From b6d3718f5af8bd5dd2cb44871529d3c290ede007 Mon Sep 17 00:00:00 2001 From: Evan Almloff Date: Thu, 11 Jun 2026 11:13:08 -0500 Subject: [PATCH 01/16] move more into the compiler --- .../src/suite/native/flash_attention_ops.rs | 14 +- fusor-ml/core/examples/bench_contraction.rs | 284 ++++ fusor-ml/core/src/access_analysis.rs | 89 ++ .../core/src/composite/flash_attention.rs | 172 ++- fusor-ml/core/src/composite/rms_norm_fused.rs | 13 +- fusor-ml/core/src/composite/rope_fused.rs | 8 +- fusor-ml/core/src/composite/softmax.rs | 30 +- fusor-ml/core/src/composite/where_cond.rs | 4 +- .../core/src/compute_graph/layout_pass.rs | 185 +-- fusor-ml/core/src/compute_graph/mod.rs | 157 +- .../compute_graph/resolve/cluster_match.rs | 218 +++ .../src/compute_graph/resolve/execution.rs | 232 +-- .../src/compute_graph/resolve/fold_views.rs | 168 ++ .../src/compute_graph/resolve/fusion_basic.rs | 185 ++- .../compute_graph/resolve/fusion_matmul.rs | 94 +- .../src/compute_graph/resolve/fusion_row.rs | 699 +++++++++ .../core/src/compute_graph/resolve/mod.rs | 100 +- .../src/compute_graph/resolve/recognize.rs | 412 +++++ .../resolve/recognize_attention.rs | 369 +++++ .../core/src/compute_graph/resolve/run.rs | 38 +- fusor-ml/core/src/compute_graph/visualize.rs | 175 +-- fusor-ml/core/src/conv.rs | 481 ------ fusor-ml/core/src/device.rs | 23 + fusor-ml/core/src/flash_attention.rs | 3 - fusor-ml/core/src/index_select.rs | 4 +- fusor-ml/core/src/lib.rs | 14 +- fusor-ml/core/src/map_layout.rs | 192 --- fusor-ml/core/src/matmul/direct.rs | 241 --- fusor-ml/core/src/matmul/kernel.rs | 519 +++---- fusor-ml/core/src/matmul/mod.rs | 111 +- fusor-ml/core/src/matmul/variants.rs | 37 +- fusor-ml/core/src/mir/kernel_backend.rs | 2 - .../kernel_backend/flash_attention/kernel.rs | 586 ------- .../mir/kernel_backend/flash_attention/mod.rs | 523 ------- .../kernel_backend/flash_attention/tests.rs | 610 -------- .../core/src/mir/kernel_backend/rms_norm.rs | 799 ---------- fusor-ml/core/src/mir/workgroup_shape.rs | 71 +- fusor-ml/core/src/nary_direct.rs | 464 +++++- fusor-ml/core/src/nary_wise.rs | 34 +- fusor-ml/core/src/quantized/dequantize.rs | 229 +-- fusor-ml/core/src/quantized/embedding.rs | 56 +- .../core/src/quantized/matmul/fallback.rs | 8 +- fusor-ml/core/src/quantized/matmul/kernel.rs | 6 + fusor-ml/core/src/quantized/matmul/mod.rs | 2 +- fusor-ml/core/src/reduce.rs | 146 +- fusor-ml/core/src/reduce_direct.rs | 644 +++++--- fusor-ml/core/src/reduce_tiled.rs | 1007 ++++++++++++ fusor-ml/core/src/resize.rs | 400 ----- fusor-ml/core/src/rms_norm.rs | 1 - fusor-ml/core/src/row_program.rs | 1247 +++++++++++++++ fusor-ml/core/src/slice_assign.rs | 136 +- fusor-ml/core/src/softmax.rs | 449 ------ fusor-ml/core/src/tensor/lazy_data.rs | 67 +- fusor-ml/core/src/tensor/mod.rs | 320 +--- fusor-ml/core/src/view.rs | 788 ++++++++++ fusor-ml/core/tests/attention.rs | 325 ++++ fusor-ml/core/tests/fused_reduce.rs | 294 ++++ fusor-ml/core/tests/keepdim_repro.rs | 29 + fusor-ml/core/tests/recognition.rs | 281 ++++ fusor-ml/fusor/src/composite/conv.rs | 16 - fusor-ml/fusor/src/gpu.rs | 22 - fusor-ml/tile-ir-kernels/src/kernels.rs | 22 +- fusor-ml/tile-ir-kernels/src/kernels/flash.rs | 1373 ----------------- .../tile-ir-kernels/src/kernels/helpers.rs | 20 - .../tile-ir-kernels/src/kernels/matmul.rs | 437 +----- .../src/kernels/qdequantize.rs | 60 - .../tile-ir-kernels/src/kernels/rms_norm.rs | 160 -- .../tile-ir-kernels/src/kernels/softmax.rs | 312 ---- fusor-ml/tile-ir-kernels/src/kernels/types.rs | 130 -- fusor-ml/tile-ir-kernels/src/lib.rs | 21 +- fusor-ml/tile-ir-kernels/tests/lowering.rs | 270 +--- fusor-ml/tile-ir/src/tile/block.rs | 6 + fusor-ml/tile-ir/src/tile/capability.rs | 9 +- .../kalosm-llama/examples/compiler_smoke.rs | 36 + 74 files changed, 8424 insertions(+), 9265 deletions(-) create mode 100644 fusor-ml/core/examples/bench_contraction.rs create mode 100644 fusor-ml/core/src/access_analysis.rs create mode 100644 fusor-ml/core/src/compute_graph/resolve/cluster_match.rs create mode 100644 fusor-ml/core/src/compute_graph/resolve/fold_views.rs create mode 100644 fusor-ml/core/src/compute_graph/resolve/fusion_row.rs create mode 100644 fusor-ml/core/src/compute_graph/resolve/recognize.rs create mode 100644 fusor-ml/core/src/compute_graph/resolve/recognize_attention.rs delete mode 100644 fusor-ml/core/src/conv.rs delete mode 100644 fusor-ml/core/src/flash_attention.rs delete mode 100644 fusor-ml/core/src/map_layout.rs delete mode 100644 fusor-ml/core/src/matmul/direct.rs delete mode 100644 fusor-ml/core/src/mir/kernel_backend/flash_attention/kernel.rs delete mode 100644 fusor-ml/core/src/mir/kernel_backend/flash_attention/mod.rs delete mode 100644 fusor-ml/core/src/mir/kernel_backend/flash_attention/tests.rs delete mode 100644 fusor-ml/core/src/mir/kernel_backend/rms_norm.rs create mode 100644 fusor-ml/core/src/reduce_tiled.rs delete mode 100644 fusor-ml/core/src/resize.rs delete mode 100644 fusor-ml/core/src/rms_norm.rs create mode 100644 fusor-ml/core/src/row_program.rs delete mode 100644 fusor-ml/core/src/softmax.rs create mode 100644 fusor-ml/core/src/view.rs create mode 100644 fusor-ml/core/tests/attention.rs create mode 100644 fusor-ml/core/tests/fused_reduce.rs create mode 100644 fusor-ml/core/tests/keepdim_repro.rs create mode 100644 fusor-ml/core/tests/recognition.rs delete mode 100644 fusor-ml/tile-ir-kernels/src/kernels/flash.rs delete mode 100644 fusor-ml/tile-ir-kernels/src/kernels/qdequantize.rs delete mode 100644 fusor-ml/tile-ir-kernels/src/kernels/rms_norm.rs delete mode 100644 fusor-ml/tile-ir-kernels/src/kernels/softmax.rs create mode 100644 models/kalosm-llama/examples/compiler_smoke.rs diff --git a/fusor-ml/conformance/src/suite/native/flash_attention_ops.rs b/fusor-ml/conformance/src/suite/native/flash_attention_ops.rs index 1d9482ac6..388ee206d 100644 --- a/fusor-ml/conformance/src/suite/native/flash_attention_ops.rs +++ b/fusor-ml/conformance/src/suite/native/flash_attention_ops.rs @@ -304,10 +304,9 @@ fn assert_flash_attention_case( /// non-deterministic workgroup-memory races. pub fn flash_attention_decode_tiled_matches_cpu_reference() -> AssertionCases { // (num_heads, num_kv_heads, kv_seq_len) - // Shapes specifically chosen to stress decode block boundaries and the - // tiled flash_decode_small_block path. head_dim=128 forces the decode-small - // kernel; kv_seq_len spans 128-wide tiled failure cases and values - // beyond the largest 1024-wide block. + // Shapes specifically chosen to stress decode tile boundaries in the + // attention row program: kv_seq_len spans ragged single tiles and + // values beyond one workgroup bucket (the split lowering). let shapes = [ (16, 2, 129), // just past one full tile (16, 2, 192), // mid-second tile @@ -349,10 +348,9 @@ pub fn flash_attention_decode_tiled_matches_cpu_reference() -> AssertionCases { /// Same as the tiled test above, but builds Q with non-canonical strides /// (reshape + transpose) — matching how the real attention path produces Q -/// in `models/kalosm-llama/src/raw/attention_layer.rs`. The kernel reads Q -/// via `index_n(meta.q_offset, meta.q_strides, ...)`, so different strides -/// hit different memory addresses and exercise different control flow paths -/// inside `flash_decode_small_block`. +/// in `models/kalosm-llama/src/raw/attention_layer.rs`. The row program +/// reads Q through its layout strides, so different strides hit different +/// memory addresses. pub fn flash_attention_decode_tiled_with_transposed_q_matches_cpu_reference() -> AssertionCases { let shapes = [(16, 2, 129), (16, 2, 257), (16, 2, 384), (16, 2, 569)]; diff --git a/fusor-ml/core/examples/bench_contraction.rs b/fusor-ml/core/examples/bench_contraction.rs new file mode 100644 index 000000000..92315ffc2 --- /dev/null +++ b/fusor-ml/core/examples/bench_contraction.rs @@ -0,0 +1,284 @@ +//! A/B kernel timing for the dense matmul routes and composed contractions. +//! Run on two builds and compare medians; wall-clock per resolve with a +//! device sync, warmup excluded. + +use fusor_core::{Device, QMatrix, StrideSpec, Tensor}; +use fusor_gguf::GgmlType; +use std::time::Instant; + +fn values(len: usize, scale: f32) -> Vec { + (0..len).map(|i| ((i as f32) * scale).sin()).collect() +} + +fn time_case(name: &str, warmup: usize, iters: usize, mut run: impl FnMut()) { + for _ in 0..warmup { + run(); + } + let mut samples = Vec::with_capacity(iters); + for _ in 0..iters { + let start = Instant::now(); + run(); + samples.push(start.elapsed().as_secs_f64() * 1e3); + } + samples.sort_by(|a, b| a.partial_cmp(b).unwrap()); + let median = samples[samples.len() / 2]; + let min = samples[0]; + println!("{name}: median {median:.3} ms, min {min:.3} ms over {iters} iters"); +} + +fn main() { + pollster::block_on(async { + let device = Device::new().await.expect("gpu device"); + + // Coop-ineligible dense matmul (1000 % 64 != 0): the workgroup-tiled + // route. + { + let a = Tensor::from_slice(&device, [1000, 1000], &values(1_000_000, 0.13)); + let b = Tensor::from_slice(&device, [1000, 1000], &values(1_000_000, 0.07)); + time_case("matmul_1000_tiled", 3, 15, || { + let out = a.mat_mul(&b); + out.materialize_sync(); + }); + } + + // Coop-eligible square (512 divisible by 64): should be unchanged. + { + let a = Tensor::from_slice(&device, [512, 512], &values(512 * 512, 0.13)); + let b = Tensor::from_slice(&device, [512, 512], &values(512 * 512, 0.07)); + time_case("matmul_512_coop", 3, 30, || { + let out = a.mat_mul(&b); + out.materialize_sync(); + }); + } + + // Low tile utilization (65/96/65): the register-tile route. + { + let a = Tensor::from_slice(&device, [16, 65, 96], &values(16 * 65 * 96, 0.13)); + let b = Tensor::from_slice(&device, [16, 96, 65], &values(16 * 96 * 65, 0.07)); + time_case("matmul_batched_65_register", 3, 30, || { + let out = a.mat_mul(&b); + out.materialize_sync(); + }); + } + + // Weighted sum: `sum_k w[k] * x[m, k]` — k-dependent inputs share no + // (row, col) pair, but `w` misses the row dim. + { + let (m, k) = (4096usize, 4096usize); + let x = Tensor::from_slice(&device, [m, k], &values(m * k, 0.13)); + let w = Tensor::from_slice(&device, [k], &values(k, 0.07)); + time_case("weighted_sum_4096", 3, 30, || { + let w2 = w.restride([StrideSpec::dim_with(0, m, 0), StrideSpec::dim(0, k)]); + let out = (&x * &w2).sum(1); + out.materialize_sync(); + }); + } + + // Broadcast scale: `x[m, n] * w[n]` — `w` is invariant along rows. + { + let (m, n) = (4096usize, 4096usize); + let x = Tensor::from_slice(&device, [m, n], &values(m * n, 0.13)); + let w = Tensor::from_slice(&device, [n], &values(n, 0.07)); + time_case("broadcast_scale_4096", 3, 30, || { + let w2 = w.restride([StrideSpec::dim_with(0, m, 0), StrideSpec::dim(0, n)]); + let out = &x * &w2; + out.materialize_sync(); + }); + } + + // Broadcast table apply: `x[b, s, h] * t[s, h]` with a table too + // large for cache — the invariant table loads hoist out of each + // thread's output run. + { + let (b, s, h) = (4usize, 2048usize, 2048usize); + let x = Tensor::from_slice(&device, [b, s, h], &values(b * s * h, 0.13)); + let t = Tensor::from_slice(&device, [s, h], &values(s * h, 0.07)); + time_case("broadcast_table_4x2048x2048", 3, 30, || { + let t3 = t.restride([ + StrideSpec::dim_with(0, b, 0), + StrideSpec::dim(0, s), + StrideSpec::dim(1, h), + ]); + let out = &x * &t3; + out.materialize_sync(); + }); + } + + // Dequantize a Q8_0 matrix to dense f32. + { + const Q8_BLOCK: usize = 32; + let (n, k) = (4096usize, 4096usize); + let scale = half::f16::from_f32(0.02); + let mut bytes = Vec::with_capacity(n * k / Q8_BLOCK * 34); + for block in 0..(n * k / Q8_BLOCK) { + bytes.extend_from_slice(&scale.to_le_bytes()); + for i in 0..Q8_BLOCK { + bytes.push((((block * 7 + i * 5) % 64) as i32 - 32) as i8 as u8); + } + } + let w = QMatrix::from_parts( + &device, + &bytes, + vec![n, k].into_boxed_slice(), + GgmlType::Q8_0, + ) + .unwrap(); + time_case("dequantize_q8_4096", 3, 30, || { + let out = w.dequantize::(); + out.materialize_sync(); + }); + } + + // Dense gemv: [m, k] x [k, 1]. + { + let (m, k) = (4096usize, 4096usize); + let a = Tensor::from_slice(&device, [m, k], &values(m * k, 0.13)); + let b = Tensor::from_slice(&device, [k, 1], &values(k, 0.07)); + time_case("gemv_4096", 3, 50, || { + let out = a.mat_mul(&b); + out.materialize_sync(); + }); + } + + // Norm shapes: decode-like single row and prefill-like many rows. + { + let hidden = 4096usize; + let x1 = Tensor::from_slice(&device, [1, hidden], &values(hidden, 0.13)); + let xs = Tensor::from_slice(&device, [512, hidden], &values(512 * hidden, 0.13)); + let w = Tensor::from_slice(&device, [hidden], &values(hidden, 0.07)); + time_case("rms_norm_1x4096", 3, 50, || { + let out = x1.rms_norm_fused(&w, None, 1e-5); + out.materialize_sync(); + }); + time_case("rms_norm_512x4096", 3, 30, || { + let out = xs.rms_norm_fused(&w, None, 1e-5); + out.materialize_sync(); + }); + time_case("softmax_512x4096", 3, 30, || { + let out = xs.softmax(1); + out.materialize_sync(); + }); + let small = Tensor::from_slice(&device, [32, 128], &values(32 * 128, 0.17)); + time_case("softmax_32x128", 3, 50, || { + let out = small.softmax(1); + out.materialize_sync(); + }); + } + + // Decode-shape attention (q_len = 1): the per-token hot path. kv + // lengths cover the single-dispatch decode buckets and (at 2048) the + // split two-dispatch route. + { + let (batch, heads, head_dim) = (1usize, 32usize, 128usize); + let scale = 1.0 / (head_dim as f32).sqrt(); + let q = Tensor::from_slice( + &device, + [batch, heads, 1, head_dim], + &values(batch * heads * head_dim, 0.13), + ); + for kv in [512usize, 1024, 2048, 4096, 8192] { + let k = Tensor::from_slice( + &device, + [batch, heads, kv, head_dim], + &values(batch * heads * kv * head_dim, 0.07), + ); + let v = Tensor::from_slice( + &device, + [batch, heads, kv, head_dim], + &values(batch * heads * kv * head_dim, 0.11), + ); + time_case(&format!("attn_decode_kv{kv}"), 3, 50, || { + let out = q.flash_attention(&k, &v, scale, None); + out.materialize_sync(); + }); + } + // Grouped-query variant: 32 query heads over 8 KV heads. + let kv_heads = 8usize; + let kv = 1024usize; + let k = Tensor::from_slice( + &device, + [batch, kv_heads, kv, head_dim], + &values(batch * kv_heads * kv * head_dim, 0.07), + ); + let v = Tensor::from_slice( + &device, + [batch, kv_heads, kv, head_dim], + &values(batch * kv_heads * kv * head_dim, 0.11), + ); + time_case("attn_decode_gqa_kv1024", 3, 50, || { + let out = q.flash_attention(&k, &v, scale, None); + out.materialize_sync(); + }); + } + + // Prefill-shape attention: causal self-attention at q == kv (the + // streaming-tiled regime) and a mid-length query block (the plain + // streaming regime). + { + let (batch, heads, head_dim) = (1usize, 32usize, 128usize); + let scale = 1.0 / (head_dim as f32).sqrt(); + let seq = 512usize; + let q = Tensor::from_slice( + &device, + [batch, heads, seq, head_dim], + &values(batch * heads * seq * head_dim, 0.13), + ); + let k = Tensor::from_slice( + &device, + [batch, heads, seq, head_dim], + &values(batch * heads * seq * head_dim, 0.07), + ); + let v = Tensor::from_slice( + &device, + [batch, heads, seq, head_dim], + &values(batch * heads * seq * head_dim, 0.11), + ); + time_case("attn_prefill_512_causal", 3, 20, || { + let out = q.flash_attention_causal(&k, &v, scale); + out.materialize_sync(); + }); + + let q64 = Tensor::from_slice( + &device, + [batch, heads, 64, head_dim], + &values(batch * heads * 64 * head_dim, 0.13), + ); + let kv = 1024usize; + let k = Tensor::from_slice( + &device, + [batch, heads, kv, head_dim], + &values(batch * heads * kv * head_dim, 0.07), + ); + let v = Tensor::from_slice( + &device, + [batch, heads, kv, head_dim], + &values(batch * heads * kv * head_dim, 0.11), + ); + time_case("attn_prefill_q64_kv1024", 3, 30, || { + let out = q64.flash_attention(&k, &v, scale, None); + out.materialize_sync(); + }); + } + + // Composed broadcast contraction: not recognized as a matmul. + { + let (m, n, k) = (256usize, 256usize, 256usize); + let a = Tensor::from_slice(&device, [m, k], &values(m * k, 0.13)); + let b = Tensor::from_slice(&device, [n, k], &values(n * k, 0.07)); + time_case("broadcast_contraction_256", 3, 30, || { + let a3 = a.restride([ + StrideSpec::dim(0, m), + StrideSpec::dim_with(0, n, 0), + StrideSpec::dim(1, k), + ]); + let b3 = b.restride([ + StrideSpec::dim_with(0, m, 0), + StrideSpec::dim(0, n), + StrideSpec::dim(1, k), + ]); + let out = (&a3 * &b3).sum(2); + out.materialize_sync(); + }); + } + }); +} diff --git a/fusor-ml/core/src/access_analysis.rs b/fusor-ml/core/src/access_analysis.rs new file mode 100644 index 000000000..74c78344c --- /dev/null +++ b/fusor-ml/core/src/access_analysis.rs @@ -0,0 +1,89 @@ +//! Access-pattern analysis over fused n-ary expressions. +//! +//! Every tiled lowering asks the same two questions about an expression's +//! inputs: *which index-space dims does each input read* (the index lists), +//! and *which of those reads actually vary* (stride-aware — a broadcast view +//! reads a dim through a zero stride, which is reuse, not dependence). The +//! answers drive tile-shape selection for reductions (which inputs stage +//! through workgroup memory or registers) and for elementwise kernels (which +//! loads hoist out of a thread's output run). + +use crate::{nary_direct::TensorMeta, nary_wise::NaryExpr}; + +/// Per-input access structure for an expression whose inputs are all read at +/// consistent, pure `DimIndex` coordinates. +pub(crate) struct InputAccesses { + /// Per input: the index-space dim read by each input dimension. + pub(crate) dims: Vec>, + /// Per input: the index-space dims the input *effectively* depends on. A + /// stride-0 or size-1 input dim is not a dependence — `layout_index` + /// drops it from addressing too — so broadcast views analyze the same as + /// index lists that never mention the dim. + pub(crate) effective: Vec>, +} + +impl InputAccesses { + /// Collect access structure, or `None` for any access a tiled kernel + /// cannot hoist or stage: non-`DimIndex` coordinates (gathers), + /// inconsistent index lists for one input, or `DimIndex` leaves in value + /// position (the value depends on the coordinate itself). + pub(crate) fn collect( + expr: &NaryExpr, + input_count: usize, + metas: &[TensorMeta], + ) -> Option { + let mut accesses = vec![None; input_count]; + collect_accesses(expr, &mut accesses)?; + let dims = accesses.into_iter().collect::>>>()?; + let effective = dims + .iter() + .zip(metas) + .map(|(dims, meta)| { + dims.iter() + .enumerate() + .filter_map(|(j, &d)| { + let stride = meta.strides.get(j).copied().unwrap_or(0); + let size = meta.shape.get(j).copied().unwrap_or(1); + (stride != 0 && size > 1).then_some(d) + }) + .collect() + }) + .collect(); + Some(Self { dims, effective }) + } + + /// Whether input `i` effectively depends on index-space dim `d`. + pub(crate) fn depends_on(&self, i: usize, d: usize) -> bool { + self.effective[i].contains(&d) + } +} + +fn collect_accesses(expr: &NaryExpr, accesses: &mut Vec>>) -> Option<()> { + match expr { + NaryExpr::Op { children, .. } => { + for child in children { + collect_accesses(child, accesses)?; + } + Some(()) + } + NaryExpr::IndexedInput { input_idx, indices } => { + let dims = indices + .iter() + .map(|index| match index { + NaryExpr::DimIndex(dim) => Some(*dim), + _ => None, + }) + .collect::>>()?; + match &accesses[*input_idx] { + Some(existing) if *existing != dims => None, + Some(_) => Some(()), + None => { + accesses[*input_idx] = Some(dims); + Some(()) + } + } + } + NaryExpr::DimIndex(_) => None, + NaryExpr::Scalar(_) => Some(()), + } +} diff --git a/fusor-ml/core/src/composite/flash_attention.rs b/fusor-ml/core/src/composite/flash_attention.rs index bfb39ebee..e36a4890e 100644 --- a/fusor-ml/core/src/composite/flash_attention.rs +++ b/fusor-ml/core/src/composite/flash_attention.rs @@ -1,46 +1,57 @@ -use crate::{DataTypeEnum, StrideSpec, Tensor}; +use crate::{ + DataTypeEnum, Layout, Tensor, + nary_wise::{ElementwiseOperation, NaryExpr, NaryFunction, NaryOp, NaryScalar}, + view::ViewOperation, +}; impl Tensor { - /// Causal flash attention: the kernel applies a strict lower-triangular - /// causal mask internally, skipping the upper-triangle Q·K work entirely. - /// `q_seq_len` must equal `kv_seq_len` (prefill self-attention); other - /// shapes fall back to [`flash_attention`] with an explicit mask. + /// A view layered directly on this tensor's node, without collapsing + /// into any underlying view chain. Composed-attention clusters use these + /// so recognition can peel the exact GQA-expand / transpose / mask + /// layouts back to the original q/k/v/mask nodes. + fn attached_view(&self, layout: Layout) -> Tensor { + Tensor::from_parts(self.data().view(ViewOperation::fully_defined( + self.key(), + layout, + self.shape(), + self.datatype(), + ))) + } + + /// Causal flash attention in its composed form: scores at kv positions + /// beyond the query position are replaced with `-inf` via an + /// index-comparison select (`kv_pos <= q_pos`), so causality is pure + /// index arithmetic — no mask tensor. The resolver recognizes the + /// cluster and routes it to the attention row program, whose axis bound + /// skips the masked upper-triangle tiles entirely. pub fn flash_attention_causal(&self, k: &Self, v: &Self, scale: f32) -> Self { - self.assert_rank::<4>(); - if let Some(output) = self.try_flash_attention_direct_causal(k, v, scale) { - return output; - } - // Fallback: materialize a causal mask and re-enter the masked path. - let q_shape = self.shape(); - let seq_len = q_shape[2]; - match self.datatype() { - DataTypeEnum::F32 => { - let mut data = vec![0.0f32; seq_len * seq_len]; - for i in 0..seq_len { - for j in (i + 1)..seq_len { - data[i * seq_len + j] = f32::NEG_INFINITY; - } - } - let mask = Tensor::from_slice::(self.device(), [seq_len, seq_len], &data); - self.flash_attention(k, v, scale, Some(&mask)) - } - DataTypeEnum::F16 => { - let mut data = vec![half::f16::from_f32(0.0); seq_len * seq_len]; - let neg_inf = half::f16::from_f32(f32::NEG_INFINITY); - for i in 0..seq_len { - for j in (i + 1)..seq_len { - data[i * seq_len + j] = neg_inf; - } - } - let mask = - Tensor::from_slice::(self.device(), [seq_len, seq_len], &data); - self.flash_attention(k, v, scale, Some(&mask)) - } - DataTypeEnum::U32 => panic!("flash_attention requires f32/f16 tensors"), - } + assert_eq!( + self.shape()[2], + k.shape()[2], + "causal flash attention requires q_seq_len == kv_seq_len \ + (self-attention prefill); use an explicit mask otherwise" + ); + self.compose_attention(k, v, scale, None, true) } + /// Scaled dot-product attention in its composed form: + /// `softmax(q · kᵀ · scale [+ mask]) · v`, with K/V expanded across query + /// heads for grouped-query attention. The resolver recognizes the + /// canonical cluster and routes it to the fused attention row program; + /// ineligible shapes lower through the recognized matmul + softmax + /// kernels (the same math). pub fn flash_attention(&self, k: &Self, v: &Self, scale: f32, mask: Option<&Tensor>) -> Self { + self.compose_attention(k, v, scale, mask, false) + } + + fn compose_attention( + &self, + k: &Self, + v: &Self, + scale: f32, + mask: Option<&Tensor>, + causal: bool, + ) -> Self { self.assert_rank::<4>(); k.assert_rank::<4>(); v.assert_rank::<4>(); @@ -50,9 +61,6 @@ impl Tensor { mask.assert_rank::<2>(); assert_eq!(self.datatype(), mask.datatype()); } - if let Some(output) = self.try_flash_attention_direct(k, v, scale, mask) { - return output; - } let q_shape = self.shape(); let k_shape = k.shape(); @@ -72,32 +80,68 @@ impl Tensor { ); let groups = num_heads / num_kv_heads; - let (k_expanded, v_expanded) = if groups > 1 { - let k_expanded = k - .reshape([batch, num_kv_heads, 1, kv_seq_len, head_dim]) - .broadcast_as([batch, num_kv_heads, groups, kv_seq_len, head_dim]) - .reshape([batch, num_heads, kv_seq_len, head_dim]); - let v_expanded = v - .reshape([batch, num_kv_heads, 1, kv_seq_len, head_dim]) - .broadcast_as([batch, num_kv_heads, groups, kv_seq_len, head_dim]) - .reshape([batch, num_heads, kv_seq_len, head_dim]); - (k_expanded, v_expanded) - } else { - (k.clone(), v.clone()) + let expanded_shape = [batch, num_heads, kv_seq_len, head_dim]; + let expand = |tensor: &Tensor| -> Tensor { + if groups == 1 { + return tensor.clone(); + } + // Two attached views: a stride-0 broadcast across the group dim, + // then a flat reinterpret down to rank 4. + let grouped = tensor.attached_view(Layout::from_parts( + 0, + [batch, num_kv_heads, groups, kv_seq_len, head_dim].into(), + [ + num_kv_heads * kv_seq_len * head_dim, + kv_seq_len * head_dim, + 0, + head_dim, + 1, + ] + .into(), + )); + grouped.attached_view(Layout::contiguous(&expanded_shape)) }; + let (k_expanded, v_expanded) = (expand(k), expand(v)); - let k_t = k_expanded.restride([ - StrideSpec::dim(0, batch), - StrideSpec::dim(1, num_heads), - StrideSpec::dim(3, head_dim), - StrideSpec::dim(2, kv_seq_len), - ]); + let k_t = k_expanded.attached_view(Layout::contiguous(&expanded_shape).transpose(2, 3)); let scores = match self.datatype() { DataTypeEnum::F32 => self.mat_mul(&k_t) * scale, DataTypeEnum::F16 => self.mat_mul(&k_t) * half::f16::from_f32(scale), DataTypeEnum::U32 => panic!("flash_attention requires f32/f16 tensors"), }; - let scores = if let Some(mask) = mask { + let scores = if causal { + // Keep kv positions at or before the query position; everything + // later contributes exp(-inf) = 0 to the softmax. + let condition = NaryExpr::Op { + children: vec![NaryExpr::DimIndex(3), NaryExpr::DimIndex(2)], + function: NaryFunction::binary( + Some("causal_bound".to_string()), + NaryOp::LessEqual, + DataTypeEnum::U32, + DataTypeEnum::U32, + DataTypeEnum::U32, + ), + }; + let datatype = self.datatype(); + let neg_inf = match datatype { + DataTypeEnum::F32 => NaryScalar::F32(f32::NEG_INFINITY), + DataTypeEnum::F16 => NaryScalar::F16(half::f16::NEG_INFINITY), + DataTypeEnum::U32 => unreachable!("attention requires f32/f16"), + }; + let expression = NaryExpr::select( + condition, + NaryExpr::input(0, 4), + NaryExpr::scalar(neg_inf), + DataTypeEnum::U32, + datatype, + ); + Tensor::from_parts(scores.data().nary(ElementwiseOperation { + inputs: vec![scores.key()], + expression, + shape: [batch, num_heads, q_seq_len, kv_seq_len].into(), + output_datatype: datatype, + })) + } else if let Some(mask) = mask { let mask_shape = mask.shape(); assert_eq!( mask_shape, @@ -107,10 +151,12 @@ impl Tensor { q_seq_len, kv_seq_len ); - scores - + mask - .reshape([1, 1, q_seq_len, kv_seq_len]) - .broadcast_as([batch, num_heads, q_seq_len, kv_seq_len]) + let mask_view = mask.attached_view(Layout::from_parts( + 0, + [batch, num_heads, q_seq_len, kv_seq_len].into(), + [0, 0, kv_seq_len, 1].into(), + )); + scores + mask_view } else { scores }; diff --git a/fusor-ml/core/src/composite/rms_norm_fused.rs b/fusor-ml/core/src/composite/rms_norm_fused.rs index 1ce50c9aa..e0a122961 100644 --- a/fusor-ml/core/src/composite/rms_norm_fused.rs +++ b/fusor-ml/core/src/composite/rms_norm_fused.rs @@ -1,6 +1,12 @@ use crate::{DataTypeEnum, Tensor}; impl Tensor { + /// RMS norm, expressed as its composed form: + /// `x / sqrt(mean(x²) + eps) * weight (+ bias)`, computed in f32 with + /// casts at the edges for f16 tensors. The resolver recognizes the + /// canonical cluster and routes it to the fused RMS norm kernel; + /// ineligible configurations lower through the generic elementwise + + /// reduce path (the same math). pub fn rms_norm_fused(&self, weight: &Tensor, bias: Option<&Tensor>, eps: f32) -> Self { let original_datatype = self.datatype(); assert!( @@ -11,9 +17,6 @@ impl Tensor { if let Some(bias) = bias { assert_eq!(bias.datatype(), original_datatype); } - if let Some(output) = self.try_rms_norm_direct(weight, bias, eps) { - return output; - } let hidden_size = *self.shape().last().unwrap() as f32; let last_dim = self.rank() - 1; @@ -47,10 +50,6 @@ impl Tensor { bias: Option<&Tensor>, eps: f32, ) -> Self { - if let Some(output) = self.try_rms_norm_residual_direct(residual, weight, bias, eps) { - return output; - } - (self.clone() + residual.clone()).rms_norm_fused(weight, bias, eps) } } diff --git a/fusor-ml/core/src/composite/rope_fused.rs b/fusor-ml/core/src/composite/rope_fused.rs index b8e2b9bee..d6a3fa074 100644 --- a/fusor-ml/core/src/composite/rope_fused.rs +++ b/fusor-ml/core/src/composite/rope_fused.rs @@ -218,8 +218,8 @@ impl RopeFusedOperation { self.shape.len() } - fn to_nary(&self) -> crate::nary_wise::NaryOperation { - crate::nary_wise::NaryOperation { + fn to_nary(&self) -> crate::nary_wise::ElementwiseOperation { + crate::nary_wise::ElementwiseOperation { inputs: vec![self.input, self.cos, self.sin], expression: self.build_expr(), shape: self.shape.clone(), @@ -410,12 +410,12 @@ fn build_sign_condition(dim_last: NaryExpr, head_dim: usize, mode: RopeMode) -> } impl RopePairFusedOperation { - fn to_nary(&self) -> crate::nary_wise::NaryOperation { + fn to_nary(&self) -> crate::nary_wise::ElementwiseOperation { let mut inputs = vec![self.q, self.k, self.cos, self.sin]; if let Some(position) = self.position { inputs.push(position); } - crate::nary_wise::NaryOperation { + crate::nary_wise::ElementwiseOperation { inputs, expression: self.build_expr(), shape: vec![self.total_elements].into_boxed_slice(), diff --git a/fusor-ml/core/src/composite/softmax.rs b/fusor-ml/core/src/composite/softmax.rs index bb73de94d..cb8200ca9 100644 --- a/fusor-ml/core/src/composite/softmax.rs +++ b/fusor-ml/core/src/composite/softmax.rs @@ -1,31 +1,13 @@ -use std::sync::Arc; - -use crate::{ - Tensor, - softmax::SoftmaxOperation, - tensor::{LazyTensorData, TensorInfo}, -}; +use crate::Tensor; impl Tensor { + /// Softmax along `axis`, expressed as its composed form: + /// `exp(x - max(x)) / sum(exp(x - max(x)))` over keepdim-broadcast + /// reductions. The resolver recognizes the canonical cluster and routes + /// it to the fused softmax kernels; ineligible shapes lower through the + /// generic elementwise + reduce path (the same math). pub fn softmax(&self, axis: usize) -> Self { assert!(axis < self.rank(), "softmax axis out of bounds"); - if let Some(operation) = SoftmaxOperation::new( - self.key(), - self.shape(), - axis, - self.datatype(), - self.device(), - ) { - let device = self.device().clone(); - let info = TensorInfo::new(self.shape().into(), self.datatype()); - let key = device.compute_graph().create_graph_op(Arc::new(operation)); - return Tensor::from_parts(LazyTensorData::from_parts(device, info, key)); - } - - self.softmax_composite(axis) - } - - fn softmax_composite(&self, axis: usize) -> Self { let mut kept_shape = self.shape().to_vec(); kept_shape[axis] = 1; let max = self diff --git a/fusor-ml/core/src/composite/where_cond.rs b/fusor-ml/core/src/composite/where_cond.rs index 7b478bdf5..a7c407976 100644 --- a/fusor-ml/core/src/composite/where_cond.rs +++ b/fusor-ml/core/src/composite/where_cond.rs @@ -1,6 +1,6 @@ use crate::{ Tensor, - nary_wise::{NaryExpr, NaryOperation}, + nary_wise::{ElementwiseOperation, NaryExpr}, }; impl Tensor { @@ -10,7 +10,7 @@ impl Tensor { assert_eq!(self.shape(), on_true.shape()); let shape: Box<[usize]> = self.shape().into(); let rank = shape.len(); - let nary = NaryOperation { + let nary = ElementwiseOperation { inputs: vec![self.key(), on_true.key(), on_false.key()], expression: NaryExpr::select( NaryExpr::input(0, rank), diff --git a/fusor-ml/core/src/compute_graph/layout_pass.rs b/fusor-ml/core/src/compute_graph/layout_pass.rs index 303874d66..2203c989d 100644 --- a/fusor-ml/core/src/compute_graph/layout_pass.rs +++ b/fusor-ml/core/src/compute_graph/layout_pass.rs @@ -1,8 +1,8 @@ use rustc_hash::FxHashMap; -use crate::{Layout, TensorLayoutInfo, nary_wise::NaryOperation}; +use crate::{Layout, TensorLayoutInfo, nary_wise::ElementwiseOperation}; -use super::{ComputeGraphNodeVariant, GraphOperation, NodeIndex, queue::ComputeQueue}; +use super::{ComputeGraphNodeVariant, NodeIndex, queue::ComputeQueue}; #[derive(Default)] pub(crate) struct LayoutPass { @@ -24,23 +24,17 @@ impl LayoutPass { continue; } match &node_data.variant { - ComputeGraphNodeVariant::Nary(op) => self.visit_nary(node, op), - ComputeGraphNodeVariant::MatMul(op) => self.visit_mat_mul(node, op), - ComputeGraphNodeVariant::QMatMul(op) => self.visit_q_mat_mul(node, op), - ComputeGraphNodeVariant::QEmbedding(op) => self.visit_q_embedding(node, op), - ComputeGraphNodeVariant::Reduce(op) => self.visit_reduce(node, op), - ComputeGraphNodeVariant::FlashAttention(op) => self.visit_flash_attention(node, op), - ComputeGraphNodeVariant::GraphOp(op) => self.visit_graph_op(node, op.as_ref()), - ComputeGraphNodeVariant::MapLayout(op) => self.visit_map_layout(node, op), - ComputeGraphNodeVariant::Resize(op) => self.visit_resize(node, op), - ComputeGraphNodeVariant::SliceAssign(op) => self.visit_slice_assign(node, op), ComputeGraphNodeVariant::Tensor(op) => self.visit_tensor(node, op), - ComputeGraphNodeVariant::Dequantize(op) => self.visit_dequantize(node, op), + ComputeGraphNodeVariant::QMatrix(op) => self.visit_dequantize(node, op), + ComputeGraphNodeVariant::Elementwise(op) => self.visit_nary(node, op), + ComputeGraphNodeVariant::Reduce(op) => self.visit_reduce(node, op), + ComputeGraphNodeVariant::View(op) => self.visit_view(node, op), + ComputeGraphNodeVariant::Assign(op) => self.visit_slice_assign(node, op), } } } - fn visit_nary(&mut self, key: NodeIndex, operation: &NaryOperation) { + fn visit_nary(&mut self, key: NodeIndex, operation: &ElementwiseOperation) { // Ensure all inputs have been visited for input in &operation.inputs { if !self.output_layout.contains_key(input) { @@ -56,169 +50,28 @@ impl LayoutPass { ); } - fn visit_mat_mul(&mut self, key: NodeIndex, operation: &crate::MatMulOperation) { - let Some(first_layout) = self.output_layout.get(&operation.first) else { - self.queue.push_back(operation.first); - self.queue.push_back(key); - return; - }; - let Some(_) = self.output_layout.get(&operation.second) else { - self.queue.push_back(operation.second); - self.queue.push_back(key); - return; - }; - let output_shape = &operation.out_shape; - let output_layout = Layout::contiguous(output_shape); - self.output_layout.insert( - key, - TensorLayoutInfo::new(output_layout, first_layout.datatype()), - ); - } - - fn visit_q_mat_mul( - &mut self, - key: NodeIndex, - operation: &crate::quantized::matmul::QMatMulOperation, - ) { - let Some(first_layout) = self.output_layout.get(&operation.input) else { - self.queue.push_back(operation.input); - self.queue.push_back(key); - return; - }; - let output_layout = Layout::contiguous(&operation.out_shape); - self.output_layout.insert( - key, - TensorLayoutInfo::new(output_layout, first_layout.datatype()), - ); - } - - fn visit_q_embedding( - &mut self, - key: NodeIndex, - operation: &crate::quantized::embedding::QEmbeddingOperation, - ) { - let Some(_) = self.output_layout.get(&operation.indexes) else { - self.queue.push_back(operation.indexes); - self.queue.push_back(key); - return; - }; - let output_layout = Layout::contiguous(&operation.out_shape); - self.output_layout.insert( - key, - TensorLayoutInfo::new(output_layout, crate::DataTypeEnum::F32), - ); - } - fn visit_reduce(&mut self, key: NodeIndex, operation: &crate::ReduceOperation) { - let dim = operation.axis; - let Some(input_layout) = self.output_layout.get(&operation.value) else { - self.queue.push_back(operation.value); - self.queue.push_back(key); - return; - }; - let new_shape = input_layout - .layout() - .shape() - .iter() - .enumerate() - .filter_map(|(i, x)| (i != dim).then_some(*x)) - .collect::>(); - let new_layout = Layout::contiguous(&new_shape); - self.output_layout.insert( - key, - TensorLayoutInfo::new(new_layout, input_layout.datatype()), - ); - } - - fn visit_graph_op(&mut self, key: NodeIndex, operation: &dyn GraphOperation) { - let mut dependencies = Vec::new(); - operation.visit_dependencies(&mut |dependency| { - dependencies.push(dependency); - }); - - let mut missing_dependency = false; - for dependency in dependencies { - if !self.output_layout.contains_key(&dependency) { - self.queue.push_back(dependency); - missing_dependency = true; - } - } - if missing_dependency { - self.queue.push_back(key); - return; - } - - let output_layout = operation - .output_layout(&self.output_layout) - .unwrap_or_else(|| { - panic!( - "graph operation {} could not infer output layout", - operation.category() - ) - }); - self.output_layout.insert(key, output_layout); - } - - fn visit_flash_attention( - &mut self, - key: NodeIndex, - operation: &crate::FlashAttentionOperation, - ) { - let Some(q_layout) = self.output_layout.get(&operation.q) else { - self.queue.push_back(operation.q); - self.queue.push_back(key); - return; - }; - if !self.output_layout.contains_key(&operation.k) { - self.queue.push_back(operation.k); - self.queue.push_back(key); - return; - } - if !self.output_layout.contains_key(&operation.v) { - self.queue.push_back(operation.v); - self.queue.push_back(key); - return; - } - if let Some(mask) = operation.mask - && !self.output_layout.contains_key(&mask) - { - self.queue.push_back(mask); - self.queue.push_back(key); - return; - } + let new_layout = Layout::contiguous(&operation.out_shape()); self.output_layout.insert( key, - TensorLayoutInfo::new( - Layout::contiguous(&operation.out_shape), - q_layout.datatype(), - ), - ); - } - - fn visit_map_layout( - &mut self, - key: NodeIndex, - operation: &crate::map_layout::MapLayoutOperation, - ) { - let Some(input_layout) = self.output_layout.get(&operation.input) else { - self.queue.push_back(operation.input); - self.queue.push_back(key); - return; - }; - let new_layout = operation.map_layout(input_layout.layout()); - self.output_layout.insert( - key, - TensorLayoutInfo::new(new_layout, input_layout.datatype()), + TensorLayoutInfo::new(new_layout, operation.out_datatype()), ); } - fn visit_resize(&mut self, key: NodeIndex, operation: &crate::resize::ResizeOperation) { + fn visit_view(&mut self, key: NodeIndex, operation: &crate::view::ViewOperation) { let Some(input_layout) = self.output_layout.get(&operation.input) else { self.queue.push_back(operation.input); self.queue.push_back(key); return; }; - let new_layout = Layout::contiguous(&operation.new_shape); + // A fully-defined view that composes with the input's layout stays a + // zero-cost buffer view; anything else materializes contiguously + // through the gather fallback. + let new_layout = operation + .is_fully_defined() + .then(|| crate::view::compose_layouts(&operation.layout, input_layout.layout())) + .flatten() + .unwrap_or_else(|| Layout::contiguous(operation.shape())); self.output_layout.insert( key, TensorLayoutInfo::new(new_layout, input_layout.datatype()), diff --git a/fusor-ml/core/src/compute_graph/mod.rs b/fusor-ml/core/src/compute_graph/mod.rs index 43833aed7..f50655375 100644 --- a/fusor-ml/core/src/compute_graph/mod.rs +++ b/fusor-ml/core/src/compute_graph/mod.rs @@ -1,10 +1,9 @@ -use std::{any::Any, sync::Arc}; +use std::sync::Arc; use parking_lot::RwLock; pub use petgraph::graph::NodeIndex; use petgraph::prelude::StableGraph; use resolve::Resolver; -use rustc_hash::FxHashMap; #[cfg(feature = "graphvis")] use tabbycat::Graph; @@ -17,18 +16,15 @@ mod tests; mod visualize; use crate::{ - DataTypeEnum, Device, FlashAttentionOperation, MatMulOperation, QMatrix, ReduceOperation, - RmsNormOperation, + DataTypeEnum, Device, QMatrix, ReduceOperation, compute_graph::resolve::ResolverResult, dequantize::DequantizeOperation, - map_layout::MapLayoutOperation, mir::{inputs::MirValue, operation::Operation}, - nary_wise::NaryOperation, - quantized::embedding::QEmbeddingOperation, + nary_wise::ElementwiseOperation, quantized::matmul::QMatMulOperation, - resize::ResizeOperation, slice_assign::SliceAssignOperation, - tensor::{TensorData, TensorLayoutInfo}, + tensor::TensorData, + view::ViewOperation, visit_tiled::MaybeQData, }; @@ -57,48 +53,37 @@ impl ComputeGraph { self.with_mut(|inner| inner.create_node(node)) } - pub(crate) fn create_nary(&self, op: NaryOperation) -> NodeIndex { - self.create_node(ComputeGraphNodeVariant::Nary(op)) - } - - pub(crate) fn create_mat_mul(&self, op: MatMulOperation) -> NodeIndex { - self.create_node(ComputeGraphNodeVariant::MatMul(op)) - } - - pub(crate) fn create_q_mat_mul(&self, op: QMatMulOperation) -> NodeIndex { - self.create_node(ComputeGraphNodeVariant::QMatMul(Box::new(op))) - } - - pub(crate) fn create_q_embedding(&self, op: QEmbeddingOperation) -> NodeIndex { - self.create_node(ComputeGraphNodeVariant::QEmbedding(op)) + pub(crate) fn create_nary(&self, op: ElementwiseOperation) -> NodeIndex { + self.create_node(ComputeGraphNodeVariant::Elementwise(op)) } pub(crate) fn create_reduce(&self, op: ReduceOperation) -> NodeIndex { self.create_node(ComputeGraphNodeVariant::Reduce(op)) } - pub(crate) fn create_graph_op(&self, op: Arc) -> NodeIndex { - self.create_node(ComputeGraphNodeVariant::GraphOp(op)) - } - - pub(crate) fn create_rms_norm(&self, op: RmsNormOperation) -> NodeIndex { - self.create_graph_op(Arc::new(op)) - } - - pub(crate) fn create_flash_attention(&self, op: FlashAttentionOperation) -> NodeIndex { - self.create_node(ComputeGraphNodeVariant::FlashAttention(op)) + pub(crate) fn create_view(&self, op: ViewOperation) -> NodeIndex { + self.create_node(ComputeGraphNodeVariant::View(op)) } - pub(crate) fn create_map_layout(&self, op: MapLayoutOperation) -> NodeIndex { - self.create_node(ComputeGraphNodeVariant::MapLayout(op)) + /// Build and submit one operation immediately against already-cached + /// inputs, bypassing the graph. Used by tuning APIs whose kernel + /// parameters cannot round-trip through the composed vocabulary. + pub(crate) fn execute_eager(&self, operation: &dyn Operation) -> Option { + self.with_mut(|inner| inner.execute_eager(operation)) } - pub(crate) fn create_resize(&self, op: ResizeOperation) -> NodeIndex { - self.create_node(ComputeGraphNodeVariant::Resize(op)) + /// Clone the view at `key` if that node is a view. Used to collapse view + /// chains at construction time. + pub(crate) fn get_view(&self, key: NodeIndex) -> Option { + let inner = self.inner.read(); + match &inner.nodes.nodes.node_weight(key)?.variant { + ComputeGraphNodeVariant::View(op) => Some(op.clone()), + _ => None, + } } pub(crate) fn create_slice_assign(&self, op: SliceAssignOperation) -> NodeIndex { - self.create_node(ComputeGraphNodeVariant::SliceAssign(op)) + self.create_node(ComputeGraphNodeVariant::Assign(op)) } pub(crate) fn create_tensor(&self, op: TensorData) -> NodeIndex { @@ -106,9 +91,9 @@ impl ComputeGraph { } pub(crate) fn dequantize(&self, matrix: QMatrix, ty: DataTypeEnum) -> NodeIndex { - self.create_node(ComputeGraphNodeVariant::Dequantize( - DequantizeOperation::new(matrix, ty), - )) + self.create_node(ComputeGraphNodeVariant::QMatrix(DequantizeOperation::new( + matrix, ty, + ))) } pub(crate) fn resolve(&self, key: NodeIndex) -> ResolverResult { @@ -310,70 +295,44 @@ impl ComputeGraphNode { } pub(crate) trait GraphOperation: Operation + Send + Sync { - fn as_any(&self) -> &dyn Any; - fn category(&self) -> &'static str; - - fn output_layout( - &self, - input_layouts: &FxHashMap, - ) -> Option; } +/// The graph vocabulary. Exactly three core operations — elementwise +/// visitation, reduction, and zero-dispatch views — over tensor and +/// quantized-matrix leaves, plus the in-place region write (pure data +/// movement, no compute dispatch). Everything else (matmul, attention, +/// normalization, embedding...) is a composition of these that the resolver +/// recognizes into fused execution regions. #[derive(Clone, Debug)] pub(crate) enum ComputeGraphNodeVariant { - Nary(NaryOperation), - SliceAssign(SliceAssignOperation), - Resize(ResizeOperation), - MapLayout(MapLayoutOperation), - Dequantize(DequantizeOperation), - QEmbedding(QEmbeddingOperation), - MatMul(MatMulOperation), - QMatMul(Box), Tensor(TensorData), + QMatrix(DequantizeOperation), + Elementwise(ElementwiseOperation), Reduce(ReduceOperation), - FlashAttention(FlashAttentionOperation), - GraphOp(Arc), + View(ViewOperation), + Assign(SliceAssignOperation), } impl ComputeGraphNodeVariant { fn visit_dependencies(&self, f: &mut dyn FnMut(NodeIndex)) { match &self { - ComputeGraphNodeVariant::Nary(op) => { + ComputeGraphNodeVariant::Elementwise(op) => { for input in &op.inputs { f(*input); } } - ComputeGraphNodeVariant::MatMul(op) => { - f(op.first); - f(op.second); - } - ComputeGraphNodeVariant::QMatMul(op) => { - f(op.input); - if let Some(epilogue) = &op.pre_element_wise_expr { - for extra in &epilogue.extras { - f(*extra); - } - } - if let Some(epilogue) = &op.post_element_wise_expr { - for extra in &epilogue.extras { - f(*extra); - } + ComputeGraphNodeVariant::Reduce(op) => { + for input in &op.inputs { + f(*input); } } - ComputeGraphNodeVariant::QEmbedding(op) => { - f(op.indexes); - } - ComputeGraphNodeVariant::Reduce(op) => f(op.value), - ComputeGraphNodeVariant::FlashAttention(op) => op.visit_dependencies(f), - ComputeGraphNodeVariant::GraphOp(op) => op.visit_dependencies(f), - ComputeGraphNodeVariant::MapLayout(op) => f(op.input), - ComputeGraphNodeVariant::Resize(op) => f(op.input), - ComputeGraphNodeVariant::SliceAssign(op) => { + ComputeGraphNodeVariant::View(op) => f(op.input), + ComputeGraphNodeVariant::Assign(op) => { f(op.input); f(op.value); } - ComputeGraphNodeVariant::Dequantize(_) => {} + ComputeGraphNodeVariant::QMatrix(_) => {} ComputeGraphNodeVariant::Tensor(_) => {} } } @@ -606,11 +565,31 @@ impl ComputeGraphInner { Some((output, total_kernels)) } - fn try_resolve_direct_qmatmul(&mut self, key: NodeIndex) -> Option { - let operation = match self.nodes.nodes.node_weight(key)?.variant.clone() { - ComputeGraphNodeVariant::QMatMul(operation) => operation, - _ => return None, + fn execute_eager(&mut self, operation: &dyn Operation) -> Option { + let device = self.device(); + let inputs = operation.inputs(self); + let workgroup_shape = operation + .workgroup_shape_constraints(&device) + .solve(device.max_subgroup_size(), &device.limits())?; + let kernel = operation.build_direct_kernel(self, &workgroup_shape, &inputs)?; + let MirValue::Tensor(output) = operation.output(self, &inputs) else { + return None; }; + + let mut command_encoder = + device + .wgpu_device() + .create_command_encoder(&wgpu::CommandEncoderDescriptor { + label: Some("Eager Operation Encoder"), + }); + kernel.run(device.kernel_cache(), &mut command_encoder); + device.wgpu_queue().submit(Some(command_encoder.finish())); + device.reset_initialized_buffers(); + Some(output) + } + + fn try_resolve_direct_qmatmul(&mut self, key: NodeIndex) -> Option { + let operation = self.match_direct_qmatmul(key)?; let (output, total_kernels) = self.try_submit_direct_qmatmul(&operation)?; self.set_cached_result(key, output.clone()); Some(ResolverResult { @@ -676,7 +655,7 @@ impl ComputeGraphInner { return Some(cached.clone().into()); } match &node.variant { - ComputeGraphNodeVariant::Dequantize(op) => Some(op.matrix.clone().into()), + ComputeGraphNodeVariant::QMatrix(op) => Some(op.matrix.clone().into()), ComputeGraphNodeVariant::Tensor(op) => Some(op.clone().into()), _ => None, } diff --git a/fusor-ml/core/src/compute_graph/resolve/cluster_match.rs b/fusor-ml/core/src/compute_graph/resolve/cluster_match.rs new file mode 100644 index 000000000..734896645 --- /dev/null +++ b/fusor-ml/core/src/compute_graph/resolve/cluster_match.rs @@ -0,0 +1,218 @@ +//! Shared expression/cluster matchers used by recognition and fusion passes. + +use crate::nary_wise::{NaryFunction, NaryOp}; +use crate::reduce::ReduceOp; + +use super::*; + +/// The keepdim-broadcast view layout a reduced tensor presents over its base: +/// row-major strides of the reduced shape with a `0` inserted at `axis`. +pub(super) fn keepdim_broadcast_layout(full_shape: &[usize], axis: usize) -> Layout { + let reduced: Vec = full_shape + .iter() + .enumerate() + .filter_map(|(dim, &size)| (dim != axis).then_some(size)) + .collect(); + let reduced_strides = Layout::continuous_strides(&reduced); + let mut strides = Vec::with_capacity(full_shape.len()); + let mut reduced_dim = 0; + for dim in 0..full_shape.len() { + if dim == axis { + strides.push(0); + } else { + strides.push(reduced_strides[reduced_dim]); + reduced_dim += 1; + } + } + Layout::from_parts(0, full_shape.into(), strides.into()) +} + +/// Layout equality modulo degenerate dimensions: strides of size-1 dims never +/// affect addressing, and view composition leaves arbitrary values there. +pub(super) fn layout_matches(actual: Option<&Layout>, expected: &Layout) -> bool { + let Some(actual) = actual else { + return false; + }; + actual.shape() == expected.shape() + && actual.offset() == expected.offset() + && actual + .shape() + .iter() + .zip(actual.strides()) + .zip(expected.strides()) + .all(|((&size, &actual), &expected)| size <= 1 || actual == expected) +} + +/// Match a single-input n-ary whose body is `f(input)` with full-rank +/// elementwise indices, returning the function and the input node. +pub(super) fn unary_elementwise(nary: &ElementwiseOperation) -> Option<(NaryFunction, NodeIndex)> { + let NaryExpr::Op { children, function } = &nary.expression else { + return None; + }; + let [ + NaryExpr::IndexedInput { + input_idx: 0, + indices, + }, + ] = children.as_slice() + else { + return None; + }; + (indices.len() == nary.shape.len() + && NaryExpr::is_elementwise_indices(indices) + && nary.inputs.len() == 1) + .then(|| (function.clone(), nary.inputs[0])) +} + +/// Match a two-input n-ary whose body is `op(input0, input1)` with full-rank +/// elementwise indices on both sides. +pub(super) fn binary_elementwise( + nary: &ElementwiseOperation, +) -> Option<(NaryOp, NodeIndex, NodeIndex)> { + let NaryExpr::Op { children, function } = &nary.expression else { + return None; + }; + let [ + NaryExpr::IndexedInput { + input_idx: 0, + indices: lhs, + }, + NaryExpr::IndexedInput { + input_idx: 1, + indices: rhs, + }, + ] = children.as_slice() + else { + return None; + }; + (lhs.len() == nary.shape.len() + && rhs.len() == nary.shape.len() + && NaryExpr::is_elementwise_indices(lhs) + && NaryExpr::is_elementwise_indices(rhs) + && nary.inputs.len() == 2) + .then(|| (function.op, nary.inputs[0], nary.inputs[1])) +} + +/// The composed softmax cluster matched without rewriting it. +pub(super) struct SoftmaxCluster { + pub(super) input: NodeIndex, + pub(super) shape: Box<[usize]>, + pub(super) axis: usize, +} + +impl Resolver { + pub(super) fn inner_nary(&self, inner: NodeIndex) -> Option<&ElementwiseOperation> { + let exec = self.get_input_node_in_exec_graph(inner)?; + match &self.execution_graph[exec].variant { + ExecutionVariant::Elementwise(nary) => Some(nary), + _ => None, + } + } + + pub(super) fn consumer_count_of(&self, inner: NodeIndex) -> Option { + let exec = self.get_input_node_in_exec_graph(inner)?; + Some( + self.execution_graph + .neighbors_directed(exec, petgraph::Direction::Outgoing) + .count(), + ) + } + + /// An intermediate node consumed by the recognized cluster alone: + /// exactly `expected` exec-graph consumers and no user-held reference or + /// cached value. + pub(super) fn exclusively_consumed( + &self, + graph: &ComputeGraphInner, + inner: NodeIndex, + expected: usize, + ) -> bool { + !graph.has_live_reference(inner) + && graph.get_cached_result(inner).is_none() + && self.consumer_count_of(inner) == Some(expected) + } + + /// A bare reduce of `op` with no fused chains: returns (axis, value). + pub(super) fn match_reduce( + &self, + inner: NodeIndex, + op: ReduceOp, + ) -> Option<(usize, NodeIndex)> { + let exec = self.get_input_node_in_exec_graph(inner)?; + let ExecutionVariant::Reduce(reduce) = &self.execution_graph[exec].variant else { + return None; + }; + let value = reduce.plain_input()?; + (reduce.function.op == op && reduce.post_element_wise.functions.is_empty()) + .then_some((reduce.axis, value)) + } + + /// A unary elementwise n-ary matching `accept`, returning its input. + pub(super) fn match_unary( + &self, + inner: NodeIndex, + accept: impl Fn(&NaryFunction) -> bool, + ) -> Option { + let nary = self.inner_nary(inner)?; + let (function, input) = unary_elementwise(nary)?; + accept(&function).then_some(input) + } + + /// Match the composed softmax cluster rooted at `probs_inner` (see + /// `Tensor::softmax`): `div(exp(sub(x, bcast(max(x, axis)))), + /// bcast(sum(exp, axis)))`, with every intermediate consumed exclusively + /// by the cluster. Used by attention recognition to see through the + /// probabilities without the cluster having been rewritten. + pub(super) fn match_softmax_cluster( + &self, + graph: &ComputeGraphInner, + probs_inner: NodeIndex, + ) -> Option { + let div = self.inner_nary(probs_inner)?; + let (div_op, exp_inner, sum_view_inner) = binary_elementwise(div)?; + if div_op != NaryOp::Div { + return None; + } + let shape = div.shape.clone(); + + let shifted_inner = self.match_unary(exp_inner, |function| function.op == NaryOp::Exp)?; + let (sub_op, x_inner, max_view_inner) = self + .inner_nary(shifted_inner) + .and_then(binary_elementwise)?; + if sub_op != NaryOp::Sub { + return None; + } + + // Denominator: bcast(sum(exp, axis)) reading the same exp node. + let (sum_base, sum_layout) = self.walk_view_chain(sum_view_inner); + let (axis, sum_value) = self.match_reduce(sum_base, ReduceOp::Sum)?; + if sum_value != exp_inner + || !layout_matches(sum_layout.as_ref(), &keepdim_broadcast_layout(&shape, axis)) + { + return None; + } + + // Shift: bcast(max(x, axis)) along the same axis, over the same x. + let (max_base, max_layout) = self.walk_view_chain(max_view_inner); + let (max_axis, max_value) = self.match_reduce(max_base, ReduceOp::Max)?; + if max_axis != axis + || max_value != x_inner + || !layout_matches(max_layout.as_ref(), &keepdim_broadcast_layout(&shape, axis)) + { + return None; + } + + // The whole cluster must exist solely for this softmax. + (self.exclusively_consumed(graph, exp_inner, 2) + && self.exclusively_consumed(graph, shifted_inner, 1) + && self.exclusively_consumed(graph, sum_view_inner, 1) + && self.exclusively_consumed(graph, max_view_inner, 1) + && self.exclusively_consumed(graph, sum_base, 1) + && self.exclusively_consumed(graph, max_base, 1)) + .then_some(SoftmaxCluster { + input: x_inner, + shape, + axis, + }) + } +} diff --git a/fusor-ml/core/src/compute_graph/resolve/execution.rs b/fusor-ml/core/src/compute_graph/resolve/execution.rs index f65dad157..54f0e185e 100644 --- a/fusor-ml/core/src/compute_graph/resolve/execution.rs +++ b/fusor-ml/core/src/compute_graph/resolve/execution.rs @@ -148,7 +148,7 @@ impl Resolver { // Add to execution graph let exec_idx = self.execution_graph.add_node(ExecutionNode { inner_idx: node, - variant: variant.clone(), + variant: variant.clone().into(), }); self.node_mapping.insert(node, exec_idx); @@ -167,39 +167,68 @@ impl Resolver { Some(exec_idx) } - pub(super) fn lower_node(&self, node: &ExecutionNode) -> Option { + pub(super) fn lower_node( + &self, + exec_idx: ExecutionNodeIndex, + node: &ExecutionNode, + ) -> Option { match &node.variant { - ComputeGraphNodeVariant::Nary(op) => { - Some(QueuedOperation::Generic(Arc::new(op.clone()))) - } - ComputeGraphNodeVariant::MatMul(op) => { - Some(QueuedOperation::Generic(Arc::new(op.clone()))) - } - ComputeGraphNodeVariant::Reduce(op) => { - Some(QueuedOperation::Generic(Arc::new(op.clone()))) - } - ComputeGraphNodeVariant::FlashAttention(op) => { - Some(QueuedOperation::Generic(Arc::new(op.clone()))) - } - ComputeGraphNodeVariant::GraphOp(op) => Some(QueuedOperation::Generic(op.clone())), - ComputeGraphNodeVariant::MapLayout(op) => { + ExecutionVariant::Elementwise(op) => { Some(QueuedOperation::Generic(Arc::new(op.clone()))) } - ComputeGraphNodeVariant::Resize(op) => { + ExecutionVariant::MatMul(op) => Some(QueuedOperation::Generic(Arc::new(op.clone()))), + ExecutionVariant::Reduce(op) => Some(QueuedOperation::Generic(Arc::new(op.clone()))), + ExecutionVariant::GraphOp(op) => Some(QueuedOperation::Generic(op.clone())), + ExecutionVariant::View(op) => Some(QueuedOperation::Generic(Arc::new(op.clone()))), + ExecutionVariant::Assign(op) => Some(QueuedOperation::Generic(Arc::new(op.clone()))), + ExecutionVariant::QEmbedding(op) => { Some(QueuedOperation::Generic(Arc::new(op.clone()))) } - ComputeGraphNodeVariant::SliceAssign(op) => { - Some(QueuedOperation::Generic(Arc::new(op.clone()))) - } - ComputeGraphNodeVariant::QEmbedding(op) => { + ExecutionVariant::QMatMul(op) => Some(QueuedOperation::QMatMul(op.clone())), + ExecutionVariant::QMatrix(op) => { + // Skip materializing the dense tensor when every consumer + // reads the block-quantized data directly (fused reduces and + // elementwise expressions decode per element; qmatmul and + // embedding kernels decode per block). + if self.qmatrix_consumed_raw(exec_idx, node.inner_idx) { + return None; + } Some(QueuedOperation::Generic(Arc::new(op.clone()))) } - ComputeGraphNodeVariant::QMatMul(op) => Some(QueuedOperation::QMatMul(op.clone())), - ComputeGraphNodeVariant::Dequantize(op) => { - Some(QueuedOperation::Generic(Arc::new(op.clone()))) + ExecutionVariant::Tensor(_) => None, // Handled in execution loop + } + } + + /// Whether every consumer of a quantized-matrix node reads the raw + /// blocks (region ops decode per block; expression reads decode per + /// element) rather than needing the dense tensor. Only custom-indexed + /// expression reads require the dense form. + fn qmatrix_consumed_raw(&self, exec_idx: ExecutionNodeIndex, inner: NodeIndex) -> bool { + let mut any = false; + for consumer in self + .execution_graph + .neighbors_directed(exec_idx, petgraph::Direction::Outgoing) + { + any = true; + let raw = match &self.execution_graph[consumer].variant { + ExecutionVariant::QMatMul(_) | ExecutionVariant::QEmbedding(_) => true, + ExecutionVariant::Elementwise(nary) => { + !nary.inputs.iter().enumerate().any(|(slot, &input)| { + input == inner && nary.expression.uses_custom_indexing_for_input(slot) + }) + } + ExecutionVariant::Reduce(reduce) => { + !reduce.inputs.iter().enumerate().any(|(slot, &input)| { + input == inner && reduce.expression.uses_custom_indexing_for_input(slot) + }) + } + _ => false, + }; + if !raw { + return false; } - ComputeGraphNodeVariant::Tensor(_) => None, // Handled in execution loop } + any } // --- Rewrite Engine --- @@ -207,32 +236,36 @@ impl Resolver { pub(super) fn optimize(&mut self, graph: &mut ComputeGraphInner) { let profile_enabled = std::env::var_os("FUSOR_TRACE_OPTIMIZE").is_some(); let mut profile = OptimizeProfile::default(); + // Rebuild composed contraction / normalization clusters into their + // specialized operations first, while they are still in the exact + // canonical form the API emitted (before view folding or fusion + // disturbs them). + self.recognize_contractions(graph); + self.recognize_embeddings(graph); + self.recognize_attention(graph); + self.fuse_row_programs(graph); // The current rewrite rules can only start from Nary nodes (nary // fusion, post-op reduce/matmul fusion) or MatMul nodes (pre-op - // unary fusion). Avoid scanning every QMatMul/RMS/attention node in + // unary fusion). Avoid scanning every QMatMul/attention node in // decode graphs with hundreds of kernels. let has_reduce = self.execution_graph.node_indices().any(|node| { matches!( self.execution_graph[node].variant, - ComputeGraphNodeVariant::Reduce(_) + ExecutionVariant::Reduce(_) ) }); let has_matmul = self.execution_graph.node_indices().any(|node| { matches!( self.execution_graph[node].variant, - ComputeGraphNodeVariant::MatMul(_) + ExecutionVariant::MatMul(_) ) }); let has_qmatmul = self.execution_graph.node_indices().any(|node| { matches!( self.execution_graph[node].variant, - ComputeGraphNodeVariant::QMatMul(_) + ExecutionVariant::QMatMul(_) ) }); - let has_rmsnorm = self - .execution_graph - .node_indices() - .any(|node| as_rms_norm(&self.execution_graph[node].variant).is_some()); let allow_qmatmul_elementwise_fusion = self.execution_graph.node_count() <= DEFAULT_OPTIMIZE_NODE_LIMIT || std::env::var_os("FUSOR_RESOLVE_QMATMUL_ELEMENTWISE_FUSION").is_some(); @@ -258,10 +291,14 @@ impl Resolver { .neighbors_directed(node_idx, petgraph::Direction::Outgoing) .collect(); - // 1. Fuse naries together (combine expression trees) - // 2. Try to fuse resulting nary into specialized ops (reduce, matmul, etc.) + // 1. Fold view inputs into the nary body so fusion sees through + // layout changes + // 2. Fuse naries together (combine expression trees) + // 3. Try to fuse resulting nary into specialized ops (reduce, matmul, etc.) + let changed = self.try_fold_view_inputs(graph, node_idx); + let start = profile_enabled.then(Instant::now); - let changed = self.try_fuse_naries(graph, node_idx); + let changed = changed | self.try_fuse_naries(graph, node_idx); if let Some(start) = start { profile.fuse_naries_count += 1; profile.fuse_naries += start.elapsed(); @@ -271,7 +308,9 @@ impl Resolver { true } else { let start = profile_enabled.then(Instant::now); - let changed = has_reduce && self.try_fuse_into_reduce(graph, node_idx); + let changed = has_reduce + && (self.try_fuse_into_reduce(graph, node_idx) + || self.try_fuse_producer_into_reduce(graph, node_idx)); if let Some(start) = start { profile.fuse_reduce_count += 1; profile.fuse_reduce += start.elapsed(); @@ -292,18 +331,6 @@ impl Resolver { changed }; - let changed = if changed { - true - } else { - let start = profile_enabled.then(Instant::now); - let changed = has_rmsnorm && self.try_fuse_into_rmsnorm(graph, node_idx); - if let Some(start) = start { - profile.fuse_rmsnorm_count += 1; - profile.fuse_rmsnorm += start.elapsed(); - } - changed - }; - if changed { profile.changed += 1; // Re-add the current node to worklist if it still exists @@ -345,10 +372,14 @@ impl Resolver { } pub(super) fn optimize_large_graph(&mut self, graph: &mut ComputeGraphInner) { + self.recognize_contractions(graph); + self.recognize_embeddings(graph); + self.recognize_attention(graph); + self.fuse_row_programs(graph); let has_qmatmul = self.execution_graph.node_indices().any(|node| { matches!( self.execution_graph[node].variant, - ComputeGraphNodeVariant::QMatMul(_) + ExecutionVariant::QMatMul(_) ) }); if !has_qmatmul { @@ -372,7 +403,8 @@ impl Resolver { .execution_graph .neighbors_directed(node_idx, petgraph::Direction::Outgoing) .collect::>(); - let mut changed = self.try_fuse_naries(graph, node_idx); + let mut changed = self.try_fold_view_inputs(graph, node_idx); + changed |= self.try_fuse_naries(graph, node_idx); if !changed && self.execution_graph.contains_node(node_idx) { changed = self.try_fuse_into_matmul(graph, node_idx, true); } @@ -419,7 +451,7 @@ impl Resolver { } } else if matches!( self.execution_graph[node].variant, - ComputeGraphNodeVariant::MapLayout(_) + ExecutionVariant::View(_) ) { stack.extend( self.execution_graph @@ -430,7 +462,7 @@ impl Resolver { } pub(super) fn is_large_graph_nary_candidate(&self, node_idx: ExecutionNodeIndex) -> bool { - let ComputeGraphNodeVariant::Nary(nary) = &self.execution_graph[node_idx].variant else { + let ExecutionVariant::Elementwise(nary) = &self.execution_graph[node_idx].variant else { return false; }; if nary.shape.last().copied().unwrap_or_default() >= LARGE_GRAPH_NARY_FUSION_MIN_LAST_DIM { @@ -438,14 +470,12 @@ impl Resolver { } nary.inputs.iter().any(|&input| { - let (base_inner, _) = self - .walk_map_layout_chain(input) - .unwrap_or((input, Vec::new())); + let (base_inner, _) = self.walk_view_chain(input); self.get_input_node_in_exec_graph(base_inner) .is_some_and(|exec_idx| { matches!( self.execution_graph[exec_idx].variant, - ComputeGraphNodeVariant::QMatMul(_) + ExecutionVariant::QMatMul(_) ) }) }) @@ -455,8 +485,7 @@ impl Resolver { let mut qmatmul_count = 0usize; let mut single_token_count = 0usize; for node in self.execution_graph.node_indices() { - let ComputeGraphNodeVariant::QMatMul(qmatmul) = &self.execution_graph[node].variant - else { + let ExecutionVariant::QMatMul(qmatmul) = &self.execution_graph[node].variant else { continue; }; qmatmul_count += 1; @@ -475,10 +504,11 @@ impl Resolver { pub(super) fn is_optimization_candidate(&self, node_idx: ExecutionNodeIndex) -> bool { matches!( self.execution_graph[node_idx].variant, - ComputeGraphNodeVariant::Nary(_) - | ComputeGraphNodeVariant::MatMul(_) - | ComputeGraphNodeVariant::QMatMul(_) - ) || as_rms_norm(&self.execution_graph[node_idx].variant).is_some() + ExecutionVariant::Elementwise(_) + | ExecutionVariant::MatMul(_) + | ExecutionVariant::QMatMul(_) + | ExecutionVariant::Reduce(_) + ) } // Helpers @@ -501,46 +531,44 @@ impl Resolver { self.node_mapping.get(&inner_input).copied() } - pub(super) fn walk_map_layout_chain( - &self, - mut inner: NodeIndex, - ) -> Option<(NodeIndex, Vec)> { - let mut chain = Vec::new(); + /// Walk through fully-defined view nodes from `inner` down to the first + /// non-view node, composing the view layouts. Returns the base node and + /// the composed view layout over the base's logical value space; the + /// layout is `None` when `inner` is not a view (identity). Views that + /// don't compose (or carry a fill region) act as chain breaks: the walk + /// stops without seeing through them. + pub(super) fn walk_view_chain(&self, mut inner: NodeIndex) -> (NodeIndex, Option) { + let mut composed: Option = None; loop { let Some(exec) = self.get_input_node_in_exec_graph(inner) else { - if chain.is_empty() { - return None; - } - chain.reverse(); - return Some((inner, chain)); + return (inner, composed); + }; + let ExecutionVariant::View(view) = &self.execution_graph[exec].variant else { + return (inner, composed); }; - let ComputeGraphNodeVariant::MapLayout(map) = - self.execution_graph[exec].variant.clone() - else { - chain.reverse(); - return Some((inner, chain)); + if !view.is_fully_defined() { + return (inner, composed); + } + let next = match &composed { + None => view.layout.clone(), + Some(outer) => match crate::view::compose_layouts(outer, &view.layout) { + Some(layout) => layout, + None => return (inner, composed), + }, }; - inner = map.input; - chain.push(map); + composed = Some(next); + inner = view.input; } } - pub(super) fn apply_map_layout_chain( - base: &Layout, - chain: &[crate::map_layout::MapLayoutOperation], - ) -> Layout { - chain - .iter() - .fold(base.clone(), |layout, map| map.map_layout(&layout)) - } - - pub(super) fn infer_layout( - graph: &ComputeGraphInner, - inner_idx: NodeIndex, - ) -> Option { - let mut pass = LayoutPass::default(); - pass.visit(graph, inner_idx); - pass.output_layout.remove(&inner_idx) + /// The layout a (possibly chained-view) node presents over its base + /// node's value space when the base materializes at `base_layout`. + /// `None` when the view does not compose with that layout. + pub(super) fn apply_view_chain(base_layout: &Layout, chain: &Option) -> Option { + match chain { + None => Some(base_layout.clone()), + Some(view) => crate::view::compose_layouts(view, base_layout), + } } pub(super) fn infer_layout_cached( @@ -548,12 +576,8 @@ impl Resolver { graph: &ComputeGraphInner, inner_idx: NodeIndex, ) -> Option { - if let Some(layout) = self.layout_cache.get(&inner_idx) { - return layout.clone(); - } - let layout = Self::infer_layout(graph, inner_idx); - self.layout_cache.insert(inner_idx, layout.clone()); - layout + self.layout_pass.visit(graph, inner_idx); + self.layout_pass.output_layout.get(&inner_idx).cloned() } pub(super) fn try_normalize_qmatmul_post_extra( @@ -580,9 +604,7 @@ impl Resolver { return Some(extra_inner); } - let (base_inner, _) = self - .walk_map_layout_chain(extra_inner) - .unwrap_or((extra_inner, Vec::new())); + let (base_inner, _) = self.walk_view_chain(extra_inner); let base_info = self.infer_layout_cached(graph, base_inner)?; let base_layout = base_info.layout(); if base_info.datatype() == DataTypeEnum::F32 diff --git a/fusor-ml/core/src/compute_graph/resolve/fold_views.rs b/fusor-ml/core/src/compute_graph/resolve/fold_views.rs new file mode 100644 index 000000000..64e791e96 --- /dev/null +++ b/fusor-ml/core/src/compute_graph/resolve/fold_views.rs @@ -0,0 +1,168 @@ +use crate::view::{AffineIndex, ViewOperation, affine_dim_indices}; + +use super::*; + +impl Resolver { + /// Fold view inputs of an n-ary node directly into its expression: each + /// `IndexedInput` through the view becomes an `IndexedInput` of the + /// view's base node with the view's coordinate mapping applied (and a + /// bounds-select around partially-defined views). This removes the view + /// node from between the producer and consumer, so the n-ary fusion + /// passes see through layout changes instead of stopping at them. + /// + /// Only affine views fold (no divmod in the rewritten indices); anything + /// else stays a view node and materializes through the gather fallback. + pub(super) fn try_fold_view_inputs( + &mut self, + graph: &mut ComputeGraphInner, + node_idx: ExecutionNodeIndex, + ) -> bool { + let ExecutionVariant::Elementwise(nary) = self.execution_graph[node_idx].variant.clone() + else { + return false; + }; + + let mut expression = nary.expression.clone(); + let mut inputs = nary.inputs.clone(); + let mut folded = Vec::new(); + + for (slot, input_inner) in nary.inputs.iter().copied().enumerate() { + if self.check_cached(graph, input_inner) { + continue; + } + let Some(input_exec) = self.get_input_node_in_exec_graph(input_inner) else { + continue; + }; + if !self.execution_graph.contains_node(input_exec) { + continue; + } + let ExecutionVariant::View(view) = &self.execution_graph[input_exec].variant else { + continue; + }; + let Some(affine) = affine_dim_indices(&view.layout, &view.input_shape) else { + continue; + }; + let view = view.clone(); + expression = Self::rewrite_view_input(&expression, slot, &view, &affine); + inputs[slot] = view.input; + folded.push((input_exec, view.input)); + } + if folded.is_empty() { + return false; + } + + let (final_inputs, final_expression) = Self::deduplicate_inputs(inputs, expression); + let new_nary = ElementwiseOperation { + inputs: final_inputs, + expression: final_expression, + shape: nary.shape.clone(), + output_datatype: nary.output_datatype, + }; + self.execution_graph[node_idx].variant = ExecutionVariant::Elementwise(new_nary.clone()); + + for (view_exec, base_inner) in &folded { + if let Some(edge) = self.execution_graph.find_edge(*view_exec, node_idx) { + self.execution_graph.remove_edge(edge); + } + if let Some(base_exec) = self.get_input_node_in_exec_graph(*base_inner) + && self + .execution_graph + .find_edge(base_exec, node_idx) + .is_none() + { + self.execution_graph.add_edge(base_exec, node_idx, ()); + } + } + self.add_physical_dependencies(graph, node_idx, &new_nary.inputs); + for (view_exec, _) in folded { + self.remove_node_if_dead(view_exec); + } + true + } + + /// Rewrite every access to input `target_idx` through `view`'s coordinate + /// mapping. The original index expressions (the view's output + /// coordinates) feed the affine per-base-dimension indices; partially + /// defined views wrap the load in `select(in_bounds, load, fill)` with + /// the load coordinates clamped in-bounds (both select branches + /// evaluate). + fn rewrite_view_input( + expr: &NaryExpr, + target_idx: usize, + view: &ViewOperation, + affine: &[AffineIndex], + ) -> NaryExpr { + match expr { + NaryExpr::Op { children, function } => NaryExpr::Op { + children: children + .iter() + .map(|child| Self::rewrite_view_input(child, target_idx, view, affine)) + .collect(), + function: function.clone(), + }, + NaryExpr::IndexedInput { input_idx, indices } => { + let indices: Vec = indices + .iter() + .map(|index| Self::rewrite_view_input(index, target_idx, view, affine)) + .collect(); + if *input_idx != target_idx { + return NaryExpr::IndexedInput { + input_idx: *input_idx, + indices, + }; + } + + let fully_defined = view.is_fully_defined(); + let base_indices = affine + .iter() + .zip(&*view.input_shape) + .map(|(index, &extent)| { + let index = index.to_expr(&indices); + if fully_defined || extent == 0 { + index + } else { + NaryExpr::unary_op( + index, + "clamp_dim", + NaryOp::MinConst(NaryScalar::U32(extent as u32 - 1)), + DataTypeEnum::U32, + DataTypeEnum::U32, + ) + } + }) + .collect(); + let loaded = NaryExpr::IndexedInput { + input_idx: *input_idx, + indices: base_indices, + }; + if fully_defined { + return loaded; + } + + let mut condition = NaryExpr::scalar(NaryScalar::U32(1)); + for (dim, (&defined, &size)) in view.defined.iter().zip(view.shape()).enumerate() { + if defined >= size { + continue; + } + let lt_defined = NaryExpr::unary_op( + indices[dim].clone(), + "lt_defined", + NaryOp::LessConst(NaryScalar::U32(defined as u32)), + DataTypeEnum::U32, + DataTypeEnum::U32, + ); + condition = NaryExpr::mul(condition, lt_defined, DataTypeEnum::U32); + } + NaryExpr::select( + condition, + loaded, + NaryExpr::scalar(view.fill), + DataTypeEnum::U32, + view.datatype, + ) + } + NaryExpr::DimIndex(dim) => NaryExpr::DimIndex(*dim), + NaryExpr::Scalar(value) => NaryExpr::Scalar(*value), + } + } +} diff --git a/fusor-ml/core/src/compute_graph/resolve/fusion_basic.rs b/fusor-ml/core/src/compute_graph/resolve/fusion_basic.rs index a398443a1..89eca64d1 100644 --- a/fusor-ml/core/src/compute_graph/resolve/fusion_basic.rs +++ b/fusor-ml/core/src/compute_graph/resolve/fusion_basic.rs @@ -8,7 +8,7 @@ impl Resolver { ) -> bool { let node_variant = self.execution_graph[node_idx].variant.clone(); - let ComputeGraphNodeVariant::Nary(nary) = node_variant else { + let ExecutionVariant::Elementwise(nary) = node_variant else { return false; }; @@ -30,7 +30,21 @@ impl Resolver { if !self.execution_graph.contains_node(input_exec) { continue; } - let ComputeGraphNodeVariant::Nary(input_nary) = + // Inlining duplicates the producer's work unless this node is its + // only consumer: the producer still materializes for everyone + // else (e.g. the residual stream feeds every later layer — fusing + // it forward would re-sum the whole prefix per layer). A user-held + // reference alone doesn't block fusion — only another consumer in + // this resolve does. + if self + .execution_graph + .neighbors_directed(input_exec, petgraph::Direction::Outgoing) + .count() + != 1 + { + continue; + } + let ExecutionVariant::Elementwise(input_nary) = &self.execution_graph[input_exec].variant else { continue; @@ -86,14 +100,14 @@ impl Resolver { // Deduplicate and remove unused inputs let (final_inputs, final_expression) = Self::deduplicate_inputs(all_inputs, expression); - let new_nary = NaryOperation { + let new_nary = ElementwiseOperation { inputs: final_inputs.clone(), expression: final_expression, shape: nary.shape.clone(), output_datatype: nary.output_datatype, }; - self.execution_graph[node_idx].variant = ComputeGraphNodeVariant::Nary(new_nary.clone()); + self.execution_graph[node_idx].variant = ExecutionVariant::Elementwise(new_nary.clone()); // Update graph edges for (input_exec, new_inputs) in fused_execs { @@ -390,11 +404,9 @@ impl Resolver { /// Try to extract a unary function chain from a node variant. /// Only Nary ops with a single input and element-wise access can be converted. - pub(super) fn try_get_unary_chain( - variant: &ComputeGraphNodeVariant, - ) -> Option { + pub(super) fn try_get_unary_chain(variant: &ExecutionVariant) -> Option { match variant { - ComputeGraphNodeVariant::Nary(nary) => nary.try_extract_unary_chain(), + ExecutionVariant::Elementwise(nary) => nary.try_extract_unary_chain(), _ => None, } } @@ -420,7 +432,7 @@ impl Resolver { }; let input_variant = self.execution_graph[input_exec_idx].variant.clone(); - let ComputeGraphNodeVariant::Reduce(reduce_op) = input_variant else { + let ExecutionVariant::Reduce(reduce_op) = input_variant else { return false; }; @@ -430,76 +442,133 @@ impl Resolver { new_reduce.post_element_wise = UnaryFunctionChain::new(existing_post, reduce_op.post_element_wise.input_datatype()); - self.execution_graph[node_idx].variant = - ComputeGraphNodeVariant::Reduce(new_reduce.clone()); + self.execution_graph[node_idx].variant = ExecutionVariant::Reduce(new_reduce.clone()); - let reduce_input_inner = reduce_op.value; - if let Some(reduce_input_exec) = self.get_input_node_in_exec_graph(reduce_input_inner) { - self.execution_graph - .add_edge(reduce_input_exec, node_idx, ()); + for &reduce_input_inner in &reduce_op.inputs { + if let Some(reduce_input_exec) = self.get_input_node_in_exec_graph(reduce_input_inner) { + self.execution_graph + .add_edge(reduce_input_exec, node_idx, ()); + } } if let Some(edge) = self.execution_graph.find_edge(input_exec_idx, node_idx) { self.execution_graph.remove_edge(edge); } - self.add_physical_dependencies(graph, node_idx, &[reduce_input_inner]); + self.add_physical_dependencies(graph, node_idx, &reduce_op.inputs); self.remove_node_if_dead(input_exec_idx); true } - pub(super) fn try_fuse_into_rmsnorm( + /// Inline elementwise producers into a reduce's fused expression: the + /// reduce evaluates the producer at every index-space coordinate, so a + /// producer consumed only by this reduce never needs to materialize. + /// Composed contractions that recognition did not claim collapse to a + /// single map-reduce kernel here, where the tiled lowering can stage + /// their reused inputs through workgroup memory. + pub(super) fn try_fuse_producer_into_reduce( &mut self, graph: &mut ComputeGraphInner, node_idx: ExecutionNodeIndex, ) -> bool { - let ComputeGraphNodeVariant::Nary(_) = &self.execution_graph[node_idx].variant else { - return false; - }; - let node_variant = self.execution_graph[node_idx].variant.clone(); - let Some(el_op) = Self::try_get_unary_chain(&node_variant) else { + let ExecutionVariant::Reduce(reduce) = self.execution_graph[node_idx].variant.clone() + else { return false; }; - let input_inner = el_op.value; - if self.check_cached(graph, input_inner) { - return false; + + let mut expression = reduce.expression.clone(); + let mut all_inputs = reduce.inputs.clone(); + let mut fused_execs = Vec::new(); + let max_fused_inputs = graph.device().nary_direct_input_binding_budget(); + + for &input_inner in reduce.inputs.iter() { + if self.check_cached(graph, input_inner) { + continue; + } + let Some(input_exec) = self.get_input_node_in_exec_graph(input_inner) else { + continue; + }; + if !self.execution_graph.contains_node(input_exec) { + continue; + } + // Same sole-consumer rule as nary fusion: inlining a shared + // producer would re-evaluate it for every other consumer. + if self + .execution_graph + .neighbors_directed(input_exec, petgraph::Direction::Outgoing) + .count() + != 1 + { + continue; + } + let ExecutionVariant::Elementwise(input_nary) = + &self.execution_graph[input_exec].variant + else { + continue; + }; + // The reduce evaluates this input across the full index space; + // a producer with any other shape reads out of range. + if input_nary.shape != reduce.shape { + continue; + } + + let offset = all_inputs.len(); + let inlined = Self::offset_input_indices(&input_nary.expression, offset); + let target_slots: Vec = all_inputs + .iter() + .enumerate() + .filter_map(|(slot, value)| (*value == input_inner).then_some(slot)) + .collect(); + let mut new_expression = expression.clone(); + let mut success = true; + for slot in target_slots { + let (next, s) = Self::substitute_input_in_expr(&new_expression, slot, &inlined); + new_expression = next; + success &= s; + } + + if success { + let unique_inputs: FxHashSet<_> = all_inputs + .iter() + .chain(input_nary.inputs.iter()) + .copied() + .collect(); + if unique_inputs.len() > max_fused_inputs { + continue; + } + + expression = new_expression; + all_inputs.extend(input_nary.inputs.iter().copied()); + fused_execs.push((input_exec, input_nary.inputs.clone())); + } } - let Some(input_exec_idx) = self.get_input_node_in_exec_graph(input_inner) else { - return false; - }; - let input_variant = self.execution_graph[input_exec_idx].variant.clone(); - let Some(rms_op) = as_rms_norm(&input_variant) else { + + if fused_execs.is_empty() { return false; - }; - let mut new_rms = rms_op.clone(); - let mut existing = new_rms.post_element_wise.functions.clone(); - existing.extend(el_op.functions.functions.iter().cloned()); - new_rms.post_element_wise = - UnaryFunctionChain::new(existing, rms_op.post_element_wise.input_datatype()); - - self.execution_graph[node_idx].variant = - ComputeGraphNodeVariant::GraphOp(Arc::new(new_rms.clone())); - - // Re-wire dependency edges to the fused RmsNorm inputs. - let mut deps = vec![new_rms.input]; - if let Some(residual) = new_rms.residual { - deps.push(residual); - } - deps.push(new_rms.weight); - if let Some(bias) = new_rms.bias { - deps.push(bias); } - for &dep in &deps { - if let Some(idx) = self.get_input_node_in_exec_graph(dep) - && self.execution_graph.find_edge(idx, node_idx).is_none() - { - self.execution_graph.add_edge(idx, node_idx, ()); + + let (final_inputs, final_expression) = Self::deduplicate_inputs(all_inputs, expression); + + let mut new_reduce = reduce.clone(); + new_reduce.inputs = final_inputs; + new_reduce.expression = final_expression; + let new_inputs = new_reduce.inputs.clone(); + self.execution_graph[node_idx].variant = ExecutionVariant::Reduce(new_reduce); + + for (input_exec, producer_inputs) in fused_execs { + if let Some(edge) = self.execution_graph.find_edge(input_exec, node_idx) { + self.execution_graph.remove_edge(edge); } + for &new_input in &producer_inputs { + if let Some(exec) = self.get_input_node_in_exec_graph(new_input) + && self.execution_graph.find_edge(exec, node_idx).is_none() + { + self.execution_graph.add_edge(exec, node_idx, ()); + } + } + self.remove_node_if_dead(input_exec); } - if let Some(edge) = self.execution_graph.find_edge(input_exec_idx, node_idx) { - self.execution_graph.remove_edge(edge); - } - self.add_physical_dependencies(graph, node_idx, &deps); - self.remove_node_if_dead(input_exec_idx); + + self.add_physical_dependencies(graph, node_idx, &new_inputs); true } } diff --git a/fusor-ml/core/src/compute_graph/resolve/fusion_matmul.rs b/fusor-ml/core/src/compute_graph/resolve/fusion_matmul.rs index f43780cba..55b6b8ae6 100644 --- a/fusor-ml/core/src/compute_graph/resolve/fusion_matmul.rs +++ b/fusor-ml/core/src/compute_graph/resolve/fusion_matmul.rs @@ -16,7 +16,7 @@ impl Resolver { && let Some(input_exec_idx) = self.get_input_node_in_exec_graph(input_inner) { let input_variant = self.execution_graph[input_exec_idx].variant.clone(); - if let ComputeGraphNodeVariant::MatMul(matmul_op) = input_variant { + if let ExecutionVariant::MatMul(matmul_op) = input_variant { let mut new_matmul = matmul_op.clone(); let mut existing_post = new_matmul.post_element_wise.functions.clone(); existing_post.extend(el_op.functions.functions.iter().cloned()); @@ -26,7 +26,7 @@ impl Resolver { ); self.execution_graph[node_idx].variant = - ComputeGraphNodeVariant::MatMul(new_matmul.clone()); + ExecutionVariant::MatMul(new_matmul.clone()); let (first_inner, second_inner) = (matmul_op.first, matmul_op.second); if let Some(idx) = self.get_input_node_in_exec_graph(first_inner) { @@ -49,7 +49,7 @@ impl Resolver { // qmatmul. This handles composite expressions like GELU and ordered // extra inputs whose layouts match the output visitation shape. if allow_qmatmul_elementwise_fusion - && let ComputeGraphNodeVariant::Nary(nary) = &node_variant + && let ExecutionVariant::Elementwise(nary) = &node_variant { // Split/gate expressions built from `narrow` views of a qmatmul // output (e.g. SwiGLU's gate/up halves) reach the qmatmul through @@ -63,19 +63,17 @@ impl Resolver { if self.get_input_node_in_exec_graph(input_inner).is_none() { continue; } - let (qmatmul_inner, map_chain) = self - .walk_map_layout_chain(input_inner) - .unwrap_or((input_inner, Vec::new())); + let (qmatmul_inner, map_chain) = self.walk_view_chain(input_inner); let Some(qmatmul_exec_idx) = self.get_input_node_in_exec_graph(qmatmul_inner) else { continue; }; - let ComputeGraphNodeVariant::QMatMul(qmatmul_op) = + let ExecutionVariant::QMatMul(qmatmul_op) = self.execution_graph[qmatmul_exec_idx].variant.clone() else { continue; }; - if map_chain.is_empty() + if map_chain.is_none() && !self.check_cached(graph, input_inner) && qmatmul_op.post_element_wise_expr.is_none() && qmatmul_op.in_shape[..qmatmul_op.in_shape.len() - 1] @@ -128,10 +126,11 @@ impl Resolver { self.commit_qmatmul_post_fusion(graph, node_idx, &nary.inputs, new_q); return true; } - let mapped_layout = Self::apply_map_layout_chain( - &Layout::contiguous(&qmatmul_op.out_shape), - &map_chain, - ); + let Some(mapped_layout) = + Self::apply_view_chain(&Layout::contiguous(&qmatmul_op.out_shape), &map_chain) + else { + continue; + }; if mapped_layout != Layout::contiguous(&nary.shape) { continue; } @@ -152,23 +151,21 @@ impl Resolver { let mut replacements = vec![None; nary.inputs.len()]; let mut valid_expression = true; for (input_idx, &nary_input) in nary.inputs.iter().enumerate() { - let (base_inner, chain) = self - .walk_map_layout_chain(nary_input) - .unwrap_or((nary_input, Vec::new())); + let (base_inner, chain) = self.walk_view_chain(nary_input); let base_qmatmul = self.get_input_node_in_exec_graph(base_inner) .and_then(|exec| match &self.execution_graph[exec].variant { - ComputeGraphNodeVariant::QMatMul(op) => Some(op.clone()), + ExecutionVariant::QMatMul(op) => Some(op.clone()), _ => None, }); if let Some(base_qmatmul) = base_qmatmul && Self::qmatmul_same_base(&qmatmul_op, &base_qmatmul) { - let alias_layout = Self::apply_map_layout_chain( + let alias_layout = Self::apply_view_chain( &Layout::contiguous(&base_qmatmul.out_shape), &chain, ); - if alias_layout == Layout::contiguous(&nary.shape) + if alias_layout == Some(Layout::contiguous(&nary.shape)) && !nary.expression.uses_custom_indexing_for_input(input_idx) { replacements[input_idx] = Self::qmatmul_output_expr( @@ -236,7 +233,7 @@ impl Resolver { // tile, so expensive expressions like GELU would be recomputed many // times. Keep those chains materialized once instead. if allow_qmatmul_elementwise_fusion - && let ComputeGraphNodeVariant::QMatMul(qmatmul_op) = &node_variant + && let ExecutionVariant::QMatMul(qmatmul_op) = &node_variant && qmatmul_op.in_shape[..qmatmul_op.in_shape.len() - 1] .iter() .product::() @@ -244,20 +241,18 @@ impl Resolver { && !self.check_cached(graph, qmatmul_op.input) && let Some(input_exec) = self.get_input_node_in_exec_graph(qmatmul_op.input) { - let (nary_inner, nary_map_chain) = self - .walk_map_layout_chain(qmatmul_op.input) - .unwrap_or((qmatmul_op.input, Vec::new())); + let (nary_inner, nary_map_chain) = self.walk_view_chain(qmatmul_op.input); let Some(nary_exec) = self.get_input_node_in_exec_graph(nary_inner) else { return false; }; - let ComputeGraphNodeVariant::Nary(nary) = + let ExecutionVariant::Elementwise(nary) = self.execution_graph[nary_exec].variant.clone() else { return false; }; let mapped_layout = - Self::apply_map_layout_chain(&Layout::contiguous(&nary.shape), &nary_map_chain); - if mapped_layout != Layout::contiguous(&qmatmul_op.in_shape) { + Self::apply_view_chain(&Layout::contiguous(&nary.shape), &nary_map_chain); + if mapped_layout != Some(Layout::contiguous(&qmatmul_op.in_shape)) { return false; } @@ -281,14 +276,15 @@ impl Resolver { continue; } - let (primary_inner, primary_chain) = self - .walk_map_layout_chain(primary_input) - .unwrap_or((primary_input, Vec::new())); + let (primary_inner, primary_chain) = self.walk_view_chain(primary_input); let Some(primary_info) = self.infer_layout_cached(graph, primary_inner) else { continue; }; - let primary_layout = - Self::apply_map_layout_chain(primary_info.layout(), &primary_chain); + let Some(primary_layout) = + Self::apply_view_chain(primary_info.layout(), &primary_chain) + else { + continue; + }; if primary_layout != Layout::contiguous(&nary.shape) { continue; } @@ -297,13 +293,10 @@ impl Resolver { let mut extras = Vec::new(); let mut valid_expression = true; for (input_idx, &nary_input) in nary.inputs.iter().enumerate() { - let (base_inner, chain) = self - .walk_map_layout_chain(nary_input) - .unwrap_or((nary_input, Vec::new())); + let (base_inner, chain) = self.walk_view_chain(nary_input); if base_inner == primary_inner { - let alias_layout = - Self::apply_map_layout_chain(primary_info.layout(), &chain); - if alias_layout == Layout::contiguous(&nary.shape) + let alias_layout = Self::apply_view_chain(primary_info.layout(), &chain); + if alias_layout == Some(Layout::contiguous(&nary.shape)) && !nary.expression.uses_custom_indexing_for_input(input_idx) { mapping[input_idx] = 0; @@ -380,8 +373,7 @@ impl Resolver { self.execution_graph.add_edge(idx, node_idx, ()); } } - self.execution_graph[node_idx].variant = - ComputeGraphNodeVariant::QMatMul(new_q.clone()); + self.execution_graph[node_idx].variant = ExecutionVariant::QMatMul(new_q.clone()); self.remove_node_if_dead(input_exec); let mut deps = vec![new_q.input]; deps.extend(deps_extras); @@ -391,7 +383,7 @@ impl Resolver { } // Pre-op: fuse elementwise before matmul inputs - if let ComputeGraphNodeVariant::MatMul(matmul_op) = &node_variant { + if let ExecutionVariant::MatMul(matmul_op) = &node_variant { let mut new_matmul = matmul_op.clone(); let mut changed = false; @@ -425,7 +417,7 @@ impl Resolver { if changed { self.execution_graph[node_idx].variant = - ComputeGraphNodeVariant::MatMul(new_matmul.clone()); + ExecutionVariant::MatMul(new_matmul.clone()); if new_matmul.first != matmul_op.first { let old = self.get_input_node_in_exec_graph(matmul_op.first).unwrap(); @@ -483,7 +475,7 @@ impl Resolver { new_q: Box, ) { let deps = Self::qmatmul_dependencies(&new_q); - self.execution_graph[node_idx].variant = ComputeGraphNodeVariant::QMatMul(new_q); + self.execution_graph[node_idx].variant = ExecutionVariant::QMatMul(new_q); for input in nary_inputs { if deps.contains(input) { @@ -525,7 +517,7 @@ impl Resolver { &mut self, graph: &mut ComputeGraphInner, node_idx: ExecutionNodeIndex, - nary: &NaryOperation, + nary: &ElementwiseOperation, ) -> bool { if nary.output_datatype != crate::DataTypeEnum::F32 { return false; @@ -535,16 +527,14 @@ impl Resolver { // (chain-less) reference is the indexed-input form handled below. let mut base = None; for &input in &nary.inputs { - let (base_inner, chain) = self - .walk_map_layout_chain(input) - .unwrap_or((input, Vec::new())); - if chain.is_empty() { + let (base_inner, chain) = self.walk_view_chain(input); + if chain.is_none() { continue; } let Some(exec) = self.get_input_node_in_exec_graph(base_inner) else { continue; }; - if let ComputeGraphNodeVariant::QMatMul(op) = &self.execution_graph[exec].variant { + if let ExecutionVariant::QMatMul(op) = &self.execution_graph[exec].variant { // A qmatmul that already carries a post epilogue isn't a clean // accumulator-offset base; leave it to the general scan. if op.post_element_wise_expr.is_some() { @@ -610,7 +600,7 @@ impl Resolver { fn try_extract_mapped_qmatmul_post_expr( &mut self, graph: &mut ComputeGraphInner, - nary: &NaryOperation, + nary: &ElementwiseOperation, qmatmul_inner: NodeIndex, qmatmul_out_shape: &[usize], ) -> Option<(NaryExpr, Vec, Vec)> { @@ -654,11 +644,9 @@ impl Resolver { if nary.expression.uses_custom_indexing_for_input(input_idx) { return None; } - let (base_inner, chain) = self - .walk_map_layout_chain(nary_input) - .unwrap_or((nary_input, Vec::new())); + let (base_inner, chain) = self.walk_view_chain(nary_input); if base_inner == qmatmul_inner { - let view = Self::apply_map_layout_chain(&qmatmul_out_layout, &chain); + let view = Self::apply_view_chain(&qmatmul_out_layout, &chain)?; let offset = Self::qmatmul_last_dim_view_offset(&view, &nary.shape, matrix_cols)?; let value_idx = *accumulator_map.entry(offset).or_insert_with(|| { let idx = accumulator_offsets.len(); @@ -725,7 +713,7 @@ impl Resolver { fn try_extract_indexed_qmatmul_post_expr( &mut self, graph: &mut ComputeGraphInner, - nary: &NaryOperation, + nary: &ElementwiseOperation, qmatmul_input_idx: usize, qmatmul_out_shape: &[usize], ) -> Option<(NaryExpr, Vec, Vec)> { diff --git a/fusor-ml/core/src/compute_graph/resolve/fusion_row.rs b/fusor-ml/core/src/compute_graph/resolve/fusion_row.rs new file mode 100644 index 000000000..c54b0c821 --- /dev/null +++ b/fusor-ml/core/src/compute_graph/resolve/fusion_row.rs @@ -0,0 +1,699 @@ +//! Row-program fusion: collapse clusters of same-axis reductions and +//! elementwise expressions into one [`RowProgramOperation`]. +//! +//! Unlike recognition, nothing here matches a named operation. The pass +//! walks upward from an elementwise root, converting everything it can into +//! expressions over a shared slot space: full-shape elementwise producers +//! inline, reductions over a single common axis become per-row scalar +//! phases (their keepdim-broadcast reads turn into scalar references), and +//! anything else becomes an external input. Composed softmax and RMS norm +//! collapse to one kernel this way — and so does any other normalization- +//! shaped cluster the tensor API emits. + +use std::sync::Arc; + +use petgraph::algo::toposort; + +use crate::{ + nary_wise::NaryFunction, + row_program::{RowOutput, RowPhase, RowProgramOperation, RowReduce}, +}; + +use super::cluster_match::{keepdim_broadcast_layout, layout_matches, unary_elementwise}; +use super::*; + +/// Scalar slot base while a cluster is under construction; remapped to +/// `externals.len() + phase` on commit. +const SCALAR_SLOT_BASE: usize = usize::MAX / 2; + +struct ClusterBuilder<'a> { + shape: Box<[usize]>, + axis: Option, + externals: Vec, + /// `(scalar chain base node, phase)` — the base node keys deduplication + /// when the same reduction is read more than once. + phases: Vec<(NodeIndex, RowReduce)>, + /// Inner nodes absorbed into the cluster: they must be consumed only by + /// other members (or the root) and die on commit. + members: Vec, + /// Memoized full-shape expressions for absorbed elementwise nodes. + full_exprs: FxHashMap, + /// The closed absorbable set: nodes outside it read as externals. + allowed: &'a FxHashSet, +} + +struct BuilderSnapshot { + externals: usize, + phases: usize, + members: usize, + full_exprs: FxHashMap, +} + +impl ClusterBuilder<'_> { + fn snapshot(&self) -> BuilderSnapshot { + BuilderSnapshot { + externals: self.externals.len(), + phases: self.phases.len(), + members: self.members.len(), + full_exprs: self.full_exprs.clone(), + } + } + + fn restore(&mut self, snapshot: BuilderSnapshot) { + self.externals.truncate(snapshot.externals); + self.phases.truncate(snapshot.phases); + self.members.truncate(snapshot.members); + self.full_exprs = snapshot.full_exprs; + } + + fn external_slot(&mut self, inner: NodeIndex) -> usize { + if let Some(slot) = self.externals.iter().position(|&node| node == inner) { + slot + } else { + self.externals.push(inner); + self.externals.len() - 1 + } + } +} + +/// Rewrite the slot references of one absorbed node's expression into the +/// cluster slot space: external slots keep their index lists, inlined +/// expressions require plain elementwise reads. +enum SlotRewrite { + External(usize), + Inline(NaryExpr), +} + +fn rewrite_slots(expr: &NaryExpr, rewrites: &[SlotRewrite]) -> Option { + match expr { + NaryExpr::Op { children, function } => Some(NaryExpr::Op { + children: children + .iter() + .map(|child| rewrite_slots(child, rewrites)) + .collect::>>()?, + function: function.clone(), + }), + NaryExpr::IndexedInput { input_idx, indices } => { + let indices = indices + .iter() + .map(|index| rewrite_slots(index, rewrites)) + .collect::>>()?; + match &rewrites[*input_idx] { + SlotRewrite::External(slot) => Some(NaryExpr::IndexedInput { + input_idx: *slot, + indices, + }), + SlotRewrite::Inline(replacement) => { + NaryExpr::is_elementwise_indices(&indices).then(|| replacement.clone()) + } + } + } + NaryExpr::DimIndex(dim) => Some(NaryExpr::DimIndex(*dim)), + NaryExpr::Scalar(value) => Some(NaryExpr::Scalar(*value)), + } +} + +/// Remap construction-time scalar slots to their final indices. +fn finalize_slots(expr: &NaryExpr, external_count: usize) -> NaryExpr { + match expr { + NaryExpr::Op { children, function } => NaryExpr::Op { + children: children + .iter() + .map(|child| finalize_slots(child, external_count)) + .collect(), + function: function.clone(), + }, + NaryExpr::IndexedInput { input_idx, indices } => { + let input_idx = if *input_idx >= SCALAR_SLOT_BASE { + external_count + (*input_idx - SCALAR_SLOT_BASE) + } else { + *input_idx + }; + NaryExpr::IndexedInput { + input_idx, + indices: indices + .iter() + .map(|index| finalize_slots(index, external_count)) + .collect(), + } + } + NaryExpr::DimIndex(dim) => NaryExpr::DimIndex(*dim), + NaryExpr::Scalar(value) => NaryExpr::Scalar(*value), + } +} + +impl Resolver { + /// Fuse row-program clusters. Candidate roots are discovered from the + /// reductions outward — a reduce's scalar flows through its keepdim + /// broadcast into some elementwise consumer, and the root is the last + /// single-consumer elementwise below it — so graphs full of elementwise + /// chains with no reductions (most of a decode graph) cost nothing. + pub(super) fn fuse_row_programs(&mut self, graph: &mut ComputeGraphInner) { + let reduces: Vec = self + .execution_graph + .node_indices() + .filter(|&node| { + matches!( + self.execution_graph[node].variant, + ExecutionVariant::Reduce(_) + ) + }) + .collect(); + if reduces.is_empty() { + return; + } + + let mut roots = Vec::new(); + let mut rootless = 0usize; + let reduce_count = reduces.len(); + for reduce in reduces { + if let Some(root) = self.row_cluster_root(reduce) { + if !roots.contains(&root) { + roots.push(root); + } + } else { + rootless += 1; + } + } + if std::env::var_os("FUSOR_TRACE_ROW_FUSION").is_some() { + eprintln!( + "row_fusion: {reduce_count} reduces, {} roots, {rootless} rootless", + roots.len() + ); + } + // Outermost roots first so each committed cluster is maximal: a + // root's cluster can contain another candidate root (softmax's two + // reductions converge on one div). + let Ok(order) = toposort(&self.execution_graph, None) else { + return; + }; + let rank: FxHashMap = order + .into_iter() + .enumerate() + .map(|(rank, node)| (node, rank)) + .collect(); + roots.sort_by_key(|root| std::cmp::Reverse(rank.get(root).copied().unwrap_or(0))); + for root in roots { + if !self.execution_graph.contains_node(root) { + continue; + } + self.try_fuse_row_program(graph, root); + } + } + + /// Walk downstream from a reduce to the cluster's natural root: through + /// any reduced-shape unary chain and keepdim views to the first + /// full-shape elementwise consumer, then down while that node feeds + /// exactly one full-shape elementwise. + fn row_cluster_root(&self, reduce: ExecutionNodeIndex) -> Option { + let ExecutionVariant::Reduce(op) = &self.execution_graph[reduce].variant else { + return None; + }; + let full_shape = op.shape.clone(); + + // Phase 1: descend through scalar-side nodes (views and elementwise + // ops over fewer elements than the full shape) to a full-shape + // elementwise consumer. + let mut node = reduce; + let mut current = loop { + let mut consumers = self + .execution_graph + .neighbors_directed(node, petgraph::Direction::Outgoing); + let first = consumers.next()?; + if consumers.next().is_some() { + return None; + } + match &self.execution_graph[first].variant { + ExecutionVariant::View(_) => node = first, + ExecutionVariant::Elementwise(nary) if nary.shape == full_shape => break first, + ExecutionVariant::Elementwise(_) => node = first, + _ => return None, + } + }; + + // Phase 2: descend while the node feeds exactly one full-shape + // elementwise (other consumers may be reductions inside the cluster). + loop { + let mut next = None; + for consumer in self + .execution_graph + .neighbors_directed(current, petgraph::Direction::Outgoing) + { + if let ExecutionVariant::Elementwise(nary) = &self.execution_graph[consumer].variant + && nary.shape == full_shape + { + if next.is_some() { + return Some(current); + } + next = Some(consumer); + } + } + match next { + Some(consumer) => current = consumer, + None => return Some(current), + } + } + } + + fn try_fuse_row_program( + &mut self, + graph: &mut ComputeGraphInner, + root_idx: ExecutionNodeIndex, + ) -> bool { + let ExecutionVariant::Elementwise(root) = self.execution_graph[root_idx].variant.clone() + else { + return false; + }; + + // Greedy absorption could claim a node another branch still needs + // (a residual stream read by the next layer, say). Compute the + // *closed* absorbable set first: collect everything reachable + // through eligible nodes, then iteratively drop anything consumed + // outside the set — what remains can fuse without anything + // escaping. Dropped nodes read as external inputs and materialize + // once, exactly as they must. + let mut allowed = self.collect_row_cluster(graph, root_idx, &root, None); + loop { + let allowed_execs: FxHashSet = allowed + .iter() + .filter_map(|&inner| self.get_input_node_in_exec_graph(inner)) + .collect(); + let mut violators = Vec::new(); + for &member in &allowed { + let Some(exec) = self.get_input_node_in_exec_graph(member) else { + violators.push(member); + continue; + }; + if self + .execution_graph + .neighbors_directed(exec, petgraph::Direction::Outgoing) + .any(|consumer| consumer != root_idx && !allowed_execs.contains(&consumer)) + { + violators.push(member); + } + } + if violators.is_empty() { + break; + } + for violator in violators { + allowed.remove(&violator); + } + // Re-walk: regions only reachable through a dropped node fall + // out of the set too. + allowed = self.collect_row_cluster(graph, root_idx, &root, Some(&allowed)); + } + if allowed.is_empty() { + return false; + } + self.build_row_cluster(graph, root_idx, &root, &allowed) + } + + /// Collect the nodes the absorption walk could claim: full-shape + /// elementwise producers and keepdim-broadcast scalar reads of same-axis + /// reductions, recursively. `within` restricts the walk to an existing + /// candidate set. + fn collect_row_cluster( + &self, + graph: &ComputeGraphInner, + _root_idx: ExecutionNodeIndex, + root: &ElementwiseOperation, + within: Option<&FxHashSet>, + ) -> FxHashSet { + let mut out = FxHashSet::default(); + let mut axis = None; + for &input in &root.inputs { + self.collect_operand(graph, &root.shape, &mut axis, input, within, &mut out); + } + out + } + + fn collect_operand( + &self, + graph: &ComputeGraphInner, + shape: &[usize], + axis: &mut Option, + inner: NodeIndex, + within: Option<&FxHashSet>, + out: &mut FxHashSet, + ) { + let eligible = |node: NodeIndex| { + within.is_none_or(|allowed| allowed.contains(&node)) && !out.contains(&node) + }; + if out.contains(&inner) { + return; + } + if !eligible(inner) && !out.contains(&inner) && within.is_some() { + return; + } + + // Scalar path: views → unary chain → same-axis reduce. + 'scalar: { + let mut nodes = Vec::new(); + let mut node = inner; + let mut layout: Option = None; + loop { + let Some(exec) = self.get_input_node_in_exec_graph(node) else { + break 'scalar; + }; + let ExecutionVariant::View(view) = &self.execution_graph[exec].variant else { + break; + }; + if !view.is_fully_defined() || !within.is_none_or(|allowed| allowed.contains(&node)) + { + break 'scalar; + } + layout = Some(match &layout { + None => view.layout.clone(), + Some(outer) => match crate::view::compose_layouts(outer, &view.layout) { + Some(layout) => layout, + None => break 'scalar, + }, + }); + nodes.push(node); + node = view.input; + } + let Some(layout) = layout else { + break 'scalar; + }; + let reduce = loop { + if self.check_cached(graph, node) + || !within.is_none_or(|allowed| allowed.contains(&node)) + { + break 'scalar; + } + let Some(exec) = self.get_input_node_in_exec_graph(node) else { + break 'scalar; + }; + match &self.execution_graph[exec].variant { + ExecutionVariant::Reduce(reduce) => break reduce.clone(), + ExecutionVariant::Elementwise(nary) => { + if unary_elementwise(nary).is_none() { + break 'scalar; + } + nodes.push(node); + let (_, input) = unary_elementwise(nary).unwrap(); + node = input; + } + _ => break 'scalar, + } + }; + let Some(value) = reduce.plain_input() else { + break 'scalar; + }; + if reduce.shape.as_ref() != shape + || axis.is_some_and(|existing| existing != reduce.axis) + || !layout_matches(Some(&layout), &keepdim_broadcast_layout(shape, reduce.axis)) + { + break 'scalar; + } + *axis = Some(reduce.axis); + out.extend(nodes); + out.insert(node); + self.collect_operand(graph, shape, axis, value, within, out); + return; + } + + // Full-shape elementwise producer. + if self.check_cached(graph, inner) { + return; + } + let Some(exec) = self.get_input_node_in_exec_graph(inner) else { + return; + }; + let ExecutionVariant::Elementwise(nary) = &self.execution_graph[exec].variant else { + return; + }; + if nary.shape.as_ref() != shape { + return; + } + let inputs = nary.inputs.clone(); + out.insert(inner); + for input in inputs { + self.collect_operand(graph, shape, axis, input, within, out); + } + } + + fn build_row_cluster( + &mut self, + graph: &mut ComputeGraphInner, + root_idx: ExecutionNodeIndex, + root: &ElementwiseOperation, + allowed: &FxHashSet, + ) -> bool { + let mut builder = ClusterBuilder { + shape: root.shape.clone(), + axis: None, + externals: Vec::new(), + phases: Vec::new(), + members: Vec::new(), + full_exprs: FxHashMap::default(), + allowed, + }; + let mut rewrites = Vec::with_capacity(root.inputs.len()); + for &input in &root.inputs { + let Some(rewrite) = self.absorb_operand(graph, &mut builder, input) else { + return false; + }; + rewrites.push(rewrite); + } + let Some(output_expr) = rewrite_slots(&root.expression, &rewrites) else { + return false; + }; + + let trace = std::env::var_os("FUSOR_TRACE_ROW_FUSION").is_some(); + // Fusing pays for itself only when at least one reduction folds in. + if builder.phases.is_empty() { + return false; + } + let axis = builder.axis.expect("phases imply a chosen axis"); + + // The fixpoint guarantees closure; verify before rewriting anyway. + let member_execs: FxHashSet = builder + .members + .iter() + .filter_map(|&inner| self.get_input_node_in_exec_graph(inner)) + .collect(); + if member_execs.len() != builder.members.len() { + return false; + } + for &exec in &member_execs { + for consumer in self + .execution_graph + .neighbors_directed(exec, petgraph::Direction::Outgoing) + { + if consumer != root_idx && !member_execs.contains(&consumer) { + return false; + } + } + } + + let external_count = builder.externals.len(); + let phases: Vec = builder + .phases + .into_iter() + .map(|(_, mut phase)| { + phase.expression = finalize_slots(&phase.expression, external_count); + RowPhase::Reduce(phase) + }) + .collect(); + let operation = RowProgramOperation { + inputs: builder.externals.clone(), + shape: builder.shape, + axis, + phases, + output: RowOutput::Map(finalize_slots(&output_expr, external_count)), + output_datatype: root.output_datatype, + dynamic_axis: None, + }; + let externals = builder.externals; + if trace { + eprintln!( + "row_fusion: committed root {:?} shape {:?} phases {} externals {}", + self.execution_graph[root_idx].inner_idx, + operation.shape, + operation.phases.len(), + externals.len() + ); + } + self.commit_recognized( + graph, + root_idx, + &externals, + ExecutionVariant::GraphOp(Arc::new(operation)), + ); + true + } + + /// Convert one operand read into the cluster slot space: a scalar + /// reference when it is a keepdim-broadcast of a same-axis reduction, an + /// inlined expression when it is an absorbable full-shape elementwise + /// node, and an external slot otherwise. + fn absorb_operand( + &self, + graph: &ComputeGraphInner, + builder: &mut ClusterBuilder<'_>, + inner: NodeIndex, + ) -> Option { + if !builder.allowed.contains(&inner) { + return Some(SlotRewrite::External(builder.external_slot(inner))); + } + let snapshot = builder.snapshot(); + if let Some(expr) = self.try_absorb_scalar(graph, builder, inner) { + return Some(SlotRewrite::Inline(expr)); + } + builder.restore(snapshot); + + let snapshot = builder.snapshot(); + if let Some(expr) = self.try_absorb_full(graph, builder, inner) { + return Some(SlotRewrite::Inline(expr)); + } + builder.restore(snapshot); + + Some(SlotRewrite::External(builder.external_slot(inner))) + } + + /// Inline a full-shape elementwise producer consumed by the cluster. + fn try_absorb_full( + &self, + graph: &ComputeGraphInner, + builder: &mut ClusterBuilder<'_>, + inner: NodeIndex, + ) -> Option { + if let Some(expr) = builder.full_exprs.get(&inner) { + return Some(expr.clone()); + } + // A cached intermediate is a hard boundary; a user-held reference + // alone is not — the inner-graph node outlives this resolve, so a + // later resolve of that handle simply recomputes it (the same rule + // n-ary fusion applies). + if self.check_cached(graph, inner) { + return None; + } + let exec = self.get_input_node_in_exec_graph(inner)?; + let ExecutionVariant::Elementwise(nary) = &self.execution_graph[exec].variant else { + return None; + }; + if nary.shape != builder.shape { + return None; + } + let nary = nary.clone(); + + let mut rewrites = Vec::with_capacity(nary.inputs.len()); + for &input in &nary.inputs { + rewrites.push(self.absorb_operand(graph, builder, input)?); + } + let expr = rewrite_slots(&nary.expression, &rewrites)?; + builder.members.push(inner); + builder.full_exprs.insert(inner, expr.clone()); + Some(expr) + } + + /// Absorb a keepdim-broadcast read of a same-axis reduction as a per-row + /// scalar phase: views compose down to the broadcast layout, an optional + /// unary chain below them becomes the phase's post chain, and the + /// reduction's producer is absorbed like any other full-shape value. + fn try_absorb_scalar( + &self, + graph: &ComputeGraphInner, + builder: &mut ClusterBuilder<'_>, + inner: NodeIndex, + ) -> Option { + // Walk the view chain, collecting the member nodes. + let mut views = Vec::new(); + let mut layout: Option = None; + let mut node = inner; + loop { + let exec = self.get_input_node_in_exec_graph(node)?; + let ExecutionVariant::View(view) = &self.execution_graph[exec].variant else { + break; + }; + if !view.is_fully_defined() { + return None; + } + layout = Some(match &layout { + None => view.layout.clone(), + Some(outer) => crate::view::compose_layouts(outer, &view.layout)?, + }); + views.push(node); + node = view.input; + } + let layout = layout?; + + // An optional unary chain on the reduced value (mean scaling, eps, + // rsqrt...) folds into the phase's post chain, innermost first. + let mut chain_nodes = Vec::new(); + let mut chain: Vec = Vec::new(); + let reduce = loop { + if self.check_cached(graph, node) { + return None; + } + let exec = self.get_input_node_in_exec_graph(node)?; + match &self.execution_graph[exec].variant { + ExecutionVariant::Reduce(reduce) => break reduce.clone(), + ExecutionVariant::Elementwise(nary) => { + let (function, input) = unary_elementwise(nary)?; + chain.push(function); + chain_nodes.push(node); + node = input; + } + _ => return None, + } + }; + chain.reverse(); + + let value = reduce.plain_input()?; + let axis = reduce.axis; + if reduce.shape != builder.shape { + return None; + } + if let Some(existing) = builder.axis + && existing != axis + { + return None; + } + if !layout_matches( + Some(&layout), + &keepdim_broadcast_layout(&builder.shape, axis), + ) { + return None; + } + + // The same reduction read again reuses its phase. + let scalar_ref = + |phase: usize, rank: usize| NaryExpr::input(SCALAR_SLOT_BASE + phase, rank); + if let Some(phase) = builder + .phases + .iter() + .position(|(base, _)| *base == chain_nodes.first().copied().unwrap_or(node)) + { + builder.members.extend(views); + return Some(scalar_ref(phase, builder.shape.len())); + } + + builder.axis = Some(axis); + let rewrite = self.absorb_operand(graph, builder, value)?; + let expression = rewrite_slots(&NaryExpr::input(0, builder.shape.len()), &[rewrite])?; + + let mut post = reduce.post_element_wise.functions.clone(); + post.extend(chain); + let post_chain = crate::nary_wise::UnaryFunctionChain::new( + post, + reduce.post_element_wise.input_datatype(), + ); + + let phase_key = chain_nodes.first().copied().unwrap_or(node); + builder.members.extend(views); + builder.members.extend(chain_nodes); + builder.members.push(node); + let phase_index = builder.phases.len(); + builder.phases.push(( + phase_key, + RowReduce { + expression, + function: reduce.function.clone(), + post_chain, + }, + )); + Some(scalar_ref(phase_index, builder.shape.len())) + } +} diff --git a/fusor-ml/core/src/compute_graph/resolve/mod.rs b/fusor-ml/core/src/compute_graph/resolve/mod.rs index b225b52e4..56c057c27 100644 --- a/fusor-ml/core/src/compute_graph/resolve/mod.rs +++ b/fusor-ml/core/src/compute_graph/resolve/mod.rs @@ -7,7 +7,7 @@ use crate::{ compute_graph::layout_pass::LayoutPass, mir::{inputs::MirValue, kernel_backend::PreparedDirectDispatch, operation::Operation}, nary_wise::{ - ExtractedUnaryChain, NaryExpr, NaryOp, NaryOperation, NaryScalar, UnaryFunctionChain, + ElementwiseOperation, ExtractedUnaryChain, NaryExpr, NaryOp, NaryScalar, UnaryFunctionChain, }, quantized::matmul::{ElementwiseEpilogue, QMatMulOperation}, tensor::TensorData, @@ -16,12 +16,24 @@ use petgraph::algo::toposort; use petgraph::stable_graph::StableGraph; use rustc_hash::{FxHashMap, FxHashSet}; -use super::{ComputeGraphInner, ComputeGraphNode, ComputeGraphNodeVariant, NodeIndex}; +use super::{ + ComputeGraphInner, ComputeGraphNode, ComputeGraphNodeVariant, GraphOperation, NodeIndex, +}; +use crate::{ + MatMulOperation, ReduceOperation, dequantize::DequantizeOperation, + quantized::embedding::QEmbeddingOperation, slice_assign::SliceAssignOperation, + view::ViewOperation, +}; +mod cluster_match; mod execution; +mod fold_views; mod fusion_basic; mod fusion_matmul; +mod fusion_row; mod plan_cache; +mod recognize; +mod recognize_attention; mod run; pub(crate) use plan_cache::structural_kernel_key; @@ -136,8 +148,6 @@ struct OptimizeProfile { fuse_reduce: Duration, fuse_matmul_count: usize, fuse_matmul: Duration, - fuse_rmsnorm_count: usize, - fuse_rmsnorm: Duration, } impl OptimizeProfile { @@ -146,8 +156,7 @@ impl OptimizeProfile { "resolve_optimize_profile iterations={} changed={} \ fuse_naries_count={} fuse_naries={:?} \ fuse_reduce_count={} fuse_reduce={:?} \ -fuse_matmul_count={} fuse_matmul={:?} \ -fuse_rmsnorm_count={} fuse_rmsnorm={:?}", +fuse_matmul_count={} fuse_matmul={:?}", self.iterations, self.changed, self.fuse_naries_count, @@ -156,8 +165,6 @@ fuse_rmsnorm_count={} fuse_rmsnorm={:?}", self.fuse_reduce, self.fuse_matmul_count, self.fuse_matmul, - self.fuse_rmsnorm_count, - self.fuse_rmsnorm, ); } } @@ -220,34 +227,69 @@ fn print_host_category_profile(profile: FxHashMap<&'static str, ResolveHostCateg tracing::info!("resolve_host_category_profile {profile:?}"); } -fn node_category(variant: &ComputeGraphNodeVariant) -> &'static str { +fn node_category_inner(variant: &ComputeGraphNodeVariant) -> &'static str { match variant { - ComputeGraphNodeVariant::Nary(_) => "nary", - ComputeGraphNodeVariant::SliceAssign(_) => "slice_assign", - ComputeGraphNodeVariant::Resize(_) => "resize", - ComputeGraphNodeVariant::MapLayout(_) => "map_layout", - ComputeGraphNodeVariant::Dequantize(_) => "dequantize", - ComputeGraphNodeVariant::QEmbedding(_) => "q_embedding", - ComputeGraphNodeVariant::MatMul(_) => "matmul", - ComputeGraphNodeVariant::QMatMul(_) => "q_matmul", ComputeGraphNodeVariant::Tensor(_) => "tensor", + ComputeGraphNodeVariant::QMatrix(_) => "q_matrix", + ComputeGraphNodeVariant::Elementwise(_) => "elementwise", ComputeGraphNodeVariant::Reduce(_) => "reduce", - ComputeGraphNodeVariant::FlashAttention(_) => "flash_attention", - ComputeGraphNodeVariant::GraphOp(op) => op.category(), + ComputeGraphNodeVariant::View(_) => "view", + ComputeGraphNodeVariant::Assign(_) => "assign", } } -fn as_rms_norm(variant: &ComputeGraphNodeVariant) -> Option<&crate::RmsNormOperation> { - let ComputeGraphNodeVariant::GraphOp(op) = variant else { - return None; - }; - op.as_any().downcast_ref::() +#[allow(dead_code, reason = "execution-side category labeling for profiling")] +fn node_category(variant: &ExecutionVariant) -> &'static str { + match variant { + ExecutionVariant::Elementwise(_) => "nary", + ExecutionVariant::Assign(_) => "slice_assign", + ExecutionVariant::View(_) => "view", + ExecutionVariant::QMatrix(_) => "dequantize", + ExecutionVariant::QEmbedding(_) => "q_embedding", + ExecutionVariant::MatMul(_) => "matmul", + ExecutionVariant::QMatMul(_) => "q_matmul", + ExecutionVariant::Tensor(_) => "tensor", + ExecutionVariant::Reduce(_) => "reduce", + ExecutionVariant::GraphOp(op) => op.category(), + } +} + +/// What an execution-graph node lowers to. The graph vocabulary (the first +/// six variants) enters verbatim; the region variants exist only here — +/// recognition rebuilds them from composed clusters, and fusion enriches +/// them with epilogues. +#[derive(Debug, Clone)] +pub(crate) enum ExecutionVariant { + Tensor(crate::tensor::TensorData), + QMatrix(DequantizeOperation), + Elementwise(ElementwiseOperation), + Reduce(ReduceOperation), + View(ViewOperation), + Assign(SliceAssignOperation), + // Recognized regions. + MatMul(MatMulOperation), + QMatMul(Box), + QEmbedding(QEmbeddingOperation), + GraphOp(Arc), +} + +impl From for ExecutionVariant { + fn from(variant: ComputeGraphNodeVariant) -> Self { + match variant { + ComputeGraphNodeVariant::Tensor(op) => Self::Tensor(op), + ComputeGraphNodeVariant::QMatrix(op) => Self::QMatrix(op), + ComputeGraphNodeVariant::Elementwise(op) => Self::Elementwise(op), + ComputeGraphNodeVariant::Reduce(op) => Self::Reduce(op), + ComputeGraphNodeVariant::View(op) => Self::View(op), + ComputeGraphNodeVariant::Assign(op) => Self::Assign(op), + } + } } #[derive(Debug, Clone)] struct ExecutionNode { inner_idx: NodeIndex, - variant: ComputeGraphNodeVariant, + variant: ExecutionVariant, } type ExecutionGraph = StableGraph; @@ -338,7 +380,11 @@ fn print_gpu_kernel_profile( pub(crate) struct Resolver { execution_graph: ExecutionGraph, node_mapping: FxHashMap, - layout_cache: FxHashMap>, + // Persistent memoized layout inference: the inner graph's node variants + // are immutable during optimization (rewrites only touch the execution + // graph and dependency edges), so layouts computed once stay valid for + // the whole resolve. + layout_pass: LayoutPass, targets: Vec, resolved_set: FxHashSet, } @@ -366,7 +412,7 @@ impl Resolver { targets, execution_graph: Default::default(), node_mapping: Default::default(), - layout_cache: Default::default(), + layout_pass: Default::default(), resolved_set, } } diff --git a/fusor-ml/core/src/compute_graph/resolve/recognize.rs b/fusor-ml/core/src/compute_graph/resolve/recognize.rs new file mode 100644 index 000000000..dfa84f560 --- /dev/null +++ b/fusor-ml/core/src/compute_graph/resolve/recognize.rs @@ -0,0 +1,412 @@ +//! Pattern recognition over the composed 3-op graph. +//! +//! The tensor API expresses contractions as `Elementwise(Mul) + Reduce(Sum)` +//! over a shared index space (see `Tensor::mat_mul` / `Tensor::q_mat_mul`). +//! This pass runs before view folding and n-ary fusion, while the composed +//! cluster is still in the exact canonical form the API emitted, and rebuilds +//! the specialized operation so the existing kernel paths (and their epilogue +//! fusion) take over. Anything it does not recognize lowers through the +//! generic elementwise + reduce kernels — slower, but correct. + +use crate::{ + MatMulOperation, ReduceOperation, dequantize::DequantizeOperation, + quantized::matmul::QMatMulOperation, reduce::ReduceOp, +}; + +use super::*; + +/// The canonical contraction cluster: `Reduce(Sum, axis = rank-1)` over a +/// two-factor multiply where each factor is a bare indexed load with pure +/// `DimIndex` coordinates. +pub(crate) struct Contraction { + /// The two factors: (inner node, output dims indexed). + factors: [(NodeIndex, Vec); 2], + shape: Box<[usize]>, + datatype: DataTypeEnum, +} + +fn pure_dim_indices(indices: &[NaryExpr]) -> Option> { + indices + .iter() + .map(|index| match index { + NaryExpr::DimIndex(dim) => Some(*dim), + _ => None, + }) + .collect() +} + +/// Match the canonical multiply + sum-reduce pair. +pub(crate) fn match_contraction( + reduce: &ReduceOperation, + nary: &ElementwiseOperation, +) -> Option { + if reduce.function.op != ReduceOp::Sum + || reduce.plain_input().is_none() + || !reduce.post_element_wise.functions.is_empty() + || reduce.axis + 1 != nary.shape.len() + { + return None; + } + + let NaryExpr::Op { children, function } = &nary.expression else { + return None; + }; + if function.op != crate::nary_wise::NaryOp::Mul || children.len() != 2 { + return None; + } + let factor = |expr: &NaryExpr| -> Option<(NodeIndex, Vec)> { + let NaryExpr::IndexedInput { input_idx, indices } = expr else { + return None; + }; + Some((nary.inputs[*input_idx], pure_dim_indices(indices)?)) + }; + + Some(Contraction { + factors: [factor(&children[0])?, factor(&children[1])?], + shape: nary.shape.clone(), + datatype: nary.output_datatype, + }) +} + +impl Contraction { + /// `input [.., K] × QMatrix [N, K] → [.., N]`: index space `[.., N, K]`, + /// activation indices `[0..r-2, k]` (skipping the `N` dim), matrix + /// indices `[n, k]`. Returns the rebuilt operation and its activation + /// input node. + pub(crate) fn to_q_mat_mul( + &self, + dequantize_for: impl Fn(NodeIndex) -> Option, + ) -> Option<(QMatMulOperation, NodeIndex)> { + let rank = self.shape.len(); + if rank < 2 { + return None; + } + let (k_dim, n_dim) = (rank - 1, rank - 2); + let expected_activation: Vec = (0..rank - 2).chain(std::iter::once(k_dim)).collect(); + + for (activation, matrix) in [ + (&self.factors[0], &self.factors[1]), + (&self.factors[1], &self.factors[0]), + ] { + let Some(matrix_op) = dequantize_for(matrix.0) else { + continue; + }; + if matrix.1 != [n_dim, k_dim] + || activation.1 != expected_activation + || matrix_op.datatype != self.datatype + || matrix_op.matrix.shape() != [self.shape[n_dim], self.shape[k_dim]] + { + continue; + } + + let in_shape: Vec = self + .shape + .iter() + .enumerate() + .filter_map(|(dim, &size)| (dim != n_dim).then_some(size)) + .collect(); + let operation = QMatMulOperation::new( + self.datatype, + &in_shape, + activation.0, + matrix_op.matrix.clone(), + ); + return Some((operation, activation.0)); + } + None + } + + /// `a [batch.., M, K] × b [batch.., K, N] → [batch.., M, N]`: index space + /// `[batch.., M, N, K]`, `a` indices `[batch.., m, k]`, `b` indices + /// `[batch.., k, n]`. Returns the rebuilt operation and its two inputs. + pub(crate) fn to_mat_mul( + &self, + device: &crate::Device, + ) -> Option<(MatMulOperation, [NodeIndex; 2])> { + let rank = self.shape.len(); + if rank < 3 { + return None; + } + let batch = rank - 3; + let (m_dim, n_dim, k_dim) = (batch, batch + 1, batch + 2); + + let expected_a: Vec = (0..batch).chain([m_dim, k_dim]).collect(); + let expected_b: Vec = (0..batch).chain([k_dim, n_dim]).collect(); + let (a, b) = if self.factors[0].1 == expected_a && self.factors[1].1 == expected_b { + (&self.factors[0], &self.factors[1]) + } else if self.factors[1].1 == expected_a && self.factors[0].1 == expected_b { + (&self.factors[1], &self.factors[0]) + } else { + return None; + }; + + let batch_shape = &self.shape[..batch]; + let first_shape: Vec = batch_shape + .iter() + .chain([&self.shape[m_dim], &self.shape[k_dim]]) + .copied() + .collect(); + let second_shape: Vec = batch_shape + .iter() + .chain([&self.shape[k_dim], &self.shape[n_dim]]) + .copied() + .collect(); + + let operation = MatMulOperation::new( + self.datatype, + a.0, + b.0, + &first_shape, + &second_shape, + None, + device, + ); + Some((operation, [a.0, b.0])) + } +} + +impl ComputeGraphInner { + pub(crate) fn dequantize_variant(&self, key: NodeIndex) -> Option { + match &self.nodes.nodes.node_weight(key)?.variant { + ComputeGraphNodeVariant::QMatrix(op) => Some(op.clone()), + _ => None, + } + } + + /// Recognize a composed contraction rooted at `key` directly on the inner + /// graph (the single-target fast path, no resolver involved). The multiply + /// must feed only this reduce. + pub(crate) fn match_direct_qmatmul(&self, key: NodeIndex) -> Option { + let ComputeGraphNodeVariant::Reduce(reduce) = &self.nodes.nodes.node_weight(key)?.variant + else { + return None; + }; + let value = reduce.plain_input()?; + if self.get_cached_result(value).is_some() || self.has_live_reference(value) { + return None; + } + let ComputeGraphNodeVariant::Elementwise(nary) = + &self.nodes.nodes.node_weight(value)?.variant + else { + return None; + }; + if self + .nodes + .nodes + .neighbors_directed(value, petgraph::Direction::Outgoing) + .count() + != 1 + { + return None; + } + let contraction = match_contraction(reduce, nary)?; + let (operation, _) = contraction.to_q_mat_mul(|key| self.dequantize_variant(key))?; + Some(operation) + } +} + +impl Resolver { + /// Recognize composed contraction clusters in the execution graph and + /// rebuild them as `MatMul` / `QMatMul` nodes. Runs as a linear sweep + /// before any other rewrite, while the clusters are still in the exact + /// form the API emitted. + pub(super) fn recognize_contractions(&mut self, graph: &mut ComputeGraphInner) { + let reduces: Vec = self + .execution_graph + .node_indices() + .filter(|&node| { + matches!( + self.execution_graph[node].variant, + ExecutionVariant::Reduce(_) + ) + }) + .collect(); + for node in reduces { + self.try_recognize_contraction(graph, node); + } + } + + fn try_recognize_contraction( + &mut self, + graph: &mut ComputeGraphInner, + node_idx: ExecutionNodeIndex, + ) -> bool { + let ExecutionVariant::Reduce(reduce) = &self.execution_graph[node_idx].variant else { + return false; + }; + let Some(value) = reduce.plain_input() else { + return false; + }; + if self.check_cached(graph, value) || graph.has_live_reference(value) { + return false; + } + let Some(nary_exec) = self.get_input_node_in_exec_graph(value) else { + return false; + }; + let ExecutionVariant::Elementwise(nary) = &self.execution_graph[nary_exec].variant else { + return false; + }; + // The multiply must exist solely for this reduce: consumed elsewhere + // it has to materialize anyway. + if self + .execution_graph + .neighbors_directed(nary_exec, petgraph::Direction::Outgoing) + .count() + != 1 + { + return false; + } + let Some(contraction) = match_contraction(reduce, nary) else { + return false; + }; + + let _ = nary_exec; + if let Some((operation, activation)) = + contraction.to_q_mat_mul(|key| graph.dequantize_variant(key)) + { + self.commit_recognized( + graph, + node_idx, + &[activation], + ExecutionVariant::QMatMul(Box::new(operation)), + ); + return true; + } + if let Some((operation, inputs)) = contraction.to_mat_mul(&graph.device()) { + self.commit_recognized( + graph, + node_idx, + &inputs, + ExecutionVariant::MatMul(operation), + ); + return true; + } + false + } + + /// Sweep elementwise nodes for quantized embedding gathers. + pub(super) fn recognize_embeddings(&mut self, graph: &mut ComputeGraphInner) { + let candidates: Vec = self + .execution_graph + .node_indices() + .filter(|&node| { + matches!( + self.execution_graph[node].variant, + ExecutionVariant::Elementwise(_) + ) + }) + .collect(); + for node in candidates { + if !self.execution_graph.contains_node(node) { + continue; + } + self.try_recognize_q_embedding(graph, node); + } + } + + /// Quantized row gather: `Elementwise([table, idx], table[idx[i], j])` + /// over a `[count, hidden]` space (see `QMatrix::index_select_rows_to`). + /// Rebuilds the block-amortized embedding kernel; dense-storage tables + /// stay on the generic elementwise path, which reads them directly. + pub(super) fn try_recognize_q_embedding( + &mut self, + graph: &mut ComputeGraphInner, + node_idx: ExecutionNodeIndex, + ) -> bool { + let ExecutionVariant::Elementwise(nary) = &self.execution_graph[node_idx].variant else { + return false; + }; + if nary.inputs.len() != 2 || nary.shape.len() != 2 { + return false; + } + // Peel the optional cast from the load type to the requested type. + let gather = match &nary.expression { + NaryExpr::Op { children, function } + if function.op == crate::nary_wise::NaryOp::Cast && children.len() == 1 => + { + &children[0] + } + expr => expr, + }; + let NaryExpr::IndexedInput { + input_idx: 0, + indices, + } = gather + else { + return false; + }; + let [row, NaryExpr::DimIndex(1)] = indices.as_slice() else { + return false; + }; + let NaryExpr::IndexedInput { + input_idx: 1, + indices: row_indices, + } = row + else { + return false; + }; + if row_indices.as_slice() != [NaryExpr::DimIndex(0)] { + return false; + } + + let Some(dequantize) = graph.dequantize_variant(nary.inputs[0]) else { + return false; + }; + if crate::quantized::dequantize::quant_format(&dequantize.matrix).is_none() { + return false; + } + if dequantize.matrix.shape().len() != 2 || dequantize.matrix.shape()[1] != nary.shape[1] { + return false; + } + + let indexes = nary.inputs[1]; + let operation = crate::quantized::embedding::QEmbeddingOperation::new( + indexes, + nary.shape[0], + dequantize.matrix.clone(), + nary.output_datatype, + ); + self.commit_recognized( + graph, + node_idx, + &[indexes], + ExecutionVariant::QEmbedding(operation), + ); + true + } + + /// Replace a recognized cluster's root with the rebuilt operation: drop + /// every edge from the cluster's intermediates, wire the operation's + /// dependencies directly, and let the now-unconsumed intermediates fall + /// out of the execution graph. + pub(super) fn commit_recognized( + &mut self, + graph: &mut ComputeGraphInner, + node_idx: ExecutionNodeIndex, + dependencies: &[NodeIndex], + variant: ExecutionVariant, + ) { + self.execution_graph[node_idx].variant = variant; + + let previous: Vec = self + .execution_graph + .neighbors_directed(node_idx, petgraph::Direction::Incoming) + .collect(); + for &prev in &previous { + if let Some(edge) = self.execution_graph.find_edge(prev, node_idx) { + self.execution_graph.remove_edge(edge); + } + } + for &dependency in dependencies { + if let Some(exec) = self.get_input_node_in_exec_graph(dependency) + && self.execution_graph.find_edge(exec, node_idx).is_none() + { + self.execution_graph.add_edge(exec, node_idx, ()); + } + } + self.add_physical_dependencies(graph, node_idx, dependencies); + for prev in previous { + self.remove_node_if_dead(prev); + } + } +} diff --git a/fusor-ml/core/src/compute_graph/resolve/recognize_attention.rs b/fusor-ml/core/src/compute_graph/resolve/recognize_attention.rs new file mode 100644 index 000000000..d6eb75998 --- /dev/null +++ b/fusor-ml/core/src/compute_graph/resolve/recognize_attention.rs @@ -0,0 +1,369 @@ +//! Recognition of composed attention clusters. +//! +//! Runs third, after contractions and normalizations: by then the canonical +//! cluster from `Tensor::flash_attention` has collapsed to +//! `MatMul(Softmax(scale·MatMul(q, kᵀ) [+ mask]), v)` with the GQA-expand / +//! transpose / mask broadcast views still attached to the original +//! q/k/v/mask nodes. Recognition rebuilds the attention row program when its +//! gates pass; otherwise the cluster runs through the recognized matmul + +//! softmax kernels. + +use crate::{ + MatMulOperation, + nary_wise::{NaryOp, NaryScalar}, + view::ViewOperation, +}; + +use super::cluster_match::{binary_elementwise, layout_matches, unary_elementwise}; +use super::*; + +struct MatchedAttention { + q: NodeIndex, + k: NodeIndex, + v: NodeIndex, + mask: Option, + q_shape: Vec, + k_shape: Vec, + v_shape: Vec, + mask_shape: Option>, + scale: f32, + datatype: DataTypeEnum, + causal: bool, +} + +/// Match `select(kv_pos <= q_pos, scores, -inf)` over a single input, +/// returning the scores node. +fn match_causal_select(nary: &ElementwiseOperation) -> Option { + let NaryExpr::Op { children, function } = &nary.expression else { + return None; + }; + if function.op != NaryOp::Select || nary.inputs.len() != 1 { + return None; + } + let [condition, on_true, on_false] = children.as_slice() else { + return None; + }; + let NaryExpr::Op { + children: bound, + function: compare, + } = condition + else { + return None; + }; + if compare.op != NaryOp::LessEqual + || bound.as_slice() != [NaryExpr::DimIndex(3), NaryExpr::DimIndex(2)] + { + return None; + } + let NaryExpr::IndexedInput { + input_idx: 0, + indices, + } = on_true + else { + return None; + }; + if indices.len() != nary.shape.len() || !NaryExpr::is_elementwise_indices(indices) { + return None; + } + let negative_infinity = match on_false { + NaryExpr::Scalar(crate::nary_wise::NaryScalar::F32(value)) => *value == f32::NEG_INFINITY, + NaryExpr::Scalar(crate::nary_wise::NaryScalar::F16(value)) => { + *value == half::f16::NEG_INFINITY + } + _ => false, + }; + negative_infinity.then_some(nary.inputs[0]) +} + +impl Resolver { + pub(super) fn recognize_attention(&mut self, graph: &mut ComputeGraphInner) { + let candidates: Vec = self + .execution_graph + .node_indices() + .filter(|&node| { + matches!( + self.execution_graph[node].variant, + ExecutionVariant::MatMul(_) + ) + }) + .collect(); + for node in candidates { + if !self.execution_graph.contains_node(node) { + continue; + } + self.try_recognize_attention(graph, node); + } + } + + fn inner_view(&self, inner: NodeIndex) -> Option<&ViewOperation> { + let exec = self.get_input_node_in_exec_graph(inner)?; + match &self.execution_graph[exec].variant { + ExecutionVariant::View(view) => Some(view), + _ => None, + } + } + + fn inner_matmul(&self, inner: NodeIndex) -> Option<&MatMulOperation> { + let exec = self.get_input_node_in_exec_graph(inner)?; + match &self.execution_graph[exec].variant { + ExecutionVariant::MatMul(matmul) => Some(matmul), + _ => None, + } + } + + fn try_recognize_attention( + &mut self, + graph: &mut ComputeGraphInner, + node_idx: ExecutionNodeIndex, + ) -> bool { + let Some(matched) = self.match_attention(graph, node_idx) else { + return false; + }; + let mut dependencies = vec![matched.q, matched.k, matched.v]; + if let Some(mask) = matched.mask { + dependencies.push(mask); + } + // The recognized cluster lowers through the generic attention row + // program; shapes it cannot host stay composed (matmul + softmax + // row program + matmul). + let Some(operation) = crate::row_program::attention_row_program( + &graph.device(), + crate::row_program::AttentionInputs { + q: matched.q, + k: matched.k, + v: matched.v, + mask: matched.mask, + q_shape: &matched.q_shape, + k_shape: &matched.k_shape, + v_shape: &matched.v_shape, + mask_shape: matched.mask_shape.as_deref(), + scale: matched.scale, + input_dtype: matched.datatype, + causal: matched.causal, + }, + ) else { + return false; + }; + self.commit_recognized( + graph, + node_idx, + &dependencies, + ExecutionVariant::GraphOp(Arc::new(operation)), + ); + true + } + + fn match_attention( + &self, + graph: &ComputeGraphInner, + node_idx: ExecutionNodeIndex, + ) -> Option { + let ExecutionVariant::MatMul(out) = &self.execution_graph[node_idx].variant else { + return None; + }; + if !out.pre_element_wise[0].functions.is_empty() + || !out.pre_element_wise[1].functions.is_empty() + || !out.post_element_wise.functions.is_empty() + { + return None; + } + let probs_inner = out.first; + let v_eff_inner = out.second; + let v_eff_shape = out.second_shape.clone(); + let datatype = out.datatype; + + // probs = softmax(scores) along the last axis of a rank-4 space — + // matched in its composed form (the compiler no longer rewrites + // standalone softmax into a named operation). + let softmax = self.match_softmax_cluster(graph, probs_inner)?; + let shape = softmax.shape.to_vec(); + if shape.len() != 4 || softmax.axis != 3 { + return None; + } + let (batch, num_heads, q_seq_len, kv_seq_len) = (shape[0], shape[1], shape[2], shape[3]); + let mut scores_inner = softmax.input; + // The composed softmax reads its input twice (the shift and the max + // reduction), so the scores node has two in-cluster consumers; every + // node below it is read once. + let expected_consumers = |node: NodeIndex| if node == softmax.input { 2 } else { 1 }; + + // Causal masking: select(kv_pos <= q_pos, scaled, -inf) — pure index + // arithmetic emitted by `flash_attention_causal`. + let mut causal = false; + if let Some(scaled) = self.inner_nary(scores_inner).and_then(match_causal_select) { + if !self.exclusively_consumed(graph, scores_inner, expected_consumers(scores_inner)) { + return None; + } + causal = true; + scores_inner = scaled; + } + + // Optional additive mask: add(scaled, bcast(mask)) where the mask + // broadcast reads a rank-2 [q_seq, kv_seq] base. + let mut mask = None; + let mut mask_shape: Option> = None; + if !causal + && let Some((NaryOp::Add, lhs, rhs)) = + self.inner_nary(scores_inner).and_then(binary_elementwise) + { + let mut matched_mask = None; + for (scaled_side, mask_side) in [(lhs, rhs), (rhs, lhs)] { + let Some(view) = self.inner_view(mask_side) else { + continue; + }; + let expected = + Layout::from_parts(0, shape.clone().into(), [0, 0, kv_seq_len, 1].into()); + if !view.is_fully_defined() + || !layout_matches(Some(&view.layout), &expected) + || view.input_shape.as_ref() != [q_seq_len, kv_seq_len] + { + continue; + } + if !self.exclusively_consumed(graph, mask_side, 1) { + continue; + } + matched_mask = Some((scaled_side, view.input, view.input_shape.to_vec())); + break; + } + let (scaled_side, mask_base, base_shape) = matched_mask?; + if !self.exclusively_consumed(graph, scores_inner, expected_consumers(scores_inner)) { + return None; + } + mask = Some(mask_base); + mask_shape = Some(base_shape); + scores_inner = scaled_side; + } + + // scaled = qk · scale + let scale = { + let nary = self.inner_nary(scores_inner)?; + let (function, _) = unary_elementwise(nary)?; + match function.op { + NaryOp::MulConst(NaryScalar::F32(scale)) => scale, + NaryOp::MulConst(NaryScalar::F16(scale)) => scale.to_f32(), + _ => return None, + } + }; + let qk_inner = self.match_unary(scores_inner, |function| { + matches!(function.op, NaryOp::MulConst(_)) + })?; + if !self.exclusively_consumed(graph, scores_inner, expected_consumers(scores_inner)) { + return None; + } + + // qk = MatMul(q, kᵀ-view) + let qk = self.inner_matmul(qk_inner)?; + if !qk.pre_element_wise[0].functions.is_empty() + || !qk.pre_element_wise[1].functions.is_empty() + || !qk.post_element_wise.functions.is_empty() + { + return None; + } + let q = qk.first; + let q_shape = qk.first_shape.to_vec(); + let kt_inner = qk.second; + if q_shape.len() != 4 + || q_shape[..3] != [batch, num_heads, q_seq_len] + || !self.exclusively_consumed(graph, qk_inner, 1) + { + return None; + } + let head_dim = q_shape[3]; + let expanded_shape = [batch, num_heads, kv_seq_len, head_dim]; + + // kᵀ: a transpose view attached to the (possibly GQA-expanded) k. + let kt = self.inner_view(kt_inner)?; + let expected_kt = Layout::contiguous(&expanded_shape).transpose(2, 3); + if !kt.is_fully_defined() + || !layout_matches(Some(&kt.layout), &expected_kt) + || kt.input_shape.as_ref() != expanded_shape + || !self.exclusively_consumed(graph, kt_inner, 1) + { + return None; + } + let (k, k_shape) = self.peel_gqa_expand(graph, kt.input, &expanded_shape)?; + let (v, v_shape) = if v_eff_shape.as_ref() == expanded_shape { + self.peel_gqa_expand(graph, v_eff_inner, &expanded_shape)? + } else { + return None; + }; + if k_shape != v_shape { + return None; + } + + // The intermediate cluster must exist solely for this attention. + if !self.exclusively_consumed(graph, probs_inner, 1) { + return None; + } + + Some(MatchedAttention { + q, + k, + v, + mask, + q_shape, + k_shape, + v_shape, + mask_shape, + scale, + datatype, + causal, + }) + } + + /// Peel the canonical GQA expansion (flat-reinterpret view over a + /// stride-0 group broadcast) back to the original K/V node. Any other + /// node — including arbitrary user views like KV-cache slices — is the + /// tensor itself with the full expanded shape (groups == 1). + fn peel_gqa_expand( + &self, + graph: &ComputeGraphInner, + inner: NodeIndex, + expanded_shape: &[usize; 4], + ) -> Option<(NodeIndex, Vec)> { + let [batch, num_heads, kv_seq_len, head_dim] = *expanded_shape; + let unexpanded = Some((inner, expanded_shape.to_vec())); + + let Some(reinterpret) = self.inner_view(inner) else { + return unexpanded; + }; + // The flat reinterpret: contiguous rank-4 over a rank-5 grouped space. + if !reinterpret.is_fully_defined() + || reinterpret.input_shape.len() != 5 + || !layout_matches( + Some(&reinterpret.layout), + &Layout::contiguous(expanded_shape), + ) + { + return unexpanded; + } + let [b, num_kv_heads, groups, s, d] = *reinterpret.input_shape.as_ref() else { + return unexpanded; + }; + if b != batch + || s != kv_seq_len + || d != head_dim + || groups <= 1 + || num_kv_heads * groups != num_heads + { + return unexpanded; + } + let Some(broadcast) = self.inner_view(reinterpret.input) else { + return unexpanded; + }; + let expected_broadcast = Layout::from_parts( + 0, + [b, num_kv_heads, groups, s, d].into(), + [num_kv_heads * s * d, s * d, 0, d, 1].into(), + ); + if !broadcast.is_fully_defined() + || !layout_matches(Some(&broadcast.layout), &expected_broadcast) + || broadcast.input_shape.as_ref() != [b, num_kv_heads, s, d] + || !self.exclusively_consumed(graph, inner, 1) + || !self.exclusively_consumed(graph, reinterpret.input, 1) + { + return unexpanded; + } + Some((broadcast.input, vec![b, num_kv_heads, s, d])) + } +} diff --git a/fusor-ml/core/src/compute_graph/resolve/run.rs b/fusor-ml/core/src/compute_graph/resolve/run.rs index 66988792b..12cc100c3 100644 --- a/fusor-ml/core/src/compute_graph/resolve/run.rs +++ b/fusor-ml/core/src/compute_graph/resolve/run.rs @@ -86,12 +86,12 @@ impl Resolver { for idx in sorted_nodes { let node = &self.execution_graph[idx]; // Handle Tensor caching explicitly here - if let ComputeGraphNodeVariant::Tensor(data) = &node.variant { + if let ExecutionVariant::Tensor(data) = &node.variant { graph.set_cached_result(node.inner_idx, data.clone()); continue; } - if let Some(op) = self.lower_node(node) { + if let Some(op) = self.lower_node(idx, node) { queued_operations.push((node.inner_idx, op)); } } @@ -151,22 +151,25 @@ impl Resolver { .nodes .nodes .node_weight(node) - .map(|node| node_category(&node.variant)) + .map(|node| node_category_inner(&node.variant)) }) .flatten(); - // Map layout isn't really a kernel. Resolve it immediately - let map_layout = if let Some(node_data) = graph.nodes.nodes.node_weight(node) { + // A view that composes with its input's buffer layout isn't a + // kernel. Resolve it immediately as a zero-cost buffer view; + // anything else (fill regions, non-composable reshapes) falls + // through to the gather kernel below. + let view_result = if let Some(node_data) = graph.nodes.nodes.node_weight(node) { match &node_data.variant { - ComputeGraphNodeVariant::MapLayout(map_layout) => Some(map_layout.clone()), - ComputeGraphNodeVariant::Resize(resize) => resize.lower(graph), + ComputeGraphNodeVariant::View(view) => graph + .get_cached_result(view.input) + .and_then(|input| view.try_map_tensor(input)), _ => None, } } else { None }; - if let Some(map_layout) = map_layout { + if let Some(result) = view_result { let start = host_trace.then(Instant::now); - let result = map_layout.run(graph); // Cache the result graph.set_cached_result(node, result); // Map-layout nodes are resolved immediately — release any @@ -184,8 +187,7 @@ impl Resolver { } } else { let slice_copy = graph.nodes.nodes.node_weight(node).and_then(|node_data| { - let ComputeGraphNodeVariant::SliceAssign(slice_assign) = &node_data.variant - else { + let ComputeGraphNodeVariant::Assign(slice_assign) = &node_data.variant else { return None; }; Self::try_prepare_in_place_slice_assign_copy(graph, slice_assign) @@ -232,7 +234,9 @@ impl Resolver { let start = host_trace.then(Instant::now); let constraints = qmatmul.workgroup_shape_constraints(&device); - let workgroup_shape = constraints.solve(max_subgroup_size).unwrap_or_else(|| { + let workgroup_shape = constraints + .solve(max_subgroup_size, &device.limits()) + .unwrap_or_else(|| { panic!( "Failed to find a valid qmatmul workgroup shape for constraints {constraints:?}" ) @@ -349,9 +353,13 @@ impl Resolver { let start = host_trace.then(Instant::now); let constraints = operation.workgroup_shape_constraints(&device); - let workgroup_shape = constraints.solve(max_subgroup_size).unwrap_or_else(|| { - panic!("Failed to find a valid workgroup shape for constraints {constraints:?}") - }); + let workgroup_shape = constraints + .solve(max_subgroup_size, &device.limits()) + .unwrap_or_else(|| { + panic!( + "Failed to find a valid workgroup shape for constraints {constraints:?}" + ) + }); if let Some(start) = start { let elapsed = start.elapsed(); host_profile.workgroup += elapsed; diff --git a/fusor-ml/core/src/compute_graph/visualize.rs b/fusor-ml/core/src/compute_graph/visualize.rs index d3af26c90..b738cd9bd 100644 --- a/fusor-ml/core/src/compute_graph/visualize.rs +++ b/fusor-ml/core/src/compute_graph/visualize.rs @@ -1,7 +1,7 @@ use rustc_hash::FxHashMap; use super::queue::ComputeQueue; -use super::{ComputeGraphInner, ComputeGraphNodeVariant, GraphOperation, NodeIndex, layout_pass}; +use super::{ComputeGraphInner, ComputeGraphNodeVariant, NodeIndex, layout_pass}; use tabbycat::Graph; use tabbycat::{Edge, GraphBuilder, GraphType, Identity, Stmt, StmtList}; @@ -14,7 +14,7 @@ struct GraphVisPass { } impl GraphVisPass { - fn visit_nary(&mut self, key: NodeIndex, operation: &crate::nary_wise::NaryOperation) { + fn visit_nary(&mut self, key: NodeIndex, operation: &crate::nary_wise::ElementwiseOperation) { let output_layout = self.layout_pass.output_layout.get(&key).unwrap(); let id = Identity::quoted(format!("nary ({}) #{:?}", output_layout, key)); self.statements.push(Stmt::Node { @@ -31,47 +31,7 @@ impl GraphVisPass { self.identities.insert(key, id.clone()); } - fn visit_mat_mul(&mut self, key: NodeIndex, operation: &crate::MatMulOperation) { - let first = self.identities.get(&operation.first).unwrap(); - let second = self.identities.get(&operation.second).unwrap(); - let output_layout = self.layout_pass.output_layout.get(&key).unwrap(); - let id = Identity::quoted(format!("matmul ({}) #{:?}", output_layout, key)); - self.statements.push(Stmt::Node { - id: id.clone(), - port: None, - attr: None, - }); - self.statements.push(Stmt::Edge( - Edge::head_node(first.clone(), None).arrow_to_node(id.clone(), None), - )); - self.statements.push(Stmt::Edge( - Edge::head_node(second.clone(), None).arrow_to_node(id.clone(), None), - )); - self.identities.insert(key, id.clone()); - } - - fn visit_q_mat_mul( - &mut self, - key: NodeIndex, - operation: &crate::quantized::matmul::QMatMulOperation, - ) { - let input = self.identities.get(&operation.input).unwrap(); - let output_layout = self.layout_pass.output_layout.get(&key).unwrap(); - let header = format!("qmatmul ({}) #{:?}", output_layout, key); - let id = Identity::quoted(header); - self.statements.push(Stmt::Node { - id: id.clone(), - port: None, - attr: None, - }); - self.statements.push(Stmt::Edge( - Edge::head_node(input.clone(), None).arrow_to_node(id.clone(), None), - )); - self.identities.insert(key, id.clone()); - } - fn visit_reduce(&mut self, key: NodeIndex, operation: &crate::ReduceOperation) { - let input = self.identities.get(&operation.value).unwrap(); let output_layout = self.layout_pass.output_layout.get(&key).unwrap(); let id = Identity::quoted(format!( "{} ({}) #{:?}", @@ -84,95 +44,19 @@ impl GraphVisPass { port: None, attr: None, }); - self.statements.push(Stmt::Edge( - Edge::head_node(input.clone(), None).arrow_to_node(id.clone(), None), - )); - self.identities.insert(key, id.clone()); - } - - fn visit_graph_op(&mut self, key: NodeIndex, operation: &dyn GraphOperation) { - let output_layout = self.layout_pass.output_layout.get(&key).unwrap(); - let id = Identity::quoted(format!( - "{} ({}) #{:?}", - operation.category(), - output_layout, - key - )); - self.statements.push(Stmt::Node { - id: id.clone(), - port: None, - attr: None, - }); - - let mut dependencies = Vec::new(); - operation.visit_dependencies(&mut |dependency| { - dependencies.push(dependency); - }); - for dependency in dependencies { - let dependency = self.identities.get(&dependency).unwrap(); - self.statements.push(Stmt::Edge( - Edge::head_node(dependency.clone(), None).arrow_to_node(id.clone(), None), - )); - } - self.identities.insert(key, id.clone()); - } - - fn visit_flash_attention( - &mut self, - key: NodeIndex, - operation: &crate::FlashAttentionOperation, - ) { - let q = self.identities.get(&operation.q).unwrap(); - let k = self.identities.get(&operation.k).unwrap(); - let v = self.identities.get(&operation.v).unwrap(); - let output_layout = self.layout_pass.output_layout.get(&key).unwrap(); - let id = Identity::quoted(format!("flash_attention ({}) #{:?}", output_layout, key)); - self.statements.push(Stmt::Node { - id: id.clone(), - port: None, - attr: None, - }); - self.statements.push(Stmt::Edge( - Edge::head_node(q.clone(), None).arrow_to_node(id.clone(), None), - )); - self.statements.push(Stmt::Edge( - Edge::head_node(k.clone(), None).arrow_to_node(id.clone(), None), - )); - self.statements.push(Stmt::Edge( - Edge::head_node(v.clone(), None).arrow_to_node(id.clone(), None), - )); - if let Some(mask) = operation.mask { - let mask = self.identities.get(&mask).unwrap(); + for input in &operation.inputs { + let input = self.identities.get(input).unwrap(); self.statements.push(Stmt::Edge( - Edge::head_node(mask.clone(), None).arrow_to_node(id.clone(), None), + Edge::head_node(input.clone(), None).arrow_to_node(id.clone(), None), )); } self.identities.insert(key, id.clone()); } - fn visit_map_layout( - &mut self, - key: NodeIndex, - operation: &crate::map_layout::MapLayoutOperation, - ) { + fn visit_view(&mut self, key: NodeIndex, operation: &crate::view::ViewOperation) { let input = self.identities.get(&operation.input).unwrap(); let output_layout = self.layout_pass.output_layout.get(&key).unwrap(); - let id = Identity::quoted(format!("map_layout ({}) #{:?}", output_layout, key)); - self.statements.push(Stmt::Node { - id: id.clone(), - port: None, - attr: None, - }); - self.statements.push(Stmt::Edge( - Edge::head_node(input.clone(), None).arrow_to_node(id.clone(), None), - )); - self.identities.insert(key, id.clone()); - } - - fn visit_resize(&mut self, key: NodeIndex, operation: &crate::resize::ResizeOperation) { - let input = self.identities.get(&operation.input).unwrap(); - let output_layout = self.layout_pass.output_layout.get(&key).unwrap(); - let id = Identity::quoted(format!("resize ({}) #{:?}", output_layout, key)); + let id = Identity::quoted(format!("view ({}) #{:?}", output_layout, key)); self.statements.push(Stmt::Node { id: id.clone(), port: None, @@ -222,25 +106,6 @@ impl GraphVisPass { self.identities.insert(key, id.clone()); } - fn visit_q_embedding( - &mut self, - key: NodeIndex, - operation: &crate::quantized::embedding::QEmbeddingOperation, - ) { - let indexes = self.identities.get(&operation.indexes).unwrap(); - let output_layout = self.layout_pass.output_layout.get(&key).unwrap(); - let id = Identity::quoted(format!("q_embedding ({}) #{:?}", output_layout, key)); - self.statements.push(Stmt::Node { - id: id.clone(), - port: None, - attr: None, - }); - self.statements.push(Stmt::Edge( - Edge::head_node(indexes.clone(), None).arrow_to_node(id.clone(), None), - )); - self.identities.insert(key, id.clone()); - } - fn visit_tensor(&mut self, key: NodeIndex, _operation: &crate::tensor::TensorData) { let output_layout = self.layout_pass.output_layout.get(&key).unwrap(); let id = Identity::quoted(format!("tensor ({}) #{:?}", output_layout, key)); @@ -295,28 +160,12 @@ impl ComputeGraphInner { let node_data = self.nodes.nodes.node_weight(node).expect("Node not found"); match &node_data.variant { - ComputeGraphNodeVariant::Nary(op) => graph_vis_pass.visit_nary(node, op), - ComputeGraphNodeVariant::MatMul(op) => graph_vis_pass.visit_mat_mul(node, op), - ComputeGraphNodeVariant::QMatMul(op) => graph_vis_pass.visit_q_mat_mul(node, op), - ComputeGraphNodeVariant::QEmbedding(op) => { - graph_vis_pass.visit_q_embedding(node, op) - } - ComputeGraphNodeVariant::Reduce(op) => graph_vis_pass.visit_reduce(node, op), - ComputeGraphNodeVariant::FlashAttention(op) => { - graph_vis_pass.visit_flash_attention(node, op) - } - ComputeGraphNodeVariant::GraphOp(op) => { - graph_vis_pass.visit_graph_op(node, op.as_ref()) - } - ComputeGraphNodeVariant::MapLayout(op) => graph_vis_pass.visit_map_layout(node, op), - ComputeGraphNodeVariant::Resize(op) => graph_vis_pass.visit_resize(node, op), - ComputeGraphNodeVariant::SliceAssign(op) => { - graph_vis_pass.visit_slice_assign(node, op) - } ComputeGraphNodeVariant::Tensor(op) => graph_vis_pass.visit_tensor(node, op), - ComputeGraphNodeVariant::Dequantize(op) => { - graph_vis_pass.visit_dequantize(node, op) - } + ComputeGraphNodeVariant::QMatrix(op) => graph_vis_pass.visit_dequantize(node, op), + ComputeGraphNodeVariant::Elementwise(op) => graph_vis_pass.visit_nary(node, op), + ComputeGraphNodeVariant::Reduce(op) => graph_vis_pass.visit_reduce(node, op), + ComputeGraphNodeVariant::View(op) => graph_vis_pass.visit_view(node, op), + ComputeGraphNodeVariant::Assign(op) => graph_vis_pass.visit_slice_assign(node, op), } } diff --git a/fusor-ml/core/src/conv.rs b/fusor-ml/core/src/conv.rs deleted file mode 100644 index ca6e80764..000000000 --- a/fusor-ml/core/src/conv.rs +++ /dev/null @@ -1,481 +0,0 @@ -use std::{any::Any, hash::Hash}; - -use fusor_tile_ir as tile_ir; -use rustc_hash::{FxHashMap, FxHasher}; - -use crate::{ - DataTypeEnum, Device, Layout, - compute_graph::{ComputeGraphInner, GraphOperation, NodeIndex}, - mir::{ - inputs::MirValue, - kernel_backend, - kernel_backend::DirectKernel, - operation::Operation, - workgroup_shape::{Constraint, WorkgroupShape, WorkgroupShapeConstraints}, - }, - nary_direct::{ - TensorMeta, ValueTile, flat_layout, layout_index, linear_group, output_dims_from_flat, - tile_u32, - }, - tensor::{TensorData, TensorLayoutInfo}, - visit_tiled::distribute_workgroups, -}; - -#[derive(Clone, Debug)] -pub(crate) struct ConvNdOperation { - input: NodeIndex, - weight: NodeIndex, - bias: Option, - input_shape: Box<[usize]>, - weight_shape: Box<[usize]>, - output_shape: Box<[usize]>, - padding: Box<[u32]>, - strides: Box<[u32]>, - datatype: DataTypeEnum, -} - -pub(crate) struct ConvNdNodes { - pub(crate) input: NodeIndex, - pub(crate) weight: NodeIndex, - pub(crate) bias: Option, -} - -pub(crate) struct ConvNdShapeSpec<'a> { - pub(crate) input_shape: &'a [usize], - pub(crate) weight_shape: &'a [usize], - pub(crate) bias_shape: Option<&'a [usize]>, - pub(crate) padding: &'a [usize], - pub(crate) strides: &'a [usize], -} - -impl ConvNdOperation { - pub(crate) fn output_shape(&self) -> &[usize] { - &self.output_shape - } - - pub(crate) fn new( - nodes: ConvNdNodes, - spec: ConvNdShapeSpec<'_>, - datatype: DataTypeEnum, - device: &Device, - ) -> Option { - let ConvNdNodes { - input, - weight, - bias, - } = nodes; - let ConvNdShapeSpec { - input_shape, - weight_shape, - bias_shape, - padding, - strides, - } = spec; - - let spatial_rank = padding.len(); - if spatial_rank == 0 - || strides.len() != spatial_rank - || input_shape.len() != spatial_rank + 2 - || weight_shape.len() != spatial_rank + 2 - || !matches!(datatype, DataTypeEnum::F32 | DataTypeEnum::F16) - || (datatype == DataTypeEnum::F16 && !device.f16_supported()) - { - return None; - } - - let batch = input_shape[0]; - let in_channels = input_shape[1]; - let out_channels = weight_shape[0]; - if weight_shape[1] != in_channels - || strides.contains(&0) - || input_shape.contains(&0) - || weight_shape.contains(&0) - { - return None; - } - if let Some(bias_shape) = bias_shape - && bias_shape != [out_channels] - { - return None; - } - - let mut output_shape = Vec::with_capacity(input_shape.len()); - output_shape.push(batch); - output_shape.push(out_channels); - for axis in 0..spatial_rank { - let input_len = input_shape[axis + 2]; - let kernel_len = weight_shape[axis + 2]; - let padded_len = input_len.checked_add(padding[axis].checked_mul(2)?)?; - let out_len = padded_len - .checked_sub(kernel_len)? - .checked_div(strides[axis])? - + 1; - if out_len == 0 { - return None; - } - output_shape.push(out_len); - } - - let kernel_volume = weight_shape[2..] - .iter() - .try_fold(1usize, |acc, dim| acc.checked_mul(*dim))?; - let total_outputs = output_shape - .iter() - .try_fold(1usize, |acc, dim| acc.checked_mul(*dim))?; - u32::try_from(total_outputs).ok()?; - u32::try_from(in_channels.checked_mul(kernel_volume)?).ok()?; - - Some(Self { - input, - weight, - bias, - input_shape: input_shape.into(), - weight_shape: weight_shape.into(), - output_shape: output_shape.into_boxed_slice(), - padding: padding - .iter() - .copied() - .map(u32::try_from) - .collect::, _>>() - .ok()? - .into_boxed_slice(), - strides: strides - .iter() - .copied() - .map(u32::try_from) - .collect::, _>>() - .ok()? - .into_boxed_slice(), - datatype, - }) - } - - #[inline] - fn spatial_rank(&self) -> usize { - self.padding.len() - } - - #[inline] - fn output_index(&self) -> usize { - 2 + usize::from(self.bias.is_some()) - } - - #[inline] - fn total_outputs(&self) -> Option { - self.output_shape - .iter() - .try_fold(1u32, |acc, dim| acc.checked_mul((*dim).try_into().ok()?)) - } - - fn storage_read( - kb: &mut tile_ir::KernelBuilder, - datatype: DataTypeEnum, - tensor: tile_ir::KernelTensorRef, - ) -> crate::nary_direct::Storage2 { - let element = crate::nary_direct::datatype_element(datatype); - let storage = kb.read(element, tensor); - match datatype { - DataTypeEnum::F32 => crate::nary_direct::Storage2::F32(storage), - DataTypeEnum::F16 => crate::nary_direct::Storage2::F16(storage), - DataTypeEnum::U32 => crate::nary_direct::Storage2::U32(storage), - } - } - - fn storage_write( - kb: &mut tile_ir::KernelBuilder, - datatype: DataTypeEnum, - tensor: tile_ir::KernelTensorRef, - ) -> crate::nary_direct::Storage2 { - let element = crate::nary_direct::datatype_element(datatype); - let storage = kb.write(element, tensor); - match datatype { - DataTypeEnum::F32 => crate::nary_direct::Storage2::F32(storage), - DataTypeEnum::F16 => crate::nary_direct::Storage2::F16(storage), - DataTypeEnum::U32 => crate::nary_direct::Storage2::U32(storage), - } - } -} - -struct ConvNdDirectKernelVariant; - -impl Operation for ConvNdOperation { - fn hash_kernel_fields(&self, state: &mut FxHasher) { - self.input_shape.hash(state); - self.weight_shape.hash(state); - self.output_shape.hash(state); - self.padding.hash(state); - self.strides.hash(state); - self.bias.is_some().hash(state); - self.datatype.hash(state); - } - - fn workgroup_shape_constraints(&self, device: &Device) -> WorkgroupShapeConstraints { - let mut constraints = WorkgroupShapeConstraints::new(); - let workgroup_size = device.limits().max_compute_workgroup_size_x.min(256); - constraints.add_constraint(0, Constraint::equals(workgroup_size)); - constraints.add_constraint(1, Constraint::equals(1)); - constraints.add_constraint(2, Constraint::equals(1)); - constraints - } - - fn dispatch_size(&self, workgroup_shape: &WorkgroupShape, inputs: &[MirValue]) -> [u32; 3] { - let Some(output) = inputs - .get(self.output_index()) - .and_then(MirValue::as_tensor) - else { - return [1, 1, 1]; - }; - let Some(total_outputs) = self.total_outputs() else { - return [1, 1, 1]; - }; - distribute_workgroups( - total_outputs.div_ceil(workgroup_shape.x()), - output - .device() - .limits() - .max_compute_workgroups_per_dimension, - ) - } - - fn visit_dependencies(&self, f: &mut dyn FnMut(NodeIndex)) { - f(self.input); - f(self.weight); - if let Some(bias) = self.bias { - f(bias); - } - } - - fn inputs(&self, nodes: &ComputeGraphInner) -> Vec { - let input = nodes.get_cached_result(self.input).unwrap(); - let weight = nodes.get_cached_result(self.weight).unwrap(); - let mut inputs = Vec::with_capacity(3 + usize::from(self.bias.is_some())); - inputs.push(input.clone().into()); - inputs.push(weight.clone().into()); - if let Some(bias) = self.bias { - inputs.push(nodes.get_cached_result(bias).unwrap().clone().into()); - } - inputs.push( - TensorData::new_for_shape(input.device(), &self.output_shape, self.datatype).into(), - ); - inputs - } - - fn output(&self, _nodes: &ComputeGraphInner, inputs: &[MirValue]) -> MirValue { - inputs[self.output_index()].clone() - } - - fn build_direct_kernel( - &self, - graph: &ComputeGraphInner, - workgroup_shape: &WorkgroupShape, - inputs: &[MirValue], - ) -> Option { - let input = inputs.first()?.as_tensor()?.clone(); - let weight = inputs.get(1)?.as_tensor()?.clone(); - let bias = self - .bias - .is_some() - .then(|| inputs.get(2)?.as_tensor().cloned()) - .flatten(); - let output = inputs.get(self.output_index())?.as_tensor()?.clone(); - if input.datatype() != self.datatype - || weight.datatype() != self.datatype - || output.datatype() != self.datatype - || bias - .as_ref() - .is_some_and(|bias| bias.datatype() != self.datatype) - || (self.datatype == DataTypeEnum::F16 && !graph.device().f16_supported()) - { - return None; - } - - let input_meta = TensorMeta::new(&input)?; - let weight_meta = TensorMeta::new(&weight)?; - let bias_meta = match &bias { - Some(bias) => Some(TensorMeta::new(bias)?), - None => None, - }; - let output_meta = TensorMeta::new(&output)?; - let input_shape = to_u32_vec(&self.input_shape)?; - let weight_shape = to_u32_vec(&self.weight_shape)?; - let output_shape_usize = self.output_shape.to_vec(); - let kernel_shape = weight_shape[2..].to_vec(); - let kernel_volume = kernel_shape - .iter() - .try_fold(1u32, |acc, dim| acc.checked_mul(*dim))?; - let reduce_len = input_shape[1].checked_mul(kernel_volume)?; - let total_outputs = self.total_outputs()?; - let dispatch_size = self.dispatch_size(workgroup_shape, inputs); - let cache_key = self.kernel_cache_key_with_dispatch( - kernel_backend::KernelVariantKey::of::(), - Some(workgroup_shape), - dispatch_size, - inputs, - ); - - let input_buffer = input.buffer().clone(); - let weight_buffer = weight.buffer().clone(); - let bias_buffer = bias.as_ref().map(|bias| bias.buffer().clone()); - let output_buffer = output.buffer().clone(); - let input_layout = flat_layout(input_meta.allocation_len); - let weight_layout = flat_layout(weight_meta.allocation_len); - let bias_layout = bias_meta - .as_ref() - .map(|meta| flat_layout(meta.allocation_len)); - let output_layout = flat_layout(output_meta.allocation_len); - let input_meta_body = input_meta.clone(); - let weight_meta_body = weight_meta.clone(); - let bias_meta_body = bias_meta.clone(); - let output_meta_body = output_meta.clone(); - let padding = self.padding.to_vec(); - let strides = self.strides.to_vec(); - let spatial_rank = self.spatial_rank(); - let datatype = self.datatype; - let block = workgroup_shape.x(); - - kernel_backend::run_kernel( - graph.device().kernel_cache(), - self.name(), - cache_key, - dispatch_size, - move |kb| { - let input_tensor = - tile_ir::KernelTensorRef::new(input_buffer.clone(), input_layout.clone()); - let weight_tensor = - tile_ir::KernelTensorRef::new(weight_buffer.clone(), weight_layout.clone()); - let output_tensor = - tile_ir::KernelTensorRef::new(output_buffer.clone(), output_layout.clone()); - let input_storage = Self::storage_read(kb, datatype, input_tensor); - let weight_storage = Self::storage_read(kb, datatype, weight_tensor); - let bias_storage = bias_buffer - .as_ref() - .zip(bias_layout.as_ref()) - .zip(bias_meta_body.as_ref()) - .map(|((buffer, layout), _)| { - let bias_tensor = - tile_ir::KernelTensorRef::new(buffer.clone(), layout.clone()); - Self::storage_read(kb, datatype, bias_tensor) - }); - let output_storage = Self::storage_write(kb, datatype, output_tensor); - - kb.program().program_grid(block, dispatch_size, |program| { - let lane = program.lane(); - let group = linear_group(program, dispatch_size); - let flat = group * block + lane; - let in_bounds = flat.lt(total_outputs); - let output_dims = output_dims_from_flat(flat.clone(), &output_shape_usize); - let batch = output_dims[0].clone(); - let out_channel = output_dims[1].clone(); - - let zero = tile_ir::tile::Tile::literal(tile_ir::TileLiteral::f32(0.0)); - let initial = if let (Some(bias_storage), Some(bias_meta)) = - (bias_storage.as_ref(), bias_meta_body.as_ref()) - { - bias_storage - .load( - program, - layout_index(bias_meta, std::slice::from_ref(&out_channel)), - in_bounds.clone(), - ) - .into_f32() - } else { - zero - }; - - let [sum] = program.fold( - tile_ir::tile::range(reduce_len), - [initial], - |program, reduce_index, [acc]| { - let in_channel = reduce_index.clone() / tile_u32(kernel_volume); - let kernel_linear = reduce_index % tile_u32(kernel_volume); - let mut active = in_bounds.clone(); - let mut input_coords = Vec::with_capacity(spatial_rank + 2); - input_coords.push(batch.clone()); - input_coords.push(in_channel.clone()); - let mut weight_coords = Vec::with_capacity(spatial_rank + 2); - weight_coords.push(out_channel.clone()); - weight_coords.push(in_channel); - - for axis in 0..spatial_rank { - let divisor = kernel_shape[axis + 1..] - .iter() - .fold(1u32, |acc, dim| acc.saturating_mul(*dim)); - let quotient = if divisor == 1 { - kernel_linear.clone() - } else { - kernel_linear.clone() / tile_u32(divisor) - }; - let kernel_coord = quotient % tile_u32(kernel_shape[axis]); - let padded_coord = output_dims[axis + 2].clone() * strides[axis] - + kernel_coord.clone(); - let source_coord = padded_coord.clone() - padding[axis]; - let valid_coord = padded_coord - .ge(padding[axis]) - .and(source_coord.clone().lt(input_shape[axis + 2])); - active = active.and(valid_coord); - input_coords.push(source_coord); - weight_coords.push(kernel_coord); - } - - let input_index = layout_index(&input_meta_body, &input_coords); - let weight_index = layout_index(&weight_meta_body, &weight_coords); - let input_value = input_storage - .load(program, input_index, active.clone()) - .into_f32(); - let input_value = tile_ir::tile::Tile::select( - active.clone(), - input_value, - tile_ir::tile::Tile::literal(tile_ir::TileLiteral::f32(0.0)), - ); - let weight_value = weight_storage - .load(program, weight_index, active) - .into_f32(); - [acc + input_value * weight_value] - }, - ); - - output_storage.store( - program, - layout_index(&output_meta_body, &output_dims), - ValueTile::F32(sum).cast_to(datatype), - in_bounds, - ); - }); - Some(()) - }, - ) - } - - fn name(&self) -> String { - format!("conv{}d_direct", self.spatial_rank()) - } -} - -impl GraphOperation for ConvNdOperation { - fn as_any(&self) -> &dyn Any { - self - } - - fn category(&self) -> &'static str { - "conv_nd" - } - - fn output_layout( - &self, - _input_layouts: &FxHashMap, - ) -> Option { - Some(TensorLayoutInfo::new( - Layout::contiguous(&self.output_shape), - self.datatype, - )) - } -} - -fn to_u32_vec(values: &[usize]) -> Option> { - values - .iter() - .copied() - .map(u32::try_from) - .collect::, _>>() - .ok() -} diff --git a/fusor-ml/core/src/device.rs b/fusor-ml/core/src/device.rs index 5c80bb5ef..e16d9a089 100644 --- a/fusor-ml/core/src/device.rs +++ b/fusor-ml/core/src/device.rs @@ -423,6 +423,29 @@ impl Device { storage_bindings.min(bind_group_bindings).saturating_sub(1) } + /// A conservative estimate of the device's last-level cache. Data reused + /// below this size is treated as cache-resident — re-reading it costs no + /// bandwidth — so the reuse-driven tilings (which trade thread-level + /// parallelism for explicit reuse) only engage above it. wgpu exposes no + /// cache size, so this is a floor per device class; override with + /// `FUSOR_LAST_LEVEL_CACHE_BYTES` when tuning a specific part. + pub fn last_level_cache_bytes(&self) -> u64 { + if let Ok(value) = std::env::var("FUSOR_LAST_LEVEL_CACHE_BYTES") + && let Ok(parsed) = value.parse::() + { + return parsed; + } + let info = &self.inner.adapter_info; + // Apple-silicon system-level cache starts at 8 MiB on the base M1 + // and only grows with tier; other GPU classes floor lower (older + // discrete L2s are 2-4 MiB, integrated parts share CPU cache). + if info.backend == wgpu::Backend::Metal && info.name.starts_with("Apple") { + 8 << 20 + } else { + 4 << 20 + } + } + pub fn features(&self) -> wgpu::Features { self.inner.device.features() } diff --git a/fusor-ml/core/src/flash_attention.rs b/fusor-ml/core/src/flash_attention.rs deleted file mode 100644 index b57be18c6..000000000 --- a/fusor-ml/core/src/flash_attention.rs +++ /dev/null @@ -1,3 +0,0 @@ -pub(crate) use crate::mir::kernel_backend::flash_attention::{ - FlashAttentionInputs, FlashAttentionOperation, -}; diff --git a/fusor-ml/core/src/index_select.rs b/fusor-ml/core/src/index_select.rs index fc590068d..1bdbe63ef 100644 --- a/fusor-ml/core/src/index_select.rs +++ b/fusor-ml/core/src/index_select.rs @@ -1,6 +1,6 @@ use crate::{ Tensor, - nary_wise::{NaryExpr, NaryOperation}, + nary_wise::{ElementwiseOperation, NaryExpr}, }; /// Compute the output shape for an index_select operation. @@ -28,7 +28,7 @@ impl Tensor { indexes.assert_datatype::(); assert!(dimension < self.rank()); let output_shape = index_select_output_shape(dimension, self.shape(), indexes.shape()); - let nary = NaryOperation { + let nary = ElementwiseOperation { inputs: vec![self.key(), indexes.key()], expression: NaryExpr::index_select(self.rank(), dimension), shape: output_shape.clone(), diff --git a/fusor-ml/core/src/lib.rs b/fusor-ml/core/src/lib.rs index 5f9f56b8d..c6c5054c9 100644 --- a/fusor-ml/core/src/lib.rs +++ b/fusor-ml/core/src/lib.rs @@ -21,21 +21,18 @@ pub use tabbycat; pub use wgpu::{WasmNotSend, WasmNotSendSync, WasmNotSync}; pub use matmul::*; -pub use resize::ShapeWithOneHole; +pub use view::ShapeWithOneHole; +mod access_analysis; mod composite; mod compute_graph; -mod conv; pub use compute_graph::NodeIndex; mod device; mod element_wise; -mod flash_attention; -pub(crate) use flash_attention::{FlashAttentionInputs, FlashAttentionOperation}; mod index_select; #[doc(hidden)] pub mod kernel_selection; mod layout; -mod map_layout; pub mod matmul; mod mir; mod nary_direct; @@ -45,14 +42,13 @@ mod quantized; mod rank; mod reduce; mod reduce_direct; -mod rms_norm; -pub(crate) use rms_norm::RmsNormOperation; -mod resize; +mod reduce_tiled; +mod row_program; mod sampling; mod slice_assign; -mod softmax; mod tensor; mod top_k; +mod view; pub use top_k::{ GpuMirostat2Sampler, GpuMirostat2SamplerParams, GpuStandardSamplerParams, PendingGpuSampledToken, diff --git a/fusor-ml/core/src/map_layout.rs b/fusor-ml/core/src/map_layout.rs deleted file mode 100644 index b77d82e30..000000000 --- a/fusor-ml/core/src/map_layout.rs +++ /dev/null @@ -1,192 +0,0 @@ -use std::{fmt::Debug, sync::Arc}; - -use crate::{Layout, Tensor, TensorData, compute_graph::NodeIndex, mir::operation::Operation}; - -type MapLayout = Arc Layout + Send + Sync>; - -#[derive(Clone)] -pub(crate) struct MapLayoutOperation { - pub(crate) input: NodeIndex, - pub(crate) map_layout_fn: MapLayout, -} - -impl Debug for MapLayoutOperation { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("MapLayoutOperation") - .field("input", &self.input) - .finish() - } -} - -impl MapLayoutOperation { - pub fn new( - input: NodeIndex, - map_layout_fn: impl Fn(&Layout) -> Layout + Send + Sync + 'static, - ) -> Self { - Self { - input, - map_layout_fn: Arc::new(map_layout_fn), - } - } - - pub fn map_tensor(&self, tensor: &TensorData) -> TensorData { - TensorData::new_from_parts( - tensor.device(), - tensor.buffer().clone(), - self.map_layout(tensor.layout()), - tensor.datatype(), - ) - } - - pub fn map_layout(&self, layout: &Layout) -> Layout { - (self.map_layout_fn)(layout) - } - - pub fn run(&self, graph: &mut crate::compute_graph::ComputeGraphInner) -> TensorData { - let input = graph.get_result(self.input).unwrap(); - self.map_tensor(&input) - } -} - -impl Operation for MapLayoutOperation { - fn hash_kernel_fields(&self, _state: &mut rustc_hash::FxHasher) {} - - fn workgroup_shape_constraints( - &self, - _: &crate::Device, - ) -> crate::mir::workgroup_shape::WorkgroupShapeConstraints { - Default::default() - } - - fn dispatch_size( - &self, - _: &crate::mir::workgroup_shape::WorkgroupShape, - _: &[crate::mir::inputs::MirValue], - ) -> [u32; 3] { - [1, 1, 1] - } - - fn visit_dependencies(&self, f: &mut dyn FnMut(NodeIndex)) { - f(self.input); - } - - fn inputs( - &self, - nodes: &crate::compute_graph::ComputeGraphInner, - ) -> Vec { - vec![nodes.get_result(self.input).unwrap().into()] - } - - fn output( - &self, - _: &crate::compute_graph::ComputeGraphInner, - inputs: &[crate::mir::inputs::MirValue], - ) -> crate::mir::inputs::MirValue { - let input = inputs[0].as_tensor().unwrap(); - self.map_tensor(input).into() - } - - fn build_direct_kernel( - &self, - _: &crate::compute_graph::ComputeGraphInner, - _: &crate::mir::workgroup_shape::WorkgroupShape, - _: &[crate::mir::inputs::MirValue], - ) -> Option { - None - } - - fn name(&self) -> String { - "map_layout".to_string() - } -} - -impl Tensor { - pub fn restride(&self, specs: impl Into>) -> Tensor { - let specs = specs.into(); - self.add_map_layout(MapLayoutOperation::new(self.key(), move |layout| { - layout.restride(&specs) - })) - } - - /// Replace the tensor's layout with `new_layout`, treating the underlying - /// buffer as a flat blob. The user-supplied offset and strides are absolute - /// (in buffer elements), so the input must itself be a contiguous view with - /// offset 0 — otherwise the user's strides would compose nonsensically with - /// the input's own strides and silently read the wrong elements. - /// - /// Callers that need to re-layout a non-contiguous view should materialize - /// it first (e.g. with `to_concrete()`), or use [`restride`] which composes - /// stride specs relative to the input's current strides. - pub fn restride_layout(&self, new_layout: Layout) -> Tensor { - self.add_map_layout(MapLayoutOperation::new(self.key(), move |input_layout| { - assert!( - input_layout.is_contiguous() && input_layout.offset() == 0, - "restride_layout requires a contiguous, offset-0 input — got \ - offset={} strides={:?} shape={:?}. Call `.to_concrete()` first, \ - or use `restride` for stride composition.", - input_layout.offset(), - input_layout.strides(), - input_layout.shape() - ); - new_layout.clone() - })) - } - - pub(crate) fn broadcast_as(&self, out_shape: impl AsRef<[usize]>) -> Tensor { - let out_shape = out_shape.as_ref(); - let shape = self.shape(); - assert!( - out_shape.len() >= shape.len(), - "The output rank must be at least the input rank" - ); - let specs: Vec = (0..out_shape.len()) - .map(|out_i| { - let in_i = out_i as isize - (out_shape.len() as isize - shape.len() as isize); - if in_i < 0 { - crate::StrideSpec::dim_with(0, out_shape[out_i], 0) - } else { - let in_i = in_i as usize; - if shape[in_i] == 1 && out_shape[out_i] > 1 { - crate::StrideSpec::dim_with(in_i, out_shape[out_i], 0) - } else { - crate::StrideSpec::dim(in_i, out_shape[out_i]) - } - } - }) - .collect(); - self.restride(specs) - } - - pub(crate) fn broadcast_together(first: &Tensor, second: &Tensor) -> (Tensor, Tensor) { - assert_eq!(first.datatype(), second.datatype()); - let first_shape = first.shape(); - let second_shape = second.shape(); - let rank = first_shape.len().max(second_shape.len()); - let shape: Vec = (0..rank) - .map(|i| { - let a = i + first_shape.len(); - let b = i + second_shape.len(); - let a = if a >= rank { first_shape[a - rank] } else { 1 }; - let b = if b >= rank { second_shape[b - rank] } else { 1 }; - assert!( - a == b || a == 1 || b == 1, - "Cannot broadcast shapes {:?} and {:?}", - first_shape, - second_shape - ); - a.max(b) - }) - .collect(); - (first.broadcast_as(&shape), second.broadcast_as(&shape)) - } - - pub(crate) fn broadcast_then_elementwise_op( - first: &Tensor, - second: &Tensor, - op: impl Fn(Tensor, Tensor) -> Tensor, - ) -> Tensor { - let (b1, b2) = Tensor::broadcast_together(first, second); - assert_eq!(b1.shape(), b2.shape()); - op(b1, b2) - } -} diff --git a/fusor-ml/core/src/matmul/direct.rs b/fusor-ml/core/src/matmul/direct.rs deleted file mode 100644 index 2813a9e4b..000000000 --- a/fusor-ml/core/src/matmul/direct.rs +++ /dev/null @@ -1,241 +0,0 @@ -use fusor_tile_ir as tile_ir; - -use crate::{ - matmul::MatMulOperation, - mir::{ - inputs::MirValue, kernel_backend, kernel_backend::DirectKernel, operation::Operation, - workgroup_shape::WorkgroupShape, - }, - nary_direct::{ - TensorMeta, ValueTile, apply_unary_function_chain, flat_layout, layout_index, linear_group, - output_dims_from_flat, - }, - tensor::DataTypeEnum, - visit_tiled::distribute_workgroups, -}; - -const BLOCK: usize = 256; - -struct MatmulSerialDirectKernelVariant; - -pub(crate) fn build_serial_matmul_direct_kernel( - operation: &MatMulOperation, - graph: &crate::compute_graph::ComputeGraphInner, - workgroup_shape: &WorkgroupShape, - inputs: &[MirValue], -) -> Option { - let [input_a, input_b, output] = inputs else { - return None; - }; - let input_a = input_a.as_tensor()?.clone(); - let input_b = input_b.as_tensor()?.clone(); - let output = output.as_tensor()?.clone(); - if input_a.layout().rank() != input_b.layout().rank() - || input_a.layout().rank() != output.layout().rank() - || input_a.layout().rank() < 2 - || input_a.datatype() != input_b.datatype() - { - return None; - } - if (input_a.datatype() == DataTypeEnum::F16 - || input_b.datatype() == DataTypeEnum::F16 - || output.datatype() == DataTypeEnum::F16) - && !graph.device().f16_supported() - { - return None; - } - - let total_outputs = output - .layout() - .shape() - .iter() - .try_fold(1u32, |acc, dim| acc.checked_mul((*dim).try_into().ok()?))?; - let dispatch_size = distribute_workgroups( - total_outputs.div_ceil(BLOCK as u32), - graph.device().limits().max_compute_workgroups_per_dimension, - ); - let cache_key = operation.kernel_cache_key_with_dispatch( - kernel_backend::KernelVariantKey::of::(), - Some(workgroup_shape), - dispatch_size, - inputs, - ); - - let a_meta = TensorMeta::new(&input_a)?; - let b_meta = TensorMeta::new(&input_b)?; - let y_meta = TensorMeta::new(&output)?; - if operation.pre_element_wise[0].input_datatype() != a_meta.datatype - || operation.pre_element_wise[1].input_datatype() != b_meta.datatype - { - return None; - } - let a_product_dtype = operation.pre_element_wise[0].out_datatype(); - let b_product_dtype = operation.pre_element_wise[1].out_datatype(); - if a_product_dtype != b_product_dtype - || operation.post_element_wise.input_datatype() != a_product_dtype - { - return None; - } - let acc_dtype = match a_product_dtype { - DataTypeEnum::U32 => DataTypeEnum::U32, - DataTypeEnum::F32 | DataTypeEnum::F16 => DataTypeEnum::F32, - }; - let result_dtype = a_product_dtype; - let out_shape = output - .layout() - .shape() - .iter() - .copied() - .map(u32::try_from) - .collect::, _>>() - .ok()?; - let rank = out_shape.len(); - let total_outputs = out_shape - .iter() - .try_fold(1u32, |acc, dim| acc.checked_mul(*dim))?; - let k: u32 = input_a - .layout() - .shape() - .last() - .copied() - .and_then(|value| value.try_into().ok())?; - - let a_buffer = input_a.buffer().clone(); - let b_buffer = input_b.buffer().clone(); - let y_buffer = output.buffer().clone(); - let a_layout = flat_layout(a_meta.allocation_len); - let b_layout = flat_layout(b_meta.allocation_len); - let y_layout = flat_layout(y_meta.allocation_len); - let a_meta_body = a_meta.clone(); - let b_meta_body = b_meta.clone(); - let y_meta_body = y_meta.clone(); - let pre_a = operation.pre_element_wise[0].clone(); - let pre_b = operation.pre_element_wise[1].clone(); - let post = operation.post_element_wise.clone(); - - kernel_backend::run_kernel( - graph.device().kernel_cache(), - operation.name(), - cache_key, - dispatch_size, - move |kb| { - let a_tensor = tile_ir::KernelTensorRef::new(a_buffer.clone(), a_layout.clone()); - let b_tensor = tile_ir::KernelTensorRef::new(b_buffer.clone(), b_layout.clone()); - let y_tensor = tile_ir::KernelTensorRef::new(y_buffer.clone(), y_layout.clone()); - let a_storage = storage_read(kb, a_meta_body.datatype, a_tensor); - let b_storage = storage_read(kb, b_meta_body.datatype, b_tensor); - let y_storage = storage_write(kb, y_meta_body.datatype, y_tensor); - - kb.program() - .program_grid(BLOCK as u32, dispatch_size, |program| { - let lane = program.lane(); - let group = linear_group(program, dispatch_size); - let flat = group * BLOCK as u32 + lane.clone(); - let in_bounds = flat.lt(total_outputs); - let dims = output_dims_from_flat_u32(flat.clone(), &out_shape); - - let value_at = - |program: &mut tile_ir::tile::TileBlock<'_>, - k_index: tile_ir::tile::Tile| { - let mut a_coords = dims[..rank - 1].to_vec(); - a_coords.push(k_index.clone()); - let mut b_coords = dims[..rank - 2].to_vec(); - b_coords.push(k_index); - b_coords.push(dims[rank - 1].clone()); - - let a = a_storage.load( - program, - layout_index(&a_meta_body, &a_coords), - in_bounds.clone(), - ); - let b = b_storage.load( - program, - layout_index(&b_meta_body, &b_coords), - in_bounds.clone(), - ); - let (a, a_ty) = apply_unary_function_chain( - a.into_f32(), - a_meta_body.datatype, - &pre_a, - ) - .expect("validated matmul pre_element_wise[0] chain"); - let (b, b_ty) = apply_unary_function_chain( - b.into_f32(), - b_meta_body.datatype, - &pre_b, - ) - .expect("validated matmul pre_element_wise[1] chain"); - let a = ValueTile::F32(a).cast_to(a_ty).cast_to(acc_dtype); - let b = ValueTile::F32(b).cast_to(b_ty).cast_to(acc_dtype); - a.binary(tile_ir::TileBinaryOp::Mul, b) - }; - - let sum = match acc_dtype { - DataTypeEnum::F32 => { - let [acc] = program.fold( - tile_ir::tile::range(k), - [tile_ir::tile::Tile::literal(tile_ir::TileLiteral::f32(0.0))], - |program, k_index, [acc]| { - [acc + value_at(program, k_index).into_f32()] - }, - ); - ValueTile::F32(acc) - } - DataTypeEnum::F16 => unreachable!("matmul accumulates f16 products in f32"), - DataTypeEnum::U32 => { - let [acc] = program.fold( - tile_ir::tile::range(k), - [tile_ir::tile::Tile::literal(tile_ir::TileLiteral::U32(0))], - |program, k_index, [acc]| { - [acc + value_at(program, k_index).into_u32()] - }, - ); - ValueTile::U32(acc) - } - }; - let sum = sum.cast_to(result_dtype); - let (sum, sum_ty) = - apply_unary_function_chain(sum.into_f32(), result_dtype, &post) - .expect("validated matmul post_element_wise chain"); - let sum = ValueTile::F32(sum) - .cast_to(sum_ty) - .cast_to(y_meta_body.datatype); - y_storage.store(program, layout_index(&y_meta_body, &dims), sum, in_bounds); - }); - Some(()) - }, - ) -} - -fn storage_read( - kb: &mut tile_ir::KernelBuilder, - datatype: DataTypeEnum, - tensor: tile_ir::KernelTensorRef, -) -> crate::nary_direct::Storage2 { - let element = crate::nary_direct::datatype_element(datatype); - let storage = kb.read(element, tensor); - match datatype { - DataTypeEnum::F32 => crate::nary_direct::Storage2::F32(storage), - DataTypeEnum::F16 => crate::nary_direct::Storage2::F16(storage), - DataTypeEnum::U32 => crate::nary_direct::Storage2::U32(storage), - } -} - -fn storage_write( - kb: &mut tile_ir::KernelBuilder, - datatype: DataTypeEnum, - tensor: tile_ir::KernelTensorRef, -) -> crate::nary_direct::Storage2 { - let element = crate::nary_direct::datatype_element(datatype); - let storage = kb.write(element, tensor); - match datatype { - DataTypeEnum::F32 => crate::nary_direct::Storage2::F32(storage), - DataTypeEnum::F16 => crate::nary_direct::Storage2::F16(storage), - DataTypeEnum::U32 => crate::nary_direct::Storage2::U32(storage), - } -} - -fn output_dims_from_flat_u32(flat: tile_ir::tile::Tile, shape: &[u32]) -> Vec { - let shape = shape.iter().map(|dim| *dim as usize).collect::>(); - output_dims_from_flat(flat, &shape) -} diff --git a/fusor-ml/core/src/matmul/kernel.rs b/fusor-ml/core/src/matmul/kernel.rs index a40a63a43..e27bd3fd9 100644 --- a/fusor-ml/core/src/matmul/kernel.rs +++ b/fusor-ml/core/src/matmul/kernel.rs @@ -11,50 +11,24 @@ use crate::{ kernel_backend::{self, DirectKernel}, operation::Operation, tile_direct::{ - DirectMatrixLayout, flatten_matrix_layout, tile_storage_read_with_direct_layout_typed, + flatten_matrix_layout, tile_storage_read_with_direct_layout_typed, tile_storage_write_with_direct_layout_typed, }, }, - nary_direct::apply_unary_function_chain, - nary_wise::UnaryFunctionChain, + nary_wise::{NaryExpr, NaryFunction, NaryOp, NaryScalar, UnaryFunctionChain}, + reduce::{ReduceFunction, ReduceOp, ReduceOperation}, tensor::{DataTypeEnum, TensorData}, }; use super::{ - MatMulOperation, MatMulParams, coop_gemm, direct, sgemm, sgemv, - variants::{ - CoopTile, DirectTileMatmulVariant, dense_coop_kinds_from_datatype, - select_dense_matmul_params, select_direct_tile_matmul_variant, - }, + MatMulOperation, MatMulParams, coop_gemm, sgemm, sgemv, + variants::{CoopTile, dense_coop_kinds_from_datatype, select_dense_matmul_params}, }; -struct MatmulTileDirectKernelVariant; - fn device_supported(value: Option) -> Result { value.ok_or(kernel_backend::DeviceNotSupported) } -#[derive(Clone, Copy)] -struct DenseCoopTokens { - config: tile_ir_kernels::SubgroupConfig, - coop: tile_ir::CoopMatrixToken, -} - -#[derive(Clone, Copy)] -enum DirectTileMatmulRoute { - Gemv(tile_ir_kernels::SubgroupConfig), - MatMul { coop: Option }, -} - -impl DirectTileMatmulRoute { - fn variant(self) -> DirectTileMatmulVariant { - match self { - Self::Gemv(_) => DirectTileMatmulVariant::Gemv, - Self::MatMul { .. } => DirectTileMatmulVariant::MatMul, - } - } -} - impl MatMulOperation { pub fn new( datatype: DataTypeEnum, @@ -124,11 +98,104 @@ impl MatMulOperation { self.out_shape.len() as u32 } - fn can_use_direct_tile_matmul(&self) -> bool { + fn can_use_hardware_matmul(&self) -> bool { matches!(self.datatype, DataTypeEnum::F32 | DataTypeEnum::F16) } - fn build_direct_tile_matmul( + /// The contraction in its composed map-reduce form: a multiply over the + /// `[batch.., m, n, k]` index space summed along `k`, with the fused + /// pre/post chains inlined and accumulation upgraded to f32 (matching + /// the dedicated kernels' accumulator). Routes that aren't + /// hardware-specialized lower through this — the same generic tiled + /// reduce any composed contraction gets. + fn as_fused_reduce(&self) -> ReduceOperation { + let rank = self.first_shape.len(); + let batch = rank - 2; + let (m_dim, n_dim, k_dim) = (batch, batch + 1, batch + 2); + let mut index_space: Vec = self.first_shape[..batch].to_vec(); + index_space.extend([ + self.first_shape[batch], + self.second_shape[batch + 1], + self.first_shape[batch + 1], + ]); + + let apply_chain = |mut expr: NaryExpr, chain: &UnaryFunctionChain| { + for function in &chain.functions { + expr = NaryExpr::Op { + children: vec![expr], + function: function.clone(), + }; + } + expr + }; + let cast_to = |expr: NaryExpr, from: DataTypeEnum, to: DataTypeEnum| { + if from == to { + expr + } else { + NaryExpr::Op { + children: vec![expr], + function: NaryFunction::unary(Some("cast".to_string()), NaryOp::Cast, from, to), + } + } + }; + + let a_indices: Vec = (0..batch) + .chain([m_dim, k_dim]) + .map(NaryExpr::DimIndex) + .collect(); + let b_indices: Vec = (0..batch) + .chain([k_dim, n_dim]) + .map(NaryExpr::DimIndex) + .collect(); + + let acc_dtype = match self.pre_element_wise[0].out_datatype() { + DataTypeEnum::U32 => DataTypeEnum::U32, + DataTypeEnum::F32 | DataTypeEnum::F16 => DataTypeEnum::F32, + }; + let a = apply_chain( + NaryExpr::indexed_input(0, a_indices), + &self.pre_element_wise[0], + ); + let a = cast_to(a, self.pre_element_wise[0].out_datatype(), acc_dtype); + let b = apply_chain( + NaryExpr::indexed_input(1, b_indices), + &self.pre_element_wise[1], + ); + let b = cast_to(b, self.pre_element_wise[1].out_datatype(), acc_dtype); + let expression = NaryExpr::mul(a, b, acc_dtype); + + let initial_value = match acc_dtype { + DataTypeEnum::U32 => NaryScalar::U32(0), + _ => NaryScalar::F32(0.0), + }; + let result_dtype = self.post_element_wise.input_datatype(); + let mut post_functions = Vec::new(); + if acc_dtype != result_dtype { + post_functions.push(NaryFunction::unary( + Some("cast".to_string()), + NaryOp::Cast, + acc_dtype, + result_dtype, + )); + } + post_functions.extend(self.post_element_wise.functions.iter().cloned()); + + ReduceOperation { + inputs: vec![self.first, self.second], + expression, + shape: index_space.into(), + function: ReduceFunction { + name: Some("sum".to_string()), + op: ReduceOp::Sum, + initial_value, + datatype: acc_dtype, + }, + post_element_wise: UnaryFunctionChain::new(post_functions, acc_dtype), + axis: k_dim, + } + } + + fn build_hardware_matmul( &self, device: &Device, input_a: &TensorData, @@ -169,124 +236,75 @@ impl MatMulOperation { } let shape = tile_ir_kernels::DenseMatmulShape { batch, m, k, n }; - // GEMV reduces through subgroup operations. Use the matmul route - // unless the device exposes a subgroup path we trust. - let subgroup_config = device.subgroup_config(); - let route_variant = if subgroup_config.is_some() { - select_direct_tile_matmul_variant(m, k, n) - } else { - DirectTileMatmulVariant::MatMul - }; - let coop_kind = match &self.parameters { - MatMulParams::CoopMatMul(params) => Some(params.kind()), - _ => None, - }; - let route = match (route_variant, subgroup_config) { - (DirectTileMatmulVariant::Gemv, Some(config)) => DirectTileMatmulRoute::Gemv(config), - (DirectTileMatmulVariant::Gemv, None) => DirectTileMatmulRoute::MatMul { coop: None }, - (DirectTileMatmulVariant::MatMul, config) => { - let coop = config.and_then(|config| { - if !config.is_fixed() { - return None; - } - let kind = coop_kind?; - Some(DenseCoopTokens { - config, - coop: device.coop_token(kind)?, - }) - }); - DirectTileMatmulRoute::MatMul { coop } - } - }; - let variant = route.variant(); - let coop_variant = if let DirectTileMatmulRoute::MatMul { coop: Some(coop) } = route { - CoopTile::select( - m, - k, - n, - device - .limits() - .max_compute_workgroup_size_x - .min(device.limits().max_compute_invocations_per_workgroup), - coop.config.max_size(), - ) - } else { - None + // Only the cooperative-matrix route stays hand-specialized; gemv + // shapes lower through the generic subgroup-per-output reduce, and + // fused chains lower through the generic tiled reduce. + if !self.pre_element_wise[0].functions.is_empty() + || !self.pre_element_wise[1].functions.is_empty() + || !self.post_element_wise.functions.is_empty() + { + return Err(kernel_backend::DeviceNotSupported); + } + let subgroup_config = device_supported(device.subgroup_config())?; + if !subgroup_config.is_fixed() { + return Err(kernel_backend::DeviceNotSupported); + } + let MatMulParams::CoopMatMul(params) = &self.parameters else { + return Err(kernel_backend::DeviceNotSupported); }; - let pre_a = self.pre_element_wise[0] - .functions - .is_empty() - .then_some(()) - .is_none() - .then(|| { - let chain = self.pre_element_wise[0].clone(); - let datatype = chain.input_datatype(); - tile_ir_kernels::UnaryEpilogue::new("matmul_pre_a_chain", move |tile| { - apply_unary_function_chain(tile, datatype, &chain) - .expect("pre-chain validated at fuse time") - .0 - }) - }); - let pre_b = self.pre_element_wise[1] - .functions - .is_empty() - .then_some(()) - .is_none() - .then(|| { - let chain = self.pre_element_wise[1].clone(); - let datatype = chain.input_datatype(); - tile_ir_kernels::UnaryEpilogue::new("matmul_pre_b_chain", move |tile| { - apply_unary_function_chain(tile, datatype, &chain) - .expect("pre-chain validated at fuse time") - .0 - }) - }); - let post = self - .post_element_wise - .functions - .is_empty() - .then_some(()) - .is_none() - .then(|| { - let chain = self.post_element_wise.clone(); - let datatype = chain.input_datatype(); - tile_ir_kernels::UnaryEpilogue::new("matmul_post_chain", move |tile| { - apply_unary_function_chain(tile, datatype, &chain) - .expect("post-chain validated at fuse time") - .0 - }) - }); - let epilogue_identity = pre_a.as_ref().map(|e| e.identity()).unwrap_or(0) - ^ pre_b.as_ref().map(|e| e.identity()).unwrap_or(0) - ^ post.as_ref().map(|e| e.identity()).unwrap_or(0); - let matmul_tile = select_matmul_tile(shape, device.limits().max_compute_workgroup_size_x); + let kind = params.kind(); + let coop = device_supported(device.coop_token(kind))?; + let tile = device_supported(CoopTile::select( + m, + k, + n, + device + .limits() + .max_compute_workgroup_size_x + .min(device.limits().max_compute_invocations_per_workgroup), + subgroup_config.max_size(), + ))?; + let max_wg_per_dim = device.limits().max_compute_workgroups_per_dimension; let datatype = self.datatype; - let ir = tile_ir::tile::build(move |phase| { - let epilogues = tile_ir_kernels::DenseMatmulEpilogues { - pre_a: pre_a.as_ref(), - pre_b: pre_b.as_ref(), - post: post.as_ref(), - }; + let used = std::cell::Cell::new(false); + let ir = tile_ir::tile::build(|phase| { let element = match datatype { DataTypeEnum::F32 => tile_ir::ElementType::F32, DataTypeEnum::F16 => tile_ir::ElementType::F16, - _ => unreachable!("direct tile matmul only supports f32/f16"), + _ => unreachable!("hardware matmul only supports f32/f16"), }; - dispatch_direct_tile_matmul( + let a = tile_storage_read_with_direct_layout_typed(phase, element, a_view.clone()); + let b = tile_storage_read_with_direct_layout_typed(phase, element, b_view.clone()); + let y = tile_storage_write_with_direct_layout_typed(phase, element, y_view.clone()); + used.set(tile_ir_kernels::try_batched_coop_matmul( phase, - element, - a_view.clone(), - b_view.clone(), - y_view.clone(), - coop_variant, - route, - matmul_tile, + tile_ir_kernels::DenseMatmulTensors { + a: &a, + b: &b, + y: &y, + }, shape, - &epilogues, + &tile_ir_kernels::DenseMatmulEpilogues { + pre_a: None, + pre_b: None, + post: None, + }, max_wg_per_dim, - ); + tile_ir_kernels::DenseCoopMatmulConfig { + coop, + subgroups: subgroup_config, + tile: tile_ir_kernels::DenseCoopMatmulTile { + bm: tile.bm, + bn: tile.bn, + bk: tile.bk, + }, + }, + )); }); + if !used.get() { + return Err(kernel_backend::DeviceNotSupported); + } let dispatch_size = ir.grid; if dispatch_size.iter().any(|dim| *dim > max_wg_per_dim) { return Err(kernel_backend::DeviceNotSupported); @@ -296,16 +314,11 @@ impl MatMulOperation { input_b.clone().into(), output.clone().into(), ]; - let variant = kernel_backend::KernelVariantKey::with_payload::( - |state| { - matmul_tile.hash(state); - variant.hash(state); - coop_variant.hash(state); - coop_kind.hash(state); + let variant = + kernel_backend::KernelVariantKey::with_payload::(|state| { + tile.hash(state); subgroup_config.hash(state); - epilogue_identity.hash(state); - }, - ); + }); let cache_key = self.kernel_cache_key_with_dispatch(variant, None, dispatch_size, &inputs); let name = self.name(); @@ -329,6 +342,8 @@ impl MatMulOperation { } } +struct HardwareMatmulVariant; + impl Operation for MatMulOperation { fn hash_kernel_fields(&self, state: &mut FxHasher) { self.datatype.hash(state); @@ -440,20 +455,29 @@ impl Operation for MatMulOperation { let input_a = input_a.as_tensor()?; let input_b = input_b.as_tensor()?; let output = output.as_tensor()?; - let direct_tile_kernel = if self.can_use_direct_tile_matmul() + if self.can_use_hardware_matmul() && input_a.datatype() == self.datatype && input_b.datatype() == self.datatype && output.datatype() == self.datatype && (self.datatype != DataTypeEnum::F16 || graph.device().f16_supported()) + && let Ok(kernel) = + self.build_hardware_matmul(&graph.device(), input_a, input_b, output) { - self.build_direct_tile_matmul(&graph.device(), input_a, input_b, output) - .ok() - } else { - None - }; - direct_tile_kernel.or_else(|| { - direct::build_serial_matmul_direct_kernel(self, graph, workgroup_shape, inputs) - }) + return Some(kernel); + } + // Everything else is the composed contraction's own lowering: the + // generic tiled (or serial) fused reduce, identical to what any + // unrecognized contraction gets. + let reduce = self.as_fused_reduce(); + crate::reduce_tiled::build_reduce_tiled_kernel(&reduce, graph, workgroup_shape, inputs) + .or_else(|| { + crate::reduce_direct::build_reduce_direct_kernel( + &reduce, + graph, + workgroup_shape, + inputs, + ) + }) } fn output( @@ -482,182 +506,3 @@ impl Operation for MatMulOperation { ) } } - -fn select_matmul_tile( - shape: tile_ir_kernels::DenseMatmulShape, - max_workgroup_size_x: u32, -) -> Option { - if shape.m < 32 || shape.n < 32 || shape.k < 8 { - return None; - } - - let candidates = [tile_ir_kernels::DenseMatmulTile::new(32, 32, 8, 4, 4, 64)]; - - candidates - .into_iter() - .filter_map(|tile| { - if tile.lanes > max_workgroup_size_x || shape.m < tile.bm || shape.n < tile.bn { - return None; - } - - let tiles_m = shape.m.div_ceil(tile.bm); - let tiles_n = shape.n.div_ceil(tile.bn); - let workgroups = shape.batch * tiles_m * tiles_n; - let actual_outputs = shape.batch * shape.m * shape.n; - let covered_outputs = workgroups * tile.bm * tile.bn; - let tile_utilization_ok = actual_outputs * 4 >= covered_outputs * 3; - if !tile_utilization_ok { - return None; - } - - // Larger workgroups reduce input rereads, but avoid collapsing small - // unbatched matmuls down to only a handful of 256-thread groups. - if (tile.lanes >= 256 && workgroups < 32) - || (tile.lanes >= 128 && shape.batch == 1 && workgroups < 32) - || (tile.lanes >= 128 && shape.batch > 1 && workgroups < 8) - { - return None; - } - - let staged_inputs_per_k_tile = workgroups * tile.bk * (tile.bm + tile.bn); - let score = (staged_inputs_per_k_tile, tile.lanes); - Some((score, tile)) - }) - .min_by_key(|(score, _)| *score) - .map(|(_, tile)| tile) -} - -#[allow(clippy::too_many_arguments)] -fn dispatch_direct_tile_matmul( - phase: &mut tile_ir::tile::Program, - element: tile_ir::ElementType, - a_view: DirectMatrixLayout, - b_view: DirectMatrixLayout, - y_view: DirectMatrixLayout, - coop_variant: Option, - route: DirectTileMatmulRoute, - matmul_tile: Option, - shape: tile_ir_kernels::DenseMatmulShape, - epilogues: &tile_ir_kernels::DenseMatmulEpilogues<'_>, - max_wg_per_dim: u32, -) { - let a = tile_storage_read_with_direct_layout_typed(phase, element, a_view); - let b = tile_storage_read_with_direct_layout_typed(phase, element, b_view); - let y = tile_storage_write_with_direct_layout_typed(phase, element, y_view); - if let (Some(tile), DirectTileMatmulRoute::MatMul { coop: Some(tokens) }) = - (coop_variant, route) - && tile_ir_kernels::try_batched_coop_matmul( - phase, - tile_ir_kernels::DenseMatmulTensors { - a: &a, - b: &b, - y: &y, - }, - shape, - epilogues, - max_wg_per_dim, - tile_ir_kernels::DenseCoopMatmulConfig { - coop: tokens.coop, - subgroups: tokens.config, - tile: tile_ir_kernels::DenseCoopMatmulTile { - bm: tile.bm, - bn: tile.bn, - bk: tile.bk, - }, - }, - ) - { - return; - } - match route { - DirectTileMatmulRoute::Gemv(tokens) => tile_ir_kernels::batched_gemv_with_epilogues( - phase, - &a, - &b, - &y, - shape, - epilogues, - max_wg_per_dim, - tokens, - ), - DirectTileMatmulRoute::MatMul { .. } => { - if let Some(tile) = matmul_tile { - tile_ir_kernels::batched_matmul_with_epilogues( - phase, - tile_ir_kernels::DenseMatmulTensors { - a: &a, - b: &b, - y: &y, - }, - shape, - epilogues, - max_wg_per_dim, - tile, - ) - } else { - tile_ir_kernels::batched_matmul_register_with_epilogues( - phase, - &a, - &b, - &y, - shape, - epilogues, - max_wg_per_dim, - ) - } - } - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn batched_full_tiles_use_shared_tile() { - assert_eq!( - select_matmul_tile( - tile_ir_kernels::DenseMatmulShape { - batch: 8, - m: 64, - k: 96, - n: 64, - }, - 256, - ), - Some(tile_ir_kernels::DenseMatmulTile::new(32, 32, 8, 4, 4, 64)), - ); - } - - #[test] - fn partial_large_tiles_need_reasonable_utilization() { - assert_eq!( - select_matmul_tile( - tile_ir_kernels::DenseMatmulShape { - batch: 2, - m: 65, - k: 96, - n: 65, - }, - 256, - ), - None, - ); - } - - #[test] - fn small_unbatched_matmul_keeps_parallelism() { - assert_eq!( - select_matmul_tile( - tile_ir_kernels::DenseMatmulShape { - batch: 1, - m: 64, - k: 96, - n: 64, - }, - 256, - ), - Some(tile_ir_kernels::DenseMatmulTile::new(32, 32, 8, 4, 4, 64)), - ); - } -} diff --git a/fusor-ml/core/src/matmul/mod.rs b/fusor-ml/core/src/matmul/mod.rs index 627b41e52..5eec4ec4d 100644 --- a/fusor-ml/core/src/matmul/mod.rs +++ b/fusor-ml/core/src/matmul/mod.rs @@ -4,7 +4,6 @@ use crate::{ }; pub mod coop_gemm; -mod direct; mod kernel; pub mod sgemm; mod sgemm_params; @@ -40,22 +39,98 @@ pub(crate) struct MatMulOperation { } impl Tensor { + /// Matrix multiply, expressed as its composed form: a broadcast multiply + /// over the `[batch.., M, N, K]` index space summed along `K`. The + /// resolver recognizes the canonical cluster and routes it to the + /// specialized matmul kernels; anything that composes differently lowers + /// through the generic elementwise + reduce path. pub fn mat_mul(&self, other: &Self) -> Self { + use crate::nary_wise::{ElementwiseOperation, NaryExpr}; + assert_eq!(self.datatype(), other.datatype()); - self.add_mat_mul(other, None) + let a_shape = self.shape(); + let b_shape = other.shape(); + let rank = a_shape.len(); + assert_eq!( + rank, + b_shape.len(), + "mat_mul requires equal ranks: {a_shape:?} x {b_shape:?}" + ); + assert!(rank >= 2, "mat_mul requires rank >= 2: {a_shape:?}"); + let batch = rank - 2; + assert_eq!( + a_shape[..batch], + b_shape[..batch], + "mat_mul batch dimensions must match: {a_shape:?} x {b_shape:?}" + ); + assert_eq!( + a_shape[rank - 1], + b_shape[rank - 2], + "mat_mul contraction dimensions must match: {a_shape:?} x {b_shape:?}" + ); + + let (m, k, n) = (a_shape[batch], a_shape[batch + 1], b_shape[batch + 1]); + let mut index_space: Vec = a_shape[..batch].to_vec(); + index_space.extend([m, n, k]); + let (m_dim, n_dim, k_dim) = (batch, batch + 1, batch + 2); + + let a_indices: Vec = (0..batch) + .chain([m_dim, k_dim]) + .map(NaryExpr::DimIndex) + .collect(); + let b_indices: Vec = (0..batch) + .chain([k_dim, n_dim]) + .map(NaryExpr::DimIndex) + .collect(); + + let datatype = self.datatype(); + let product = Tensor::from_parts(self.data.nary(ElementwiseOperation { + inputs: vec![self.key(), other.key()], + expression: NaryExpr::mul( + NaryExpr::indexed_input(0, a_indices), + NaryExpr::indexed_input(1, b_indices), + datatype, + ), + shape: index_space.into(), + output_datatype: datatype, + })); + product.sum(k_dim) } + /// Matrix multiply with explicit kernel parameters: a tuning/benchmark + /// API. The parameters cannot round-trip through the composed graph, so + /// the operation builds directly against materialized inputs and + /// executes eagerly, returning a fresh leaf tensor. pub fn mat_mul_with_parameters(&self, other: &Self, parameters: MatMulParams) -> Self { assert_eq!(self.datatype(), other.datatype()); - self.add_mat_mul(other, Some(parameters)) + self.data.materialize(); + other.data.materialize(); + let operation = MatMulOperation::new( + self.datatype(), + self.key(), + other.key(), + self.shape(), + other.shape(), + Some(parameters), + self.device(), + ); + let output = self + .device() + .compute_graph() + .execute_eager(&operation) + .unwrap_or_else(|| { + panic!( + "mat_mul_with_parameters could not build a kernel for the requested parameters" + ) + }); + Tensor::from(output) } } #[cfg(test)] mod selection_tests { use super::variants::{ - CoopTile, DenseMatmulCtx, DenseMatmulVariant, DirectTileMatmulVariant, - dense_matmul_selector, direct_tile_matmul_selector, select_coop_kind, + CoopTile, DenseMatmulCtx, DenseMatmulVariant, dense_matmul_selector, select_coop_kind, }; use crate::kernel_selection::{ CooperativeMatrixCaps, CooperativeMatrixKind, DeterministicShapeRng, KernelDeviceCaps, @@ -204,32 +279,6 @@ mod selection_tests { ); } - #[test] - fn direct_tile_selector_generates_each_variant() { - let selector = direct_tile_matmul_selector(); - let caps = KernelDeviceCaps { - subgroups_supported: false, - cooperative_matrix: CooperativeMatrixCaps::default(), - min_subgroup_size: 0, - max_subgroup_size: 0, - max_compute_invocations_per_workgroup: 0, - max_compute_workgroup_storage_size: 0, - max_compute_workgroup_size_x: 0, - backend: wgpu::Backend::Noop, - }; - let mut rng = DeterministicShapeRng::default(); - - for variant in [ - DirectTileMatmulVariant::Gemv, - DirectTileMatmulVariant::MatMul, - ] { - let shape = selector - .generate_for(variant, &(), caps, &mut rng) - .expect("variant should generate"); - assert_eq!(selector.select(shape, &(), caps), Some(variant)); - } - } - #[test] fn direct_tile_coop_selector_prefers_largest_supported_tile() { let select = diff --git a/fusor-ml/core/src/matmul/variants.rs b/fusor-ml/core/src/matmul/variants.rs index b1ed4b49e..15516cecf 100644 --- a/fusor-ml/core/src/matmul/variants.rs +++ b/fusor-ml/core/src/matmul/variants.rs @@ -1,8 +1,7 @@ use crate::{ Device, kernel_selection::{ - Axis, CooperativeMatrixCaps, CooperativeMatrixKind, KernelDeviceCaps, KernelShape, - ShapeRule, ShapeSelector, eq, range, + Axis, CooperativeMatrixKind, KernelDeviceCaps, KernelShape, ShapeRule, ShapeSelector, range, }, matmul::sgemm_params::gemm_parameters, matmul::sgemv_params::gemv_parameters, @@ -135,12 +134,6 @@ pub(super) fn select_coop_kind( .expect("coop selector called with no supported cooperative matrix kind") } -#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] -pub(super) enum DirectTileMatmulVariant { - Gemv, - MatMul, -} - /// (BM, BN, BK) tile dimensions for a cooperative-matrix matmul tile. The /// `select` helper below returns `Option` (`None` = no coop variant /// fits the shape); the kernel layer uses the tuple to look up the matching @@ -239,31 +232,3 @@ impl CoopTile { None } } - -pub(super) fn direct_tile_matmul_selector() -> ShapeSelector<3, (), DirectTileMatmulVariant> { - ShapeSelector::new() - .rule( - DirectTileMatmulVariant::Gemv, - ShapeRule::new().axis(DENSE_N, eq(1)), - ) - .rule(DirectTileMatmulVariant::MatMul, ShapeRule::new()) -} - -pub(super) fn select_direct_tile_matmul_variant(m: u32, k: u32, n: u32) -> DirectTileMatmulVariant { - direct_tile_matmul_selector() - .select( - KernelShape::new([m as usize, k as usize, n as usize]), - &(), - KernelDeviceCaps { - subgroups_supported: false, - cooperative_matrix: CooperativeMatrixCaps::default(), - min_subgroup_size: 0, - max_subgroup_size: 0, - max_compute_invocations_per_workgroup: 0, - max_compute_workgroup_storage_size: 0, - max_compute_workgroup_size_x: 0, - backend: wgpu::Backend::Noop, - }, - ) - .expect("direct tile matmul selector has a catch-all rule") -} diff --git a/fusor-ml/core/src/mir/kernel_backend.rs b/fusor-ml/core/src/mir/kernel_backend.rs index 316fadceb..447170764 100644 --- a/fusor-ml/core/src/mir/kernel_backend.rs +++ b/fusor-ml/core/src/mir/kernel_backend.rs @@ -16,8 +16,6 @@ impl std::fmt::Display for DeviceNotSupported { impl std::error::Error for DeviceNotSupported {} -pub(crate) mod flash_attention; pub(crate) mod mirostat; -pub(crate) mod rms_norm; pub(crate) mod sampling_topk; pub(crate) mod standard_sampler; diff --git a/fusor-ml/core/src/mir/kernel_backend/flash_attention/kernel.rs b/fusor-ml/core/src/mir/kernel_backend/flash_attention/kernel.rs deleted file mode 100644 index 8d50e89d7..000000000 --- a/fusor-ml/core/src/mir/kernel_backend/flash_attention/kernel.rs +++ /dev/null @@ -1,586 +0,0 @@ -use std::{any::TypeId, hash::Hash, sync::Arc}; - -use fusor_tile_ir as tile_ir; -use fusor_tile_ir_kernels as tile_ir_kernels; - -use crate::{ - DataTypeEnum, - compute_graph::{ComputeGraphInner, NodeIndex}, - kernel_selection::KernelDeviceCaps, - mir::{ - inputs::MirValue, - kernel_backend, - kernel_backend::DirectKernel, - operation::Operation, - workgroup_shape::{Constraint, WorkgroupShape, WorkgroupShapeConstraints}, - }, - tensor::TensorData, -}; - -use super::{ - DECODE_SMALL_BLOCK, FLASH_STREAMING_SUBGROUP_SIZES, FLASH_STREAMING_TILED_Q_BLOCK, - FlashAttentionDirectKernelVariant, FlashAttentionKernelVariant, FlashAttentionOperation, - FlashDecodeSmallMeta, FlashDecodeSmallTensors, TensorMeta, build_flash_decode_small_meta, - dispatch_streaming_flash_attention, dispatch_streaming_tiled_flash_attention, - flash_streaming_tiled_eligible, select_flash_attention_variant, streaming_dispatch_size, -}; - -fn flash_decode_cache_variant( - variant: FlashAttentionKernelVariant, - scale: f32, - meta: &FlashDecodeSmallMeta, -) -> kernel_backend::KernelVariantKey { - kernel_backend::KernelVariantKey::with_payload::(|state| { - variant.hash(state); - scale.to_bits().hash(state); - meta.decode_block.hash(state); - meta.tiled.hash(state); - meta.split_blocks.hash(state); - }) -} - -fn hash_flash_decode_dims( - state: &mut rustc_hash::FxHasher, - dims: &tile_ir_kernels::FlashAttentionDims, -) { - dims.batch.hash(state); - dims.num_heads.hash(state); - dims.num_kv_heads.hash(state); - dims.q_seq_len.hash(state); - dims.kv_seq_len.hash(state); - dims.head_dim.hash(state); -} - -pub(super) fn flash_decode_cache_key( - workgroup_shape: Option<&WorkgroupShape>, - dispatch_size: [u32; 3], - input_dtype: DataTypeEnum, - scale: f32, - meta: &FlashDecodeSmallMeta, -) -> kernel_backend::KernelCacheKey { - flash_decode_cache_key_for_variant( - FlashAttentionKernelVariant::DecodeSmall, - workgroup_shape, - dispatch_size, - input_dtype, - scale, - meta, - ) -} - -fn flash_decode_cache_key_for_variant( - decode_variant: FlashAttentionKernelVariant, - workgroup_shape: Option<&WorkgroupShape>, - dispatch_size: [u32; 3], - input_dtype: DataTypeEnum, - scale: f32, - meta: &FlashDecodeSmallMeta, -) -> kernel_backend::KernelCacheKey { - let variant = flash_decode_cache_variant(decode_variant, scale, meta); - kernel_backend::KernelCacheKey::from_hash_inputs(|state| { - // Decode kernels take the active KV length from a params buffer. Do - // not hash `active_kv_len`, or every generated token would miss the - // kernel cache even though the IR is otherwise bucketed by block size. - 2u64.hash(state); - variant.hash(state); - TypeId::of::().hash(state); - workgroup_shape - .map(|workgroup_shape| workgroup_shape.shape()) - .hash(state); - dispatch_size.hash(state); - input_dtype.hash(state); - hash_flash_decode_dims(state, &meta.dims); - meta.decode_block.hash(state); - meta.tiled.hash(state); - meta.split_blocks.hash(state); - meta.groups.hash(state); - meta.q_offset.hash(state); - meta.k_offset.hash(state); - meta.v_offset.hash(state); - meta.output_offset.hash(state); - meta.q_strides.hash(state); - meta.k_strides.hash(state); - meta.v_strides.hash(state); - meta.output_strides.hash(state); - }) -} - -impl Operation for FlashAttentionOperation { - fn hash_kernel_fields(&self, state: &mut rustc_hash::FxHasher) { - self.out_shape.hash(state); - self.q_shape.hash(state); - self.k_shape.hash(state); - self.mask.is_some().hash(state); - self.scale.to_bits().hash(state); - self.input_dtype.hash(state); - self.causal.hash(state); - } - - fn workgroup_shape_constraints(&self, _device: &crate::Device) -> WorkgroupShapeConstraints { - let mut constraints = WorkgroupShapeConstraints::new(); - constraints.add_constraint(0, Constraint::Equals(1)); - constraints.add_constraint(1, Constraint::Equals(1)); - constraints.add_constraint(2, Constraint::Equals(1)); - constraints - } - - fn dispatch_size(&self, _workgroup_shape: &WorkgroupShape, _inputs: &[MirValue]) -> [u32; 3] { - // The streaming kernel's per-axis dispatch depends on the device's - // hardware subgroup size, which isn't known at this layer; the - // workgroup-shape pass only uses the row axis (Y/Z) to size local - // work. Conservatively report the smallest streaming variant's - // per-axis dispatch; `build_direct_kernel` recomputes the real - // dispatch with the correct subgroup width before launch. - let dims = self.dims().expect("flash attention dimensions fit in u32"); - let outputs_per_workgroup = tile_ir_kernels::flash_outputs_per_workgroup( - *FLASH_STREAMING_SUBGROUP_SIZES.last().expect("non-empty"), - ); - streaming_dispatch_size(dims, outputs_per_workgroup) - } - - fn visit_dependencies(&self, f: &mut dyn FnMut(NodeIndex)) { - f(self.q); - f(self.k); - f(self.v); - if let Some(mask) = self.mask { - f(mask); - } - } - - fn inputs(&self, nodes: &ComputeGraphInner) -> Vec { - let q = nodes.get_result(self.q).unwrap(); - let k = nodes.get_result(self.k).unwrap(); - let v = nodes.get_result(self.v).unwrap(); - let output = TensorData::new_for_shape(q.device(), &self.out_shape, self.input_dtype); - - let mut inputs = vec![q.into(), k.into(), v.into()]; - if let Some(mask) = self.mask { - inputs.push(nodes.get_result(mask).unwrap().into()); - } - inputs.push(output.into()); - inputs - } - - fn build_direct_kernel( - &self, - graph: &ComputeGraphInner, - workgroup_shape: &WorkgroupShape, - inputs: &[MirValue], - ) -> Option { - let q = inputs.first()?.as_tensor()?.clone(); - let k = inputs.get(1)?.as_tensor()?.clone(); - let v = inputs.get(2)?.as_tensor()?.clone(); - let (mask, output_index) = if self.mask.is_some() { - (Some(inputs.get(3)?.as_tensor()?.clone()), 4) - } else { - (None, 3) - }; - let output = inputs.get(output_index)?.as_tensor()?.clone(); - let device = graph.device(); - - let input_dtype = self.input_dtype; - if !matches!(input_dtype, DataTypeEnum::F32 | DataTypeEnum::F16) { - return None; - } - if q.datatype() != input_dtype - || k.datatype() != input_dtype - || v.datatype() != input_dtype - || output.datatype() != input_dtype - || mask - .as_ref() - .is_some_and(|mask| mask.datatype() != input_dtype) - { - return None; - } - - let dims = self.dims()?; - if dims.batch == 0 - || dims.num_heads == 0 - || dims.num_kv_heads == 0 - || dims.q_seq_len == 0 - || dims.kv_seq_len == 0 - || dims.head_dim == 0 - { - return None; - } - - let q_meta = TensorMeta::new(&q)?; - let k_meta = TensorMeta::new(&k)?; - let v_meta = TensorMeta::new(&v)?; - let mask_meta = if let Some(mask) = mask.as_ref() { - Some(TensorMeta::new(mask)?) - } else { - None - }; - let output_meta = TensorMeta::new(&output)?; - - if q_meta.datatype != input_dtype - || k_meta.datatype != input_dtype - || v_meta.datatype != input_dtype - || output_meta.datatype != input_dtype - || mask_meta - .as_ref() - .is_some_and(|mask| mask.datatype != input_dtype) - { - return None; - } - - let caps = KernelDeviceCaps::from_device(&device); - let selected_variant = select_flash_attention_variant(dims, mask_meta.is_some(), caps); - // Decode-small kernel only supports f32; force streaming for other dtypes. - let decode_eligible = - input_dtype == DataTypeEnum::F32 && selected_variant.decode_block().is_some(); - let decode_candidate = mask_meta.is_none() - && dims.q_seq_len == 1 - && dims.head_dim > 0 - && input_dtype == DataTypeEnum::F32; - assert!( - !decode_candidate || selected_variant.decode_block().is_some(), - "decode attention refused slow fallback: device must support at least {DECODE_SMALL_BLOCK} workgroup invocations on x" - ); - let decode_meta = if decode_eligible { - let meta = build_flash_decode_small_meta( - dims, - self.scale, - caps, - FlashDecodeSmallTensors { - q: q_meta.clone(), - k: k_meta.clone(), - v: v_meta.clone(), - mask: mask_meta.as_ref(), - output: output_meta.clone(), - }, - )?; - assert_eq!( - Some(meta.decode_block), - selected_variant.decode_block(), - "flash attention selector and decode meta disagree" - ); - Some(meta) - } else { - None - }; - let variant = if decode_eligible { - selected_variant.kernel_variant() - } else { - // Decode-small uses workgroup reductions only. The streaming - // kernels partition lanes by the hardware subgroup width, so - // construct those variants only with a trusted fixed subgroup. - let subgroup = device.subgroup_token()?; - let subgroup_size = device.fixed_width_subgroup_size()?; - if !FLASH_STREAMING_SUBGROUP_SIZES.contains(&subgroup_size) { - return None; - } - if flash_streaming_tiled_eligible(dims) { - FlashAttentionKernelVariant::StreamingTiled { - subgroup, - subgroup_size, - } - } else { - FlashAttentionKernelVariant::Streaming { - subgroup, - subgroup_size, - } - } - }; - let dispatch_size = match variant { - FlashAttentionKernelVariant::Streaming { subgroup_size, .. } => { - streaming_dispatch_size( - dims, - tile_ir_kernels::flash_outputs_per_workgroup(subgroup_size), - ) - } - FlashAttentionKernelVariant::StreamingTiled { subgroup_size, .. } => { - tile_ir_kernels::flash_tiled_dispatch_size( - dims, - tile_ir_kernels::flash_tiled_outputs_per_workgroup(subgroup_size), - FLASH_STREAMING_TILED_Q_BLOCK, - ) - } - FlashAttentionKernelVariant::DecodeSmall => [ - dims.batch - .checked_mul(dims.num_heads) - .expect("flash decode dispatch overflow"), - 1, - 1, - ], - FlashAttentionKernelVariant::DecodeSplitPartials - | FlashAttentionKernelVariant::DecodeSplitReduce => { - unreachable!("split decode kernels are built as a sequence") - } - }; - if dispatch_size - .iter() - .any(|dim| *dim > device.limits().max_compute_workgroups_per_dimension) - { - return None; - } - - if let Some(meta) = decode_meta.filter(|meta| meta.tiled) { - let rows = meta - .dims - .batch - .checked_mul(meta.dims.num_heads) - .expect("flash split decode row overflow"); - let scratch_elements = - rows as u64 * meta.split_blocks as u64 * (meta.dims.head_dim as u64 + 2); - let scratch_buffer = device.create_buffer( - scratch_elements * std::mem::size_of::() as u64, - wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC, - ); - let params = [meta.active_kv_len, 0, 0, 0]; - let params_buffer = device.create_buffer_init( - bytemuck::cast_slice(¶ms), - wgpu::BufferUsages::STORAGE - | wgpu::BufferUsages::COPY_DST - | wgpu::BufferUsages::COPY_SRC, - ); - let partial_dispatch = [ - rows.checked_mul(meta.split_blocks) - .expect("flash split decode partial dispatch overflow"), - 1, - 1, - ]; - let reduce_dispatch = [rows, 1, 1]; - let layout = tile_ir_kernels::linear_storage_layout(); - - let partial_key = flash_decode_cache_key_for_variant( - FlashAttentionKernelVariant::DecodeSplitPartials, - Some(workgroup_shape), - partial_dispatch, - input_dtype, - self.scale, - &meta, - ); - let partial_buffers = vec![ - q.buffer().clone(), - k.buffer().clone(), - v.buffer().clone(), - scratch_buffer.clone(), - params_buffer, - ]; - let partial_layout = layout.clone(); - let partial_kernel = kernel_backend::dynamic_kernel_from_ir( - device.kernel_cache(), - "flash_attention_decode_split_partials", - partial_key, - move || { - let mut kb = tile_ir::KernelBuilder::<()>::new(); - let q_ref = tile_ir::KernelTensorRef::new((), partial_layout.clone()); - let k_ref = tile_ir::KernelTensorRef::new((), partial_layout.clone()); - let v_ref = tile_ir::KernelTensorRef::new((), partial_layout.clone()); - let scratch_ref = tile_ir::KernelTensorRef::new((), partial_layout.clone()); - let params_ref = tile_ir::KernelTensorRef::new((), partial_layout); - tile_ir_kernels::flash_decode_split_partials( - &mut kb, - q_ref, - k_ref, - v_ref, - scratch_ref, - params_ref, - meta, - )?; - Some(kb.finish().0) - }, - partial_buffers, - partial_dispatch, - )?; - - let reduce_key = flash_decode_cache_key_for_variant( - FlashAttentionKernelVariant::DecodeSplitReduce, - Some(workgroup_shape), - reduce_dispatch, - input_dtype, - self.scale, - &meta, - ); - let reduce_buffers = vec![scratch_buffer, output.buffer().clone()]; - let reduce_layout = layout.clone(); - let reduce_kernel = kernel_backend::dynamic_kernel_from_ir( - device.kernel_cache(), - "flash_attention_decode_split_reduce", - reduce_key, - move || { - let mut kb = tile_ir::KernelBuilder::<()>::new(); - let scratch_ref = tile_ir::KernelTensorRef::new((), reduce_layout.clone()); - let output_ref = tile_ir::KernelTensorRef::new((), reduce_layout); - tile_ir_kernels::flash_decode_split_reduce( - &mut kb, - scratch_ref, - output_ref, - meta, - )?; - Some(kb.finish().0) - }, - reduce_buffers, - reduce_dispatch, - )?; - - return Some(kernel_backend::DirectKernel::sequence( - "flash_attention_decode_split", - vec![partial_kernel, reduce_kernel], - )); - } - - let kernel_label = match variant { - FlashAttentionKernelVariant::Streaming { .. } => "flash_attention", - FlashAttentionKernelVariant::StreamingTiled { .. } => "flash_attention_tiled", - FlashAttentionKernelVariant::DecodeSmall => "flash_attention_decode", - FlashAttentionKernelVariant::DecodeSplitPartials => { - "flash_attention_decode_split_partials" - } - FlashAttentionKernelVariant::DecodeSplitReduce => "flash_attention_decode_split_reduce", - }; - let cache_key = if let Some(meta) = decode_meta.as_ref() { - flash_decode_cache_key( - Some(workgroup_shape), - dispatch_size, - input_dtype, - self.scale, - meta, - ) - } else { - let cache_variant = kernel_backend::KernelVariantKey::with_payload::< - FlashAttentionDirectKernelVariant, - >(|state| { - variant.hash(state); - self.scale.to_bits().hash(state); - self.causal.hash(state); - }); - self.kernel_cache_key_with_dispatch( - cache_variant, - Some(workgroup_shape), - dispatch_size, - inputs, - ) - }; - - let _ = output_index; // Bindings are derived from the kernel IR. - let layout = tile_ir_kernels::linear_storage_layout(); - let q_buffer = q.buffer().clone(); - let k_buffer = k.buffer().clone(); - let v_buffer = v.buffer().clone(); - let mask_buffer = mask.as_ref().map(|m| m.buffer().clone()); - let output_buffer = output.buffer().clone(); - let scale = self.scale; - let causal = self.causal; - let q_tile_meta = q_meta.tile.clone(); - let k_tile_meta = k_meta.tile.clone(); - let v_tile_meta = v_meta.tile.clone(); - let mask_tile_meta = mask_meta.clone().map(|meta| meta.tile); - let output_tile_meta = output_meta.tile.clone(); - - // Hoist params-buffer upload (decode path only) and the buffer-list - // collection OUTSIDE the IR-build closure so cache hits skip the - // entire IR construction. The closure runs only on cache miss. - let mut buffers: Vec> = Vec::with_capacity(6); - buffers.push(q_buffer.clone()); - buffers.push(k_buffer.clone()); - buffers.push(v_buffer.clone()); - if let Some(mask_buf) = mask_buffer.as_ref() { - buffers.push(mask_buf.clone()); - } - buffers.push(output_buffer.clone()); - if let Some(meta) = decode_meta { - let params = [meta.active_kv_len, 0, 0, 0]; - let params_buffer = device.create_buffer_init( - bytemuck::cast_slice(¶ms), - wgpu::BufferUsages::STORAGE - | wgpu::BufferUsages::COPY_DST - | wgpu::BufferUsages::COPY_SRC, - ); - buffers.push(params_buffer); - } - - kernel_backend::dynamic_kernel_from_ir( - device.kernel_cache(), - kernel_label, - cache_key, - move || { - let mut kb = tile_ir::KernelBuilder::<()>::new(); - let q_ref = tile_ir::KernelTensorRef::new((), layout.clone()); - let k_ref = tile_ir::KernelTensorRef::new((), layout.clone()); - let v_ref = tile_ir::KernelTensorRef::new((), layout.clone()); - let mask_ref = mask_buffer - .as_ref() - .map(|_| tile_ir::KernelTensorRef::new((), layout.clone())); - let output_ref = tile_ir::KernelTensorRef::new((), layout.clone()); - if let Some(meta) = decode_meta { - let params_ref = tile_ir::KernelTensorRef::new((), layout.clone()); - tile_ir_kernels::flash_decode_small( - &mut kb, q_ref, k_ref, v_ref, output_ref, params_ref, meta, - ) - } else { - let stream_meta = tile_ir_kernels::FlashAttentionMeta { - dims, - scale: tile_ir::F32Bits::new(scale), - q_meta: q_tile_meta, - k_meta: k_tile_meta, - v_meta: v_tile_meta, - mask_meta: mask_tile_meta, - output_meta: output_tile_meta, - dispatch_size, - causal, - }; - match variant { - FlashAttentionKernelVariant::StreamingTiled { - subgroup, - subgroup_size, - } => dispatch_streaming_tiled_flash_attention( - &mut kb, - q_ref, - k_ref, - v_ref, - mask_ref, - output_ref, - stream_meta, - input_dtype, - subgroup, - subgroup_size, - ), - FlashAttentionKernelVariant::Streaming { - subgroup, - subgroup_size, - } => dispatch_streaming_flash_attention( - &mut kb, - q_ref, - k_ref, - v_ref, - mask_ref, - output_ref, - stream_meta, - input_dtype, - subgroup, - subgroup_size, - ), - FlashAttentionKernelVariant::DecodeSmall - | FlashAttentionKernelVariant::DecodeSplitPartials - | FlashAttentionKernelVariant::DecodeSplitReduce => { - unreachable!("decode variants are handled before streaming IR build") - } - } - }?; - Some(kb.finish().0) - }, - buffers, - dispatch_size, - ) - } - - fn output(&self, _nodes: &ComputeGraphInner, inputs: &[MirValue]) -> MirValue { - inputs.last().unwrap().clone() - } - - fn name(&self) -> String { - format!( - "flash_attention_{}_{}x{}x{}x{}_by_{}x{}", - self.input_dtype, - self.q_shape[0], - self.q_shape[1], - self.q_shape[2], - self.q_shape[3], - self.k_shape[1], - self.k_shape[2], - ) - } -} diff --git a/fusor-ml/core/src/mir/kernel_backend/flash_attention/mod.rs b/fusor-ml/core/src/mir/kernel_backend/flash_attention/mod.rs deleted file mode 100644 index 15c666401..000000000 --- a/fusor-ml/core/src/mir/kernel_backend/flash_attention/mod.rs +++ /dev/null @@ -1,523 +0,0 @@ -use fusor_tile_ir as tile_ir; -use fusor_tile_ir_kernels as tile_ir_kernels; - -use crate::{ - DataTypeEnum, - compute_graph::NodeIndex, - kernel_selection::{ - Axis, DimConstraint, KernelDeviceCaps, KernelShape, ShapeRule, ShapeSelector, eq, range, - }, - tensor::TensorData, -}; - -mod kernel; -#[cfg(test)] -mod tests; - -const DECODE_SMALL_BLOCK: u32 = 128; -const DECODE_MID_BLOCK: u32 = 256; -const DECODE_MEDIUM_BLOCK: u32 = 512; -const DECODE_LARGE_BLOCK: u32 = 1024; -pub(crate) const MIN_DECODE_KV_SEQ: usize = 32; -/// Hardware subgroup sizes the streaming flash kernel emits IR for. The -/// kernel layout assumes one subgroup per output dim and uses subgroup -/// reductions across a `SIZE`-wide KV chunk, so the runtime subgroup width -/// must match one of these exactly. -const FLASH_STREAMING_SUBGROUP_SIZES: &[u32] = &[4, 8, 16, 32, 64]; - -#[allow(clippy::too_many_arguments)] -fn dispatch_streaming_flash_attention( - kb: &mut tile_ir::KernelBuilder<()>, - q: tile_ir::KernelTensorRef<()>, - k: tile_ir::KernelTensorRef<()>, - v: tile_ir::KernelTensorRef<()>, - mask: Option>, - output: tile_ir::KernelTensorRef<()>, - meta: tile_ir_kernels::FlashAttentionMeta, - input_dtype: DataTypeEnum, - subgroup: tile_ir::SubgroupToken, - subgroup_size: u32, -) -> Option<()> { - let element = match input_dtype { - DataTypeEnum::F32 => tile_ir::ElementType::F32, - DataTypeEnum::F16 => tile_ir::ElementType::F16, - _ => return None, - }; - tile_ir_kernels::flash_attention( - kb, - element, - tile_ir_kernels::FlashAttentionTensors { - q, - k, - v, - mask, - output, - }, - meta, - subgroup, - subgroup_size, - ) -} - -#[allow(clippy::too_many_arguments)] -fn dispatch_streaming_tiled_flash_attention( - kb: &mut tile_ir::KernelBuilder<()>, - q: tile_ir::KernelTensorRef<()>, - k: tile_ir::KernelTensorRef<()>, - v: tile_ir::KernelTensorRef<()>, - mask: Option>, - output: tile_ir::KernelTensorRef<()>, - meta: tile_ir_kernels::FlashAttentionMeta, - input_dtype: DataTypeEnum, - subgroup: tile_ir::SubgroupToken, - subgroup_size: u32, -) -> Option<()> { - let element = match input_dtype { - DataTypeEnum::F32 => tile_ir::ElementType::F32, - DataTypeEnum::F16 => tile_ir::ElementType::F16, - _ => return None, - }; - tile_ir_kernels::flash_attention_tiled( - kb, - element, - tile_ir_kernels::FlashAttentionTensors { - q, - k, - v, - mask, - output, - }, - meta, - subgroup, - subgroup_size, - FLASH_STREAMING_TILED_Q_BLOCK, - ) -} - -/// Returns true when the per-shape gating for the tiled (Q-batched) streaming -/// kernel is satisfied. Decode (q_seq_len < threshold) keeps the existing -/// streaming kernel because the per-workgroup K-cache load offers no reuse -/// when there's only one query. -pub(crate) fn flash_streaming_tiled_eligible(dims: FlashAttentionDims) -> bool { - dims.q_seq_len >= FLASH_STREAMING_TILED_MIN_Q - && dims - .head_dim - .is_multiple_of(FLASH_STREAMING_TILED_HEAD_DIM_ALIGN) -} - -pub(crate) fn flash_decode_direct_candidate( - q_shape: &[usize], - k_shape: &[usize], - has_mask: bool, - causal: bool, - input_dtype: DataTypeEnum, -) -> bool { - q_shape.get(2).copied() == Some(1) - && q_shape.get(3).copied().is_some_and(|head_dim| head_dim > 0) - && k_shape - .get(2) - .copied() - .is_some_and(|kv| kv >= MIN_DECODE_KV_SEQ) - && !has_mask - && !causal - && input_dtype == DataTypeEnum::F32 -} - -fn streaming_dispatch_size(dims: FlashAttentionDims, outputs_per_workgroup: u32) -> [u32; 3] { - [ - dims.head_dim.div_ceil(outputs_per_workgroup), - dims.batch - .checked_mul(dims.num_heads) - .and_then(|value| value.checked_mul(dims.q_seq_len)) - .expect("flash attention row dispatch overflow"), - 1, - ] -} - -#[derive(Clone, Copy, Debug, Hash, PartialEq, Eq)] -enum FlashAttentionKernelVariant { - Streaming { - subgroup: tile_ir::SubgroupToken, - subgroup_size: u32, - }, - StreamingTiled { - subgroup: tile_ir::SubgroupToken, - subgroup_size: u32, - }, - DecodeSmall, - DecodeSplitPartials, - DecodeSplitReduce, -} - -/// Q-block size used by the tiled (Q-batched) flash attention kernel: each -/// workgroup processes this many contiguous query rows, sharing one K/V -/// workgroup-memory load across them. -pub(crate) const FLASH_STREAMING_TILED_Q_BLOCK: u32 = 8; -/// Minimum query sequence length to switch to the tiled kernel. Below this -/// the per-workgroup setup overhead dominates the K/V reuse win. -pub(crate) const FLASH_STREAMING_TILED_MIN_Q: u32 = 128; -/// Required head_dim alignment for the tiled kernel. Apple subgroup width is -/// 32, but we round-robin K loads in groups of 8 (the subgroup-warp width on -/// Metal); enforce 8-alignment so the cooperative K-cache load issues stay -/// well-formed without per-lane tail handling. -pub(crate) const FLASH_STREAMING_TILED_HEAD_DIM_ALIGN: u32 = 8; - -struct FlashAttentionDirectKernelVariant; - -const FLASH_Q_SEQ: Axis<3> = Axis; -const FLASH_KV_SEQ: Axis<4> = Axis; -const FLASH_HEAD_DIM: Axis<5> = Axis; - -type FlashAttentionDims = tile_ir_kernels::FlashAttentionDims; -type FlashDecodeSmallMeta = tile_ir_kernels::FlashDecodeSmallMeta; - -/// Workgroup sizes the decode kernel monomorphizes over, smallest first. -/// `choose_decode_block` picks the smallest entry that covers the given KV -/// length and that the device's workgroup-size limits support. -const DECODE_BLOCKS: [u32; 4] = [ - DECODE_SMALL_BLOCK, - DECODE_MID_BLOCK, - DECODE_MEDIUM_BLOCK, - DECODE_LARGE_BLOCK, -]; - -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -enum FlashAttentionSelectedVariant { - Streaming, - /// Decode kernel parameterized by its workgroup BLOCK (128/512/1024). - DecodeSmall(u32), -} - -impl FlashAttentionSelectedVariant { - fn kernel_variant(self) -> FlashAttentionKernelVariant { - match self { - Self::Streaming => { - unreachable!("streaming variants are constructed with device subgroup tokens") - } - Self::DecodeSmall(_) => FlashAttentionKernelVariant::DecodeSmall, - } - } - - fn decode_block(self) -> Option { - match self { - Self::Streaming => None, - Self::DecodeSmall(block) => Some(block), - } - } -} - -#[derive(Clone, Copy, Debug)] -struct FlashAttentionSelectionCtx { - has_mask: bool, -} - -fn decode_block_supported(block: u32, caps: KernelDeviceCaps) -> bool { - let size = block; - size <= caps.max_compute_invocations_per_workgroup && size <= caps.max_compute_workgroup_size_x -} - -fn choose_decode_block(kv_seq_len: u32, head_dim: u32, caps: KernelDeviceCaps) -> Option { - if kv_seq_len == 0 || head_dim == 0 { - return None; - } - if kv_seq_len > DECODE_LARGE_BLOCK - && head_dim <= DECODE_MID_BLOCK - && decode_block_supported(DECODE_MID_BLOCK, caps) - { - return Some(DECODE_MID_BLOCK); - } - - let required_block = kv_seq_len.min(DECODE_LARGE_BLOCK).max(head_dim); - let mut largest_supported = None; - for block in DECODE_BLOCKS { - if !decode_block_supported(block, caps) { - continue; - } - largest_supported = Some(block); - if required_block <= block { - return Some(block); - } - } - largest_supported.filter(|block| head_dim <= *block) -} - -fn selected_decode_block_for_shape( - shape: KernelShape<6>, - ctx: &FlashAttentionSelectionCtx, - caps: KernelDeviceCaps, -) -> Option { - if ctx.has_mask || shape[FLASH_KV_SEQ] == 0 { - return None; - } - choose_decode_block( - shape[FLASH_KV_SEQ].try_into().ok()?, - shape[FLASH_HEAD_DIM].try_into().ok()?, - caps, - ) -} - -fn decode_shape_rule( - block: u32, - kv_constraint: Option, -) -> ShapeRule<6, FlashAttentionSelectionCtx> { - let mut rule = ShapeRule::new().axis(FLASH_Q_SEQ, eq(1)); - if let Some(kv_constraint) = kv_constraint { - rule = rule.axis(FLASH_KV_SEQ, kv_constraint); - } - rule.when(move |shape, ctx, caps| { - selected_decode_block_for_shape(shape, ctx, caps) == Some(block) - }) -} - -fn flash_attention_selector() --> ShapeSelector<6, FlashAttentionSelectionCtx, FlashAttentionSelectedVariant> { - ShapeSelector::new() - .rule( - FlashAttentionSelectedVariant::DecodeSmall(DECODE_SMALL_BLOCK), - decode_shape_rule( - DECODE_SMALL_BLOCK, - Some(range(1..=DECODE_SMALL_BLOCK as usize)), - ), - ) - .rule( - FlashAttentionSelectedVariant::DecodeSmall(DECODE_MID_BLOCK), - decode_shape_rule( - DECODE_MID_BLOCK, - Some(range( - (DECODE_SMALL_BLOCK as usize + 1)..=DECODE_MID_BLOCK as usize, - )), - ), - ) - .rule( - FlashAttentionSelectedVariant::DecodeSmall(DECODE_MEDIUM_BLOCK), - decode_shape_rule( - DECODE_MEDIUM_BLOCK, - Some(range( - (DECODE_MID_BLOCK as usize + 1)..=DECODE_MEDIUM_BLOCK as usize, - )), - ), - ) - .rule( - FlashAttentionSelectedVariant::DecodeSmall(DECODE_LARGE_BLOCK), - decode_shape_rule( - DECODE_LARGE_BLOCK, - Some(range( - (DECODE_MEDIUM_BLOCK as usize + 1)..=DECODE_LARGE_BLOCK as usize, - )), - ), - ) - .rule( - FlashAttentionSelectedVariant::DecodeSmall(DECODE_SMALL_BLOCK), - decode_shape_rule(DECODE_SMALL_BLOCK, None), - ) - .rule( - FlashAttentionSelectedVariant::DecodeSmall(DECODE_MID_BLOCK), - decode_shape_rule(DECODE_MID_BLOCK, None), - ) - .rule( - FlashAttentionSelectedVariant::DecodeSmall(DECODE_MEDIUM_BLOCK), - decode_shape_rule(DECODE_MEDIUM_BLOCK, None), - ) - .rule( - FlashAttentionSelectedVariant::DecodeSmall(DECODE_LARGE_BLOCK), - decode_shape_rule(DECODE_LARGE_BLOCK, None), - ) - .rule(FlashAttentionSelectedVariant::Streaming, ShapeRule::new()) -} - -fn select_flash_attention_variant( - dims: FlashAttentionDims, - has_mask: bool, - caps: KernelDeviceCaps, -) -> FlashAttentionSelectedVariant { - let shape = KernelShape::new([ - dims.batch as usize, - dims.num_heads as usize, - dims.num_kv_heads as usize, - dims.q_seq_len as usize, - dims.kv_seq_len as usize, - dims.head_dim as usize, - ]); - let ctx = FlashAttentionSelectionCtx { has_mask }; - flash_attention_selector() - .select(shape, &ctx, caps) - .expect("flash attention selector has a catch-all rule") -} - -#[derive(Clone, Debug)] -pub(crate) struct FlashAttentionOperation { - pub(crate) q: NodeIndex, - pub(crate) k: NodeIndex, - pub(crate) v: NodeIndex, - pub(crate) mask: Option, - pub(crate) out_shape: Box<[usize]>, - q_shape: Box<[usize]>, - k_shape: Box<[usize]>, - scale: f32, - input_dtype: DataTypeEnum, - pub(crate) causal: bool, -} - -pub(crate) struct FlashAttentionInputs<'a> { - pub(crate) q: NodeIndex, - pub(crate) k: NodeIndex, - pub(crate) v: NodeIndex, - pub(crate) mask: Option, - pub(crate) q_shape: &'a [usize], - pub(crate) k_shape: &'a [usize], - pub(crate) v_shape: &'a [usize], - pub(crate) scale: f32, - pub(crate) input_dtype: DataTypeEnum, - pub(crate) causal: bool, -} - -impl FlashAttentionOperation { - pub(crate) fn new(inputs: FlashAttentionInputs<'_>) -> Self { - let FlashAttentionInputs { - q, - k, - v, - mask, - q_shape, - k_shape, - v_shape, - scale, - input_dtype, - causal, - } = inputs; - assert_eq!(q_shape.len(), 4, "Q must be rank-4"); - assert_eq!(k_shape.len(), 4, "K must be rank-4"); - assert_eq!(v_shape.len(), 4, "V must be rank-4"); - assert_eq!(q_shape[0], k_shape[0], "Q and K batch dimensions differ"); - assert_eq!(q_shape[0], v_shape[0], "Q and V batch dimensions differ"); - assert_eq!(k_shape[1], v_shape[1], "K and V head dimensions differ"); - assert_eq!(k_shape[2], v_shape[2], "K and V sequence dimensions differ"); - assert_eq!(q_shape[3], k_shape[3], "Q and K head dimensions differ"); - assert_eq!(q_shape[3], v_shape[3], "Q and V head dimensions differ"); - assert!( - q_shape[1].is_multiple_of(k_shape[1]), - "Number of Q heads ({}) must be divisible by number of K/V heads ({})", - q_shape[1], - k_shape[1] - ); - - if causal { - assert!( - mask.is_none(), - "causal flash attention cannot accept an additive mask" - ); - assert_eq!( - q_shape[2], k_shape[2], - "causal flash attention requires q_seq_len == kv_seq_len, got {} vs {}", - q_shape[2], k_shape[2] - ); - } - - Self { - q, - k, - v, - mask, - out_shape: q_shape.into(), - q_shape: q_shape.into(), - k_shape: k_shape.into(), - scale, - input_dtype, - causal, - } - } - - fn dims(&self) -> Option { - Some(FlashAttentionDims { - batch: self.q_shape[0].try_into().ok()?, - num_heads: self.q_shape[1].try_into().ok()?, - num_kv_heads: self.k_shape[1].try_into().ok()?, - q_seq_len: self.q_shape[2].try_into().ok()?, - kv_seq_len: self.k_shape[2].try_into().ok()?, - head_dim: self.q_shape[3].try_into().ok()?, - }) - } -} - -#[derive(Clone)] -pub(crate) struct TensorMeta { - datatype: DataTypeEnum, - tile: tile_ir_kernels::TensorMeta, -} - -impl TensorMeta { - fn new(tensor: &TensorData) -> Option { - let strides = tensor - .layout() - .strides() - .iter() - .copied() - .map(u32::try_from) - .collect::, _>>() - .ok()?; - let offset = tensor.layout().offset().try_into().ok()?; - Some(Self { - datatype: tensor.datatype(), - tile: tile_ir_kernels::TensorMeta::new(strides, offset), - }) - } - - fn stride4(&self) -> Option<[u32; 4]> { - self.tile.strides.as_slice().try_into().ok() - } -} - -struct FlashDecodeSmallTensors<'a> { - q: TensorMeta, - k: TensorMeta, - v: TensorMeta, - mask: Option<&'a TensorMeta>, - output: TensorMeta, -} - -fn build_flash_decode_small_meta( - dims: FlashAttentionDims, - scale: f32, - caps: KernelDeviceCaps, - tensors: FlashDecodeSmallTensors<'_>, -) -> Option { - let FlashDecodeSmallTensors { - q: q_meta, - k: k_meta, - v: v_meta, - mask: mask_meta, - output: output_meta, - } = tensors; - if mask_meta.is_some() || dims.q_seq_len != 1 || dims.head_dim == 0 || dims.kv_seq_len == 0 { - return None; - } - let decode_block = choose_decode_block(dims.kv_seq_len, dims.head_dim, caps)?; - let tiled = dims.kv_seq_len > decode_block; - let split_blocks = dims.kv_seq_len.div_ceil(decode_block); - - let groups = dims.num_heads.checked_div(dims.num_kv_heads)?; - if groups == 0 { - return None; - } - - let mut module_dims = dims; - module_dims.kv_seq_len = decode_block; - - Some(FlashDecodeSmallMeta { - dims: module_dims, - scale: tile_ir::F32Bits::new(scale), - active_kv_len: dims.kv_seq_len, - decode_block, - tiled, - split_blocks, - groups, - q_offset: q_meta.tile.offset, - k_offset: k_meta.tile.offset, - v_offset: v_meta.tile.offset, - output_offset: output_meta.tile.offset, - q_strides: q_meta.stride4()?, - k_strides: k_meta.stride4()?, - v_strides: v_meta.stride4()?, - output_strides: output_meta.stride4()?, - }) -} diff --git a/fusor-ml/core/src/mir/kernel_backend/flash_attention/tests.rs b/fusor-ml/core/src/mir/kernel_backend/flash_attention/tests.rs deleted file mode 100644 index f9606a327..000000000 --- a/fusor-ml/core/src/mir/kernel_backend/flash_attention/tests.rs +++ /dev/null @@ -1,610 +0,0 @@ -use super::*; -use crate::{ - Device, Tensor, - kernel_selection::{CooperativeMatrixCaps, assert_selector_generates}, - mir::workgroup_shape::WorkgroupShape, -}; -use std::sync::{Mutex, MutexGuard, OnceLock}; - -const TEST_HEAD_DIM: usize = 128; - -fn caps(max_compute_invocations_per_workgroup: u32) -> KernelDeviceCaps { - KernelDeviceCaps { - cooperative_matrix: CooperativeMatrixCaps::default(), - max_compute_invocations_per_workgroup, - ..KernelDeviceCaps::test_caps() - } -} - -fn tensor_meta4() -> TensorMeta { - TensorMeta { - datatype: DataTypeEnum::F32, - tile: tile_ir_kernels::TensorMeta::new(vec![65_536, 8_192, 128, 1], 0), - } -} - -fn decode_dims(kv_seq_len: u32) -> FlashAttentionDims { - FlashAttentionDims { - batch: 1, - num_heads: 32, - num_kv_heads: 8, - q_seq_len: 1, - kv_seq_len, - head_dim: TEST_HEAD_DIM as u32, - } -} - -fn decode_shape(kv_seq_len: usize) -> KernelShape<6> { - KernelShape::new([1, 32, 8, 1, kv_seq_len, TEST_HEAD_DIM]) -} - -fn decode_shape_with_head_dim(kv_seq_len: usize, head_dim: usize) -> KernelShape<6> { - KernelShape::new([1, 32, 8, 1, kv_seq_len, head_dim]) -} - -fn gpu_test_guard() -> MutexGuard<'static, ()> { - static LOCK: OnceLock> = OnceLock::new(); - LOCK.get_or_init(|| Mutex::new(())) - .lock() - .unwrap_or_else(|poisoned| poisoned.into_inner()) -} - -#[test] -fn decode_block_choice_uses_smallest_covering_supported_block() { - assert_eq!( - choose_decode_block(64, TEST_HEAD_DIM as u32, caps(DECODE_LARGE_BLOCK)), - Some(DECODE_SMALL_BLOCK) - ); - assert_eq!( - choose_decode_block(200, TEST_HEAD_DIM as u32, caps(DECODE_LARGE_BLOCK)), - Some(DECODE_MID_BLOCK) - ); - assert_eq!( - choose_decode_block(600, TEST_HEAD_DIM as u32, caps(DECODE_LARGE_BLOCK)), - Some(DECODE_LARGE_BLOCK) - ); - assert_eq!( - choose_decode_block(600, TEST_HEAD_DIM as u32, caps(DECODE_MEDIUM_BLOCK)), - Some(DECODE_MEDIUM_BLOCK) - ); - assert_eq!( - choose_decode_block( - DECODE_LARGE_BLOCK + 1, - TEST_HEAD_DIM as u32, - caps(DECODE_LARGE_BLOCK) - ), - Some(DECODE_MID_BLOCK) - ); - assert_eq!( - choose_decode_block( - DECODE_SMALL_BLOCK, - TEST_HEAD_DIM as u32, - caps(DECODE_SMALL_BLOCK - 1) - ), - None - ); - assert_eq!( - choose_decode_block(64, 64, caps(DECODE_LARGE_BLOCK)), - Some(DECODE_SMALL_BLOCK) - ); - - assert_eq!( - choose_decode_block( - 200, - TEST_HEAD_DIM as u32, - KernelDeviceCaps { - backend: wgpu::Backend::Metal, - ..caps(DECODE_LARGE_BLOCK) - }, - ), - Some(DECODE_MID_BLOCK) - ); -} - -#[test] -fn decode_small_meta_buckets_dynamic_kv_len() { - let meta = build_flash_decode_small_meta( - decode_dims(DECODE_SMALL_BLOCK + 1), - 1.0, - caps(DECODE_LARGE_BLOCK), - FlashDecodeSmallTensors { - q: tensor_meta4(), - k: tensor_meta4(), - v: tensor_meta4(), - mask: None, - output: tensor_meta4(), - }, - ) - .unwrap(); - - assert_eq!(meta.active_kv_len, DECODE_SMALL_BLOCK + 1); - assert_eq!(meta.decode_block, DECODE_MID_BLOCK); - assert_eq!(meta.dims.kv_seq_len, DECODE_MID_BLOCK); - assert_eq!(meta.split_blocks, 1); - assert!(!meta.tiled); -} - -#[test] -fn decode_cache_key_uses_bucketed_kv_len() { - let meta = build_flash_decode_small_meta( - decode_dims(DECODE_SMALL_BLOCK + 1), - 1.0, - caps(DECODE_LARGE_BLOCK), - FlashDecodeSmallTensors { - q: tensor_meta4(), - k: tensor_meta4(), - v: tensor_meta4(), - mask: None, - output: tensor_meta4(), - }, - ) - .unwrap(); - let mut next_token_meta = meta; - next_token_meta.active_kv_len += 1; - - let workgroup_shape = WorkgroupShape::new(1, 1, 1); - let dispatch_size = [meta.dims.batch * meta.dims.num_heads, 1, 1]; - let key = kernel::flash_decode_cache_key( - Some(&workgroup_shape), - dispatch_size, - DataTypeEnum::F32, - 1.0, - &meta, - ); - let next_token_key = kernel::flash_decode_cache_key( - Some(&workgroup_shape), - dispatch_size, - DataTypeEnum::F32, - 1.0, - &next_token_meta, - ); - - assert_eq!(key, next_token_key); - - let mut different_stride_meta = meta; - different_stride_meta.k_strides[1] += TEST_HEAD_DIM as u32; - let different_stride_key = kernel::flash_decode_cache_key( - Some(&workgroup_shape), - dispatch_size, - DataTypeEnum::F32, - 1.0, - &different_stride_meta, - ); - - assert_ne!(key, different_stride_key); -} - -#[test] -fn decode_small_meta_tiles_with_largest_supported_block() { - let meta = build_flash_decode_small_meta( - decode_dims(DECODE_MEDIUM_BLOCK + 1), - 1.0, - caps(DECODE_MEDIUM_BLOCK), - FlashDecodeSmallTensors { - q: tensor_meta4(), - k: tensor_meta4(), - v: tensor_meta4(), - mask: None, - output: tensor_meta4(), - }, - ); - - let meta = meta.unwrap(); - assert_eq!(meta.active_kv_len, DECODE_MEDIUM_BLOCK + 1); - assert_eq!(meta.decode_block, DECODE_MEDIUM_BLOCK); - assert_eq!(meta.dims.kv_seq_len, DECODE_MEDIUM_BLOCK); - assert_eq!(meta.split_blocks, 2); - assert!(meta.tiled); -} - -#[test] -fn decode_small_meta_requires_minimum_workgroup_limit() { - let meta = build_flash_decode_small_meta( - decode_dims(DECODE_SMALL_BLOCK), - 1.0, - caps(DECODE_SMALL_BLOCK - 1), - FlashDecodeSmallTensors { - q: tensor_meta4(), - k: tensor_meta4(), - v: tensor_meta4(), - mask: None, - output: tensor_meta4(), - }, - ); - - assert!(meta.is_none()); -} - -#[test] -fn flash_attention_selector_selects_decode_block_buckets() { - let selector = flash_attention_selector(); - let decode_ctx = FlashAttentionSelectionCtx { has_mask: false }; - let masked_ctx = FlashAttentionSelectionCtx { has_mask: true }; - - assert_eq!( - selector.select(decode_shape(64), &decode_ctx, caps(DECODE_LARGE_BLOCK)), - Some(FlashAttentionSelectedVariant::DecodeSmall( - DECODE_SMALL_BLOCK - )) - ); - assert_eq!( - selector.select(decode_shape(200), &decode_ctx, caps(DECODE_LARGE_BLOCK)), - Some(FlashAttentionSelectedVariant::DecodeSmall(DECODE_MID_BLOCK)) - ); - assert_eq!( - selector.select( - decode_shape_with_head_dim(200, 64), - &decode_ctx, - caps(DECODE_LARGE_BLOCK) - ), - Some(FlashAttentionSelectedVariant::DecodeSmall(DECODE_MID_BLOCK)) - ); - assert_eq!( - selector.select(decode_shape(600), &decode_ctx, caps(DECODE_LARGE_BLOCK)), - Some(FlashAttentionSelectedVariant::DecodeSmall( - DECODE_LARGE_BLOCK - )) - ); - assert_eq!( - selector.select(decode_shape(600), &decode_ctx, caps(DECODE_MEDIUM_BLOCK)), - Some(FlashAttentionSelectedVariant::DecodeSmall( - DECODE_MEDIUM_BLOCK - )) - ); - assert_eq!( - selector.select( - decode_shape(DECODE_LARGE_BLOCK as usize + 1), - &decode_ctx, - caps(DECODE_LARGE_BLOCK) - ), - Some(FlashAttentionSelectedVariant::DecodeSmall(DECODE_MID_BLOCK)) - ); - assert_eq!( - selector.select(decode_shape(200), &decode_ctx, caps(DECODE_SMALL_BLOCK - 1)), - Some(FlashAttentionSelectedVariant::Streaming) - ); - assert_eq!( - selector.select(decode_shape(200), &masked_ctx, caps(DECODE_LARGE_BLOCK)), - Some(FlashAttentionSelectedVariant::Streaming) - ); -} - -#[test] -fn decode_direct_candidate_does_not_require_subgroups() { - let q_shape = [1, 32, 1, 64]; - let k_shape = [1, 8, MIN_DECODE_KV_SEQ, 64]; - assert!(flash_decode_direct_candidate( - &q_shape, - &k_shape, - false, - false, - DataTypeEnum::F32, - )); - - let caps_without_subgroups = KernelDeviceCaps { - subgroups_supported: false, - min_subgroup_size: 0, - max_subgroup_size: 0, - ..caps(DECODE_LARGE_BLOCK) - }; - assert_eq!( - select_flash_attention_variant( - decode_dims(MIN_DECODE_KV_SEQ as u32), - false, - caps_without_subgroups - ), - FlashAttentionSelectedVariant::DecodeSmall(DECODE_SMALL_BLOCK), - ); -} - -#[test] -fn flash_attention_selector_generates_each_variant() { - let selector = flash_attention_selector(); - let decode_ctx = FlashAttentionSelectionCtx { has_mask: false }; - let streaming_ctx = FlashAttentionSelectionCtx { has_mask: true }; - let cases = [ - ( - FlashAttentionSelectedVariant::DecodeSmall(DECODE_SMALL_BLOCK), - decode_ctx, - caps(DECODE_SMALL_BLOCK), - ), - ( - FlashAttentionSelectedVariant::DecodeSmall(DECODE_MID_BLOCK), - decode_ctx, - caps(DECODE_MID_BLOCK), - ), - ( - FlashAttentionSelectedVariant::DecodeSmall(DECODE_MEDIUM_BLOCK), - decode_ctx, - caps(DECODE_MEDIUM_BLOCK), - ), - ( - FlashAttentionSelectedVariant::DecodeSmall(DECODE_LARGE_BLOCK), - decode_ctx, - caps(DECODE_LARGE_BLOCK), - ), - ( - FlashAttentionSelectedVariant::Streaming, - streaming_ctx, - caps(DECODE_LARGE_BLOCK), - ), - ]; - assert_selector_generates(&selector, cases); -} - -type AttentionFixture = Vec>>>; - -fn attention_fixture( - heads: usize, - tokens: usize, - f: impl Fn(usize, usize, usize) -> f32, -) -> AttentionFixture { - vec![ - (0..heads) - .map(|head| { - (0..tokens) - .map(|token| (0..TEST_HEAD_DIM).map(|dim| f(head, token, dim)).collect()) - .collect() - }) - .collect(), - ] -} - -fn decode_q() -> AttentionFixture { - attention_fixture(1, 1, |_, _, dim| ((dim % 17) as f32 - 8.0) * 0.0075) -} - -fn decode_k(kv_len: usize) -> AttentionFixture { - attention_fixture(1, kv_len, |_, token, dim| { - let value = ((token * 13 + dim * 7) % 31) as f32 - 15.0; - value * 0.004 - }) -} - -fn decode_v(kv_len: usize) -> AttentionFixture { - attention_fixture(1, kv_len, |_, token, dim| { - let value = ((token * 5 + dim * 11) % 37) as f32 - 18.0; - 0.25 + value * 0.01 - }) -} - -fn decode_q_gqa(num_heads: usize) -> AttentionFixture { - attention_fixture(num_heads, 1, |head, _, dim| { - let value = ((head * 19 + dim * 7) % 43) as f32 - 21.0; - value * 0.003 - }) -} - -fn decode_k_gqa(num_kv_heads: usize, kv_len: usize) -> AttentionFixture { - attention_fixture(num_kv_heads, kv_len, |kv_head, token, dim| { - let value = ((kv_head * 23 + token * 13 + dim * 5) % 47) as f32 - 23.0; - value * 0.0025 - }) -} - -fn decode_v_gqa(num_kv_heads: usize, kv_len: usize) -> AttentionFixture { - attention_fixture(num_kv_heads, kv_len, |kv_head, token, dim| { - let value = ((kv_head * 29 + token * 3 + dim * 11) % 53) as f32 - 26.0; - 0.05 + value * 0.004 - }) -} - -fn cpu_decode_reference(q: &[f32], k: &[Vec], v: &[Vec], scale: f32) -> Vec { - let scores = k - .iter() - .map(|key| { - q.iter() - .zip(key) - .map(|(q, k)| (*q as f64) * (*k as f64)) - .sum::() - * scale as f64 - }) - .collect::>(); - let max_score = scores.iter().copied().fold(f64::NEG_INFINITY, f64::max); - let denom = scores - .iter() - .map(|score| (score - max_score).exp()) - .sum::(); - let mut output = vec![0.0; TEST_HEAD_DIM]; - for (token, score) in scores.iter().copied().enumerate() { - let prob = (score - max_score).exp() / denom; - for (dim, output) in output.iter_mut().enumerate() { - *output += prob * v[token][dim] as f64; - } - } - output.into_iter().map(|value| value as f32).collect() -} - -fn decode_max_error( - num_heads: usize, - groups: usize, - q_data: &AttentionFixture, - k_data: &AttentionFixture, - v_data: &AttentionFixture, - scale: f32, - actual: impl Fn(usize, usize) -> f32, -) -> (f32, usize, usize, f32, f32) { - let mut max_error = 0.0f32; - let mut max_head = 0usize; - let mut max_dim = 0usize; - let mut max_actual = 0.0f32; - let mut max_expected = 0.0f32; - for (head, q_head) in q_data[0].iter().enumerate().take(num_heads) { - let kv_head = head / groups; - let expected = - cpu_decode_reference(&q_head[0], &k_data[0][kv_head], &v_data[0][kv_head], scale); - for (dim, expected) in expected.into_iter().enumerate() { - let actual = actual(head, dim); - let error = (actual - expected).abs(); - if error > max_error { - max_error = error; - max_head = head; - max_dim = dim; - max_actual = actual; - max_expected = expected; - } - } - } - (max_error, max_head, max_dim, max_actual, max_expected) -} - -#[test] -fn tiled_decode_attention_matches_cpu_reference() { - let _guard = gpu_test_guard(); - pollster::block_on(async { - let Ok(device) = Device::new().await else { - return; - }; - - let kv_len = DECODE_LARGE_BLOCK as usize + 1; - let q_data = decode_q(); - let k_data = decode_k(kv_len); - let v_data = decode_v(kv_len); - let scale = 1.0 / f32::sqrt(TEST_HEAD_DIM as f32); - - let q = Tensor::new(&device, &q_data); - let k = Tensor::new(&device, &k_data); - let v = Tensor::new(&device, &v_data); - let output = q.try_flash_attention_direct(&k, &v, scale, None).unwrap(); - let output = output.as_slice::<4, f32>().await.unwrap(); - let (max_error, _, max_dim, max_actual, max_expected) = - decode_max_error(1, 1, &q_data, &k_data, &v_data, scale, |_, dim| { - output[[0, 0, 0, dim]] - }); - assert!( - max_error < 2.0e-4, - "dim {max_dim}: actual={max_actual} expected={max_expected} error={max_error}" - ); - }); -} - -#[test] -fn tiled_decode_attention_gqa_matches_cpu_reference() { - let _guard = gpu_test_guard(); - pollster::block_on(async { - let Ok(device) = Device::new().await else { - return; - }; - - let num_heads = 32; - let num_kv_heads = 8; - let groups = num_heads / num_kv_heads; - let kv_len = DECODE_LARGE_BLOCK as usize + 1; - let q_data = decode_q_gqa(num_heads); - let k_data = decode_k_gqa(num_kv_heads, kv_len); - let v_data = decode_v_gqa(num_kv_heads, kv_len); - let scale = 1.0 / f32::sqrt(TEST_HEAD_DIM as f32); - - let q = Tensor::new(&device, &q_data); - let k = Tensor::new(&device, &k_data); - let v = Tensor::new(&device, &v_data); - let output = q.try_flash_attention_direct(&k, &v, scale, None).unwrap(); - let output = output.as_slice::<4, f32>().await.unwrap(); - - let (max_error, max_head, max_dim, max_actual, max_expected) = decode_max_error( - num_heads, - groups, - &q_data, - &k_data, - &v_data, - scale, - |head, dim| output[[0, head, 0, dim]], - ); - assert!( - max_error < 3.0e-4, - "head {max_head} dim {max_dim}: actual={max_actual} expected={max_expected} error={max_error}" - ); - }); -} - -/// Regression test for the non-tiled 512/1024-thread decode blocks. -/// Before the fix, the per-thread score loop folded its 128 q*k -/// accumulations into a single deeply-nested Naga expression, which -/// miscompiled on Metal once the kernel's `workgroup_size` exceeded 128; -/// the kernel produced all-zero output. The fix emits the dot product as a -/// shader loop with a function-scope accumulator. -#[test] -fn decode_gqa_non_tiled_large_blocks_match_cpu_reference() { - let _guard = gpu_test_guard(); - pollster::block_on(async { - let Ok(device) = Device::new().await else { - return; - }; - - let num_heads = 32; - let num_kv_heads = 8; - let groups = num_heads / num_kv_heads; - let caps = KernelDeviceCaps::from_device(&device); - - // On devices that support the larger workgroups, 200 uses the 512 - // block and 600 uses the 1024 block. - for (kv_len, expected_block) in [(200usize, DECODE_MID_BLOCK), (600, DECODE_LARGE_BLOCK)] { - if choose_decode_block(kv_len as u32, TEST_HEAD_DIM as u32, caps) - != Some(expected_block) - { - continue; - } - let q_data = decode_q_gqa(num_heads); - let k_data = decode_k_gqa(num_kv_heads, kv_len); - let v_data = decode_v_gqa(num_kv_heads, kv_len); - let scale = 1.0 / f32::sqrt(TEST_HEAD_DIM as f32); - - let q = Tensor::new(&device, &q_data); - let k = Tensor::new(&device, &k_data); - let v = Tensor::new(&device, &v_data); - let output = q.try_flash_attention_direct(&k, &v, scale, None).unwrap(); - let output = output.as_slice::<4, f32>().await.unwrap(); - - let (max_error, max_head, max_dim, max_actual, max_expected) = decode_max_error( - num_heads, - groups, - &q_data, - &k_data, - &v_data, - scale, - |head, dim| output[[0, head, 0, dim]], - ); - assert!( - max_error < 5.0e-4, - "kv_len={kv_len} head={max_head} dim={max_dim}: actual={max_actual} expected={max_expected} error={max_error}" - ); - } - }); -} - -#[test] -fn streaming_gqa_regression_shape_builds_direct_kernel() { - let _guard = gpu_test_guard(); - pollster::block_on(async { - let Ok(device) = Device::new().await else { - return; - }; - - let q_data = vec![ - (0..32) - .map(|head| { - (0..48) - .map(|token| { - (0..TEST_HEAD_DIM) - .map(|dim| { - let value = - ((head * 17 + token * 11 + dim * 5) % 41) as f32 - 20.0; - value * 0.002 - }) - .collect::>() - }) - .collect::>() - }) - .collect::>(), - ]; - let k_data = decode_k_gqa(8, 48); - let v_data = decode_v_gqa(8, 48); - let scale = 1.0 / f32::sqrt(TEST_HEAD_DIM as f32); - - let q = Tensor::new(&device, &q_data); - let k = Tensor::new(&device, &k_data); - let v = Tensor::new(&device, &v_data); - let output = q.try_flash_attention_direct(&k, &v, scale, None).unwrap(); - output.materialize().await; - }); -} diff --git a/fusor-ml/core/src/mir/kernel_backend/rms_norm.rs b/fusor-ml/core/src/mir/kernel_backend/rms_norm.rs deleted file mode 100644 index d6a808fe2..000000000 --- a/fusor-ml/core/src/mir/kernel_backend/rms_norm.rs +++ /dev/null @@ -1,799 +0,0 @@ -use std::{any::Any, hash::Hash, sync::Arc}; - -use crate::{ - DataTypeEnum, Layout, - compute_graph::{ComputeGraphInner, GraphOperation, NodeIndex}, - kernel_selection::{ - Axis, KernelDeviceCaps, KernelShape, ShapeRule, ShapeSelector, multiple_of, - }, - mir::{ - inputs::MirValue, - kernel_backend, - kernel_backend::DirectKernel, - operation::Operation, - tile_direct::{ - flatten_matrix_layout, tile_storage_read_with_direct_layout_typed, - tile_storage_write_with_direct_layout_typed, - }, - workgroup_shape::{Constraint, WorkgroupShape, WorkgroupShapeConstraints}, - }, - nary_direct::apply_unary_function_chain, - nary_wise::UnaryFunctionChain, - tensor::{TensorData, TensorLayoutInfo}, -}; -use fusor_tile_ir as tile_ir; -use fusor_tile_ir_kernels as tile_ir_kernels; -use rustc_hash::FxHashMap; -use rustc_hash::FxHasher; - -// 512 instead of 1024 because Windows wgpu (DX12 / WARP) caps -// `max_compute_invocations_per_workgroup` at 768, so a 1024-wide -// workgroup fails pipeline creation on those adapters with -// `Shader entry point's workgroup size [1024, 1, 1] must be less or -// equal to the per-dimension limit […] and the total invocation limit -// […] of 768`. 512 fits every backend wgpu supports (the cooperative -// matmul also uses 512 for the same reason) and is still wide enough -// to feed a 128-block subgroup-reduce on Apple Silicon / NVIDIA. -const BLOCK: usize = 512; - -fn collect_rms_buffers( - input: &TensorData, - residual: Option<&TensorData>, - weight: &TensorData, - bias: Option<&TensorData>, - output: &TensorData, -) -> Vec> { - let mut buffers = Vec::with_capacity(5); - buffers.push(input.buffer().clone()); - if let Some(residual) = residual { - buffers.push(residual.buffer().clone()); - } - buffers.push(weight.buffer().clone()); - if let Some(bias) = bias { - buffers.push(bias.buffer().clone()); - } - buffers.push(output.buffer().clone()); - buffers -} - -#[derive(Clone, Copy, Debug, Hash, PartialEq, Eq)] -enum RmsNormKernelVariant { - Tile, - Vec4, -} - -struct RmsNormDirectKernelVariant; - -const RMS_COLS: Axis<1> = Axis; - -#[derive(Clone, Copy, Debug)] -struct RmsNormSelectionCtx { - vec4_supported: bool, -} - -fn rms_norm_selector() -> ShapeSelector<2, RmsNormSelectionCtx, RmsNormKernelVariant> { - ShapeSelector::new() - .rule( - RmsNormKernelVariant::Vec4, - ShapeRule::new() - .axis(RMS_COLS, multiple_of(4)) - .when_ctx(|ctx: &RmsNormSelectionCtx| ctx.vec4_supported), - ) - .rule(RmsNormKernelVariant::Tile, ShapeRule::new()) -} - -fn select_rms_norm_variant( - rows: u32, - cols: u32, - ctx: &RmsNormSelectionCtx, - caps: KernelDeviceCaps, -) -> RmsNormKernelVariant { - rms_norm_selector() - .select(KernelShape::new([rows as usize, cols as usize]), ctx, caps) - .expect("rms norm selector has a catch-all rule") -} - -#[derive(Clone, Debug)] -pub(crate) struct RmsNormOperation { - pub(crate) input: NodeIndex, - pub(crate) residual: Option, - pub(crate) weight: NodeIndex, - pub(crate) bias: Option, - shape: Box<[usize]>, - eps: f32, - /// Unary chain applied to each normalized output element in-register - /// before the store. Populated by `try_fuse_into_rmsnorm` when a - /// downstream `Nary` matches the element-wise fusion pattern. - pub(crate) post_element_wise: UnaryFunctionChain, -} - -impl RmsNormOperation { - pub(crate) fn new( - input: NodeIndex, - weight: NodeIndex, - bias: Option, - shape: &[usize], - eps: f32, - ) -> Self { - Self { - input, - residual: None, - weight, - bias, - shape: shape.into(), - eps, - post_element_wise: UnaryFunctionChain::empty(DataTypeEnum::F32), - } - } - - pub(crate) fn new_with_residual( - input: NodeIndex, - residual: NodeIndex, - weight: NodeIndex, - bias: Option, - shape: &[usize], - eps: f32, - ) -> Self { - Self { - input, - residual: Some(residual), - weight, - bias, - shape: shape.into(), - eps, - post_element_wise: UnaryFunctionChain::empty(DataTypeEnum::F32), - } - } - - fn rows_cols(&self) -> Option<(u32, u32)> { - let cols = *self.shape.last()?; - let rows = self.shape[..self.shape.len().saturating_sub(1)] - .iter() - .try_fold(1usize, |acc, dim| acc.checked_mul(*dim))?; - Some((rows.try_into().ok()?, cols.try_into().ok()?)) - } -} - -impl Operation for RmsNormOperation { - fn hash_kernel_fields(&self, state: &mut FxHasher) { - self.residual.is_some().hash(state); - self.bias.is_some().hash(state); - self.shape.hash(state); - self.eps.to_bits().hash(state); - self.post_element_wise.hash(state); - } - - fn workgroup_shape_constraints(&self, _device: &crate::Device) -> WorkgroupShapeConstraints { - let mut constraints = WorkgroupShapeConstraints::new(); - constraints.add_constraint(0, Constraint::Equals(1)); - constraints.add_constraint(1, Constraint::Equals(1)); - constraints.add_constraint(2, Constraint::Equals(1)); - constraints - } - - fn dispatch_size(&self, _workgroup_shape: &WorkgroupShape, _inputs: &[MirValue]) -> [u32; 3] { - let (rows, _) = self - .rows_cols() - .expect("rms norm requires a non-empty shape"); - [rows, 1, 1] - } - - fn visit_dependencies(&self, f: &mut dyn FnMut(NodeIndex)) { - f(self.input); - if let Some(residual) = self.residual { - f(residual); - } - f(self.weight); - if let Some(bias) = self.bias { - f(bias); - } - } - - fn inputs(&self, nodes: &ComputeGraphInner) -> Vec { - let input = nodes.get_cached_result(self.input).unwrap(); - let residual = self - .residual - .map(|residual| nodes.get_cached_result(residual).unwrap()); - let weight = nodes.get_cached_result(self.weight).unwrap(); - let output = - TensorData::new_for_shape(input.device(), input.layout().shape(), input.datatype()); - - let mut inputs = vec![input.clone().into()]; - if let Some(residual) = residual { - inputs.push(residual.clone().into()); - } - inputs.push(weight.clone().into()); - if let Some(bias) = self.bias { - inputs.push(nodes.get_cached_result(bias).unwrap().clone().into()); - } - inputs.push(output.into()); - inputs - } - - fn build_direct_kernel( - &self, - graph: &ComputeGraphInner, - workgroup_shape: &WorkgroupShape, - inputs: &[MirValue], - ) -> Option { - let input = inputs.first()?.as_tensor()?; - let (residual, weight_index) = if self.residual.is_some() { - (Some(inputs.get(1)?.as_tensor()?), 2) - } else { - (None, 1) - }; - let weight = inputs.get(weight_index)?.as_tensor()?; - let (bias, output_index) = if self.bias.is_some() { - ( - Some(inputs.get(weight_index + 1)?.as_tensor()?), - weight_index + 2, - ) - } else { - (None, weight_index + 1) - }; - let output = inputs.get(output_index)?.as_tensor()?; - - let storage_datatype = input.datatype(); - if !matches!(storage_datatype, DataTypeEnum::F32 | DataTypeEnum::F16) - || (storage_datatype == DataTypeEnum::F16 && !graph.device().f16_supported()) - || residual.is_some_and(|residual| residual.datatype() != storage_datatype) - || weight.datatype() != storage_datatype - || output.datatype() != storage_datatype - || bias.is_some_and(|bias| bias.datatype() != storage_datatype) - { - return None; - } - - let input_view = flatten_matrix_layout(input.layout())?; - let residual_view = match residual { - Some(residual) => Some(flatten_matrix_layout(residual.layout())?), - None => None, - }; - let output_view = flatten_matrix_layout(output.layout())?; - let rows = input_view.rows; - let cols = input_view.cols; - if rows != output_view.rows || cols != output_view.cols { - return None; - } - if let Some(residual_view) = residual_view.as_ref() - && (rows != residual_view.rows || cols != residual_view.cols) - { - return None; - } - if weight.layout().shape() != [cols as usize] { - return None; - } - if let Some(bias) = bias - && bias.layout().shape() != [cols as usize] - { - return None; - } - - // The vec4 path is a pre-built kernel that doesn't accept an arbitrary - // post-element-wise chain; force the tile path when fusion has - // attached one. (The tile path applies the chain inline below.) - let post_chain_nonempty = !self.post_element_wise.functions.is_empty(); - let vec4_meta = (graph.device().subgroups_supported() - && storage_datatype == DataTypeEnum::F32) - .then(|| { - if post_chain_nonempty { - return None; - } - build_vec4_rms_norm_meta( - input_view.clone(), - residual_view.clone(), - weight, - bias, - output_view.clone(), - self.eps, - ) - }) - .flatten(); - let selection_ctx = RmsNormSelectionCtx { - vec4_supported: vec4_meta.is_some(), - }; - let variant = select_rms_norm_variant( - rows, - cols, - &selection_ctx, - KernelDeviceCaps::from_device(&graph.device()), - ); - let dispatch_size = [rows, 1, 1]; - let kernel_label = match variant { - RmsNormKernelVariant::Tile => "rms_norm", - RmsNormKernelVariant::Vec4 => "rms_norm_vec4", - }; - let cache_variant = - kernel_backend::KernelVariantKey::with_payload::(|state| { - variant.hash(state); - }); - let cache_key = self.kernel_cache_key_with_dispatch( - cache_variant, - Some(workgroup_shape), - dispatch_size, - inputs, - ); - - // Collect buffers in the SAME order as the IR builder declares - // them (input, residual?, weight, bias?, output), so the IR-build - // closure can be deferred to cache-miss only. - let buffers = collect_rms_buffers(input, residual, weight, bias, output); - - if let Some(meta) = vec4_meta { - let has_residual = residual.is_some(); - let has_bias = bias.is_some(); - kernel_backend::dynamic_kernel_from_ir( - graph.device().kernel_cache(), - kernel_label, - cache_key, - move || { - let vec_layout = tile_ir::Layout::strided( - tile_ir::MemoryLevel::Storage, - tile_ir::Shape::new([1]), - &[1], - ); - let mut kb = tile_ir::KernelBuilder::<()>::new(); - let input_ref = tile_ir::KernelTensorRef::with_offset( - (), - vec_layout.clone(), - meta.input_offset_vec, - ); - let residual_ref = if has_residual { - meta.residual_offset_vec.map(|offset| { - tile_ir::KernelTensorRef::with_offset((), vec_layout.clone(), offset) - }) - } else { - None - }; - let weight_ref = tile_ir::KernelTensorRef::with_offset( - (), - vec_layout.clone(), - meta.weight_offset_vec, - ); - let bias_ref = if has_bias { - meta.bias_offset_vec.map(|offset| { - tile_ir::KernelTensorRef::with_offset((), vec_layout.clone(), offset) - }) - } else { - None - }; - let output_ref = tile_ir::KernelTensorRef::with_offset( - (), - vec_layout, - meta.output_offset_vec, - ); - tile_ir_kernels::rms_norm_vec4( - &mut kb, - tile_ir_kernels::RmsNormVec4 { - input: input_ref, - residual: residual_ref, - weight: weight_ref, - bias: bias_ref, - output: output_ref, - meta, - rows, - }, - )?; - Some(kb.finish().0) - }, - buffers, - dispatch_size, - ) - } else { - let post_chain = self.post_element_wise.clone(); - kernel_backend::dynamic_kernel_from_ir( - graph.device().kernel_cache(), - kernel_label, - cache_key, - || { - build_rms_norm_tile_ir(RmsNormTileIrParams { - input_view, - residual_view, - weight, - bias, - output_view, - eps: self.eps, - post_chain, - storage_datatype, - }) - }, - buffers, - dispatch_size, - ) - } - } - - fn output(&self, _nodes: &ComputeGraphInner, inputs: &[MirValue]) -> MirValue { - inputs.last().unwrap().clone() - } - - fn name(&self) -> String { - let op = if self.residual.is_some() { - "rms_norm_residual" - } else { - "rms_norm" - }; - format!( - "{op}_f32_{}", - self.shape - .iter() - .map(|dim| dim.to_string()) - .collect::>() - .join("x") - ) - } -} - -impl GraphOperation for RmsNormOperation { - fn as_any(&self) -> &dyn Any { - self - } - - fn category(&self) -> &'static str { - "rms_norm" - } - - fn output_layout( - &self, - input_layouts: &FxHashMap, - ) -> Option { - let input_layout = input_layouts.get(&self.input)?; - Some(TensorLayoutInfo::new( - Layout::contiguous(input_layout.shape()), - input_layout.datatype(), - )) - } -} - -fn build_vec4_rms_norm_meta( - input_view: crate::mir::tile_direct::DirectMatrixLayout, - residual_view: Option, - weight: &TensorData, - bias: Option<&TensorData>, - output_view: crate::mir::tile_direct::DirectMatrixLayout, - eps: f32, -) -> Option { - if !input_view.layout.is_affine() - || !output_view.layout.is_affine() - || residual_view - .as_ref() - .is_some_and(|residual| !residual.layout.is_affine()) - || !input_view.cols.is_multiple_of(4) - { - return None; - } - - let [input_row_stride, input_col_stride] = matrix_strides(input_view.layout.affine_strides())?; - let [output_row_stride, output_col_stride] = - matrix_strides(output_view.layout.affine_strides())?; - if input_col_stride != 1 - || output_col_stride != 1 - || !input_view.offset.is_multiple_of(4) - || !output_view.offset.is_multiple_of(4) - || !input_row_stride.is_multiple_of(4) - || !output_row_stride.is_multiple_of(4) - { - return None; - } - - let (residual_offset_vec, residual_row_stride_vec) = if let Some(residual_view) = residual_view - { - let [residual_row_stride, residual_col_stride] = - matrix_strides(residual_view.layout.affine_strides())?; - if residual_col_stride != 1 - || !residual_view.offset.is_multiple_of(4) - || !residual_row_stride.is_multiple_of(4) - { - return None; - } - (Some(residual_view.offset / 4), residual_row_stride / 4) - } else { - (None, 0) - }; - - let weight_stride = *weight.layout().strides().first()?; - if weight.layout().shape() != [input_view.cols as usize] - || weight_stride != 1 - || !weight.layout().offset().is_multiple_of(4) - { - return None; - } - let bias_offset_vec = if let Some(bias) = bias { - let bias_stride = *bias.layout().strides().first()?; - if bias.layout().shape() != [input_view.cols as usize] - || bias_stride != 1 - || !bias.layout().offset().is_multiple_of(4) - { - return None; - } - Some((bias.layout().offset() / 4).try_into().ok()?) - } else { - None - }; - - Some(tile_ir_kernels::RmsNormVec4Meta { - cols: input_view.cols, - cols_vec: input_view.cols / 4, - eps: tile_ir::F32Bits::new(eps), - input_offset_vec: input_view.offset / 4, - input_row_stride_vec: input_row_stride / 4, - residual_offset_vec, - residual_row_stride_vec, - weight_offset_vec: (weight.layout().offset() / 4).try_into().ok()?, - bias_offset_vec, - output_offset_vec: output_view.offset / 4, - output_row_stride_vec: output_row_stride / 4, - }) -} - -fn matrix_strides(strides: Vec) -> Option<[u32; 2]> { - strides.try_into().ok() -} - -struct RmsNormTileIrParams<'a> { - input_view: crate::mir::tile_direct::DirectMatrixLayout, - residual_view: Option, - weight: &'a TensorData, - bias: Option<&'a TensorData>, - output_view: crate::mir::tile_direct::DirectMatrixLayout, - eps: f32, - post_chain: UnaryFunctionChain, - storage_datatype: DataTypeEnum, -} - -fn build_rms_norm_tile_ir(params: RmsNormTileIrParams<'_>) -> Option { - let element = match params.storage_datatype { - DataTypeEnum::F32 => tile_ir::ElementType::F32, - DataTypeEnum::F16 => tile_ir::ElementType::F16, - DataTypeEnum::U32 => return None, - }; - build_rms_norm_tile_ir_typed( - element, - params.input_view, - params.residual_view, - params.weight, - params.bias, - params.output_view, - params.eps, - params.post_chain, - ) -} - -#[allow(clippy::too_many_arguments)] -fn build_rms_norm_tile_ir_typed( - element: tile_ir::ElementType, - input_view: crate::mir::tile_direct::DirectMatrixLayout, - residual_view: Option, - weight: &TensorData, - bias: Option<&TensorData>, - output_view: crate::mir::tile_direct::DirectMatrixLayout, - eps: f32, - post_chain: UnaryFunctionChain, -) -> Option { - let zero_fill = rms_norm_zero_fill(element); - let rows = input_view.rows; - let cols = input_view.cols; - let input_storage_layout = input_view.layout.clone(); - let residual_storage_layout = residual_view - .as_ref() - .map(|residual_view| residual_view.layout.clone()); - let residual_offset = residual_view.as_ref().map(|residual| residual.offset); - let output_storage_layout = output_view.layout.clone(); - let weight_layout = vector_as_row_layout(weight.layout())?; - let bias_layout = match bias { - Some(bias) => Some(vector_as_row_layout(bias.layout())?), - None => None, - }; - let weight_offset = weight.layout().offset().try_into().ok()?; - let bias_offset = match bias { - Some(bias) => Some(bias.layout().offset().try_into().ok()?), - None => None, - }; - - Some(tile_ir::tile::build(move |phase| { - let input = tile_storage_read_with_direct_layout_typed( - phase, - element, - crate::mir::tile_direct::DirectMatrixLayout { - rows, - cols, - offset: input_view.offset, - layout: input_storage_layout, - }, - ); - let residual = residual_storage_layout.map(|layout| { - tile_storage_read_with_direct_layout_typed( - phase, - element, - crate::mir::tile_direct::DirectMatrixLayout { - rows, - cols, - offset: residual_offset.expect("residual offset exists with layout"), - layout, - }, - ) - }); - let weight = phase.storage_read_with_layout_offset(element, weight_layout, weight_offset); - let bias = bias_layout.map(|layout| { - phase.storage_read_with_layout_offset( - element, - layout, - bias_offset.expect("bias offset exists when bias layout exists"), - ) - }); - let output = tile_storage_write_with_direct_layout_typed( - phase, - element, - crate::mir::tile_direct::DirectMatrixLayout { - rows, - cols, - offset: output_view.offset, - layout: output_storage_layout, - }, - ); - - // Values are loaded in their storage element type and accumulated in - // f32. The cast is a no-op when the storage element is already f32. - let into_accum = |value: tile_ir::tile::Tile| -> tile_ir::tile::Tile { - if element == tile_ir::ElementType::F32 { - value - } else { - value.cast(tile_ir::ElementType::F32) - } - }; - let from_accum = |value: tile_ir::tile::Tile| -> tile_ir::tile::Tile { - if element == tile_ir::ElementType::F32 { - value - } else { - value.cast(element) - } - }; - - let chunks = cols.div_ceil(BLOCK as u32); - phase.program_grid(BLOCK as u32, [rows, 1, 1], |program| { - let row = program.program_id(tile_ir::WorkgroupAxis::X); - let lane = program.lane(); - let [sum_square] = program.fold( - tile_ir::tile::range(chunks), - [tile_ir::tile::Tile::literal(0.0f32)], - |program, loop_index, [acc]| { - let reduce_col = loop_index * BLOCK as u32 + lane.clone(); - let reduce_mask = reduce_col.clone().lt(cols); - let mut value = into_accum(program.load( - input.at((&row, &reduce_col)), - reduce_mask.clone(), - zero_fill, - )); - if let Some(residual) = &residual { - value = value - + into_accum(program.load( - residual.at((&row, &reduce_col)), - reduce_mask, - zero_fill, - )); - } - [acc + value.clone() * value] - }, - ); - let total_sum = program.group_reduce_sum(BLOCK as u32, sum_square); - let rms = (total_sum / tile_ir::tile::Tile::literal(cols as f32) - + tile_ir::tile::Tile::literal(eps)) - .unary(tile_ir::TileUnaryOp::Sqrt); - for chunk in 0..chunks { - let col = lane.clone() + chunk * BLOCK as u32; - let mask = col.lt(cols); - let mut value = - into_accum(program.load(input.at((&row, &col)), mask.clone(), zero_fill)); - if let Some(residual) = &residual { - value = value - + into_accum(program.load( - residual.at((&row, &col)), - mask.clone(), - zero_fill, - )); - } - let weight = - into_accum(program.load(weight.at((0, &col)), mask.clone(), zero_fill)); - let mut normalized = value / rms.clone() * weight; - if let Some(bias) = &bias { - let bias_value = - into_accum(program.load(bias.at((0, &col)), mask.clone(), zero_fill)); - normalized = normalized + bias_value; - } - // Apply any fused post-element-wise chain in-register before - // the store. Empty chain is a no-op (the loop below short- - // circuits and returns the value unchanged). - if !post_chain.functions.is_empty() { - let (val, _) = - apply_unary_function_chain(normalized, DataTypeEnum::F32, &post_chain) - .expect("rms_norm post-chain validated at fuse time"); - normalized = val; - } - program.store(output.at((&row, col)), from_accum(normalized), mask); - } - }); - })) -} - -/// The storage-typed zero used as the masked-out fill for rms-norm loads. -fn rms_norm_zero_fill(element: tile_ir::ElementType) -> tile_ir::TileLiteral { - match element { - tile_ir::ElementType::F16 => tile_ir::TileLiteral::F16(0), - _ => tile_ir::TileLiteral::f32(0.0), - } -} - -fn vector_as_row_layout(layout: &crate::Layout) -> Option { - let shape = layout.shape(); - let strides = layout.strides(); - if shape.len() != 1 { - return None; - } - Some(tile_ir::Layout::strided( - tile_ir::MemoryLevel::Storage, - tile_ir::Shape::new([1, (*shape.first()?).try_into().ok()?]), - &[0, (*strides.first()?).try_into().ok()?], - )) -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::{ - Device, Tensor, - kernel_selection::{CooperativeMatrixCaps, assert_selector_generates}, - }; - - fn caps() -> KernelDeviceCaps { - KernelDeviceCaps { - cooperative_matrix: CooperativeMatrixCaps::default(), - ..KernelDeviceCaps::test_caps() - } - } - - #[test] - fn rms_norm_selector_generates_each_variant() { - let selector = rms_norm_selector(); - let cases = [ - ( - RmsNormKernelVariant::Vec4, - RmsNormSelectionCtx { - vec4_supported: true, - }, - ), - ( - RmsNormKernelVariant::Tile, - RmsNormSelectionCtx { - vec4_supported: false, - }, - ), - ]; - assert_selector_generates( - &selector, - cases.map(|(variant, ctx)| (variant, ctx, caps())), - ); - } - - #[test] - fn rms_norm_direct_matches_reference() { - pollster::block_on(async { - let Ok(device) = Device::new().await else { - return; - }; - - let input = Tensor::new::(&device, &vec![vec![1.0f32, 2.0, 3.0, 4.0]]); - let weight = Tensor::new::(&device, &vec![0.5f32, 1.0, 1.5, 2.0]); - let output = input.try_rms_norm_direct(&weight, None, 1e-5).unwrap(); - let output = output.as_slice::<2, f32>().await.unwrap(); - - let mean_square = (1.0 + 4.0 + 9.0 + 16.0) / 4.0; - let rms = f32::sqrt(mean_square + 1e-5); - let expected = [1.0 / rms * 0.5, 2.0 / rms, 3.0 / rms * 1.5, 4.0 / rms * 2.0]; - - for (i, expected) in expected.into_iter().enumerate() { - let actual = output[[0, i]]; - assert!((actual - expected).abs() < 1e-5, "{actual} != {expected}"); - } - }); - } -} diff --git a/fusor-ml/core/src/mir/workgroup_shape.rs b/fusor-ml/core/src/mir/workgroup_shape.rs index 534a2eb1c..89ddcf537 100644 --- a/fusor-ml/core/src/mir/workgroup_shape.rs +++ b/fusor-ml/core/src/mir/workgroup_shape.rs @@ -1,4 +1,5 @@ -// The shape of the workgroup. [x, y, z] where their product is <= 256. +// The shape of the workgroup. [x, y, z] where their product is bounded by +// the device's workgroup limits. // // Kernels can be fused if their workgroup shape can be coerced. Coercion can happen if // the biggest linearized workgroup shape is a multiple of all smaller workgroup shapes. @@ -8,8 +9,6 @@ use parking_lot::RwLock; use rustc_hash::FxBuildHasher; use std::{num::NonZeroUsize, sync::OnceLock}; -const MAX_WORKGROUP_SIZE: u32 = 256; - #[derive(Debug, Clone, Copy)] pub struct WorkgroupShape { shape: [u32; 3], @@ -27,10 +26,6 @@ impl WorkgroupShape { x > 0 && y > 0 && z > 0, "Workgroup shape dimensions must be greater than zero" ); - assert!( - x * y * z <= 256, - "Workgroup shape dimensions must be less than or equal to 256" - ); Self { shape: [x, y, z] } } @@ -88,24 +83,38 @@ impl WorkgroupShapeConstraints { }) } - fn possible(&self) -> impl Iterator { - possible_workgroup_shapes().filter(move |shape| self.is_valid(shape)) + fn possible(&self, limits: &wgpu::Limits) -> impl Iterator { + possible_workgroup_shapes(limits).filter(move |shape| self.is_valid(shape)) } - pub(crate) fn solve(&self, max_subgroup_size: u32) -> Option { - static CACHE: OnceLock< - RwLock, FxBuildHasher>>, - > = OnceLock::new(); + pub(crate) fn solve( + &self, + max_subgroup_size: u32, + limits: &wgpu::Limits, + ) -> Option { + type CacheKey = (WorkgroupShapeConstraints, u32, [u32; 4]); + static CACHE: OnceLock, FxBuildHasher>>> = + OnceLock::new(); let cache = CACHE.get_or_init(|| { RwLock::new(LruCache::with_hasher( const { NonZeroUsize::new(2048).unwrap() }, Default::default(), )) }); + let key = ( + self.clone(), + max_subgroup_size, + [ + limits.max_compute_workgroup_size_x, + limits.max_compute_workgroup_size_y, + limits.max_compute_workgroup_size_z, + limits.max_compute_invocations_per_workgroup, + ], + ); let mut write = cache.write(); - *write.get_or_insert_ref(self, || { + *write.get_or_insert_ref(&key, || { // Find the smallest valid shape that matches the max subgroup size - self.possible().min_by_key(|shape| { + self.possible(limits).min_by_key(|shape| { let linearized = shape.linearized(); (linearized as i64) + if max_subgroup_size == 0 || shape.x() % max_subgroup_size == 0 { @@ -118,10 +127,16 @@ impl WorkgroupShapeConstraints { } } -fn possible_workgroup_shapes() -> impl Iterator { - (1..=MAX_WORKGROUP_SIZE).flat_map(move |x| { - (1..=(MAX_WORKGROUP_SIZE / x)).flat_map(move |y| { - (1..=(MAX_WORKGROUP_SIZE / (x * y))).map(move |z| WorkgroupShape::new(x, y, z)) +/// Every workgroup shape the device can run: each dimension within its +/// per-axis limit and the product within the invocation limit. +fn possible_workgroup_shapes(limits: &wgpu::Limits) -> impl Iterator { + let total = limits.max_compute_invocations_per_workgroup; + let max_x = limits.max_compute_workgroup_size_x.min(total); + let max_y = limits.max_compute_workgroup_size_y.min(total); + let max_z = limits.max_compute_workgroup_size_z.min(total); + (1..=max_x).flat_map(move |x| { + (1..=max_y.min(total / x)).flat_map(move |y| { + (1..=max_z.min(total / (x * y))).map(move |z| WorkgroupShape::new(x, y, z)) }) }) } @@ -165,23 +180,35 @@ mod tests { const TEST_MAX_SUBGROUP_SIZE: u32 = 64; + fn test_limits(size: u32) -> wgpu::Limits { + wgpu::Limits { + max_compute_workgroup_size_x: size, + max_compute_workgroup_size_y: size, + max_compute_workgroup_size_z: size, + max_compute_invocations_per_workgroup: size, + ..wgpu::Limits::default() + } + } + #[test] fn test_all_possible_workgroup_shapes() { - assert_eq!(possible_workgroup_shapes().count(), 5136); + assert_eq!(possible_workgroup_shapes(&test_limits(256)).count(), 5136); + assert_eq!(possible_workgroup_shapes(&test_limits(1024)).count(), 30343); } #[test] fn test_workgroup_shape_constraints() { + let limits = test_limits(256); let mut constraints = WorkgroupShapeConstraints::new(); constraints.add_constraint(0, Constraint::Equals(4)); constraints.add_constraint(1, Constraint::LessThan(3)); - for shape in constraints.possible() { + for shape in constraints.possible(&limits) { assert_eq!(shape.shape()[0], 4); assert!(shape.shape()[1] < 3); } - let valid_shape = constraints.solve(TEST_MAX_SUBGROUP_SIZE).unwrap(); + let valid_shape = constraints.solve(TEST_MAX_SUBGROUP_SIZE, &limits).unwrap(); assert_eq!(valid_shape.shape(), [4, 1, 1]); assert_eq!(valid_shape.linearized(), 4); } diff --git a/fusor-ml/core/src/nary_direct.rs b/fusor-ml/core/src/nary_direct.rs index ac644bd2b..3c33a0dcd 100644 --- a/fusor-ml/core/src/nary_direct.rs +++ b/fusor-ml/core/src/nary_direct.rs @@ -9,8 +9,12 @@ use crate::{ operation::Operation, workgroup_shape::WorkgroupShape, }, - nary_wise::{NaryExpr, NaryFunction, NaryOp, NaryOperation, NaryScalar, UnaryFunctionChain}, + nary_wise::{ + ElementwiseOperation, NaryExpr, NaryFunction, NaryOp, NaryScalar, UnaryFunctionChain, + }, + quantized::QMatrix, tensor::{DataTypeEnum, TensorData}, + visit_tiled::MaybeQData, }; const BLOCK: usize = 256; @@ -19,7 +23,7 @@ const SMALL_BLOCK: usize = 1; struct NaryDirectKernelVariant; pub(crate) fn build_nary_direct_kernel( - operation: &NaryOperation, + operation: &ElementwiseOperation, graph: &crate::compute_graph::ComputeGraphInner, workgroup_shape: &WorkgroupShape, inputs: &[MirValue], @@ -28,7 +32,7 @@ pub(crate) fn build_nary_direct_kernel( } pub(crate) fn build_nary_direct_kernel_to_output( - operation: &NaryOperation, + operation: &ElementwiseOperation, graph: &crate::compute_graph::ComputeGraphInner, workgroup_shape: &WorkgroupShape, inputs: &[MirValue], @@ -44,29 +48,52 @@ pub(crate) fn build_nary_direct_kernel_to_output( } fn build_nary_direct_kernel_with_output_index( - operation: &NaryOperation, + operation: &ElementwiseOperation, graph: &crate::compute_graph::ComputeGraphInner, workgroup_shape: &WorkgroupShape, inputs: &[MirValue], forced_output_index: Option, ) -> Option { let output_index = forced_output_index.or_else(|| operation.output_tensor_index(inputs))?; - let tensors = inputs + let values = inputs .iter() - .map(|input| input.as_tensor().cloned()) + .map(|input| MaybeQData::try_from(input.clone()).ok()) .collect::>>()?; - tensors.get(output_index)?; + let MaybeQData::Tensor(_) = values.get(output_index)? else { + return None; + }; - if tensors - .iter() - .any(|tensor| tensor.datatype() == DataTypeEnum::F16 && !graph.device().f16_supported()) - { + if values.iter().any(|value| { + let datatype = match value { + MaybeQData::Tensor(tensor) => tensor.datatype(), + MaybeQData::QMatrix(matrix) => match matrix.datatype() { + fusor_gguf::GgmlType::F16 => DataTypeEnum::F16, + _ => return false, + }, + }; + datatype == DataTypeEnum::F16 && !graph.device().f16_supported() + }) { return None; } let total_elements = total_elements(&operation.shape)?; + let plan = plan_nary_tiling(operation, &graph.device(), &values, output_index); + if let Some(plan) = &plan + && std::env::var_os("FUSOR_TRACE_REDUCE_TILED").is_some() + { + eprintln!( + "nary_tiled dim={} invariant={:?} threads={} shape={:?}", + plan.dim, plan.invariant, plan.total_threads, operation.shape, + ); + } let small_dispatch = total_elements < BLOCK as u32; - let dispatch_size = if small_dispatch { + let dispatch_size = if let Some(plan) = &plan { + let max_per_dim = graph.device().limits().max_compute_workgroups_per_dimension; + crate::visit_tiled::distribute_workgroups( + plan.total_threads.div_ceil(BLOCK as u32), + max_per_dim, + ) + } else if small_dispatch { [total_elements, 1, 1] } else { operation.dispatch_size(workgroup_shape, inputs) @@ -74,6 +101,10 @@ fn build_nary_direct_kernel_with_output_index( let variant = kernel_backend::KernelVariantKey::with_payload::(|state| { output_index.hash(state); + if let Some(plan) = &plan { + plan.dim.hash(state); + plan.invariant.hash(state); + } }); let cache_key = operation.kernel_cache_key_with_dispatch( variant, @@ -86,33 +117,218 @@ fn build_nary_direct_kernel_with_output_index( } else { format!("nary_direct_out_{output_index}") }; - let buffers = tensors - .iter() - .map(|tensor| tensor.buffer().clone()) - .collect::>(); - kernel_backend::dynamic_kernel_from_ir( + kernel_backend::run_kernel( graph.device().kernel_cache(), name, cache_key, - move || { - if small_dispatch { - build_nary_tile_ir::(operation, &tensors, output_index, dispatch_size) + dispatch_size, + |kb| { + if let Some(plan) = &plan { + build_nary_tiled_ir(operation, &values, output_index, plan, dispatch_size, kb) + } else if small_dispatch { + build_nary_tile_ir::( + operation, + &values, + output_index, + dispatch_size, + kb, + ) } else { - build_nary_tile_ir::(operation, &tensors, output_index, dispatch_size) + build_nary_tile_ir::(operation, &values, output_index, dispatch_size, kb) } }, - buffers, - dispatch_size, ) } +/// Outputs per thread along the tiled dim of a reuse-tiled elementwise +/// kernel. +const NARY_TM: u32 = 4; +/// Floor on post-tiling thread count: trading threads for register reuse +/// must leave the device saturated. +const MIN_TILED_THREADS: u32 = 65536; + +/// A register-reuse tiling for an elementwise kernel: each thread covers +/// `NARY_TM` outputs along `dim`, loading the inputs that are invariant +/// along `dim` once instead of per output. +struct NaryTilePlan { + dim: usize, + /// Per input: invariant along `dim` (its loads hoist out of the run). + invariant: Vec, + /// Per input: index-space dim read by each input dimension. + dims: Vec>, + /// The output shape with `dim` divided by `NARY_TM`. + thread_shape: Vec, + total_threads: u32, +} + +/// Tile only when the hoisted loads buy real bandwidth: the invariant +/// inputs must exceed the device's cache-residency threshold (cache-resident +/// re-reads are free, and the tiling costs thread-level parallelism), the +/// tiled dim must not be the innermost output dim (thread-local runs there +/// break inter-thread store coalescing), and enough threads must remain. +fn plan_nary_tiling( + operation: &ElementwiseOperation, + device: &crate::Device, + values: &[MaybeQData], + output_index: usize, +) -> Option { + let input_count = operation.inputs.len(); + let rank = operation.shape.len(); + if rank < 2 || output_index != input_count || values.len() != input_count + 1 { + return None; + } + // Quantized inputs decode through a dedicated load path the value-tile + // evaluator can't hoist. + if values[..input_count] + .iter() + .any(|value| matches!(value, MaybeQData::QMatrix(_))) + { + return None; + } + let metas: Vec = values[..input_count] + .iter() + .map(|value| match value { + MaybeQData::Tensor(tensor) => TensorMeta::new(tensor), + MaybeQData::QMatrix(_) => None, + }) + .collect::>()?; + let access = + crate::access_analysis::InputAccesses::collect(&operation.expression, input_count, &metas)?; + + let mut best: Option<(u64, usize)> = None; + for dim in 0..rank.saturating_sub(1) { + if operation.shape[dim] < NARY_TM as usize { + continue; + } + let invariant_bytes: u64 = (0..input_count) + .filter(|&i| !access.depends_on(i, dim)) + .map(|i| crate::reduce_tiled::input_allocation_bytes(&metas[i], &values[i])) + .sum(); + if invariant_bytes < device.last_level_cache_bytes() { + continue; + } + if best + .as_ref() + .is_none_or(|(best_bytes, _)| invariant_bytes > *best_bytes) + { + best = Some((invariant_bytes, dim)); + } + } + let (_, dim) = best?; + let invariant: Vec = (0..input_count) + .map(|i| !access.depends_on(i, dim)) + .collect(); + + let mut thread_shape = operation.shape.to_vec(); + thread_shape[dim] = thread_shape[dim].div_ceil(NARY_TM as usize); + let total_threads = total_elements(&thread_shape)?; + if total_threads < MIN_TILED_THREADS { + return None; + } + Some(NaryTilePlan { + dim, + invariant, + dims: access.dims, + thread_shape, + total_threads, + }) +} + +fn build_nary_tiled_ir( + operation: &ElementwiseOperation, + values: &[MaybeQData], + output_index: usize, + plan: &NaryTilePlan, + dispatch_size: [u32; 3], + kb: &mut tile_ir::KernelBuilder>, +) -> Option<()> { + let mut storages = Vec::with_capacity(values.len()); + let mut metas = Vec::with_capacity(values.len()); + for (binding, value) in values.iter().enumerate() { + let (storage, meta) = declare_value(kb, value, binding == output_index)?; + storages.push(storage); + metas.push(meta); + } + let extent = operation.shape[plan.dim] as u32; + let total_threads = plan.total_threads; + let input_count = operation.inputs.len(); + + let input_coords = |coords: &[tile_ir::tile::Tile], i: usize| -> Vec { + plan.dims[i].iter().map(|&d| coords[d].clone()).collect() + }; + + kb.program() + .program_grid(BLOCK as u32, dispatch_size, |program| { + let lane = program.lane(); + let group = linear_group(program, dispatch_size); + let flat = group * BLOCK as u32 + lane; + let in_bounds = flat.clone().lt(total_threads); + let mut coords = output_dims_from_flat(flat, &plan.thread_shape); + let base = program.bind(coords[plan.dim].clone() * NARY_TM); + + // Invariant loads hoist out of the per-output run; the base + // coordinate is always in range for an in-bounds thread. + coords[plan.dim] = base.clone(); + let hoisted: Vec> = (0..input_count) + .map(|i| { + if !plan.invariant[i] { + return None; + } + let index = layout_index(&metas[i], &input_coords(&coords, i)); + let loaded = storages[i].load(program, index, in_bounds.clone()); + let native = match loaded { + ValueTile::F32(v) | ValueTile::F16(v) | ValueTile::U32(v) => v, + ValueTile::Bool(_) => unreachable!("tensor inputs are f32/f16/u32"), + }; + Some(( + value_tile_of(metas[i].datatype, program.bind(native)), + metas[i].datatype, + )) + }) + .collect(); + + for j in 0..NARY_TM { + let coord = base.clone() + j; + let in_bounds_j = in_bounds.clone().and(coord.clone().lt(extent)); + coords[plan.dim] = coord; + let slot_values: Vec<(ValueTile, DataTypeEnum)> = (0..input_count) + .map(|i| match &hoisted[i] { + Some(value) => value.clone(), + None => { + let index = layout_index(&metas[i], &input_coords(&coords, i)); + ( + storages[i].load(program, index, in_bounds_j.clone()), + metas[i].datatype, + ) + } + }) + .collect(); + let (value, value_ty) = + eval_nary_expr_on_value_tiles(&operation.expression, &slot_values); + let value = value.cast_to(operation.output_datatype); + debug_assert_eq!(value_ty, operation.output_datatype); + let output_index_value = layout_index(&metas[output_index], &coords); + storages[output_index].store(program, output_index_value, value, in_bounds_j); + } + }); + Some(()) +} + +fn value_tile_of(datatype: DataTypeEnum, value: tile_ir::tile::Tile) -> ValueTile { + match datatype { + DataTypeEnum::F32 => ValueTile::F32(value), + DataTypeEnum::F16 => ValueTile::F16(value), + DataTypeEnum::U32 => ValueTile::U32(value), + } +} + fn total_elements(shape: &[usize]) -> Option { shape .iter() .try_fold(1u32, |acc, dim| acc.checked_mul((*dim).try_into().ok()?)) } -impl NaryOperation { +impl ElementwiseOperation { pub(crate) fn output_tensor_index(&self, inputs: &[MirValue]) -> Option { inputs.len().checked_sub(1) } @@ -247,6 +463,10 @@ pub(crate) enum Storage2 { F32(tile_ir::tile::Storage), F16(tile_ir::tile::Storage), U32(tile_ir::tile::Storage), + /// A block-quantized matrix loaded per element (block decode in the + /// lowerer). Loads through [`eval_nary_expr`]'s dedicated path — `load` + /// on this variant is unreachable. + Quantized(tile_ir::QuantizedMatrix), } impl Storage2 { @@ -274,6 +494,9 @@ impl Storage2 { mask, zero_literal(DataTypeEnum::U32), )), + Self::Quantized(_) => { + unreachable!("quantized inputs load through the row/col path") + } } } @@ -300,52 +523,93 @@ impl Storage2 { program.store(storage.at((0u32, index)), value, mask); } } + Self::Quantized(_) => { + unreachable!("quantized matrices are read-only inputs") + } } } } -pub(crate) fn declare_storage( - phase: &mut tile_ir::tile::Program, - meta: &TensorMeta, +/// Declare one n-ary input or output as a kernel binding: dense tensors and +/// dense-storage matrices as flat strided reads/writes, block-quantized +/// matrices through the format-aware quantized binding. +pub(crate) fn declare_value( + kb: &mut tile_ir::KernelBuilder>, + value: &MaybeQData, write: bool, -) -> Storage2 { - let layout = tile_ir::Layout::strided( - tile_ir::MemoryLevel::Storage, - tile_ir::Shape::new([1, meta.allocation_len]), - &[0, 1], - ); - let element = datatype_element(meta.datatype); - let storage = if write { - phase.storage_write_with_layout_offset(element, layout, 0) - } else { - phase.storage_read_with_layout_offset(element, layout, 0) +) -> Option<(Storage2, TensorMeta)> { + let declare_dense = |kb: &mut tile_ir::KernelBuilder>, + buffer: std::sync::Arc, + meta: &TensorMeta| { + let tensor = tile_ir::KernelTensorRef::new(buffer, flat_layout(meta.allocation_len)); + let element = datatype_element(meta.datatype); + let storage = if write { + kb.write(element, tensor) + } else { + kb.read(element, tensor) + }; + match meta.datatype { + DataTypeEnum::F32 => Storage2::F32(storage), + DataTypeEnum::F16 => Storage2::F16(storage), + DataTypeEnum::U32 => Storage2::U32(storage), + } }; - match meta.datatype { - DataTypeEnum::F32 => Storage2::F32(storage), - DataTypeEnum::F16 => Storage2::F16(storage), - DataTypeEnum::U32 => Storage2::U32(storage), + match value { + MaybeQData::Tensor(tensor) => { + let meta = TensorMeta::new(tensor)?; + let storage = declare_dense(kb, tensor.buffer().clone(), &meta); + Some((storage, meta)) + } + MaybeQData::QMatrix(matrix) => { + if write { + return None; + } + let meta = TensorMeta::for_matrix(matrix)?; + match crate::quantized::dequantize::quant_format(matrix) { + Some(format) => { + // `QuantizedMatrix::rows` is the dense row *length* (the + // contiguous K axis); `cols` counts the rows. + let row_len = *matrix.shape().last()? as u32; + let row_count = matrix.shape()[..matrix.shape().len() - 1] + .iter() + .try_fold(1u32, |acc, dim| acc.checked_mul((*dim).try_into().ok()?))?; + let storage = fusor_tile_ir_kernels::quantized_matrix_for( + kb, + matrix.buffer().clone(), + format, + row_len, + row_count, + ); + Some((Storage2::Quantized(storage), meta)) + } + // Dense f16/f32 storage reads like a plain row-major tensor. + None => { + let storage = declare_dense(kb, matrix.buffer().clone(), &meta); + Some((storage, meta)) + } + } + } } } fn build_nary_tile_ir( - operation: &NaryOperation, - tensors: &[TensorData], + operation: &ElementwiseOperation, + values: &[MaybeQData], output_index: usize, dispatch_size: [u32; 3], -) -> Option { + kb: &mut tile_ir::KernelBuilder>, +) -> Option<()> { let total_elements = total_elements(&operation.shape)?; - let tensor_metas = tensors - .iter() - .map(TensorMeta::new) - .collect::>>()?; - - Some(tile_ir::tile::build(move |phase| { - let storages = tensor_metas - .iter() - .enumerate() - .map(|(binding, meta)| declare_storage(phase, meta, binding == output_index)) - .collect::>(); + let mut storages = Vec::with_capacity(values.len()); + let mut tensor_metas = Vec::with_capacity(values.len()); + for (binding, value) in values.iter().enumerate() { + let (storage, meta) = declare_value(kb, value, binding == output_index)?; + storages.push(storage); + tensor_metas.push(meta); + } + { + let phase = kb.program(); phase.program_grid(BLOCK_SIZE as u32, dispatch_size, |program| { let lane = program.lane(); let group = linear_group(program, dispatch_size); @@ -359,25 +623,31 @@ fn build_nary_tile_ir( &storages, &tensor_metas, in_bounds.clone(), + &[], ); let value = value.cast_to(operation.output_datatype); debug_assert_eq!(value_ty, operation.output_datatype); let output_index_value = layout_index(&tensor_metas[output_index], &dims); storages[output_index].store(program, output_index_value, value, in_bounds); }); - })) + } + Some(()) } -fn eval_nary_expr( +/// Evaluate an expression at one coordinate. Slot indices beyond the bound +/// storages resolve from `extras` — per-row scalars a surrounding kernel has +/// already computed (row-program phases). +pub(crate) fn eval_nary_expr( program: &mut tile_ir::tile::TileBlock<'_>, expr: &NaryExpr, dims: &[tile_ir::tile::Tile], storages: &[Storage2], metas: &[TensorMeta], mask: tile_ir::tile::Mask, + extras: &[(ValueTile, DataTypeEnum)], ) -> (ValueTile, DataTypeEnum) { if let Some(value) = - eval_associative_binary_tree(program, expr, dims, storages, metas, mask.clone()) + eval_associative_binary_tree(program, expr, dims, storages, metas, mask.clone(), extras) { return value; } @@ -389,25 +659,39 @@ fn eval_nary_expr( .zip(&function.input_types) .map(|(child, expected)| { let (value, ty) = - eval_nary_expr(program, child, dims, storages, metas, mask.clone()); + eval_nary_expr(program, child, dims, storages, metas, mask.clone(), extras); (value.cast_to(*expected), ty) }) .collect::>(); (emit_function(function, &mut values), function.output_type) } NaryExpr::IndexedInput { input_idx, indices } => { + if *input_idx >= storages.len() { + return extras[*input_idx - storages.len()].clone(); + } let meta = &metas[*input_idx]; let coords = indices .iter() .map(|index| { let (value, _) = - eval_nary_expr(program, index, dims, storages, metas, mask.clone()); + eval_nary_expr(program, index, dims, storages, metas, mask.clone(), extras); match value.cast_to(DataTypeEnum::U32) { ValueTile::U32(value) => value, _ => unreachable!(), } }) .collect::>(); + if let Storage2::Quantized(matrix) = &storages[*input_idx] { + // Block-quantized loads address by (position within the row, + // which row): the position is the last coordinate, and the + // row is the row-major flattening of every leading + // coordinate (the meta carries those strides with a 0 in the + // final slot). + let along_row = coords.last().cloned().unwrap_or_else(|| tile_u32(0)); + let which_row = layout_index(meta, &coords); + let value = program.load_quantized(matrix, along_row, which_row, mask, 0.0); + return (ValueTile::F32(value), DataTypeEnum::F32); + } let index = layout_index(meta, &coords); let value = storages[*input_idx].load(program, index, mask); (value, meta.datatype) @@ -417,6 +701,7 @@ fn eval_nary_expr( } } +#[allow(clippy::too_many_arguments)] fn eval_associative_binary_tree( program: &mut tile_ir::tile::TileBlock<'_>, expr: &NaryExpr, @@ -424,14 +709,15 @@ fn eval_associative_binary_tree( storages: &[Storage2], metas: &[TensorMeta], mask: tile_ir::tile::Mask, + extras: &[(ValueTile, DataTypeEnum)], ) -> Option<(ValueTile, DataTypeEnum)> { let (op, datatype, terms) = flatten_associative_binary_terms(expr)?; let mut terms = terms.into_iter(); let first = terms.next()?; - let (value, _) = eval_nary_expr(program, first, dims, storages, metas, mask.clone()); + let (value, _) = eval_nary_expr(program, first, dims, storages, metas, mask.clone(), extras); let mut value = value.cast_to(datatype); for term in terms { - let (rhs, _) = eval_nary_expr(program, term, dims, storages, metas, mask.clone()); + let (rhs, _) = eval_nary_expr(program, term, dims, storages, metas, mask.clone(), extras); value = value.binary(op, rhs.cast_to(datatype)); } @@ -572,6 +858,11 @@ fn emit_function(function: &NaryFunction, values: &mut [(ValueTile, DataTypeEnum NaryOp::Acosh => values[0].0.clone().unary(tile_ir::TileUnaryOp::Acosh), NaryOp::Atanh => values[0].0.clone().unary(tile_ir::TileUnaryOp::Atanh), NaryOp::Abs => values[0].0.clone().unary(tile_ir::TileUnaryOp::Abs), + NaryOp::LessEqual => values[0].0.clone().compare( + tile_ir::TileCompareOp::Le, + values[1].0.clone(), + function.output_type, + ), NaryOp::AddConst(scalar) => values[0].0.clone().binary( tile_ir::TileBinaryOp::Add, tile_literal(scalar).cast_to(values[0].1), @@ -646,7 +937,7 @@ fn emit_function(function: &NaryFunction, values: &mut [(ValueTile, DataTypeEnum } } -fn eval_nary_expr_on_value_tiles( +pub(crate) fn eval_nary_expr_on_value_tiles( expr: &NaryExpr, inputs: &[(ValueTile, DataTypeEnum)], ) -> (ValueTile, DataTypeEnum) { @@ -853,6 +1144,51 @@ pub(crate) struct TensorMeta { } impl TensorMeta { + /// Metadata for a quantized-or-dense matrix input. Quantized matrices + /// load by (row, col), so the strides flatten the leading dims row-major + /// and zero out the final (column) slot; dense-storage matrices read as + /// plain row-major tensors of their storage element type. + pub(crate) fn for_matrix(matrix: &QMatrix) -> Option { + let shape: Vec = matrix + .shape() + .iter() + .copied() + .map(u32::try_from) + .collect::>() + .ok()?; + let allocation_len = shape + .iter() + .try_fold(1u32, |acc, dim| acc.checked_mul(*dim))?; + let quantized = crate::quantized::dequantize::quant_format(matrix).is_some(); + let datatype = if quantized { + DataTypeEnum::F32 + } else { + match matrix.datatype() { + fusor_gguf::GgmlType::F32 => DataTypeEnum::F32, + fusor_gguf::GgmlType::F16 => DataTypeEnum::F16, + _ => return None, + } + }; + let mut strides = vec![0u32; shape.len()]; + let mut acc = 1u32; + let stride_dims = if quantized { + shape.len().saturating_sub(1) + } else { + shape.len() + }; + for dim in (0..stride_dims).rev() { + strides[dim] = acc; + acc = acc.checked_mul(shape[dim])?; + } + Some(Self { + datatype, + shape, + strides, + offset: 0, + allocation_len, + }) + } + pub(crate) fn new(tensor: &TensorData) -> Option { Some(Self { datatype: tensor.datatype(), diff --git a/fusor-ml/core/src/nary_wise.rs b/fusor-ml/core/src/nary_wise.rs index e3455e9c1..859799931 100644 --- a/fusor-ml/core/src/nary_wise.rs +++ b/fusor-ml/core/src/nary_wise.rs @@ -82,6 +82,10 @@ pub(crate) enum NaryOp { Abs, ApproximateExp, LessApproximateExp, + /// Binary `a <= b` comparison producing 1/0 in the output type. Used for + /// index-dependent masking (e.g. composed causal attention compares the + /// kv position against the query position). + LessEqual, AddConst(NaryScalar), SubConst(NaryScalar), RSubConst(NaryScalar), @@ -195,7 +199,7 @@ impl UnaryFunctionChain { } } -/// Result of extracting a unary function chain from an NaryOperation. +/// Result of extracting a unary function chain from an ElementwiseOperation. /// Used by the resolver to fuse unary ops into reduce/matmul/dequantize. pub(crate) struct ExtractedUnaryChain { pub(crate) value: crate::compute_graph::NodeIndex, @@ -459,7 +463,7 @@ impl NaryExpr { /// N-ary operation combining multiple inputs with arbitrary operations. /// Can fuse chains of element-wise and pair-wise operations into a single kernel. #[derive(Clone, Debug)] -pub(crate) struct NaryOperation { +pub(crate) struct ElementwiseOperation { /// Input tensors (leaves of expression tree) pub(crate) inputs: Vec, /// Expression tree describing computation (includes all operations) @@ -468,8 +472,8 @@ pub(crate) struct NaryOperation { pub(crate) output_datatype: DataTypeEnum, } -impl NaryOperation { - /// Attempt to extract a unary function chain from this NaryOperation. +impl ElementwiseOperation { + /// Attempt to extract a unary function chain from this ElementwiseOperation. /// This will only succeed if there is only a single input to the operation. pub(crate) fn try_extract_unary_chain(&self) -> Option { if self.inputs.len() != 1 { @@ -477,7 +481,7 @@ impl NaryOperation { } let mut functions = Vec::new(); - if !Self::collect_unary_chain(&self.expression, &mut functions)? { + if !Self::collect_unary_chain(&self.expression, &mut functions, self.shape.len())? { return None; } if functions.is_empty() { @@ -502,16 +506,28 @@ impl NaryOperation { }) } - fn collect_unary_chain(expr: &NaryExpr, functions: &mut Vec) -> Option { + fn collect_unary_chain( + expr: &NaryExpr, + functions: &mut Vec, + rank: usize, + ) -> Option { match expr { NaryExpr::IndexedInput { input_idx, indices } => { - Some(*input_idx == 0 && NaryExpr::is_elementwise_indices(indices)) + // The index list must cover the full output rank: a shorter + // prefix (e.g. a folded keepdim view reading a lower-rank + // input) is still pointwise but not shape-preserving, so it + // must not be treated as a fusible unary chain. + Some( + *input_idx == 0 + && indices.len() == rank + && NaryExpr::is_elementwise_indices(indices), + ) } NaryExpr::Op { children, function } => { if children.len() != 1 || function.input_types.len() != 1 { return None; } - if !Self::collect_unary_chain(&children[0], functions)? { + if !Self::collect_unary_chain(&children[0], functions, rank)? { return Some(false); } functions.push(function.clone()); @@ -522,7 +538,7 @@ impl NaryOperation { } } -impl Operation for NaryOperation { +impl Operation for ElementwiseOperation { fn hash_kernel_fields(&self, state: &mut FxHasher) { self.expression.hash(state); self.shape.hash(state); diff --git a/fusor-ml/core/src/quantized/dequantize.rs b/fusor-ml/core/src/quantized/dequantize.rs index 5c45dcc96..858814760 100644 --- a/fusor-ml/core/src/quantized/dequantize.rs +++ b/fusor-ml/core/src/quantized/dequantize.rs @@ -2,63 +2,22 @@ use std::hash::Hash; use fusor_gguf::GgmlType; use fusor_tile_ir as tile_ir; -use fusor_tile_ir_kernels as tile_ir_kernels; +use crate::compute_graph::NodeIndex; use crate::mir::inputs::MirValue; use crate::mir::operation::Operation; use crate::{ CastTensor, DataType, DataTypeEnum, Device, Layout, LazyTensorData, Tensor, TensorData, TensorInfo, mir::{ - kernel_backend, kernel_backend::DirectKernel, workgroup_shape::{Constraint, WorkgroupShapeConstraints}, }, - nary_wise::UnaryFunctionChain, + nary_wise::{ElementwiseOperation, NaryExpr, NaryFunction, UnaryFunctionChain}, }; use super::{QMatrix, QMatrixStorageLayout}; -struct DequantizeDirectKernelVariant; - -struct QDequantizeKernelParams { - matrix_buffer: std::sync::Arc, - output_buffer: std::sync::Arc, - output_layout: tile_ir::Layout, - output_element: tile_ir::ElementType, - format: tile_ir::GgmlQuantFormat, - k: u32, - n: u32, - dispatch_x: u32, -} - -fn emit_qdequantize_kernel( - kb: &mut tile_ir::KernelBuilder>, - params: QDequantizeKernelParams, -) -> Option<()> { - let q = tile_ir_kernels::quantized_matrix_for( - kb, - params.matrix_buffer, - params.format, - params.k, - params.n, - ); - let y = kb.write( - params.output_element, - tile_ir::KernelTensorRef::new(params.output_buffer, params.output_layout), - ); - tile_ir_kernels::qdequantize(kb.program(), &q, &y, params.dispatch_x); - Some(()) -} - -fn datatype_element(datatype: DataTypeEnum) -> Option { - Some(match datatype { - DataTypeEnum::F32 => tile_ir::ElementType::F32, - DataTypeEnum::F16 => tile_ir::ElementType::F16, - DataTypeEnum::U32 => return None, - }) -} - #[derive(Debug, Clone)] pub(crate) struct DequantizeOperation { pub(crate) matrix: QMatrix, @@ -74,42 +33,44 @@ impl DequantizeOperation { post_dequantize: UnaryFunctionChain::empty(datatype), } } +} - fn direct_quant_format(&self) -> Option { - Some(match self.matrix.datatype { - GgmlType::Q4_0 if self.matrix.storage_layout() == QMatrixStorageLayout::Native => { - tile_ir::GgmlQuantFormat::Q4_0Native - } - GgmlType::Q4_0 => tile_ir::GgmlQuantFormat::Q4_0, - GgmlType::Q4_1 => tile_ir::GgmlQuantFormat::Q4_1, - GgmlType::Q5_0 if self.matrix.storage_layout() == QMatrixStorageLayout::Native => { - tile_ir::GgmlQuantFormat::Q5_0Native - } - GgmlType::Q5_0 => tile_ir::GgmlQuantFormat::Q5_0, - GgmlType::Q5_1 => tile_ir::GgmlQuantFormat::Q5_1, - GgmlType::Q8_0 if self.matrix.storage_layout() == QMatrixStorageLayout::Native => { - tile_ir::GgmlQuantFormat::Q8_0Native - } - GgmlType::Q8_0 => tile_ir::GgmlQuantFormat::Q8_0, - GgmlType::Q8_1 => tile_ir::GgmlQuantFormat::Q8_1, - GgmlType::Q2K => tile_ir::GgmlQuantFormat::Q2K, - GgmlType::Q3K => tile_ir::GgmlQuantFormat::Q3K, - GgmlType::Q4K if self.matrix.storage_layout() == QMatrixStorageLayout::Native => { - tile_ir::GgmlQuantFormat::Q4KNative - } - GgmlType::Q4K => tile_ir::GgmlQuantFormat::Q4K, - GgmlType::Q5K if self.matrix.storage_layout() == QMatrixStorageLayout::Native => { - tile_ir::GgmlQuantFormat::Q5KNative - } - GgmlType::Q5K => tile_ir::GgmlQuantFormat::Q5K, - GgmlType::Q6K if self.matrix.storage_layout() == QMatrixStorageLayout::Native => { - tile_ir::GgmlQuantFormat::Q6KNative - } - GgmlType::Q6K => tile_ir::GgmlQuantFormat::Q6K, - GgmlType::Q8K => tile_ir::GgmlQuantFormat::Q8K, - GgmlType::F16 | GgmlType::F32 => return None, - }) - } +/// The tile-ir quantization format for a matrix, accounting for its storage +/// layout. `None` for dense f16/f32 storage. +pub(crate) fn quant_format(matrix: &QMatrix) -> Option { + Some(match matrix.datatype { + GgmlType::Q4_0 if matrix.storage_layout() == QMatrixStorageLayout::Native => { + tile_ir::GgmlQuantFormat::Q4_0Native + } + GgmlType::Q4_0 => tile_ir::GgmlQuantFormat::Q4_0, + GgmlType::Q4_1 => tile_ir::GgmlQuantFormat::Q4_1, + GgmlType::Q5_0 if matrix.storage_layout() == QMatrixStorageLayout::Native => { + tile_ir::GgmlQuantFormat::Q5_0Native + } + GgmlType::Q5_0 => tile_ir::GgmlQuantFormat::Q5_0, + GgmlType::Q5_1 => tile_ir::GgmlQuantFormat::Q5_1, + GgmlType::Q8_0 if matrix.storage_layout() == QMatrixStorageLayout::Native => { + tile_ir::GgmlQuantFormat::Q8_0Native + } + GgmlType::Q8_0 => tile_ir::GgmlQuantFormat::Q8_0, + GgmlType::Q8_1 => tile_ir::GgmlQuantFormat::Q8_1, + GgmlType::Q2K => tile_ir::GgmlQuantFormat::Q2K, + GgmlType::Q3K => tile_ir::GgmlQuantFormat::Q3K, + GgmlType::Q4K if matrix.storage_layout() == QMatrixStorageLayout::Native => { + tile_ir::GgmlQuantFormat::Q4KNative + } + GgmlType::Q4K => tile_ir::GgmlQuantFormat::Q4K, + GgmlType::Q5K if matrix.storage_layout() == QMatrixStorageLayout::Native => { + tile_ir::GgmlQuantFormat::Q5KNative + } + GgmlType::Q5K => tile_ir::GgmlQuantFormat::Q5K, + GgmlType::Q6K if matrix.storage_layout() == QMatrixStorageLayout::Native => { + tile_ir::GgmlQuantFormat::Q6KNative + } + GgmlType::Q6K => tile_ir::GgmlQuantFormat::Q6K, + GgmlType::Q8K => tile_ir::GgmlQuantFormat::Q8K, + GgmlType::F16 | GgmlType::F32 => return None, + }) } impl Operation for DequantizeOperation { @@ -163,82 +124,54 @@ impl Operation for DequantizeOperation { fn build_direct_kernel( &self, graph: &crate::compute_graph::ComputeGraphInner, - _workgroup_shape: &crate::mir::workgroup_shape::WorkgroupShape, + workgroup_shape: &crate::mir::workgroup_shape::WorkgroupShape, inputs: &[MirValue], ) -> Option { - if !matches!(self.datatype, DataTypeEnum::F32 | DataTypeEnum::F16) - || !self.post_dequantize.functions.is_empty() - || (self.datatype == DataTypeEnum::F16 && !graph.device().f16_supported()) + // Dequantization is the identity expression over one block-quantized + // input: the generic n-ary kernel decodes per element through the + // format-aware load, and any fused chain rides along. + let rank = self.matrix.shape.len(); + let mut expression = NaryExpr::input(0, rank); + let mut current = DataTypeEnum::F32; + if self.post_dequantize.input_datatype() != current + || self.post_dequantize.out_datatype() != self.datatype { - return None; + expression = NaryExpr::Op { + children: vec![expression], + function: NaryFunction::unary( + Some("cast".to_string()), + crate::nary_wise::NaryOp::Cast, + current, + self.post_dequantize.input_datatype(), + ), + }; + current = self.post_dequantize.input_datatype(); } - let [matrix, output] = inputs else { - return None; - }; - let MirValue::QMatrix(matrix) = matrix else { - return None; - }; - let output = output.as_tensor()?; - if output.datatype() != self.datatype { - return None; + for function in &self.post_dequantize.functions { + expression = NaryExpr::Op { + children: vec![expression], + function: function.clone(), + }; + current = function.output_type; } - - let format = self.direct_quant_format()?; - let k = *self.matrix.shape.last()? as u32; - let n: u32 = self - .matrix - .shape - .iter() - .rev() - .skip(1) - .try_fold(1u32, |acc, dim| acc.checked_mul((*dim).try_into().ok()?))?; - let total = k.checked_mul(n)?; - let workgroups = total.div_ceil(256); - let max_workgroups = graph - .device() - .limits() - .max_compute_workgroups_per_dimension - .max(1); - let dispatch_x = workgroups.min(max_workgroups); - let dispatch_y = workgroups.div_ceil(dispatch_x.max(1)).max(1); - if dispatch_y > max_workgroups { - return None; + if current != self.datatype { + expression = NaryExpr::Op { + children: vec![expression], + function: NaryFunction::unary( + Some("cast".to_string()), + crate::nary_wise::NaryOp::Cast, + current, + self.datatype, + ), + }; } - let cache_key = self.kernel_cache_key_with_dispatch( - kernel_backend::KernelVariantKey::of::(), - Some(_workgroup_shape), - [dispatch_x, dispatch_y, 1], - inputs, - ); - let matrix_buffer = matrix.buffer().clone(); - let output_buffer = output.buffer().clone(); - let output_layout = tile_ir::Layout::contiguous( - tile_ir::MemoryLevel::Storage, - tile_ir::Shape::new([total.max(1)]), - ); - let output_datatype = self.datatype; - kernel_backend::run_kernel( - graph.device().kernel_cache(), - self.name(), - cache_key, - [dispatch_x, dispatch_y, 1], - move |kb| { - emit_qdequantize_kernel( - kb, - QDequantizeKernelParams { - matrix_buffer, - output_buffer, - output_layout, - output_element: datatype_element(output_datatype)?, - format, - k, - n, - dispatch_x, - }, - )?; - Some(()) - }, - ) + let synthesized = ElementwiseOperation { + inputs: vec![NodeIndex::new(0)], + expression, + shape: self.matrix.shape.clone(), + output_datatype: self.datatype, + }; + crate::nary_direct::build_nary_direct_kernel(&synthesized, graph, workgroup_shape, inputs) } fn name(&self) -> String { diff --git a/fusor-ml/core/src/quantized/embedding.rs b/fusor-ml/core/src/quantized/embedding.rs index 757a11fb1..457b4ebc3 100644 --- a/fusor-ml/core/src/quantized/embedding.rs +++ b/fusor-ml/core/src/quantized/embedding.rs @@ -328,7 +328,13 @@ impl QMatrix { self.index_select_rows_to(indexes, DataTypeEnum::F32) } + /// Row gather in its composed form: an elementwise read of the quantized + /// table at `[indexes[i], j]`. The resolver recognizes the canonical + /// cluster and routes it to the block-amortized embedding kernel for + /// quantized tables; dense-storage tables read directly. pub fn index_select_rows_to(&self, indexes: &Tensor, datatype: DataTypeEnum) -> Tensor { + use crate::nary_wise::{ElementwiseOperation, NaryExpr, NaryFunction, NaryOp}; + indexes.assert_rank::<1>(); indexes.assert_datatype::(); assert_eq!( @@ -337,25 +343,51 @@ impl QMatrix { "quantized row index_select requires a 2D table, got {}D", self.shape.len() ); - if matches!(self.datatype, GgmlType::F32 | GgmlType::F16) { - let dense = if datatype == DataTypeEnum::F16 { - self.dequantize::() - } else { - self.dequantize::() - }; - return dense.index_select(0, indexes).cast_to(datatype); - } let index_count = indexes.shape()[0]; + let hidden = self.shape[1]; let device = self.device.clone(); let datatype = if datatype == DataTypeEnum::F16 && !self.device.f16_supported() { DataTypeEnum::F32 } else { datatype }; - let operation = - QEmbeddingOperation::new(indexes.key(), index_count, self.clone(), datatype); - let info = TensorInfo::new(operation.out_shape.clone(), datatype); - let key = device.compute_graph().create_q_embedding(operation); + + // Quantized loads decode to f32; dense-storage tables read at their + // storage type. Cast to the requested type when they differ. + let loaded_datatype = match self.datatype { + GgmlType::F16 => DataTypeEnum::F16, + _ => DataTypeEnum::F32, + }; + let row = NaryExpr::indexed_input(1, vec![NaryExpr::DimIndex(0)]); + let gather = NaryExpr::indexed_input(0, vec![row, NaryExpr::DimIndex(1)]); + let body = if loaded_datatype == datatype { + gather + } else { + NaryExpr::Op { + children: vec![gather], + function: NaryFunction::unary( + Some("cast".to_string()), + NaryOp::Cast, + loaded_datatype, + datatype, + ), + } + }; + + let matrix_key = device.compute_graph().dequantize(self.clone(), datatype); + let matrix = Tensor::from_parts(LazyTensorData::from_parts( + device.clone(), + TensorInfo::new(self.shape.clone(), datatype), + matrix_key, + )); + let operation = ElementwiseOperation { + inputs: vec![matrix.key(), indexes.key()], + expression: body, + shape: [index_count, hidden].into(), + output_datatype: datatype, + }; + let info = TensorInfo::new(operation.shape.clone(), datatype); + let key = device.compute_graph().create_nary(operation); Tensor::from_parts(LazyTensorData::from_parts(device, info, key)) } } diff --git a/fusor-ml/core/src/quantized/matmul/fallback.rs b/fusor-ml/core/src/quantized/matmul/fallback.rs index 50d3a8996..835a6bd01 100644 --- a/fusor-ml/core/src/quantized/matmul/fallback.rs +++ b/fusor-ml/core/src/quantized/matmul/fallback.rs @@ -78,7 +78,7 @@ impl QMatMulOperation { graph, &dense_matmul .workgroup_shape_constraints(&device) - .solve(device.max_subgroup_size())?, + .solve(device.max_subgroup_size(), &device.limits())?, &[ input.clone().into(), dense_weight_t.into(), @@ -109,7 +109,7 @@ impl QMatMulOperation { let dequantize_inputs = vec![matrix.clone().into(), dense_weight.clone().into()]; let dequantize_workgroup = dequantize .workgroup_shape_constraints(&graph.device()) - .solve(graph.device().max_subgroup_size())?; + .solve(graph.device().max_subgroup_size(), &graph.device().limits())?; let dequantize_kernel = dequantize.build_direct_kernel(graph, &dequantize_workgroup, &dequantize_inputs)?; let dense_matrix = QMatrix { @@ -133,7 +133,7 @@ impl QMatMulOperation { shape: &[usize], output_datatype: DataTypeEnum, ) -> Option { - let operation = NaryOperation { + let operation = ElementwiseOperation { inputs: (0..inputs.len()).map(NodeIndex::new).collect(), expression, shape: shape.into(), @@ -146,7 +146,7 @@ impl QMatMulOperation { mir_inputs.push(output.clone().into()); let workgroup_shape = operation .workgroup_shape_constraints(&graph.device()) - .solve(graph.device().max_subgroup_size())?; + .solve(graph.device().max_subgroup_size(), &graph.device().limits())?; operation.build_direct_kernel(graph, &workgroup_shape, &mir_inputs) } diff --git a/fusor-ml/core/src/quantized/matmul/kernel.rs b/fusor-ml/core/src/quantized/matmul/kernel.rs index 8d545e0d8..ccd08998c 100644 --- a/fusor-ml/core/src/quantized/matmul/kernel.rs +++ b/fusor-ml/core/src/quantized/matmul/kernel.rs @@ -740,6 +740,12 @@ impl Operation for QMatMulOperation { } } if matches!(matrix.datatype(), GgmlType::F32 | GgmlType::F16) { + // The dense kernel has no epilogue slots; declining here routes + // fused epilogues to the dequantize+dense fallback, which applies + // them as separate elementwise kernels. + if self.pre_element_wise_expr.is_some() || self.post_element_wise_expr.is_some() { + return None; + } return self.build_dense_direct_kernel(graph, input, matrix, output); } Self::direct_kernel_for_tensors( diff --git a/fusor-ml/core/src/quantized/matmul/mod.rs b/fusor-ml/core/src/quantized/matmul/mod.rs index eabafd7e2..15a300a41 100644 --- a/fusor-ml/core/src/quantized/matmul/mod.rs +++ b/fusor-ml/core/src/quantized/matmul/mod.rs @@ -21,7 +21,7 @@ use crate::{ workgroup_shape::{Constraint, WorkgroupShapeConstraints}, }, nary_direct::{apply_multi_input_elementwise_expr, apply_single_input_elementwise_expr}, - nary_wise::{NaryExpr, NaryOp, NaryOperation}, + nary_wise::{ElementwiseOperation, NaryExpr, NaryOp}, }; use fusor_gguf::GgmlType; use fusor_tile_ir as tile_ir; diff --git a/fusor-ml/core/src/reduce.rs b/fusor-ml/core/src/reduce.rs index 8854fab66..604a4e249 100644 --- a/fusor-ml/core/src/reduce.rs +++ b/fusor-ml/core/src/reduce.rs @@ -3,8 +3,9 @@ use std::hash::Hash; use rustc_hash::FxHasher; use crate::{ - Layout, Tensor, + Tensor, compute_graph::NodeIndex, + nary_wise::NaryExpr, tensor::{DataTypeEnum, TensorData}, }; use crate::{ @@ -34,21 +35,34 @@ fn unsqueeze_dim(tensor: &Tensor, dim_idx: usize) -> Tensor { tensor.reshape(new_shape) } +/// A reduction over one axis of an index space, with a fused n-ary producer. +/// +/// `expression` is evaluated at every coordinate of `shape` (the full +/// pre-reduce index space, including the reduced `axis`) and folded with +/// `function` along `axis`. A plain tensor reduction is the trivial producer +/// `input(0, rank)`; the resolver widens it by inlining upstream elementwise +/// expressions, so composed map-reduce clusters (contractions included) lower +/// as a single kernel without materializing the intermediate. #[derive(Debug, Clone)] pub(crate) struct ReduceOperation { - pub(crate) value: NodeIndex, - pub(crate) pre_element_wise: UnaryFunctionChain, + /// Producer inputs referenced by `expression`. + pub(crate) inputs: Vec, + /// Fused producer over the full index space `shape`. + pub(crate) expression: NaryExpr, + /// The full pre-reduce index space, including the reduced axis. + pub(crate) shape: Box<[usize]>, pub(crate) function: ReduceFunction, pub(crate) post_element_wise: UnaryFunctionChain, pub(crate) axis: usize, } impl ReduceOperation { - pub fn new(value: NodeIndex, function: ReduceFunction, axis: usize, _shape: &[usize]) -> Self { + pub fn new(value: NodeIndex, function: ReduceFunction, axis: usize, shape: &[usize]) -> Self { let datatype = function.datatype(); Self { - value, - pre_element_wise: UnaryFunctionChain::empty(datatype), + inputs: vec![value], + expression: NaryExpr::input(0, shape.len()), + shape: shape.into(), function, post_element_wise: UnaryFunctionChain::empty(datatype), axis, @@ -58,11 +72,34 @@ impl ReduceOperation { pub fn out_datatype(&self) -> DataTypeEnum { self.post_element_wise.out_datatype() } + + /// The single input of a trivial (un-fused) reduction: the producer is + /// still the bare `input(0, rank)` the tensor API emitted. Recognition + /// runs before fusion, so the canonical clusters it matches always take + /// this form. + pub(crate) fn plain_input(&self) -> Option { + (self.inputs.len() == 1 && self.expression == NaryExpr::input(0, self.shape.len())) + .then(|| self.inputs[0]) + } + + /// The output shape: the index space with the reduced axis removed. + pub(crate) fn out_shape(&self) -> Vec { + self.shape + .iter() + .enumerate() + .filter_map(|(i, x)| (i != self.axis).then_some(*x)) + .collect() + } + + pub(crate) fn reduce_size(&self) -> usize { + self.shape[self.axis] + } } impl Operation for ReduceOperation { fn hash_kernel_fields(&self, state: &mut FxHasher) { - self.pre_element_wise.hash(state); + self.expression.hash(state); + self.shape.hash(state); self.function.hash(state); self.post_element_wise.hash(state); self.axis.hash(state); @@ -85,12 +122,9 @@ impl Operation for ReduceOperation { workgroup_shape: &crate::mir::workgroup_shape::WorkgroupShape, inputs: &[MirValue], ) -> [u32; 3] { - let output_tensor: TensorData = inputs[1].as_tensor().unwrap().clone(); + let output_tensor: TensorData = inputs.last().unwrap().as_tensor().unwrap().clone(); let total_outputs = output_tensor.layout().shape().iter().product::() as u32; - let reduce_size = match inputs.get(2) { - Some(MirValue::Integer(value)) => *value, - _ => 1, - }; + let reduce_size = self.reduce_size() as u32; let serial_workgroups = total_outputs.div_ceil(workgroup_shape.x()); let total_workgroups = if use_cooperative_reduce(total_outputs, reduce_size, workgroup_shape.x()) { @@ -109,53 +143,37 @@ impl Operation for ReduceOperation { } fn visit_dependencies(&self, f: &mut dyn FnMut(NodeIndex)) { - f(self.value); + for input in &self.inputs { + f(*input); + } } fn inputs(&self, nodes: &crate::compute_graph::ComputeGraphInner) -> Vec { - let dim = self.axis; - let tensor = nodes.get_cached_result(self.value).unwrap(); - assert_eq!(self.pre_element_wise.input_datatype(), tensor.datatype()); - let layout = tensor.layout(); - let shape = layout.shape(); - let new_tensor_shape = shape + let mut mir_inputs: Vec = self + .inputs .iter() .enumerate() - .filter_map(|(i, x)| (i != dim).then_some(*x)) - .collect::>(); - let output_type = self.out_datatype(); + .map(|(i, idx)| { + // Custom-indexed inputs need the dense (dequantized) tensor; + // block-quantized data only supports the plain row/col path. + if self.expression.uses_custom_indexing_for_input(i) + && let Some(cached) = nodes.get_result(*idx) + { + return cached.into(); + } + nodes.get_result_or_qmatrix(*idx).unwrap().into() + }) + .collect(); + + let device = match &mir_inputs[0] { + MirValue::Tensor(tensor) => tensor.device().clone(), + MirValue::QMatrix(matrix) => matrix.device().clone(), + _ => unreachable!("reduce inputs are tensors or quantized matrices"), + }; let output_tensor = - TensorData::new_for_shape(tensor.device(), &new_tensor_shape, output_type); - - let trimmed_tensor_layout = Layout::from_parts( - tensor.layout().offset(), - tensor - .layout() - .shape() - .iter() - .enumerate() - .filter_map(|(i, x)| (i != dim).then_some(*x)) - .collect(), - tensor - .layout() - .strides() - .iter() - .enumerate() - .filter_map(|(i, x)| (i != dim).then_some(*x)) - .collect(), - ); - let trimmed_tensor = TensorData::new_from_parts( - tensor.device(), - tensor.buffer().clone(), - trimmed_tensor_layout, - tensor.datatype(), - ); - vec![ - MirValue::Tensor(trimmed_tensor.clone()), - MirValue::Tensor(output_tensor.clone()), - MirValue::Integer(tensor.layout().shape()[dim] as u32), - MirValue::Integer(tensor.layout().strides()[dim] as u32), - ] + TensorData::new_for_shape(&device, &self.out_shape(), self.out_datatype()); + mir_inputs.push(output_tensor.into()); + mir_inputs } fn build_direct_kernel( @@ -164,16 +182,28 @@ impl Operation for ReduceOperation { workgroup_shape: &crate::mir::workgroup_shape::WorkgroupShape, inputs: &[MirValue], ) -> Option { - crate::reduce_direct::build_reduce_direct_kernel(self, graph, workgroup_shape, inputs) + crate::reduce_tiled::build_reduce_tiled_kernel(self, graph, workgroup_shape, inputs) + .or_else(|| { + crate::reduce_direct::build_reduce_direct_kernel( + self, + graph, + workgroup_shape, + inputs, + ) + }) } fn output(&self, _: &crate::compute_graph::ComputeGraphInner, inputs: &[MirValue]) -> MirValue { - let output_tensor: TensorData = inputs[1].as_tensor().unwrap().clone(); + let output_tensor: TensorData = inputs.last().unwrap().as_tensor().unwrap().clone(); output_tensor.into() } fn name(&self) -> String { - format!("reduce_{}", self.function.name()) + if self.plain_input().is_some() { + format!("reduce_{}", self.function.name()) + } else { + format!("reduce_{}_fused", self.function.name()) + } } } @@ -233,7 +263,7 @@ impl Tensor { } } -fn sum_fn(datatype: DataTypeEnum) -> ReduceFunction { +pub(crate) fn sum_fn(datatype: DataTypeEnum) -> ReduceFunction { ReduceFunction::new(ReduceOp::Sum, zero_for_dtype(datatype), datatype).with_name("sum") } @@ -248,7 +278,7 @@ impl Tensor { } } -fn max_fn(datatype: DataTypeEnum) -> ReduceFunction { +pub(crate) fn max_fn(datatype: DataTypeEnum) -> ReduceFunction { ReduceFunction::new(ReduceOp::Max, min_scalar_for_dtype(datatype), datatype).with_name("max") } diff --git a/fusor-ml/core/src/reduce_direct.rs b/fusor-ml/core/src/reduce_direct.rs index f0dc219b7..4d18d4e4d 100644 --- a/fusor-ml/core/src/reduce_direct.rs +++ b/fusor-ml/core/src/reduce_direct.rs @@ -1,87 +1,232 @@ use fusor_tile_ir as tile_ir; +use fusor_tile_ir_kernels as tile_ir_kernels; +use std::hash::Hash; use crate::{ + access_analysis::InputAccesses, mir::{ inputs::MirValue, kernel_backend, kernel_backend::DirectKernel, operation::Operation, workgroup_shape::WorkgroupShape, }, nary_direct::{ - TensorMeta, ValueTile, apply_unary_function_chain, flat_layout, layout_index, linear_group, - output_dims_from_flat, + ValueTile, apply_unary_function_chain, declare_value, eval_nary_expr, linear_group, + output_dims_from_flat, tile_u32, }, nary_wise::NaryScalar, reduce::{ReduceOp, ReduceOperation}, + reduce_tiled::{datatype_scalar, input_allocation_bytes}, tensor::DataTypeEnum, + visit_tiled::{MaybeQData, distribute_workgroups}, }; const BLOCK: usize = 256; +/// Output rows per workgroup on the subgroup-per-output route. +const SUBGROUP_ROWS: u32 = 4; +/// Consecutive reduce-axis values each lane folds per iteration on the +/// subgroup-per-output route (matches the dense gemv kernel's run length). +const SUBGROUP_VALUES_PER_LANE: u32 = 4; struct ReduceDirectKernelVariant; -pub(crate) fn build_reduce_direct_kernel( +/// How the fused reduce distributes its fold across threads. +#[derive(Clone, Copy)] +enum ReduceRoute { + /// One workgroup per output, 256 lanes splitting the reduce axis. The + /// final cross-lane combine uses subgroup collectives (one reduction per + /// subgroup, partials through workgroup memory, one more reduction) when + /// the device has a fixed subgroup width; otherwise the scratch tree. + Cooperative { + subgroup_finish: Option, + }, + /// One subgroup per output: lanes walk the reduce axis in + /// `SUBGROUP_VALUES_PER_LANE`-long runs (consecutive lanes read adjacent + /// runs — coalesced when the dominant input is contiguous along the + /// axis) and a single subgroup collective combines them. + SubgroupRow(tile_ir_kernels::SubgroupConfig), + /// One thread per output with a per-thread fold. + Serial, +} + +/// The subgroup-per-output route only pays when adjacent lanes read adjacent +/// addresses: the byte-dominant reduce-axis-dependent input must be +/// contiguous along the axis. Otherwise the serial layout (adjacent threads +/// = adjacent outputs) is the coalesced one. +fn dominant_k_dep_input_is_k_contiguous( operation: &ReduceOperation, - graph: &crate::compute_graph::ComputeGraphInner, - workgroup_shape: &WorkgroupShape, - inputs: &[MirValue], -) -> Option { - let input = inputs[0].as_tensor()?.clone(); - let output = inputs[1].as_tensor()?.clone(); - let reduce_size = match inputs.get(2)? { - MirValue::Integer(value) => *value, - _ => return None, + values: &[MaybeQData], +) -> bool { + let Some(metas) = values + .iter() + .map(crate::reduce_tiled::analysis_meta) + .collect::>>() + else { + return false; }; - let reduce_stride = match inputs.get(3)? { - MirValue::Integer(value) => *value, - _ => return None, + let Some(access) = + InputAccesses::collect(&operation.expression, operation.inputs.len(), &metas) + else { + return false; }; - - if (input.datatype() == DataTypeEnum::F16 || output.datatype() == DataTypeEnum::F16) - && !graph.device().f16_supported() - { - return None; + let axis = operation.axis; + let mut best: Option<(u64, bool)> = None; + for i in 0..operation.inputs.len() { + if !access.depends_on(i, axis) { + continue; + } + let contiguous = access.dims[i] + .iter() + .enumerate() + .any(|(j, &d)| d == axis && metas[i].strides.get(j).copied() == Some(1)); + let bytes = input_allocation_bytes(&metas[i], &values[i]); + if best.as_ref().is_none_or(|(b, _)| bytes > *b) { + best = Some((bytes, contiguous)); + } } + best.is_some_and(|(_, contiguous)| contiguous) +} - let input_meta = TensorMeta::new(&input)?; - let output_meta = TensorMeta::new(&output)?; - if operation.pre_element_wise.input_datatype() != input_meta.datatype { - return None; - } - let reduce_dtype = operation.pre_element_wise.out_datatype(); - if reduce_dtype != operation.function.datatype() - || operation.post_element_wise.input_datatype() != reduce_dtype - { - return None; +/// The shared preamble for every fused-reduce kernel builder: split the +/// MirValues into producer values + output, gate on device support, and +/// declare the kernel bindings. +pub(crate) struct ReduceKernelInputs { + pub(crate) values: Vec, + pub(crate) output_shape: Vec, + pub(crate) total_outputs: u32, +} + +impl ReduceKernelInputs { + pub(crate) fn parse( + operation: &ReduceOperation, + graph: &crate::compute_graph::ComputeGraphInner, + inputs: &[MirValue], + ) -> Option { + let (output, producers) = inputs.split_last()?; + let output = output.as_tensor()?; + let values = producers + .iter() + .map(|input| MaybeQData::try_from(input.clone()).ok()) + .collect::>>()?; + + let f16_unsupported = !graph.device().f16_supported(); + if f16_unsupported { + let uses_f16 = output.datatype() == DataTypeEnum::F16 + || operation.function.datatype() == DataTypeEnum::F16 + || values.iter().any(|value| match value { + MaybeQData::Tensor(tensor) => tensor.datatype() == DataTypeEnum::F16, + MaybeQData::QMatrix(matrix) => { + matches!(matrix.datatype(), fusor_gguf::GgmlType::F16) + } + }); + if uses_f16 { + return None; + } + } + + if operation.post_element_wise.input_datatype() != operation.function.datatype() { + return None; + } + + let output_shape = output + .layout() + .shape() + .iter() + .copied() + .map(u32::try_from) + .collect::, _>>() + .ok()?; + let total_outputs = output_shape + .iter() + .try_fold(1u32, |acc, dim| acc.checked_mul(*dim))?; + + Some(Self { + values, + output_shape, + total_outputs, + }) } +} - let output_shape = output - .layout() - .shape() - .iter() - .copied() - .map(u32::try_from) - .collect::, _>>() - .ok()?; - let total_outputs = output_shape - .iter() - .try_fold(1u32, |acc, dim| acc.checked_mul(*dim))?; +pub(crate) fn build_reduce_direct_kernel( + operation: &ReduceOperation, + graph: &crate::compute_graph::ComputeGraphInner, + workgroup_shape: &WorkgroupShape, + inputs: &[MirValue], +) -> Option { + let parsed = ReduceKernelInputs::parse(operation, graph, inputs)?; + let output = inputs.last()?.as_tensor()?.clone(); + let reduce_size: u32 = operation.reduce_size().try_into().ok()?; + + let reduce_dtype = operation.function.datatype(); let reduce_op = tile_reduce_op(operation.function.op); let initial = operation.function.initial_value; - let dispatch_size = operation.dispatch_size(workgroup_shape, inputs); + + let device = graph.device(); + let limits = device.limits(); + let subgroup_config = device.subgroup_config(); + let route = + if crate::reduce::use_cooperative_reduce(parsed.total_outputs, reduce_size, BLOCK as u32) { + // The two-stage subgroup finish needs a compile-time subgroup count + // small enough for the second collective to cover the partials. + let subgroup_finish = subgroup_config.filter(|config| { + config.is_fixed() + && (BLOCK as u32).is_multiple_of(config.max_size()) + && BLOCK as u32 / config.max_size() <= config.max_size() + }); + ReduceRoute::Cooperative { subgroup_finish } + } else { + let row_config = subgroup_config.filter(|config| { + let block = config.block_for_subgroups(SUBGROUP_ROWS); + reduce_size >= 32 + && block.is_power_of_two() + && block + <= limits + .max_compute_workgroup_size_x + .min(limits.max_compute_invocations_per_workgroup) + && dominant_k_dep_input_is_k_contiguous(operation, &parsed.values) + }); + match row_config { + Some(config) => ReduceRoute::SubgroupRow(config), + None => ReduceRoute::Serial, + } + }; + let dispatch_size = match route { + ReduceRoute::SubgroupRow(_) => distribute_workgroups( + parsed.total_outputs.div_ceil(SUBGROUP_ROWS), + limits.max_compute_workgroups_per_dimension, + ), + _ => operation.dispatch_size(workgroup_shape, inputs), + }; + let variant = + kernel_backend::KernelVariantKey::with_payload::(|state| { + match route { + ReduceRoute::Cooperative { subgroup_finish } => { + 0u8.hash(state); + subgroup_finish.hash(state); + } + ReduceRoute::SubgroupRow(config) => { + 1u8.hash(state); + config.hash(state); + } + ReduceRoute::Serial => 2u8.hash(state), + } + }); let cache_key = operation.kernel_cache_key_with_dispatch( - kernel_backend::KernelVariantKey::of::(), + variant, Some(workgroup_shape), dispatch_size, inputs, ); - let input_buffer = input.buffer().clone(); - let output_buffer = output.buffer().clone(); - let input_layout = flat_layout(input_meta.allocation_len); - let output_layout = flat_layout(output_meta.allocation_len); - let input_meta_body = input_meta.clone(); - let output_meta_body = output_meta.clone(); - let pre_chain = operation.pre_element_wise.clone(); + let axis = operation.axis; + let expression = operation.expression.clone(); let post_chain = operation.post_element_wise.clone(); + let output_dtype = output.datatype(); + let output_value = MaybeQData::Tensor(output); + let ReduceKernelInputs { + values, + output_shape, + total_outputs, + } = parsed; kernel_backend::run_kernel( graph.device().kernel_cache(), @@ -89,220 +234,227 @@ pub(crate) fn build_reduce_direct_kernel( cache_key, dispatch_size, move |kb| { - let input_tensor = - tile_ir::KernelTensorRef::new(input_buffer.clone(), input_layout.clone()); - let output_tensor = - tile_ir::KernelTensorRef::new(output_buffer.clone(), output_layout.clone()); - let input_element = crate::nary_direct::datatype_element(input_meta_body.datatype); - let input_storage = match input_meta_body.datatype { - DataTypeEnum::F32 => { - crate::nary_direct::Storage2::F32(kb.read(input_element, input_tensor)) - } - DataTypeEnum::F16 => { - crate::nary_direct::Storage2::F16(kb.read(input_element, input_tensor)) - } - DataTypeEnum::U32 => { - crate::nary_direct::Storage2::U32(kb.read(input_element, input_tensor)) - } + let mut storages = Vec::with_capacity(values.len()); + let mut metas = Vec::with_capacity(values.len()); + for value in &values { + let (storage, meta) = declare_value(kb, value, false)?; + storages.push(storage); + metas.push(meta); + } + let (output_storage, output_meta) = declare_value(kb, &output_value, true)?; + + // Evaluate the fused producer at one index-space coordinate: the + // output dims with the running reduce index inserted at `axis`. + let value_at = |program: &mut tile_ir::tile::TileBlock<'_>, + dims_out: &[tile_ir::tile::Tile], + reduce_index: tile_ir::tile::Tile, + mask: tile_ir::tile::Mask| { + let mut coords = dims_out.to_vec(); + coords.insert(axis, reduce_index); + let (value, _) = + eval_nary_expr(program, &expression, &coords, &storages, &metas, mask, &[]); + value.cast_to(reduce_dtype) }; - let output_element = crate::nary_direct::datatype_element(output_meta_body.datatype); - let output_storage = match output_meta_body.datatype { - DataTypeEnum::F32 => { - crate::nary_direct::Storage2::F32(kb.write(output_element, output_tensor)) - } - DataTypeEnum::F16 => { - crate::nary_direct::Storage2::F16(kb.write(output_element, output_tensor)) - } - DataTypeEnum::U32 => { - crate::nary_direct::Storage2::U32(kb.write(output_element, output_tensor)) - } + + let store_reduced = |program: &mut tile_ir::tile::TileBlock<'_>, + reduced: ValueTile, + dims_out: &[tile_ir::tile::Tile], + mask: tile_ir::tile::Mask| { + let (reduced, reduced_ty) = + apply_unary_function_chain(reduced.into_f32(), reduce_dtype, &post_chain) + .expect("validated reduce post_element_wise chain"); + let reduced = ValueTile::F32(reduced) + .cast_to(reduced_ty) + .cast_to(output_dtype); + let output_index = crate::nary_direct::layout_index(&output_meta, dims_out); + output_storage.store(program, output_index, reduced, mask); + }; + + let identity = || tile_ir::tile::Tile::literal(tile_literal_for(initial, reduce_dtype)); + let untag = |value: ValueTile| match reduce_dtype { + DataTypeEnum::F32 => value.into_f32(), + DataTypeEnum::F16 => value.into_f16(), + DataTypeEnum::U32 => value.into_u32(), }; + let tag = |value: tile_ir::tile::Tile| match reduce_dtype { + DataTypeEnum::F32 => ValueTile::F32(value), + DataTypeEnum::F16 => ValueTile::F16(value), + DataTypeEnum::U32 => ValueTile::U32(value), + }; + + // The per-lane fold every route shares: `k_of(loop_index, run)` + // gives this lane's reduce coordinate for each of `runs` values + // per iteration, and out-of-range slots collapse to the identity. + // A route whose coordinates never leave the axis (serial) skips + // the select by passing `mask_oob: false`. + let fold_reduce_lane = + |program: &mut tile_ir::tile::TileBlock<'_>, + dims_out: &[tile_ir::tile::Tile], + iterations: tile_ir::tile::FoldIter, + in_bounds: &tile_ir::tile::Mask, + runs: u32, + mask_oob: bool, + k_of: &dyn Fn(&tile_ir::tile::Tile, u32) -> tile_ir::tile::Tile| + -> tile_ir::tile::Tile { + let reduce_binary = reduce_op.binary(); + let [acc] = + program.fold(iterations, [identity()], |program, loop_index, [acc]| { + let mut acc = acc; + for run in 0..runs { + let k_index = k_of(&loop_index, run); + let active = in_bounds.clone().and(k_index.clone().lt(reduce_size)); + let value = value_at(program, dims_out, k_index, active.clone()); + let value = if mask_oob { + tile_ir::tile::Tile::select(active, untag(value), identity()) + } else { + untag(value) + }; + acc = acc.binary(reduce_binary, value); + } + [acc] + }); + acc + }; - let cooperative = - crate::reduce::use_cooperative_reduce(total_outputs, reduce_size, BLOCK as u32); - if cooperative { - let chunks = reduce_size.div_ceil(BLOCK as u32); - kb.program() - .program_grid(BLOCK as u32, dispatch_size, |program| { + match route { + ReduceRoute::Cooperative { subgroup_finish } => { + let chunks = reduce_size.div_ceil(BLOCK as u32); + let phase = kb.program(); + let scratch = subgroup_finish + .map(|config| BLOCK as u32 / config.max_size()) + .filter(|partials| *partials > 1) + .map(|partials| { + phase.alloc_workgroup_array(datatype_scalar(reduce_dtype), partials) + }); + phase.program_grid(BLOCK as u32, dispatch_size, |program| { let lane = program.lane(); let output_flat = linear_group(program, dispatch_size); let in_bounds = output_flat.lt(total_outputs); - let dims = output_dims_from_flat_usize(output_flat.clone(), &output_shape); - let base = layout_index(&input_meta_body, &dims); - let reduce_binary = reduce_op.binary(); - - let value_at = - |program: &mut tile_ir::tile::TileBlock<'_>, - reduce_index: tile_ir::tile::Tile, - active: tile_ir::tile::Mask| { - let value_index = base.clone() + reduce_index * reduce_stride; - let value = - input_storage.load(program, value_index, active.clone()); - let (value, value_ty) = apply_unary_function_chain( - value.into_f32(), - input_meta_body.datatype, - &pre_chain, - ) - .expect("validated reduce pre_element_wise chain"); - let value = ValueTile::F32(value) - .cast_to(value_ty) - .cast_to(reduce_dtype); - let identity = tile_ir::tile::Tile::literal(tile_literal_for( - initial, - reduce_dtype, - )); - match reduce_dtype { - DataTypeEnum::F32 => tile_ir::tile::Tile::select( - active, - value.into_f32(), - identity, - ), - DataTypeEnum::F16 => tile_ir::tile::Tile::select( - active, - value.into_f16(), - identity, - ), - DataTypeEnum::U32 => tile_ir::tile::Tile::select( - active, - value.into_u32(), - identity, - ), - } - }; + let dims_out = + output_dims_from_flat_usize(output_flat.clone(), &output_shape); - let initial_acc = - tile_ir::tile::Tile::literal(tile_literal_for(initial, reduce_dtype)); - let [partial] = program.fold( + let partial = fold_reduce_lane( + program, + &dims_out, tile_ir::tile::range(chunks), - [initial_acc], - |program, loop_index, [acc]| { - let reduce_index = loop_index * BLOCK as u32 + lane.clone(); - let active = in_bounds.clone().and(reduce_index.lt(reduce_size)); - [ - acc.binary( - reduce_binary, - value_at(program, reduce_index, active), - ), - ] - }, + &in_bounds, + 1, + true, + &|loop_index, _| loop_index.clone() * BLOCK as u32 + lane.clone(), ); - let reduced = program.group_reduce(reduce_op, BLOCK as u32, partial); - - let reduced = match reduce_dtype { - DataTypeEnum::F32 => ValueTile::F32(reduced), - DataTypeEnum::F16 => ValueTile::F16(reduced), - DataTypeEnum::U32 => ValueTile::U32(reduced), + let (reduced, store_mask) = match subgroup_finish { + // One collective per subgroup, partials through + // workgroup memory, one more collective: two + // barriers and two subgroup ops instead of the + // log2(block) scratch-tree rounds. + Some(config) => { + let token = config.token(); + let first = token.subgroup_reduce(program, reduce_op, partial); + let partials = BLOCK as u32 / config.max_size(); + let sg_lane = token.subgroup_lane(program); + let sg_id = token.subgroup_id(program); + let reduced = if partials <= 1 { + first + } else { + let scratch = scratch.as_ref().unwrap(); + let first = program.bind(first); + program.if_then(sg_lane.clone().eq(0u32), |program| { + program.store_workgroup( + scratch, + sg_id.clone(), + first.clone(), + ); + }); + program.workgroup_barrier(); + let slot = tile_ir::tile::Tile::select( + lane.clone().lt(partials), + lane.clone(), + tile_u32(0), + ); + let loaded = program.load_workgroup(scratch, slot); + let masked = tile_ir::tile::Tile::select( + lane.clone().lt(partials), + loaded, + identity(), + ); + token.subgroup_reduce(program, reduce_op, masked) + }; + // The combined value is uniform across + // subgroup 0; store from its first lane + // rather than assuming workgroup lane 0 + // belongs to it. + let store_mask = + in_bounds.clone().and(sg_id.eq(0u32)).and(sg_lane.eq(0u32)); + (reduced, store_mask) + } + None => ( + program.group_reduce(reduce_op, BLOCK as u32, partial), + in_bounds.clone().and(lane.eq(0u32)), + ), }; - let (reduced, reduced_ty) = apply_unary_function_chain( - reduced.into_f32(), - reduce_dtype, - &post_chain, - ) - .expect("validated reduce post_element_wise chain"); - let reduced = ValueTile::F32(reduced) - .cast_to(reduced_ty) - .cast_to(output_meta_body.datatype); - let output_index = layout_index(&output_meta_body, &dims); - let store_mask = in_bounds.and(lane.eq(0u32)); - output_storage.store(program, output_index, reduced, store_mask); + store_reduced(program, tag(reduced), &dims_out, store_mask); }); - } else { - kb.program() - .program_grid(BLOCK as u32, dispatch_size, |program| { - let lane = program.lane(); + } + ReduceRoute::SubgroupRow(config) => { + let block = config.block_for_subgroups(SUBGROUP_ROWS); + let token = config.token(); + kb.program().program_grid(block, dispatch_size, |program| { let group = linear_group(program, dispatch_size); - let flat = group * BLOCK as u32 + lane.clone(); - let in_bounds = flat.lt(total_outputs); - let dims = output_dims_from_flat_usize(flat.clone(), &output_shape); - let base = layout_index(&input_meta_body, &dims); - let value_at = - |program: &mut tile_ir::tile::TileBlock<'_>, - loop_index: tile_ir::tile::Tile| { - let value_index = base.clone() + loop_index * reduce_stride; - let value = - input_storage.load(program, value_index, in_bounds.clone()); - let (value, value_ty) = apply_unary_function_chain( - value.into_f32(), - input_meta_body.datatype, - &pre_chain, - ) - .expect("validated reduce pre_element_wise chain"); - ValueTile::F32(value) - .cast_to(value_ty) - .cast_to(reduce_dtype) - }; + let row = group * SUBGROUP_ROWS + token.subgroup_id(program); + let in_bounds = row.clone().lt(total_outputs); + let dims_out = output_dims_from_flat_usize(row, &output_shape); + let sg_lane = token.subgroup_lane(program); + let k_per_iter = + program.bind(token.subgroup_size(program) * SUBGROUP_VALUES_PER_LANE); + let k_iterations = (tile_u32(reduce_size) + k_per_iter.clone() - 1u32) + / k_per_iter.clone(); - let reduce_binary = reduce_op.binary(); - let reduced = match reduce_dtype { - DataTypeEnum::F32 => { - let [acc] = program.fold( - tile_ir::tile::range(reduce_size), - [tile_ir::tile::Tile::literal(tile_literal_for( - initial, - DataTypeEnum::F32, - ))], - |program, loop_index, [acc]| { - [acc.binary( - reduce_binary, - value_at(program, loop_index).into_f32(), - )] - }, - ); - ValueTile::F32(acc) - } - DataTypeEnum::F16 => { - let [acc] = program.fold( - tile_ir::tile::range(reduce_size), - [tile_ir::tile::Tile::literal(tile_literal_for( - initial, - DataTypeEnum::F16, - ))], - |program, loop_index, [acc]| { - [acc.binary( - reduce_binary, - value_at(program, loop_index).into_f16(), - )] - }, - ); - ValueTile::F16(acc) - } - DataTypeEnum::U32 => { - let [acc] = program.fold( - tile_ir::tile::range(reduce_size), - [tile_ir::tile::Tile::literal(tile_literal_for( - initial, - DataTypeEnum::U32, - ))], - |program, loop_index, [acc]| { - [acc.binary( - reduce_binary, - value_at(program, loop_index).into_u32(), - )] - }, - ); - ValueTile::U32(acc) - } - }; - - let (reduced, reduced_ty) = apply_unary_function_chain( - reduced.into_f32(), - reduce_dtype, - &post_chain, - ) - .expect("validated reduce post_element_wise chain"); - let reduced = ValueTile::F32(reduced) - .cast_to(reduced_ty) - .cast_to(output_meta_body.datatype); - let output_index = layout_index(&output_meta_body, &dims); - output_storage.store(program, output_index, reduced, in_bounds); + let acc = fold_reduce_lane( + program, + &dims_out, + tile_ir::tile::range(k_iterations), + &in_bounds, + SUBGROUP_VALUES_PER_LANE, + true, + &|loop_index, run| { + loop_index.clone() * k_per_iter.clone() + + sg_lane.clone() * SUBGROUP_VALUES_PER_LANE + + run + }, + ); + let reduced = token.subgroup_reduce(program, reduce_op, acc); + let store_mask = in_bounds.and(sg_lane.eq(0u32)); + store_reduced(program, tag(reduced), &dims_out, store_mask); }); + } + ReduceRoute::Serial => { + kb.program() + .program_grid(BLOCK as u32, dispatch_size, |program| { + let lane = program.lane(); + let group = linear_group(program, dispatch_size); + let flat = group * BLOCK as u32 + lane.clone(); + let in_bounds = flat.lt(total_outputs); + let dims_out = output_dims_from_flat_usize(flat.clone(), &output_shape); + + let acc = fold_reduce_lane( + program, + &dims_out, + tile_ir::tile::range(reduce_size), + &in_bounds, + 1, + false, + &|loop_index, _| loop_index.clone(), + ); + store_reduced(program, tag(acc), &dims_out, in_bounds); + }); + } } Some(()) }, ) } -fn tile_literal_for(value: NaryScalar, target: DataTypeEnum) -> tile_ir::TileLiteral { +pub(crate) fn tile_literal_for(value: NaryScalar, target: DataTypeEnum) -> tile_ir::TileLiteral { match target { DataTypeEnum::F32 => match value { NaryScalar::F32(value) => tile_ir::TileLiteral::f32(value), @@ -326,7 +478,7 @@ fn tile_literal_for(value: NaryScalar, target: DataTypeEnum) -> tile_ir::TileLit } } -fn output_dims_from_flat_usize( +pub(crate) fn output_dims_from_flat_usize( flat: tile_ir::tile::Tile, shape: &[u32], ) -> Vec { @@ -334,7 +486,7 @@ fn output_dims_from_flat_usize( output_dims_from_flat(flat, &shape) } -fn tile_reduce_op(op: ReduceOp) -> tile_ir::TileReduceOp { +pub(crate) fn tile_reduce_op(op: ReduceOp) -> tile_ir::TileReduceOp { match op { ReduceOp::Sum => tile_ir::TileReduceOp::Sum, ReduceOp::Product => tile_ir::TileReduceOp::Product, diff --git a/fusor-ml/core/src/reduce_tiled.rs b/fusor-ml/core/src/reduce_tiled.rs new file mode 100644 index 000000000..64c7fe3ee --- /dev/null +++ b/fusor-ml/core/src/reduce_tiled.rs @@ -0,0 +1,1007 @@ +//! Tiled lowering for fused map-reduce operations. +//! +//! When the fused producer's index structure shows 2D reuse — every +//! reduce-axis-dependent input misses at least one of two parallel dims — the +//! reduction lowers as a workgroup-tiled kernel: each k-tile of every staged +//! input is loaded cooperatively into workgroup memory once and reused by all +//! lanes, and each thread accumulates a TM×TN register tile. A composed +//! matmul is exactly this shape (`a[.., m, k]` misses `n`, `b[.., k, n]` +//! misses `m`), so the tiled matmul kernel falls out of the caching the tiled +//! loads provide — as does any other contraction the tensor API composes, +//! whatever its dim order, pre-scaling, or surrounding expression. + +use fusor_tile_ir as tile_ir; +use std::hash::Hash; +use tile_ir::tile::{Mask, Tile}; + +use crate::{ + access_analysis::InputAccesses, + mir::{ + inputs::MirValue, kernel_backend, kernel_backend::DirectKernel, operation::Operation, + workgroup_shape::WorkgroupShape, + }, + nary_direct::{ + TensorMeta, ValueTile, apply_unary_function_chain, declare_value, + eval_nary_expr_on_value_tiles, layout_index, linear_group, tile_u32, zero_literal, + }, + nary_wise::{NaryExpr, NaryOp, NaryScalar}, + reduce::{ReduceOp, ReduceOperation}, + reduce_direct::{ReduceKernelInputs, tile_literal_for, tile_reduce_op}, + tensor::DataTypeEnum, + visit_tiled::{MaybeQData, distribute_workgroups}, +}; + +/// Workgroup tile geometry. One fixed shape keeps plan selection +/// deterministic; the gates below reject shapes it cannot cover profitably. +const BM: u32 = 32; +const BN: u32 = 32; +const BK: u32 = 8; +const TM: u32 = 4; +const TN: u32 = 4; +const LANES: u32 = (BM / TM) * (BN / TN); + +/// Metadata for dependence analysis. Block-quantized matrices address by +/// (row, col) with a zero in the kernel meta's column slot; the *analysis* +/// must still see the column as a real dependence, so it gets full row-major +/// strides here. +pub(crate) fn analysis_meta(value: &MaybeQData) -> Option { + match value { + MaybeQData::Tensor(tensor) => TensorMeta::new(tensor), + MaybeQData::QMatrix(matrix) => { + let mut meta = TensorMeta::for_matrix(matrix)?; + let mut acc = 1u32; + for dim in (0..meta.shape.len()).rev() { + meta.strides[dim] = acc; + acc = acc.checked_mul(meta.shape[dim])?; + } + Some(meta) + } + } +} + +/// Load one input value at `coords`, dequantizing block-quantized inputs +/// through the format-aware per-element path. +fn load_input_value( + program: &mut tile_ir::tile::TileBlock<'_>, + storage: &crate::nary_direct::Storage2, + meta: &TensorMeta, + coords: &[Tile], + mask: Mask, +) -> ValueTile { + if let crate::nary_direct::Storage2::Quantized(matrix) = storage { + let along_row = coords.last().cloned().unwrap_or_else(|| tile_u32(0)); + let which_row = layout_index(meta, coords); + return ValueTile::F32(program.load_quantized(matrix, along_row, which_row, mask, 0.0)); + } + storage.load(program, layout_index(meta, coords), mask) +} + +/// The allocation footprint one input re-streams per redundant read. +pub(crate) fn input_allocation_bytes(meta: &TensorMeta, value: &MaybeQData) -> u64 { + let element_bytes = match value { + MaybeQData::Tensor(tensor) => match tensor.datatype() { + DataTypeEnum::F16 => 2, + DataTypeEnum::F32 | DataTypeEnum::U32 => 4, + }, + MaybeQData::QMatrix(_) => 4, + }; + meta.allocation_len as u64 * element_bytes +} + +struct ReduceTiledKernelVariant; + +/// How a staged input's workgroup tile is addressed. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] +enum StagedKind { + /// References the row dim and the reduce axis: `[BM, BK]` tile. + RowK, + /// References the col dim and the reduce axis: `[BK, BN]` tile. + KCol, + /// References only the reduce axis (plus grid dims): `[BK]` tile. + KOnly, +} + +impl StagedKind { + fn tile_elements(self) -> u32 { + match self { + StagedKind::RowK => BM * BK, + StagedKind::KCol => BK * BN, + StagedKind::KOnly => BK, + } + } +} + +#[derive(Clone, Debug)] +enum InputRole { + /// Reduce-axis-dependent: staged through workgroup memory each k-tile. + Staged(StagedKind), + /// Independent of the reduce axis: loaded once per thread, outside the + /// k loop. + Direct, +} + +#[derive(Clone, Debug)] +struct PlannedInput { + /// Index-space dim read by each input dimension (pure `DimIndex` only). + dims: Vec, + role: InputRole, +} + +/// Where staged inputs are cached between reuses. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] +enum Staging { + /// Full tiles big enough to amortize barriers: k-tiles staged + /// cooperatively through workgroup memory. + Workgroup, + /// Partial or skinny tiles: each thread keeps its TM row values and TN + /// column values in registers per k step, so the reuse across its + /// register tile survives without barriers. + Register, +} + +/// Per-workgroup output-tile geometry. Always 64 lanes: +/// `(bm / tm) * (bn / tn) == LANES`. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] +struct TileGeom { + bm: u32, + bn: u32, + tm: u32, + tn: u32, +} + +impl TileGeom { + /// The 2D contraction tile: matches the dense matmul geometry. + const PAIR: Self = Self { + bm: BM, + bn: BN, + tm: TM, + tn: TN, + }; + /// The 1D tile: one parallel dim, each thread covering `tm` outputs so + /// row-missing inputs load once per k step instead of `tm` times. + const SINGLE: Self = Self { + bm: 256, + bn: 1, + tm: 4, + tn: 1, + }; + + fn outs(self) -> u32 { + self.tm * self.tn + } +} + +struct TiledReducePlan { + /// The parallel dim(s) tiled across the workgroup. `col_dim` is `None` + /// for 1D tiling: a single reuse dim with no second axis to pair it + /// with. + row_dim: usize, + col_dim: Option, + geom: TileGeom, + /// Every other non-reduce dim: one coordinate per workgroup. + grid_dims: Vec, + inputs: Vec, + staging: Staging, + /// Out-of-bounds index-space slots evaluate to the reduction identity + /// with zero-filled tiles (a bare multiply chain summed), so the inner + /// loop can skip per-element masking — the tiled matmul fast path. + fill_neutral: bool, +} + +/// Out-of-bounds slots collapse to the reduction identity without masking iff +/// the expression is a flat multiply folded with `Sum` whose factors map a +/// zero-filled load to zero (and stay finite): every staged tile zero-fills +/// its out-of-bounds slots, every out-of-bounds index-space slot is out of +/// bounds for at least one staged input, and one zero factor kills the +/// product. Factors may wrap their load in zero-preserving unary ops (casts, +/// negation, finite constant scales) — the dense matmul's f16→f32 accumulate +/// casts and scale epilogues keep the fast path. +fn is_fill_neutral(expr: &NaryExpr, op: ReduceOp) -> bool { + if op != ReduceOp::Sum { + return false; + } + fn scalar_is_finite(scalar: NaryScalar) -> bool { + match scalar { + NaryScalar::F32(value) => value.is_finite(), + NaryScalar::F16(value) => value.is_finite(), + NaryScalar::U32(_) => true, + } + } + fn factor_is_zero_preserving(expr: &NaryExpr) -> bool { + match expr { + NaryExpr::IndexedInput { .. } => true, + NaryExpr::Scalar(scalar) => scalar_is_finite(*scalar), + NaryExpr::Op { children, function } if children.len() == 1 => { + let zero_preserving = matches!( + function.op, + NaryOp::Cast | NaryOp::Neg | NaryOp::Abs + ) || matches!(function.op, NaryOp::MulConst(scalar) if scalar_is_finite(scalar)); + zero_preserving && factor_is_zero_preserving(&children[0]) + } + _ => false, + } + } + fn flat_mul_factors_collapse(expr: &NaryExpr) -> bool { + match expr { + NaryExpr::Op { children, function } + if function.op == NaryOp::Mul && children.len() == 2 => + { + children.iter().all(flat_mul_factors_collapse) + } + other => factor_is_zero_preserving(other), + } + } + matches!(expr, NaryExpr::Op { function, .. } if function.op == NaryOp::Mul) + && flat_mul_factors_collapse(expr) +} + +fn plan_tiled_reduce( + operation: &ReduceOperation, + values: &[MaybeQData], + device: &crate::Device, +) -> Option { + let shape = &operation.shape; + let axis = operation.axis; + let rank = shape.len(); + // One parallel dim + the reduce axis is enough for 1D tiling; the 2D + // pair search simply finds nothing below rank 3. + if rank < 2 { + return None; + } + let metas: Vec = values.iter().map(analysis_meta).collect::>()?; + let access = InputAccesses::collect(&operation.expression, operation.inputs.len(), &metas)?; + + // Pick the (row, col) pair: both must be fed by at least one staged + // input, and no staged input may reference both (its tile would need all + // three tiled dims). Among valid pairs prefer the most square, largest + // coverage; ties resolve to the lowest dim indices for determinism. + let k_dep: Vec = (0..access.effective.len()) + .map(|i| access.depends_on(i, axis)) + .collect(); + if !k_dep.iter().any(|dep| *dep) { + return None; + } + let mut best: Option<((usize, usize), (usize, usize))> = None; + for row_dim in 0..rank { + for col_dim in 0..rank { + if row_dim == col_dim || row_dim == axis || col_dim == axis { + continue; + } + let mut row_fed = false; + let mut col_fed = false; + let mut valid = true; + for (i, k_dep) in k_dep.iter().enumerate() { + if !k_dep { + continue; + } + let has_row = access.depends_on(i, row_dim); + let has_col = access.depends_on(i, col_dim); + if has_row && has_col { + valid = false; + break; + } + row_fed |= has_row; + col_fed |= has_col; + } + if !valid || !row_fed || !col_fed { + continue; + } + let (m, n) = (shape[row_dim], shape[col_dim]); + let score = (m.min(n), m * n); + if best + .as_ref() + .is_none_or(|(best_score, _)| score > *best_score) + { + best = Some((score, (row_dim, col_dim))); + } + } + } + // No pair: fall back to a single reuse dim — a parallel dim some + // k-dependent input misses. Each thread then covers `tm` outputs along + // it, so the missing input streams once per thread instead of `tm` + // times. Trading `tm`× thread count for that reuse only pays when the + // missed inputs are too large for cache to absorb the re-streams + // (cache-resident reuse is free, and the occupancy loss is not), so the + // missed inputs must clear the cache-thrash threshold. + let (row_dim, col_dim) = match best { + Some((_, (row_dim, col_dim))) => (row_dim, Some(col_dim)), + None => { + let mut best_single: Option<((u64, usize), usize)> = None; + for d in 0..rank { + if d == axis || shape[d] < TileGeom::SINGLE.bm as usize { + continue; + } + let missed_bytes: u64 = (0..k_dep.len()) + .filter(|&i| k_dep[i] && !access.depends_on(i, d)) + .map(|i| input_allocation_bytes(&metas[i], &values[i])) + .sum(); + if missed_bytes < device.last_level_cache_bytes() { + continue; + } + let score = (missed_bytes, shape[d]); + if best_single + .as_ref() + .is_none_or(|(best_score, _)| score > *best_score) + { + best_single = Some((score, d)); + } + } + let (_, d) = best_single?; + (d, None) + } + }; + let geom = match col_dim { + Some(_) => TileGeom::PAIR, + None => TileGeom::SINGLE, + }; + + let m: u32 = shape[row_dim].try_into().ok()?; + let n: u32 = match col_dim { + Some(col_dim) => shape[col_dim].try_into().ok()?, + None => 1, + }; + let k: u32 = shape[axis].try_into().ok()?; + let limits = device.limits(); + if LANES > limits.max_compute_workgroup_size_x + || LANES > limits.max_compute_invocations_per_workgroup + { + return None; + } + let grid_dims: Vec = (0..rank) + .filter(|&d| d != axis && d != row_dim && Some(d) != col_dim) + .collect(); + let batch: u32 = grid_dims.iter().try_fold(1u32, |acc, &d| { + acc.checked_mul(u32::try_from(shape[d]).ok()?) + })?; + let tiles_m = m.div_ceil(geom.bm); + let tiles_n = n.div_ceil(geom.bn); + let workgroups = batch.checked_mul(tiles_m)?.checked_mul(tiles_n)?; + let actual_outputs = batch.checked_mul(m)?.checked_mul(n)?; + let covered_outputs = workgroups.checked_mul(geom.bm * geom.bn)?; + + let mut workgroup_bytes = 0u32; + let mut inputs = Vec::with_capacity(access.dims.len()); + for (i, (dims, k_dep)) in access.dims.iter().zip(&k_dep).enumerate() { + let value = &values[i]; + let role = if *k_dep { + let has_col = col_dim.is_some_and(|col_dim| access.depends_on(i, col_dim)); + let kind = match (access.depends_on(i, row_dim), has_col) { + (true, false) => StagedKind::RowK, + (false, true) => StagedKind::KCol, + (false, false) => StagedKind::KOnly, + (true, true) => unreachable!("pair selection rejects row+col staged inputs"), + }; + let element_bytes = match value { + MaybeQData::Tensor(tensor) => match tensor.datatype() { + DataTypeEnum::F16 => 2, + DataTypeEnum::F32 | DataTypeEnum::U32 => 4, + }, + MaybeQData::QMatrix(_) => 4, + }; + workgroup_bytes = + workgroup_bytes.checked_add(kind.tile_elements().checked_mul(element_bytes)?)?; + InputRole::Staged(kind) + } else { + InputRole::Direct + }; + inputs.push(PlannedInput { + dims: dims.clone(), + role, + }); + } + + // Workgroup staging needs full-enough 2D tiles to amortize its barriers: + // the same shape and 75%-utilization gates as the dedicated dense matmul + // tile selection. Anything else keeps the register tile, whose reuse + // needs no barriers and tolerates partial coverage. + let workgroup_eligible = col_dim.is_some() + && m >= BM + && n >= BN + && k >= BK + && (actual_outputs as u64) * 4 >= (covered_outputs as u64) * 3 + && workgroup_bytes <= limits.max_compute_workgroup_storage_size; + let staging = if workgroup_eligible { + Staging::Workgroup + } else { + Staging::Register + }; + + let fill_neutral = is_fill_neutral(&operation.expression, operation.function.op); + + Some(TiledReducePlan { + row_dim, + col_dim, + geom, + grid_dims, + inputs, + staging, + fill_neutral, + }) +} + +/// The full index-space coordinate vector for one (batch, row, col, k) +/// position, in dim order. +fn full_coords( + plan: &TiledReducePlan, + rank: usize, + axis: usize, + batch_coords: &[Tile], + row: &Tile, + col: &Tile, + k: &Tile, +) -> Vec { + (0..rank) + .map(|d| { + if d == axis { + k.clone() + } else if d == plan.row_dim { + row.clone() + } else if Some(d) == plan.col_dim { + col.clone() + } else { + let slot = plan.grid_dims.iter().position(|&g| g == d).unwrap(); + batch_coords[slot].clone() + } + }) + .collect() +} + +fn input_coords(input: &PlannedInput, coords: &[Tile]) -> Vec { + input.dims.iter().map(|&d| coords[d].clone()).collect() +} + +pub(crate) fn datatype_scalar(datatype: DataTypeEnum) -> tile_ir::ScalarElement { + match datatype { + DataTypeEnum::F32 => tile_ir::ScalarElement::F32, + DataTypeEnum::F16 => tile_ir::ScalarElement::F16, + DataTypeEnum::U32 => tile_ir::ScalarElement::U32, + } +} + +fn value_tile(datatype: DataTypeEnum, value: Tile) -> ValueTile { + match datatype { + DataTypeEnum::F32 => ValueTile::F32(value), + DataTypeEnum::F16 => ValueTile::F16(value), + DataTypeEnum::U32 => ValueTile::U32(value), + } +} + +pub(crate) fn build_reduce_tiled_kernel( + operation: &ReduceOperation, + graph: &crate::compute_graph::ComputeGraphInner, + workgroup_shape: &WorkgroupShape, + inputs: &[MirValue], +) -> Option { + let parsed = ReduceKernelInputs::parse(operation, graph, inputs)?; + let device = graph.device(); + let plan = plan_tiled_reduce(operation, &parsed.values, &device)?; + if std::env::var_os("FUSOR_TRACE_REDUCE_TILED").is_some() { + eprintln!( + "reduce_tiled row={} col={:?} axis={} staging={:?} fill_neutral={} staged={} name={}", + plan.row_dim, + plan.col_dim, + operation.axis, + plan.staging, + plan.fill_neutral, + plan.inputs + .iter() + .filter(|input| matches!(input.role, InputRole::Staged(_))) + .count(), + operation.name(), + ); + } + + let rank = operation.shape.len(); + let axis = operation.axis; + let geom = plan.geom; + let m: u32 = operation.shape[plan.row_dim].try_into().ok()?; + let n: u32 = match plan.col_dim { + Some(col_dim) => operation.shape[col_dim].try_into().ok()?, + None => 1, + }; + let k: u32 = operation.shape[axis].try_into().ok()?; + let batch_sizes: Vec = plan + .grid_dims + .iter() + .map(|&d| u32::try_from(operation.shape[d]).ok()) + .collect::>()?; + let batch: u32 = batch_sizes.iter().product(); + let tiles_m = m.div_ceil(geom.bm); + let tiles_n = n.div_ceil(geom.bn); + let total_tiles = batch * tiles_m * tiles_n; + let k_tiles = k.div_ceil(BK); + + let dispatch_size = distribute_workgroups( + total_tiles, + device.limits().max_compute_workgroups_per_dimension, + ); + let variant = + kernel_backend::KernelVariantKey::with_payload::(|state| { + plan.row_dim.hash(state); + plan.col_dim.hash(state); + plan.geom.hash(state); + plan.staging.hash(state); + plan.fill_neutral.hash(state); + for input in &plan.inputs { + input.dims.hash(state); + match input.role { + InputRole::Staged(kind) => kind.hash(state), + InputRole::Direct => "direct".hash(state), + } + } + }); + let cache_key = operation.kernel_cache_key_with_dispatch( + variant, + Some(workgroup_shape), + dispatch_size, + inputs, + ); + + let reduce_dtype = operation.function.datatype(); + let reduce_op = tile_reduce_op(operation.function.op); + let initial = operation.function.initial_value; + let expression = operation.expression.clone(); + let post_chain = operation.post_element_wise.clone(); + let output = inputs.last()?.as_tensor()?.clone(); + let output_dtype = output.datatype(); + let output_value = MaybeQData::Tensor(output); + let values = parsed.values; + let plan = std::sync::Arc::new(plan); + + kernel_backend::run_kernel( + device.kernel_cache(), + format!("{}_tiled", operation.name()), + cache_key, + dispatch_size, + move |kb| { + let mut storages = Vec::with_capacity(values.len()); + let mut metas = Vec::with_capacity(values.len()); + for value in &values { + let (storage, meta) = declare_value(kb, value, false)?; + storages.push(storage); + metas.push(meta); + } + let (output_storage, output_meta) = declare_value(kb, &output_value, true)?; + + let phase = kb.program(); + let staged_tiles: Vec> = plan + .inputs + .iter() + .zip(&metas) + .map(|(input, meta)| match (plan.staging, &input.role) { + (Staging::Workgroup, InputRole::Staged(kind)) => { + let element = datatype_scalar(meta.datatype); + Some(match kind { + StagedKind::RowK => phase.alloc_workgroup_tile(element, BM, BK), + StagedKind::KCol => phase.alloc_workgroup_tile(element, BK, BN), + StagedKind::KOnly => phase.alloc_workgroup_array(element, BK), + }) + } + _ => None, + }) + .collect(); + + phase.program_grid(LANES, dispatch_size, |program| { + let tile_id = linear_group(program, dispatch_size); + let tile_active = tile_id.clone().lt(total_tiles); + let batch_id = tile_id.clone() / (tiles_m * tiles_n); + let local_tile = tile_id % (tiles_m * tiles_n); + let m_tile = local_tile.clone() / tiles_n; + let n_tile = local_tile % tiles_n; + + // Mixed-radix batch decomposition over the grid dims, in dim + // order (earlier dims vary slowest). + let batch_coords: Vec = (0..batch_sizes.len()) + .map(|i| { + let size = batch_sizes[i]; + if size == 1 { + return tile_u32(0); + } + let divisor: u32 = batch_sizes[i + 1..].iter().product(); + let quotient = if divisor == 1 { + batch_id.clone() + } else { + batch_id.clone() / divisor + }; + quotient % size + }) + .collect(); + + let lane = program.lane(); + let lane_row = program.bind(lane.clone() / (geom.bn / geom.tn)); + let lane_col = program.bind(lane % (geom.bn / geom.tn)); + let m_tile_base = program.bind(m_tile * geom.bm); + let n_tile_base = program.bind(n_tile * geom.bn); + let row_base = program.bind(m_tile_base.clone() + lane_row.clone() * geom.tm); + let col_base = program.bind(n_tile_base.clone() + lane_col.clone() * geom.tn); + + let identity = || Tile::literal(tile_literal_for(initial, reduce_dtype)); + + // Per-thread output-slot bounds, hoisted out of the k loop. + let rc_masks: Vec = (0..geom.outs()) + .map(|idx| { + let (r, c) = (idx / geom.tn, idx % geom.tn); + let row = row_base.clone() + r; + let col = col_base.clone() + c; + tile_active.clone().and(row.lt(m)).and(col.lt(n)) + }) + .collect(); + + // Direct (k-independent) inputs: one bound load per output + // slot, reused across the whole k loop. + let direct_values: Vec>> = plan + .inputs + .iter() + .enumerate() + .map(|(i, input)| { + if !matches!(input.role, InputRole::Direct) { + return None; + } + Some( + (0..geom.outs()) + .map(|idx| { + let (r, c) = (idx / geom.tn, idx % geom.tn); + let row = row_base.clone() + r; + let col = col_base.clone() + c; + let coords = full_coords( + &plan, + rank, + axis, + &batch_coords, + &row, + &col, + &tile_u32(0), + ); + let index = + layout_index(&metas[i], &input_coords(input, &coords)); + let loaded = storages[i].load( + program, + index, + rc_masks[idx as usize].clone(), + ); + // Bind to a local so the k loop reuses + // the value instead of re-issuing the + // storage load every iteration. + let native = match loaded { + ValueTile::F32(v) + | ValueTile::F16(v) + | ValueTile::U32(v) => v, + ValueTile::Bool(_) => { + unreachable!("tensor inputs are f32/f16/u32") + } + }; + ( + value_tile(metas[i].datatype, program.bind(native)), + metas[i].datatype, + ) + }) + .collect(), + ) + }) + .collect(); + + let initial_accs: Vec = (0..geom.outs()).map(|_| identity()).collect(); + let sums = match plan.staging { + Staging::Register => program.fold_vec( + tile_ir::tile::range(k), + initial_accs, + |program, k_index, accs| { + // Per staged input, the register slice this thread + // reuses across its TM×TN tile: TM row values, TN + // column values, or one k-only value. + let reg_values: Vec>> = plan + .inputs + .iter() + .enumerate() + .map(|(i, input)| { + let InputRole::Staged(kind) = input.role else { + return None; + }; + let count = match kind { + StagedKind::RowK => geom.tm, + StagedKind::KCol => geom.tn, + StagedKind::KOnly => 1, + }; + Some( + (0..count) + .map(|j| { + let (row, col, mask) = match kind { + StagedKind::RowK => { + let row = row_base.clone() + j; + let mask = tile_active + .clone() + .and(row.clone().lt(m)); + (row, tile_u32(0), mask) + } + StagedKind::KCol => { + let col = col_base.clone() + j; + let mask = tile_active + .clone() + .and(col.clone().lt(n)); + (tile_u32(0), col, mask) + } + StagedKind::KOnly => ( + tile_u32(0), + tile_u32(0), + tile_active.clone(), + ), + }; + let coords = full_coords( + &plan, + rank, + axis, + &batch_coords, + &row, + &col, + &k_index, + ); + let index = layout_index( + &metas[i], + &input_coords(input, &coords), + ); + let loaded = + storages[i].load(program, index, mask.clone()); + let value = match loaded { + ValueTile::F32(v) + | ValueTile::F16(v) + | ValueTile::U32(v) => Tile::select( + mask, + v, + Tile::literal(zero_literal( + metas[i].datatype, + )), + ), + ValueTile::Bool(_) => unreachable!( + "tensor inputs are f32/f16/u32" + ), + }; + ( + value_tile(metas[i].datatype, value), + metas[i].datatype, + ) + }) + .collect(), + ) + }) + .collect(); + + // Every k step is in range and out-of-bounds + // output slots never store, so no accumulate + // masking is needed in register mode. + accs.into_iter() + .enumerate() + .map(|(idx, acc)| { + let (r, c) = (idx as u32 / geom.tn, idx as u32 % geom.tn); + let slot_values: Vec<(ValueTile, DataTypeEnum)> = plan + .inputs + .iter() + .enumerate() + .map(|(i, input)| match input.role { + InputRole::Staged(kind) => { + let j = match kind { + StagedKind::RowK => r, + StagedKind::KCol => c, + StagedKind::KOnly => 0, + }; + reg_values[i].as_ref().unwrap()[j as usize].clone() + } + InputRole::Direct => { + direct_values[i].as_ref().unwrap()[idx].clone() + } + }) + .collect(); + let (value, _) = + eval_nary_expr_on_value_tiles(&expression, &slot_values); + let value = match value.cast_to(reduce_dtype) { + ValueTile::F32(v) + | ValueTile::F16(v) + | ValueTile::U32(v) => v, + ValueTile::Bool(_) => { + unreachable!("reduce dtype is f32/f16/u32") + } + }; + acc.binary(reduce_op.binary(), value) + }) + .collect() + }, + ), + Staging::Workgroup => program.fold_vec( + tile_ir::tile::range(k_tiles), + initial_accs, + |program, k_tile, accs| { + let k_base = program.bind(k_tile * BK); + + // Cooperative staging: every lane copies its share of + // each staged input's k-tile into workgroup memory. + for (i, input) in plan.inputs.iter().enumerate() { + let InputRole::Staged(kind) = input.role else { + continue; + }; + let tile = staged_tiles[i].as_ref().unwrap(); + let elements = kind.tile_elements(); + for pass in 0..elements.div_ceil(LANES) { + let flat = program.lane() + pass * LANES; + let (row, col, k_index, flat_ok) = match kind { + StagedKind::RowK => { + let local_row = flat.clone() / BK; + let local_k = flat.clone() % BK; + ( + m_tile_base.clone() + local_row, + tile_u32(0), + k_base.clone() + local_k, + Mask::all(), + ) + } + StagedKind::KCol => { + let local_k = flat.clone() / BN; + let local_col = flat.clone() % BN; + ( + tile_u32(0), + n_tile_base.clone() + local_col, + k_base.clone() + local_k, + Mask::all(), + ) + } + StagedKind::KOnly => ( + tile_u32(0), + tile_u32(0), + k_base.clone() + flat.clone(), + flat.clone().lt(elements), + ), + }; + let mut in_bounds = + tile_active.clone().and(flat_ok).and(k_index.clone().lt(k)); + if matches!(kind, StagedKind::RowK) { + in_bounds = in_bounds.and(row.clone().lt(m)); + } + if matches!(kind, StagedKind::KCol) { + in_bounds = in_bounds.and(col.clone().lt(n)); + } + let coords = full_coords( + &plan, + rank, + axis, + &batch_coords, + &row, + &col, + &k_index, + ); + let loaded = load_input_value( + program, + &storages[i], + &metas[i], + &input_coords(input, &coords), + in_bounds.clone(), + ); + // Zero-fill out-of-bounds slots: in the + // fill-neutral fast path that zero is what + // collapses the product, elsewhere the inner + // mask discards the slot anyway. + let value = match loaded.cast_to(metas[i].datatype) { + ValueTile::F32(v) + | ValueTile::F16(v) + | ValueTile::U32(v) => Tile::select( + in_bounds, + v, + Tile::literal(zero_literal(metas[i].datatype)), + ), + ValueTile::Bool(_) => { + unreachable!("tensor inputs are f32/f16/u32") + } + }; + program.store_workgroup(tile, flat, value); + } + } + program.workgroup_barrier(); + + let chunk_sums: Vec = (0..geom.outs()) + .map(|idx| { + let (r, c) = (idx / geom.tn, idx % geom.tn); + let local_row = lane_row.clone() * geom.tm + r; + let local_col = lane_col.clone() * geom.tn + c; + let mut chunk = identity(); + for kk in 0..BK { + let slot_values: Vec<(ValueTile, DataTypeEnum)> = plan + .inputs + .iter() + .enumerate() + .map(|(i, input)| match input.role { + InputRole::Staged(kind) => { + let tile = staged_tiles[i].as_ref().unwrap(); + let flat = match kind { + StagedKind::RowK => { + local_row.clone() * BK + kk + } + StagedKind::KCol => { + local_col.clone() + kk * BN + } + StagedKind::KOnly => tile_u32(kk), + }; + ( + value_tile( + metas[i].datatype, + program.load_workgroup(tile, flat), + ), + metas[i].datatype, + ) + } + InputRole::Direct => { + direct_values[i].as_ref().unwrap()[idx as usize] + .clone() + } + }) + .collect(); + let (value, _) = eval_nary_expr_on_value_tiles( + &expression, + &slot_values, + ); + let value = match value.cast_to(reduce_dtype) { + ValueTile::F32(v) + | ValueTile::F16(v) + | ValueTile::U32(v) => v, + ValueTile::Bool(_) => { + unreachable!("reduce dtype is f32/f16/u32") + } + }; + let value = if plan.fill_neutral { + value + } else { + let valid = rc_masks[idx as usize] + .clone() + .and((k_base.clone() + kk).lt(k)); + Tile::select(valid, value, identity()) + }; + chunk = chunk.binary(reduce_op.binary(), value); + } + chunk + }) + .collect(); + // Bind every chunk before the trailing barrier: the + // next iteration's staging overwrites the tiles these + // reads source from. + let chunk_sums: Vec = chunk_sums + .into_iter() + .map(|chunk| program.bind(chunk)) + .collect(); + program.workgroup_barrier(); + accs.into_iter() + .zip(chunk_sums) + .map(|(acc, chunk)| acc.binary(reduce_op.binary(), chunk)) + .collect() + }, + ), + }; + + for (idx, sum) in sums.into_iter().enumerate() { + let idx = idx as u32; + let (r, c) = (idx / geom.tn, idx % geom.tn); + let row = row_base.clone() + r; + let col = col_base.clone() + c; + let (value, value_ty) = apply_unary_function_chain( + value_tile(reduce_dtype, sum).into_f32(), + reduce_dtype, + &post_chain, + ) + .expect("validated reduce post_element_wise chain"); + let value = ValueTile::F32(value) + .cast_to(value_ty) + .cast_to(output_dtype); + let coords = + full_coords(&plan, rank, axis, &batch_coords, &row, &col, &tile_u32(0)); + let out_coords: Vec = (0..rank) + .filter(|&d| d != axis) + .map(|d| coords[d].clone()) + .collect(); + let output_index = layout_index(&output_meta, &out_coords); + output_storage.store( + program, + output_index, + value, + rc_masks[idx as usize].clone(), + ); + } + }); + Some(()) + }, + ) +} diff --git a/fusor-ml/core/src/resize.rs b/fusor-ml/core/src/resize.rs deleted file mode 100644 index 93c722266..000000000 --- a/fusor-ml/core/src/resize.rs +++ /dev/null @@ -1,400 +0,0 @@ -use std::hash::Hash; - -use crate::{ - DataTypeEnum, Layout, TILE_SIZE, Tensor, TensorData, - compute_graph::NodeIndex, - map_layout::MapLayoutOperation, - mir::{ - inputs::MirValue, - kernel_backend::DirectKernel, - operation::Operation, - workgroup_shape::{Constraint, WorkgroupShape, WorkgroupShapeConstraints}, - }, - nary_wise::{NaryExpr, NaryOp, NaryOperation, NaryScalar}, - visit_tiled::distribute_workgroups, -}; - -const BLOCKSIZE: u32 = 256; - -#[derive(Debug, Clone)] -pub(crate) struct ResizeOperation { - pub(crate) input: NodeIndex, - pub(crate) current_shape: Box<[usize]>, - pub(crate) new_shape: Box<[usize]>, - pub(crate) fill_shape: Box<[usize]>, -} - -impl ResizeOperation { - pub fn new( - input: NodeIndex, - current_shape: Box<[usize]>, - new_shape: Box<[usize]>, - fill_shape: Box<[usize]>, - ) -> Self { - Self { - input, - current_shape, - new_shape, - fill_shape, - } - } -} - -impl ResizeOperation { - pub(crate) fn lower( - &self, - graph: &crate::compute_graph::ComputeGraphInner, - ) -> Option { - let full_fill = self.fill_shape == self.new_shape; - let matching_size = self.current_shape.iter().product::() - == self.new_shape.iter().product::(); - if !full_fill || !matching_size { - return None; - } - - let input = graph.get_cached_result(self.input)?; - let input_layout = input.layout(); - if !is_row_major_contiguous(input_layout) { - return None; - } - - // Find the chunks of strides that are contiguous in the input - let mut contiguous_stride_chunks = Vec::new(); - for (stride, len) in input_layout - .strides() - .iter() - .rev() - .zip(input_layout.shape().iter().rev()) - { - let Some((last_stride, last_len)) = contiguous_stride_chunks.last_mut() else { - contiguous_stride_chunks.push((*stride, *len)); - continue; - }; - let contiguous_stride = *last_stride * *last_len; - if *stride == contiguous_stride { - *last_len *= len; - } else { - contiguous_stride_chunks.push((*stride, *len)); - } - } - - // Check if the new shape can be formed by combining contiguous chunks - // of the input shape - let new_shape = self.new_shape.clone(); - let mut new_strides = Vec::new(); - let offset = input_layout.offset(); - let mut contiguous_stride_chunks_iter = contiguous_stride_chunks.iter_mut().peekable(); - for shape in new_shape.iter().rev() { - // If we've used up this chunk and the current shape dimension is more than 1, - // move to the next chunk - while contiguous_stride_chunks_iter - .next_if(|(_, len)| *len == 1 && *shape > 1) - .is_some() - {} - let (stride, len) = contiguous_stride_chunks_iter.peek_mut()?; - // Make sure the current chunk can be divided to form the new shape - if *len % *shape != 0 { - return None; - } - *len /= *shape; - new_strides.push(*stride); - *stride *= *shape; - } - new_strides.reverse(); - - Some(MapLayoutOperation::new(self.input, move |_layout| { - Layout::from_parts(offset, new_shape.clone(), new_strides.as_slice().into()) - })) - } - - fn copy_expression(&self) -> Option { - let flat = row_major_flat_expr(&self.fill_shape)?; - let input_indices = row_major_indices_from_flat(flat, &self.current_shape)?; - Some(NaryExpr::indexed_input(0, input_indices)) - } - - fn in_fill_bounds_expression(&self) -> NaryExpr { - let mut condition = NaryExpr::scalar(NaryScalar::U32(1)); - for (dim, &fill) in self.fill_shape.iter().enumerate() { - let lt_fill = NaryExpr::unary_op( - NaryExpr::DimIndex(dim), - "lt_fill", - NaryOp::LessConst(NaryScalar::U32(fill as u32)), - DataTypeEnum::U32, - DataTypeEnum::U32, - ); - condition = NaryExpr::mul(condition, lt_fill, DataTypeEnum::U32); - } - condition - } - - fn zero_expression(datatype: DataTypeEnum) -> NaryExpr { - let zero = match datatype { - DataTypeEnum::F32 => NaryScalar::F32(0.0), - DataTypeEnum::F16 => NaryScalar::F16(half::f16::from_f32(0.0)), - DataTypeEnum::U32 => NaryScalar::U32(0), - }; - NaryExpr::scalar(zero) - } - - fn expression(&self, datatype: DataTypeEnum) -> Option { - let copied = self.copy_expression()?; - if self.fill_shape == self.new_shape { - return Some(copied); - } - Some(NaryExpr::select( - self.in_fill_bounds_expression(), - copied, - Self::zero_expression(datatype), - DataTypeEnum::U32, - datatype, - )) - } -} - -fn is_row_major_contiguous(layout: &Layout) -> bool { - let mut expected_stride = 1usize; - for (dim, stride) in layout.shape().iter().zip(layout.strides()).rev() { - if *dim > 1 && *stride != expected_stride { - return false; - } - expected_stride = expected_stride.saturating_mul(*dim); - } - true -} - -impl Operation for ResizeOperation { - fn hash_kernel_fields(&self, state: &mut rustc_hash::FxHasher) { - self.current_shape.hash(state); - self.new_shape.hash(state); - self.fill_shape.hash(state); - } - - fn workgroup_shape_constraints( - &self, - _: &crate::Device, - ) -> crate::mir::workgroup_shape::WorkgroupShapeConstraints { - let mut constraints = WorkgroupShapeConstraints::new(); - constraints.add_constraint(0, Constraint::equals(BLOCKSIZE)); - constraints.add_constraint(1, Constraint::equals(1)); - constraints.add_constraint(2, Constraint::equals(1)); - constraints - } - - fn dispatch_size( - &self, - _: &crate::mir::workgroup_shape::WorkgroupShape, - inputs: &[crate::mir::inputs::MirValue], - ) -> [u32; 3] { - let output = inputs[1].as_tensor().unwrap(); - let total_workgroups = (output.layout().shape().iter().product::() as u32) - .div_ceil(TILE_SIZE * BLOCKSIZE); - distribute_workgroups( - total_workgroups, - output - .device() - .limits() - .max_compute_workgroups_per_dimension, - ) - } - - fn visit_dependencies(&self, f: &mut dyn FnMut(NodeIndex)) { - f(self.input); - } - - fn inputs( - &self, - nodes: &crate::compute_graph::ComputeGraphInner, - ) -> Vec { - let input = nodes.get_cached_result(self.input).unwrap().clone(); - let output = TensorData::new_for_shape(input.device(), &self.new_shape, input.datatype()); - vec![input.into(), output.into()] - } - - fn build_direct_kernel( - &self, - graph: &crate::compute_graph::ComputeGraphInner, - workgroup_shape: &WorkgroupShape, - inputs: &[MirValue], - ) -> Option { - let input = inputs[0].as_tensor()?; - let operation = NaryOperation { - inputs: vec![self.input], - expression: self.expression(input.datatype())?, - shape: self.new_shape.clone(), - output_datatype: input.datatype(), - }; - crate::nary_direct::build_nary_direct_kernel_to_output( - &operation, - graph, - workgroup_shape, - inputs, - 1, - ) - } - - fn output( - &self, - _: &crate::compute_graph::ComputeGraphInner, - inputs: &[crate::mir::inputs::MirValue], - ) -> crate::mir::inputs::MirValue { - let output = inputs[1].as_tensor().unwrap(); - TensorData::new_from_buffer( - output.device(), - output.buffer().clone(), - &self.new_shape, - output.datatype(), - ) - .into() - } - - fn name(&self) -> String { - format!( - "resize_from_{}_to_{}", - self.current_shape - .iter() - .map(|x| x.to_string()) - .collect::>() - .join("x"), - self.new_shape - .iter() - .map(|x| x.to_string()) - .collect::>() - .join("x") - ) - } -} - -fn row_major_flat_expr(shape: &[usize]) -> Option { - let mut flat = NaryExpr::scalar(NaryScalar::U32(0)); - for axis in 0..shape.len() { - let stride = shape[axis + 1..] - .iter() - .try_fold(1u32, |acc, dim| acc.checked_mul((*dim).try_into().ok()?))?; - let dim = NaryExpr::DimIndex(axis); - let term = if stride == 1 { - dim - } else { - NaryExpr::unary_op( - dim, - "mul_const", - NaryOp::MulConst(NaryScalar::U32(stride)), - DataTypeEnum::U32, - DataTypeEnum::U32, - ) - }; - flat = NaryExpr::add(flat, term, DataTypeEnum::U32); - } - Some(flat) -} - -fn row_major_indices_from_flat(flat: NaryExpr, shape: &[usize]) -> Option> { - let mut indices = Vec::with_capacity(shape.len()); - for axis in 0..shape.len() { - let divisor = shape[axis + 1..] - .iter() - .try_fold(1u32, |acc, dim| acc.checked_mul((*dim).try_into().ok()?))?; - let dim = u32::try_from(shape[axis]).ok()?; - let quotient = if divisor == 1 { - flat.clone() - } else { - NaryExpr::unary_op( - flat.clone(), - "div_const", - NaryOp::DivConst(NaryScalar::U32(divisor)), - DataTypeEnum::U32, - DataTypeEnum::U32, - ) - }; - indices.push(if dim == 1 { - NaryExpr::scalar(NaryScalar::U32(0)) - } else { - NaryExpr::unary_op( - quotient, - "rem_const", - NaryOp::RemConst(NaryScalar::U32(dim)), - DataTypeEnum::U32, - DataTypeEnum::U32, - ) - }); - } - Some(indices) -} - -impl Tensor { - pub fn resize(&self, new_shape: impl AsRef<[usize]>) -> Tensor { - let new_shape: Box<[usize]> = new_shape.as_ref().into(); - let input = self.key(); - self.add_resize(ResizeOperation::new( - input, - self.shape().into(), - new_shape, - self.shape().into(), - )) - } - - pub fn reshape(&self, new_shape: impl AsRef<[usize]>) -> Tensor { - let new_shape = new_shape.as_ref(); - assert_eq!( - new_shape.iter().product::(), - self.shape().iter().product::(), - "Reshape requires the number of elements to be the same. \ - Current shape: {:?}, target shape: {:?}", - self.shape(), - new_shape - ); - let new_shape: Box<[usize]> = new_shape.into(); - let input = self.key(); - self.add_resize(ResizeOperation::new( - input, - self.shape().into(), - new_shape.clone(), - new_shape.clone(), - )) - } - - pub fn flatten_last_n(&self, from_end: usize) -> Tensor { - assert!( - from_end < self.rank(), - "flatten_last_n FROM_END must be less than input rank" - ); - let out_rank = self.rank() - from_end; - let new_shape: Vec = (0..out_rank) - .map(|i| { - if i < self.rank() - 1 - from_end { - self.shape()[i] - } else if i == self.rank() - 1 - from_end { - self.shape()[i..].iter().product() - } else { - 1 - } - }) - .collect(); - self.reshape(new_shape) - } - - pub fn flatten_first_n(&self, from_start: usize) -> Tensor { - assert!( - from_start < self.rank(), - "flatten_first_n FROM_START must be less than input rank" - ); - let out_rank = self.rank() - from_start; - let new_shape: Vec = (0..out_rank) - .map(|i| { - if i == 0 { - self.shape()[..=from_start].iter().product() - } else { - self.shape()[i + from_start] - } - }) - .collect(); - self.reshape(new_shape) - } - - pub fn flatten_all(&self) -> Tensor { - let size = self.shape().iter().product(); - self.reshape([size]) - } -} - -pub use fusor_types::ShapeWithOneHole; diff --git a/fusor-ml/core/src/rms_norm.rs b/fusor-ml/core/src/rms_norm.rs deleted file mode 100644 index 326831f34..000000000 --- a/fusor-ml/core/src/rms_norm.rs +++ /dev/null @@ -1 +0,0 @@ -pub(crate) use crate::mir::kernel_backend::rms_norm::RmsNormOperation; diff --git a/fusor-ml/core/src/row_program.rs b/fusor-ml/core/src/row_program.rs new file mode 100644 index 000000000..f300ecf3a --- /dev/null +++ b/fusor-ml/core/src/row_program.rs @@ -0,0 +1,1247 @@ +//! Generic lowering for fused row programs. +//! +//! A row program is a cluster of same-axis reductions and elementwise +//! expressions over one axis of an index space. Each phase is either a +//! scalar reduction (folding an expression over the axis into a per-row +//! value visible to later phases) or a staged per-element value (evaluated +//! once at each axis position, optionally folding a private inner axis — an +//! inline dot product). The output either maps one element per index-space +//! position (softmax, RMS norm) or folds the axis once more with a free +//! output dimension appended to the row shape (attention's `Σ p·v`). +//! +//! One workgroup runs one row. Map-style programs stride the axis in chunks; +//! programs with element phases or a reducing output pin one lane per axis +//! position (the axis must fit one workgroup), which lets the score of a +//! decode-attention row live in a register across every phase. The axis +//! length may be dynamic: kernels compile per block bucket and read the +//! active length from a trailing u32 params input, so a growing KV cache +//! stays in one cached kernel. + +use std::{any::TypeId, hash::Hash}; + +use fusor_tile_ir as tile_ir; +use rustc_hash::FxHasher; +use tile_ir::{ + ElementType, ScalarElement, + tile::{Mask, Tile, WorkgroupTile}, +}; + +use crate::{ + compute_graph::{GraphOperation, NodeIndex}, + mir::{ + inputs::MirValue, + kernel_backend::{self, DirectKernel}, + operation::Operation, + workgroup_shape::{Constraint, WorkgroupShape, WorkgroupShapeConstraints}, + }, + nary_direct::{ + ValueTile, apply_unary_function_chain, declare_value, eval_nary_expr, layout_index, + linear_group, output_dims_from_flat, tile_u32, + }, + nary_wise::{NaryExpr, NaryFunction, NaryOp, NaryScalar, UnaryFunctionChain}, + reduce::{ReduceFunction, ReduceOp, max_fn, sum_fn}, + reduce_direct::{tile_literal_for, tile_reduce_op}, + tensor::{DataTypeEnum, TensorData}, + visit_tiled::{MaybeQData, distribute_workgroups}, +}; + +const BLOCK: u32 = 256; + +/// Below this many rows a long axis is split across workgroups (one tile +/// each) with a combine kernel folding the spans — decode has too few rows +/// to fill the device with one workgroup per row. +const SPLIT_ROWS_TARGET: u32 = 256; + +/// Workgroup buckets for dynamic-axis row programs, smallest first. The +/// kernel monomorphizes per bucket; the active axis length rides in the +/// params input. +const ROW_DYNAMIC_BLOCKS: [u32; 4] = [128, 256, 512, 1024]; + +/// One reduction phase: `expression` (over the external inputs and the +/// slots of earlier phases) folded along the row axis, then `post_chain` +/// applied to the combined value once per row. +#[derive(Debug, Clone, Hash)] +pub(crate) struct RowReduce { + pub(crate) expression: NaryExpr, + pub(crate) function: ReduceFunction, + pub(crate) post_chain: UnaryFunctionChain, +} + +/// The private inner axis of an element phase: the expression is evaluated +/// `len` times with the fold coordinate as `DimIndex(rank)` and accumulated +/// with `function` — an inline dot product per axis position. +#[derive(Debug, Clone, Hash)] +pub(crate) struct RowFold { + pub(crate) len: usize, + pub(crate) function: ReduceFunction, +} + +#[derive(Debug, Clone, Hash)] +pub(crate) enum RowPhase { + /// Fold over the row axis into a per-row scalar slot. + Reduce(RowReduce), + /// A staged per-element value: evaluated once at each axis position, + /// with one lane pinned per position (the axis must fit the workgroup). + Element { + expression: NaryExpr, + fold: Option, + datatype: DataTypeEnum, + }, +} + +#[derive(Debug, Clone, Hash)] +pub(crate) enum RowOutput { + /// One output element per index-space position; output shape == `shape`. + Map(NaryExpr), + /// Fold the axis once more with one free output dimension appended to + /// the row shape: `combine` sees the free coordinate as `DimIndex(rank)` + /// and element-phase slots at the folded axis position. Output shape = + /// row dims ++ `[free_dim]`. + Reduce { + combine: NaryExpr, + function: ReduceFunction, + free_dim: usize, + }, +} + +/// Dynamic-axis configuration: the kernel is compiled for the `block` +/// capacity bucket and reads the active axis length from a trailing u32 +/// params input, so per-token axis growth (the KV cache) reuses one kernel. +#[derive(Debug, Clone, Hash)] +pub(crate) struct DynamicAxis { + /// Workgroup size and per-tile capacity; axis lengths beyond it stream + /// through the online tile loop. + pub(crate) block: u32, + /// For each input, the input dimension whose extent tracks the axis + /// (the KV dim of K/V) — normalized out of the kernel cache key. + pub(crate) input_axis_dims: Vec>, + /// A row dimension whose coordinate bounds the active axis length: + /// `effective_len = min(axis_len, coord + 1)` — causal attention skips + /// every tile past the query position. + pub(crate) axis_bound_dim: Option, +} + +/// The algebraic shape the online-streaming lowering requires: a staged +/// score element, a max phase over a scaled score `e`, an exp-sum phase +/// shifted by that max, a staged probability element, and a linear combine. +/// The exp shift is what licenses streaming — rescaling the running sum and +/// accumulators by `exp(M_old − M_new)` keeps them exact across tiles. +struct OnlineSoftmax<'a> { + /// Element phase: the raw per-position score (slot `n`). + score: (&'a NaryExpr, &'a Option), + /// The scaled/masked score expression `e` over slot `n` (max phase body). + scaled: &'a NaryExpr, + max_identity: NaryScalar, + /// The per-position weight in `combine = p · weight`. + weight: &'a NaryExpr, +} + +fn slot_expr(input_count: usize, phase: usize) -> NaryExpr { + NaryExpr::IndexedInput { + input_idx: input_count + phase, + indices: vec![], + } +} + +fn unary_op_child<'a>(expr: &'a NaryExpr, op: &NaryOp) -> Option<&'a NaryExpr> { + match expr { + NaryExpr::Op { children, function } if function.op == *op && children.len() == 1 => { + Some(&children[0]) + } + _ => None, + } +} + +fn binary_op_children<'a>(expr: &'a NaryExpr, op: &NaryOp) -> Option<(&'a NaryExpr, &'a NaryExpr)> { + match expr { + NaryExpr::Op { children, function } if function.op == *op && children.len() == 2 => { + Some((&children[0], &children[1])) + } + _ => None, + } +} + +/// Match the four-phase online-softmax shape (see [`OnlineSoftmax`]). The +/// attention constructor builds exactly this; the lowering re-derives it so +/// the phase expressions stay the single source of truth. +fn match_online_softmax<'a>( + phases: &'a [RowPhase], + output: &'a RowOutput, + input_count: usize, +) -> Option> { + let [ + RowPhase::Element { + expression: score, + fold: score_fold, + .. + }, + RowPhase::Reduce(max_phase), + RowPhase::Reduce(sum_phase), + RowPhase::Element { + expression: prob, + fold: None, + .. + }, + ] = phases + else { + return None; + }; + if max_phase.function.op != ReduceOp::Max + || sum_phase.function.op != ReduceOp::Sum + || !max_phase.post_chain.functions.is_empty() + || !sum_phase.post_chain.functions.is_empty() + { + return None; + } + let scaled = &max_phase.expression; + // sum phase: exp(e − m) + let (shift_lhs, shift_rhs) = binary_op_children( + unary_op_child(&sum_phase.expression, &NaryOp::Exp)?, + &NaryOp::Sub, + )?; + if shift_lhs != scaled || *shift_rhs != slot_expr(input_count, 1) { + return None; + } + // prob element: exp(e − m) / l + let (num, denom) = binary_op_children(prob, &NaryOp::Div)?; + if num != &sum_phase.expression || *denom != slot_expr(input_count, 2) { + return None; + } + // combine: p · weight, summed + let RowOutput::Reduce { + combine, function, .. + } = output + else { + return None; + }; + if function.op != ReduceOp::Sum { + return None; + } + let (p_ref, weight) = binary_op_children(combine, &NaryOp::Mul)?; + if *p_ref != slot_expr(input_count, 3) { + return None; + } + // The weight and scaled score may only reference tensor inputs, dims, + // and the score slot — never later phase slots (those are consumed by + // the streaming structure itself). + for later in 1..4 { + if weight.uses_input(input_count + later) || scaled.uses_input(input_count + later) { + return None; + } + } + if weight.uses_input(input_count) { + return None; + } + Some(OnlineSoftmax { + score: (score, score_fold), + scaled, + max_identity: max_phase.function.initial_value, + weight, + }) +} + +/// Slot convention for every expression in the program: indices below +/// `inputs.len()` are tensor reads; index `inputs.len() + p` is phase `p`'s +/// value (a per-row scalar for reduce phases, a per-element value for +/// element phases). +#[derive(Debug, Clone)] +pub(crate) struct RowProgramOperation { + pub(crate) inputs: Vec, + /// The full row-parallel index space (including the axis). + pub(crate) shape: Box<[usize]>, + pub(crate) axis: usize, + pub(crate) phases: Vec, + pub(crate) output: RowOutput, + pub(crate) output_datatype: DataTypeEnum, + /// `Some` exactly when the program pins one lane per axis position + /// (element phases / reducing output); `None` for chunked map programs. + pub(crate) dynamic_axis: Option, +} + +impl RowProgramOperation { + pub(crate) fn rows(&self) -> usize { + self.shape + .iter() + .enumerate() + .filter_map(|(dim, &size)| (dim != self.axis).then_some(size)) + .product() + } + + fn row_shape(&self) -> Vec { + self.shape + .iter() + .enumerate() + .filter_map(|(dim, &size)| (dim != self.axis).then_some(size)) + .collect() + } + + pub(crate) fn out_shape(&self) -> Vec { + match &self.output { + RowOutput::Map(_) => self.shape.to_vec(), + RowOutput::Reduce { free_dim, .. } => { + let mut shape = self.row_shape(); + shape.push(*free_dim); + shape + } + } + } + + fn block(&self, device: &crate::Device) -> u32 { + match &self.dynamic_axis { + Some(dynamic) => dynamic.block, + None => device.limits().max_compute_workgroup_size_x.min(BLOCK), + } + } +} + +impl Operation for RowProgramOperation { + fn hash_kernel_fields(&self, state: &mut FxHasher) { + // With a dynamic axis the kernel is bucketed by `block`; the actual + // axis extent rides in the params input and must stay out of the + // key, or every generated token would recompile. + match &self.dynamic_axis { + Some(dynamic) => { + 1u8.hash(state); + dynamic.hash(state); + for (dim, size) in self.shape.iter().enumerate() { + if dim != self.axis { + size.hash(state); + } + } + } + None => { + 0u8.hash(state); + self.shape.hash(state); + } + } + self.axis.hash(state); + self.phases.hash(state); + self.output.hash(state); + self.output_datatype.hash(state); + } + + fn workgroup_shape_constraints(&self, device: &crate::Device) -> WorkgroupShapeConstraints { + let mut constraints = WorkgroupShapeConstraints::new(); + constraints.add_constraint(0, Constraint::equals(self.block(device))); + constraints.add_constraint(1, Constraint::equals(1)); + constraints.add_constraint(2, Constraint::equals(1)); + constraints + } + + fn dispatch_size(&self, _workgroup_shape: &WorkgroupShape, inputs: &[MirValue]) -> [u32; 3] { + let output: TensorData = inputs.last().unwrap().as_tensor().unwrap().clone(); + distribute_workgroups( + self.rows() as u32, + output + .device() + .limits() + .max_compute_workgroups_per_dimension, + ) + } + + fn visit_dependencies(&self, f: &mut dyn FnMut(NodeIndex)) { + for input in &self.inputs { + f(*input); + } + } + + fn inputs(&self, nodes: &crate::compute_graph::ComputeGraphInner) -> Vec { + let mut mir_inputs: Vec = self + .inputs + .iter() + .map(|idx| nodes.get_result_or_qmatrix(*idx).unwrap().into()) + .collect(); + let device = match &mir_inputs[0] { + MirValue::Tensor(tensor) => tensor.device().clone(), + MirValue::QMatrix(matrix) => matrix.device().clone(), + _ => unreachable!("row program inputs are tensors or quantized matrices"), + }; + if self.dynamic_axis.is_some() { + mir_inputs + .push(TensorData::new_splat(&device, &[1], self.shape[self.axis] as u32).into()); + } + let output_tensor = + TensorData::new_for_shape(&device, &self.out_shape(), self.output_datatype); + mir_inputs.push(output_tensor.into()); + mir_inputs + } + + fn build_direct_kernel( + &self, + graph: &crate::compute_graph::ComputeGraphInner, + workgroup_shape: &WorkgroupShape, + inputs: &[MirValue], + ) -> Option { + build_row_program_kernel(self, graph, workgroup_shape, inputs) + } + + fn output(&self, _: &crate::compute_graph::ComputeGraphInner, inputs: &[MirValue]) -> MirValue { + inputs.last().unwrap().clone() + } + + fn name(&self) -> String { + format!( + "row_program_{}p_{}", + self.phases.len(), + self.shape + .iter() + .map(|x| x.to_string()) + .collect::>() + .join("x") + ) + } +} + +impl GraphOperation for RowProgramOperation { + fn category(&self) -> &'static str { + "row_program" + } +} + +struct RowProgramKernelVariant; + +/// The shared cache-key recipe, except inputs with a dynamic-axis dimension +/// hash a normalized extent: the KV cache's strides and offset are stable +/// across tokens, so bucketing out the growing dim keeps one kernel per +/// block size. +fn row_program_cache_key( + operation: &RowProgramOperation, + workgroup_shape: &WorkgroupShape, + dispatch_size: [u32; 3], + inputs: &[MirValue], + role: u64, + splits: u32, +) -> kernel_backend::KernelCacheKey { + kernel_backend::KernelCacheKey::from_hash_inputs(|state| { + 11u64.hash(state); + role.hash(state); + splits.hash(state); + kernel_backend::KernelVariantKey::of::().hash(state); + TypeId::of::().hash(state); + operation.hash_kernel_fields(state); + workgroup_shape.shape().hash(state); + dispatch_size.hash(state); + inputs.len().hash(state); + for (i, input) in inputs.iter().enumerate() { + let dynamic_dim = operation + .dynamic_axis + .as_ref() + .and_then(|dynamic| dynamic.input_axis_dims.get(i).copied().flatten()); + std::mem::discriminant(input).hash(state); + match input { + MirValue::Tensor(tensor) => { + tensor.datatype().hash(state); + let layout = tensor.layout(); + layout.offset().hash(state); + for (dim, (size, stride)) in layout + .shape() + .iter() + .zip(layout.strides().iter()) + .enumerate() + { + if Some(dim) == dynamic_dim { + 0usize.hash(state); + } else { + size.hash(state); + } + stride.hash(state); + } + } + MirValue::QMatrix(matrix) => { + matrix.datatype().hash(state); + matrix.storage_layout().hash(state); + matrix.shape().hash(state); + } + MirValue::Integer(value) => value.hash(state), + MirValue::Float(value) => value.to_bits().hash(state), + } + } + }) +} + +fn raw_tile(value: ValueTile) -> Tile { + match value { + ValueTile::F32(tile) | ValueTile::F16(tile) | ValueTile::U32(tile) => tile, + ValueTile::Bool(_) => unreachable!("row program values are f32/f16/u32"), + } +} + +fn f32_literal(value: f32) -> Tile { + Tile::literal(tile_ir::TileLiteral::f32(value)) +} + +fn build_row_program_kernel( + operation: &RowProgramOperation, + graph: &crate::compute_graph::ComputeGraphInner, + workgroup_shape: &WorkgroupShape, + inputs: &[MirValue], +) -> Option { + let (output, producers) = inputs.split_last()?; + let output = output.as_tensor()?.clone(); + let (params, producers) = if operation.dynamic_axis.is_some() { + let (params, producers) = producers.split_last()?; + (Some(params.as_tensor()?.clone()), producers) + } else { + (None, producers) + }; + let values = producers + .iter() + .map(|input| MaybeQData::try_from(input.clone()).ok()) + .collect::>>()?; + if !graph.device().f16_supported() { + let uses_f16 = output.datatype() == DataTypeEnum::F16 + || values.iter().any(|value| match value { + MaybeQData::Tensor(tensor) => tensor.datatype() == DataTypeEnum::F16, + MaybeQData::QMatrix(matrix) => { + matches!(matrix.datatype(), fusor_gguf::GgmlType::F16) + } + }); + if uses_f16 { + return None; + } + } + + let rows: u32 = operation.rows().try_into().ok()?; + let k: u32 = operation.shape[operation.axis].try_into().ok()?; + let row_shape = operation.row_shape(); + let axis = operation.axis; + let rank = operation.shape.len(); + let block = workgroup_shape.x(); + let lanes_own_axis = operation.dynamic_axis.is_some(); + + // Long axes with few rows fan out across workgroups: each split runs + // the online body over one tile, writing its unnormalized accumulator + // and softmax statistics to scratch; a combine kernel folds the spans + // with the online monoid. + let free_dim_out = match &operation.output { + RowOutput::Reduce { free_dim, .. } => Some(*free_dim), + RowOutput::Map(_) => None, + }; + let tiles = k.div_ceil(block); + let splits: u32 = match free_dim_out { + Some(free) + if lanes_own_axis + && rows < SPLIT_ROWS_TARGET + && tiles > 1 + && tiles <= block + && (free as u32 + 2) <= block => + { + tiles + } + _ => 1, + }; + let max_dispatch_dim = graph.device().limits().max_compute_workgroups_per_dimension; + let dispatch_size = distribute_workgroups(rows * splits, max_dispatch_dim); + let cache_key = + row_program_cache_key(operation, workgroup_shape, dispatch_size, inputs, 1, splits); + + let input_count = operation.inputs.len(); + let axis_bound_dim = operation + .dynamic_axis + .as_ref() + .and_then(|dynamic| dynamic.axis_bound_dim); + let phases = operation.phases.clone(); + let output_kind = operation.output.clone(); + let output_dtype = output.datatype(); + let output_value = MaybeQData::Tensor(output); + let params_value = params.map(MaybeQData::Tensor); + + let scratch_value = (splits > 1).then(|| { + MaybeQData::Tensor(TensorData::new_for_shape( + &graph.device(), + &[ + rows as usize, + splits as usize, + free_dim_out.expect("split row programs reduce") + 2, + ], + DataTypeEnum::F32, + )) + }); + let online_max_identity = match phases.get(1) { + Some(RowPhase::Reduce(reduce)) => Some(reduce.function.initial_value), + _ => None, + }; + let combine = if splits > 1 { + let scratch_b = scratch_value.clone().expect("split scratch"); + let output_b = output_value.clone(); + let row_shape_b = row_shape.clone(); + let dispatch_b = distribute_workgroups(rows, max_dispatch_dim); + let free = free_dim_out.expect("split row programs reduce") as u32; + let max_identity_scalar = + online_max_identity.expect("split row programs carry a max phase"); + let key_b = + row_program_cache_key(operation, workgroup_shape, dispatch_b, inputs, 2, splits); + let combine = kernel_backend::run_kernel( + graph.device().kernel_cache(), + format!("{}_combine", operation.name()), + key_b, + dispatch_b, + move |kb| { + let (scratch_storage, scratch_meta) = declare_value(kb, &scratch_b, false)?; + let (output_storage, output_meta) = declare_value(kb, &output_b, true)?; + kb.program().program_grid(block, dispatch_b, |program| { + let lane = program.lane(); + let row_flat = linear_group(program, dispatch_b); + let in_bounds = row_flat.clone().lt(rows); + let row_dims = output_dims_from_flat(row_flat.clone(), &row_shape_b); + let max_identity = + Tile::literal(tile_literal_for(max_identity_scalar, DataTypeEnum::F32)); + + // Lane = one span: fold the spans' maxima and rescaled + // sums into the row's softmax statistics. + let j_active = in_bounds.clone().and(lane.clone().lt(splits)); + let m_index = layout_index( + &scratch_meta, + &[row_flat.clone(), lane.clone(), tile_u32(free + 1)], + ); + let m_j = program.bind(raw_tile(scratch_storage.load( + program, + m_index, + j_active.clone(), + ))); + let masked_m = Tile::select(j_active.clone(), m_j.clone(), max_identity); + let global_max = + program.group_reduce(tile_reduce_op(ReduceOp::Max), block, masked_m); + let global_max = program.bind(global_max); + let l_index = layout_index( + &scratch_meta, + &[row_flat.clone(), lane.clone(), tile_u32(free)], + ); + let l_j = raw_tile(scratch_storage.load(program, l_index, j_active.clone())); + let weighted = Tile::select( + j_active, + l_j * (m_j - global_max.clone()).exp(), + f32_literal(0.0), + ); + let denom = + program.group_reduce(tile_reduce_op(ReduceOp::Sum), block, weighted); + let denom = program.bind(denom); + + // Lane = one free-dim position: rescale and fold the + // spans' accumulators. + let acc = program.private(ElementType::F32); + program.store_local(&acc, f32_literal(0.0)); + let out_active = in_bounds.and(lane.clone().lt(free)); + program.if_then(out_active, |program| { + program.loop_range(splits, |program, j| { + let m_index = layout_index( + &scratch_meta, + &[row_flat.clone(), j.clone(), tile_u32(free + 1)], + ); + let m = raw_tile(scratch_storage.load(program, m_index, Mask::all())); + let o_index = + layout_index(&scratch_meta, &[row_flat.clone(), j, lane.clone()]); + let o = raw_tile(scratch_storage.load(program, o_index, Mask::all())); + let current = program.load_local(&acc); + program.store_local(&acc, current + o * (m - global_max.clone()).exp()); + }); + let value = program.load_local(&acc) / denom.clone(); + let mut out_coords = row_dims.clone(); + out_coords.push(lane.clone()); + let output_index = layout_index(&output_meta, &out_coords); + output_storage.store( + program, + output_index, + ValueTile::F32(value).cast_to(output_dtype), + Mask::all(), + ); + }); + }); + Some(()) + }, + )?; + Some(combine) + } else { + None + }; + + let partials = kernel_backend::run_kernel( + graph.device().kernel_cache(), + operation.name(), + cache_key, + dispatch_size, + move |kb| { + let mut storages = Vec::with_capacity(values.len()); + let mut metas = Vec::with_capacity(values.len()); + for value in &values { + let (storage, meta) = declare_value(kb, value, false)?; + storages.push(storage); + metas.push(meta); + } + let params_storage = match ¶ms_value { + Some(value) => Some(declare_value(kb, value, false)?), + None => None, + }; + let (output_storage, output_meta) = match &scratch_value { + Some(scratch) => declare_value(kb, scratch, true)?, + None => declare_value(kb, &output_value, true)?, + }; + + let phase_handle = kb.program(); + // Element-phase values the reducing output reads across lanes are + // staged through workgroup memory (the probs of decode attention). + let staged: Vec> = phases + .iter() + .enumerate() + .map(|(p, phase)| match (&output_kind, phase) { + (RowOutput::Reduce { combine, .. }, RowPhase::Element { .. }) + if combine.uses_input(input_count + p) => + { + Some(phase_handle.alloc_workgroup_array(ScalarElement::F32, block)) + } + _ => None, + }) + .collect(); + + phase_handle.program_grid(block, dispatch_size, |program| { + let lane = program.lane(); + let wg_flat = linear_group(program, dispatch_size); + let (row_flat, split_idx) = if splits > 1 { + ( + program.bind(wg_flat.clone() / splits), + program.bind(wg_flat % splits), + ) + } else { + (wg_flat, tile_u32(0)) + }; + let in_bounds = row_flat.clone().lt(rows); + let row_dims = output_dims_from_flat(row_flat.clone(), &row_shape); + let full_coords = |k_index: Tile| -> Vec { + let mut coords = Vec::with_capacity(rank); + let mut row_dim = 0; + for dim in 0..rank { + if dim == axis { + coords.push(k_index.clone()); + } else { + coords.push(row_dims[row_dim].clone()); + row_dim += 1; + } + } + coords + }; + + // The active axis length: a params read for dynamic-axis + // programs, the compiled extent otherwise. + let axis_len: Tile = match ¶ms_storage { + Some((storage, _)) => raw_tile(storage.load(program, tile_u32(0), Mask::all())), + None => tile_u32(k), + }; + + let mut slots: Vec<(ValueTile, DataTypeEnum)> = Vec::new(); + + if lanes_own_axis { + // Online streaming over axis tiles of `block`: lanes own + // one axis position per tile; the running max, sum, and + // free-dim accumulator are rescaled by + // `exp(M_old − M_new)` each tile, so any axis length + // streams through one workgroup with exact results. + let RowOutput::Reduce { free_dim, .. } = &output_kind else { + unreachable!("dynamic-axis row programs have a reducing output") + }; + let online = match_online_softmax(&phases, &output_kind, input_count) + .expect("dynamic-axis row programs are built in online-softmax shape"); + let probs = staged[3] + .as_ref() + .expect("the probability element phase is staged"); + let free = *free_dim as u32; + + // Causal bound: tiles past the bounding row coordinate + // hold no live positions, so the loop ends there. + let effective_len: Tile = match axis_bound_dim { + Some(dim) => { + let row_index = dim - usize::from(axis < dim); + program.bind( + axis_len + .clone() + .min(row_dims[row_index].clone() + tile_u32(1)), + ) + } + None => axis_len, + }; + + // A split workgroup owns one `block`-wide span of the + // axis; the single-workgroup form owns it all. + let (span_start, span_end) = if splits > 1 { + let start = program.bind(split_idx.clone() * block); + let end = program + .bind((start.clone() + tile_u32(block)).min(effective_len.clone())); + (start, end) + } else { + (tile_u32(0), effective_len.clone()) + }; + + let max_identity = + || Tile::literal(tile_literal_for(online.max_identity, DataTypeEnum::F32)); + let running_max = program.private(ElementType::F32); + let running_sum = program.private(ElementType::F32); + let acc = program.private(ElementType::F32); + let tile_base = program.private(ElementType::U32); + let item = program.private(ElementType::U32); + program.store_local(&running_max, max_identity()); + program.store_local(&running_sum, f32_literal(0.0)); + program.store_local(&acc, f32_literal(0.0)); + program.store_local(&tile_base, span_start); + let out_active = in_bounds.clone().and(lane.clone().lt(free)); + + program.loop_forever(|program| { + let base = program.load_local(&tile_base); + program.break_if(base.clone().ge(span_end.clone())); + let kv = program.bind(base.clone() + lane.clone()); + let kv_active = in_bounds.clone().and(kv.clone().lt(effective_len.clone())); + let coords = full_coords(kv); + + // Per-position score, optionally an inline fold (the + // q·k dot over the head dim). + let (score_expr, score_fold) = online.score; + let score = match score_fold { + Some(RowFold { + len, + function: fold_fn, + }) => { + let fold_dtype = fold_fn.datatype(); + let identity = Tile::literal(tile_literal_for( + fold_fn.initial_value, + fold_dtype, + )); + let reduce_op = tile_reduce_op(fold_fn.op); + let [folded] = program.fold( + tile_ir::tile::range(*len as u32), + [identity], + |program, fold_idx, [fold_acc]| { + let mut fold_coords = coords.clone(); + fold_coords.push(fold_idx); + let (value, _) = eval_nary_expr( + program, + score_expr, + &fold_coords, + &storages, + &metas, + kv_active.clone(), + &slots, + ); + let value = raw_tile(value.cast_to(fold_dtype)); + [fold_acc.binary(reduce_op.binary(), value)] + }, + ); + folded + } + None => { + let (value, _) = eval_nary_expr( + program, + score_expr, + &coords, + &storages, + &metas, + kv_active.clone(), + &slots, + ); + raw_tile(value.cast_to(DataTypeEnum::F32)) + } + }; + let score = program.bind(score); + let score_slot = [(ValueTile::F32(score), DataTypeEnum::F32)]; + + // Scaled/masked score, tile max, and the online + // rescale of the running state. + let (scaled, _) = eval_nary_expr( + program, + online.scaled, + &coords, + &storages, + &metas, + kv_active.clone(), + &score_slot, + ); + let scaled = program.bind(Tile::select( + kv_active.clone(), + raw_tile(scaled.cast_to(DataTypeEnum::F32)), + max_identity(), + )); + let tile_max = program.group_reduce( + tile_reduce_op(ReduceOp::Max), + block, + scaled.clone(), + ); + let old_max = program.load_local(&running_max); + let new_max = program.bind(old_max.clone().max(tile_max)); + program.store_local(&running_max, new_max.clone()); + let factor = program.bind((old_max - new_max.clone()).exp()); + + let prob = program.bind(Tile::select( + kv_active.clone(), + (scaled - new_max).exp(), + f32_literal(0.0), + )); + let tile_sum = program.group_reduce( + tile_reduce_op(ReduceOp::Sum), + block, + prob.clone(), + ); + let sum = program.load_local(&running_sum); + program.store_local(&running_sum, sum * factor.clone() + tile_sum); + program.store_workgroup(probs, lane.clone(), prob); + program.workgroup_barrier(); + + // Lanes switch to owning one free-dim position each + // and fold this tile's staged probabilities against + // the weight input. + program.if_then(out_active.clone(), |program| { + let rescaled = program.load_local(&acc) * factor.clone(); + program.store_local(&acc, rescaled); + program.store_local(&item, tile_u32(0)); + program.loop_forever(|program| { + let j = program.load_local(&item); + let kv_j = program.bind(base.clone() + j.clone()); + program.break_if( + j.clone() + .ge(block) + .or(kv_j.clone().ge(effective_len.clone())), + ); + let prob_j = program.load_workgroup(probs, j.clone()); + let mut weight_coords = full_coords(kv_j); + weight_coords.push(lane.clone()); + let (weight, _) = eval_nary_expr( + program, + online.weight, + &weight_coords, + &storages, + &metas, + Mask::all(), + &slots, + ); + let weight = raw_tile(weight.cast_to(DataTypeEnum::F32)); + let current = program.load_local(&acc); + program.store_local(&acc, current + prob_j * weight); + program.store_local(&item, j + tile_u32(1)); + }); + }); + program.workgroup_barrier(); + program.store_local(&tile_base, base + tile_u32(block)); + }); + + if splits > 1 { + // Partials: the unnormalized accumulator plus this + // span's softmax statistics — the combine kernel + // folds the spans with the online monoid. + program.if_then(out_active, |program| { + let index = layout_index( + &output_meta, + &[row_flat.clone(), split_idx.clone(), lane.clone()], + ); + let value = ValueTile::F32(program.load_local(&acc)); + output_storage.store(program, index, value, Mask::all()); + }); + let stat_active = |offset: u32| { + in_bounds + .clone() + .and(lane.clone().ge(free + offset)) + .and(lane.clone().lt(free + offset + 1)) + }; + program.if_then(stat_active(0), |program| { + let index = layout_index( + &output_meta, + &[row_flat.clone(), split_idx.clone(), tile_u32(free)], + ); + let value = ValueTile::F32(program.load_local(&running_sum)); + output_storage.store(program, index, value, Mask::all()); + }); + program.if_then(stat_active(1), |program| { + let index = layout_index( + &output_meta, + &[row_flat.clone(), split_idx.clone(), tile_u32(free + 1)], + ); + let value = ValueTile::F32(program.load_local(&running_max)); + output_storage.store(program, index, value, Mask::all()); + }); + } else { + // out = O / L — the division the probability phase + // declares, applied once at the end. + program.if_then(out_active, |program| { + let value = program.load_local(&acc) / program.load_local(&running_sum); + let mut out_coords = row_dims.clone(); + out_coords.push(lane.clone()); + let output_index = layout_index(&output_meta, &out_coords); + let value = ValueTile::F32(value).cast_to(output_dtype); + output_storage.store(program, output_index, value, Mask::all()); + }); + } + return; + } + + // Chunked map program: lanes stride the axis for each phase's + // fold; full-shape intermediates referenced by several phases + // are recomputed per phase — the same trade every multi-pass + // normalization kernel makes. + let chunks = k.div_ceil(block); + for phase in &phases { + let RowPhase::Reduce(reduce) = phase else { + unreachable!("element phases require a dynamic-axis row program") + }; + let phase_dtype = reduce.function.datatype(); + let identity = || { + Tile::literal(tile_literal_for(reduce.function.initial_value, phase_dtype)) + }; + let reduce_op = tile_reduce_op(reduce.function.op); + let [partial] = program.fold( + tile_ir::tile::range(chunks), + [identity()], + |program, chunk, [acc]| { + let k_index = chunk * block + lane.clone(); + let active = in_bounds.clone().and(k_index.clone().lt(k)); + let coords = full_coords(k_index); + let (value, _) = eval_nary_expr( + program, + &reduce.expression, + &coords, + &storages, + &metas, + active.clone(), + &slots, + ); + let value = raw_tile(value.cast_to(phase_dtype)); + let masked = Tile::select(active, value, identity()); + [acc.binary(reduce_op.binary(), masked)] + }, + ); + // The workgroup reduction broadcasts: every lane reads the + // combined value, so later phases can use it directly. + let combined = program.group_reduce(reduce_op, block, partial); + let (combined, combined_ty) = + apply_unary_function_chain(combined, phase_dtype, &reduce.post_chain) + .expect("validated row program post chain"); + let scalar = ValueTile::F32(program.bind(combined)).cast_to(combined_ty); + slots.push((scalar, combined_ty)); + } + + let RowOutput::Map(output_expr) = &output_kind else { + unreachable!("reducing output requires a dynamic-axis row program") + }; + program.loop_range(chunks, |program, chunk| { + let k_index = chunk * block + lane.clone(); + let active = in_bounds.clone().and(k_index.clone().lt(k)); + let coords = full_coords(k_index); + let (value, _) = eval_nary_expr( + program, + output_expr, + &coords, + &storages, + &metas, + active.clone(), + &slots, + ); + let value = value.cast_to(output_dtype); + let output_index = layout_index(&output_meta, &coords); + output_storage.store(program, output_index, value, active); + }); + }); + Some(()) + }, + )?; + + match combine { + Some(combine) => Some(DirectKernel::sequence( + "row_program_split", + vec![partials, combine], + )), + None => Some(partials), + } +} + +/// Scaled dot-product attention as a row program over the KV axis: an +/// element phase stages the q·k score per lane via an inline head-dim fold, +/// a max phase and an exp-sum phase form the softmax statistics, a second +/// element phase stages the probabilities, and the reducing output folds +/// `Σ p·v` with the head dim as the free output dimension. Causal masking +/// is an index-compare select inside the scaled score (plus an axis bound +/// that skips tiles past the query position); an additive mask is a fourth +/// input read at `[q, kv]`. Each query row is one workgroup; KV histories +/// beyond one workgroup bucket stream through the online-softmax tile loop. +/// Returns `None` for shapes the row program cannot host (head dim beyond +/// the largest workgroup bucket, non-float dtypes). +pub(crate) struct AttentionInputs<'a> { + pub(crate) q: NodeIndex, + pub(crate) k: NodeIndex, + pub(crate) v: NodeIndex, + pub(crate) mask: Option, + pub(crate) q_shape: &'a [usize], + pub(crate) k_shape: &'a [usize], + pub(crate) v_shape: &'a [usize], + pub(crate) mask_shape: Option<&'a [usize]>, + pub(crate) scale: f32, + pub(crate) input_dtype: DataTypeEnum, + pub(crate) causal: bool, +} + +pub(crate) fn attention_row_program( + device: &crate::Device, + inputs: AttentionInputs<'_>, +) -> Option { + let AttentionInputs { + q, + k, + v, + mask, + q_shape, + k_shape, + v_shape, + mask_shape, + scale, + input_dtype, + causal, + } = inputs; + let [batch, num_heads, q_seq_len, head_dim] = *q_shape else { + return None; + }; + let [k_batch, num_kv_heads, kv_len, k_head_dim] = *k_shape else { + return None; + }; + if !matches!(input_dtype, DataTypeEnum::F32 | DataTypeEnum::F16) + || (input_dtype == DataTypeEnum::F16 && !device.f16_supported()) + || q_seq_len == 0 + || kv_len == 0 + || head_dim == 0 + || k_batch != batch + || k_head_dim != head_dim + || v_shape != k_shape + || num_kv_heads == 0 + || !num_heads.is_multiple_of(num_kv_heads) + || (causal && mask.is_some()) + { + return None; + } + if let Some(mask_shape) = mask_shape + && mask_shape != [q_seq_len, kv_len] + { + return None; + } + if mask.is_some() != mask_shape.is_some() { + return None; + } + let limits = device.limits(); + // The workgroup bucket is one axis tile: small tiles let the split + // lowering fan decode across workgroups and stream longer axes through + // the online loop with good occupancy. + let needed = head_dim.max(256) as u32; + let block = ROW_DYNAMIC_BLOCKS.iter().copied().find(|&candidate| { + candidate >= needed + && candidate <= limits.max_compute_workgroup_size_x + && candidate <= limits.max_compute_invocations_per_workgroup + })?; + + let groups = num_heads / num_kv_heads; + let f32 = DataTypeEnum::F32; + let dim = NaryExpr::DimIndex; + let kv_head = || { + if groups == 1 { + dim(1) + } else { + NaryExpr::Op { + children: vec![dim(1)], + function: NaryFunction::unary( + Some("kv_head".to_string()), + NaryOp::DivConst(NaryScalar::U32(groups as u32)), + DataTypeEnum::U32, + DataTypeEnum::U32, + ), + } + } + }; + let binary = |op, a, b| NaryExpr::Op { + children: vec![a, b], + function: NaryFunction::binary(None, op, f32, f32, f32), + }; + let unary = |op, a| NaryExpr::Op { + children: vec![a], + function: NaryFunction::unary(None, op, f32, f32), + }; + let graph_inputs: Vec = match mask { + Some(mask) => vec![q, k, v, mask], + None => vec![q, k, v], + }; + let input_count = graph_inputs.len(); + let slot = |p: usize| slot_expr(input_count, p); + // The fold/free coordinate is one past the rank-4 index space. Loads + // are cast to the f32 expression types by evaluation, so f16 tensors + // need no explicit casts. + let q_read = NaryExpr::indexed_input(0, vec![dim(0), dim(1), dim(2), dim(4)]); + let k_read = NaryExpr::indexed_input(1, vec![dim(0), kv_head(), dim(3), dim(4)]); + let v_read = NaryExpr::indexed_input(2, vec![dim(0), kv_head(), dim(3), dim(4)]); + + // The scaled score, with causality or the additive mask folded in. The + // causal arm masks to the max-phase identity so masked positions vanish + // from both the max and the exp sum. + let scaled_score = || { + let base = unary(NaryOp::MulConst(NaryScalar::F32(scale)), slot(0)); + if causal { + let bound = NaryExpr::Op { + children: vec![dim(3), dim(2)], + function: NaryFunction::binary( + Some("causal_bound".to_string()), + NaryOp::LessEqual, + DataTypeEnum::U32, + DataTypeEnum::U32, + DataTypeEnum::U32, + ), + }; + NaryExpr::select( + bound, + base, + NaryExpr::Scalar(max_fn(f32).initial_value), + DataTypeEnum::U32, + f32, + ) + } else if mask.is_some() { + let mask_read = NaryExpr::indexed_input(3, vec![dim(2), dim(3)]); + binary(NaryOp::Add, base, mask_read) + } else { + base + } + }; + let shifted_exp = || unary(NaryOp::Exp, binary(NaryOp::Sub, scaled_score(), slot(1))); + + let mut input_axis_dims = vec![None, Some(2), Some(2)]; + if mask.is_some() { + input_axis_dims.push(Some(1)); + } + Some(RowProgramOperation { + inputs: graph_inputs, + shape: [batch, num_heads, q_seq_len, kv_len].into(), + axis: 3, + phases: vec![ + RowPhase::Element { + expression: binary(NaryOp::Mul, q_read, k_read), + fold: Some(RowFold { + len: head_dim, + function: sum_fn(f32), + }), + datatype: f32, + }, + RowPhase::Reduce(RowReduce { + expression: scaled_score(), + function: max_fn(f32), + post_chain: UnaryFunctionChain::empty(f32), + }), + RowPhase::Reduce(RowReduce { + expression: shifted_exp(), + function: sum_fn(f32), + post_chain: UnaryFunctionChain::empty(f32), + }), + RowPhase::Element { + expression: binary(NaryOp::Div, shifted_exp(), slot(2)), + fold: None, + datatype: f32, + }, + ], + output: RowOutput::Reduce { + combine: binary(NaryOp::Mul, slot(3), v_read), + function: sum_fn(f32), + free_dim: head_dim, + }, + output_datatype: input_dtype, + dynamic_axis: Some(DynamicAxis { + block, + input_axis_dims, + axis_bound_dim: causal.then_some(2), + }), + }) +} diff --git a/fusor-ml/core/src/slice_assign.rs b/fusor-ml/core/src/slice_assign.rs index 9a99fe89b..82596129d 100644 --- a/fusor-ml/core/src/slice_assign.rs +++ b/fusor-ml/core/src/slice_assign.rs @@ -9,7 +9,7 @@ use crate::{ operation::Operation, workgroup_shape::{WorkgroupShape, WorkgroupShapeConstraints}, }, - nary_wise::{NaryExpr, NaryOp, NaryOperation, NaryScalar}, + nary_wise::{ElementwiseOperation, NaryExpr, NaryOp, NaryScalar}, visit_tiled::{titled_map_dispatch_size, titled_map_workgroup_size_constraints}, }; @@ -22,22 +22,67 @@ pub(crate) struct SliceAssignOperation { pub(crate) in_place: bool, } -impl SliceAssignOperation { - pub fn new( - input: NodeIndex, - value: NodeIndex, - slices: Box<[Range]>, - input_shape: Box<[usize]>, - ) -> Self { - Self { - input, - value, - slices, - input_shape, - in_place: false, - } +/// The composed slice-assign body: per output coordinate, read the assigned +/// value inside the slice region and the original input outside it. Inputs: +/// 0 = the original tensor, 1 = the assigned value. +pub(crate) fn slice_assign_expression(slices: &[Range], datatype: DataTypeEnum) -> NaryExpr { + let rank = slices.len(); + let mut condition = NaryExpr::scalar(NaryScalar::U32(1)); + for (dim, slice) in slices.iter().enumerate() { + let dim_index = NaryExpr::DimIndex(dim); + let ge_start = NaryExpr::unary_op( + dim_index.clone(), + "ge_start", + NaryOp::GreaterEqualConst(NaryScalar::U32(slice.start as u32)), + DataTypeEnum::U32, + DataTypeEnum::U32, + ); + let lt_end = NaryExpr::unary_op( + dim_index, + "lt_end", + NaryOp::LessConst(NaryScalar::U32(slice.end as u32)), + DataTypeEnum::U32, + DataTypeEnum::U32, + ); + condition = NaryExpr::mul(condition, ge_start, DataTypeEnum::U32); + condition = NaryExpr::mul(condition, lt_end, DataTypeEnum::U32); } + let value_indices = slices + .iter() + .enumerate() + .map(|(dim, slice)| { + let shifted_index = if slice.start == 0 { + NaryExpr::DimIndex(dim) + } else { + NaryExpr::unary_op( + NaryExpr::DimIndex(dim), + "slice_offset", + NaryOp::SubConst(NaryScalar::U32(slice.start as u32)), + DataTypeEnum::U32, + DataTypeEnum::U32, + ) + }; + NaryExpr::select( + condition.clone(), + shifted_index, + NaryExpr::scalar(NaryScalar::U32(0)), + DataTypeEnum::U32, + DataTypeEnum::U32, + ) + }) + .collect(); + + NaryExpr::select( + condition, + NaryExpr::indexed_input(1, value_indices), + NaryExpr::input(0, rank), + DataTypeEnum::U32, + datatype, + ) +} + +impl SliceAssignOperation { pub fn new_in_place( input: NodeIndex, value: NodeIndex, @@ -72,62 +117,7 @@ impl SliceAssignOperation { if self.in_place { return NaryExpr::input(0, self.slices.len()); } - - let rank = self.slices.len(); - let mut condition = NaryExpr::scalar(NaryScalar::U32(1)); - for (dim, slice) in self.slices.iter().enumerate() { - let dim_index = NaryExpr::DimIndex(dim); - let ge_start = NaryExpr::unary_op( - dim_index.clone(), - "ge_start", - NaryOp::GreaterEqualConst(NaryScalar::U32(slice.start as u32)), - DataTypeEnum::U32, - DataTypeEnum::U32, - ); - let lt_end = NaryExpr::unary_op( - dim_index, - "lt_end", - NaryOp::LessConst(NaryScalar::U32(slice.end as u32)), - DataTypeEnum::U32, - DataTypeEnum::U32, - ); - condition = NaryExpr::mul(condition, ge_start, DataTypeEnum::U32); - condition = NaryExpr::mul(condition, lt_end, DataTypeEnum::U32); - } - - let value_indices = self - .slices - .iter() - .enumerate() - .map(|(dim, slice)| { - let shifted_index = if slice.start == 0 { - NaryExpr::DimIndex(dim) - } else { - NaryExpr::unary_op( - NaryExpr::DimIndex(dim), - "slice_offset", - NaryOp::SubConst(NaryScalar::U32(slice.start as u32)), - DataTypeEnum::U32, - DataTypeEnum::U32, - ) - }; - NaryExpr::select( - condition.clone(), - shifted_index, - NaryExpr::scalar(NaryScalar::U32(0)), - DataTypeEnum::U32, - DataTypeEnum::U32, - ) - }) - .collect(); - - NaryExpr::select( - condition, - NaryExpr::indexed_input(1, value_indices), - NaryExpr::input(0, rank), - DataTypeEnum::U32, - datatype, - ) + slice_assign_expression(&self.slices, datatype) } } @@ -187,7 +177,7 @@ impl Operation for SliceAssignOperation { ) -> Option { if self.in_place { let value = inputs[1].as_tensor()?; - let operation = NaryOperation { + let operation = ElementwiseOperation { inputs: vec![self.value], expression: self.expression(value.datatype()), shape: value.layout().shape().into(), @@ -203,7 +193,7 @@ impl Operation for SliceAssignOperation { } let input = inputs[0].as_tensor()?; - let operation = NaryOperation { + let operation = ElementwiseOperation { inputs: vec![self.input, self.value], expression: self.expression(input.datatype()), shape: self.input_shape.clone(), diff --git a/fusor-ml/core/src/softmax.rs b/fusor-ml/core/src/softmax.rs deleted file mode 100644 index 8438dd6e1..000000000 --- a/fusor-ml/core/src/softmax.rs +++ /dev/null @@ -1,449 +0,0 @@ -use std::{any::Any, hash::Hash}; - -use fusor_tile_ir as tile_ir; -use fusor_tile_ir_kernels as tile_ir_kernels; -use rustc_hash::{FxHashMap, FxHasher}; - -use crate::{ - DataTypeEnum, Device, Layout, - compute_graph::{ComputeGraphInner, GraphOperation, NodeIndex}, - kernel_selection::KernelDeviceCaps, - mir::{ - inputs::MirValue, - kernel_backend, - kernel_backend::DirectKernel, - operation::Operation, - workgroup_shape::{Constraint, WorkgroupShape, WorkgroupShapeConstraints}, - }, - tensor::{TensorData, TensorLayoutInfo}, - visit_tiled::distribute_workgroups, -}; - -const SOFTMAX_BLOCKS: [u32; 3] = [128, 512, 1024]; - -#[derive(Clone, Copy, Debug, Hash, PartialEq, Eq)] -enum SoftmaxKernelVariant { - Single, - Partials, - Reduce, - Write, -} - -struct SoftmaxDirectKernelVariant; - -fn block_supported(block: u32, caps: KernelDeviceCaps) -> bool { - block <= caps.max_compute_invocations_per_workgroup - && block <= caps.max_compute_workgroup_size_x -} - -fn choose_softmax_block(axis_len: u32, caps: KernelDeviceCaps) -> Option { - for block in SOFTMAX_BLOCKS { - if axis_len <= block && block_supported(block, caps) { - return Some(block); - } - } - - SOFTMAX_BLOCKS - .iter() - .rev() - .copied() - .find(|block| block_supported(*block, caps)) -} - -fn total_elements(shape: &[usize]) -> Option { - shape - .iter() - .try_fold(1u32, |acc, dim| acc.checked_mul((*dim).try_into().ok()?)) -} - -fn tensor_meta(tensor: &TensorData) -> Option { - let strides = tensor - .layout() - .strides() - .iter() - .copied() - .map(u32::try_from) - .collect::, _>>() - .ok()?; - let offset = tensor.layout().offset().try_into().ok()?; - Some(tile_ir_kernels::TensorMeta::new(strides, offset)) -} - -#[derive(Clone, Debug)] -pub(crate) struct SoftmaxOperation { - input: NodeIndex, - shape: Box<[usize]>, - axis: usize, - datatype: DataTypeEnum, -} - -impl SoftmaxOperation { - pub(crate) fn new( - input: NodeIndex, - shape: &[usize], - axis: usize, - datatype: DataTypeEnum, - device: &Device, - ) -> Option { - if axis >= shape.len() || shape.contains(&0) { - return None; - } - if axis + 1 != shape.len() { - return None; - } - if !matches!(datatype, DataTypeEnum::F32 | DataTypeEnum::F16) { - return None; - } - if datatype == DataTypeEnum::F16 && !device.f16_supported() { - return None; - } - - let axis_len: u32 = shape[axis].try_into().ok()?; - let _total = total_elements(shape)?; - choose_softmax_block(axis_len, KernelDeviceCaps::from_device(device))?; - - Some(Self { - input, - shape: shape.into(), - axis, - datatype, - }) - } - - fn meta( - &self, - input: &TensorData, - output: &TensorData, - dispatch_size: [u32; 3], - caps: KernelDeviceCaps, - ) -> Option { - let shape = self - .shape - .iter() - .copied() - .map(u32::try_from) - .collect::, _>>() - .ok()?; - let axis_len = shape[self.axis]; - let rows = total_elements(&self.shape)?.checked_div(axis_len)?; - let block = choose_softmax_block(axis_len, caps)?; - let split_blocks = axis_len.div_ceil(block); - - Some(tile_ir_kernels::SoftmaxMeta { - shape, - axis: self.axis.try_into().ok()?, - rows, - axis_len, - block, - split_blocks, - input_meta: tensor_meta(input)?, - output_meta: tensor_meta(output)?, - dispatch_size, - }) - } - - fn dispatch_for(&self, total_groups: u32, device: &Device) -> [u32; 3] { - distribute_workgroups( - total_groups, - device.limits().max_compute_workgroups_per_dimension, - ) - } - - fn split_blocks(&self, device: &Device) -> Option { - let axis_len: u32 = self.shape[self.axis].try_into().ok()?; - let block = choose_softmax_block(axis_len, KernelDeviceCaps::from_device(device))?; - Some(axis_len.div_ceil(block)) - } - - fn dispatch_softmax( - &self, - element: tile_ir::ElementType, - device: &Device, - input: &TensorData, - output: &TensorData, - meta: tile_ir_kernels::SoftmaxMeta, - ) -> Option { - let dispatch_size = meta.dispatch_size; - let variant = - kernel_backend::KernelVariantKey::with_payload::(|state| { - SoftmaxKernelVariant::Single.hash(state); - meta.block.hash(state); - meta.split_blocks.hash(state); - self.datatype.hash(state); - }); - let inputs = vec![input.clone().into(), output.clone().into()]; - let key = self.kernel_cache_key_with_dispatch(variant, None, dispatch_size, &inputs); - let buffers = vec![input.buffer().clone(), output.buffer().clone()]; - let layout = tile_ir_kernels::linear_storage_layout(); - - kernel_backend::dynamic_kernel_from_ir( - device.kernel_cache(), - "softmax", - key, - move || { - let mut kb = tile_ir::KernelBuilder::<()>::new(); - let input_ref = tile_ir::KernelTensorRef::new((), layout.clone()); - let output_ref = tile_ir::KernelTensorRef::new((), layout); - tile_ir_kernels::softmax(&mut kb, element, input_ref, output_ref, meta)?; - Some(kb.finish().0) - }, - buffers, - dispatch_size, - ) - } - - fn dispatch_split_softmax( - &self, - element: tile_ir::ElementType, - device: &Device, - input: &TensorData, - output: &TensorData, - meta: tile_ir_kernels::SoftmaxMeta, - ) -> Option { - let scratch_elements = meta.rows as u64 * meta.split_blocks as u64 * 2; - let scratch = device.create_buffer( - scratch_elements * std::mem::size_of::() as u64, - wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC, - ); - let global_scratch = device.create_buffer( - meta.rows as u64 * 2 * std::mem::size_of::() as u64, - wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC, - ); - let partial_dispatch_size = meta.dispatch_size; - let reduce_dispatch_size = self.dispatch_for(meta.rows, device); - let inputs = vec![input.clone().into(), output.clone().into()]; - let layout = tile_ir_kernels::linear_storage_layout(); - - let partial_variant = - kernel_backend::KernelVariantKey::with_payload::(|state| { - SoftmaxKernelVariant::Partials.hash(state); - meta.block.hash(state); - meta.split_blocks.hash(state); - self.datatype.hash(state); - }); - let partial_key = self.kernel_cache_key_with_dispatch( - partial_variant, - None, - partial_dispatch_size, - &inputs, - ); - let partial_buffers = vec![input.buffer().clone(), scratch.clone()]; - let partial_layout = layout.clone(); - let partial_meta = meta.clone(); - let partial = kernel_backend::dynamic_kernel_from_ir( - device.kernel_cache(), - "softmax_partials", - partial_key, - move || { - let mut kb = tile_ir::KernelBuilder::<()>::new(); - let input_ref = tile_ir::KernelTensorRef::new((), partial_layout.clone()); - let scratch_ref = tile_ir::KernelTensorRef::new((), partial_layout); - tile_ir_kernels::softmax_partials( - &mut kb, - element, - input_ref, - scratch_ref, - partial_meta, - )?; - Some(kb.finish().0) - }, - partial_buffers, - partial_dispatch_size, - )?; - - let reduce_variant = - kernel_backend::KernelVariantKey::with_payload::(|state| { - SoftmaxKernelVariant::Reduce.hash(state); - meta.block.hash(state); - meta.split_blocks.hash(state); - }); - let reduce_key = self.kernel_cache_key_with_dispatch( - reduce_variant, - None, - reduce_dispatch_size, - &inputs, - ); - let reduce_buffers = vec![scratch, global_scratch.clone()]; - let reduce_layout = layout.clone(); - let mut reduce_meta = meta.clone(); - reduce_meta.dispatch_size = reduce_dispatch_size; - let reduce = kernel_backend::dynamic_kernel_from_ir( - device.kernel_cache(), - "softmax_reduce", - reduce_key, - move || { - let mut kb = tile_ir::KernelBuilder::<()>::new(); - let scratch_ref = tile_ir::KernelTensorRef::new((), reduce_layout.clone()); - let global_ref = tile_ir::KernelTensorRef::new((), reduce_layout); - tile_ir_kernels::softmax_reduce::<_>( - &mut kb, - scratch_ref, - global_ref, - reduce_meta, - )?; - Some(kb.finish().0) - }, - reduce_buffers, - reduce_dispatch_size, - )?; - - let write_variant = - kernel_backend::KernelVariantKey::with_payload::(|state| { - SoftmaxKernelVariant::Write.hash(state); - meta.block.hash(state); - meta.split_blocks.hash(state); - self.datatype.hash(state); - }); - let write_key = self.kernel_cache_key_with_dispatch( - write_variant, - None, - partial_dispatch_size, - &inputs, - ); - let write_buffers = vec![ - input.buffer().clone(), - global_scratch, - output.buffer().clone(), - ]; - let write_layout = layout.clone(); - let write = kernel_backend::dynamic_kernel_from_ir( - device.kernel_cache(), - "softmax_write", - write_key, - move || { - let mut kb = tile_ir::KernelBuilder::<()>::new(); - let input_ref = tile_ir::KernelTensorRef::new((), write_layout.clone()); - let scratch_ref = tile_ir::KernelTensorRef::new((), write_layout.clone()); - let output_ref = tile_ir::KernelTensorRef::new((), write_layout); - tile_ir_kernels::softmax_write( - &mut kb, - element, - input_ref, - scratch_ref, - output_ref, - meta, - )?; - Some(kb.finish().0) - }, - write_buffers, - partial_dispatch_size, - )?; - - Some(DirectKernel::sequence( - "softmax_split", - vec![partial, reduce, write], - )) - } -} - -impl Operation for SoftmaxOperation { - fn hash_kernel_fields(&self, state: &mut FxHasher) { - self.shape.hash(state); - self.axis.hash(state); - self.datatype.hash(state); - } - - fn workgroup_shape_constraints(&self, _device: &Device) -> WorkgroupShapeConstraints { - let mut constraints = WorkgroupShapeConstraints::new(); - constraints.add_constraint(0, Constraint::Equals(1)); - constraints.add_constraint(1, Constraint::Equals(1)); - constraints.add_constraint(2, Constraint::Equals(1)); - constraints - } - - fn dispatch_size(&self, _workgroup_shape: &WorkgroupShape, inputs: &[MirValue]) -> [u32; 3] { - let Some(input) = inputs.first().and_then(MirValue::as_tensor) else { - return [1, 1, 1]; - }; - let Some(axis_len) = self - .shape - .get(self.axis) - .and_then(|dim| u32::try_from(*dim).ok()) - else { - return [1, 1, 1]; - }; - let Some(rows) = total_elements(&self.shape).and_then(|total| total.checked_div(axis_len)) - else { - return [1, 1, 1]; - }; - let Some(split_blocks) = self.split_blocks(input.device()) else { - return [1, 1, 1]; - }; - self.dispatch_for(rows.saturating_mul(split_blocks), input.device()) - } - - fn visit_dependencies(&self, f: &mut dyn FnMut(NodeIndex)) { - f(self.input); - } - - fn inputs(&self, nodes: &ComputeGraphInner) -> Vec { - let input = nodes.get_cached_result(self.input).unwrap(); - let output = TensorData::new_for_shape(input.device(), &self.shape, self.datatype); - vec![input.clone().into(), output.into()] - } - - fn output(&self, _nodes: &ComputeGraphInner, inputs: &[MirValue]) -> MirValue { - inputs[1].clone() - } - - fn build_direct_kernel( - &self, - graph: &ComputeGraphInner, - _workgroup_shape: &WorkgroupShape, - inputs: &[MirValue], - ) -> Option { - let input = inputs.first()?.as_tensor()?; - let output = inputs.get(1)?.as_tensor()?; - if input.datatype() != self.datatype || output.datatype() != self.datatype { - return None; - } - if self.datatype == DataTypeEnum::F16 && !graph.device().f16_supported() { - return None; - } - - let device = graph.device(); - let caps = KernelDeviceCaps::from_device(&device); - let dispatch_size = self.dispatch_size(&WorkgroupShape::new(1, 1, 1), inputs); - let meta = self.meta(input, output, dispatch_size, caps)?; - - match self.datatype { - DataTypeEnum::F32 if meta.split_blocks == 1 => { - self.dispatch_softmax(tile_ir::ElementType::F32, &device, input, output, meta) - } - DataTypeEnum::F16 if meta.split_blocks == 1 => { - self.dispatch_softmax(tile_ir::ElementType::F16, &device, input, output, meta) - } - DataTypeEnum::F32 => { - self.dispatch_split_softmax(tile_ir::ElementType::F32, &device, input, output, meta) - } - DataTypeEnum::F16 => { - self.dispatch_split_softmax(tile_ir::ElementType::F16, &device, input, output, meta) - } - DataTypeEnum::U32 => None, - } - } - - fn name(&self) -> String { - format!("softmax_axis_{}", self.axis) - } -} - -impl GraphOperation for SoftmaxOperation { - fn as_any(&self) -> &dyn Any { - self - } - - fn category(&self) -> &'static str { - "softmax" - } - - fn output_layout( - &self, - _input_layouts: &FxHashMap, - ) -> Option { - Some(TensorLayoutInfo::new( - Layout::contiguous(&self.shape), - self.datatype, - )) - } -} diff --git a/fusor-ml/core/src/tensor/lazy_data.rs b/fusor-ml/core/src/tensor/lazy_data.rs index b780ccf9d..b1e36089e 100644 --- a/fusor-ml/core/src/tensor/lazy_data.rs +++ b/fusor-ml/core/src/tensor/lazy_data.rs @@ -2,14 +2,11 @@ use tabbycat::Graph; use crate::{ - Device, FlashAttentionOperation, Layout, MatMulOperation, ReduceOperation, + Device, ReduceOperation, compute_graph::NodeIndex, - map_layout::MapLayoutOperation, - nary_wise::{NaryExpr, NaryFunction, NaryOperation}, - quantized::matmul::QMatMulOperation, - resize::ResizeOperation, - rms_norm::RmsNormOperation, + nary_wise::{ElementwiseOperation, NaryExpr, NaryFunction}, slice_assign::SliceAssignOperation, + view::ViewOperation, }; use super::{TensorData, TensorInfo}; @@ -54,7 +51,7 @@ impl LazyTensorData { Self { device, info, key } } - pub(crate) fn nary(&self, nary: NaryOperation) -> Self { + pub(crate) fn nary(&self, nary: ElementwiseOperation) -> Self { let device = self.device.clone(); let mut info = self.info.clone(); info.shape = nary.shape.clone(); @@ -69,7 +66,7 @@ impl LazyTensorData { let mut info = self.info.clone(); info.datatype = function.output_type; let rank = info.rank(); - let nary = NaryOperation { + let nary = ElementwiseOperation { inputs: vec![self.key], expression: NaryExpr::Op { children: vec![NaryExpr::input(0, rank)], @@ -93,7 +90,7 @@ impl LazyTensorData { let mut info = self.info.clone(); info.datatype = function.output_type; let rank = shape.len(); - let nary = NaryOperation { + let nary = ElementwiseOperation { inputs: vec![self.key, other_key], expression: NaryExpr::Op { children: vec![NaryExpr::input(0, rank), NaryExpr::input(1, rank)], @@ -107,24 +104,6 @@ impl LazyTensorData { Self::from_parts(device, info, key) } - pub(crate) fn mat_mul(&self, function: MatMulOperation) -> Self { - let device = self.device.clone(); - let mut info = self.info.clone(); - info.shape = function.out_shape.clone(); - let key = device.compute_graph().create_mat_mul(function); - - Self::from_parts(device, info, key) - } - - pub(crate) fn q_mat_mul(&self, function: QMatMulOperation) -> Self { - let device = self.device.clone(); - let mut info = self.info.clone(); - info.shape = function.out_shape.clone(); - let key = device.compute_graph().create_q_mat_mul(function); - - Self::from_parts(device, info, key) - } - pub(crate) fn reduce(&self, function: ReduceOperation) -> Self { let device = self.device.clone(); let mut info = self.info.clone(); @@ -152,38 +131,10 @@ impl LazyTensorData { Self::from_parts(device, info, key) } - pub(crate) fn rms_norm(&self, function: RmsNormOperation) -> Self { - let device = self.device.clone(); - let info = self.info.clone(); - let key = device.compute_graph().create_rms_norm(function); - - Self::from_parts(device, info, key) - } - - pub(crate) fn flash_attention(&self, function: FlashAttentionOperation) -> Self { - let device = self.device.clone(); - let mut info = self.info.clone(); - info.shape = function.out_shape.clone(); - let key = device.compute_graph().create_flash_attention(function); - - Self::from_parts(device, info, key) - } - - pub(crate) fn map_layout(&self, op: MapLayoutOperation) -> Self { - let device = self.device.clone(); - // Compute output shape by applying the layout transformation to a temporary layout - let temp_layout = Layout::contiguous(self.info.shape()); - let new_layout = op.map_layout(&temp_layout); - let info = TensorInfo::new(new_layout.shape().into(), self.info.datatype()); - let key = device.compute_graph().create_map_layout(op); - - Self::from_parts(device, info, key) - } - - pub(crate) fn resize(&self, op: ResizeOperation) -> Self { + pub(crate) fn view(&self, op: ViewOperation) -> Self { let device = self.device.clone(); - let info = TensorInfo::new(op.new_shape.clone(), self.info.datatype()); - let key = device.compute_graph().create_resize(op); + let info = TensorInfo::new(op.shape().into(), op.datatype); + let key = device.compute_graph().create_view(op); Self::from_parts(device, info, key) } diff --git a/fusor-ml/core/src/tensor/mod.rs b/fusor-ml/core/src/tensor/mod.rs index 133aa419b..485434f09 100644 --- a/fusor-ml/core/src/tensor/mod.rs +++ b/fusor-ml/core/src/tensor/mod.rs @@ -8,16 +8,10 @@ use tabbycat::Graph; use wgpu::COPY_BUFFER_ALIGNMENT; use crate::{ - Device, FlashAttentionInputs, FlashAttentionOperation, MatMulOperation, MatMulParams, - ReduceFunction, ReduceOperation, + Device, ReduceFunction, ReduceOperation, compute_graph::NodeIndex, - conv::ConvNdOperation, - map_layout::MapLayoutOperation, - nary_wise::{NaryExpr, NaryFunction, NaryOperation}, + nary_wise::{ElementwiseOperation, NaryExpr, NaryFunction}, quantized::QMatrix, - quantized::matmul::QMatMulOperation, - resize::ResizeOperation, - rms_norm::RmsNormOperation, slice_assign::SliceAssignOperation, }; @@ -324,7 +318,7 @@ impl Tensor { let mut info = self.data.info.clone(); info.datatype = function.output_type; let rank = self.shape().len(); - let nary = NaryOperation { + let nary = ElementwiseOperation { inputs: vec![self.data.key], expression: NaryExpr::Op { children: vec![NaryExpr::input(0, rank), NaryExpr::input(0, rank)], @@ -344,40 +338,77 @@ impl Tensor { ) } - pub(crate) fn add_mat_mul(&self, other: &Self, parameters: Option) -> Self { - let operation = MatMulOperation::new( - self.datatype(), - self.data.key, - other.data.key, - self.shape(), - other.shape(), - parameters, - &self.data.device, - ); - - Self::from_parts(self.data.mat_mul(operation)) - } - + /// Quantized matrix multiply in its composed form: the activation + /// `[.., K]` and the dequantized matrix `[N, K]` multiply over the + /// `[.., N, K]` index space and sum along `K`. The resolver recognizes + /// the canonical cluster and routes it to the quantized matmul kernels. pub(crate) fn add_q_mat_mul(&self, other: &QMatrix) -> Self { - let operation = - QMatMulOperation::new(self.datatype(), self.shape(), self.data.key, other.clone()); - - Self::from_parts(self.data.q_mat_mul(operation)) - } - - pub(crate) fn add_resize(&self, op: ResizeOperation) -> Tensor { - Tensor::from_parts(self.data.resize(op)) - } + let in_shape = self.shape(); + let rank = in_shape.len(); + assert!(rank >= 1, "q_mat_mul requires rank >= 1"); + assert_eq!( + in_shape[rank - 1], + other.shape()[1], + "q_mat_mul contraction dimensions must match: {in_shape:?} x {:?}", + other.shape() + ); + let datatype = self.datatype(); + let device = self.device().clone(); + let matrix_key = device.compute_graph().dequantize(other.clone(), datatype); + let matrix = Tensor::from_parts(LazyTensorData::from_parts( + device, + TensorInfo::new(other.shape().into(), datatype), + matrix_key, + )); + + let n = other.shape()[0]; + // Index space [.., N, K]: K stays last so the reduce axis is the + // final dimension. + let mut index_space = in_shape.to_vec(); + index_space.insert(rank - 1, n); + let (n_dim, k_dim) = (rank - 1, rank); + + let activation_indices: Vec = (0..rank - 1) + .chain(std::iter::once(k_dim)) + .map(NaryExpr::DimIndex) + .collect(); + let matrix_indices: Vec = [n_dim, k_dim].map(NaryExpr::DimIndex).to_vec(); + + let product = Tensor::from_parts(self.data.nary(ElementwiseOperation { + inputs: vec![self.key(), matrix.key()], + expression: NaryExpr::mul( + NaryExpr::indexed_input(0, activation_indices), + NaryExpr::indexed_input(1, matrix_indices), + datatype, + ), + shape: index_space.into(), + output_datatype: datatype, + })); + product.sum(k_dim) + } + + /// Slice assignment in its composed form: per output coordinate, read + /// the assigned value inside the slice region and this tensor outside + /// it. A plain elementwise op — no specialized kernel. pub(crate) fn add_slice_assign( &self, other: &Self, slices: impl Into]>>, ) -> Self { - let input_shape: Box<[usize]> = self.shape().to_vec().into_boxed_slice(); - let op = - SliceAssignOperation::new(self.data.key, other.data.key, slices.into(), input_shape); - Self::from_parts(self.data.slice_assign(op)) + let slices: Box<[Range]> = slices.into(); + assert_eq!( + slices.len(), + self.rank(), + "slice_assign requires one range per dimension" + ); + let expression = crate::slice_assign::slice_assign_expression(&slices, self.datatype()); + Self::from_parts(self.data.nary(ElementwiseOperation { + inputs: vec![self.data.key, other.data.key], + expression, + shape: self.shape().into(), + output_datatype: self.datatype(), + })) } #[doc(hidden)] @@ -405,10 +436,6 @@ impl Tensor { ))) } - pub(crate) fn add_map_layout(&self, op: MapLayoutOperation) -> Tensor { - Tensor::from_parts(self.data.map_layout(op)) - } - /// Return the compute-graph node index for this tensor. pub fn key(&self) -> NodeIndex { self.data.key @@ -443,219 +470,6 @@ impl Tensor { self.data.info.datatype() } - pub(crate) fn try_rms_norm_direct( - &self, - weight: &Tensor, - bias: Option<&Tensor>, - eps: f32, - ) -> Option { - if !matches!(self.datatype(), DataTypeEnum::F32 | DataTypeEnum::F16) - || self.datatype() != weight.datatype() - || bias.is_some_and(|bias| bias.datatype() != self.datatype()) - || (self.datatype() == DataTypeEnum::F16 && !self.device().f16_supported()) - { - return None; - } - let operation = RmsNormOperation::new( - self.data.key, - weight.data.key, - bias.map(|bias| bias.data.key), - self.shape(), - eps, - ); - Some(Self::from_parts(self.data.rms_norm(operation))) - } - - #[doc(hidden)] - pub fn try_conv_nd_direct( - &self, - weight: &Tensor, - bias: Option<&Tensor>, - padding: &[usize], - strides: &[usize], - ) -> Option { - if self.datatype() != weight.datatype() - || bias.is_some_and(|bias| bias.datatype() != self.datatype()) - { - return None; - } - let operation = ConvNdOperation::new( - crate::conv::ConvNdNodes { - input: self.data.key, - weight: weight.data.key, - bias: bias.map(|bias| bias.data.key), - }, - crate::conv::ConvNdShapeSpec { - input_shape: self.shape(), - weight_shape: weight.shape(), - bias_shape: bias.map(|bias| bias.shape()), - padding, - strides, - }, - self.datatype(), - self.device(), - )?; - let device = self.device().clone(); - let info = TensorInfo::new(operation.output_shape().into(), self.datatype()); - let key = device - .compute_graph() - .create_graph_op(std::sync::Arc::new(operation)); - Some(Self::from_parts(LazyTensorData::from_parts( - device, info, key, - ))) - } - - pub(crate) fn try_rms_norm_residual_direct( - &self, - residual: &Self, - weight: &Tensor, - bias: Option<&Tensor>, - eps: f32, - ) -> Option { - if !matches!(self.datatype(), DataTypeEnum::F32 | DataTypeEnum::F16) - || self.shape() != residual.shape() - || residual.datatype() != self.datatype() - || weight.datatype() != self.datatype() - || bias.is_some_and(|bias| bias.datatype() != self.datatype()) - || (self.datatype() == DataTypeEnum::F16 && !self.device().f16_supported()) - { - return None; - } - let operation = RmsNormOperation::new_with_residual( - self.data.key, - residual.data.key, - weight.data.key, - bias.map(|bias| bias.data.key), - self.shape(), - eps, - ); - Some(Self::from_parts(self.data.rms_norm(operation))) - } - - pub(crate) fn try_flash_attention_direct( - &self, - k: &Self, - v: &Self, - scale: f32, - mask: Option<&Tensor>, - ) -> Option { - self.try_flash_attention_direct_inner(k, v, scale, mask, false) - } - - pub(crate) fn try_flash_attention_direct_causal( - &self, - k: &Self, - v: &Self, - scale: f32, - ) -> Option { - self.try_flash_attention_direct_inner(k, v, scale, None, true) - } - - fn try_flash_attention_direct_inner( - &self, - k: &Self, - v: &Self, - scale: f32, - mask: Option<&Tensor>, - causal: bool, - ) -> Option { - if self.rank() != 4 || !matches!(self.datatype(), DataTypeEnum::F32 | DataTypeEnum::F16) { - return None; - } - if causal && mask.is_some() { - return None; - } - let q_shape = self.shape(); - let k_shape = k.shape(); - let is_decode_shape = q_shape[2] == 1 - && q_shape[3] > 0 - && mask.is_none() - && !causal - && self.datatype() == DataTypeEnum::F32; - if is_decode_shape - && k_shape[2] < crate::mir::kernel_backend::flash_attention::MIN_DECODE_KV_SEQ - { - return None; - } - let is_decode_candidate = - crate::mir::kernel_backend::flash_attention::flash_decode_direct_candidate( - q_shape, - k_shape, - mask.is_some(), - causal, - self.datatype(), - ); - // Decode (q_seq_len == 1) now uses the fused flash-attention decode - // kernel on wasm as well. The browser path was previously disabled here - // over WebGPU hang/crash reports on the split decode workgroup kernel; - // re-enabled to measure whether that's still an issue. - // The streaming flash attention kernels emit a separate - // monomorphization per hardware subgroup width and rely on - // `subgroup_reduce_*`, so they can only target devices where we know - // the effective subgroup width. Decode-small uses workgroup reductions - // and should not be blocked by browser adapters that support subgroups - // but report a subgroup-width range. - if !is_decode_candidate { - self.data.device.fixed_width_subgroup_size()?; - } - let v_shape = v.shape(); - if q_shape[0] != k_shape[0] - || q_shape[0] != v_shape[0] - || k_shape[1] != v_shape[1] - || k_shape[2] != v_shape[2] - || q_shape[3] != k_shape[3] - || q_shape[3] != v_shape[3] - || q_shape[0] == 0 - || q_shape[1] == 0 - || q_shape[2] == 0 - || k_shape[1] == 0 - || !q_shape[1].is_multiple_of(k_shape[1]) - || q_shape[3] == 0 - || k_shape[2] == 0 - { - return None; - } - if let Some(mask) = mask - && mask.shape() != [q_shape[2], k_shape[2]] - { - return None; - } - if causal && q_shape[2] != k_shape[2] { - // Causal optimisation only kicks in for self-attention prefill - // where q_seq_len == kv_seq_len. Other shapes (e.g. cached decode) - // fall back to the masked path. - return None; - } - let batch = u32::try_from(q_shape[0]).ok()?; - let num_heads = u32::try_from(q_shape[1]).ok()?; - let q_seq_len = u32::try_from(q_shape[2]).ok()?; - let head_dim = u32::try_from(q_shape[3]).ok()?; - let row_dispatch = batch.checked_mul(num_heads)?.checked_mul(q_seq_len)?; - let x_dispatch = head_dim.div_ceil(8); - let max_dispatch = self - .data - .device - .limits() - .max_compute_workgroups_per_dimension; - if x_dispatch > max_dispatch || row_dispatch > max_dispatch { - return None; - } - - let operation = FlashAttentionOperation::new(FlashAttentionInputs { - q: self.data.key, - k: k.data.key, - v: v.data.key, - mask: mask.map(|mask| mask.data.key), - q_shape, - k_shape, - v_shape, - scale, - input_dtype: self.datatype(), - causal, - }); - Some(Self::from_parts(self.data.flash_attention(operation))) - } - pub fn device(&self) -> &Device { &self.data.device } diff --git a/fusor-ml/core/src/view.rs b/fusor-ml/core/src/view.rs new file mode 100644 index 000000000..2591ea55f --- /dev/null +++ b/fusor-ml/core/src/view.rs @@ -0,0 +1,788 @@ +use std::hash::Hash; + +use crate::{ + DataTypeEnum, Layout, Tensor, TensorData, + compute_graph::NodeIndex, + mir::{ + inputs::MirValue, + kernel_backend::DirectKernel, + operation::Operation, + workgroup_shape::{Constraint, WorkgroupShape, WorkgroupShapeConstraints}, + }, + nary_wise::{ElementwiseOperation, NaryExpr, NaryOp, NaryScalar}, + visit_tiled::distribute_workgroups, +}; + +const BLOCKSIZE: u32 = 256; + +/// A zero-dispatch view of a node's logical value space. +/// +/// `layout` maps output coordinates to flat indices in the input's logical +/// row-major value space (`flat = offset + Σ coord_i * strides[i]`). Because +/// it indexes the *logical* space — not a concrete buffer — every producer's +/// output is contiguous by definition, so restride, transpose, broadcast, +/// slice, reshape, and resize are all plain stride arithmetic here. +/// +/// `defined` is a prefix box of `layout.shape()`: coordinates with +/// `coord_i < defined[i]` for every axis read input data; anything outside +/// reads `fill`. A fully-defined view (`defined == shape`) is a pure +/// relayout; a partially-defined view is a clip + pad (resize). +#[derive(Clone, Debug)] +pub(crate) struct ViewOperation { + pub(crate) input: NodeIndex, + pub(crate) layout: Layout, + /// Logical shape of the input node's value space — the space `layout` + /// indexes into. Used to recover per-dimension input coordinates when the + /// view cannot stay a flat index (gather fallback) and to bounds-check. + pub(crate) input_shape: Box<[usize]>, + pub(crate) defined: Box<[usize]>, + pub(crate) fill: NaryScalar, + pub(crate) datatype: DataTypeEnum, +} + +impl ViewOperation { + /// A fully-defined view: pure relayout of the input's logical space. + pub(crate) fn fully_defined( + input: NodeIndex, + layout: Layout, + input_shape: impl Into>, + datatype: DataTypeEnum, + ) -> Self { + let defined = layout.shape().into(); + Self { + input, + layout, + input_shape: input_shape.into(), + defined, + fill: zero_scalar(datatype), + datatype, + } + } + + pub(crate) fn is_fully_defined(&self) -> bool { + self.defined.as_ref() == self.layout.shape() + } + + pub(crate) fn shape(&self) -> &[usize] { + self.layout.shape() + } + + /// Resolve as a zero-cost view over the input's concrete buffer, if the + /// view's logical layout composes with the buffer's layout as a single + /// strided layout. Partially-defined views never qualify — their fill + /// region has no backing memory. + pub(crate) fn try_map_tensor(&self, input: &TensorData) -> Option { + if !self.is_fully_defined() { + return None; + } + let composed = compose_layouts(&self.layout, input.layout())?; + Some(TensorData::new_from_parts( + input.device(), + input.buffer().clone(), + composed, + input.datatype(), + )) + } + + /// The gather expression materializing this view: per output coordinate, + /// load the input at the mapped logical coordinates, or `fill` outside + /// the defined box. + fn copy_expression(&self) -> Option { + let flat = self.flat_logical_expression()?; + let indices = row_major_indices_from_flat(flat, &self.input_shape)?; + let copied = NaryExpr::indexed_input(0, indices); + if self.is_fully_defined() { + return Some(copied); + } + Some(NaryExpr::select( + self.in_defined_bounds_expression(), + copied, + NaryExpr::scalar(self.fill), + DataTypeEnum::U32, + self.datatype, + )) + } + + /// `offset + Σ DimIndex(d) * strides[d]` as a u32 expression. + fn flat_logical_expression(&self) -> Option { + let mut flat = NaryExpr::scalar(NaryScalar::U32(self.layout.offset().try_into().ok()?)); + for (axis, (&stride, &dim)) in self + .layout + .strides() + .iter() + .zip(self.layout.shape()) + .enumerate() + { + if stride == 0 || dim == 1 { + continue; + } + let stride: u32 = stride.try_into().ok()?; + let dim_index = NaryExpr::DimIndex(axis); + let term = if stride == 1 { + dim_index + } else { + NaryExpr::unary_op( + dim_index, + "mul_const", + NaryOp::MulConst(NaryScalar::U32(stride)), + DataTypeEnum::U32, + DataTypeEnum::U32, + ) + }; + flat = NaryExpr::add(flat, term, DataTypeEnum::U32); + } + Some(flat) + } + + fn in_defined_bounds_expression(&self) -> NaryExpr { + let mut condition = NaryExpr::scalar(NaryScalar::U32(1)); + for (dim, (&defined, &size)) in self.defined.iter().zip(self.layout.shape()).enumerate() { + if defined >= size { + continue; + } + let lt_defined = NaryExpr::unary_op( + NaryExpr::DimIndex(dim), + "lt_defined", + NaryOp::LessConst(NaryScalar::U32(defined as u32)), + DataTypeEnum::U32, + DataTypeEnum::U32, + ); + condition = NaryExpr::mul(condition, lt_defined, DataTypeEnum::U32); + } + condition + } +} + +fn zero_scalar(datatype: DataTypeEnum) -> NaryScalar { + match datatype { + DataTypeEnum::F32 => NaryScalar::F32(0.0), + DataTypeEnum::F16 => NaryScalar::F16(half::f16::from_f32(0.0)), + DataTypeEnum::U32 => NaryScalar::U32(0), + } +} + +/// Re-express `outer` over `inner`'s index space. +/// +/// `outer` maps its output coordinates to flat indices in the row-major space +/// of `inner.shape()`; `inner` maps its own coordinates to some target space +/// (a logical input space, or a concrete buffer). The result maps `outer`'s +/// coordinates directly to `inner`'s target space, or `None` when the +/// composition is not expressible as a single strided layout (e.g. a reshape +/// that regroups elements across non-contiguous strides). +/// +/// The check is exact: each outer stride and the outer offset are decomposed +/// as mixed-radix digits over `inner`'s contiguity chunks, and the per-chunk +/// digit spans must never carry for any in-range coordinate. +pub(crate) fn compose_layouts(outer: &Layout, inner: &Layout) -> Option { + // Merge inner dims into contiguity chunks (right to left): a run of dims + // whose strides chain as stride[i] == stride[i+1] * shape[i+1] acts as one + // flat axis. Size-1 dims are transparent. + let mut chunks: Vec<(usize, usize)> = Vec::new(); // (extent, target stride) + for (&dim, &stride) in inner.shape().iter().zip(inner.strides()).rev() { + if dim == 1 { + continue; + } + match chunks.last_mut() { + Some((extent, last_stride)) if stride == *last_stride * *extent => { + *extent *= dim; + } + _ => chunks.push((dim, stride)), + } + } + chunks.reverse(); + if chunks.iter().any(|(extent, _)| *extent == 0) { + return None; + } + + // Row-major radix strides over the chunk extents: chunk k covers flat + // positions in steps of radix[k]. + let mut radix = vec![0usize; chunks.len()]; + let mut acc = 1usize; + for (k, (extent, _)) in chunks.iter().enumerate().rev() { + radix[k] = acc; + acc = acc.checked_mul(*extent)?; + } + + // Mixed-radix decomposition over the chunk extents. Oversized digits are + // allowed here; the no-carry span check below is the correctness gate. + let digits = |mut value: usize| -> Option> { + let mut digits = vec![0usize; chunks.len()]; + for (k, radix) in radix.iter().enumerate() { + digits[k] = value / radix; + value %= radix; + } + (value == 0).then_some(digits) + }; + + let offset_digits = digits(outer.offset())?; + let stride_digits = outer + .strides() + .iter() + .zip(outer.shape()) + .map(|(&stride, &dim)| { + if dim <= 1 { + // A dim that never steps contributes nothing; its stride is + // irrelevant (and may be a degenerate placeholder). + Some(vec![0; chunks.len()]) + } else { + digits(stride) + } + }) + .collect::>>()?; + + // No-carry check: the offset digit plus every dim's maximum travel along + // each chunk must stay within the chunk extent. + for (k, (extent, _)) in chunks.iter().enumerate() { + let mut max_coord = offset_digits[k]; + for (digits, &dim) in stride_digits.iter().zip(outer.shape()) { + max_coord = max_coord.checked_add(digits[k].checked_mul(dim.saturating_sub(1))?)?; + } + if max_coord >= *extent { + return None; + } + } + + let offset = inner.offset() + + offset_digits + .iter() + .zip(&chunks) + .map(|(digit, (_, stride))| digit * stride) + .sum::(); + let strides: Box<[usize]> = stride_digits + .iter() + .map(|digits| { + digits + .iter() + .zip(&chunks) + .map(|(digit, (_, stride))| digit * stride) + .sum() + }) + .collect(); + Some(Layout::from_parts(offset, outer.shape().into(), strides)) +} + +/// One base-dimension coordinate of an affine view: +/// `base_coord = constant + Σ coefficient * out_coord[dim]`. +#[derive(Clone, Debug, PartialEq, Eq)] +pub(crate) struct AffineIndex { + pub(crate) constant: u32, + pub(crate) terms: Vec<(usize, u32)>, +} + +impl AffineIndex { + /// Render as an index expression, substituting `out[j]` for the view's + /// output coordinate `j`. Collapses to the bare coordinate when this is + /// an identity mapping. + pub(crate) fn to_expr(&self, out: &[NaryExpr]) -> NaryExpr { + if self.constant == 0 + && let [(dim, 1)] = self.terms.as_slice() + { + return out[*dim].clone(); + } + let mut expr = NaryExpr::scalar(NaryScalar::U32(self.constant)); + for &(dim, coefficient) in &self.terms { + let term = if coefficient == 1 { + out[dim].clone() + } else { + NaryExpr::unary_op( + out[dim].clone(), + "mul_const", + NaryOp::MulConst(NaryScalar::U32(coefficient)), + DataTypeEnum::U32, + DataTypeEnum::U32, + ) + }; + expr = NaryExpr::add(expr, term, DataTypeEnum::U32); + } + expr + } +} + +/// Decompose a view layout into per-base-dimension affine coordinate +/// expressions, when every output dimension's stride and the offset split +/// into mixed-radix digits over `base_shape` without carries. This is the +/// divmod-free form of the view: reshape/restride/broadcast/slice over a +/// dense base all qualify. +pub(crate) fn affine_dim_indices( + layout: &Layout, + base_shape: &[usize], +) -> Option> { + if base_shape.contains(&0) { + return None; + } + let radix = Layout::continuous_strides(base_shape); + + let digits = |mut value: usize| -> Option> { + let mut digits = vec![0u32; base_shape.len()]; + for (k, radix) in radix.iter().enumerate() { + digits[k] = u32::try_from(value / radix).ok()?; + value %= radix; + } + (value == 0).then_some(digits) + }; + + let offset_digits = digits(layout.offset())?; + let stride_digits = layout + .strides() + .iter() + .zip(layout.shape()) + .map(|(&stride, &dim)| { + if dim <= 1 { + Some(vec![0; base_shape.len()]) + } else { + digits(stride) + } + }) + .collect::>>()?; + + // No-carry check: the maximum coordinate reached along each base dim must + // stay inside that dim. + for (k, &extent) in base_shape.iter().enumerate() { + let mut max_coord = offset_digits[k] as usize; + for (digits, &dim) in stride_digits.iter().zip(layout.shape()) { + max_coord = + max_coord.checked_add((digits[k] as usize).checked_mul(dim.saturating_sub(1))?)?; + } + if max_coord >= extent { + return None; + } + } + + Some( + (0..base_shape.len()) + .map(|k| AffineIndex { + constant: offset_digits[k], + terms: stride_digits + .iter() + .enumerate() + .filter(|(_, digits)| digits[k] != 0) + .map(|(j, digits)| (j, digits[k])) + .collect(), + }) + .collect(), + ) +} + +fn row_major_indices_from_flat(flat: NaryExpr, shape: &[usize]) -> Option> { + let mut indices = Vec::with_capacity(shape.len()); + for axis in 0..shape.len() { + let divisor = shape[axis + 1..] + .iter() + .try_fold(1u32, |acc, dim| acc.checked_mul((*dim).try_into().ok()?))?; + let dim = u32::try_from(shape[axis]).ok()?; + let quotient = if divisor == 1 { + flat.clone() + } else { + NaryExpr::unary_op( + flat.clone(), + "div_const", + NaryOp::DivConst(NaryScalar::U32(divisor)), + DataTypeEnum::U32, + DataTypeEnum::U32, + ) + }; + indices.push(if dim == 1 { + NaryExpr::scalar(NaryScalar::U32(0)) + } else if axis == 0 { + quotient + } else { + NaryExpr::unary_op( + quotient, + "rem_const", + NaryOp::RemConst(NaryScalar::U32(dim)), + DataTypeEnum::U32, + DataTypeEnum::U32, + ) + }); + } + Some(indices) +} + +impl Operation for ViewOperation { + fn hash_kernel_fields(&self, state: &mut rustc_hash::FxHasher) { + self.layout.offset().hash(state); + self.layout.shape().hash(state); + self.layout.strides().hash(state); + self.input_shape.hash(state); + self.defined.hash(state); + self.fill.hash(state); + } + + fn workgroup_shape_constraints(&self, _: &crate::Device) -> WorkgroupShapeConstraints { + let mut constraints = WorkgroupShapeConstraints::new(); + constraints.add_constraint(0, Constraint::equals(BLOCKSIZE)); + constraints.add_constraint(1, Constraint::equals(1)); + constraints.add_constraint(2, Constraint::equals(1)); + constraints + } + + fn dispatch_size(&self, _: &WorkgroupShape, inputs: &[MirValue]) -> [u32; 3] { + let output = inputs[1].as_tensor().unwrap(); + let total_workgroups = (output.layout().shape().iter().product::() as u32) + .div_ceil(crate::TILE_SIZE * BLOCKSIZE); + distribute_workgroups( + total_workgroups, + output + .device() + .limits() + .max_compute_workgroups_per_dimension, + ) + } + + fn visit_dependencies(&self, f: &mut dyn FnMut(NodeIndex)) { + f(self.input); + } + + fn inputs(&self, nodes: &crate::compute_graph::ComputeGraphInner) -> Vec { + let input = nodes.get_cached_result(self.input).unwrap().clone(); + let output = + TensorData::new_for_shape(input.device(), self.layout.shape(), input.datatype()); + vec![input.into(), output.into()] + } + + fn output(&self, _: &crate::compute_graph::ComputeGraphInner, inputs: &[MirValue]) -> MirValue { + inputs[1].clone() + } + + fn build_direct_kernel( + &self, + graph: &crate::compute_graph::ComputeGraphInner, + workgroup_shape: &WorkgroupShape, + inputs: &[MirValue], + ) -> Option { + let operation = ElementwiseOperation { + inputs: vec![self.input], + expression: self.copy_expression()?, + shape: self.layout.shape().into(), + output_datatype: self.datatype, + }; + crate::nary_direct::build_nary_direct_kernel_to_output( + &operation, + graph, + workgroup_shape, + inputs, + 1, + ) + } + + fn name(&self) -> String { + format!( + "view_{}", + self.layout + .shape() + .iter() + .map(|x| x.to_string()) + .collect::>() + .join("x") + ) + } +} + +impl Tensor { + /// The view spec to layer a new view on top of this tensor: composes with + /// an existing fully-defined view (so chains collapse at construction) or + /// starts from this tensor's logical space. + fn view_base(&self) -> (NodeIndex, Layout, Box<[usize]>) { + if let Some(view) = self.device().compute_graph().get_view(self.key()) + && view.is_fully_defined() + { + return (view.input, view.layout.clone(), view.input_shape.clone()); + } + ( + self.key(), + Layout::contiguous(self.shape()), + self.shape().into(), + ) + } + + fn add_view_op(&self, op: ViewOperation) -> Tensor { + Tensor::from_parts(self.data.view(op)) + } + + pub fn restride(&self, specs: impl Into>) -> Tensor { + let specs = specs.into(); + let (input, base, input_shape) = self.view_base(); + self.add_view_op(ViewOperation::fully_defined( + input, + base.restride(&specs), + input_shape, + self.datatype(), + )) + } + + /// Replace the tensor's layout with `new_layout` over its logical value + /// space (the row-major order of its elements). The offset and strides + /// are in logical elements, independent of how the data is laid out in + /// any concrete buffer. + pub fn restride_layout(&self, new_layout: Layout) -> Tensor { + let numel = self.shape().iter().product::(); + let max_index = new_layout + .shape() + .iter() + .zip(new_layout.strides()) + .fold(new_layout.offset(), |acc, (dim, stride)| { + acc + dim.saturating_sub(1) * stride + }); + assert!( + numel == 0 || max_index < numel, + "restride_layout out of bounds: layout reaches element {max_index} \ + but the input has only {numel} elements" + ); + let (input, base, input_shape) = self.view_base(); + let layout = compose_layouts(&new_layout, &base).unwrap_or_else(|| { + panic!( + "restride_layout could not compose {new_layout:?} with the existing view {base:?}" + ) + }); + self.add_view_op(ViewOperation::fully_defined( + input, + layout, + input_shape, + self.datatype(), + )) + } + + pub fn broadcast_as(&self, out_shape: impl AsRef<[usize]>) -> Tensor { + let out_shape = out_shape.as_ref(); + let shape = self.shape(); + assert!( + out_shape.len() >= shape.len(), + "The output rank must be at least the input rank" + ); + let specs: Vec = (0..out_shape.len()) + .map(|out_i| { + let in_i = out_i as isize - (out_shape.len() as isize - shape.len() as isize); + if in_i < 0 { + crate::StrideSpec::dim_with(0, out_shape[out_i], 0) + } else { + let in_i = in_i as usize; + if shape[in_i] == 1 && out_shape[out_i] > 1 { + crate::StrideSpec::dim_with(in_i, out_shape[out_i], 0) + } else { + crate::StrideSpec::dim(in_i, out_shape[out_i]) + } + } + }) + .collect(); + self.restride(specs) + } + + pub(crate) fn broadcast_together(first: &Tensor, second: &Tensor) -> (Tensor, Tensor) { + assert_eq!(first.datatype(), second.datatype()); + let first_shape = first.shape(); + let second_shape = second.shape(); + let rank = first_shape.len().max(second_shape.len()); + let shape: Vec = (0..rank) + .map(|i| { + let a = i + first_shape.len(); + let b = i + second_shape.len(); + let a = if a >= rank { first_shape[a - rank] } else { 1 }; + let b = if b >= rank { second_shape[b - rank] } else { 1 }; + assert!( + a == b || a == 1 || b == 1, + "Cannot broadcast shapes {:?} and {:?}", + first_shape, + second_shape + ); + a.max(b) + }) + .collect(); + (first.broadcast_as(&shape), second.broadcast_as(&shape)) + } + + pub(crate) fn broadcast_then_elementwise_op( + first: &Tensor, + second: &Tensor, + op: impl Fn(Tensor, Tensor) -> Tensor, + ) -> Tensor { + let (b1, b2) = Tensor::broadcast_together(first, second); + assert_eq!(b1.shape(), b2.shape()); + op(b1, b2) + } + + pub fn reshape(&self, new_shape: impl AsRef<[usize]>) -> Tensor { + let new_shape = new_shape.as_ref(); + assert_eq!( + new_shape.iter().product::(), + self.shape().iter().product::(), + "Reshape requires the number of elements to be the same. \ + Current shape: {:?}, target shape: {:?}", + self.shape(), + new_shape + ); + let reinterpret = Layout::contiguous(new_shape); + let (input, base, input_shape) = self.view_base(); + let op = match compose_layouts(&reinterpret, &base) { + Some(layout) => { + ViewOperation::fully_defined(input, layout, input_shape, self.datatype()) + } + // The reshape regroups elements across the view's non-contiguous + // strides: keep it as a flat reinterpret of this tensor's own + // logical space (a chained view). + None => { + ViewOperation::fully_defined(self.key(), reinterpret, self.shape(), self.datatype()) + } + }; + self.add_view_op(op) + } + + /// Resize to `new_shape`, clipping or zero-padding each axis: coordinates + /// inside `min(old, new)` per axis keep their values, anything beyond is + /// zero. + pub fn resize(&self, new_shape: impl AsRef<[usize]>) -> Tensor { + let new_shape = new_shape.as_ref(); + let old_shape = self.shape(); + assert_eq!( + new_shape.len(), + old_shape.len(), + "resize requires matching ranks (got {old_shape:?} -> {new_shape:?}); use reshape \ + to change rank" + ); + let defined: Box<[usize]> = new_shape + .iter() + .zip(old_shape) + .map(|(new, old)| (*new).min(*old)) + .collect(); + // Within the defined box the old row-major strides address the input + // exactly; outside it the load is masked to `fill`. + let resize = Layout::from_parts(0, new_shape.into(), Layout::continuous_strides(old_shape)); + let (input, base, input_shape) = self.view_base(); + let op = match compose_layouts(&resize, &base) { + Some(layout) => ViewOperation { + input, + layout, + input_shape, + defined, + fill: zero_scalar(self.datatype()), + datatype: self.datatype(), + }, + None => ViewOperation { + input: self.key(), + layout: resize, + input_shape: old_shape.into(), + defined, + fill: zero_scalar(self.datatype()), + datatype: self.datatype(), + }, + }; + self.add_view_op(op) + } + + pub fn flatten_last_n(&self, from_end: usize) -> Tensor { + assert!( + from_end < self.rank(), + "flatten_last_n FROM_END must be less than input rank" + ); + let out_rank = self.rank() - from_end; + let new_shape: Vec = (0..out_rank) + .map(|i| { + if i < self.rank() - 1 - from_end { + self.shape()[i] + } else if i == self.rank() - 1 - from_end { + self.shape()[i..].iter().product() + } else { + 1 + } + }) + .collect(); + self.reshape(new_shape) + } + + pub fn flatten_first_n(&self, from_start: usize) -> Tensor { + assert!( + from_start < self.rank(), + "flatten_first_n FROM_START must be less than input rank" + ); + let out_rank = self.rank() - from_start; + let new_shape: Vec = (0..out_rank) + .map(|i| { + if i == 0 { + self.shape()[..=from_start].iter().product() + } else { + self.shape()[i + from_start] + } + }) + .collect(); + self.reshape(new_shape) + } + + pub fn flatten_all(&self) -> Tensor { + let size = self.shape().iter().product(); + self.reshape([size]) + } +} + +pub use fusor_types::ShapeWithOneHole; + +#[cfg(test)] +mod tests { + use super::*; + + fn strided(offset: usize, shape: &[usize], strides: &[usize]) -> Layout { + Layout::from_parts(offset, shape.into(), strides.into()) + } + + #[test] + fn compose_with_contiguous_is_identity() { + let inner = Layout::contiguous(&[4, 6]); + let outer = strided(3, &[2, 6], &[12, 1]); + let composed = compose_layouts(&outer, &inner).unwrap(); + assert_eq!(composed.offset(), 3); + assert_eq!(composed.shape(), &[2, 6]); + assert_eq!(composed.strides(), &[12, 1]); + } + + #[test] + fn compose_transpose_then_slice() { + // inner: transpose of a [4, 6] tensor -> shape [6, 4], strides [1, 6] + let inner = strided(0, &[6, 4], &[1, 6]); + // outer: narrow rows 2..5 of the transposed view + let outer = strided(2 * 4, &[3, 4], &[4, 1]); + let composed = compose_layouts(&outer, &inner).unwrap(); + assert_eq!(composed.shape(), &[3, 4]); + // flat index f over inner's [6,4] output decomposes as (r, c) with + // target = r * 1 + c * 6 + assert_eq!(composed.offset(), 2); + assert_eq!(composed.strides(), &[1, 6]); + } + + #[test] + fn compose_reshape_of_transpose_fails() { + // Flat reinterpret of a transposed (non-contiguous) view regroups + // elements: not expressible as one strided layout. + let inner = strided(0, &[6, 4], &[1, 6]); + let outer = Layout::contiguous(&[8, 3]); + assert!(compose_layouts(&outer, &inner).is_none()); + } + + #[test] + fn compose_reshape_merges_contiguous_chunks() { + // inner: a sliced batch of contiguous rows: [2, 12] with strides + // [24, 1] inside a larger allocation (row padding -> chunked). + let inner = strided(0, &[2, 12], &[24, 1]); + // outer: reshape each row into [3, 4] -> [2, 3, 4] flat over [2, 12] + let outer = Layout::contiguous(&[2, 3, 4]); + let composed = compose_layouts(&outer, &inner).unwrap(); + assert_eq!(composed.shape(), &[2, 3, 4]); + assert_eq!(composed.strides(), &[24, 4, 1]); + } + + #[test] + fn compose_broadcast_strides() { + let inner = Layout::contiguous(&[4, 6]); + // outer broadcasts a [6] row across 5: shape [5, 6], strides [0, 1] + let outer = strided(6, &[5, 6], &[0, 1]); + let composed = compose_layouts(&outer, &inner).unwrap(); + assert_eq!(composed.offset(), 6); + assert_eq!(composed.strides(), &[0, 1]); + } + + #[test] + fn compose_rejects_out_of_bounds_span() { + let inner = strided(0, &[2, 12], &[24, 1]); + // Steps of 12 over 4 elements span 36 flat positions, crossing the + // row chunk boundary (each chunk holds 12). + let outer = strided(0, &[4], &[12]); + assert!(compose_layouts(&outer, &inner).is_none()); + } +} diff --git a/fusor-ml/core/tests/attention.rs b/fusor-ml/core/tests/attention.rs new file mode 100644 index 000000000..2aee33924 --- /dev/null +++ b/fusor-ml/core/tests/attention.rs @@ -0,0 +1,325 @@ +//! Correctness gates for the generic attention row program across its +//! lowering regimes: single-tile decode, the online multi-tile streaming +//! loop (KV beyond one workgroup bucket), causal prefill (axis bound), +//! additive masks, GQA head mapping, and f16 IO. + +use fusor_core::{Device, Tensor}; + +fn values(len: usize, scale: f32) -> Vec { + (0..len).map(|i| ((i as f32) * scale).sin()).collect() +} + +struct AttentionCase { + batch: usize, + heads: usize, + kv_heads: usize, + q_len: usize, + kv_len: usize, + head_dim: usize, + causal: bool, + masked: bool, +} + +/// CPU reference for `softmax(q·kᵀ·scale [+ mask])·v` with GQA expansion. +#[allow(clippy::too_many_arguments)] +fn cpu_attention( + case: &AttentionCase, + q: &[f32], + k: &[f32], + v: &[f32], + mask: Option<&[f32]>, + scale: f32, +) -> Vec { + let AttentionCase { + batch, + heads, + kv_heads, + q_len, + kv_len, + head_dim, + causal, + .. + } = *case; + let groups = heads / kv_heads; + let mut out = vec![0.0f32; batch * heads * q_len * head_dim]; + for b in 0..batch { + for h in 0..heads { + let kv_h = h / groups; + for qi in 0..q_len { + let q_base = ((b * heads + h) * q_len + qi) * head_dim; + let mut scores = vec![f32::NEG_INFINITY; kv_len]; + for (pos, score) in scores.iter_mut().enumerate() { + if causal && pos > qi { + continue; + } + let k_base = ((b * kv_heads + kv_h) * kv_len + pos) * head_dim; + let mut dot = 0.0f32; + for d in 0..head_dim { + dot += q[q_base + d] * k[k_base + d]; + } + let mut value = dot * scale; + if let Some(mask) = mask { + value += mask[qi * kv_len + pos]; + } + *score = value; + } + let max = scores.iter().fold(f32::NEG_INFINITY, |a, &b| a.max(b)); + let weights: Vec = scores.iter().map(|s| (s - max).exp()).collect(); + let total: f32 = weights.iter().sum(); + for d in 0..head_dim { + let mut acc = 0.0f32; + for (pos, weight) in weights.iter().enumerate() { + let v_base = ((b * kv_heads + kv_h) * kv_len + pos) * head_dim; + acc += weight / total * v[v_base + d]; + } + out[q_base + d] = acc; + } + } + } + } + out +} + +fn check_attention(case: AttentionCase, tolerance: f32) { + pollster::block_on(async { + let Ok(device) = Device::new().await else { + return; + }; + let AttentionCase { + batch, + heads, + kv_heads, + q_len, + kv_len, + head_dim, + causal, + masked, + } = case; + let q_data = values(batch * heads * q_len * head_dim, 0.13); + let k_data = values(batch * kv_heads * kv_len * head_dim, 0.07); + let v_data = values(batch * kv_heads * kv_len * head_dim, 0.11); + let mask_data = masked.then(|| { + (0..q_len * kv_len) + .map(|i| if i % 7 == 0 { -1.5 } else { 0.25 }) + .collect::>() + }); + let scale = 1.0 / (head_dim as f32).sqrt(); + + let q = Tensor::from_slice(&device, [batch, heads, q_len, head_dim], &q_data); + let k = Tensor::from_slice(&device, [batch, kv_heads, kv_len, head_dim], &k_data); + let v = Tensor::from_slice(&device, [batch, kv_heads, kv_len, head_dim], &v_data); + let mask_tensor = mask_data + .as_ref() + .map(|data| Tensor::from_slice(&device, [q_len, kv_len], data.as_slice())); + + let out = if causal { + q.flash_attention_causal(&k, &v, scale) + } else { + q.flash_attention(&k, &v, scale, mask_tensor.as_ref()) + }; + assert_eq!( + out.count_kernels_to_resolve(), + 1, + "attention should lower as one row-program kernel" + ); + let actual = out.as_slice::<4, f32>().await.unwrap(); + + let expected = cpu_attention( + &case, + &q_data, + &k_data, + &v_data, + mask_data.as_deref(), + scale, + ); + for b in 0..batch { + for h in [0, heads - 1] { + for qi in [0, q_len / 2, q_len - 1] { + for d in [0, head_dim / 2, head_dim - 1] { + let want = expected[((b * heads + h) * q_len + qi) * head_dim + d]; + let got = actual[[b, h, qi, d]]; + assert!( + (got - want).abs() < tolerance, + "b={b} h={h} q={qi} d={d}: got {got}, expected {want}" + ); + } + } + } + } + }); +} + +#[test] +fn attention_decode_single_tile() { + check_attention( + AttentionCase { + batch: 1, + heads: 8, + kv_heads: 2, + q_len: 1, + kv_len: 100, + head_dim: 64, + causal: false, + masked: false, + }, + 1e-4, + ); +} + +#[test] +fn attention_decode_streams_long_kv() { + // KV beyond the largest workgroup bucket: the online tile loop with + // rescaling, including a ragged final tile. + check_attention( + AttentionCase { + batch: 1, + heads: 4, + kv_heads: 4, + q_len: 1, + kv_len: 3000, + head_dim: 64, + causal: false, + masked: false, + }, + 1e-4, + ); +} + +#[test] +fn attention_causal_prefill() { + check_attention( + AttentionCase { + batch: 1, + heads: 4, + kv_heads: 2, + q_len: 96, + kv_len: 96, + head_dim: 32, + causal: true, + masked: false, + }, + 1e-4, + ); +} + +#[test] +fn attention_causal_prefill_streams_long_kv() { + // Causal with q == kv beyond one bucket: the per-row axis bound must + // stop each query row at its own position while later rows stream on. + check_attention( + AttentionCase { + batch: 1, + heads: 2, + kv_heads: 2, + q_len: 1536, + kv_len: 1536, + head_dim: 32, + causal: true, + masked: false, + }, + 1e-4, + ); +} + +#[test] +fn attention_odd_sized_causal_prefill() { + // Odd extents that don't align with any tile or bucket boundary. + check_attention( + AttentionCase { + batch: 1, + heads: 4, + kv_heads: 2, + q_len: 100, + kv_len: 100, + head_dim: 32, + causal: true, + masked: false, + }, + 1e-4, + ); +} + +#[test] +fn attention_masked_prefill() { + check_attention( + AttentionCase { + batch: 1, + heads: 4, + kv_heads: 4, + q_len: 48, + kv_len: 80, + head_dim: 32, + causal: false, + masked: true, + }, + 1e-4, + ); +} + +#[test] +fn attention_f16_io() { + pollster::block_on(async { + let Ok(device) = Device::new().await else { + return; + }; + if !device.f16_supported() { + return; + } + let (heads, kv_len, head_dim) = (4usize, 64usize, 32usize); + let q_data = values(heads * head_dim, 0.13); + let k_data = values(heads * kv_len * head_dim, 0.07); + let v_data = values(heads * kv_len * head_dim, 0.11); + let scale = 1.0 / (head_dim as f32).sqrt(); + let to_f16 = |data: &[f32]| { + data.iter() + .map(|&x| half::f16::from_f32(x)) + .collect::>() + }; + let q = Tensor::from_slice(&device, [1, heads, 1, head_dim], &to_f16(&q_data)); + let k = Tensor::from_slice(&device, [1, heads, kv_len, head_dim], &to_f16(&k_data)); + let v = Tensor::from_slice(&device, [1, heads, kv_len, head_dim], &to_f16(&v_data)); + + let out = q.flash_attention(&k, &v, scale, None); + let actual = out.as_slice::<4, half::f16>().await.unwrap(); + + let case = AttentionCase { + batch: 1, + heads, + kv_heads: heads, + q_len: 1, + kv_len, + head_dim, + causal: false, + masked: false, + }; + let expected = cpu_attention(&case, &q_data, &k_data, &v_data, None, scale); + for h in [0, heads - 1] { + for d in [0, head_dim - 1] { + let want = expected[h * head_dim + d]; + let got = actual[[0, h, 0, d]].to_f32(); + assert!( + (got - want).abs() < 1e-2, + "h={h} d={d}: got {got}, expected {want}" + ); + } + } + }); +} + +#[test] +fn attention_many_rows_streams_tiles() { + // Prefill-shaped rows with the axis spanning several online tiles. + check_attention( + AttentionCase { + batch: 1, + heads: 2, + kv_heads: 2, + q_len: 256, + kv_len: 1536, + head_dim: 32, + causal: false, + masked: false, + }, + 1e-4, + ); +} diff --git a/fusor-ml/core/tests/fused_reduce.rs b/fusor-ml/core/tests/fused_reduce.rs new file mode 100644 index 000000000..f819f6fe7 --- /dev/null +++ b/fusor-ml/core/tests/fused_reduce.rs @@ -0,0 +1,294 @@ +//! Fused map-reduce gates: composed reduce clusters that recognition does +//! NOT claim must still collapse to a single kernel. The resolver inlines the +//! elementwise producer into the reduce, and contraction-shaped clusters — +//! whatever their dim order or broadcast structure — lower through the tiled +//! (workgroup-cached) path. A fusion miss materializes the full index space, +//! which these tests turn into a kernel-count failure. + +use fusor_core::{Device, QMatrix, Tensor}; +use fusor_gguf::GgmlType; + +/// `a [M, K]` and `b [N, K]` broadcast into the `[M, N, K]` index space. +fn broadcast_factors( + device: &Device, + m: usize, + n: usize, + k: usize, + a_values: &[f32], + b_values: &[f32], +) -> (Tensor, Tensor) { + let a = Tensor::from_slice(device, [m, k], a_values); + let b = Tensor::from_slice(device, [n, k], b_values); + let a3 = a.reshape([m, 1, k]).broadcast_as([m, n, k]); + let b3 = b.reshape([1, n, k]).broadcast_as([m, n, k]); + (a3, b3) +} + +fn pattern(len: usize, scale: f32) -> Vec { + (0..len).map(|i| ((i as f32) * scale).sin()).collect() +} + +#[test] +fn broadcast_composed_contraction_fuses_to_single_kernel() { + pollster::block_on(async { + let Ok(device) = Device::new().await else { + return; + }; + + // Edge shapes on every tiled dim: M not a multiple of the 32-row + // tile, K not a multiple of the 8-deep tile. + let (m, n, k) = (72usize, 64usize, 20usize); + let a_values = pattern(m * k, 0.13); + let b_values = pattern(n * k, 0.07); + let (a3, b3) = broadcast_factors(&device, m, n, k, &a_values, &b_values); + + let out = (&a3 * &b3).sum(2); + assert_eq!( + out.count_kernels_to_resolve(), + 1, + "broadcast-composed contraction must fuse into one map-reduce kernel" + ); + + let slice = out.as_slice::<2, f32>().await.unwrap(); + for row in [0usize, 31, 32, m - 1] { + for col in [0usize, 33, n - 1] { + let expected: f32 = (0..k) + .map(|kk| a_values[row * k + kk] * b_values[col * k + kk]) + .sum(); + let actual = slice[[row, col]]; + assert!( + (actual - expected).abs() < 1e-3, + "[{row}, {col}]: got {actual}, expected {expected}" + ); + } + } + }); +} + +#[test] +fn small_composed_contraction_fuses_through_serial_path() { + pollster::block_on(async { + let Ok(device) = Device::new().await else { + return; + }; + + // Below the tiled gates (m < 32): the fused reduce still collapses + // to one kernel through the serial per-output path. + let (m, n, k) = (8usize, 6usize, 16usize); + let a_values = pattern(m * k, 0.21); + let b_values = pattern(n * k, 0.17); + let (a3, b3) = broadcast_factors(&device, m, n, k, &a_values, &b_values); + + let out = (&a3 * &b3).sum(2); + assert_eq!(out.count_kernels_to_resolve(), 1, "small fused contraction"); + + let slice = out.as_slice::<2, f32>().await.unwrap(); + for row in 0..m { + for col in 0..n { + let expected: f32 = (0..k) + .map(|kk| a_values[row * k + kk] * b_values[col * k + kk]) + .sum(); + let actual = slice[[row, col]]; + assert!( + (actual - expected).abs() < 1e-4, + "[{row}, {col}]: got {actual}, expected {expected}" + ); + } + } + }); +} + +#[test] +fn max_contraction_fuses_with_masked_tiles() { + pollster::block_on(async { + let Ok(device) = Device::new().await else { + return; + }; + + // A non-sum reduction over the same contraction cluster: the tiled + // path cannot rely on zero fills collapsing out-of-bounds slots, so + // this exercises the masked accumulate with edge tiles on M and K. + let (m, n, k) = (72usize, 64usize, 20usize); + let a_values = pattern(m * k, 0.31); + let b_values = pattern(n * k, 0.23); + let (a3, b3) = broadcast_factors(&device, m, n, k, &a_values, &b_values); + + let out = (&a3 * &b3).max(2); + assert_eq!(out.count_kernels_to_resolve(), 1, "fused max contraction"); + + let slice = out.as_slice::<2, f32>().await.unwrap(); + for row in [0usize, 39, m - 1] { + for col in [0usize, 40, n - 1] { + let expected = (0..k) + .map(|kk| a_values[row * k + kk] * b_values[col * k + kk]) + .fold(f32::NEG_INFINITY, f32::max); + let actual = slice[[row, col]]; + assert!( + (actual - expected).abs() < 1e-4, + "[{row}, {col}]: got {actual}, expected {expected}" + ); + } + } + }); +} + +#[test] +fn quantized_weighted_reduce_fuses_to_single_kernel() { + pollster::block_on(async { + let Ok(device) = Device::new().await else { + return; + }; + + // `sum_k w_q8[n, k] * x[k]`: the dequantize node feeds the fused + // reduce directly as a block-quantized input — no dense + // materialization kernel, one fused dispatch decoding per element. + const Q8_BLOCK: usize = 32; + let (n, k) = (64usize, 64usize); + let scale = half::f16::from_f32(0.02); + let mut bytes = Vec::new(); + let mut dequantized = Vec::with_capacity(n * k); + for block in 0..(n * k / Q8_BLOCK) { + bytes.extend_from_slice(&scale.to_le_bytes()); + for i in 0..Q8_BLOCK { + let q = (((block * 7 + i * 5) % 64) as i32 - 32) as i8; + bytes.push(q as u8); + dequantized.push(scale.to_f32() * q as f32); + } + } + let w = QMatrix::from_parts( + &device, + &bytes, + vec![n, k].into_boxed_slice(), + GgmlType::Q8_0, + ) + .unwrap(); + let x_values = pattern(k, 0.19); + let x = Tensor::from_slice(&device, [k], &x_values); + + let wd = w.dequantize::(); + let xb = x.reshape([1, k]).broadcast_as([n, k]); + let out = (&wd * &xb).sum(1); + assert_eq!(out.count_kernels_to_resolve(), 1, "fused quantized reduce"); + + let slice = out.as_slice::<1, f32>().await.unwrap(); + for row in [0usize, 17, n - 1] { + let expected: f32 = (0..k) + .map(|kk| dequantized[row * k + kk] * x_values[kk]) + .sum(); + let actual = slice[[row]]; + assert!( + (actual - expected).abs() < 1e-3, + "[{row}]: got {actual}, expected {expected}" + ); + } + }); +} + +#[test] +fn weighted_sum_fuses_to_single_kernel() { + pollster::block_on(async { + let Ok(device) = Device::new().await else { + return; + }; + + // `sum_k w[k] * x[m, k]`: the k-dependent inputs share no (row, col) + // pair. Lowers serially by default; with the cache threshold forced + // down (FUSOR_LAST_LEVEL_CACHE_BYTES) it takes the 1D register + // tiling — correct either way. + let (m, k) = (512usize, 64usize); + let x_values = pattern(m * k, 0.13); + let w_values = pattern(k, 0.07); + let x = Tensor::from_slice(&device, [m, k], &x_values); + let w = Tensor::from_slice(&device, [k], &w_values); + let w2 = w.reshape([1, k]).broadcast_as([m, k]); + + let out = (&x * &w2).sum(1); + assert_eq!(out.count_kernels_to_resolve(), 1, "fused weighted sum"); + + let slice = out.as_slice::<1, f32>().await.unwrap(); + for row in [0usize, 100, 255, m - 1] { + let expected: f32 = (0..k).map(|kk| x_values[row * k + kk] * w_values[kk]).sum(); + let actual = slice[[row]]; + assert!( + (actual - expected).abs() < 1e-4, + "[{row}]: got {actual}, expected {expected}" + ); + } + }); +} + +#[test] +fn broadcast_table_elementwise_reuses_invariant_loads() { + pollster::block_on(async { + let Ok(device) = Device::new().await else { + return; + }; + + // The table is 8 MiB — at the cache-residency threshold — so the + // elementwise lowering tiles the batch dim and hoists the table + // loads out of each thread's run of outputs. + let (b, s, h) = (4usize, 1024usize, 2048usize); + let x_values = pattern(b * s * h, 0.13); + let t_values = pattern(s * h, 0.07); + let x = Tensor::from_slice(&device, [b, s, h], &x_values); + let t = Tensor::from_slice(&device, [s, h], &t_values); + let t3 = t.reshape([1, s, h]).broadcast_as([b, s, h]); + + let out = &x * &t3; + assert_eq!(out.count_kernels_to_resolve(), 1, "broadcast table apply"); + + let slice = out.as_slice::<3, f32>().await.unwrap(); + for batch in 0..b { + for (row, col) in [(0usize, 0usize), (511, 1023), (s - 1, h - 1)] { + let expected = x_values[batch * s * h + row * h + col] * t_values[row * h + col]; + let actual = slice[[batch, row, col]]; + assert!( + (actual - expected).abs() < 1e-5, + "[{batch}, {row}, {col}]: got {actual}, expected {expected}" + ); + } + } + }); +} + +#[test] +fn contraction_with_k_independent_factor_fuses() { + pollster::block_on(async { + let Ok(device) = Device::new().await else { + return; + }; + + // `c [M, N]` joins the product but never varies along K: it loads + // once per output slot instead of being staged per k-tile. + let (m, n, k) = (64usize, 64usize, 32usize); + let a_values = pattern(m * k, 0.19); + let b_values = pattern(n * k, 0.29); + let c_values = pattern(m * n, 0.11); + let (a3, b3) = broadcast_factors(&device, m, n, k, &a_values, &b_values); + let c = Tensor::from_slice(&device, [m, n], &c_values); + let c3 = c.reshape([m, n, 1]).broadcast_as([m, n, k]); + + let out = (&(&a3 * &b3) * &c3).sum(2); + assert_eq!( + out.count_kernels_to_resolve(), + 1, + "contraction with k-independent factor" + ); + + let slice = out.as_slice::<2, f32>().await.unwrap(); + for row in [0usize, 17, m - 1] { + for col in [0usize, 25, n - 1] { + let expected: f32 = (0..k) + .map(|kk| { + a_values[row * k + kk] * b_values[col * k + kk] * c_values[row * n + col] + }) + .sum(); + let actual = slice[[row, col]]; + assert!( + (actual - expected).abs() < 1e-3, + "[{row}, {col}]: got {actual}, expected {expected}" + ); + } + } + }); +} diff --git a/fusor-ml/core/tests/keepdim_repro.rs b/fusor-ml/core/tests/keepdim_repro.rs new file mode 100644 index 000000000..e7f6bad0a --- /dev/null +++ b/fusor-ml/core/tests/keepdim_repro.rs @@ -0,0 +1,29 @@ +//! Regression tests for view folding interacting with reduce fusion: a +//! folded keepdim view reads a lower-rank input pointwise, which must not be +//! mistaken for a shape-preserving unary chain. + +#[test] +fn keepdim_chain_keeps_rank() { + pollster::block_on(async { + let Ok(device) = fusor_core::Device::new().await else { + return; + }; + let t = + fusor_core::Tensor::new::(&device, &[[1.0f32, 2.0, 3.0], [4.0, 5.0, 6.0]]); + + let s = t.sum_keepdim(1); + let slice = s.as_slice::<2, f32>().await.unwrap(); + assert_eq!(slice.shape(), &[2, 1]); + assert_eq!(slice[[0, 0]], 6.0); + assert_eq!(slice[[1, 0]], 15.0); + + // mean-style chain: an elementwise op on top of the keepdim view, so + // the view folds into the nary and the nary must NOT fuse into the + // reduce (that would drop the keepdim rank). + let m = t.sum_keepdim(1) / 3.0; + let slice = m.as_slice::<2, f32>().await.unwrap(); + assert_eq!(slice.shape(), &[2, 1]); + assert_eq!(slice[[0, 0]], 2.0); + assert_eq!(slice[[1, 0]], 5.0); + }); +} diff --git a/fusor-ml/core/tests/recognition.rs b/fusor-ml/core/tests/recognition.rs new file mode 100644 index 000000000..c0cc06d15 --- /dev/null +++ b/fusor-ml/core/tests/recognition.rs @@ -0,0 +1,281 @@ +//! Recognition fidelity gates: composed contraction clusters must resolve to +//! the same kernel counts as the specialized operations they replace. A +//! recognition miss falls back to the generic elementwise + reduce path, +//! which is correct but slow — these tests turn that silent regression into +//! a failure. + +use fusor_core::{Device, QMatrix, Tensor}; +use fusor_gguf::GgmlType; + +fn f32_weight(device: &Device, n: usize, k: usize) -> QMatrix { + let bytes: Vec = (0..n * k) + .map(|i| 0.1 + (i as f32) * 0.05) + .flat_map(f32::to_le_bytes) + .collect(); + QMatrix::from_parts(device, &bytes, vec![n, k].into_boxed_slice(), GgmlType::F32).unwrap() +} + +const Q8_BLOCK: usize = 32; + +/// Patterned Q8_0 blocks: one f16 scale + 32 i8 weights per block, blocks +/// row-major over `[n, k]`. +fn q8_0_weight(device: &Device, n: usize, k: usize) -> (QMatrix, Vec) { + let block_count = n * k / Q8_BLOCK; + let scale = half::f16::from_f32(0.01); + let mut bytes = Vec::new(); + let mut dequantized = Vec::with_capacity(n * k); + for block in 0..block_count { + bytes.extend_from_slice(&scale.to_le_bytes()); + for i in 0..Q8_BLOCK { + let q = (((block * 5 + i * 3) % 64) as i32 - 32) as i8; + bytes.push(q as u8); + dequantized.push(scale.to_f32() * q as f32); + } + } + let matrix = QMatrix::from_parts( + device, + &bytes, + vec![n, k].into_boxed_slice(), + GgmlType::Q8_0, + ) + .unwrap(); + (matrix, dequantized) +} + +#[test] +fn composed_dense_matmul_resolves_to_single_kernel() { + pollster::block_on(async { + let Ok(device) = Device::new().await else { + return; + }; + + let a = Tensor::new::(&device, &[[1.0f32, 2.0], [3.0, 4.0]]); + let b = Tensor::new::(&device, &[[5.0f32, 6.0], [7.0, 8.0]]); + let out = a.mat_mul(&b); + assert_eq!(out.count_kernels_to_resolve(), 1, "dense matmul"); + let slice = out.as_slice::<2, f32>().await.unwrap(); + assert_eq!(slice[[0, 0]], 19.0); + assert_eq!(slice[[0, 1]], 22.0); + assert_eq!(slice[[1, 0]], 43.0); + assert_eq!(slice[[1, 1]], 50.0); + + let a = Tensor::new::( + &device, + &[[[1.0f32, 2.0], [3.0, 4.0]], [[5.0, 6.0], [7.0, 8.0]]], + ); + let b = Tensor::new::( + &device, + &[[[1.0f32, 0.0], [0.0, 1.0]], [[2.0, 0.0], [0.0, 2.0]]], + ); + let out = a.mat_mul(&b); + assert_eq!(out.count_kernels_to_resolve(), 1, "batched matmul"); + let slice = out.as_slice::<3, f32>().await.unwrap(); + assert_eq!(slice[[0, 0, 0]], 1.0); + assert_eq!(slice[[1, 0, 0]], 10.0); + assert_eq!(slice[[1, 1, 1]], 16.0); + }); +} + +#[test] +fn composed_qmatmul_resolves_to_single_kernel() { + pollster::block_on(async { + let Ok(device) = Device::new().await else { + return; + }; + + const K: usize = 8; + let w = f32_weight(&device, 4, K); + let x = Tensor::new::(&device, &[[1.0f32, 0.5, -1.0, 2.0, 0.0, 1.0, -0.5, 3.0]]); + let out = x.q_mat_mul(&w); + assert_eq!(out.count_kernels_to_resolve(), 1, "bare quantized matmul"); + }); +} + +#[test] +fn composed_q8_qmatmul_with_epilogue_resolves_to_single_kernel() { + pollster::block_on(async { + let Ok(device) = Device::new().await else { + return; + }; + + const K: usize = 32; + const N: usize = 4; + let (w, dequantized) = q8_0_weight(&device, N, K); + let x_values: Vec = (0..K).map(|i| ((i as f32) * 0.37).sin()).collect(); + let bias_values = [0.5f32, -1.0, 2.0, -0.25]; + + let x = Tensor::from_slice(&device, [1, K], &x_values); + let bias = Tensor::from_slice(&device, [1, N], &bias_values); + let out = x.q_mat_mul(&w) + &bias; + assert_eq!( + out.count_kernels_to_resolve(), + 1, + "q8_0 matmul + bias epilogue" + ); + + let slice = out.as_slice::<2, f32>().await.unwrap(); + for col in 0..N { + let expected: f32 = (0..K) + .map(|k| x_values[k] * dequantized[col * K + k]) + .sum::() + + bias_values[col]; + assert!( + (slice[[0, col]] - expected).abs() < 1e-3, + "column {col}: got {}, expected {expected}", + slice[[0, col]] + ); + } + }); +} + +#[test] +fn f32_weight_epilogue_takes_correct_fallback() { + pollster::block_on(async { + let Ok(device) = Device::new().await else { + return; + }; + + // F32-native weights lower through the dense matmul kernel, which has + // no epilogue slots: the epilogue must run as its own kernel rather + // than being dropped (regression: the dense shortcut used to discard + // fused epilogues). + const K: usize = 8; + const N: usize = 4; + let w = f32_weight(&device, N, K); + let x_values = [2.0f32, 1.5, -1.0, 2.0, 1.0, 1.0, -0.5, 3.0]; + let bias_values = [0.5f32, -1.0, 2.0, -0.25]; + let x = Tensor::from_slice(&device, [1, K], &x_values); + let bias = Tensor::from_slice(&device, [1, N], &bias_values); + let out = x.q_mat_mul(&w) + &bias; + + let slice = out.as_slice::<2, f32>().await.unwrap(); + for col in 0..N { + let expected: f32 = (0..K) + .map(|k| x_values[k] * (0.1 + ((col * K + k) as f32) * 0.05)) + .sum::() + + bias_values[col]; + assert!( + (slice[[0, col]] - expected).abs() < 1e-3, + "column {col}: got {}, expected {expected}", + slice[[0, col]] + ); + } + }); +} + +#[test] +fn composed_softmax_resolves_to_fused_kernel() { + pollster::block_on(async { + let Ok(device) = Device::new().await else { + return; + }; + let values: Vec = (0..256).map(|i| ((i as f32) * 0.1).sin()).collect(); + let x = Tensor::from_slice(&device, [2, 128], &values); + let out = x.softmax(1); + assert_eq!(out.count_kernels_to_resolve(), 1, "single-pass softmax"); + + let slice = out.as_slice::<2, f32>().await.unwrap(); + for row in 0..2 { + let max = (0..128) + .map(|col| values[row * 128 + col]) + .fold(f32::NEG_INFINITY, f32::max); + let sum: f32 = (0..128) + .map(|col| (values[row * 128 + col] - max).exp()) + .sum(); + for col in [0, 63, 127] { + let expected = (values[row * 128 + col] - max).exp() / sum; + assert!( + (slice[[row, col]] - expected).abs() < 1e-5, + "row {row} col {col}: got {}, expected {expected}", + slice[[row, col]] + ); + } + } + }); +} + +#[test] +fn composed_rms_norm_with_bias_resolves_to_fused_kernel() { + pollster::block_on(async { + let Ok(device) = Device::new().await else { + return; + }; + let x_values = [1.0f32, 2.0, 3.0, 4.0]; + let x = Tensor::from_slice(&device, [1, 4], &x_values); + let weight = Tensor::from_slice(&device, [4], &[0.5f32, 1.0, 1.5, 2.0]); + let bias = Tensor::from_slice(&device, [4], &[0.1f32, -0.2, 0.3, -0.4]); + let out = x.rms_norm_fused(&weight, Some(&bias), 1e-5); + assert_eq!(out.count_kernels_to_resolve(), 1, "rms norm with bias"); + + let slice = out.as_slice::<2, f32>().await.unwrap(); + let mean_square = (1.0 + 4.0 + 9.0 + 16.0) / 4.0; + let rms = f32::sqrt(mean_square + 1e-5); + let weights = [0.5f32, 1.0, 1.5, 2.0]; + let biases = [0.1f32, -0.2, 0.3, -0.4]; + for col in 0..4 { + let expected = x_values[col] / rms * weights[col] + biases[col]; + assert!( + (slice[[0, col]] - expected).abs() < 1e-5, + "col {col}: got {}, expected {expected}", + slice[[0, col]] + ); + } + }); +} + +#[test] +fn composed_attention_resolves_to_flash_kernel() { + pollster::block_on(async { + let Ok(device) = Device::new().await else { + return; + }; + + // Decode shape: q [1, heads, 1, d] against a kv history, GQA 4:1. + let (heads, kv_heads, kv_len, d) = (8usize, 2usize, 64usize, 16usize); + let q_data: Vec = (0..heads * d).map(|i| ((i as f32) * 0.13).sin()).collect(); + let k_data: Vec = (0..kv_heads * kv_len * d) + .map(|i| ((i as f32) * 0.07).cos()) + .collect(); + let v_data: Vec = (0..kv_heads * kv_len * d) + .map(|i| ((i as f32) * 0.11).sin()) + .collect(); + let q = Tensor::from_slice(&device, [1, heads, 1, d], &q_data); + let k = Tensor::from_slice(&device, [1, kv_heads, kv_len, d], &k_data); + let v = Tensor::from_slice(&device, [1, kv_heads, kv_len, d], &v_data); + let scale = 1.0 / (d as f32).sqrt(); + + let out = q.flash_attention(&k, &v, scale, None); + assert_eq!( + out.count_kernels_to_resolve(), + 1, + "gqa decode attention should recognize as one fused flash kernel" + ); + let slice = out.as_slice::<4, f32>().await.unwrap(); + + // Host reference. + for head in [0usize, 5] { + let kv_head = head / (heads / kv_heads); + let scores: Vec = (0..kv_len) + .map(|pos| { + (0..d) + .map(|i| q_data[head * d + i] * k_data[kv_head * kv_len * d + pos * d + i]) + .sum::() + * scale + }) + .collect(); + let max = scores.iter().fold(f32::NEG_INFINITY, |a, &b| a.max(b)); + let weights: Vec = scores.iter().map(|s| (s - max).exp()).collect(); + let total: f32 = weights.iter().sum(); + for dim in [0usize, d - 1] { + let expected: f32 = (0..kv_len) + .map(|pos| weights[pos] / total * v_data[kv_head * kv_len * d + pos * d + dim]) + .sum(); + let actual = slice[[0, head, 0, dim]]; + assert!( + (actual - expected).abs() < 1e-4, + "head {head} dim {dim}: got {actual}, expected {expected}" + ); + } + } + }); +} diff --git a/fusor-ml/fusor/src/composite/conv.rs b/fusor-ml/fusor/src/composite/conv.rs index a9a918695..f6ae09706 100644 --- a/fusor-ml/fusor/src/composite/conv.rs +++ b/fusor-ml/fusor/src/composite/conv.rs @@ -123,22 +123,6 @@ where "Weight in_channels must match input in_channels" ); - if WEIGHT_RANK == 2 + DIFF - && R2 == R + DIFF - && let (Tensor::Gpu(input), Tensor::Gpu(weight)) = (self, weight) - && (bias.is_none() || matches!(bias, Some(Tensor::Gpu(_)))) - { - let bias = bias.and_then(|bias| match bias { - Tensor::Gpu(bias) => Some(bias), - Tensor::Cpu(_) => None, - }); - if let Some(output) = - input.try_conv_nd_direct::(weight, bias, padding, strides) - { - return Tensor::Gpu(output); - } - } - // Step 1: Apply padding to the spatial dimensions (last DIFF dimensions) let padded = if padding.iter().any(|&p| p > 0) { let mut result = self.clone(); diff --git a/fusor-ml/fusor/src/gpu.rs b/fusor-ml/fusor/src/gpu.rs index abce5bffe..6dc51b3ce 100644 --- a/fusor-ml/fusor/src/gpu.rs +++ b/fusor-ml/fusor/src/gpu.rs @@ -269,28 +269,6 @@ impl Tensor { ) } - #[inline] - pub(crate) fn try_conv_nd_direct( - &self, - weight: &Tensor, - bias: Option<&Tensor<1, D>>, - padding: [usize; DIFF], - strides: [usize; DIFF], - ) -> Option { - if R != 2 + DIFF || WEIGHT_RANK != 2 + DIFF { - return None; - } - - self.inner - .try_conv_nd_direct( - weight.as_core(), - bias.map(|bias| bias.as_core()), - &padding, - &strides, - ) - .map(Tensor::from_core) - } - #[inline] pub fn sum(&self, dim: impl Dim) -> Tensor { Tensor::from_core(self.inner.sum(dim.resolve())) diff --git a/fusor-ml/tile-ir-kernels/src/kernels.rs b/fusor-ml/tile-ir-kernels/src/kernels.rs index fc7b31f86..de70d5b4c 100644 --- a/fusor-ml/tile-ir-kernels/src/kernels.rs +++ b/fusor-ml/tile-ir-kernels/src/kernels.rs @@ -1,34 +1,23 @@ use fusor_tile_ir::{Layout, MemoryLevel, Shape}; -mod flash; mod helpers; mod matmul; mod mirostat; -mod qdequantize; mod qgemv; mod qgemv_q4k_ggml; mod qgemv_q6k; mod qmatmul; mod qmatmul_workgroup; mod quantized_matrix; -mod rms_norm; -mod softmax; mod top_k; mod types; -pub use flash::{ - flash_attention, flash_attention_tiled, flash_decode_small, flash_decode_split_partials, - flash_decode_split_reduce, flash_outputs_per_workgroup, flash_tiled_dispatch_size, - flash_tiled_outputs_per_workgroup, FlashAttentionTensors, -}; pub use helpers::AccumCast; pub use matmul::{ - batched_gemv_with_epilogues, batched_matmul_register_with_epilogues, - batched_matmul_with_epilogues, try_batched_coop_matmul, DenseCoopMatmulConfig, - DenseCoopMatmulTile, DenseMatmulShape, DenseMatmulTensors, DenseMatmulTile, + try_batched_coop_matmul, DenseCoopMatmulConfig, DenseCoopMatmulTile, DenseMatmulShape, + DenseMatmulTensors, }; pub use mirostat::{mirostat2, standard_sampler, Mirostat2, StandardSampler}; -pub use qdequantize::qdequantize; pub use qgemv::{qgemv_with_epilogue, IntoQgemvEpilogues}; pub use qmatmul::qmatmul_with_epilogue; pub use qmatmul_workgroup::{ @@ -37,13 +26,8 @@ pub use qmatmul_workgroup::{ qmatmul_workgroup_storage_f16_with_epilogues, qmatmul_workgroup_with_epilogues, }; pub use quantized_matrix::{quantized_matrix, quantized_matrix_for}; -pub use rms_norm::{rms_norm_vec4, RmsNormVec4}; -pub use softmax::{softmax, softmax_partials, softmax_reduce, softmax_write}; pub use top_k::{top_k_chunk, top_k_exactness, top_k_merge}; -pub use types::{ - FlashAttentionDims, FlashAttentionMeta, FlashDecodeSmallMeta, MergeTopKMeta, Mirostat2Meta, - RmsNormVec4Meta, SoftmaxMeta, TensorMeta, TopKChunkMeta, TopKExactnessMeta, -}; +pub use types::{MergeTopKMeta, Mirostat2Meta, TensorMeta, TopKChunkMeta, TopKExactnessMeta}; /// The default rank-1 unit-stride layout used by tile-ir's pre-built kernels /// for tensors whose offset/stride is encoded in the `Meta` struct itself. diff --git a/fusor-ml/tile-ir-kernels/src/kernels/flash.rs b/fusor-ml/tile-ir-kernels/src/kernels/flash.rs deleted file mode 100644 index 73dfd28c4..000000000 --- a/fusor-ml/tile-ir-kernels/src/kernels/flash.rs +++ /dev/null @@ -1,1373 +0,0 @@ -use fusor_tile_ir::{ - tile::{range, Mask, PrivateLocal, Storage, Tile, TileBlock, WorkgroupTile}, - ElementType, ScalarElement, SubgroupToken, TileLiteral, WorkgroupAxis, -}; - -use super::helpers::{index_n, reduce_workgroup, supports_float, u32_tile, zero_fill, NEG_MAX_F32}; -use super::softmax::{softmax_partial_scale, workgroup_softmax_block}; -use super::types::{FlashAttentionDims, FlashAttentionMeta, FlashDecodeSmallMeta}; - -const FLASH_BLOCK: u32 = 256; -const TILED_OUTS_PER_SUBGROUP: u32 = 4; - -/// Runtime tensor bindings consumed by the streaming flash-attention kernels. -pub struct FlashAttentionTensors { - pub q: fusor_tile_ir::KernelTensorRef, - pub k: fusor_tile_ir::KernelTensorRef, - pub v: fusor_tile_ir::KernelTensorRef, - pub mask: Option>, - pub output: fusor_tile_ir::KernelTensorRef, -} - -/// Storage-typed `0.0` literal for the kernel's `element` (F32 or F16). -fn f32_tile(value: f32) -> Tile { - Tile::literal(TileLiteral::f32(value)) -} - -/// Number of output dimensions a single workgroup writes per KV-axis pass, -/// given the device's hardware subgroup size. Each subgroup handles one -/// output dim; the workgroup contains `FLASH_BLOCK / subgroup_size` subgroups. -pub const fn flash_outputs_per_workgroup(subgroup_size: u32) -> u32 { - FLASH_BLOCK / subgroup_size -} - -/// Number of output dimensions the tiled prefill kernel writes per workgroup. -/// -/// It keeps the physical workgroup at `FLASH_BLOCK` lanes, but each subgroup -/// serially accumulates a small run of output dimensions while reusing the -/// same QK score. -pub const fn flash_tiled_outputs_per_workgroup(subgroup_size: u32) -> u32 { - flash_outputs_per_workgroup(subgroup_size) * TILED_OUTS_PER_SUBGROUP -} - -/// Build a streaming flash-attention kernel for F32 or F16 tensors. -/// -/// `element` is the runtime storage element (F32 or F16); the accumulator math -/// is always F32. -/// -/// `subgroup_size` must match the runtime hardware subgroup width on the -/// target device — the kernel layout assigns one subgroup per output dim and -/// uses `subgroup_reduce_*` to fold a `subgroup_size`-wide chunk of KV. Pick -/// the size by reading the device's subgroup caps; pinning is not exposed by -/// wgpu, so this must come from a fixed `(min == max)` adapter range. -/// -/// The metadata supplies tensor strides, offsets, dimensions, scale, and the -/// dispatch grid. Returns `None` if the tensor ranks or optional mask binding -/// are inconsistent with the metadata. -pub fn flash_attention( - kb: &mut fusor_tile_ir::KernelBuilder, - element: ElementType, - tensors: FlashAttentionTensors, - meta: FlashAttentionMeta, - subgroup: SubgroupToken, - subgroup_size: u32, -) -> Option<()> { - let FlashAttentionTensors { - q, - k, - v, - mask, - output, - } = tensors; - if !supports_float(element) { - return None; - } - if subgroup_size == 0 || !FLASH_BLOCK.is_multiple_of(subgroup_size) { - return None; - } - let outputs_per_workgroup = flash_outputs_per_workgroup(subgroup_size); - let q_strides: [u32; 4] = meta.q_meta.strides.as_slice().try_into().ok()?; - let k_strides: [u32; 4] = meta.k_meta.strides.as_slice().try_into().ok()?; - let v_strides: [u32; 4] = meta.v_meta.strides.as_slice().try_into().ok()?; - let output_strides: [u32; 4] = meta.output_meta.strides.as_slice().try_into().ok()?; - let mask_strides: Option<[u32; 2]> = if let Some(mask) = meta.mask_meta.as_ref() { - Some(mask.strides.as_slice().try_into().ok()?) - } else { - None - }; - if meta.dims.batch == 0 - || meta.dims.num_heads == 0 - || meta.dims.num_kv_heads == 0 - || meta.dims.q_seq_len == 0 - || meta.dims.kv_seq_len == 0 - || meta.dims.head_dim == 0 - { - return None; - } - if meta.mask_meta.is_some() != mask.is_some() { - return None; - } - if meta.causal && meta.mask_meta.is_some() { - // Causal mode is mutually exclusive with an explicit additive mask: - // the kernel skips out-of-bound kv positions and emits no mask add. - return None; - } - let groups = meta.dims.num_heads.checked_div(meta.dims.num_kv_heads)?; - if groups == 0 { - return None; - } - let elem_fill = zero_fill(element); - let q = kb.read(element, q); - let k = kb.read(element, k); - let v = kb.read(element, v); - let mask = mask.map(|m| kb.read(element, m)); - let output = kb.write(element, output); - let phase = kb.program(); - { - phase.program_grid(FLASH_BLOCK, meta.dispatch_size, |program| { - let lane = program.lane(); - let workgroup_x = program.program_id(WorkgroupAxis::X); - let row = program.program_id(WorkgroupAxis::Y); - let q_idx = program.bind(row.clone() % meta.dims.q_seq_len); - let row_over_q = row.clone() / meta.dims.q_seq_len; - let head_idx = program.bind(row_over_q.clone() % meta.dims.num_heads); - let batch_idx = program.bind(row / (meta.dims.q_seq_len * meta.dims.num_heads)); - let kv_head_idx = program.bind(head_idx.clone() / u32_tile(groups)); - let kv_lane = lane.clone() % subgroup_size; - let out_dim = - program.bind(workgroup_x * outputs_per_workgroup + (lane.clone() / subgroup_size)); - let out_valid = program.bind(out_dim.clone().lt(u32_tile(meta.dims.head_dim))); - // Per-iteration scratch locals — used to bridge values across - // `if_then` branches inside the body. Not loop-carried. - let score_local = program.private(ElementType::F32); - let weighted_local = program.private(ElementType::F32); - - let kv_chunks = meta.dims.kv_seq_len.div_ceil(subgroup_size); - let causal = meta.causal; - let [_final_m, final_s, final_o] = program.fold( - range(kv_chunks), - [f32_tile(NEG_MAX_F32), f32_tile(0.0), f32_tile(0.0)], - |program, chunk_idx, [m_state, s_state, o_state]| { - let chunk = chunk_idx; - let kv_idx = - program.bind(chunk.clone() * u32_tile(subgroup_size) + kv_lane.clone()); - let bound_valid = - program.bind(kv_idx.clone().lt(u32_tile(meta.dims.kv_seq_len))); - // For causal attention we additionally restrict to kv <= q. - // We can also skip the per-dim Q·K work for an entire - // chunk when `chunk_start > q_idx`: in that case every - // lane's kv_idx > q_idx and the chunk's contribution is - // zero. We still need to fold (so the post-loop - // accumulator state is correct), but we gate the heavy - // load+dot under `chunk_in_range`. - let (kv_valid, chunk_in_range) = if causal { - let chunk_start = program.bind(chunk.clone() * u32_tile(subgroup_size)); - let chunk_in_range = program.bind(chunk_start.le(q_idx.clone())); - let kv_le_q = kv_idx.clone().le(q_idx.clone()); - let kv_valid = program.bind(bound_valid.clone().and(kv_le_q)); - (kv_valid, Some(chunk_in_range)) - } else { - (bound_valid.clone(), None) - }; - program.store_local(&score_local, f32_tile(NEG_MAX_F32)); - let compute_guard = match &chunk_in_range { - Some(in_range) => kv_valid.clone().and(in_range.clone()), - None => kv_valid.clone(), - }; - program.if_then(compute_guard, |program| { - let mut products = Vec::with_capacity(meta.dims.head_dim as usize); - for dim in 0..meta.dims.head_dim { - let q_index = index_n( - meta.q_meta.offset, - q_strides, - (batch_idx.clone(), head_idx.clone(), q_idx.clone(), dim), - ); - let k_index = index_n( - meta.k_meta.offset, - k_strides, - (batch_idx.clone(), kv_head_idx.clone(), kv_idx.clone(), dim), - ); - let q_value = program - .load(q.at(q_index), Mask::all(), elem_fill) - .cast(ElementType::F32); - let k_value = program - .load(k.at(k_index), Mask::all(), elem_fill) - .cast(ElementType::F32); - products.push(q_value * k_value); - } - let mut score = - program.sum(products, ElementType::F32) * f32_tile(meta.scale.get()); - if let (Some(mask), Some(mask_meta), Some(mask_strides)) = - (&mask, meta.mask_meta.as_ref(), mask_strides) - { - let mask_index = index_n( - mask_meta.offset, - mask_strides, - (q_idx.clone(), kv_idx.clone()), - ); - let mask_value = program - .load(mask.at(mask_index), Mask::all(), elem_fill) - .cast(ElementType::F32); - score = score + mask_value; - } - program.store_local(&score_local, score); - }); - - let score = program.bind(program.load_local(&score_local)); - let block_max = - program.bind(subgroup.subgroup_reduce_max(program, score.clone())); - let old_m = program.bind(m_state); - let new_m = program.bind(old_m.clone().max(block_max.clone())); - let raw_exp = (score.clone() - new_m.clone()).exp(); - let exp_score = - program.bind(Tile::select(kv_valid.clone(), raw_exp, f32_tile(0.0))); - let block_sum = - program.bind(subgroup.subgroup_reduce_sum(program, exp_score.clone())); - - program.store_local(&weighted_local, f32_tile(0.0)); - let valid_value = kv_valid.clone().and(out_valid.clone()); - program.if_then(valid_value, |program| { - let v_index = index_n( - meta.v_meta.offset, - v_strides, - ( - batch_idx.clone(), - kv_head_idx.clone(), - kv_idx.clone(), - out_dim.clone(), - ), - ); - let v_value = program - .load(v.at(v_index), Mask::all(), elem_fill) - .cast(ElementType::F32); - program.store_local(&weighted_local, exp_score.clone() * v_value); - }); - let weighted = program.load_local(&weighted_local); - let block_out = program.bind(subgroup.subgroup_reduce_sum(program, weighted)); - - let old_m_scale = program.bind((old_m.clone() - new_m.clone()).exp()); - let new_s = s_state * old_m_scale.clone() + block_sum; - let new_o = o_state * old_m_scale + block_out; - [new_m, new_s, new_o] - }, - ); - - let store_valid = kv_lane.eq(u32_tile(0)).and(out_valid.clone()); - // Bind the fold results so the divide and the `if_then` body share - // the same SSA values rather than re-emitting the loop materialize. - let final_o_bound = program.bind(final_o); - let final_s_bound = program.bind(final_s); - // Evaluate `final_m` once so the loop fires even though we don't - // need its value in the post-loop stage. - program.if_then(store_valid, |program| { - let output_value = (final_o_bound.clone() / final_s_bound.clone()).cast(element); - let output_index = index_n( - meta.output_meta.offset, - output_strides, - ( - batch_idx.clone(), - head_idx.clone(), - q_idx.clone(), - out_dim.clone(), - ), - ); - program.store(output.at(output_index), output_value, Mask::all()); - }); - }); - } - Some(()) -} - -/// Q-batched streaming flash attention. Same online-softmax algorithm as -/// [`flash_attention`], but each workgroup handles `q_block` contiguous q_idx -/// values, caching Q, K, and V slices in workgroup memory so the big -/// `(kv_seq_len * head_dim)` K traffic and the per-chunk Q reads are reused -/// across `q_block` queries instead of being re-fetched per query. -/// -/// Layout: -/// - One workgroup per (batch, head, q_block). -/// - `FLASH_BLOCK` lanes split into subgroups of `subgroup_size` lanes each. -/// Each subgroup is pinned to one kv position per lane and accumulates a -/// short run of output dims serially, reusing the same QK score. -/// - Once per workgroup: load `Q[q_block, head_dim]` into workgroup memory. -/// - Per chunk: cooperatively load `K[chunk_kv_positions, head_dim]` and -/// `V[chunk_kv_positions, workgroup_out_dims]` into workgroup memory once, -/// then loop over `q_block` queries reusing those loads. -pub fn flash_attention_tiled( - kb: &mut fusor_tile_ir::KernelBuilder, - element: ElementType, - tensors: FlashAttentionTensors, - meta: FlashAttentionMeta, - subgroup: SubgroupToken, - subgroup_size: u32, - q_block: u32, -) -> Option<()> { - let FlashAttentionTensors { - q, - k, - v, - mask, - output, - } = tensors; - if !supports_float(element) { - return None; - } - if subgroup_size == 0 || !FLASH_BLOCK.is_multiple_of(subgroup_size) || q_block == 0 { - return None; - } - let outputs_per_workgroup = flash_tiled_outputs_per_workgroup(subgroup_size); - let q_strides: [u32; 4] = meta.q_meta.strides.as_slice().try_into().ok()?; - let k_strides: [u32; 4] = meta.k_meta.strides.as_slice().try_into().ok()?; - let v_strides: [u32; 4] = meta.v_meta.strides.as_slice().try_into().ok()?; - let output_strides: [u32; 4] = meta.output_meta.strides.as_slice().try_into().ok()?; - let mask_strides: Option<[u32; 2]> = if let Some(mask) = meta.mask_meta.as_ref() { - Some(mask.strides.as_slice().try_into().ok()?) - } else { - None - }; - if meta.dims.batch == 0 - || meta.dims.num_heads == 0 - || meta.dims.num_kv_heads == 0 - || meta.dims.q_seq_len == 0 - || meta.dims.kv_seq_len == 0 - || meta.dims.head_dim == 0 - { - return None; - } - if meta.mask_meta.is_some() != mask.is_some() { - return None; - } - if meta.causal && meta.mask_meta.is_some() { - return None; - } - let groups = meta.dims.num_heads.checked_div(meta.dims.num_kv_heads)?; - if groups == 0 { - return None; - } - // Workgroup-memory budgets — keep totals well under the 32KB Apple cap. - // K cache: subgroup_size rows × head_dim cols. - let k_cache_elems = subgroup_size.checked_mul(meta.dims.head_dim)?; - if k_cache_elems > 4096 { - return None; - } - // V cache: subgroup_size rows × outputs_per_workgroup cols. - let v_cache_elems = subgroup_size.checked_mul(outputs_per_workgroup)?; - if v_cache_elems > 4096 { - return None; - } - // Q cache: q_block rows × head_dim cols. - let q_cache_elems = q_block.checked_mul(meta.dims.head_dim)?; - if q_cache_elems > 4096 { - return None; - } - let elem_fill = zero_fill(element); - let q = kb.read(element, q); - let k = kb.read(element, k); - let v = kb.read(element, v); - let mask = mask.map(|m| kb.read(element, m)); - let output = kb.write(element, output); - let phase = kb.program(); - { - // Allocate workgroup caches once. - let k_cache = phase.alloc_workgroup_array(ScalarElement::F32, k_cache_elems); - let v_cache = phase.alloc_workgroup_array(ScalarElement::F32, v_cache_elems); - let q_cache = phase.alloc_workgroup_array(ScalarElement::F32, q_cache_elems); - let q_blocks = meta.dims.q_seq_len.div_ceil(q_block); - let dispatch_size = [ - meta.dims.head_dim.div_ceil(outputs_per_workgroup), - meta.dims - .batch - .checked_mul(meta.dims.num_heads)? - .checked_mul(q_blocks)?, - 1, - ]; - - phase.program_grid(FLASH_BLOCK, dispatch_size, |program| { - let lane = program.lane(); - let workgroup_x = program.program_id(WorkgroupAxis::X); - let row = program.program_id(WorkgroupAxis::Y); - let q_block_idx = program.bind(row.clone() % q_blocks); - let row_over_q = row.clone() / q_blocks; - let head_idx = program.bind(row_over_q.clone() % meta.dims.num_heads); - let batch_idx = program.bind(row_over_q / meta.dims.num_heads); - let kv_head_idx = program.bind(head_idx.clone() / u32_tile(groups)); - let subgroup_idx = program.bind(lane.clone() / subgroup_size); - let kv_lane = program.bind(lane.clone() % subgroup_size); - let out_dim_base = program.bind( - workgroup_x * outputs_per_workgroup - + subgroup_idx.clone() * TILED_OUTS_PER_SUBGROUP, - ); - let mut out_dims = Vec::with_capacity(TILED_OUTS_PER_SUBGROUP as usize); - let mut out_valids = Vec::with_capacity(TILED_OUTS_PER_SUBGROUP as usize); - for out_offset in 0..TILED_OUTS_PER_SUBGROUP { - let out_dim = program.bind(out_dim_base.clone() + u32_tile(out_offset)); - let out_valid = program.bind(out_dim.clone().lt(u32_tile(meta.dims.head_dim))); - out_dims.push(out_dim); - out_valids.push(out_valid); - } - - // Per-query accumulator state, stored in private locals indexed by - // q-offset within the workgroup's Q-block. Initialised below. - let mut m_locals: Vec = Vec::with_capacity(q_block as usize); - let mut s_locals: Vec = Vec::with_capacity(q_block as usize); - let mut o_locals: Vec = - Vec::with_capacity((q_block * TILED_OUTS_PER_SUBGROUP) as usize); - for _ in 0..q_block { - m_locals.push(program.private(ElementType::F32)); - s_locals.push(program.private(ElementType::F32)); - for _ in 0..TILED_OUTS_PER_SUBGROUP { - o_locals.push(program.private(ElementType::F32)); - } - } - for q_offset in 0..q_block { - program.store_local(&m_locals[q_offset as usize], f32_tile(NEG_MAX_F32)); - program.store_local(&s_locals[q_offset as usize], f32_tile(0.0)); - for out_offset in 0..TILED_OUTS_PER_SUBGROUP { - let o_idx = (q_offset * TILED_OUTS_PER_SUBGROUP + out_offset) as usize; - program.store_local(&o_locals[o_idx], f32_tile(0.0)); - } - } - - let kv_chunks = meta.dims.kv_seq_len.div_ceil(subgroup_size); - let causal = meta.causal; - - // Q-block base index in the q dimension. - let q_block_base = program.bind(q_block_idx.clone() * u32_tile(q_block)); - - // ----------------------------------------------------------- - // One-shot Q cache load. - // - // Layout: q_cache[q_offset * head_dim + dim]. - // Loaded once per workgroup, indexed by lane in strided passes - // of FLASH_BLOCK lanes each. Total = q_block * head_dim ≤ 4096. - // - // The original streaming kernel re-loaded Q from global per - // (kv_chunk, q_idx, dim, lane) — 256× duplicate reads per - // (q_idx, dim) within a workgroup. Caching Q removes that. - // ----------------------------------------------------------- - let total_q_loads = q_block * meta.dims.head_dim; - let q_passes = total_q_loads.div_ceil(FLASH_BLOCK); - for pass in 0..q_passes { - let pass_base = pass * FLASH_BLOCK; - let idx = program.bind(lane.clone() + u32_tile(pass_base)); - let q_offset_local = program.bind(idx.clone() / meta.dims.head_dim); - let dim_local = program.bind(idx.clone() % meta.dims.head_dim); - let q_pos = program.bind(q_block_base.clone() + q_offset_local.clone()); - let in_bounds = idx - .clone() - .lt(u32_tile(total_q_loads)) - .and(q_pos.clone().lt(u32_tile(meta.dims.q_seq_len))); - let q_index = index_n( - meta.q_meta.offset, - q_strides, - ( - batch_idx.clone(), - head_idx.clone(), - q_pos.clone(), - dim_local.clone(), - ), - ); - let q_val = program - .load(q.at(q_index), in_bounds.clone(), elem_fill) - .cast(ElementType::F32); - let store_in_bounds = idx.clone().lt(u32_tile(total_q_loads)); - program.if_then(store_in_bounds, |program| { - program.store_workgroup(&q_cache, idx.clone(), q_val); - }); - } - program.workgroup_barrier(); - - // Per-iteration scratch local used to bridge values across if_then. - let score_local = program.private(ElementType::F32); - - // Runtime chunk counter; collapses the kv-chunk dimension into a - // single IR loop rather than unrolling `kv_chunks` copies of the - // body (which blew up code size for big prefills). - let chunk_local = program.private(ElementType::U32); - program.store_local(&chunk_local, u32_tile(0)); - program.loop_forever(|program| { - let chunk_tile = program.bind(program.load_local(&chunk_local)); - program.break_if(chunk_tile.clone().ge(u32_tile(kv_chunks))); - let chunk_start = program.bind(chunk_tile.clone() * u32_tile(subgroup_size)); - if causal { - let q_block_last = program.bind( - (q_block_base.clone() + u32_tile(q_block.saturating_sub(1))) - .min(u32_tile(meta.dims.q_seq_len - 1)), - ); - program.break_if(chunk_start.clone().gt(q_block_last)); - } - let kv_idx = program.bind(chunk_start.clone() + kv_lane.clone()); - let bound_valid = program.bind(kv_idx.clone().lt(u32_tile(meta.dims.kv_seq_len))); - - // ------------------------------------------------------- - // Cooperative load of K[chunk kv rows, all head_dim] into - // workgroup memory. Indexed by `kv_local * head_dim + dim`. - // Strided across FLASH_BLOCK lanes. The k_cache_elems check - // above guarantees total ≤ 4096, and we use mask-load + an - // if_then-guarded store to handle the last partial pass - // (only fires when total_k_loads % FLASH_BLOCK != 0). - // ------------------------------------------------------- - let total_k_loads = subgroup_size * meta.dims.head_dim; - let k_passes = total_k_loads.div_ceil(FLASH_BLOCK); - let k_aligned = total_k_loads.is_multiple_of(FLASH_BLOCK); - for pass in 0..k_passes { - let pass_base = pass * FLASH_BLOCK; - let idx = program.bind(lane.clone() + u32_tile(pass_base)); - let kv_local = program.bind(idx.clone() / meta.dims.head_dim); - let dim_local = program.bind(idx.clone() % meta.dims.head_dim); - let kv_pos = program.bind(chunk_start.clone() + kv_local.clone()); - let in_bounds = if k_aligned { - kv_pos.clone().lt(u32_tile(meta.dims.kv_seq_len)) - } else { - idx.clone() - .lt(u32_tile(total_k_loads)) - .and(kv_pos.clone().lt(u32_tile(meta.dims.kv_seq_len))) - }; - let k_index = index_n( - meta.k_meta.offset, - k_strides, - ( - batch_idx.clone(), - kv_head_idx.clone(), - kv_pos.clone(), - dim_local.clone(), - ), - ); - let k_val = program - .load(k.at(k_index), in_bounds.clone(), elem_fill) - .cast(ElementType::F32); - if k_aligned { - program.store_workgroup(&k_cache, idx.clone(), k_val); - } else { - let store_in_bounds = idx.clone().lt(u32_tile(total_k_loads)); - program.if_then(store_in_bounds, |program| { - program.store_workgroup(&k_cache, idx.clone(), k_val.clone()); - }); - } - } - - // Cooperative load of V[chunk kv rows, out_dims handled by - // this workgroup] into workgroup memory. Each subgroup loads - // a short run of output dims for the same kv lane. - for out_offset in 0..TILED_OUTS_PER_SUBGROUP { - let out_idx = out_offset as usize; - let v_cell_idx = program.bind( - kv_lane.clone() * outputs_per_workgroup - + subgroup_idx.clone() * TILED_OUTS_PER_SUBGROUP - + u32_tile(out_offset), - ); - let v_in_bounds = bound_valid.clone().and(out_valids[out_idx].clone()); - let v_index = index_n( - meta.v_meta.offset, - v_strides, - ( - batch_idx.clone(), - kv_head_idx.clone(), - kv_idx.clone(), - out_dims[out_idx].clone(), - ), - ); - let v_val = program - .load(v.at(v_index), v_in_bounds.clone(), elem_fill) - .cast(ElementType::F32); - program.store_workgroup(&v_cache, v_cell_idx.clone(), v_val); - } - - program.workgroup_barrier(); - - // ------------------------------------------------------- - // For each query in the Q block, compute score, update the - // online-softmax stats, accumulate weighted V. Q and K both - // come from workgroup memory; only the optional mask still - // comes from global. - // ------------------------------------------------------- - for q_offset in 0..q_block { - let q_idx = program.bind(q_block_base.clone() + u32_tile(q_offset)); - let q_in_range = program.bind(q_idx.clone().lt(u32_tile(meta.dims.q_seq_len))); - - let (kv_valid, chunk_in_range) = if causal { - let chunk_in_range = program.bind(chunk_start.clone().le(q_idx.clone())); - let kv_le_q = kv_idx.clone().le(q_idx.clone()); - let kv_valid = - program.bind(bound_valid.clone().and(kv_le_q).and(q_in_range.clone())); - (kv_valid, Some(chunk_in_range)) - } else { - let kv_valid = program.bind(bound_valid.clone().and(q_in_range.clone())); - (kv_valid, None) - }; - - program.store_local(&score_local, f32_tile(NEG_MAX_F32)); - let compute_guard = match &chunk_in_range { - Some(in_range) => kv_valid.clone().and(in_range.clone()), - None => kv_valid.clone(), - }; - program.if_then(compute_guard, |program| { - let mut products = Vec::with_capacity(meta.dims.head_dim as usize); - for dim in 0..meta.dims.head_dim { - // Q from workgroup cache: q_cache[q_offset*head_dim + dim]. - let q_cache_idx = u32_tile(q_offset * meta.dims.head_dim + dim); - let q_value = program.load_workgroup(&q_cache, q_cache_idx); - // K from workgroup cache: k_cache[kv_lane*head_dim + dim]. - let k_cache_idx = kv_lane.clone() * meta.dims.head_dim + u32_tile(dim); - let k_value = program.load_workgroup(&k_cache, k_cache_idx); - products.push(q_value * k_value); - } - let mut score = - program.sum(products, ElementType::F32) * f32_tile(meta.scale.get()); - if let (Some(mask), Some(mask_meta), Some(mask_strides)) = - (&mask, meta.mask_meta.as_ref(), mask_strides) - { - let mask_index = index_n( - mask_meta.offset, - mask_strides, - (q_idx.clone(), kv_idx.clone()), - ); - let mask_value = program - .load(mask.at(mask_index), Mask::all(), elem_fill) - .cast(ElementType::F32); - score = score + mask_value; - } - program.store_local(&score_local, score); - }); - - let score = program.bind(program.load_local(&score_local)); - let block_max = - program.bind(subgroup.subgroup_reduce_max(program, score.clone())); - let old_m = program.bind(program.load_local(&m_locals[q_offset as usize])); - let new_m = program.bind(old_m.clone().max(block_max.clone())); - let raw_exp = (score.clone() - new_m.clone()).exp(); - let exp_score = - program.bind(Tile::select(kv_valid.clone(), raw_exp, f32_tile(0.0))); - let block_sum = - program.bind(subgroup.subgroup_reduce_sum(program, exp_score.clone())); - - let old_m_scale = program.bind((old_m.clone() - new_m.clone()).exp()); - let old_s = program.load_local(&s_locals[q_offset as usize]); - let new_s = old_s * old_m_scale.clone() + block_sum; - for out_offset in 0..TILED_OUTS_PER_SUBGROUP { - let out_idx = out_offset as usize; - let v_cache_idx = kv_lane.clone() * outputs_per_workgroup - + subgroup_idx.clone() * TILED_OUTS_PER_SUBGROUP - + u32_tile(out_offset); - let v_cached = program.bind(program.load_workgroup(&v_cache, v_cache_idx)); - let weighted = Tile::select( - kv_valid.clone().and(out_valids[out_idx].clone()), - exp_score.clone() * v_cached, - f32_tile(0.0), - ); - let block_out = - program.bind(subgroup.subgroup_reduce_sum(program, weighted)); - let o_idx = (q_offset * TILED_OUTS_PER_SUBGROUP + out_offset) as usize; - let old_o = program.load_local(&o_locals[o_idx]); - let new_o = old_o * old_m_scale.clone() + block_out; - program.store_local(&o_locals[o_idx], new_o); - } - program.store_local(&m_locals[q_offset as usize], new_m); - program.store_local(&s_locals[q_offset as usize], new_s); - } - - // Barrier before next chunk's K/V load overwrites the cache. - program.workgroup_barrier(); - - program.store_local(&chunk_local, chunk_tile + u32_tile(1)); - }); - - // Write each query's accumulated output. - for q_offset in 0..q_block { - let q_idx = program.bind(q_block_base.clone() + u32_tile(q_offset)); - let q_in_range = program.bind(q_idx.clone().lt(u32_tile(meta.dims.q_seq_len))); - let final_s = program.bind(program.load_local(&s_locals[q_offset as usize])); - for out_offset in 0..TILED_OUTS_PER_SUBGROUP { - let out_idx = out_offset as usize; - let store_valid = kv_lane - .clone() - .eq(u32_tile(0)) - .and(out_valids[out_idx].clone()) - .and(q_in_range.clone()); - let o_idx = (q_offset * TILED_OUTS_PER_SUBGROUP + out_offset) as usize; - let final_o = program.bind(program.load_local(&o_locals[o_idx])); - program.if_then(store_valid, |program| { - let output_value = (final_o.clone() / final_s.clone()).cast(element); - let output_index = index_n( - meta.output_meta.offset, - output_strides, - ( - batch_idx.clone(), - head_idx.clone(), - q_idx.clone(), - out_dims[out_idx].clone(), - ), - ); - program.store(output.at(output_index), output_value, Mask::all()); - }); - } - } - }); - } - Some(()) -} - -/// Dispatch grid for the tiled (Q-batched) flash-attention kernel. -pub fn flash_tiled_dispatch_size( - dims: FlashAttentionDims, - outputs_per_workgroup: u32, - q_block: u32, -) -> [u32; 3] { - [ - dims.head_dim.div_ceil(outputs_per_workgroup), - dims.batch - .checked_mul(dims.num_heads) - .and_then(|value| value.checked_mul(dims.q_seq_len.div_ceil(q_block))) - .expect("flash attention tiled dispatch overflow"), - 1, - ] -} - -struct DecodeScoreForKv<'a> { - q: &'a Storage, - k: &'a Storage, - meta: FlashDecodeSmallMeta, - batch_idx: Tile, - head_idx: Tile, - kv_head_idx: Tile, - kv: Tile, - score_acc: &'a PrivateLocal, - dim_local: &'a PrivateLocal, -} - -fn decode_score_for_kv(program: &mut TileBlock<'_>, request: DecodeScoreForKv<'_>) -> Tile { - let DecodeScoreForKv { - q, - k, - meta, - batch_idx, - head_idx, - kv_head_idx, - kv, - score_acc, - dim_local, - } = request; - program.store_local(score_acc, f32_tile(0.0)); - program.store_local(dim_local, u32_tile(0)); - program.loop_forever(|program| { - let dim = program.load_local(dim_local); - program.break_if(dim.clone().ge(u32_tile(meta.dims.head_dim))); - let q_index = index_n( - meta.q_offset, - meta.q_strides, - (batch_idx.clone(), head_idx.clone(), 0, dim.clone()), - ); - let k_index = index_n( - meta.k_offset, - meta.k_strides, - ( - batch_idx.clone(), - kv_head_idx.clone(), - kv.clone(), - dim.clone(), - ), - ); - let q_value = program.load(q.at(q_index), Mask::all(), TileLiteral::f32(0.0)); - let k_value = program.load(k.at(k_index), Mask::all(), TileLiteral::f32(0.0)); - let acc = program.load_local(score_acc); - program.store_local(score_acc, acc + q_value * k_value); - program.store_local(dim_local, dim + u32_tile(1)); - }); - program.load_local(score_acc) * f32_tile(meta.scale.get()) -} - -struct DecodeOutputLoop<'a> { - v: &'a Storage, - output: &'a Storage, - probs: &'a WorkgroupTile, - meta: FlashDecodeSmallMeta, - batch_idx: Tile, - head_idx: Tile, - kv_head_idx: Tile, - out_dim: Tile, - active_kv_len: Tile, - acc: &'a PrivateLocal, - kv_local: &'a PrivateLocal, -} - -struct FlashDecodeSmallBlock { - q: fusor_tile_ir::KernelTensorRef, - k: fusor_tile_ir::KernelTensorRef, - v: fusor_tile_ir::KernelTensorRef, - output: fusor_tile_ir::KernelTensorRef, - params: fusor_tile_ir::KernelTensorRef, - meta: FlashDecodeSmallMeta, -} - -struct FlashDecodeSplitPartialsBlock { - q: fusor_tile_ir::KernelTensorRef, - k: fusor_tile_ir::KernelTensorRef, - v: fusor_tile_ir::KernelTensorRef, - scratch: fusor_tile_ir::KernelTensorRef, - params: fusor_tile_ir::KernelTensorRef, - meta: FlashDecodeSmallMeta, -} - -fn append_decode_output_loop(program: &mut TileBlock<'_>, request: DecodeOutputLoop<'_>) { - let DecodeOutputLoop { - v, - output, - probs, - meta, - batch_idx, - head_idx, - kv_head_idx, - out_dim, - active_kv_len, - acc, - kv_local, - } = request; - program.loop_forever(|program| { - let kv = program.load_local(kv_local); - program.break_if(kv.clone().ge(active_kv_len.clone())); - let prob = program.load_workgroup(probs, kv.clone()); - let v_index = index_n( - meta.v_offset, - meta.v_strides, - ( - batch_idx.clone(), - kv_head_idx.clone(), - kv.clone(), - out_dim.clone(), - ), - ); - let v_value = program.load(v.at(v_index), Mask::all(), TileLiteral::f32(0.0)); - let current = program.load_local(acc); - program.store_local(acc, current + prob * v_value); - program.store_local(kv_local, kv + u32_tile(1)); - }); - - let output_value = program.load_local(acc); - let output_index = index_n( - meta.output_offset, - meta.output_strides, - (batch_idx, head_idx, 0, out_dim), - ); - program.store(output.at(output_index), output_value, Mask::all()); -} - -fn flash_decode_small_block( - kb: &mut fusor_tile_ir::KernelBuilder, - block: u32, - args: FlashDecodeSmallBlock, -) { - let FlashDecodeSmallBlock { - q, - k, - v, - output, - params, - meta, - } = args; - let q = kb.read(ElementType::F32, q); - let k = kb.read(ElementType::F32, k); - let v = kb.read(ElementType::F32, v); - let output = kb.write(ElementType::F32, output); - let params = kb.read(ElementType::U32, params); - let phase = kb.program(); - let probs = phase.alloc_workgroup_array(ScalarElement::F32, block); - let reduce = phase.alloc_workgroup_array(ScalarElement::F32, block); - - phase.program_grid( - block, - [meta.dims.batch * meta.dims.num_heads, 1, 1], - |program| { - let lane = program.lane(); - let row = program.program_id(WorkgroupAxis::X); - let active_kv_len = program.load( - params.at(0), - Mask::all(), - TileLiteral::U32(meta.active_kv_len), - ); - let head_idx = row.clone() % meta.dims.num_heads; - let batch_idx = row / meta.dims.num_heads; - let kv_head_idx = head_idx.clone() / u32_tile(meta.groups); - let lane_value = lane.clone(); - let acc = program.private(ElementType::F32); - let kv_local = program.private(ElementType::U32); - let item = program.private(ElementType::U32); - let dim = program.private(ElementType::U32); - let score_acc = program.private(ElementType::F32); - let score_local = program.private(ElementType::F32); - let max_score_local = program.private(ElementType::F32); - - if meta.tiled { - program.store_workgroup(&reduce, lane.clone(), f32_tile(NEG_MAX_F32)); - program.store_local(&kv_local, u32_tile(0)); - program.loop_forever(|program| { - let tile_base = program.load_local(&kv_local); - program.break_if(tile_base.clone().ge(active_kv_len.clone())); - let kv = tile_base.clone() + lane_value.clone(); - let kv_valid = kv.clone().lt(active_kv_len.clone()); - program.if_then(kv_valid, |program| { - let score = decode_score_for_kv( - program, - DecodeScoreForKv { - q: &q, - k: &k, - meta, - batch_idx: batch_idx.clone(), - head_idx: head_idx.clone(), - kv_head_idx: kv_head_idx.clone(), - kv: kv.clone(), - score_acc: &score_acc, - dim_local: &dim, - }, - ); - let current = program.load_workgroup(&reduce, lane.clone()); - program.store_workgroup(&reduce, lane.clone(), current.max(score)); - }); - program.store_local(&kv_local, tile_base + u32_tile(block)); - }); - program.workgroup_barrier(); - reduce_workgroup(program, &reduce, lane.clone(), |lhs, rhs| lhs.max(rhs)); - let max_score = program.load_workgroup(&reduce, 0u32); - program.store_local(&max_score_local, max_score); - let max_score = program.load_local(&max_score_local); - - // All lanes load `reduce[0]` (the max) above. Before any lane - // overwrites `reduce[lane]` for the denominator accumulator we - // need a barrier — without it, lane 0's store to slot 0 races - // with other lanes still loading slot 0, which surfaces as - // intermittent wrong values for individual heads. - program.workgroup_barrier(); - program.store_workgroup(&reduce, lane.clone(), f32_tile(0.0)); - program.store_local(&kv_local, u32_tile(0)); - program.loop_forever(|program| { - let tile_base = program.load_local(&kv_local); - program.break_if(tile_base.clone().ge(active_kv_len.clone())); - let kv = tile_base.clone() + lane_value.clone(); - let kv_valid = kv.clone().lt(active_kv_len.clone()); - program.if_then(kv_valid, |program| { - let score = decode_score_for_kv( - program, - DecodeScoreForKv { - q: &q, - k: &k, - meta, - batch_idx: batch_idx.clone(), - head_idx: head_idx.clone(), - kv_head_idx: kv_head_idx.clone(), - kv: kv.clone(), - score_acc: &score_acc, - dim_local: &dim, - }, - ); - let prob = (score - max_score.clone()).exp(); - let current = program.load_workgroup(&reduce, lane.clone()); - program.store_workgroup(&reduce, lane.clone(), current + prob); - }); - program.store_local(&kv_local, tile_base + u32_tile(block)); - }); - program.workgroup_barrier(); - reduce_workgroup(program, &reduce, lane.clone(), |lhs, rhs| lhs + rhs); - let denom = program.load_workgroup(&reduce, 0u32); - - program.store_local(&acc, f32_tile(0.0)); - program.store_local(&kv_local, u32_tile(0)); - program.loop_forever(|program| { - let tile_base = program.load_local(&kv_local); - program.break_if(tile_base.clone().ge(active_kv_len.clone())); - let kv = tile_base.clone() + lane_value.clone(); - let kv_valid = kv.clone().lt(active_kv_len.clone()); - program.if_else( - kv_valid, - |program| { - let score = decode_score_for_kv( - program, - DecodeScoreForKv { - q: &q, - k: &k, - meta, - batch_idx: batch_idx.clone(), - head_idx: head_idx.clone(), - kv_head_idx: kv_head_idx.clone(), - kv: kv.clone(), - score_acc: &score_acc, - dim_local: &dim, - }, - ); - let prob = (score - max_score.clone()).exp() / denom.clone(); - program.store_workgroup(&probs, lane.clone(), prob); - }, - |program| { - program.store_workgroup(&probs, lane.clone(), f32_tile(0.0)); - }, - ); - program.workgroup_barrier(); - program.store_local(&item, u32_tile(0)); - let out_condition = lane_value.clone().lt(u32_tile(meta.dims.head_dim)); - program.if_then(out_condition, |program| { - program.loop_forever(|program| { - let item_value = program.load_local(&item); - let block_done = item_value.clone().ge(u32_tile(block)); - let kv = tile_base.clone() + item_value.clone(); - let kv_done = kv.clone().ge(active_kv_len.clone()); - program.break_if(block_done.or(kv_done)); - let prob = program.load_workgroup(&probs, item_value.clone()); - let v_index = index_n( - meta.v_offset, - meta.v_strides, - ( - batch_idx.clone(), - kv_head_idx.clone(), - kv, - lane_value.clone(), - ), - ); - let v_value = - program.load(v.at(v_index), Mask::all(), TileLiteral::f32(0.0)); - let current = program.load_local(&acc); - program.store_local(&acc, current + prob * v_value); - program.store_local(&item, item_value + u32_tile(1)); - }); - }); - program.workgroup_barrier(); - program.store_local(&kv_local, tile_base + u32_tile(block)); - }); - let out_condition = lane_value.clone().lt(u32_tile(meta.dims.head_dim)); - program.if_then(out_condition, |program| { - let output_value = program.load_local(&acc); - let output_index = index_n( - meta.output_offset, - meta.output_strides, - (batch_idx.clone(), head_idx.clone(), 0, lane_value.clone()), - ); - program.store(output.at(output_index), output_value, Mask::all()); - }); - return; - } - - let kv_valid = lane_value.clone().lt(active_kv_len.clone()); - program.store_local(&score_local, f32_tile(NEG_MAX_F32)); - program.if_then(kv_valid.clone(), |program| { - let score = decode_score_for_kv( - program, - DecodeScoreForKv { - q: &q, - k: &k, - meta, - batch_idx: batch_idx.clone(), - head_idx: head_idx.clone(), - kv_head_idx: kv_head_idx.clone(), - kv: lane_value.clone(), - score_acc: &score_acc, - dim_local: &dim, - }, - ); - program.store_local(&score_local, score); - }); - let stats = workgroup_softmax_block( - program, - lane.clone(), - program.load_local(&score_local), - kv_valid.clone(), - &reduce, - Some(&probs), - ); - program.if_then(kv_valid, |program| { - program.store_workgroup( - &probs, - lane.clone(), - stats.prob.clone() / stats.denom.clone(), - ); - }); - program.workgroup_barrier(); - - let out_condition = lane_value.clone().lt(u32_tile(meta.dims.head_dim)); - program.if_then(out_condition, |program| { - program.store_local(&acc, f32_tile(0.0)); - program.store_local(&kv_local, u32_tile(0)); - append_decode_output_loop( - program, - DecodeOutputLoop { - v: &v, - output: &output, - probs: &probs, - meta, - batch_idx: batch_idx.clone(), - head_idx: head_idx.clone(), - kv_head_idx: kv_head_idx.clone(), - out_dim: lane_value.clone(), - active_kv_len: active_kv_len.clone(), - acc: &acc, - kv_local: &kv_local, - }, - ); - }); - }, - ); -} - -fn flash_decode_split_partials_block( - kb: &mut fusor_tile_ir::KernelBuilder, - block: u32, - args: FlashDecodeSplitPartialsBlock, -) { - let FlashDecodeSplitPartialsBlock { - q, - k, - v, - scratch, - params, - meta, - } = args; - let q = kb.read(ElementType::F32, q); - let k = kb.read(ElementType::F32, k); - let v = kb.read(ElementType::F32, v); - let scratch = kb.write(ElementType::F32, scratch); - let params = kb.read(ElementType::U32, params); - let phase = kb.program(); - let probs = phase.alloc_workgroup_array(ScalarElement::F32, block); - let reduce = phase.alloc_workgroup_array(ScalarElement::F32, block); - - phase.program_grid( - block, - [ - meta.dims.batch * meta.dims.num_heads * meta.split_blocks, - 1, - 1, - ], - |program| { - let lane = program.lane(); - let tile = program.program_id(WorkgroupAxis::X); - let rows = u32_tile(meta.dims.batch * meta.dims.num_heads); - let row = program.bind(tile.clone() % rows.clone()); - let split_block = program.bind(tile / rows); - let active_kv_len = program.load( - params.at(0), - Mask::all(), - TileLiteral::U32(meta.active_kv_len), - ); - let head_idx = row.clone() % meta.dims.num_heads; - let batch_idx = row.clone() / meta.dims.num_heads; - let kv_head_idx = head_idx.clone() / u32_tile(meta.groups); - let lane_value = lane.clone(); - let tile_base = program.bind(split_block.clone() * u32_tile(meta.decode_block)); - let tile_end = program - .bind((tile_base.clone() + u32_tile(meta.decode_block)).min(active_kv_len.clone())); - let kv = program.bind(tile_base.clone() + lane_value.clone()); - let kv_valid = kv.clone().lt(tile_end.clone()); - let item = program.private(ElementType::U32); - let acc = program.private(ElementType::F32); - let dim = program.private(ElementType::U32); - let score_acc = program.private(ElementType::F32); - let score_local = program.private(ElementType::F32); - - program.store_local(&score_local, f32_tile(NEG_MAX_F32)); - program.if_then(kv_valid.clone(), |program| { - let score = decode_score_for_kv( - program, - DecodeScoreForKv { - q: &q, - k: &k, - meta, - batch_idx: batch_idx.clone(), - head_idx: head_idx.clone(), - kv_head_idx: kv_head_idx.clone(), - kv: kv.clone(), - score_acc: &score_acc, - dim_local: &dim, - }, - ); - program.store_local(&score_local, score); - }); - let stats = workgroup_softmax_block( - program, - lane.clone(), - program.load_local(&score_local), - kv_valid.clone(), - &reduce, - Some(&probs), - ); - - let scratch_base = program.bind( - (row.clone() * u32_tile(meta.split_blocks) + split_block.clone()) - * u32_tile(meta.dims.head_dim + 2), - ); - program.if_then(lane_value.clone().eq(u32_tile(0)), |program| { - program.store( - scratch.at(scratch_base.clone() + meta.dims.head_dim), - stats.denom.clone(), - Mask::all(), - ); - program.store( - scratch.at(scratch_base.clone() + meta.dims.head_dim + 1), - stats.max.clone(), - Mask::all(), - ); - }); - - let out_condition = lane_value.clone().lt(u32_tile(meta.dims.head_dim)); - program.if_then(out_condition, |program| { - program.store_local(&acc, f32_tile(0.0)); - program.store_local(&item, u32_tile(0)); - program.loop_forever(|program| { - let item_value = program.load_local(&item); - let block_done = item_value.clone().ge(u32_tile(block)); - let kv = tile_base.clone() + item_value.clone(); - let kv_done = kv.clone().ge(tile_end.clone()); - program.break_if(block_done.or(kv_done)); - let prob = program.load_workgroup(&probs, item_value.clone()); - let v_index = index_n( - meta.v_offset, - meta.v_strides, - ( - batch_idx.clone(), - kv_head_idx.clone(), - kv, - lane_value.clone(), - ), - ); - let v_value = program.load(v.at(v_index), Mask::all(), TileLiteral::f32(0.0)); - let current = program.load_local(&acc); - program.store_local(&acc, current + prob * v_value); - program.store_local(&item, item_value + u32_tile(1)); - }); - program.store( - scratch.at(scratch_base.clone() + lane_value.clone()), - program.load_local(&acc), - Mask::all(), - ); - }); - }, - ); -} - -pub fn flash_decode_split_partials( - kb: &mut fusor_tile_ir::KernelBuilder, - q: fusor_tile_ir::KernelTensorRef, - k: fusor_tile_ir::KernelTensorRef, - v: fusor_tile_ir::KernelTensorRef, - scratch: fusor_tile_ir::KernelTensorRef, - params: fusor_tile_ir::KernelTensorRef, - meta: FlashDecodeSmallMeta, -) -> Option<()> { - if meta.dims.head_dim == 0 - || meta.dims.head_dim > meta.decode_block - || meta.decode_block == 0 - || meta.groups == 0 - || meta.split_blocks < 2 - { - return None; - } - if !matches!(meta.decode_block, 128 | 256 | 512 | 1024) { - return None; - } - flash_decode_split_partials_block( - kb, - meta.decode_block, - FlashDecodeSplitPartialsBlock { - q, - k, - v, - scratch, - params, - meta, - }, - ); - Some(()) -} - -pub fn flash_decode_split_reduce( - kb: &mut fusor_tile_ir::KernelBuilder, - scratch: fusor_tile_ir::KernelTensorRef, - output: fusor_tile_ir::KernelTensorRef, - meta: FlashDecodeSmallMeta, -) -> Option<()> { - if meta.dims.head_dim == 0 || meta.groups == 0 || meta.split_blocks < 2 { - return None; - } - let output_strides: [u32; 4] = meta.output_strides; - let scratch = kb.read(ElementType::F32, scratch); - let output = kb.write(ElementType::F32, output); - let phase = kb.program(); - phase.program_grid( - meta.dims.head_dim, - [meta.dims.batch * meta.dims.num_heads, 1, 1], - |program| { - let lane = program.lane(); - let row = program.program_id(WorkgroupAxis::X); - let out_dim = lane.clone(); - let head_idx = row.clone() % meta.dims.num_heads; - let batch_idx = row.clone() / meta.dims.num_heads; - let scratch_stride = meta.dims.head_dim + 2; - let row_base = program.bind(row.clone() * u32_tile(meta.split_blocks * scratch_stride)); - let mut max_score = f32_tile(NEG_MAX_F32); - for split_block in 0..meta.split_blocks { - let block_base = row_base.clone() + split_block * scratch_stride; - let block_max = program.load( - scratch.at(block_base + meta.dims.head_dim + 1), - Mask::all(), - TileLiteral::f32(NEG_MAX_F32), - ); - max_score = max_score.max(block_max); - } - let max_score = program.bind(max_score); - - let mut denom = f32_tile(0.0); - let mut acc = f32_tile(0.0); - for split_block in 0..meta.split_blocks { - let block_base = row_base.clone() + split_block * scratch_stride; - let block_denom = program.load( - scratch.at(block_base.clone() + meta.dims.head_dim), - Mask::all(), - TileLiteral::f32(0.0), - ); - let block_max = program.load( - scratch.at(block_base.clone() + meta.dims.head_dim + 1), - Mask::all(), - TileLiteral::f32(NEG_MAX_F32), - ); - let scale = softmax_partial_scale(block_max, max_score.clone()); - denom = denom + block_denom * scale.clone(); - let partial = program.load( - scratch.at(block_base + out_dim.clone()), - Mask::all(), - TileLiteral::f32(0.0), - ); - acc = acc + partial * scale; - } - let output_index = index_n( - meta.output_offset, - output_strides, - (batch_idx, head_idx, 0, out_dim), - ); - program.store(output.at(output_index), acc / denom, Mask::all()); - }, - ); - Some(()) -} - -/// Build the small F32 decode-attention kernel. -/// -/// Supports head dimensions no larger than the decode block size and the decode block sizes accepted by -/// [`FlashDecodeSmallMeta::decode_block`](crate::FlashDecodeSmallMeta::decode_block). -pub fn flash_decode_small( - kb: &mut fusor_tile_ir::KernelBuilder, - q: fusor_tile_ir::KernelTensorRef, - k: fusor_tile_ir::KernelTensorRef, - v: fusor_tile_ir::KernelTensorRef, - output: fusor_tile_ir::KernelTensorRef, - params: fusor_tile_ir::KernelTensorRef, - meta: FlashDecodeSmallMeta, -) -> Option<()> { - if meta.dims.head_dim == 0 - || meta.dims.head_dim > meta.decode_block - || meta.decode_block == 0 - || meta.groups == 0 - { - return None; - } - if !matches!(meta.decode_block, 128 | 256 | 512 | 1024) { - return None; - } - flash_decode_small_block( - kb, - meta.decode_block, - FlashDecodeSmallBlock { - q, - k, - v, - output, - params, - meta, - }, - ); - Some(()) -} diff --git a/fusor-ml/tile-ir-kernels/src/kernels/helpers.rs b/fusor-ml/tile-ir-kernels/src/kernels/helpers.rs index b5553da33..5763afe97 100644 --- a/fusor-ml/tile-ir-kernels/src/kernels/helpers.rs +++ b/fusor-ml/tile-ir-kernels/src/kernels/helpers.rs @@ -499,23 +499,3 @@ pub(super) fn scalar_of(element: ElementType) -> ScalarElement { ElementType::CoopMatrix { scalar, .. } => scalar, } } - -/// Zero literal for a float `element` (F32 or F16). Panics on other types — -/// callers (flash / softmax) only ever stage F32 and F16. -pub(super) fn zero_fill(element: ElementType) -> TileLiteral { - match scalar_of(element) { - ScalarElement::F32 => TileLiteral::f32(0.0), - ScalarElement::F16 => TileLiteral::F16(0), - _ => panic!("only F32 and F16 element types are supported"), - } -} - -/// Whether `element`'s scalar is a float (F32 or F16). -pub(super) fn supports_float(element: ElementType) -> bool { - matches!(scalar_of(element), ScalarElement::F32 | ScalarElement::F16) -} - -/// A `u32` literal tile. -pub(super) fn u32_tile(value: u32) -> Tile { - Tile::literal(TileLiteral::U32(value)) -} diff --git a/fusor-ml/tile-ir-kernels/src/kernels/matmul.rs b/fusor-ml/tile-ir-kernels/src/kernels/matmul.rs index ef3b33202..4f749ed69 100644 --- a/fusor-ml/tile-ir-kernels/src/kernels/matmul.rs +++ b/fusor-ml/tile-ir-kernels/src/kernels/matmul.rs @@ -1,19 +1,15 @@ //! Dense matrix multiply program kernels. -use fusor_tile_ir::tile::{range, CoopAcc, Program, Storage, Tile, TileBlock}; +use fusor_tile_ir::tile::{CoopAcc, Program, Storage, Tile, TileBlock}; use fusor_tile_ir::{CoopMatrixToken, ScalarElement, SubgroupToken, TileLiteral, WorkgroupAxis}; use crate::{ dispatch::SubgroupConfig, - grid::dot4_sum, kernels::helpers::{ coop_load_a_fragments, coop_load_b_fragments, coop_mma_grid, coop_store_acc_grid, - dispatch_grid_1d, scalar_of, zero_coop_acc_grid, AccumCast, - }, - types::{ - apply_optional_epilogue, cooperative_store_layout_supported, matrix_shape, - DenseMatmulEpilogues, + dispatch_grid_1d, scalar_of, zero_coop_acc_grid, }, + types::{cooperative_store_layout_supported, DenseMatmulEpilogues}, }; /// Logical shape for flattened batched dense matmul views. @@ -29,40 +25,6 @@ pub struct DenseMatmulShape { pub n: u32, } -/// Workgroup tile geometry for the direct dense matmul kernel. -#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] -pub struct DenseMatmulTile { - pub bm: u32, - pub bn: u32, - pub bk: u32, - pub tm: u32, - pub tn: u32, - pub lanes: u32, -} - -impl DenseMatmulTile { - pub const fn new(bm: u32, bn: u32, bk: u32, tm: u32, tn: u32, lanes: u32) -> Self { - Self { - bm, - bn, - bk, - tm, - tn, - lanes, - } - } - - pub fn validate(self) { - assert!(self.bm >= self.tm && self.bm.is_multiple_of(self.tm)); - assert!(self.bn >= self.tn && self.bn.is_multiple_of(self.tn)); - assert_eq!( - self.lanes, - (self.bm / self.tm) * (self.bn / self.tn), - "dense matmul lanes must cover one thread tile per lane" - ); - } -} - /// Direct storage bindings for dense matrix multiplication kernels. /// /// The storage element travels in each [`Storage`] view, so this bundle is not @@ -105,399 +67,6 @@ impl CoopTileEntry { } } -/// The accumulator element for every dense matmul kernel is F32. The storage -/// element (F32 or F16) travels in the [`Storage`] view; the runtime -/// [`AccumCast`] inserts the F16↔F32 cast pair on load/store and is the -/// identity for F32 storage. -fn accum_cast(storage: ScalarElement) -> AccumCast { - AccumCast::new(storage, ScalarElement::F32) -} - -/// Batched dense GEMV over flattened direct views: -/// A is `[batch * m, k]`, B is `[batch * k, 1]`, Y is `[batch * m, 1]`. -/// -/// The storage element (F32 or F16) is recovered at runtime from the bound -/// [`Storage`] views; accumulation is in F32 via the [`AccumCast`], which -/// inserts the F16→F32 cast on load and F32→F16 cast on store. F32 storage has -/// identity casts. -/// -/// Each subgroup computes one output row. Lanes cooperatively walk K in -/// `VALUES_PER_LANE` chunks and then reduce the partial sums inside the -/// subgroup, avoiding the scalar-lane behavior of the generic edge matmul. -#[allow(clippy::too_many_arguments)] -pub fn batched_gemv_with_epilogues( - program: &mut Program, - a: &Storage, - b: &Storage, - y: &Storage, - shape: DenseMatmulShape, - epilogues: &DenseMatmulEpilogues<'_>, - max_workgroups_per_dimension: u32, - subgroups: SubgroupConfig, -) { - // Runtime subgroup width × rows per workgroup = logical row coverage. - // Each lane folds VALUES_PER_LANE elements of K via dot4. - const ROWS_PER_WORKGROUP: u32 = 4; - const VALUES_PER_LANE: u32 = 8; - let block = subgroups.block_for_subgroups(ROWS_PER_WORKGROUP); - let subgroup = subgroups.token(); - let rows_per_workgroup = ROWS_PER_WORKGROUP; - let values_per_lane = VALUES_PER_LANE; - assert_eq!(shape.n, 1, "batched_gemv expects a single RHS column"); - - let cast = accum_cast(scalar_of(a.element())); - - let [a_rows, a_k] = matrix_shape(a.layout()); - let [b_rows, b_n] = matrix_shape(b.layout()); - let [y_rows, y_n] = matrix_shape(y.layout()); - assert_eq!(shape.batch * shape.m, a_rows); - assert_eq!(shape.k, a_k); - assert_eq!(shape.batch * shape.k, b_rows); - assert_eq!(1, b_n); - assert_eq!(shape.batch * shape.m, y_rows); - assert_eq!(1, y_n); - - let row_groups = shape.m.div_ceil(rows_per_workgroup); - let total_groups = shape.batch * row_groups; - let grid = dispatch_grid_1d(total_groups, max_workgroups_per_dimension); - - program.program_grid(block, grid, |program| { - let group_id = program.program_id(WorkgroupAxis::X) - + program.program_id(WorkgroupAxis::Y) * grid[0] - + program.program_id(WorkgroupAxis::Z) * grid[0] * grid[1]; - let group_active = group_id.clone().lt(total_groups); - let batch_tile = group_id.clone() / row_groups; - let row_group = group_id % row_groups; - let row = row_group * rows_per_workgroup + subgroup.subgroup_id(program); - let lane = subgroup.subgroup_lane(program); - let k_per_iter = subgroup.subgroup_size(program) * values_per_lane; - let k_iterations = (Tile::literal(TileLiteral::U32(shape.k)) + k_per_iter.clone() - 1u32) - / k_per_iter.clone(); - let row_in_bounds = group_active.clone().and(row.clone().lt(shape.m)); - let a_batch_base = batch_tile.clone() * shape.m; - let b_batch_base = batch_tile.clone() * shape.k; - let y_batch_base = batch_tile * shape.m; - - let [sum] = program.fold( - range(k_iterations), - [Tile::literal(TileLiteral::f32(0.0))], - |program, loop_index, [acc]| { - let k_base = loop_index * k_per_iter.clone() + lane.clone() * values_per_lane; - let a_values: Vec = (0..values_per_lane) - .map(|i| { - let k_index = k_base.clone() + i; - let mask = row_in_bounds.clone().and(k_index.clone().lt(shape.k)); - let loaded = program.load( - a.at((a_batch_base.clone() + row.clone(), k_index)), - mask.clone(), - cast.zero_storage(), - ); - Tile::select( - mask, - apply_optional_epilogue(epilogues.pre_a, cast.into_accum(loaded)), - Tile::literal(TileLiteral::f32(0.0)), - ) - }) - .collect(); - let b_values: Vec = (0..values_per_lane) - .map(|i| { - let k_index = k_base.clone() + i; - let mask = group_active.clone().and(k_index.clone().lt(shape.k)); - let loaded = program.load( - b.at((b_batch_base.clone() + k_index, 0)), - mask.clone(), - cast.zero_storage(), - ); - Tile::select( - mask, - apply_optional_epilogue(epilogues.pre_b, cast.into_accum(loaded)), - Tile::literal(TileLiteral::f32(0.0)), - ) - }) - .collect(); - [acc + dot4_sum(program, &a_values, &b_values)] - }, - ); - let reduced = subgroup.subgroup_reduce_sum(program, sum); - let value = cast.from_accum(apply_optional_epilogue(epilogues.post, reduced)); - let mask = lane.eq(0).and(row_in_bounds); - program.store(y.at((y_batch_base + row, 0)), value, mask); - }); -} - -/// Batched dense matmul over flattened direct views. The storage element -/// (F32 or F16) is recovered at runtime from the bound [`Storage`] views; -/// accumulation is in F32 via the [`AccumCast`]. -/// A is `[batch * m, k]`, B is `[batch * k, n]`, Y is `[batch * m, n]`. -pub fn batched_matmul_with_epilogues( - program: &mut Program, - tensors: DenseMatmulTensors<'_>, - shape: DenseMatmulShape, - epilogues: &DenseMatmulEpilogues<'_>, - max_workgroups_per_dimension: u32, - tile: DenseMatmulTile, -) { - let DenseMatmulTensors { a, b, y } = tensors; - tile.validate(); - let DenseMatmulTile { - bm, - bn, - bk, - tm, - tn, - lanes, - } = tile; - let outs = (tm * tn) as usize; - - let scalar = scalar_of(a.element()); - let cast = accum_cast(scalar); - - let [a_rows, a_k] = matrix_shape(a.layout()); - let [b_rows, b_n] = matrix_shape(b.layout()); - let [y_rows, y_n] = matrix_shape(y.layout()); - assert_eq!(shape.batch * shape.m, a_rows); - assert_eq!(shape.k, a_k); - assert_eq!(shape.batch * shape.k, b_rows); - assert_eq!(shape.n, b_n); - assert_eq!(shape.batch * shape.m, y_rows); - assert_eq!(shape.n, y_n); - - let tiles_m = shape.m.div_ceil(bm); - let tiles_n = shape.n.div_ceil(bn); - let total_tiles = shape.batch * tiles_m * tiles_n; - let k_tiles = shape.k.div_ceil(bk); - let grid = dispatch_grid_1d(total_tiles, max_workgroups_per_dimension); - let a_tile = program.alloc_workgroup_tile(scalar, bm, bk); - let b_tile = program.alloc_workgroup_tile(scalar, bk, bn); - - program.program_grid(lanes, grid, |program| { - let tile_id = program.program_id(WorkgroupAxis::X) - + program.program_id(WorkgroupAxis::Y) * grid[0] - + program.program_id(WorkgroupAxis::Z) * grid[0] * grid[1]; - let tile_active = tile_id.clone().lt(total_tiles); - let batch_tile = tile_id.clone() / (tiles_m * tiles_n); - let local_tile = tile_id % (tiles_m * tiles_n); - let m_tile = local_tile.clone() / tiles_n; - let n_tile = local_tile % tiles_n; - - let lane = program.lane(); - let lane_row = lane.clone() / (bn / tn); - let lane_col = lane % (bn / tn); - let m_tile_base = m_tile * bm; - let n_tile_base = n_tile * bn; - let row_base = m_tile_base.clone() + lane_row.clone() * tm; - let col_base = n_tile_base.clone() + lane_col.clone() * tn; - let a_batch_base = batch_tile.clone() * shape.m; - let b_batch_base = batch_tile.clone() * shape.k; - let y_batch_base = batch_tile * shape.m; - - let sums = program.fold_vec( - range(k_tiles), - (0..outs) - .map(|_| Tile::literal(TileLiteral::f32(0.0))) - .collect(), - |program, k_tile, accs| { - let k_base = k_tile * bk; - for pass in 0..(bm * bk).div_ceil(lanes) { - let flat = program.lane() + pass * lanes; - let local_row = flat.clone() / bk; - let local_k = flat.clone() % bk; - let global_row = m_tile_base.clone() + local_row.clone(); - let global_k = k_base.clone() + local_k.clone(); - let in_bounds = tile_active - .clone() - .and(flat.clone().lt(bm * bk)) - .and(global_row.clone().lt(shape.m)) - .and(global_k.clone().lt(shape.k)); - let loaded = program.load( - a.at((a_batch_base.clone() + global_row, &global_k)), - in_bounds.clone(), - cast.zero_storage(), - ); - let value = cast.from_accum(Tile::select( - in_bounds, - apply_optional_epilogue(epilogues.pre_a, cast.into_accum(loaded)), - Tile::literal(TileLiteral::f32(0.0)), - )); - program.store_workgroup(&a_tile, flat, value); - } - for pass in 0..(bk * bn).div_ceil(lanes) { - let flat = program.lane() + pass * lanes; - let local_k = flat.clone() / bn; - let local_col = flat.clone() % bn; - let global_k = k_base.clone() + local_k.clone(); - let global_col = n_tile_base.clone() + local_col.clone(); - let in_bounds = tile_active - .clone() - .and(flat.clone().lt(bk * bn)) - .and(global_k.clone().lt(shape.k)) - .and(global_col.clone().lt(shape.n)); - let loaded = program.load( - b.at((b_batch_base.clone() + global_k, global_col)), - in_bounds.clone(), - cast.zero_storage(), - ); - let value = cast.from_accum(Tile::select( - in_bounds, - apply_optional_epilogue(epilogues.pre_b, cast.into_accum(loaded)), - Tile::literal(TileLiteral::f32(0.0)), - )); - program.store_workgroup(&b_tile, flat, value); - } - program.workgroup_barrier(); - - // Each chunk starts from a fresh `0.0` base, is bound to a - // local, and is added to the carried accumulator after the - // per-chunk dot is complete. - let chunk_sums: Vec<_> = (0..outs as u32) - .map(|idx| { - let r = idx / tn; - let c = idx % tn; - let local_row = lane_row.clone() * tm + r; - let local_col = lane_col.clone() * tn + c; - let mut sum = Tile::literal(TileLiteral::f32(0.0)); - for kk in 0..bk { - let a_value = cast.into_accum( - program.load_workgroup(&a_tile, local_row.clone() * bk + kk), - ); - let b_value = cast.into_accum( - program.load_workgroup(&b_tile, local_col.clone() + kk * bn), - ); - sum = sum + a_value * b_value; - } - sum - }) - .collect(); - let chunk_sums: Vec<_> = chunk_sums - .into_iter() - .map(|sum| program.bind(sum)) - .collect(); - program.workgroup_barrier(); - accs.into_iter() - .zip(chunk_sums) - .map(|(acc, chunk)| acc + chunk) - .collect() - }, - ); - - for (idx, sum) in sums.into_iter().enumerate() { - let idx = idx as u32; - let r = idx / tn; - let c = idx % tn; - let row = row_base.clone() + r; - let col = col_base.clone() + c; - let value = cast.from_accum(apply_optional_epilogue(epilogues.post, sum)); - let mask = tile_active - .clone() - .and(row.clone().lt(shape.m)) - .and(col.clone().lt(shape.n)); - program.store(y.at((y_batch_base.clone() + row, col)), value, mask); - } - }); -} - -/// Batched dense matmul fallback for partial tiles. This keeps the 4x4 -/// register tile but reads directly from storage so skinny/edge shapes avoid -/// workgroup-tile corner cases. The storage element (F32 or F16) is recovered -/// at runtime from the bound [`Storage`] views with F32 accumulation. -pub fn batched_matmul_register_with_epilogues( - program: &mut Program, - a: &Storage, - b: &Storage, - y: &Storage, - shape: DenseMatmulShape, - epilogues: &DenseMatmulEpilogues<'_>, - max_workgroups_per_dimension: u32, -) { - // BM/BN are pinned to the register tile geometry (4x4 lanes × 8x8 = 32x32). - const BM: u32 = 32; - const BN: u32 = 32; - const TM: u32 = 4; - const TN: u32 = 4; - const OUTS: usize = (TM * TN) as usize; - const LANES: u32 = 64; - - let cast = accum_cast(scalar_of(a.element())); - - let tiles_m = shape.m.div_ceil(BM); - let tiles_n = shape.n.div_ceil(BN); - let total_tiles = shape.batch * tiles_m * tiles_n; - let grid = dispatch_grid_1d(total_tiles, max_workgroups_per_dimension); - - program.program_grid(LANES, grid, |program| { - let tile_id = program.program_id(WorkgroupAxis::X) - + program.program_id(WorkgroupAxis::Y) * grid[0] - + program.program_id(WorkgroupAxis::Z) * grid[0] * grid[1]; - let tile_active = tile_id.clone().lt(total_tiles); - let batch_tile = tile_id.clone() / (tiles_m * tiles_n); - let local_tile = tile_id % (tiles_m * tiles_n); - let m_tile = local_tile.clone() / tiles_n; - let n_tile = local_tile % tiles_n; - - let lane = program.lane(); - let lane_row = lane.clone() / (BN / TN); - let lane_col = lane % (BN / TN); - let row_base = m_tile * BM + lane_row * TM; - let col_base = n_tile * BN + lane_col * TN; - let a_batch_base = batch_tile.clone() * shape.m; - let b_batch_base = batch_tile.clone() * shape.k; - let y_batch_base = batch_tile * shape.m; - - let sums: [Tile; OUTS] = program.fold( - range(shape.k), - std::array::from_fn(|_| Tile::literal(TileLiteral::f32(0.0))), - |program, k_index, accs| { - let a_values: [Tile; TM as usize] = std::array::from_fn(|r| { - let row = row_base.clone() + r as u32; - let in_bounds = tile_active.clone().and(row.clone().lt(shape.m)); - let loaded = program.load( - a.at((a_batch_base.clone() + row, &k_index)), - in_bounds.clone(), - cast.zero_storage(), - ); - Tile::select( - in_bounds, - apply_optional_epilogue(epilogues.pre_a, cast.into_accum(loaded)), - Tile::literal(TileLiteral::f32(0.0)), - ) - }); - let b_values: [Tile; TN as usize] = std::array::from_fn(|c| { - let col = col_base.clone() + c as u32; - let in_bounds = tile_active.clone().and(col.clone().lt(shape.n)); - let loaded = program.load( - b.at((b_batch_base.clone() + k_index.clone(), col)), - in_bounds.clone(), - cast.zero_storage(), - ); - Tile::select( - in_bounds, - apply_optional_epilogue(epilogues.pre_b, cast.into_accum(loaded)), - Tile::literal(TileLiteral::f32(0.0)), - ) - }); - std::array::from_fn(|idx| { - let r = idx / TN as usize; - let c = idx % TN as usize; - accs[idx].clone() + a_values[r].clone() * b_values[c].clone() - }) - }, - ); - - for (idx, sum) in sums.into_iter().enumerate() { - let r = idx / TN as usize; - let c = idx % TN as usize; - let row = row_base.clone() + r as u32; - let col = col_base.clone() + c as u32; - let value = cast.from_accum(apply_optional_epilogue(epilogues.post, sum)); - let mask = tile_active - .clone() - .and(row.clone().lt(shape.m)) - .and(col.clone().lt(shape.n)); - program.store(y.at((y_batch_base.clone() + row, col)), value, mask); - } - }); -} - /// Try to emit a fast cooperative-matrix batched matmul. Returns false /// when shape/layout/epilogues require the generic path. The storage element /// travels in the bound [`Storage`] views, so both F32 and F16 use the same diff --git a/fusor-ml/tile-ir-kernels/src/kernels/qdequantize.rs b/fusor-ml/tile-ir-kernels/src/kernels/qdequantize.rs deleted file mode 100644 index 6597617d8..000000000 --- a/fusor-ml/tile-ir-kernels/src/kernels/qdequantize.rs +++ /dev/null @@ -1,60 +0,0 @@ -//! Quantized dequantization program kernels. - -use fusor_tile_ir::tile::{Program, Storage}; -use fusor_tile_ir::{ - ElementType, Layout, MemoryLevel, QuantizedMatrix, Shape, StorageView, WorkgroupAxis, -}; - -/// Lane-per-element dequantization. -/// -/// Emits one dense f32/f16 element per quantized element of `b` and writes it -/// to a row-major `y` of `b.rows * b.cols` elements. -pub fn qdequantize(program: &mut Program, b: &QuantizedMatrix, y: &Storage, workgroups_x: u32) { - const BLOCK: u32 = 256; - assert!( - workgroups_x > 0, - "qdequantize workgroups_x must be non-zero" - ); - assert_eq!( - y.view().layout.element_count().get(), - b.rows - .checked_mul(b.cols) - .expect("qdequantize output element count overflow"), - "qdequantize output must contain one dense element per quantized element" - ); - assert!( - y.view().layout.is_row_major(), - "qdequantize output must be row-major" - ); - assert!( - matches!(y.view().buffer.element, ElementType::F32 | ElementType::F16), - "qdequantize output must be f32 or f16" - ); - - let total = b - .rows - .checked_mul(b.cols) - .expect("qdequantize output element count overflow"); - let workgroups = total.div_ceil(BLOCK); - let dispatch_y = workgroups.div_ceil(workgroups_x); - let y = Storage::from_view(StorageView { - buffer: y.view().buffer.clone(), - offset: y.view().offset, - layout: Layout::contiguous(MemoryLevel::Storage, Shape::new([1, total])), - }); - program.program_grid(BLOCK, [workgroups_x, dispatch_y, 1], |program| { - let lane = program.lane(); - let linear_group = program.program_id(WorkgroupAxis::X) - + program.program_id(WorkgroupAxis::Y) * workgroups_x; - let flat = linear_group * BLOCK + lane; - let mask = flat.lt(total); - let value = program.load_quantized( - b, - flat.clone() % b.rows, - flat.clone() / b.rows, - mask.clone(), - 0.0, - ); - program.store_cast(y.at((0, flat)), value, mask); - }); -} diff --git a/fusor-ml/tile-ir-kernels/src/kernels/rms_norm.rs b/fusor-ml/tile-ir-kernels/src/kernels/rms_norm.rs deleted file mode 100644 index 5c8d15431..000000000 --- a/fusor-ml/tile-ir-kernels/src/kernels/rms_norm.rs +++ /dev/null @@ -1,160 +0,0 @@ -use fusor_tile_ir::tile::{range, Tile}; -use fusor_tile_ir::{ElementType, ScalarElement, TileLiteral, WorkgroupAxis}; - -use super::types::RmsNormVec4Meta; - -const RMS_NORM_VEC4_BLOCK: u32 = 128; - -/// Tensor bindings and shape metadata for [`rms_norm_vec4`]. -/// -/// The offsets and strides live in [`RmsNormVec4Meta`]; these tensor refs only -/// describe the bound buffers and base layouts. -/// -/// ```no_run -/// # use fusor_tile_ir::{F32Bits, KernelBuilder, KernelTensorRef}; -/// # use fusor_tile_ir_kernels::{linear_storage_layout, rms_norm_vec4, RmsNormVec4, RmsNormVec4Meta}; -/// let layout = linear_storage_layout(); -/// let mut kb = KernelBuilder::<()>::new(); -/// let input = KernelTensorRef::new((), layout.clone()); -/// let weight = KernelTensorRef::new((), layout.clone()); -/// let output = KernelTensorRef::new((), layout); -/// rms_norm_vec4( -/// &mut kb, -/// RmsNormVec4 { -/// input, -/// residual: None, -/// weight, -/// bias: None, -/// output, -/// meta: RmsNormVec4Meta { -/// cols: 4, -/// cols_vec: 1, -/// eps: F32Bits::new(1e-5), -/// input_offset_vec: 0, -/// input_row_stride_vec: 1, -/// residual_offset_vec: None, -/// residual_row_stride_vec: 0, -/// weight_offset_vec: 0, -/// bias_offset_vec: None, -/// output_offset_vec: 0, -/// output_row_stride_vec: 1, -/// }, -/// rows: 1, -/// }, -/// ); -/// ``` -pub struct RmsNormVec4 { - /// Input tensor, read as packed `vec4` values. - pub input: fusor_tile_ir::KernelTensorRef, - /// Optional residual tensor added before normalization. - pub residual: Option>, - /// Per-column weight tensor, read as packed `vec4` values. - pub weight: fusor_tile_ir::KernelTensorRef, - /// Optional per-column bias tensor. - pub bias: Option>, - /// Output tensor, written as packed `vec4` values. - pub output: fusor_tile_ir::KernelTensorRef, - /// Column counts and vec4 offsets/strides. - pub meta: RmsNormVec4Meta, - /// Number of rows to normalize. - pub rows: u32, -} - -/// Build a packed-vec4 F32 RMS-norm kernel. -/// -/// Returns `None` when the row/column metadata is empty or the optional -/// residual/bias bindings do not match their optional metadata offsets. -pub fn rms_norm_vec4( - kb: &mut fusor_tile_ir::KernelBuilder, - spec: RmsNormVec4, -) -> Option<()> { - let RmsNormVec4 { - input, - residual, - weight, - bias, - output, - meta, - rows, - } = spec; - if rows == 0 || meta.cols == 0 || meta.cols_vec == 0 { - return None; - } - if meta.residual_offset_vec.is_some() != residual.is_some() - || meta.bias_offset_vec.is_some() != bias.is_some() - { - return None; - } - - let chunks = meta.cols_vec.div_ceil(RMS_NORM_VEC4_BLOCK); - let eps = meta.eps.get(); - - let vec4 = ElementType::vector(ScalarElement::F32, 4); - let input = kb.read(vec4, input); - let residual = residual.map(|r| kb.read(vec4, r)); - let weight = kb.read(vec4, weight); - let bias = bias.map(|b| kb.read(vec4, b)); - let output = kb.write(vec4, output); - let phase = kb.program(); - - phase.program_grid(RMS_NORM_VEC4_BLOCK, [rows, 1, 1], |program| { - let row = program.program_id(WorkgroupAxis::X); - let lane = program.lane(); - let [partial_sum] = program.fold( - range(chunks), - [Tile::literal(TileLiteral::f32(0.0))], - |program, loop_index, [acc]| { - let reduce_col = loop_index * RMS_NORM_VEC4_BLOCK + lane.clone(); - let reduce_mask = reduce_col.lt(meta.cols_vec); - let input_index = row.clone() * meta.input_row_stride_vec + reduce_col.clone(); - let mut value = program.load( - input.at(input_index), - reduce_mask.clone(), - TileLiteral::f32(0.0), - ); - if let Some(residual) = &residual { - let residual_index = - row.clone() * meta.residual_row_stride_vec + reduce_col.clone(); - value = value - + program.load( - residual.at(residual_index), - reduce_mask, - TileLiteral::f32(0.0), - ); - } - [acc + program.vector_dot(value.clone(), value)] - }, - ); - let total_sum = program.group_reduce_sum(RMS_NORM_VEC4_BLOCK, partial_sum); - let mean = total_sum / Tile::literal(TileLiteral::f32(meta.cols as f32)); - let scale = (mean + Tile::literal(TileLiteral::f32(eps))).inverse_sqrt(); - let scale = program.bind(scale); - - program.loop_range(chunks, |program, chunk| { - let col = lane.clone() + chunk * RMS_NORM_VEC4_BLOCK; - let mask = col.lt(meta.cols_vec); - let input_index = row.clone() * meta.input_row_stride_vec + col.clone(); - let mut value = - program.load(input.at(input_index), mask.clone(), TileLiteral::f32(0.0)); - if let Some(residual) = &residual { - let residual_index = row.clone() * meta.residual_row_stride_vec + col.clone(); - value = value - + program.load( - residual.at(residual_index), - mask.clone(), - TileLiteral::f32(0.0), - ); - } - let scale = program.vector_splat::<4>(ScalarElement::F32, scale.clone()); - let weight = program.load(weight.at(col.clone()), mask.clone(), TileLiteral::f32(0.0)); - let mut normalized = value * scale * weight; - if let Some(bias) = &bias { - normalized = normalized - + program.load(bias.at(col.clone()), mask.clone(), TileLiteral::f32(0.0)); - } - let output_index = row.clone() * meta.output_row_stride_vec + col; - program.store(output.at(output_index), normalized, mask); - }); - }); - Some(()) -} diff --git a/fusor-ml/tile-ir-kernels/src/kernels/softmax.rs b/fusor-ml/tile-ir-kernels/src/kernels/softmax.rs deleted file mode 100644 index ad119a0b3..000000000 --- a/fusor-ml/tile-ir-kernels/src/kernels/softmax.rs +++ /dev/null @@ -1,312 +0,0 @@ -use fusor_tile_ir::tile::{Mask, Tile, TileBlock, WorkgroupTile}; -use fusor_tile_ir::{ElementType, ScalarElement, TileLiteral, WorkgroupAxis}; - -use super::{ - helpers::{reduce_workgroup, scalar_of, supports_float, u32_tile, zero_fill, NEG_MAX_F32}, - types::SoftmaxMeta, -}; - -/// The per-workgroup softmax statistics produced by [`workgroup_softmax_block`]. -/// Every tile carries its `ElementType` as data; these values are all F32 -/// accumulators. -#[derive(Clone)] -pub(super) struct WorkgroupSoftmaxBlock { - pub max: Tile, - pub denom: Tile, - pub prob: Tile, -} - -/// Storage-typed `-f32::MAX` literal for the kernel's `element` (F32 or F16). -fn neg_max_fill(element: ElementType) -> TileLiteral { - match scalar_of(element) { - ScalarElement::F32 => TileLiteral::f32(NEG_MAX_F32), - ScalarElement::F16 => TileLiteral::F16(half::f16::from_f32(NEG_MAX_F32).to_bits()), - _ => panic!("softmax only supports F32 and F16 element types"), - } -} - -/// Accept the block sizes supported by the softmax kernels before -/// `program_grid` bakes `block` into `@workgroup_size`. -fn supported_block(block: u32) -> bool { - matches!(block, 128 | 512 | 1024) -} - -fn linear_group(program: &TileBlock<'_>, dispatch_size: [u32; 3]) -> Tile { - let x = program.program_id(WorkgroupAxis::X); - let y = program.program_id(WorkgroupAxis::Y); - let z = program.program_id(WorkgroupAxis::Z); - x + y * u32_tile(dispatch_size[0]) + z * u32_tile(dispatch_size[0] * dispatch_size[1]) -} - -fn storage_index( - program: &mut TileBlock<'_>, - meta: &SoftmaxMeta, - row: Tile, - axis_value: Tile, - output: bool, -) -> Tile { - let tensor_meta = if output { - &meta.output_meta - } else { - &meta.input_meta - }; - let strides = tensor_meta.strides.as_slice(); - let axis = meta.axis as usize; - let mut remaining = row; - let mut index = u32_tile(tensor_meta.offset); - - for dim in (0..meta.shape.len()).rev() { - let coord = if dim == axis { - axis_value.clone() - } else { - let size = meta.shape[dim]; - let coord = if size == 1 { - u32_tile(0) - } else { - remaining.clone() % u32_tile(size) - }; - if size != 1 { - remaining = program.bind(remaining / u32_tile(size)); - } - coord - }; - match strides[dim] { - 0 => {} - 1 => { - index = index + coord; - } - stride => { - index = index + coord * u32_tile(stride); - } - } - } - - program.bind(index) -} - -pub(super) fn softmax_partial_scale(block_max: Tile, global_max: Tile) -> Tile { - (block_max - global_max).exp() -} - -pub(super) fn workgroup_softmax_block( - program: &mut TileBlock<'_>, - lane: Tile, - score: Tile, - valid: Mask, - reduce: &WorkgroupTile, - probs: Option<&WorkgroupTile>, -) -> WorkgroupSoftmaxBlock { - let score = Tile::select( - valid.clone(), - score, - Tile::literal(TileLiteral::f32(NEG_MAX_F32)), - ); - program.store_workgroup(reduce, lane.clone(), score.clone()); - program.workgroup_barrier(); - reduce_workgroup(program, reduce, lane.clone(), |lhs, rhs| lhs.max(rhs)); - - let max_local = program.private(ElementType::F32); - let max_score = program.load_workgroup(reduce, 0u32); - program.store_local(&max_local, max_score); - let max_score = program.load_local(&max_local); - - program.workgroup_barrier(); - let raw_prob = (score - max_score.clone()).exp(); - let prob = Tile::select(valid, raw_prob, Tile::literal(TileLiteral::f32(0.0))); - if let Some(probs) = probs { - program.store_workgroup(probs, lane.clone(), prob.clone()); - } - program.store_workgroup(reduce, lane.clone(), prob.clone()); - program.workgroup_barrier(); - reduce_workgroup(program, reduce, lane, |lhs, rhs| lhs + rhs); - let denom = program.load_workgroup(reduce, 0u32); - - WorkgroupSoftmaxBlock { - max: max_score, - denom, - prob, - } -} - -pub fn softmax( - kb: &mut fusor_tile_ir::KernelBuilder, - element: ElementType, - input: fusor_tile_ir::KernelTensorRef, - output: fusor_tile_ir::KernelTensorRef, - meta: SoftmaxMeta, -) -> Option<()> { - if !supports_float(element) || !supported_block(meta.block) || meta.split_blocks != 1 { - return None; - } - let input = kb.read(element, input); - let output = kb.write(element, output); - let phase = kb.program(); - let reduce = phase.alloc_workgroup_array(ScalarElement::F32, meta.block); - - phase.program_grid(meta.block, meta.dispatch_size, |program| { - let lane = program.lane(); - let row = linear_group(program, meta.dispatch_size); - let axis_value = lane.clone(); - let valid = row - .clone() - .lt(u32_tile(meta.rows)) - .and(axis_value.clone().lt(u32_tile(meta.axis_len))); - let input_index = storage_index(program, &meta, row.clone(), axis_value.clone(), false); - let score = program - .load(input.at(input_index), valid.clone(), neg_max_fill(element)) - .cast(ElementType::F32); - let stats = workgroup_softmax_block(program, lane, score, valid.clone(), &reduce, None); - let output_index = storage_index(program, &meta, row, axis_value, true); - let value = (stats.prob / stats.denom).cast(element); - program.store(output.at(output_index), value, valid); - }); - Some(()) -} - -pub fn softmax_partials( - kb: &mut fusor_tile_ir::KernelBuilder, - element: ElementType, - input: fusor_tile_ir::KernelTensorRef, - scratch: fusor_tile_ir::KernelTensorRef, - meta: SoftmaxMeta, -) -> Option<()> { - if !supports_float(element) || !supported_block(meta.block) || meta.split_blocks < 2 { - return None; - } - let input = kb.read(element, input); - let scratch = kb.write(ElementType::F32, scratch); - let phase = kb.program(); - let reduce = phase.alloc_workgroup_array(ScalarElement::F32, meta.block); - - phase.program_grid(meta.block, meta.dispatch_size, |program| { - let lane = program.lane(); - let group = linear_group(program, meta.dispatch_size); - let total_groups = u32_tile(meta.rows * meta.split_blocks); - let group_valid = group.clone().lt(total_groups); - let row = program.bind(group.clone() % u32_tile(meta.rows)); - let split = program.bind(group / u32_tile(meta.rows)); - let axis_value = program.bind(split.clone() * u32_tile(meta.block) + lane.clone()); - let valid = group_valid - .clone() - .and(axis_value.clone().lt(u32_tile(meta.axis_len))); - let input_index = storage_index(program, &meta, row.clone(), axis_value, false); - let score = program - .load(input.at(input_index), valid.clone(), neg_max_fill(element)) - .cast(ElementType::F32); - let stats = workgroup_softmax_block(program, lane.clone(), score, valid, &reduce, None); - let partial_base = program.bind((row * u32_tile(meta.split_blocks) + split) * u32_tile(2)); - program.if_then(group_valid.and(lane.eq(u32_tile(0))), |program| { - program.store(scratch.at(partial_base.clone()), stats.denom, Mask::all()); - program.store( - scratch.at(partial_base + u32_tile(1)), - stats.max, - Mask::all(), - ); - }); - }); - Some(()) -} - -pub fn softmax_write( - kb: &mut fusor_tile_ir::KernelBuilder, - element: ElementType, - input: fusor_tile_ir::KernelTensorRef, - global: fusor_tile_ir::KernelTensorRef, - output: fusor_tile_ir::KernelTensorRef, - meta: SoftmaxMeta, -) -> Option<()> { - if !supports_float(element) || !supported_block(meta.block) || meta.split_blocks < 2 { - return None; - } - let input = kb.read(element, input); - let global = kb.read(ElementType::F32, global); - let output = kb.write(element, output); - let phase = kb.program(); - - phase.program_grid(meta.block, meta.dispatch_size, |program| { - let lane = program.lane(); - let group = linear_group(program, meta.dispatch_size); - let total_groups = u32_tile(meta.rows * meta.split_blocks); - let group_valid = group.clone().lt(total_groups); - let row = program.bind(group.clone() % u32_tile(meta.rows)); - let split = program.bind(group / u32_tile(meta.rows)); - let axis_value = program.bind(split * u32_tile(meta.block) + lane.clone()); - - let row_base = program.bind(row.clone() * u32_tile(2)); - let denom = program.load( - global.at(row_base.clone()), - group_valid.clone(), - TileLiteral::f32(0.0), - ); - let max_score = program.load( - global.at(row_base + u32_tile(1)), - group_valid.clone(), - TileLiteral::f32(NEG_MAX_F32), - ); - - let valid = group_valid.and(axis_value.clone().lt(u32_tile(meta.axis_len))); - let input_index = storage_index(program, &meta, row.clone(), axis_value.clone(), false); - let output_index = storage_index(program, &meta, row, axis_value, true); - let value = program - .load(input.at(input_index), valid.clone(), zero_fill(element)) - .cast(ElementType::F32); - let prob = (value - max_score).exp() / denom; - program.store(output.at(output_index), prob.cast(element), valid); - }); - Some(()) -} - -pub fn softmax_reduce( - kb: &mut fusor_tile_ir::KernelBuilder, - scratch: fusor_tile_ir::KernelTensorRef, - global: fusor_tile_ir::KernelTensorRef, - meta: SoftmaxMeta, -) -> Option<()> { - if !supported_block(meta.block) || meta.split_blocks < 2 { - return None; - } - let scratch = kb.read(ElementType::F32, scratch); - let global = kb.write(ElementType::F32, global); - let phase = kb.program(); - let block = meta.block; - - phase.program_grid(block, meta.dispatch_size, |program| { - let lane = program.lane(); - let row = linear_group(program, meta.dispatch_size); - let row_valid = row.clone().lt(u32_tile(meta.rows)); - let partial_row_base = program.bind(row.clone() * u32_tile(meta.split_blocks * 2)); - - program.if_then(row_valid.and(lane.eq(u32_tile(0))), |program| { - let mut max_score = Tile::literal(TileLiteral::f32(NEG_MAX_F32)); - for split in 0..meta.split_blocks { - let block_max = program.load( - scratch.at(partial_row_base.clone() + u32_tile(split * 2 + 1)), - Mask::all(), - TileLiteral::f32(NEG_MAX_F32), - ); - max_score = max_score.max(block_max); - } - - let mut denom = Tile::literal(TileLiteral::f32(0.0)); - for split in 0..meta.split_blocks { - let block_base = partial_row_base.clone() + u32_tile(split * 2); - let block_denom = program.load( - scratch.at(block_base.clone()), - Mask::all(), - TileLiteral::f32(0.0), - ); - let block_max = program.load( - scratch.at(block_base + u32_tile(1)), - Mask::all(), - TileLiteral::f32(NEG_MAX_F32), - ); - denom = denom + block_denom * softmax_partial_scale(block_max, max_score.clone()); - } - - let global_base = program.bind(row * u32_tile(2)); - program.store(global.at(global_base.clone()), denom, Mask::all()); - program.store(global.at(global_base + u32_tile(1)), max_score, Mask::all()); - }); - }); - Some(()) -} diff --git a/fusor-ml/tile-ir-kernels/src/kernels/types.rs b/fusor-ml/tile-ir-kernels/src/kernels/types.rs index 100e8bd64..952666b80 100644 --- a/fusor-ml/tile-ir-kernels/src/kernels/types.rs +++ b/fusor-ml/tile-ir-kernels/src/kernels/types.rs @@ -1,22 +1,3 @@ -use fusor_tile_ir::F32Bits; - -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -/// Flash-attention tensor dimensions. -pub struct FlashAttentionDims { - /// Batch size. - pub batch: u32, - /// Number of query heads. - pub num_heads: u32, - /// Number of key/value heads. - pub num_kv_heads: u32, - /// Query sequence length. - pub q_seq_len: u32, - /// Key/value sequence length. - pub kv_seq_len: u32, - /// Per-head embedding dimension. - pub head_dim: u32, -} - #[derive(Clone, Debug, PartialEq, Eq)] /// Runtime tensor strides and base offset. pub struct TensorMeta { @@ -33,117 +14,6 @@ impl TensorMeta { } } -#[derive(Clone, Debug, PartialEq, Eq)] -/// Metadata for streaming flash attention. -pub struct FlashAttentionMeta { - /// Logical attention dimensions. - pub dims: FlashAttentionDims, - /// Attention scale applied to QK scores. - pub scale: F32Bits, - /// Query tensor metadata. - pub q_meta: TensorMeta, - /// Key tensor metadata. - pub k_meta: TensorMeta, - /// Value tensor metadata. - pub v_meta: TensorMeta, - /// Optional additive mask metadata. - pub mask_meta: Option, - /// Output tensor metadata. - pub output_meta: TensorMeta, - /// Dispatch grid used for the generated tile program. - pub dispatch_size: [u32; 3], - /// When `true`, the kernel applies a strict lower-triangular causal mask - /// (kv_idx <= q_idx) by skipping out-of-bound KV chunks/lanes. The - /// `mask_meta` field must be `None` in this case — no additive mask is - /// loaded. - pub causal: bool, -} - -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -/// Metadata for the small F32 flash-decode kernel. -pub struct FlashDecodeSmallMeta { - /// Logical decode dimensions. - pub dims: FlashAttentionDims, - /// Attention scale applied to QK scores. - pub scale: F32Bits, - /// Default active KV length. Runtime params may override this. - pub active_kv_len: u32, - /// Workgroup block size to use for decode. - pub decode_block: u32, - /// Whether to use the tiled decode path for long active KV lengths. - pub tiled: bool, - /// Number of KV tiles used by the split decode path. - pub split_blocks: u32, - /// Query-heads per KV-head group. - pub groups: u32, - /// Query tensor element offset. - pub q_offset: u32, - /// Key tensor element offset. - pub k_offset: u32, - /// Value tensor element offset. - pub v_offset: u32, - /// Output tensor element offset. - pub output_offset: u32, - /// Query strides. - pub q_strides: [u32; 4], - /// Key strides. - pub k_strides: [u32; 4], - /// Value strides. - pub v_strides: [u32; 4], - /// Output strides. - pub output_strides: [u32; 4], -} - -#[derive(Clone, Debug, PartialEq, Eq)] -/// Metadata for a direct softmax kernel over one tensor axis. -pub struct SoftmaxMeta { - /// Logical tensor shape. - pub shape: Vec, - /// Softmax axis. - pub axis: u32, - /// Number of logical rows after removing the softmax axis. - pub rows: u32, - /// Length of the softmax axis. - pub axis_len: u32, - /// Workgroup block size used for one axis tile. - pub block: u32, - /// Number of axis tiles for split softmax. - pub split_blocks: u32, - /// Input tensor metadata. - pub input_meta: TensorMeta, - /// Output tensor metadata. - pub output_meta: TensorMeta, - /// Dispatch grid used by the generated tile program. - pub dispatch_size: [u32; 3], -} - -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -/// Vec4 RMS-norm offsets, strides, and scalar parameters. -pub struct RmsNormVec4Meta { - /// Dense column count. - pub cols: u32, - /// Column count in packed vec4 elements. - pub cols_vec: u32, - /// Epsilon added before reciprocal square root. - pub eps: F32Bits, - /// Input vec4 offset. - pub input_offset_vec: u32, - /// Input row stride in vec4 elements. - pub input_row_stride_vec: u32, - /// Optional residual vec4 offset. - pub residual_offset_vec: Option, - /// Residual row stride in vec4 elements. - pub residual_row_stride_vec: u32, - /// Weight vec4 offset. - pub weight_offset_vec: u32, - /// Optional bias vec4 offset. - pub bias_offset_vec: Option, - /// Output vec4 offset. - pub output_offset_vec: u32, - /// Output row stride in vec4 elements. - pub output_row_stride_vec: u32, -} - #[derive(Clone, Copy, Debug, PartialEq, Eq)] /// Metadata for one top-k chunk pass. pub struct TopKChunkMeta { diff --git a/fusor-ml/tile-ir-kernels/src/lib.rs b/fusor-ml/tile-ir-kernels/src/lib.rs index db9c43f1e..6afd85f81 100644 --- a/fusor-ml/tile-ir-kernels/src/lib.rs +++ b/fusor-ml/tile-ir-kernels/src/lib.rs @@ -1,8 +1,8 @@ //! Pre-built kernels for `fusor-tile-ir`. //! //! `fusor-tile-ir` contains the IR, lowerer, and generic tile builder. This -//! crate contains concrete kernels: dense matmul/GEMV, quantized matmul/GEMV, -//! dequantization, flash attention, top-k, RMS norm, and Mirostat sampling. +//! crate contains concrete kernels: cooperative dense matmul, quantized +//! matmul/GEMV, top-k, and Mirostat sampling. mod dispatch; mod grid; @@ -14,21 +14,14 @@ pub use dispatch::{ qgemv_subgroups_per_workgroup_for_shape, SubgroupConfig, }; pub use kernels::{ - batched_gemv_with_epilogues, batched_matmul_register_with_epilogues, - batched_matmul_with_epilogues, flash_attention, flash_attention_tiled, flash_decode_small, - flash_decode_split_partials, flash_decode_split_reduce, flash_outputs_per_workgroup, - flash_tiled_dispatch_size, flash_tiled_outputs_per_workgroup, linear_storage_layout, mirostat2, - qdequantize, qgemv_with_epilogue, qgemv_workgroup_f16_with_epilogue, + linear_storage_layout, mirostat2, qgemv_with_epilogue, qgemv_workgroup_f16_with_epilogue, qgemv_workgroup_storage_f16_with_epilogue, qgemv_workgroup_with_epilogue, qmatmul_with_epilogue, qmatmul_workgroup_f16_with_epilogues, qmatmul_workgroup_storage_f16_with_epilogues, qmatmul_workgroup_with_epilogues, - quantized_matrix, quantized_matrix_for, rms_norm_vec4, softmax, softmax_partials, - softmax_reduce, softmax_write, standard_sampler, top_k_chunk, top_k_exactness, top_k_merge, - try_batched_coop_matmul, AccumCast, DenseCoopMatmulConfig, DenseCoopMatmulTile, - DenseMatmulShape, DenseMatmulTensors, DenseMatmulTile, FlashAttentionDims, FlashAttentionMeta, - FlashAttentionTensors, FlashDecodeSmallMeta, IntoQgemvEpilogues, MergeTopKMeta, Mirostat2, - Mirostat2Meta, RmsNormVec4, RmsNormVec4Meta, SoftmaxMeta, StandardSampler, TensorMeta, - TopKChunkMeta, TopKExactnessMeta, + quantized_matrix, quantized_matrix_for, standard_sampler, top_k_chunk, top_k_exactness, + top_k_merge, try_batched_coop_matmul, AccumCast, DenseCoopMatmulConfig, DenseCoopMatmulTile, + DenseMatmulShape, DenseMatmulTensors, IntoQgemvEpilogues, MergeTopKMeta, Mirostat2, + Mirostat2Meta, StandardSampler, TensorMeta, TopKChunkMeta, TopKExactnessMeta, }; pub use types::{ cooperative_store_layout_supported, DenseMatmulEpilogues, QmatmulEpilogues, QmatmulExtra, diff --git a/fusor-ml/tile-ir-kernels/tests/lowering.rs b/fusor-ml/tile-ir-kernels/tests/lowering.rs index 6cfd3df9f..b8eb3cbd3 100644 --- a/fusor-ml/tile-ir-kernels/tests/lowering.rs +++ b/fusor-ml/tile-ir-kernels/tests/lowering.rs @@ -1,16 +1,10 @@ -use fusor_tile_ir::{ - tile, ElementType, F32Bits, GgmlQuantFormat, KernelBuilder, KernelTensorRef, Layout, - MemoryLevel, NagaKernel, ScalarElement, Shape, -}; +use fusor_tile_ir::{tile, GgmlQuantFormat, NagaKernel, ScalarElement, Shape}; use fusor_tile_ir_kernels::{ - batched_gemv_with_epilogues, batched_matmul_with_epilogues, flash_attention, - linear_storage_layout, qdequantize, qgemv_with_epilogue, qgemv_workgroup_f16_with_epilogue, - qgemv_workgroup_with_epilogue, qmatmul_with_epilogue, qmatmul_workgroup_f16_with_epilogues, - qmatmul_workgroup_with_epilogues, quantized_matrix, rms_norm_vec4, try_batched_coop_matmul, - DenseCoopMatmulConfig, DenseCoopMatmulTile, DenseMatmulEpilogues, DenseMatmulShape, - DenseMatmulTensors, DenseMatmulTile, FlashAttentionDims, FlashAttentionMeta, - FlashAttentionTensors, QmatmulEpilogues, RmsNormVec4, RmsNormVec4Meta, SubgroupConfig, - TensorMeta, UnaryEpilogue, UnaryEpilogueWithExtras, + qgemv_with_epilogue, qgemv_workgroup_f16_with_epilogue, qgemv_workgroup_with_epilogue, + qmatmul_with_epilogue, qmatmul_workgroup_f16_with_epilogues, qmatmul_workgroup_with_epilogues, + quantized_matrix, try_batched_coop_matmul, DenseCoopMatmulConfig, DenseCoopMatmulTile, + DenseMatmulEpilogues, DenseMatmulShape, DenseMatmulTensors, QmatmulEpilogues, SubgroupConfig, + UnaryEpilogue, UnaryEpilogueWithExtras, }; fn lower_or_fail(ir: &fusor_tile_ir::KernelIr, label: &str) -> NagaKernel { @@ -30,84 +24,6 @@ fn coop_token() -> fusor_tile_ir::CoopMatrixToken { fusor_tile_ir::CoopMatrixToken::new_unchecked() } -#[test] -fn streaming_flash_attention_regression_shape_lowers_to_naga() { - let layout = linear_storage_layout(); - let mut kb = KernelBuilder::<()>::new(); - flash_attention::<()>( - &mut kb, - ScalarElement::F32.element(), - FlashAttentionTensors { - q: KernelTensorRef::new((), layout.clone()), - k: KernelTensorRef::new((), layout.clone()), - v: KernelTensorRef::new((), layout.clone()), - mask: None, - output: KernelTensorRef::new((), layout), - }, - FlashAttentionMeta { - dims: FlashAttentionDims { - batch: 1, - num_heads: 32, - num_kv_heads: 8, - q_seq_len: 48, - kv_seq_len: 48, - head_dim: 128, - }, - scale: F32Bits::new(1.0 / 128.0f32.sqrt()), - q_meta: TensorMeta::new(vec![196_608, 6_144, 128, 1], 0), - k_meta: TensorMeta::new(vec![49_152, 6_144, 128, 1], 0), - v_meta: TensorMeta::new(vec![49_152, 6_144, 128, 1], 0), - mask_meta: None, - output_meta: TensorMeta::new(vec![196_608, 6_144, 128, 1], 0), - dispatch_size: [16, 1536, 1], - causal: false, - }, - subgroup_token(), - 32, - ) - .expect("streaming flash attention should build"); - let (ir, _) = kb.finish(); - - lower_or_fail(&ir, "streaming flash attention"); -} - -#[test] -fn rms_norm_vec4_minimal_lowers() { - let layout = Layout::strided(MemoryLevel::Storage, Shape::new([1]), &[1]); - let mut kb = KernelBuilder::<()>::new(); - let input = KernelTensorRef::with_offset((), layout.clone(), 0); - let weight = KernelTensorRef::with_offset((), layout.clone(), 0); - let output = KernelTensorRef::with_offset((), layout.clone(), 0); - let meta = RmsNormVec4Meta { - cols: 4, - cols_vec: 1, - eps: F32Bits::new(1e-5), - input_offset_vec: 0, - input_row_stride_vec: 1, - residual_offset_vec: None, - residual_row_stride_vec: 0, - weight_offset_vec: 0, - bias_offset_vec: None, - output_offset_vec: 0, - output_row_stride_vec: 1, - }; - rms_norm_vec4( - &mut kb, - RmsNormVec4 { - input, - residual: None, - weight, - bias: None, - output, - meta, - rows: 1, - }, - ) - .unwrap(); - let (ir, _) = kb.finish(); - lower_or_fail(&ir, "rms_norm_vec4"); -} - fn qgemv_ir_with_subgroup_size( format: GgmlQuantFormat, rows: u32, @@ -220,150 +136,6 @@ fn cooperative_qmatmul_lowers() { lower_or_fail(&ir, "cooperative qmatmul"); } -#[test] -fn batched_dense_f32_matmul_lowers() { - let ir = tile::build(|program| { - let shape = DenseMatmulShape { - batch: 3, - m: 8, - k: 256, - n: 4, - }; - let a = program.storage_read( - ScalarElement::F32.element(), - Shape::new([shape.batch * shape.m, shape.k]), - ); - let b = program.storage_read( - ScalarElement::F32.element(), - Shape::new([shape.batch * shape.k, shape.n]), - ); - let y = program.storage_write( - ScalarElement::F32.element(), - Shape::new([shape.batch * shape.m, shape.n]), - ); - batched_matmul_with_epilogues( - program, - DenseMatmulTensors { - a: &a, - b: &b, - y: &y, - }, - shape, - &DenseMatmulEpilogues::empty(), - 65_535, - DenseMatmulTile::new(32, 32, 8, 4, 4, 64), - ); - }); - lower_or_fail(&ir, "batched dense f32 matmul"); -} - -#[test] -fn batched_dense_f32_gemv_lowers() { - let ir = tile::build(|program| { - let shape = DenseMatmulShape { - batch: 3, - m: 5, - k: 256, - n: 1, - }; - let a = program.storage_read( - ScalarElement::F32.element(), - Shape::new([shape.batch * shape.m, shape.k]), - ); - let b = program.storage_read( - ScalarElement::F32.element(), - Shape::new([shape.batch * shape.k, shape.n]), - ); - let y = program.storage_write( - ScalarElement::F32.element(), - Shape::new([shape.batch * shape.m, shape.n]), - ); - batched_gemv_with_epilogues( - program, - &a, - &b, - &y, - shape, - &DenseMatmulEpilogues::empty(), - 65_535, - subgroup_config(32), - ); - }); - lower_or_fail(&ir, "batched dense f32 gemv"); -} - -#[test] -fn batched_dense_f16_matmul_lowers() { - let ir = tile::build(|program| { - let shape = DenseMatmulShape { - batch: 2, - m: 8, - k: 128, - n: 4, - }; - let a = program.storage_read( - ScalarElement::F16.element(), - Shape::new([shape.batch * shape.m, shape.k]), - ); - let b = program.storage_read( - ScalarElement::F16.element(), - Shape::new([shape.batch * shape.k, shape.n]), - ); - let y = program.storage_write( - ScalarElement::F16.element(), - Shape::new([shape.batch * shape.m, shape.n]), - ); - batched_matmul_with_epilogues( - program, - DenseMatmulTensors { - a: &a, - b: &b, - y: &y, - }, - shape, - &DenseMatmulEpilogues::empty(), - 65_535, - DenseMatmulTile::new(32, 32, 8, 4, 4, 64), - ); - }); - lower_or_fail(&ir, "batched dense f16 matmul"); -} - -#[test] -fn batched_dense_f16_gemv_lowers() { - let ir = tile::build(|program| { - let shape = DenseMatmulShape { - batch: 2, - m: 5, - k: 128, - n: 1, - }; - let a = program.storage_read( - ScalarElement::F16.element(), - Shape::new([shape.batch * shape.m, shape.k]), - ); - let b = program.storage_read( - ScalarElement::F16.element(), - Shape::new([shape.batch * shape.k, shape.n]), - ); - let y = program.storage_write( - ScalarElement::F16.element(), - Shape::new([shape.batch * shape.m, shape.n]), - ); - batched_gemv_with_epilogues( - program, - &a, - &b, - &y, - shape, - &DenseMatmulEpilogues::empty(), - 65_535, - subgroup_config(32), - ); - }); - lower_or_fail(&ir, "batched dense f16 gemv"); -} - #[test] fn cooperative_dense_f32_matmul_lowers() { let ir = tile::build(|program| { @@ -591,36 +363,6 @@ fn cooperative_dense_f32_matmul_128x256_npass_lowers() { lower_or_fail(&ir, "cooperative dense f32 128x256 N_PASSES=4 matmul"); } -#[test] -fn qdequantize_lowers() { - let ir = tile::build(|program| { - let b = quantized_matrix(program, GgmlQuantFormat::Q4K, 256, 4); - let y = program.storage_write(ElementType::F32, Shape::new([1024])); - qdequantize(program, &b, &y, 1); - }); - lower_or_fail(&ir, "qdequantize"); -} - -#[test] -fn q4k_native_qdequantize_lowers() { - let ir = tile::build(|program| { - let b = quantized_matrix(program, GgmlQuantFormat::Q4KNative, 256, 4); - let y = program.storage_write(ElementType::F32, Shape::new([1024])); - qdequantize(program, &b, &y, 1); - }); - lower_or_fail(&ir, "q4k native qdequantize"); -} - -#[test] -fn qdequantize_f16_output_lowers() { - let ir = tile::build(|program| { - let b = quantized_matrix(program, GgmlQuantFormat::Q4KNative, 256, 4); - let y = program.storage_write(ElementType::F16, Shape::new([1024])); - qdequantize(program, &b, &y, 1); - }); - lower_or_fail(&ir, "qdequantize f16 output"); -} - /// Regression for the fallback branch in `qmatmul_tile_with_epilogue`. When /// `BM*BN*BK != 256` (true for every caller in core's `quantized/matmul`), /// the fallback used to drop the epilogue. With a non-identity `post` diff --git a/fusor-ml/tile-ir/src/tile/block.rs b/fusor-ml/tile-ir/src/tile/block.rs index 08376816a..7b341669a 100644 --- a/fusor-ml/tile-ir/src/tile/block.rs +++ b/fusor-ml/tile-ir/src/tile/block.rs @@ -70,6 +70,12 @@ impl TileBlock<'_> { /// Masked dequantizing load of one f32 value from a quantized matrix at a /// `(row, col)` coordinate. + /// Dequantize one element of a block-quantized matrix. + /// + /// Addressing follows the matrix's flat block order, where + /// [`QuantizedMatrix::rows`] is the *length* of one dense row: `row` is + /// the position along that contiguous axis and `col` selects which row + /// (`flat = col * matrix.rows + row`). pub fn load_quantized( &self, matrix: &QuantizedMatrix, diff --git a/fusor-ml/tile-ir/src/tile/capability.rs b/fusor-ml/tile-ir/src/tile/capability.rs index aca46e3b2..b1562efc3 100644 --- a/fusor-ml/tile-ir/src/tile/capability.rs +++ b/fusor-ml/tile-ir/src/tile/capability.rs @@ -40,13 +40,8 @@ impl SubgroupToken { program.num_subgroups() } - /// Reduction across one subgroup. - pub(crate) fn subgroup_reduce( - self, - program: &TileBlock<'_>, - op: TileReduceOp, - value: Tile, - ) -> Tile { + /// Reduction across one subgroup with an explicit operator. + pub fn subgroup_reduce(self, program: &TileBlock<'_>, op: TileReduceOp, value: Tile) -> Tile { program.subgroup_reduce(op, value) } diff --git a/models/kalosm-llama/examples/compiler_smoke.rs b/models/kalosm-llama/examples/compiler_smoke.rs new file mode 100644 index 000000000..47ffd53eb --- /dev/null +++ b/models/kalosm-llama/examples/compiler_smoke.rs @@ -0,0 +1,36 @@ +use kalosm_llama::*; +use kalosm_model_types::ModelLoadingProgress; +use prelude::{StreamExt, TextCompletionModelExt}; + +fn main() { + let _ = tracing_subscriber::fmt::try_init(); + pollster::block_on(async { + let model = Llama::builder() + .with_source(LlamaSource::llama_3_1_8b_chat()) + .build_with_loading_handler(|_: ModelLoadingProgress| {}) + .await + .unwrap(); + + let prompt = "The capital of France is"; + let mut stream = model.complete(prompt).take(72); + let mut text = String::new(); + // Warmup: first tokens pay one-time kernel compilation. + for _ in 0..8 { + if let Some(token) = stream.next().await { + text.push_str(&token); + } + } + let start = std::time::Instant::now(); + let mut tokens = 0usize; + while let Some(token) = stream.next().await { + text.push_str(&token); + tokens += 1; + } + let elapsed = start.elapsed(); + println!("OUTPUT: {}", text.replace('\n', " | ")); + println!( + "steady-state tokens={tokens} elapsed={elapsed:?} tps={:.2}", + tokens as f64 / elapsed.as_secs_f64() + ); + }); +} From 75420a57612ef75a2dd1fbde4d225b5106636a9f Mon Sep 17 00:00:00 2001 From: Evan Almloff Date: Thu, 11 Jun 2026 17:40:44 -0500 Subject: [PATCH 02/16] remove matmul fallback --- fusor-ml/conformance/src/suite/webgpu.rs | 7 +- fusor-ml/core/src/compute_graph/mod.rs | 28 +- .../src/compute_graph/resolve/execution.rs | 17 +- .../compute_graph/resolve/fusion_matmul.rs | 2 + .../core/src/compute_graph/resolve/mod.rs | 1 + .../compute_graph/resolve/recognize_cat.rs | 588 +++++++ fusor-ml/core/src/compute_graph/tests.rs | 216 ++- .../core/src/quantized/matmul/fallback.rs | 239 --- fusor-ml/core/src/quantized/matmul/kernel.rs | 219 ++- fusor-ml/core/src/quantized/matmul/mod.rs | 120 +- fusor-ml/core/src/sampling/mod.rs | 7 +- fusor-ml/core/src/sampling/pipeline.rs | 1457 ++++------------- fusor-ml/core/src/sampling/qmat_topk.rs | 104 -- fusor-ml/core/src/sampling/tests.rs | 40 +- fusor-ml/core/src/sampling/topk.rs | 2 +- fusor-ml/core/src/slice_assign.rs | 106 +- fusor-ml/core/src/tensor/mod.rs | 8 +- fusor-ml/core/src/tensor/sampling.rs | 257 ++- fusor-ml/core/src/top_k.rs | 6 +- fusor-ml/core/src/view.rs | 2 +- fusor-ml/cpu/src/matmul.rs | 12 +- fusor-ml/cpu/src/quantized.rs | 4 +- fusor-ml/fusor/src/gpu.rs | 39 +- fusor-ml/fusor/src/lib.rs | 77 +- models/kalosm-llama/examples/qwen_probe.rs | 38 + models/kalosm-llama/examples/vision.rs | 44 +- models/kalosm-llama/src/model/forward.rs | 47 +- models/kalosm-llama/src/model/mod.rs | 3 +- 28 files changed, 1791 insertions(+), 1899 deletions(-) create mode 100644 fusor-ml/core/src/compute_graph/resolve/recognize_cat.rs delete mode 100644 fusor-ml/core/src/quantized/matmul/fallback.rs delete mode 100644 fusor-ml/core/src/sampling/qmat_topk.rs create mode 100644 models/kalosm-llama/examples/qwen_probe.rs diff --git a/fusor-ml/conformance/src/suite/webgpu.rs b/fusor-ml/conformance/src/suite/webgpu.rs index 173ca8484..53071eb81 100644 --- a/fusor-ml/conformance/src/suite/webgpu.rs +++ b/fusor-ml/conformance/src/suite/webgpu.rs @@ -338,8 +338,8 @@ async fn check_q4k_fused_sampler(device: &Device) -> Result<(), SuiteError> { let gpu_device = gpu.device().clone(); let mut sampler = Mirostat2Sampler::new(&gpu_device, 10.0); let token = hidden - .try_sample_mirostat2_token_q_mat( - &matrix, + .q_mat_mul(&matrix) + .sample_mirostat2_token( &mut sampler, &[], Mirostat2SamplerParams { @@ -352,8 +352,7 @@ async fn check_q4k_fused_sampler(device: &Device) -> Result<(), SuiteError> { }, ) .await - .map_err(|err| SuiteError::case(case, err))? - .ok_or_else(|| SuiteError::case(case, "fused sampler returned None"))?; + .map_err(|err| SuiteError::case(case, err))?; if token >= weight_shape[0] as u32 { return Err(SuiteError::case( case, diff --git a/fusor-ml/core/src/compute_graph/mod.rs b/fusor-ml/core/src/compute_graph/mod.rs index f50655375..f28a0c735 100644 --- a/fusor-ml/core/src/compute_graph/mod.rs +++ b/fusor-ml/core/src/compute_graph/mod.rs @@ -72,11 +72,19 @@ impl ComputeGraph { self.with_mut(|inner| inner.execute_eager(operation)) } - /// Clone the view at `key` if that node is a view. Used to collapse view - /// chains at construction time. + /// Clone the view at `key` if that node is an unresolved view. Used to + /// collapse view chains at construction time. Cached views are excluded: + /// a resolved view no longer keeps its base alive, so the base node may + /// already be culled from the graph — and a new view over the cached + /// buffer is a zero-cost map anyway, while composing past it would force + /// a recompute from the (possibly released) base. pub(crate) fn get_view(&self, key: NodeIndex) -> Option { let inner = self.inner.read(); - match &inner.nodes.nodes.node_weight(key)?.variant { + let node = inner.nodes.nodes.node_weight(key)?; + if node.cached.is_some() { + return None; + } + match &node.variant { ComputeGraphNodeVariant::View(op) => Some(op.clone()), _ => None, } @@ -726,15 +734,21 @@ impl ComputeGraphInner { } } - // Check that all dependencies of non-cached nodes exist + // Check that all dependencies of non-cached nodes that could still + // resolve exist. A dead uncached node (no references and no alive + // descendants — e.g. an intermediate whose buffer was released after + // its handle dropped) is unreachable by any future resolve: resolves + // start from a live handle, and everything alive transitively keeps + // its dependencies' `live_descendant_count` positive. Its + // dependencies may therefore be legitimately removed from under it. for key in self.nodes.nodes.node_indices() { - let is_cached = self + let resolvable = self .nodes .nodes .node_weight(key) - .map(|n| n.cached.is_some()) + .map(|n| n.cached.is_none() && n.should_keep_cached()) .unwrap_or(false); - if is_cached { + if !resolvable { continue; } self.visit_dependencies(key, &mut |dependency| { diff --git a/fusor-ml/core/src/compute_graph/resolve/execution.rs b/fusor-ml/core/src/compute_graph/resolve/execution.rs index 54f0e185e..93039d67c 100644 --- a/fusor-ml/core/src/compute_graph/resolve/execution.rs +++ b/fusor-ml/core/src/compute_graph/resolve/execution.rs @@ -11,11 +11,19 @@ impl Resolver { op.visit_dependencies(&mut |dep| { if let Some(count) = remaining_consumers.get_mut(&dep) { *count = count.saturating_sub(1); - if *count == 0 && !targets.contains(&dep) && !graph.has_live_reference(dep) { + if *count == 0 + && !targets.contains(&dep) + && !graph.has_live_lazy_descendant(dep) + { // All consumers within this execution have been // processed and no user-held lazy tensor still // transitively depends on `dep` — free the cached - // buffer. + // buffer. The descendant check must include + // `live_descendant_count`, not just direct + // references: clearing `cached` on a node that still + // has an alive-uncached descendant flips it back to + // alive-uncached without propagating the transition, + // undercounting every ancestor's descendant counter. if let Some(node) = graph.nodes.nodes.node_weight_mut(dep) { node.cached = None; } @@ -59,9 +67,6 @@ impl Resolver { graph: &ComputeGraphInner, operation: &crate::slice_assign::SliceAssignOperation, ) -> Option<(TensorData, Vec)> { - if !operation.in_place { - return None; - } let input = graph.get_cached_result(operation.input)?; let value = graph.get_cached_result(operation.value)?; if input.datatype() != value.datatype() || operation.slices.len() != input.layout().rank() { @@ -244,6 +249,7 @@ impl Resolver { self.recognize_embeddings(graph); self.recognize_attention(graph); self.fuse_row_programs(graph); + self.recognize_assign_chains(graph); // The current rewrite rules can only start from Nary nodes (nary // fusion, post-op reduce/matmul fusion) or MatMul nodes (pre-op // unary fusion). Avoid scanning every QMatMul/attention node in @@ -376,6 +382,7 @@ impl Resolver { self.recognize_embeddings(graph); self.recognize_attention(graph); self.fuse_row_programs(graph); + self.recognize_assign_chains(graph); let has_qmatmul = self.execution_graph.node_indices().any(|node| { matches!( self.execution_graph[node].variant, diff --git a/fusor-ml/core/src/compute_graph/resolve/fusion_matmul.rs b/fusor-ml/core/src/compute_graph/resolve/fusion_matmul.rs index 55b6b8ae6..5e42ae802 100644 --- a/fusor-ml/core/src/compute_graph/resolve/fusion_matmul.rs +++ b/fusor-ml/core/src/compute_graph/resolve/fusion_matmul.rs @@ -200,6 +200,7 @@ impl Resolver { if self.check_cached(graph, input_inner) || input_datatype != crate::DataTypeEnum::F32 || nary.output_datatype != crate::DataTypeEnum::F32 + || !qmatmul_op.supports_elementwise_epilogue_fusion(&graph.device()) { continue; } @@ -238,6 +239,7 @@ impl Resolver { .iter() .product::() == 1 + && qmatmul_op.supports_elementwise_epilogue_fusion(&graph.device()) && !self.check_cached(graph, qmatmul_op.input) && let Some(input_exec) = self.get_input_node_in_exec_graph(qmatmul_op.input) { diff --git a/fusor-ml/core/src/compute_graph/resolve/mod.rs b/fusor-ml/core/src/compute_graph/resolve/mod.rs index 56c057c27..89c7213be 100644 --- a/fusor-ml/core/src/compute_graph/resolve/mod.rs +++ b/fusor-ml/core/src/compute_graph/resolve/mod.rs @@ -34,6 +34,7 @@ mod fusion_row; mod plan_cache; mod recognize; mod recognize_attention; +mod recognize_cat; mod run; pub(crate) use plan_cache::structural_kernel_key; diff --git a/fusor-ml/core/src/compute_graph/resolve/recognize_cat.rs b/fusor-ml/core/src/compute_graph/resolve/recognize_cat.rs new file mode 100644 index 000000000..82fec4ffa --- /dev/null +++ b/fusor-ml/core/src/compute_graph/resolve/recognize_cat.rs @@ -0,0 +1,588 @@ +//! Recognize composed slice-assign chains (`Tensor::cat`, sequential +//! `slice_assign` calls) while the graph is still in canonical form — before +//! view folding smears the narrow offsets into index arithmetic — and rewrite +//! each chain into a single elementwise kernel over the destination index +//! space. +//! +//! Every branch expression lifts into its region's select arm with chunk +//! coordinates rewritten to destination coordinates. Reads through `narrow` +//! views compose at the integer level (`AffineIndex`), so an aligned +//! split + op + cat cancels exactly: `(c - slice_start) + narrow_start = c`. +//! When all arms end up structurally identical and the slices tile the +//! destination, the selects disappear entirely and the chain becomes the op +//! applied to the larger tensor. + +use std::ops::Range; + +use crate::slice_assign::{slice_assign_expression, slice_region_condition}; +use crate::view::{AffineIndex, affine_dim_indices}; + +use super::*; + +/// Every destination element evaluates all lifted arms (both sides of a +/// select execute), so a K-way lift multiplies branch arithmetic by K. +/// Chains longer than this keep their materialized branches and fuse through +/// the regular n-ary path instead. +const MAX_LIFTED_BRANCHES: usize = 4; +/// Per-arm expression size cap for the non-collapsed (select chain) form. +const MAX_ARM_OPS: usize = 32; + +/// One matched slice-assign link: an `Elementwise` node whose expression is +/// exactly `slice_assign_expression(slices)` over `[destination, value]`. +struct AssignLink { + exec: ExecutionNodeIndex, + inner: NodeIndex, + destination: NodeIndex, + value: NodeIndex, + slices: Box<[Range]>, + shape: Box<[usize]>, + datatype: DataTypeEnum, +} + +/// Inputs of the rewritten kernel, deduplicated eagerly so that identical +/// branches produce identical arm expressions (slot-for-slot), which is what +/// the equal-arms collapse compares. +#[derive(Default)] +struct LiftState { + inputs: Vec, + slots: FxHashMap, + /// Whether any branch expression was inlined or any view was folded — + /// when nothing lifts, the rewrite would only replicate what n-ary + /// fusion already does, so the chain is left alone. + lifted: bool, +} + +impl LiftState { + fn slot(&mut self, inner: NodeIndex) -> usize { + if let Some(&slot) = self.slots.get(&inner) { + return slot; + } + let slot = self.inputs.len(); + self.inputs.push(inner); + self.slots.insert(inner, slot); + slot + } +} + +/// How a branch input slot is rewritten during chunk-space composition. +enum Rep { + /// Producer expression inlined at element-wise accesses. + Inline(NaryExpr), + /// Opaque input, remapped to a slot in the combined input list. + Slot(usize), +} + +fn apply_reps(expr: &NaryExpr, reps: &[Rep]) -> NaryExpr { + match expr { + NaryExpr::Op { children, function } => NaryExpr::Op { + children: children.iter().map(|c| apply_reps(c, reps)).collect(), + function: function.clone(), + }, + NaryExpr::IndexedInput { input_idx, indices } => { + let indices: Vec = indices.iter().map(|c| apply_reps(c, reps)).collect(); + match &reps[*input_idx] { + // Inline replacements are only built for slots accessed + // element-wise, where the whole load is the producer's value. + Rep::Inline(expr) => expr.clone(), + Rep::Slot(slot) => NaryExpr::IndexedInput { + input_idx: *slot, + indices, + }, + } + } + NaryExpr::DimIndex(dim) => NaryExpr::DimIndex(*dim), + NaryExpr::Scalar(value) => NaryExpr::Scalar(*value), + } +} + +/// Recover the slice ranges from a candidate slice-assign expression, then +/// verify the match by regenerating the canonical expression and comparing. +fn match_slice_assign(nary: &ElementwiseOperation) -> Option]>> { + if nary.inputs.len() != 2 { + return None; + } + let rank = nary.shape.len(); + if rank == 0 { + return None; + } + let NaryExpr::Op { children, function } = &nary.expression else { + return None; + }; + if !matches!(function.op, NaryOp::Select) || children.len() != 3 { + return None; + } + let mut starts = vec![None; rank]; + let mut ends = vec![None; rank]; + if !collect_region_bounds(&children[0], &mut starts, &mut ends) { + return None; + } + let slices: Box<[Range]> = starts + .iter() + .zip(&ends) + .map(|(&start, &end)| Some(start? as usize..end? as usize)) + .collect::>()?; + (slice_assign_expression(&slices, nary.output_datatype) == nary.expression).then_some(slices) +} + +fn collect_region_bounds( + expr: &NaryExpr, + starts: &mut [Option], + ends: &mut [Option], +) -> bool { + match expr { + NaryExpr::Scalar(NaryScalar::U32(1)) => true, + NaryExpr::Op { children, function } => match function.op { + NaryOp::Mul if children.len() == 2 => { + collect_region_bounds(&children[0], starts, ends) + && collect_region_bounds(&children[1], starts, ends) + } + NaryOp::GreaterEqualConst(NaryScalar::U32(value)) if children.len() == 1 => { + let NaryExpr::DimIndex(dim) = children[0] else { + return false; + }; + match starts.get_mut(dim) { + Some(slot) => { + *slot = Some(value); + true + } + None => false, + } + } + NaryOp::LessConst(NaryScalar::U32(value)) if children.len() == 1 => { + let NaryExpr::DimIndex(dim) = children[0] else { + return false; + }; + match ends.get_mut(dim) { + Some(slot) => { + *slot = Some(value); + true + } + None => false, + } + } + _ => false, + }, + _ => false, + } +} + +/// Whether the chain's slices cover every destination coordinate: all dims +/// full except at most one, whose ranges merge to the full extent. +fn slices_tile(out_shape: &[usize], chain: &[&AssignLink]) -> bool { + let mut partial_dims: Vec = Vec::new(); + for (dim, &extent) in out_shape.iter().enumerate() { + if chain.iter().any(|link| link.slices[dim] != (0..extent)) { + partial_dims.push(dim); + } + } + match partial_dims.as_slice() { + [] => true, + [dim] => { + let mut ranges: Vec> = + chain.iter().map(|link| link.slices[*dim].clone()).collect(); + ranges.sort_by_key(|range| range.start); + let mut covered = 0; + for range in ranges { + if range.start > covered { + return false; + } + covered = covered.max(range.end); + } + covered >= out_shape[*dim] + } + _ => false, + } +} + +fn expr_op_count(expr: &NaryExpr) -> usize { + match expr { + NaryExpr::Op { children, .. } => 1 + children.iter().map(expr_op_count).sum::(), + NaryExpr::IndexedInput { indices, .. } => { + 1 + indices.iter().map(expr_op_count).sum::() + } + NaryExpr::DimIndex(_) | NaryExpr::Scalar(_) => 0, + } +} + +impl Resolver { + pub(super) fn recognize_assign_chains(&mut self, graph: &mut ComputeGraphInner) { + if std::env::var_os("FUSOR_RESOLVE_SKIP_ASSIGN_CHAINS").is_some() { + return; + } + let mut links: FxHashMap = FxHashMap::default(); + for exec in self.execution_graph.node_indices() { + let node = &self.execution_graph[exec]; + let ExecutionVariant::Elementwise(nary) = &node.variant else { + continue; + }; + let Some(slices) = match_slice_assign(nary) else { + continue; + }; + links.insert( + node.inner_idx, + AssignLink { + exec, + inner: node.inner_idx, + destination: nary.inputs[0], + value: nary.inputs[1], + slices, + shape: nary.shape.clone(), + datatype: nary.output_datatype, + }, + ); + } + if links.is_empty() { + return; + } + + // A link is interior when its sole consumer is the next link of the + // same chain, reading it as the destination. Everything else is a + // chain tail (its consumers are ordinary readers of the cat result). + let tails: Vec = links + .values() + .filter(|link| !self.is_interior_link(&links, link)) + .map(|link| link.inner) + .collect(); + for tail in tails { + self.rewrite_assign_chain(graph, &links, tail); + } + } + + fn is_interior_link( + &self, + links: &FxHashMap, + link: &AssignLink, + ) -> bool { + let mut consumers = self + .execution_graph + .neighbors_directed(link.exec, petgraph::Direction::Outgoing); + let (Some(consumer), None) = (consumers.next(), consumers.next()) else { + return false; + }; + let consumer = &self.execution_graph[consumer]; + links.get(&consumer.inner_idx).is_some_and(|next| { + next.destination == link.inner + && next.shape == link.shape + && next.datatype == link.datatype + }) + } + + fn rewrite_assign_chain( + &mut self, + graph: &mut ComputeGraphInner, + links: &FxHashMap, + tail_inner: NodeIndex, + ) { + let tail = &links[&tail_inner]; + if !self.execution_graph.contains_node(tail.exec) { + return; + } + + // Walk destinations back to the chain base (base → tail order after + // the reverse). + let mut chain: Vec<&AssignLink> = vec![tail]; + loop { + let cur = chain.last().unwrap(); + let Some(prev) = links.get(&cur.destination) else { + break; + }; + if !self.execution_graph.contains_node(prev.exec) + || prev.shape != tail.shape + || prev.datatype != tail.datatype + || self + .execution_graph + .neighbors_directed(prev.exec, petgraph::Direction::Outgoing) + .count() + != 1 + { + break; + } + chain.push(prev); + } + chain.reverse(); + if chain.len() > MAX_LIFTED_BRANCHES { + return; + } + let base_inner = chain[0].destination; + + let out_shape = tail.shape.clone(); + let rank = out_shape.len(); + let mut state = LiftState::default(); + let base_slot = state.slot(base_inner); + let mut arms = Vec::with_capacity(chain.len()); + for link in &chain { + let condition = slice_region_condition(&link.slices); + let arm = self.lift_branch(graph, link, &condition, &out_shape, &mut state); + arms.push((condition, arm)); + } + if !state.lifted { + return; + } + + let collapsed = + arms.windows(2).all(|pair| pair[0].1 == pair[1].1) && slices_tile(&out_shape, &chain); + let expression = if collapsed { + arms.swap_remove(0).1 + } else { + if arms.iter().any(|(_, arm)| expr_op_count(arm) > MAX_ARM_OPS) { + return; + } + let mut expression = NaryExpr::input(base_slot, rank); + for (condition, arm) in arms { + expression = + NaryExpr::select(condition, arm, expression, DataTypeEnum::U32, tail.datatype); + } + expression + }; + + let (final_inputs, final_expression) = Self::deduplicate_inputs(state.inputs, expression); + if final_inputs.len() > graph.device().nary_direct_input_binding_budget() { + return; + } + + let new_nary = ElementwiseOperation { + inputs: final_inputs.clone(), + expression: final_expression, + shape: out_shape, + output_datatype: tail.datatype, + }; + let tail_exec = tail.exec; + self.execution_graph[tail_exec].variant = ExecutionVariant::Elementwise(new_nary); + + let incoming: Vec = self + .execution_graph + .neighbors_directed(tail_exec, petgraph::Direction::Incoming) + .collect(); + for &source in &incoming { + if let Some(edge) = self.execution_graph.find_edge(source, tail_exec) { + self.execution_graph.remove_edge(edge); + } + } + for &input in &final_inputs { + if let Some(exec) = self.get_input_node_in_exec_graph(input) + && self.execution_graph.contains_node(exec) + && self.execution_graph.find_edge(exec, tail_exec).is_none() + { + self.execution_graph.add_edge(exec, tail_exec, ()); + } + } + self.add_physical_dependencies(graph, tail_exec, &final_inputs); + for source in incoming { + self.remove_node_if_dead(source); + } + } + + /// Build one region's arm: compose the branch's elementwise cluster in + /// chunk coordinate space, then rewrite it into destination coordinates. + fn lift_branch( + &self, + graph: &ComputeGraphInner, + link: &AssignLink, + condition: &NaryExpr, + out_shape: &[usize], + state: &mut LiftState, + ) -> NaryExpr { + let chunk_shape: Box<[usize]> = link + .slices + .iter() + .map(|slice| slice.end - slice.start) + .collect(); + let chunk_expr = self.compose_chunk_expr(graph, link.value, &chunk_shape, state); + self.lift_expr(&chunk_expr, link, condition, out_shape, state, false) + } + + /// Inline the branch's elementwise producers within the shared chunk + /// index space. Producers stay opaque (a plain load) when they are + /// cached, shared, shaped differently, accessed with custom indices, or + /// not elementwise. + fn compose_chunk_expr( + &self, + graph: &ComputeGraphInner, + inner: NodeIndex, + chunk_shape: &[usize], + state: &mut LiftState, + ) -> NaryExpr { + let rank = chunk_shape.len(); + let inlinable = !self.check_cached(graph, inner) + && self + .get_input_node_in_exec_graph(inner) + .filter(|&exec| self.execution_graph.contains_node(exec)) + .is_some_and(|exec| { + matches!( + &self.execution_graph[exec].variant, + ExecutionVariant::Elementwise(nary) if *nary.shape == *chunk_shape + ) && self + .execution_graph + .neighbors_directed(exec, petgraph::Direction::Outgoing) + .count() + == 1 + }); + if !inlinable { + return NaryExpr::input(state.slot(inner), rank); + } + let exec = self.get_input_node_in_exec_graph(inner).unwrap(); + let ExecutionVariant::Elementwise(nary) = self.execution_graph[exec].variant.clone() else { + unreachable!("inlinable check matched an elementwise variant"); + }; + + let reps: Vec = nary + .inputs + .iter() + .enumerate() + .map(|(slot, &input)| { + if nary.expression.uses_custom_indexing_for_input(slot) { + Rep::Slot(state.slot(input)) + } else { + Rep::Inline(self.compose_chunk_expr(graph, input, chunk_shape, state)) + } + }) + .collect(); + state.lifted = true; + apply_reps(&nary.expression, &reps) + } + + /// Rewrite a chunk-space expression into destination coordinates. + /// `DimIndex(d)` in value position becomes the unguarded shifted + /// coordinate (wrapping arithmetic in dead lanes is discarded by the + /// region select); in index position it becomes the guarded form so + /// every load stays in bounds on both sides of the select. Element-wise + /// loads through affine views fold at the integer level when the shifted + /// map stays in bounds for every destination coordinate — that is where + /// aligned narrow offsets cancel. + fn lift_expr( + &self, + expr: &NaryExpr, + link: &AssignLink, + condition: &NaryExpr, + out_shape: &[usize], + state: &mut LiftState, + index_position: bool, + ) -> NaryExpr { + match expr { + NaryExpr::Op { children, function } => NaryExpr::Op { + children: children + .iter() + .map(|c| self.lift_expr(c, link, condition, out_shape, state, index_position)) + .collect(), + function: function.clone(), + }, + NaryExpr::IndexedInput { input_idx, indices } => { + if NaryExpr::is_elementwise_indices(indices) + && let Some(folded) = + self.try_fold_shifted_view(state.inputs[*input_idx], link, out_shape, state) + { + state.lifted = true; + return folded; + } + NaryExpr::IndexedInput { + input_idx: *input_idx, + indices: indices + .iter() + .map(|c| self.lift_expr(c, link, condition, out_shape, state, true)) + .collect(), + } + } + NaryExpr::DimIndex(dim) => { + if index_position { + self.guarded_coord(*dim, link, condition, out_shape) + } else { + Self::unguarded_coord(*dim, link) + } + } + NaryExpr::Scalar(value) => NaryExpr::Scalar(*value), + } + } + + /// The chunk coordinate `c_d - slice_start` as a raw value. Out-of-region + /// lanes wrap, which only feeds dead select arms. + fn unguarded_coord(dim: usize, link: &AssignLink) -> NaryExpr { + let start = link.slices[dim].start; + if start == 0 { + return NaryExpr::DimIndex(dim); + } + NaryExpr::unary_op( + NaryExpr::DimIndex(dim), + "slice_offset", + NaryOp::SubConst(NaryScalar::U32(start as u32)), + DataTypeEnum::U32, + DataTypeEnum::U32, + ) + } + + /// The chunk coordinate clamped to 0 outside the region, so loads through + /// it stay in bounds even in dead lanes. + fn guarded_coord( + &self, + dim: usize, + link: &AssignLink, + condition: &NaryExpr, + out_shape: &[usize], + ) -> NaryExpr { + let slice = &link.slices[dim]; + if slice.start == 0 && slice.end == out_shape[dim] { + return NaryExpr::DimIndex(dim); + } + NaryExpr::select( + condition.clone(), + Self::unguarded_coord(dim, link), + NaryExpr::scalar(NaryScalar::U32(0)), + DataTypeEnum::U32, + DataTypeEnum::U32, + ) + } + + /// Fold an element-wise load through a fully-defined affine view by + /// composing the chunk shift into the view's affine map: + /// `constant' = constant - Σ coeff·slice_start`. Only fires when the + /// shifted map provably stays inside the base for every destination + /// coordinate; aligned narrows compose to the identity and read the base + /// at the bare destination coordinates. + fn try_fold_shifted_view( + &self, + view_inner: NodeIndex, + link: &AssignLink, + out_shape: &[usize], + state: &mut LiftState, + ) -> Option { + let exec = self.get_input_node_in_exec_graph(view_inner)?; + if !self.execution_graph.contains_node(exec) { + return None; + } + let ExecutionVariant::View(view) = &self.execution_graph[exec].variant else { + return None; + }; + if !view.is_fully_defined() || view.shape().len() != link.slices.len() { + return None; + } + for (extent, slice) in view.shape().iter().zip(&*link.slices) { + if *extent != slice.end - slice.start { + return None; + } + } + let affine = affine_dim_indices(&view.layout, &view.input_shape)?; + + let mut shifted = Vec::with_capacity(affine.len()); + for (index, &extent) in affine.iter().zip(&*view.input_shape) { + let mut constant = index.constant as i64; + let mut max_offset = 0i64; + for &(dim, coefficient) in &index.terms { + constant -= coefficient as i64 * link.slices[dim].start as i64; + max_offset += coefficient as i64 * (out_shape[dim] as i64 - 1); + } + if constant < 0 || constant + max_offset >= extent as i64 { + return None; + } + shifted.push(AffineIndex { + constant: constant as u32, + terms: index.terms.clone(), + }); + } + + let base_slot = state.slot(view.input); + let coords: Vec = (0..out_shape.len()).map(NaryExpr::DimIndex).collect(); + Some(NaryExpr::IndexedInput { + input_idx: base_slot, + indices: shifted.iter().map(|index| index.to_expr(&coords)).collect(), + }) + } +} diff --git a/fusor-ml/core/src/compute_graph/tests.rs b/fusor-ml/core/src/compute_graph/tests.rs index e244e420f..b8c722aea 100644 --- a/fusor-ml/core/src/compute_graph/tests.rs +++ b/fusor-ml/core/src/compute_graph/tests.rs @@ -1,4 +1,4 @@ -use crate::{Device, Tensor}; +use crate::{Device, StrideSpec, Tensor}; // Build a small intermediate that requires a real kernel (not a Tensor input). // `x` materializes via `(input * 2.0) + 1.0`, which fuses to a single nary. @@ -201,3 +201,217 @@ fn auto_flush_resolves_pending_siblings() { } }); } + +// --- split + op + cat lowering (resolve/recognize_cat.rs) --- + +/// Narrow a 2D tensor along a dimension as a view. +fn narrow2(tensor: &Tensor, dim: usize, start: usize, length: usize) -> Tensor { + let shape = tensor.shape(); + let specs: Vec = (0..2) + .map(|i| { + if i == dim { + StrideSpec::dim(i, length).with_offset(start) + } else { + StrideSpec::dim(i, shape[i]) + } + }) + .collect(); + tensor.restride(specs) +} + +fn cat_test_input(device: &Device) -> (Tensor, Vec>) { + let rows: Vec> = (0..4) + .map(|r| (0..8).map(|c| (r * 8 + c) as f32 * 0.1).collect()) + .collect(); + (Tensor::new::(device, &rows), rows) +} + +#[test] +fn cat_of_same_op_collapses_to_op_on_whole_tensor() { + pollster::block_on(async { + let Ok(device) = Device::new().await else { + return; + }; + let (x, rows) = cat_test_input(&device); + + let first = narrow2(&x, 1, 0, 4).sin(); + let second = narrow2(&x, 1, 4, 4).sin(); + let dest = Tensor::splat(&device, 0.0f32, [4, 8]); + let out = dest + .slice_assign([0..4, 0..4], &first) + .slice_assign([0..4, 4..8], &second); + + let (_, kernels) = out.data.materialize(); + assert_eq!( + kernels, 1, + "same op over chunks tiling the tensor should collapse to one kernel" + ); + + let result = out.as_slice::<2, f32>().await.unwrap(); + for (r, row) in rows.iter().enumerate() { + for (c, &value) in row.iter().enumerate() { + assert!( + (result[[r, c]] - value.sin()).abs() < 1e-5, + "mismatch at [{r}, {c}]" + ); + } + } + }); +} + +#[test] +fn cat_of_different_ops_fuses_to_single_select_kernel() { + pollster::block_on(async { + let Ok(device) = Device::new().await else { + return; + }; + let (x, rows) = cat_test_input(&device); + + let first = narrow2(&x, 1, 0, 4).sin(); + let second = narrow2(&x, 1, 4, 4).cos(); + let dest = Tensor::splat(&device, 0.0f32, [4, 8]); + let out = dest + .slice_assign([0..4, 0..4], &first) + .slice_assign([0..4, 4..8], &second); + + let (_, kernels) = out.data.materialize(); + assert_eq!(kernels, 1, "branch ops should lift into the select arms"); + + let result = out.as_slice::<2, f32>().await.unwrap(); + for (r, row) in rows.iter().enumerate() { + for (c, &value) in row.iter().enumerate() { + let expected = if c < 4 { value.sin() } else { value.cos() }; + assert!( + (result[[r, c]] - expected).abs() < 1e-5, + "mismatch at [{r}, {c}]" + ); + } + } + }); +} + +#[test] +fn reordered_cat_with_op_fuses_to_single_kernel() { + pollster::block_on(async { + let Ok(device) = Device::new().await else { + return; + }; + let (x, rows) = cat_test_input(&device); + + // rotate_half: cat([-x2, x1], last_dim) + let negated_second = &narrow2(&x, 1, 4, 4) * -1.0; + let first = narrow2(&x, 1, 0, 4); + let dest = Tensor::splat(&device, 0.0f32, [4, 8]); + let out = dest + .slice_assign([0..4, 0..4], &negated_second) + .slice_assign([0..4, 4..8], &first); + + let (_, kernels) = out.data.materialize(); + assert_eq!(kernels, 1, "reordered cat should still fuse to one kernel"); + + let result = out.as_slice::<2, f32>().await.unwrap(); + for (r, row) in rows.iter().enumerate() { + for c in 0..8 { + let expected = if c < 4 { -row[c + 4] } else { row[c - 4] }; + assert!( + (result[[r, c]] - expected).abs() < 1e-5, + "mismatch at [{r}, {c}]" + ); + } + } + }); +} + +#[test] +fn partial_slice_assign_keeps_destination_outside_region() { + pollster::block_on(async { + let Ok(device) = Device::new().await else { + return; + }; + let (x, rows) = cat_test_input(&device); + + // Only cover the left half: the right half must keep the splat value. + let first = narrow2(&x, 1, 0, 4).sin(); + let dest = Tensor::splat(&device, 7.0f32, [4, 8]); + let out = dest.slice_assign([0..4, 0..4], &first); + + let (_, kernels) = out.data.materialize(); + assert_eq!(kernels, 1); + + let result = out.as_slice::<2, f32>().await.unwrap(); + for (r, row) in rows.iter().enumerate() { + for c in 0..8 { + let expected = if c < 4 { row[c].sin() } else { 7.0 }; + assert!( + (result[[r, c]] - expected).abs() < 1e-5, + "mismatch at [{r}, {c}]" + ); + } + } + }); +} + +#[test] +fn three_way_chunk_cat_collapses() { + pollster::block_on(async { + let Ok(device) = Device::new().await else { + return; + }; + let rows: Vec> = (0..6) + .map(|r| (0..4).map(|c| (r * 4 + c) as f32 * 0.1).collect()) + .collect(); + let x = Tensor::new::(&device, &rows); + + let dest = Tensor::splat(&device, 0.0f32, [6, 4]); + let mut out = dest; + for chunk in 0..3 { + let start = chunk * 2; + let part = &narrow2(&x, 0, start, 2) * 2.0; + out = out.slice_assign([start..start + 2, 0..4], &part); + } + + let (_, kernels) = out.data.materialize(); + assert_eq!(kernels, 1, "3-way same-op chunk cat should collapse"); + + let result = out.as_slice::<2, f32>().await.unwrap(); + for (r, row) in rows.iter().enumerate() { + for (c, &value) in row.iter().enumerate() { + assert!( + (result[[r, c]] - value * 2.0).abs() < 1e-5, + "mismatch at [{r}, {c}]" + ); + } + } + }); +} + +#[test] +fn deep_branch_chains_collapse_through_cat() { + pollster::block_on(async { + let Ok(device) = Device::new().await else { + return; + }; + let (x, rows) = cat_test_input(&device); + + // Multi-node branches: ((narrow * 2) + 1).sin() on each half. + let branch = |start: usize| ((&narrow2(&x, 1, start, 4) * 2.0) + 1.0).sin(); + let dest = Tensor::splat(&device, 0.0f32, [4, 8]); + let out = dest + .slice_assign([0..4, 0..4], &branch(0)) + .slice_assign([0..4, 4..8], &branch(4)); + + let (_, kernels) = out.data.materialize(); + assert_eq!(kernels, 1, "whole branch chains should inline and collapse"); + + let result = out.as_slice::<2, f32>().await.unwrap(); + for (r, row) in rows.iter().enumerate() { + for (c, &value) in row.iter().enumerate() { + let expected = (value * 2.0 + 1.0).sin(); + assert!( + (result[[r, c]] - expected).abs() < 1e-5, + "mismatch at [{r}, {c}]" + ); + } + } + }); +} diff --git a/fusor-ml/core/src/quantized/matmul/fallback.rs b/fusor-ml/core/src/quantized/matmul/fallback.rs deleted file mode 100644 index 835a6bd01..000000000 --- a/fusor-ml/core/src/quantized/matmul/fallback.rs +++ /dev/null @@ -1,239 +0,0 @@ -use super::*; - -impl QMatMulOperation { - pub(crate) fn build_direct_kernels( - &self, - graph: &crate::compute_graph::ComputeGraphInner, - workgroup_shape: &crate::mir::workgroup_shape::WorkgroupShape, - inputs: &[MirValue], - ) -> Result { - if inputs - .last() - .and_then(MirValue::as_tensor) - .is_some_and(|output| output.layout().shape().contains(&0)) - { - return Ok(QMatMulKernelPlan::EmptyOutput); - } - - if let Some(kernel) = self.build_direct_kernel(graph, workgroup_shape, inputs) { - return Ok(QMatMulKernelPlan::Kernels(vec![kernel])); - } - - self.build_dequantize_dense_fallback_direct_kernels(graph, inputs) - .and_then(QMatMulKernelPlan::from_kernels) - .ok_or_else(|| QMatMulLoweringError::new(self.name())) - } - - pub(super) fn build_dense_direct_kernel( - &self, - graph: &crate::compute_graph::ComputeGraphInner, - input: &TensorData, - matrix: &QMatrix, - output: &TensorData, - ) -> Option { - let [n, k] = matrix.shape() else { - return None; - }; - let (n, k) = (*n, *k); - let input_shape = input.layout().shape(); - let rank = input_shape.len(); - if rank < 2 { - return None; - } - let mut dense_shape = input_shape.to_vec(); - dense_shape[rank - 2] = k; - dense_shape[rank - 1] = n; - let mut dense_strides = vec![0; rank]; - dense_strides[rank - 2] = 1; - dense_strides[rank - 1] = k; - let matrix_datatype = match matrix.datatype() { - GgmlType::F32 => DataTypeEnum::F32, - GgmlType::F16 => DataTypeEnum::F16, - _ => return None, - }; - if input.datatype() != matrix_datatype || output.datatype() != matrix_datatype { - return None; - } - let dense_weight_t = TensorData::new_from_parts( - matrix.device(), - matrix.buffer().clone(), - Layout::from_parts( - 0, - dense_shape.into_boxed_slice(), - dense_strides.into_boxed_slice(), - ), - matrix_datatype, - ); - let device = graph.device(); - let dense_matmul = MatMulOperation::new( - matrix_datatype, - self.input, - self.input, - input.layout().shape(), - dense_weight_t.layout().shape(), - None, - &device, - ); - dense_matmul.build_direct_kernel( - graph, - &dense_matmul - .workgroup_shape_constraints(&device) - .solve(device.max_subgroup_size(), &device.limits())?, - &[ - input.clone().into(), - dense_weight_t.into(), - output.clone().into(), - ], - ) - } - - fn build_dense_qmatmul_fallback_direct_kernels( - &self, - graph: &crate::compute_graph::ComputeGraphInner, - input: &TensorData, - matrix: &QMatrix, - output: &TensorData, - ) -> Option> { - if input.datatype() != output.datatype() { - return None; - } - if matches!(matrix.datatype(), GgmlType::F32 | GgmlType::F16) { - return self - .build_dense_direct_kernel(graph, input, matrix, output) - .map(|kernel| vec![kernel]); - } - - let dense_weight = - TensorData::new_for_shape(&graph.device(), matrix.shape(), DataTypeEnum::F32); - let dequantize = DequantizeOperation::new(matrix.clone(), DataTypeEnum::F32); - let dequantize_inputs = vec![matrix.clone().into(), dense_weight.clone().into()]; - let dequantize_workgroup = dequantize - .workgroup_shape_constraints(&graph.device()) - .solve(graph.device().max_subgroup_size(), &graph.device().limits())?; - let dequantize_kernel = - dequantize.build_direct_kernel(graph, &dequantize_workgroup, &dequantize_inputs)?; - let dense_matrix = QMatrix { - device: graph.device(), - shape: matrix.shape.clone(), - buffer: dense_weight.buffer().clone(), - datatype: GgmlType::F32, - storage_layout: QMatrixStorageLayout::Native, - direct_pipeline_cache: matrix.direct_pipeline_cache.clone(), - }; - let matmul_kernel = self.build_dense_direct_kernel(graph, input, &dense_matrix, output)?; - Some(vec![dequantize_kernel, matmul_kernel]) - } - - fn build_nary_fallback_direct_kernel( - &self, - graph: &crate::compute_graph::ComputeGraphInner, - expression: NaryExpr, - inputs: &[&TensorData], - output: &TensorData, - shape: &[usize], - output_datatype: DataTypeEnum, - ) -> Option { - let operation = ElementwiseOperation { - inputs: (0..inputs.len()).map(NodeIndex::new).collect(), - expression, - shape: shape.into(), - output_datatype, - }; - let mut mir_inputs = inputs - .iter() - .map(|input| (*input).clone().into()) - .collect::>(); - mir_inputs.push(output.clone().into()); - let workgroup_shape = operation - .workgroup_shape_constraints(&graph.device()) - .solve(graph.device().max_subgroup_size(), &graph.device().limits())?; - operation.build_direct_kernel(graph, &workgroup_shape, &mir_inputs) - } - - fn build_dequantize_dense_fallback_direct_kernels( - &self, - graph: &crate::compute_graph::ComputeGraphInner, - inputs: &[MirValue], - ) -> Option> { - if !self.post_accumulator_offsets.is_empty() { - return None; - } - let [input, matrix, rest @ .., output] = inputs else { - return None; - }; - let mut input = input.as_tensor()?.clone(); - let MirValue::QMatrix(matrix) = matrix else { - return None; - }; - let output = output.as_tensor()?.clone(); - - let pre_extra_count = self - .pre_element_wise_expr - .as_ref() - .map(|epilogue| epilogue.extras.len()) - .unwrap_or(0); - let post_extra_count = self - .post_element_wise_expr - .as_ref() - .map(|epilogue| epilogue.extras.len()) - .unwrap_or(0); - if rest.len() != pre_extra_count + post_extra_count { - return None; - } - let extra_tensors = rest - .iter() - .map(MirValue::as_tensor) - .collect::>>()?; - let pre_extra_tensors = &extra_tensors[..pre_extra_count]; - let post_extra_tensors = &extra_tensors[pre_extra_count..]; - - let mut kernels = Vec::new(); - if let Some(pre) = &self.pre_element_wise_expr { - let pre_output = TensorData::new_for_shape( - &graph.device(), - input.layout().shape(), - pre.output_datatype, - ); - let mut nary_inputs = Vec::with_capacity(1 + pre_extra_tensors.len()); - nary_inputs.push(&input); - nary_inputs.extend(pre_extra_tensors.iter().copied()); - kernels.push(self.build_nary_fallback_direct_kernel( - graph, - pre.expression.clone(), - &nary_inputs, - &pre_output, - input.layout().shape(), - pre.output_datatype, - )?); - input = pre_output; - } - - let matmul_output = if self.post_element_wise_expr.is_some() { - TensorData::new_for_shape(&graph.device(), output.layout().shape(), DataTypeEnum::F32) - } else { - output.clone() - }; - kernels.extend(self.build_dense_qmatmul_fallback_direct_kernels( - graph, - &input, - matrix, - &matmul_output, - )?); - - if let Some(post) = &self.post_element_wise_expr { - let mut nary_inputs = Vec::with_capacity(1 + post_extra_tensors.len()); - nary_inputs.push(&matmul_output); - nary_inputs.extend(post_extra_tensors.iter().copied()); - kernels.push(self.build_nary_fallback_direct_kernel( - graph, - post.expression.clone(), - &nary_inputs, - &output, - output.layout().shape(), - post.output_datatype, - )?); - } - - Some(kernels) - } -} diff --git a/fusor-ml/core/src/quantized/matmul/kernel.rs b/fusor-ml/core/src/quantized/matmul/kernel.rs index ccd08998c..442d5c891 100644 --- a/fusor-ml/core/src/quantized/matmul/kernel.rs +++ b/fusor-ml/core/src/quantized/matmul/kernel.rs @@ -13,6 +13,183 @@ enum QmatmulDirectTokens { } impl QMatMulOperation { + /// Lower this operation to its kernel plan. Recognition and epilogue + /// fusion only build operations the direct paths can lower (see + /// `supports_elementwise_epilogue_fusion`), so a `None` from every path + /// here is an invariant violation, not a recoverable state. + pub(crate) fn build_direct_kernels( + &self, + graph: &crate::compute_graph::ComputeGraphInner, + workgroup_shape: &crate::mir::workgroup_shape::WorkgroupShape, + inputs: &[MirValue], + ) -> Result { + if inputs + .last() + .and_then(MirValue::as_tensor) + .is_some_and(|output| output.layout().shape().contains(&0)) + { + return Ok(QMatMulKernelPlan::EmptyOutput); + } + + if let Some(kernels) = self.build_m_padded_kernels(graph, inputs) { + return Ok(QMatMulKernelPlan::Kernels(kernels)); + } + if let Some(kernel) = self.build_direct_kernel(graph, workgroup_shape, inputs) { + return Ok(QMatMulKernelPlan::Kernels(vec![kernel])); + } + Err(QMatMulLoweringError::new(self.name())) + } + + /// Lower an M-padded matmul: copy the activation into a zero-padded + /// scratch tensor (the same kernel a `resize` view lowers to) and run the + /// matmul over the padded views. The output buffer's slack rows were + /// allocated by `qmatmul_operation_inputs`, which makes the same + /// `m_pad_target` decision. + fn build_m_padded_kernels( + &self, + graph: &crate::compute_graph::ComputeGraphInner, + inputs: &[MirValue], + ) -> Option> { + let device = graph.device(); + let padded_m = self.m_pad_target(KernelDeviceCaps::from_device(&device))?; + // A pad target implies no epilogues: inputs are [input, matrix, output]. + let [input, MirValue::QMatrix(matrix), output] = inputs else { + return None; + }; + let input = input.as_tensor()?; + let output = output.as_tensor()?; + let in_shape = input.layout().shape(); + let out_shape = output.layout().shape(); + if in_shape.len() < 2 || out_shape.len() != in_shape.len() { + return None; + } + let m_axis = in_shape.len() - 2; + + let mut padded_in_shape = in_shape.to_vec(); + padded_in_shape[m_axis] = padded_m; + let scratch = TensorData::new_for_shape(&device, &padded_in_shape, input.datatype()); + let pad_copy = crate::view::ViewOperation { + input: self.input, + layout: Layout::from_parts( + 0, + padded_in_shape.into(), + Layout::continuous_strides(in_shape), + ), + input_shape: in_shape.into(), + defined: in_shape.into(), + fill: crate::view::zero_scalar(input.datatype()), + datatype: input.datatype(), + }; + let pad_workgroup = pad_copy + .workgroup_shape_constraints(&device) + .solve(device.max_subgroup_size(), &device.limits())?; + let pad_inputs = [input.clone().into(), scratch.clone().into()]; + let pad_kernel = pad_copy.build_direct_kernel(graph, &pad_workgroup, &pad_inputs)?; + + let mut padded_out_shape = out_shape.to_vec(); + padded_out_shape[m_axis] = padded_m; + let padded_out_layout = Layout::contiguous(&padded_out_shape); + let padded_bytes = + padded_out_shape.iter().product::() * output.datatype().element_size(); + if output.layout().offset() != 0 || (padded_bytes as u64) > output.buffer().size() { + return None; + } + let padded_output = TensorData::new_from_parts( + &device, + output.buffer().clone(), + padded_out_layout, + output.datatype(), + ); + + let matmul = if matches!(matrix.datatype(), GgmlType::F32 | GgmlType::F16) { + self.build_dense_direct_kernel(graph, &scratch, matrix, &padded_output)? + } else { + Self::direct_kernel_for_tensors( + &device, + DirectKernelTensors { + input: &scratch, + matrix, + pre_extra_tensors: &[], + post_extra_tensors: &[], + output: &padded_output, + }, + self.name(), + DirectKernelChains { + pre_expr: None, + post_expr: None, + post_accumulator_offsets: &[], + }, + Some((self, inputs)), + )? + }; + Some(vec![pad_kernel, matmul]) + } + + /// F32/F16 ggml storage is dense values: read the matrix buffer as a + /// transposed dense weight and run a regular matmul kernel. + fn build_dense_direct_kernel( + &self, + graph: &crate::compute_graph::ComputeGraphInner, + input: &TensorData, + matrix: &QMatrix, + output: &TensorData, + ) -> Option { + let [n, k] = matrix.shape() else { + return None; + }; + let (n, k) = (*n, *k); + let input_shape = input.layout().shape(); + let rank = input_shape.len(); + if rank < 2 { + return None; + } + let mut dense_shape = input_shape.to_vec(); + dense_shape[rank - 2] = k; + dense_shape[rank - 1] = n; + let mut dense_strides = vec![0; rank]; + dense_strides[rank - 2] = 1; + dense_strides[rank - 1] = k; + let matrix_datatype = match matrix.datatype() { + GgmlType::F32 => DataTypeEnum::F32, + GgmlType::F16 => DataTypeEnum::F16, + _ => return None, + }; + if input.datatype() != matrix_datatype || output.datatype() != matrix_datatype { + return None; + } + let dense_weight_t = TensorData::new_from_parts( + matrix.device(), + matrix.buffer().clone(), + Layout::from_parts( + 0, + dense_shape.into_boxed_slice(), + dense_strides.into_boxed_slice(), + ), + matrix_datatype, + ); + let device = graph.device(); + let dense_matmul = MatMulOperation::new( + matrix_datatype, + self.input, + self.input, + input.layout().shape(), + dense_weight_t.layout().shape(), + None, + &device, + ); + dense_matmul.build_direct_kernel( + graph, + &dense_matmul + .workgroup_shape_constraints(&device) + .solve(device.max_subgroup_size(), &device.limits())?, + &[ + input.clone().into(), + dense_weight_t.into(), + output.clone().into(), + ], + ) + } + /// Build a direct quantized-matmul kernel for the supplied tensors. /// `pre_chain`/`post_chain` are pre- and post-element-wise unary chains /// to fuse into the kernel; pass `None` to skip. `operation_key` ties the @@ -638,7 +815,9 @@ impl Operation for QMatMulOperation { } fn inputs(&self, nodes: &crate::compute_graph::ComputeGraphInner) -> Vec { - let base = qmatmul_operation_inputs(self.input, &self.matrix, &self.out_shape, nodes); + let m_pad = self.m_pad_target(KernelDeviceCaps::from_device(&nodes.device())); + let base = + qmatmul_operation_inputs(self.input, &self.matrix, &self.out_shape, m_pad, nodes); let pre_extras = self .pre_element_wise_expr .as_ref() @@ -676,6 +855,38 @@ impl Operation for QMatMulOperation { return None; }; let output = inputs.last()?.as_tensor()?; + // A rank-1 activation is a single matrix row: lower it through the + // same [1, K] -> [1, N] views the rank-2 path uses. + let input_row; + let output_row; + let (input, output) = if input.layout().rank() == 1 && output.layout().rank() == 1 { + let in_len = input.layout().shape()[0]; + let out_len = output.layout().shape()[0]; + let out_stride = output.layout().strides()[0]; + input_row = TensorData::new_from_parts( + input.device(), + input.buffer().clone(), + Layout::from_parts( + input.layout().offset(), + Box::new([1, in_len]), + Box::new([0, input.layout().strides()[0]]), + ), + input.datatype(), + ); + output_row = TensorData::new_from_parts( + output.device(), + output.buffer().clone(), + Layout::from_parts( + output.layout().offset(), + Box::new([1, out_len]), + Box::new([out_len * out_stride, out_stride]), + ), + output.datatype(), + ); + (&input_row, &output_row) + } else { + (input, output) + }; let pre_extra_count = self .pre_element_wise_expr .as_ref() @@ -740,9 +951,9 @@ impl Operation for QMatMulOperation { } } if matches!(matrix.datatype(), GgmlType::F32 | GgmlType::F16) { - // The dense kernel has no epilogue slots; declining here routes - // fused epilogues to the dequantize+dense fallback, which applies - // them as separate elementwise kernels. + // The dense kernel has no epilogue slots; fusion never attaches + // epilogues to dense-storage operations (see + // `supports_elementwise_epilogue_fusion`). if self.pre_element_wise_expr.is_some() || self.post_element_wise_expr.is_some() { return None; } diff --git a/fusor-ml/core/src/quantized/matmul/mod.rs b/fusor-ml/core/src/quantized/matmul/mod.rs index 15a300a41..96fd9bb95 100644 --- a/fusor-ml/core/src/quantized/matmul/mod.rs +++ b/fusor-ml/core/src/quantized/matmul/mod.rs @@ -21,18 +21,15 @@ use crate::{ workgroup_shape::{Constraint, WorkgroupShapeConstraints}, }, nary_direct::{apply_multi_input_elementwise_expr, apply_single_input_elementwise_expr}, - nary_wise::{ElementwiseOperation, NaryExpr, NaryOp}, + nary_wise::{NaryExpr, NaryOp}, }; use fusor_gguf::GgmlType; use fusor_tile_ir as tile_ir; use fusor_tile_ir_kernels as tile_ir_kernels; use rustc_hash::FxHasher; -use super::{ - QMatMulDirectPipelineKey, QMatrix, QMatrixStorageLayout, dequantize::DequantizeOperation, -}; +use super::{QMatMulDirectPipelineKey, QMatrix, QMatrixStorageLayout}; -mod fallback; mod kernel; #[cfg(test)] mod tests; @@ -343,6 +340,10 @@ fn select_qmatmul_direct_variant( } fn matmul_m_size(shape: &[usize]) -> u32 { + // Rank-1 activations are a single matrix row. + if shape.len() < 2 { + return 1; + } shape[shape.len() - 2] as u32 } @@ -350,12 +351,34 @@ fn qmatmul_operation_inputs( input: NodeIndex, matrix: &QMatrix, out_shape: &[usize], + m_pad: Option, nodes: &crate::compute_graph::ComputeGraphInner, ) -> Vec { let input = nodes.get_result(input).unwrap(); let q_matrix = matrix.clone(); let device = input.device(); - let output_tensor = TensorData::new_for_shape(device, out_shape, input.datatype()); + let output_tensor = match m_pad { + // M-padded lowering writes `padded_m` rows: allocate the padded + // buffer but expose it cropped to the real `M` (the trailing rows + // are slack the kernel scribbles into). + Some(padded_m) => { + let m_axis = out_shape.len() - 2; + let mut padded_shape = out_shape.to_vec(); + padded_shape[m_axis] = padded_m; + let padded = TensorData::new_for_shape(device, &padded_shape, input.datatype()); + TensorData::new_from_parts( + device, + padded.buffer().clone(), + Layout::from_parts( + 0, + out_shape.into(), + Layout::continuous_strides(&padded_shape), + ), + input.datatype(), + ) + } + None => TensorData::new_for_shape(device, out_shape, input.datatype()), + }; vec![input.into(), q_matrix.into(), output_tensor.into()] } @@ -444,6 +467,44 @@ impl QMatMulOperation { matmul_m_size(&self.in_shape) } + /// The padded `M` this operation's lowering will use, when M-padding + /// applies: the activation is copied into a zero-padded scratch tensor so + /// the selector can pick a coop tile variant, and the output buffer is + /// over-allocated to hold the padded rows (exposed cropped to the real + /// `M`). Epilogue-carrying operations never pad — the padded views don't + /// carry epilogue extras, so fusion declines those ops instead (see + /// [`Self::supports_elementwise_epilogue_fusion`]). + pub(crate) fn m_pad_target(&self, caps: KernelDeviceCaps) -> Option { + if self.pre_element_wise_expr.is_some() + || self.post_element_wise_expr.is_some() + || !self.post_accumulator_offsets.is_empty() + { + return None; + } + if self.in_shape.len() < 2 { + return None; + } + let m = self.in_shape[self.in_shape.len() - 2]; + let n = self.out_shape[self.out_shape.len() - 1]; + qmatmul_m_pad_target_for_caps(m, n, caps) + } + + /// Whether the resolver may fold a general element-wise expression into + /// this operation's pre/post epilogue. Dense-storage (F32/F16) matrices + /// lower through the epilogue-less dense kernel, and an operation whose + /// lowering will M-pad keeps element-wise work as separate kernels. + pub(crate) fn supports_elementwise_epilogue_fusion(&self, device: &Device) -> bool { + if qmatrix_direct_quant_format(&self.matrix).is_none() { + return false; + } + if self.in_shape.len() < 2 { + return true; + } + let m = self.in_shape[self.in_shape.len() - 2]; + let n = self.out_shape[self.out_shape.len() - 1]; + qmatmul_m_pad_target_for_caps(m, n, KernelDeviceCaps::from_device(device)).is_none() + } + fn n_size(&self) -> u32 { self.out_shape[self.out_shape.len() - 1] as u32 } @@ -536,10 +597,6 @@ pub(crate) enum QMatMulKernelPlan { } impl QMatMulKernelPlan { - fn from_kernels(kernels: Vec) -> Option { - (!kernels.is_empty()).then_some(Self::Kernels(kernels)) - } - pub(crate) fn dispatch_count(&self) -> usize { match self { Self::EmptyOutput => 0, @@ -913,44 +970,9 @@ impl Tensor { DataTypeEnum::F16 | DataTypeEnum::F32 => {} DataTypeEnum::U32 => panic!("q_mat_mul requires f32/f16 tensors"), } - - if self.rank() < 2 { - return self.add_q_mat_mul(other); - } - let in_shape = self.shape(); - let m_axis = self.rank() - 2; - let m = in_shape[m_axis]; - let n = other.shape()[0]; - let Some(padded_m) = - qmatmul_m_pad_target_for_caps(m, n, KernelDeviceCaps::from_device(self.device())) - else { - return self.add_q_mat_mul(other); - }; - - // Build padded input shape: replace the M dim with padded_m. - let mut padded_shape = in_shape.to_vec(); - padded_shape[m_axis] = padded_m; - - // Resize writes zeros outside the copied region so the trailing - // `padded_m - m` rows contribute nothing to the dot product. - let padded_input = self.resize(padded_shape); - - // Run the aligned matmul. - let padded_out = padded_input.add_q_mat_mul(other); - - // Narrow the output back to the caller's M along dim R-2 via - // a restride view. All other dims are full-size, so this is a - // pure layout change (no copy). - let out_shape = padded_out.shape(); - let specs: Vec = (0..padded_out.rank()) - .map(|i| { - if i == m_axis { - crate::StrideSpec::dim(i, m) - } else { - crate::StrideSpec::dim(i, out_shape[i]) - } - }) - .collect(); - padded_out.restride(specs) + // M-padding for the coop tile variants is a lowering decision: the + // recognized operation pads its own activation scratch and output + // slack (see `QMatMulOperation::m_pad_target`). + self.add_q_mat_mul(other) } } diff --git a/fusor-ml/core/src/sampling/mod.rs b/fusor-ml/core/src/sampling/mod.rs index 5d2c5c094..e86cfb66d 100644 --- a/fusor-ml/core/src/sampling/mod.rs +++ b/fusor-ml/core/src/sampling/mod.rs @@ -3,18 +3,13 @@ use crate::{Device, Tensor, tensor::TensorData}; mod mirostat; mod pipeline; pub(crate) mod processors; -mod qmat_topk; mod standard_sampler; mod topk; #[cfg(test)] mod tests; -pub(crate) use pipeline::{ - mirostat2_sample_token_to_host, qmat_mirostat2_sample_lazy_token_pending, - qmat_mirostat2_sample_lazy_token_to_host, qmat_standard_sample_lazy_token_pending, - qmat_standard_sample_lazy_token_to_host, standard_sample_token_to_host, -}; +pub(crate) use pipeline::{GpuSamplerRequest, sample_token_pending, sample_token_to_host}; pub(crate) use topk::{ MergeSortedChunkTopKParams, chunk_top_k_pair_data_with_encoder, merge_sorted_chunk_top_k_pair_data_with_encoder, diff --git a/fusor-ml/core/src/sampling/pipeline.rs b/fusor-ml/core/src/sampling/pipeline.rs index 66961821c..889d500ac 100644 --- a/fusor-ml/core/src/sampling/pipeline.rs +++ b/fusor-ml/core/src/sampling/pipeline.rs @@ -1,6 +1,5 @@ use crate::{ Layout, Tensor, - quantized::QMatrix, tensor::{DataTypeEnum, LazyTensorData, TensorData}, }; use web_time::Instant; @@ -10,977 +9,282 @@ use wgpu::CommandEncoder; use super::{ GPU_SAMPLE_RESULT_WORDS, GPU_SAMPLE_STATUS_INVALID, GPU_SAMPLE_STATUS_RETRY_NEEDED, GPU_SAMPLE_STATUS_SAMPLED, GpuMirostat2Sampler, GpuMirostat2SamplerParams, - GpuStandardSamplerParams, PendingGpuSampledToken, TOP_K_CHUNK, + GpuStandardSamplerParams, PendingGpuSampledToken, TOP_K_CHUNK, min_top_k_candidates_per_chunk, mirostat::sample_from_sorted_top_k_data_with_encoder, - qmat_topk::{ - initial_sampler_candidate_count, next_sampler_candidate_count, - qmat_logits_data_with_encoder, sampler_output_per_chunk, - }, standard_sampler::sample_from_sorted_top_k_data_with_encoder as sample_standard_from_sorted_top_k_data_with_encoder, topk::{ ProcessorSettings, chunk_top_k_pair_data_with_processors_and_gpu_tail_with_encoder, - chunk_top_k_pair_data_with_processors_with_encoder, merge_sorted_chunk_top_k_pair_data_with_encoder, top_k_exactness_flag_data_with_encoder, }, }; -pub(crate) async fn mirostat2_sample_token_to_host( - input: &TensorData, - sampler: &mut GpuMirostat2Sampler, - previous_tokens: &[u32], - params: GpuMirostat2SamplerParams, -) -> Result, wgpu::BufferAsyncError> { - sample_processed_logits_to_host( - input, - sampler, - previous_tokens, - params, - None, - "mirostat2 sampled token download", - ) - .await +/// Which sampler kernel terminates the top-k tail, along with its parameters. +/// The chunked top-k, merge, and exactness stages are shared between kinds. +pub(crate) enum GpuSamplerRequest<'a> { + Mirostat2 { + sampler: &'a mut GpuMirostat2Sampler, + params: GpuMirostat2SamplerParams, + }, + Standard { + params: GpuStandardSamplerParams, + }, } -pub(crate) async fn qmat_mirostat2_sample_token_to_host( - hidden: &TensorData, - matrix: &QMatrix, - sampler: &mut GpuMirostat2Sampler, - previous_tokens: &[u32], - params: GpuMirostat2SamplerParams, -) -> Result, wgpu::BufferAsyncError> { - if hidden.datatype() != DataTypeEnum::F32 || hidden.layout().rank() != 1 { - return Ok(None); - } - let hidden_len = hidden.layout().shape()[0]; - let [vocab_len, hidden_matrix_len] = matrix.shape() else { - return Ok(None); - }; - if hidden_len != *hidden_matrix_len || *vocab_len == 0 { - return Ok(None); - } - if !hidden.device().is_same_device(matrix.device()) { - return Ok(None); +impl GpuSamplerRequest<'_> { + fn top_k(&self) -> usize { + match self { + Self::Mirostat2 { params, .. } => params.top_k, + Self::Standard { params } => params.top_k, + } } - let device = hidden.device(); - let mut encoder = - device - .wgpu_device() - .create_command_encoder(&wgpu::CommandEncoderDescriptor { - label: Some("qmat_mirostat2_sample_token_to_host encoder"), - }); - - let trace = cfg!(target_arch = "wasm32") - || std::env::var_os("FUSOR_TRACE_DECODE").is_some() - || std::env::var_os("FUSOR_TRACE_SAMPLER").is_some(); - let qmat_start = trace.then(Instant::now); - let Some(logits) = qmat_logits_data_with_encoder(hidden, matrix, &mut encoder) else { - return Ok(None); - }; - if let Some(start) = qmat_start { - tracing::info!( - "sampler_trace qmat_logits_setup elapsed={:?}", - start.elapsed() - ); + fn processor_settings(&self) -> ProcessorSettings { + match self { + Self::Mirostat2 { params, .. } => ProcessorSettings { + temperature: params.temperature, + repetition_penalty: params.repetition_penalty, + }, + Self::Standard { params } => ProcessorSettings { + temperature: params.temperature, + repetition_penalty: params.repetition_penalty, + }, + } } - let hidden_dump_buffer = if std::env::var_os("FUSOR_DEBUG_SAMPLER").is_some() { - let hidden_bytes = (std::mem::size_of::() * hidden_len) as u64; - let hidden_dl = device.wgpu_device().create_buffer(&wgpu::BufferDescriptor { - size: hidden_bytes, - usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ, - mapped_at_creation: false, - label: Some("debug sampler hidden download"), - }); - encoder.copy_buffer_to_buffer(hidden.buffer(), 0, &hidden_dl, 0, hidden_bytes); - Some(hidden_dl) - } else { - None - }; - - let result = sample_processed_logits_to_host( - &logits, - sampler, - previous_tokens, - params, - Some(encoder), - "qmat mirostat2 sampled token download", - ) - .await?; - - if result.is_none() - && let Some(hidden_dl) = hidden_dump_buffer - { - let (tx, rx) = futures_channel::oneshot::channel(); - hidden_dl - .slice(..) - .map_async(wgpu::MapMode::Read, move |r| { - _ = tx.send(r); - }); - #[cfg(not(target_arch = "wasm32"))] - device.poll_wait(); - let _ = rx.await; - let view = hidden_dl.slice(..).get_mapped_range(); - let hidden_vec: Vec = bytemuck::cast_slice(&view).to_vec(); - drop(view); - hidden_dl.unmap(); - - let mut nan = 0usize; - let mut pos_inf = 0usize; - let mut neg_inf = 0usize; - let mut finite = 0usize; - let mut min_h = f32::INFINITY; - let mut max_h = f32::NEG_INFINITY; - for &v in &hidden_vec { - if v.is_nan() { - nan += 1; - } else if v == f32::INFINITY { - pos_inf += 1; - } else if v == f32::NEG_INFINITY { - neg_inf += 1; - } else { - finite += 1; - if v < min_h { - min_h = v; - } - if v > max_h { - max_h = v; - } - } + fn encode_sample( + &mut self, + ids: &TensorData, + values: &TensorData, + exactness_flag: Option<&TensorData>, + encoder: &mut CommandEncoder, + ) -> Option { + match self { + Self::Mirostat2 { sampler, params } => sample_from_sorted_top_k_data_with_encoder( + ids, + values, + sampler, + *params, + exactness_flag, + Some(encoder), + ), + Self::Standard { params } => sample_standard_from_sorted_top_k_data_with_encoder( + ids, + values, + *params, + exactness_flag, + Some(encoder), + ), } - tracing::warn!( - "sampler_debug HIDDEN len={} nan={} +inf={} -inf={} finite={} min={} max={} first8={:?}", - hidden_vec.len(), - nan, - pos_inf, - neg_inf, - finite, - min_h, - max_h, - hidden_vec.iter().take(8).collect::>() - ); } - - Ok(result) } -pub(crate) async fn qmat_mirostat2_sample_lazy_token_to_host( - hidden: &LazyTensorData, - matrix: &QMatrix, - sampler: &mut GpuMirostat2Sampler, - previous_tokens: &[u32], - params: GpuMirostat2SamplerParams, -) -> Result, wgpu::BufferAsyncError> { - if hidden.info.datatype() != DataTypeEnum::F32 || hidden.info.rank() != 1 { - return Ok(None); - } - let hidden_len = hidden.info.shape()[0]; - let [vocab_len, hidden_matrix_len] = matrix.shape() else { - return Ok(None); - }; - if hidden_len != *hidden_matrix_len || *vocab_len == 0 { - return Ok(None); - } - if !hidden.device.is_same_device(matrix.device()) { - return Ok(None); - } - - let trace = cfg!(target_arch = "wasm32") - || std::env::var_os("FUSOR_TRACE_DECODE").is_some() - || std::env::var_os("FUSOR_TRACE_SAMPLER").is_some(); - let qmat_start = trace.then(Instant::now); - let (materialized_hidden, _, logits) = hidden.materialize_with_tail(|hidden_data, encoder| { - qmat_logits_data_with_encoder(hidden_data, matrix, encoder) - }); - if let Some(start) = qmat_start { - tracing::info!( - "sampler_trace qmat_logits_tail elapsed={:?}", - start.elapsed() - ); - } - - let Some(logits) = logits else { - return qmat_mirostat2_sample_token_to_host( - &materialized_hidden, - matrix, - sampler, - previous_tokens, - params, - ) - .await; - }; - - sample_processed_logits_to_host( - &logits, - sampler, - previous_tokens, - params, - None, - "qmat mirostat2 sampled token download", - ) - .await +#[derive(Clone, Copy)] +struct SampleAttemptDims { + top_k: usize, + candidate_count: usize, + chunks: usize, + input_len: usize, } -pub(crate) async fn standard_sample_token_to_host( - input: &TensorData, - previous_tokens: &[u32], - params: GpuStandardSamplerParams, -) -> Result, wgpu::BufferAsyncError> { - sample_processed_standard_logits_to_host( - input, - previous_tokens, - params, - None, - "standard sampled token download", - ) - .await +fn initial_sampler_candidate_count(top_k: usize, chunks: usize) -> usize { + top_k + .div_ceil(chunks) + .max(min_top_k_candidates_per_chunk()) + .min(top_k) + .min(TOP_K_CHUNK) } -pub(crate) async fn qmat_standard_sample_lazy_token_to_host( - hidden: &LazyTensorData, - matrix: &QMatrix, - previous_tokens: &[u32], - params: GpuStandardSamplerParams, -) -> Result, wgpu::BufferAsyncError> { - if hidden.info.datatype() != DataTypeEnum::F32 || hidden.info.rank() != 1 { - return Ok(None); - } - let hidden_len = hidden.info.shape()[0]; - let [vocab_len, hidden_matrix_len] = matrix.shape() else { - return Ok(None); - }; - if hidden_len != *hidden_matrix_len || *vocab_len == 0 { - return Ok(None); - } - if !hidden.device.is_same_device(matrix.device()) { - return Ok(None); - } - - let trace = cfg!(target_arch = "wasm32") - || std::env::var_os("FUSOR_TRACE_DECODE").is_some() - || std::env::var_os("FUSOR_TRACE_SAMPLER").is_some(); - let qmat_start = trace.then(Instant::now); - let (materialized_hidden, _, logits) = hidden.materialize_with_tail(|hidden_data, encoder| { - qmat_logits_data_with_encoder(hidden_data, matrix, encoder) - }); - if let Some(start) = qmat_start { - tracing::info!( - "sampler_trace qmat_standard_logits_tail elapsed={:?}", - start.elapsed() - ); +fn sampler_output_per_chunk(candidate_count: usize) -> usize { + if candidate_count >= TOP_K_CHUNK { + TOP_K_CHUNK + } else { + candidate_count + 1 } - - let Some(logits) = logits else { - return qmat_standard_sample_token_to_host( - &materialized_hidden, - matrix, - previous_tokens, - params, - ) - .await; - }; - - sample_processed_standard_logits_to_host( - &logits, - previous_tokens, - params, - None, - "qmat standard sampled token download", - ) - .await } -async fn qmat_standard_sample_token_to_host( - hidden: &TensorData, - matrix: &QMatrix, - previous_tokens: &[u32], - params: GpuStandardSamplerParams, -) -> Result, wgpu::BufferAsyncError> { - if hidden.datatype() != DataTypeEnum::F32 || hidden.layout().rank() != 1 { - return Ok(None); - } - let hidden_len = hidden.layout().shape()[0]; - let [vocab_len, hidden_matrix_len] = matrix.shape() else { - return Ok(None); - }; - if hidden_len != *hidden_matrix_len || *vocab_len == 0 { - return Ok(None); - } - if !hidden.device().is_same_device(matrix.device()) { - return Ok(None); - } - - let device = hidden.device(); - let mut encoder = - device - .wgpu_device() - .create_command_encoder(&wgpu::CommandEncoderDescriptor { - label: Some("qmat_standard_sample_token_to_host encoder"), - }); - let Some(logits) = qmat_logits_data_with_encoder(hidden, matrix, &mut encoder) else { - return Ok(None); - }; - sample_processed_standard_logits_to_host( - &logits, - previous_tokens, - params, - Some(encoder), - "qmat standard sampled token download", - ) - .await +fn next_sampler_candidate_count(candidate_count: usize, top_k: usize) -> usize { + candidate_count + .saturating_mul(2) + .min(top_k) + .min(TOP_K_CHUNK) } -pub(crate) fn qmat_mirostat2_sample_lazy_token_pending( - hidden: &LazyTensorData, - matrix: &QMatrix, - sampler: &mut GpuMirostat2Sampler, - previous_tokens: &[u32], - previous_gpu_token: Option<&Tensor>, - params: GpuMirostat2SamplerParams, -) -> Option { - if hidden.info.datatype() != DataTypeEnum::F32 || hidden.info.rank() != 1 { - return None; - } - let hidden_len = hidden.info.shape()[0]; - let [vocab_len, hidden_matrix_len] = matrix.shape() else { - return None; - }; - if hidden_len != *hidden_matrix_len || *vocab_len == 0 { - return None; - } - if !hidden.device.is_same_device(matrix.device()) { - return None; - } - - let previous_gpu_token = previous_gpu_token.and_then(materialize_gpu_previous_token); - if previous_gpu_token - .as_ref() - .is_some_and(|token| !hidden.device.is_same_device(token.device())) - { - return None; - } - - let qmat_start = (cfg!(target_arch = "wasm32") +fn sampler_trace_enabled() -> bool { + cfg!(target_arch = "wasm32") || std::env::var_os("FUSOR_TRACE_DECODE").is_some() - || std::env::var_os("FUSOR_TRACE_SAMPLER").is_some()) - .then(Instant::now); - let (materialized_hidden, _) = hidden.materialize(); - if let Some(start) = qmat_start { - tracing::info!( - "sampler_trace hidden_materialize_pending elapsed={:?}", - start.elapsed() - ); - } - - qmat_mirostat2_sample_token_pending( - &materialized_hidden, - matrix, - sampler, - previous_tokens, - previous_gpu_token.as_ref(), - params, - ) + || std::env::var_os("FUSOR_TRACE_SAMPLER").is_some() } -pub(crate) fn qmat_standard_sample_lazy_token_pending( - hidden: &LazyTensorData, - matrix: &QMatrix, - previous_tokens: &[u32], - previous_gpu_token: Option<&Tensor>, - params: GpuStandardSamplerParams, -) -> Option { - if hidden.info.datatype() != DataTypeEnum::F32 || hidden.info.rank() != 1 { - return None; - } - let hidden_len = hidden.info.shape()[0]; - let [vocab_len, hidden_matrix_len] = matrix.shape() else { - return None; - }; - if hidden_len != *hidden_matrix_len || *vocab_len == 0 { - return None; - } - if !hidden.device.is_same_device(matrix.device()) { - return None; - } - - let previous_gpu_token = previous_gpu_token.and_then(materialize_gpu_previous_token); - if previous_gpu_token - .as_ref() - .is_some_and(|token| !hidden.device.is_same_device(token.device())) - { - return None; - } - - let qmat_start = (cfg!(target_arch = "wasm32") - || std::env::var_os("FUSOR_TRACE_DECODE").is_some() - || std::env::var_os("FUSOR_TRACE_SAMPLER").is_some()) - .then(Instant::now); - let (materialized_hidden, _) = hidden.materialize(); - if let Some(start) = qmat_start { - tracing::info!( - "sampler_trace standard_hidden_materialize_pending elapsed={:?}", - start.elapsed() - ); - } - - qmat_standard_sample_token_pending( - &materialized_hidden, - matrix, - previous_tokens, - previous_gpu_token.as_ref(), - params, - ) -} - -fn qmat_mirostat2_sample_token_pending( - hidden: &TensorData, - matrix: &QMatrix, - sampler: &mut GpuMirostat2Sampler, - previous_tokens: &[u32], - previous_gpu_token: Option<&TensorData>, - params: GpuMirostat2SamplerParams, -) -> Option { - if hidden.datatype() != DataTypeEnum::F32 || hidden.layout().rank() != 1 { - return None; - } - let hidden_len = hidden.layout().shape()[0]; - let [vocab_len, hidden_matrix_len] = matrix.shape() else { - return None; - }; - if hidden_len != *hidden_matrix_len || *vocab_len == 0 { - return None; - } - if !hidden.device().is_same_device(matrix.device()) { - return None; - } - - let device = hidden.device(); - let mut encoder = - device - .wgpu_device() - .create_command_encoder(&wgpu::CommandEncoderDescriptor { - label: Some("qmat_mirostat2_sample_token_pending encoder"), - }); - let logits = qmat_logits_data_with_encoder(hidden, matrix, &mut encoder)?; - sample_processed_logits_pending( - &logits, - sampler, - previous_tokens, - previous_gpu_token, - params, - Some(encoder), - "qmat mirostat2 pending sampled token download", - ) -} - -fn qmat_standard_sample_token_pending( - hidden: &TensorData, - matrix: &QMatrix, +/// Encode one full sampling attempt — processed chunk top-k, merge, +/// optional exactness proof, and the sampler kernel — into `encoder`. +/// Returns the `[status, token]` output along with the sorted top-k +/// ids/values (kept for `FUSOR_DEBUG_SAMPLER` dumps). +/// +/// This runs inside the resolver tail while the graph lock is held, so it +/// must only touch raw buffers — no compute-graph access. +fn encode_sample_attempt( + logits: &TensorData, previous_tokens: &[u32], previous_gpu_token: Option<&TensorData>, - params: GpuStandardSamplerParams, -) -> Option { - if hidden.datatype() != DataTypeEnum::F32 || hidden.layout().rank() != 1 { - return None; - } - let hidden_len = hidden.layout().shape()[0]; - let [vocab_len, hidden_matrix_len] = matrix.shape() else { - return None; - }; - if hidden_len != *hidden_matrix_len || *vocab_len == 0 { - return None; - } - if !hidden.device().is_same_device(matrix.device()) { - return None; - } + request: &mut GpuSamplerRequest<'_>, + dims: SampleAttemptDims, + encoder: &mut CommandEncoder, +) -> Option<(TensorData, TensorData, TensorData)> { + let output_per_chunk = sampler_output_per_chunk(dims.candidate_count); + let (chunk_ids, chunk_values) = + chunk_top_k_pair_data_with_processors_and_gpu_tail_with_encoder( + logits, + previous_tokens, + previous_gpu_token, + request.processor_settings(), + dims.candidate_count, + output_per_chunk, + Some(encoder), + )?; - let device = hidden.device(); - let mut encoder = - device - .wgpu_device() - .create_command_encoder(&wgpu::CommandEncoderDescriptor { - label: Some("qmat_standard_sample_token_pending encoder"), - }); - let logits = qmat_logits_data_with_encoder(hidden, matrix, &mut encoder)?; - sample_processed_standard_logits_pending( - &logits, - previous_tokens, - previous_gpu_token, - params, + let (ids, values) = merge_sorted_chunk_top_k_pair_data_with_encoder( + &chunk_ids, + &chunk_values, + crate::sampling::topk::MergeSortedChunkTopKParams { + chunks: dims.chunks, + chunk_len: dims.candidate_count, + chunk_stride: output_per_chunk, + input_len: dims.input_len, + k: dims.top_k, + }, Some(encoder), - "qmat standard pending sampled token download", - ) -} + )?; -fn materialize_gpu_previous_token(token: &Tensor) -> Option { - if token.datatype() != DataTypeEnum::U32 - || token.rank() != 1 - || token.shape().first().copied().unwrap_or_default() == 0 + let exactness_flag = if dims.candidate_count < dims.top_k && dims.candidate_count < TOP_K_CHUNK { - return None; - } - let (data, _) = token.data.materialize(); - Some(data) -} - -fn sample_processed_logits_pending( - input: &TensorData, - sampler: &mut GpuMirostat2Sampler, - previous_tokens: &[u32], - previous_gpu_token: Option<&TensorData>, - params: GpuMirostat2SamplerParams, - mut initial_encoder: Option, - download_label: &'static str, -) -> Option { - if input.datatype() != DataTypeEnum::F32 || input.layout().rank() != 1 { - return None; - } - - let input_len = input.layout().shape()[0]; - let top_k = params.top_k.min(input_len); - if top_k == 0 { - return None; - } - - let chunks = input_len.div_ceil(TOP_K_CHUNK); - if top_k > TOP_K_CHUNK { - return None; - } - let candidate_count = top_k; - let (output, _, _, encoder) = build_sample_attempt( - input, - sampler, - previous_tokens, - PreviousTokenSource::GpuTail(previous_gpu_token), - params, - SampleAttemptConfig { - top_k, - chunks, - input_len, - candidate_count, - trace: cfg!(target_arch = "wasm32"), - encoder_label: "mirostat2_sample_token_pending encoder", - }, - &mut initial_encoder, - )?; + Some(top_k_exactness_flag_data_with_encoder( + &values, + &chunk_values, + dims.chunks, + dims.candidate_count, + output_per_chunk, + dims.top_k, + Some(encoder), + )?) + } else { + None + }; - Some(submit_pending_sample_output( - &output, - encoder, - download_label, - )) + let output = request.encode_sample(&ids, &values, exactness_flag.as_ref(), encoder)?; + Some((output, ids, values)) } -fn submit_pending_sample_output( +fn encode_token_download( output: &TensorData, - mut encoder: CommandEncoder, - download_label: &'static str, -) -> PendingGpuSampledToken { + encoder: &mut CommandEncoder, + label: &'static str, +) -> wgpu::Buffer { let device = output.device(); let download_size = (std::mem::size_of::() * GPU_SAMPLE_RESULT_WORDS) as u64; let download = device.wgpu_device().create_buffer(&wgpu::BufferDescriptor { size: download_size, usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ, mapped_at_creation: false, - label: Some(download_label), + label: Some(label), }); encoder.copy_buffer_to_buffer(output.buffer(), 0, &download, 0, download_size); - device.wgpu_queue().submit(Some(encoder.finish())); + download +} +async fn read_sample_result( + device: &crate::Device, + download: wgpu::Buffer, + trace: bool, +) -> Result<(u32, u32), wgpu::BufferAsyncError> { + #[cfg(target_arch = "wasm32")] + let _ = device; + let map_start = trace.then(Instant::now); let (sender, receiver) = futures_channel::oneshot::channel(); download .slice(..) .map_async(wgpu::MapMode::Read, move |result| { _ = sender.send(result); }); - - let token = Tensor::from(TensorData::new_from_parts( - device, - output.buffer().clone(), - Layout::from_parts(1, Box::new([1]), Box::new([1])), - DataTypeEnum::U32, - )); - PendingGpuSampledToken::new(token, download, receiver) -} - -enum PreviousTokenSource<'a> { - HostOnly, - GpuTail(Option<&'a TensorData>), -} - -struct SampleAttemptConfig { - top_k: usize, - chunks: usize, - input_len: usize, - candidate_count: usize, - trace: bool, - encoder_label: &'static str, + #[cfg(not(target_arch = "wasm32"))] + device.poll_wait(); + receiver.await.map_err(|_| wgpu::BufferAsyncError)??; + if let Some(start) = map_start { + tracing::info!("sampler_trace map_wait elapsed={:?}", start.elapsed()); + } + + let view = download.slice(..).get_mapped_range(); + let word_size = std::mem::size_of::(); + let status = view + .get(..word_size) + .map(bytemuck::from_bytes::) + .copied() + .unwrap_or(GPU_SAMPLE_STATUS_INVALID); + let token = view + .get(word_size..word_size * GPU_SAMPLE_RESULT_WORDS) + .map(bytemuck::from_bytes::) + .copied() + .unwrap_or_default(); + drop(view); + download.unmap(); + Ok((status, token)) } -fn build_sample_attempt( - input: &TensorData, - sampler: &mut GpuMirostat2Sampler, +/// Sample a token from a lazy 1-D logits tensor, downloading the result. +/// +/// The first attempt's kernels ride the resolver's command encoder, so the +/// graph that produces the logits (including a recognized lm_head qgemv) and +/// the sampling tail land in a single submission. Retries with escalated +/// candidate counts re-encode over the already-materialized logits. +pub(crate) async fn sample_token_to_host( + logits: &LazyTensorData, + mut request: GpuSamplerRequest<'_>, previous_tokens: &[u32], - previous_token_source: PreviousTokenSource<'_>, - params: GpuMirostat2SamplerParams, - config: SampleAttemptConfig, - initial_encoder: &mut Option, -) -> Option<(TensorData, TensorData, TensorData, CommandEncoder)> { - let device = input.device(); - let mut encoder = initial_encoder.take().unwrap_or_else(|| { - device - .wgpu_device() - .create_command_encoder(&wgpu::CommandEncoderDescriptor { - label: Some(config.encoder_label), - }) - }); - - let output_per_chunk = sampler_output_per_chunk(config.candidate_count); - let topk_start = config.trace.then(Instant::now); - let (chunk_ids, chunk_values) = match previous_token_source { - PreviousTokenSource::HostOnly => chunk_top_k_pair_data_with_processors_with_encoder( - input, - previous_tokens, - ProcessorSettings { - temperature: params.temperature, - repetition_penalty: params.repetition_penalty, - }, - config.candidate_count, - output_per_chunk, - Some(&mut encoder), - )?, - PreviousTokenSource::GpuTail(previous_gpu_token) => { - chunk_top_k_pair_data_with_processors_and_gpu_tail_with_encoder( - input, - previous_tokens, - previous_gpu_token, - ProcessorSettings { - temperature: params.temperature, - repetition_penalty: params.repetition_penalty, - }, - config.candidate_count, - output_per_chunk, - Some(&mut encoder), - )? - } - }; - if let Some(start) = topk_start { - tracing::info!("sampler_trace topk_setup elapsed={:?}", start.elapsed()); - } - - let merge_start = config.trace.then(Instant::now); - let (ids, values) = merge_sorted_chunk_top_k_pair_data_with_encoder( - &chunk_ids, - &chunk_values, - crate::sampling::topk::MergeSortedChunkTopKParams { - chunks: config.chunks, - chunk_len: config.candidate_count, - chunk_stride: output_per_chunk, - input_len: config.input_len, - k: config.top_k, - }, - Some(&mut encoder), - )?; - if let Some(start) = merge_start { - tracing::info!("sampler_trace merge_setup elapsed={:?}", start.elapsed()); - } - - let exactness_start = config.trace.then(Instant::now); - let exactness_flag = - if config.candidate_count < config.top_k && config.candidate_count < TOP_K_CHUNK { - Some(top_k_exactness_flag_data_with_encoder( - &values, - &chunk_values, - config.chunks, - config.candidate_count, - output_per_chunk, - config.top_k, - Some(&mut encoder), - )?) - } else { - None - }; - if let Some(start) = exactness_start { - tracing::info!( - "sampler_trace exactness_setup elapsed={:?}", - start.elapsed() - ); - } - - let sample_start = config.trace.then(Instant::now); - let output = sample_from_sorted_top_k_data_with_encoder( - &ids, - &values, - sampler, - params, - exactness_flag.as_ref(), - Some(&mut encoder), - )?; - if let Some(start) = sample_start { - tracing::info!("sampler_trace sample_setup elapsed={:?}", start.elapsed()); - } - - Some((output, ids, values, encoder)) -} - -fn build_standard_sample_attempt( - input: &TensorData, - previous_tokens: &[u32], - previous_token_source: PreviousTokenSource<'_>, - params: GpuStandardSamplerParams, - config: SampleAttemptConfig, - initial_encoder: &mut Option, -) -> Option<(TensorData, CommandEncoder)> { - let device = input.device(); - let mut encoder = initial_encoder.take().unwrap_or_else(|| { - device - .wgpu_device() - .create_command_encoder(&wgpu::CommandEncoderDescriptor { - label: Some(config.encoder_label), - }) - }); - - let output_per_chunk = sampler_output_per_chunk(config.candidate_count); - let topk_start = config.trace.then(Instant::now); - let (chunk_ids, chunk_values) = match previous_token_source { - PreviousTokenSource::HostOnly => chunk_top_k_pair_data_with_processors_with_encoder( - input, - previous_tokens, - ProcessorSettings { - temperature: params.temperature, - repetition_penalty: params.repetition_penalty, - }, - config.candidate_count, - output_per_chunk, - Some(&mut encoder), - )?, - PreviousTokenSource::GpuTail(previous_gpu_token) => { - chunk_top_k_pair_data_with_processors_and_gpu_tail_with_encoder( - input, - previous_tokens, - previous_gpu_token, - ProcessorSettings { - temperature: params.temperature, - repetition_penalty: params.repetition_penalty, - }, - config.candidate_count, - output_per_chunk, - Some(&mut encoder), - )? - } - }; - if let Some(start) = topk_start { - tracing::info!( - "sampler_trace standard_topk_setup elapsed={:?}", - start.elapsed() - ); - } - - let merge_start = config.trace.then(Instant::now); - let (ids, values) = merge_sorted_chunk_top_k_pair_data_with_encoder( - &chunk_ids, - &chunk_values, - crate::sampling::topk::MergeSortedChunkTopKParams { - chunks: config.chunks, - chunk_len: config.candidate_count, - chunk_stride: output_per_chunk, - input_len: config.input_len, - k: config.top_k, - }, - Some(&mut encoder), - )?; - if let Some(start) = merge_start { - tracing::info!( - "sampler_trace standard_merge_setup elapsed={:?}", - start.elapsed() - ); - } - - let exactness_start = config.trace.then(Instant::now); - let exactness_flag = - if config.candidate_count < config.top_k && config.candidate_count < TOP_K_CHUNK { - Some(top_k_exactness_flag_data_with_encoder( - &values, - &chunk_values, - config.chunks, - config.candidate_count, - output_per_chunk, - config.top_k, - Some(&mut encoder), - )?) - } else { - None - }; - if let Some(start) = exactness_start { - tracing::info!( - "sampler_trace standard_exactness_setup elapsed={:?}", - start.elapsed() - ); - } - - let sample_start = config.trace.then(Instant::now); - let output = sample_standard_from_sorted_top_k_data_with_encoder( - &ids, - &values, - params, - exactness_flag.as_ref(), - Some(&mut encoder), - )?; - if let Some(start) = sample_start { - tracing::info!( - "sampler_trace standard_sample_setup elapsed={:?}", - start.elapsed() - ); - } - - Some((output, encoder)) -} - -fn sample_processed_standard_logits_pending( - input: &TensorData, - previous_tokens: &[u32], - previous_gpu_token: Option<&TensorData>, - params: GpuStandardSamplerParams, - mut initial_encoder: Option, - download_label: &'static str, -) -> Option { - if input.datatype() != DataTypeEnum::F32 || input.layout().rank() != 1 { - return None; - } - - let input_len = input.layout().shape()[0]; - let top_k = params.top_k.min(input_len); - if top_k == 0 { - return None; - } - - let chunks = input_len.div_ceil(TOP_K_CHUNK); - if top_k > TOP_K_CHUNK { - return None; - } - let candidate_count = top_k; - let (output, encoder) = build_standard_sample_attempt( - input, - previous_tokens, - PreviousTokenSource::GpuTail(previous_gpu_token), - params, - SampleAttemptConfig { - top_k, - chunks, - input_len, - candidate_count, - trace: cfg!(target_arch = "wasm32"), - encoder_label: "standard_sample_token_pending encoder", - }, - &mut initial_encoder, - )?; - - Some(submit_pending_sample_output( - &output, - encoder, - download_label, - )) -} - -async fn sample_processed_standard_logits_to_host( - input: &TensorData, - previous_tokens: &[u32], - params: GpuStandardSamplerParams, - mut initial_encoder: Option, - download_label: &'static str, ) -> Result, wgpu::BufferAsyncError> { - if input.datatype() != DataTypeEnum::F32 || input.layout().rank() != 1 { + if logits.info.datatype() != DataTypeEnum::F32 || logits.info.rank() != 1 { return Ok(None); } - - let input_len = input.layout().shape()[0]; - let top_k = params.top_k.min(input_len); + let input_len = logits.info.shape()[0]; + let top_k = request.top_k().min(input_len); if top_k == 0 { return Ok(None); } let chunks = input_len.div_ceil(TOP_K_CHUNK); let mut candidate_count = initial_sampler_candidate_count(top_k, chunks); - let trace = cfg!(target_arch = "wasm32") - || std::env::var_os("FUSOR_TRACE_DECODE").is_some() - || std::env::var_os("FUSOR_TRACE_SAMPLER").is_some(); - let mut attempt = 0usize; - loop { - attempt += 1; - let device = input.device(); - let Some((output, mut encoder)) = build_standard_sample_attempt( - input, + let trace = sampler_trace_enabled(); + let debug_dump = std::env::var_os("FUSOR_DEBUG_SAMPLER").is_some(); + + let (logits_data, _, mut attempt) = logits.materialize_with_tail(|logits_data, encoder| { + let attempt = encode_sample_attempt( + logits_data, previous_tokens, - PreviousTokenSource::HostOnly, - params, - SampleAttemptConfig { + None, + &mut request, + SampleAttemptDims { top_k, + candidate_count, chunks, input_len, - candidate_count, - trace, - encoder_label: "standard_sample_token_to_host encoder", }, - &mut initial_encoder, - ) else { + encoder, + )?; + let download = encode_token_download(&attempt.0, encoder, "sampled token download"); + Some((attempt, download)) + }); + let device = logits_data.device().clone(); + + let mut attempt_index = 0usize; + loop { + attempt_index += 1; + let Some(((_, ids, values), download)) = attempt else { return Ok(None); }; - - let download_size = (std::mem::size_of::() * GPU_SAMPLE_RESULT_WORDS) as u64; - let download = device.wgpu_device().create_buffer(&wgpu::BufferDescriptor { - size: download_size, - usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ, - mapped_at_creation: false, - label: Some(download_label), - }); - encoder.copy_buffer_to_buffer(output.buffer(), 0, &download, 0, download_size); - - let submit_start = trace.then(Instant::now); - device.wgpu_queue().submit(Some(encoder.finish())); - if let Some(start) = submit_start { - tracing::info!( - "sampler_trace standard_submit elapsed={:?}", - start.elapsed() - ); - } - - let map_start = trace.then(Instant::now); - let (sender, receiver) = futures_channel::oneshot::channel(); - download - .slice(..) - .map_async(wgpu::MapMode::Read, move |result| { - _ = sender.send(result); - }); - #[cfg(not(target_arch = "wasm32"))] - device.poll_wait(); - receiver.await.map_err(|_| wgpu::BufferAsyncError)??; - if let Some(start) = map_start { - tracing::info!( - "sampler_trace standard_map_wait elapsed={:?}", - start.elapsed() - ); - } - - let view = download.slice(..).get_mapped_range(); - let word_size = std::mem::size_of::(); - let status = view - .get(..word_size) - .map(bytemuck::from_bytes::) - .copied() - .unwrap_or(GPU_SAMPLE_STATUS_INVALID); - let token = view - .get(word_size..word_size * GPU_SAMPLE_RESULT_WORDS) - .map(bytemuck::from_bytes::) - .copied() - .unwrap_or_default(); - drop(view); - download.unmap(); - + let (status, token) = read_sample_result(&device, download, trace).await?; match status { GPU_SAMPLE_STATUS_SAMPLED => { if trace { tracing::info!( - "sampler_trace standard_sampled attempt={attempt} top_k={top_k} chunks={chunks} candidate_count={candidate_count} token={token}" + "sampler_trace sampled attempt={attempt_index} top_k={top_k} chunks={chunks} candidate_count={candidate_count} token={token}" ); } return Ok(Some(token)); @@ -988,16 +292,19 @@ async fn sample_processed_standard_logits_to_host( GPU_SAMPLE_STATUS_RETRY_NEEDED => { if trace { tracing::info!( - "sampler_trace standard_retry attempt={attempt} top_k={top_k} chunks={chunks} candidate_count={candidate_count}" + "sampler_trace retry attempt={attempt_index} top_k={top_k} chunks={chunks} candidate_count={candidate_count}" ); } } _ => { if trace { tracing::warn!( - "sampler_trace standard_invalid attempt={attempt} top_k={top_k} chunks={chunks} candidate_count={candidate_count} status={status}" + "sampler_trace invalid attempt={attempt_index} top_k={top_k} chunks={chunks} candidate_count={candidate_count} status={status}" ); } + if debug_dump { + debug_dump_invalid_sample(&logits_data, &ids, &values, previous_tokens).await; + } return Ok(None); } } @@ -1007,236 +314,166 @@ async fn sample_processed_standard_logits_to_host( return Ok(None); } candidate_count = next; + + let mut encoder = + device + .wgpu_device() + .create_command_encoder(&wgpu::CommandEncoderDescriptor { + label: Some("sampler retry encoder"), + }); + attempt = encode_sample_attempt( + &logits_data, + previous_tokens, + None, + &mut request, + SampleAttemptDims { + top_k, + candidate_count, + chunks, + input_len, + }, + &mut encoder, + ) + .map(|attempt| { + let download = + encode_token_download(&attempt.0, &mut encoder, "sampled token download"); + (attempt, download) + }); + device.wgpu_queue().submit(Some(encoder.finish())); } } -async fn sample_processed_logits_to_host( - input: &TensorData, - sampler: &mut GpuMirostat2Sampler, +/// Sample a token from a lazy 1-D logits tensor without waiting for the +/// result. The whole tail — including the optional copy of the previously +/// sampled GPU token into the repetition-penalty window — is appended to the +/// resolver's submission, and the token stays on the GPU as a 1-element +/// tensor for the next decode step. +pub(crate) fn sample_token_pending( + logits: &LazyTensorData, + mut request: GpuSamplerRequest<'_>, previous_tokens: &[u32], - params: GpuMirostat2SamplerParams, - mut initial_encoder: Option, - download_label: &'static str, -) -> Result, wgpu::BufferAsyncError> { - if input.datatype() != DataTypeEnum::F32 || input.layout().rank() != 1 { - return Ok(None); + previous_gpu_token: Option<&Tensor>, +) -> Option { + if logits.info.datatype() != DataTypeEnum::F32 || logits.info.rank() != 1 { + return None; + } + let input_len = logits.info.shape()[0]; + let top_k = request.top_k().min(input_len); + if top_k == 0 || top_k > TOP_K_CHUNK { + return None; } - let input_len = input.layout().shape()[0]; - let top_k = params.top_k.min(input_len); - if top_k == 0 { - return Ok(None); + let previous_gpu_token = previous_gpu_token.and_then(materialize_gpu_previous_token); + if previous_gpu_token + .as_ref() + .is_some_and(|token| !logits.device.is_same_device(token.device())) + { + return None; } - let chunks = input_len.div_ceil(TOP_K_CHUNK); - let mut candidate_count = initial_sampler_candidate_count(top_k, chunks); - let trace = cfg!(target_arch = "wasm32") - || std::env::var_os("FUSOR_TRACE_DECODE").is_some() - || std::env::var_os("FUSOR_TRACE_SAMPLER").is_some(); - let debug_dump = std::env::var_os("FUSOR_DEBUG_SAMPLER").is_some(); - let mut attempt = 0usize; - loop { - attempt += 1; - let device = input.device(); - let Some((output, ids, values, mut encoder)) = build_sample_attempt( - input, - sampler, + let dims = SampleAttemptDims { + top_k, + candidate_count: top_k, + chunks: input_len.div_ceil(TOP_K_CHUNK), + input_len, + }; + let (_, _, tail) = logits.materialize_with_tail(|logits_data, encoder| { + let attempt = encode_sample_attempt( + logits_data, previous_tokens, - PreviousTokenSource::HostOnly, - params, - SampleAttemptConfig { - top_k, - chunks, - input_len, - candidate_count, - trace, - encoder_label: "mirostat2_sample_token_to_host encoder", - }, - &mut initial_encoder, - ) else { - return Ok(None); - }; + previous_gpu_token.as_ref(), + &mut request, + dims, + encoder, + )?; + let download = encode_token_download(&attempt.0, encoder, "pending sampled token download"); + Some((attempt.0, download)) + }); + let (output, download) = tail?; - let download_size = (std::mem::size_of::() * GPU_SAMPLE_RESULT_WORDS) as u64; - let download = device.wgpu_device().create_buffer(&wgpu::BufferDescriptor { - size: download_size, - usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ, - mapped_at_creation: false, - label: Some(download_label), + let (sender, receiver) = futures_channel::oneshot::channel(); + download + .slice(..) + .map_async(wgpu::MapMode::Read, move |result| { + _ = sender.send(result); }); - encoder.copy_buffer_to_buffer(output.buffer(), 0, &download, 0, download_size); - - let debug_buffers = if debug_dump { - let ids_bytes = (std::mem::size_of::() * top_k) as u64; - let values_bytes = (std::mem::size_of::() * top_k) as u64; - let logits_bytes = (std::mem::size_of::() * input_len) as u64; - let ids_dl = device.wgpu_device().create_buffer(&wgpu::BufferDescriptor { - size: ids_bytes, - usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ, - mapped_at_creation: false, - label: Some("debug sampler ids download"), - }); - let values_dl = device.wgpu_device().create_buffer(&wgpu::BufferDescriptor { - size: values_bytes, - usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ, - mapped_at_creation: false, - label: Some("debug sampler values download"), - }); - let logits_dl = device.wgpu_device().create_buffer(&wgpu::BufferDescriptor { - size: logits_bytes, - usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ, - mapped_at_creation: false, - label: Some("debug sampler logits download"), - }); - encoder.copy_buffer_to_buffer(ids.buffer(), 0, &ids_dl, 0, ids_bytes); - encoder.copy_buffer_to_buffer(values.buffer(), 0, &values_dl, 0, values_bytes); - encoder.copy_buffer_to_buffer(input.buffer(), 0, &logits_dl, 0, logits_bytes); - Some((ids_dl, values_dl, logits_dl)) - } else { - None - }; - let submit_start = trace.then(Instant::now); - device.wgpu_queue().submit(Some(encoder.finish())); - if let Some(start) = submit_start { - tracing::info!("sampler_trace submit elapsed={:?}", start.elapsed()); - } - - let map_start = trace.then(Instant::now); - let (sender, receiver) = futures_channel::oneshot::channel(); - download - .slice(..) - .map_async(wgpu::MapMode::Read, move |result| { - _ = sender.send(result); - }); - #[cfg(not(target_arch = "wasm32"))] - device.poll_wait(); - receiver.await.map_err(|_| wgpu::BufferAsyncError)??; - if let Some(start) = map_start { - tracing::info!("sampler_trace map_wait elapsed={:?}", start.elapsed()); - } + // The token lives at word 1 of the `[status, token]` output buffer. + let token = Tensor::from(TensorData::new_from_parts( + output.device(), + output.buffer().clone(), + Layout::from_parts(1, Box::new([1]), Box::new([1])), + DataTypeEnum::U32, + )); + Some(PendingGpuSampledToken::new(token, download, receiver)) +} - let view = download.slice(..).get_mapped_range(); - let word_size = std::mem::size_of::(); - let status = view - .get(..word_size) - .map(bytemuck::from_bytes::) - .copied() - .unwrap_or(GPU_SAMPLE_STATUS_INVALID); - let token = view - .get(word_size..word_size * GPU_SAMPLE_RESULT_WORDS) - .map(bytemuck::from_bytes::) - .copied() - .unwrap_or_default(); - drop(view); - download.unmap(); +fn materialize_gpu_previous_token(token: &Tensor) -> Option { + if token.datatype() != DataTypeEnum::U32 + || token.rank() != 1 + || token.shape().first().copied().unwrap_or_default() == 0 + { + return None; + } + let (data, _) = token.data.materialize(); + Some(data) +} - match status { - GPU_SAMPLE_STATUS_SAMPLED => { - if trace { - tracing::info!( - "sampler_trace sampled attempt={attempt} top_k={top_k} chunks={chunks} candidate_count={candidate_count} token={token}" - ); - } - return Ok(Some(token)); - } - GPU_SAMPLE_STATUS_RETRY_NEEDED => { - if trace { - tracing::info!( - "sampler_trace retry attempt={attempt} top_k={top_k} chunks={chunks} candidate_count={candidate_count}" - ); - } +async fn debug_dump_invalid_sample( + logits: &TensorData, + ids: &TensorData, + values: &TensorData, + previous_tokens: &[u32], +) { + let Ok(ids_slice) = Tensor::as_slice_from_tensor_data::<1, u32>(ids).await else { + return; + }; + let Ok(values_slice) = Tensor::as_slice_from_tensor_data::<1, f32>(values).await else { + return; + }; + let Ok(logits_slice) = Tensor::as_slice_from_tensor_data::<1, f32>(logits).await else { + return; + }; + let logits_vec = logits_slice.as_slice(); + + let mut nan_count = 0usize; + let mut inf_pos = 0usize; + let mut inf_neg = 0usize; + let mut finite_count = 0usize; + let mut min_f = f32::INFINITY; + let mut max_f = f32::NEG_INFINITY; + let mut argmax = 0usize; + for (i, &v) in logits_vec.iter().enumerate() { + if v.is_nan() { + nan_count += 1; + } else if v == f32::INFINITY { + inf_pos += 1; + } else if v == f32::NEG_INFINITY { + inf_neg += 1; + } else { + finite_count += 1; + if v < min_f { + min_f = v; } - _ => { - if trace { - tracing::warn!( - "sampler_trace invalid attempt={attempt} top_k={top_k} chunks={chunks} candidate_count={candidate_count} status={status}" - ); - } - if let Some((ids_dl, values_dl, logits_dl)) = debug_buffers { - let (id_tx, id_rx) = futures_channel::oneshot::channel(); - ids_dl.slice(..).map_async(wgpu::MapMode::Read, move |r| { - _ = id_tx.send(r); - }); - let (val_tx, val_rx) = futures_channel::oneshot::channel(); - values_dl - .slice(..) - .map_async(wgpu::MapMode::Read, move |r| { - _ = val_tx.send(r); - }); - let (log_tx, log_rx) = futures_channel::oneshot::channel(); - logits_dl - .slice(..) - .map_async(wgpu::MapMode::Read, move |r| { - _ = log_tx.send(r); - }); - #[cfg(not(target_arch = "wasm32"))] - device.poll_wait(); - let _ = id_rx.await; - let _ = val_rx.await; - let _ = log_rx.await; - let ids_view = ids_dl.slice(..).get_mapped_range(); - let vals_view = values_dl.slice(..).get_mapped_range(); - let logits_view = logits_dl.slice(..).get_mapped_range(); - let ids_vec: Vec = bytemuck::cast_slice(&ids_view).to_vec(); - let vals_vec: Vec = bytemuck::cast_slice(&vals_view).to_vec(); - let logits_vec: Vec = bytemuck::cast_slice(&logits_view).to_vec(); - drop(ids_view); - drop(vals_view); - drop(logits_view); - ids_dl.unmap(); - values_dl.unmap(); - logits_dl.unmap(); - - let mut nan_count = 0usize; - let mut inf_pos = 0usize; - let mut inf_neg = 0usize; - let mut finite_count = 0usize; - let mut min_f = f32::INFINITY; - let mut max_f = f32::NEG_INFINITY; - let mut argmax = 0usize; - for (i, &v) in logits_vec.iter().enumerate() { - if v.is_nan() { - nan_count += 1; - } else if v == f32::INFINITY { - inf_pos += 1; - } else if v == f32::NEG_INFINITY { - inf_neg += 1; - } else { - finite_count += 1; - if v < min_f { - min_f = v; - } - if v > max_f { - max_f = v; - argmax = i; - } - } - } - tracing::warn!( - "sampler_debug INVALID ids={:?} values={:?} logits_len={} nan={} +inf={} -inf={} finite={} min={} max={} argmax={} first8={:?} previous_tokens_last={:?}", - ids_vec, - vals_vec, - logits_vec.len(), - nan_count, - inf_pos, - inf_neg, - finite_count, - min_f, - max_f, - argmax, - logits_vec.iter().take(8).collect::>(), - previous_tokens.iter().rev().take(8).collect::>() - ); - } - return Ok(None); + if v > max_f { + max_f = v; + argmax = i; } } - - let next = next_sampler_candidate_count(candidate_count, top_k); - if next == candidate_count { - return Ok(None); - } - candidate_count = next; } + tracing::warn!( + "sampler_debug INVALID ids={:?} values={:?} logits_len={} nan={} +inf={} -inf={} finite={} min={} max={} argmax={} first8={:?} previous_tokens_last={:?}", + ids_slice.as_slice(), + values_slice.as_slice(), + logits_vec.len(), + nan_count, + inf_pos, + inf_neg, + finite_count, + min_f, + max_f, + argmax, + logits_vec.iter().take(8).collect::>(), + previous_tokens.iter().rev().take(8).collect::>() + ); } diff --git a/fusor-ml/core/src/sampling/qmat_topk.rs b/fusor-ml/core/src/sampling/qmat_topk.rs deleted file mode 100644 index da449ffb9..000000000 --- a/fusor-ml/core/src/sampling/qmat_topk.rs +++ /dev/null @@ -1,104 +0,0 @@ -use crate::{ - Layout, - quantized::{QMatrix, matmul::QMatMulOperation}, - tensor::{DataTypeEnum, TensorData}, -}; -use wgpu::CommandEncoder; - -use super::{TOP_K_CHUNK, min_top_k_candidates_per_chunk}; - -pub(super) fn qmat_logits_data_with_encoder( - hidden: &TensorData, - matrix: &QMatrix, - encoder: &mut CommandEncoder, -) -> Option { - if hidden.datatype() != DataTypeEnum::F32 || hidden.layout().rank() != 1 { - return None; - } - let hidden_len = hidden.layout().shape()[0]; - let hidden_stride = hidden.layout().strides()[0]; - let [vocab_len, matrix_hidden_len] = matrix.shape() else { - return None; - }; - if hidden_len != *matrix_hidden_len || *vocab_len == 0 { - return None; - } - - let device = hidden.device(); - let compact_hidden; - let hidden = if hidden.layout().offset() == 0 && hidden_stride == 1 { - hidden - } else if hidden_stride == 1 { - compact_hidden = TensorData::new_for_shape(device, &[hidden_len], DataTypeEnum::F32); - encoder.copy_buffer_to_buffer( - hidden.buffer(), - (hidden.layout().offset() * std::mem::size_of::()) as u64, - compact_hidden.buffer(), - 0, - (hidden_len * std::mem::size_of::()) as u64, - ); - &compact_hidden - } else { - return None; - }; - let logits = TensorData::new_for_shape(device, &[*vocab_len], DataTypeEnum::F32); - let hidden_2d = TensorData::new_from_parts( - device, - hidden.buffer().clone(), - Layout::from_parts( - hidden.layout().offset(), - Box::new([1, hidden_len]), - Box::new([0, hidden_stride]), - ), - DataTypeEnum::F32, - ); - let logits_2d = TensorData::new_from_parts( - device, - logits.buffer().clone(), - Layout::from_parts(0, Box::new([1, *vocab_len]), Box::new([*vocab_len, 1])), - DataTypeEnum::F32, - ); - let kernel = QMatMulOperation::direct_kernel_for_tensors( - device, - crate::quantized::matmul::DirectKernelTensors { - input: &hidden_2d, - matrix, - pre_extra_tensors: &[], - post_extra_tensors: &[], - output: &logits_2d, - }, - "q_mat_logits_for_sampler", - crate::quantized::matmul::DirectKernelChains { - pre_expr: None, - post_expr: None, - post_accumulator_offsets: &[], - }, - None, - )?; - kernel.run(device.kernel_cache(), encoder); - - Some(logits) -} - -pub(super) fn initial_sampler_candidate_count(top_k: usize, chunks: usize) -> usize { - top_k - .div_ceil(chunks) - .max(min_top_k_candidates_per_chunk()) - .min(top_k) - .min(TOP_K_CHUNK) -} - -pub(super) fn sampler_output_per_chunk(candidate_count: usize) -> usize { - if candidate_count >= TOP_K_CHUNK { - TOP_K_CHUNK - } else { - candidate_count + 1 - } -} - -pub(super) fn next_sampler_candidate_count(candidate_count: usize, top_k: usize) -> usize { - candidate_count - .saturating_mul(2) - .min(top_k) - .min(TOP_K_CHUNK) -} diff --git a/fusor-ml/core/src/sampling/tests.rs b/fusor-ml/core/src/sampling/tests.rs index bbc4764d3..9c77305a0 100644 --- a/fusor-ml/core/src/sampling/tests.rs +++ b/fusor-ml/core/src/sampling/tests.rs @@ -3,13 +3,14 @@ use std::mem::size_of; use crate::{DataTypeEnum, Device, Tensor, TensorData, quantized::QMatrix}; use fusor_gguf::{BlockQ4_0, GgmlType}; +use crate::mir::kernel_backend::sampling_topk::chunk_top_k_pair_data_with_processors_with_encoder; + use super::{ GPU_SAMPLE_STATUS_RETRY_NEEDED, GPU_SAMPLE_STATUS_SAMPLED, GpuMirostat2Sampler, - GpuMirostat2SamplerParams, GpuStandardSamplerParams, - mirostat::sample_from_sorted_top_k_data_with_encoder, - mirostat2_sample_token_to_host, + GpuMirostat2SamplerParams, GpuSamplerRequest, GpuStandardSamplerParams, + mirostat::sample_from_sorted_top_k_data_with_encoder, sample_token_to_host, standard_sampler::sample_from_sorted_top_k_data_with_encoder as sample_standard_from_sorted_top_k_data_with_encoder, - topk::{ProcessorSettings, chunk_top_k_pair_data_with_processors_with_encoder}, + topk::ProcessorSettings, }; #[test] @@ -459,9 +460,17 @@ fn mirostat2_sampler_uses_exact_top_k_when_candidates_cluster() { ); let mut sampler = GpuMirostat2Sampler::new(&device, mu); - let token = mirostat2_sample_token_to_host(&data, &mut sampler, &[], params) - .await - .unwrap(); + let logits = Tensor::from(data); + let token = sample_token_to_host( + &logits.data, + GpuSamplerRequest::Mirostat2 { + sampler: &mut sampler, + params, + }, + &[], + ) + .await + .unwrap(); assert_eq!(token, Some(expected)); }); @@ -558,7 +567,7 @@ fn top_k_pairs_large_vocab_merge_path_matches_cpu_sorted_order() { } #[test] -fn qmat_mirostat2_sample_token_uses_direct_sampler_path() { +fn qmat_logits_sample_token_through_graph() { pollster::block_on(async { let device = Device::new().await.unwrap(); let hidden = Tensor::new(&device, vec![1.0f32; 32].as_slice()); @@ -577,10 +586,17 @@ fn qmat_mirostat2_sample_token_uses_direct_sampler_path() { random: 0.0, }; - let token = hidden - .try_sample_mirostat2_token_q_mat(&matrix, &mut sampler, &[], params) - .await - .unwrap(); + let logits = hidden.q_mat_mul(&matrix); + let token = sample_token_to_host( + &logits.data, + GpuSamplerRequest::Mirostat2 { + sampler: &mut sampler, + params, + }, + &[], + ) + .await + .unwrap(); assert_eq!(token, Some(7)); }); diff --git a/fusor-ml/core/src/sampling/topk.rs b/fusor-ml/core/src/sampling/topk.rs index 0f43ac390..95c4a55d6 100644 --- a/fusor-ml/core/src/sampling/topk.rs +++ b/fusor-ml/core/src/sampling/topk.rs @@ -5,5 +5,5 @@ pub(crate) use crate::mir::kernel_backend::sampling_topk::{ pub(super) use crate::mir::kernel_backend::sampling_topk::{ ProcessorSettings, chunk_top_k_pair_data_with_processors_and_gpu_tail_with_encoder, - chunk_top_k_pair_data_with_processors_with_encoder, top_k_exactness_flag_data_with_encoder, + top_k_exactness_flag_data_with_encoder, }; diff --git a/fusor-ml/core/src/slice_assign.rs b/fusor-ml/core/src/slice_assign.rs index 82596129d..b25720d67 100644 --- a/fusor-ml/core/src/slice_assign.rs +++ b/fusor-ml/core/src/slice_assign.rs @@ -1,7 +1,7 @@ use std::{hash::Hash, ops::Range}; use crate::{ - DataTypeEnum, TILE_SIZE, Tensor, TensorData, + DataTypeEnum, TILE_SIZE, Tensor, compute_graph::{ComputeGraphInner, NodeIndex}, mir::{ inputs::MirValue, @@ -13,20 +13,20 @@ use crate::{ visit_tiled::{titled_map_dispatch_size, titled_map_workgroup_size_constraints}, }; +/// The in-place region write of the graph vocabulary: dispatch over the +/// slice region only, writing the value into the input's buffer. The +/// out-of-place form has no operation — `Tensor::slice_assign` composes it +/// as a plain elementwise select (see [`slice_assign_expression`]). #[derive(Clone, Debug)] pub(crate) struct SliceAssignOperation { pub(crate) input: NodeIndex, pub(crate) value: NodeIndex, pub(crate) slices: Box<[Range]>, - pub(crate) input_shape: Box<[usize]>, - pub(crate) in_place: bool, } -/// The composed slice-assign body: per output coordinate, read the assigned -/// value inside the slice region and the original input outside it. Inputs: -/// 0 = the original tensor, 1 = the assigned value. -pub(crate) fn slice_assign_expression(slices: &[Range], datatype: DataTypeEnum) -> NaryExpr { - let rank = slices.len(); +/// The region predicate of a slice assign: 1 when every output coordinate +/// lies inside its slice range, 0 otherwise. +pub(crate) fn slice_region_condition(slices: &[Range]) -> NaryExpr { let mut condition = NaryExpr::scalar(NaryScalar::U32(1)); for (dim, slice) in slices.iter().enumerate() { let dim_index = NaryExpr::DimIndex(dim); @@ -47,6 +47,15 @@ pub(crate) fn slice_assign_expression(slices: &[Range], datatype: DataTyp condition = NaryExpr::mul(condition, ge_start, DataTypeEnum::U32); condition = NaryExpr::mul(condition, lt_end, DataTypeEnum::U32); } + condition +} + +/// The composed slice-assign body: per output coordinate, read the assigned +/// value inside the slice region and the original input outside it. Inputs: +/// 0 = the original tensor, 1 = the assigned value. +pub(crate) fn slice_assign_expression(slices: &[Range], datatype: DataTypeEnum) -> NaryExpr { + let rank = slices.len(); + let condition = slice_region_condition(slices); let value_indices = slices .iter() @@ -83,18 +92,11 @@ pub(crate) fn slice_assign_expression(slices: &[Range], datatype: DataTyp } impl SliceAssignOperation { - pub fn new_in_place( - input: NodeIndex, - value: NodeIndex, - slices: Box<[Range]>, - input_shape: Box<[usize]>, - ) -> Self { + pub fn new_in_place(input: NodeIndex, value: NodeIndex, slices: Box<[Range]>) -> Self { Self { input, value, slices, - input_shape, - in_place: true, } } @@ -104,32 +106,15 @@ impl SliceAssignOperation { .map(|slice| slice.end - slice.start) .collect() } - - fn operation_shape(&self) -> Box<[usize]> { - if self.in_place { - self.value_shape() - } else { - self.input_shape.clone() - } - } - - fn expression(&self, datatype: DataTypeEnum) -> NaryExpr { - if self.in_place { - return NaryExpr::input(0, self.slices.len()); - } - slice_assign_expression(&self.slices, datatype) - } } impl Operation for SliceAssignOperation { fn hash_kernel_fields(&self, state: &mut rustc_hash::FxHasher) { self.slices.hash(state); - self.input_shape.hash(state); - self.in_place.hash(state); } fn workgroup_shape_constraints(&self, device: &crate::Device) -> WorkgroupShapeConstraints { - titled_map_workgroup_size_constraints(&self.operation_shape(), device) + titled_map_workgroup_size_constraints(&self.value_shape(), device) } fn dispatch_size(&self, workgroup_shape: &WorkgroupShape, inputs: &[MirValue]) -> [u32; 3] { @@ -142,7 +127,7 @@ impl Operation for SliceAssignOperation { titled_map_dispatch_size( TILE_SIZE, *workgroup_shape, - &self.operation_shape(), + &self.value_shape(), max_per_dim, ) } @@ -156,16 +141,7 @@ impl Operation for SliceAssignOperation { // Pass the ORIGINAL input tensor (not sliced) and the value tensor let input = nodes.get_cached_result(self.input).unwrap(); let value = nodes.get_cached_result(self.value).unwrap(); - - if self.in_place { - let output = input.slice(&self.slices); - return vec![input.clone().into(), value.clone().into(), output.into()]; - } - - // Create output buffer with the same shape as input - let output = - TensorData::new_for_shape(input.device(), input.layout().shape(), input.datatype()); - + let output = input.slice(&self.slices); vec![input.clone().into(), value.clone().into(), output.into()] } @@ -175,46 +151,26 @@ impl Operation for SliceAssignOperation { workgroup_shape: &WorkgroupShape, inputs: &[MirValue], ) -> Option { - if self.in_place { - let value = inputs[1].as_tensor()?; - let operation = ElementwiseOperation { - inputs: vec![self.value], - expression: self.expression(value.datatype()), - shape: value.layout().shape().into(), - output_datatype: value.datatype(), - }; - return crate::nary_direct::build_nary_direct_kernel_to_output( - &operation, - graph, - workgroup_shape, - &[inputs[1].clone(), inputs[2].clone()], - 1, - ); - } - - let input = inputs[0].as_tensor()?; + // A copy kernel over the slice region: read the value, write into + // the sliced view of the input's buffer. + let value = inputs[1].as_tensor()?; let operation = ElementwiseOperation { - inputs: vec![self.input, self.value], - expression: self.expression(input.datatype()), - shape: self.input_shape.clone(), - output_datatype: input.datatype(), + inputs: vec![self.value], + expression: NaryExpr::input(0, self.slices.len()), + shape: value.layout().shape().into(), + output_datatype: value.datatype(), }; crate::nary_direct::build_nary_direct_kernel_to_output( &operation, graph, workgroup_shape, - inputs, - 2, + &[inputs[1].clone(), inputs[2].clone()], + 1, ) } fn output(&self, _nodes: &ComputeGraphInner, inputs: &[MirValue]) -> MirValue { - if self.in_place { - return inputs[0].clone(); - } - - // Return the output tensor (last input) - inputs[2].clone() + inputs[0].clone() } fn name(&self) -> String { diff --git a/fusor-ml/core/src/tensor/mod.rs b/fusor-ml/core/src/tensor/mod.rs index 485434f09..03d2de2fa 100644 --- a/fusor-ml/core/src/tensor/mod.rs +++ b/fusor-ml/core/src/tensor/mod.rs @@ -417,13 +417,7 @@ impl Tensor { slices: impl Into]>>, value: &Self, ) -> Self { - let input_shape: Box<[usize]> = self.shape().to_vec().into_boxed_slice(); - let op = SliceAssignOperation::new_in_place( - self.data.key, - value.data.key, - slices.into(), - input_shape, - ); + let op = SliceAssignOperation::new_in_place(self.data.key, value.data.key, slices.into()); Self::from_parts(self.data.slice_assign(op)) } diff --git a/fusor-ml/core/src/tensor/sampling.rs b/fusor-ml/core/src/tensor/sampling.rs index 8eacb1ca2..5d7a14643 100644 --- a/fusor-ml/core/src/tensor/sampling.rs +++ b/fusor-ml/core/src/tensor/sampling.rs @@ -1,119 +1,83 @@ use super::{DataTypeEnum, Tensor, TensorData}; -use crate::quantized::QMatrix; +use crate::top_k::GpuSamplerRequest; impl Tensor { - pub async fn try_sample_mirostat2_token_q_mat( + pub async fn sample_mirostat2_token( &self, - matrix: &QMatrix, sampler: &mut crate::top_k::GpuMirostat2Sampler, previous_tokens: &[u32], params: crate::top_k::GpuMirostat2SamplerParams, - ) -> Result, wgpu::BufferAsyncError> { + ) -> Result { self.assert_rank::<1>(); self.assert_datatype::(); - crate::top_k::qmat_mirostat2_sample_lazy_token_to_host( + if let Some(token) = crate::top_k::sample_token_to_host( &self.data, - matrix, - sampler, + GpuSamplerRequest::Mirostat2 { sampler, params }, previous_tokens, - params, ) - .await + .await? + { + return Ok(token); + } + + let (ids, _) = self.top_k_pairs(params.top_k).await?; + Ok(ids.first().copied().unwrap_or_default()) } - pub fn try_sample_mirostat2_token_q_mat_pending( + pub async fn sample_standard_token( &self, - matrix: &QMatrix, - sampler: &mut crate::top_k::GpuMirostat2Sampler, previous_tokens: &[u32], - previous_gpu_token: Option<&Tensor>, - params: crate::top_k::GpuMirostat2SamplerParams, - ) -> Option { + params: crate::top_k::GpuStandardSamplerParams, + ) -> Result { self.assert_rank::<1>(); self.assert_datatype::(); - crate::top_k::qmat_mirostat2_sample_lazy_token_pending( + if let Some(token) = crate::top_k::sample_token_to_host( &self.data, - matrix, - sampler, + GpuSamplerRequest::Standard { params }, previous_tokens, - previous_gpu_token, - params, ) + .await? + { + return Ok(token); + } + + let (ids, _) = self.top_k_pairs(params.top_k).await?; + Ok(ids.first().copied().unwrap_or_default()) } - pub async fn try_sample_standard_token_q_mat( + pub fn sample_mirostat2_token_pending( &self, - matrix: &QMatrix, + sampler: &mut crate::top_k::GpuMirostat2Sampler, previous_tokens: &[u32], - params: crate::top_k::GpuStandardSamplerParams, - ) -> Result, wgpu::BufferAsyncError> { + previous_gpu_token: Option<&Tensor>, + params: crate::top_k::GpuMirostat2SamplerParams, + ) -> Option { self.assert_rank::<1>(); self.assert_datatype::(); - crate::top_k::qmat_standard_sample_lazy_token_to_host( + crate::top_k::sample_token_pending( &self.data, - matrix, + GpuSamplerRequest::Mirostat2 { sampler, params }, previous_tokens, - params, + previous_gpu_token, ) - .await } - pub fn try_sample_standard_token_q_mat_pending( + pub fn sample_standard_token_pending( &self, - matrix: &QMatrix, previous_tokens: &[u32], previous_gpu_token: Option<&Tensor>, params: crate::top_k::GpuStandardSamplerParams, ) -> Option { self.assert_rank::<1>(); self.assert_datatype::(); - crate::top_k::qmat_standard_sample_lazy_token_pending( + crate::top_k::sample_token_pending( &self.data, - matrix, + GpuSamplerRequest::Standard { params }, previous_tokens, previous_gpu_token, - params, ) } - pub async fn sample_mirostat2_token( - &self, - sampler: &mut crate::top_k::GpuMirostat2Sampler, - previous_tokens: &[u32], - params: crate::top_k::GpuMirostat2SamplerParams, - ) -> Result { - self.assert_rank::<1>(); - self.assert_datatype::(); - let (input, _) = self.data.materialize(); - if let Some(token) = - crate::top_k::mirostat2_sample_token_to_host(&input, sampler, previous_tokens, params) - .await? - { - return Ok(token); - } - - let (ids, _) = self.top_k_pairs(params.top_k).await?; - Ok(ids.first().copied().unwrap_or_default()) - } - - pub async fn sample_standard_token( - &self, - previous_tokens: &[u32], - params: crate::top_k::GpuStandardSamplerParams, - ) -> Result { - self.assert_rank::<1>(); - self.assert_datatype::(); - let (input, _) = self.data.materialize(); - if let Some(token) = - crate::top_k::standard_sample_token_to_host(&input, previous_tokens, params).await? - { - return Ok(token); - } - - let (ids, _) = self.top_k_pairs(params.top_k).await?; - Ok(ids.first().copied().unwrap_or_default()) - } - pub async fn top_k_pairs( &self, k: usize, @@ -122,12 +86,12 @@ impl Tensor { return Ok((Vec::new(), Vec::new())); } - let (input, _) = self.data.materialize(); - if input.datatype() != DataTypeEnum::F32 || input.layout().rank() != 1 { + if self.datatype() != DataTypeEnum::F32 || self.rank() != 1 { + let (input, _) = self.data.materialize(); return cpu_top_k_pairs_from_tensor_data(&input, k).await; } - let input_len = input.layout().shape()[0]; + let input_len = self.shape()[0]; let k = k.min(input_len); if k == 0 { return Ok((Vec::new(), Vec::new())); @@ -140,67 +104,34 @@ impl Tensor { .min(k) .min(crate::top_k::TOP_K_CHUNK); + // The first attempt's chunk + merge kernels ride the resolver's + // encoder, so the graph producing the logits and the top-k selection + // land in one submission. Retries with escalated candidate counts + // re-encode over the materialized input. + let (input, _, mut attempt) = self.data.materialize_with_tail(|input, encoder| { + encode_top_k_attempt(input, k, input_len, chunks, candidate_count, encoder) + }); + loop { - let output_per_chunk = if candidate_count >= crate::top_k::TOP_K_CHUNK { - crate::top_k::TOP_K_CHUNK - } else { - candidate_count + 1 - }; - let mut encoder = input.device().wgpu_device().create_command_encoder( - &wgpu::CommandEncoderDescriptor { - label: Some("top_k_pairs encoder"), - }, - ); - let Some((ids, values)) = crate::top_k::chunk_top_k_pair_data_with_encoder( - &input, - candidate_count, - output_per_chunk, - Some(&mut encoder), - ) else { + let Some(TopKAttempt { + chunk_ids, + chunk_values, + merged_ids, + merged_values, + exhaustive, + }) = attempt + else { return cpu_top_k_pairs_from_tensor_data(&input, k).await; }; - if candidate_count >= crate::top_k::TOP_K_CHUNK { - let Some((ids, values)) = - crate::top_k::merge_sorted_chunk_top_k_pair_data_with_encoder( - &ids, - &values, - crate::top_k::MergeSortedChunkTopKParams { - chunks, - chunk_len: crate::top_k::TOP_K_CHUNK, - chunk_stride: crate::top_k::TOP_K_CHUNK, - input_len, - k, - }, - Some(&mut encoder), - ) - else { - return cpu_top_k_pairs_from_tensor_data(&input, k).await; - }; - input.device().wgpu_queue().submit(Some(encoder.finish())); - let ids = Tensor::as_slice_from_tensor_data::<1, u32>(&ids).await?; - let values = Tensor::as_slice_from_tensor_data::<1, f32>(&values).await?; + if exhaustive { + let ids = Tensor::as_slice_from_tensor_data::<1, u32>(&merged_ids).await?; + let values = Tensor::as_slice_from_tensor_data::<1, f32>(&merged_values).await?; return Ok((ids.as_slice().to_vec(), values.as_slice().to_vec())); } - let Some((merged_ids, merged_values)) = - crate::top_k::merge_sorted_chunk_top_k_pair_data_with_encoder( - &ids, - &values, - crate::top_k::MergeSortedChunkTopKParams { - chunks, - chunk_len: candidate_count, - chunk_stride: output_per_chunk, - input_len, - k, - }, - Some(&mut encoder), - ) - else { - return cpu_top_k_pairs_from_tensor_data(&input, k).await; - }; - input.device().wgpu_queue().submit(Some(encoder.finish())); + let output_per_chunk = candidate_count + 1; let merged_ids = Tensor::as_slice_from_tensor_data::<1, u32>(&merged_ids).await?; let merged_values = Tensor::as_slice_from_tensor_data::<1, f32>(&merged_values).await?; - let chunk_values = Tensor::as_slice_from_tensor_data::<1, f32>(&values).await?; + let chunk_values = Tensor::as_slice_from_tensor_data::<1, f32>(&chunk_values).await?; let exact = top_k_chunk_bounds_prove_exact( merged_values.as_slice(), chunk_values.as_slice(), @@ -216,7 +147,7 @@ impl Tensor { )); } - let ids = Tensor::as_slice_from_tensor_data::<1, u32>(&ids).await?; + let ids = Tensor::as_slice_from_tensor_data::<1, u32>(&chunk_ids).await?; if let Some(top) = top_k_from_chunk_candidates( ids.as_slice(), chunk_values.as_slice(), @@ -233,10 +164,78 @@ impl Tensor { return cpu_top_k_pairs_from_tensor_data(&input, k).await; } candidate_count = (candidate_count * 2).min(crate::top_k::TOP_K_CHUNK); + + let mut encoder = input.device().wgpu_device().create_command_encoder( + &wgpu::CommandEncoderDescriptor { + label: Some("top_k_pairs retry encoder"), + }, + ); + attempt = + encode_top_k_attempt(&input, k, input_len, chunks, candidate_count, &mut encoder); + input.device().wgpu_queue().submit(Some(encoder.finish())); } } } +struct TopKAttempt { + chunk_ids: TensorData, + chunk_values: TensorData, + merged_ids: TensorData, + merged_values: TensorData, + /// Every chunk contributed all of its elements: the merge result is the + /// exact top-k with no host-side exactness check needed. + exhaustive: bool, +} + +/// Encode one chunked top-k + merge attempt into `encoder`. Runs inside the +/// resolver tail while the graph lock is held, so it must only touch raw +/// buffers — no compute-graph access. +fn encode_top_k_attempt( + input: &TensorData, + k: usize, + input_len: usize, + chunks: usize, + candidate_count: usize, + encoder: &mut wgpu::CommandEncoder, +) -> Option { + let exhaustive = candidate_count >= crate::top_k::TOP_K_CHUNK; + let output_per_chunk = if exhaustive { + crate::top_k::TOP_K_CHUNK + } else { + candidate_count + 1 + }; + let (chunk_ids, chunk_values) = crate::top_k::chunk_top_k_pair_data_with_encoder( + input, + candidate_count, + output_per_chunk, + Some(encoder), + )?; + let (merged_ids, merged_values) = + crate::top_k::merge_sorted_chunk_top_k_pair_data_with_encoder( + &chunk_ids, + &chunk_values, + crate::top_k::MergeSortedChunkTopKParams { + chunks, + chunk_len: if exhaustive { + crate::top_k::TOP_K_CHUNK + } else { + candidate_count + }, + chunk_stride: output_per_chunk, + input_len, + k, + }, + Some(encoder), + )?; + Some(TopKAttempt { + chunk_ids, + chunk_values, + merged_ids, + merged_values, + exhaustive, + }) +} + fn top_k_chunk_bounds_prove_exact( top_values: &[f32], chunk_values: &[f32], diff --git a/fusor-ml/core/src/top_k.rs b/fusor-ml/core/src/top_k.rs index 85e3e896c..76cfab446 100644 --- a/fusor-ml/core/src/top_k.rs +++ b/fusor-ml/core/src/top_k.rs @@ -4,11 +4,9 @@ pub use crate::sampling::{ }; pub(crate) use crate::sampling::{ - MergeSortedChunkTopKParams, TOP_K_CHUNK, chunk_top_k_pair_data_with_encoder, + GpuSamplerRequest, MergeSortedChunkTopKParams, TOP_K_CHUNK, chunk_top_k_pair_data_with_encoder, merge_sorted_chunk_top_k_pair_data_with_encoder, min_top_k_candidates_per_chunk, - mirostat2_sample_token_to_host, qmat_mirostat2_sample_lazy_token_pending, - qmat_mirostat2_sample_lazy_token_to_host, qmat_standard_sample_lazy_token_pending, - qmat_standard_sample_lazy_token_to_host, standard_sample_token_to_host, + sample_token_pending, sample_token_to_host, }; #[cfg(test)] diff --git a/fusor-ml/core/src/view.rs b/fusor-ml/core/src/view.rs index 2591ea55f..d09a9564b 100644 --- a/fusor-ml/core/src/view.rs +++ b/fusor-ml/core/src/view.rs @@ -153,7 +153,7 @@ impl ViewOperation { } } -fn zero_scalar(datatype: DataTypeEnum) -> NaryScalar { +pub(crate) fn zero_scalar(datatype: DataTypeEnum) -> NaryScalar { match datatype { DataTypeEnum::F32 => NaryScalar::F32(0.0), DataTypeEnum::F16 => NaryScalar::F16(half::f16::from_f32(0.0)), diff --git a/fusor-ml/cpu/src/matmul.rs b/fusor-ml/cpu/src/matmul.rs index 4d4938598..a53336e9e 100644 --- a/fusor-ml/cpu/src/matmul.rs +++ b/fusor-ml/cpu/src/matmul.rs @@ -227,12 +227,12 @@ pub fn batched_matmul( lhs: &ConcreteTensor, rhs: &ConcreteTensor, ) -> ConcreteTensor { - const { - assert!( - R >= 2, - "Matrix multiplication requires at least 2 dimensions" - ) - }; + // Runtime (not const) so rank-generic callers can guard rank-1 themselves + // without this monomorphization failing. + assert!( + R >= 2, + "Matrix multiplication requires at least 2 dimensions" + ); let lhs_shape = lhs.layout().shape(); let rhs_shape = rhs.layout().shape(); diff --git a/fusor-ml/cpu/src/quantized.rs b/fusor-ml/cpu/src/quantized.rs index bd33dea91..caae5057a 100644 --- a/fusor-ml/cpu/src/quantized.rs +++ b/fusor-ml/cpu/src/quantized.rs @@ -248,7 +248,9 @@ impl ConcreteTensor { B::Dequantized: AsRef<[f32]>, B::ActivationBlock: Pod + Send + Sync, { - const { assert!(R >= 2, "q_mat_mul requires at least 2 dimensions") }; + // Runtime (not const) so rank-generic callers can guard rank-1 + // themselves without this monomorphization failing. + assert!(R >= 2, "q_mat_mul requires at least 2 dimensions"); let rhs_shape = rhs.element_shape(); assert_eq!( diff --git a/fusor-ml/fusor/src/gpu.rs b/fusor-ml/fusor/src/gpu.rs index 6dc51b3ce..615368c84 100644 --- a/fusor-ml/fusor/src/gpu.rs +++ b/fusor-ml/fusor/src/gpu.rs @@ -565,23 +565,8 @@ impl Tensor { impl Tensor<1, f32> { #[inline] - pub async fn try_sample_mirostat2_token_q_mat( + pub fn sample_mirostat2_token_pending( &self, - matrix: &QMatrix, - sampler: &mut GpuMirostat2Sampler, - previous_tokens: &[u32], - params: GpuMirostat2SamplerParams, - ) -> Result> { - self.inner - .try_sample_mirostat2_token_q_mat(matrix, sampler, previous_tokens, params) - .await - .map_err(Error::from) - } - - #[inline] - pub fn try_sample_mirostat2_token_q_mat_pending( - &self, - matrix: &QMatrix, sampler: &mut GpuMirostat2Sampler, previous_tokens: &[u32], previous_gpu_token: Option<&GpuSampledToken>, @@ -589,8 +574,7 @@ impl Tensor<1, f32> { ) -> Result> { Ok(self .inner - .try_sample_mirostat2_token_q_mat_pending( - matrix, + .sample_mirostat2_token_pending( sampler, previous_tokens, previous_gpu_token.map(|token| token.as_core_token()), @@ -600,30 +584,15 @@ impl Tensor<1, f32> { } #[inline] - pub async fn try_sample_standard_token_q_mat( - &self, - matrix: &QMatrix, - previous_tokens: &[u32], - params: GpuStandardSamplerParams, - ) -> Result> { - self.inner - .try_sample_standard_token_q_mat(matrix, previous_tokens, params) - .await - .map_err(Error::from) - } - - #[inline] - pub fn try_sample_standard_token_q_mat_pending( + pub fn sample_standard_token_pending( &self, - matrix: &QMatrix, previous_tokens: &[u32], previous_gpu_token: Option<&GpuSampledToken>, params: GpuStandardSamplerParams, ) -> Result> { Ok(self .inner - .try_sample_standard_token_q_mat_pending( - matrix, + .sample_standard_token_pending( previous_tokens, previous_gpu_token.map(|token| token.as_core_token()), params, diff --git a/fusor-ml/fusor/src/lib.rs b/fusor-ml/fusor/src/lib.rs index 90f5ea05c..1d3360641 100644 --- a/fusor-ml/fusor/src/lib.rs +++ b/fusor-ml/fusor/src/lib.rs @@ -493,34 +493,16 @@ where } } - pub async fn try_sample_mirostat2_token_q_mat( + pub fn sample_mirostat2_token_pending( &self, - weights: &crate::QMatrix, - sampler: &mut Mirostat2Sampler, - previous_tokens: &[u32], - params: Mirostat2SamplerParams, - ) -> Result, Error> { - match (self, weights) { - (Tensor::Gpu(t), crate::QMatrix::Gpu(weights)) => t - .try_sample_mirostat2_token_q_mat(weights, sampler, previous_tokens, params) - .await - .map_err(Error::Gpu), - _ => Ok(None), - } - } - - pub fn try_sample_mirostat2_token_q_mat_pending( - &self, - weights: &crate::QMatrix, sampler: &mut Mirostat2Sampler, previous_tokens: &[u32], previous_gpu_token: Option<&GpuSampledToken>, params: Mirostat2SamplerParams, ) -> Result, Error> { - match (self, weights) { - (Tensor::Gpu(t), crate::QMatrix::Gpu(weights)) => t - .try_sample_mirostat2_token_q_mat_pending( - weights, + match self { + Tensor::Gpu(t) => t + .sample_mirostat2_token_pending( sampler, previous_tokens, previous_gpu_token, @@ -531,36 +513,15 @@ where } } - pub async fn try_sample_standard_token_q_mat( - &self, - weights: &crate::QMatrix, - previous_tokens: &[u32], - params: StandardSamplerParams, - ) -> Result, Error> { - match (self, weights) { - (Tensor::Gpu(t), crate::QMatrix::Gpu(weights)) => t - .try_sample_standard_token_q_mat(weights, previous_tokens, params) - .await - .map_err(Error::Gpu), - _ => Ok(None), - } - } - - pub fn try_sample_standard_token_q_mat_pending( + pub fn sample_standard_token_pending( &self, - weights: &crate::QMatrix, previous_tokens: &[u32], previous_gpu_token: Option<&GpuSampledToken>, params: StandardSamplerParams, ) -> Result, Error> { - match (self, weights) { - (Tensor::Gpu(t), crate::QMatrix::Gpu(weights)) => t - .try_sample_standard_token_q_mat_pending( - weights, - previous_tokens, - previous_gpu_token, - params, - ) + match self { + Tensor::Gpu(t) => t + .sample_standard_token_pending(previous_tokens, previous_gpu_token, params) .map_err(Error::Gpu), _ => Ok(None), } @@ -1176,6 +1137,28 @@ where weights.shape().len() ); + // A rank-1 activation is a single matrix row. The GPU graph lowers + // quantized formats natively; the CPU and dense F16/F32 paths only + // handle rank >= 2, so route them through a [1, K] view. + if R == 1 { + let gpu_native = matches!((self, weights), (Tensor::Gpu(_), QMatrix::Gpu(_))) + && !matches!( + weights.ggml_type(), + fusor_gguf::GgmlType::F16 | fusor_gguf::GgmlType::F32 + ); + if !gpu_native { + let k = self.shape()[0]; + let n = weights.shape()[0]; + let out_shape: [usize; R] = std::array::from_fn(|_| n); + return self + .reshape([1, k]) + .to_concrete() + .q_mat_mul(weights) + .reshape(out_shape) + .to_concrete(); + } + } + match (self, weights) { // CPU path - dispatch based on block type // eval() returns Tensor, so we need .inner() to get ConcreteTensor diff --git a/models/kalosm-llama/examples/qwen_probe.rs b/models/kalosm-llama/examples/qwen_probe.rs new file mode 100644 index 000000000..c1e659d39 --- /dev/null +++ b/models/kalosm-llama/examples/qwen_probe.rs @@ -0,0 +1,38 @@ +//! Temporary perf probe for the Qwen2.5-VL mrope split+op+cat path. + +use kalosm_llama::*; +use kalosm_model_types::ModelLoadingProgress; +use prelude::{StreamExt, TextCompletionModelExt}; + +fn main() { + let _ = tracing_subscriber::fmt::try_init(); + pollster::block_on(async { + let model = Llama::builder() + .with_source(LlamaSource::qwen_2_5_3b_vl_chat_q4()) + .build_with_loading_handler(|_: ModelLoadingProgress| {}) + .await + .unwrap(); + + let prompt = "The capital of France is"; + let mut stream = model.complete(prompt).take(72); + let mut text = String::new(); + // Warmup: first tokens pay one-time kernel compilation. + for _ in 0..8 { + if let Some(token) = stream.next().await { + text.push_str(&token); + } + } + let start = std::time::Instant::now(); + let mut tokens = 0usize; + while let Some(token) = stream.next().await { + text.push_str(&token); + tokens += 1; + } + let elapsed = start.elapsed(); + println!("OUTPUT: {}", text.replace('\n', " | ")); + println!( + "steady-state tokens={tokens} elapsed={elapsed:?} tps={:.2}", + tokens as f64 / elapsed.as_secs_f64() + ); + }); +} diff --git a/models/kalosm-llama/examples/vision.rs b/models/kalosm-llama/examples/vision.rs index 6556daf8d..421634e4b 100644 --- a/models/kalosm-llama/examples/vision.rs +++ b/models/kalosm-llama/examples/vision.rs @@ -1,29 +1,27 @@ -use fusor::Device; use kalosm_llama::prelude::*; use kalosm_streams::text_stream::TextStream; -fn main() { - pollster::block_on(async { - tracing_subscriber::fmt::init(); - let model = Llama::builder() - .with_source(LlamaSource::qwen_2_5_3b_vl_chat_q4()) - .with_device(Device::cpu()) - .build() - .await - .unwrap(); +// The demo image is fetched over HTTP, which needs a tokio reactor. +#[tokio::main] +async fn main() { + tracing_subscriber::fmt::init(); + let model = Llama::builder() + .with_source(LlamaSource::qwen_2_5_3b_vl_chat_q4()) + .build() + .await + .unwrap(); - let mut chat = model.chat(); - let mut response = chat(&( - MediaChunk::new( - MediaSource::url( - "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-VL/assets/demo.jpeg", - ), - MediaType::Image, + let mut chat = model.chat(); + let mut response = chat(&( + MediaChunk::new( + MediaSource::url( + "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-VL/assets/demo.jpeg", ), - "Describe this image.", - )); - response.to_std_out().await.unwrap(); - response.await.unwrap(); - println!("\n"); - }); + MediaType::Image, + ), + "Describe this image.", + )); + response.to_std_out().await.unwrap(); + response.await.unwrap(); + println!("\n"); } diff --git a/models/kalosm-llama/src/model/forward.rs b/models/kalosm-llama/src/model/forward.rs index 33db557fd..fc93bf559 100644 --- a/models/kalosm-llama/src/model/forward.rs +++ b/models/kalosm-llama/src/model/forward.rs @@ -377,16 +377,18 @@ where previous_gpu_token: Option<&fusor::GpuSampledToken>, top_k: usize, ) -> Result, LlamaModelError> { - let hidden = hidden.squeeze(0).to_concrete(); + let logits = hidden + .squeeze(0) + .to_concrete() + .q_mat_mul(model.output_matrix()); match sampler.config.sampling_strategy { kalosm_language_model::SamplingStrategy::Mirostat2 => { let params = sampler.mirostat_params(top_k); let mirostat = sampler.mirostat.as_mut().ok_or_else(|| { LlamaModelError::SamplerError("missing Mirostat GPU sampler".into()) })?; - hidden - .try_sample_mirostat2_token_q_mat_pending( - model.output_matrix(), + logits + .sample_mirostat2_token_pending( mirostat, &previous_tokens, previous_gpu_token, @@ -396,13 +398,8 @@ where } kalosm_language_model::SamplingStrategy::Standard => { let params = sampler.standard_params(top_k); - hidden - .try_sample_standard_token_q_mat_pending( - model.output_matrix(), - &previous_tokens, - previous_gpu_token, - params, - ) + logits + .sample_standard_token_pending(&previous_tokens, previous_gpu_token, params) .map_err(LlamaModelError::from) } } @@ -446,13 +443,15 @@ where Ok(hidden) => hidden, Err(err) => return Box::pin(async move { Err(err.into()) }), }; - let hidden = hidden.squeeze(0).to_concrete(); - let output_matrix = model.output_matrix().clone(); + let logits = hidden + .squeeze(0) + .to_concrete() + .q_mat_mul(model.output_matrix()); let mut kernels = 0; if trace { - if let Some(gpu_hidden) = hidden.as_gpu() { + if let Some(gpu_logits) = logits.as_gpu() { let resolve_start = Some(Instant::now()); - kernels = gpu_hidden.count_kernels_to_resolve(); + kernels = gpu_logits.count_kernels_to_resolve(); if let Some(start) = resolve_start { tracing::info!( "forward_resolve path={path} decode_eligible={decode_eligible} kernels={kernels} elapsed={:?}", @@ -469,25 +468,17 @@ where let mirostat = sampler.mirostat.as_mut().ok_or_else(|| { LlamaModelError::SamplerError("missing Mirostat GPU sampler".into()) })?; - hidden - .try_sample_mirostat2_token_q_mat( - &output_matrix, - mirostat, - &previous_tokens, - params, - ) + logits + .sample_mirostat2_token(mirostat, &previous_tokens, params) .await? } kalosm_language_model::SamplingStrategy::Standard => { let params = sampler.standard_params(top_k); - hidden - .try_sample_standard_token_q_mat(&output_matrix, &previous_tokens, params) + logits + .sample_standard_token(&previous_tokens, params) .await? } - } - .ok_or_else(|| { - LlamaModelError::SamplerError("fused logits sampler refused slow fallback".into()) - })?; + }; if let Some(start) = sample_start { tracing::info!( "forward_sample_token_download path={path} decode_eligible={decode_eligible} fused_logits=1 k={} elapsed={:?}", diff --git a/models/kalosm-llama/src/model/mod.rs b/models/kalosm-llama/src/model/mod.rs index 9fd7c78a9..def415db7 100644 --- a/models/kalosm-llama/src/model/mod.rs +++ b/models/kalosm-llama/src/model/mod.rs @@ -659,7 +659,8 @@ where }; let _ = hidden - .try_sample_mirostat2_token_q_mat(model.output_matrix(), &mut sampler, &[], params) + .q_mat_mul(model.output_matrix()) + .sample_mirostat2_token(&mut sampler, &[], params) .await; if let Some(start) = start { tracing::info!( From 5cb7007fde497d190cccd2641979c2966c16a152 Mon Sep 17 00:00:00 2001 From: Evan Almloff Date: Thu, 11 Jun 2026 18:13:32 -0500 Subject: [PATCH 03/16] remove output field --- .../src/compute_graph/resolve/fusion_row.rs | 15 +- fusor-ml/core/src/row_program.rs | 146 ++++++++++-------- 2 files changed, 90 insertions(+), 71 deletions(-) diff --git a/fusor-ml/core/src/compute_graph/resolve/fusion_row.rs b/fusor-ml/core/src/compute_graph/resolve/fusion_row.rs index c54b0c821..29c92e2ab 100644 --- a/fusor-ml/core/src/compute_graph/resolve/fusion_row.rs +++ b/fusor-ml/core/src/compute_graph/resolve/fusion_row.rs @@ -16,7 +16,7 @@ use petgraph::algo::toposort; use crate::{ nary_wise::NaryFunction, - row_program::{RowOutput, RowPhase, RowProgramOperation, RowReduce}, + row_program::{RowOutput, RowProgramOperation, RowReduce, RowStep}, }; use super::cluster_match::{keepdim_broadcast_layout, layout_matches, unary_elementwise}; @@ -488,20 +488,23 @@ impl Resolver { } let external_count = builder.externals.len(); - let phases: Vec = builder + let mut steps: Vec = builder .phases .into_iter() .map(|(_, mut phase)| { phase.expression = finalize_slots(&phase.expression, external_count); - RowPhase::Reduce(phase) + RowStep::Reduce(phase) }) .collect(); + steps.push(RowStep::Output(RowOutput::Map(finalize_slots( + &output_expr, + external_count, + )))); let operation = RowProgramOperation { inputs: builder.externals.clone(), shape: builder.shape, axis, - phases, - output: RowOutput::Map(finalize_slots(&output_expr, external_count)), + steps, output_datatype: root.output_datatype, dynamic_axis: None, }; @@ -511,7 +514,7 @@ impl Resolver { "row_fusion: committed root {:?} shape {:?} phases {} externals {}", self.execution_graph[root_idx].inner_idx, operation.shape, - operation.phases.len(), + operation.phase_count(), externals.len() ); } diff --git a/fusor-ml/core/src/row_program.rs b/fusor-ml/core/src/row_program.rs index f300ecf3a..dfbdcba28 100644 --- a/fusor-ml/core/src/row_program.rs +++ b/fusor-ml/core/src/row_program.rs @@ -76,19 +76,6 @@ pub(crate) struct RowFold { pub(crate) function: ReduceFunction, } -#[derive(Debug, Clone, Hash)] -pub(crate) enum RowPhase { - /// Fold over the row axis into a per-row scalar slot. - Reduce(RowReduce), - /// A staged per-element value: evaluated once at each axis position, - /// with one lane pinned per position (the axis must fit the workgroup). - Element { - expression: NaryExpr, - fold: Option, - datatype: DataTypeEnum, - }, -} - #[derive(Debug, Clone, Hash)] pub(crate) enum RowOutput { /// One output element per index-space position; output shape == `shape`. @@ -121,15 +108,32 @@ pub(crate) struct DynamicAxis { pub(crate) axis_bound_dim: Option, } +/// One ordered row-program step. All non-output steps produce slots for later +/// expressions; the final step must be [`RowStep::Output`] and defines the +/// tensor write contract. +#[derive(Debug, Clone, Hash)] +pub(crate) enum RowStep { + /// Fold over the row axis into a per-row scalar slot. + Reduce(RowReduce), + /// A staged per-element value: evaluated once at each axis position, + /// with one lane pinned per position (the axis must fit the workgroup). + Element { + expression: NaryExpr, + fold: Option, + datatype: DataTypeEnum, + }, + Output(RowOutput), +} + /// The algebraic shape the online-streaming lowering requires: a staged -/// score element, a max phase over a scaled score `e`, an exp-sum phase +/// score element, a max step over a scaled score `e`, an exp-sum step /// shifted by that max, a staged probability element, and a linear combine. /// The exp shift is what licenses streaming — rescaling the running sum and /// accumulators by `exp(M_old − M_new)` keeps them exact across tiles. struct OnlineSoftmax<'a> { - /// Element phase: the raw per-position score (slot `n`). + /// Element step: the raw per-position score (slot `n`). score: (&'a NaryExpr, &'a Option), - /// The scaled/masked score expression `e` over slot `n` (max phase body). + /// The scaled/masked score expression `e` over slot `n` (max step body). scaled: &'a NaryExpr, max_identity: NaryScalar, /// The per-position weight in `combine = p · weight`. @@ -161,28 +165,27 @@ fn binary_op_children<'a>(expr: &'a NaryExpr, op: &NaryOp) -> Option<(&'a NaryEx } } -/// Match the four-phase online-softmax shape (see [`OnlineSoftmax`]). The +/// Match the terminal-output online-softmax shape (see [`OnlineSoftmax`]). The /// attention constructor builds exactly this; the lowering re-derives it so -/// the phase expressions stay the single source of truth. -fn match_online_softmax<'a>( - phases: &'a [RowPhase], - output: &'a RowOutput, - input_count: usize, -) -> Option> { +/// the step expressions stay the single source of truth. +fn match_online_softmax<'a>(steps: &'a [RowStep], input_count: usize) -> Option> { let [ - RowPhase::Element { + RowStep::Element { expression: score, fold: score_fold, .. }, - RowPhase::Reduce(max_phase), - RowPhase::Reduce(sum_phase), - RowPhase::Element { + RowStep::Reduce(max_phase), + RowStep::Reduce(sum_phase), + RowStep::Element { expression: prob, fold: None, .. }, - ] = phases + RowStep::Output(RowOutput::Reduce { + combine, function, .. + }), + ] = steps else { return None; }; @@ -208,12 +211,6 @@ fn match_online_softmax<'a>( return None; } // combine: p · weight, summed - let RowOutput::Reduce { - combine, function, .. - } = output - else { - return None; - }; if function.op != ReduceOp::Sum { return None; } @@ -241,24 +238,38 @@ fn match_online_softmax<'a>( } /// Slot convention for every expression in the program: indices below -/// `inputs.len()` are tensor reads; index `inputs.len() + p` is phase `p`'s -/// value (a per-row scalar for reduce phases, a per-element value for -/// element phases). +/// `inputs.len()` are tensor reads; index `inputs.len() + p` is step `p`'s +/// value for non-output steps (a per-row scalar for reduce steps, a +/// per-element value for element steps). The final output step never produces +/// a reusable slot. #[derive(Debug, Clone)] pub(crate) struct RowProgramOperation { pub(crate) inputs: Vec, /// The full row-parallel index space (including the axis). pub(crate) shape: Box<[usize]>, pub(crate) axis: usize, - pub(crate) phases: Vec, - pub(crate) output: RowOutput, + pub(crate) steps: Vec, pub(crate) output_datatype: DataTypeEnum, /// `Some` exactly when the program pins one lane per axis position - /// (element phases / reducing output); `None` for chunked map programs. + /// (element steps / reducing output); `None` for chunked map programs. pub(crate) dynamic_axis: Option, } impl RowProgramOperation { + fn output_step(&self) -> &RowOutput { + match self.steps.last() { + Some(RowStep::Output(output)) => output, + _ => panic!("row program must end with an output step"), + } + } + + fn phase_steps(&self) -> &[RowStep] { + match self.steps.last() { + Some(RowStep::Output(_)) => &self.steps[..self.steps.len() - 1], + _ => panic!("row program must end with an output step"), + } + } + pub(crate) fn rows(&self) -> usize { self.shape .iter() @@ -276,7 +287,7 @@ impl RowProgramOperation { } pub(crate) fn out_shape(&self) -> Vec { - match &self.output { + match self.output_step() { RowOutput::Map(_) => self.shape.to_vec(), RowOutput::Reduce { free_dim, .. } => { let mut shape = self.row_shape(); @@ -286,6 +297,10 @@ impl RowProgramOperation { } } + pub(crate) fn phase_count(&self) -> usize { + self.phase_steps().len() + } + fn block(&self, device: &crate::Device) -> u32 { match &self.dynamic_axis { Some(dynamic) => dynamic.block, @@ -315,8 +330,7 @@ impl Operation for RowProgramOperation { } } self.axis.hash(state); - self.phases.hash(state); - self.output.hash(state); + self.steps.hash(state); self.output_datatype.hash(state); } @@ -381,8 +395,8 @@ impl Operation for RowProgramOperation { fn name(&self) -> String { format!( - "row_program_{}p_{}", - self.phases.len(), + "row_program_{}s_{}", + self.steps.len(), self.shape .iter() .map(|x| x.to_string()) @@ -513,7 +527,9 @@ fn build_row_program_kernel( // the online body over one tile, writing its unnormalized accumulator // and softmax statistics to scratch; a combine kernel folds the spans // with the online monoid. - let free_dim_out = match &operation.output { + let output_kind = operation.output_step().clone(); + let phase_steps = operation.phase_steps().to_vec(); + let free_dim_out = match &output_kind { RowOutput::Reduce { free_dim, .. } => Some(*free_dim), RowOutput::Map(_) => None, }; @@ -540,8 +556,6 @@ fn build_row_program_kernel( .dynamic_axis .as_ref() .and_then(|dynamic| dynamic.axis_bound_dim); - let phases = operation.phases.clone(); - let output_kind = operation.output.clone(); let output_dtype = output.datatype(); let output_value = MaybeQData::Tensor(output); let params_value = params.map(MaybeQData::Tensor); @@ -557,8 +571,8 @@ fn build_row_program_kernel( DataTypeEnum::F32, )) }); - let online_max_identity = match phases.get(1) { - Some(RowPhase::Reduce(reduce)) => Some(reduce.function.initial_value), + let online_max_identity = match phase_steps.get(1) { + Some(RowStep::Reduce(reduce)) => Some(reduce.function.initial_value), _ => None, }; let combine = if splits > 1 { @@ -680,11 +694,11 @@ fn build_row_program_kernel( let phase_handle = kb.program(); // Element-phase values the reducing output reads across lanes are // staged through workgroup memory (the probs of decode attention). - let staged: Vec> = phases + let staged: Vec> = phase_steps .iter() .enumerate() .map(|(p, phase)| match (&output_kind, phase) { - (RowOutput::Reduce { combine, .. }, RowPhase::Element { .. }) + (RowOutput::Reduce { combine, .. }, RowStep::Element { .. }) if combine.uses_input(input_count + p) => { Some(phase_handle.alloc_workgroup_array(ScalarElement::F32, block)) @@ -738,7 +752,9 @@ fn build_row_program_kernel( let RowOutput::Reduce { free_dim, .. } = &output_kind else { unreachable!("dynamic-axis row programs have a reducing output") }; - let online = match_online_softmax(&phases, &output_kind, input_count) + let mut online_steps = phase_steps.clone(); + online_steps.push(RowStep::Output(output_kind.clone())); + let online = match_online_softmax(&online_steps, input_count) .expect("dynamic-axis row programs are built in online-softmax shape"); let probs = staged[3] .as_ref() @@ -973,8 +989,8 @@ fn build_row_program_kernel( // are recomputed per phase — the same trade every multi-pass // normalization kernel makes. let chunks = k.div_ceil(block); - for phase in &phases { - let RowPhase::Reduce(reduce) = phase else { + for phase in &phase_steps { + let RowStep::Reduce(reduce) = phase else { unreachable!("element phases require a dynamic-axis row program") }; let phase_dtype = reduce.function.datatype(); @@ -1207,8 +1223,8 @@ pub(crate) fn attention_row_program( inputs: graph_inputs, shape: [batch, num_heads, q_seq_len, kv_len].into(), axis: 3, - phases: vec![ - RowPhase::Element { + steps: vec![ + RowStep::Element { expression: binary(NaryOp::Mul, q_read, k_read), fold: Some(RowFold { len: head_dim, @@ -1216,27 +1232,27 @@ pub(crate) fn attention_row_program( }), datatype: f32, }, - RowPhase::Reduce(RowReduce { + RowStep::Reduce(RowReduce { expression: scaled_score(), function: max_fn(f32), post_chain: UnaryFunctionChain::empty(f32), }), - RowPhase::Reduce(RowReduce { + RowStep::Reduce(RowReduce { expression: shifted_exp(), function: sum_fn(f32), post_chain: UnaryFunctionChain::empty(f32), }), - RowPhase::Element { + RowStep::Element { expression: binary(NaryOp::Div, shifted_exp(), slot(2)), fold: None, datatype: f32, }, + RowStep::Output(RowOutput::Reduce { + combine: binary(NaryOp::Mul, slot(3), v_read), + function: sum_fn(f32), + free_dim: head_dim, + }), ], - output: RowOutput::Reduce { - combine: binary(NaryOp::Mul, slot(3), v_read), - function: sum_fn(f32), - free_dim: head_dim, - }, output_datatype: input_dtype, dynamic_axis: Some(DynamicAxis { block, From 3a80cb7187c43687e5278d7e56e3dcc138555c78 Mon Sep 17 00:00:00 2001 From: Evan Almloff Date: Thu, 11 Jun 2026 20:26:14 -0500 Subject: [PATCH 04/16] remove specialized reduce kernels. everything goes through row program --- .../core/examples/bench_conv_implicit_gemm.rs | 186 +++ .../compute_graph/resolve/fusion_matmul.rs | 17 +- .../src/compute_graph/resolve/recognize.rs | 84 +- .../resolve/recognize_attention.rs | 8 +- fusor-ml/core/src/compute_graph/tests.rs | 88 ++ fusor-ml/core/src/lib.rs | 2 - fusor-ml/core/src/matmul/kernel.rs | 180 ++- fusor-ml/core/src/matmul/mod.rs | 83 +- fusor-ml/core/src/mir/tile_direct.rs | 79 +- fusor-ml/core/src/nary_direct.rs | 27 +- fusor-ml/core/src/reduce.rs | 38 +- fusor-ml/core/src/reduce_direct.rs | 496 -------- fusor-ml/core/src/reduce_tiled.rs | 1007 ----------------- fusor-ml/core/src/row_program.rs | 151 ++- fusor-ml/core/src/view.rs | 5 +- 15 files changed, 773 insertions(+), 1678 deletions(-) create mode 100644 fusor-ml/core/examples/bench_conv_implicit_gemm.rs delete mode 100644 fusor-ml/core/src/reduce_direct.rs delete mode 100644 fusor-ml/core/src/reduce_tiled.rs diff --git a/fusor-ml/core/examples/bench_conv_implicit_gemm.rs b/fusor-ml/core/examples/bench_conv_implicit_gemm.rs new file mode 100644 index 000000000..70407c0d5 --- /dev/null +++ b/fusor-ml/core/examples/bench_conv_implicit_gemm.rs @@ -0,0 +1,186 @@ +// A/B microbench: conv-as-im2col matmul with the gather materialization +// (the old pipeline) vs the implicit-GEMM unflatten (the new pipeline). +// +// Holding the flat [M, K] view tensor across materialization trips the +// resolver's live-reference guard, which declines the unflatten — giving +// exactly the old gather + matmul pipeline in the same binary. +// +// Run with: +// cargo run --package fusor-core --example bench_conv_implicit_gemm --release + +use std::time::{Duration, Instant}; + +use fusor_core::{Device, StrideSpec, Tensor}; + +const WARMUP: usize = 5; +const MEASURED: usize = 30; + +struct ConvCase { + name: &'static str, + b: usize, + c: usize, + h: usize, + w: usize, + n: usize, + kh: usize, + kw: usize, +} + +fn main() -> Result<(), Box> { + pollster::block_on(async { + let device = Device::new().await?; + println!("bench_conv_implicit_gemm (warmup {WARMUP}, measured {MEASURED})"); + println!(); + + let cases = [ + // Coop-tile-unaligned M: both variants run the generic reduce. + ConvCase { + name: "small_unaligned", + b: 2, + c: 8, + h: 16, + w: 16, + n: 16, + kh: 3, + kw: 3, + }, + ConvCase { + name: "mid_unaligned_n256", + b: 1, + c: 128, + h: 35, + w: 35, + n: 256, + kh: 3, + kw: 3, + }, + ConvCase { + name: "mid_unaligned_n128", + b: 1, + c: 128, + h: 35, + w: 35, + n: 128, + kh: 3, + kw: 3, + }, + ConvCase { + name: "mid_unaligned_n64", + b: 1, + c: 128, + h: 35, + w: 35, + n: 64, + kh: 3, + kw: 3, + }, + // Coop-tile-aligned M/K/N: the matmul runs the hardware kernel. + ConvCase { + name: "large_aligned", + b: 2, + c: 64, + h: 34, + w: 34, + n: 128, + kh: 3, + kw: 3, + }, + ConvCase { + name: "vision_aligned", + b: 1, + c: 256, + h: 66, + w: 66, + n: 256, + kh: 3, + kw: 3, + }, + ]; + + for case in &cases { + bench_case(&device, case); + } + Ok(()) + }) +} + +fn bench_case(device: &Device, case: &ConvCase) { + let &ConvCase { + name, + b, + c, + h, + w, + n, + kh, + kw, + } = case; + let (oh, ow) = (h - kh + 1, w - kw + 1); + let (m, k) = (b * oh * ow, c * kh * kw); + + let input_host: Vec = (0..b * c * h * w).map(|i| (i % 13) as f32 * 0.1).collect(); + let weight_host: Vec = (0..n * k).map(|i| (i % 7) as f32 * 0.01).collect(); + let input = Tensor::from_slice(device, [b, c, h, w], &input_host); + let weight = Tensor::from_slice(device, [n, k], &weight_host); + input.materialize_sync(); + weight.materialize_sync(); + + let build = |hold_flat: bool| -> (Tensor, Option) { + let windows = input.restride([ + StrideSpec::dim(0, b), + StrideSpec::dim_with(2, oh, 1), + StrideSpec::dim_with(3, ow, 1), + StrideSpec::dim(1, c), + StrideSpec::dim(2, kh), + StrideSpec::dim(3, kw), + ]); + let a = windows.reshape([m, k]); + let b_mat = weight.restride([StrideSpec::dim(1, k), StrideSpec::dim(0, n)]); + let out = a.mat_mul(&b_mat); + (out, hold_flat.then_some(a)) + }; + + let run = |hold_flat: bool| -> (Vec, usize) { + let mut kernels = 0; + for _ in 0..WARMUP { + let (out, held) = build(hold_flat); + kernels = out.count_kernels_to_resolve(); + device.poll_wait(); + drop(held); + drop(out); + } + let mut samples = Vec::with_capacity(MEASURED); + for _ in 0..MEASURED { + let (out, held) = build(hold_flat); + let start = Instant::now(); + let _ = out.count_kernels_to_resolve(); + device.poll_wait(); + samples.push(start.elapsed()); + drop(held); + drop(out); + } + samples.sort_unstable(); + (samples, kernels) + }; + + let (gather, gather_kernels) = run(true); + let (implicit, implicit_kernels) = run(false); + + let stats = |samples: &[Duration]| { + let mean = samples.iter().sum::() / samples.len() as u32; + let p50 = samples[samples.len() / 2]; + ( + mean.as_secs_f64() * 1000.0, + p50.as_secs_f64() * 1000.0, + samples[0].as_secs_f64() * 1000.0, + ) + }; + let (g_mean, g_p50, g_min) = stats(&gather); + let (i_mean, i_p50, i_min) = stats(&implicit); + + println!("{name}: conv {b}x{c}x{h}x{w} k{kh}x{kw} -> matmul {m}x{k} @ {k}x{n}"); + println!(" gather+matmul ({gather_kernels} dispatches): mean {g_mean:.3} ms, p50 {g_p50:.3} ms, min {g_min:.3} ms"); + println!(" implicit-GEMM ({implicit_kernels} dispatches): mean {i_mean:.3} ms, p50 {i_p50:.3} ms, min {i_min:.3} ms"); + println!(" speedup (p50): {:.2}x", g_p50 / i_p50); + println!(); +} diff --git a/fusor-ml/core/src/compute_graph/resolve/fusion_matmul.rs b/fusor-ml/core/src/compute_graph/resolve/fusion_matmul.rs index 5e42ae802..d1436ffa5 100644 --- a/fusor-ml/core/src/compute_graph/resolve/fusion_matmul.rs +++ b/fusor-ml/core/src/compute_graph/resolve/fusion_matmul.rs @@ -16,7 +16,13 @@ impl Resolver { && let Some(input_exec_idx) = self.get_input_node_in_exec_graph(input_inner) { let input_variant = self.execution_graph[input_exec_idx].variant.clone(); - if let ExecutionVariant::MatMul(matmul_op) = input_variant { + // An un-flattened operand was chosen for the coop kernel, + // which hosts no element-wise chains: fusing one here would + // demote the matmul to the generic divmod-per-load reduce. + if let ExecutionVariant::MatMul(matmul_op) = input_variant + && matmul_op.a.is_plain() + && matmul_op.b.is_plain() + { let mut new_matmul = matmul_op.clone(); let mut existing_post = new_matmul.post_element_wise.functions.clone(); existing_post.extend(el_op.functions.functions.iter().cloned()); @@ -384,8 +390,13 @@ impl Resolver { } } - // Pre-op: fuse elementwise before matmul inputs - if let ExecutionVariant::MatMul(matmul_op) = &node_variant { + // Pre-op: fuse elementwise before matmul inputs. Skipped for + // un-flattened operands: pre chains would demote the matmul off the + // coop kernel they were chosen for. + if let ExecutionVariant::MatMul(matmul_op) = &node_variant + && matmul_op.a.is_plain() + && matmul_op.b.is_plain() + { let mut new_matmul = matmul_op.clone(); let mut changed = false; diff --git a/fusor-ml/core/src/compute_graph/resolve/recognize.rs b/fusor-ml/core/src/compute_graph/resolve/recognize.rs index dfa84f560..ca89a31b3 100644 --- a/fusor-ml/core/src/compute_graph/resolve/recognize.rs +++ b/fusor-ml/core/src/compute_graph/resolve/recognize.rs @@ -272,7 +272,9 @@ impl Resolver { ); return true; } - if let Some((operation, inputs)) = contraction.to_mat_mul(&graph.device()) { + if let Some((mut operation, _)) = contraction.to_mat_mul(&graph.device()) { + self.try_unflatten_matmul_input(graph, &mut operation); + let inputs = [operation.first, operation.second]; self.commit_recognized( graph, node_idx, @@ -284,6 +286,86 @@ impl Resolver { false } + /// Read a recognized matmul's A operand through its un-flattened + /// producer. Conv's im2col flatten regroups a windowed view's dims + /// across overlapping strides, which no single strided layout can + /// express: the reshape becomes a chained view that would materialize + /// through the gather fallback. When the flat `[M, K]` operand is + /// exactly such a reinterpret, point the matmul at the producer and let + /// the kernels divmod the flat coordinates back apart per load — an + /// implicit GEMM with no gather dispatch. + fn try_unflatten_matmul_input( + &mut self, + graph: &ComputeGraphInner, + operation: &mut crate::MatMulOperation, + ) { + if !operation.a.is_plain() || operation.a.batch_dims != 0 { + return; + } + // Only the cooperative-matrix kernel reads an un-flattened operand + // faster than gather-then-matmul: its tile staging amortizes the + // per-load coordinate decomposition. The generic reduce re-derives + // coordinates for every load and measures slower than the gather at + // every meaningful size, so anything bound for it keeps the + // materialized matrix. + if !operation.hardware_matmul_statically_viable(&graph.device()) { + return; + } + let (m, k) = (operation.a.rows(), operation.a.cols()); + // An already-materialized (or externally held) operand is cheaper to + // read flat than to re-derive coordinates for. + if self.check_cached(graph, operation.first) || graph.has_live_reference(operation.first) { + return; + } + let Some(node) = graph.nodes.nodes.node_weight(operation.first) else { + return; + }; + let ComputeGraphNodeVariant::View(view) = &node.variant else { + return; + }; + if !view.is_fully_defined() + || !view.layout.is_contiguous() + || view.layout.offset() != 0 + || view.layout.shape() != [m, k] + { + return; + } + let inner_shape = &view.input_shape; + // The producer's dims must split cleanly into an `M` prefix and a + // `K` suffix for the per-side flat-coordinate decomposition. + let mut product = 1usize; + let mut k_start = inner_shape.len(); + while k_start > 0 && product < k { + k_start -= 1; + let Some(next) = product.checked_mul(inner_shape[k_start]) else { + return; + }; + product = next; + } + if product != k + || k_start == 0 + || k_start == inner_shape.len() + || inner_shape[..k_start].iter().product::() != m + { + return; + } + // The kernels decompose with u32 arithmetic; validate both sides + // here so the lowering can rely on it. + let probe = NaryExpr::DimIndex(0); + if crate::view::row_major_indices_from_flat(probe.clone(), &inner_shape[..k_start]) + .is_none() + || crate::view::row_major_indices_from_flat(probe, &inner_shape[k_start..]).is_none() + { + return; + } + operation.first = view.input; + operation.a = crate::matmul::MatrixOperand { + shape: inner_shape.clone(), + batch_dims: 0, + row_dims: k_start, + }; + } + /// Sweep elementwise nodes for quantized embedding gathers. pub(super) fn recognize_embeddings(&mut self, graph: &mut ComputeGraphInner) { let candidates: Vec = self diff --git a/fusor-ml/core/src/compute_graph/resolve/recognize_attention.rs b/fusor-ml/core/src/compute_graph/resolve/recognize_attention.rs index d6eb75998..d60fc99bc 100644 --- a/fusor-ml/core/src/compute_graph/resolve/recognize_attention.rs +++ b/fusor-ml/core/src/compute_graph/resolve/recognize_attention.rs @@ -164,12 +164,14 @@ impl Resolver { if !out.pre_element_wise[0].functions.is_empty() || !out.pre_element_wise[1].functions.is_empty() || !out.post_element_wise.functions.is_empty() + || !out.a.is_plain() + || !out.b.is_plain() { return None; } let probs_inner = out.first; let v_eff_inner = out.second; - let v_eff_shape = out.second_shape.clone(); + let v_eff_shape = out.b.shape.clone(); let datatype = out.datatype; // probs = softmax(scores) along the last axis of a rank-4 space — @@ -256,11 +258,13 @@ impl Resolver { if !qk.pre_element_wise[0].functions.is_empty() || !qk.pre_element_wise[1].functions.is_empty() || !qk.post_element_wise.functions.is_empty() + || !qk.a.is_plain() + || !qk.b.is_plain() { return None; } let q = qk.first; - let q_shape = qk.first_shape.to_vec(); + let q_shape = qk.a.shape.to_vec(); let kt_inner = qk.second; if q_shape.len() != 4 || q_shape[..3] != [batch, num_heads, q_seq_len] diff --git a/fusor-ml/core/src/compute_graph/tests.rs b/fusor-ml/core/src/compute_graph/tests.rs index b8c722aea..916572c64 100644 --- a/fusor-ml/core/src/compute_graph/tests.rs +++ b/fusor-ml/core/src/compute_graph/tests.rs @@ -385,6 +385,94 @@ fn three_way_chunk_cat_collapses() { }); } +// --- implicit-GEMM conv (resolve/recognize.rs unflatten) --- + +/// Conv's im2col composition: a sliding-window restride whose rank-2 flatten +/// regroups dims across overlapping strides, feeding the composed matmul. +/// When the contraction will reach the cooperative-matrix kernel, the +/// recognizer reads it through the un-flattened windows and the whole conv +/// is exactly one dispatch — no gather materializing the im2col matrix. +/// Shapes bound for the generic reduce keep the gather (measured faster +/// there), so the coop-unaligned case asserts two dispatches. +#[test] +fn conv_im2col_matmul_runs_without_gather() { + pollster::block_on(async { + let Ok(device) = Device::new().await else { + return; + }; + let coop_viable = device + .coop_token(crate::kernel_selection::CooperativeMatrixKind::F32F32M8N8K8) + .is_some() + && device + .subgroup_config() + .is_some_and(|config| config.is_fixed()); + + for (b, c, h, w, n, kh, kw, implicit) in [ + (2usize, 8usize, 16usize, 16usize, 16usize, 3usize, 3usize, false), + (2, 64, 34, 34, 128, 3, 3, coop_viable), + ] { + let (oh, ow) = (h - kh + 1, w - kw + 1); + let (m, k) = (b * oh * ow, c * kh * kw); + let input_host: Vec = (0..b * c * h * w).map(|i| (i % 13) as f32 * 0.1).collect(); + let weight_host: Vec = (0..n * k).map(|i| (i % 7) as f32 * 0.01).collect(); + let input = Tensor::from_slice(&device, [b, c, h, w], &input_host); + let weight = Tensor::from_slice(&device, [n, k], &weight_host); + + // The window views drop at the end of the block, like a conv + // layer's locals dropping before anything materializes. + let out = { + // sliding_window_view + window permutation in one restride: + // dims [b, oh, ow, c, kh, kw] over the (B, C, H, W) buffer. + let windows = input.restride([ + StrideSpec::dim(0, b), + StrideSpec::dim_with(2, oh, 1), + StrideSpec::dim_with(3, ow, 1), + StrideSpec::dim(1, c), + StrideSpec::dim(2, kh), + StrideSpec::dim(3, kw), + ]); + let a = windows.reshape([m, k]); + // Weight read through a transposed [K, N] view, like conv's + // weight.t(). + let b_mat = weight.restride([StrideSpec::dim(1, k), StrideSpec::dim(0, n)]); + a.mat_mul(&b_mat) + }; + let (_, kernels) = out.data.materialize(); + let expected = if implicit { 1 } else { 2 }; + assert_eq!( + kernels, expected, + "conv should be {expected} dispatch(es) (got {kernels})", + ); + + let result = out.as_slice::<2, f32>().await.unwrap(); + // Spot-check a strided sample against the host reference (the + // full reference is O(M·N·K) in a debug build). + for mi in (0..m).step_by(37) { + let (bi, rest) = (mi / (oh * ow), mi % (oh * ow)); + let (ohi, owi) = (rest / ow, rest % ow); + for ni in (0..n).step_by(13) { + let mut acc = 0f32; + for ci in 0..c { + for khi in 0..kh { + for kwi in 0..kw { + let iv = + input_host[((bi * c + ci) * h + ohi + khi) * w + owi + kwi]; + let wv = weight_host[ni * k + (ci * kh + khi) * kw + kwi]; + acc += iv * wv; + } + } + } + let got = result[[mi, ni]]; + assert!( + (got - acc).abs() < 1e-3 * acc.abs().max(1.0), + "mismatch at [{mi}, {ni}]: got {got}, expected {acc}", + ); + } + } + } + }); +} + #[test] fn deep_branch_chains_collapse_through_cat() { pollster::block_on(async { diff --git a/fusor-ml/core/src/lib.rs b/fusor-ml/core/src/lib.rs index c6c5054c9..1442515fe 100644 --- a/fusor-ml/core/src/lib.rs +++ b/fusor-ml/core/src/lib.rs @@ -41,8 +41,6 @@ mod pair_wise; mod quantized; mod rank; mod reduce; -mod reduce_direct; -mod reduce_tiled; mod row_program; mod sampling; mod slice_assign; diff --git a/fusor-ml/core/src/matmul/kernel.rs b/fusor-ml/core/src/matmul/kernel.rs index e27bd3fd9..d87502ae5 100644 --- a/fusor-ml/core/src/matmul/kernel.rs +++ b/fusor-ml/core/src/matmul/kernel.rs @@ -11,7 +11,8 @@ use crate::{ kernel_backend::{self, DirectKernel}, operation::Operation, tile_direct::{ - flatten_matrix_layout, tile_storage_read_with_direct_layout_typed, + flatten_matrix_layout, flatten_matrix_layout_split, + tile_storage_read_with_direct_layout_typed, tile_storage_write_with_direct_layout_typed, }, }, @@ -21,7 +22,7 @@ use crate::{ }; use super::{ - MatMulOperation, MatMulParams, coop_gemm, sgemm, sgemv, + MatMulOperation, MatMulParams, MatrixOperand, coop_gemm, sgemm, sgemv, variants::{CoopTile, dense_coop_kinds_from_datatype, select_dense_matmul_params}, }; @@ -81,8 +82,8 @@ impl MatMulOperation { Self { first, second, - first_shape: first_shape.into(), - second_shape: second_shape.into(), + a: MatrixOperand::plain(first_shape), + b: MatrixOperand::plain(second_shape), out_shape: out_shape.into(), datatype, pre_element_wise: [ @@ -94,10 +95,6 @@ impl MatMulOperation { } } - pub fn rank(&self) -> u32 { - self.out_shape.len() as u32 - } - fn can_use_hardware_matmul(&self) -> bool { matches!(self.datatype, DataTypeEnum::F32 | DataTypeEnum::F16) } @@ -109,15 +106,10 @@ impl MatMulOperation { /// hardware-specialized lower through this — the same generic tiled /// reduce any composed contraction gets. fn as_fused_reduce(&self) -> ReduceOperation { - let rank = self.first_shape.len(); - let batch = rank - 2; + let batch = self.a.batch_dims; let (m_dim, n_dim, k_dim) = (batch, batch + 1, batch + 2); - let mut index_space: Vec = self.first_shape[..batch].to_vec(); - index_space.extend([ - self.first_shape[batch], - self.second_shape[batch + 1], - self.first_shape[batch + 1], - ]); + let mut index_space: Vec = self.a.batch_shape().to_vec(); + index_space.extend([self.a.rows(), self.b.cols(), self.a.cols()]); let apply_chain = |mut expr: NaryExpr, chain: &UnaryFunctionChain| { for function in &chain.functions { @@ -139,14 +131,8 @@ impl MatMulOperation { } }; - let a_indices: Vec = (0..batch) - .chain([m_dim, k_dim]) - .map(NaryExpr::DimIndex) - .collect(); - let b_indices: Vec = (0..batch) - .chain([k_dim, n_dim]) - .map(NaryExpr::DimIndex) - .collect(); + let a_indices = self.a.index_expressions(m_dim, k_dim); + let b_indices = self.b.index_expressions(k_dim, n_dim); let acc_dtype = match self.pre_element_wise[0].out_datatype() { DataTypeEnum::U32 => DataTypeEnum::U32, @@ -195,6 +181,73 @@ impl MatMulOperation { } } + /// The static half of [`Self::build_hardware_matmul`]'s gates: whether + /// this contraction will reach the cooperative-matrix kernel on this + /// device. The resolver uses it to decide if reading an operand through + /// its un-flattened producer is profitable — the coop kernel's tile + /// staging amortizes the per-load coordinate decomposition, while the + /// generic reduce re-derives it for every load and loses to a one-time + /// gather. + pub(crate) fn hardware_matmul_statically_viable(&self, device: &Device) -> bool { + if !self.can_use_hardware_matmul() + || (self.datatype == DataTypeEnum::F16 && !device.f16_supported()) + || !self.pre_element_wise[0].functions.is_empty() + || !self.pre_element_wise[1].functions.is_empty() + || !self.post_element_wise.functions.is_empty() + { + return false; + } + let MatMulParams::CoopMatMul(params) = &self.parameters else { + return false; + }; + if device.coop_token(params.kind()).is_none() + || !device + .subgroup_config() + .is_some_and(|config| config.is_fixed()) + { + return false; + } + let (Ok(m), Ok(k), Ok(n)): (Result, Result, Result) = ( + self.a.rows().try_into(), + self.a.cols().try_into(), + self.b.cols().try_into(), + ) else { + return false; + }; + let Some(batch) = self + .a + .batch_shape() + .iter() + .try_fold(1u32, |acc, &dim| acc.checked_mul(u32::try_from(dim).ok()?)) + else { + return false; + }; + let limits = device.limits(); + let Some(tile) = CoopTile::select( + m, + k, + n, + limits + .max_compute_workgroup_size_x + .min(limits.max_compute_invocations_per_workgroup), + device + .subgroup_config() + .expect("checked above") + .max_size(), + ) else { + return false; + }; + // `CoopTile::select` guarantees the divisibility the kernel needs; + // only the dispatch-grid bound remains. + let Some(total_tiles) = (m / tile.bm) + .checked_mul(n / tile.bn) + .and_then(|tiles| tiles.checked_mul(batch)) + else { + return false; + }; + total_tiles <= limits.max_compute_workgroups_per_dimension + } + fn build_hardware_matmul( &self, device: &Device, @@ -202,22 +255,34 @@ impl MatMulOperation { input_b: &TensorData, output: &TensorData, ) -> Result { - let a_view = device_supported(flatten_matrix_layout(input_a.layout()))?; - let b_view = device_supported(flatten_matrix_layout(input_b.layout()))?; + let a_view = device_supported(flatten_matrix_layout_split( + input_a.layout(), + self.a.split(), + ))?; + let b_view = device_supported(flatten_matrix_layout_split( + input_b.layout(), + self.b.split(), + ))?; let y_view = device_supported(flatten_matrix_layout(output.layout()))?; - let rank = self.first_shape.len(); - let m: u32 = self.first_shape[rank - 2] + let m: u32 = self + .a + .rows() .try_into() .map_err(|_| kernel_backend::DeviceNotSupported)?; - let k: u32 = self.first_shape[rank - 1] + let k: u32 = self + .a + .cols() .try_into() .map_err(|_| kernel_backend::DeviceNotSupported)?; - let n: u32 = self.second_shape[rank - 1] + let n: u32 = self + .b + .cols() .try_into() .map_err(|_| kernel_backend::DeviceNotSupported)?; let batch: u32 = device_supported( - self.first_shape[..rank - 2] + self.a + .batch_shape() .iter() .try_fold(1usize, |acc, dim| acc.checked_mul(*dim)), )? @@ -347,8 +412,8 @@ struct HardwareMatmulVariant; impl Operation for MatMulOperation { fn hash_kernel_fields(&self, state: &mut FxHasher) { self.datatype.hash(state); - self.first_shape.hash(state); - self.second_shape.hash(state); + self.a.hash(state); + self.b.hash(state); self.out_shape.hash(state); self.pre_element_wise.hash(state); self.post_element_wise.hash(state); @@ -377,18 +442,15 @@ impl Operation for MatMulOperation { workgroup_shape: &crate::mir::workgroup_shape::WorkgroupShape, inputs: &[crate::mir::inputs::MirValue], ) -> [u32; 3] { - let [input_a, input_b, _output] = inputs else { + let [input_a, _input_b, _output] = inputs else { panic!("MatMulOperation requires 3 inputs"); }; + // The logical contraction shape: an un-flattened operand's runtime + // layout has a different rank, so the runtime layouts can't be used. let input_a = input_a.as_tensor().unwrap(); - let input_b = input_b.as_tensor().unwrap(); - let a_shape = input_a.layout().shape(); - let b_shape = input_b.layout().shape(); - let last_dim = self.rank() as usize - 1; - let last_dim_size = b_shape[last_dim]; - let second_to_last_dim = self.rank() as usize - 2; - let second_to_last_dim_size = a_shape[second_to_last_dim]; - let batch_size = a_shape.iter().rev().skip(2).product::(); + let last_dim_size = self.b.cols(); + let second_to_last_dim_size = self.a.rows(); + let batch_size = self.a.batch_shape().iter().product::(); match &self.parameters { MatMulParams::Vector(sgemv_params) => sgemv::dispatch_size( @@ -430,16 +492,12 @@ impl Operation for MatMulOperation { ) -> Vec { let a = nodes.get_result(self.first).unwrap(); let b = nodes.get_result(self.second).unwrap(); - let last_dim = self.rank() as usize - 1; - let second_to_last_dim = self.rank() as usize - 2; let device = a.device(); - let a_shape = a.layout().shape(); - let b_shape = b.layout().shape(); - let mut out_shape = a_shape.to_vec(); - out_shape[second_to_last_dim] = a_shape[second_to_last_dim]; - out_shape[last_dim] = b_shape[last_dim]; - let output_tensor = - TensorData::new_for_shape(device, &out_shape, self.post_element_wise.out_datatype()); + let output_tensor = TensorData::new_for_shape( + device, + &self.out_shape, + self.post_element_wise.out_datatype(), + ); vec![a.into(), b.into(), output_tensor.into()] } @@ -469,15 +527,11 @@ impl Operation for MatMulOperation { // generic tiled (or serial) fused reduce, identical to what any // unrecognized contraction gets. let reduce = self.as_fused_reduce(); - crate::reduce_tiled::build_reduce_tiled_kernel(&reduce, graph, workgroup_shape, inputs) - .or_else(|| { - crate::reduce_direct::build_reduce_direct_kernel( - &reduce, - graph, - workgroup_shape, - inputs, - ) - }) + crate::row_program::RowProgramOperation::from_reduce(&reduce).build_direct_kernel( + graph, + workgroup_shape, + inputs, + ) } fn output( @@ -493,12 +547,14 @@ impl Operation for MatMulOperation { format!( "matmul_{}_{}_by_{}", self.datatype, - self.first_shape + self.a + .shape .iter() .map(|s| s.to_string()) .collect::>() .join("x"), - self.second_shape + self.b + .shape .iter() .map(|s| s.to_string()) .collect::>() diff --git a/fusor-ml/core/src/matmul/mod.rs b/fusor-ml/core/src/matmul/mod.rs index 5eec4ec4d..0c01c18a4 100644 --- a/fusor-ml/core/src/matmul/mod.rs +++ b/fusor-ml/core/src/matmul/mod.rs @@ -25,13 +25,92 @@ pub enum MatMulParams { CoopMatMul(coop_gemm::CoopGemmParams), } +/// One matmul operand's dim grouping: the producer node's logical dims +/// split into `batch_dims` leading batch dims, `row_dims` row dims, and +/// column dims for the rest. The row and column groups flatten to the two +/// matrix axes. A plain `[batch.., rows, cols]` operand has one dim per +/// group; conv's im2col operand keeps the windowed view's dims, and the +/// kernels divmod the flat matrix coordinates back apart per load. +#[derive(Debug, Clone, Hash)] +pub(crate) struct MatrixOperand { + pub(crate) shape: Box<[usize]>, + pub(crate) batch_dims: usize, + pub(crate) row_dims: usize, +} + +impl MatrixOperand { + pub(crate) fn plain(shape: &[usize]) -> Self { + assert!(shape.len() >= 2, "matrix operands are at least rank 2"); + Self { + shape: shape.into(), + batch_dims: shape.len() - 2, + row_dims: 1, + } + } + + /// One dim per group: the operand's shape is the logical matmul shape. + pub(crate) fn is_plain(&self) -> bool { + self.row_dims == 1 && self.batch_dims + 2 == self.shape.len() + } + + pub(crate) fn batch_shape(&self) -> &[usize] { + &self.shape[..self.batch_dims] + } + + pub(crate) fn row_shape(&self) -> &[usize] { + &self.shape[self.batch_dims..self.batch_dims + self.row_dims] + } + + pub(crate) fn col_shape(&self) -> &[usize] { + &self.shape[self.batch_dims + self.row_dims..] + } + + pub(crate) fn rows(&self) -> usize { + self.row_shape().iter().product() + } + + pub(crate) fn cols(&self) -> usize { + self.col_shape().iter().product() + } + + /// Leading dim count flattening to the 2-D matrix row axis (batch + /// included): the split for [`flatten_matrix_layout_split`]. + /// + /// [`flatten_matrix_layout_split`]: crate::mir::tile_direct::flatten_matrix_layout_split + pub(crate) fn split(&self) -> usize { + self.batch_dims + self.row_dims + } + + /// Index expressions reading the operand at the contraction coordinates: + /// batch dims index through, and the flat row/column coordinates + /// decompose over the row/column groups (the identity for single-dim + /// groups, so plain operands load with bare `DimIndex` coordinates). + pub(crate) fn index_expressions( + &self, + row_dim: usize, + col_dim: usize, + ) -> Vec { + use crate::nary_wise::NaryExpr; + let mut indices: Vec = (0..self.batch_dims).map(NaryExpr::DimIndex).collect(); + indices.extend( + crate::view::row_major_indices_from_flat(NaryExpr::DimIndex(row_dim), self.row_shape()) + .expect("operand dims fit u32"), + ); + indices.extend( + crate::view::row_major_indices_from_flat(NaryExpr::DimIndex(col_dim), self.col_shape()) + .expect("operand dims fit u32"), + ); + indices + } +} + #[derive(Debug, Clone)] pub(crate) struct MatMulOperation { pub(crate) datatype: DataTypeEnum, pub(crate) first: NodeIndex, pub(crate) second: NodeIndex, - pub(crate) first_shape: Box<[usize]>, - pub(crate) second_shape: Box<[usize]>, + pub(crate) a: MatrixOperand, + pub(crate) b: MatrixOperand, pub(crate) out_shape: Box<[usize]>, pub(crate) pre_element_wise: [UnaryFunctionChain; 2], pub(crate) post_element_wise: UnaryFunctionChain, diff --git a/fusor-ml/core/src/mir/tile_direct.rs b/fusor-ml/core/src/mir/tile_direct.rs index 4b9ef42a4..d052df39d 100644 --- a/fusor-ml/core/src/mir/tile_direct.rs +++ b/fusor-ml/core/src/mir/tile_direct.rs @@ -10,24 +10,42 @@ pub(crate) struct DirectMatrixLayout { } pub(crate) fn flatten_matrix_layout(layout: &Layout) -> Option { + flatten_matrix_layout_split(layout, layout.shape().len().checked_sub(1)?) +} + +/// Flatten a strided layout into a 2-D matrix view: `shape[..row_dims]` +/// flattens to rows, `shape[row_dims..]` to columns. Sides whose dims merge +/// affinely use a plain strided layout; anything else (a conv im2col window, +/// a non-affine batch prefix) becomes a `MultiFlattenMap` whose sub-axes +/// divmod the flat coordinate back apart per load. +pub(crate) fn flatten_matrix_layout_split( + layout: &Layout, + row_dims: usize, +) -> Option { let shape = layout.shape(); let strides = layout.strides(); - if shape.len() < 2 || shape.contains(&0) { + if row_dims == 0 || row_dims >= shape.len() || shape.contains(&0) { return None; } - let rows = shape[..shape.len() - 1] + let rows = shape[..row_dims] + .iter() + .try_fold(1usize, |acc, dim| acc.checked_mul(*dim))?; + let cols = shape[row_dims..] .iter() .try_fold(1usize, |acc, dim| acc.checked_mul(*dim))?; - let cols = *shape.last()?; let rows_u32 = rows.try_into().ok()?; let cols_u32 = cols.try_into().ok()?; let offset = layout.offset().try_into().ok()?; - let prefix_is_affine = (0..shape.len().saturating_sub(2)) - .all(|axis| strides[axis] == strides[axis + 1].saturating_mul(shape[axis + 1])); + let side_is_affine = |range: std::ops::Range| { + range + .clone() + .zip(range.skip(1)) + .all(|(axis, next)| strides[axis] == strides[next].saturating_mul(shape[next])) + }; - let layout = if prefix_is_affine { - let row_stride: u32 = strides[shape.len() - 2].try_into().ok()?; + let layout = if side_is_affine(0..row_dims) && side_is_affine(row_dims..shape.len()) { + let row_stride: u32 = strides[row_dims - 1].try_into().ok()?; let col_stride: u32 = strides[shape.len() - 1].try_into().ok()?; tile_ir::Layout::strided( tile_ir::MemoryLevel::Storage, @@ -35,37 +53,32 @@ pub(crate) fn flatten_matrix_layout(layout: &Layout) -> Option = shape[..shape.len() - 1] - .iter() - .copied() - .map(u32::try_from) - .collect::, _>>() - .ok()?; - let prefix_strides: Vec = strides[..strides.len() - 1] - .iter() - .copied() - .map(u32::try_from) - .collect::, _>>() - .ok()?; - let column_stride: u32 = strides[strides.len() - 1].try_into().ok()?; - let m_group = tile_ir::AxisGroup { - sub_axes: prefix_shape - .iter() - .zip(prefix_strides.iter()) - .map(|(&extent, &stride)| tile_ir::SubAxis { extent, stride }) - .collect(), - }; - let k_group = tile_ir::AxisGroup { - sub_axes: vec![tile_ir::SubAxis { - extent: cols_u32, - stride: column_stride, - }], + let axis_group = |range: std::ops::Range| -> Option { + let mut sub_axes = Vec::with_capacity(range.len()); + for axis in range { + // Extent-1 axes contribute nothing to the flat coordinate + // decomposition; dropping them saves a divmod per load. + if shape[axis] == 1 { + continue; + } + sub_axes.push(tile_ir::SubAxis { + extent: shape[axis].try_into().ok()?, + stride: strides[axis].try_into().ok()?, + }); + } + if sub_axes.is_empty() { + sub_axes.push(tile_ir::SubAxis { + extent: 1, + stride: 0, + }); + } + Some(tile_ir::AxisGroup { sub_axes }) }; tile_ir::Layout::with_indexing( tile_ir::MemoryLevel::Storage, tile_ir::Shape::new([rows_u32, cols_u32]), tile_ir::MultiFlattenMap { - groups: vec![m_group, k_group], + groups: vec![axis_group(0..row_dims)?, axis_group(row_dims..shape.len())?], }, ) }; diff --git a/fusor-ml/core/src/nary_direct.rs b/fusor-ml/core/src/nary_direct.rs index 3c33a0dcd..687a87c3c 100644 --- a/fusor-ml/core/src/nary_direct.rs +++ b/fusor-ml/core/src/nary_direct.rs @@ -202,7 +202,7 @@ fn plan_nary_tiling( } let invariant_bytes: u64 = (0..input_count) .filter(|&i| !access.depends_on(i, dim)) - .map(|i| crate::reduce_tiled::input_allocation_bytes(&metas[i], &values[i])) + .map(|i| input_allocation_bytes(&metas[i], &values[i])) .sum(); if invariant_bytes < device.last_level_cache_bytes() { continue; @@ -382,20 +382,6 @@ impl ValueTile { } } - pub(crate) fn into_f16(self) -> tile_ir::tile::Tile { - match self.cast_to(DataTypeEnum::F16) { - Self::F16(v) => v, - _ => unreachable!(), - } - } - - pub(crate) fn into_u32(self) -> tile_ir::tile::Tile { - match self.cast_to(DataTypeEnum::U32) { - Self::U32(v) => v, - _ => unreachable!(), - } - } - pub(crate) fn into_mask(self) -> tile_ir::tile::Mask { match self { Self::Bool(v) => v, @@ -1143,6 +1129,17 @@ pub(crate) struct TensorMeta { pub(crate) allocation_len: u32, } +pub(crate) fn input_allocation_bytes(meta: &TensorMeta, value: &MaybeQData) -> u64 { + let element_bytes = match value { + MaybeQData::Tensor(tensor) => match tensor.datatype() { + DataTypeEnum::F16 => 2, + DataTypeEnum::F32 | DataTypeEnum::U32 => 4, + }, + MaybeQData::QMatrix(_) => 4, + }; + meta.allocation_len as u64 * element_bytes +} + impl TensorMeta { /// Metadata for a quantized-or-dense matrix input. Quantized matrices /// load by (row, col), so the strides flatten the leading dims row-major diff --git a/fusor-ml/core/src/reduce.rs b/fusor-ml/core/src/reduce.rs index 604a4e249..de97aa444 100644 --- a/fusor-ml/core/src/reduce.rs +++ b/fusor-ml/core/src/reduce.rs @@ -90,10 +90,6 @@ impl ReduceOperation { .filter_map(|(i, x)| (i != self.axis).then_some(*x)) .collect() } - - pub(crate) fn reduce_size(&self) -> usize { - self.shape[self.axis] - } } impl Operation for ReduceOperation { @@ -119,22 +115,13 @@ impl Operation for ReduceOperation { fn dispatch_size( &self, - workgroup_shape: &crate::mir::workgroup_shape::WorkgroupShape, + _workgroup_shape: &crate::mir::workgroup_shape::WorkgroupShape, inputs: &[MirValue], ) -> [u32; 3] { let output_tensor: TensorData = inputs.last().unwrap().as_tensor().unwrap().clone(); - let total_outputs = output_tensor.layout().shape().iter().product::() as u32; - let reduce_size = self.reduce_size() as u32; - let serial_workgroups = total_outputs.div_ceil(workgroup_shape.x()); - let total_workgroups = - if use_cooperative_reduce(total_outputs, reduce_size, workgroup_shape.x()) { - total_outputs - } else { - serial_workgroups - }; - + let rows = output_tensor.layout().shape().iter().product::() as u32; distribute_workgroups( - total_workgroups, + rows, output_tensor .device() .limits() @@ -182,15 +169,11 @@ impl Operation for ReduceOperation { workgroup_shape: &crate::mir::workgroup_shape::WorkgroupShape, inputs: &[MirValue], ) -> Option { - crate::reduce_tiled::build_reduce_tiled_kernel(self, graph, workgroup_shape, inputs) - .or_else(|| { - crate::reduce_direct::build_reduce_direct_kernel( - self, - graph, - workgroup_shape, - inputs, - ) - }) + crate::row_program::RowProgramOperation::from_reduce(self).build_direct_kernel( + graph, + workgroup_shape, + inputs, + ) } fn output(&self, _: &crate::compute_graph::ComputeGraphInner, inputs: &[MirValue]) -> MirValue { @@ -207,11 +190,6 @@ impl Operation for ReduceOperation { } } -pub(crate) fn use_cooperative_reduce(total_outputs: u32, reduce_size: u32, block: u32) -> bool { - let serial_workgroups = total_outputs.div_ceil(block); - reduce_size >= block && serial_workgroups <= 4 -} - #[derive(Clone, Debug, Hash)] pub struct ReduceFunction { pub(crate) name: Option, diff --git a/fusor-ml/core/src/reduce_direct.rs b/fusor-ml/core/src/reduce_direct.rs deleted file mode 100644 index 4d18d4e4d..000000000 --- a/fusor-ml/core/src/reduce_direct.rs +++ /dev/null @@ -1,496 +0,0 @@ -use fusor_tile_ir as tile_ir; -use fusor_tile_ir_kernels as tile_ir_kernels; -use std::hash::Hash; - -use crate::{ - access_analysis::InputAccesses, - mir::{ - inputs::MirValue, kernel_backend, kernel_backend::DirectKernel, operation::Operation, - workgroup_shape::WorkgroupShape, - }, - nary_direct::{ - ValueTile, apply_unary_function_chain, declare_value, eval_nary_expr, linear_group, - output_dims_from_flat, tile_u32, - }, - nary_wise::NaryScalar, - reduce::{ReduceOp, ReduceOperation}, - reduce_tiled::{datatype_scalar, input_allocation_bytes}, - tensor::DataTypeEnum, - visit_tiled::{MaybeQData, distribute_workgroups}, -}; - -const BLOCK: usize = 256; -/// Output rows per workgroup on the subgroup-per-output route. -const SUBGROUP_ROWS: u32 = 4; -/// Consecutive reduce-axis values each lane folds per iteration on the -/// subgroup-per-output route (matches the dense gemv kernel's run length). -const SUBGROUP_VALUES_PER_LANE: u32 = 4; - -struct ReduceDirectKernelVariant; - -/// How the fused reduce distributes its fold across threads. -#[derive(Clone, Copy)] -enum ReduceRoute { - /// One workgroup per output, 256 lanes splitting the reduce axis. The - /// final cross-lane combine uses subgroup collectives (one reduction per - /// subgroup, partials through workgroup memory, one more reduction) when - /// the device has a fixed subgroup width; otherwise the scratch tree. - Cooperative { - subgroup_finish: Option, - }, - /// One subgroup per output: lanes walk the reduce axis in - /// `SUBGROUP_VALUES_PER_LANE`-long runs (consecutive lanes read adjacent - /// runs — coalesced when the dominant input is contiguous along the - /// axis) and a single subgroup collective combines them. - SubgroupRow(tile_ir_kernels::SubgroupConfig), - /// One thread per output with a per-thread fold. - Serial, -} - -/// The subgroup-per-output route only pays when adjacent lanes read adjacent -/// addresses: the byte-dominant reduce-axis-dependent input must be -/// contiguous along the axis. Otherwise the serial layout (adjacent threads -/// = adjacent outputs) is the coalesced one. -fn dominant_k_dep_input_is_k_contiguous( - operation: &ReduceOperation, - values: &[MaybeQData], -) -> bool { - let Some(metas) = values - .iter() - .map(crate::reduce_tiled::analysis_meta) - .collect::>>() - else { - return false; - }; - let Some(access) = - InputAccesses::collect(&operation.expression, operation.inputs.len(), &metas) - else { - return false; - }; - let axis = operation.axis; - let mut best: Option<(u64, bool)> = None; - for i in 0..operation.inputs.len() { - if !access.depends_on(i, axis) { - continue; - } - let contiguous = access.dims[i] - .iter() - .enumerate() - .any(|(j, &d)| d == axis && metas[i].strides.get(j).copied() == Some(1)); - let bytes = input_allocation_bytes(&metas[i], &values[i]); - if best.as_ref().is_none_or(|(b, _)| bytes > *b) { - best = Some((bytes, contiguous)); - } - } - best.is_some_and(|(_, contiguous)| contiguous) -} - -/// The shared preamble for every fused-reduce kernel builder: split the -/// MirValues into producer values + output, gate on device support, and -/// declare the kernel bindings. -pub(crate) struct ReduceKernelInputs { - pub(crate) values: Vec, - pub(crate) output_shape: Vec, - pub(crate) total_outputs: u32, -} - -impl ReduceKernelInputs { - pub(crate) fn parse( - operation: &ReduceOperation, - graph: &crate::compute_graph::ComputeGraphInner, - inputs: &[MirValue], - ) -> Option { - let (output, producers) = inputs.split_last()?; - let output = output.as_tensor()?; - let values = producers - .iter() - .map(|input| MaybeQData::try_from(input.clone()).ok()) - .collect::>>()?; - - let f16_unsupported = !graph.device().f16_supported(); - if f16_unsupported { - let uses_f16 = output.datatype() == DataTypeEnum::F16 - || operation.function.datatype() == DataTypeEnum::F16 - || values.iter().any(|value| match value { - MaybeQData::Tensor(tensor) => tensor.datatype() == DataTypeEnum::F16, - MaybeQData::QMatrix(matrix) => { - matches!(matrix.datatype(), fusor_gguf::GgmlType::F16) - } - }); - if uses_f16 { - return None; - } - } - - if operation.post_element_wise.input_datatype() != operation.function.datatype() { - return None; - } - - let output_shape = output - .layout() - .shape() - .iter() - .copied() - .map(u32::try_from) - .collect::, _>>() - .ok()?; - let total_outputs = output_shape - .iter() - .try_fold(1u32, |acc, dim| acc.checked_mul(*dim))?; - - Some(Self { - values, - output_shape, - total_outputs, - }) - } -} - -pub(crate) fn build_reduce_direct_kernel( - operation: &ReduceOperation, - graph: &crate::compute_graph::ComputeGraphInner, - workgroup_shape: &WorkgroupShape, - inputs: &[MirValue], -) -> Option { - let parsed = ReduceKernelInputs::parse(operation, graph, inputs)?; - let output = inputs.last()?.as_tensor()?.clone(); - let reduce_size: u32 = operation.reduce_size().try_into().ok()?; - - let reduce_dtype = operation.function.datatype(); - let reduce_op = tile_reduce_op(operation.function.op); - let initial = operation.function.initial_value; - - let device = graph.device(); - let limits = device.limits(); - let subgroup_config = device.subgroup_config(); - let route = - if crate::reduce::use_cooperative_reduce(parsed.total_outputs, reduce_size, BLOCK as u32) { - // The two-stage subgroup finish needs a compile-time subgroup count - // small enough for the second collective to cover the partials. - let subgroup_finish = subgroup_config.filter(|config| { - config.is_fixed() - && (BLOCK as u32).is_multiple_of(config.max_size()) - && BLOCK as u32 / config.max_size() <= config.max_size() - }); - ReduceRoute::Cooperative { subgroup_finish } - } else { - let row_config = subgroup_config.filter(|config| { - let block = config.block_for_subgroups(SUBGROUP_ROWS); - reduce_size >= 32 - && block.is_power_of_two() - && block - <= limits - .max_compute_workgroup_size_x - .min(limits.max_compute_invocations_per_workgroup) - && dominant_k_dep_input_is_k_contiguous(operation, &parsed.values) - }); - match row_config { - Some(config) => ReduceRoute::SubgroupRow(config), - None => ReduceRoute::Serial, - } - }; - let dispatch_size = match route { - ReduceRoute::SubgroupRow(_) => distribute_workgroups( - parsed.total_outputs.div_ceil(SUBGROUP_ROWS), - limits.max_compute_workgroups_per_dimension, - ), - _ => operation.dispatch_size(workgroup_shape, inputs), - }; - let variant = - kernel_backend::KernelVariantKey::with_payload::(|state| { - match route { - ReduceRoute::Cooperative { subgroup_finish } => { - 0u8.hash(state); - subgroup_finish.hash(state); - } - ReduceRoute::SubgroupRow(config) => { - 1u8.hash(state); - config.hash(state); - } - ReduceRoute::Serial => 2u8.hash(state), - } - }); - let cache_key = operation.kernel_cache_key_with_dispatch( - variant, - Some(workgroup_shape), - dispatch_size, - inputs, - ); - - let axis = operation.axis; - let expression = operation.expression.clone(); - let post_chain = operation.post_element_wise.clone(); - let output_dtype = output.datatype(); - let output_value = MaybeQData::Tensor(output); - let ReduceKernelInputs { - values, - output_shape, - total_outputs, - } = parsed; - - kernel_backend::run_kernel( - graph.device().kernel_cache(), - operation.name(), - cache_key, - dispatch_size, - move |kb| { - let mut storages = Vec::with_capacity(values.len()); - let mut metas = Vec::with_capacity(values.len()); - for value in &values { - let (storage, meta) = declare_value(kb, value, false)?; - storages.push(storage); - metas.push(meta); - } - let (output_storage, output_meta) = declare_value(kb, &output_value, true)?; - - // Evaluate the fused producer at one index-space coordinate: the - // output dims with the running reduce index inserted at `axis`. - let value_at = |program: &mut tile_ir::tile::TileBlock<'_>, - dims_out: &[tile_ir::tile::Tile], - reduce_index: tile_ir::tile::Tile, - mask: tile_ir::tile::Mask| { - let mut coords = dims_out.to_vec(); - coords.insert(axis, reduce_index); - let (value, _) = - eval_nary_expr(program, &expression, &coords, &storages, &metas, mask, &[]); - value.cast_to(reduce_dtype) - }; - - let store_reduced = |program: &mut tile_ir::tile::TileBlock<'_>, - reduced: ValueTile, - dims_out: &[tile_ir::tile::Tile], - mask: tile_ir::tile::Mask| { - let (reduced, reduced_ty) = - apply_unary_function_chain(reduced.into_f32(), reduce_dtype, &post_chain) - .expect("validated reduce post_element_wise chain"); - let reduced = ValueTile::F32(reduced) - .cast_to(reduced_ty) - .cast_to(output_dtype); - let output_index = crate::nary_direct::layout_index(&output_meta, dims_out); - output_storage.store(program, output_index, reduced, mask); - }; - - let identity = || tile_ir::tile::Tile::literal(tile_literal_for(initial, reduce_dtype)); - let untag = |value: ValueTile| match reduce_dtype { - DataTypeEnum::F32 => value.into_f32(), - DataTypeEnum::F16 => value.into_f16(), - DataTypeEnum::U32 => value.into_u32(), - }; - let tag = |value: tile_ir::tile::Tile| match reduce_dtype { - DataTypeEnum::F32 => ValueTile::F32(value), - DataTypeEnum::F16 => ValueTile::F16(value), - DataTypeEnum::U32 => ValueTile::U32(value), - }; - - // The per-lane fold every route shares: `k_of(loop_index, run)` - // gives this lane's reduce coordinate for each of `runs` values - // per iteration, and out-of-range slots collapse to the identity. - // A route whose coordinates never leave the axis (serial) skips - // the select by passing `mask_oob: false`. - let fold_reduce_lane = - |program: &mut tile_ir::tile::TileBlock<'_>, - dims_out: &[tile_ir::tile::Tile], - iterations: tile_ir::tile::FoldIter, - in_bounds: &tile_ir::tile::Mask, - runs: u32, - mask_oob: bool, - k_of: &dyn Fn(&tile_ir::tile::Tile, u32) -> tile_ir::tile::Tile| - -> tile_ir::tile::Tile { - let reduce_binary = reduce_op.binary(); - let [acc] = - program.fold(iterations, [identity()], |program, loop_index, [acc]| { - let mut acc = acc; - for run in 0..runs { - let k_index = k_of(&loop_index, run); - let active = in_bounds.clone().and(k_index.clone().lt(reduce_size)); - let value = value_at(program, dims_out, k_index, active.clone()); - let value = if mask_oob { - tile_ir::tile::Tile::select(active, untag(value), identity()) - } else { - untag(value) - }; - acc = acc.binary(reduce_binary, value); - } - [acc] - }); - acc - }; - - match route { - ReduceRoute::Cooperative { subgroup_finish } => { - let chunks = reduce_size.div_ceil(BLOCK as u32); - let phase = kb.program(); - let scratch = subgroup_finish - .map(|config| BLOCK as u32 / config.max_size()) - .filter(|partials| *partials > 1) - .map(|partials| { - phase.alloc_workgroup_array(datatype_scalar(reduce_dtype), partials) - }); - phase.program_grid(BLOCK as u32, dispatch_size, |program| { - let lane = program.lane(); - let output_flat = linear_group(program, dispatch_size); - let in_bounds = output_flat.lt(total_outputs); - let dims_out = - output_dims_from_flat_usize(output_flat.clone(), &output_shape); - - let partial = fold_reduce_lane( - program, - &dims_out, - tile_ir::tile::range(chunks), - &in_bounds, - 1, - true, - &|loop_index, _| loop_index.clone() * BLOCK as u32 + lane.clone(), - ); - let (reduced, store_mask) = match subgroup_finish { - // One collective per subgroup, partials through - // workgroup memory, one more collective: two - // barriers and two subgroup ops instead of the - // log2(block) scratch-tree rounds. - Some(config) => { - let token = config.token(); - let first = token.subgroup_reduce(program, reduce_op, partial); - let partials = BLOCK as u32 / config.max_size(); - let sg_lane = token.subgroup_lane(program); - let sg_id = token.subgroup_id(program); - let reduced = if partials <= 1 { - first - } else { - let scratch = scratch.as_ref().unwrap(); - let first = program.bind(first); - program.if_then(sg_lane.clone().eq(0u32), |program| { - program.store_workgroup( - scratch, - sg_id.clone(), - first.clone(), - ); - }); - program.workgroup_barrier(); - let slot = tile_ir::tile::Tile::select( - lane.clone().lt(partials), - lane.clone(), - tile_u32(0), - ); - let loaded = program.load_workgroup(scratch, slot); - let masked = tile_ir::tile::Tile::select( - lane.clone().lt(partials), - loaded, - identity(), - ); - token.subgroup_reduce(program, reduce_op, masked) - }; - // The combined value is uniform across - // subgroup 0; store from its first lane - // rather than assuming workgroup lane 0 - // belongs to it. - let store_mask = - in_bounds.clone().and(sg_id.eq(0u32)).and(sg_lane.eq(0u32)); - (reduced, store_mask) - } - None => ( - program.group_reduce(reduce_op, BLOCK as u32, partial), - in_bounds.clone().and(lane.eq(0u32)), - ), - }; - store_reduced(program, tag(reduced), &dims_out, store_mask); - }); - } - ReduceRoute::SubgroupRow(config) => { - let block = config.block_for_subgroups(SUBGROUP_ROWS); - let token = config.token(); - kb.program().program_grid(block, dispatch_size, |program| { - let group = linear_group(program, dispatch_size); - let row = group * SUBGROUP_ROWS + token.subgroup_id(program); - let in_bounds = row.clone().lt(total_outputs); - let dims_out = output_dims_from_flat_usize(row, &output_shape); - let sg_lane = token.subgroup_lane(program); - let k_per_iter = - program.bind(token.subgroup_size(program) * SUBGROUP_VALUES_PER_LANE); - let k_iterations = (tile_u32(reduce_size) + k_per_iter.clone() - 1u32) - / k_per_iter.clone(); - - let acc = fold_reduce_lane( - program, - &dims_out, - tile_ir::tile::range(k_iterations), - &in_bounds, - SUBGROUP_VALUES_PER_LANE, - true, - &|loop_index, run| { - loop_index.clone() * k_per_iter.clone() - + sg_lane.clone() * SUBGROUP_VALUES_PER_LANE - + run - }, - ); - let reduced = token.subgroup_reduce(program, reduce_op, acc); - let store_mask = in_bounds.and(sg_lane.eq(0u32)); - store_reduced(program, tag(reduced), &dims_out, store_mask); - }); - } - ReduceRoute::Serial => { - kb.program() - .program_grid(BLOCK as u32, dispatch_size, |program| { - let lane = program.lane(); - let group = linear_group(program, dispatch_size); - let flat = group * BLOCK as u32 + lane.clone(); - let in_bounds = flat.lt(total_outputs); - let dims_out = output_dims_from_flat_usize(flat.clone(), &output_shape); - - let acc = fold_reduce_lane( - program, - &dims_out, - tile_ir::tile::range(reduce_size), - &in_bounds, - 1, - false, - &|loop_index, _| loop_index.clone(), - ); - store_reduced(program, tag(acc), &dims_out, in_bounds); - }); - } - } - Some(()) - }, - ) -} - -pub(crate) fn tile_literal_for(value: NaryScalar, target: DataTypeEnum) -> tile_ir::TileLiteral { - match target { - DataTypeEnum::F32 => match value { - NaryScalar::F32(value) => tile_ir::TileLiteral::f32(value), - NaryScalar::F16(value) => tile_ir::TileLiteral::f32(value.to_f32()), - NaryScalar::U32(value) => tile_ir::TileLiteral::f32(value as f32), - }, - DataTypeEnum::F16 => match value { - NaryScalar::F32(value) => { - tile_ir::TileLiteral::F16(half::f16::from_f32(value).to_bits()) - } - NaryScalar::F16(value) => tile_ir::TileLiteral::F16(value.to_bits()), - NaryScalar::U32(value) => { - tile_ir::TileLiteral::F16(half::f16::from_f32(value as f32).to_bits()) - } - }, - DataTypeEnum::U32 => match value { - NaryScalar::F32(value) => tile_ir::TileLiteral::U32(value as u32), - NaryScalar::F16(value) => tile_ir::TileLiteral::U32(value.to_f32() as u32), - NaryScalar::U32(value) => tile_ir::TileLiteral::U32(value), - }, - } -} - -pub(crate) fn output_dims_from_flat_usize( - flat: tile_ir::tile::Tile, - shape: &[u32], -) -> Vec { - let shape = shape.iter().map(|dim| *dim as usize).collect::>(); - output_dims_from_flat(flat, &shape) -} - -pub(crate) fn tile_reduce_op(op: ReduceOp) -> tile_ir::TileReduceOp { - match op { - ReduceOp::Sum => tile_ir::TileReduceOp::Sum, - ReduceOp::Product => tile_ir::TileReduceOp::Product, - ReduceOp::Max => tile_ir::TileReduceOp::Max, - ReduceOp::Min => tile_ir::TileReduceOp::Min, - } -} diff --git a/fusor-ml/core/src/reduce_tiled.rs b/fusor-ml/core/src/reduce_tiled.rs deleted file mode 100644 index 64c7fe3ee..000000000 --- a/fusor-ml/core/src/reduce_tiled.rs +++ /dev/null @@ -1,1007 +0,0 @@ -//! Tiled lowering for fused map-reduce operations. -//! -//! When the fused producer's index structure shows 2D reuse — every -//! reduce-axis-dependent input misses at least one of two parallel dims — the -//! reduction lowers as a workgroup-tiled kernel: each k-tile of every staged -//! input is loaded cooperatively into workgroup memory once and reused by all -//! lanes, and each thread accumulates a TM×TN register tile. A composed -//! matmul is exactly this shape (`a[.., m, k]` misses `n`, `b[.., k, n]` -//! misses `m`), so the tiled matmul kernel falls out of the caching the tiled -//! loads provide — as does any other contraction the tensor API composes, -//! whatever its dim order, pre-scaling, or surrounding expression. - -use fusor_tile_ir as tile_ir; -use std::hash::Hash; -use tile_ir::tile::{Mask, Tile}; - -use crate::{ - access_analysis::InputAccesses, - mir::{ - inputs::MirValue, kernel_backend, kernel_backend::DirectKernel, operation::Operation, - workgroup_shape::WorkgroupShape, - }, - nary_direct::{ - TensorMeta, ValueTile, apply_unary_function_chain, declare_value, - eval_nary_expr_on_value_tiles, layout_index, linear_group, tile_u32, zero_literal, - }, - nary_wise::{NaryExpr, NaryOp, NaryScalar}, - reduce::{ReduceOp, ReduceOperation}, - reduce_direct::{ReduceKernelInputs, tile_literal_for, tile_reduce_op}, - tensor::DataTypeEnum, - visit_tiled::{MaybeQData, distribute_workgroups}, -}; - -/// Workgroup tile geometry. One fixed shape keeps plan selection -/// deterministic; the gates below reject shapes it cannot cover profitably. -const BM: u32 = 32; -const BN: u32 = 32; -const BK: u32 = 8; -const TM: u32 = 4; -const TN: u32 = 4; -const LANES: u32 = (BM / TM) * (BN / TN); - -/// Metadata for dependence analysis. Block-quantized matrices address by -/// (row, col) with a zero in the kernel meta's column slot; the *analysis* -/// must still see the column as a real dependence, so it gets full row-major -/// strides here. -pub(crate) fn analysis_meta(value: &MaybeQData) -> Option { - match value { - MaybeQData::Tensor(tensor) => TensorMeta::new(tensor), - MaybeQData::QMatrix(matrix) => { - let mut meta = TensorMeta::for_matrix(matrix)?; - let mut acc = 1u32; - for dim in (0..meta.shape.len()).rev() { - meta.strides[dim] = acc; - acc = acc.checked_mul(meta.shape[dim])?; - } - Some(meta) - } - } -} - -/// Load one input value at `coords`, dequantizing block-quantized inputs -/// through the format-aware per-element path. -fn load_input_value( - program: &mut tile_ir::tile::TileBlock<'_>, - storage: &crate::nary_direct::Storage2, - meta: &TensorMeta, - coords: &[Tile], - mask: Mask, -) -> ValueTile { - if let crate::nary_direct::Storage2::Quantized(matrix) = storage { - let along_row = coords.last().cloned().unwrap_or_else(|| tile_u32(0)); - let which_row = layout_index(meta, coords); - return ValueTile::F32(program.load_quantized(matrix, along_row, which_row, mask, 0.0)); - } - storage.load(program, layout_index(meta, coords), mask) -} - -/// The allocation footprint one input re-streams per redundant read. -pub(crate) fn input_allocation_bytes(meta: &TensorMeta, value: &MaybeQData) -> u64 { - let element_bytes = match value { - MaybeQData::Tensor(tensor) => match tensor.datatype() { - DataTypeEnum::F16 => 2, - DataTypeEnum::F32 | DataTypeEnum::U32 => 4, - }, - MaybeQData::QMatrix(_) => 4, - }; - meta.allocation_len as u64 * element_bytes -} - -struct ReduceTiledKernelVariant; - -/// How a staged input's workgroup tile is addressed. -#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] -enum StagedKind { - /// References the row dim and the reduce axis: `[BM, BK]` tile. - RowK, - /// References the col dim and the reduce axis: `[BK, BN]` tile. - KCol, - /// References only the reduce axis (plus grid dims): `[BK]` tile. - KOnly, -} - -impl StagedKind { - fn tile_elements(self) -> u32 { - match self { - StagedKind::RowK => BM * BK, - StagedKind::KCol => BK * BN, - StagedKind::KOnly => BK, - } - } -} - -#[derive(Clone, Debug)] -enum InputRole { - /// Reduce-axis-dependent: staged through workgroup memory each k-tile. - Staged(StagedKind), - /// Independent of the reduce axis: loaded once per thread, outside the - /// k loop. - Direct, -} - -#[derive(Clone, Debug)] -struct PlannedInput { - /// Index-space dim read by each input dimension (pure `DimIndex` only). - dims: Vec, - role: InputRole, -} - -/// Where staged inputs are cached between reuses. -#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] -enum Staging { - /// Full tiles big enough to amortize barriers: k-tiles staged - /// cooperatively through workgroup memory. - Workgroup, - /// Partial or skinny tiles: each thread keeps its TM row values and TN - /// column values in registers per k step, so the reuse across its - /// register tile survives without barriers. - Register, -} - -/// Per-workgroup output-tile geometry. Always 64 lanes: -/// `(bm / tm) * (bn / tn) == LANES`. -#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] -struct TileGeom { - bm: u32, - bn: u32, - tm: u32, - tn: u32, -} - -impl TileGeom { - /// The 2D contraction tile: matches the dense matmul geometry. - const PAIR: Self = Self { - bm: BM, - bn: BN, - tm: TM, - tn: TN, - }; - /// The 1D tile: one parallel dim, each thread covering `tm` outputs so - /// row-missing inputs load once per k step instead of `tm` times. - const SINGLE: Self = Self { - bm: 256, - bn: 1, - tm: 4, - tn: 1, - }; - - fn outs(self) -> u32 { - self.tm * self.tn - } -} - -struct TiledReducePlan { - /// The parallel dim(s) tiled across the workgroup. `col_dim` is `None` - /// for 1D tiling: a single reuse dim with no second axis to pair it - /// with. - row_dim: usize, - col_dim: Option, - geom: TileGeom, - /// Every other non-reduce dim: one coordinate per workgroup. - grid_dims: Vec, - inputs: Vec, - staging: Staging, - /// Out-of-bounds index-space slots evaluate to the reduction identity - /// with zero-filled tiles (a bare multiply chain summed), so the inner - /// loop can skip per-element masking — the tiled matmul fast path. - fill_neutral: bool, -} - -/// Out-of-bounds slots collapse to the reduction identity without masking iff -/// the expression is a flat multiply folded with `Sum` whose factors map a -/// zero-filled load to zero (and stay finite): every staged tile zero-fills -/// its out-of-bounds slots, every out-of-bounds index-space slot is out of -/// bounds for at least one staged input, and one zero factor kills the -/// product. Factors may wrap their load in zero-preserving unary ops (casts, -/// negation, finite constant scales) — the dense matmul's f16→f32 accumulate -/// casts and scale epilogues keep the fast path. -fn is_fill_neutral(expr: &NaryExpr, op: ReduceOp) -> bool { - if op != ReduceOp::Sum { - return false; - } - fn scalar_is_finite(scalar: NaryScalar) -> bool { - match scalar { - NaryScalar::F32(value) => value.is_finite(), - NaryScalar::F16(value) => value.is_finite(), - NaryScalar::U32(_) => true, - } - } - fn factor_is_zero_preserving(expr: &NaryExpr) -> bool { - match expr { - NaryExpr::IndexedInput { .. } => true, - NaryExpr::Scalar(scalar) => scalar_is_finite(*scalar), - NaryExpr::Op { children, function } if children.len() == 1 => { - let zero_preserving = matches!( - function.op, - NaryOp::Cast | NaryOp::Neg | NaryOp::Abs - ) || matches!(function.op, NaryOp::MulConst(scalar) if scalar_is_finite(scalar)); - zero_preserving && factor_is_zero_preserving(&children[0]) - } - _ => false, - } - } - fn flat_mul_factors_collapse(expr: &NaryExpr) -> bool { - match expr { - NaryExpr::Op { children, function } - if function.op == NaryOp::Mul && children.len() == 2 => - { - children.iter().all(flat_mul_factors_collapse) - } - other => factor_is_zero_preserving(other), - } - } - matches!(expr, NaryExpr::Op { function, .. } if function.op == NaryOp::Mul) - && flat_mul_factors_collapse(expr) -} - -fn plan_tiled_reduce( - operation: &ReduceOperation, - values: &[MaybeQData], - device: &crate::Device, -) -> Option { - let shape = &operation.shape; - let axis = operation.axis; - let rank = shape.len(); - // One parallel dim + the reduce axis is enough for 1D tiling; the 2D - // pair search simply finds nothing below rank 3. - if rank < 2 { - return None; - } - let metas: Vec = values.iter().map(analysis_meta).collect::>()?; - let access = InputAccesses::collect(&operation.expression, operation.inputs.len(), &metas)?; - - // Pick the (row, col) pair: both must be fed by at least one staged - // input, and no staged input may reference both (its tile would need all - // three tiled dims). Among valid pairs prefer the most square, largest - // coverage; ties resolve to the lowest dim indices for determinism. - let k_dep: Vec = (0..access.effective.len()) - .map(|i| access.depends_on(i, axis)) - .collect(); - if !k_dep.iter().any(|dep| *dep) { - return None; - } - let mut best: Option<((usize, usize), (usize, usize))> = None; - for row_dim in 0..rank { - for col_dim in 0..rank { - if row_dim == col_dim || row_dim == axis || col_dim == axis { - continue; - } - let mut row_fed = false; - let mut col_fed = false; - let mut valid = true; - for (i, k_dep) in k_dep.iter().enumerate() { - if !k_dep { - continue; - } - let has_row = access.depends_on(i, row_dim); - let has_col = access.depends_on(i, col_dim); - if has_row && has_col { - valid = false; - break; - } - row_fed |= has_row; - col_fed |= has_col; - } - if !valid || !row_fed || !col_fed { - continue; - } - let (m, n) = (shape[row_dim], shape[col_dim]); - let score = (m.min(n), m * n); - if best - .as_ref() - .is_none_or(|(best_score, _)| score > *best_score) - { - best = Some((score, (row_dim, col_dim))); - } - } - } - // No pair: fall back to a single reuse dim — a parallel dim some - // k-dependent input misses. Each thread then covers `tm` outputs along - // it, so the missing input streams once per thread instead of `tm` - // times. Trading `tm`× thread count for that reuse only pays when the - // missed inputs are too large for cache to absorb the re-streams - // (cache-resident reuse is free, and the occupancy loss is not), so the - // missed inputs must clear the cache-thrash threshold. - let (row_dim, col_dim) = match best { - Some((_, (row_dim, col_dim))) => (row_dim, Some(col_dim)), - None => { - let mut best_single: Option<((u64, usize), usize)> = None; - for d in 0..rank { - if d == axis || shape[d] < TileGeom::SINGLE.bm as usize { - continue; - } - let missed_bytes: u64 = (0..k_dep.len()) - .filter(|&i| k_dep[i] && !access.depends_on(i, d)) - .map(|i| input_allocation_bytes(&metas[i], &values[i])) - .sum(); - if missed_bytes < device.last_level_cache_bytes() { - continue; - } - let score = (missed_bytes, shape[d]); - if best_single - .as_ref() - .is_none_or(|(best_score, _)| score > *best_score) - { - best_single = Some((score, d)); - } - } - let (_, d) = best_single?; - (d, None) - } - }; - let geom = match col_dim { - Some(_) => TileGeom::PAIR, - None => TileGeom::SINGLE, - }; - - let m: u32 = shape[row_dim].try_into().ok()?; - let n: u32 = match col_dim { - Some(col_dim) => shape[col_dim].try_into().ok()?, - None => 1, - }; - let k: u32 = shape[axis].try_into().ok()?; - let limits = device.limits(); - if LANES > limits.max_compute_workgroup_size_x - || LANES > limits.max_compute_invocations_per_workgroup - { - return None; - } - let grid_dims: Vec = (0..rank) - .filter(|&d| d != axis && d != row_dim && Some(d) != col_dim) - .collect(); - let batch: u32 = grid_dims.iter().try_fold(1u32, |acc, &d| { - acc.checked_mul(u32::try_from(shape[d]).ok()?) - })?; - let tiles_m = m.div_ceil(geom.bm); - let tiles_n = n.div_ceil(geom.bn); - let workgroups = batch.checked_mul(tiles_m)?.checked_mul(tiles_n)?; - let actual_outputs = batch.checked_mul(m)?.checked_mul(n)?; - let covered_outputs = workgroups.checked_mul(geom.bm * geom.bn)?; - - let mut workgroup_bytes = 0u32; - let mut inputs = Vec::with_capacity(access.dims.len()); - for (i, (dims, k_dep)) in access.dims.iter().zip(&k_dep).enumerate() { - let value = &values[i]; - let role = if *k_dep { - let has_col = col_dim.is_some_and(|col_dim| access.depends_on(i, col_dim)); - let kind = match (access.depends_on(i, row_dim), has_col) { - (true, false) => StagedKind::RowK, - (false, true) => StagedKind::KCol, - (false, false) => StagedKind::KOnly, - (true, true) => unreachable!("pair selection rejects row+col staged inputs"), - }; - let element_bytes = match value { - MaybeQData::Tensor(tensor) => match tensor.datatype() { - DataTypeEnum::F16 => 2, - DataTypeEnum::F32 | DataTypeEnum::U32 => 4, - }, - MaybeQData::QMatrix(_) => 4, - }; - workgroup_bytes = - workgroup_bytes.checked_add(kind.tile_elements().checked_mul(element_bytes)?)?; - InputRole::Staged(kind) - } else { - InputRole::Direct - }; - inputs.push(PlannedInput { - dims: dims.clone(), - role, - }); - } - - // Workgroup staging needs full-enough 2D tiles to amortize its barriers: - // the same shape and 75%-utilization gates as the dedicated dense matmul - // tile selection. Anything else keeps the register tile, whose reuse - // needs no barriers and tolerates partial coverage. - let workgroup_eligible = col_dim.is_some() - && m >= BM - && n >= BN - && k >= BK - && (actual_outputs as u64) * 4 >= (covered_outputs as u64) * 3 - && workgroup_bytes <= limits.max_compute_workgroup_storage_size; - let staging = if workgroup_eligible { - Staging::Workgroup - } else { - Staging::Register - }; - - let fill_neutral = is_fill_neutral(&operation.expression, operation.function.op); - - Some(TiledReducePlan { - row_dim, - col_dim, - geom, - grid_dims, - inputs, - staging, - fill_neutral, - }) -} - -/// The full index-space coordinate vector for one (batch, row, col, k) -/// position, in dim order. -fn full_coords( - plan: &TiledReducePlan, - rank: usize, - axis: usize, - batch_coords: &[Tile], - row: &Tile, - col: &Tile, - k: &Tile, -) -> Vec { - (0..rank) - .map(|d| { - if d == axis { - k.clone() - } else if d == plan.row_dim { - row.clone() - } else if Some(d) == plan.col_dim { - col.clone() - } else { - let slot = plan.grid_dims.iter().position(|&g| g == d).unwrap(); - batch_coords[slot].clone() - } - }) - .collect() -} - -fn input_coords(input: &PlannedInput, coords: &[Tile]) -> Vec { - input.dims.iter().map(|&d| coords[d].clone()).collect() -} - -pub(crate) fn datatype_scalar(datatype: DataTypeEnum) -> tile_ir::ScalarElement { - match datatype { - DataTypeEnum::F32 => tile_ir::ScalarElement::F32, - DataTypeEnum::F16 => tile_ir::ScalarElement::F16, - DataTypeEnum::U32 => tile_ir::ScalarElement::U32, - } -} - -fn value_tile(datatype: DataTypeEnum, value: Tile) -> ValueTile { - match datatype { - DataTypeEnum::F32 => ValueTile::F32(value), - DataTypeEnum::F16 => ValueTile::F16(value), - DataTypeEnum::U32 => ValueTile::U32(value), - } -} - -pub(crate) fn build_reduce_tiled_kernel( - operation: &ReduceOperation, - graph: &crate::compute_graph::ComputeGraphInner, - workgroup_shape: &WorkgroupShape, - inputs: &[MirValue], -) -> Option { - let parsed = ReduceKernelInputs::parse(operation, graph, inputs)?; - let device = graph.device(); - let plan = plan_tiled_reduce(operation, &parsed.values, &device)?; - if std::env::var_os("FUSOR_TRACE_REDUCE_TILED").is_some() { - eprintln!( - "reduce_tiled row={} col={:?} axis={} staging={:?} fill_neutral={} staged={} name={}", - plan.row_dim, - plan.col_dim, - operation.axis, - plan.staging, - plan.fill_neutral, - plan.inputs - .iter() - .filter(|input| matches!(input.role, InputRole::Staged(_))) - .count(), - operation.name(), - ); - } - - let rank = operation.shape.len(); - let axis = operation.axis; - let geom = plan.geom; - let m: u32 = operation.shape[plan.row_dim].try_into().ok()?; - let n: u32 = match plan.col_dim { - Some(col_dim) => operation.shape[col_dim].try_into().ok()?, - None => 1, - }; - let k: u32 = operation.shape[axis].try_into().ok()?; - let batch_sizes: Vec = plan - .grid_dims - .iter() - .map(|&d| u32::try_from(operation.shape[d]).ok()) - .collect::>()?; - let batch: u32 = batch_sizes.iter().product(); - let tiles_m = m.div_ceil(geom.bm); - let tiles_n = n.div_ceil(geom.bn); - let total_tiles = batch * tiles_m * tiles_n; - let k_tiles = k.div_ceil(BK); - - let dispatch_size = distribute_workgroups( - total_tiles, - device.limits().max_compute_workgroups_per_dimension, - ); - let variant = - kernel_backend::KernelVariantKey::with_payload::(|state| { - plan.row_dim.hash(state); - plan.col_dim.hash(state); - plan.geom.hash(state); - plan.staging.hash(state); - plan.fill_neutral.hash(state); - for input in &plan.inputs { - input.dims.hash(state); - match input.role { - InputRole::Staged(kind) => kind.hash(state), - InputRole::Direct => "direct".hash(state), - } - } - }); - let cache_key = operation.kernel_cache_key_with_dispatch( - variant, - Some(workgroup_shape), - dispatch_size, - inputs, - ); - - let reduce_dtype = operation.function.datatype(); - let reduce_op = tile_reduce_op(operation.function.op); - let initial = operation.function.initial_value; - let expression = operation.expression.clone(); - let post_chain = operation.post_element_wise.clone(); - let output = inputs.last()?.as_tensor()?.clone(); - let output_dtype = output.datatype(); - let output_value = MaybeQData::Tensor(output); - let values = parsed.values; - let plan = std::sync::Arc::new(plan); - - kernel_backend::run_kernel( - device.kernel_cache(), - format!("{}_tiled", operation.name()), - cache_key, - dispatch_size, - move |kb| { - let mut storages = Vec::with_capacity(values.len()); - let mut metas = Vec::with_capacity(values.len()); - for value in &values { - let (storage, meta) = declare_value(kb, value, false)?; - storages.push(storage); - metas.push(meta); - } - let (output_storage, output_meta) = declare_value(kb, &output_value, true)?; - - let phase = kb.program(); - let staged_tiles: Vec> = plan - .inputs - .iter() - .zip(&metas) - .map(|(input, meta)| match (plan.staging, &input.role) { - (Staging::Workgroup, InputRole::Staged(kind)) => { - let element = datatype_scalar(meta.datatype); - Some(match kind { - StagedKind::RowK => phase.alloc_workgroup_tile(element, BM, BK), - StagedKind::KCol => phase.alloc_workgroup_tile(element, BK, BN), - StagedKind::KOnly => phase.alloc_workgroup_array(element, BK), - }) - } - _ => None, - }) - .collect(); - - phase.program_grid(LANES, dispatch_size, |program| { - let tile_id = linear_group(program, dispatch_size); - let tile_active = tile_id.clone().lt(total_tiles); - let batch_id = tile_id.clone() / (tiles_m * tiles_n); - let local_tile = tile_id % (tiles_m * tiles_n); - let m_tile = local_tile.clone() / tiles_n; - let n_tile = local_tile % tiles_n; - - // Mixed-radix batch decomposition over the grid dims, in dim - // order (earlier dims vary slowest). - let batch_coords: Vec = (0..batch_sizes.len()) - .map(|i| { - let size = batch_sizes[i]; - if size == 1 { - return tile_u32(0); - } - let divisor: u32 = batch_sizes[i + 1..].iter().product(); - let quotient = if divisor == 1 { - batch_id.clone() - } else { - batch_id.clone() / divisor - }; - quotient % size - }) - .collect(); - - let lane = program.lane(); - let lane_row = program.bind(lane.clone() / (geom.bn / geom.tn)); - let lane_col = program.bind(lane % (geom.bn / geom.tn)); - let m_tile_base = program.bind(m_tile * geom.bm); - let n_tile_base = program.bind(n_tile * geom.bn); - let row_base = program.bind(m_tile_base.clone() + lane_row.clone() * geom.tm); - let col_base = program.bind(n_tile_base.clone() + lane_col.clone() * geom.tn); - - let identity = || Tile::literal(tile_literal_for(initial, reduce_dtype)); - - // Per-thread output-slot bounds, hoisted out of the k loop. - let rc_masks: Vec = (0..geom.outs()) - .map(|idx| { - let (r, c) = (idx / geom.tn, idx % geom.tn); - let row = row_base.clone() + r; - let col = col_base.clone() + c; - tile_active.clone().and(row.lt(m)).and(col.lt(n)) - }) - .collect(); - - // Direct (k-independent) inputs: one bound load per output - // slot, reused across the whole k loop. - let direct_values: Vec>> = plan - .inputs - .iter() - .enumerate() - .map(|(i, input)| { - if !matches!(input.role, InputRole::Direct) { - return None; - } - Some( - (0..geom.outs()) - .map(|idx| { - let (r, c) = (idx / geom.tn, idx % geom.tn); - let row = row_base.clone() + r; - let col = col_base.clone() + c; - let coords = full_coords( - &plan, - rank, - axis, - &batch_coords, - &row, - &col, - &tile_u32(0), - ); - let index = - layout_index(&metas[i], &input_coords(input, &coords)); - let loaded = storages[i].load( - program, - index, - rc_masks[idx as usize].clone(), - ); - // Bind to a local so the k loop reuses - // the value instead of re-issuing the - // storage load every iteration. - let native = match loaded { - ValueTile::F32(v) - | ValueTile::F16(v) - | ValueTile::U32(v) => v, - ValueTile::Bool(_) => { - unreachable!("tensor inputs are f32/f16/u32") - } - }; - ( - value_tile(metas[i].datatype, program.bind(native)), - metas[i].datatype, - ) - }) - .collect(), - ) - }) - .collect(); - - let initial_accs: Vec = (0..geom.outs()).map(|_| identity()).collect(); - let sums = match plan.staging { - Staging::Register => program.fold_vec( - tile_ir::tile::range(k), - initial_accs, - |program, k_index, accs| { - // Per staged input, the register slice this thread - // reuses across its TM×TN tile: TM row values, TN - // column values, or one k-only value. - let reg_values: Vec>> = plan - .inputs - .iter() - .enumerate() - .map(|(i, input)| { - let InputRole::Staged(kind) = input.role else { - return None; - }; - let count = match kind { - StagedKind::RowK => geom.tm, - StagedKind::KCol => geom.tn, - StagedKind::KOnly => 1, - }; - Some( - (0..count) - .map(|j| { - let (row, col, mask) = match kind { - StagedKind::RowK => { - let row = row_base.clone() + j; - let mask = tile_active - .clone() - .and(row.clone().lt(m)); - (row, tile_u32(0), mask) - } - StagedKind::KCol => { - let col = col_base.clone() + j; - let mask = tile_active - .clone() - .and(col.clone().lt(n)); - (tile_u32(0), col, mask) - } - StagedKind::KOnly => ( - tile_u32(0), - tile_u32(0), - tile_active.clone(), - ), - }; - let coords = full_coords( - &plan, - rank, - axis, - &batch_coords, - &row, - &col, - &k_index, - ); - let index = layout_index( - &metas[i], - &input_coords(input, &coords), - ); - let loaded = - storages[i].load(program, index, mask.clone()); - let value = match loaded { - ValueTile::F32(v) - | ValueTile::F16(v) - | ValueTile::U32(v) => Tile::select( - mask, - v, - Tile::literal(zero_literal( - metas[i].datatype, - )), - ), - ValueTile::Bool(_) => unreachable!( - "tensor inputs are f32/f16/u32" - ), - }; - ( - value_tile(metas[i].datatype, value), - metas[i].datatype, - ) - }) - .collect(), - ) - }) - .collect(); - - // Every k step is in range and out-of-bounds - // output slots never store, so no accumulate - // masking is needed in register mode. - accs.into_iter() - .enumerate() - .map(|(idx, acc)| { - let (r, c) = (idx as u32 / geom.tn, idx as u32 % geom.tn); - let slot_values: Vec<(ValueTile, DataTypeEnum)> = plan - .inputs - .iter() - .enumerate() - .map(|(i, input)| match input.role { - InputRole::Staged(kind) => { - let j = match kind { - StagedKind::RowK => r, - StagedKind::KCol => c, - StagedKind::KOnly => 0, - }; - reg_values[i].as_ref().unwrap()[j as usize].clone() - } - InputRole::Direct => { - direct_values[i].as_ref().unwrap()[idx].clone() - } - }) - .collect(); - let (value, _) = - eval_nary_expr_on_value_tiles(&expression, &slot_values); - let value = match value.cast_to(reduce_dtype) { - ValueTile::F32(v) - | ValueTile::F16(v) - | ValueTile::U32(v) => v, - ValueTile::Bool(_) => { - unreachable!("reduce dtype is f32/f16/u32") - } - }; - acc.binary(reduce_op.binary(), value) - }) - .collect() - }, - ), - Staging::Workgroup => program.fold_vec( - tile_ir::tile::range(k_tiles), - initial_accs, - |program, k_tile, accs| { - let k_base = program.bind(k_tile * BK); - - // Cooperative staging: every lane copies its share of - // each staged input's k-tile into workgroup memory. - for (i, input) in plan.inputs.iter().enumerate() { - let InputRole::Staged(kind) = input.role else { - continue; - }; - let tile = staged_tiles[i].as_ref().unwrap(); - let elements = kind.tile_elements(); - for pass in 0..elements.div_ceil(LANES) { - let flat = program.lane() + pass * LANES; - let (row, col, k_index, flat_ok) = match kind { - StagedKind::RowK => { - let local_row = flat.clone() / BK; - let local_k = flat.clone() % BK; - ( - m_tile_base.clone() + local_row, - tile_u32(0), - k_base.clone() + local_k, - Mask::all(), - ) - } - StagedKind::KCol => { - let local_k = flat.clone() / BN; - let local_col = flat.clone() % BN; - ( - tile_u32(0), - n_tile_base.clone() + local_col, - k_base.clone() + local_k, - Mask::all(), - ) - } - StagedKind::KOnly => ( - tile_u32(0), - tile_u32(0), - k_base.clone() + flat.clone(), - flat.clone().lt(elements), - ), - }; - let mut in_bounds = - tile_active.clone().and(flat_ok).and(k_index.clone().lt(k)); - if matches!(kind, StagedKind::RowK) { - in_bounds = in_bounds.and(row.clone().lt(m)); - } - if matches!(kind, StagedKind::KCol) { - in_bounds = in_bounds.and(col.clone().lt(n)); - } - let coords = full_coords( - &plan, - rank, - axis, - &batch_coords, - &row, - &col, - &k_index, - ); - let loaded = load_input_value( - program, - &storages[i], - &metas[i], - &input_coords(input, &coords), - in_bounds.clone(), - ); - // Zero-fill out-of-bounds slots: in the - // fill-neutral fast path that zero is what - // collapses the product, elsewhere the inner - // mask discards the slot anyway. - let value = match loaded.cast_to(metas[i].datatype) { - ValueTile::F32(v) - | ValueTile::F16(v) - | ValueTile::U32(v) => Tile::select( - in_bounds, - v, - Tile::literal(zero_literal(metas[i].datatype)), - ), - ValueTile::Bool(_) => { - unreachable!("tensor inputs are f32/f16/u32") - } - }; - program.store_workgroup(tile, flat, value); - } - } - program.workgroup_barrier(); - - let chunk_sums: Vec = (0..geom.outs()) - .map(|idx| { - let (r, c) = (idx / geom.tn, idx % geom.tn); - let local_row = lane_row.clone() * geom.tm + r; - let local_col = lane_col.clone() * geom.tn + c; - let mut chunk = identity(); - for kk in 0..BK { - let slot_values: Vec<(ValueTile, DataTypeEnum)> = plan - .inputs - .iter() - .enumerate() - .map(|(i, input)| match input.role { - InputRole::Staged(kind) => { - let tile = staged_tiles[i].as_ref().unwrap(); - let flat = match kind { - StagedKind::RowK => { - local_row.clone() * BK + kk - } - StagedKind::KCol => { - local_col.clone() + kk * BN - } - StagedKind::KOnly => tile_u32(kk), - }; - ( - value_tile( - metas[i].datatype, - program.load_workgroup(tile, flat), - ), - metas[i].datatype, - ) - } - InputRole::Direct => { - direct_values[i].as_ref().unwrap()[idx as usize] - .clone() - } - }) - .collect(); - let (value, _) = eval_nary_expr_on_value_tiles( - &expression, - &slot_values, - ); - let value = match value.cast_to(reduce_dtype) { - ValueTile::F32(v) - | ValueTile::F16(v) - | ValueTile::U32(v) => v, - ValueTile::Bool(_) => { - unreachable!("reduce dtype is f32/f16/u32") - } - }; - let value = if plan.fill_neutral { - value - } else { - let valid = rc_masks[idx as usize] - .clone() - .and((k_base.clone() + kk).lt(k)); - Tile::select(valid, value, identity()) - }; - chunk = chunk.binary(reduce_op.binary(), value); - } - chunk - }) - .collect(); - // Bind every chunk before the trailing barrier: the - // next iteration's staging overwrites the tiles these - // reads source from. - let chunk_sums: Vec = chunk_sums - .into_iter() - .map(|chunk| program.bind(chunk)) - .collect(); - program.workgroup_barrier(); - accs.into_iter() - .zip(chunk_sums) - .map(|(acc, chunk)| acc.binary(reduce_op.binary(), chunk)) - .collect() - }, - ), - }; - - for (idx, sum) in sums.into_iter().enumerate() { - let idx = idx as u32; - let (r, c) = (idx / geom.tn, idx % geom.tn); - let row = row_base.clone() + r; - let col = col_base.clone() + c; - let (value, value_ty) = apply_unary_function_chain( - value_tile(reduce_dtype, sum).into_f32(), - reduce_dtype, - &post_chain, - ) - .expect("validated reduce post_element_wise chain"); - let value = ValueTile::F32(value) - .cast_to(value_ty) - .cast_to(output_dtype); - let coords = - full_coords(&plan, rank, axis, &batch_coords, &row, &col, &tile_u32(0)); - let out_coords: Vec = (0..rank) - .filter(|&d| d != axis) - .map(|d| coords[d].clone()) - .collect(); - let output_index = layout_index(&output_meta, &out_coords); - output_storage.store( - program, - output_index, - value, - rc_masks[idx as usize].clone(), - ); - } - }); - Some(()) - }, - ) -} diff --git a/fusor-ml/core/src/row_program.rs b/fusor-ml/core/src/row_program.rs index dfbdcba28..445e3ac49 100644 --- a/fusor-ml/core/src/row_program.rs +++ b/fusor-ml/core/src/row_program.rs @@ -39,8 +39,7 @@ use crate::{ linear_group, output_dims_from_flat, tile_u32, }, nary_wise::{NaryExpr, NaryFunction, NaryOp, NaryScalar, UnaryFunctionChain}, - reduce::{ReduceFunction, ReduceOp, max_fn, sum_fn}, - reduce_direct::{tile_literal_for, tile_reduce_op}, + reduce::{ReduceFunction, ReduceOp, ReduceOperation, max_fn, sum_fn}, tensor::{DataTypeEnum, TensorData}, visit_tiled::{MaybeQData, distribute_workgroups}, }; @@ -80,6 +79,8 @@ pub(crate) struct RowFold { pub(crate) enum RowOutput { /// One output element per index-space position; output shape == `shape`. Map(NaryExpr), + /// One output element per row; output shape == row dims. + Scalar(NaryExpr), /// Fold the axis once more with one free output dimension appended to /// the row shape: `combine` sees the free coordinate as `DimIndex(rank)` /// and element-phase slots at the folded axis position. Output shape = @@ -256,6 +257,25 @@ pub(crate) struct RowProgramOperation { } impl RowProgramOperation { + pub(crate) fn from_reduce(reduce: &ReduceOperation) -> Self { + let scalar_slot = slot_expr(reduce.inputs.len(), 0); + Self { + inputs: reduce.inputs.clone(), + shape: reduce.shape.clone(), + axis: reduce.axis, + steps: vec![ + RowStep::Reduce(RowReduce { + expression: reduce.expression.clone(), + function: reduce.function.clone(), + post_chain: reduce.post_element_wise.clone(), + }), + RowStep::Output(RowOutput::Scalar(scalar_slot)), + ], + output_datatype: reduce.out_datatype(), + dynamic_axis: None, + } + } + fn output_step(&self) -> &RowOutput { match self.steps.last() { Some(RowStep::Output(output)) => output, @@ -289,6 +309,7 @@ impl RowProgramOperation { pub(crate) fn out_shape(&self) -> Vec { match self.output_step() { RowOutput::Map(_) => self.shape.to_vec(), + RowOutput::Scalar(_) => self.row_shape(), RowOutput::Reduce { free_dim, .. } => { let mut shape = self.row_shape(); shape.push(*free_dim); @@ -307,6 +328,25 @@ impl RowProgramOperation { None => device.limits().max_compute_workgroup_size_x.min(BLOCK), } } + + fn uses_custom_indexing_for_input(&self, input_idx: usize) -> bool { + let output_uses_custom_indexing = match self.output_step() { + RowOutput::Map(expr) | RowOutput::Scalar(expr) => { + expr.uses_custom_indexing_for_input(input_idx) + } + RowOutput::Reduce { combine, .. } => combine.uses_custom_indexing_for_input(input_idx), + }; + output_uses_custom_indexing + || self.phase_steps().iter().any(|step| match step { + RowStep::Reduce(reduce) => { + reduce.expression.uses_custom_indexing_for_input(input_idx) + } + RowStep::Element { expression, .. } => { + expression.uses_custom_indexing_for_input(input_idx) + } + RowStep::Output(_) => false, + }) + } } impl Operation for RowProgramOperation { @@ -363,7 +403,17 @@ impl Operation for RowProgramOperation { let mut mir_inputs: Vec = self .inputs .iter() - .map(|idx| nodes.get_result_or_qmatrix(*idx).unwrap().into()) + .enumerate() + .map(|(i, idx)| { + // Custom-indexed inputs need dense tensor addressing; the + // quantized matrix path only supports plain row/col reads. + if self.uses_custom_indexing_for_input(i) + && let Some(cached) = nodes.get_result(*idx) + { + return cached.into(); + } + nodes.get_result_or_qmatrix(*idx).unwrap().into() + }) .collect(); let device = match &mir_inputs[0] { MirValue::Tensor(tensor) => tensor.device().clone(), @@ -484,6 +534,39 @@ fn f32_literal(value: f32) -> Tile { Tile::literal(tile_ir::TileLiteral::f32(value)) } +fn tile_literal_for(value: NaryScalar, target: DataTypeEnum) -> tile_ir::TileLiteral { + match target { + DataTypeEnum::F32 => match value { + NaryScalar::F32(value) => tile_ir::TileLiteral::f32(value), + NaryScalar::F16(value) => tile_ir::TileLiteral::f32(value.to_f32()), + NaryScalar::U32(value) => tile_ir::TileLiteral::f32(value as f32), + }, + DataTypeEnum::F16 => match value { + NaryScalar::F32(value) => { + tile_ir::TileLiteral::F16(half::f16::from_f32(value).to_bits()) + } + NaryScalar::F16(value) => tile_ir::TileLiteral::F16(value.to_bits()), + NaryScalar::U32(value) => { + tile_ir::TileLiteral::F16(half::f16::from_f32(value as f32).to_bits()) + } + }, + DataTypeEnum::U32 => match value { + NaryScalar::F32(value) => tile_ir::TileLiteral::U32(value as u32), + NaryScalar::F16(value) => tile_ir::TileLiteral::U32(value.to_f32() as u32), + NaryScalar::U32(value) => tile_ir::TileLiteral::U32(value), + }, + } +} + +fn tile_reduce_op(op: ReduceOp) -> tile_ir::TileReduceOp { + match op { + ReduceOp::Sum => tile_ir::TileReduceOp::Sum, + ReduceOp::Product => tile_ir::TileReduceOp::Product, + ReduceOp::Max => tile_ir::TileReduceOp::Max, + ReduceOp::Min => tile_ir::TileReduceOp::Min, + } +} + fn build_row_program_kernel( operation: &RowProgramOperation, graph: &crate::compute_graph::ComputeGraphInner, @@ -531,7 +614,7 @@ fn build_row_program_kernel( let phase_steps = operation.phase_steps().to_vec(); let free_dim_out = match &output_kind { RowOutput::Reduce { free_dim, .. } => Some(*free_dim), - RowOutput::Map(_) => None, + RowOutput::Map(_) | RowOutput::Scalar(_) => None, }; let tiles = k.div_ceil(block); let splits: u32 = match free_dim_out { @@ -1029,26 +1112,46 @@ fn build_row_program_kernel( slots.push((scalar, combined_ty)); } - let RowOutput::Map(output_expr) = &output_kind else { - unreachable!("reducing output requires a dynamic-axis row program") - }; - program.loop_range(chunks, |program, chunk| { - let k_index = chunk * block + lane.clone(); - let active = in_bounds.clone().and(k_index.clone().lt(k)); - let coords = full_coords(k_index); - let (value, _) = eval_nary_expr( - program, - output_expr, - &coords, - &storages, - &metas, - active.clone(), - &slots, - ); - let value = value.cast_to(output_dtype); - let output_index = layout_index(&output_meta, &coords); - output_storage.store(program, output_index, value, active); - }); + match &output_kind { + RowOutput::Map(output_expr) => { + program.loop_range(chunks, |program, chunk| { + let k_index = chunk * block + lane.clone(); + let active = in_bounds.clone().and(k_index.clone().lt(k)); + let coords = full_coords(k_index); + let (value, _) = eval_nary_expr( + program, + output_expr, + &coords, + &storages, + &metas, + active.clone(), + &slots, + ); + let value = value.cast_to(output_dtype); + let output_index = layout_index(&output_meta, &coords); + output_storage.store(program, output_index, value, active); + }); + } + RowOutput::Scalar(output_expr) => { + let active = in_bounds.and(lane.eq(0u32)); + let coords = full_coords(tile_u32(0)); + let (value, _) = eval_nary_expr( + program, + output_expr, + &coords, + &storages, + &metas, + active.clone(), + &slots, + ); + let value = value.cast_to(output_dtype); + let output_index = layout_index(&output_meta, &row_dims); + output_storage.store(program, output_index, value, active); + } + RowOutput::Reduce { .. } => { + unreachable!("reducing output requires a dynamic-axis row program") + } + } }); Some(()) }, diff --git a/fusor-ml/core/src/view.rs b/fusor-ml/core/src/view.rs index d09a9564b..1c9d8b0d8 100644 --- a/fusor-ml/core/src/view.rs +++ b/fusor-ml/core/src/view.rs @@ -363,7 +363,10 @@ pub(crate) fn affine_dim_indices( ) } -fn row_major_indices_from_flat(flat: NaryExpr, shape: &[usize]) -> Option> { +pub(crate) fn row_major_indices_from_flat( + flat: NaryExpr, + shape: &[usize], +) -> Option> { let mut indices = Vec::with_capacity(shape.len()); for axis in 0..shape.len() { let divisor = shape[axis + 1..] From 14c313649befc607e910fcdbbedd2147eab8d924 Mon Sep 17 00:00:00 2001 From: Evan Almloff Date: Thu, 11 Jun 2026 22:15:30 -0500 Subject: [PATCH 05/16] generic pseudo-afline shapes --- .../core/src/compute_graph/layout_pass.rs | 13 +- .../src/compute_graph/resolve/execution.rs | 23 +- .../src/compute_graph/resolve/fold_views.rs | 160 +++--- .../src/compute_graph/resolve/fusion_row.rs | 18 +- .../src/compute_graph/resolve/recognize.rs | 49 +- .../resolve/recognize_attention.rs | 42 +- .../compute_graph/resolve/recognize_cat.rs | 8 +- fusor-ml/core/src/compute_graph/tests.rs | 51 ++ fusor-ml/core/src/matmul/kernel.rs | 216 +++++--- fusor-ml/core/src/matmul/mod.rs | 68 ++- fusor-ml/core/src/matmul/variants.rs | 24 +- fusor-ml/core/src/quantized/matmul/kernel.rs | 18 +- fusor-ml/core/src/view.rs | 469 ++++++++++++------ .../tile-ir-kernels/src/kernels/matmul.rs | 71 ++- fusor-ml/tile-ir/src/ir/layout.rs | 17 + fusor-ml/tile-ir/src/ir/program.rs | 5 + fusor-ml/tile-ir/src/lower/analysis.rs | 5 +- fusor-ml/tile-ir/src/lower/coop.rs | 293 ++++++++++- fusor-ml/tile-ir/src/tile/block.rs | 18 + 19 files changed, 1129 insertions(+), 439 deletions(-) diff --git a/fusor-ml/core/src/compute_graph/layout_pass.rs b/fusor-ml/core/src/compute_graph/layout_pass.rs index 2203c989d..8ec503ffa 100644 --- a/fusor-ml/core/src/compute_graph/layout_pass.rs +++ b/fusor-ml/core/src/compute_graph/layout_pass.rs @@ -64,12 +64,17 @@ impl LayoutPass { self.queue.push_back(key); return; }; - // A fully-defined view that composes with the input's layout stays a - // zero-cost buffer view; anything else materializes contiguously - // through the gather fallback. + // A fully-defined view whose stages compose with the input's layout + // stays a zero-cost buffer view; anything else materializes + // contiguously through the gather fallback. let new_layout = operation .is_fully_defined() - .then(|| crate::view::compose_layouts(&operation.layout, input_layout.layout())) + .then(|| { + operation.stages.iter().try_fold( + input_layout.layout().clone(), + |composed, stage| crate::view::compose_layouts(&stage.layout, &composed), + ) + }) .flatten() .unwrap_or_else(|| Layout::contiguous(operation.shape())); self.output_layout.insert( diff --git a/fusor-ml/core/src/compute_graph/resolve/execution.rs b/fusor-ml/core/src/compute_graph/resolve/execution.rs index 93039d67c..26c7cfde3 100644 --- a/fusor-ml/core/src/compute_graph/resolve/execution.rs +++ b/fusor-ml/core/src/compute_graph/resolve/execution.rs @@ -538,12 +538,15 @@ impl Resolver { self.node_mapping.get(&inner_input).copied() } - /// Walk through fully-defined view nodes from `inner` down to the first - /// non-view node, composing the view layouts. Returns the base node and - /// the composed view layout over the base's logical value space; the - /// layout is `None` when `inner` is not a view (identity). Views that - /// don't compose (or carry a fill region) act as chain breaks: the walk - /// stops without seeing through them. + /// Walk through view nodes from `inner` down to the first non-view + /// node, composing each view's collapsed stage stack. Public tensor ops + /// collapse into single view nodes at construction, but composed + /// clusters (attention's attached GQA/transpose views) still layer view + /// nodes deliberately. Returns the base node and the composed layout + /// over the base's logical value space; the layout is `None` when + /// `inner` is not a view (identity). Views that don't collapse or + /// compose (or carry a fill region) act as chain breaks: the walk stops + /// without seeing through them. pub(super) fn walk_view_chain(&self, mut inner: NodeIndex) -> (NodeIndex, Option) { let mut composed: Option = None; loop { @@ -553,12 +556,12 @@ impl Resolver { let ExecutionVariant::View(view) = &self.execution_graph[exec].variant else { return (inner, composed); }; - if !view.is_fully_defined() { + let Some(collapsed) = view.composed_layout() else { return (inner, composed); - } + }; let next = match &composed { - None => view.layout.clone(), - Some(outer) => match crate::view::compose_layouts(outer, &view.layout) { + None => collapsed, + Some(outer) => match crate::view::compose_layouts(outer, &collapsed) { Some(layout) => layout, None => return (inner, composed), }, diff --git a/fusor-ml/core/src/compute_graph/resolve/fold_views.rs b/fusor-ml/core/src/compute_graph/resolve/fold_views.rs index 64e791e96..9de390b0c 100644 --- a/fusor-ml/core/src/compute_graph/resolve/fold_views.rs +++ b/fusor-ml/core/src/compute_graph/resolve/fold_views.rs @@ -1,17 +1,21 @@ -use crate::view::{AffineIndex, ViewOperation, affine_dim_indices}; +use crate::view::ViewOperation; use super::*; impl Resolver { /// Fold view inputs of an n-ary node directly into its expression: each - /// `IndexedInput` through the view becomes an `IndexedInput` of the - /// view's base node with the view's coordinate mapping applied (and a - /// bounds-select around partially-defined views). This removes the view - /// node from between the producer and consumer, so the n-ary fusion - /// passes see through layout changes instead of stopping at them. + /// `IndexedInput` through the view becomes a load of the view's base + /// node with the view's coordinate map applied (and a bounds-select + /// around partially-defined stages). This removes the view node from + /// between the producer and consumer, so the n-ary fusion passes see + /// through layout changes instead of stopping at them. /// - /// Only affine views fold (no divmod in the rewritten indices); anything - /// else stays a view node and materializes through the gather fallback. + /// Affine maps always fold — they rewrite to plain index arithmetic. + /// Maps that need delinearization (divmod address arithmetic, from a + /// reshape regrouping non-mergeable strides) re-derive coordinates on + /// every load, so they only fold where each element is loaded once; a + /// load re-read across unindexed dims (a contraction operand) keeps the + /// view node and materializes through the gather fallback instead. pub(super) fn try_fold_view_inputs( &mut self, graph: &mut ComputeGraphInner, @@ -39,11 +43,17 @@ impl Resolver { let ExecutionVariant::View(view) = &self.execution_graph[input_exec].variant else { continue; }; - let Some(affine) = affine_dim_indices(&view.layout, &view.input_shape) else { + let needs_delinearize = view.stages.iter().any(|stage| { + crate::view::affine_dim_indices(&stage.layout, &stage.input_shape).is_none() + }); + if needs_delinearize && input_reread_factor(&expression, &nary.shape, slot) > 1 { continue; - }; + } let view = view.clone(); - expression = Self::rewrite_view_input(&expression, slot, &view, &affine); + let Some(rewritten) = Self::rewrite_view_input(&expression, slot, &view) else { + continue; + }; + expression = rewritten; inputs[slot] = view.input; folded.push((input_exec, view.input)); } @@ -80,89 +90,93 @@ impl Resolver { true } - /// Rewrite every access to input `target_idx` through `view`'s coordinate - /// mapping. The original index expressions (the view's output - /// coordinates) feed the affine per-base-dimension indices; partially - /// defined views wrap the load in `select(in_bounds, load, fill)` with - /// the load coordinates clamped in-bounds (both select branches - /// evaluate). + /// Rewrite every access to input `target_idx` through `view`'s + /// coordinate map: the original index expressions (the view's output + /// coordinates) walk down the stage stack to base coordinates, with + /// fill selects and in-bounds clamps where stages are partially defined + /// (both select branches evaluate). fn rewrite_view_input( expr: &NaryExpr, target_idx: usize, view: &ViewOperation, - affine: &[AffineIndex], - ) -> NaryExpr { - match expr { + ) -> Option { + Some(match expr { NaryExpr::Op { children, function } => NaryExpr::Op { children: children .iter() - .map(|child| Self::rewrite_view_input(child, target_idx, view, affine)) - .collect(), + .map(|child| Self::rewrite_view_input(child, target_idx, view)) + .collect::>>()?, function: function.clone(), }, NaryExpr::IndexedInput { input_idx, indices } => { let indices: Vec = indices .iter() - .map(|index| Self::rewrite_view_input(index, target_idx, view, affine)) - .collect(); + .map(|index| Self::rewrite_view_input(index, target_idx, view)) + .collect::>>()?; if *input_idx != target_idx { - return NaryExpr::IndexedInput { + NaryExpr::IndexedInput { input_idx: *input_idx, indices, - }; + } + } else { + view.value_expression(*input_idx, &indices)?.0 } + } + NaryExpr::DimIndex(dim) => NaryExpr::DimIndex(*dim), + NaryExpr::Scalar(value) => NaryExpr::Scalar(*value), + }) + } +} - let fully_defined = view.is_fully_defined(); - let base_indices = affine - .iter() - .zip(&*view.input_shape) - .map(|(index, &extent)| { - let index = index.to_expr(&indices); - if fully_defined || extent == 0 { - index - } else { - NaryExpr::unary_op( - index, - "clamp_dim", - NaryOp::MinConst(NaryScalar::U32(extent as u32 - 1)), - DataTypeEnum::U32, - DataTypeEnum::U32, - ) - } - }) - .collect(); - let loaded = NaryExpr::IndexedInput { - input_idx: *input_idx, - indices: base_indices, - }; - if fully_defined { - return loaded; +/// The worst re-read factor across this slot's loads: the product of +/// index-space dims a load's coordinates never reference — each such dim +/// re-reads the same element once per step. +fn input_reread_factor(expr: &NaryExpr, shape: &[usize], slot: usize) -> usize { + fn collect_dims(expr: &NaryExpr, referenced: &mut [bool]) { + match expr { + NaryExpr::Op { children, .. } => { + for child in children { + collect_dims(child, referenced); } - - let mut condition = NaryExpr::scalar(NaryScalar::U32(1)); - for (dim, (&defined, &size)) in view.defined.iter().zip(view.shape()).enumerate() { - if defined >= size { - continue; + } + NaryExpr::IndexedInput { indices, .. } => { + for index in indices { + collect_dims(index, referenced); + } + } + NaryExpr::DimIndex(dim) => referenced[*dim] = true, + NaryExpr::Scalar(_) => {} + } + } + fn visit_loads(expr: &NaryExpr, shape: &[usize], slot: usize, worst: &mut usize) { + match expr { + NaryExpr::Op { children, .. } => { + for child in children { + visit_loads(child, shape, slot, worst); + } + } + NaryExpr::IndexedInput { input_idx, indices } => { + for index in indices { + visit_loads(index, shape, slot, worst); + } + if *input_idx == slot { + let mut referenced = vec![false; shape.len()]; + for index in indices { + collect_dims(index, &mut referenced); } - let lt_defined = NaryExpr::unary_op( - indices[dim].clone(), - "lt_defined", - NaryOp::LessConst(NaryScalar::U32(defined as u32)), - DataTypeEnum::U32, - DataTypeEnum::U32, - ); - condition = NaryExpr::mul(condition, lt_defined, DataTypeEnum::U32); + let factor: usize = shape + .iter() + .zip(&referenced) + .filter(|(_, referenced)| !**referenced) + .map(|(size, _)| *size) + .product(); + *worst = (*worst).max(factor); } - NaryExpr::select( - condition, - loaded, - NaryExpr::scalar(view.fill), - DataTypeEnum::U32, - view.datatype, - ) } - NaryExpr::DimIndex(dim) => NaryExpr::DimIndex(*dim), - NaryExpr::Scalar(value) => NaryExpr::Scalar(*value), + NaryExpr::DimIndex(_) | NaryExpr::Scalar(_) => {} } } + let mut worst = 1; + visit_loads(expr, shape, slot, &mut worst); + worst } diff --git a/fusor-ml/core/src/compute_graph/resolve/fusion_row.rs b/fusor-ml/core/src/compute_graph/resolve/fusion_row.rs index 29c92e2ab..8d0d58ecd 100644 --- a/fusor-ml/core/src/compute_graph/resolve/fusion_row.rs +++ b/fusor-ml/core/src/compute_graph/resolve/fusion_row.rs @@ -358,13 +358,15 @@ impl Resolver { let ExecutionVariant::View(view) = &self.execution_graph[exec].variant else { break; }; - if !view.is_fully_defined() || !within.is_none_or(|allowed| allowed.contains(&node)) - { + if !within.is_none_or(|allowed| allowed.contains(&node)) { break 'scalar; } + let Some(collapsed) = view.composed_layout() else { + break 'scalar; + }; layout = Some(match &layout { - None => view.layout.clone(), - Some(outer) => match crate::view::compose_layouts(outer, &view.layout) { + None => collapsed, + Some(outer) => match crate::view::compose_layouts(outer, &collapsed) { Some(layout) => layout, None => break 'scalar, }, @@ -610,12 +612,10 @@ impl Resolver { let ExecutionVariant::View(view) = &self.execution_graph[exec].variant else { break; }; - if !view.is_fully_defined() { - return None; - } + let collapsed = view.composed_layout()?; layout = Some(match &layout { - None => view.layout.clone(), - Some(outer) => crate::view::compose_layouts(outer, &view.layout)?, + None => collapsed, + Some(outer) => crate::view::compose_layouts(outer, &collapsed)?, }); views.push(node); node = view.input; diff --git a/fusor-ml/core/src/compute_graph/resolve/recognize.rs b/fusor-ml/core/src/compute_graph/resolve/recognize.rs index ca89a31b3..f1083b32d 100644 --- a/fusor-ml/core/src/compute_graph/resolve/recognize.rs +++ b/fusor-ml/core/src/compute_graph/resolve/recognize.rs @@ -289,10 +289,11 @@ impl Resolver { /// Read a recognized matmul's A operand through its un-flattened /// producer. Conv's im2col flatten regroups a windowed view's dims /// across overlapping strides, which no single strided layout can - /// express: the reshape becomes a chained view that would materialize + /// express: the view keeps a stage boundary and would materialize /// through the gather fallback. When the flat `[M, K]` operand is - /// exactly such a reinterpret, point the matmul at the producer and let - /// the kernels divmod the flat coordinates back apart per load — an + /// exactly such a reinterpret over an affine stage, point the matmul at + /// the producer, carry the affine stage as the operand's base map, and + /// let the kernels divmod the flat coordinates back apart per load — an /// implicit GEMM with no gather dispatch. fn try_unflatten_matmul_input( &mut self, @@ -323,46 +324,60 @@ impl Resolver { let ComputeGraphNodeVariant::View(view) = &node.variant else { return; }; - if !view.is_fully_defined() - || !view.layout.is_contiguous() - || view.layout.offset() != 0 - || view.layout.shape() != [m, k] + // The stack must be an affine relayout under a flat [M, K] + // reinterpret, both pure relayouts (no fill regions). + let [windowed, flat] = view.stages.as_slice() else { + return; + }; + if !windowed.is_fully_defined() + || !flat.is_fully_defined() + || !flat.layout.is_contiguous() + || flat.layout.offset() != 0 + || flat.layout.shape() != [m, k] { return; } - let inner_shape = &view.input_shape; + // The kernels substitute the windowed map as affine per-dim index + // arithmetic; validate it here so the lowering can rely on it. + if crate::view::affine_dim_indices(&windowed.layout, &windowed.input_shape).is_none() { + return; + } + let operand_shape = windowed.shape(); // The producer's dims must split cleanly into an `M` prefix and a // `K` suffix for the per-side flat-coordinate decomposition. let mut product = 1usize; - let mut k_start = inner_shape.len(); + let mut k_start = operand_shape.len(); while k_start > 0 && product < k { k_start -= 1; - let Some(next) = product.checked_mul(inner_shape[k_start]) else { + let Some(next) = product.checked_mul(operand_shape[k_start]) else { return; }; product = next; } if product != k || k_start == 0 - || k_start == inner_shape.len() - || inner_shape[..k_start].iter().product::() != m + || k_start == operand_shape.len() + || operand_shape[..k_start].iter().product::() != m { return; } - // The kernels decompose with u32 arithmetic; validate both sides - // here so the lowering can rely on it. + // The flat row/column coordinates decompose with u32 arithmetic. let probe = NaryExpr::DimIndex(0); - if crate::view::row_major_indices_from_flat(probe.clone(), &inner_shape[..k_start]) + if crate::view::row_major_indices_from_flat(probe.clone(), &operand_shape[..k_start]) .is_none() - || crate::view::row_major_indices_from_flat(probe, &inner_shape[k_start..]).is_none() + || crate::view::row_major_indices_from_flat(probe, &operand_shape[k_start..]).is_none() { return; } operation.first = view.input; operation.a = crate::matmul::MatrixOperand { - shape: inner_shape.clone(), + shape: operand_shape.into(), batch_dims: 0, row_dims: k_start, + base_map: Some(crate::matmul::OperandBaseMap { + layout: windowed.layout.clone(), + base_shape: windowed.input_shape.clone(), + }), }; } diff --git a/fusor-ml/core/src/compute_graph/resolve/recognize_attention.rs b/fusor-ml/core/src/compute_graph/resolve/recognize_attention.rs index d60fc99bc..5f096f78b 100644 --- a/fusor-ml/core/src/compute_graph/resolve/recognize_attention.rs +++ b/fusor-ml/core/src/compute_graph/resolve/recognize_attention.rs @@ -213,18 +213,20 @@ impl Resolver { let Some(view) = self.inner_view(mask_side) else { continue; }; + let Some(stage) = view.plain().filter(|stage| stage.is_fully_defined()) else { + continue; + }; let expected = Layout::from_parts(0, shape.clone().into(), [0, 0, kv_seq_len, 1].into()); - if !view.is_fully_defined() - || !layout_matches(Some(&view.layout), &expected) - || view.input_shape.as_ref() != [q_seq_len, kv_seq_len] + if !layout_matches(Some(&stage.layout), &expected) + || stage.input_shape.as_ref() != [q_seq_len, kv_seq_len] { continue; } if !self.exclusively_consumed(graph, mask_side, 1) { continue; } - matched_mask = Some((scaled_side, view.input, view.input_shape.to_vec())); + matched_mask = Some((scaled_side, view.input, stage.input_shape.to_vec())); break; } let (scaled_side, mask_base, base_shape) = matched_mask?; @@ -277,10 +279,10 @@ impl Resolver { // kᵀ: a transpose view attached to the (possibly GQA-expanded) k. let kt = self.inner_view(kt_inner)?; + let kt_stage = kt.plain().filter(|stage| stage.is_fully_defined())?; let expected_kt = Layout::contiguous(&expanded_shape).transpose(2, 3); - if !kt.is_fully_defined() - || !layout_matches(Some(&kt.layout), &expected_kt) - || kt.input_shape.as_ref() != expanded_shape + if !layout_matches(Some(&kt_stage.layout), &expected_kt) + || kt_stage.input_shape.as_ref() != expanded_shape || !self.exclusively_consumed(graph, kt_inner, 1) { return None; @@ -328,12 +330,17 @@ impl Resolver { let [batch, num_heads, kv_seq_len, head_dim] = *expanded_shape; let unexpanded = Some((inner, expanded_shape.to_vec())); - let Some(reinterpret) = self.inner_view(inner) else { + let Some(reinterpret_view) = self.inner_view(inner) else { + return unexpanded; + }; + let Some(reinterpret) = reinterpret_view + .plain() + .filter(|stage| stage.is_fully_defined()) + else { return unexpanded; }; // The flat reinterpret: contiguous rank-4 over a rank-5 grouped space. - if !reinterpret.is_fully_defined() - || reinterpret.input_shape.len() != 5 + if reinterpret.input_shape.len() != 5 || !layout_matches( Some(&reinterpret.layout), &Layout::contiguous(expanded_shape), @@ -352,7 +359,13 @@ impl Resolver { { return unexpanded; } - let Some(broadcast) = self.inner_view(reinterpret.input) else { + let Some(broadcast_view) = self.inner_view(reinterpret_view.input) else { + return unexpanded; + }; + let Some(broadcast) = broadcast_view + .plain() + .filter(|stage| stage.is_fully_defined()) + else { return unexpanded; }; let expected_broadcast = Layout::from_parts( @@ -360,14 +373,13 @@ impl Resolver { [b, num_kv_heads, groups, s, d].into(), [num_kv_heads * s * d, s * d, 0, d, 1].into(), ); - if !broadcast.is_fully_defined() - || !layout_matches(Some(&broadcast.layout), &expected_broadcast) + if !layout_matches(Some(&broadcast.layout), &expected_broadcast) || broadcast.input_shape.as_ref() != [b, num_kv_heads, s, d] || !self.exclusively_consumed(graph, inner, 1) - || !self.exclusively_consumed(graph, reinterpret.input, 1) + || !self.exclusively_consumed(graph, reinterpret_view.input, 1) { return unexpanded; } - Some((broadcast.input, vec![b, num_kv_heads, s, d])) + Some((broadcast_view.input, vec![b, num_kv_heads, s, d])) } } diff --git a/fusor-ml/core/src/compute_graph/resolve/recognize_cat.rs b/fusor-ml/core/src/compute_graph/resolve/recognize_cat.rs index 82fec4ffa..cc5efae09 100644 --- a/fusor-ml/core/src/compute_graph/resolve/recognize_cat.rs +++ b/fusor-ml/core/src/compute_graph/resolve/recognize_cat.rs @@ -551,7 +551,7 @@ impl Resolver { let ExecutionVariant::View(view) = &self.execution_graph[exec].variant else { return None; }; - if !view.is_fully_defined() || view.shape().len() != link.slices.len() { + if view.shape().len() != link.slices.len() { return None; } for (extent, slice) in view.shape().iter().zip(&*link.slices) { @@ -559,10 +559,12 @@ impl Resolver { return None; } } - let affine = affine_dim_indices(&view.layout, &view.input_shape)?; + let base_shape = &view.stages[0].input_shape; + let collapsed = view.composed_layout()?; + let affine = affine_dim_indices(&collapsed, base_shape)?; let mut shifted = Vec::with_capacity(affine.len()); - for (index, &extent) in affine.iter().zip(&*view.input_shape) { + for (index, &extent) in affine.iter().zip(&**base_shape) { let mut constant = index.constant as i64; let mut max_offset = 0i64; for &(dim, coefficient) in &index.terms { diff --git a/fusor-ml/core/src/compute_graph/tests.rs b/fusor-ml/core/src/compute_graph/tests.rs index 916572c64..197152e46 100644 --- a/fusor-ml/core/src/compute_graph/tests.rs +++ b/fusor-ml/core/src/compute_graph/tests.rs @@ -385,6 +385,57 @@ fn three_way_chunk_cat_collapses() { }); } +/// A staged (divmod) view feeding a single-read elementwise folds into the +/// consumer's loads: one dispatch, no gather. The same composition used to +/// cost a gather plus the elementwise. +#[test] +fn staged_view_folds_into_single_read_elementwise() { + pollster::block_on(async { + let Ok(device) = Device::new().await else { + return; + }; + + let (b, c, h, w) = (2usize, 4usize, 8usize, 8usize); + let (kh, kw) = (3usize, 3usize); + let (oh, ow) = (h - kh + 1, w - kw + 1); + let (m, k) = (b * oh * ow, c * kh * kw); + let input_host: Vec = (0..b * c * h * w).map(|i| (i % 13) as f32 * 0.1).collect(); + let input = Tensor::from_slice(&device, [b, c, h, w], &input_host); + + let out = { + let windows = input.restride([ + StrideSpec::dim(0, b), + StrideSpec::dim_with(2, oh, 1), + StrideSpec::dim_with(3, ow, 1), + StrideSpec::dim(1, c), + StrideSpec::dim(2, kh), + StrideSpec::dim(3, kw), + ]); + &windows.reshape([m, k]) + 1.0 + }; + let (_, kernels) = out.data.materialize(); + assert_eq!( + kernels, 1, + "the staged view should fold into the elementwise (got {kernels} dispatches)", + ); + + let result = out.as_slice::<2, f32>().await.unwrap(); + for mi in 0..m { + let (bi, rest) = (mi / (oh * ow), mi % (oh * ow)); + let (ohi, owi) = (rest / ow, rest % ow); + for ki in 0..k { + let (ci, rest) = (ki / (kh * kw), ki % (kh * kw)); + let (khi, kwi) = (rest / kw, rest % kw); + let expected = input_host[((bi * c + ci) * h + ohi + khi) * w + owi + kwi] + 1.0; + assert!( + (result[[mi, ki]] - expected).abs() < 1e-6, + "mismatch at [{mi}, {ki}]", + ); + } + } + }); +} + // --- implicit-GEMM conv (resolve/recognize.rs unflatten) --- /// Conv's im2col composition: a sliding-window restride whose rank-2 flatten diff --git a/fusor-ml/core/src/matmul/kernel.rs b/fusor-ml/core/src/matmul/kernel.rs index d87502ae5..a9c87652f 100644 --- a/fusor-ml/core/src/matmul/kernel.rs +++ b/fusor-ml/core/src/matmul/kernel.rs @@ -11,8 +11,7 @@ use crate::{ kernel_backend::{self, DirectKernel}, operation::Operation, tile_direct::{ - flatten_matrix_layout, flatten_matrix_layout_split, - tile_storage_read_with_direct_layout_typed, + flatten_matrix_layout_split, tile_storage_read_with_direct_layout_typed, tile_storage_write_with_direct_layout_typed, }, }, @@ -189,63 +188,75 @@ impl MatMulOperation { /// generic reduce re-derives it for every load and loses to a one-time /// gather. pub(crate) fn hardware_matmul_statically_viable(&self, device: &Device) -> bool { + self.coop_tile(device).is_some() + } + + /// The tile geometry the cooperative-matrix kernel would run with on + /// this device, `None` when any static gate fails and the contraction is + /// bound for the generic path. Shapes need not divide the tile: edge + /// tiles mask their fills and the output allocation pads to whole tiles. + pub(crate) fn coop_tile(&self, device: &Device) -> Option { if !self.can_use_hardware_matmul() || (self.datatype == DataTypeEnum::F16 && !device.f16_supported()) || !self.pre_element_wise[0].functions.is_empty() || !self.pre_element_wise[1].functions.is_empty() || !self.post_element_wise.functions.is_empty() { - return false; + return None; } let MatMulParams::CoopMatMul(params) = &self.parameters else { - return false; + return None; }; - if device.coop_token(params.kind()).is_none() - || !device - .subgroup_config() - .is_some_and(|config| config.is_fixed()) - { - return false; + device.coop_token(params.kind())?; + let subgroup_config = device.subgroup_config()?; + if !subgroup_config.is_fixed() { + return None; } - let (Ok(m), Ok(k), Ok(n)): (Result, Result, Result) = ( - self.a.rows().try_into(), - self.a.cols().try_into(), - self.b.cols().try_into(), - ) else { - return false; - }; - let Some(batch) = self + let (m, k, n): (u32, u32, u32) = ( + self.a.rows().try_into().ok()?, + self.a.cols().try_into().ok()?, + self.b.cols().try_into().ok()?, + ); + let batch = self .a .batch_shape() .iter() - .try_fold(1u32, |acc, &dim| acc.checked_mul(u32::try_from(dim).ok()?)) - else { - return false; - }; + .try_fold(1u32, |acc, &dim| acc.checked_mul(u32::try_from(dim).ok()?))?; let limits = device.limits(); - let Some(tile) = CoopTile::select( + let tile = CoopTile::select( m, k, n, limits .max_compute_workgroup_size_x .min(limits.max_compute_invocations_per_workgroup), - device - .subgroup_config() - .expect("checked above") - .max_size(), - ) else { - return false; - }; - // `CoopTile::select` guarantees the divisibility the kernel needs; - // only the dispatch-grid bound remains. - let Some(total_tiles) = (m / tile.bm) - .checked_mul(n / tile.bn) - .and_then(|tiles| tiles.checked_mul(batch)) - else { - return false; - }; - total_tiles <= limits.max_compute_workgroups_per_dimension + subgroup_config.max_size(), + )?; + let total_tiles = m + .div_ceil(tile.bm) + .checked_mul(n.div_ceil(tile.bn)) + .and_then(|tiles| tiles.checked_mul(batch))?; + (total_tiles <= limits.max_compute_workgroups_per_dimension).then_some(tile) + } + + /// Row-major strides of the logical output over its padded backing: + /// rows step `n_padded`, each batch block spans `m_padded * n_padded`. + fn padded_out_strides( + out_shape: &[usize], + m_padded: usize, + n_padded: usize, + ) -> Box<[usize]> { + let rank = out_shape.len(); + let mut strides = vec![0usize; rank]; + strides[rank - 1] = 1; + strides[rank - 2] = n_padded; + if rank >= 3 { + strides[rank - 3] = m_padded * n_padded; + for axis in (0..rank - 3).rev() { + strides[axis] = strides[axis + 1] * out_shape[axis + 1]; + } + } + strides.into() } fn build_hardware_matmul( @@ -255,15 +266,20 @@ impl MatMulOperation { input_b: &TensorData, output: &TensorData, ) -> Result { - let a_view = device_supported(flatten_matrix_layout_split( - input_a.layout(), - self.a.split(), - ))?; - let b_view = device_supported(flatten_matrix_layout_split( - input_b.layout(), - self.b.split(), - ))?; - let y_view = device_supported(flatten_matrix_layout(output.layout()))?; + // Operands with a base map read their producer through it: compose + // with the runtime buffer layout, then flatten with the operand's + // dim grouping. + let operand_layout = + |operand: &MatrixOperand, input: &TensorData| -> Option { + match &operand.base_map { + Some(map) => crate::view::compose_layouts(&map.layout, input.layout()), + None => Some(input.layout().clone()), + } + }; + let a_layout = device_supported(operand_layout(&self.a, input_a))?; + let b_layout = device_supported(operand_layout(&self.b, input_b))?; + let a_view = device_supported(flatten_matrix_layout_split(&a_layout, self.a.split()))?; + let b_view = device_supported(flatten_matrix_layout_split(&b_layout, self.b.split()))?; let m: u32 = self .a @@ -290,12 +306,7 @@ impl MatMulOperation { .map_err(|_| kernel_backend::DeviceNotSupported)?; let batch_m = device_supported(batch.checked_mul(m))?; let batch_k = device_supported(batch.checked_mul(k))?; - if a_view.rows != batch_m - || a_view.cols != k - || b_view.rows != batch_k - || b_view.cols != n - || y_view.rows != batch_m - || y_view.cols != n + if a_view.rows != batch_m || a_view.cols != k || b_view.rows != batch_k || b_view.cols != n { return Err(kernel_backend::DeviceNotSupported); } @@ -304,31 +315,49 @@ impl MatMulOperation { // Only the cooperative-matrix route stays hand-specialized; gemv // shapes lower through the generic subgroup-per-output reduce, and // fused chains lower through the generic tiled reduce. - if !self.pre_element_wise[0].functions.is_empty() - || !self.pre_element_wise[1].functions.is_empty() - || !self.post_element_wise.functions.is_empty() - { - return Err(kernel_backend::DeviceNotSupported); - } + let tile = device_supported(self.coop_tile(device))?; let subgroup_config = device_supported(device.subgroup_config())?; - if !subgroup_config.is_fixed() { - return Err(kernel_backend::DeviceNotSupported); - } let MatMulParams::CoopMatMul(params) = &self.parameters else { return Err(kernel_backend::DeviceNotSupported); }; - let kind = params.kind(); - let coop = device_supported(device.coop_token(kind))?; - let tile = device_supported(CoopTile::select( - m, - k, - n, - device - .limits() - .max_compute_workgroup_size_x - .min(device.limits().max_compute_invocations_per_workgroup), - subgroup_config.max_size(), - ))?; + let coop = device_supported(device.coop_token(params.kind()))?; + + // The store covers whole tiles, so `y` is the padded matrix: rows + // padded to `ceil(m / bm) * bm` per batch and columns to + // `ceil(n / bn) * bn`, allocated by `inputs()` with the logical + // output viewing it. Verify the output really has that geometry — + // a mismatch (the allocation predicted a different tile) falls back + // to the generic path, which writes through the logical layout. + let m_padded = m.div_ceil(tile.bm) * tile.bm; + let n_padded = n.div_ceil(tile.bn) * tile.bn; + let expected_strides = Self::padded_out_strides( + &self.out_shape, + m_padded as usize, + n_padded as usize, + ); + let padded_elements = device_supported( + (batch as usize) + .checked_mul(m_padded as usize) + .and_then(|rows| rows.checked_mul(n_padded as usize)), + )?; + let padded_bytes = padded_elements as u64 * self.datatype.element_size() as u64; + if output.layout().offset() != 0 + || output.layout().strides() != &*expected_strides + || padded_bytes > output.buffer().size() + { + return Err(kernel_backend::DeviceNotSupported); + } + let batch_m_padded = device_supported(batch.checked_mul(m_padded))?; + let y_view = crate::mir::tile_direct::DirectMatrixLayout { + rows: batch_m_padded, + cols: n_padded, + offset: 0, + layout: tile_ir::Layout::strided( + tile_ir::MemoryLevel::Storage, + tile_ir::Shape::new([batch_m_padded, n_padded]), + &[n_padded, 1], + ), + }; let max_wg_per_dim = device.limits().max_compute_workgroups_per_dimension; let datatype = self.datatype; @@ -493,11 +522,38 @@ impl Operation for MatMulOperation { let a = nodes.get_result(self.first).unwrap(); let b = nodes.get_result(self.second).unwrap(); let device = a.device(); - let output_tensor = TensorData::new_for_shape( - device, - &self.out_shape, - self.post_element_wise.out_datatype(), - ); + let datatype = self.post_element_wise.out_datatype(); + // The coop kernel stores whole tiles: pad the backing to tile + // multiples and view the logical shape over it (consumers never + // read the pad region). Shapes that already divide the tile — and + // anything bound for the generic path — allocate exactly. + let (m, n) = (self.a.rows(), self.b.cols()); + let padded = self.coop_tile(&device).and_then(|tile| { + let m_padded = m.div_ceil(tile.bm as usize) * tile.bm as usize; + let n_padded = n.div_ceil(tile.bn as usize) * tile.bn as usize; + (m_padded != m || n_padded != n).then_some((m_padded, n_padded)) + }); + let output_tensor = match padded { + Some((m_padded, n_padded)) => { + let batch: usize = self.a.batch_shape().iter().product(); + let backing = TensorData::new_for_shape( + &device, + &[batch, m_padded, n_padded], + datatype, + ); + TensorData::new_from_parts( + &device, + backing.buffer().clone(), + crate::Layout::from_parts( + 0, + self.out_shape.clone(), + Self::padded_out_strides(&self.out_shape, m_padded, n_padded), + ), + datatype, + ) + } + None => TensorData::new_for_shape(&device, &self.out_shape, datatype), + }; vec![a.into(), b.into(), output_tensor.into()] } diff --git a/fusor-ml/core/src/matmul/mod.rs b/fusor-ml/core/src/matmul/mod.rs index 0c01c18a4..e6f888fef 100644 --- a/fusor-ml/core/src/matmul/mod.rs +++ b/fusor-ml/core/src/matmul/mod.rs @@ -1,5 +1,5 @@ use crate::{ - Device, Tensor, compute_graph::NodeIndex, kernel_selection::CooperativeMatrixKind, + Device, Layout, Tensor, compute_graph::NodeIndex, kernel_selection::CooperativeMatrixKind, nary_wise::UnaryFunctionChain, tensor::DataTypeEnum, }; @@ -25,17 +25,38 @@ pub enum MatMulParams { CoopMatMul(coop_gemm::CoopGemmParams), } -/// One matmul operand's dim grouping: the producer node's logical dims -/// split into `batch_dims` leading batch dims, `row_dims` row dims, and -/// column dims for the rest. The row and column groups flatten to the two -/// matrix axes. A plain `[batch.., rows, cols]` operand has one dim per -/// group; conv's im2col operand keeps the windowed view's dims, and the -/// kernels divmod the flat matrix coordinates back apart per load. +/// An affine relayout between an operand's dims and its node's logical +/// space: conv's sliding windows. The kernels concretize it lazily — the +/// coop path composes it with the node's runtime buffer layout, the generic +/// reduce substitutes it into the load coordinates. +#[derive(Debug, Clone)] +pub(crate) struct OperandBaseMap { + pub(crate) layout: Layout, + pub(crate) base_shape: Box<[usize]>, +} + +impl std::hash::Hash for OperandBaseMap { + fn hash(&self, state: &mut H) { + self.layout.offset().hash(state); + self.layout.shape().hash(state); + self.layout.strides().hash(state); + self.base_shape.hash(state); + } +} + +/// One matmul operand's dim grouping: the operand's dims split into +/// `batch_dims` leading batch dims, `row_dims` row dims, and column dims +/// for the rest. The row and column groups flatten to the two matrix axes. +/// A plain `[batch.., rows, cols]` operand has one dim per group; conv's +/// im2col operand keeps the windowed view's dims (mapped onto the node by +/// `base_map`), and the kernels divmod the flat matrix coordinates back +/// apart per load. #[derive(Debug, Clone, Hash)] pub(crate) struct MatrixOperand { pub(crate) shape: Box<[usize]>, pub(crate) batch_dims: usize, pub(crate) row_dims: usize, + pub(crate) base_map: Option, } impl MatrixOperand { @@ -45,12 +66,14 @@ impl MatrixOperand { shape: shape.into(), batch_dims: shape.len() - 2, row_dims: 1, + base_map: None, } } - /// One dim per group: the operand's shape is the logical matmul shape. + /// One dim per group, reading the node's own dims directly: the + /// operand's shape is the logical matmul shape. pub(crate) fn is_plain(&self) -> bool { - self.row_dims == 1 && self.batch_dims + 2 == self.shape.len() + self.row_dims == 1 && self.batch_dims + 2 == self.shape.len() && self.base_map.is_none() } pub(crate) fn batch_shape(&self) -> &[usize] { @@ -82,9 +105,10 @@ impl MatrixOperand { } /// Index expressions reading the operand at the contraction coordinates: - /// batch dims index through, and the flat row/column coordinates - /// decompose over the row/column groups (the identity for single-dim - /// groups, so plain operands load with bare `DimIndex` coordinates). + /// batch dims index through, the flat row/column coordinates decompose + /// over the row/column groups (the identity for single-dim groups, so + /// plain operands load with bare `DimIndex` coordinates), and the + /// operand coordinates map through `base_map` to the node's own dims. pub(crate) fn index_expressions( &self, row_dim: usize, @@ -100,7 +124,14 @@ impl MatrixOperand { crate::view::row_major_indices_from_flat(NaryExpr::DimIndex(col_dim), self.col_shape()) .expect("operand dims fit u32"), ); - indices + match &self.base_map { + None => indices, + Some(map) => crate::view::affine_dim_indices(&map.layout, &map.base_shape) + .expect("validated when the resolver attached the base map") + .iter() + .map(|index| index.to_expr(&indices)) + .collect(), + } } } @@ -416,6 +447,15 @@ mod selection_tests { select(1024, 1024, 1024, 128), Some(CoopTile::new(64, 64, 16)) ); - assert_eq!(select(1000, 1024, 1024, 512), None); + // M=1000 divides nothing: the masked-edge fallback picks the tile + // with the least padded work (all candidates pad M to 1024; the + // preference order breaks the tie toward the biggest tile). + assert_eq!( + select(1000, 1024, 1024, 512), + Some(CoopTile::new(128, 128, 16)) + ); + // Tiny N inflates every candidate's padding past the waste bound: + // gemv-shaped contractions stay on the generic path. + assert_eq!(select(1000, 1024, 16, 512), None); } } diff --git a/fusor-ml/core/src/matmul/variants.rs b/fusor-ml/core/src/matmul/variants.rs index 15516cecf..cb6c9d95f 100644 --- a/fusor-ml/core/src/matmul/variants.rs +++ b/fusor-ml/core/src/matmul/variants.rs @@ -183,7 +183,7 @@ impl CoopTile { max_subgroup_size: u32, ) -> Option { let tiles_for = |bm: u32, bn: u32| -> u32 { (m / bm) * (n / bn) }; - if !k.is_multiple_of(16) { + if m == 0 || n == 0 || k == 0 { return None; } // Tile256x256 single-buffer has lower memory traffic (sqrt-min) but @@ -229,6 +229,26 @@ impl CoopTile { return Some(tile); } } - None + + // Shapes that divide no tile run with masked edge tiles: pick the + // candidate minimizing padded work, in preference order on ties. + // Selections whose padding inflates the output by more than a + // quarter stay on the generic path — that bound also keeps + // gemv-shaped contractions (tiny M or N) off the tile kernels. + let mut best: Option<(u64, Self)> = None; + for (bm, bn) in [(128, 128), (128, 64), (64, 128), (64, 64)] { + let tile = Self::new(bm, bn, 16); + if !tile.workgroup_size_supported(max_workgroup_size_x, max_subgroup_size) { + continue; + } + let padded = u64::from(m.div_ceil(bm) * bm) * u64::from(n.div_ceil(bn) * bn); + if padded * 4 > u64::from(m) * u64::from(n) * 5 { + continue; + } + if best.is_none_or(|(best_padded, _)| padded < best_padded) { + best = Some((padded, tile)); + } + } + best.map(|(_, tile)| tile) } } diff --git a/fusor-ml/core/src/quantized/matmul/kernel.rs b/fusor-ml/core/src/quantized/matmul/kernel.rs index 442d5c891..d104e93ea 100644 --- a/fusor-ml/core/src/quantized/matmul/kernel.rs +++ b/fusor-ml/core/src/quantized/matmul/kernel.rs @@ -70,14 +70,16 @@ impl QMatMulOperation { let scratch = TensorData::new_for_shape(&device, &padded_in_shape, input.datatype()); let pad_copy = crate::view::ViewOperation { input: self.input, - layout: Layout::from_parts( - 0, - padded_in_shape.into(), - Layout::continuous_strides(in_shape), - ), - input_shape: in_shape.into(), - defined: in_shape.into(), - fill: crate::view::zero_scalar(input.datatype()), + stages: vec![crate::view::ViewStage { + layout: Layout::from_parts( + 0, + padded_in_shape.into(), + Layout::continuous_strides(in_shape), + ), + input_shape: in_shape.into(), + defined: in_shape.into(), + fill: crate::view::zero_scalar(input.datatype()), + }], datatype: input.datatype(), }; let pad_workgroup = pad_copy diff --git a/fusor-ml/core/src/view.rs b/fusor-ml/core/src/view.rs index 1c9d8b0d8..82c97bc52 100644 --- a/fusor-ml/core/src/view.rs +++ b/fusor-ml/core/src/view.rs @@ -15,67 +15,193 @@ use crate::{ const BLOCKSIZE: u32 = 256; -/// A zero-dispatch view of a node's logical value space. +/// One stage of a view: a quasi-affine relayout of the space below it. /// -/// `layout` maps output coordinates to flat indices in the input's logical -/// row-major value space (`flat = offset + Σ coord_i * strides[i]`). Because -/// it indexes the *logical* space — not a concrete buffer — every producer's -/// output is contiguous by definition, so restride, transpose, broadcast, -/// slice, reshape, and resize are all plain stride arithmetic here. -/// -/// `defined` is a prefix box of `layout.shape()`: coordinates with -/// `coord_i < defined[i]` for every axis read input data; anything outside -/// reads `fill`. A fully-defined view (`defined == shape`) is a pure -/// relayout; a partially-defined view is a clip + pad (resize). +/// `layout` maps the stage's output coordinates to flat indices in the +/// row-major space of `input_shape` (the stage below, or the input node's +/// logical value space). `defined` is a prefix box of `layout.shape()`: +/// coordinates with `coord_i < defined[i]` for every axis read data; +/// anything outside reads `fill`. #[derive(Clone, Debug)] -pub(crate) struct ViewOperation { - pub(crate) input: NodeIndex, +pub(crate) struct ViewStage { pub(crate) layout: Layout, - /// Logical shape of the input node's value space — the space `layout` - /// indexes into. Used to recover per-dimension input coordinates when the - /// view cannot stay a flat index (gather fallback) and to bounds-check. pub(crate) input_shape: Box<[usize]>, pub(crate) defined: Box<[usize]>, pub(crate) fill: NaryScalar, +} + +impl ViewStage { + pub(crate) fn fully_defined(layout: Layout, input_shape: impl Into>) -> Self { + let defined = layout.shape().into(); + Self { + layout, + input_shape: input_shape.into(), + defined, + fill: NaryScalar::U32(0), + } + } + + pub(crate) fn is_fully_defined(&self) -> bool { + self.defined.as_ref() == self.layout.shape() + } + + pub(crate) fn shape(&self) -> &[usize] { + self.layout.shape() + } + + /// Coordinate expressions of the space below, given this stage's output + /// coordinate expressions: the divmod-free affine per-dim form when one + /// exists, otherwise the flat index delinearized over `input_shape`. + /// Returns the coordinates and whether delinearization was needed. + /// Coordinates are clamped in-bounds when `clamp` is set (callers set it + /// for partially-defined stages, where out-of-box coordinates still + /// evaluate the data branch of the fill select). + pub(crate) fn coords_below( + &self, + out: &[NaryExpr], + clamp: bool, + ) -> Option<(Vec, bool)> { + let (coords, delinearized) = + match affine_dim_indices(&self.layout, &self.input_shape) { + Some(affine) => ( + affine.iter().map(|index| index.to_expr(out)).collect(), + false, + ), + None => { + let flat = flat_index_expression(&self.layout, out)?; + (row_major_indices_from_flat(flat, &self.input_shape)?, true) + } + }; + if !clamp { + return Some((coords, delinearized)); + } + let clamped = coords + .into_iter() + .zip(&*self.input_shape) + .map(|(index, &extent)| { + if extent == 0 { + index + } else { + NaryExpr::unary_op( + index, + "clamp_dim", + NaryOp::MinConst(NaryScalar::U32(extent as u32 - 1)), + DataTypeEnum::U32, + DataTypeEnum::U32, + ) + } + }) + .collect(); + Some((clamped, delinearized)) + } + + /// `Some(condition)` over this stage's output coordinate expressions when + /// partially defined; `None` for a pure relayout. + pub(crate) fn bounds_condition(&self, out: &[NaryExpr]) -> Option { + if self.is_fully_defined() { + return None; + } + let mut condition = NaryExpr::scalar(NaryScalar::U32(1)); + for (dim, (&defined, &size)) in self.defined.iter().zip(self.layout.shape()).enumerate() { + if defined >= size { + continue; + } + let lt_defined = NaryExpr::unary_op( + out[dim].clone(), + "lt_defined", + NaryOp::LessConst(NaryScalar::U32(defined as u32)), + DataTypeEnum::U32, + DataTypeEnum::U32, + ); + condition = NaryExpr::mul(condition, lt_defined, DataTypeEnum::U32); + } + Some(condition) + } +} + +/// A zero-dispatch view of a node's logical value space, as a stack of +/// quasi-affine stages. +/// +/// `stages[0]` indexes the input node's logical row-major value space; +/// each later stage indexes the output space of the stage below it. Most +/// relayouts merge into the top stage at construction; a stage boundary +/// remains only where the composition is not expressible as one strided +/// layout (a reshape regrouping elements across non-mergeable strides) or +/// where a fill region intervenes. The compiler chooses how a consumer +/// concretizes the map: a zero-cost buffer view when the stack collapses +/// over the buffer's layout, index expressions folded into the consumer's +/// loads, or a gather dispatch materializing the result. +/// +/// `input` is never itself a view node: view-over-view always collapses +/// into one node at construction. +#[derive(Clone, Debug)] +pub(crate) struct ViewOperation { + pub(crate) input: NodeIndex, + /// Innermost first; never empty. + pub(crate) stages: Vec, pub(crate) datatype: DataTypeEnum, } impl ViewOperation { - /// A fully-defined view: pure relayout of the input's logical space. + /// A fully-defined single-stage view: pure relayout of the input's + /// logical space. pub(crate) fn fully_defined( input: NodeIndex, layout: Layout, input_shape: impl Into>, datatype: DataTypeEnum, ) -> Self { - let defined = layout.shape().into(); Self { input, - layout, - input_shape: input_shape.into(), - defined, - fill: zero_scalar(datatype), + stages: vec![ViewStage::fully_defined(layout, input_shape)], datatype, } } pub(crate) fn is_fully_defined(&self) -> bool { - self.defined.as_ref() == self.layout.shape() + self.stages.iter().all(ViewStage::is_fully_defined) } pub(crate) fn shape(&self) -> &[usize] { - self.layout.shape() + self.stages.last().expect("views have stages").shape() + } + + /// The single stage, when this view is one relayout deep. Pattern + /// matchers that need a specific affine layout use this and bail on + /// staged views. + pub(crate) fn plain(&self) -> Option<&ViewStage> { + match self.stages.as_slice() { + [stage] => Some(stage), + _ => None, + } + } + + /// Collapse the stack into one strided layout over the input's logical + /// space, when every stage is fully defined and the compositions hold. + pub(crate) fn composed_layout(&self) -> Option { + if !self.is_fully_defined() { + return None; + } + let mut stages = self.stages.iter(); + let mut layout = stages.next()?.layout.clone(); + for stage in stages { + layout = compose_layouts(&stage.layout, &layout)?; + } + Some(layout) } /// Resolve as a zero-cost view over the input's concrete buffer, if the - /// view's logical layout composes with the buffer's layout as a single - /// strided layout. Partially-defined views never qualify — their fill - /// region has no backing memory. + /// stack composes with the buffer's layout as a single strided layout. + /// Partially-defined views never qualify — their fill region has no + /// backing memory. pub(crate) fn try_map_tensor(&self, input: &TensorData) -> Option { if !self.is_fully_defined() { return None; } - let composed = compose_layouts(&self.layout, input.layout())?; + let mut composed = input.layout().clone(); + for stage in &self.stages { + composed = compose_layouts(&stage.layout, &composed)?; + } Some(TensorData::new_from_parts( input.device(), input.buffer().clone(), @@ -84,73 +210,80 @@ impl ViewOperation { )) } - /// The gather expression materializing this view: per output coordinate, - /// load the input at the mapped logical coordinates, or `fill` outside - /// the defined box. - fn copy_expression(&self) -> Option { - let flat = self.flat_logical_expression()?; - let indices = row_major_indices_from_flat(flat, &self.input_shape)?; - let copied = NaryExpr::indexed_input(0, indices); - if self.is_fully_defined() { - return Some(copied); + /// The value of this view at the given output coordinate expressions, as + /// a load of input `input_slot` walked through every stage (with fill + /// selects where stages are partially defined). Returns the expression + /// and whether any stage needed delinearization (divmod address + /// arithmetic) — the signal consumers use to decide whether folding this + /// view into their loads beats materializing it. + pub(crate) fn value_expression( + &self, + input_slot: usize, + out: &[NaryExpr], + ) -> Option<(NaryExpr, bool)> { + let mut coords: Vec = out.to_vec(); + // (condition, fill) per partially-defined stage, outermost first. + let mut fills: Vec<(NaryExpr, NaryScalar)> = Vec::new(); + let mut delinearized = false; + for stage in self.stages.iter().rev() { + if let Some(condition) = stage.bounds_condition(&coords) { + fills.push((condition, stage.fill)); + } + let (below, stage_delinearized) = + stage.coords_below(&coords, !stage.is_fully_defined())?; + delinearized |= stage_delinearized; + coords = below; } - Some(NaryExpr::select( - self.in_defined_bounds_expression(), - copied, - NaryExpr::scalar(self.fill), - DataTypeEnum::U32, - self.datatype, - )) + let mut value = NaryExpr::IndexedInput { + input_idx: input_slot, + indices: coords, + }; + // Innermost select wraps first so each stage's fill masks everything + // below it. + for (condition, fill) in fills.into_iter().rev() { + value = NaryExpr::select( + condition, + value, + NaryExpr::scalar(fill), + DataTypeEnum::U32, + self.datatype, + ); + } + Some((value, delinearized)) } - /// `offset + Σ DimIndex(d) * strides[d]` as a u32 expression. - fn flat_logical_expression(&self) -> Option { - let mut flat = NaryExpr::scalar(NaryScalar::U32(self.layout.offset().try_into().ok()?)); - for (axis, (&stride, &dim)) in self - .layout - .strides() - .iter() - .zip(self.layout.shape()) - .enumerate() - { - if stride == 0 || dim == 1 { - continue; - } - let stride: u32 = stride.try_into().ok()?; - let dim_index = NaryExpr::DimIndex(axis); - let term = if stride == 1 { - dim_index - } else { - NaryExpr::unary_op( - dim_index, - "mul_const", - NaryOp::MulConst(NaryScalar::U32(stride)), - DataTypeEnum::U32, - DataTypeEnum::U32, - ) - }; - flat = NaryExpr::add(flat, term, DataTypeEnum::U32); - } - Some(flat) + /// The gather expression materializing this view: per output coordinate, + /// load the input at the mapped coordinates, or fill outside the defined + /// boxes. + fn copy_expression(&self) -> Option { + let out: Vec = (0..self.shape().len()).map(NaryExpr::DimIndex).collect(); + Some(self.value_expression(0, &out)?.0) } +} - fn in_defined_bounds_expression(&self) -> NaryExpr { - let mut condition = NaryExpr::scalar(NaryScalar::U32(1)); - for (dim, (&defined, &size)) in self.defined.iter().zip(self.layout.shape()).enumerate() { - if defined >= size { - continue; - } - let lt_defined = NaryExpr::unary_op( - NaryExpr::DimIndex(dim), - "lt_defined", - NaryOp::LessConst(NaryScalar::U32(defined as u32)), +/// `offset + Σ coords[d] * strides[d]` as a u32 expression. +fn flat_index_expression(layout: &Layout, coords: &[NaryExpr]) -> Option { + let mut flat = NaryExpr::scalar(NaryScalar::U32(layout.offset().try_into().ok()?)); + for (axis, (&stride, &dim)) in layout.strides().iter().zip(layout.shape()).enumerate() { + if stride == 0 || dim == 1 { + continue; + } + let stride: u32 = stride.try_into().ok()?; + let coord = coords[axis].clone(); + let term = if stride == 1 { + coord + } else { + NaryExpr::unary_op( + coord, + "mul_const", + NaryOp::MulConst(NaryScalar::U32(stride)), DataTypeEnum::U32, DataTypeEnum::U32, - ); - condition = NaryExpr::mul(condition, lt_defined, DataTypeEnum::U32); - } - condition + ) + }; + flat = NaryExpr::add(flat, term, DataTypeEnum::U32); } + Some(flat) } pub(crate) fn zero_scalar(datatype: DataTypeEnum) -> NaryScalar { @@ -403,12 +536,15 @@ pub(crate) fn row_major_indices_from_flat( impl Operation for ViewOperation { fn hash_kernel_fields(&self, state: &mut rustc_hash::FxHasher) { - self.layout.offset().hash(state); - self.layout.shape().hash(state); - self.layout.strides().hash(state); - self.input_shape.hash(state); - self.defined.hash(state); - self.fill.hash(state); + self.stages.len().hash(state); + for stage in &self.stages { + stage.layout.offset().hash(state); + stage.layout.shape().hash(state); + stage.layout.strides().hash(state); + stage.input_shape.hash(state); + stage.defined.hash(state); + stage.fill.hash(state); + } } fn workgroup_shape_constraints(&self, _: &crate::Device) -> WorkgroupShapeConstraints { @@ -438,8 +574,7 @@ impl Operation for ViewOperation { fn inputs(&self, nodes: &crate::compute_graph::ComputeGraphInner) -> Vec { let input = nodes.get_cached_result(self.input).unwrap().clone(); - let output = - TensorData::new_for_shape(input.device(), self.layout.shape(), input.datatype()); + let output = TensorData::new_for_shape(input.device(), self.shape(), input.datatype()); vec![input.into(), output.into()] } @@ -456,7 +591,7 @@ impl Operation for ViewOperation { let operation = ElementwiseOperation { inputs: vec![self.input], expression: self.copy_expression()?, - shape: self.layout.shape().into(), + shape: self.shape().into(), output_datatype: self.datatype, }; crate::nary_direct::build_nary_direct_kernel_to_output( @@ -471,8 +606,7 @@ impl Operation for ViewOperation { fn name(&self) -> String { format!( "view_{}", - self.layout - .shape() + self.shape() .iter() .map(|x| x.to_string()) .collect::>() @@ -482,35 +616,80 @@ impl Operation for ViewOperation { } impl Tensor { - /// The view spec to layer a new view on top of this tensor: composes with - /// an existing fully-defined view (so chains collapse at construction) or - /// starts from this tensor's logical space. - fn view_base(&self) -> (NodeIndex, Layout, Box<[usize]>) { - if let Some(view) = self.device().compute_graph().get_view(self.key()) - && view.is_fully_defined() - { - return (view.input, view.layout.clone(), view.input_shape.clone()); + /// The stage stack to extend when layering a new view on this tensor: + /// the existing view's parts when this tensor is an unresolved view, or + /// an empty stack over this tensor's own logical space. + fn view_parts(&self) -> (NodeIndex, Vec) { + match self.device().compute_graph().get_view(self.key()) { + Some(view) => (view.input, view.stages), + None => (self.key(), Vec::new()), } - ( - self.key(), - Layout::contiguous(self.shape()), - self.shape().into(), - ) } fn add_view_op(&self, op: ViewOperation) -> Tensor { Tensor::from_parts(self.data.view(op)) } + /// Layer `stage` — a relayout of this tensor's current output space — + /// onto the stack: merged into the top stage when the composition stays + /// one strided layout over a pure relayout, pushed as a new stage + /// otherwise. Either way the result is a single view node over the + /// non-view base. + fn add_view_stage(&self, stage: ViewStage) -> Tensor { + let (input, mut stages) = self.view_parts(); + let stage = match stages.last_mut() { + Some(top) if top.is_fully_defined() => { + match compose_layouts(&stage.layout, &top.layout) { + Some(composed) => { + top.layout = composed; + top.defined = stage.defined; + top.fill = stage.fill; + None + } + None => Some(stage), + } + } + _ => Some(stage), + }; + if let Some(stage) = stage { + stages.push(stage); + } + self.add_view_op(ViewOperation { + input, + stages, + datatype: self.datatype(), + }) + } + pub fn restride(&self, specs: impl Into>) -> Tensor { let specs = specs.into(); - let (input, base, input_shape) = self.view_base(); - self.add_view_op(ViewOperation::fully_defined( + let (input, mut stages) = self.view_parts(); + match stages.last_mut() { + // Restride is direct stride arithmetic on the top stage — the + // composition is infallible, no merge check needed. + Some(top) if top.is_fully_defined() => { + top.layout = top.layout.restride(&specs); + top.defined = top.layout.shape().into(); + } + Some(top) => { + let below: Box<[usize]> = top.shape().into(); + stages.push(ViewStage::fully_defined( + Layout::contiguous(&below).restride(&specs), + below, + )); + } + None => { + stages.push(ViewStage::fully_defined( + Layout::contiguous(self.shape()).restride(&specs), + self.shape(), + )); + } + } + self.add_view_op(ViewOperation { input, - base.restride(&specs), - input_shape, - self.datatype(), - )) + stages, + datatype: self.datatype(), + }) } /// Replace the tensor's layout with `new_layout` over its logical value @@ -531,18 +710,7 @@ impl Tensor { "restride_layout out of bounds: layout reaches element {max_index} \ but the input has only {numel} elements" ); - let (input, base, input_shape) = self.view_base(); - let layout = compose_layouts(&new_layout, &base).unwrap_or_else(|| { - panic!( - "restride_layout could not compose {new_layout:?} with the existing view {base:?}" - ) - }); - self.add_view_op(ViewOperation::fully_defined( - input, - layout, - input_shape, - self.datatype(), - )) + self.add_view_stage(ViewStage::fully_defined(new_layout, self.shape())) } pub fn broadcast_as(&self, out_shape: impl AsRef<[usize]>) -> Tensor { @@ -613,20 +781,11 @@ impl Tensor { self.shape(), new_shape ); + // When the reshape regroups elements across non-mergeable strides + // the reinterpret stays its own stage; otherwise it merges into the + // top stage. let reinterpret = Layout::contiguous(new_shape); - let (input, base, input_shape) = self.view_base(); - let op = match compose_layouts(&reinterpret, &base) { - Some(layout) => { - ViewOperation::fully_defined(input, layout, input_shape, self.datatype()) - } - // The reshape regroups elements across the view's non-contiguous - // strides: keep it as a flat reinterpret of this tensor's own - // logical space (a chained view). - None => { - ViewOperation::fully_defined(self.key(), reinterpret, self.shape(), self.datatype()) - } - }; - self.add_view_op(op) + self.add_view_stage(ViewStage::fully_defined(reinterpret, self.shape())) } /// Resize to `new_shape`, clipping or zero-padding each axis: coordinates @@ -649,26 +808,12 @@ impl Tensor { // Within the defined box the old row-major strides address the input // exactly; outside it the load is masked to `fill`. let resize = Layout::from_parts(0, new_shape.into(), Layout::continuous_strides(old_shape)); - let (input, base, input_shape) = self.view_base(); - let op = match compose_layouts(&resize, &base) { - Some(layout) => ViewOperation { - input, - layout, - input_shape, - defined, - fill: zero_scalar(self.datatype()), - datatype: self.datatype(), - }, - None => ViewOperation { - input: self.key(), - layout: resize, - input_shape: old_shape.into(), - defined, - fill: zero_scalar(self.datatype()), - datatype: self.datatype(), - }, - }; - self.add_view_op(op) + self.add_view_stage(ViewStage { + layout: resize, + input_shape: old_shape.into(), + defined, + fill: zero_scalar(self.datatype()), + }) } pub fn flatten_last_n(&self, from_end: usize) -> Tensor { diff --git a/fusor-ml/tile-ir-kernels/src/kernels/matmul.rs b/fusor-ml/tile-ir-kernels/src/kernels/matmul.rs index 4f749ed69..022b211a5 100644 --- a/fusor-ml/tile-ir-kernels/src/kernels/matmul.rs +++ b/fusor-ml/tile-ir-kernels/src/kernels/matmul.rs @@ -87,18 +87,20 @@ pub fn try_batched_coop_matmul( } = config; let subgroup = subgroups.token(); let DenseCoopMatmulTile { bm, bn, bk } = tile; + // Shapes need not divide the tile geometry: edge tiles fill zero past + // the logical extents, and the caller provides `y` with its rows padded + // to `ceil(m / bm) * bm` per batch and its columns to `ceil(n / bn) * bn` + // (the stores cover whole tiles; the pad region holds garbage the + // logical view never reads). if !subgroups.is_fixed() || epilogues.pre_a.is_some() || epilogues.pre_b.is_some() || epilogues.post.is_some() - || !shape.m.is_multiple_of(bm) - || !shape.n.is_multiple_of(bn) - || !shape.k.is_multiple_of(bk) || !cooperative_store_layout_supported(y.layout()) { return false; } - let total_tiles = shape.batch * (shape.m / bm) * (shape.n / bn); + let total_tiles = shape.batch * shape.m.div_ceil(bm) * shape.n.div_ceil(bn); if total_tiles > max_workgroups_per_dimension { return false; } @@ -255,6 +257,8 @@ fn coop_stage_and_mma( k_base: &Tile, sg_row_base: &Tile, sg_col_base_in_pass: &Tile, + a_bounds: &[Option; 2], + b_bounds: &[Option; 2], accs: &[Vec], tile_rows_per_sg: u32, tile_cols_per_sg: u32, @@ -262,12 +266,19 @@ fn coop_stage_and_mma( coop_dim: u32, scalar: ScalarElement, ) { - program.fill_tile(a_tile, a, a_batch_base.clone() + row_base.clone(), k_base); - program.fill_tile( + program.fill_tile_bounded( + a_tile, + a, + a_batch_base.clone() + row_base.clone(), + k_base, + a_bounds.clone(), + ); + program.fill_tile_bounded( b_tile, b, b_batch_base.clone() + k_base.clone(), pass_col_base, + b_bounds.clone(), ); program.workgroup_barrier(); @@ -383,10 +394,14 @@ fn batched_coop_matmul_perf_single( let scalar = scalar_of(a.element()); - let tiles_m = shape.m / bm; - let tiles_n = shape.n / bn; + let tiles_m = shape.m.div_ceil(bm); + let tiles_n = shape.n.div_ceil(bn); let total_tiles = shape.batch * tiles_m * tiles_n; - let k_iterations = shape.k / bk; + let k_iterations = shape.k.div_ceil(bk); + // The y rows are padded to whole tiles per batch (the caller allocates + // the pad region); A/B are logical and edge tiles fill zero past the + // extents. + let m_padded = tiles_m * bm; let a_tile = program.alloc_workgroup_tile_padded(scalar, bm, bk, 1); let b_tile = program.alloc_workgroup_tile_padded(scalar, bk, bn_pass, 1); @@ -404,7 +419,15 @@ fn batched_coop_matmul_perf_single( let col_base = n_tile * bn; let a_batch_base = batch.clone() * shape.m; let b_batch_base = batch.clone() * shape.k; - let y_batch_base = batch * shape.m; + let y_batch_base = batch * m_padded; + let a_bounds: [Option; 2] = [ + (!shape.m.is_multiple_of(bm)).then(|| a_batch_base.clone() + shape.m), + (!shape.k.is_multiple_of(bk)).then(|| Tile::literal(TileLiteral::U32(shape.k))), + ]; + let b_bounds: [Option; 2] = [ + (!shape.k.is_multiple_of(bk)).then(|| b_batch_base.clone() + shape.k), + (!shape.n.is_multiple_of(bn)).then(|| Tile::literal(TileLiteral::U32(shape.n))), + ]; let subgroup_id = subgroup.subgroup_id(program); let sg_row = subgroup_id.clone() / col_groups; @@ -443,6 +466,8 @@ fn batched_coop_matmul_perf_single( &k_base, &sg_row_base, &sg_col_base_in_pass, + &a_bounds, + &b_bounds, accs, tile_rows_per_sg, tile_cols_per_sg, @@ -507,12 +532,16 @@ fn batched_coop_matmul_perf( let scalar = scalar_of(a.element()); - let tiles_m = shape.m / bm; - let tiles_n = shape.n / bn; + let tiles_m = shape.m.div_ceil(bm); + let tiles_n = shape.n.div_ceil(bn); let total_tiles = shape.batch * tiles_m * tiles_n; - let k_iterations = shape.k / bk; + let k_iterations = shape.k.div_ceil(bk); let k_pairs = k_iterations / 2; let k_remainder = k_iterations % 2; + // The y rows are padded to whole tiles per batch (the caller allocates + // the pad region); A/B are logical and edge tiles fill zero past the + // extents. + let m_padded = tiles_m * bm; // +1 inner padding on workgroup tiles avoids Apple shared-memory bank // conflicts on the inner stride (matches `stride_a = block_k + 1` in @@ -536,7 +565,15 @@ fn batched_coop_matmul_perf( let col_base = n_tile * bn; let a_batch_base = batch.clone() * shape.m; let b_batch_base = batch.clone() * shape.k; - let y_batch_base = batch * shape.m; + let y_batch_base = batch * m_padded; + let a_bounds: [Option; 2] = [ + (!shape.m.is_multiple_of(bm)).then(|| a_batch_base.clone() + shape.m), + (!shape.k.is_multiple_of(bk)).then(|| Tile::literal(TileLiteral::U32(shape.k))), + ]; + let b_bounds: [Option; 2] = [ + (!shape.k.is_multiple_of(bk)).then(|| b_batch_base.clone() + shape.k), + (!shape.n.is_multiple_of(bn)).then(|| Tile::literal(TileLiteral::U32(shape.n))), + ]; let subgroup_id = subgroup.subgroup_id(program); let sg_row = subgroup_id.clone() / col_groups; @@ -584,6 +621,8 @@ fn batched_coop_matmul_perf( &k_base_0, &sg_row_base, &sg_col_base_in_pass, + &a_bounds, + &b_bounds, accs, tile_rows_per_sg, tile_cols_per_sg, @@ -606,6 +645,8 @@ fn batched_coop_matmul_perf( &k_base_1, &sg_row_base, &sg_col_base_in_pass, + &a_bounds, + &b_bounds, accs, tile_rows_per_sg, tile_cols_per_sg, @@ -637,6 +678,8 @@ fn batched_coop_matmul_perf( &k_base_epi, &sg_row_base, &sg_col_base_in_pass, + &a_bounds, + &b_bounds, accs, tile_rows_per_sg, tile_cols_per_sg, diff --git a/fusor-ml/tile-ir/src/ir/layout.rs b/fusor-ml/tile-ir/src/ir/layout.rs index 207e5cf87..e6463d861 100644 --- a/fusor-ml/tile-ir/src/ir/layout.rs +++ b/fusor-ml/tile-ir/src/ir/layout.rs @@ -86,6 +86,23 @@ impl MultiFlattenMap { .collect() } + /// The length of contiguous (unit-stride) storage runs as this logical + /// axis's coordinate increments — the coalescing measure for choosing + /// which axis lanes should advance along. Walks the axis's sub-axes from + /// the innermost out, extending the run while each sub-axis continues + /// the previous one's span; 1 when the innermost sub-axis is not + /// unit-stride. + pub fn axis_unit_run(&self, axis: usize) -> u32 { + let mut run = 1u32; + for sub_axis in self.groups[axis].sub_axes.iter().rev() { + if sub_axis.stride != run { + break; + } + run = run.saturating_mul(sub_axis.extent); + } + run + } + /// True if this affine map is row-major contiguous for `shape`. pub fn is_row_major(&self, shape: &Shape) -> bool { self.is_affine() && self.affine_strides() == row_major_strides(shape) diff --git a/fusor-ml/tile-ir/src/ir/program.rs b/fusor-ml/tile-ir/src/ir/program.rs index 2dccd89ee..3cea67412 100644 --- a/fusor-ml/tile-ir/src/ir/program.rs +++ b/fusor-ml/tile-ir/src/ir/program.rs @@ -88,6 +88,11 @@ pub enum Stmt { dst: Tile, /// Per-element source value (typically a masked `Load`). value: Expr, + /// Exclusive global (row, col) limits: elements at or beyond a limit + /// load zero instead of touching storage. `None` per axis when the + /// tile is known in-bounds along it (the common aligned case, which + /// lowers without any guard). + bounds: [Option; 2], }, /// Cooperatively store an accumulator to a global storage view. A distinct /// subgroup-collective primitive — never lowered as a per-lane `Store`. diff --git a/fusor-ml/tile-ir/src/lower/analysis.rs b/fusor-ml/tile-ir/src/lower/analysis.rs index 02b83fe17..96359c75f 100644 --- a/fusor-ml/tile-ir/src/lower/analysis.rs +++ b/fusor-ml/tile-ir/src/lower/analysis.rs @@ -117,9 +117,12 @@ impl Analysis { self.visit_expr(index); self.visit_expr(value); } - Stmt::FillTile { dst, value } => { + Stmt::FillTile { dst, value, bounds } => { self.note_tile(dst); self.visit_expr(value); + for bound in bounds.iter().flatten() { + self.visit_expr(bound); + } } Stmt::CoopStore { acc, dst, addr } => { self.caps.uses_coop = true; diff --git a/fusor-ml/tile-ir/src/lower/coop.rs b/fusor-ml/tile-ir/src/lower/coop.rs index 26f62c445..44dcf4207 100644 --- a/fusor-ml/tile-ir/src/lower/coop.rs +++ b/fusor-ml/tile-ir/src/lower/coop.rs @@ -38,7 +38,9 @@ impl<'a> Lowerer<'a> { body.push(Statement::Store { pointer, value }, Span::default()); Ok(()) } - Stmt::FillTile { dst, value } => self.lower_fill_tile(expressions, body, dst, value), + Stmt::FillTile { dst, value, bounds } => { + self.lower_fill_tile(expressions, body, dst, value, bounds) + } Stmt::CoopStore { acc, dst, addr } => { self.lower_store_coop_acc(expressions, body, acc, dst, addr) } @@ -311,6 +313,7 @@ impl<'a> Lowerer<'a> { body: &mut Block, dst: &Tile, value: &Expr, + bounds: &[Option; 2], ) -> Result<(), LowerError> { let ExprKind::Load { src, @@ -329,10 +332,23 @@ impl<'a> Lowerer<'a> { )); }; match src { - Source::Storage(view) => { - self.lower_copy_to_tile(expressions, body, dst, view, row, col, mask, fill) - } + Source::Storage(view) => self.lower_copy_to_tile( + expressions, + body, + dst, + view, + row, + col, + bounds, + mask, + fill, + ), Source::Quantized(matrix) => { + if bounds.iter().any(Option::is_some) { + return Err(LowerError::UnsupportedOperation( + "bounded fills are not supported for quantized sources", + )); + } self.lower_copy_quant_to_tile(expressions, body, dst, matrix, row, col) } } @@ -347,6 +363,7 @@ impl<'a> Lowerer<'a> { src: &StorageView, row_offset: &Expr, col_offset: &Expr, + bounds: &[Option; 2], _mask: &Expr, _fill: &Expr, ) -> Result<(), LowerError> { @@ -356,9 +373,30 @@ impl<'a> Lowerer<'a> { let local = Self::function_arg(expressions, LOCAL_INVOCATION_INDEX_ARG); let row_base = self.lower_expr(expressions, body, row_offset)?; let col_base = self.lower_expr(expressions, body, col_offset)?; + let bounded = bounds.iter().any(Option::is_some); + + // The source layout drives the lane enumeration: lanes advance along + // the axis whose unit-stride storage runs they can follow, so a + // transposed matrix visits column-fastest down its unit-stride rows + // instead of scattering every lane across row strides. Divmod maps + // (im2col windows) keep the column order: their column-fastest + // blocks are compact in memory (a few cache lines per simdgroup), + // and Apple coalesces at line granularity — visiting along the long + // row runs measured slower by spreading each pass over the whole + // window footprint. + let cols_fastest = { + let indexing = src.layout.indexing(); + !indexing.is_affine() + || indexing.rank() != 2 + || indexing.axis_unit_run(1) >= indexing.axis_unit_run(0) + }; const VEC: u32 = 4; - if cols.is_multiple_of(VEC) && Self::storage_has_unit_inner_stride(&src.layout) { + if !bounded + && cols_fastest + && cols.is_multiple_of(VEC) + && Self::storage_axis_unit_stride(&src.layout, 1) + { let groups_per_row = cols / VEC; let total_groups = rows.checked_mul(groups_per_row) @@ -445,16 +483,120 @@ impl<'a> Lowerer<'a> { ); } + // Vectorized fast path down unit-stride rows (a transposed source): + // each lane loads four row-consecutive elements of one column. + if !bounded + && !cols_fastest + && rows.is_multiple_of(VEC) + && Self::storage_axis_unit_stride(&src.layout, 0) + { + let groups_per_col = rows / VEC; + let total_groups = + cols.checked_mul(groups_per_col) + .ok_or(LowerError::UnsupportedOperation( + "workgroup tile size overflow", + ))?; + return self.lower_copy_passes( + expressions, + body, + local, + total_groups, + |expressions, flat| { + let mut accept = Block::new(); + let local_col = self.div_literal_u32_emitted( + expressions, + flat, + groups_per_col, + &mut accept, + ); + let local_row_group = self.mod_literal_u32_emitted( + expressions, + flat, + groups_per_col, + &mut accept, + ); + let local_row_base = self.mul_literal_u32_emitted( + expressions, + local_row_group, + VEC, + &mut accept, + ); + let global_row_base = + self.add(expressions, &mut accept, row_base, local_row_base); + let global_col = self.add(expressions, &mut accept, col_base, local_col); + let storage_index_base = self.storage_index_from_coords( + expressions, + src, + &[global_row_base, global_col], + &mut accept, + )?; + let tile_index_base = self.tile_matrix_index_inline( + expressions, + &mut accept, + local_row_base, + local_col, + stride, + ); + let mut values = [None; VEC as usize]; + for i in 0..VEC { + let storage_index = self.add_literal_u32_emitted( + expressions, + storage_index_base, + i, + &mut accept, + ); + let storage_ptr = self.storage_dynamic_pointer( + expressions, + src, + storage_index, + &mut accept, + )?; + values[i as usize] = + Some(Self::emit_load(expressions, &mut accept, storage_ptr)); + } + for i in 0..VEC { + let tile_index = self.add_literal_u32_emitted( + expressions, + tile_index_base, + i * stride, + &mut accept, + ); + let tile_ptr = + self.tile_dynamic_pointer(expressions, dst, tile_index, &mut accept)?; + accept.push( + Statement::Store { + pointer: tile_ptr, + value: values[i as usize].expect("loaded above"), + }, + Span::default(), + ); + } + Ok(accept) + }, + ); + } + let total = rows .checked_mul(cols) .ok_or(LowerError::UnsupportedOperation( "workgroup tile size overflow", ))?; let lane_layout = CopyLaneLayout { + rows, cols, stride, row_base, col_base, + cols_fastest, + }; + + let row_limit = match &bounds[0] { + Some(bound) => Some(self.lower_expr(expressions, body, bound)?), + None => None, + }; + let col_limit = match &bounds[1] { + Some(bound) => Some(self.lower_expr(expressions, body, bound)?), + None => None, }; self.lower_copy_passes(expressions, body, local, total, |expressions, flat| { @@ -470,33 +612,111 @@ impl<'a> Lowerer<'a> { dst, lane_layout, )?; - let storage_index = self.storage_index_from_coords( - expressions, - src, - &[global_row, global_col], - &mut accept, - )?; - let storage_ptr = - self.storage_dynamic_pointer(expressions, src, storage_index, &mut accept)?; - let value = Self::emit_load(expressions, &mut accept, storage_ptr); - accept.push( - Statement::Store { - pointer: tile_ptr, - value, - }, - Span::default(), - ); + // Edge tiles: elements past a bound store zero without touching + // storage, so partial M/N/K tails stay garbage-free for the MMA. + let mut in_bounds: Option> = None; + for (coord, limit) in [(global_row, row_limit), (global_col, col_limit)] { + let Some(limit) = limit else { continue }; + let check = self.bin( + expressions, + &mut accept, + BinaryOperator::Less, + coord, + limit, + ); + in_bounds = Some(match in_bounds { + None => check, + Some(previous) => self.bin( + expressions, + &mut accept, + BinaryOperator::LogicalAnd, + previous, + check, + ), + }); + } + match in_bounds { + None => { + let storage_index = self.storage_index_from_coords( + expressions, + src, + &[global_row, global_col], + &mut accept, + )?; + let storage_ptr = self.storage_dynamic_pointer( + expressions, + src, + storage_index, + &mut accept, + )?; + let value = Self::emit_load(expressions, &mut accept, storage_ptr); + accept.push( + Statement::Store { + pointer: tile_ptr, + value, + }, + Span::default(), + ); + } + Some(in_bounds) => { + let mut in_bounds_body = Block::new(); + let storage_index = self.storage_index_from_coords( + expressions, + src, + &[global_row, global_col], + &mut in_bounds_body, + )?; + let storage_ptr = self.storage_dynamic_pointer( + expressions, + src, + storage_index, + &mut in_bounds_body, + )?; + let value = Self::emit_load(expressions, &mut in_bounds_body, storage_ptr); + in_bounds_body.push( + Statement::Store { + pointer: tile_ptr, + value, + }, + Span::default(), + ); + + let zero_f32 = self.f32(expressions, 0.0); + let mut out_of_bounds_body = Block::new(); + let zero = self.cast_tile_value( + expressions, + &mut out_of_bounds_body, + zero_f32, + ElementType::F32, + dst.element, + ); + out_of_bounds_body.push( + Statement::Store { + pointer: tile_ptr, + value: zero, + }, + Span::default(), + ); + accept.push( + Statement::If { + condition: in_bounds, + accept: in_bounds_body, + reject: out_of_bounds_body, + }, + Span::default(), + ); + } + } Ok(accept) }) } - /// True if the storage view's innermost axis is contiguous (stride 1). - fn storage_has_unit_inner_stride(layout: &Layout) -> bool { + /// True if the storage view is affine with a unit stride along `axis`. + fn storage_axis_unit_stride(layout: &Layout, axis: usize) -> bool { if !layout.is_affine() { return false; } - let strides = layout.affine_strides(); - strides.last().copied() == Some(1) + layout.affine_strides().get(axis).copied() == Some(1) } fn lower_copy_quant_to_tile( @@ -678,11 +898,15 @@ impl<'a> Lowerer<'a> { } let total = rows * cols; + // Quantized matrices store row-major blocks; lanes keep the column + // order. let lane_layout = CopyLaneLayout { + rows, cols, stride, row_base, col_base, + cols_fastest: true, }; self.lower_copy_passes(expressions, body, local, total, |expressions, flat| { let mut accept = Block::new(); @@ -881,13 +1105,25 @@ impl<'a> Lowerer<'a> { layout: CopyLaneLayout, ) -> Result { let CopyLaneLayout { + rows, cols, stride, row_base, col_base, + cols_fastest, } = layout; - let local_row = self.div_literal_u32_emitted(expressions, flat, cols, body); - let local_col = self.mod_literal_u32_emitted(expressions, flat, cols, body); + // Lanes advance along the axis the source layout coalesces best. + let (local_row, local_col) = if cols_fastest { + ( + self.div_literal_u32_emitted(expressions, flat, cols, body), + self.mod_literal_u32_emitted(expressions, flat, cols, body), + ) + } else { + ( + self.mod_literal_u32_emitted(expressions, flat, rows, body), + self.div_literal_u32_emitted(expressions, flat, rows, body), + ) + }; let global_row = self.add(expressions, body, row_base, local_row); let global_col = self.add(expressions, body, col_base, local_col); let tile_index = @@ -917,8 +1153,11 @@ struct CopyLaneCoords { #[derive(Clone, Copy)] struct CopyLaneLayout { + rows: u32, cols: u32, stride: u32, row_base: Handle, col_base: Handle, + /// Lane enumeration order, chosen from the source layout's unit runs. + cols_fastest: bool, } diff --git a/fusor-ml/tile-ir/src/tile/block.rs b/fusor-ml/tile-ir/src/tile/block.rs index 7b341669a..24a0f9d65 100644 --- a/fusor-ml/tile-ir/src/tile/block.rs +++ b/fusor-ml/tile-ir/src/tile/block.rs @@ -187,6 +187,22 @@ impl TileBlock<'_> { src: &Storage, row: impl Into, col: impl Into, + ) { + self.fill_tile_bounded(dst, src, row, col, [None, None]); + } + + /// [`Self::fill_tile`] with optional exclusive global (row, col) limits: + /// elements at or beyond a limit fill zero instead of reading storage. + /// Edge tiles of a matmul whose shape doesn't divide the tile geometry + /// pass the logical extents here; `None` per axis keeps the unguarded + /// fast path. + pub fn fill_tile_bounded( + &mut self, + dst: &WorkgroupTile, + src: &Storage, + row: impl Into, + col: impl Into, + bounds: [Option; 2], ) { let value = Expr::new( ExprKind::Load { @@ -203,6 +219,7 @@ impl TileBlock<'_> { self.push_stmt(Stmt::FillTile { dst: dst.decl().clone(), value, + bounds: bounds.map(|bound| bound.map(Tile::into_expr)), }); } @@ -233,6 +250,7 @@ impl TileBlock<'_> { self.push_stmt(Stmt::FillTile { dst: dst.decl().clone(), value, + bounds: [None, None], }); } From fc452728191d2d4cc6dbcaca400e5a0dfa91c35f Mon Sep 17 00:00:00 2001 From: Evan Almloff Date: Fri, 12 Jun 2026 06:12:53 -0500 Subject: [PATCH 06/16] adjust CoopTile selection for alignment --- fusor-ml/core/src/compute_graph/tests.rs | 3 +++ fusor-ml/core/src/matmul/mod.rs | 2 +- fusor-ml/core/src/matmul/variants.rs | 4 +++- 3 files changed, 7 insertions(+), 2 deletions(-) diff --git a/fusor-ml/core/src/compute_graph/tests.rs b/fusor-ml/core/src/compute_graph/tests.rs index 197152e46..9268c3ca8 100644 --- a/fusor-ml/core/src/compute_graph/tests.rs +++ b/fusor-ml/core/src/compute_graph/tests.rs @@ -461,6 +461,9 @@ fn conv_im2col_matmul_runs_without_gather() { for (b, c, h, w, n, kh, kw, implicit) in [ (2usize, 8usize, 16usize, 16usize, 16usize, 3usize, 3usize, false), (2, 64, 34, 34, 128, 3, 3, coop_viable), + // Coop-tile-unaligned M (1089 divides nothing): the masked-edge + // kernel with a padded output, still one dispatch. + (1, 128, 35, 35, 128, 3, 3, coop_viable), ] { let (oh, ow) = (h - kh + 1, w - kw + 1); let (m, k) = (b * oh * ow, c * kh * kw); diff --git a/fusor-ml/core/src/matmul/mod.rs b/fusor-ml/core/src/matmul/mod.rs index e6f888fef..8447ee6b9 100644 --- a/fusor-ml/core/src/matmul/mod.rs +++ b/fusor-ml/core/src/matmul/mod.rs @@ -452,7 +452,7 @@ mod selection_tests { // preference order breaks the tie toward the biggest tile). assert_eq!( select(1000, 1024, 1024, 512), - Some(CoopTile::new(128, 128, 16)) + Some(CoopTile::new(128, 64, 16)) ); // Tiny N inflates every candidate's padding past the waste bound: // gemv-shaped contractions stay on the generic path. diff --git a/fusor-ml/core/src/matmul/variants.rs b/fusor-ml/core/src/matmul/variants.rs index cb6c9d95f..cf91ec12d 100644 --- a/fusor-ml/core/src/matmul/variants.rs +++ b/fusor-ml/core/src/matmul/variants.rs @@ -235,8 +235,10 @@ impl CoopTile { // Selections whose padding inflates the output by more than a // quarter stay on the generic path — that bound also keeps // gemv-shaped contractions (tiny M or N) off the tile kernels. + // Candidates stick to geometries the aligned rules already reach + // (the (128, 128) table entry was never selectable and miscomputes). let mut best: Option<(u64, Self)> = None; - for (bm, bn) in [(128, 128), (128, 64), (64, 128), (64, 64)] { + for (bm, bn) in [(128, 64), (64, 128), (64, 64)] { let tile = Self::new(bm, bn, 16); if !tile.workgroup_size_supported(max_workgroup_size_x, max_subgroup_size) { continue; From 8688b3a3db1ed3190a346eedf8723ce146d77f15 Mon Sep 17 00:00:00 2001 From: Evan Almloff Date: Fri, 12 Jun 2026 20:02:09 -0500 Subject: [PATCH 07/16] fix clippy, fmt, and burn dx12 init on windows --- fusor-ml/conformance/src/bench/burn.rs | 15 +++++++--- .../core/examples/bench_conv_implicit_gemm.rs | 8 +++-- .../core/src/compute_graph/layout_pass.rs | 10 ++++--- fusor-ml/core/src/compute_graph/tests.rs | 4 ++- fusor-ml/core/src/matmul/kernel.rs | 26 +++++----------- fusor-ml/core/src/view.rs | 21 +++++++------ fusor-ml/tile-ir/src/lower/coop.rs | 30 ++++--------------- 7 files changed, 50 insertions(+), 64 deletions(-) diff --git a/fusor-ml/conformance/src/bench/burn.rs b/fusor-ml/conformance/src/bench/burn.rs index 0bdefb2d2..f2bd3e057 100644 --- a/fusor-ml/conformance/src/bench/burn.rs +++ b/fusor-ml/conformance/src/bench/burn.rs @@ -17,27 +17,34 @@ use super::time_samples; type BurnTensor = Tensor; -#[cfg(target_arch = "wasm32")] +#[cfg(any(target_arch = "wasm32", windows))] static BURN_WGPU_RUNTIME_INITIALIZED: core::sync::atomic::AtomicBool = core::sync::atomic::AtomicBool::new(false); -#[cfg(target_arch = "wasm32")] +#[cfg(any(target_arch = "wasm32", windows))] pub(crate) async fn initialized_device() -> WgpuDevice { use core::sync::atomic::Ordering; use burn::backend::wgpu::{RuntimeOptions, graphics, init_setup_async}; + #[cfg(target_arch = "wasm32")] + type Api = graphics::WebGpu; + // Cubecl's auto selection requests Vulkan on Windows, which is missing on + // GPU-less runners; DX12 always has at least the WARP software adapter. + #[cfg(windows)] + type Api = graphics::Dx12; + let device = WgpuDevice::default(); if BURN_WGPU_RUNTIME_INITIALIZED .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire) .is_ok() { - init_setup_async::(&device, RuntimeOptions::default()).await; + init_setup_async::(&device, RuntimeOptions::default()).await; } device } -#[cfg(not(target_arch = "wasm32"))] +#[cfg(not(any(target_arch = "wasm32", windows)))] pub(crate) async fn initialized_device() -> WgpuDevice { WgpuDevice::default() } diff --git a/fusor-ml/core/examples/bench_conv_implicit_gemm.rs b/fusor-ml/core/examples/bench_conv_implicit_gemm.rs index 70407c0d5..e53c0be47 100644 --- a/fusor-ml/core/examples/bench_conv_implicit_gemm.rs +++ b/fusor-ml/core/examples/bench_conv_implicit_gemm.rs @@ -179,8 +179,12 @@ fn bench_case(device: &Device, case: &ConvCase) { let (i_mean, i_p50, i_min) = stats(&implicit); println!("{name}: conv {b}x{c}x{h}x{w} k{kh}x{kw} -> matmul {m}x{k} @ {k}x{n}"); - println!(" gather+matmul ({gather_kernels} dispatches): mean {g_mean:.3} ms, p50 {g_p50:.3} ms, min {g_min:.3} ms"); - println!(" implicit-GEMM ({implicit_kernels} dispatches): mean {i_mean:.3} ms, p50 {i_p50:.3} ms, min {i_min:.3} ms"); + println!( + " gather+matmul ({gather_kernels} dispatches): mean {g_mean:.3} ms, p50 {g_p50:.3} ms, min {g_min:.3} ms" + ); + println!( + " implicit-GEMM ({implicit_kernels} dispatches): mean {i_mean:.3} ms, p50 {i_p50:.3} ms, min {i_min:.3} ms" + ); println!(" speedup (p50): {:.2}x", g_p50 / i_p50); println!(); } diff --git a/fusor-ml/core/src/compute_graph/layout_pass.rs b/fusor-ml/core/src/compute_graph/layout_pass.rs index 8ec503ffa..8b27fb2e2 100644 --- a/fusor-ml/core/src/compute_graph/layout_pass.rs +++ b/fusor-ml/core/src/compute_graph/layout_pass.rs @@ -70,10 +70,12 @@ impl LayoutPass { let new_layout = operation .is_fully_defined() .then(|| { - operation.stages.iter().try_fold( - input_layout.layout().clone(), - |composed, stage| crate::view::compose_layouts(&stage.layout, &composed), - ) + operation + .stages + .iter() + .try_fold(input_layout.layout().clone(), |composed, stage| { + crate::view::compose_layouts(&stage.layout, &composed) + }) }) .flatten() .unwrap_or_else(|| Layout::contiguous(operation.shape())); diff --git a/fusor-ml/core/src/compute_graph/tests.rs b/fusor-ml/core/src/compute_graph/tests.rs index 9268c3ca8..73bf7f53e 100644 --- a/fusor-ml/core/src/compute_graph/tests.rs +++ b/fusor-ml/core/src/compute_graph/tests.rs @@ -459,7 +459,9 @@ fn conv_im2col_matmul_runs_without_gather() { .is_some_and(|config| config.is_fixed()); for (b, c, h, w, n, kh, kw, implicit) in [ - (2usize, 8usize, 16usize, 16usize, 16usize, 3usize, 3usize, false), + ( + 2usize, 8usize, 16usize, 16usize, 16usize, 3usize, 3usize, false, + ), (2, 64, 34, 34, 128, 3, 3, coop_viable), // Coop-tile-unaligned M (1089 divides nothing): the masked-edge // kernel with a padded output, still one dispatch. diff --git a/fusor-ml/core/src/matmul/kernel.rs b/fusor-ml/core/src/matmul/kernel.rs index a9c87652f..8f1325d8c 100644 --- a/fusor-ml/core/src/matmul/kernel.rs +++ b/fusor-ml/core/src/matmul/kernel.rs @@ -241,11 +241,7 @@ impl MatMulOperation { /// Row-major strides of the logical output over its padded backing: /// rows step `n_padded`, each batch block spans `m_padded * n_padded`. - fn padded_out_strides( - out_shape: &[usize], - m_padded: usize, - n_padded: usize, - ) -> Box<[usize]> { + fn padded_out_strides(out_shape: &[usize], m_padded: usize, n_padded: usize) -> Box<[usize]> { let rank = out_shape.len(); let mut strides = vec![0usize; rank]; strides[rank - 1] = 1; @@ -330,11 +326,8 @@ impl MatMulOperation { // to the generic path, which writes through the logical layout. let m_padded = m.div_ceil(tile.bm) * tile.bm; let n_padded = n.div_ceil(tile.bn) * tile.bn; - let expected_strides = Self::padded_out_strides( - &self.out_shape, - m_padded as usize, - n_padded as usize, - ); + let expected_strides = + Self::padded_out_strides(&self.out_shape, m_padded as usize, n_padded as usize); let padded_elements = device_supported( (batch as usize) .checked_mul(m_padded as usize) @@ -528,7 +521,7 @@ impl Operation for MatMulOperation { // read the pad region). Shapes that already divide the tile — and // anything bound for the generic path — allocate exactly. let (m, n) = (self.a.rows(), self.b.cols()); - let padded = self.coop_tile(&device).and_then(|tile| { + let padded = self.coop_tile(device).and_then(|tile| { let m_padded = m.div_ceil(tile.bm as usize) * tile.bm as usize; let n_padded = n.div_ceil(tile.bn as usize) * tile.bn as usize; (m_padded != m || n_padded != n).then_some((m_padded, n_padded)) @@ -536,13 +529,10 @@ impl Operation for MatMulOperation { let output_tensor = match padded { Some((m_padded, n_padded)) => { let batch: usize = self.a.batch_shape().iter().product(); - let backing = TensorData::new_for_shape( - &device, - &[batch, m_padded, n_padded], - datatype, - ); + let backing = + TensorData::new_for_shape(device, &[batch, m_padded, n_padded], datatype); TensorData::new_from_parts( - &device, + device, backing.buffer().clone(), crate::Layout::from_parts( 0, @@ -552,7 +542,7 @@ impl Operation for MatMulOperation { datatype, ) } - None => TensorData::new_for_shape(&device, &self.out_shape, datatype), + None => TensorData::new_for_shape(device, &self.out_shape, datatype), }; vec![a.into(), b.into(), output_tensor.into()] } diff --git a/fusor-ml/core/src/view.rs b/fusor-ml/core/src/view.rs index 82c97bc52..03411b7aa 100644 --- a/fusor-ml/core/src/view.rs +++ b/fusor-ml/core/src/view.rs @@ -61,17 +61,16 @@ impl ViewStage { out: &[NaryExpr], clamp: bool, ) -> Option<(Vec, bool)> { - let (coords, delinearized) = - match affine_dim_indices(&self.layout, &self.input_shape) { - Some(affine) => ( - affine.iter().map(|index| index.to_expr(out)).collect(), - false, - ), - None => { - let flat = flat_index_expression(&self.layout, out)?; - (row_major_indices_from_flat(flat, &self.input_shape)?, true) - } - }; + let (coords, delinearized) = match affine_dim_indices(&self.layout, &self.input_shape) { + Some(affine) => ( + affine.iter().map(|index| index.to_expr(out)).collect(), + false, + ), + None => { + let flat = flat_index_expression(&self.layout, out)?; + (row_major_indices_from_flat(flat, &self.input_shape)?, true) + } + }; if !clamp { return Some((coords, delinearized)); } diff --git a/fusor-ml/tile-ir/src/lower/coop.rs b/fusor-ml/tile-ir/src/lower/coop.rs index 44dcf4207..b1414a13c 100644 --- a/fusor-ml/tile-ir/src/lower/coop.rs +++ b/fusor-ml/tile-ir/src/lower/coop.rs @@ -332,17 +332,9 @@ impl<'a> Lowerer<'a> { )); }; match src { - Source::Storage(view) => self.lower_copy_to_tile( - expressions, - body, - dst, - view, - row, - col, - bounds, - mask, - fill, - ), + Source::Storage(view) => { + self.lower_copy_to_tile(expressions, body, dst, view, row, col, bounds, mask, fill) + } Source::Quantized(matrix) => { if bounds.iter().any(Option::is_some) { return Err(LowerError::UnsupportedOperation( @@ -617,13 +609,7 @@ impl<'a> Lowerer<'a> { let mut in_bounds: Option> = None; for (coord, limit) in [(global_row, row_limit), (global_col, col_limit)] { let Some(limit) = limit else { continue }; - let check = self.bin( - expressions, - &mut accept, - BinaryOperator::Less, - coord, - limit, - ); + let check = self.bin(expressions, &mut accept, BinaryOperator::Less, coord, limit); in_bounds = Some(match in_bounds { None => check, Some(previous) => self.bin( @@ -643,12 +629,8 @@ impl<'a> Lowerer<'a> { &[global_row, global_col], &mut accept, )?; - let storage_ptr = self.storage_dynamic_pointer( - expressions, - src, - storage_index, - &mut accept, - )?; + let storage_ptr = + self.storage_dynamic_pointer(expressions, src, storage_index, &mut accept)?; let value = Self::emit_load(expressions, &mut accept, storage_ptr); accept.push( Statement::Store { From 8b2e0e424a3f9a12aa6803717476bd0c6e1caf4d Mon Sep 17 00:00:00 2001 From: Evan Almloff Date: Sat, 20 Jun 2026 15:34:55 -0500 Subject: [PATCH 08/16] add gemma support --- fusor-ml/core/src/quantized/matmul/mod.rs | 10 +- fusor-ml/core/src/quantized/mod.rs | 59 +- fusor-ml/core/src/row_program.rs | 21 +- fusor-ml/core/tests/attention.rs | 276 +++- fusor-ml/fusor/src/cache/mask_cache.rs | 7 +- .../fusor/src/composite/flash_attention.rs | 20 +- fusor-ml/fusor/src/layers/rms_norm.rs | 15 + fusor-ml/tile-ir-kernels/src/dispatch.rs | 42 +- fusor-ml/tile-ir-kernels/src/kernels/qgemv.rs | 47 +- .../tile-ir/src/lower/quantized/helpers.rs | 13 + .../tile-ir/src/lower/quantized/values.rs | 76 ++ .../src/lower/tile_program/quantized.rs | 5 +- fusor-ml/tile-ir/src/quantized.rs | 4 + interfaces/kalosm/examples/vision.rs | 2 +- models/kalosm-llama/examples/vision.rs | 144 +- models/kalosm-llama/src/chat_template.rs | 76 +- models/kalosm-llama/src/gguf_tokenizer.rs | 40 +- models/kalosm-llama/src/language_model.rs | 14 +- models/kalosm-llama/src/model/forward.rs | 82 +- models/kalosm-llama/src/model/inference.rs | 613 ++++++++- models/kalosm-llama/src/model/mod.rs | 166 ++- .../kalosm-llama/src/raw/attention_layer.rs | 250 +++- models/kalosm-llama/src/raw/cache.rs | 8 +- models/kalosm-llama/src/raw/mod.rs | 1187 ++++++++++++++--- models/kalosm-llama/src/raw/mtp.rs | 427 ++++++ models/kalosm-llama/src/raw/rope.rs | 12 +- models/kalosm-llama/src/raw/vision/gemma.rs | 744 +++++++++++ models/kalosm-llama/src/raw/vision/mod.rs | 172 +++ models/kalosm-llama/src/source.rs | 27 + models/kalosm-tokenizer/src/lib.rs | 200 ++- 30 files changed, 4417 insertions(+), 342 deletions(-) create mode 100644 models/kalosm-llama/src/raw/mtp.rs create mode 100644 models/kalosm-llama/src/raw/vision/gemma.rs diff --git a/fusor-ml/core/src/quantized/matmul/mod.rs b/fusor-ml/core/src/quantized/matmul/mod.rs index 96fd9bb95..44cd6340b 100644 --- a/fusor-ml/core/src/quantized/matmul/mod.rs +++ b/fusor-ml/core/src/quantized/matmul/mod.rs @@ -911,8 +911,9 @@ fn qgemv_cols_per_workgroup_for_direct(format: tile_ir::GgmlQuantFormat, k: u32, if format.is_q6k_family() && n <= 4096 && k > 4096 { return 4; // was Q6KLargeNarrow4 } - // Q8_0 wide. - if format.is_q8_0_family() && k <= 1024 && n >= 8192 { + // Q8_0 wide. Keep this aligned with qgemv_tile_with_epilogue, which + // switches Q8_0/Q8_0Native to 4x8 output columns for wide outputs. + if format.is_q8_0_family() && n >= 8192 { return 32; // was Q8WideAccelerated32 } // FormatAccelerated: Q5_0 mid (K,N both 2048..=4096), Q4K/Q6K general, @@ -929,7 +930,10 @@ fn qgemv_cols_per_workgroup_for_direct(format: tile_ir::GgmlQuantFormat, k: u32, if format.is_q5_0_family() && k <= 1024 && n <= 4096 { return 8; // was Q5Small8 } - 4 // was Default4 + // The shader-side qgemv shape table is the source of truth for generic + // formats. Returning a smaller value here over-dispatches masked columns, + // and every extra workgroup still pays the full K-loop cost. + tile_ir_kernels::qgemv_cols_per_workgroup_for_shape(format, k, n) } fn qmatmul_m_pad_target_for_caps(m: usize, n: usize, caps: KernelDeviceCaps) -> Option { diff --git a/fusor-ml/core/src/quantized/mod.rs b/fusor-ml/core/src/quantized/mod.rs index a63d3f20e..050005840 100644 --- a/fusor-ml/core/src/quantized/mod.rs +++ b/fusor-ml/core/src/quantized/mod.rs @@ -255,21 +255,26 @@ fn qmatrix_storage_layout_for_parts_with_env( // the raw-word ggml qgemv dot (`qgemv_q6k_ggml`) cannot address. The // f32-scale layout is 212 bytes (word-aligned) and feeds that amortized // decode; the +2 bytes/block (~0.9%) is paid back many times over by the - // faster kernel. Q4K/Q5K keep their native f16-scale layout (their native - // blocks are word-aligned, and native f16 scales read less memory). + // faster kernel. if ty == GgmlType::Q6K { return QMatrixStorageLayout::GpuF32Scales; } - if matches!( - ty, - GgmlType::Q4_0 | GgmlType::Q5_0 | GgmlType::Q8_0 | GgmlType::Q4K | GgmlType::Q5K - ) && native_half_scale_storage_enabled(shader_f16_supported, env_override) + // The small 32-element block formats save only two bytes/block with native + // f16 scales, but native Q4_0/Q5_0/Q8_0 blocks are byte-addressed. Decode is + // much faster with word-aligned f32-scale storage on Apple GPUs. + if matches!(ty, GgmlType::Q4_0 | GgmlType::Q5_0 | GgmlType::Q8_0) { + if env_override.is_some() + && native_half_scale_storage_enabled(shader_f16_supported, env_override) + { + QMatrixStorageLayout::Native + } else { + QMatrixStorageLayout::GpuF32Scales + } + } else if matches!(ty, GgmlType::Q4K | GgmlType::Q5K) + && native_half_scale_storage_enabled(shader_f16_supported, env_override) { QMatrixStorageLayout::Native - } else if matches!( - ty, - GgmlType::Q4_0 | GgmlType::Q5_0 | GgmlType::Q8_0 | GgmlType::Q4K | GgmlType::Q5K - ) { + } else if matches!(ty, GgmlType::Q4K | GgmlType::Q5K) { QMatrixStorageLayout::GpuF32Scales } else { QMatrixStorageLayout::Native @@ -536,14 +541,22 @@ mod tests { use super::*; #[test] - fn native_capable_quant_storage_defaults_to_native_when_shader_f16_is_supported() { - for ty in [ - GgmlType::Q4_0, - GgmlType::Q5_0, - GgmlType::Q8_0, - GgmlType::Q4K, - GgmlType::Q5K, - ] { + fn small_block_quant_storage_defaults_to_f32_scales() { + for ty in [GgmlType::Q4_0, GgmlType::Q5_0, GgmlType::Q8_0] { + assert_eq!( + qmatrix_storage_layout_for_parts_with_env(ty, true, None), + QMatrixStorageLayout::GpuF32Scales + ); + assert_eq!( + qmatrix_storage_layout_for_parts_with_env(ty, false, None), + QMatrixStorageLayout::GpuF32Scales + ); + } + } + + #[test] + fn k_quant_storage_defaults_to_native_when_shader_f16_is_supported() { + for ty in [GgmlType::Q4K, GgmlType::Q5K] { assert_eq!( qmatrix_storage_layout_for_parts_with_env(ty, true, None), QMatrixStorageLayout::Native @@ -564,7 +577,15 @@ mod tests { } #[test] - fn q4k_storage_env_can_force_native_or_f32_expanded() { + fn quant_storage_env_can_force_native_or_f32_expanded() { + assert_eq!( + qmatrix_storage_layout_for_parts_with_env(GgmlType::Q4_0, true, Some("1")), + QMatrixStorageLayout::Native + ); + assert_eq!( + qmatrix_storage_layout_for_parts_with_env(GgmlType::Q4_0, true, Some("0")), + QMatrixStorageLayout::GpuF32Scales + ); assert_eq!( qmatrix_storage_layout_for_parts_with_env(GgmlType::Q4K, false, Some("1")), QMatrixStorageLayout::Native diff --git a/fusor-ml/core/src/row_program.rs b/fusor-ml/core/src/row_program.rs index 445e3ac49..713eb4715 100644 --- a/fusor-ml/core/src/row_program.rs +++ b/fusor-ml/core/src/row_program.rs @@ -46,11 +46,6 @@ use crate::{ const BLOCK: u32 = 256; -/// Below this many rows a long axis is split across workgroups (one tile -/// each) with a combine kernel folding the spans — decode has too few rows -/// to fill the device with one workgroup per row. -const SPLIT_ROWS_TARGET: u32 = 256; - /// Workgroup buckets for dynamic-axis row programs, smallest first. The /// kernel monomorphizes per bucket; the active axis length rides in the /// params input. @@ -606,10 +601,12 @@ fn build_row_program_kernel( let block = workgroup_shape.x(); let lanes_own_axis = operation.dynamic_axis.is_some(); - // Long axes with few rows fan out across workgroups: each split runs - // the online body over one tile, writing its unnormalized accumulator - // and softmax statistics to scratch; a combine kernel folds the spans - // with the online monoid. + // Long reducing outputs fan out across workgroups: each split owns one + // axis tile, writes its unnormalized accumulator and softmax statistics + // to scratch, then a combine kernel folds the spans with the online + // monoid. This avoids the single-workgroup multi-tile loop for attention + // prefill, which is both slower to fill the GPU and less robust at Gemma's + // unscaled vision-attention logits. let output_kind = operation.output_step().clone(); let phase_steps = operation.phase_steps().to_vec(); let free_dim_out = match &output_kind { @@ -619,11 +616,7 @@ fn build_row_program_kernel( let tiles = k.div_ceil(block); let splits: u32 = match free_dim_out { Some(free) - if lanes_own_axis - && rows < SPLIT_ROWS_TARGET - && tiles > 1 - && tiles <= block - && (free as u32 + 2) <= block => + if lanes_own_axis && tiles > 1 && tiles <= block && (free as u32 + 2) <= block => { tiles } diff --git a/fusor-ml/core/tests/attention.rs b/fusor-ml/core/tests/attention.rs index 2aee33924..fd114d2a3 100644 --- a/fusor-ml/core/tests/attention.rs +++ b/fusor-ml/core/tests/attention.rs @@ -9,6 +9,15 @@ fn values(len: usize, scale: f32) -> Vec { (0..len).map(|i| ((i as f32) * scale).sin()).collect() } +fn high_variance_values(len: usize, scale: f32) -> Vec { + (0..len) + .map(|i| { + let x = i as f32; + ((x * 0.173).sin() * 1.7 + (x * 0.071).cos() * 0.9) * scale + }) + .collect() +} + struct AttentionCase { batch: usize, heads: usize, @@ -81,6 +90,24 @@ fn cpu_attention( } fn check_attention(case: AttentionCase, tolerance: f32) { + check_attention_with_scale(case, tolerance, None); +} + +fn check_attention_with_scale(case: AttentionCase, tolerance: f32, scale_override: Option) { + check_attention_impl(case, tolerance, scale_override, false, false); +} + +fn check_attention_full_high_variance(case: AttentionCase, tolerance: f32, scale_override: f32) { + check_attention_impl(case, tolerance, Some(scale_override), true, true); +} + +fn check_attention_impl( + case: AttentionCase, + tolerance: f32, + scale_override: Option, + full_compare: bool, + high_variance: bool, +) { pollster::block_on(async { let Ok(device) = Device::new().await else { return; @@ -95,15 +122,27 @@ fn check_attention(case: AttentionCase, tolerance: f32) { causal, masked, } = case; - let q_data = values(batch * heads * q_len * head_dim, 0.13); - let k_data = values(batch * kv_heads * kv_len * head_dim, 0.07); - let v_data = values(batch * kv_heads * kv_len * head_dim, 0.11); + let q_data = if high_variance { + high_variance_values(batch * heads * q_len * head_dim, 1.0) + } else { + values(batch * heads * q_len * head_dim, 0.13) + }; + let k_data = if high_variance { + high_variance_values(batch * kv_heads * kv_len * head_dim, 0.8) + } else { + values(batch * kv_heads * kv_len * head_dim, 0.07) + }; + let v_data = if high_variance { + high_variance_values(batch * kv_heads * kv_len * head_dim, 1.1) + } else { + values(batch * kv_heads * kv_len * head_dim, 0.11) + }; let mask_data = masked.then(|| { (0..q_len * kv_len) .map(|i| if i % 7 == 0 { -1.5 } else { 0.25 }) .collect::>() }); - let scale = 1.0 / (head_dim as f32).sqrt(); + let scale = scale_override.unwrap_or_else(|| 1.0 / (head_dim as f32).sqrt()); let q = Tensor::from_slice(&device, [batch, heads, q_len, head_dim], &q_data); let k = Tensor::from_slice(&device, [batch, kv_heads, kv_len, head_dim], &k_data); @@ -132,16 +171,44 @@ fn check_attention(case: AttentionCase, tolerance: f32) { mask_data.as_deref(), scale, ); - for b in 0..batch { - for h in [0, heads - 1] { - for qi in [0, q_len / 2, q_len - 1] { - for d in [0, head_dim / 2, head_dim - 1] { - let want = expected[((b * heads + h) * q_len + qi) * head_dim + d]; - let got = actual[[b, h, qi, d]]; - assert!( - (got - want).abs() < tolerance, - "b={b} h={h} q={qi} d={d}: got {got}, expected {want}" - ); + if full_compare { + let mut worst = 0.0f32; + let mut worst_index = (0, 0, 0, 0); + for b in 0..batch { + for h in 0..heads { + for qi in 0..q_len { + for d in 0..head_dim { + let want = expected[((b * heads + h) * q_len + qi) * head_dim + d]; + let got = actual[[b, h, qi, d]]; + let err = (got - want).abs(); + if err > worst { + worst = err; + worst_index = (b, h, qi, d); + } + } + } + } + } + assert!( + worst < tolerance, + "worst attention diff {worst} at {worst_index:?}: got {}, expected {}, tolerance {tolerance}", + actual[[worst_index.0, worst_index.1, worst_index.2, worst_index.3]], + expected[((worst_index.0 * heads + worst_index.1) * q_len + worst_index.2) + * head_dim + + worst_index.3] + ); + } else { + for b in 0..batch { + for h in [0, heads - 1] { + for qi in [0, q_len / 2, q_len - 1] { + for d in [0, head_dim / 2, head_dim - 1] { + let want = expected[((b * heads + h) * q_len + qi) * head_dim + d]; + let got = actual[[b, h, qi, d]]; + assert!( + (got - want).abs() < tolerance, + "b={b} h={h} q={qi} d={d}: got {got}, expected {want}" + ); + } } } } @@ -221,6 +288,80 @@ fn attention_causal_prefill_streams_long_kv() { ); } +#[test] +fn attention_unmasked_prefill_streams_long_kv() { + // Gemma 4 vision uses full, non-causal image self-attention over thousands + // of patch tokens. This exercises the streaming row-program path without + // the causal axis bound. + check_attention( + AttentionCase { + batch: 1, + heads: 4, + kv_heads: 4, + q_len: 1024, + kv_len: 1024, + head_dim: 64, + causal: false, + masked: false, + }, + 1e-4, + ); +} + +#[test] +fn attention_unmasked_prefill_streams_long_kv_scale_one() { + check_attention_full_high_variance( + AttentionCase { + batch: 1, + heads: 4, + kv_heads: 4, + q_len: 1024, + kv_len: 1024, + head_dim: 64, + causal: false, + masked: false, + }, + 3e-4, + 1.0, + ); +} + +#[test] +fn attention_unmasked_prefill_many_tiles_scale_one() { + check_attention_full_high_variance( + AttentionCase { + batch: 1, + heads: 2, + kv_heads: 2, + q_len: 384, + kv_len: 2304, + head_dim: 64, + causal: false, + masked: false, + }, + 3e-4, + 1.0, + ); +} + +#[test] +fn attention_masked_prefill_streams_long_kv_scale_one() { + check_attention_full_high_variance( + AttentionCase { + batch: 1, + heads: 4, + kv_heads: 4, + q_len: 512, + kv_len: 768, + head_dim: 64, + causal: false, + masked: true, + }, + 3e-4, + 1.0, + ); +} + #[test] fn attention_odd_sized_causal_prefill() { // Odd extents that don't align with any tile or bucket boundary. @@ -256,6 +397,113 @@ fn attention_masked_prefill() { ); } +#[test] +fn attention_offset_causal_mask_prefill() { + check_offset_causal_mask_prefill(20, 280, 4, 4); +} + +#[test] +fn attention_offset_causal_mask_prefill_single_tile() { + check_offset_causal_mask_prefill(20, 256, 4, 4); +} + +#[test] +fn attention_offset_causal_mask_prefill_gqa() { + check_offset_causal_mask_prefill(20, 256, 8, 4); +} + +fn check_offset_causal_mask_prefill(q_len: usize, kv_len: usize, heads: usize, kv_heads: usize) { + pollster::block_on(async { + let Ok(device) = Device::new().await else { + return; + }; + let case = AttentionCase { + batch: 1, + heads, + kv_heads, + q_len, + kv_len, + head_dim: 64, + causal: false, + masked: true, + }; + let q_data = + high_variance_values(case.batch * case.heads * case.q_len * case.head_dim, 1.0); + let k_data = high_variance_values( + case.batch * case.kv_heads * case.kv_len * case.head_dim, + 0.8, + ); + let v_data = high_variance_values( + case.batch * case.kv_heads * case.kv_len * case.head_dim, + 1.1, + ); + let prefix = case.kv_len - case.q_len; + let mask_data = (0..case.q_len * case.kv_len) + .map(|i| { + let q = i / case.kv_len; + let kv = i % case.kv_len; + if kv <= prefix + q { + 0.0 + } else { + f32::NEG_INFINITY + } + }) + .collect::>(); + let q = Tensor::from_slice( + &device, + [case.batch, case.heads, case.q_len, case.head_dim], + &q_data, + ); + let k = Tensor::from_slice( + &device, + [case.batch, case.kv_heads, case.kv_len, case.head_dim], + &k_data, + ); + let v = Tensor::from_slice( + &device, + [case.batch, case.kv_heads, case.kv_len, case.head_dim], + &v_data, + ); + let mask = Tensor::from_slice(&device, [case.q_len, case.kv_len], &mask_data); + let scale = 1.0; + + let out = q.flash_attention(&k, &v, scale, Some(&mask)); + assert_eq!( + out.count_kernels_to_resolve(), + 1, + "offset-causal mask attention should lower as one row-program kernel" + ); + let actual = out.as_slice::<4, f32>().await.unwrap(); + let expected = cpu_attention(&case, &q_data, &k_data, &v_data, Some(&mask_data), scale); + let mut worst = 0.0f32; + let mut worst_index = (0, 0, 0, 0); + for b in 0..case.batch { + for h in 0..case.heads { + for qi in 0..case.q_len { + for d in 0..case.head_dim { + let want = + expected[((b * case.heads + h) * case.q_len + qi) * case.head_dim + d]; + let got = actual[[b, h, qi, d]]; + let err = (got - want).abs(); + if err > worst { + worst = err; + worst_index = (b, h, qi, d); + } + } + } + } + } + assert!( + worst < 3e-4, + "worst offset-causal attention diff {worst} at {worst_index:?}: got {}, expected {}", + actual[[worst_index.0, worst_index.1, worst_index.2, worst_index.3]], + expected[((worst_index.0 * case.heads + worst_index.1) * case.q_len + worst_index.2) + * case.head_dim + + worst_index.3] + ); + }); +} + #[test] fn attention_f16_io() { pollster::block_on(async { diff --git a/fusor-ml/fusor/src/cache/mask_cache.rs b/fusor-ml/fusor/src/cache/mask_cache.rs index f05db019f..4b1707061 100644 --- a/fusor-ml/fusor/src/cache/mask_cache.rs +++ b/fusor-ml/fusor/src/cache/mask_cache.rs @@ -60,7 +60,12 @@ where } else { // Create the mask based on whether we have a sliding window let mask = if let Some(sliding_window_size) = sliding_window_size { - Self::create_sliding_window_mask(device, seq_len, sliding_window_size) + let mask = Self::create_sliding_window_mask(device, seq_len, sliding_window_size); + if seq_len <= sliding_window_size { + mask.mark_strict_causal() + } else { + mask + } } else { AttentionMask::causal(device, seq_len) }; diff --git a/fusor-ml/fusor/src/composite/flash_attention.rs b/fusor-ml/fusor/src/composite/flash_attention.rs index 32568638c..00aaba64f 100644 --- a/fusor-ml/fusor/src/composite/flash_attention.rs +++ b/fusor-ml/fusor/src/composite/flash_attention.rs @@ -68,12 +68,20 @@ where (Tensor::Gpu(q), Tensor::Gpu(k_gpu), Tensor::Gpu(v_gpu)) if !matches!(mask, Some((_, MaskKind::BatchKeyMask))) => { - // Decode (q_seq_len == 1) runs the DecodeSmall attention kernel, - // which uses workgroup reductions and needs no subgroups, so it - // works on browser adapters that report no subgroup support. Only - // the prefill/streaming kernels require subgroups — keep the - // fallback for those (q_seq_len > 1). - if !q.device().subgroups_supported() && self.shape()[2] != 1 { + let q_seq_len = self.shape()[2]; + let use_cpu_qk_mask_prefill = q_seq_len != 1 + && matches!(mask, Some((_, MaskKind::QKMask))) + && std::env::var_os("FUSOR_ALLOW_GPU_QK_MASK_PREFILL").is_none(); + let force_cpu_prefill = q_seq_len != 1 + && (std::env::var_os("FUSOR_FORCE_CPU_PREFILL_ATTENTION").is_some() + || (mask.is_none() + && std::env::var_os("FUSOR_FORCE_CPU_UNMASKED_PREFILL_ATTENTION") + .is_some()) + || (matches!(mask, Some((_, MaskKind::QKMask))) + && std::env::var_os("FUSOR_FORCE_CPU_QK_MASK_PREFILL_ATTENTION") + .is_some()) + || use_cpu_qk_mask_prefill); + if force_cpu_prefill { #[cfg(target_arch = "wasm32")] { return self.flash_attention_composite_impl(k, v, scale, mask); diff --git a/fusor-ml/fusor/src/layers/rms_norm.rs b/fusor-ml/fusor/src/layers/rms_norm.rs index e902f77ed..2f5bc29ea 100644 --- a/fusor-ml/fusor/src/layers/rms_norm.rs +++ b/fusor-ml/fusor/src/layers/rms_norm.rs @@ -98,6 +98,21 @@ where T: CastTo + CastTensor, f32: CastTo + CastTensor, { + /// Forward pass for 2D input with generic type. + /// Converts input to f32 for computation, then converts back. + pub fn forward_generic_2d(&self, input: &Tensor<2, T, B>) -> Tensor<2, T> + where + B: Fusion<2, T>, + { + let input_f32 = input.cast::(); + let weight_f32: Tensor<1, f32> = self.weight.cast(); + let bias_f32: Option> = self.bias.as_ref().map(|b| b.cast()); + + let result_f32 = input_f32.rms_norm_fused::<1, 1>(&weight_f32, bias_f32.as_ref(), self.eps); + + result_f32.cast() + } + /// Forward pass for 3D input with generic type. /// Converts input to f32 for computation, then converts back. pub fn forward_generic(&self, input: &Tensor<3, T, B>) -> Tensor<3, T> diff --git a/fusor-ml/tile-ir-kernels/src/dispatch.rs b/fusor-ml/tile-ir-kernels/src/dispatch.rs index 03e4d6a5e..63d167152 100644 --- a/fusor-ml/tile-ir-kernels/src/dispatch.rs +++ b/fusor-ml/tile-ir-kernels/src/dispatch.rs @@ -42,6 +42,10 @@ pub const fn qgemv_cols_per_workgroup(format: GgmlQuantFormat) -> u32 { /// This includes the Q4K/Q6K GGML specializations whose column grouping /// depends on both K (`rows`) and N (`cols`). pub fn qgemv_cols_per_workgroup_for_shape(format: GgmlQuantFormat, rows: u32, cols: u32) -> u32 { + if matches!(format, GgmlQuantFormat::Q4_0 | GgmlQuantFormat::Q4_0Native) { + return q4_0_override(q4_0_default(rows, cols)).cols_per_workgroup(); + } + if is_q4k_family(format) && rows <= 4096 && (4096..8192).contains(&cols) { return q4k_mid_override(q4k_default_mid(rows, cols)).cols_per_workgroup(); } @@ -95,12 +99,15 @@ pub(crate) const fn qgemv_subgroups_per_workgroup(format: GgmlQuantFormat) -> u3 } /// Shape-aware subgroup count used by the qgemv dispatch policy. -pub const fn qgemv_subgroups_per_workgroup_for_shape( +pub fn qgemv_subgroups_per_workgroup_for_shape( format: GgmlQuantFormat, rows: u32, - _cols: u32, + cols: u32, ) -> u32 { match format { + GgmlQuantFormat::Q4_0 | GgmlQuantFormat::Q4_0Native => { + q4_0_override(q4_0_default(rows, cols)).subgroups + } format if format.is_q6k_family() && rows > 4096 => 8, _ => qgemv_subgroups_per_workgroup(format), } @@ -235,6 +242,20 @@ const STANDARD_8_TILES: &[(&str, QgemvShape)] = &[ ("ggml_8x4", QgemvShape::new(8, 4)), ]; +const Q4_0_TILES: &[(&str, QgemvShape)] = &[ + ("ggml_1x4", QgemvShape::new(1, 4)), + ("ggml_1x8", QgemvShape::new(1, 8)), + ("ggml_2x2", QgemvShape::new(2, 2)), + ("ggml_2x4", QgemvShape::new(2, 4)), + ("ggml_2x8", QgemvShape::new(2, 8)), + ("ggml_4x2", QgemvShape::new(4, 2)), + ("ggml_4x4", QgemvShape::new(4, 4)), + ("ggml_4x8", QgemvShape::new(4, 8)), + ("ggml_8x1", QgemvShape::new(8, 1)), + ("ggml_8x2", QgemvShape::new(8, 2)), + ("ggml_8x4", QgemvShape::new(8, 4)), +]; + fn env_tile_override(var: &str, table: &[(&str, QgemvShape)], default: QgemvShape) -> QgemvShape { let Ok(value) = std::env::var(var) else { return default; @@ -250,6 +271,23 @@ pub(crate) fn q4k_mid_override(default: QgemvShape) -> QgemvShape { env_tile_override("FUSOR_Q4K_MID_TILE", Q4K_MID_TILES, default) } +pub(crate) const fn q4_0_default(_rows: u32, _cols: u32) -> QgemvShape { + QgemvShape::new(4, 4) +} + +pub(crate) fn q4_0_override(default: QgemvShape) -> QgemvShape { + env_tile_override("FUSOR_Q4_0_TILE", Q4_0_TILES, default) +} + +pub(crate) fn q4_0_values_per_lane(default: u32) -> u32 { + match std::env::var("FUSOR_Q4_0_VALUES_PER_LANE").ok().as_deref() { + Some("8") => 8, + Some("16") => 16, + Some("32") => 32, + _ => default, + } +} + // ----- Q4K large (rows<=4096, cols>=8192) ----- /// Default Q4K large-shape: cols<=16_384 → 8x4, else 2x4. diff --git a/fusor-ml/tile-ir-kernels/src/kernels/qgemv.rs b/fusor-ml/tile-ir-kernels/src/kernels/qgemv.rs index 4cbb9f24d..6b69d9f59 100644 --- a/fusor-ml/tile-ir-kernels/src/kernels/qgemv.rs +++ b/fusor-ml/tile-ir-kernels/src/kernels/qgemv.rs @@ -19,8 +19,9 @@ use fusor_tile_ir::tile::{range, Mask, Program, Storage, Tile, TileBlock}; use fusor_tile_ir::{GgmlQuantFormat, QuantizedMatrix, TileLiteral}; use crate::dispatch::{ - q4k_default_large, q4k_default_mid, q4k_default_tall, q4k_large_override, q4k_mid_override, - q4k_tall_override, q6k_default_large, q6k_default_tall, q6k_large_override, q6k_tall_override, + q4_0_default, q4_0_override, q4_0_values_per_lane, q4k_default_large, q4k_default_mid, + q4k_default_tall, q4k_large_override, q4k_mid_override, q4k_tall_override, q6k_default_large, + q6k_default_tall, q6k_large_override, q6k_tall_override, qgemv_subgroups_per_workgroup_for_shape, QgemvShape, SubgroupConfig, }; use crate::grid::{ @@ -216,19 +217,30 @@ pub(crate) fn qgemv_tile_with_epilogue( qgemv_shape(2, 4), 16, ), - GgmlQuantFormat::Q4_0 - | GgmlQuantFormat::Q4_0Native - | GgmlQuantFormat::Q4_1 - | GgmlQuantFormat::Q5_1 - | GgmlQuantFormat::Q2K => qgemv_perf_with_epilogue( - program, - tensors, - workgroups_x, - subgroups, - ep, - qgemv_shape(2, 4), - 8, - ), + GgmlQuantFormat::Q4_0 | GgmlQuantFormat::Q4_0Native => { + let shape = q4_0_override(q4_0_default(b.rows, output_cols)); + let values_per_lane = q4_0_values_per_lane(8); + qgemv_perf_with_epilogue( + program, + tensors, + workgroups_x, + subgroups, + ep, + shape, + values_per_lane, + ) + } + GgmlQuantFormat::Q4_1 | GgmlQuantFormat::Q5_1 | GgmlQuantFormat::Q2K => { + qgemv_perf_with_epilogue( + program, + tensors, + workgroups_x, + subgroups, + ep, + qgemv_shape(2, 4), + 8, + ) + } GgmlQuantFormat::Q3K | GgmlQuantFormat::Q8K => qgemv_perf_with_epilogue( program, tensors, @@ -336,6 +348,11 @@ fn select_qgemv_dot( if format.is_q8_0_family() && values_per_lane == 8 { return QgemvDot::F32Vec; } + if format.is_q4_0_family() + && (values_per_lane == 8 || values_per_lane == 16 || values_per_lane == 32) + { + return QgemvDot::F32Vec; + } if format.is_q4k_family() && (values_per_lane == 8 || values_per_lane == 16 || values_per_lane == 32) { diff --git a/fusor-ml/tile-ir/src/lower/quantized/helpers.rs b/fusor-ml/tile-ir/src/lower/quantized/helpers.rs index a5a708be4..98154c2aa 100644 --- a/fusor-ml/tile-ir/src/lower/quantized/helpers.rs +++ b/fusor-ml/tile-ir/src/lower/quantized/helpers.rs @@ -139,6 +139,19 @@ impl<'a> Lowerer<'a> { } } + pub(in crate::lower) fn q4_0_data_byte_offset( + &self, + format: GgmlQuantFormat, + ) -> Result { + match format { + GgmlQuantFormat::Q4_0 => Ok(4), + GgmlQuantFormat::Q4_0Native => Ok(2), + _ => Err(LowerError::UnsupportedOperation( + "q4_0 data offset requires a Q4_0 format", + )), + } + } + pub(in crate::lower) fn q4k_data_word_offset( &self, format: GgmlQuantFormat, diff --git a/fusor-ml/tile-ir/src/lower/quantized/values.rs b/fusor-ml/tile-ir/src/lower/quantized/values.rs index 8df0c445f..1cdb6f259 100644 --- a/fusor-ml/tile-ir/src/lower/quantized/values.rs +++ b/fusor-ml/tile-ir/src/lower/quantized/values.rs @@ -103,6 +103,82 @@ impl<'a> Lowerer<'a> { Ok(self.mul(expressions, body, sum, parts.scale)) } + pub(in crate::lower) fn q4_0_f32_dot( + &self, + expressions: &mut Arena, + matrix: &QuantizedMatrix, + k_base: Handle, + col: Handle, + a: &[Handle], + body: &mut Block, + ) -> Result, LowerError> { + if !matrix.format.is_q4_0_family() || !matches!(a.len(), 8 | 16 | 32) { + return Err(LowerError::UnsupportedOperation( + "q4_0 f32 dot requires a Q4_0 format and 8, 16, or 32 activation values", + )); + } + + let (base, q_base) = self.quantized_flat_block_base_and_q( + expressions, + matrix, + k_base, + col, + matrix.format.block_words(), + body, + ); + let scale = self.load_affine_scale_f32(expressions, matrix, base, 0, body)?; + let data_offset = self.q4_0_data_byte_offset(matrix.format)?; + let q_local = self.and_lit(expressions, body, q_base, 15); + let high = self.cmp_lit(expressions, body, BinaryOperator::GreaterEqual, q_base, 16); + let word_count = if a.len() == 32 { 4 } else { a.len() / 4 }; + let mut words = Vec::with_capacity(word_count); + for word_index in 0..word_count { + let byte_offset = self.add_lit( + expressions, + body, + q_local, + data_offset + (word_index * 4) as u32, + ); + words.push(self.load_word_at_block_dynamic_byte_offset( + expressions, + matrix, + base, + byte_offset, + body, + )?); + } + + let mut quants = Vec::with_capacity(a.len()); + for lane in 0..a.len() { + let byte_lane = self.u32(expressions, (lane % 4) as u32); + let word = if a.len() == 32 { + words[(lane % 16) / 4] + } else { + words[lane / 4] + }; + let byte = self.byte_at(expressions, body, word, byte_lane); + let low = self.and_lit(expressions, body, byte, 0x0f); + let high4 = self.shr_lit(expressions, body, byte, 4); + let quant = if a.len() == 32 { + if lane >= 16 { + high4 + } else { + low + } + } else { + self.select(expressions, body, high, high4, low) + }; + quants.push(quant); + } + + let weighted_sum = self.dot_quant_vec4_chunks(expressions, body, a, &quants); + let activation_sum = self.sum_values(expressions, body, a); + let center = self.f32(expressions, 8.0); + let center_term = self.mul(expressions, body, activation_sum, center); + let centered = self.sub(expressions, body, weighted_sum, center_term); + Ok(self.mul(expressions, body, centered, scale)) + } + pub(in crate::lower) fn q8_0_block_parts8( &self, expressions: &mut Arena, diff --git a/fusor-ml/tile-ir/src/lower/tile_program/quantized.rs b/fusor-ml/tile-ir/src/lower/tile_program/quantized.rs index ef7b9bf66..8c432f427 100644 --- a/fusor-ml/tile-ir/src/lower/tile_program/quantized.rs +++ b/fusor-ml/tile-ir/src/lower/tile_program/quantized.rs @@ -135,6 +135,9 @@ impl<'a> Lowerer<'a> { let c = self.lower_expr_lane(expressions, block, col, spill_depth)?; match packing { QuantActivation::F32 => match (src.format, block_n) { + (GgmlQuantFormat::Q4_0 | GgmlQuantFormat::Q4_0Native, 8 | 16 | 32) => { + self.q4_0_f32_dot(expressions, src, k, c, &a_handles, block) + } (GgmlQuantFormat::Q8_0 | GgmlQuantFormat::Q8_0Native, 8) => { let a8 = Self::expect_dot8(&a_handles)?; self.dequantize_q8_0_dot8(expressions, src, k, c, &a8, block) @@ -147,7 +150,7 @@ impl<'a> Lowerer<'a> { self.q4k_f32_dot(expressions, src, k, c, &a_handles, block) } _ => Err(LowerError::UnsupportedOperation( - "f32 activation dot only supports Q8_0/Q6K dot8 or Q4K dot8/16/32", + "f32 activation dot only supports Q4_0 dot8/16, Q8_0/Q6K dot8, or Q4K dot8/16/32", )), }, QuantActivation::Q8 => { diff --git a/fusor-ml/tile-ir/src/quantized.rs b/fusor-ml/tile-ir/src/quantized.rs index b9fdebe24..eebf4ed44 100644 --- a/fusor-ml/tile-ir/src/quantized.rs +++ b/fusor-ml/tile-ir/src/quantized.rs @@ -111,6 +111,10 @@ impl GgmlQuantFormat { matches!(self, Self::Q4K | Self::Q4KNative) } + pub const fn is_q4_0_family(self) -> bool { + matches!(self, Self::Q4_0 | Self::Q4_0Native) + } + pub const fn is_q8_0_family(self) -> bool { matches!(self, Self::Q8_0 | Self::Q8_0Native) } diff --git a/interfaces/kalosm/examples/vision.rs b/interfaces/kalosm/examples/vision.rs index 3ccdc50a0..7ec920b37 100644 --- a/interfaces/kalosm/examples/vision.rs +++ b/interfaces/kalosm/examples/vision.rs @@ -6,7 +6,7 @@ async fn main() { tracing_subscriber::fmt::init(); let t_load_start = Instant::now(); let model = Llama::builder() - .with_source(LlamaSource::qwen_2_5_3b_vl_chat_q4()) + .with_source(LlamaSource::gemma_4_e2b_it_qat_chat()) .build() .await .unwrap(); diff --git a/models/kalosm-llama/examples/vision.rs b/models/kalosm-llama/examples/vision.rs index 421634e4b..425ccb572 100644 --- a/models/kalosm-llama/examples/vision.rs +++ b/models/kalosm-llama/examples/vision.rs @@ -1,27 +1,141 @@ +#![recursion_limit = "256"] + use kalosm_llama::prelude::*; use kalosm_streams::text_stream::TextStream; +use std::time::Instant; + +async fn collect_or_bench(mut response: impl TextStream + Unpin, max_tokens: u32) { + let bench = std::env::var_os("KALOSM_VISION_BENCH").is_some(); + let total_start = Instant::now(); + let mut first = None; + let mut last = None; + let mut tokens = 0usize; + let mut output = String::new(); + while let Some(token) = response.next().await { + let now = Instant::now(); + first.get_or_insert(now); + last = Some(now); + tokens += 1; + output.push_str(&token); + } + + if !bench { + print!("{output}"); + println!("\n"); + return; + } -// The demo image is fetched over HTTP, which needs a tokio reactor. + let total_elapsed = total_start.elapsed().as_secs_f64(); + let steady_elapsed = first + .zip(last) + .map(|(first, last)| (last - first).as_secs_f64()) + .unwrap_or_default(); + let steady_tokens = tokens.saturating_sub(1); + let total_tps = if total_elapsed > 0.0 { + tokens as f64 / total_elapsed + } else { + 0.0 + }; + let steady_tps = if steady_elapsed > 0.0 { + steady_tokens as f64 / steady_elapsed + } else { + 0.0 + }; + + print!("{output}"); + println!("\n"); + eprintln!( + "vision_bench tokens={tokens} max_tokens={max_tokens} total_s={total_elapsed:.3} total_tok_s={total_tps:.2} steady_tok_s={steady_tps:.2}" + ); +} + +// The demo image may be fetched over HTTP, which needs a tokio reactor. #[tokio::main] async fn main() { tracing_subscriber::fmt::init(); - let model = Llama::builder() - .with_source(LlamaSource::qwen_2_5_3b_vl_chat_q4()) - .build() - .await - .unwrap(); + let source = match std::env::var("KALOSM_VISION_SOURCE").as_deref() { + Ok("gemma4") => { + let mut source = LlamaSource::gemma_4_e2b_it_qat_chat().with_vision_model( + kalosm_model_types::FileSource::HuggingFace { + model_id: "unsloth/gemma-4-E2B-it-qat-GGUF".into(), + revision: "main".into(), + file: "mmproj-F16.gguf".into(), + }, + ); + if std::env::var_os("KALOSM_LLAMA_MTP_DRAFT").is_some() { + let file = std::env::var("KALOSM_LLAMA_MTP_FILE") + .unwrap_or_else(|_| "MTP/gemma-4-E2B-it-Q4_0-MTP.gguf".into()); + source = source.with_mtp_model(kalosm_model_types::FileSource::HuggingFace { + model_id: "unsloth/gemma-4-E2B-it-qat-GGUF".into(), + revision: "main".into(), + file, + }); + } + source + } + _ => LlamaSource::qwen_2_5_3b_vl_chat_q4(), + }; + let mut builder = Llama::builder().with_source(source); + if std::env::var_os("KALOSM_VISION_CPU").is_some() { + builder = builder.with_device(kalosm_llama::Device::Cpu); + } + let model = builder.build().await.unwrap(); let mut chat = model.chat(); + let max_tokens = std::env::var("KALOSM_VISION_MAX_TOKENS") + .ok() + .and_then(|value| value.parse::().ok()) + .unwrap_or(64); + let mut sampler = GenerationParameters::new().with_max_length(max_tokens); + if let Ok(seed) = std::env::var("KALOSM_VISION_SEED") { + if let Ok(seed) = seed.parse::() { + sampler = sampler.with_seed(seed); + } + } + if std::env::var_os("KALOSM_VISION_STANDARD").is_some() { + sampler = sampler.with_standard_sampler(); + } + if let Ok(temperature) = std::env::var("KALOSM_VISION_TEMPERATURE") { + if let Ok(temperature) = temperature.parse::() { + sampler = sampler.with_temperature(temperature); + if temperature <= 0.0 { + sampler = sampler.with_standard_sampler(); + } + } + } + if std::env::var_os("KALOSM_VISION_TEXT_ONLY").is_some() { + let prompt = std::env::var("KALOSM_VISION_TEXT_PROMPT").unwrap_or_else(|_| { + "Answer in one short sentence: what color is a ripe banana?".into() + }); + let mut response = chat(&prompt).with_sampler(sampler); + if std::env::var_os("KALOSM_VISION_COLLECT").is_some() { + collect_or_bench(response, max_tokens).await; + return; + } + response.to_std_out().await.unwrap(); + println!("\n"); + return; + } + + let image_source = if let Ok(url) = std::env::var("KALOSM_VISION_URL") { + MediaSource::url(url) + } else if let Ok(path) = std::env::var("KALOSM_VISION_IMAGE") { + MediaSource::file(path).unwrap() + } else { + MediaSource::url("https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-VL/assets/demo.jpeg") + }; + let prompt = + std::env::var("KALOSM_VISION_PROMPT").unwrap_or_else(|_| "Describe this image.".into()); + let mut response = chat(&( - MediaChunk::new( - MediaSource::url( - "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-VL/assets/demo.jpeg", - ), - MediaType::Image, - ), - "Describe this image.", - )); + MediaChunk::new(image_source, MediaType::Image), + prompt.as_str(), + )) + .with_sampler(sampler); + if std::env::var_os("KALOSM_VISION_COLLECT").is_some() { + collect_or_bench(response, max_tokens).await; + return; + } response.to_std_out().await.unwrap(); - response.await.unwrap(); println!("\n"); } diff --git a/models/kalosm-llama/src/chat_template.rs b/models/kalosm-llama/src/chat_template.rs index bfb11b54a..59f1648ae 100644 --- a/models/kalosm-llama/src/chat_template.rs +++ b/models/kalosm-llama/src/chat_template.rs @@ -1,6 +1,6 @@ -use std::fmt::Display; +use std::{collections::BTreeMap, fmt::Display}; -use kalosm_language_model::{ChatMessage, ContentChunk}; +use kalosm_language_model::{ChatMessage, ContentChunk, MediaType}; use minijinja::{context, Environment, ErrorKind, Value}; use minijinja_contrib::pycompat; @@ -11,11 +11,14 @@ use pretty_assertions::assert_eq; pub(crate) struct HuggingFaceChatTemplate { environment: Environment<'static>, + enable_thinking_by_default: bool, } impl HuggingFaceChatTemplate { pub(crate) fn create(chat_template: impl Display) -> Result { let chat_template = chat_template.to_string(); + let enable_thinking_by_default = + chat_template.contains("enable_thinking") && chat_template.contains("<|think|>"); let mut environment = Environment::new(); // enable python compatibility methods because most models are tested with python @@ -40,7 +43,10 @@ impl HuggingFaceChatTemplate { // compile the template expression in the environment environment.add_template_owned("main", chat_template)?; - Ok(Self { environment }) + Ok(Self { + environment, + enable_thinking_by_default, + }) } pub(crate) fn format( @@ -64,11 +70,17 @@ impl HuggingFaceChatTemplate { .iter() .map(|chunk| match chunk { ContentChunk::Text(text) => { - context! { text } - } - ContentChunk::Media(_) => { - context! { image => "" } + chunk_context([("type", "text"), ("text", text)]) } + ContentChunk::Media(media) => match media.media_type() { + MediaType::Image => { + chunk_context([("type", "image"), ("image", "")]) + } + MediaType::Video => { + chunk_context([("type", "video"), ("video", "")]) + } + _ => chunk_context([("type", "media")]), + }, }) .collect::>(); chunks.into() @@ -76,13 +88,18 @@ impl HuggingFaceChatTemplate { context! { role, content } }) .collect::>(); - let ctx = context! { bos_token, eos_token, messages, add_generation_prompt, tools }; + let enable_thinking = self.enable_thinking_by_default; + let ctx = context! { bos_token, eos_token, messages, add_generation_prompt, tools, enable_thinking }; let template = self.environment.get_template("main")?; let result = template.render(&ctx)?; Ok(result) } } +fn chunk_context(entries: [(&str, &str); N]) -> Value { + Value::from_serialize(BTreeMap::from(entries)) +} + #[test] fn test_qwen_chat_template() { let template = r#"{%- if tools %} @@ -238,6 +255,49 @@ I'd like to show off how chat templating works!<|vision_start|><|image_pad|><|vi ); } +#[test] +fn test_gemma_vl_chat_template_media_type() { + use kalosm_language_model::{MediaChunk, MediaSource}; + + let template = "{% for message in messages %}{% for item in message['content'] %}{% if item['type'] == 'image' %}<|image|>{% elif item['type'] == 'text' %}{{ item['text'] }}{% endif %}{% endfor %}{% endfor %}"; + let template = HuggingFaceChatTemplate::create(template).unwrap(); + let inputs = [ChatMessage::new( + MessageType::UserMessage, + ( + MediaChunk::new( + MediaSource::url("https://example.com/image.png"), + kalosm_language_model::MediaType::Image, + ), + "Describe this image.", + ), + )]; + let result = template.format("", "", &inputs, false).unwrap(); + assert_eq!(result, "<|image|>Describe this image."); +} + +#[test] +fn test_gemma_4_chat_template_enables_thinking() { + use kalosm_language_model::{MediaChunk, MediaSource}; + + let template = "{{- bos_token -}}{%- if enable_thinking is defined and enable_thinking -%}{{- '<|turn>system\n<|think|>\n\n' -}}{%- endif -%}{%- for message in messages -%}{{- '<|turn>' + message['role'] + '\n' -}}{%- for item in message['content'] -%}{%- if item['type'] == 'image' -%}{{- '<|image|>' -}}{%- elif item['type'] == 'text' -%}{{- item['text'] | trim -}}{%- endif -%}{%- endfor -%}{{- '\n' -}}{%- endfor -%}{%- if add_generation_prompt -%}{{- '<|turn>model\n' -}}{%- endif -%}"; + let template = HuggingFaceChatTemplate::create(template).unwrap(); + let inputs = [ChatMessage::new( + MessageType::UserMessage, + ( + MediaChunk::new( + MediaSource::url("https://example.com/image.png"), + kalosm_language_model::MediaType::Image, + ), + "Describe this image.", + ), + )]; + let result = template.format("", "", &inputs, true).unwrap(); + assert_eq!( + result, + "<|turn>system\n<|think|>\n\n<|turn>user\n<|image|>Describe this image.\n<|turn>model\n" + ); +} + #[test] fn test_llama_chat_template() { let template = "{% set loop_messages = messages %}{% for message in loop_messages %}{% set content = '<|start_header_id|>' + message['role'] + '<|end_header_id|>\n\n'+ message['content'] | trim + '<|eot_id|>' %}{% if loop.index0 == 0 %}{% set content = bos_token + content %}{% endif %}{{ content }}{% endfor %}{% if add_generation_prompt %}{{ '<|start_header_id|>assistant<|end_header_id|>\n\n' }}{% endif %}"; diff --git a/models/kalosm-llama/src/gguf_tokenizer.rs b/models/kalosm-llama/src/gguf_tokenizer.rs index 95d2c037a..a962782ae 100644 --- a/models/kalosm-llama/src/gguf_tokenizer.rs +++ b/models/kalosm-llama/src/gguf_tokenizer.rs @@ -21,6 +21,7 @@ enum PreTokenizerType { Default, Exaone, Falcon, + Gemma4, Gpt2, Gpt3Finnish, Jais, @@ -69,6 +70,7 @@ impl PreTokenizerType { "'s|'t|'re|'ve|'m|'ll|'d| ?\\p{L}+| ?\\p{N}+| ?[^\\s\\p{L}\\p{N}]+|\\s+(?!\\S)", "[0-9][0-9][0-9]", ], + Self::Gemma4 => &["[^\\n]+|[\\n]+"], Self::Starcoder | Self::Refact | Self::CommandR @@ -261,6 +263,8 @@ fn sanitize_regex(regex: &str) -> String { #[derive(Clone, Copy)] pub(crate) struct GGUFPreTokenizerConfig { add_bos: bool, + byte_level: bool, + escape_whitespaces: bool, ignore_merges: bool, ty: PreTokenizerType, } @@ -298,7 +302,11 @@ impl GGUFPreTokenizerConfig { let merges = merges .into_iter() .map(|(left, right)| format!("{left} {right}")); - let bpe = FastBpe::from_vocab_and_merges(vocab, merges, self.ignore_merges)?; + let bpe = if self.byte_level { + FastBpe::from_vocab_and_merges(vocab, merges, self.ignore_merges)? + } else { + FastBpe::from_raw_utf8_vocab_and_merges(vocab, merges, self.ignore_merges)? + }; let pre_tokenizer = PreTokenizer::new(self.ty)?; Ok(GgufTokenizer { @@ -308,6 +316,7 @@ impl GGUFPreTokenizerConfig { special_token_matches, bos_token, add_bos: self.add_bos, + escape_whitespaces: self.escape_whitespaces, }) } } @@ -316,6 +325,8 @@ impl Default for GGUFPreTokenizerConfig { fn default() -> Self { Self { add_bos: true, + byte_level: true, + escape_whitespaces: false, ignore_merges: false, ty: PreTokenizerType::Default, } @@ -330,6 +341,7 @@ pub(crate) struct GgufTokenizer { special_token_matches: Vec, u32)>>, bos_token: Option, add_bos: bool, + escape_whitespaces: bool, } impl GgufTokenizer { @@ -394,7 +406,12 @@ impl GgufTokenizer { bytes.extend_from_slice(token_bytes); } } - String::from_utf8_lossy(&bytes).into_owned() + let text = String::from_utf8_lossy(&bytes).into_owned(); + if self.escape_whitespaces { + text.replace('▁', " ") + } else { + text + } } pub(crate) fn is_special_token(&self, token: u32) -> bool { @@ -419,6 +436,14 @@ impl GgufTokenizer { tokenization, } = buffers; + let normalized; + let text = if self.escape_whitespaces { + normalized = text.replace(' ', "▁"); + normalized.as_str() + } else { + text + }; + self.pre_tokenizer.split_into_ranges(text, pre_tokenizer); for piece in pre_tokenizer.pieces.iter().copied() { if piece.is_empty() { @@ -460,6 +485,14 @@ pub(crate) fn get_pre_tokenizer( ignore_merges: true, add_bos: true, ty: PreTokenizerType::Llama3, + ..Default::default() + }, + "gemma4" => GGUFPreTokenizerConfig { + add_bos: true, + byte_level: false, + escape_whitespaces: true, + ty: PreTokenizerType::Gemma4, + ..Default::default() }, "deepseek-llm" => GGUFPreTokenizerConfig { ty: PreTokenizerType::DeepseekLlm, @@ -534,6 +567,7 @@ pub(crate) fn get_pre_tokenizer( ignore_merges: true, add_bos: true, ty: PreTokenizerType::Tekken, + ..Default::default() }, "smollm" => GGUFPreTokenizerConfig { ty: PreTokenizerType::Smollm, @@ -569,7 +603,7 @@ pub(crate) fn get_pre_tokenizer( }; if let Some(add_bos) = add_bos { - tokenizer.add_bos = add_bos; + tokenizer.add_bos = add_bos || pre_tokenizer_type == "gemma4"; } tokenizer diff --git a/models/kalosm-llama/src/language_model.rs b/models/kalosm-llama/src/language_model.rs index ac8d71a8c..c659530b4 100644 --- a/models/kalosm-llama/src/language_model.rs +++ b/models/kalosm-llama/src/language_model.rs @@ -54,7 +54,12 @@ where .tokenizer .as_ref() .is_some_and(|tokenizer| !cache.exists(tokenizer)); - model_missing || tokenizer_missing + let mtp_missing = self + .source + .mtp_model + .as_ref() + .is_some_and(|mtp| !cache.exists(mtp)); + model_missing || tokenizer_missing || mtp_missing } } @@ -118,6 +123,13 @@ where } Vec::new() }; + if std::env::var_os("KALOSM_TRACE_PROMPT").is_some() { + eprintln!( + "[stream_text] prompt_len={} image_count={}", + text.len(), + images.len() + ); + } self.inner .sender .unbounded_send(Task::UnstructuredGeneration(UnstructuredGenerationTask::< diff --git a/models/kalosm-llama/src/model/forward.rs b/models/kalosm-llama/src/model/forward.rs index fc93bf559..d2224a354 100644 --- a/models/kalosm-llama/src/model/forward.rs +++ b/models/kalosm-llama/src/model/forward.rs @@ -46,7 +46,11 @@ where }; let token_start = trace_enabled.then(Instant::now); let build_start = trace_enabled.then(Instant::now); - let logits = model.forward(tokens, images, device, cache); + let logits = if !images.is_empty() && model.should_chunk_multimodal_prompt(device) { + model.forward_chunked_multimodal(tokens, images, device, cache) + } else { + model.forward(tokens, images, device, cache) + }; if let Some(start) = build_start { tracing::info!( "forward_graph_build path={path} decode_eligible={decode_eligible} elapsed={:?}", @@ -229,7 +233,9 @@ where ); } - if gpu_fused_logits_sampling_enabled() { + if gpu_fused_logits_sampling_enabled() + && (images.is_empty() || !model.should_chunk_multimodal_prompt(device)) + { return Self::forward_sample_token_fused_logits( ForwardInputs { model, @@ -327,6 +333,18 @@ where return Ok(None); } + if !images.is_empty() && model.should_chunk_multimodal_prompt(device) { + let logits = model.forward_chunked_multimodal(tokens, images, device, cache)?; + let logits: fusor::Tensor<1, f32> = logits.squeeze(0).cast(); + return Self::forward_sample_token_pending_from_logits( + logits, + sampler, + previous_tokens, + None, + top_k, + ); + } + Self::forward_sample_token_pending_from_hidden( model.forward_last_hidden_f32(tokens, images, device, cache)?, model, @@ -337,6 +355,37 @@ where ) } + fn forward_sample_token_pending_from_logits( + logits: fusor::Tensor<1, f32>, + sampler: &mut LlamaGpuSamplerState, + previous_tokens: Vec, + previous_gpu_token: Option<&fusor::GpuSampledToken>, + top_k: usize, + ) -> Result, LlamaModelError> { + match sampler.config.sampling_strategy { + kalosm_language_model::SamplingStrategy::Mirostat2 => { + let params = sampler.mirostat_params(top_k); + let mirostat = sampler.mirostat.as_mut().ok_or_else(|| { + LlamaModelError::SamplerError("missing Mirostat GPU sampler".into()) + })?; + logits + .sample_mirostat2_token_pending( + mirostat, + &previous_tokens, + previous_gpu_token, + params, + ) + .map_err(LlamaModelError::from) + } + kalosm_language_model::SamplingStrategy::Standard => { + let params = sampler.standard_params(top_k); + logits + .sample_standard_token_pending(&previous_tokens, previous_gpu_token, params) + .map_err(LlamaModelError::from) + } + } + } + pub(crate) fn forward_sample_token_from_gpu_token_pending( model: &Model, device: &Device, @@ -381,28 +430,13 @@ where .squeeze(0) .to_concrete() .q_mat_mul(model.output_matrix()); - match sampler.config.sampling_strategy { - kalosm_language_model::SamplingStrategy::Mirostat2 => { - let params = sampler.mirostat_params(top_k); - let mirostat = sampler.mirostat.as_mut().ok_or_else(|| { - LlamaModelError::SamplerError("missing Mirostat GPU sampler".into()) - })?; - logits - .sample_mirostat2_token_pending( - mirostat, - &previous_tokens, - previous_gpu_token, - params, - ) - .map_err(LlamaModelError::from) - } - kalosm_language_model::SamplingStrategy::Standard => { - let params = sampler.standard_params(top_k); - logits - .sample_standard_token_pending(&previous_tokens, previous_gpu_token, params) - .map_err(LlamaModelError::from) - } - } + Self::forward_sample_token_pending_from_logits( + logits, + sampler, + previous_tokens, + previous_gpu_token, + top_k, + ) } fn forward_sample_token_fused_logits<'a>( diff --git a/models/kalosm-llama/src/model/inference.rs b/models/kalosm-llama/src/model/inference.rs index ae36773a1..ae9ede11a 100644 --- a/models/kalosm-llama/src/model/inference.rs +++ b/models/kalosm-llama/src/model/inference.rs @@ -28,6 +28,16 @@ where .tokenizer .encode(&prompt, false) .map_err(LlamaModelError::Tokenizer)?; + if std::env::var_os("KALOSM_TRACE_PROMPT").is_some() { + let decoded = self + .tokenizer + .decode(&tokens, false) + .unwrap_or_else(|err| format!("")); + eprintln!( + "[prompt] len={} text={prompt:?} tokens={tokens:?} decoded={decoded:?}", + tokens.len() + ); + } let mut text_stream = TokenOutputStream::new(self.tokenizer.clone()); for &token in &tokens { text_stream @@ -35,13 +45,33 @@ where .map_err(LlamaModelError::TokenOutputStreamError)?; } + if mtp_speculative_enabled() + && stop_on.is_none() + && sampler.sampling_strategy == kalosm_language_model::SamplingStrategy::Standard + && sampler.temperature <= 0.0 + && gpu_token_sampling_enabled() + && self.mtp.is_some() + { + if let Some(gpu_sampler) = LlamaGpuSamplerState::new(&self.device, sampler, seed) { + return self + .infer_mtp_speculative( + &tokens, + &images, + text_stream, + session, + max_tokens, + gpu_sampler, + on_token, + finished, + ) + .await; + } + } + if gpu_token_sampling_enabled() && stop_on.is_none() { if let Some(mut gpu_sampler) = LlamaGpuSamplerState::new(&self.device, sampler, seed) { let top_k = gpu_sample_top_k(&gpu_sampler.config); - if gpu_run_ahead_enabled() - && images.is_empty() - && self.model.supports_gpu_token_run_ahead() - { + if gpu_run_ahead_enabled() && self.model.supports_gpu_token_run_ahead() { let next_token = { let previous_tokens = gpu_sampler.previous_tokens(&text_stream); let mut session_lock = session @@ -64,7 +94,7 @@ where }; if let Some(mut next_token) = next_token { - let stop_token = self.model.config.stop_token; + let stop_tokens = &self.model.config.stop_tokens; let mut tokens_generated = 0; while !finished.is_canceled() && tokens_generated < max_tokens { let scheduled_next = if tokens_generated + 1 < max_tokens { @@ -102,7 +132,16 @@ where "pending GPU sampler refused slow fallback".into(), ) })?; - if new_token == stop_token { + if std::env::var_os("KALOSM_TRACE_SAMPLED_TOKEN").is_some() { + let decoded = self + .tokenizer + .decode(&[new_token], false) + .unwrap_or_else(|err| format!("")); + eprintln!( + "[sampled_token] index={tokens_generated} id={new_token} text={decoded:?} stop_tokens={stop_tokens:?}" + ); + } + if stop_tokens.contains(&new_token) { tracing::trace!("Stopping on stop token"); break; } @@ -196,11 +235,20 @@ where } .await?; - let stop_token = self.model.config.stop_token; + let stop_tokens = &self.model.config.stop_tokens; let mut tokens_generated = 0; while !finished.is_canceled() && tokens_generated < max_tokens { let new_token = next_token; - if new_token == stop_token { + if std::env::var_os("KALOSM_TRACE_SAMPLED_TOKEN").is_some() { + let decoded = self + .tokenizer + .decode(&[new_token], false) + .unwrap_or_else(|err| format!("")); + eprintln!( + "[sampled_token] index={tokens_generated} id={new_token} text={decoded:?} stop_tokens={stop_tokens:?}" + ); + } + if stop_tokens.contains(&new_token) { tracing::trace!("Stopping on stop token"); break; } @@ -282,14 +330,23 @@ where let mut queued_text_matching_stop_on = String::new(); let stop_on_lowercase = stop_on.as_ref().map(|s| s.to_lowercase()); let stop_on_lowercase = stop_on_lowercase.as_deref(); - let stop_token = self.model.config.stop_token; + let stop_tokens = &self.model.config.stop_tokens; let mut tokens_generated = 0; 'generate: while !finished.is_canceled() && tokens_generated < max_tokens { let new_token = text_stream .sample_token(&mut cpu_sampler, logits, stop_on.as_deref(), sample_top_k) .map_err(LlamaModelError::TokenOutputStreamError)?; - if new_token == stop_token { + if std::env::var_os("KALOSM_TRACE_SAMPLED_TOKEN").is_some() { + let decoded = self + .tokenizer + .decode(&[new_token], false) + .unwrap_or_else(|err| format!("")); + eprintln!( + "[sampled_token] index={tokens_generated} id={new_token} text={decoded:?} stop_tokens={stop_tokens:?}" + ); + } + if stop_tokens.contains(&new_token) { tracing::trace!("Stopping on stop token"); break; } @@ -394,4 +451,540 @@ where Ok(()) } + + async fn infer_mtp_speculative( + &self, + prompt_tokens: &[u32], + images: &[LlamaImage], + mut text_stream: TokenOutputStream, + session: crate::LlamaSession, + max_tokens: u32, + mut gpu_sampler: LlamaGpuSamplerState, + mut on_token: crate::BoxedTokenCallback, + finished: &futures_channel::oneshot::Sender>, + ) -> Result<(), LlamaModelError> { + let mtp = self.mtp.as_ref().ok_or_else(|| { + LlamaModelError::SamplerError("Gemma4 MTP assistant is not loaded".into()) + })?; + let top_k = gpu_sample_top_k(&gpu_sampler.config); + let target = { + let mut session_lock = session + .cache + .write() + .map_err(|err| LlamaModelError::Session(err.to_string()))?; + self.model + .forward_logits_and_nextn_f32( + prompt_tokens, + images, + &self.device, + Some(&mut session_lock), + ) + .map_err(LlamaModelError::from)? + }; + let last_row = target.logits.shape()[0].saturating_sub(1); + let mut pending_h = Self::h_nextn_row(&target.h_nextn, last_row); + let history = Self::mtp_previous_tokens(&gpu_sampler.config, text_stream.tokens(), &[]); + let mut next_token = Self::sample_standard_logits_row( + target.logits, + last_row, + &mut gpu_sampler, + history, + top_k, + ) + .await?; + + let stop_tokens = &self.model.config.stop_tokens; + let mut tokens_generated = 0usize; + let mut drafted_total = 0usize; + let mut accepted_total = 0usize; + while !finished.is_canceled() && tokens_generated < max_tokens as usize { + let new_token = next_token; + if std::env::var_os("KALOSM_TRACE_SAMPLED_TOKEN").is_some() { + let decoded = self + .tokenizer + .decode(&[new_token], false) + .unwrap_or_else(|err| format!("")); + eprintln!( + "[sampled_token] index={tokens_generated} id={new_token} text={decoded:?} stop_tokens={stop_tokens:?}" + ); + } + if stop_tokens.contains(&new_token) { + tracing::trace!("Stopping on stop token"); + break; + } + tokens_generated += 1; + if let Some(new_text) = text_stream + .next_token(new_token) + .map_err(LlamaModelError::TokenOutputStreamError)? + { + on_token(new_text)?; + } + if finished.is_canceled() || tokens_generated >= max_tokens as usize { + break; + } + + let remaining = max_tokens as usize - tokens_generated; + let mut draft_limit = mtp_draft_limit(mtp.draft_n()).min(remaining); + if mtp_auto_fallback_enabled() && drafted_total < mtp_fallback_probe_drafts() { + draft_limit = + draft_limit.min(mtp_fallback_probe_drafts().saturating_sub(drafted_total)); + } + let mut draft_tokens = Vec::with_capacity(draft_limit); + let mut assistant_h = pending_h.clone(); + let mut assistant_token = new_token; + let draft_position = { + let session_lock = session + .cache + .read() + .map_err(|err| LlamaModelError::Session(err.to_string()))?; + session_lock.tokens.len() + }; + for _ in 0..draft_limit { + let step = { + let session_lock = session + .cache + .read() + .map_err(|err| LlamaModelError::Session(err.to_string()))?; + mtp.draft_step( + &self.model, + assistant_token, + &assistant_h, + &session_lock, + &self.device, + draft_position, + ) + .map_err(LlamaModelError::from)? + }; + let history = Self::mtp_previous_tokens( + &gpu_sampler.config, + text_stream.tokens(), + &draft_tokens, + ); + assistant_token = + Self::sample_standard_logits(step.logits, &mut gpu_sampler, history, top_k) + .await?; + assistant_h = step.h_nextn; + draft_tokens.push(assistant_token); + } + drafted_total += draft_tokens.len(); + + let mut accepted = 0usize; + let mut verify_input = new_token; + let mut mismatched = false; + for draft in draft_tokens.iter().copied() { + let history = + Self::mtp_previous_tokens(&gpu_sampler.config, text_stream.tokens(), &[]); + let (verified, h_nextn) = self + .mtp_target_step(verify_input, &session, &mut gpu_sampler, history, top_k) + .await?; + pending_h = h_nextn; + if verified != draft { + next_token = verified; + mismatched = true; + break; + } + + accepted += 1; + if stop_tokens.contains(&draft) { + tracing::trace!("Stopping on accepted MTP stop token"); + if std::env::var_os("KALOSM_TRACE_MTP").is_some() { + tracing::info!( + "mtp_summary drafted={drafted_total} accepted={accepted_total}" + ); + } + return Ok(()); + } + + tokens_generated += 1; + if let Some(new_text) = text_stream + .next_token(draft) + .map_err(LlamaModelError::TokenOutputStreamError)? + { + on_token(new_text)?; + } + + verify_input = draft; + if tokens_generated >= max_tokens as usize || finished.is_canceled() { + break; + } + } + accepted_total += accepted; + + if !mismatched + && accepted == draft_tokens.len() + && tokens_generated < max_tokens as usize + && !finished.is_canceled() + { + let history = + Self::mtp_previous_tokens(&gpu_sampler.config, text_stream.tokens(), &[]); + let (bonus, h_nextn) = self + .mtp_target_step(verify_input, &session, &mut gpu_sampler, history, top_k) + .await?; + pending_h = h_nextn; + next_token = bonus; + } + + if Self::should_mtp_fallback(drafted_total, accepted_total) { + if std::env::var_os("KALOSM_TRACE_MTP").is_some() { + tracing::info!( + "mtp_auto_fallback drafted={drafted_total} accepted={accepted_total}" + ); + } + return self + .infer_target_only_from_pending( + next_token, + text_stream, + session, + max_tokens, + tokens_generated, + gpu_sampler, + on_token, + finished, + ) + .await; + } + + { + use std::sync::atomic::{AtomicBool, Ordering}; + let yielded = AtomicBool::new(false); + std::future::poll_fn(|cx| { + if yielded.load(Ordering::Relaxed) { + std::task::Poll::Ready(()) + } else { + yielded.store(true, Ordering::Relaxed); + cx.waker().wake_by_ref(); + std::task::Poll::Pending + } + }) + .await; + } + } + if std::env::var_os("KALOSM_TRACE_MTP").is_some() { + tracing::info!("mtp_summary drafted={drafted_total} accepted={accepted_total}"); + } + Ok(()) + } + + async fn infer_target_only_from_pending( + &self, + mut next_token: u32, + mut text_stream: TokenOutputStream, + session: crate::LlamaSession, + max_tokens: u32, + mut tokens_generated: usize, + mut gpu_sampler: LlamaGpuSamplerState, + mut on_token: crate::BoxedTokenCallback, + finished: &futures_channel::oneshot::Sender>, + ) -> Result<(), LlamaModelError> { + let stop_tokens = &self.model.config.stop_tokens; + let top_k = gpu_sample_top_k(&gpu_sampler.config); + while !finished.is_canceled() && tokens_generated < max_tokens as usize { + let new_token = next_token; + if std::env::var_os("KALOSM_TRACE_SAMPLED_TOKEN").is_some() { + let decoded = self + .tokenizer + .decode(&[new_token], false) + .unwrap_or_else(|err| format!("")); + eprintln!( + "[sampled_token] index={tokens_generated} id={new_token} text={decoded:?} stop_tokens={stop_tokens:?}" + ); + } + if stop_tokens.contains(&new_token) { + tracing::trace!("Stopping on stop token"); + break; + } + tokens_generated += 1; + if let Some(new_text) = text_stream + .next_token(new_token) + .map_err(LlamaModelError::TokenOutputStreamError)? + { + on_token(new_text)?; + } + if finished.is_canceled() || tokens_generated >= max_tokens as usize { + break; + } + + let previous_tokens = gpu_sampler.previous_tokens(&text_stream); + let pending = { + let mut session_lock = session + .cache + .write() + .map_err(|err| LlamaModelError::Session(err.to_string()))?; + Self::forward_sample_token_pending( + ForwardInputs { + model: &self.model, + device: &self.device, + tokens: &[new_token], + images: &[], + cache: Some(&mut session_lock), + tokenizer: &self.tokenizer, + }, + &mut gpu_sampler, + previous_tokens.clone(), + top_k, + )? + }; + if let Some(next_pending) = pending { + return self + .infer_target_only_from_gpu_pending( + next_pending, + text_stream, + session, + max_tokens, + tokens_generated, + gpu_sampler, + on_token, + finished, + ) + .await; + } + + next_token = { + let mut session_lock = session + .cache + .write() + .map_err(|err| LlamaModelError::Session(err.to_string()))?; + Self::forward_sample_token( + ForwardInputs { + model: &self.model, + device: &self.device, + tokens: &[new_token], + images: &[], + cache: Some(&mut session_lock), + tokenizer: &self.tokenizer, + }, + &mut gpu_sampler, + previous_tokens, + top_k, + ) + } + .await?; + + { + use std::sync::atomic::{AtomicBool, Ordering}; + let yielded = AtomicBool::new(false); + std::future::poll_fn(|cx| { + if yielded.load(Ordering::Relaxed) { + std::task::Poll::Ready(()) + } else { + yielded.store(true, Ordering::Relaxed); + cx.waker().wake_by_ref(); + std::task::Poll::Pending + } + }) + .await; + } + } + Ok(()) + } + + async fn infer_target_only_from_gpu_pending( + &self, + mut next_token: fusor::GpuSampledToken, + mut text_stream: TokenOutputStream, + session: crate::LlamaSession, + max_tokens: u32, + mut tokens_generated: usize, + mut gpu_sampler: LlamaGpuSamplerState, + mut on_token: crate::BoxedTokenCallback, + finished: &futures_channel::oneshot::Sender>, + ) -> Result<(), LlamaModelError> { + let stop_tokens = &self.model.config.stop_tokens; + let top_k = gpu_sample_top_k(&gpu_sampler.config); + while !finished.is_canceled() && tokens_generated < max_tokens as usize { + let scheduled_next = if tokens_generated + 1 < max_tokens as usize { + let previous_tokens = gpu_sampler.previous_tokens(&text_stream); + let mut speculative_cache = session + .cache + .read() + .map_err(|err| LlamaModelError::Session(err.to_string()))? + .clone(); + if speculative_cache.tokens.len() < self.model.config.context_length { + Self::forward_sample_token_from_gpu_token_pending( + &self.model, + &self.device, + &next_token, + &mut speculative_cache, + &mut gpu_sampler, + previous_tokens, + top_k, + )? + .map(|(next, cache_slot)| (next, speculative_cache, cache_slot)) + } else { + None + } + } else { + None + }; + + let new_token = next_token + .read_token() + .await + .map_err(|err| LlamaModelError::Fusor(fusor::Error::Gpu(err)))? + .ok_or_else(|| { + LlamaModelError::SamplerError( + "pending GPU sampler refused slow fallback".into(), + ) + })?; + if std::env::var_os("KALOSM_TRACE_SAMPLED_TOKEN").is_some() { + let decoded = self + .tokenizer + .decode(&[new_token], false) + .unwrap_or_else(|err| format!("")); + eprintln!( + "[sampled_token] index={tokens_generated} id={new_token} text={decoded:?} stop_tokens={stop_tokens:?}" + ); + } + if stop_tokens.contains(&new_token) { + tracing::trace!("Stopping on stop token"); + break; + } + + tokens_generated += 1; + if let Some(new_text) = text_stream + .next_token(new_token) + .map_err(LlamaModelError::TokenOutputStreamError)? + { + on_token(new_text)?; + } + + if let Some((scheduled_token, mut speculative_cache, cache_slot)) = scheduled_next { + if let Some(slot) = speculative_cache.tokens.get_mut(cache_slot) { + *slot = new_token; + } + *session + .cache + .write() + .map_err(|err| LlamaModelError::Session(err.to_string()))? = speculative_cache; + next_token = scheduled_token; + } else if !finished.is_canceled() && tokens_generated < max_tokens as usize { + let previous_tokens = gpu_sampler.previous_tokens(&text_stream); + let pending = { + let mut session_lock = session + .cache + .write() + .map_err(|err| LlamaModelError::Session(err.to_string()))?; + Self::forward_sample_token_pending( + ForwardInputs { + model: &self.model, + device: &self.device, + tokens: &[new_token], + images: &[], + cache: Some(&mut session_lock), + tokenizer: &self.tokenizer, + }, + &mut gpu_sampler, + previous_tokens, + top_k, + )? + }; + match pending { + Some(sampled) => next_token = sampled, + None => break, + } + } else { + break; + } + + { + use std::sync::atomic::{AtomicBool, Ordering}; + let yielded = AtomicBool::new(false); + std::future::poll_fn(|cx| { + if yielded.load(Ordering::Relaxed) { + std::task::Poll::Ready(()) + } else { + yielded.store(true, Ordering::Relaxed); + cx.waker().wake_by_ref(); + std::task::Poll::Pending + } + }) + .await; + } + } + Ok(()) + } + + async fn mtp_target_step( + &self, + token: u32, + session: &crate::LlamaSession, + sampler: &mut LlamaGpuSamplerState, + previous_tokens: Vec, + top_k: usize, + ) -> Result<(u32, fusor::Tensor<2, f32>), LlamaModelError> { + let target = { + let mut session_lock = session + .cache + .write() + .map_err(|err| LlamaModelError::Session(err.to_string()))?; + self.model + .forward_logits_and_nextn_f32(&[token], &[], &self.device, Some(&mut session_lock)) + .map_err(LlamaModelError::from)? + }; + let h_nextn = Self::h_nextn_row(&target.h_nextn, 0); + let token = + Self::sample_standard_logits_row(target.logits, 0, sampler, previous_tokens, top_k) + .await?; + Ok((token, h_nextn)) + } + + fn should_mtp_fallback(drafted_total: usize, accepted_total: usize) -> bool { + if !mtp_auto_fallback_enabled() { + return false; + } + if drafted_total < mtp_fallback_probe_drafts() { + return false; + } + accepted_total * 100 < drafted_total * mtp_fallback_min_accept_percent() + } + + fn h_nextn_row(h_nextn: &fusor::Tensor<2, f32>, row: usize) -> fusor::Tensor<2, f32> { + let row: fusor::Tensor<1, f32> = h_nextn.i((row, ..)).to_concrete(); + row.unsqueeze(0).to_concrete() + } + + fn mtp_previous_tokens( + config: &GpuSamplerConfig, + base_tokens: &[u32], + extra_tokens: &[u32], + ) -> Vec { + let range = config.repetition_penalty_range; + if range == 0 { + return Vec::new(); + } + let total_len = base_tokens.len() + extra_tokens.len(); + let keep_from = total_len.saturating_sub(range); + let mut result = Vec::with_capacity(range.min(total_len)); + if keep_from < base_tokens.len() { + result.extend_from_slice(&base_tokens[keep_from..]); + result.extend_from_slice(extra_tokens); + } else { + result.extend_from_slice(&extra_tokens[keep_from - base_tokens.len()..]); + } + result + } + + async fn sample_standard_logits( + logits: fusor::Tensor<1, f32>, + sampler: &mut LlamaGpuSamplerState, + previous_tokens: Vec, + top_k: usize, + ) -> Result { + let params = sampler.standard_params(top_k); + logits + .sample_standard_token(&previous_tokens, params) + .await + .map_err(LlamaModelError::from) + } + + async fn sample_standard_logits_row( + logits: fusor::Tensor<2, f32>, + row: usize, + sampler: &mut LlamaGpuSamplerState, + previous_tokens: Vec, + top_k: usize, + ) -> Result { + let row_logits: fusor::Tensor<1, f32> = logits.i((row, ..)).to_concrete(); + Self::sample_standard_logits(row_logits, sampler, previous_tokens, top_k).await + } } diff --git a/models/kalosm-llama/src/model/mod.rs b/models/kalosm-llama/src/model/mod.rs index def415db7..55420cc8b 100644 --- a/models/kalosm-llama/src/model/mod.rs +++ b/models/kalosm-llama/src/model/mod.rs @@ -1,6 +1,6 @@ use crate::gguf_tokenizer::get_pre_tokenizer; use crate::raw::cache::LlamaCache; -use crate::raw::{LlamaVarSource, Model, RopeScalingConfig}; +use crate::raw::{Gemma4MtpAssistant, LlamaVarSource, Model, RopeScalingConfig}; use crate::token_stream::TokenOutputStream; use crate::token_stream::TokenOutputStreamError; use crate::tokenizer::{LlamaTokenizer, LlamaTokenizerError}; @@ -109,6 +109,42 @@ fn gpu_run_ahead_enabled() -> bool { .unwrap_or(true) } +fn mtp_speculative_enabled() -> bool { + std::env::var_os("KALOSM_LLAMA_MTP") + .map(|value| value != "0") + .unwrap_or(true) +} + +fn mtp_draft_limit(default: usize) -> usize { + std::env::var("KALOSM_LLAMA_MTP_DRAFT_N") + .ok() + .and_then(|value| value.parse::().ok()) + .unwrap_or(default) + .max(1) +} + +fn mtp_auto_fallback_enabled() -> bool { + std::env::var_os("KALOSM_LLAMA_MTP_AUTO_FALLBACK") + .map(|value| value != "0") + .unwrap_or(true) +} + +fn mtp_fallback_probe_drafts() -> usize { + std::env::var("KALOSM_LLAMA_MTP_FALLBACK_PROBE_DRAFTS") + .ok() + .and_then(|value| value.parse::().ok()) + .unwrap_or(1) + .max(1) +} + +fn mtp_fallback_min_accept_percent() -> usize { + std::env::var("KALOSM_LLAMA_MTP_FALLBACK_MIN_ACCEPT_PERCENT") + .ok() + .and_then(|value| value.parse::().ok()) + .unwrap_or(60) + .min(100) +} + fn decode_trace_enabled() -> bool { std::env::var_os("KALOSM_TRACE_DECODE_TIMING").is_some() || std::env::var_os("FUSOR_TRACE_DECODE").is_some() @@ -116,16 +152,22 @@ fn decode_trace_enabled() -> bool { } fn gpu_sample_top_k(config: &GpuSamplerConfig) -> usize { + if let Some(top_k) = std::env::var("KALOSM_LLAMA_GPU_SAMPLE_TOP_K") + .ok() + .and_then(|value| value.parse::().ok()) + { + return top_k.max(1); + } + if config.sampling_strategy == kalosm_language_model::SamplingStrategy::Standard + && config.temperature <= 0.0 + { + return 1; + } let default_top_k = match config.sampling_strategy { kalosm_language_model::SamplingStrategy::Mirostat2 => 16, kalosm_language_model::SamplingStrategy::Standard => 64, }; - std::env::var("KALOSM_LLAMA_GPU_SAMPLE_TOP_K") - .ok() - .and_then(|value| value.parse().ok()) - .or(config.top_k) - .unwrap_or(default_top_k) - .max(1) + config.top_k.unwrap_or(default_top_k).max(1) } fn parse_external_tokenizer( @@ -179,15 +221,17 @@ fn tokenizer_from_gguf_source( .clone() .try_into() .map_err(|_| LlamaSourceError::NoTokenizer)?; - if &*tokenizer_model != "gpt2" { + if !matches!(&*tokenizer_model, "gpt2" | "gemma4") { return Err(LlamaSourceError::NoTokenizer); } let pre: Box = source .get("tokenizer.ggml.pre") + .ok() + .cloned() + .map(|v| v.try_into()) + .transpose() .map_err(|_| LlamaSourceError::NoTokenizer)? - .clone() - .try_into() - .map_err(|_| LlamaSourceError::NoTokenizer)?; + .unwrap_or_else(|| tokenizer_model.clone()); let add_bos_token = source .get("tokenizer.ggml.add_bos_token") .ok() @@ -425,7 +469,35 @@ fn trim_previous_tokens_for_gpu_tail( #[cfg(test)] mod tests { - use super::trim_previous_tokens_for_gpu_tail; + use super::{gpu_sample_top_k, trim_previous_tokens_for_gpu_tail}; + use crate::GpuSamplerConfig; + use kalosm_language_model::SamplingStrategy; + use std::sync::{Mutex, OnceLock}; + + fn env_lock() -> &'static Mutex<()> { + static LOCK: OnceLock> = OnceLock::new(); + LOCK.get_or_init(|| Mutex::new(())) + } + + fn with_gpu_sample_top_k_env(value: Option<&str>, f: impl FnOnce() -> R) -> R { + let _guard = env_lock().lock().unwrap_or_else(|err| err.into_inner()); + let prior = std::env::var("KALOSM_LLAMA_GPU_SAMPLE_TOP_K").ok(); + // SAFETY: these tests serialize access to this process-wide env var. + unsafe { + match value { + Some(value) => std::env::set_var("KALOSM_LLAMA_GPU_SAMPLE_TOP_K", value), + None => std::env::remove_var("KALOSM_LLAMA_GPU_SAMPLE_TOP_K"), + } + } + let result = f(); + unsafe { + match prior { + Some(value) => std::env::set_var("KALOSM_LLAMA_GPU_SAMPLE_TOP_K", value), + None => std::env::remove_var("KALOSM_LLAMA_GPU_SAMPLE_TOP_K"), + } + } + result + } #[test] fn gpu_tail_history_preserves_repetition_window() { @@ -446,6 +518,31 @@ mod tests { (vec![1, 2], true) ); } + + #[test] + fn greedy_standard_gpu_sampling_uses_top_one() { + let mut config = GpuSamplerConfig::new(0.0, 5.0, 0.1, 10.0, 1.0, 64, Some(64)); + config.sampling_strategy = SamplingStrategy::Standard; + with_gpu_sample_top_k_env(None, || { + assert_eq!(gpu_sample_top_k(&config), 1); + }); + + config.temperature = 0.8; + with_gpu_sample_top_k_env(None, || { + assert_eq!(gpu_sample_top_k(&config), 64); + }); + with_gpu_sample_top_k_env(Some("8"), || { + assert_eq!(gpu_sample_top_k(&config), 8); + }); + } + + #[test] + fn greedy_top_k_keeps_mirostat_default() { + let config = GpuSamplerConfig::new(0.0, 5.0, 0.1, 10.0, 1.0, 64, None); + with_gpu_sample_top_k_env(None, || { + assert_eq!(gpu_sample_top_k(&config), 16); + }); + } } struct ForwardTrace { @@ -610,6 +707,7 @@ impl From for LlamaModelError { /// The inner, synchronous Llama model. pub(crate) struct LlamaModel { pub(crate) model: Model, + pub(crate) mtp: Option>, pub(crate) device: Device, pub(crate) tokenizer: Arc, } @@ -755,6 +853,29 @@ where None }; + #[cfg(not(target_arch = "wasm32"))] + let mtp_model_bytes = match &builder.source.mtp_model { + Some(mtp_model) => { + let mtp_model_source = format!("MTP Model ({mtp_model})"); + let mut create_progress = + ModelLoadingProgress::downloading_progress(mtp_model_source); + let mtp_model_bytes = builder + .source + .cache + .get_bytes(mtp_model, |progress| handler(create_progress(progress))) + .await?; + Some(mtp_model_bytes) + } + None => None, + }; + #[cfg(target_arch = "wasm32")] + let mtp_model_bytes: Option> = { + if builder.source.mtp_model.is_some() { + tracing::warn!("Gemma4 MTP speculative decoding is not loaded on wasm targets"); + } + None + }; + let source = format!("Model ({})", builder.source.model[0]); let mut create_progress = ModelLoadingProgress::downloading_progress(source); #[cfg(target_arch = "wasm32")] @@ -784,7 +905,7 @@ where let override_chat_template = builder.source.override_chat_template.clone(); #[cfg(target_arch = "wasm32")] - let (model, tokenizer) = { + let (model, tokenizer, mtp) = { let files_with_metadata = match model_source { WasmModelSource::Opfs(model_files) => { if model_files.is_empty() { @@ -840,13 +961,14 @@ where ) .await?; - (model, tokenizer) + (model, tokenizer, None) }; #[cfg(not(target_arch = "wasm32"))] - let (model, tokenizer) = { + let (model, tokenizer, mtp) = { let device = device.clone(); - let load_model = move || -> Result<(Model, LlamaTokenizer), LlamaSourceError> { + let load_model = + move || -> Result<(Model, LlamaTokenizer, Option>), LlamaSourceError> { let tokenizer = parse_external_tokenizer(tokenizer_source)?; let config = parse_external_config(config_bytes)?; if model_bytes.is_empty() { @@ -885,7 +1007,16 @@ where override_chat_template, config, )?; - Ok((model, tokenizer)) + let mtp = match mtp_model_bytes { + Some(bytes) => { + let mut cursor = std::io::Cursor::new(bytes); + let metadata = GgufMetadata::read(&mut cursor)?; + let mut mtp_source = ShardedVarBuilder::new(vec![(metadata, cursor)]); + Some(Gemma4MtpAssistant::from_gguf(&mut mtp_source, &device)?) + } + None => None, + }; + Ok((model, tokenizer, mtp)) }; load_model()? @@ -894,6 +1025,7 @@ where Ok(Self { model, + mtp, tokenizer: Arc::new(tokenizer), device, }) diff --git a/models/kalosm-llama/src/raw/attention_layer.rs b/models/kalosm-llama/src/raw/attention_layer.rs index 104a3721b..69825fc5d 100644 --- a/models/kalosm-llama/src/raw/attention_layer.rs +++ b/models/kalosm-llama/src/raw/attention_layer.rs @@ -16,6 +16,12 @@ pub enum FeedForwardVariant { Phi(PhiFeedForward), } +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum FeedForwardActivation { + Silu, + Gelu, +} + impl FeedForwardVariant where F: CastTo + CastTensor, @@ -83,6 +89,7 @@ impl PhiFeedForward { } pub struct LlamaFeedForward { + activation: FeedForwardActivation, gate: QMatrix, gate_up: Option, gate_bias: Option>, @@ -93,9 +100,15 @@ pub struct LlamaFeedForward { } impl LlamaFeedForward { - pub(crate) fn new(gate: QMatrix, down: QMatrix, up: QMatrix) -> Self { + pub(crate) fn new_with_activation( + gate: QMatrix, + down: QMatrix, + up: QMatrix, + activation: FeedForwardActivation, + ) -> Self { let gate_up = QMatrix::concat_rows(&[&gate, &up]); Self { + activation, gate, gate_up, down, @@ -117,6 +130,7 @@ impl LlamaFeedForward { ) -> Self { let gate_up = QMatrix::concat_rows(&[&gate, &up]); Self { + activation: FeedForwardActivation::Silu, gate, gate_up, gate_bias, @@ -192,7 +206,7 @@ impl LlamaFeedForward { let up = projected .narrow(D::Minus1, pair_len, pair_len) .to_concrete(); - (gate.silu() * up).to_concrete() + (self.activate(gate) * up).to_concrete() } Some(gate_up) => { let gate_width = self.gate.shape()[0]; @@ -215,7 +229,7 @@ impl LlamaFeedForward { up_states = up_states.add_(&bias_f32); } - (gate_states.silu() * up_states).to_concrete() + (self.activate(gate_states) * up_states).to_concrete() } None => { let mut w1 = x_f32.q_mat_mul(&self.gate); @@ -223,7 +237,7 @@ impl LlamaFeedForward { let bias_f32: Tensor<1, f32> = bias.cast(); w1 = w1.add_(&bias_f32); } - let w1 = w1.silu(); + let w1 = self.activate(w1); let mut w3 = x_f32.q_mat_mul(&self.up); if let Some(ref bias) = self.up_bias { @@ -235,6 +249,13 @@ impl LlamaFeedForward { } } } + + fn activate(&self, x: Tensor<3, f32>) -> Tensor<3, f32> { + match self.activation { + FeedForwardActivation::Silu => x.silu(), + FeedForwardActivation::Gelu => x.gelu(), + } + } } pub enum AttentionVariant { @@ -265,9 +286,10 @@ pub struct SeparateAttention { pub attention_wq: QMatrix, pub attention_qkv: Option, pub attention_q_norm: Option>, - pub attention_wk: QMatrix, + pub attention_wk: Option, pub attention_k_norm: Option>, - pub attention_wv: QMatrix, + pub attention_wv: Option, + pub attention_v_norm: Option>, pub bias: Option>, pub interleaved_rope: bool, } @@ -343,11 +365,16 @@ where .narrow(D::Minus1, query_width + key_width, value_width) .to_concrete(); - value_states + let value_states: Tensor<4, F> = value_states .reshape([b_sz, seq_len, num_key_value_heads, head_dim]) .transpose(1, 2) .to_concrete() - .cast() + .cast(); + if let Some(norm) = &self.attention_v_norm { + norm.forward_generic_4d(&value_states) + } else { + value_states + } }; let (query_states, key_states) = rope_cache.forward( @@ -381,7 +408,11 @@ where } }; let key_states: Tensor<4, F> = { - let mut key_states = hidden_f32.q_mat_mul(&self.attention_wk); + let attention_wk = self + .attention_wk + .as_ref() + .expect("separate attention without K weights must use a shared KV cache"); + let mut key_states = hidden_f32.q_mat_mul(attention_wk); if let Some(bias) = &self.bias { let bias_f32: Tensor<1, f32> = bias.bias_k.cast(); @@ -401,18 +432,27 @@ where } }; let value_states: Tensor<4, F> = { - let mut value_states = hidden_f32.q_mat_mul(&self.attention_wv); + let attention_wv = self + .attention_wv + .as_ref() + .expect("separate attention without V weights must use a shared KV cache"); + let mut value_states = hidden_f32.q_mat_mul(attention_wv); if let Some(bias) = &self.bias { let bias_f32: Tensor<1, f32> = bias.bias_v.cast(); value_states = value_states.add_(&bias_f32); } - value_states + let value_states: Tensor<4, F> = value_states .reshape([b_sz, seq_len, num_key_value_heads, head_dim]) .transpose(1, 2) .to_concrete() - .cast() + .cast(); + if let Some(norm) = &self.attention_v_norm { + norm.forward_generic_4d(&value_states) + } else { + value_states + } }; let (query_states, key_states) = rope_cache.forward( @@ -424,6 +464,68 @@ where ); (query_states, key_states, value_states) } + + #[allow(clippy::too_many_arguments)] + fn forward_query( + &self, + num_heads: usize, + head_dim: usize, + hidden_states: &Tensor<3, F, B>, + rope_cache: &RopeImplementation, + start_pos: usize, + pos_ids: Option<&Tensor<2, F>>, + ) -> Tensor<4, F> + where + B: Fusion<3, F>, + { + let [b_sz, seq_len, _] = hidden_states.shape(); + let hidden_f32 = hidden_states.cast::(); + + let query_states: Tensor<4, F> = if let Some(attention_qkv) = &self.attention_qkv { + let query_width = num_heads * head_dim; + let mut qkv = hidden_f32.q_mat_mul(attention_qkv); + if let Some(bias) = &self.bias { + let bias_f32: Tensor<1, f32> = bias.bias_qkv.cast(); + qkv = qkv.add_(&bias_f32); + } + let query_states = qkv.narrow(D::Minus1, 0, query_width).to_concrete(); + let query = query_states + .reshape([b_sz, seq_len, num_heads, head_dim]) + .transpose(1, 2) + .to_concrete(); + let query: Tensor<4, F> = query.cast(); + if let Some(norm) = &self.attention_q_norm { + norm.forward_generic_4d(&query) + } else { + query + } + } else { + let mut query_states = hidden_f32.q_mat_mul(&self.attention_wq); + if let Some(bias) = &self.bias { + let bias_f32: Tensor<1, f32> = bias.bias_q.cast(); + query_states = query_states.add_(&bias_f32); + } + let query = query_states + .reshape([b_sz, seq_len, num_heads, head_dim]) + .transpose(1, 2) + .to_concrete(); + let query: Tensor<4, F> = query.cast(); + if let Some(norm) = &self.attention_q_norm { + norm.forward_generic_4d(&query) + } else { + query + } + }; + + let (query_states, _) = rope_cache.forward( + &query_states, + &query_states, + start_pos, + pos_ids, + self.interleaved_rope, + ); + query_states + } } pub struct GroupedAttention { @@ -504,6 +606,12 @@ pub struct LlamaAttention { pub hidden_size: usize, pub rope_cache: RopeImplementation, pub(crate) sliding_window_size: Option, + pub(crate) attention_scale: f32, + pub(crate) shared_kv_layer: Option, + pub(crate) per_layer_inp_gate: Option, + pub(crate) per_layer_proj: Option, + pub(crate) per_layer_post_norm: Option>, + pub(crate) layer_output_scale: Option>, } impl LlamaAttention @@ -511,6 +619,29 @@ where F: CastTo + CastTensor, f32: CastTo + CastTensor, { + fn logical_kv_len(&self, start_pos: usize, q_len: usize) -> usize { + let len = start_pos + q_len; + self.sliding_window_size + .map(|window| len.min(window)) + .unwrap_or(len) + } + + fn narrow_kv_to_logical_len( + &self, + key_states: Tensor<4, f32>, + value_states: Tensor<4, f32>, + logical_kv_len: usize, + ) -> (Tensor<4, f32>, Tensor<4, f32>) { + if key_states.shape()[2] <= logical_kv_len { + return (key_states, value_states); + } + + ( + key_states.narrow(2, 0, logical_kv_len).to_concrete(), + value_states.narrow(2, 0, logical_kv_len).to_concrete(), + ) + } + pub(crate) fn forward( &self, hidden_states: &Tensor<3, F, B>, @@ -558,6 +689,11 @@ where None => (key_f32, value_f32), Some(cache) => cache.append(&query_f32.device(), &key_f32, &value_f32), }; + let (key_f32, value_f32) = self.narrow_kv_to_logical_len( + key_f32, + value_f32, + self.logical_kv_len(start_pos, q_len), + ); forward_attention_qkv_f32( &query_f32, @@ -565,10 +701,56 @@ where &value_f32, &self.attention_wo, attention_mask, - head_dim, b_sz, q_len, hidden_size, + self.attention_scale, + ) + } + + pub(crate) fn forward_with_shared_kv( + &self, + hidden_states: &Tensor<3, F, B>, + attention_mask: Option<&AttentionMask>, + start_pos: usize, + pos_ids: Option<&Tensor<2, F>>, + key_states: &Tensor<4, f32>, + value_states: &Tensor<4, f32>, + ) -> Tensor<3, F> + where + B: Fusion<3, F>, + { + let [b_sz, q_len, _] = hidden_states.shape(); + let query_states = match self.attention_variant { + AttentionVariant::Separate(ref attention) => attention.forward_query( + self.n_head, + self.head_dim, + hidden_states, + &self.rope_cache, + start_pos, + pos_ids, + ), + AttentionVariant::Grouped(_) => { + panic!("grouped attention cannot reuse a shared KV cache") + } + }; + let query_f32: Tensor<4, f32> = query_states.cast(); + let (key_states, value_states) = self.narrow_kv_to_logical_len( + key_states.clone(), + value_states.clone(), + self.logical_kv_len(start_pos, q_len), + ); + + forward_attention_qkv_f32( + &query_f32, + &key_states, + &value_states, + &self.attention_wo, + attention_mask, + b_sz, + q_len, + self.hidden_size, + self.attention_scale, ) } @@ -624,15 +806,23 @@ where None => (key_f32, value_f32), Some(cache) => cache.append(&query_f32.device(), &key_f32, &value_f32), }; + let (key_f32, value_f32) = self.narrow_kv_to_logical_len( + key_f32, + value_f32, + self.logical_kv_len(start_pos, q_len), + ); crate::raw::debug_check_nan_f32(&key_f32, layer_idx, "K_cache_view", start_pos); crate::raw::debug_check_nan_f32(&value_f32, layer_idx, "V_cache_view", start_pos); - let scale = 1. / (head_dim as f64).sqrt(); + let scale = self.attention_scale; + let padded_attention_mask = + attention_mask.and_then(|m| pad_attention_mask_to_kv_len(m, key_f32.shape()[2])); + let attention_mask = padded_attention_mask.as_ref().or(attention_mask); let attn_raw = query_f32.flash_attention( &key_f32, &value_f32, - scale as f32, + scale, attention_mask.map(|m| { let kind = if m.is_strict_causal() { fusor::MaskKind::Causal @@ -653,6 +843,28 @@ where } } +fn pad_attention_mask_to_kv_len( + attention_mask: &AttentionMask, + kv_seq_len: usize, +) -> Option> { + let mask = attention_mask.mask(); + let [rows, cols] = mask.shape(); + if cols == kv_seq_len { + return None; + } + + if cols > kv_seq_len { + return Some(AttentionMask::new( + mask.narrow(1, 0, kv_seq_len).to_concrete(), + )); + } + + let padded = Tensor::full(&mask.device(), [rows, kv_seq_len], f32::NEG_INFINITY); + Some(AttentionMask::new( + padded.slice_assign([0..rows, 0..cols], mask), + )) +} + /// Forward attention QKV computation in f32 for SIMD compatibility. /// All intermediate computation happens in f32, with the final result cast back to F. #[allow(clippy::too_many_arguments)] @@ -662,20 +874,22 @@ pub(crate) fn forward_attention_qkv_f32( value_states: &Tensor<4, f32>, attention_wo: &Linear, attention_mask: Option<&AttentionMask>, - head_dim: usize, b_sz: usize, q_len: usize, hidden_size: usize, + attention_scale: f32, ) -> Tensor<3, F> where F: FloatDataType + SimdElement + Default + CastTo + CastTensor, f32: CastTo + CastTensor, { - let scale = 1. / (head_dim as f64).sqrt(); + let padded_attention_mask = + attention_mask.and_then(|m| pad_attention_mask_to_kv_len(m, key_states.shape()[2])); + let attention_mask = padded_attention_mask.as_ref().or(attention_mask); let attn_output = query_states.flash_attention( key_states, value_states, - scale as f32, + attention_scale, attention_mask.map(|m| { let kind = if m.is_strict_causal() { fusor::MaskKind::Causal diff --git a/models/kalosm-llama/src/raw/cache.rs b/models/kalosm-llama/src/raw/cache.rs index dd26510d6..94627496f 100644 --- a/models/kalosm-llama/src/raw/cache.rs +++ b/models/kalosm-llama/src/raw/cache.rs @@ -24,7 +24,13 @@ impl LlamaCache { let max_seq_len = config.context_length; let mut blocks = Vec::with_capacity(config.n_layer); for layer_idx in 0..config.n_layer { - let max_seq_len = if let (Some(sliding_window_type), Some(sliding_window_size)) = + let max_seq_len = if let Some(layer_window_size) = config + .layer_sliding_window_sizes + .as_ref() + .and_then(|sizes| sizes.get(layer_idx).copied().flatten()) + { + layer_window_size + } else if let (Some(sliding_window_type), Some(sliding_window_size)) = (config.sliding_window_type, config.sliding_window_size) { let is_sliding = (layer_idx + 1) % sliding_window_type != 0; diff --git a/models/kalosm-llama/src/raw/mod.rs b/models/kalosm-llama/src/raw/mod.rs index c2d8bbfe3..fb9063b61 100644 --- a/models/kalosm-llama/src/raw/mod.rs +++ b/models/kalosm-llama/src/raw/mod.rs @@ -1,4 +1,6 @@ +use std::collections::HashMap; use std::future::Future; +use std::ops::Range; use std::pin::Pin; use std::sync::Arc; #[cfg(not(target_arch = "wasm32"))] @@ -57,6 +59,68 @@ pub(crate) fn debug_check_nan_f32( } } +#[cfg(feature = "vision")] +pub(crate) fn debug_tensor_stats_f32(t: &fusor::Tensor, label: &str) { + #[cfg(target_arch = "wasm32")] + { + let _ = (t, label); + return; + } + #[cfg(not(target_arch = "wasm32"))] + { + if std::env::var_os("KALOSM_TRACE_VISION_STATS").is_none() { + return; + } + let Ok(slice) = pollster::block_on(t.as_slice()) else { + tracing::warn!( + "[vision_stats] {label} shape={:?} readback failed", + t.shape() + ); + return; + }; + let values = slice.as_slice(); + let mut nan = 0usize; + let mut pos_inf = 0usize; + let mut neg_inf = 0usize; + let mut min = f32::INFINITY; + let mut max = f32::NEG_INFINITY; + let mut sum = 0.0f64; + for value in values.iter().copied() { + if value.is_nan() { + nan += 1; + } else if value == f32::INFINITY { + pos_inf += 1; + } else if value == f32::NEG_INFINITY { + neg_inf += 1; + } else { + min = min.min(value); + max = max.max(value); + sum += value as f64; + } + } + let finite = values.len().saturating_sub(nan + pos_inf + neg_inf); + let mean = if finite == 0 { + 0.0 + } else { + (sum / finite as f64) as f32 + }; + let sample = values.iter().copied().take(8).collect::>(); + tracing::info!( + "[vision_stats] {label} shape={:?} len={} finite={} nan={} +inf={} -inf={} min={} max={} mean={} sample={:?}", + t.shape(), + values.len(), + finite, + nan, + pos_inf, + neg_inf, + min, + max, + mean, + sample + ); + } +} + #[cfg(not(feature = "vision"))] pub(crate) fn debug_check_nan_f32( _: &fusor::Tensor, @@ -66,18 +130,61 @@ pub(crate) fn debug_check_nan_f32( ) { } +#[cfg(not(feature = "vision"))] +pub(crate) fn debug_tensor_stats_f32(_: &fusor::Tensor, _: &str) {} + +#[cfg(not(target_arch = "wasm32"))] +fn resolve_intermediate_hidden_f32(tensor: &fusor::Tensor<2, f32>) { + let marker = tensor.clone().mul_scalar(1.0).to_concrete(); + if let Some(gpu_marker) = marker.as_gpu() { + gpu_marker.materialize_sync(); + } else { + std::mem::drop(marker.to_concrete()); + } +} + +#[cfg(target_arch = "wasm32")] +fn resolve_intermediate_hidden_f32(_: &fusor::Tensor<2, f32>) {} + +#[cfg(all(feature = "vision", not(target_arch = "wasm32")))] +fn copy_image_embeddings_to_device( + tensor: Tensor<2, f32>, + device: &Device, +) -> Result> { + let shape = tensor.shape(); + let expected_len = shape.iter().product::(); + let values = pollster::block_on(tensor.as_slice())?; + let values = values.as_slice(); + if values.len() != expected_len { + return Err(fusor::Error::msg(format!( + "Image embedding transfer expected {expected_len} values for shape {shape:?}, got {}", + values.len() + ))); + } + Ok(Tensor::from_slice(device, shape, values)) +} + +#[cfg(all(feature = "vision", target_arch = "wasm32"))] +fn copy_image_embeddings_to_device( + tensor: Tensor<2, f32>, + _device: &Device, +) -> Result> { + Ok(tensor) +} + use crate::chat_template::HuggingFaceChatTemplate; use crate::raw::attention_layer::LlamaAttention; use crate::raw::rope::RopeImplementation; use crate::LlamaSourceError; use attention_layer::AttentionBias; use attention_layer::AttentionVariant; +use attention_layer::FeedForwardActivation; use attention_layer::FeedForwardVariant; use attention_layer::GroupedAttention; use attention_layer::LlamaFeedForward; use attention_layer::PhiFeedForward; use attention_layer::SeparateAttention; -use fusor::cache::MaskCache; +use fusor::cache::{AttentionMask, MaskCache}; use fusor::layers::Embedding; use fusor::layers::Linear; use fusor::layers::RmsNorm; @@ -94,12 +201,14 @@ use fusor_gguf::GgufValue; mod attention_layer; pub mod cache; +mod mtp; mod rope; #[cfg(feature = "vision")] mod vision; use crate::LlamaImage; use cache::LlamaCache; +pub(crate) use mtp::Gemma4MtpAssistant; pub const DEFAULT_ROPE_FREQUENCY: f32 = 1_000_000.; pub const GEMMA_DEFAULT_SLIDING_WINDOW_TYPE: usize = 6; @@ -111,30 +220,32 @@ pub struct LlamaConfig { pub(crate) rope_theta: f32, pub(crate) context_length: usize, pub(crate) head_dimension: usize, - n_head: usize, pub(crate) n_layer: usize, pub(crate) start_token_string: String, - pub(crate) stop_token: u32, + pub(crate) stop_tokens: Vec, pub(crate) stop_token_string: String, pub(crate) chat_template: Option, pub(crate) rope_scaling: Option, pub(crate) sliding_window_type: Option, pub(crate) sliding_window_size: Option, + pub(crate) layer_sliding_window_sizes: Option>>, + pub(crate) final_logit_softcapping: Option, + pub(crate) per_layer_embedding_length: Option, #[cfg_attr(not(feature = "vision"), allow(dead_code))] pub(crate) vision_start_token: Option, pub(crate) _vision_end_token: Option, #[cfg_attr(not(feature = "vision"), allow(dead_code))] pub(crate) image_pad_token: Option, #[cfg_attr(not(feature = "vision"), allow(dead_code))] + pub(crate) image_start_token: Option, + #[cfg_attr(not(feature = "vision"), allow(dead_code))] + pub(crate) image_end_token: Option, + #[cfg_attr(not(feature = "vision"), allow(dead_code))] pub(crate) video_pad_token: Option, pub(crate) mrope_sections: Option>, } impl LlamaConfig { - fn hidden_size(&self) -> usize { - self.head_dimension * self.n_head - } - #[cfg(test)] pub(crate) fn mock_test() -> Self { Self { @@ -142,18 +253,22 @@ impl LlamaConfig { rope_theta: 5000., context_length: 6, head_dimension: 2, - n_head: 0, n_layer: 0, start_token_string: "<|startoftext|>".to_string(), - stop_token: 0, + stop_tokens: vec![0], stop_token_string: "<|endoftext|>".to_string(), sliding_window_type: None, sliding_window_size: None, + layer_sliding_window_sizes: None, + final_logit_softcapping: None, + per_layer_embedding_length: None, chat_template: None, rope_scaling: None, vision_start_token: None, _vision_end_token: None, image_pad_token: None, + image_start_token: None, + image_end_token: None, video_pad_token: None, mrope_sections: None, } @@ -171,9 +286,12 @@ pub struct RopeScalingConfig { pub struct Model { pub(crate) config: Arc>, #[cfg(feature = "vision")] - vision_encoder: Option>, + vision_encoder: Option>, tok_embeddings: Embedding, tok_embedding_scale: Option, + per_layer_tok_embeddings: Option>, + per_layer_model_proj: Option, + per_layer_proj_norm: Option>, layers: Vec>, norm: RmsNorm<1, F>, output: QMatrix, @@ -181,13 +299,25 @@ pub struct Model { masks: MaskCache, } +pub(crate) struct TargetBatchOutput { + pub(crate) logits: Tensor<2, f32>, + pub(crate) h_nextn: Tensor<2, f32>, +} + +struct PreNormForwardOutput { + hidden: Tensor<3, F>, + seq_len: usize, +} + /// The embedded token inputs produced by [`Model::encode_tokens`], ready to be /// run through the transformer layers. pub(crate) struct EncodedTokens { embeddings: Tensor<3, F>, + per_layer_inputs: Option>, seq_len: usize, index_pos: usize, pos_ids: Option>, + non_causal_token_ranges: Vec>, } pub(crate) trait LlamaVarSource { @@ -342,17 +472,22 @@ where .get("tokenizer.ggml.bos_token_id") .ok() .and_then(|v| v.try_into().ok()); + let eos_token: u32 = source + .get("tokenizer.ggml.eos_token_id")? + .clone() + .try_into()?; let stop_token = if let Some(override_stop_token_string) = override_stop_token_string { tokens .iter() .position(|v| **v == override_stop_token_string) .unwrap_or(0) as u32 } else { - source - .get("tokenizer.ggml.eos_token_id")? - .clone() - .try_into()? + eos_token }; + let mut stop_tokens = vec![stop_token]; + if eos_token != stop_token { + stop_tokens.push(eos_token); + } let start_token_string = start_token .map(|v| tokens[v as usize].to_string()) .unwrap_or_default(); @@ -375,12 +510,22 @@ where // Parameter extraction from metadata. let architecture = source.get("general.architecture")?.to_string()?.clone(); + let is_gemma4 = architecture.as_ref() == "gemma4"; let head_count = source.get(".attention.head_count")?.to_u32()? as usize; let head_count_kv = source.get(".attention.head_count_kv")?.to_u32()? as usize; let block_count = source.get(".block_count")?.to_u32()? as usize; let embedding_length = source.get(".embedding_length")?.to_u32()? as usize; + let per_layer_embedding_length = source + .get(".embedding_length_per_layer_input") + .and_then(|m| Ok(m.to_u32()?)) + .ok() + .map(|x| x as usize); // Strangely this value is generally 1e-6 in GGUF file but used to be 1e-5 by default. let rms_norm_eps = source.get(".attention.layer_norm_rms_epsilon")?.to_f32()? as f64; + let final_logit_softcapping = source + .get(".final_logit_softcapping") + .and_then(|m| Ok(m.to_f32()?)) + .ok(); let rope_freq_base = source .get(".rope.freq_base") @@ -399,20 +544,68 @@ where .or_else(|| (&*architecture == "gemma3").then_some(GEMMA_DEFAULT_SLIDING_WINDOW_TYPE)); let rope_freq_base_sliding = source - .get(".rope.local_freq_base") + .get(".rope.freq_base_swa") .and_then(|m| Ok(m.to_f32()?)) .ok() .or_else(|| { - (&*architecture == "gemma3").then_some(GEMMA_DEFAULT_ROPE_FREQUENCY_SLIDING) + source + .get(".rope.local_freq_base") + .and_then(|m| Ok(m.to_f32()?)) + .ok() + }) + .or_else(|| { + (&*architecture == "gemma3" || is_gemma4) + .then_some(GEMMA_DEFAULT_ROPE_FREQUENCY_SLIDING) }); - let context_length = source.get(".context_length")?.to_u32()? as usize; + let sliding_window_pattern = source + .get(".attention.sliding_window_pattern") + .ok() + .and_then(|m| { + let values = m.to_array().ok()?; + values + .iter() + .map(|value| value.to_bool().ok()) + .collect::>>() + }); + let layer_is_sliding: Vec = if let Some(pattern) = sliding_window_pattern { + pattern + } else if let Some(sliding_window_type) = sliding_window_type { + (0..block_count) + .map(|layer_idx| (layer_idx + 1) % sliding_window_type != 0) + .collect() + } else { + vec![false; block_count] + }; + let layer_sliding_window_sizes = sliding_window_size.map(|sliding_window_size| { + layer_is_sliding + .iter() + .map(|is_sliding| is_sliding.then_some(sliding_window_size)) + .collect::>() + }); + + let shared_kv_layers = source + .get(".attention.shared_kv_layers") + .and_then(|m| Ok(m.to_u32()?)) + .ok() + .map(|x| x as usize) + .unwrap_or_default(); + let n_layer_kv_from_start = block_count.saturating_sub(shared_kv_layers); + let head_dim = source .get(".attention.key_length") .and_then(|v| Ok(v.to_u32()?)) .ok() .map(|x| x as usize) .unwrap_or_else(|| embedding_length / head_count); + let head_dim_swa = source + .get(".attention.key_length_swa") + .and_then(|v| Ok(v.to_u32()?)) + .ok() + .map(|x| x as usize) + .unwrap_or(head_dim); + + let context_length = source.get(".context_length")?.to_u32()? as usize; let rope_freq_weight: Option> = source .tensor("rope_freqs.weight", device) @@ -425,15 +618,17 @@ where rope_theta: rope_freq_base, context_length, head_dimension: head_dim, - n_head: head_count, n_layer: block_count, start_token_string, - stop_token, + stop_tokens, stop_token_string, chat_template, rope_scaling, sliding_window_type, sliding_window_size, + layer_sliding_window_sizes, + final_logit_softcapping, + per_layer_embedding_length, vision_start_token: tokens .iter() .position(|v| &**v == "<|vision_start|>") @@ -445,6 +640,15 @@ where image_pad_token: tokens .iter() .position(|v| &**v == "<|image_pad|>") + .or_else(|| tokens.iter().position(|v| &**v == "<|image|>")) + .map(|v| v as u32), + image_start_token: tokens + .iter() + .position(|v| &**v == "<|image>") + .map(|v| v as u32), + image_end_token: tokens + .iter() + .position(|v| &**v == "") .map(|v| v as u32), video_pad_token: tokens .iter() @@ -462,19 +666,48 @@ where }; let config = Arc::new(config); - let rope: RopeImplementation = - rope::RopeImplementation::new(&config, config.rope_theta, device)?; + let rope: RopeImplementation = rope::RopeImplementation::new_with_head_dimension( + &config, + head_dim, + config.rope_freq_weight.as_ref(), + config.rope_theta, + device, + )?; let sliding_rope: Option> = rope_freq_base_sliding + .filter(|_| layer_is_sliding.iter().any(|is_sliding| *is_sliding)) .map(|rope_freq_base_sliding| { - RopeImplementation::new(&config, rope_freq_base_sliding, device) + let rope_freq_weight = (!is_gemma4) + .then_some(config.rope_freq_weight.as_ref()) + .flatten(); + RopeImplementation::new_with_head_dimension( + &config, + head_dim_swa, + rope_freq_weight, + rope_freq_base_sliding, + device, + ) }) .transpose()?; let tok_embeddings_q = source.tensor("token_embd.weight", device).await?; let tok_embedding_scale = - (&*architecture == "gemma3").then(|| (embedding_length as f32).sqrt()); + (&*architecture == "gemma3" || is_gemma4).then(|| (embedding_length as f32).sqrt()); let tok_embeddings = Embedding::new(tok_embeddings_q.clone()); + let (per_layer_tok_embeddings, per_layer_model_proj, per_layer_proj_norm) = + if per_layer_embedding_length.is_some() { + let embeddings = source.tensor("per_layer_token_embd.weight", device).await?; + let model_proj = source.tensor("per_layer_model_proj.weight", device).await?; + let proj_norm = source.tensor("per_layer_proj_norm.weight", device).await?; + ( + Some(Embedding::new(embeddings)), + Some(model_proj), + Some(decode_norm(proj_norm, rms_norm_eps)?), + ) + } else { + (None, None, None) + }; + let norm = source.tensor("output_norm.weight", device).await?; let norm = decode_norm(norm, rms_norm_eps)?; let output = match source.tensor("output.weight", device).await { @@ -484,11 +717,43 @@ where tok_embeddings_q.clone() } }; + let mut layers = Vec::with_capacity(block_count); let interleaved_rope = architecture.as_ref() != "qwen2" && architecture.as_ref() != "qwen3" - && architecture.as_ref() != "gemma3"; + && architecture.as_ref() != "gemma3" + && !is_gemma4; + for layer_idx in 0..block_count { + let layer_is_sliding = layer_is_sliding.get(layer_idx).copied().unwrap_or(false); + let layer_head_dim = if is_gemma4 && layer_is_sliding { + head_dim_swa + } else { + head_dim + }; + let layer_attention_width = head_count * layer_head_dim; + let layer_sliding_window_size = config + .layer_sliding_window_sizes + .as_ref() + .and_then(|sizes| sizes.get(layer_idx).copied().flatten()); + + let has_kv = !is_gemma4 || layer_idx < n_layer_kv_from_start; + let shared_kv_layer = if is_gemma4 && !has_kv { + let offset = if layer_is_sliding { 2 } else { 1 }; + Some(n_layer_kv_from_start.saturating_sub(offset)) + } else { + None + }; + + let rope_cache = if layer_is_sliding { + sliding_rope + .as_ref() + .cloned() + .unwrap_or_else(|| rope.clone()) + } else { + rope.clone() + }; + let prefix = format!("blk.{layer_idx}"); let attention_variant = if let Ok(attention_qkv) = source .tensor(&format!("{prefix}.attn_qkv.weight"), device) @@ -502,13 +767,29 @@ where let q = source .tensor(&format!("{prefix}.attn_q.weight"), device) .await?; - let k = source - .tensor(&format!("{prefix}.attn_k.weight"), device) - .await?; - let v = source - .tensor(&format!("{prefix}.attn_v.weight"), device) - .await?; - let qkv = QMatrix::concat_rows(&[&q, &k, &v]); + let k = if has_kv { + Some( + source + .tensor(&format!("{prefix}.attn_k.weight"), device) + .await?, + ) + } else { + None + }; + let v = if has_kv { + Some( + source + .tensor(&format!("{prefix}.attn_v.weight"), device) + .await?, + ) + } else { + None + }; + let qkv = if let (Some(k), Some(v)) = (&k, &v) { + QMatrix::concat_rows(&[&q, k, v]) + } else { + None + }; let bias_q = source .tensor(&format!("{prefix}.attn_q.bias"), device) .await; @@ -535,6 +816,15 @@ where .tensor(&format!("{prefix}.attn_k_norm.weight"), device) .await .ok(); + let attention_v_norm = if is_gemma4 && has_kv { + Some(RmsNorm::new( + Tensor::ones(device, [layer_head_dim]), + None, + rms_norm_eps as f32, + )) + } else { + None + }; let separate = SeparateAttention { attention_wq: q, attention_qkv: qkv, @@ -546,6 +836,7 @@ where .map(|norm| decode_norm(norm, rms_norm_eps)) .transpose()?, attention_wv: v, + attention_v_norm, interleaved_rope, bias, }; @@ -566,10 +857,16 @@ where let feed_forward_w3 = source .tensor(&format!("{prefix}.ffn_up.weight"), device) .await?; - FeedForwardVariant::Llama(Box::new(LlamaFeedForward::new( + let activation = if is_gemma4 { + FeedForwardActivation::Gelu + } else { + FeedForwardActivation::Silu + }; + FeedForwardVariant::Llama(Box::new(LlamaFeedForward::new_with_activation( feed_forward_w1, feed_forward_w2, feed_forward_w3, + activation, ))) } else { // Otherwise, try to read from the up, and down weights @@ -603,26 +900,41 @@ where .await .ok(); - let mut layer_sliding_window_size = None; - - let rope_cache = if let ( - Some(rope_sliding), - Some(sliding_window_type), - Some(sliding_window_size), - ) = ( - sliding_rope.as_ref(), - sliding_window_type, - sliding_window_size, - ) { - let is_sliding = (layer_idx + 1) % sliding_window_type != 0; - if is_sliding { - layer_sliding_window_size = Some(sliding_window_size); - rope_sliding.clone() - } else { - rope.clone() - } + let per_layer_inp_gate = if per_layer_embedding_length.is_some() { + Some( + source + .tensor(&format!("{prefix}.inp_gate.weight"), device) + .await?, + ) } else { - rope.clone() + None + }; + let per_layer_proj = if per_layer_embedding_length.is_some() { + Some( + source + .tensor(&format!("{prefix}.proj.weight"), device) + .await?, + ) + } else { + None + }; + let per_layer_post_norm = if per_layer_embedding_length.is_some() { + let norm = source + .tensor(&format!("{prefix}.post_norm.weight"), device) + .await?; + Some(decode_norm(norm, rms_norm_eps)?) + } else { + None + }; + let layer_output_scale = source + .tensor(&format!("{prefix}.layer_output_scale.weight"), device) + .await + .ok() + .map(&dequantize_1d); + let attention_scale = if is_gemma4 { + 1.0 + } else { + 1.0 / (layer_head_dim as f32).sqrt() }; layers.push(LlamaAttention { @@ -639,10 +951,16 @@ where .transpose()?, n_head: head_count, n_kv_head: head_count_kv, - head_dim, - hidden_size: config.hidden_size(), + head_dim: layer_head_dim, + hidden_size: layer_attention_width, rope_cache, sliding_window_size: layer_sliding_window_size, + attention_scale, + shared_kv_layer, + per_layer_inp_gate, + per_layer_proj, + per_layer_post_norm, + layer_output_scale, }) } @@ -650,11 +968,11 @@ where #[cfg(feature = "vision")] let vision_encoder = if let (Some(vision_ct), Some(vision_bytes)) = (vision_ct, vision_bytes) { - Some(vision::QwenVisionTransformer::from_gguf( + Some(vision::VisionTransformer::from_gguf( vision_ct, &vision_bytes, device, - )) + )?) } else { None }; @@ -662,12 +980,15 @@ where config, tok_embeddings, tok_embedding_scale, + per_layer_tok_embeddings, + per_layer_model_proj, + per_layer_proj_norm, layers, norm, output, masks: Default::default(), #[cfg(feature = "vision")] - vision_encoder: vision_encoder.transpose()?, + vision_encoder, }) } } @@ -684,6 +1005,10 @@ where #[cfg(feature = "vision")] { self.vision_encoder.is_none() + || matches!( + self.vision_encoder, + Some(vision::VisionTransformer::Gemma(_)) + ) } #[cfg(not(feature = "vision"))] @@ -716,41 +1041,19 @@ where } } - // Add any image padding tokens to the tokens if needed - let tokens = if let (Some(image_pad_token), Some(vision_start_token), Some(vision)) = ( - self.config.image_pad_token, - self.config.vision_start_token, - &self.vision_encoder, - ) { - let mut tokens = Vec::new(); - let mut token_iter = raw_tokens.iter().copied(); - let mut image_iter = grid_thw.iter(); - while let Some(token) = token_iter.next() { - tokens.push(token); - let start_index = tokens.len(); - if token == vision_start_token { - match token_iter.next() { - Some(next) if next == image_pad_token => { - // Push a pad token for every image token - let grid = image_iter.next().ok_or_else(|| { - fusor::Error::msg( - "Image pad token found without matching image.", - ) - })?; - for _ in 0..grid.iter().product::() - / (vision.spacial_merge_size as u32).pow(2) - { - tokens.push(image_pad_token); - } - image_token_ranges.push(start_index..tokens.len()); - } - Some(next) => { - tokens.push(next); - } - None => break, - } - } - } + // Add image padding tokens for any placeholders in the prompt. + let tokens = if let (Some(image_pad_token), Some(vision)) = + (self.config.image_pad_token, &self.vision_encoder) + { + let (tokens, ranges) = vision.expand_image_tokens( + raw_tokens, + image_pad_token, + self.config.vision_start_token, + self.config.image_start_token, + self.config.image_end_token, + &grid_thw, + )?; + image_token_ranges = ranges; tokens } else { raw_tokens.to_vec() @@ -809,42 +1112,113 @@ where embeddings_f32 = (embeddings_f32 * scale).to_concrete(); } #[cfg(feature = "vision")] - let mut embeddings: Tensor<3, F> = embeddings_f32.cast(); - #[cfg(not(feature = "vision"))] - let embeddings: Tensor<3, F> = embeddings_f32.cast(); - #[cfg(feature = "vision")] let mut pos_ids = None; #[cfg(not(feature = "vision"))] let pos_ids = None; - #[cfg(feature = "vision")] - let batch_size = embeddings.shape()[0]; - #[cfg(feature = "vision")] - let embed_dim = embeddings.shape()[2]; #[cfg(feature = "vision")] if let Some(vision_encoder) = &self.vision_encoder { - for ((pixels, grid), range) in images.iter().zip(&grid_thw).zip(image_token_ranges) { + let batch_size = embeddings_f32.shape()[0]; + let embed_dim = embeddings_f32.shape()[2]; + for ((pixels, grid), range) in + images.iter().zip(&grid_thw).zip(image_token_ranges.iter()) + { let pixels_f: Tensor<2, F> = pixels.cast(); let image_embeds = vision_encoder.forward_image(&pixels_f, *grid)?; - let image_embeds_3d = image_embeds.unsqueeze(0); - embeddings = - embeddings.slice_assign([0..batch_size, range, 0..embed_dim], &image_embeds_3d); + let mut image_embeds_f32: Tensor<2, f32> = image_embeds.cast(); + if vision_encoder.outputs_on_isolated_device() { + image_embeds_f32 = copy_image_embeddings_to_device(image_embeds_f32, device)?; + } + debug_tensor_stats_f32(&image_embeds_f32, "image_embeds_projected"); + let image_embeds_3d: Tensor<3, f32> = image_embeds_f32.unsqueeze(0).to_concrete(); + embeddings_f32 = embeddings_f32.slice_assign( + [0..batch_size, range.clone(), 0..embed_dim], + &image_embeds_3d, + ); } - let (new_pos_ids, new_start_time) = - vision_encoder.get_rope_index(&tokens, &grid_thw, &self.config, start_time)?; - if let Some(cache) = cache.as_mut() { - cache.start_time = new_start_time; + if let Some((new_pos_ids, new_start_time)) = + vision_encoder.get_rope_index(&tokens, &grid_thw, &self.config, start_time)? + { + if let Some(cache) = cache.as_mut() { + cache.start_time = new_start_time; + } + let pos_f32: Tensor<2, f32> = new_pos_ids.cast(); + let pos_f: Tensor<2, F> = pos_f32.cast(); + pos_ids = Some(pos_f); } - let pos_f32: Tensor<2, f32> = new_pos_ids.cast(); - let pos_f: Tensor<2, F> = pos_f32.cast(); - pos_ids = Some(pos_f); } + let per_layer_inputs = if let ( + Some(per_layer_tok_embeddings), + Some(per_layer_model_proj), + Some(per_layer_proj_norm), + Some(per_layer_embedding_length), + ) = ( + &self.per_layer_tok_embeddings, + &self.per_layer_model_proj, + &self.per_layer_proj_norm, + self.config.per_layer_embedding_length, + ) { + let [batch, seq, embedding_dim] = embeddings_f32.shape(); + #[cfg(feature = "vision")] + let per_layer_x_base = { + let mut per_layer_tokens = tokens.clone(); + for range in &image_token_ranges { + per_layer_tokens[range.clone()].fill(0); + } + if let Some(image_start_token) = self.config.image_start_token { + for token in &mut per_layer_tokens { + if *token == image_start_token { + *token = 0; + } + } + } + if let Some(image_end_token) = self.config.image_end_token { + for token in &mut per_layer_tokens { + if *token == image_end_token { + *token = 0; + } + } + } + Tensor::new(device, per_layer_tokens.as_slice()) + }; + #[cfg(feature = "vision")] + let per_layer_x = per_layer_x_base.unsqueeze(0); + #[cfg(not(feature = "vision"))] + let per_layer_x = x.clone(); + let token_inputs = per_layer_tok_embeddings.forward::<2, 3, _>(&per_layer_x) + * (per_layer_embedding_length as f32).sqrt(); + let token_inputs = token_inputs + .reshape([batch, seq, self.config.n_layer, per_layer_embedding_length]) + .to_concrete(); + + let projected_inputs = embeddings_f32.q_mat_mul(per_layer_model_proj) + * (1.0 / (embedding_dim as f32).sqrt()); + let projected_inputs = projected_inputs + .reshape([batch, seq, self.config.n_layer, per_layer_embedding_length]) + .to_concrete(); + let projected_inputs: Tensor<4, F> = + per_layer_proj_norm.forward_generic_4d(&projected_inputs.cast()); + Some( + ((projected_inputs.cast::() + token_inputs) * (1.0 / 2.0_f32.sqrt())) + .to_concrete() + .cast(), + ) + } else { + None + }; + let embeddings: Tensor<3, F> = embeddings_f32.cast(); + Ok(EncodedTokens { embeddings, + per_layer_inputs, seq_len, index_pos, pos_ids, + #[cfg(feature = "vision")] + non_causal_token_ranges: image_token_ranges, + #[cfg(not(feature = "vision"))] + non_causal_token_ranges: Vec::new(), }) } @@ -860,10 +1234,275 @@ where f32: CastTo + CastTensor, { let x_f32 = self.forward_last_hidden_f32(tokens, images, device, cache)?; - let result_f32 = x_f32.q_mat_mul(&self.output); + self.forward_logits_from_hidden_f32(x_f32) + } + + fn forward_logits_from_hidden_f32(&self, x_f32: Tensor<2, f32>) -> Result> + where + f32: CastTo + CastTensor, + { + let mut result_f32 = x_f32.q_mat_mul(&self.output); + if let Some(softcap) = self.config.final_logit_softcapping { + result_f32 = result_f32 + .mul_scalar(1.0 / softcap) + .tanh() + .mul_scalar(softcap); + } Ok(result_f32.cast()) } + pub(crate) fn should_chunk_multimodal_prompt(&self, device: &Device) -> bool { + if std::env::var_os("KALOSM_LLAMA_DISABLE_MULTIMODAL_CHUNK").is_some() { + return false; + } + if std::env::var_os("KALOSM_LLAMA_FORCE_MULTIMODAL_CHUNK").is_some() { + return self.config.image_pad_token.is_some(); + } + #[cfg(feature = "vision")] + { + device.is_gpu() + && matches!( + self.vision_encoder, + Some(vision::VisionTransformer::Gemma(_)) + ) + && self.config.image_pad_token.is_some() + } + #[cfg(not(feature = "vision"))] + { + let _ = device; + false + } + } + + pub(crate) fn forward_chunked_multimodal( + &self, + tokens: &[u32], + images: &[LlamaImage], + device: &Device, + mut cache: Option<&mut LlamaCache>, + ) -> Result> + where + F: CastTo + CastTensor + Default, + f32: CastTo + CastTensor, + { + let Some(image_pad_token) = self.config.image_pad_token else { + return self.forward(tokens, images, device, cache); + }; + if images.is_empty() { + return self.forward(tokens, images, device, cache); + } + + let mut image_index = 0; + let mut segment_start = 0; + let mut text_prefix = Vec::new(); + let mut last_logits = None; + for (index, token) in tokens.iter().copied().enumerate() { + if token != image_pad_token || image_index >= images.len() { + continue; + } + let mut text_tokens = std::mem::take(&mut text_prefix); + text_tokens.extend_from_slice(&tokens[segment_start..index]); + if let Some(image_start_token) = self.config.image_start_token { + text_tokens.push(image_start_token); + } + if !text_tokens.is_empty() { + self.forward_text_chunk_for_multimodal( + &text_tokens, + device, + cache.as_deref_mut(), + false, + )?; + } + let image_hidden = self.forward_image_embeddings_only_hidden_f32( + &images[image_index..image_index + 1], + device, + cache.as_deref_mut(), + )?; + image_index += 1; + segment_start = index + 1; + if let Some(image_end_token) = self.config.image_end_token { + text_prefix.push(image_end_token); + } + if segment_start < tokens.len() || !text_prefix.is_empty() { + resolve_intermediate_hidden_f32(&image_hidden); + } else { + last_logits = Some(self.forward_logits_from_hidden_f32(image_hidden)?); + } + } + + if segment_start < tokens.len() || !text_prefix.is_empty() { + let mut text_tokens = text_prefix; + text_tokens.extend_from_slice(&tokens[segment_start..]); + last_logits = self.forward_text_chunk_for_multimodal( + &text_tokens, + device, + cache.as_deref_mut(), + true, + )?; + } + + last_logits.ok_or_else(|| fusor::Error::msg("No tokens to forward")) + } + + fn forward_text_chunk_for_multimodal( + &self, + text_tokens: &[u32], + device: &Device, + mut cache: Option<&mut LlamaCache>, + return_logits: bool, + ) -> Result>> + where + F: CastTo + CastTensor + Default, + f32: CastTo + CastTensor, + { + if text_tokens.is_empty() { + return Ok(None); + } + + let has_cache_prefix = cache + .as_ref() + .map(|cache| !cache.tokens.is_empty()) + .unwrap_or(false); + let incremental = device.is_gpu() + && has_cache_prefix + && text_tokens.len() > 1 + && std::env::var_os("KALOSM_LLAMA_ENABLE_MULTIMODAL_TEXT_INCREMENTAL").is_some(); + + if incremental { + let last = text_tokens.len() - 1; + for (index, token) in text_tokens.iter().copied().enumerate() { + let one = [token]; + if return_logits && index == last { + return Ok(Some(self.forward( + &one, + &[], + device, + cache.as_deref_mut(), + )?)); + } + let hidden = + self.forward_last_hidden_f32(&one, &[], device, cache.as_deref_mut())?; + resolve_intermediate_hidden_f32(&hidden); + } + return Ok(None); + } + + if return_logits { + Ok(Some(self.forward(text_tokens, &[], device, cache)?)) + } else { + let hidden = self.forward_last_hidden_f32(text_tokens, &[], device, cache)?; + resolve_intermediate_hidden_f32(&hidden); + Ok(None) + } + } + + #[cfg(feature = "vision")] + fn forward_image_embeddings_only_hidden_f32( + &self, + raw_images: &[LlamaImage], + device: &Device, + mut cache: Option<&mut LlamaCache>, + ) -> Result> + where + F: CastTo + CastTensor + Default, + f32: CastTo + CastTensor, + { + let vision_encoder = self + .vision_encoder + .as_ref() + .ok_or_else(|| fusor::Error::msg("Image chunk requires a vision encoder"))?; + let image_pad_token = self + .config + .image_pad_token + .ok_or_else(|| fusor::Error::msg("Image chunk requires an image token"))?; + let (image, hints) = raw_images + .first() + .ok_or_else(|| fusor::Error::msg("Image chunk requires an image"))?; + let t_encode = Instant::now(); + let (pixels, grid) = + vision_encoder.preprocess_image(image, hints.min_tokens(), hints.max_tokens())?; + let pixels_f: Tensor<2, F> = pixels.cast(); + let image_embeds = vision_encoder.forward_image(&pixels_f, grid)?; + let mut embeddings_f32: Tensor<2, f32> = image_embeds.cast(); + if vision_encoder.outputs_on_isolated_device() { + embeddings_f32 = copy_image_embeddings_to_device(embeddings_f32, device)?; + } + debug_tensor_stats_f32(&embeddings_f32, "image_embeds_projected"); + let seq_len = embeddings_f32.shape()[0]; + let embeddings_f32 = embeddings_f32.unsqueeze(0).to_concrete(); + + let index_pos = cache.as_ref().map(|c| c.tokens.len()).unwrap_or_default(); + if let Some(cache) = cache.as_mut() { + cache + .tokens + .extend(std::iter::repeat_n(image_pad_token, seq_len)); + } + + let per_layer_inputs = if let ( + Some(per_layer_tok_embeddings), + Some(per_layer_model_proj), + Some(per_layer_proj_norm), + Some(per_layer_embedding_length), + ) = ( + &self.per_layer_tok_embeddings, + &self.per_layer_model_proj, + &self.per_layer_proj_norm, + self.config.per_layer_embedding_length, + ) { + let [batch, seq, embedding_dim] = embeddings_f32.shape(); + let per_layer_tokens = [0u32]; + let per_layer_x_base = Tensor::new(device, per_layer_tokens.as_slice()); + let per_layer_x = per_layer_x_base.unsqueeze(0); + let token_inputs = per_layer_tok_embeddings.forward::<2, 3, _>(&per_layer_x) + * (per_layer_embedding_length as f32).sqrt(); + let token_inputs = token_inputs + .reshape([batch, 1, self.config.n_layer, per_layer_embedding_length]) + .broadcast_as([batch, seq, self.config.n_layer, per_layer_embedding_length]) + .to_concrete(); + + let projected_inputs = embeddings_f32.q_mat_mul(per_layer_model_proj) + * (1.0 / (embedding_dim as f32).sqrt()); + let projected_inputs = projected_inputs + .reshape([batch, seq, self.config.n_layer, per_layer_embedding_length]) + .to_concrete(); + let projected_inputs: Tensor<4, F> = + per_layer_proj_norm.forward_generic_4d(&projected_inputs.cast()); + Some( + ((projected_inputs.cast::() + token_inputs) * (1.0 / 2.0_f32.sqrt())) + .to_concrete() + .cast(), + ) + } else { + None + }; + + let encoded = EncodedTokens { + embeddings: embeddings_f32.cast(), + per_layer_inputs, + seq_len, + index_pos, + pos_ids: None, + non_causal_token_ranges: vec![0..seq_len], + }; + self.forward_last_hidden_from_embeddings(encoded, device, cache, Some(t_encode.elapsed())) + } + + #[cfg(not(feature = "vision"))] + fn forward_image_embeddings_only_hidden_f32( + &self, + _raw_images: &[LlamaImage], + _device: &Device, + _cache: Option<&mut LlamaCache>, + ) -> Result> + where + F: CastTo + CastTensor + Default, + f32: CastTo + CastTensor, + { + Err(fusor::Error::msg( + "Image chunks require the `vision` feature", + )) + } + pub(crate) fn forward_last_hidden_f32( &self, tokens: &[u32], @@ -880,6 +1519,35 @@ where self.forward_last_hidden_from_embeddings(encoded, device, cache, Some(t_encode.elapsed())) } + pub(crate) fn forward_logits_and_nextn_f32( + &self, + tokens: &[u32], + images: &[LlamaImage], + device: &Device, + mut cache: Option<&mut LlamaCache>, + ) -> Result + where + F: CastTo + CastTensor + Default, + f32: CastTo + CastTensor, + { + let t_encode = Instant::now(); + let encoded = self.encode_tokens(tokens, images, device, cache.as_deref_mut())?; + let pre_norm = self.forward_pre_norm_hidden_from_embeddings( + encoded, + device, + cache, + Some(t_encode.elapsed()), + )?; + let normed = self.norm.forward_generic(&pre_norm.hidden); + let h_nextn: Tensor<2, f32> = normed.clone().cast::().squeeze(0).to_concrete(); + let mut logits = normed.cast::().q_mat_mul(&self.output); + if let Some(softcap) = self.config.final_logit_softcapping { + logits = logits.mul_scalar(1.0 / softcap).tanh().mul_scalar(softcap); + } + let logits: Tensor<2, f32> = logits.squeeze(0).to_concrete(); + Ok(TargetBatchOutput { logits, h_nextn }) + } + pub(crate) fn forward_last_hidden_f32_gpu_token( &self, token: &Tensor<1, u32>, @@ -891,9 +1559,9 @@ where f32: CastTo + CastTensor, { #[cfg(feature = "vision")] - if self.vision_encoder.is_some() { + if !self.supports_gpu_token_run_ahead() { return Err(fusor::Error::msg( - "GPU token run-ahead is only available for text-only models", + "GPU token run-ahead is not available for this vision model", )); } @@ -910,34 +1578,149 @@ where if let Some(scale) = self.tok_embedding_scale { embeddings_f32 = (embeddings_f32 * scale).to_concrete(); } + let per_layer_inputs = if let ( + Some(per_layer_tok_embeddings), + Some(per_layer_model_proj), + Some(per_layer_proj_norm), + Some(per_layer_embedding_length), + ) = ( + &self.per_layer_tok_embeddings, + &self.per_layer_model_proj, + &self.per_layer_proj_norm, + self.config.per_layer_embedding_length, + ) { + let [batch, seq, embedding_dim] = embeddings_f32.shape(); + let token_inputs = per_layer_tok_embeddings.forward::<2, 3, _>(&x) + * (per_layer_embedding_length as f32).sqrt(); + let token_inputs = token_inputs + .reshape([batch, seq, self.config.n_layer, per_layer_embedding_length]) + .to_concrete(); + + let projected_inputs = embeddings_f32.q_mat_mul(per_layer_model_proj) + * (1.0 / (embedding_dim as f32).sqrt()); + let projected_inputs = projected_inputs + .reshape([batch, seq, self.config.n_layer, per_layer_embedding_length]) + .to_concrete(); + let projected_inputs: Tensor<4, F> = + per_layer_proj_norm.forward_generic_4d(&projected_inputs.cast()); + Some( + ((projected_inputs.cast::() + token_inputs) * (1.0 / 2.0_f32.sqrt())) + .to_concrete() + .cast(), + ) + } else { + None + }; let embeddings: Tensor<3, F> = embeddings_f32.cast(); let encoded = EncodedTokens { embeddings, + per_layer_inputs, seq_len: 1, index_pos: cache_slot, pos_ids: None, + non_causal_token_ranges: Vec::new(), }; let hidden = self.forward_last_hidden_from_embeddings(encoded, device, Some(cache), None)?; Ok((hidden, cache_slot)) } + fn get_attention_mask( + &self, + seq_len: usize, + index_pos: usize, + sliding_window_size: Option, + non_causal_token_ranges: &[Range], + device: &Device, + ) -> AttentionMask { + if non_causal_token_ranges.is_empty() { + return self + .masks + .get_mask(seq_len, index_pos, sliding_window_size, device); + } + + let cols = index_pos + seq_len; + let mut mask_data = vec![0.0_f32; seq_len * cols]; + for row in 0..seq_len { + let global_row = index_pos + row; + for col in 0..cols { + let same_non_causal_range = col >= index_pos + && non_causal_token_ranges + .iter() + .any(|range| range.contains(&row) && range.contains(&(col - index_pos))); + let future = col > global_row && !same_non_causal_range; + let outside_window = sliding_window_size + .map(|window| col + window <= global_row) + .unwrap_or(false); + if future || outside_window { + mask_data[row * cols + col] = f32::NEG_INFINITY; + } + } + } + let mask: Tensor<2, f32> = Tensor::new(device, mask_data.as_slice()) + .reshape([seq_len, cols]) + .to_concrete(); + AttentionMask::::new(mask) + } + + fn can_skip_attention_mask( + seq_len: usize, + index_pos: usize, + sliding_window_size: Option, + non_causal_token_ranges: &[Range], + ) -> bool { + if std::env::var_os("KALOSM_LLAMA_DISABLE_NON_CAUSAL_MASK_SKIP").is_some() { + return false; + } + let all_query_tokens_are_non_causal = non_causal_token_ranges + .iter() + .any(|range| range.start == 0 && range.end >= seq_len); + if !all_query_tokens_are_non_causal { + return false; + } + + match sliding_window_size { + Some(window) => window >= index_pos + seq_len, + None => true, + } + } + fn forward_last_hidden_from_embeddings( &self, encoded: EncodedTokens, device: &Device, - mut cache: Option<&mut LlamaCache>, + cache: Option<&mut LlamaCache>, encode_elapsed: Option, ) -> Result> + where + F: CastTo + CastTensor + Default, + f32: CastTo + CastTensor, + { + let output = + self.forward_pre_norm_hidden_from_embeddings(encoded, device, cache, encode_elapsed)?; + let x = output.hidden.i((.., output.seq_len - 1, ..)); + let x = self.norm.forward_generic_2d(&x); + Ok(x.cast::()) + } + + fn forward_pre_norm_hidden_from_embeddings( + &self, + encoded: EncodedTokens, + device: &Device, + mut cache: Option<&mut LlamaCache>, + encode_elapsed: Option, + ) -> Result> where F: CastTo + CastTensor + Default, f32: CastTo + CastTensor, { let EncodedTokens { embeddings: mut layer_in, + per_layer_inputs, seq_len, index_pos, pos_ids, + non_causal_token_ranges, } = encoded; let _trace_text_prefill = seq_len > 1 && std::env::var_os("KALOSM_TRACE_TEXT").is_some(); let trace_forward_timing = @@ -958,6 +1741,7 @@ where debug_check_nan_f32(&probe, usize::MAX, "embed", index_pos); } + let mut non_causal_masks: HashMap, AttentionMask> = HashMap::new(); for (i, layer) in self.layers.iter().enumerate() { let x = layer_in; let residual: Tensor<3, f32> = x.cast(); @@ -966,23 +1750,95 @@ where let probe: fusor::Tensor<3, f32> = x.clone().cast(); debug_check_nan_f32(&probe, i, "post_attn_norm", index_pos); } - let mask = (seq_len > 1).then(|| { - self.masks - .get_mask(seq_len, index_pos, layer.sliding_window_size, device) - }); + let mask = if seq_len > 1 { + if Self::can_skip_attention_mask( + seq_len, + index_pos, + layer.sliding_window_size, + &non_causal_token_ranges, + ) { + None + } else if non_causal_token_ranges.is_empty() { + Some(self.get_attention_mask( + seq_len, + index_pos, + layer.sliding_window_size, + &non_causal_token_ranges, + device, + )) + } else { + Some( + non_causal_masks + .entry(layer.sliding_window_size) + .or_insert_with(|| { + self.get_attention_mask( + seq_len, + index_pos, + layer.sliding_window_size, + &non_causal_token_ranges, + device, + ) + }) + .clone(), + ) + } + } else { + None + }; + let shared_kv = if let Some(shared_kv_layer) = layer.shared_kv_layer { + let cache_ref = cache.as_deref().ok_or_else(|| { + fusor::Error::msg("Gemma 4 shared KV attention requires a populated cache") + })?; + let key = cache_ref.blocks[shared_kv_layer] + .k() + .cloned() + .ok_or_else(|| { + fusor::Error::msg("Gemma 4 shared KV source key cache is empty") + })?; + let value = cache_ref.blocks[shared_kv_layer] + .v() + .cloned() + .ok_or_else(|| { + fusor::Error::msg("Gemma 4 shared KV source value cache is empty") + })?; + Some((key, value)) + } else { + None + }; let mut attn = { - #[cfg(feature = "vision")] - { - if trace_layer_nan { - layer.forward_with_trace( - &x, - mask.as_ref(), - index_pos, - pos_ids.as_ref(), - cache.as_mut().map(|c| &mut c.blocks[i]), - i, - ) - } else { + if let Some((shared_key, shared_value)) = shared_kv.as_ref() { + layer.forward_with_shared_kv( + &x, + mask.as_ref(), + index_pos, + pos_ids.as_ref(), + shared_key, + shared_value, + ) + } else { + #[cfg(feature = "vision")] + { + if trace_layer_nan { + layer.forward_with_trace( + &x, + mask.as_ref(), + index_pos, + pos_ids.as_ref(), + cache.as_mut().map(|c| &mut c.blocks[i]), + i, + ) + } else { + layer.forward( + &x, + mask.as_ref(), + index_pos, + pos_ids.as_ref(), + cache.as_mut().map(|c| &mut c.blocks[i]), + ) + } + } + #[cfg(not(feature = "vision"))] + { layer.forward( &x, mask.as_ref(), @@ -992,16 +1848,6 @@ where ) } } - #[cfg(not(feature = "vision"))] - { - layer.forward( - &x, - mask.as_ref(), - index_pos, - pos_ids.as_ref(), - cache.as_mut().map(|c| &mut c.blocks[i]), - ) - } }; if trace_layer_nan { let probe: fusor::Tensor<3, f32> = attn.clone().cast(); @@ -1021,19 +1867,52 @@ where .forward_add_residuals(&x, &attn_f32, &residual) { layer_in = layer_out; - if trace_layer_nan { - let probe: fusor::Tensor<3, f32> = layer_in.cast(); - debug_check_nan_f32(&probe, i, "ffn_fused", index_pos); - } - continue; + } else { + let x = layer.feed_forward_variant.forward(&x); + let x_f32: Tensor<3, f32> = x.cast(); + layer_in = (x_f32 + attn_f32 + residual).cast(); } + } else { + let mut x = layer.feed_forward_variant.forward(&x); + if let Some(post_ffn_norm) = &layer.post_ffn_norm { + x = post_ffn_norm.forward_generic(&x); + } + let x_f32: Tensor<3, f32> = x.cast(); + layer_in = (x_f32 + attn_f32 + residual).cast(); } - let mut x = layer.feed_forward_variant.forward(&x); - if let Some(post_ffn_norm) = &layer.post_ffn_norm { - x = post_ffn_norm.forward_generic(&x); + if let ( + Some(per_layer_inputs), + Some(per_layer_inp_gate), + Some(per_layer_proj), + Some(per_layer_post_norm), + ) = ( + per_layer_inputs.as_ref(), + layer.per_layer_inp_gate.as_ref(), + layer.per_layer_proj.as_ref(), + layer.per_layer_post_norm.as_ref(), + ) { + let pe_in: Tensor<3, f32> = layer_in.cast(); + let gate = pe_in.q_mat_mul(per_layer_inp_gate).gelu(); + let layer_input: Tensor<3, F> = per_layer_inputs + .narrow(2, i, 1) + .squeeze::<3>(2) + .to_concrete(); + let projected = (gate.cast::() * layer_input) + .to_concrete() + .cast::() + .q_mat_mul(per_layer_proj) + .cast(); + let projected = per_layer_post_norm.forward_generic(&projected); + let projected_f32: Tensor<3, f32> = projected.cast(); + layer_in = (pe_in + projected_f32).cast(); + } + if let Some(layer_output_scale) = &layer.layer_output_scale { + let scale = layer_output_scale + .reshape([1, 1, 1]) + .broadcast_as(layer_in.shape()) + .to_concrete(); + layer_in = (layer_in * scale).to_concrete(); } - let x_f32: Tensor<3, f32> = x.cast(); - layer_in = (x_f32 + attn_f32 + residual).cast(); if trace_layer_nan { let probe: fusor::Tensor<3, f32> = layer_in.cast(); debug_check_nan_f32(&probe, i, "ffn_unfused", index_pos); @@ -1042,10 +1921,10 @@ where if trace_forward_timing { tracing::info!("[timing] text layer loop: {:.2?}", t_text_layers.elapsed()); } - let x = self.norm.forward_generic(&layer_in); - let x = x.i((.., seq_len - 1, ..)); - let out = x.cast::(); - Ok(out) + Ok(PreNormForwardOutput { + hidden: layer_in, + seq_len, + }) } pub(crate) fn output_matrix(&self) -> &QMatrix { diff --git a/models/kalosm-llama/src/raw/mtp.rs b/models/kalosm-llama/src/raw/mtp.rs new file mode 100644 index 000000000..719d950f8 --- /dev/null +++ b/models/kalosm-llama/src/raw/mtp.rs @@ -0,0 +1,427 @@ +use super::*; + +pub(crate) struct Gemma4MtpAssistant { + config: Arc>, + pre_projection: QMatrix, + post_projection: QMatrix, + layers: Vec>, + norm: RmsNorm<1, F>, + output: QMatrix, + layer_is_sliding: Vec, +} + +pub(crate) struct Gemma4MtpStep { + pub(crate) logits: Tensor<1, f32>, + pub(crate) h_nextn: Tensor<2, f32>, +} + +impl Gemma4MtpAssistant +where + MulOp: SimdBinaryOp, + AddOp: SimdBinaryOp, + SumOp: SimdReduceOp, +{ + #[cfg(not(target_arch = "wasm32"))] + pub(crate) fn from_gguf( + source: &mut ShardedVarBuilder, + device: &Device, + ) -> std::result::Result + where + f32: CastTensor + CastTo, + F: CastTensor + CastTo, + { + super::block_on_ready(Self::from_var_source(source, device)) + } + + pub(crate) async fn from_var_source( + source: &mut S, + device: &Device, + ) -> std::result::Result + where + f32: CastTensor + CastTo, + F: CastTensor + CastTo, + { + let dequantize_1d = |qmatrix: QMatrix| -> Tensor<1, F> { + let shape = qmatrix.shape(); + if shape.len() == 1 { + let w1d: Tensor<1, f32> = qmatrix.dequantize(); + w1d.cast() + } else if shape.len() == 2 { + let w2d: Tensor<2, f32> = qmatrix.dequantize(); + w2d.reshape([w2d.shape()[0] * w2d.shape()[1]]) + .to_concrete() + .cast() + } else { + panic!( + "Expected 1D or 2D tensor for dequantize_1d, got {}D", + shape.len() + ) + } + }; + let decode_norm = |qmatrix: QMatrix, eps: f64| -> Result> { + let weight = dequantize_1d(qmatrix); + Ok(RmsNorm::new(weight, None, eps as f32)) + }; + + let architecture = source.get("general.architecture")?.to_string()?.clone(); + if architecture.as_ref() != "gemma4-assistant" { + return Err(fusor::Error::msg(format!( + "MTP assistant architecture must be gemma4-assistant, got {architecture}" + )) + .into()); + } + + let block_count = source.get(".block_count")?.to_u32()? as usize; + let context_length = source.get(".context_length")?.to_u32()? as usize; + let embedding_length = source.get(".embedding_length")?.to_u32()? as usize; + let target_embedding_length = source.get(".embedding_length_out")?.to_u32()? as usize; + let head_count = source.get(".attention.head_count")?.to_u32()? as usize; + let head_count_kv = source.get(".attention.head_count_kv")?.to_u32()? as usize; + let rms_norm_eps = source.get(".attention.layer_norm_rms_epsilon")?.to_f32()? as f64; + let rope_freq_base = source + .get(".rope.freq_base") + .and_then(|m| Ok(m.to_f32()?)) + .unwrap_or(DEFAULT_ROPE_FREQUENCY); + let rope_freq_base_sliding = source + .get(".rope.freq_base_swa") + .and_then(|m| Ok(m.to_f32()?)) + .ok() + .or_else(|| { + source + .get(".rope.local_freq_base") + .and_then(|m| Ok(m.to_f32()?)) + .ok() + }) + .unwrap_or(GEMMA_DEFAULT_ROPE_FREQUENCY_SLIDING); + let sliding_window_size = source + .get(".attention.sliding_window") + .and_then(|m| Ok(m.to_u32()?)) + .ok() + .map(|x| x as usize); + let head_dim = source + .get(".attention.key_length") + .and_then(|m| Ok(m.to_u32()?)) + .ok() + .map(|x| x as usize) + .unwrap_or_else(|| target_embedding_length / head_count); + let head_dim_swa = source + .get(".attention.key_length_swa") + .and_then(|m| Ok(m.to_u32()?)) + .ok() + .map(|x| x as usize) + .unwrap_or(head_dim); + let sliding_window_pattern = source + .get(".attention.sliding_window_pattern") + .ok() + .and_then(|m| { + let values = m.to_array().ok()?; + values + .iter() + .map(|value| value.to_bool().ok()) + .collect::>>() + }); + let layer_is_sliding = sliding_window_pattern.unwrap_or_else(|| { + (0..block_count) + .map(|idx| sliding_window_size.is_some() && idx + 1 < block_count) + .collect() + }); + let layer_sliding_window_sizes = sliding_window_size.map(|window| { + layer_is_sliding + .iter() + .map(|is_sliding| is_sliding.then_some(window)) + .collect::>() + }); + let rope_freq_weight: Option> = source + .tensor("rope_freqs.weight", device) + .await + .ok() + .map(&dequantize_1d); + + let config = Arc::new(LlamaConfig { + rope_freq_weight, + rope_theta: rope_freq_base, + context_length, + head_dimension: head_dim, + n_layer: block_count, + start_token_string: String::new(), + stop_tokens: Vec::new(), + stop_token_string: String::new(), + chat_template: None, + rope_scaling: None, + sliding_window_type: None, + sliding_window_size, + layer_sliding_window_sizes, + final_logit_softcapping: None, + per_layer_embedding_length: None, + vision_start_token: None, + _vision_end_token: None, + image_pad_token: None, + image_start_token: None, + image_end_token: None, + video_pad_token: None, + mrope_sections: None, + }); + + let rope = RopeImplementation::new_with_head_dimension( + &config, + head_dim, + config.rope_freq_weight.as_ref(), + rope_freq_base, + device, + )?; + let sliding_rope = RopeImplementation::new_with_head_dimension( + &config, + head_dim_swa, + None, + rope_freq_base_sliding, + device, + )?; + + let pre_projection = source.tensor("nextn.pre_projection.weight", device).await?; + let post_projection = source + .tensor("nextn.post_projection.weight", device) + .await?; + let output_norm = source.tensor("output_norm.weight", device).await?; + let norm = decode_norm(output_norm, rms_norm_eps)?; + let token_embd = source.tensor("token_embd.weight", device).await?; + let output = source + .tensor("output.weight", device) + .await + .unwrap_or_else(|_| token_embd.clone()); + + let mut layers = Vec::with_capacity(block_count); + for layer_idx in 0..block_count { + let layer_is_sliding = layer_is_sliding.get(layer_idx).copied().unwrap_or(false); + let layer_head_dim = if layer_is_sliding { + head_dim_swa + } else { + head_dim + }; + let layer_attention_width = head_count * layer_head_dim; + let layer_sliding_window_size = config + .layer_sliding_window_sizes + .as_ref() + .and_then(|sizes| sizes.get(layer_idx).copied().flatten()); + let rope_cache = if layer_is_sliding { + sliding_rope.clone() + } else { + rope.clone() + }; + let prefix = format!("blk.{layer_idx}"); + let q = source + .tensor(&format!("{prefix}.attn_q.weight"), device) + .await?; + let q_norm = source + .tensor(&format!("{prefix}.attn_q_norm.weight"), device) + .await + .ok(); + let attention_variant = AttentionVariant::Separate(Box::new(SeparateAttention { + attention_wq: q, + attention_qkv: None, + attention_q_norm: q_norm + .map(|norm| decode_norm(norm, rms_norm_eps)) + .transpose()?, + attention_wk: None, + attention_k_norm: None, + attention_wv: None, + attention_v_norm: None, + interleaved_rope: false, + bias: None, + })); + + let attention_wo = source + .tensor(&format!("{prefix}.attn_output.weight"), device) + .await?; + let feed_forward_w1 = source + .tensor(&format!("{prefix}.ffn_gate.weight"), device) + .await?; + let feed_forward_w2 = source + .tensor(&format!("{prefix}.ffn_down.weight"), device) + .await?; + let feed_forward_w3 = source + .tensor(&format!("{prefix}.ffn_up.weight"), device) + .await?; + let attention_norm = source + .tensor(&format!("{prefix}.attn_norm.weight"), device) + .await?; + let post_attention_norm = source + .tensor(&format!("{prefix}.post_attention_norm.weight"), device) + .await + .ok(); + let ffn_norm = source + .tensor(&format!("{prefix}.ffn_norm.weight"), device) + .await?; + let ffn_post_norm = source + .tensor(&format!("{prefix}.post_ffw_norm.weight"), device) + .await + .ok(); + let layer_output_scale = source + .tensor(&format!("{prefix}.layer_output_scale.weight"), device) + .await + .ok() + .map(&dequantize_1d); + + layers.push(LlamaAttention { + attention_variant, + attention_wo: Linear::new(attention_wo, None), + attention_norm: decode_norm(attention_norm, rms_norm_eps)?, + post_attention_norm: post_attention_norm + .map(|norm| decode_norm(norm, rms_norm_eps)) + .transpose()?, + feed_forward_variant: FeedForwardVariant::Llama(Box::new( + LlamaFeedForward::new_with_activation( + feed_forward_w1, + feed_forward_w2, + feed_forward_w3, + FeedForwardActivation::Gelu, + ), + )), + ffn_norm: decode_norm(ffn_norm, rms_norm_eps)?, + post_ffn_norm: ffn_post_norm + .map(|norm| decode_norm(norm, rms_norm_eps)) + .transpose()?, + n_head: head_count, + n_kv_head: head_count_kv, + head_dim: layer_head_dim, + hidden_size: layer_attention_width, + rope_cache, + sliding_window_size: layer_sliding_window_size, + attention_scale: 1.0, + shared_kv_layer: None, + per_layer_inp_gate: None, + per_layer_proj: None, + per_layer_post_norm: None, + layer_output_scale, + }); + } + + if pre_projection.shape().get(1).copied() != Some(target_embedding_length * 2) { + return Err(fusor::Error::msg(format!( + "unexpected Gemma4 MTP pre_projection input width {:?}, target hidden {target_embedding_length}", + pre_projection.shape() + )) + .into()); + } + if post_projection.shape().first().copied() != Some(target_embedding_length) { + return Err(fusor::Error::msg(format!( + "unexpected Gemma4 MTP post_projection output width {:?}, target hidden {target_embedding_length}", + post_projection.shape() + )) + .into()); + } + if embedding_length == 0 { + return Err(fusor::Error::msg("Gemma4 MTP assistant hidden size is zero").into()); + } + + Ok(Self { + config, + pre_projection, + post_projection, + layers, + norm, + output, + layer_is_sliding, + }) + } +} + +impl Gemma4MtpAssistant +where + F: CastTo + CastTensor, + f32: CastTo + CastTensor, + MulOp: SimdBinaryOp, + AddOp: SimdBinaryOp, + SumOp: SimdReduceOp, +{ + pub(crate) fn draft_step( + &self, + target: &Model, + token: u32, + h_nextn: &Tensor<2, f32>, + target_cache: &LlamaCache, + device: &Device, + position: usize, + ) -> Result { + let token_tensor: Tensor<2, u32> = + Tensor::from_slice(device, [1, 1], &[token]).to_concrete(); + let mut token_embedding = target.tok_embeddings.forward::<2, 3, _>(&token_tensor); + if let Some(scale) = target.tok_embedding_scale { + token_embedding = (token_embedding * scale).to_concrete(); + } + let h_nextn = h_nextn.unsqueeze(0).to_concrete(); + let projected = fusor::cat([token_embedding, h_nextn], 2) + .to_concrete() + .q_mat_mul(&self.pre_projection); + let mut layer_in: Tensor<3, F> = projected.cast(); + + for (layer_idx, layer) in self.layers.iter().enumerate() { + let target_kv_layer = self.target_kv_layer(target, layer_idx).ok_or_else(|| { + fusor::Error::msg("Gemma4 MTP could not find a matching target KV cache layer") + })?; + let key = target_cache.blocks[target_kv_layer] + .k() + .cloned() + .ok_or_else(|| fusor::Error::msg("Gemma4 MTP source key cache is empty"))?; + let value = target_cache.blocks[target_kv_layer] + .v() + .cloned() + .ok_or_else(|| fusor::Error::msg("Gemma4 MTP source value cache is empty"))?; + + let x = layer_in; + let residual: Tensor<3, f32> = x.cast(); + let x = layer.attention_norm.forward_generic(&x); + let mut attn = layer.forward_with_shared_kv(&x, None, position, None, &key, &value); + if let Some(post_attention_norm) = &layer.post_attention_norm { + attn = post_attention_norm.forward_generic(&attn); + } + let attn_f32: Tensor<3, f32> = attn.cast(); + let x = layer.ffn_norm.forward_residual_f32(&attn_f32, &residual); + let mut x = layer.feed_forward_variant.forward(&x); + if let Some(post_ffn_norm) = &layer.post_ffn_norm { + x = post_ffn_norm.forward_generic(&x); + } + let x_f32: Tensor<3, f32> = x.cast(); + layer_in = (x_f32 + attn_f32 + residual).cast(); + if let Some(layer_output_scale) = &layer.layer_output_scale { + let scale = layer_output_scale + .reshape([1, 1, 1]) + .broadcast_as(layer_in.shape()) + .to_concrete(); + layer_in = (layer_in * scale).to_concrete(); + } + } + + let normed = self.norm.forward_generic(&layer_in); + let logits: Tensor<1, f32> = normed + .cast::() + .q_mat_mul(&self.output) + .squeeze(0) + .squeeze(0) + .to_concrete(); + let h_nextn: Tensor<2, f32> = normed + .cast::() + .q_mat_mul(&self.post_projection) + .squeeze(0) + .to_concrete(); + Ok(Gemma4MtpStep { logits, h_nextn }) + } + + fn target_kv_layer(&self, target: &Model, assistant_layer_idx: usize) -> Option { + let wants_sliding = self + .layer_is_sliding + .get(assistant_layer_idx) + .copied() + .unwrap_or(false); + target + .layers + .iter() + .enumerate() + .rev() + .find(|(_, layer)| layer.sliding_window_size.is_some() == wants_sliding) + .map(|(idx, layer)| layer.shared_kv_layer.unwrap_or(idx)) + } + + pub(crate) fn draft_n(&self) -> usize { + self.config.n_layer.max(1) + } +} diff --git a/models/kalosm-llama/src/raw/rope.rs b/models/kalosm-llama/src/raw/rope.rs index 6bbfe717b..0128bb0cf 100644 --- a/models/kalosm-llama/src/raw/rope.rs +++ b/models/kalosm-llama/src/raw/rope.rs @@ -60,15 +60,21 @@ where F: CastTo + CastTensor, f32: CastTo + CastTensor, { - pub fn new(config: &LlamaConfig, rope_theta: f32, device: &Device) -> fusor::Result { + pub fn new_with_head_dimension( + config: &LlamaConfig, + head_dimension: usize, + rope_freq_weight: Option<&Tensor<1, F>>, + rope_theta: f32, + device: &Device, + ) -> fusor::Result { if let Some(mrope_sections) = &config.mrope_sections { let cache = QwenVLRopeCache::new(config, rope_theta, mrope_sections, device)?; Ok(Self::QwenVL(cache)) } else { let inverse_frequency: Tensor<2, F> = create_inverse_frequency( config.rope_scaling.as_ref(), - config.rope_freq_weight.as_ref(), - config.head_dimension, + rope_freq_weight, + head_dimension, rope_theta, device, ); diff --git a/models/kalosm-llama/src/raw/vision/gemma.rs b/models/kalosm-llama/src/raw/vision/gemma.rs new file mode 100644 index 000000000..d9a7913be --- /dev/null +++ b/models/kalosm-llama/src/raw/vision/gemma.rs @@ -0,0 +1,744 @@ +use fusor::{ + layers::RmsNorm, AddOp, CastTensor, CastTo, Device, DivOp, ExpOp, FloatDataType, FloatOps, + Fusion, MatmulImpl, MulOp, NegOp, QMatrix, Result, SimdBinaryOp, SimdElement, SimdUnaryOp, + Tensor, VarBuilder, +}; +use fusor_gguf::GgufMetadata; + +use crate::raw::rope::create_inverse_frequency; + +pub(crate) struct GemmaVisionTransformer { + patch_size: usize, + merge_size: usize, + patch_embed: GemmaVisionPatchEmbed, + position_embeddings: Tensor<3, F>, + blocks: Vec>, + projector_norm: RmsNorm<1, F>, + projector: GemmaClippedLinear, + std_bias: Option>, + std_scale: Option>, + image_mean: Vec, + image_std: Vec, + rope_theta: f32, + device: Device, + uses_isolated_device: bool, +} + +impl GemmaVisionTransformer +where + F: FloatDataType + + SimdElement + + FloatOps + + MatmulImpl + + Default + + CastTo + + CastTensor, + f32: CastTo + CastTensor, + fusor::MulOp: fusor::SimdBinaryOp, + fusor::AddOp: fusor::SimdBinaryOp, + fusor::SumOp: fusor::SimdReduceOp, +{ + pub(crate) fn from_gguf( + vision_ct: GgufMetadata, + vision_bytes: &[u8], + device: &Device, + ) -> Result { + let uses_isolated_device = cfg!(not(target_arch = "wasm32")) + && device.is_gpu() + && std::env::var_os("KALOSM_GEMMA4_VISION_DISABLE_SUBGROUPS").is_some(); + let vision_device = if uses_isolated_device { + // Debug fallback: use the no-subgroup sibling while sharing the + // same GPU adapter. The normal path is faster and covered by the + // split/combine long-attention tests. + device.without_subgroups() + } else { + device.clone() + }; + let device = &vision_device; + let block_count = metadata_usize(&vision_ct, "clip.vision.block_count", 16); + let head_count = metadata_usize(&vision_ct, "clip.vision.attention.head_count", 12); + let patch_size = metadata_usize(&vision_ct, "clip.vision.patch_size", 16); + let hidden_size = metadata_usize(&vision_ct, "clip.vision.embedding_length", 768); + let merge_size = metadata_usize(&vision_ct, "clip.vision.projector.scale_factor", 3); + let rope_theta = vision_ct + .metadata + .get("clip.vision.rope_theta") + .and_then(|x| x.to_f64().ok()) + .unwrap_or(100.0) as f32; + let layer_norm_eps = vision_ct + .metadata + .get("clip.vision.attention.layer_norm_epsilon") + .and_then(|x| x.to_f64().ok()) + .unwrap_or(1e-6) as f32; + let image_mean = metadata_f32_array(&vision_ct, "clip.vision.image_mean") + .unwrap_or_else(|| vec![0.0, 0.0, 0.0]); + let image_std = metadata_f32_array(&vision_ct, "clip.vision.image_std") + .unwrap_or_else(|| vec![1.0, 1.0, 1.0]); + + let mut cursor = std::io::Cursor::new(vision_bytes); + let mut vb = VarBuilder::from_gguf(&mut cursor)?; + let patch_embed = GemmaVisionPatchEmbed::new( + patch_size, + hidden_size, + &mut vb.pp("v.patch_embd"), + device, + )?; + let position_embeddings: Tensor<3, F> = vb + .get("v.position_embd.weight", device)? + .dequantize() + .cast(); + let head_dim = hidden_size / head_count; + let blocks = (0..block_count) + .map(|i| { + GemmaVisionBlock::new( + &mut vb.pp(format!("v.blk.{i}")), + device, + head_count, + head_dim, + hidden_size, + layer_norm_eps, + ) + }) + .collect::>>()?; + let projector_norm = + RmsNorm::new(Tensor::ones(device, [hidden_size]), None, layer_norm_eps); + let projector = clipped_linear(&mut vb, device, "mm.input_projection")?; + let std_bias: Option> = + vb.get("v.std_bias", device).ok().map(|x| x.dequantize()); + let std_scale: Option> = + vb.get("v.std_scale", device).ok().map(|x| x.dequantize()); + + Ok(Self { + patch_size, + merge_size, + patch_embed, + position_embeddings, + blocks, + projector_norm, + projector, + std_bias, + std_scale, + image_mean, + image_std, + rope_theta, + device: device.clone(), + uses_isolated_device, + }) + } + + pub(crate) fn uses_isolated_device(&self) -> bool { + self.uses_isolated_device + } + + pub(crate) fn preprocess_image( + &self, + image: &image::DynamicImage, + _min_pixels: Option, + _max_pixels: Option, + ) -> Result<(Tensor<2, f32>, [u32; 3])> { + let (target_width, target_height) = self.target_image_size(image); + let resized = image.resize_exact( + target_width as u32, + target_height as u32, + image::imageops::FilterType::Triangle, + ); + let rgb = image_to_rgb(&resized, &self.device)?; + let mean_tensor = Tensor::new(&self.device, &self.image_mean); + let mean = mean_tensor.reshape([1, 3, 1, 1]); + let std_tensor = Tensor::new(&self.device, &self.image_std); + let std = std_tensor.reshape([1, 3, 1, 1]); + let rgb = rgb + .sub_(&mean) + .div_(&std) + .mul_scalar(2.0) + .add_scalar(-1.0) + .to_concrete(); + let grid_h = target_height / self.patch_size; + let grid_w = target_width / self.patch_size; + if std::env::var_os("KALOSM_TRACE_VISION_STATS").is_some() { + tracing::info!( + "[vision_stats] preprocess original={}x{} target={}x{} grid=[1,{grid_h},{grid_w}] pooled_tokens={}", + image.width(), + image.height(), + target_width, + target_height, + (grid_h / self.merge_size) * (grid_w / self.merge_size) + ); + } + let patches = rgb + .reshape([1, 3, grid_h, self.patch_size, grid_w, self.patch_size]) + .permute([0, 2, 4, 1, 3, 5]) + .reshape([grid_h * grid_w, 3 * self.patch_size * self.patch_size]) + .to_concrete(); + crate::raw::debug_tensor_stats_f32(&patches, "image_patches"); + + Ok((patches, [1, grid_h as u32, grid_w as u32])) + } + + pub(crate) fn image_token_count(&self, grid: [u32; 3]) -> u32 { + grid[0] * (grid[1] / self.merge_size as u32) * (grid[2] / self.merge_size as u32) + } + + pub(crate) fn forward_image( + &self, + pixels: &Tensor<2, F>, + grid: [u32; 3], + ) -> Result> { + let [pos_x, pos_y] = self.patch_positions(grid)?; + let (cos_x, sin_x) = self.rope_sin_cos(&pos_x)?; + let (cos_y, sin_y) = self.rope_sin_cos(&pos_y)?; + let mut hidden_states = self.patch_embed.forward(pixels); + let hidden_f32: Tensor<2, f32> = hidden_states.cast(); + crate::raw::debug_tensor_stats_f32(&hidden_f32, "patch_embeds"); + hidden_states = self.add_position_embeddings(&hidden_states, &pos_x, &pos_y, grid)?; + let hidden_f32: Tensor<2, f32> = hidden_states.cast(); + crate::raw::debug_tensor_stats_f32(&hidden_f32, "patch_plus_position"); + let mut hidden_states = hidden_states.unsqueeze(0).to_concrete(); + for block in &self.blocks { + hidden_states = block.forward(&hidden_states, &cos_x, &sin_x, &cos_y, &sin_y); + } + let hidden_f32: Tensor<3, f32> = hidden_states.cast(); + crate::raw::debug_tensor_stats_f32(&hidden_f32, "vision_blocks_out"); + let mut hidden_states = self.pool(hidden_states, grid)?; + let hidden_f32: Tensor<3, f32> = hidden_states.cast(); + crate::raw::debug_tensor_stats_f32(&hidden_f32, "vision_pooled"); + if let (Some(std_bias), Some(std_scale)) = (&self.std_bias, &self.std_scale) { + let hidden_f32: Tensor<3, f32> = hidden_states.cast(); + hidden_states = hidden_f32.sub_(std_bias).mul_(std_scale).cast(); + } + let hidden_states = self.projector_norm.forward_generic(&hidden_states); + let hidden_f32: Tensor<3, f32> = hidden_states.cast(); + crate::raw::debug_tensor_stats_f32(&hidden_f32, "vision_projector_norm"); + Ok(self + .projector + .forward(&hidden_states) + .squeeze::<2>(0) + .to_concrete()) + } + + fn target_image_size(&self, image: &image::DynamicImage) -> (usize, usize) { + let align = self.patch_size * self.merge_size; + let min_pixels = 40 * align * align; + let max_pixels = 280 * align * align; + smart_resize( + image.width() as usize, + image.height() as usize, + align, + min_pixels, + max_pixels, + ) + } + + fn add_position_embeddings( + &self, + hidden_states: &Tensor<2, F, B>, + x_ids: &Tensor<1, u32>, + y_ids: &Tensor<1, u32>, + grid: [u32; 3], + ) -> Result> + where + B: Fusion<2, F>, + { + let [grid_t, _grid_h, _grid_w] = grid; + if grid_t != 1 { + return Err(fusor::Error::msg( + "Gemma 4 vision currently supports image inputs, not video frames.", + )); + } + + let pos_x_table = self.position_embeddings.i((0, .., ..)).to_concrete(); + let pos_y_table = self.position_embeddings.i((1, .., ..)).to_concrete(); + let pos_x = pos_x_table.index_select(0, x_ids); + let pos_y = pos_y_table.index_select(0, y_ids); + let position_embeddings = (pos_x + pos_y).to_concrete(); + Ok((hidden_states.to_concrete() + position_embeddings).to_concrete()) + } + + fn patch_positions(&self, grid: [u32; 3]) -> Result<[Tensor<1, u32>; 2]> { + let [grid_t, grid_h, grid_w] = grid; + if grid_t != 1 { + return Err(fusor::Error::msg( + "Gemma 4 vision currently supports image inputs, not video frames.", + )); + } + let mut y_ids = Vec::with_capacity((grid_h * grid_w) as usize); + let mut x_ids = Vec::with_capacity((grid_h * grid_w) as usize); + for y in 0..grid_h { + for x in 0..grid_w { + y_ids.push(y); + x_ids.push(x); + } + } + Ok([ + Tensor::new(&self.device, &x_ids), + Tensor::new(&self.device, &y_ids), + ]) + } + + fn rope_sin_cos(&self, positions: &Tensor<1, u32>) -> Result<(Tensor<2, f32>, Tensor<2, f32>)> { + let half_head_dim = self + .blocks + .first() + .map(|block| block.head_dim() / 2) + .ok_or_else(|| { + fusor::Error::msg("Gemma 4 vision transformer must have at least one block") + })?; + let positions: Tensor<2, f32> = positions + .cast::() + .reshape([positions.shape()[0], 1]) + .to_concrete(); + let inv_freq: Tensor<2, f32> = create_inverse_frequency::( + None, + None, + half_head_dim, + self.rope_theta, + &self.device, + ); + let freqs = positions.matmul(&inv_freq); + Ok((freqs.cos().to_concrete(), freqs.sin().to_concrete())) + } + + fn pool(&self, hidden_states: Tensor<3, F>, grid: [u32; 3]) -> Result> { + let [batch, seq, hidden] = hidden_states.shape(); + let grid_h = grid[1] as usize; + let grid_w = grid[2] as usize; + if batch != 1 || seq != grid_h * grid_w { + return Err(fusor::Error::msg( + "Gemma 4 vision grid does not match hidden states", + )); + } + if grid_h % self.merge_size != 0 || grid_w % self.merge_size != 0 { + return Err(fusor::Error::msg( + "Gemma 4 vision grid must be divisible by the merge size", + )); + } + + let out_h = grid_h / self.merge_size; + let out_w = grid_w / self.merge_size; + let pooled = hidden_states + .reshape([ + batch, + out_h, + self.merge_size, + out_w, + self.merge_size, + hidden, + ]) + .sum::<5>(4) + .mul_scalar(F::from_f32(1.0 / self.merge_size as f32)) + .sum::<4>(2) + .mul_scalar(F::from_f32(1.0 / self.merge_size as f32)) + .reshape([batch, out_h * out_w, hidden]) + .mul_scalar(F::from_f32((hidden as f32).sqrt())) + .to_concrete(); + Ok(pooled) + } +} + +struct GemmaVisionPatchEmbed { + weight: Tensor<2, F>, +} + +impl GemmaVisionPatchEmbed +where + F: FloatDataType + SimdElement + FloatOps + MatmulImpl + CastTo + CastTensor, + f32: CastTo + CastTensor, +{ + fn new( + patch_size: usize, + hidden_size: usize, + vb: &mut VarBuilder, + device: &Device, + ) -> Result { + let weight: Tensor<4, F> = vb.get("weight", device)?.dequantize().cast(); + let weight = weight + .permute([1, 2, 3, 0]) + .reshape([3 * patch_size * patch_size, hidden_size]) + .to_concrete(); + Ok(Self { weight }) + } + + fn forward(&self, pixels: &Tensor<2, F, B>) -> Tensor<2, F> + where + B: Fusion<2, F>, + { + pixels.matmul(&self.weight) + } +} + +struct GemmaVisionBlock { + norm1: RmsNorm<1, F>, + norm2: RmsNorm<1, F>, + attn: GemmaVisionAttention, + attn_post_norm: RmsNorm<1, F>, + mlp: GemmaVisionFeedForward, + ffn_post_norm: RmsNorm<1, F>, +} + +impl GemmaVisionBlock +where + F: FloatDataType + SimdElement + Default + CastTo + CastTensor, + f32: CastTo + CastTensor, +{ + fn new( + vb: &mut VarBuilder, + device: &Device, + head_count: usize, + head_dim: usize, + hidden_size: usize, + layer_norm_eps: f32, + ) -> Result { + let norm1 = rms_norm(vb.get("ln1.weight", device)?, layer_norm_eps); + let norm2 = rms_norm(vb.get("ln2.weight", device)?, layer_norm_eps); + let attn_post_norm = rms_norm(vb.get("attn_post_norm.weight", device)?, layer_norm_eps); + let ffn_post_norm = rms_norm(vb.get("ffn_post_norm.weight", device)?, layer_norm_eps); + let attn = GemmaVisionAttention::new(vb, device, head_count, head_dim, hidden_size)?; + let mlp = GemmaVisionFeedForward::new(vb, device)?; + + Ok(Self { + norm1, + norm2, + attn, + attn_post_norm, + mlp, + ffn_post_norm, + }) + } + + fn head_dim(&self) -> usize { + self.attn.head_dim + } + + fn forward( + &self, + xs: &Tensor<3, F, B>, + cos_x: &Tensor<2, f32>, + sin_x: &Tensor<2, f32>, + cos_y: &Tensor<2, f32>, + sin_y: &Tensor<2, f32>, + ) -> Tensor<3, F> + where + B: Fusion<3, F>, + { + let residual: Tensor<3, f32> = xs.cast(); + let x = self.norm1.forward_generic(xs); + let attn = self.attn.forward(&x, cos_x, sin_x, cos_y, sin_y); + let attn = self.attn_post_norm.forward_generic(&attn); + let attn_f32: Tensor<3, f32> = attn.cast(); + let x = self.norm2.forward_residual_f32(&attn_f32, &residual); + let ffn = self.mlp.forward(&x); + let ffn = self.ffn_post_norm.forward_generic(&ffn); + let ffn_f32: Tensor<3, f32> = ffn.cast(); + (ffn_f32 + attn_f32 + residual).cast() + } +} + +struct GemmaVisionAttention { + q: GemmaClippedLinear, + k: GemmaClippedLinear, + v: GemmaClippedLinear, + out: GemmaClippedLinear, + q_norm: RmsNorm<1, F>, + k_norm: RmsNorm<1, F>, + v_norm: RmsNorm<1, F>, + head_count: usize, + head_dim: usize, + hidden_size: usize, +} + +impl GemmaVisionAttention +where + F: FloatDataType + SimdElement + Default + CastTo + CastTensor, + f32: CastTo + CastTensor, +{ + fn new( + vb: &mut VarBuilder, + device: &Device, + head_count: usize, + head_dim: usize, + hidden_size: usize, + ) -> Result { + Ok(Self { + q: clipped_linear(vb, device, "attn_q")?, + k: clipped_linear(vb, device, "attn_k")?, + v: clipped_linear(vb, device, "attn_v")?, + out: clipped_linear(vb, device, "attn_out")?, + q_norm: rms_norm(vb.get("attn_q_norm.weight", device)?, 1e-6), + k_norm: rms_norm(vb.get("attn_k_norm.weight", device)?, 1e-6), + v_norm: RmsNorm::new(Tensor::ones(device, [head_dim]), None, 1e-6), + head_count, + head_dim, + hidden_size, + }) + } + + fn forward( + &self, + xs: &Tensor<3, F, B>, + cos_x: &Tensor<2, f32>, + sin_x: &Tensor<2, f32>, + cos_y: &Tensor<2, f32>, + sin_y: &Tensor<2, f32>, + ) -> Tensor<3, F> + where + B: Fusion<3, F>, + { + let [batch, seq_len, _] = xs.shape(); + let q = self + .q + .forward(xs) + .reshape([batch, seq_len, self.head_count, self.head_dim]) + .transpose(1, 2) + .to_concrete(); + let k = self + .k + .forward(xs) + .reshape([batch, seq_len, self.head_count, self.head_dim]) + .transpose(1, 2) + .to_concrete(); + let v = self + .v + .forward(xs) + .reshape([batch, seq_len, self.head_count, self.head_dim]) + .transpose(1, 2) + .to_concrete(); + let q: Tensor<4, f32> = self.q_norm.forward_generic_4d(&q).cast(); + let k: Tensor<4, f32> = self.k_norm.forward_generic_4d(&k).cast(); + let v: Tensor<4, f32> = self.v_norm.forward_generic_4d(&v).cast(); + let half = self.head_dim / 2; + let q_x = q.narrow(3, 0, half).to_concrete(); + let k_x = k.narrow(3, 0, half).to_concrete(); + let q_y = q.narrow(3, half, half).to_concrete(); + let k_y = k.narrow(3, half, half).to_concrete(); + let (q_x, k_x) = q_x.rope_normal_pair_fused(&k_x, cos_x, sin_x); + let (q_y, k_y) = q_y.rope_normal_pair_fused(&k_y, cos_y, sin_y); + let q = Tensor::cat([q_x, q_y], 3).to_concrete(); + let k = Tensor::cat([k_x, k_y], 3).to_concrete(); + let attn = q.flash_attention(&k, &v, 1.0, None); + let attn = attn + .transpose(1, 2) + .reshape([batch, seq_len, self.hidden_size]) + .cast(); + self.out.forward(&attn) + } +} + +struct GemmaVisionFeedForward { + gate: GemmaClippedLinear, + down: GemmaClippedLinear, + up: GemmaClippedLinear, +} + +impl GemmaVisionFeedForward { + fn new(vb: &mut VarBuilder, device: &Device) -> Result { + Ok(Self { + gate: clipped_linear(vb, device, "ffn_gate")?, + down: clipped_linear(vb, device, "ffn_down")?, + up: clipped_linear(vb, device, "ffn_up")?, + }) + } + + fn forward(&self, x: &Tensor<3, F, B>) -> Tensor<3, F> + where + F: FloatDataType + SimdElement + Default + CastTo + CastTensor, + f32: CastTo + CastTensor, + B: Fusion<3, F>, + { + let gate = quick_gelu(&self.gate.forward_f32(x)); + let up = self.up.forward_f32(x); + let hidden = (gate * up).to_concrete(); + self.down.forward_from_f32(&hidden).cast() + } +} + +fn quick_gelu(x: &Tensor) -> Tensor +where + B: Fusion, + AddOp: SimdBinaryOp, + DivOp: SimdBinaryOp, + ExpOp: SimdUnaryOp, + MulOp: SimdBinaryOp, + NegOp: SimdUnaryOp, +{ + let x = x.to_concrete(); + let scaled = (x.clone() * 1.702).to_concrete(); + let sigmoid = scaled.sigmoid(); + (x * sigmoid).to_concrete() +} + +#[derive(Clone, Copy)] +struct ClampInfo { + input_min: f32, + input_max: f32, + output_min: f32, + output_max: f32, +} + +impl ClampInfo { + fn from_tensors(vb: &mut VarBuilder, device: &Device, prefix: &str) -> Self { + Self { + input_min: scalar_tensor(vb, device, &format!("{prefix}.input_min"), -f32::MAX), + input_max: scalar_tensor(vb, device, &format!("{prefix}.input_max"), f32::MAX), + output_min: scalar_tensor(vb, device, &format!("{prefix}.output_min"), -f32::MAX), + output_max: scalar_tensor(vb, device, &format!("{prefix}.output_max"), f32::MAX), + } + } + + fn has_input_clamp(self) -> bool { + self.input_min > -f32::MAX || self.input_max < f32::MAX + } + + fn has_output_clamp(self) -> bool { + self.output_min > -f32::MAX || self.output_max < f32::MAX + } +} + +struct GemmaClippedLinear { + weight: QMatrix, + clamp: ClampInfo, +} + +impl GemmaClippedLinear { + fn new(weight: QMatrix, clamp: ClampInfo) -> Self { + Self { weight, clamp } + } + + fn forward(&self, input: &Tensor<3, F, B>) -> Tensor<3, F> + where + F: FloatDataType + SimdElement + Default + CastTo + CastTensor, + f32: CastTo + CastTensor, + B: Fusion<3, F>, + { + self.forward_f32(input).cast() + } + + fn forward_f32(&self, input: &Tensor<3, F, B>) -> Tensor<3, f32> + where + F: FloatDataType + SimdElement + Default + CastTo + CastTensor, + B: Fusion<3, F>, + { + self.forward_from_f32(&input.cast::()) + } + + fn forward_from_f32(&self, input: &Tensor<3, f32, B>) -> Tensor<3, f32> + where + B: Fusion<3, f32>, + { + let input = if self.clamp.has_input_clamp() { + input + .clamp(self.clamp.input_min, self.clamp.input_max) + .to_concrete() + } else { + input.to_concrete() + }; + let output = input.q_mat_mul(&self.weight); + if self.clamp.has_output_clamp() { + output + .clamp(self.clamp.output_min, self.clamp.output_max) + .to_concrete() + } else { + output + } + } +} + +fn clipped_linear( + vb: &mut VarBuilder, + device: &Device, + prefix: &str, +) -> Result { + let weight = vb.get(&format!("{prefix}.weight"), device)?; + let clamp = ClampInfo::from_tensors(vb, device, prefix); + Ok(GemmaClippedLinear::new(weight, clamp)) +} + +fn scalar_tensor(vb: &mut VarBuilder, device: &Device, name: &str, default: f32) -> f32 { + let Ok(tensor) = vb.get(name, device) else { + return default; + }; + let tensor: Tensor<1, f32> = tensor.dequantize(); + first_f32(&tensor).unwrap_or(default) +} + +#[cfg(not(target_arch = "wasm32"))] +fn first_f32(tensor: &Tensor<1, f32>) -> Option { + let slice = pollster::block_on(tensor.as_slice()).ok()?; + slice.as_slice().first().copied() +} + +#[cfg(target_arch = "wasm32")] +fn first_f32(_tensor: &Tensor<1, f32>) -> Option { + None +} + +fn rms_norm(weight: QMatrix, eps: f32) -> RmsNorm<1, F> +where + F: FloatDataType + SimdElement + Default + CastTo + CastTensor, + f32: CastTo + CastTensor, +{ + let weight: Tensor<1, F> = weight.dequantize().cast(); + RmsNorm::new(weight, None, eps) +} + +fn metadata_usize(metadata: &GgufMetadata, key: &str, default: usize) -> usize { + metadata + .metadata + .get(key) + .and_then(|x| x.to_u64().ok()) + .map(|x| x as usize) + .unwrap_or(default) +} + +fn smart_resize( + width: usize, + height: usize, + align: usize, + min_pixels: usize, + max_pixels: usize, +) -> (usize, usize) { + let round_by_factor = |x: f64| ((x / align as f64).round() as usize).max(1) * align; + let ceil_by_factor = |x: f64| ((x / align as f64).ceil() as usize).max(1) * align; + let floor_by_factor = |x: f64| ((x / align as f64).floor() as usize).max(1) * align; + + let mut target_h = round_by_factor(height as f64); + let mut target_w = round_by_factor(width as f64); + let pixels = (width * height) as f64; + if target_h * target_w > max_pixels { + let beta = (pixels / max_pixels as f64).sqrt(); + target_h = floor_by_factor(height as f64 / beta); + target_w = floor_by_factor(width as f64 / beta); + } else if target_h * target_w < min_pixels { + let beta = (min_pixels as f64 / pixels).sqrt(); + target_h = ceil_by_factor(height as f64 * beta); + target_w = ceil_by_factor(width as f64 * beta); + } + + (target_w, target_h) +} + +fn metadata_f32_array(metadata: &GgufMetadata, key: &str) -> Option> { + metadata.metadata.get(key).and_then(|value| { + value.to_array().ok().map(|values| { + values + .iter() + .filter_map(|value| value.to_f32().ok()) + .collect() + }) + }) +} + +fn image_to_rgb(image: &image::DynamicImage, device: &Device) -> Result> { + let height = image.height() as usize; + let width = image.width() as usize; + let rgb = image.to_rgb8(); + let as_u32 = rgb + .into_raw() + .into_iter() + .map(|x| x as u32) + .collect::>(); + let data_tensor = Tensor::new(device, &as_u32); + let data = data_tensor.reshape([height, width, 3]); + let img = data.permute([2, 0, 1]).cast::() * (1.0 / 255.0); + + Ok(img.unsqueeze(0).to_concrete()) +} diff --git a/models/kalosm-llama/src/raw/vision/mod.rs b/models/kalosm-llama/src/raw/vision/mod.rs index f3671fb95..f0ef4c77c 100644 --- a/models/kalosm-llama/src/raw/vision/mod.rs +++ b/models/kalosm-llama/src/raw/vision/mod.rs @@ -1,3 +1,4 @@ +mod gemma; mod qwen; mod qwen_image_processing; mod qwen_patch_merger; @@ -6,6 +7,177 @@ mod qwen_vision; mod qwen_vision_block; mod qwen_vision_embed; +use std::ops::Range; + +use fusor::{CastTensor, CastTo, Device, FloatDataType, Result, SimdElement, Tensor}; +use fusor_gguf::GgufMetadata; + +pub(crate) use gemma::GemmaVisionTransformer; pub(crate) use qwen::QwenVisionTransformer; pub const QWEN_EPS: f64 = 1e-6; + +pub(crate) enum VisionTransformer { + Qwen(QwenVisionTransformer), + Gemma(GemmaVisionTransformer), +} + +impl VisionTransformer +where + F: CastTo + CastTensor + fusor::FloatOps + fusor::MatmulImpl + Default, + f32: CastTo + CastTensor, + fusor::MulOp: fusor::SimdBinaryOp, + fusor::AddOp: fusor::SimdBinaryOp, + fusor::SumOp: fusor::SimdReduceOp, +{ + pub(crate) fn from_gguf( + vision_ct: GgufMetadata, + vision_bytes: &[u8], + device: &Device, + ) -> Result { + let projector_type = vision_ct + .metadata + .get("clip.vision.projector_type") + .and_then(|x| x.to_string().ok()) + .map(|x| x.to_string()); + + match projector_type.as_deref() { + Some("gemma4v") => Ok(Self::Gemma(GemmaVisionTransformer::from_gguf( + vision_ct, + vision_bytes, + device, + )?)), + _ => Ok(Self::Qwen(QwenVisionTransformer::from_gguf( + vision_ct, + vision_bytes, + device, + )?)), + } + } + + pub(crate) fn preprocess_image( + &self, + image: &image::DynamicImage, + min_pixels: Option, + max_pixels: Option, + ) -> Result<(Tensor<2, f32>, [u32; 3])> { + match self { + Self::Qwen(vision) => vision.preprocess_image(image, min_pixels, max_pixels), + Self::Gemma(vision) => vision.preprocess_image(image, min_pixels, max_pixels), + } + } + + pub(crate) fn image_token_count(&self, grid: [u32; 3]) -> u32 { + match self { + Self::Qwen(vision) => { + grid.iter().product::() / (vision.spacial_merge_size as u32).pow(2) + } + Self::Gemma(vision) => vision.image_token_count(grid), + } + } + + pub(crate) fn expand_image_tokens( + &self, + raw_tokens: &[u32], + image_pad_token: u32, + vision_start_token: Option, + image_start_token: Option, + image_end_token: Option, + grid_thw: &[[u32; 3]], + ) -> Result<(Vec, Vec>)> { + match self { + Self::Qwen(_) => { + let Some(vision_start_token) = vision_start_token else { + return Ok((raw_tokens.to_vec(), Vec::new())); + }; + let mut tokens = Vec::new(); + let mut token_iter = raw_tokens.iter().copied(); + let mut image_iter = grid_thw.iter(); + let mut image_token_ranges = Vec::new(); + while let Some(token) = token_iter.next() { + tokens.push(token); + let start_index = tokens.len(); + if token == vision_start_token { + match token_iter.next() { + Some(next) if next == image_pad_token => { + let grid = *image_iter.next().ok_or_else(|| { + fusor::Error::msg( + "Image pad token found without matching image.", + ) + })?; + for _ in 0..self.image_token_count(grid) { + tokens.push(image_pad_token); + } + image_token_ranges.push(start_index..tokens.len()); + } + Some(next) => { + tokens.push(next); + } + None => break, + } + } + } + Ok((tokens, image_token_ranges)) + } + Self::Gemma(_) => { + let mut tokens = Vec::new(); + let mut image_iter = grid_thw.iter(); + let mut image_token_ranges = Vec::new(); + for token in raw_tokens.iter().copied() { + if token == image_pad_token { + let grid = *image_iter.next().ok_or_else(|| { + fusor::Error::msg("Image token found without matching image.") + })?; + if let Some(image_start_token) = image_start_token { + tokens.push(image_start_token); + } + let start_index = tokens.len(); + for _ in 0..self.image_token_count(grid) { + tokens.push(image_pad_token); + } + image_token_ranges.push(start_index..tokens.len()); + if let Some(image_end_token) = image_end_token { + tokens.push(image_end_token); + } + } else { + tokens.push(token); + } + } + Ok((tokens, image_token_ranges)) + } + } + } + + pub(crate) fn get_rope_index( + &self, + input_ids: &[u32], + grid_thw: &[[u32; 3]], + config: &crate::raw::LlamaConfig, + start_time: u32, + ) -> Result, u32)>> { + match self { + Self::Qwen(vision) => vision + .get_rope_index(input_ids, grid_thw, config, start_time) + .map(Some), + Self::Gemma(_) => Ok(None), + } + } + + pub(crate) fn forward_image( + &self, + pixels: &Tensor<2, F>, + grid: [u32; 3], + ) -> Result> { + match self { + Self::Qwen(vision) => vision.forward_image(pixels, grid), + Self::Gemma(vision) => vision.forward_image(pixels, grid), + } + } + + pub(crate) fn outputs_on_isolated_device(&self) -> bool { + match self { + Self::Qwen(_) => false, + Self::Gemma(vision) => vision.uses_isolated_device(), + } + } +} diff --git a/models/kalosm-llama/src/source.rs b/models/kalosm-llama/src/source.rs index 610990eee..dacdb6443 100644 --- a/models/kalosm-llama/src/source.rs +++ b/models/kalosm-llama/src/source.rs @@ -56,6 +56,7 @@ pub(crate) struct LlamaConfigJson { pub struct LlamaSource { pub(crate) model: Vec, pub(crate) vision_model: Option, + pub(crate) mtp_model: Option, pub(crate) tokenizer: Option, pub(crate) config: Option, pub(crate) group_query_attention: u8, @@ -125,6 +126,7 @@ impl LlamaSource { override_stop_token_string: None, override_chat_template: None, vision_model: None, + mtp_model: None, } } @@ -139,6 +141,7 @@ impl LlamaSource { override_stop_token_string: None, override_chat_template: None, vision_model: None, + mtp_model: None, } } @@ -223,6 +226,13 @@ impl LlamaSource { self } + /// Set the Gemma4 MTP assistant model to use for opt-in speculative decoding. + pub fn with_mtp_model(mut self, model: FileSource) -> Self { + self.mtp_model = Some(model); + + self + } + #[cfg(not(target_arch = "wasm32"))] pub(crate) async fn model( &self, @@ -976,6 +986,23 @@ impl LlamaSource { .with_override_stop_token_string("") } + /// A preset for Unsloth's Gemma 4 E2B instruction QAT GGUF. + /// + /// Note: The gemma model series does not support system prompts. + pub fn gemma_4_e2b_it_qat_chat() -> Self { + Self::new(FileSource::huggingface( + "unsloth/gemma-4-E2B-it-qat-GGUF".to_string(), + "main".to_string(), + "gemma-4-E2B-it-qat-UD-Q4_K_XL.gguf".to_string(), + )) + .with_mtp_model(FileSource::huggingface( + "unsloth/gemma-4-E2B-it-qat-GGUF".to_string(), + "main".to_string(), + "MTP/gemma-4-E2B-it-Q4_0-MTP.gguf".to_string(), + )) + .with_override_stop_token_string("") + } + /// A preset for qwen 2.5 3b VL chat in f16 precision pub fn qwen_2_5_3b_vl_chat_f16() -> Self { Self::new(kalosm_model_types::FileSource::HuggingFace { diff --git a/models/kalosm-tokenizer/src/lib.rs b/models/kalosm-tokenizer/src/lib.rs index a42dd0671..8a30e01aa 100644 --- a/models/kalosm-tokenizer/src/lib.rs +++ b/models/kalosm-tokenizer/src/lib.rs @@ -57,6 +57,7 @@ pub struct FastBpe { byte_to_token: [u32; 256], all_bytes_present: bool, byte_token_mapping: ByteTokenMapping, + raw_utf8_initial_tokens: bool, token_bytes: FxHashMap, u32>, id_to_bytes: Vec>>, merges: FxHashMap, @@ -71,25 +72,76 @@ impl FastBpe { vocab: impl IntoIterator, merges: impl IntoIterator, ignore_merges: bool, + ) -> Result { + Self::from_vocab_and_merges_with_decoder( + vocab, + merges, + ignore_merges, + TokenByteDecoder::ByteLevel, + ) + } + + /// Build a BPE tokenizer whose vocabulary tokens are raw UTF-8 strings. + /// + /// Some GGUF tokenizers, including Gemma 4, are BPE tokenizers but do not + /// use the GPT-2 byte-level character mapping. Their byte fallback tokens + /// are encoded as `<0xNN>` strings in the vocabulary. + pub fn from_raw_utf8_vocab_and_merges( + vocab: impl IntoIterator, + merges: impl IntoIterator, + ignore_merges: bool, + ) -> Result { + Self::from_vocab_and_merges_with_decoder( + vocab, + merges, + ignore_merges, + TokenByteDecoder::RawUtf8WithHexBytes, + ) + } + + fn from_vocab_and_merges_with_decoder( + vocab: impl IntoIterator, + merges: impl IntoIterator, + ignore_merges: bool, + decoder: TokenByteDecoder, ) -> Result { let mut byte_to_token = [MISSING_BYTE_TOKEN; 256]; let mut token_bytes = FxHashMap::default(); + let mut token_byte_preference = FxHashMap::default(); + let mut exact_token_ids = FxHashMap::default(); let mut id_to_bytes = Vec::new(); + let mut vocab = vocab.into_iter().collect::>(); + vocab.sort_unstable_by_key(|(_, id)| *id); for (token, id) in vocab { - let bytes = decode_token_bytes(&token); + let bytes = decoder.decode_token_bytes(&token); + let is_byte_fallback = decoder.is_byte_fallback_token(&token); if bytes.len() == 1 { if id == MISSING_BYTE_TOKEN { return Err(TokenizerError::ReservedByteTokenId(id)); } - byte_to_token[bytes[0] as usize] = id; + if decoder == TokenByteDecoder::RawUtf8WithHexBytes { + if is_byte_fallback { + byte_to_token[bytes[0] as usize] = id; + } + } else { + byte_to_token[bytes[0] as usize] = id; + } } let id = id as usize; if id >= id_to_bytes.len() { id_to_bytes.resize(id + 1, None); } id_to_bytes[id] = Some(bytes.clone()); - token_bytes.insert(bytes, id as u32); + exact_token_ids.insert(token, id as u32); + let preference = TokenBytePreference::new(is_byte_fallback, id as u32); + if token_byte_preference + .get(&bytes) + .is_none_or(|existing: &TokenBytePreference| preference < *existing) + { + token_byte_preference.insert(bytes.clone(), preference); + token_bytes.insert(bytes, id as u32); + } } let all_bytes_present = byte_to_token @@ -100,7 +152,9 @@ impl FastBpe { let raw_merges = merges .into_iter() .enumerate() - .map(|(rank, merge)| parse_merge(rank as u32, &merge, &token_bytes)) + .map(|(rank, merge)| { + parse_merge(rank as u32, &merge, &token_bytes, &exact_token_ids, decoder) + }) .collect::, _>>()?; let merge_levels = assign_merge_levels(&raw_merges); let level_count = merge_levels @@ -154,6 +208,7 @@ impl FastBpe { byte_to_token, all_bytes_present, byte_token_mapping, + raw_utf8_initial_tokens: decoder == TokenByteDecoder::RawUtf8WithHexBytes, token_bytes, id_to_bytes, merges: merges_by_pair, @@ -271,6 +326,10 @@ impl FastBpe { } fn encode_bytes_into(&self, input: &[u8], out: &mut Vec) -> Result<(), TokenizerError> { + if self.raw_utf8_initial_tokens { + return self.encode_raw_utf8_chars_into(input, out); + } + out.clear(); out.resize(input.len(), 0); @@ -283,6 +342,28 @@ impl FastBpe { Ok(()) } + fn encode_raw_utf8_chars_into( + &self, + input: &[u8], + out: &mut Vec, + ) -> Result<(), TokenizerError> { + out.clear(); + let text = std::str::from_utf8(input).map_err(|_| TokenizerError::InvalidUtf8)?; + out.reserve(text.len()); + for ch in text.chars() { + let mut buf = [0; 4]; + let bytes = ch.encode_utf8(&mut buf).as_bytes(); + if let Some(token) = self.token_bytes.get(bytes) { + out.push(*token); + } else { + for byte in bytes { + out.push(lookup_byte_token(&self.byte_to_token, *byte)?); + } + } + } + Ok(()) + } + fn encode_bytes_and_apply_single_merge( &self, input: &[u8], @@ -457,6 +538,9 @@ pub enum TokenizerError { /// A byte was missing from the vocabulary. #[error("byte 0x{0:02x} is missing from the vocabulary")] MissingByteToken(u8), + /// Raw UTF-8 tokenization received invalid UTF-8. + #[error("raw UTF-8 input is invalid")] + InvalidUtf8, } #[inline(always)] @@ -625,27 +709,52 @@ struct RawMerge { new_token: u32, } +#[derive(Clone, Copy, Debug, Eq, PartialEq, Ord, PartialOrd)] +struct TokenBytePreference { + byte_fallback_rank: u8, + id: u32, +} + +impl TokenBytePreference { + fn new(is_byte_fallback: bool, id: u32) -> Self { + Self { + byte_fallback_rank: u8::from(is_byte_fallback), + id, + } + } +} + fn parse_merge( rank: u32, merge: &str, token_bytes: &FxHashMap, u32>, + exact_token_ids: &FxHashMap, + decoder: TokenByteDecoder, ) -> Result { - let (left, right) = merge - .split_once(' ') + let split_at = merge + .char_indices() + .skip(1) + .find_map(|(index, ch)| (ch == ' ').then_some(index)) .ok_or_else(|| TokenizerError::InvalidMerge(merge.to_owned()))?; - let left_bytes = decode_token_bytes(left); - let right_bytes = decode_token_bytes(right); + let left = &merge[..split_at]; + let right = &merge[split_at + 1..]; + let left_bytes = decoder.decode_token_bytes(left); + let right_bytes = decoder.decode_token_bytes(right); let new_bytes = left_bytes .iter() .chain(right_bytes.iter()) .copied() .collect::>(); - let left = *token_bytes - .get(&left_bytes) + let left = exact_token_ids + .get(left) + .or_else(|| token_bytes.get(&left_bytes)) + .copied() .ok_or_else(|| TokenizerError::MissingToken(left.to_owned()))?; - let right = *token_bytes - .get(&right_bytes) + let right = exact_token_ids + .get(right) + .or_else(|| token_bytes.get(&right_bytes)) + .copied() .ok_or_else(|| TokenizerError::MissingToken(right.to_owned()))?; let new_token = *token_bytes .get(&new_bytes) @@ -1381,6 +1490,59 @@ fn decode_token_bytes(token: &str) -> Vec { bytes } +#[derive(Clone, Copy, Eq, PartialEq)] +enum TokenByteDecoder { + ByteLevel, + RawUtf8WithHexBytes, +} + +impl TokenByteDecoder { + fn decode_token_bytes(self, token: &str) -> Vec { + match self { + Self::ByteLevel => decode_token_bytes(token), + Self::RawUtf8WithHexBytes => decode_raw_utf8_token_bytes(token), + } + } + + fn is_byte_fallback_token(self, token: &str) -> bool { + self == Self::RawUtf8WithHexBytes && decode_raw_utf8_byte_fallback(token).is_some() + } +} + +fn decode_raw_utf8_token_bytes(token: &str) -> Vec { + if let Some(byte) = decode_raw_utf8_byte_fallback(token) { + return vec![byte]; + } + + token.as_bytes().to_vec() +} + +fn decode_raw_utf8_byte_fallback(token: &str) -> Option { + if token.len() == 6 + && token.as_bytes()[0] == b'<' + && token.as_bytes()[1] == b'0' + && matches!(token.as_bytes()[2], b'x' | b'X') + && token.as_bytes()[5] == b'>' + { + if let (Some(high), Some(low)) = ( + hex_value(token.as_bytes()[3]), + hex_value(token.as_bytes()[4]), + ) { + return Some((high << 4) | low); + } + } + None +} + +fn hex_value(byte: u8) -> Option { + match byte { + b'0'..=b'9' => Some(byte - b'0'), + b'a'..=b'f' => Some(byte - b'a' + 10), + b'A'..=b'F' => Some(byte - b'A' + 10), + _ => None, + } +} + fn byte_level_char_to_byte(ch: char) -> Option { let codepoint = ch as u32; if (33..=126).contains(&codepoint) @@ -1474,6 +1636,20 @@ mod tests { assert_eq!(decode_token_bytes("Ā"), &[0]); } + #[test] + fn raw_utf8_hex_byte_tokens_decode_to_original_bytes() { + let vocab = [("<0x20>", 0), ("<0xC3>", 1), ("<0xA9>", 2), ("é", 3)] + .into_iter() + .map(|(token, id)| (token.to_owned(), id)); + let tokenizer = + FastBpe::from_raw_utf8_vocab_and_merges(vocab, ["<0xC3> <0xA9>".to_owned()], false) + .unwrap(); + + assert_eq!(tokenizer.tokenize(" é".as_bytes()).unwrap(), vec![0, 3]); + assert_eq!(tokenizer.token_bytes(0), Some(b" ".as_slice())); + assert_eq!(tokenizer.token_bytes(3), Some("é".as_bytes())); + } + #[test] fn levelized_tokenization_matches_reference_for_long_input() { let tokenizer = small_tokenizer(false); From b167e83d2df68409219c4eb24d4fff5bea8d5fdb Mon Sep 17 00:00:00 2001 From: Evan Almloff Date: Sat, 20 Jun 2026 16:27:42 -0500 Subject: [PATCH 09/16] cleanup --- interfaces/kalosm/examples/vision.rs | 36 +++++++++++++++------- models/kalosm-llama/src/model/forward.rs | 6 ++-- models/kalosm-llama/src/model/inference.rs | 1 + models/kalosm-llama/src/raw/mod.rs | 17 +++++----- 4 files changed, 38 insertions(+), 22 deletions(-) diff --git a/interfaces/kalosm/examples/vision.rs b/interfaces/kalosm/examples/vision.rs index 7ec920b37..a8b2571dc 100644 --- a/interfaces/kalosm/examples/vision.rs +++ b/interfaces/kalosm/examples/vision.rs @@ -5,11 +5,21 @@ use std::time::Instant; async fn main() { tracing_subscriber::fmt::init(); let t_load_start = Instant::now(); - let model = Llama::builder() - .with_source(LlamaSource::gemma_4_e2b_it_qat_chat()) - .build() - .await - .unwrap(); + let mut builder = Llama::builder().with_source( + LlamaSource::gemma_4_e2b_it_qat_chat().with_vision_model(FileSource::HuggingFace { + model_id: "unsloth/gemma-4-E2B-it-qat-GGUF".into(), + revision: "main".into(), + file: "mmproj-F16.gguf".into(), + }), + ); + builder = if std::env::var_os("KALOSM_VISION_CPU").is_some() { + builder.with_device(Device::Cpu) + } else { + builder.with_device(Device::gpu().await.expect( + "The vision example requires a GPU by default; set KALOSM_VISION_CPU=1 to run the slow CPU path.", + )) + }; + let model = builder.build().await.unwrap(); tracing::info!("[timing] model load: {:.2?}", t_load_start.elapsed()); let mut chat = model.chat(); @@ -23,18 +33,22 @@ async fn main() { } else if let Ok(path) = std::env::var("KALOSM_VISION_IMAGE") { MediaSource::file(path).unwrap() } else { - MediaSource::url("https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-VL/assets/demo.jpeg") + MediaSource::bytes(include_bytes!("landscape.jpg").as_slice()) }; - let mut response = chat(&( - MediaChunk::new(image_source, MediaType::Image), - "Describe this image.", - )); + let mut sampler = GenerationParameters::new() + .with_standard_sampler() + .with_temperature(0.0); if let Some(seed) = std::env::var("KALOSM_VISION_SEED") .ok() .and_then(|seed| seed.parse::().ok()) { - response = response.with_sampler(GenerationParameters::new().with_seed(seed)); + sampler = sampler.with_seed(seed); } + let mut response = chat(&( + MediaChunk::new(image_source, MediaType::Image), + "Describe this image.", + )) + .with_sampler(sampler); let mut first_token_at: Option = None; let mut token_count = 0u64; let t_prefill = Instant::now(); diff --git a/models/kalosm-llama/src/model/forward.rs b/models/kalosm-llama/src/model/forward.rs index d2224a354..fc6dd6a4c 100644 --- a/models/kalosm-llama/src/model/forward.rs +++ b/models/kalosm-llama/src/model/forward.rs @@ -46,7 +46,7 @@ where }; let token_start = trace_enabled.then(Instant::now); let build_start = trace_enabled.then(Instant::now); - let logits = if !images.is_empty() && model.should_chunk_multimodal_prompt(device) { + let logits = if !images.is_empty() && model.should_chunk_multimodal_prompt() { model.forward_chunked_multimodal(tokens, images, device, cache) } else { model.forward(tokens, images, device, cache) @@ -234,7 +234,7 @@ where } if gpu_fused_logits_sampling_enabled() - && (images.is_empty() || !model.should_chunk_multimodal_prompt(device)) + && (images.is_empty() || !model.should_chunk_multimodal_prompt()) { return Self::forward_sample_token_fused_logits( ForwardInputs { @@ -333,7 +333,7 @@ where return Ok(None); } - if !images.is_empty() && model.should_chunk_multimodal_prompt(device) { + if !images.is_empty() && model.should_chunk_multimodal_prompt() { let logits = model.forward_chunked_multimodal(tokens, images, device, cache)?; let logits: fusor::Tensor<1, f32> = logits.squeeze(0).cast(); return Self::forward_sample_token_pending_from_logits( diff --git a/models/kalosm-llama/src/model/inference.rs b/models/kalosm-llama/src/model/inference.rs index ae9ede11a..68f685179 100644 --- a/models/kalosm-llama/src/model/inference.rs +++ b/models/kalosm-llama/src/model/inference.rs @@ -46,6 +46,7 @@ where } if mtp_speculative_enabled() + && images.is_empty() && stop_on.is_none() && sampler.sampling_strategy == kalosm_language_model::SamplingStrategy::Standard && sampler.temperature <= 0.0 diff --git a/models/kalosm-llama/src/raw/mod.rs b/models/kalosm-llama/src/raw/mod.rs index fb9063b61..9bda8ac5a 100644 --- a/models/kalosm-llama/src/raw/mod.rs +++ b/models/kalosm-llama/src/raw/mod.rs @@ -1039,6 +1039,10 @@ where images.push(image); grid_thw.push(thw) } + } else if !raw_images.is_empty() { + return Err(fusor::Error::msg( + "Media inputs require a loaded vision encoder.", + )); } // Add image padding tokens for any placeholders in the prompt. @@ -1251,7 +1255,7 @@ where Ok(result_f32.cast()) } - pub(crate) fn should_chunk_multimodal_prompt(&self, device: &Device) -> bool { + pub(crate) fn should_chunk_multimodal_prompt(&self) -> bool { if std::env::var_os("KALOSM_LLAMA_DISABLE_MULTIMODAL_CHUNK").is_some() { return false; } @@ -1260,16 +1264,13 @@ where } #[cfg(feature = "vision")] { - device.is_gpu() - && matches!( - self.vision_encoder, - Some(vision::VisionTransformer::Gemma(_)) - ) - && self.config.image_pad_token.is_some() + matches!( + self.vision_encoder, + Some(vision::VisionTransformer::Gemma(_)) + ) && self.config.image_pad_token.is_some() } #[cfg(not(feature = "vision"))] { - let _ = device; false } } From 01f6491381b13666f27794aa9a58bfeb5cb80f52 Mon Sep 17 00:00:00 2001 From: Evan Almloff Date: Sat, 20 Jun 2026 16:39:12 -0500 Subject: [PATCH 10/16] fix vision example --- interfaces/kalosm/examples/vision.rs | 9 +++------ models/kalosm-llama/src/source.rs | 5 +++++ 2 files changed, 8 insertions(+), 6 deletions(-) diff --git a/interfaces/kalosm/examples/vision.rs b/interfaces/kalosm/examples/vision.rs index a8b2571dc..0cea4f116 100644 --- a/interfaces/kalosm/examples/vision.rs +++ b/interfaces/kalosm/examples/vision.rs @@ -6,11 +6,7 @@ async fn main() { tracing_subscriber::fmt::init(); let t_load_start = Instant::now(); let mut builder = Llama::builder().with_source( - LlamaSource::gemma_4_e2b_it_qat_chat().with_vision_model(FileSource::HuggingFace { - model_id: "unsloth/gemma-4-E2B-it-qat-GGUF".into(), - revision: "main".into(), - file: "mmproj-F16.gguf".into(), - }), + LlamaSource::gemma_4_e2b_it_qat_chat(), ); builder = if std::env::var_os("KALOSM_VISION_CPU").is_some() { builder.with_device(Device::Cpu) @@ -26,7 +22,7 @@ async fn main() { let max_tokens = std::env::var("KALOSM_VISION_MAX_TOKENS") .ok() .and_then(|value| value.parse::().ok()) - .unwrap_or(64); + .unwrap_or(1024); let t_total = Instant::now(); let image_source = if let Ok(url) = std::env::var("KALOSM_VISION_URL") { MediaSource::url(url) @@ -62,6 +58,7 @@ async fn main() { } token_count += 1; print!("{}", token); + std::io::Write::flush(&mut std::io::stdout()).unwrap(); if token_count >= max_tokens { break; } diff --git a/models/kalosm-llama/src/source.rs b/models/kalosm-llama/src/source.rs index dacdb6443..8c23fb820 100644 --- a/models/kalosm-llama/src/source.rs +++ b/models/kalosm-llama/src/source.rs @@ -1000,6 +1000,11 @@ impl LlamaSource { "main".to_string(), "MTP/gemma-4-E2B-it-Q4_0-MTP.gguf".to_string(), )) + .with_vision_model(FileSource::huggingface( + "unsloth/gemma-4-E2B-it-qat-GGUF".to_string(), + "main".to_string(), + "mmproj-F16.gguf".to_string(), + )) .with_override_stop_token_string("") } From 2dca83de852b6328411d53abb5c4ee34a1146b59 Mon Sep 17 00:00:00 2001 From: Evan Almloff Date: Sat, 20 Jun 2026 17:18:26 -0500 Subject: [PATCH 11/16] pull out mask --- models/kalosm-llama/src/raw/mod.rs | 60 +++++++++++++++++++++--------- 1 file changed, 42 insertions(+), 18 deletions(-) diff --git a/models/kalosm-llama/src/raw/mod.rs b/models/kalosm-llama/src/raw/mod.rs index 9bda8ac5a..10502eef9 100644 --- a/models/kalosm-llama/src/raw/mod.rs +++ b/models/kalosm-llama/src/raw/mod.rs @@ -214,6 +214,40 @@ pub const DEFAULT_ROPE_FREQUENCY: f32 = 1_000_000.; pub const GEMMA_DEFAULT_SLIDING_WINDOW_TYPE: usize = 6; pub const GEMMA_DEFAULT_ROPE_FREQUENCY_SLIDING: f32 = 10_000.; +/// Build the additive attention-mask values (`0.0` allowed, `-inf` blocked) +/// for a `[seq_len, index_pos + seq_len]` score matrix. +/// +/// Tokens are causal by default. Any query/key position that falls inside the +/// same entry of `non_causal_token_ranges` may attend to its peers regardless +/// of ordering (this is how image-token blocks attend bidirectionally), and an +/// optional sliding `window` blocks keys older than `window` positions. +fn non_causal_mask_data( + seq_len: usize, + index_pos: usize, + sliding_window_size: Option, + non_causal_token_ranges: &[Range], +) -> Vec { + let cols = index_pos + seq_len; + let mut mask_data = vec![0.0_f32; seq_len * cols]; + for row in 0..seq_len { + let global_row = index_pos + row; + for col in 0..cols { + let same_non_causal_range = col >= index_pos + && non_causal_token_ranges + .iter() + .any(|range| range.contains(&row) && range.contains(&(col - index_pos))); + let future = col > global_row && !same_non_causal_range; + let outside_window = sliding_window_size + .map(|window| col + window <= global_row) + .unwrap_or(false); + if future || outside_window { + mask_data[row * cols + col] = f32::NEG_INFINITY; + } + } + } + mask_data +} + /// The configuration of a Llama model. pub struct LlamaConfig { pub(crate) rope_freq_weight: Option>, @@ -1483,7 +1517,8 @@ where seq_len, index_pos, pos_ids: None, - non_causal_token_ranges: vec![0..seq_len], + // The whole image chunk attends bidirectionally. + non_causal_token_ranges: std::iter::once(0..seq_len).collect(), }; self.forward_last_hidden_from_embeddings(encoded, device, cache, Some(t_encode.elapsed())) } @@ -1641,23 +1676,12 @@ where } let cols = index_pos + seq_len; - let mut mask_data = vec![0.0_f32; seq_len * cols]; - for row in 0..seq_len { - let global_row = index_pos + row; - for col in 0..cols { - let same_non_causal_range = col >= index_pos - && non_causal_token_ranges - .iter() - .any(|range| range.contains(&row) && range.contains(&(col - index_pos))); - let future = col > global_row && !same_non_causal_range; - let outside_window = sliding_window_size - .map(|window| col + window <= global_row) - .unwrap_or(false); - if future || outside_window { - mask_data[row * cols + col] = f32::NEG_INFINITY; - } - } - } + let mask_data = non_causal_mask_data( + seq_len, + index_pos, + sliding_window_size, + non_causal_token_ranges, + ); let mask: Tensor<2, f32> = Tensor::new(device, mask_data.as_slice()) .reshape([seq_len, cols]) .to_concrete(); From ce870c9b96b320612603ab9ec02245cea89562c2 Mon Sep 17 00:00:00 2001 From: Evan Almloff Date: Sat, 20 Jun 2026 17:32:36 -0500 Subject: [PATCH 12/16] Fix row program merge import --- fusor-ml/core/src/row_program.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/fusor-ml/core/src/row_program.rs b/fusor-ml/core/src/row_program.rs index 270657433..6d8ee5e09 100644 --- a/fusor-ml/core/src/row_program.rs +++ b/fusor-ml/core/src/row_program.rs @@ -41,7 +41,7 @@ use crate::{ nary_wise::{NaryExpr, NaryFunction, NaryOp, NaryScalar, UnaryFunctionChain}, reduce::{ReduceFunction, ReduceOp, ReduceOperation, max_fn, sum_fn}, tensor::{DataTypeEnum, TensorData}, - visit_tiled::MaybeQData, + visit_tiled::{MaybeQData, distribute_workgroups}, }; const BLOCK: u32 = 256; From 29b68a144538f37263e81384e88b965ccd5b3a22 Mon Sep 17 00:00:00 2001 From: Evan Almloff Date: Sat, 20 Jun 2026 18:02:51 -0500 Subject: [PATCH 13/16] fix formatting --- interfaces/kalosm/examples/vision.rs | 4 +- models/kalosm-llama/src/model/forward.rs | 10 +--- models/kalosm-llama/src/model/mod.rs | 4 +- .../kalosm-llama/src/raw/attention_layer.rs | 3 +- models/kalosm-llama/src/raw/mod.rs | 28 ++++++---- models/kalosm-llama/src/raw/mtp.rs | 5 +- models/kalosm-tokenizer/src/lib.rs | 51 +++++++++++++++++++ 7 files changed, 78 insertions(+), 27 deletions(-) diff --git a/interfaces/kalosm/examples/vision.rs b/interfaces/kalosm/examples/vision.rs index 0cea4f116..0e5ac5cff 100644 --- a/interfaces/kalosm/examples/vision.rs +++ b/interfaces/kalosm/examples/vision.rs @@ -5,9 +5,7 @@ use std::time::Instant; async fn main() { tracing_subscriber::fmt::init(); let t_load_start = Instant::now(); - let mut builder = Llama::builder().with_source( - LlamaSource::gemma_4_e2b_it_qat_chat(), - ); + let mut builder = Llama::builder().with_source(LlamaSource::gemma_4_e2b_it_qat_chat()); builder = if std::env::var_os("KALOSM_VISION_CPU").is_some() { builder.with_device(Device::Cpu) } else { diff --git a/models/kalosm-llama/src/model/forward.rs b/models/kalosm-llama/src/model/forward.rs index fc6dd6a4c..fcbb6e7b1 100644 --- a/models/kalosm-llama/src/model/forward.rs +++ b/models/kalosm-llama/src/model/forward.rs @@ -426,10 +426,7 @@ where previous_gpu_token: Option<&fusor::GpuSampledToken>, top_k: usize, ) -> Result, LlamaModelError> { - let logits = hidden - .squeeze(0) - .to_concrete() - .q_mat_mul(model.output_matrix()); + let logits = model.logits_from_hidden_f32(hidden.squeeze(0).to_concrete()); Self::forward_sample_token_pending_from_logits( logits, sampler, @@ -477,10 +474,7 @@ where Ok(hidden) => hidden, Err(err) => return Box::pin(async move { Err(err.into()) }), }; - let logits = hidden - .squeeze(0) - .to_concrete() - .q_mat_mul(model.output_matrix()); + let logits = model.logits_from_hidden_f32(hidden.squeeze(0).to_concrete()); let mut kernels = 0; if trace { if let Some(gpu_logits) = logits.as_gpu() { diff --git a/models/kalosm-llama/src/model/mod.rs b/models/kalosm-llama/src/model/mod.rs index 55420cc8b..20fcb18ee 100644 --- a/models/kalosm-llama/src/model/mod.rs +++ b/models/kalosm-llama/src/model/mod.rs @@ -756,8 +756,8 @@ where random: 0.5, }; - let _ = hidden - .q_mat_mul(model.output_matrix()) + let _ = model + .logits_from_hidden_f32(hidden) .sample_mirostat2_token(&mut sampler, &[], params) .await; if let Some(start) = start { diff --git a/models/kalosm-llama/src/raw/attention_layer.rs b/models/kalosm-llama/src/raw/attention_layer.rs index 69825fc5d..7044b95ae 100644 --- a/models/kalosm-llama/src/raw/attention_layer.rs +++ b/models/kalosm-llama/src/raw/attention_layer.rs @@ -854,8 +854,9 @@ fn pad_attention_mask_to_kv_len( } if cols > kv_seq_len { + let start_col = cols - kv_seq_len; return Some(AttentionMask::new( - mask.narrow(1, 0, kv_seq_len).to_concrete(), + mask.narrow(1, start_col, kv_seq_len).to_concrete(), )); } diff --git a/models/kalosm-llama/src/raw/mod.rs b/models/kalosm-llama/src/raw/mod.rs index 10502eef9..e18071100 100644 --- a/models/kalosm-llama/src/raw/mod.rs +++ b/models/kalosm-llama/src/raw/mod.rs @@ -1279,14 +1279,25 @@ where where f32: CastTo + CastTensor, { - let mut result_f32 = x_f32.q_mat_mul(&self.output); + Ok(self.logits_from_hidden_f32(x_f32).cast()) + } + + pub(crate) fn logits_from_hidden_f32( + &self, + x_f32: Tensor, + ) -> Tensor { + self.apply_final_logit_softcap(x_f32.q_mat_mul(&self.output)) + } + + pub(crate) fn apply_final_logit_softcap( + &self, + logits: Tensor, + ) -> Tensor { if let Some(softcap) = self.config.final_logit_softcapping { - result_f32 = result_f32 - .mul_scalar(1.0 / softcap) - .tanh() - .mul_scalar(softcap); + logits.mul_scalar(1.0 / softcap).tanh().mul_scalar(softcap) + } else { + logits } - Ok(result_f32.cast()) } pub(crate) fn should_chunk_multimodal_prompt(&self) -> bool { @@ -1576,10 +1587,7 @@ where )?; let normed = self.norm.forward_generic(&pre_norm.hidden); let h_nextn: Tensor<2, f32> = normed.clone().cast::().squeeze(0).to_concrete(); - let mut logits = normed.cast::().q_mat_mul(&self.output); - if let Some(softcap) = self.config.final_logit_softcapping { - logits = logits.mul_scalar(1.0 / softcap).tanh().mul_scalar(softcap); - } + let logits = self.logits_from_hidden_f32(normed.cast::()); let logits: Tensor<2, f32> = logits.squeeze(0).to_concrete(); Ok(TargetBatchOutput { logits, h_nextn }) } diff --git a/models/kalosm-llama/src/raw/mtp.rs b/models/kalosm-llama/src/raw/mtp.rs index 719d950f8..a1e802870 100644 --- a/models/kalosm-llama/src/raw/mtp.rs +++ b/models/kalosm-llama/src/raw/mtp.rs @@ -392,9 +392,8 @@ where } let normed = self.norm.forward_generic(&layer_in); - let logits: Tensor<1, f32> = normed - .cast::() - .q_mat_mul(&self.output) + let logits: Tensor<1, f32> = target + .apply_final_logit_softcap(normed.cast::().q_mat_mul(&self.output)) .squeeze(0) .squeeze(0) .to_concrete(); diff --git a/models/kalosm-tokenizer/src/lib.rs b/models/kalosm-tokenizer/src/lib.rs index 8a30e01aa..ee72fbfed 100644 --- a/models/kalosm-tokenizer/src/lib.rs +++ b/models/kalosm-tokenizer/src/lib.rs @@ -372,6 +372,12 @@ impl FastBpe { right: u32, new_token: u32, ) -> Result<(), TokenizerError> { + if self.raw_utf8_initial_tokens { + self.encode_raw_utf8_chars_into(input, out)?; + apply_single_merge(out, left, right, new_token); + return Ok(()); + } + if self.all_bytes_present { self.encode_bytes_and_apply_single_merge_unchecked(input, out, left, right, new_token); Ok(()) @@ -774,6 +780,25 @@ fn apply_greedy_merges(merges: &FxHashMap, tokens: &mut Vec< } } +fn apply_single_merge(tokens: &mut Vec, left: u32, right: u32, new_token: u32) { + let len = tokens.len(); + let mut read = 0; + let mut write = 0; + + while read < len { + if read + 1 < len && tokens[read] == left && tokens[read + 1] == right { + tokens[write] = new_token; + read += 2; + } else { + tokens[write] = tokens[read]; + read += 1; + } + write += 1; + } + + tokens.truncate(write); +} + fn rebuild_merge_candidates( merges: &MergeLookup, tokens: &[u32], @@ -1650,6 +1675,32 @@ mod tests { assert_eq!(tokenizer.token_bytes(3), Some("é".as_bytes())); } + #[test] + fn raw_utf8_single_merge_fast_path_keeps_direct_char_tokens() { + let vocab = [ + ("<0x20>", 0), + ("<0xC3>", 1), + ("<0xA9>", 2), + ("é", 3), + ("a", 4), + ("b", 5), + ("ab", 6), + ] + .into_iter() + .map(|(token, id)| (token.to_owned(), id)); + let tokenizer = + FastBpe::from_raw_utf8_vocab_and_merges(vocab, ["a b".to_owned()], false).unwrap(); + + assert_eq!( + tokenizer.tokenize(" éab".as_bytes()).unwrap(), + vec![0, 3, 6] + ); + assert_eq!( + tokenizer.tokenize(" éab".as_bytes()).unwrap(), + tokenizer.tokenize_reference(" éab".as_bytes()).unwrap() + ); + } + #[test] fn levelized_tokenization_matches_reference_for_long_input() { let tokenizer = small_tokenizer(false); From 0efbdc37b2284556d6eae62a277226f35cbbd2c9 Mon Sep 17 00:00:00 2001 From: Evan Almloff Date: Sat, 20 Jun 2026 18:32:33 -0500 Subject: [PATCH 14/16] more refactoring --- models/kalosm-llama/src/model/inference.rs | 182 ++++------------ models/kalosm-llama/src/model/mod.rs | 4 + .../kalosm-llama/src/raw/attention_layer.rs | 5 + models/kalosm-llama/src/raw/mod.rs | 198 ++++++++---------- models/kalosm-llama/src/raw/mtp.rs | 3 + 5 files changed, 141 insertions(+), 251 deletions(-) diff --git a/models/kalosm-llama/src/model/inference.rs b/models/kalosm-llama/src/model/inference.rs index 68f685179..2e5e87113 100644 --- a/models/kalosm-llama/src/model/inference.rs +++ b/models/kalosm-llama/src/model/inference.rs @@ -1,5 +1,22 @@ use super::*; +/// Yield once to the async runtime so long generation loops stay cooperative +/// between GPU dispatches without spinning the executor. +async fn yield_once() { + use std::sync::atomic::{AtomicBool, Ordering}; + let yielded = AtomicBool::new(false); + std::future::poll_fn(|cx| { + if yielded.load(Ordering::Relaxed) { + std::task::Poll::Ready(()) + } else { + yielded.store(true, Ordering::Relaxed); + cx.waker().wake_by_ref(); + std::task::Poll::Pending + } + }) + .await; +} + impl LlamaModel where F: CastTo + CastTensor + WasmNotSend + WasmNotSync + 'static, @@ -8,6 +25,21 @@ where AddOp: SimdBinaryOp, SumOp: SimdReduceOp, { + /// Emit a `[sampled_token]` trace line when `KALOSM_TRACE_SAMPLED_TOKEN` is + /// set. No-op (beyond an env lookup) otherwise. + fn trace_sampled_token(&self, index: impl std::fmt::Display, token: u32, stop_tokens: &[u32]) { + if std::env::var_os("KALOSM_TRACE_SAMPLED_TOKEN").is_none() { + return; + } + let decoded = self + .tokenizer + .decode(&[token], false) + .unwrap_or_else(|err| format!("")); + eprintln!( + "[sampled_token] index={index} id={token} text={decoded:?} stop_tokens={stop_tokens:?}" + ); + } + pub(crate) async fn _infer( &mut self, settings: InferenceSettings, @@ -133,15 +165,7 @@ where "pending GPU sampler refused slow fallback".into(), ) })?; - if std::env::var_os("KALOSM_TRACE_SAMPLED_TOKEN").is_some() { - let decoded = self - .tokenizer - .decode(&[new_token], false) - .unwrap_or_else(|err| format!("")); - eprintln!( - "[sampled_token] index={tokens_generated} id={new_token} text={decoded:?} stop_tokens={stop_tokens:?}" - ); - } + self.trace_sampled_token(tokens_generated, new_token, stop_tokens); if stop_tokens.contains(&new_token) { tracing::trace!("Stopping on stop token"); break; @@ -193,20 +217,7 @@ where break; } - { - use std::sync::atomic::{AtomicBool, Ordering}; - let yielded = AtomicBool::new(false); - std::future::poll_fn(|cx| { - if yielded.load(Ordering::Relaxed) { - std::task::Poll::Ready(()) - } else { - yielded.store(true, Ordering::Relaxed); - cx.waker().wake_by_ref(); - std::task::Poll::Pending - } - }) - .await; - } + yield_once().await; } return Ok(()); @@ -240,15 +251,7 @@ where let mut tokens_generated = 0; while !finished.is_canceled() && tokens_generated < max_tokens { let new_token = next_token; - if std::env::var_os("KALOSM_TRACE_SAMPLED_TOKEN").is_some() { - let decoded = self - .tokenizer - .decode(&[new_token], false) - .unwrap_or_else(|err| format!("")); - eprintln!( - "[sampled_token] index={tokens_generated} id={new_token} text={decoded:?} stop_tokens={stop_tokens:?}" - ); - } + self.trace_sampled_token(tokens_generated, new_token, stop_tokens); if stop_tokens.contains(&new_token) { tracing::trace!("Stopping on stop token"); break; @@ -288,20 +291,7 @@ where } .await?; - { - use std::sync::atomic::{AtomicBool, Ordering}; - let yielded = AtomicBool::new(false); - std::future::poll_fn(|cx| { - if yielded.load(Ordering::Relaxed) { - std::task::Poll::Ready(()) - } else { - yielded.store(true, Ordering::Relaxed); - cx.waker().wake_by_ref(); - std::task::Poll::Pending - } - }) - .await; - } + yield_once().await; } return Ok(()); @@ -338,15 +328,7 @@ where let new_token = text_stream .sample_token(&mut cpu_sampler, logits, stop_on.as_deref(), sample_top_k) .map_err(LlamaModelError::TokenOutputStreamError)?; - if std::env::var_os("KALOSM_TRACE_SAMPLED_TOKEN").is_some() { - let decoded = self - .tokenizer - .decode(&[new_token], false) - .unwrap_or_else(|err| format!("")); - eprintln!( - "[sampled_token] index={tokens_generated} id={new_token} text={decoded:?} stop_tokens={stop_tokens:?}" - ); - } + self.trace_sampled_token(tokens_generated, new_token, stop_tokens); if stop_tokens.contains(&new_token) { tracing::trace!("Stopping on stop token"); break; @@ -427,20 +409,7 @@ where .await?; logits = logits_from_sorted_top_k(logit_probs); // Yield control to allow the stream to deliver tokens - { - use std::sync::atomic::{AtomicBool, Ordering}; - let yielded = AtomicBool::new(false); - std::future::poll_fn(|cx| { - if yielded.load(Ordering::Relaxed) { - std::task::Poll::Ready(()) - } else { - yielded.store(true, Ordering::Relaxed); - cx.waker().wake_by_ref(); - std::task::Poll::Pending - } - }) - .await; - } + yield_once().await; } // Flush the queued text @@ -500,15 +469,7 @@ where let mut accepted_total = 0usize; while !finished.is_canceled() && tokens_generated < max_tokens as usize { let new_token = next_token; - if std::env::var_os("KALOSM_TRACE_SAMPLED_TOKEN").is_some() { - let decoded = self - .tokenizer - .decode(&[new_token], false) - .unwrap_or_else(|err| format!("")); - eprintln!( - "[sampled_token] index={tokens_generated} id={new_token} text={decoded:?} stop_tokens={stop_tokens:?}" - ); - } + self.trace_sampled_token(tokens_generated, new_token, stop_tokens); if stop_tokens.contains(&new_token) { tracing::trace!("Stopping on stop token"); break; @@ -645,20 +606,7 @@ where .await; } - { - use std::sync::atomic::{AtomicBool, Ordering}; - let yielded = AtomicBool::new(false); - std::future::poll_fn(|cx| { - if yielded.load(Ordering::Relaxed) { - std::task::Poll::Ready(()) - } else { - yielded.store(true, Ordering::Relaxed); - cx.waker().wake_by_ref(); - std::task::Poll::Pending - } - }) - .await; - } + yield_once().await; } if std::env::var_os("KALOSM_TRACE_MTP").is_some() { tracing::info!("mtp_summary drafted={drafted_total} accepted={accepted_total}"); @@ -681,15 +629,7 @@ where let top_k = gpu_sample_top_k(&gpu_sampler.config); while !finished.is_canceled() && tokens_generated < max_tokens as usize { let new_token = next_token; - if std::env::var_os("KALOSM_TRACE_SAMPLED_TOKEN").is_some() { - let decoded = self - .tokenizer - .decode(&[new_token], false) - .unwrap_or_else(|err| format!("")); - eprintln!( - "[sampled_token] index={tokens_generated} id={new_token} text={decoded:?} stop_tokens={stop_tokens:?}" - ); - } + self.trace_sampled_token(tokens_generated, new_token, stop_tokens); if stop_tokens.contains(&new_token) { tracing::trace!("Stopping on stop token"); break; @@ -761,20 +701,7 @@ where } .await?; - { - use std::sync::atomic::{AtomicBool, Ordering}; - let yielded = AtomicBool::new(false); - std::future::poll_fn(|cx| { - if yielded.load(Ordering::Relaxed) { - std::task::Poll::Ready(()) - } else { - yielded.store(true, Ordering::Relaxed); - cx.waker().wake_by_ref(); - std::task::Poll::Pending - } - }) - .await; - } + yield_once().await; } Ok(()) } @@ -827,15 +754,7 @@ where "pending GPU sampler refused slow fallback".into(), ) })?; - if std::env::var_os("KALOSM_TRACE_SAMPLED_TOKEN").is_some() { - let decoded = self - .tokenizer - .decode(&[new_token], false) - .unwrap_or_else(|err| format!("")); - eprintln!( - "[sampled_token] index={tokens_generated} id={new_token} text={decoded:?} stop_tokens={stop_tokens:?}" - ); - } + self.trace_sampled_token(tokens_generated, new_token, stop_tokens); if stop_tokens.contains(&new_token) { tracing::trace!("Stopping on stop token"); break; @@ -887,20 +806,7 @@ where break; } - { - use std::sync::atomic::{AtomicBool, Ordering}; - let yielded = AtomicBool::new(false); - std::future::poll_fn(|cx| { - if yielded.load(Ordering::Relaxed) { - std::task::Poll::Ready(()) - } else { - yielded.store(true, Ordering::Relaxed); - cx.waker().wake_by_ref(); - std::task::Poll::Pending - } - }) - .await; - } + yield_once().await; } Ok(()) } diff --git a/models/kalosm-llama/src/model/mod.rs b/models/kalosm-llama/src/model/mod.rs index 20fcb18ee..f6c331d42 100644 --- a/models/kalosm-llama/src/model/mod.rs +++ b/models/kalosm-llama/src/model/mod.rs @@ -158,6 +158,10 @@ fn gpu_sample_top_k(config: &GpuSamplerConfig) -> usize { { return top_k.max(1); } + // Greedy standard sampling (temperature <= 0) is argmax, so a top-1 cut is + // both exact and cheaper. This intentionally ignores any configured `top_k`, + // which has no effect once the distribution collapses to its max; the + // `KALOSM_LLAMA_GPU_SAMPLE_TOP_K` override above still wins if set. if config.sampling_strategy == kalosm_language_model::SamplingStrategy::Standard && config.temperature <= 0.0 { diff --git a/models/kalosm-llama/src/raw/attention_layer.rs b/models/kalosm-llama/src/raw/attention_layer.rs index 7044b95ae..c3c94ec11 100644 --- a/models/kalosm-llama/src/raw/attention_layer.rs +++ b/models/kalosm-llama/src/raw/attention_layer.rs @@ -482,6 +482,11 @@ where let hidden_f32 = hidden_states.cast::(); let query_states: Tensor<4, F> = if let Some(attention_qkv) = &self.attention_qkv { + // Shared-KV callers only need Q, but a fused QKV weight forces us to + // project K/V as well and discard them. No current shared-KV model + // (Gemma 4 or the MTP assistant) uses a fused QKV weight, so this + // branch is effectively unreachable today; if a future one does, it + // pays a ~3x projection here and should grow a Q-only weight slice. let query_width = num_heads * head_dim; let mut qkv = hidden_f32.q_mat_mul(attention_qkv); if let Some(bias) = &self.bias { diff --git a/models/kalosm-llama/src/raw/mod.rs b/models/kalosm-llama/src/raw/mod.rs index e18071100..30ee2cb1a 100644 --- a/models/kalosm-llama/src/raw/mod.rs +++ b/models/kalosm-llama/src/raw/mod.rs @@ -130,9 +130,6 @@ pub(crate) fn debug_check_nan_f32( ) { } -#[cfg(not(feature = "vision"))] -pub(crate) fn debug_tensor_stats_f32(_: &fusor::Tensor, _: &str) {} - #[cfg(not(target_arch = "wasm32"))] fn resolve_intermediate_hidden_f32(tensor: &fusor::Tensor<2, f32>) { let marker = tensor.clone().mul_scalar(1.0).to_concrete(); @@ -191,7 +188,7 @@ use fusor::layers::RmsNorm; use fusor::QMatrix; use fusor::ShardedVarBuilder; use fusor::{ - AddOp, CastTensor, CastTo, FloatDataType, FloatOps, MatmulImpl, MulOp, SimdBinaryOp, + AddOp, CastTensor, CastTo, FloatDataType, FloatOps, Fusion, MatmulImpl, MulOp, SimdBinaryOp, SimdElement, SimdReduceOp, SumOp, }; use fusor::{AsyncReadRange, AsyncShardedVarBuilder}; @@ -965,6 +962,10 @@ where .await .ok() .map(&dequantize_1d); + // Gemma 4 folds the query pre-attention scaling into the exported + // `attn_q_norm` weights, so the softmax logits are already scaled and + // flash-attention must run with a unit scale. Every other supported + // architecture applies the usual 1/sqrt(head_dim) here. let attention_scale = if is_gemma4 { 1.0 } else { @@ -1051,6 +1052,73 @@ where } } + /// Compute the Gemma "per-layer input" embeddings that are blended into each + /// decoder layer (the `inp_gate`/`proj`/`post_norm` path). Returns `None` + /// for models without per-layer embeddings. + /// + /// `per_layer_token_ids` is invoked lazily, so models without per-layer + /// embeddings never pay for building the token id tensor. It yields the + /// `[batch, positions]` ids used for the per-layer token lookup, with + /// image/control tokens already zeroed by the caller. A single position is + /// broadcast across the whole sequence (image chunks share one zeroed + /// per-layer token). + fn compute_per_layer_inputs( + &self, + embeddings_f32: &Tensor<3, f32>, + per_layer_token_ids: impl FnOnce() -> Tensor<2, u32, B>, + ) -> Option> + where + B: Fusion<2, u32>, + { + let ( + per_layer_tok_embeddings, + per_layer_model_proj, + per_layer_proj_norm, + per_layer_embedding_length, + ) = match ( + &self.per_layer_tok_embeddings, + &self.per_layer_model_proj, + &self.per_layer_proj_norm, + self.config.per_layer_embedding_length, + ) { + (Some(embeddings), Some(model_proj), Some(proj_norm), Some(length)) => { + (embeddings, model_proj, proj_norm, length) + } + _ => return None, + }; + + let [batch, seq, embedding_dim] = embeddings_f32.shape(); + let n_layer = self.config.n_layer; + let per_layer_token_ids = per_layer_token_ids(); + let positions = per_layer_token_ids.shape()[1]; + + let token_inputs = per_layer_tok_embeddings.forward::<2, 3, _>(&per_layer_token_ids) + * (per_layer_embedding_length as f32).sqrt(); + let token_inputs = + token_inputs.reshape([batch, positions, n_layer, per_layer_embedding_length]); + let token_inputs: Tensor<4, f32> = if positions == seq { + token_inputs.to_concrete() + } else { + token_inputs + .broadcast_as([batch, seq, n_layer, per_layer_embedding_length]) + .to_concrete() + }; + + let projected_inputs = embeddings_f32.q_mat_mul(per_layer_model_proj) + * (1.0 / (embedding_dim as f32).sqrt()); + let projected_inputs = projected_inputs + .reshape([batch, seq, n_layer, per_layer_embedding_length]) + .to_concrete(); + let projected_inputs: Tensor<4, F> = + per_layer_proj_norm.forward_generic_4d(&projected_inputs.cast()); + + Some( + ((projected_inputs.cast::() + token_inputs) * (1.0 / 2.0_f32.sqrt())) + .to_concrete() + .cast(), + ) + } + pub fn encode_tokens( &self, raw_tokens: &[u32], @@ -1186,20 +1254,9 @@ where } } - let per_layer_inputs = if let ( - Some(per_layer_tok_embeddings), - Some(per_layer_model_proj), - Some(per_layer_proj_norm), - Some(per_layer_embedding_length), - ) = ( - &self.per_layer_tok_embeddings, - &self.per_layer_model_proj, - &self.per_layer_proj_norm, - self.config.per_layer_embedding_length, - ) { - let [batch, seq, embedding_dim] = embeddings_f32.shape(); + let per_layer_inputs = self.compute_per_layer_inputs(&embeddings_f32, || { #[cfg(feature = "vision")] - let per_layer_x_base = { + { let mut per_layer_tokens = tokens.clone(); for range in &image_token_ranges { per_layer_tokens[range.clone()].fill(0); @@ -1218,33 +1275,13 @@ where } } } - Tensor::new(device, per_layer_tokens.as_slice()) - }; - #[cfg(feature = "vision")] - let per_layer_x = per_layer_x_base.unsqueeze(0); + Tensor::from_slice(device, [1, per_layer_tokens.len()], per_layer_tokens.as_slice()) + } #[cfg(not(feature = "vision"))] - let per_layer_x = x.clone(); - let token_inputs = per_layer_tok_embeddings.forward::<2, 3, _>(&per_layer_x) - * (per_layer_embedding_length as f32).sqrt(); - let token_inputs = token_inputs - .reshape([batch, seq, self.config.n_layer, per_layer_embedding_length]) - .to_concrete(); - - let projected_inputs = embeddings_f32.q_mat_mul(per_layer_model_proj) - * (1.0 / (embedding_dim as f32).sqrt()); - let projected_inputs = projected_inputs - .reshape([batch, seq, self.config.n_layer, per_layer_embedding_length]) - .to_concrete(); - let projected_inputs: Tensor<4, F> = - per_layer_proj_norm.forward_generic_4d(&projected_inputs.cast()); - Some( - ((projected_inputs.cast::() + token_inputs) * (1.0 / 2.0_f32.sqrt())) - .to_concrete() - .cast(), - ) - } else { - None - }; + { + x.clone() + } + }); let embeddings: Tensor<3, F> = embeddings_f32.cast(); Ok(EncodedTokens { @@ -1484,43 +1521,10 @@ where .extend(std::iter::repeat_n(image_pad_token, seq_len)); } - let per_layer_inputs = if let ( - Some(per_layer_tok_embeddings), - Some(per_layer_model_proj), - Some(per_layer_proj_norm), - Some(per_layer_embedding_length), - ) = ( - &self.per_layer_tok_embeddings, - &self.per_layer_model_proj, - &self.per_layer_proj_norm, - self.config.per_layer_embedding_length, - ) { - let [batch, seq, embedding_dim] = embeddings_f32.shape(); - let per_layer_tokens = [0u32]; - let per_layer_x_base = Tensor::new(device, per_layer_tokens.as_slice()); - let per_layer_x = per_layer_x_base.unsqueeze(0); - let token_inputs = per_layer_tok_embeddings.forward::<2, 3, _>(&per_layer_x) - * (per_layer_embedding_length as f32).sqrt(); - let token_inputs = token_inputs - .reshape([batch, 1, self.config.n_layer, per_layer_embedding_length]) - .broadcast_as([batch, seq, self.config.n_layer, per_layer_embedding_length]) - .to_concrete(); - - let projected_inputs = embeddings_f32.q_mat_mul(per_layer_model_proj) - * (1.0 / (embedding_dim as f32).sqrt()); - let projected_inputs = projected_inputs - .reshape([batch, seq, self.config.n_layer, per_layer_embedding_length]) - .to_concrete(); - let projected_inputs: Tensor<4, F> = - per_layer_proj_norm.forward_generic_4d(&projected_inputs.cast()); - Some( - ((projected_inputs.cast::() + token_inputs) * (1.0 / 2.0_f32.sqrt())) - .to_concrete() - .cast(), - ) - } else { - None - }; + // Image chunks carry no text tokens, so every position shares a single + // zeroed per-layer token that the helper broadcasts across the chunk. + let per_layer_inputs = self + .compute_per_layer_inputs(&embeddings_f32, || Tensor::from_slice(device, [1, 1], &[0u32])); let encoded = EncodedTokens { embeddings: embeddings_f32.cast(), @@ -1622,39 +1626,7 @@ where if let Some(scale) = self.tok_embedding_scale { embeddings_f32 = (embeddings_f32 * scale).to_concrete(); } - let per_layer_inputs = if let ( - Some(per_layer_tok_embeddings), - Some(per_layer_model_proj), - Some(per_layer_proj_norm), - Some(per_layer_embedding_length), - ) = ( - &self.per_layer_tok_embeddings, - &self.per_layer_model_proj, - &self.per_layer_proj_norm, - self.config.per_layer_embedding_length, - ) { - let [batch, seq, embedding_dim] = embeddings_f32.shape(); - let token_inputs = per_layer_tok_embeddings.forward::<2, 3, _>(&x) - * (per_layer_embedding_length as f32).sqrt(); - let token_inputs = token_inputs - .reshape([batch, seq, self.config.n_layer, per_layer_embedding_length]) - .to_concrete(); - - let projected_inputs = embeddings_f32.q_mat_mul(per_layer_model_proj) - * (1.0 / (embedding_dim as f32).sqrt()); - let projected_inputs = projected_inputs - .reshape([batch, seq, self.config.n_layer, per_layer_embedding_length]) - .to_concrete(); - let projected_inputs: Tensor<4, F> = - per_layer_proj_norm.forward_generic_4d(&projected_inputs.cast()); - Some( - ((projected_inputs.cast::() + token_inputs) * (1.0 / 2.0_f32.sqrt())) - .to_concrete() - .cast(), - ) - } else { - None - }; + let per_layer_inputs = self.compute_per_layer_inputs(&embeddings_f32, || x.clone()); let embeddings: Tensor<3, F> = embeddings_f32.cast(); let encoded = EncodedTokens { embeddings, diff --git a/models/kalosm-llama/src/raw/mtp.rs b/models/kalosm-llama/src/raw/mtp.rs index a1e802870..aa5a9eaff 100644 --- a/models/kalosm-llama/src/raw/mtp.rs +++ b/models/kalosm-llama/src/raw/mtp.rs @@ -286,6 +286,9 @@ where hidden_size: layer_attention_width, rope_cache, sliding_window_size: layer_sliding_window_size, + // Unit scale: like the Gemma 4 target model, the assistant's + // query pre-attention scaling is baked into its `attn_q_norm` + // weights (see `Model::from_gguf`). attention_scale: 1.0, shared_kv_layer: None, per_layer_inp_gate: None, From 462fcb492916309223103be3e85cb651f6a38721 Mon Sep 17 00:00:00 2001 From: Evan Almloff Date: Sun, 28 Jun 2026 08:54:10 -0500 Subject: [PATCH 15/16] fix formatting --- .devin/config.local.json | 12 +++++++++++ models/kalosm-llama/src/model/inference.rs | 3 +++ models/kalosm-llama/src/model/mod.rs | 1 + models/kalosm-llama/src/raw/mod.rs | 23 +++++++++++----------- 4 files changed, 28 insertions(+), 11 deletions(-) create mode 100644 .devin/config.local.json diff --git a/.devin/config.local.json b/.devin/config.local.json new file mode 100644 index 000000000..3054c4bf5 --- /dev/null +++ b/.devin/config.local.json @@ -0,0 +1,12 @@ +{ + "permissions": { + "allow": [ + "Exec(git status)", + "Exec(gh pr)", + "Exec(gh run)", + "Exec(cargo fmt)", + "Exec(git diff)", + "Exec(cargo clippy)" + ] + } +} \ No newline at end of file diff --git a/models/kalosm-llama/src/model/inference.rs b/models/kalosm-llama/src/model/inference.rs index 2e5e87113..cddff99a0 100644 --- a/models/kalosm-llama/src/model/inference.rs +++ b/models/kalosm-llama/src/model/inference.rs @@ -422,6 +422,7 @@ where Ok(()) } + #[allow(clippy::too_many_arguments)] async fn infer_mtp_speculative( &self, prompt_tokens: &[u32], @@ -614,6 +615,7 @@ where Ok(()) } + #[allow(clippy::too_many_arguments)] async fn infer_target_only_from_pending( &self, mut next_token: u32, @@ -706,6 +708,7 @@ where Ok(()) } + #[allow(clippy::too_many_arguments)] async fn infer_target_only_from_gpu_pending( &self, mut next_token: fusor::GpuSampledToken, diff --git a/models/kalosm-llama/src/model/mod.rs b/models/kalosm-llama/src/model/mod.rs index f6c331d42..ace2570b8 100644 --- a/models/kalosm-llama/src/model/mod.rs +++ b/models/kalosm-llama/src/model/mod.rs @@ -971,6 +971,7 @@ where #[cfg(not(target_arch = "wasm32"))] let (model, tokenizer, mtp) = { let device = device.clone(); + #[allow(clippy::type_complexity)] let load_model = move || -> Result<(Model, LlamaTokenizer, Option>), LlamaSourceError> { let tokenizer = parse_external_tokenizer(tokenizer_source)?; diff --git a/models/kalosm-llama/src/raw/mod.rs b/models/kalosm-llama/src/raw/mod.rs index 30ee2cb1a..db0f88c02 100644 --- a/models/kalosm-llama/src/raw/mod.rs +++ b/models/kalosm-llama/src/raw/mod.rs @@ -1104,8 +1104,8 @@ where .to_concrete() }; - let projected_inputs = embeddings_f32.q_mat_mul(per_layer_model_proj) - * (1.0 / (embedding_dim as f32).sqrt()); + let projected_inputs = + embeddings_f32.q_mat_mul(per_layer_model_proj) * (1.0 / (embedding_dim as f32).sqrt()); let projected_inputs = projected_inputs .reshape([batch, seq, n_layer, per_layer_embedding_length]) .to_concrete(); @@ -1275,7 +1275,11 @@ where } } } - Tensor::from_slice(device, [1, per_layer_tokens.len()], per_layer_tokens.as_slice()) + Tensor::from_slice( + device, + [1, per_layer_tokens.len()], + per_layer_tokens.as_slice(), + ) } #[cfg(not(feature = "vision"))] { @@ -1416,12 +1420,8 @@ where if segment_start < tokens.len() || !text_prefix.is_empty() { let mut text_tokens = text_prefix; text_tokens.extend_from_slice(&tokens[segment_start..]); - last_logits = self.forward_text_chunk_for_multimodal( - &text_tokens, - device, - cache.as_deref_mut(), - true, - )?; + last_logits = + self.forward_text_chunk_for_multimodal(&text_tokens, device, cache, true)?; } last_logits.ok_or_else(|| fusor::Error::msg("No tokens to forward")) @@ -1523,8 +1523,9 @@ where // Image chunks carry no text tokens, so every position shares a single // zeroed per-layer token that the helper broadcasts across the chunk. - let per_layer_inputs = self - .compute_per_layer_inputs(&embeddings_f32, || Tensor::from_slice(device, [1, 1], &[0u32])); + let per_layer_inputs = self.compute_per_layer_inputs(&embeddings_f32, || { + Tensor::from_slice(device, [1, 1], &[0u32]) + }); let encoded = EncodedTokens { embeddings: embeddings_f32.cast(), From 89403162a1a7b691594af953a23f029c62182b9a Mon Sep 17 00:00:00 2001 From: Evan Almloff Date: Sun, 28 Jun 2026 19:44:31 -0500 Subject: [PATCH 16/16] fix sliding window prefill --- models/kalosm-llama/src/model/inference.rs | 24 ++++++++ .../kalosm-llama/src/raw/attention_layer.rs | 59 ++++++++++++++++++- models/kalosm-llama/src/raw/vision/gemma.rs | 53 +++++++++++++++-- 3 files changed, 128 insertions(+), 8 deletions(-) diff --git a/models/kalosm-llama/src/model/inference.rs b/models/kalosm-llama/src/model/inference.rs index cddff99a0..8713c0ea1 100644 --- a/models/kalosm-llama/src/model/inference.rs +++ b/models/kalosm-llama/src/model/inference.rs @@ -503,6 +503,7 @@ where session_lock.tokens.len() }; for _ in 0..draft_limit { + let draft_position = mtp_draft_position(draft_position, draft_tokens.len()); let step = { let session_lock = session .cache @@ -898,3 +899,26 @@ where Self::sample_standard_logits(row_logits, sampler, previous_tokens, top_k).await } } + +fn mtp_draft_position(base_position: usize, draft_offset: usize) -> usize { + base_position + draft_offset +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn multi_token_mtp_draft_positions_advance_by_draft_offset() { + let base_position = 10; + let positions: Vec<_> = (0..3) + .map(|draft_offset| mtp_draft_position(base_position, draft_offset)) + .collect(); + + assert_eq!( + positions, + vec![10, 11, 12], + "each MTP draft token must be evaluated at its future position" + ); + } +} diff --git a/models/kalosm-llama/src/raw/attention_layer.rs b/models/kalosm-llama/src/raw/attention_layer.rs index c3c94ec11..0a1ad1ade 100644 --- a/models/kalosm-llama/src/raw/attention_layer.rs +++ b/models/kalosm-llama/src/raw/attention_layer.rs @@ -642,8 +642,20 @@ where } ( - key_states.narrow(2, 0, logical_kv_len).to_concrete(), - value_states.narrow(2, 0, logical_kv_len).to_concrete(), + key_states + .narrow( + 2, + kv_narrow_start(key_states.shape()[2], logical_kv_len), + logical_kv_len, + ) + .to_concrete(), + value_states + .narrow( + 2, + kv_narrow_start(value_states.shape()[2], logical_kv_len), + logical_kv_len, + ) + .to_concrete(), ) } @@ -848,6 +860,10 @@ where } } +fn kv_narrow_start(kv_seq_len: usize, logical_kv_len: usize) -> usize { + kv_seq_len.saturating_sub(logical_kv_len) +} + fn pad_attention_mask_to_kv_len( attention_mask: &AttentionMask, kv_seq_len: usize, @@ -912,3 +928,42 @@ where attention_wo.forward_generic(&attn_output.cast()) } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn sliding_window_prefill_kv_slice_keeps_newest_entries() { + let kv_seq_len = 8; + let logical_kv_len = 4; + + assert_eq!( + kv_narrow_start(kv_seq_len, logical_kv_len), + kv_seq_len - logical_kv_len, + "sliding-window prefill must keep the newest KV entries" + ); + } + + #[test] + fn sliding_window_prefill_mask_crop_targets_newest_columns() { + let device = fusor::Device::Cpu; + let mask = AttentionMask::new( + Tensor::from_slice( + &device, + [1, 8], + &[-7.0f32, -6.0, -5.0, -4.0, -3.0, -2.0, -1.0, 0.0], + ) + .to_concrete(), + ); + + let cropped = pad_attention_mask_to_kv_len(&mask, 4).expect("mask should be cropped"); + let values = pollster::block_on(cropped.mask().as_slice()).unwrap(); + + assert_eq!( + values.as_slice(), + &[-3.0, -2.0, -1.0, 0.0], + "attention mask cropping keeps the newest columns, so KV slicing must match it" + ); + } +} diff --git a/models/kalosm-llama/src/raw/vision/gemma.rs b/models/kalosm-llama/src/raw/vision/gemma.rs index d9a7913be..05b7ed03b 100644 --- a/models/kalosm-llama/src/raw/vision/gemma.rs +++ b/models/kalosm-llama/src/raw/vision/gemma.rs @@ -133,10 +133,10 @@ where pub(crate) fn preprocess_image( &self, image: &image::DynamicImage, - _min_pixels: Option, - _max_pixels: Option, + min_pixels: Option, + max_pixels: Option, ) -> Result<(Tensor<2, f32>, [u32; 3])> { - let (target_width, target_height) = self.target_image_size(image); + let (target_width, target_height) = self.target_image_size(image, min_pixels, max_pixels); let resized = image.resize_exact( target_width as u32, target_height as u32, @@ -216,10 +216,14 @@ where .to_concrete()) } - fn target_image_size(&self, image: &image::DynamicImage) -> (usize, usize) { + fn target_image_size( + &self, + image: &image::DynamicImage, + min_pixels: Option, + max_pixels: Option, + ) -> (usize, usize) { let align = self.patch_size * self.merge_size; - let min_pixels = 40 * align * align; - let max_pixels = 280 * align * align; + let (min_pixels, max_pixels) = gemma_image_pixel_bounds(align, min_pixels, max_pixels); smart_resize( image.width() as usize, image.height() as usize, @@ -335,6 +339,21 @@ where } } +fn gemma_image_pixel_bounds( + align: usize, + min_pixels: Option, + max_pixels: Option, +) -> (usize, usize) { + ( + min_pixels + .map(|pixels| pixels as usize) + .unwrap_or(40 * align * align), + max_pixels + .map(|pixels| pixels as usize) + .unwrap_or(280 * align * align), + ) +} + struct GemmaVisionPatchEmbed { weight: Tensor<2, F>, } @@ -716,6 +735,28 @@ fn smart_resize( (target_w, target_h) } +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn gemma_image_pixel_bounds_honor_media_hints() { + let align = 8; + let requested_min = 10 * align * align; + let requested_max = 20 * align * align; + + assert_eq!( + gemma_image_pixel_bounds( + align, + Some(requested_min as u32), + Some(requested_max as u32) + ), + (requested_min, requested_max), + "Gemma preprocessing must honor caller-provided MediaHints image budgets" + ); + } +} + fn metadata_f32_array(metadata: &GgufMetadata, key: &str) -> Option> { metadata.metadata.get(key).and_then(|value| { value.to_array().ok().map(|values| {