diff --git a/.claude/settings.local.json b/.claude/settings.local.json new file mode 100644 index 000000000..49fb87fe2 --- /dev/null +++ b/.claude/settings.local.json @@ -0,0 +1,159 @@ +{ + "permissions": { + "allow": [ + "Bash(cargo test:*)", + "Bash(cargo check:*)", + "Bash(find:*)", + "Bash(cargo bench:*)", + "Bash(cargo build:*)", + "Bash(RUST_LOG=debug cargo test test_matmul)", + "Bash(git stash:*)", + "Bash(RUST_BACKTRACE=1 cargo test fuzz_matmul)", + "Bash(timeout 60 cargo bench --package fusor-core matmul -- --output-format=quiet \"matmul-wgpu/1024x1\")", + "Bash(RUSTFLAGS=\"-C opt-level=3\" cargo run --release --bin test_sgemv_fix test_sgemv_fix.rs)", + "Bash(timeout 30 cargo run --example test_sgemv_fix ../../test_sgemv_fix.rs)", + "Bash(timeout 120 cargo bench --package fusor-core matmul)", + "Bash(cargo run:*)", + "Bash(cargo clippy:*)", + "Bash(timeout 30 cargo run --example gemv_tuner)", + "Bash(timeout 30 cargo run --package fusor-core --example gemv_tuner)", + "Bash(timeout 120 cargo run --package fusor-core --example gemv_tuner)", + "Bash(sed:*)", + "Bash(RUST_BACKTRACE=1 cargo run --example embed --package rbert)", + "Bash(RUST_BACKTRACE=1 cargo test test_transposed_matmul)", + "Bash(timeout 180 cargo bench --package fusor-core --bench elementwise bench_vec4_optimization)", + "Bash(cat:*)", + "Bash(timeout 120 cargo bench --bench elementwise -- --noplot vec4)", + "Bash(tee:*)", + "Bash(timeout 300 cargo bench --bench elementwise -- --noplot vec4)", + "Bash(timeout 360 cargo bench --bench elementwise -- --noplot vec4)", + "Bash(cargo clean:*)", + "Bash(timeout 30 cargo run --example embed --package rbert)", + "Bash(RUST_BACKTRACE=1 cargo test test_fuzz_q_mat_mul_q6k --package fusor-core -- --nocapture)", + "Bash(RUST_BACKTRACE=1 cargo test test_fuzz_q_mat_mul_q6k --package fusor-core)", + "Bash(for i in {1..3})", + "Bash(do echo \"=== Run $i ===\")", + "Bash(done)", + "Bash(RUST_BACKTRACE=full timeout 45 cargo run --example embed --package rbert)", + "Bash(RUST_BACKTRACE=1 timeout 30 cargo run --example embed --package rbert)", + "Bash(timeout:*)", + "Bash(RUST_BACKTRACE=1 cargo test:*)", + "Bash(python3:*)", + "Bash(RUST_BACKTRACE=1 timeout 60 cargo run --example embed --package rbert)", + "Bash(RUST_LOG=debug timeout 60 cargo run --example embed_with_tracing --package rbert)", + "Bash(RUST_BACKTRACE=1 cargo run:*)", + "Bash(for i in 1 2 3)", + "Bash(for:*)", + "Bash(do cargo test --package fusor-core fuzz_matmul -- --nocapture)", + "Bash(break)", + "Bash(git checkout:*)", + "Bash(git log:*)", + "Bash(cargo doc:*)", + "Bash(echo:*)", + "Bash(RUST_BACKTRACE=1 timeout 30 cargo run:*)", + "Bash(1)", + "Bash(RUST_BACKTRACE=1 timeout 60 cargo run:*)", + "Bash(git restore:*)", + "Bash(grep:*)", + "Bash(RUST_LOG=kalosm_llama=debug cargo run:*)", + "Bash(cargo fmt:*)", + "Bash(pkill:*)", + "Bash(RUST_BACKTRACE=1 timeout 90 cargo run:*)", + "WebSearch", + "WebFetch(domain:github.com)", + "Bash(cargo update:*)", + "Bash(do cargo test --package fusor-core test_flash_attention_masked_fuzz)", + "Bash(git -C /Users/evanalmloff/Desktop/Github/floneum log --oneline --all)", + "Bash(git -C /Users/evanalmloff/Desktop/Github/floneum log --oneline -20)", + "Bash(git -C /Users/evanalmloff/Desktop/Github/floneum show:*)", + "WebFetch(domain:docs.rs)", + "WebFetch(domain:web.dev)", + "Bash(RUSTFLAGS='--cfg=web_sys_unstable_apis' cargo check:*)", + "Bash(RUST_LOG=warn cargo run:*)", + "Bash(git show:*)", + "Bash(curl:*)", + "Bash(RUST_BACKTRACE=1 timeout 300 cargo run:*)", + "Bash(RUST_BACKTRACE=1 RUST_LOG=debug timeout 180 cargo run:*)", + "Bash(VISION_GGUF=\"/Users/evanalmloff/Library/Application Support/kalosm/cache/ggml-org/pixtral-12b-GGUF/main/pixtral-12b-Q4_K_M.gguf\" cargo run:*)", + "Bash(RUST_BACKTRACE=1 timeout 120 cargo run:*)", + "Bash(VISION_GGUF=\"/Users/evanalmloff/Library/Application Support/kalosm/cache/ggml-org/pixtral-12b-GGUF/main/mmproj-pixtral-12b-f16.gguf\" cargo run:*)", + "Bash(RUST_BACKTRACE=full timeout 300 cargo run:*)", + "Bash(GGUF_PATH=\"/Users/evanalmloff/Library/Application Support/kalosm/cache/ggml-org/pixtral-12b-GGUF/main/pixtral-12b-Q4_K_M.gguf\" python3:*)", + "Bash(RUST_LOG=debug timeout 300 cargo run:*)", + "Bash(RUST_BACKTRACE=1 timeout 400 cargo run:*)", + "Bash(RUST_BACKTRACE=1 timeout 600 cargo run:*)", + "WebFetch(domain:raw.githubusercontent.com)", + "WebFetch(domain:towardsdatascience.com)", + "WebFetch(domain:docs.anthropic.com)", + "Bash(gh pr:*)", + "Bash(RUST_BACKTRACE=1 cargo build:*)", + "Bash(./target/debug/fusor-nanochat:*)", + "Bash(md5:*)", + "Bash(rm:*)", + "Bash(RUST_BACKTRACE=1 ./target/debug/fusor-nanochat:*)", + "Bash(git add:*)", + "Bash(git -C /Users/evanalmloff/Desktop/Github/floneum log --oneline --all -- fusor-ml/cpu/src/concrete_tensor.rs)", + "Bash(git -C /Users/evanalmloff/Desktop/Github/floneum diff -- fusor-ml/cpu/src/quantized.rs)", + "Bash(git -C /Users/evanalmloff/Desktop/Github/floneum checkout -- fusor-ml/cpu/src/quantized.rs)", + "Bash(git -C /Users/evanalmloff/Desktop/Github/floneum diff --stat)", + "Bash(git -C /Users/evanalmloff/Desktop/Github/floneum diff --name-only -- fusor-ml/nanochat/)", + "Bash(git -C /Users/evanalmloff/Desktop/Github/floneum diff -- fusor-ml/fusor/src/lib.rs)", + "Bash(system_profiler:*)", + "Bash(sysctl:*)", + "Bash(RUST_BACKTRACE=1 timeout 180 cargo run:*)", + "Bash(git reset:*)", + "Bash(git rm:*)", + "Bash(gh api:*)", + "Bash(xargs git rm --cached)", + "Bash(xargs -I {} sh -c 'echo \"\"\"\"=== {} ===\"\"\"\" && grep -c \"\"\"\"pub fn\"\"\"\" {}')", + "Bash(xargs -I {} sh -c 'echo \"\"\"\"=== {} ===\"\"\"\" && grep -n \"\"\"\"pub fn\"\"\"\" {} | grep -v test')", + "Bash(git -C /Users/evanalmloff/Desktop/Github/floneum diff --name-only)", + "Bash(git -C /Users/evanalmloff/Desktop/Github/floneum diff --stat -- fusor-ml/core/src/compute_graph/resolve.rs fusor-ml/core/src/device.rs fusor-ml/core/src/mir/kernel.rs)", + "Bash(git -C /Users/evanalmloff/Desktop/Github/floneum stash)", + "Bash(git -C /Users/evanalmloff/Desktop/Github/floneum stash pop)", + "Bash(git -C /Users/evanalmloff/Desktop/Github/floneum stash drop)", + "Bash(git -C /Users/evanalmloff/Desktop/Github/floneum diff --stat -- fusor-ml/core/src/compute_graph/resolve.rs)", + "Bash(git -C /Users/evanalmloff/Desktop/Github/floneum stash list:*)", + "Bash(git -C /Users/evanalmloff/Desktop/Github/floneum diff --name-only --cached)", + "Bash(git -C /Users/evanalmloff/Desktop/Github/floneum status --short -- fusor-ml/core/src/device.rs fusor-ml/core/src/mir/kernel.rs)", + "Bash(git -C /Users/evanalmloff/Desktop/Github/floneum fsck --unreachable --no-reflogs)", + "Bash(git -C /Users/evanalmloff/Desktop/Github/floneum status --short)", + "Bash(git -C /Users/evanalmloff/Desktop/Github/floneum diff -- fusor-ml/types/src/layout.rs fusor-ml/cpu/src/tensor.rs fusor-ml/core/src/map_layout.rs fusor-ml/core/src/lib.rs fusor-ml/fusor/src/composite/shape.rs)", + "Bash(git -C /Users/evanalmloff/Desktop/Github/floneum log --oneline -5)", + "Bash(rt:*)", + "Bash(mkdir:*)", + "Bash(RUST_BACKTRACE=full cargo test:*)", + "Bash(pip3 show:*)", + "Bash(huggingface-cli download:*)", + "Bash(readlink:*)", + "Bash(RWHISPER_COHERE_DIR=/private/tmp/cohere-transcribe-03-2026 cargo run:*)", + "Bash(huggingface-cli whoami:*)", + "Bash(huggingface-cli repo create:*)", + "Bash(RWHISPER_COHERE=1 cargo run:*)", + "Bash(do echo \"=== $commit ===\")", + "Bash(hf whoami:*)", + "Bash(hf auth:*)", + "Bash(git -s ls-files models/rwhisper/artifacts/)", + "Bash(hf download:*)", + "Bash(mv moonshine-streaming-__TRACKED_VAR__/model.gguf moonshine-streaming-__TRACKED_VAR__.gguf)", + "Bash(rm -rf moonshine-streaming-__TRACKED_VAR__)", + "Bash(FUSOR_USE_GPU=1 cargo test -p rwhisper --test jfk_regression -r)", + "Bash(FUSOR_USE_GPU=1 RUST_BACKTRACE=1 cargo test -p rwhisper --test jfk_regression -r moonshine_tiny_matches_the_jfk_reference)", + "Bash(FUSOR_USE_GPU=1 cargo test -p rwhisper --test jfk_regression -r moonshine_tiny_matches_the_jfk_reference)", + "Bash(mkdir -p ~/.cache/huggingface/hub/models--Demonthos--moonshot-gguf/refs)", + "Bash(echo -n \"main\")", + "Read(//Users/evanalmloff/.cache/huggingface/hub/models--Demonthos--moonshot-gguf/**)", + "Bash(FUSOR_USE_GPU=1 FUSOR_DUMP_SHADERS=1 cargo test -p rwhisper --test jfk_regression -r moonshine_tiny_matches_the_jfk_reference)", + "Bash(awk '/=== fusor shader ===/{n++; out=\"\"; next} /=== end shader ===/{last=out; next} /invalid accessor/{print last; exit} {out=out \"\\\\n\" $0}')", + "Bash(FUSOR_USE_GPU=1 cargo test -p rwhisper --test jfk_regression -r moonshine_tiny_streaming_matches_the_jfk_reference)", + "Bash(FUSOR_USE_GPU=1 cargo test -p rwhisper --lib quantized::moonshine::tests::cached_decode_matches_full_decode -- --nocapture)", + "Bash(FUSOR_USE_GPU=1 cargo test -p rwhisper --lib quantized::moonshine::tests::streaming_encoder_matches_full_encoder_on_jfk -- --nocapture)", + "Bash(awk 'NR<1357 || NR>1544' models/rwhisper/src/quantized/moonshine.rs)", + "Bash(awk 'NR<=1352 || NR>=1469' models/rwhisper/src/quantized/moonshine.rs)" + ], + "deny": [], + "additionalDirectories": [ + "/private/tmp" + ] + } +} diff --git a/fusor-ml/core/.claude/settings.local.json b/fusor-ml/core/.claude/settings.local.json new file mode 100644 index 000000000..b025c9c23 --- /dev/null +++ b/fusor-ml/core/.claude/settings.local.json @@ -0,0 +1,39 @@ +{ + "permissions": { + "allow": [ + "Bash(cargo check:*)", + "Bash(cargo test:*)", + "Bash(RUST_BACKTRACE=1 cargo test test_fuzz_q_mat_mul -- --nocapture)", + "Bash(RUST_BACKTRACE=1 cargo test test_fuzz_q_mat_mul_q8_0 -- --nocapture)", + "Bash(RUST_BACKTRACE=1 cargo test test_fuzz_q_mat_mul_sgemv -- --nocapture)", + "Bash(git log:*)", + "Read(//Users/evanalmloff/Desktop/Github/floneum/models/rwhisper/src/**)", + "Read(//Users/evanalmloff/Desktop/Github/floneum/models/rwhisper/**)", + "Read(//Users/evanalmloff/.cargo/registry/src/**)", + "Bash(git merge:*)", + "Bash(git checkout:*)", + "Bash(git add:*)", + "Bash(git rm:*)", + "Bash(git commit:*)", + "Bash(RUST_BACKTRACE=1 cargo test:*)", + "Bash(cargo build:*)", + "Bash(RUST_BACKTRACE=1 cargo run:*)", + "Bash(cargo run:*)", + "Bash(FUSOR_DEBUG_SHADERS=1 cargo run:*)", + "Bash(tee:*)", + "Bash(timeout 120 bash:*)", + "Bash(timeout 120 cargo run:*)", + "Bash(timeout 60 cargo run:*)", + "Bash(FUSOR_DEBUG_SHADERS=1 timeout 45 cargo run:*)", + "Bash(cat:*)", + "Bash(timeout 300 cargo run:*)", + "Bash(FUSOR_DEBUG_SHADERS=1 timeout 120 cargo run:*)", + "Bash(RUST_BACKTRACE=1 timeout 120 cargo run:*)", + "Bash(git stash:*)", + "Bash(FUSOR_DEBUG_SHADERS=1 cargo test:*)", + "Bash(git show:*)" + ], + "deny": [], + "ask": [] + } +} diff --git a/fusor-ml/core/src/composite/argmax.rs b/fusor-ml/core/src/composite/argmax.rs new file mode 100644 index 000000000..bab90acec --- /dev/null +++ b/fusor-ml/core/src/composite/argmax.rs @@ -0,0 +1,414 @@ +use std::{fmt::Write, sync::Arc}; + +use crate::{ + DataTypeEnum, Layout, Tensor, + compute_graph::NodeIndex, + min_for_dtype, + mir::{ + globals::KernelGlobalSpace, + inputs::MirValue, + kernel::GenericKernel, + operation::Operation, + workgroup_shape::{Constraint, WorkgroupShape, WorkgroupShapeConstraints}, + }, + tensor::{LazyTensorData, TensorData, TensorInfo}, + visit_tiled::distribute_workgroups, +}; + +impl Tensor { + pub fn argmax_last_dim(&self) -> Tensor { + assert_eq!(R, OUT_RANK + 1); + + let operation = ArgmaxLastDimOperation::new(self.key(), self.shape()); + let device = self.device().clone(); + let output_shape = self.shape()[..R - 1].to_vec().into_boxed_slice(); + let info = TensorInfo::new(output_shape, DataTypeEnum::U32); + let key = device.compute_graph().create_custom(Arc::new(operation)); + + Tensor::from_parts(LazyTensorData::from_parts(device, info, key)) + } +} + +#[derive(Debug, Clone)] +struct ArgmaxLastDimOperation { + input: NodeIndex, + shape: Box<[usize]>, +} + +impl ArgmaxLastDimOperation { + fn new(input: NodeIndex, shape: &[usize]) -> Self { + Self { + input, + shape: shape.into(), + } + } + + fn rank(&self) -> u32 { + self.shape.len() as u32 + } + + fn output_rank(&self) -> u32 { + self.rank() - 1 + } + + fn kernel( + &self, + workgroup_shape: &WorkgroupShape, + blocksize: u32, + kernel: &mut GenericKernel, + device: &crate::Device, + ) { + let output_rank = self.output_rank(); + let input_tensor = kernel.add_tensor_input(output_rank, false, DataTypeEnum::F32); + let output_tensor = kernel.add_tensor_input(output_rank, true, DataTypeEnum::U32); + let reduce_size = kernel.add_integer_input(); + let reduce_stride = kernel.add_integer_input(); + let workgroup_local_index = kernel.workgroup_local_index(); + + let linearized_workgroup = workgroup_shape.linearized_workgroup_index(kernel); + writeln!( + kernel, + "var workgroup_index_remainder = {};", + linearized_workgroup + ) + .unwrap(); + for i in (0..output_rank).rev() { + let out_shape_i = output_tensor.shape_binding(i); + writeln!( + kernel, + "let index_{i} = workgroup_index_remainder % {out_shape_i};" + ) + .unwrap(); + writeln!(kernel, "workgroup_index_remainder /= {out_shape_i};").unwrap(); + } + + writeln!(kernel, "var in_start_offset = ").unwrap(); + input_tensor.strided_index(kernel, (0..).map(|i| format!("index_{i}"))); + writeln!(kernel, ";").unwrap(); + writeln!(kernel, "var out_start_offset = ").unwrap(); + output_tensor.strided_index(kernel, (0..).map(|i| format!("index_{i}"))); + writeln!(kernel, ";").unwrap(); + + writeln!( + kernel, + "var best_value = f32({});", + min_for_dtype(DataTypeEnum::F32) + ) + .unwrap(); + writeln!(kernel, "var best_index = 0u;").unwrap(); + writeln!( + kernel, + "let bucket_size = ({reduce_size} + {blocksize}u - 1u) / {blocksize}u;" + ) + .unwrap(); + writeln!( + kernel, + "let base_axis_index = {workgroup_local_index} * bucket_size;" + ) + .unwrap(); + writeln!( + kernel, + "let end_axis_index = min(base_axis_index + bucket_size, {reduce_size});" + ) + .unwrap(); + writeln!(kernel, "var index = base_axis_index;").unwrap(); + writeln!(kernel, "while (index < end_axis_index) {{").unwrap(); + writeln!( + kernel, + "let value = {input_tensor}[in_start_offset + index * {reduce_stride}];" + ) + .unwrap(); + writeln!( + kernel, + "if value > best_value || (value == best_value && index < best_index) {{" + ) + .unwrap(); + writeln!(kernel, "best_value = value;").unwrap(); + writeln!(kernel, "best_index = index;").unwrap(); + writeln!(kernel, "}}").unwrap(); + writeln!(kernel, "index += 1u;").unwrap(); + writeln!(kernel, "}}").unwrap(); + + if device.subgroups_supported() { + let max_subgroup_size = device.max_subgroup_size(); + let local_values = kernel.add_global_array( + KernelGlobalSpace::Workgroup, + DataTypeEnum::F32, + max_subgroup_size.to_string(), + ); + let local_indexes = kernel.add_global_array( + KernelGlobalSpace::Workgroup, + DataTypeEnum::U32, + max_subgroup_size.to_string(), + ); + let subgroup_id = kernel.subgroup_index(); + let subgroup_local_id = kernel.subgroup_local_index(); + let subgroups_per_workgroup = kernel.subgroups_per_workgroup(); + let subgroup_size = kernel.subgroup_size(); + + let mut offset = max_subgroup_size; + while offset > 1 { + writeln!(kernel, "if {subgroup_size} >= {offset}u {{").unwrap(); + offset /= 2; + writeln!( + kernel, + "let neighbor_value = subgroupShuffleDown(best_value, {offset}u);" + ) + .unwrap(); + writeln!( + kernel, + "let neighbor_index = subgroupShuffleDown(best_index, {offset}u);" + ) + .unwrap(); + writeln!( + kernel, + "if neighbor_value > best_value || (neighbor_value == best_value && neighbor_index < best_index) {{" + ) + .unwrap(); + writeln!(kernel, "best_value = neighbor_value;").unwrap(); + writeln!(kernel, "best_index = neighbor_index;").unwrap(); + writeln!(kernel, "}}").unwrap(); + writeln!(kernel, "}}").unwrap(); + } + + writeln!(kernel, "if {subgroup_local_id} == 0u {{").unwrap(); + writeln!(kernel, "{local_values}[{subgroup_id}] = best_value;").unwrap(); + writeln!(kernel, "{local_indexes}[{subgroup_id}] = best_index;").unwrap(); + writeln!(kernel, "}}").unwrap(); + writeln!(kernel, "workgroupBarrier();").unwrap(); + + writeln!( + kernel, + "if {subgroup_local_id} < {subgroups_per_workgroup} {{" + ) + .unwrap(); + writeln!(kernel, "best_value = {local_values}[{subgroup_local_id}];").unwrap(); + writeln!(kernel, "best_index = {local_indexes}[{subgroup_local_id}];").unwrap(); + writeln!(kernel, "}} else {{").unwrap(); + writeln!( + kernel, + "best_value = f32({});", + min_for_dtype(DataTypeEnum::F32) + ) + .unwrap(); + writeln!(kernel, "best_index = 0u;").unwrap(); + writeln!(kernel, "}}").unwrap(); + + offset = max_subgroup_size; + while offset > 1 { + writeln!(kernel, "if {subgroup_size} >= {offset}u {{").unwrap(); + offset /= 2; + writeln!( + kernel, + "let neighbor_value = subgroupShuffleDown(best_value, {offset}u);" + ) + .unwrap(); + writeln!( + kernel, + "let neighbor_index = subgroupShuffleDown(best_index, {offset}u);" + ) + .unwrap(); + writeln!( + kernel, + "if neighbor_value > best_value || (neighbor_value == best_value && neighbor_index < best_index) {{" + ) + .unwrap(); + writeln!(kernel, "best_value = neighbor_value;").unwrap(); + writeln!(kernel, "best_index = neighbor_index;").unwrap(); + writeln!(kernel, "}}").unwrap(); + writeln!(kernel, "}}").unwrap(); + } + } else { + let local_values = kernel.add_global_array( + KernelGlobalSpace::Workgroup, + DataTypeEnum::F32, + blocksize.to_string(), + ); + let local_indexes = kernel.add_global_array( + KernelGlobalSpace::Workgroup, + DataTypeEnum::U32, + blocksize.to_string(), + ); + let mut offset = blocksize; + while offset > 1 { + writeln!( + kernel, + "{local_values}[{workgroup_local_index}] = best_value;" + ) + .unwrap(); + writeln!( + kernel, + "{local_indexes}[{workgroup_local_index}] = best_index;" + ) + .unwrap(); + writeln!(kernel, "workgroupBarrier();").unwrap(); + offset /= 2; + writeln!(kernel, "if {workgroup_local_index} < {offset}u {{").unwrap(); + writeln!( + kernel, + "let neighbor_value = {local_values}[{workgroup_local_index} + {offset}u];" + ) + .unwrap(); + writeln!( + kernel, + "let neighbor_index = {local_indexes}[{workgroup_local_index} + {offset}u];" + ) + .unwrap(); + writeln!( + kernel, + "if neighbor_value > best_value || (neighbor_value == best_value && neighbor_index < best_index) {{" + ) + .unwrap(); + writeln!(kernel, "best_value = neighbor_value;").unwrap(); + writeln!(kernel, "best_index = neighbor_index;").unwrap(); + writeln!(kernel, "}}").unwrap(); + writeln!(kernel, "}}").unwrap(); + } + } + + writeln!(kernel, "if {workgroup_local_index} == 0u {{").unwrap(); + writeln!(kernel, "{output_tensor}[out_start_offset] = best_index;").unwrap(); + writeln!(kernel, "}}").unwrap(); + } +} + +impl Operation for ArgmaxLastDimOperation { + fn workgroup_shape_constraints(&self, device: &crate::Device) -> WorkgroupShapeConstraints { + let mut constraints = WorkgroupShapeConstraints::new(); + constraints.add_constraint( + 0, + Constraint::less_than(device.limits().max_compute_workgroup_size_x + 1), + ); + if device.subgroups_supported() { + constraints.add_constraint( + 0, + Constraint::more_than_or_equals(device.min_subgroup_size()), + ); + constraints.add_constraint( + 0, + Constraint::less_than_or_equals(device.max_subgroup_size()), + ); + } + 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_tensor: TensorData = inputs[1].as_tensor().unwrap().clone(); + let total_workgroups = output_tensor.layout().shape().iter().product::() as u32; + + distribute_workgroups(total_workgroups) + } + + fn visit_dependencies(&self, f: &mut dyn FnMut(NodeIndex)) { + f(self.input); + } + + fn inputs(&self, nodes: &crate::compute_graph::ComputeGraphInner) -> Vec { + let tensor = nodes.get_cached_result(self.input).unwrap(); + let last_dim = tensor.layout().shape().len() - 1; + let output_shape = tensor.layout().shape()[..last_dim].to_vec(); + let output_tensor = + TensorData::new_for_shape(tensor.device(), &output_shape, DataTypeEnum::U32); + let trimmed_layout = Layout::from_parts( + tensor.layout().offset(), + tensor.layout().shape()[..last_dim].to_vec().into(), + tensor.layout().strides()[..last_dim].to_vec().into(), + ); + let trimmed_tensor = TensorData::new_from_parts( + tensor.device(), + tensor.buffer().clone(), + trimmed_layout, + DataTypeEnum::F32, + ); + + vec![ + MirValue::Tensor(trimmed_tensor), + MirValue::Tensor(output_tensor), + MirValue::Integer(tensor.layout().shape()[last_dim] as u32), + MirValue::Integer(tensor.layout().strides()[last_dim] as u32), + ] + } + + fn build_kernel( + &self, + graph: &crate::compute_graph::ComputeGraphInner, + workgroup_shape: &WorkgroupShape, + _: &[MirValue], + kernel: &mut GenericKernel, + ) { + self.kernel( + workgroup_shape, + workgroup_shape.x(), + kernel, + &graph.device(), + ); + } + + fn output(&self, _: &crate::compute_graph::ComputeGraphInner, inputs: &[MirValue]) -> MirValue { + let output_tensor: TensorData = inputs[1].as_tensor().unwrap().clone(); + output_tensor.into() + } + + fn name(&self) -> String { + format!( + "argmax_last_dim_f32_{}", + self.shape + .iter() + .map(|dim| dim.to_string()) + .collect::>() + .join("x") + ) + } +} + +#[cfg(test)] +#[tokio::test] +async fn test_argmax_last_dim_1d() { + use crate::Device; + + let device = Device::test_instance(); + let tensor = Tensor::new(&device, &[1.0, 3.0, 2.0, 3.0]); + let output: Tensor<0, u32> = tensor.argmax_last_dim(); + + assert_eq!(output.to_scalar().await.unwrap(), 1); +} + +#[cfg(test)] +#[tokio::test] +async fn test_argmax_last_dim_2d() { + use crate::Device; + + let device = Device::test_instance(); + let tensor = Tensor::new(&device, &[[1.0, 5.0, 2.0], [3.0, 0.0, 4.0]]); + let output: Tensor<1, u32> = tensor.argmax_last_dim(); + let output = output.as_slice().await.unwrap(); + + assert_eq!(output[[0]], 1); + assert_eq!(output[[1]], 2); +} + +#[cfg(test)] +#[tokio::test] +async fn test_argmax_last_dim_non_contiguous() { + use crate::Device; + + let device = Device::test_instance(); + let tensor = Tensor::new( + &device, + &[ + [1.0, 10.0, 3.0, 8.0], + [7.0, 4.0, 9.0, 2.0], + [5.0, 6.0, 0.0, 11.0], + ], + ); + let transposed = tensor.t(); + let output: Tensor<1, u32> = transposed.argmax_last_dim(); + let output = output.as_slice().await.unwrap(); + + assert_eq!(output[[0]], 1); + assert_eq!(output[[1]], 0); + assert_eq!(output[[2]], 1); + assert_eq!(output[[3]], 2); +} diff --git a/fusor-ml/core/src/composite/cat.rs b/fusor-ml/core/src/composite/cat.rs index 52d74dd62..a429cce06 100644 --- a/fusor-ml/core/src/composite/cat.rs +++ b/fusor-ml/core/src/composite/cat.rs @@ -4,6 +4,7 @@ impl Tensor { pub fn cat(vectors: impl IntoIterator, dim: impl Dim) -> Self { let dim = dim.resolve(); let vectors = vectors.into_iter().collect::>(); + assert!(!vectors.is_empty(), "cat requires at least one tensor"); let mut shape = [0; R]; for (i, v) in vectors[0].shape().iter().enumerate() { if i != dim { @@ -20,14 +21,15 @@ impl Tensor { } } } - let mut iter = vectors.into_iter(); - + let device = vectors[0].device().clone(); let mut index = 0; - let first = iter.next().unwrap(); - let mut larger = first.resize(shape); - index += first.shape()[dim]; - for vector in iter { + let mut larger = Tensor::zeros(&device, shape); + for vector in vectors { let length = vector.shape()[dim]; + if vector.is_zero_splat() { + index += length; + continue; + } let slice = std::array::from_fn(|i| { if i == dim { index..(index + length) @@ -120,3 +122,40 @@ async fn test_multi_dim_cat() { } } } + +#[cfg(test)] +#[tokio::test] +async fn test_rel_shift_cat_pattern_matches_cpu() { + use crate::Device; + + let gpu = Device::test_instance(); + + let data = vec![vec![vec![ + vec![1.0, 2.0, 3.0, 4.0, 5.0], + vec![6.0, 7.0, 8.0, 9.0, 10.0], + ]]]; + let gpu_x = Tensor::new(&gpu, &data); + + let run = |x: Tensor<4, f32>| -> Tensor<4, f32> { + let [batch, heads, q_len, pos_len] = *x.shape(); + let zeros = Tensor::zeros(&x.device(), [batch, heads, q_len, 1]); + let padded = Tensor::cat([zeros, x], 3); + let reshaped = padded.reshape([batch, heads, pos_len + 1, q_len]); + reshaped + .narrow(2, 1, pos_len) + .reshape([batch, heads, q_len, pos_len]) + }; + + let gpu_out = run(gpu_x).as_slice().await.unwrap(); + let expected = [[[[2.0f32, 3.0, 4.0, 5.0, 0.0], [6.0, 7.0, 8.0, 9.0, 10.0]]]]; + + for b in 0..1 { + for h in 0..1 { + for q in 0..2 { + for p in 0..5 { + assert_eq!(gpu_out[[b, h, q, p]], expected[b][h][q][p]); + } + } + } + } +} diff --git a/fusor-ml/core/src/composite/conv.rs b/fusor-ml/core/src/composite/conv.rs index 3be60fb28..85ade9c97 100644 --- a/fusor-ml/core/src/composite/conv.rs +++ b/fusor-ml/core/src/composite/conv.rs @@ -1,4 +1,322 @@ -use crate::{DataType, LargerRank, Tensor}; +use std::{fmt::Write, sync::Arc}; + +use crate::{ + DataType, DataTypeEnum, LargerRank, Layout, LazyTensorData, TILE_SIZE, Tensor, TensorData, + TensorInfo, TensorLayoutInfo, + compute_graph::{ComputeGraphInner, NodeIndex}, + mir::{ + inputs::MirValue, + kernel::GenericKernel, + operation::Operation, + workgroup_shape::{WorkgroupShape, WorkgroupShapeConstraints}, + }, + visit_tiled::{ + MaybeQTensorInput, VisitTiledInput, build_visit_tiled_kernel, titled_map_dispatch_size, + titled_map_workgroup_size_constraints, + }, +}; + +#[derive(Debug, Clone)] +struct GroupedConv1dOperation { + input: NodeIndex, + weight: NodeIndex, + bias: Option, + datatype: DataTypeEnum, + input_shape: [usize; 3], + weight_shape: [usize; 3], + output_shape: [usize; 3], + padding: usize, + stride: usize, + groups: usize, +} + +impl GroupedConv1dOperation { + #[allow(clippy::too_many_arguments)] + fn new( + input: NodeIndex, + weight: NodeIndex, + bias: Option, + datatype: DataTypeEnum, + input_shape: [usize; 3], + weight_shape: [usize; 3], + output_shape: [usize; 3], + padding: usize, + stride: usize, + groups: usize, + ) -> Self { + Self { + input, + weight, + bias, + datatype, + input_shape, + weight_shape, + output_shape, + padding, + stride, + groups, + } + } + + fn zero_literal(&self) -> &'static str { + match self.datatype { + DataTypeEnum::F32 => "0.0", + DataTypeEnum::F16 => "f16(0.0)", + DataTypeEnum::U32 => "0u", + } + } +} + +impl Operation for GroupedConv1dOperation { + fn workgroup_shape_constraints(&self, device: &crate::Device) -> WorkgroupShapeConstraints { + titled_map_workgroup_size_constraints(&self.output_shape, device) + } + + fn dispatch_size(&self, workgroup_shape: &WorkgroupShape, _inputs: &[MirValue]) -> [u32; 3] { + titled_map_dispatch_size(TILE_SIZE, *workgroup_shape, &self.output_shape) + } + + 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 output = TensorData::new_for_shape(input.device(), &self.output_shape, self.datatype); + + let mut inputs = Vec::with_capacity(if self.bias.is_some() { 4 } else { 3 }); + 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(output.into()); + inputs + } + + fn build_kernel( + &self, + graph: &ComputeGraphInner, + _workgroup_shape: &WorkgroupShape, + inputs: &[MirValue], + kernel: &mut GenericKernel, + ) { + let output_tensor_idx = inputs.len() - 1; + let has_bias = self.bias.is_some(); + let input_rank = self.input_shape.len() as u32; + let weight_rank = self.weight_shape.len() as u32; + let output_rank = self.output_shape.len() as u32; + let input_len = self.input_shape[2]; + let in_channels_per_group = self.weight_shape[1]; + let out_channels_per_group = self.output_shape[1] / self.groups; + let kernel_size = self.weight_shape[2]; + let dtype = self.datatype; + let zero = self.zero_literal(); + + let mut tiled_inputs = vec![ + VisitTiledInput::new(dtype.into(), input_rank), + VisitTiledInput::new(dtype.into(), weight_rank), + ]; + if has_bias { + tiled_inputs.push(VisitTiledInput::new(dtype.into(), 1)); + } + tiled_inputs.push(VisitTiledInput::new(dtype.into(), output_rank)); + + build_visit_tiled_kernel( + &graph.device(), + &self.output_shape, + TILE_SIZE, + tiled_inputs, + output_tensor_idx, + |kernel, indexes, tensors, _values| { + let input_tensor = match &tensors[0] { + MaybeQTensorInput::Tensor(tensor) => tensor, + MaybeQTensorInput::QTensor(_) => { + panic!("Grouped conv input cannot be quantized") + } + }; + let weight_tensor = match &tensors[1] { + MaybeQTensorInput::Tensor(tensor) => tensor, + MaybeQTensorInput::QTensor(_) => { + panic!("Grouped conv weight cannot be quantized") + } + }; + let bias_tensor = if has_bias { + match &tensors[2] { + MaybeQTensorInput::Tensor(tensor) => Some(tensor), + MaybeQTensorInput::QTensor(_) => { + panic!("Grouped conv bias cannot be quantized") + } + } + } else { + None + }; + let output_tensor = match &tensors[output_tensor_idx] { + MaybeQTensorInput::Tensor(tensor) => tensor, + MaybeQTensorInput::QTensor(_) => { + panic!("Grouped conv output cannot be quantized") + } + }; + + if let Some(bias_tensor) = bias_tensor { + write!(kernel, "var acc: {dtype} = {bias_tensor}[").unwrap(); + bias_tensor.strided_index(kernel, ["dim_1"]); + writeln!(kernel, "];").unwrap(); + } else { + writeln!(kernel, "var acc: {dtype} = {zero};").unwrap(); + } + + writeln!( + kernel, + "let conv_group = dim_1 / {out_channels_per_group}u;" + ) + .unwrap(); + writeln!( + kernel, + "for (var ic_local = 0u; ic_local < {in_channels_per_group}u; ic_local++) {{" + ) + .unwrap(); + writeln!( + kernel, + " let input_channel = conv_group * {in_channels_per_group}u + ic_local;" + ) + .unwrap(); + writeln!( + kernel, + " for (var kernel_index = 0u; kernel_index < {kernel_size}u; kernel_index++) {{" + ) + .unwrap(); + writeln!( + kernel, + " let padded_pos = dim_2 * {}u + kernel_index;", + self.stride + ) + .unwrap(); + writeln!( + kernel, + " if (padded_pos >= {}u && padded_pos < {}u) {{", + self.padding, + self.padding + input_len + ) + .unwrap(); + writeln!( + kernel, + " let input_pos = padded_pos - {}u;", + self.padding + ) + .unwrap(); + write!(kernel, " let input_value = {input_tensor}[").unwrap(); + input_tensor.strided_index(kernel, ["dim_0", "input_channel", "input_pos"]); + writeln!(kernel, "];").unwrap(); + write!(kernel, " let weight_value = {weight_tensor}[").unwrap(); + weight_tensor.strided_index(kernel, ["dim_1", "ic_local", "kernel_index"]); + writeln!(kernel, "];").unwrap(); + writeln!(kernel, " acc += input_value * weight_value;").unwrap(); + writeln!(kernel, " }}").unwrap(); + writeln!(kernel, " }}").unwrap(); + writeln!(kernel, "}}").unwrap(); + + format!("{output_tensor}[{}] = acc;", indexes[output_tensor_idx]) + }, + kernel, + ); + } + + fn output(&self, _nodes: &ComputeGraphInner, inputs: &[MirValue]) -> MirValue { + inputs.last().unwrap().clone() + } + + fn name(&self) -> String { + format!( + "conv1d_grouped_{}_{}x{}x{}_groups_{}", + self.datatype, + self.input_shape[1], + self.output_shape[1], + self.weight_shape[2], + self.groups + ) + } + + fn output_layout( + &self, + _map: &rustc_hash::FxHashMap, + ) -> TensorLayoutInfo { + TensorLayoutInfo::new(Layout::contiguous(&self.output_shape), self.datatype) + } +} + +impl Tensor<3, D> { + pub fn conv1d_grouped( + &self, + weight: &Tensor<3, D>, + bias: Option<&Tensor<1, D>>, + padding: usize, + stride: usize, + groups: usize, + ) -> Self { + assert!(groups > 0, "groups must be greater than zero"); + assert!(stride > 0, "stride must be greater than zero"); + + let input_shape = *self.shape(); + let weight_shape = *weight.shape(); + let in_channels = input_shape[1]; + let out_channels = weight_shape[0]; + let in_channels_per_group = weight_shape[1]; + let kernel_size = weight_shape[2]; + + assert_eq!( + in_channels, + in_channels_per_group * groups, + "weight in_channels per group must match input channels / groups" + ); + assert_eq!( + out_channels % groups, + 0, + "out_channels ({out_channels}) must be divisible by groups ({groups})" + ); + assert!( + input_shape[2] + 2 * padding >= kernel_size, + "kernel size ({kernel_size}) cannot exceed padded input length ({})", + input_shape[2] + 2 * padding + ); + + if let Some(bias) = bias { + assert_eq!( + bias.shape()[0], + out_channels, + "bias shape must match out_channels" + ); + } + + let output_length = (input_shape[2] + 2 * padding - kernel_size) / stride + 1; + let output_shape = [input_shape[0], out_channels, output_length]; + let operation = GroupedConv1dOperation::new( + self.key(), + weight.key(), + bias.map(|bias| bias.key()), + self.datatype(), + input_shape, + weight_shape, + output_shape, + padding, + stride, + groups, + ); + let device = self.device().clone(); + let key = device.compute_graph().create_custom(Arc::new(operation)); + + Tensor::from_parts(LazyTensorData::from_parts( + device, + TensorInfo::new(output_shape.to_vec().into_boxed_slice(), self.datatype()), + key, + )) + } +} impl Tensor { /// Pad a specific axis with zeros on both sides @@ -351,4 +669,189 @@ mod tests { } } } + + #[tokio::test] + async fn test_grouped_conv_1d_depthwise_single_kernel() { + let device = Device::test_instance(); + + let input_data = [[[1.0f32, 2.0, 3.0, 4.0], [10.0f32, 20.0, 30.0, 40.0]]]; + let input = Tensor::new(&device, &input_data); + let weight_data = [[[1.0f32, 0.0, -1.0]], [[0.5f32, 0.25, -0.5]]]; + let weight = Tensor::new(&device, &weight_data); + let bias = Tensor::new(&device, &[0.5f32, -1.0]); + + let output = input.conv1d_grouped(&weight, Some(&bias), 1, 1, 2); + assert_eq!(output.count_kernels_to_resolve(), 1); + + let result = output.as_slice().await.unwrap(); + assert_eq!(result.shape(), &[1, 2, 4]); + let expected = [[[-1.5f32, -1.5, -1.5, 3.5], [-8.5f32, -6.0, -3.5, 24.0]]]; + for batch in 0..1 { + for channel in 0..2 { + for position in 0..4 { + let actual = result[[batch, channel, position]]; + let expected = expected[batch][channel][position]; + assert!( + (actual - expected).abs() < 1e-5, + "Mismatch at [{batch}, {channel}, {position}]: got {actual}, expected {expected}" + ); + } + } + } + } + + #[tokio::test] + async fn test_conv_1d_moonshine_conv2_shape_matches_cpu_reference() { + let device = Device::test_instance(); + let batch = 1; + let in_channels = 640; + let out_channels = 320; + let input_len = 204; + let kernel = 5; + let stride = 2; + + let input_data: Vec = (0..batch * in_channels * input_len) + .map(|i| ((i % 101) as f32 - 50.0) * 0.002) + .collect(); + let weight_data: Vec = (0..out_channels * in_channels * kernel) + .map(|i| ((i % 103) as f32 - 51.0) * 0.001) + .collect(); + let bias_data: Vec = (0..out_channels) + .map(|i| ((i % 17) as f32 - 8.0) * 0.001) + .collect(); + + let input_nested: Vec>> = input_data + .chunks_exact(in_channels * input_len) + .map(|batch| { + batch + .chunks_exact(input_len) + .map(|channel| channel.to_vec()) + .collect() + }) + .collect(); + let weight_nested: Vec>> = weight_data + .chunks_exact(in_channels * kernel) + .map(|out_channel| { + out_channel + .chunks_exact(kernel) + .map(|channel| channel.to_vec()) + .collect() + }) + .collect(); + + let input = Tensor::new(&device, &input_nested); + let weight = Tensor::new(&device, &weight_nested); + let bias = Tensor::new(&device, &bias_data); + let result = input + .conv(&weight, Some(&bias), [0], [stride]) + .as_slice() + .await + .unwrap(); + + let check_positions = [ + (0, 0, 0), + (0, 0, result.shape()[2] - 1), + (0, out_channels / 2, result.shape()[2] / 2), + (0, out_channels - 1, 0), + (0, out_channels - 1, result.shape()[2] - 1), + ]; + + for (batch_idx, out_channel, out_pos) in check_positions { + let mut expected = bias_data[out_channel]; + for in_channel in 0..in_channels { + for kernel_idx in 0..kernel { + let input_pos = out_pos * stride + kernel_idx; + expected += input_data + [batch_idx * in_channels * input_len + in_channel * input_len + input_pos] + * weight_data + [out_channel * in_channels * kernel + in_channel * kernel + kernel_idx]; + } + } + let actual = result[[batch_idx, out_channel, out_pos]]; + assert!( + (actual - expected).abs() < 1e-3, + "Mismatch at [{batch_idx}, {out_channel}, {out_pos}]: actual={actual}, expected={expected}" + ); + } + } + + #[tokio::test] + async fn test_conv_1d_after_zero_cat_uses_materialized_layout() { + let device = Device::test_instance(); + let batch = 1; + let in_channels = 64; + let out_channels = 32; + let input_len = 17; + let pad = 4; + let padded_len = input_len + pad; + let kernel = 5; + let stride = 2; + + let input_data: Vec = (0..batch * in_channels * input_len) + .map(|i| ((i % 101) as f32 - 50.0) * 0.002) + .collect(); + let weight_data: Vec = (0..out_channels * in_channels * kernel) + .map(|i| ((i % 103) as f32 - 51.0) * 0.001) + .collect(); + let bias_data: Vec = (0..out_channels) + .map(|i| ((i % 17) as f32 - 8.0) * 0.001) + .collect(); + + let input_nested: Vec>> = input_data + .chunks_exact(in_channels * input_len) + .map(|batch| { + batch + .chunks_exact(input_len) + .map(|channel| channel.to_vec()) + .collect() + }) + .collect(); + let weight_nested: Vec>> = weight_data + .chunks_exact(in_channels * kernel) + .map(|out_channel| { + out_channel + .chunks_exact(kernel) + .map(|channel| channel.to_vec()) + .collect() + }) + .collect(); + + let input = Tensor::new(&device, &input_nested); + let padded = Tensor::cat( + [Tensor::zeros(&device, [batch, in_channels, pad]), input], + 2, + ); + let weight = Tensor::new(&device, &weight_nested); + let bias = Tensor::new(&device, &bias_data); + let result = padded + .conv(&weight, Some(&bias), [0], [stride]) + .as_slice() + .await + .unwrap(); + + for out_channel in [0, out_channels / 2, out_channels - 1] { + for out_pos in [0, result.shape()[2] / 2, result.shape()[2] - 1] { + let mut expected = bias_data[out_channel]; + for in_channel in 0..in_channels { + for kernel_idx in 0..kernel { + let padded_pos = out_pos * stride + kernel_idx; + let input_value = if padded_pos < pad { + 0.0 + } else { + input_data[in_channel * input_len + padded_pos - pad] + }; + expected += input_value + * weight_data[out_channel * in_channels * kernel + + in_channel * kernel + + kernel_idx]; + } + } + let actual = result[[0, out_channel, out_pos]]; + assert!( + (actual - expected).abs() < 1e-3, + "Mismatch at [0, {out_channel}, {out_pos}] with padded_len={padded_len}: actual={actual}, expected={expected}" + ); + } + } + } } diff --git a/fusor-ml/core/src/composite/mod.rs b/fusor-ml/core/src/composite/mod.rs index f98c30842..b286f92a5 100644 --- a/fusor-ml/core/src/composite/mod.rs +++ b/fusor-ml/core/src/composite/mod.rs @@ -1,4 +1,5 @@ mod arange; +mod argmax; mod cat; mod chunk; mod conv; @@ -19,6 +20,7 @@ mod sliding_window_view; mod softmax; mod sqr; mod squeeze; +mod swiglu; mod to_vec; pub use to_vec::{ToVec1, ToVec2, ToVec3}; mod unsqueeze; diff --git a/fusor-ml/core/src/composite/rms_norm_fused.rs b/fusor-ml/core/src/composite/rms_norm_fused.rs index 5f7577119..f1e9ee4c6 100644 --- a/fusor-ml/core/src/composite/rms_norm_fused.rs +++ b/fusor-ml/core/src/composite/rms_norm_fused.rs @@ -40,6 +40,37 @@ impl Tensor { self.shape(), weight.shape(), eps, + false, + ); + let data = self.data(); + + Self::from_parts(data.custom(Arc::new(operation))) + } + + /// Fused LayerNorm kernel that performs mean, variance, scale, and bias in one launch. + /// + /// Formula: output = (input - mean(input)) / sqrt(var(input) + eps) * weight + bias + pub fn layer_norm_fused( + &self, + weight: &Tensor, + bias: Option<&Tensor>, + eps: f32, + ) -> Self + where + T: CastTensor, + f32: CastTensor, + (Tensor, Tensor): MaxRank, + { + let operation = RmsNormOperation::new( + self.key(), + weight.key(), + bias.map(|b| b.key()), + self.datatype(), + weight.datatype(), + self.shape(), + weight.shape(), + eps, + true, ); let data = self.data(); @@ -73,6 +104,8 @@ struct RmsNormOperation { input_shape: Box<[usize]>, /// Epsilon for numerical stability eps: f32, + /// Whether to subtract the mean and normalize by variance instead of RMS. + remove_mean: bool, } impl RmsNormOperation { @@ -86,6 +119,7 @@ impl RmsNormOperation { input_shape: &[usize], _weight_shape: &[usize], eps: f32, + remove_mean: bool, ) -> Self { Self { input, @@ -95,6 +129,7 @@ impl RmsNormOperation { weight_dtype, input_shape: input_shape.into(), eps, + remove_mean, } } @@ -119,6 +154,7 @@ impl RmsNormOperation { let hidden_size = self.hidden_size(); let large_reduction = hidden_size > 256; let has_bias = self.bias.is_some(); + let remove_mean = self.remove_mean; // Input tensor (without the last dimension in the layout for workgroup indexing) let input_tensor = kernel.add_tensor_input(output_rank, false, input_dtype); @@ -165,7 +201,10 @@ impl RmsNormOperation { writeln!(kernel, ";").unwrap(); writeln!(kernel).unwrap(); - // Phase 1: Compute sum of squares + // Phase 1: Compute row statistics + if remove_mean { + writeln!(kernel, "var sum = f32(0.0);").unwrap(); + } writeln!(kernel, "var sum_sq = f32(0.0);").unwrap(); // Divide work among threads in the workgroup @@ -204,6 +243,13 @@ impl RmsNormOperation { // Convert to f32 and compute squared values writeln!(kernel, "let f32_data = vec4(data);").unwrap(); + if remove_mean { + writeln!( + kernel, + "sum += f32_data.x + f32_data.y + f32_data.z + f32_data.w;" + ) + .unwrap(); + } writeln!(kernel, "let sq_data = f32_data * f32_data;").unwrap(); writeln!( kernel, @@ -223,13 +269,21 @@ impl RmsNormOperation { "let data = f32({input_tensor}[in_start_offset + index * {reduce_stride}]);" ) .unwrap(); + if remove_mean { + writeln!(kernel, "sum += data;").unwrap(); + } writeln!(kernel, "sum_sq += data * data;").unwrap(); writeln!(kernel, "index += 1u;").unwrap(); writeln!(kernel, "}}").unwrap(); writeln!(kernel).unwrap(); - // Phase 2: Reduce sum_sq across the workgroup + // Phase 2: Reduce row statistics across the workgroup let global_rms = kernel.add_global_value(KernelGlobalSpace::Workgroup, DataTypeEnum::F32); + let global_mean = if remove_mean { + Some(kernel.add_global_value(KernelGlobalSpace::Workgroup, DataTypeEnum::F32)) + } else { + None + }; if device.subgroups_supported() { let max_subgroup_size = device.max_subgroup_size(); let local_data = kernel.add_global_array( @@ -237,15 +291,30 @@ impl RmsNormOperation { DataTypeEnum::F32, max_subgroup_size.to_string(), ); + let local_sum_data = if remove_mean { + Some(kernel.add_global_array( + KernelGlobalSpace::Workgroup, + DataTypeEnum::F32, + max_subgroup_size.to_string(), + )) + } else { + None + }; let subgroup_id = kernel.subgroup_index(); let subgroup_local_id = kernel.subgroup_local_index(); let subgroups_per_workgroup = kernel.subgroups_per_workgroup(); // First: reduce within each subgroup writeln!(kernel, "sum_sq = subgroupAdd(sum_sq);").unwrap(); + if remove_mean { + writeln!(kernel, "sum = subgroupAdd(sum);").unwrap(); + } // Write subgroup results to shared memory writeln!(kernel, "{local_data}[{subgroup_id}] = sum_sq;").unwrap(); + if let Some(local_sum_data) = &local_sum_data { + writeln!(kernel, "{local_sum_data}[{subgroup_id}] = sum;").unwrap(); + } writeln!(kernel, "workgroupBarrier();").unwrap(); // Final reduction across subgroups (only first subgroup participates) @@ -255,20 +324,39 @@ impl RmsNormOperation { ) .unwrap(); writeln!(kernel, "sum_sq = {local_data}[{subgroup_local_id}];").unwrap(); + if let Some(local_sum_data) = &local_sum_data { + writeln!(kernel, "sum = {local_sum_data}[{subgroup_local_id}];").unwrap(); + } writeln!(kernel, "}}").unwrap(); writeln!(kernel, "else {{").unwrap(); writeln!(kernel, "sum_sq = f32(0.0);").unwrap(); + if remove_mean { + writeln!(kernel, "sum = f32(0.0);").unwrap(); + } writeln!(kernel, "}}").unwrap(); writeln!(kernel, "sum_sq = subgroupAdd(sum_sq);").unwrap(); + if remove_mean { + writeln!(kernel, "sum = subgroupAdd(sum);").unwrap(); + } - // Thread 0 now has the final sum_sq, compute RMS + // Thread 0 now has the final row statistics. writeln!(kernel, "if {subgroup_id} == 0u {{").unwrap(); - // Store to shared memory for other subgroups to read - writeln!( - kernel, - "{global_rms} = sqrt(sum_sq / f32({reduce_size}) + {eps_input});" - ) - .unwrap(); + if let Some(global_mean) = &global_mean { + writeln!(kernel, "let mean = sum / f32({reduce_size});").unwrap(); + writeln!( + kernel, + "let variance = max(sum_sq / f32({reduce_size}) - mean * mean, f32(0.0));" + ) + .unwrap(); + writeln!(kernel, "{global_mean} = mean;").unwrap(); + writeln!(kernel, "{global_rms} = sqrt(variance + {eps_input});").unwrap(); + } else { + writeln!( + kernel, + "{global_rms} = sqrt(sum_sq / f32({reduce_size}) + {eps_input});" + ) + .unwrap(); + } writeln!(kernel, "}}").unwrap(); } else { // Fallback: shared memory reduction @@ -277,9 +365,21 @@ impl RmsNormOperation { DataTypeEnum::F32, blocksize.to_string(), ); + let local_sum_data = if remove_mean { + Some(kernel.add_global_array( + KernelGlobalSpace::Workgroup, + DataTypeEnum::F32, + blocksize.to_string(), + )) + } else { + None + }; let mut offset = blocksize; while offset > 1 { writeln!(kernel, "{local_data}[{workgroup_local_index}] = sum_sq;").unwrap(); + if let Some(local_sum_data) = &local_sum_data { + writeln!(kernel, "{local_sum_data}[{workgroup_local_index}] = sum;").unwrap(); + } writeln!(kernel, "workgroupBarrier();").unwrap(); offset /= 2; // Only threads in the first half do the reduction to avoid OOB reads @@ -291,17 +391,156 @@ impl RmsNormOperation { ) .unwrap(); writeln!(kernel, "sum_sq += neighbor;").unwrap(); + if let Some(local_sum_data) = &local_sum_data { + writeln!( + kernel, + "let neighbor_sum = {local_sum_data}[{workgroup_local_index} + {offset}u];" + ) + .unwrap(); + writeln!(kernel, "sum += neighbor_sum;").unwrap(); + } writeln!(kernel, "}}").unwrap(); } - // Compute RMS and store in shared memory + // Compute denominator and store it in shared memory. writeln!(kernel, "if {workgroup_local_index} == 0u {{").unwrap(); - writeln!(kernel, "let mean_sq = sum_sq / f32({reduce_size});").unwrap(); - writeln!(kernel, "{global_rms} = sqrt(mean_sq + {eps_input});").unwrap(); + if let Some(global_mean) = &global_mean { + writeln!(kernel, "let mean = sum / f32({reduce_size});").unwrap(); + writeln!( + kernel, + "let variance = max(sum_sq / f32({reduce_size}) - mean * mean, f32(0.0));" + ) + .unwrap(); + writeln!(kernel, "{global_mean} = mean;").unwrap(); + writeln!(kernel, "{global_rms} = sqrt(variance + {eps_input});").unwrap(); + } else { + writeln!(kernel, "let mean_sq = sum_sq / f32({reduce_size});").unwrap(); + writeln!(kernel, "{global_rms} = sqrt(mean_sq + {eps_input});").unwrap(); + } writeln!(kernel, "}}").unwrap(); } writeln!(kernel, "workgroupBarrier();").unwrap(); - // Read RMS from shared memory + // Read normalization parameters from shared memory. + if let Some(global_mean) = &global_mean { + writeln!(kernel, "let mean = {global_mean};").unwrap(); + + // Recompute variance from centered values. This preserves LayerNorm's + // two-pass numerical behavior while still keeping the whole op in one kernel. + writeln!(kernel, "var variance_sum = f32(0.0);").unwrap(); + writeln!(kernel, "var variance_index = base_axis_index;").unwrap(); + if large_reduction { + writeln!(kernel, "while (variance_index + 4u <= end_axis_index) {{").unwrap(); + write!(kernel, "let variance_data = vec4<{input_dtype}>(").unwrap(); + for i in 0..4 { + if i > 0 { + write!(kernel, ", ").unwrap(); + } + write!( + kernel, + "{input_tensor}[in_start_offset + (variance_index + {i}u) * {reduce_stride}]" + ) + .unwrap(); + } + writeln!(kernel, ");").unwrap(); + writeln!( + kernel, + "let variance_diff = vec4(variance_data) - vec4(mean);" + ) + .unwrap(); + writeln!(kernel, "let variance_sq = variance_diff * variance_diff;").unwrap(); + writeln!( + kernel, + "variance_sum += variance_sq.x + variance_sq.y + variance_sq.z + variance_sq.w;" + ) + .unwrap(); + writeln!(kernel, "variance_index += 4u;").unwrap(); + writeln!(kernel, "}}").unwrap(); + } + writeln!(kernel, "while (variance_index < end_axis_index) {{").unwrap(); + writeln!( + kernel, + "let variance_data = f32({input_tensor}[in_start_offset + variance_index * {reduce_stride}]);" + ) + .unwrap(); + writeln!(kernel, "let variance_diff = variance_data - mean;").unwrap(); + writeln!(kernel, "variance_sum += variance_diff * variance_diff;").unwrap(); + writeln!(kernel, "variance_index += 1u;").unwrap(); + writeln!(kernel, "}}").unwrap(); + + if device.subgroups_supported() { + let max_subgroup_size = device.max_subgroup_size(); + let local_variance_data = kernel.add_global_array( + KernelGlobalSpace::Workgroup, + DataTypeEnum::F32, + max_subgroup_size.to_string(), + ); + let subgroup_id = kernel.subgroup_index(); + let subgroup_local_id = kernel.subgroup_local_index(); + let subgroups_per_workgroup = kernel.subgroups_per_workgroup(); + + writeln!(kernel, "variance_sum = subgroupAdd(variance_sum);").unwrap(); + writeln!( + kernel, + "{local_variance_data}[{subgroup_id}] = variance_sum;" + ) + .unwrap(); + writeln!(kernel, "workgroupBarrier();").unwrap(); + writeln!( + kernel, + "if {subgroup_local_id} < {subgroups_per_workgroup} {{" + ) + .unwrap(); + writeln!( + kernel, + "variance_sum = {local_variance_data}[{subgroup_local_id}];" + ) + .unwrap(); + writeln!(kernel, "}}").unwrap(); + writeln!(kernel, "else {{").unwrap(); + writeln!(kernel, "variance_sum = f32(0.0);").unwrap(); + writeln!(kernel, "}}").unwrap(); + writeln!(kernel, "variance_sum = subgroupAdd(variance_sum);").unwrap(); + writeln!(kernel, "if {subgroup_id} == 0u {{").unwrap(); + writeln!( + kernel, + "{global_rms} = sqrt(variance_sum / f32({reduce_size}) + {eps_input});" + ) + .unwrap(); + writeln!(kernel, "}}").unwrap(); + } else { + let local_variance_data = kernel.add_global_array( + KernelGlobalSpace::Workgroup, + DataTypeEnum::F32, + blocksize.to_string(), + ); + let mut offset = blocksize; + while offset > 1 { + writeln!( + kernel, + "{local_variance_data}[{workgroup_local_index}] = variance_sum;" + ) + .unwrap(); + writeln!(kernel, "workgroupBarrier();").unwrap(); + offset /= 2; + writeln!(kernel, "if {workgroup_local_index} < {offset}u {{").unwrap(); + writeln!( + kernel, + "let variance_neighbor = {local_variance_data}[{workgroup_local_index} + {offset}u];" + ) + .unwrap(); + writeln!(kernel, "variance_sum += variance_neighbor;").unwrap(); + writeln!(kernel, "}}").unwrap(); + } + writeln!(kernel, "if {workgroup_local_index} == 0u {{").unwrap(); + writeln!( + kernel, + "{global_rms} = sqrt(variance_sum / f32({reduce_size}) + {eps_input});" + ) + .unwrap(); + writeln!(kernel, "}}").unwrap(); + } + writeln!(kernel, "workgroupBarrier();").unwrap(); + } writeln!(kernel, "let rms = {global_rms};").unwrap(); // Phase 3: Normalize and apply weight/bias @@ -335,8 +574,15 @@ impl RmsNormOperation { } writeln!(kernel, ");").unwrap(); - // Normalize: (input / rms) * weight - writeln!(kernel, "let normalized = vec4(data) / rms;").unwrap(); + if remove_mean { + writeln!( + kernel, + "let normalized = (vec4(data) - vec4(mean)) / rms;" + ) + .unwrap(); + } else { + writeln!(kernel, "let normalized = vec4(data) / rms;").unwrap(); + } writeln!(kernel, "var result = normalized * vec4(w);").unwrap(); // Add bias if present @@ -375,7 +621,11 @@ impl RmsNormOperation { ) .unwrap(); writeln!(kernel, "let w = f32({weight_tensor}[out_index]);").unwrap(); - writeln!(kernel, "var result = (data / rms) * w;").unwrap(); + if remove_mean { + writeln!(kernel, "var result = ((data - mean) / rms) * w;").unwrap(); + } else { + writeln!(kernel, "var result = (data / rms) * w;").unwrap(); + } if let Some(bias_tensor) = &bias_tensor { writeln!(kernel, "result += f32({bias_tensor}[out_index]);").unwrap(); @@ -529,7 +779,12 @@ impl Operation for RmsNormOperation { fn name(&self) -> String { format!( - "rms_norm_fused_{}_{}{}", + "{}_{}_{}{}", + if self.remove_mean { + "layer_norm_fused" + } else { + "rms_norm_fused" + }, self.rank(), self.input_dtype, if self.bias.is_some() { "_bias" } else { "" } @@ -541,7 +796,7 @@ impl Operation for RmsNormOperation { map: &rustc_hash::FxHashMap, ) -> crate::TensorLayoutInfo { let input_layout = map.get(&self.input).unwrap(); - input_layout.clone() + crate::TensorLayoutInfo::new(Layout::contiguous(input_layout.shape()), self.input_dtype) } } @@ -618,6 +873,195 @@ async fn test_rms_norm_fused_vs_composite() { } } +#[cfg(test)] +#[tokio::test] +async fn test_layer_norm_fused_vs_composite() { + use crate::Device; + + let device = Device::test_instance(); + + let input = Tensor::new(&device, &[[1.0f32, 2.0, 3.0, 4.0], [5.0, 7.0, 11.0, 13.0]]); + let weight = Tensor::new(&device, &[1.0f32, 0.5, 2.0, -1.0]); + let bias = Tensor::new(&device, &[0.0f32, 0.25, -0.5, 1.0]); + let eps = 1e-5; + + let composite_result = input.layer_norm(&weight, Some(&bias), eps, true); + let composite_output = composite_result.as_slice().await.unwrap(); + + let fused_result = input.layer_norm_fused(&weight, Some(&bias), eps); + let fused_output = fused_result.as_slice().await.unwrap(); + + for i in 0..2 { + for j in 0..4 { + let diff = (composite_output[[i, j]] - fused_output[[i, j]]).abs(); + assert!( + diff < 0.001, + "Mismatch at [{}, {}]: composite={}, fused={}, diff={}", + i, + j, + composite_output[[i, j]], + fused_output[[i, j]], + diff + ); + } + } +} + +#[cfg(test)] +#[tokio::test] +async fn test_layer_norm_fused_large() { + use crate::Device; + + let device = Device::test_instance(); + + let hidden_size = 320; + let rows = 5; + let input_data: Vec> = (0..rows) + .map(|row| { + (0..hidden_size) + .map(|col| (((row * 17 + col * 13) % 29) as f32 - 14.0) * 0.125) + .collect() + }) + .collect(); + let weight_data: Vec = (0..hidden_size) + .map(|col| 0.75 + (col % 7) as f32 * 0.05) + .collect(); + + let input: Tensor<2, f32> = Tensor::new(&device, &input_data); + let weight: Tensor<1, f32> = Tensor::new(&device, &weight_data); + let eps = 1e-5; + + let composite_result = input.layer_norm(&weight, None, eps, true); + let composite_output = composite_result.as_slice().await.unwrap(); + + let fused_result = input.layer_norm_fused(&weight, None, eps); + let fused_output = fused_result.as_slice().await.unwrap(); + + for row in 0..rows { + for col in 0..hidden_size { + let diff = (composite_output[[row, col]] - fused_output[[row, col]]).abs(); + assert!( + diff < 0.001, + "Mismatch at [{}, {}]: composite={}, fused={}, diff={}", + row, + col, + composite_output[[row, col]], + fused_output[[row, col]], + diff + ); + } + } +} + +#[cfg(test)] +#[tokio::test] +async fn test_layer_norm_fused_non_contiguous_input_layout() { + use crate::Device; + + let device = Device::test_instance(); + + let input = Tensor::new( + &device, + &[ + [ + [1.0f32, 2.0, 3.0, 4.0], + [5.0, 7.0, 11.0, 13.0], + [-4.0, -1.0, 0.5, 3.0], + [8.0, 6.0, 4.0, 2.0], + ], + [ + [0.25, 0.5, 1.0, 2.0], + [3.0, 1.5, 0.75, 0.25], + [9.0, 10.0, 12.0, 15.0], + [-8.0, -4.0, -2.0, -1.0], + ], + ], + ) + .narrow(1, 1, 2); + let weight = Tensor::new(&device, &[1.0f32, 0.5, 2.0, -1.0]); + let eps = 1e-5; + + let composite_result = input.layer_norm(&weight, None, eps, true); + let composite_output = composite_result.as_slice().await.unwrap(); + + let fused_result = input.layer_norm_fused(&weight, None, eps); + let fused_output = fused_result.as_slice().await.unwrap(); + + assert_eq!(fused_output.shape(), &[2, 2, 4]); + for batch in 0..2 { + for row in 0..2 { + for col in 0..4 { + let diff = + (composite_output[[batch, row, col]] - fused_output[[batch, row, col]]).abs(); + assert!( + diff < 0.001, + "Mismatch at [{}, {}, {}]: composite={}, fused={}, diff={}", + batch, + row, + col, + composite_output[[batch, row, col]], + fused_output[[batch, row, col]], + diff + ); + } + } + } +} + +#[cfg(test)] +#[tokio::test] +async fn test_layer_norm_fused_matmul_chain_matches_composite() { + use crate::Device; + + let device = Device::test_instance(); + + let rows = 7; + let hidden = 320; + let out = 13; + let input_data: Vec> = (0..rows) + .map(|row| { + (0..hidden + 3) + .map(|col| (((row * 19 + col * 11) % 37) as f32 - 18.0) * 0.03125) + .collect() + }) + .collect(); + let proj_data: Vec> = (0..hidden) + .map(|row| { + (0..out) + .map(|col| (((row * 7 + col * 5) % 23) as f32 - 11.0) * 0.015625) + .collect() + }) + .collect(); + let norm_weight: Vec = (0..hidden) + .map(|col| 0.9 + (col % 5) as f32 * 0.025) + .collect(); + + let input = Tensor::new(&device, &input_data).narrow(1, 0, hidden); + let weight = Tensor::new(&device, &norm_weight); + let proj = Tensor::new(&device, &proj_data); + let eps = 1e-5; + + let composite = input.layer_norm(&weight, None, eps, true).mat_mul(&proj); + let fused = input.layer_norm_fused(&weight, None, eps).mat_mul(&proj); + let composite = composite.as_slice().await.unwrap(); + let fused = fused.as_slice().await.unwrap(); + + for row in 0..rows { + for col in 0..out { + let diff = (composite[[row, col]] - fused[[row, col]]).abs(); + assert!( + diff < 0.001, + "Mismatch at [{}, {}]: composite={}, fused={}, diff={}", + row, + col, + composite[[row, col]], + fused[[row, col]], + diff + ); + } + } +} + #[cfg(test)] #[tokio::test] async fn test_rms_norm_fused_with_bias() { diff --git a/fusor-ml/core/src/composite/rope_fused.rs b/fusor-ml/core/src/composite/rope_fused.rs index 31f0ccf9d..a5b30e7b8 100644 --- a/fusor-ml/core/src/composite/rope_fused.rs +++ b/fusor-ml/core/src/composite/rope_fused.rs @@ -8,13 +8,23 @@ impl Tensor<4, D> { /// Apply fused interleaved RoPE (rotary position embedding). /// This pairs adjacent elements: (0, 1), (2, 3), etc. pub fn rope_fused(&self, cos: &Tensor<2, D>, sin: &Tensor<2, D>) -> Tensor<4, D> { - self.rope_fused_impl(cos, sin, RopeMode::Interleaved) + self.rope_fused_impl(cos, sin, RopeMode::Interleaved, None) + } + + /// Apply fused interleaved RoPE to the first `rotary_dim` elements and pass the rest through. + pub fn rope_partial_fused( + &self, + cos: &Tensor<2, D>, + sin: &Tensor<2, D>, + rotary_dim: usize, + ) -> Tensor<4, D> { + self.rope_fused_impl(cos, sin, RopeMode::Interleaved, Some(rotary_dim)) } /// Apply fused normal RoPE (rotary position embedding). /// This pairs first half with second half: (0, head_dim/2), (1, head_dim/2+1), etc. pub fn rope_normal_fused(&self, cos: &Tensor<2, D>, sin: &Tensor<2, D>) -> Tensor<4, D> { - self.rope_fused_impl(cos, sin, RopeMode::Normal) + self.rope_fused_impl(cos, sin, RopeMode::Normal, None) } fn rope_fused_impl( @@ -22,8 +32,12 @@ impl Tensor<4, D> { cos: &Tensor<2, D>, sin: &Tensor<2, D>, mode: RopeMode, + rotary_dim: Option, ) -> Tensor<4, D> { let [_, _, sequence_length, head_dim] = *self.shape(); + let rotary_dim = rotary_dim.unwrap_or(head_dim); + assert!(rotary_dim <= head_dim); + assert_eq!(rotary_dim % 2, 0); // Narrow cos/sin to the sequence length let cos = cos.narrow(0, 0, sequence_length); let sin = sin.narrow(0, 0, sequence_length); @@ -36,6 +50,7 @@ impl Tensor<4, D> { shape: (*self.shape()).into(), mode, head_dim, + rotary_dim, } .to_nary(); @@ -61,6 +76,7 @@ struct RopeFusedOperation { shape: Box<[usize]>, mode: RopeMode, head_dim: usize, + rotary_dim: usize, } impl RopeFusedOperation { @@ -108,9 +124,27 @@ impl RopeFusedOperation { ); // Final expression: input * cos + neighbor * sin_with_sign - let input_times_cos = NaryExpr::mul(input_val, cos_val, self.datatype); + let input_times_cos = NaryExpr::mul(input_val.clone(), cos_val, self.datatype); let neighbor_times_sin = NaryExpr::mul(neighbor_val, sin_with_sign, self.datatype); - NaryExpr::add(input_times_cos, neighbor_times_sin, self.datatype) + let rotated = NaryExpr::add(input_times_cos, neighbor_times_sin, self.datatype); + if self.rotary_dim == self.head_dim { + rotated + } else { + let in_rotary = NaryExpr::unary_op( + NaryExpr::DimIndex(dim_last_idx), + "lt_rotary", + format!("let output = select(0u, 1u, input < {}u);", self.rotary_dim), + DataTypeEnum::U32, + DataTypeEnum::U32, + ); + NaryExpr::select( + in_rotary, + rotated, + input_val, + DataTypeEnum::U32, + self.datatype, + ) + } } /// Build the index expressions for accessing cos/sin values (returns Vec for indexed_input) @@ -120,13 +154,25 @@ impl RopeFusedOperation { let cos_sin_dim = match self.mode { // Interleaved: index = dim_last / 2 - RopeMode::Interleaved => NaryExpr::unary_op( - dim_last, - "div2", - "let output = input / 2u;", - DataTypeEnum::U32, - DataTypeEnum::U32, - ), + RopeMode::Interleaved => { + if self.rotary_dim == self.head_dim { + NaryExpr::unary_op( + dim_last, + "div2", + "let output = input / 2u;", + DataTypeEnum::U32, + DataTypeEnum::U32, + ) + } else { + NaryExpr::unary_op( + dim_last, + "partial_div2", + format!("let output = min(input, {}u) / 2u;", self.rotary_dim - 1), + DataTypeEnum::U32, + DataTypeEnum::U32, + ) + } + } // Normal: index = dim_last % half RopeMode::Normal => { let half = self.head_dim / 2; @@ -295,6 +341,39 @@ mod tests { } } + #[tokio::test] + async fn test_rope_partial_fused_interleaved() { + let device = Device::test_instance(); + let cos = Tensor::new( + &device, + &[[1.0, 0.5], [0.75, 0.25], [0.25, 0.125], [0.9, 0.33]], + ); + let sin = Tensor::new( + &device, + &[[0.0, 0.25], [0.5, 0.75], [0.125, 0.875], [0.1, 0.66]], + ); + let x = Tensor::new( + &device, + &[[[ + [1.0, 2.0, 3.0, 4.0, 50.0, 60.0], + [5.0, 6.0, 7.0, 8.0, 70.0, 80.0], + [9.0, 10.0, 11.0, 12.0, 90.0, 100.0], + [13.0, 14.0, 15.0, 16.0, 110.0, 120.0], + ]]], + ); + + let rotated = x.narrow(3, 0, 4).rope_fused(&cos, &sin); + let pass = x.narrow(3, 4, 2); + let expected = Tensor::cat([rotated, pass], 3); + let fused = x.rope_partial_fused(&cos, &sin, 4); + + let expected = expected.as_slice().await.unwrap(); + let fused = fused.as_slice().await.unwrap(); + for (actual, expected) in fused.as_slice().iter().zip(expected.as_slice()) { + assert!((actual - expected).abs() < 1e-5); + } + } + #[tokio::test] async fn test_rope_fused_normal() { let device = Device::test_instance(); diff --git a/fusor-ml/core/src/composite/swiglu.rs b/fusor-ml/core/src/composite/swiglu.rs new file mode 100644 index 000000000..631cca516 --- /dev/null +++ b/fusor-ml/core/src/composite/swiglu.rs @@ -0,0 +1,79 @@ +use std::sync::Arc; + +use crate::{ + DataTypeEnum, Tensor, + nary_wise::NaryExpr, + tensor::{LazyTensorData, TensorInfo}, +}; + +impl Tensor<3, f32> { + pub fn swiglu_split(&self, intermediate_size: usize) -> Self { + assert_eq!( + self.shape()[2], + intermediate_size * 2, + "swiglu_split expects the last dimension to be 2 * intermediate_size" + ); + let shape = [self.shape()[0], self.shape()[1], intermediate_size]; + let gate_index = NaryExpr::unary_op( + NaryExpr::DimIndex(2), + "add_intermediate", + format!("let output = input + {intermediate_size}u;"), + DataTypeEnum::U32, + DataTypeEnum::U32, + ); + let state = NaryExpr::indexed_input( + 0, + vec![ + NaryExpr::DimIndex(0), + NaryExpr::DimIndex(1), + NaryExpr::DimIndex(2), + ], + ); + let gate = NaryExpr::indexed_input( + 0, + vec![NaryExpr::DimIndex(0), NaryExpr::DimIndex(1), gate_index], + ); + let gate = NaryExpr::unary_op( + gate, + "silu", + "let output = input / (1.0 + exp(-input));", + DataTypeEnum::F32, + DataTypeEnum::F32, + ); + let expression = NaryExpr::mul(state, gate, DataTypeEnum::F32); + let operation = crate::nary_wise::NaryOperation { + inputs: vec![self.key()], + expression, + shape: shape.into(), + output_datatype: DataTypeEnum::F32, + }; + + let device = self.device().clone(); + let info = TensorInfo::new(shape.into(), DataTypeEnum::F32); + let key = device.compute_graph().create_custom(Arc::new(operation)); + Tensor::from_parts(LazyTensorData::from_parts(device, info, key)) + } +} + +#[cfg(test)] +#[tokio::test] +async fn test_swiglu_split_matches_composite() { + use crate::Device; + + let device = Device::test_instance(); + let data = [[ + [1.0, -2.0, 0.5, -0.25, 0.1, 1.5], + [3.0, 0.25, -1.0, 2.0, -0.5, 0.75], + ]]; + let tensor = Tensor::new(&device, &data); + let fused = tensor.swiglu_split(3); + let state = tensor.narrow(2, 0, 3); + let gate = tensor.narrow(2, 3, 3).silu(); + let expected = state * gate; + + let fused = fused.as_slice().await.unwrap(); + let expected = expected.as_slice().await.unwrap(); + for (actual, expected) in fused.as_slice().iter().zip(expected.as_slice()) { + assert!((actual - expected).abs() < 1e-6); + } +} diff --git a/fusor-ml/core/src/compute_graph/layout_pass.rs b/fusor-ml/core/src/compute_graph/layout_pass.rs index ecac29cc3..513ee544d 100644 --- a/fusor-ml/core/src/compute_graph/layout_pass.rs +++ b/fusor-ml/core/src/compute_graph/layout_pass.rs @@ -112,7 +112,7 @@ impl LayoutPass { key: NodeIndex, operation: &crate::quantized::matmul::QMatMulOperation, ) { - let Some(first_layout) = self.output_layout.get(&operation.input) else { + let Some(_first_layout) = self.output_layout.get(&operation.input) else { self.queue.push_back(operation.input); self.queue.push_back(key); return; @@ -120,7 +120,7 @@ impl LayoutPass { let output_layout = Layout::contiguous(&operation.out_shape); self.output_layout.insert( key, - TensorLayoutInfo::new(output_layout, first_layout.datatype()), + TensorLayoutInfo::new(output_layout, operation.post_element_wise.out_datatype()), ); } diff --git a/fusor-ml/core/src/compute_graph/mod.rs b/fusor-ml/core/src/compute_graph/mod.rs index d89c6f76b..6d7d184b5 100644 --- a/fusor-ml/core/src/compute_graph/mod.rs +++ b/fusor-ml/core/src/compute_graph/mod.rs @@ -5,7 +5,6 @@ pub(crate) use petgraph::graph::NodeIndex; use petgraph::prelude::StableGraph; use resolve::Resolver; use tabbycat::Graph; -use wgpu::CommandEncoderDescriptor; mod layout_pass; mod queue; @@ -14,11 +13,19 @@ mod visualize; use crate::{ DataTypeEnum, Device, ElementWiseOperation, MatMulOperation, PairWiseOperation, QMatrix, - ReduceOperation, composite::where_cond::WhereCondOperation, - compute_graph::resolve::ResolverResult, dequantize::DequantizeOperation, - index_select::IndexSelectOperation, map_layout::MapLayoutOperation, mir::operation::Operation, - nary_wise::NaryOperation, quantized::matmul::QMatMulOperation, resize::ResizeOperation, - slice_assign::SliceAssignOperation, tensor::TensorData, visit_tiled::MaybeQData, + ReduceOperation, + composite::where_cond::WhereCondOperation, + compute_graph::resolve::{ResolverManyResult, ResolverResult}, + dequantize::DequantizeOperation, + index_select::IndexSelectOperation, + map_layout::MapLayoutOperation, + mir::operation::Operation, + nary_wise::NaryOperation, + quantized::matmul::QMatMulOperation, + resize::ResizeOperation, + slice_assign::SliceAssignOperation, + tensor::TensorData, + visit_tiled::MaybeQData, }; #[derive(Clone)] @@ -101,16 +108,10 @@ impl ComputeGraph { } pub(crate) fn resolve(&self, key: NodeIndex, device: &Device) -> ResolverResult { - let mut encoder = device - .wgpu_device() - .create_command_encoder(&CommandEncoderDescriptor { - label: Some("ComputeGraph Encoder"), - }); let data = self.with_mut(|inner| { - let mut resolver = Resolver::new(inner, key, &mut encoder); + let mut resolver = Resolver::new(inner, key, device); resolver.run(inner) }); - device.wgpu_queue().submit(Some(encoder.finish())); // Reset the written flag on all buffers device.reset_initialized_buffers(); @@ -129,10 +130,58 @@ impl ComputeGraph { data } + pub(crate) fn resolve_many(&self, keys: &[NodeIndex], device: &Device) -> ResolverManyResult { + let data = self.with_mut(|inner| { + let mut resolver = Resolver::new_many(inner, keys, device); + resolver.run_many(inner) + }); + device.reset_initialized_buffers(); + + if let (Some(pipeline_cache), Some(cache_file)) = + (device.wgpu_cache(), device.wgpu_cache_file()) + { + let data = pipeline_cache.get_data(); + if let Some(data) = data { + let temp_file = cache_file.with_extension("temp"); + std::fs::write(&temp_file, &data).unwrap(); + std::fs::rename(&temp_file, cache_file).unwrap(); + } + } + + data + } + pub(crate) fn graphvis(&self, root: NodeIndex) -> Graph { self.with_mut(|inner| inner.graphvis(root)) } + #[cfg(test)] + pub(crate) fn node_variant_name(&self, key: NodeIndex) -> &'static str { + self.with_mut(|inner| { + let node = inner + .nodes + .nodes + .node_weight(key) + .expect("Node not found in graph"); + match &node.variant { + ComputeGraphNodeVariant::ElementWise(_) => "ElementWise", + ComputeGraphNodeVariant::PairWise(_) => "PairWise", + ComputeGraphNodeVariant::Nary(_) => "Nary", + ComputeGraphNodeVariant::SliceAssign(_) => "SliceAssign", + ComputeGraphNodeVariant::Resize(_) => "Resize", + ComputeGraphNodeVariant::MapLayout(_) => "MapLayout", + ComputeGraphNodeVariant::Dequantize(_) => "Dequantize", + ComputeGraphNodeVariant::MatMul(_) => "MatMul", + ComputeGraphNodeVariant::QMatMul(_) => "QMatMul", + ComputeGraphNodeVariant::Tensor(_) => "Tensor", + ComputeGraphNodeVariant::Reduce(_) => "Reduce", + ComputeGraphNodeVariant::IndexSelect(_) => "IndexSelect", + ComputeGraphNodeVariant::WhereCond(_) => "WhereCond", + ComputeGraphNodeVariant::Custom(_) => "Custom", + } + }) + } + pub(crate) fn add_reference(&self, key: NodeIndex) { self.with_mut(|inner| inner.add_reference(key)); } @@ -140,6 +189,10 @@ impl ComputeGraph { pub(crate) fn remove_reference(&self, key: NodeIndex) { self.with_mut(|inner| inner.remove_reference(key)); } + + pub(crate) fn is_cached(&self, key: NodeIndex) -> bool { + self.with_mut(|inner| inner.get_cached_result(key).is_some()) + } } #[derive(Default)] @@ -189,7 +242,7 @@ impl ComputeGraphNodeVariant { f(op.second); } ComputeGraphNodeVariant::QMatMul(op) => { - f(op.input); + op.visit_dependencies(f); } ComputeGraphNodeVariant::Reduce(op) => f(op.value), ComputeGraphNodeVariant::MapLayout(op) => f(op.input), diff --git a/fusor-ml/core/src/compute_graph/resolve.rs b/fusor-ml/core/src/compute_graph/resolve.rs index 6c5547d8c..4f22643b6 100644 --- a/fusor-ml/core/src/compute_graph/resolve.rs +++ b/fusor-ml/core/src/compute_graph/resolve.rs @@ -5,10 +5,10 @@ use std::sync::Arc; use petgraph::algo::toposort; use petgraph::stable_graph::StableGraph; use rustc_hash::{FxHashMap, FxHashSet}; -use wgpu::CommandEncoder; +use wgpu::{CommandEncoder, CommandEncoderDescriptor}; use crate::{ - ElementWiseFunctions, + Device, ElementWiseFunctions, mir::{ inputs::{KernelInputValue, MirValue}, kernel::GenericKernel, @@ -16,7 +16,6 @@ use crate::{ workgroup_shape::{self, WorkgroupShapeConstraints}, }, nary_wise::{NaryExpr, NaryOperation}, - quantized::matmul::QMatMulOperation, tensor::TensorData, }; @@ -27,6 +26,11 @@ pub(crate) struct ResolverResult { pub(crate) total_kernels: usize, } +pub(crate) struct ResolverManyResult { + pub(crate) data: Vec, + pub(crate) total_kernels: usize, +} + #[derive(Debug, Clone)] struct ExecutionNode { inner_idx: NodeIndex, @@ -36,19 +40,29 @@ struct ExecutionNode { type ExecutionGraph = StableGraph; type ExecutionNodeIndex = petgraph::graph::NodeIndex; -pub(crate) struct Resolver<'a> { - command_encoder: &'a mut CommandEncoder, +const MAX_FUSED_NARY_INPUTS: usize = 8; +const MAX_FUSED_NARY_EXPR_NODES: usize = 128; +const MAX_BATCHED_STORAGE_BUFFERS: usize = 16; + +pub(crate) struct Resolver { + device: Device, + command_encoder: CommandEncoder, execution_graph: ExecutionGraph, node_mapping: FxHashMap, - target: NodeIndex, + targets: Vec, resolved_set: FxHashSet, + dispatched_kernels_since_submit: usize, } -impl<'a> Resolver<'a> { - pub(crate) fn new( +impl Resolver { + pub(crate) fn new(graph: &mut ComputeGraphInner, target: NodeIndex, device: &Device) -> Self { + Self::new_many(graph, &[target], device) + } + + pub(crate) fn new_many( graph: &mut ComputeGraphInner, - target: NodeIndex, - command_encoder: &'a mut CommandEncoder, + targets: &[NodeIndex], + device: &Device, ) -> Self { let resolved_set = graph .nodes @@ -64,20 +78,30 @@ impl<'a> Resolver<'a> { }) .collect(); Self { - command_encoder, - target, + device: device.clone(), + command_encoder: device.wgpu_device().create_command_encoder( + &CommandEncoderDescriptor { + label: Some("ComputeGraph Resolver Encoder"), + }, + ), + targets: targets.to_vec(), execution_graph: Default::default(), node_mapping: Default::default(), resolved_set, + dispatched_kernels_since_submit: 0, } } - pub(crate) fn run(&mut self, graph: &mut ComputeGraphInner) -> ResolverResult { + fn execute(&mut self, graph: &mut ComputeGraphInner) -> usize { let device = graph.device(); let max_subgroup_size = device.max_subgroup_size(); + let max_storage_buffers = (device.limits().max_storage_buffers_per_shader_stage as usize) + .min(MAX_BATCHED_STORAGE_BUFFERS); // Pass 1: Build execution graph - self.build_execution_graph(graph, self.target); + for target in self.targets.clone() { + self.build_execution_graph(graph, target); + } // Pass 2: Apply Rewrite Rules self.optimize(graph); @@ -106,12 +130,14 @@ impl<'a> Resolver<'a> { let mut current_constraints = WorkgroupShapeConstraints::new(); let mut pending_operations = Vec::new(); let mut inputs = Vec::new(); + let mut outputs = Vec::new(); let mut all_input_values = Vec::new(); let mut kernel = GenericKernel::new(); let mut total_kernels = 0; for (node, operation) in queued_operations { let new_inputs = operation.inputs(graph); + let new_output = operation.output(graph, &new_inputs); let constraint = operation.workgroup_shape_constraints(&device); let mut new_merged = current_constraints.clone(); new_merged.merge(&constraint); @@ -120,7 +146,9 @@ impl<'a> Resolver<'a> { "Failed to find a valid workgroup shape for constraints {current_constraints:?}" ) }); - let mut extend = self.should_extend_kernel(new_inputs.clone(), &inputs); + let mut extend = self.should_extend_kernel(&new_inputs, &new_output, &inputs, &outputs); + extend &= Self::combined_storage_buffer_count(&all_input_values, &new_inputs) + <= max_storage_buffers; extend &= new_merged.solve(max_subgroup_size).is_some(); if extend { current_constraints = new_merged; @@ -136,9 +164,11 @@ impl<'a> Resolver<'a> { old_best, ); pending_operations.clear(); + outputs.clear(); all_input_values.clear(); inputs.clear(); kernel.clear(); + self.dispatched_kernels_since_submit += 1; } current_constraints = constraint; } @@ -163,7 +193,9 @@ impl<'a> Resolver<'a> { &mut kernel, node, operation, + new_output, &mut inputs, + &mut outputs, &mut all_input_values, &mut pending_operations, ); @@ -185,10 +217,22 @@ impl<'a> Resolver<'a> { &all_input_values, old_best, ); + self.dispatched_kernels_since_submit += 1; } + self.submit_encoder(); + let data = graph - .get_result(self.target) + .get_result(self.targets[0]) + .expect("Target result not cached"); + graph.set_cached_result(self.targets[0], data.clone()); + total_kernels + } + + pub(crate) fn run(&mut self, graph: &mut ComputeGraphInner) -> ResolverResult { + let total_kernels = self.execute(graph); + let data = graph + .get_result(self.targets[0]) .expect("Target result not cached"); ResolverResult { data, @@ -196,6 +240,35 @@ impl<'a> Resolver<'a> { } } + pub(crate) fn run_many(&mut self, graph: &mut ComputeGraphInner) -> ResolverManyResult { + let total_kernels = self.execute(graph); + let data = self + .targets + .iter() + .map(|target| graph.get_result(*target).expect("Target result not cached")) + .collect(); + ResolverManyResult { + data, + total_kernels, + } + } + + fn submit_encoder(&mut self) { + if self.dispatched_kernels_since_submit == 0 { + return; + } + let encoder = std::mem::replace( + &mut self.command_encoder, + self.device + .wgpu_device() + .create_command_encoder(&CommandEncoderDescriptor { + label: Some("ComputeGraph Resolver Encoder"), + }), + ); + self.device.wgpu_queue().submit(Some(encoder.finish())); + self.dispatched_kernels_since_submit = 0; + } + fn build_execution_graph( &mut self, graph: &ComputeGraphInner, @@ -296,12 +369,7 @@ impl<'a> Resolver<'a> { }; Some(Arc::new(nary)) } - ComputeGraphNodeVariant::QMatMul(op) => Some(Arc::new(QMatMulOperation::new( - op.input_datatype, - &op.in_shape, - op.input, - op.matrix.clone(), - ))), + ComputeGraphNodeVariant::QMatMul(op) => Some(Arc::new(op.clone())), ComputeGraphNodeVariant::Dequantize(op) => Some(Arc::new(op.clone())), ComputeGraphNodeVariant::WhereCond(op) => Some(Arc::new(op.to_nary())), ComputeGraphNodeVariant::Tensor(_) => None, // Handled in execution loop @@ -341,6 +409,7 @@ impl<'a> Resolver<'a> { || self.try_fuse_naries(graph, node_idx) || self.try_fuse_into_reduce(graph, node_idx) || self.try_fuse_into_matmul(graph, node_idx) + || self.try_fuse_into_q_matmul(graph, node_idx) || self.try_fuse_into_dequantize(graph, node_idx); if changed { @@ -566,6 +635,12 @@ impl<'a> Resolver<'a> { // Skip fusion - would exceed GPU binding limit continue; } + if unique_inputs.len() > MAX_FUSED_NARY_INPUTS { + continue; + } + if Self::nary_expr_node_count(&new_expression) > MAX_FUSED_NARY_EXPR_NODES { + continue; + } expression = new_expression; all_inputs.extend(input_nary.inputs.iter().copied()); @@ -784,6 +859,24 @@ impl<'a> Resolver<'a> { } } + fn nary_expr_node_count(expr: &NaryExpr) -> usize { + match expr { + NaryExpr::Op { children, .. } => { + 1 + children + .iter() + .map(Self::nary_expr_node_count) + .sum::() + } + NaryExpr::IndexedInput { indices, .. } => { + 1 + indices + .iter() + .map(Self::nary_expr_node_count) + .sum::() + } + NaryExpr::DimIndex(_) => 1, + } + } + fn remap_input_indices(expr: &NaryExpr, mapping: &FxHashMap) -> NaryExpr { match expr { NaryExpr::Op { children, function } => NaryExpr::Op { @@ -1006,9 +1099,160 @@ impl<'a> Resolver<'a> { true } - fn should_extend_kernel(&mut self, _: Vec, _: &[Vec]) -> bool { - // TODO: Restore with better testing. This passes all tests in fusor, but breaks rbert and rwhisper - false + fn try_fuse_into_q_matmul( + &mut self, + graph: &mut ComputeGraphInner, + node_idx: ExecutionNodeIndex, + ) -> bool { + let node_variant = self.execution_graph[node_idx].variant.clone(); + + // Post-op: fuse elementwise after qmatmul + if let Some(el_op) = Self::try_get_elementwise(&node_variant) { + let input_inner = el_op.value; + if !self.check_cached(graph, input_inner) + && 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::QMatMul(q_mat_mul_op) = input_variant { + let mut new_q_mat_mul = q_mat_mul_op.clone(); + let mut existing_post = new_q_mat_mul.post_element_wise.functions.clone(); + existing_post.extend(el_op.functions.functions.iter().cloned()); + new_q_mat_mul.post_element_wise = ElementWiseFunctions::new( + existing_post, + q_mat_mul_op.post_element_wise.input_datatype(), + ); + + self.execution_graph[node_idx].variant = + ComputeGraphNodeVariant::QMatMul(new_q_mat_mul.clone()); + + let input_inner = q_mat_mul_op.input; + if let Some(idx) = self.get_input_node_in_exec_graph(input_inner) { + self.execution_graph.add_edge(idx, 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, &[input_inner]); + self.remove_node_if_dead(input_exec_idx); + return true; + } + } + } + + // Pre-op: fuse elementwise before qmatmul input + let ComputeGraphNodeVariant::QMatMul(q_mat_mul_op) = &node_variant else { + return false; + }; + + if self.check_cached(graph, q_mat_mul_op.input) { + return false; + } + + let Some(input_exec_idx) = self.get_input_node_in_exec_graph(q_mat_mul_op.input) else { + return false; + }; + let Some(el_op) = Self::try_get_elementwise(&self.execution_graph[input_exec_idx].variant) + else { + return false; + }; + + let mut new_q_mat_mul = q_mat_mul_op.clone(); + new_q_mat_mul.input = el_op.value; + new_q_mat_mul.pre_element_wise = el_op.functions.clone(); + + self.execution_graph[node_idx].variant = + ComputeGraphNodeVariant::QMatMul(new_q_mat_mul.clone()); + + if let Some(edge) = self.execution_graph.find_edge(input_exec_idx, node_idx) { + self.execution_graph.remove_edge(edge); + } + if let Some(new_input_exec_idx) = self.get_input_node_in_exec_graph(new_q_mat_mul.input) { + self.execution_graph + .add_edge(new_input_exec_idx, node_idx, ()); + } + self.add_physical_dependencies(graph, node_idx, &[new_q_mat_mul.input]); + self.remove_node_if_dead(input_exec_idx); + true + } + + fn should_extend_kernel( + &mut self, + new_inputs: &[MirValue], + new_output: &MirValue, + inputs: &[Vec], + outputs: &[MirValue], + ) -> bool { + let new_output_tensor = new_output.as_tensor(); + + for input in new_inputs { + let Some(input_tensor) = input.as_tensor() else { + continue; + }; + for pending_output in outputs { + let Some(pending_output_tensor) = pending_output.as_tensor() else { + continue; + }; + + // A later read of a pending write needs a dispatch boundary; + // otherwise workgroups can observe incomplete output data. + if std::sync::Arc::ptr_eq(input_tensor.buffer(), pending_output_tensor.buffer()) { + return false; + } + } + } + + let Some(new_output_tensor) = new_output_tensor else { + return true; + }; + + for pending_output in outputs { + let Some(pending_output_tensor) = pending_output.as_tensor() else { + continue; + }; + // Two writes to the same storage need to be ordered. + if std::sync::Arc::ptr_eq(new_output_tensor.buffer(), pending_output_tensor.buffer()) { + return false; + } + } + + for pending in inputs { + for input in pending { + let Some(input_tensor) = input.as_tensor() else { + continue; + }; + + // A later write that aliases an earlier read also needs a boundary. + if std::sync::Arc::ptr_eq(new_output_tensor.buffer(), input_tensor.buffer()) { + return false; + } + } + } + + true + } + + fn combined_storage_buffer_count( + all_input_values: &[KernelInputValue], + new_inputs: &[MirValue], + ) -> usize { + let mut combined = all_input_values.to_vec(); + for input in new_inputs { + input.visit_input_values(|value| { + if !combined.iter().any(|existing| *existing == value) { + combined.push(value); + } + }); + } + + combined + .iter() + .filter(|value| { + matches!( + value, + KernelInputValue::TensorBuffer(_) | KernelInputValue::QBuffer(_) + ) + }) + .count() } #[allow(clippy::too_many_arguments)] @@ -1019,7 +1263,9 @@ impl<'a> Resolver<'a> { kernel: &mut GenericKernel, key: NodeIndex, operation: Arc, + result: MirValue, inputs: &mut Vec>, + outputs: &mut Vec, all_input_values: &mut Vec, queued_operations: &mut Vec<(NodeIndex, Arc)>, ) { @@ -1033,13 +1279,14 @@ impl<'a> Resolver<'a> { } }); } - let result = operation.output(graph, &new_inputs); let MirValue::Tensor(resolved) = result else { panic!("Kernel input value is not a tensor"); }; + let output = resolved.clone(); // Cache the result graph.set_cached_result(key, resolved); inputs.push(new_inputs); + outputs.push(MirValue::Tensor(output)); queued_operations.push((key, operation)); } @@ -1092,12 +1339,24 @@ impl<'a> Resolver<'a> { } kernel.set_workgroup_size(workgroup_shape); let device = graph.device(); + let profile = std::env::var_os("FUSOR_RESOLVER_PROFILE").is_some(); + if profile { + let names = queued_operations + .iter() + .map(|(key, operation)| format!("{}#{key:?}", operation.name())) + .collect::>() + .join(" | "); + eprintln!("resolver kernel start: {names}"); + } kernel.run( &device, all_input_values, - self.command_encoder, + &mut self.command_encoder, max_dispatch_size, ); + if profile { + eprintln!("resolver kernel queued"); + } } /// Wrap an expression with element-wise functions (each becomes a unary Op node) @@ -1117,3 +1376,126 @@ impl<'a> Resolver<'a> { expr } } + +#[cfg(test)] +mod tests { + use crate::{DataTypeEnum, Device, Tensor, TensorData}; + + use super::*; + + #[test] + fn test_should_extend_kernel_allows_disjoint_tensor_buffers() { + let device = Device::test_instance(); + let mut graph = ComputeGraphInner::new(&device); + let mut resolver = Resolver::new_many(&mut graph, &[], &device); + + let pending_tensor = TensorData::new_for_shape(&device, &[4, 4], DataTypeEnum::F32); + let new_tensor = TensorData::new_for_shape(&device, &[4, 4], DataTypeEnum::F32); + let pending_output = TensorData::new_for_shape(&device, &[4, 4], DataTypeEnum::F32); + let new_output = TensorData::new_for_shape(&device, &[4, 4], DataTypeEnum::F32); + + assert!(resolver.should_extend_kernel( + &[ + MirValue::Tensor(new_tensor), + MirValue::Tensor(new_output.clone()) + ], + &MirValue::Tensor(new_output), + &[vec![ + MirValue::Tensor(pending_tensor), + MirValue::Tensor(pending_output.clone()), + ]], + &[MirValue::Tensor(pending_output)], + )); + } + + #[test] + fn test_should_extend_kernel_allows_shared_read_only_inputs() { + let device = Device::test_instance(); + let mut graph = ComputeGraphInner::new(&device); + let mut resolver = Resolver::new_many(&mut graph, &[], &device); + + let shared_input = TensorData::new_for_shape(&device, &[4, 4], DataTypeEnum::F32); + let pending_output = TensorData::new_for_shape(&device, &[4, 4], DataTypeEnum::F32); + let new_output = TensorData::new_for_shape(&device, &[4, 4], DataTypeEnum::F32); + + assert!(resolver.should_extend_kernel( + &[ + MirValue::Tensor(shared_input.clone()), + MirValue::Tensor(new_output.clone()), + ], + &MirValue::Tensor(new_output), + &[vec![ + MirValue::Tensor(shared_input), + MirValue::Tensor(pending_output.clone()), + ]], + &[MirValue::Tensor(pending_output)], + )); + } + + #[test] + fn test_should_extend_kernel_blocks_output_to_view_dependency() { + let device = Device::test_instance(); + let mut graph = ComputeGraphInner::new(&device); + let mut resolver = Resolver::new_many(&mut graph, &[], &device); + + let pending_input = TensorData::new_for_shape(&device, &[4, 4], DataTypeEnum::F32); + let pending_output = TensorData::new_for_shape(&device, &[4, 4], DataTypeEnum::F32); + let aliased_view = pending_output.slice(&[1..3, 0..4]); + let new_output = TensorData::new_for_shape(&device, &[2, 4], DataTypeEnum::F32); + + assert!(!resolver.should_extend_kernel( + &[ + MirValue::Tensor(aliased_view), + MirValue::Tensor(new_output.clone()) + ], + &MirValue::Tensor(new_output), + &[vec![ + MirValue::Tensor(pending_input), + MirValue::Tensor(pending_output.clone()), + ]], + &[MirValue::Tensor(pending_output)], + )); + } + + #[tokio::test] + async fn test_nary_fusion_caps_large_input_accumulator() { + let device = Device::test_instance(); + let tensors = (0..(MAX_FUSED_NARY_INPUTS + 4)) + .map(|_| Tensor::splat(&device, 1.0f32, [4])) + .collect::>(); + let mut sum = tensors[0].clone(); + for tensor in &tensors[1..] { + sum = &sum + tensor; + } + + let kernels = sum.count_kernels_to_resolve(); + assert!( + kernels > 1, + "large n-ary accumulators should be split before they become oversized kernels" + ); + + let result = sum.as_slice().await.unwrap(); + for index in 0..4 { + assert_eq!(result[[index]], tensors.len() as f32); + } + } + + #[tokio::test] + async fn test_many_independent_slice_assigns_respect_storage_binding_limit() { + let device = Device::test_instance(); + let mut outputs = Vec::new(); + for index in 0..12 { + let input = Tensor::splat(&device, index as f32, [1, 1]); + let value = Tensor::splat(&device, (100 + index) as f32, [1, 1]); + outputs.push(input.slice_assign([0..1, 0..1], &value)); + } + + let output_refs = outputs.iter().collect::>(); + let resolved = Tensor::materialized_many(&output_refs).await; + + for (index, output) in resolved.iter().enumerate() { + let result = output.as_slice().await.unwrap(); + assert_eq!(result[[0, 0]], (100 + index) as f32); + } + } +} diff --git a/fusor-ml/core/src/device.rs b/fusor-ml/core/src/device.rs index aaa89601c..b146d3a7b 100644 --- a/fusor-ml/core/src/device.rs +++ b/fusor-ml/core/src/device.rs @@ -211,12 +211,16 @@ impl Device { &self, source: impl Into>, ) -> wgpu::ShaderModule { + let source: Cow<'a, str> = source.into(); + if std::env::var("FUSOR_DUMP_SHADERS").ok().as_deref() == Some("1") { + eprintln!("=== fusor shader ===\n{source}\n=== end shader ==="); + } // SAFTEY: All kernels don't access memory outside of bounds and don't have unbounded loops unsafe { self.inner.device.create_shader_module_trusted( wgpu::ShaderModuleDescriptor { label: Some("Fusor ML Shader Module"), - source: wgpu::ShaderSource::Wgsl(source.into()), + source: wgpu::ShaderSource::Wgsl(source), }, wgpu::ShaderRuntimeChecks::unchecked(), ) @@ -333,9 +337,41 @@ impl Device { usage: wgpu::BufferUsages, to_initilize: bool, ) -> Arc { + // wgpu's copy / write APIs require the buffer size to respect + // `COPY_BUFFER_ALIGNMENT` (4 bytes) and reject zero-sized buffers. + // Round up at allocation time so small tensors (e.g. scalar F16 params + // or empty-iter uniforms) don't trip validation later. + let size = size + .max(wgpu::COPY_BUFFER_ALIGNMENT) + .next_multiple_of(wgpu::COPY_BUFFER_ALIGNMENT); + let disable_cache = + std::env::var("FUSOR_DISABLE_BUFFER_CACHE").ok().as_deref() == Some("1"); + let log_allocations = std::env::var("FUSOR_LOG_BUFFER_ALLOCATIONS") + .ok() + .as_deref() + == Some("1"); + + if disable_cache { + if log_allocations { + eprintln!( + "fusor buffer alloc uncached: {} bytes usage={usage:?}", + size + ); + } + return Arc::new(self.wgpu_device().create_buffer(&wgpu::BufferDescriptor { + label: Some("Tensor Buffer"), + size, + usage, + mapped_at_creation: false, + })); + } + // Try to get a buffer from the cache first self.get_cached_buffer(size, usage, to_initilize) .unwrap_or_else(|| { + if log_allocations { + eprintln!("fusor buffer alloc cached: {} bytes usage={usage:?}", size); + } let new_buffer = self.wgpu_device().create_buffer(&wgpu::BufferDescriptor { label: Some("Tensor Buffer"), size, @@ -361,7 +397,18 @@ impl Device { /// Get or create a buffer of the specified size. pub fn create_buffer_init(&self, data: &[u8], usage: wgpu::BufferUsages) -> Arc { let buffer = self.create_buffer_inner(data.len() as u64, usage, true); - self.wgpu_queue().write_buffer(&buffer, 0, data); + let align = wgpu::COPY_BUFFER_ALIGNMENT as usize; + let padded_len = data.len().next_multiple_of(align); + if padded_len == data.len() { + self.wgpu_queue().write_buffer(&buffer, 0, data); + } else { + // `write_buffer` copy size must respect `COPY_BUFFER_ALIGNMENT`; + // the underlying buffer is already padded to that size above. + let mut padded = Vec::with_capacity(padded_len); + padded.extend_from_slice(data); + padded.resize(padded_len, 0); + self.wgpu_queue().write_buffer(&buffer, 0, &padded); + } buffer } diff --git a/fusor-ml/core/src/map_layout.rs b/fusor-ml/core/src/map_layout.rs index ec75e8a50..b320e8833 100644 --- a/fusor-ml/core/src/map_layout.rs +++ b/fusor-ml/core/src/map_layout.rs @@ -33,12 +33,16 @@ impl MapLayoutOperation { } pub fn map_tensor(&self, tensor: &TensorData) -> TensorData { - TensorData::new_from_parts( + let mut out = TensorData::new_from_parts( tensor.device(), tensor.buffer().clone(), self.map_layout(tensor.layout()), tensor.datatype(), - ) + ); + if tensor.is_zero_splat() { + out.mark_zero_splat(); + } + out } pub fn map_layout(&self, layout: &Layout) -> Layout { diff --git a/fusor-ml/core/src/matmul/mod.rs b/fusor-ml/core/src/matmul/mod.rs index 638d619f2..8d450e9fb 100644 --- a/fusor-ml/core/src/matmul/mod.rs +++ b/fusor-ml/core/src/matmul/mod.rs @@ -319,6 +319,87 @@ async fn test_matrix_vector_mul_non_contiguous() { assert_eq!(as_slice[[1, 0]], 122.); } +#[cfg(test)] +#[tokio::test] +async fn test_large_skinny_k_matmul_matches_cpu_reference() { + let device = Device::test_instance(); + let m = 100; + let k = 3200; + let n = 320; + + let lhs_data: Vec = (0..m * k) + .map(|i| ((i % 97) as f32 - 48.0) * 0.001) + .collect(); + let rhs_data: Vec = (0..k * n) + .map(|i| ((i % 89) as f32 - 44.0) * 0.0015) + .collect(); + let lhs_nested: Vec> = lhs_data.chunks_exact(k).map(|row| row.to_vec()).collect(); + let rhs_nested: Vec> = rhs_data.chunks_exact(n).map(|row| row.to_vec()).collect(); + let lhs = Tensor::new(&device, &lhs_nested); + let rhs = Tensor::new(&device, &rhs_nested); + let result = lhs.mat_mul(&rhs).as_slice().await.unwrap(); + + let check_positions = [ + (0, 0), + (0, n - 1), + (m / 2, n / 2), + (m - 1, 0), + (m - 1, n - 1), + ]; + + for (row, col) in check_positions { + let expected = (0..k) + .map(|idx| lhs_data[row * k + idx] * rhs_data[idx * n + col]) + .sum::(); + let actual = result[[row, col]]; + assert!( + (actual - expected).abs() < 1e-3, + "Mismatch at [{row}, {col}]: actual={actual}, expected={expected}" + ); + } +} + +#[cfg(test)] +#[tokio::test] +async fn test_moonshine_attention_projection_matmul_matches_cpu_reference() { + let device = Device::test_instance(); + let batch = 1; + let m = 100; + let k = 320; + let n = 320; + + let lhs_data: Vec = (0..batch * m * k) + .map(|i| ((i % 97) as f32 - 48.0) * 0.002) + .collect(); + let rhs_data: Vec = (0..batch * k * n) + .map(|i| ((i % 89) as f32 - 44.0) * 0.0015) + .collect(); + let lhs_nested: Vec>> = lhs_data + .chunks_exact(m * k) + .map(|batch| batch.chunks_exact(k).map(|row| row.to_vec()).collect()) + .collect(); + let rhs_nested: Vec>> = rhs_data + .chunks_exact(k * n) + .map(|batch| batch.chunks_exact(n).map(|row| row.to_vec()).collect()) + .collect(); + let lhs = Tensor::new(&device, &lhs_nested); + let rhs = Tensor::new(&device, &rhs_nested); + let result = lhs.mat_mul(&rhs).as_slice().await.unwrap(); + + for row in [0, m / 2, m - 1] { + for col in [0, n / 2, n - 1] { + let expected = (0..k) + .map(|idx| lhs_data[row * k + idx] * rhs_data[idx * n + col]) + .sum::(); + let actual = result[[0, row, col]]; + assert!( + (actual - expected).abs() < 1e-3, + "Mismatch at [0, {row}, {col}]: actual={actual}, expected={expected}" + ); + } + } +} + #[cfg(test)] #[tokio::test] async fn test_multi_row_matrix_vector_mul() { @@ -495,6 +576,129 @@ async fn test_transposed_matmul() { } } +#[cfg(test)] +#[tokio::test] +async fn test_attention_chain_single_resolution_matches_staged() { + let device = Device::test_instance(); + let batch = 1; + let seq_len = 20; + let heads = 8; + let head_dim = 40; + let hidden = heads * head_dim; + + fn tensor3( + device: &Device, + batch: usize, + seq_len: usize, + hidden: usize, + seed: usize, + ) -> Tensor<3, f32> { + let data = (0..batch) + .map(|batch_idx| { + (0..seq_len) + .map(|seq_idx| { + (0..hidden) + .map(|hidden_idx| { + let value = + (batch_idx * 131 + seq_idx * 17 + hidden_idx * 7 + seed) % 97; + (value as f32 - 48.0) * 0.01 + }) + .collect::>() + }) + .collect::>() + }) + .collect::>(); + Tensor::new(device, &data) + } + + fn tensor4_mask(device: &Device, batch: usize, heads: usize, seq_len: usize) -> Tensor<4, f32> { + let data = (0..batch) + .map(|_| { + (0..heads) + .map(|_| { + (0..seq_len) + .map(|q_idx| { + (0..seq_len) + .map(|k_idx| if k_idx <= q_idx { 0.0 } else { -1.0e9 }) + .collect::>() + }) + .collect::>() + }) + .collect::>() + }) + .collect::>(); + Tensor::new(device, &data) + } + + fn attention( + q: &Tensor<3, f32>, + k: &Tensor<3, f32>, + v: &Tensor<3, f32>, + mask: &Tensor<4, f32>, + heads: usize, + head_dim: usize, + ) -> Tensor<3, f32> { + let [batch, seq_len, _] = *q.shape(); + let q = q.reshape([batch, seq_len, heads, head_dim]).transpose(1, 2); + let k = k.reshape([batch, seq_len, heads, head_dim]).transpose(1, 2); + let v = v.reshape([batch, seq_len, heads, head_dim]).transpose(1, 2); + let k_t = k.transpose(2, 3); + let scores = q.mat_mul(&k_t) * (head_dim as f32).powf(-0.5); + let scores = scores.add_(mask); + let weights = scores.softmax_last_dim::<3>(); + weights + .mat_mul(&v) + .transpose(1, 2) + .reshape([batch, seq_len, heads * head_dim]) + } + + let q = tensor3(&device, batch, seq_len, hidden, 3); + let k = tensor3(&device, batch, seq_len, hidden, 11); + let v = tensor3(&device, batch, seq_len, hidden, 19); + let mask = tensor4_mask(&device, batch, heads, seq_len); + + let single_pass = attention(&q, &k, &v, &mask, heads, head_dim); + let single_pass = single_pass.as_slice().await.unwrap(); + + let q_heads = q.reshape([batch, seq_len, heads, head_dim]).transpose(1, 2); + drop(q_heads.as_slice().await.unwrap()); + let k_heads = k.reshape([batch, seq_len, heads, head_dim]).transpose(1, 2); + drop(k_heads.as_slice().await.unwrap()); + let v_heads = v.reshape([batch, seq_len, heads, head_dim]).transpose(1, 2); + drop(v_heads.as_slice().await.unwrap()); + let k_t = k_heads.transpose(2, 3); + drop(k_t.as_slice().await.unwrap()); + let scores = q_heads.mat_mul(&k_t) * (head_dim as f32).powf(-0.5); + drop(scores.as_slice().await.unwrap()); + let scores = scores.add_(&mask); + drop(scores.as_slice().await.unwrap()); + let weights = scores.softmax_last_dim::<3>(); + drop(weights.as_slice().await.unwrap()); + let staged = weights + .mat_mul(&v_heads) + .transpose(1, 2) + .reshape([batch, seq_len, hidden]); + let staged = staged.as_slice().await.unwrap(); + + assert_eq!(single_pass.shape(), staged.shape()); + let mut max_diff = 0.0f32; + for batch_idx in 0..batch { + for seq_idx in 0..seq_len { + for hidden_idx in 0..hidden { + max_diff = max_diff.max( + (single_pass[[batch_idx, seq_idx, hidden_idx]] + - staged[[batch_idx, seq_idx, hidden_idx]]) + .abs(), + ); + } + } + } + assert!( + max_diff < 1e-4, + "single-pass attention chain diverged from staged execution: max_diff={max_diff}" + ); +} + #[cfg(test)] #[tokio::test] async fn test_batched_matmul() { diff --git a/fusor-ml/core/src/matmul/sgemv.rs b/fusor-ml/core/src/matmul/sgemv.rs index a7d20ff8f..0a29200b2 100644 --- a/fusor-ml/core/src/matmul/sgemv.rs +++ b/fusor-ml/core/src/matmul/sgemv.rs @@ -13,6 +13,10 @@ use crate::{ }, }; +fn matmul_sgemv_subgroups_supported(device: &crate::Device) -> bool { + device.subgroups_supported() && device.wgpu_adapter().get_info().backend != wgpu::Backend::Metal +} + #[allow(clippy::too_many_arguments)] pub(crate) fn sgemv( op: &MatMulOperation, @@ -165,7 +169,7 @@ pub(crate) fn sgemv( writeln!(kernel, "}}").unwrap(); // If subgroups are supported, perform a reduction with subgroup operations - if device.subgroups_supported() { + if matmul_sgemv_subgroups_supported(&device) { // Get the sum among all threads in the subgroup writeln!( kernel, @@ -332,19 +336,23 @@ pub(crate) fn workgroup_shape_constraints( device.limits().max_compute_workgroup_size_x + 1, ), ); - constraints.add_constraint( - 0, - crate::mir::workgroup_shape::Constraint::more_than_or_equals(device.min_subgroup_size()), - ); - constraints.add_constraint( - 0, - crate::mir::workgroup_shape::Constraint::less_than_or_equals( - device.max_subgroup_size() - * params - .subgroups_per_workgroup - .min(device.max_subgroup_size()), - ), - ); + if matmul_sgemv_subgroups_supported(device) { + constraints.add_constraint( + 0, + crate::mir::workgroup_shape::Constraint::more_than_or_equals( + device.min_subgroup_size(), + ), + ); + constraints.add_constraint( + 0, + crate::mir::workgroup_shape::Constraint::less_than_or_equals( + device.max_subgroup_size() + * params + .subgroups_per_workgroup + .min(device.max_subgroup_size()), + ), + ); + } constraints.add_constraint(1, crate::mir::workgroup_shape::Constraint::Equals(1)); constraints.add_constraint(2, crate::mir::workgroup_shape::Constraint::Equals(1)); constraints diff --git a/fusor-ml/core/src/mir/inputs/mod.rs b/fusor-ml/core/src/mir/inputs/mod.rs index 5affeffb1..b2d543ca2 100644 --- a/fusor-ml/core/src/mir/inputs/mod.rs +++ b/fusor-ml/core/src/mir/inputs/mod.rs @@ -107,8 +107,15 @@ impl Display for KernelInput { KernelInputType::QInfo(matrix) => { let info_index = matrix.info_binding; writeln!(f, "struct Tensor{info_index}Info {{")?; - for i in 0..matrix.rank { - writeln!(f, " shape_{i}: u32,")?; + if matrix.rank == 0 { + // WGSL structs must have at least one member; the field is + // unused for scalar tensors but keeps the type definition + // valid. + writeln!(f, " _unused: u32,")?; + } else { + for i in 0..matrix.rank { + writeln!(f, " shape_{i}: u32,")?; + } } writeln!(f, "}};")?; diff --git a/fusor-ml/core/src/mir/inputs/qmatrix.rs b/fusor-ml/core/src/mir/inputs/qmatrix.rs index 506427f88..bd7344004 100644 --- a/fusor-ml/core/src/mir/inputs/qmatrix.rs +++ b/fusor-ml/core/src/mir/inputs/qmatrix.rs @@ -28,6 +28,10 @@ impl QMatrixInput { write: &mut impl std::fmt::Write, indexes: impl IntoIterator, ) { + if self.rank == 0 { + write!(write, "0u").unwrap(); + return; + } let mut strides = Vec::new(); let mut product = "1".to_string(); for i in (0..self.rank).rev() { diff --git a/fusor-ml/core/src/mir/kernel.rs b/fusor-ml/core/src/mir/kernel.rs index 63fa434b2..b2de79524 100644 --- a/fusor-ml/core/src/mir/kernel.rs +++ b/fusor-ml/core/src/mir/kernel.rs @@ -118,6 +118,29 @@ impl GenericKernel { }; if binding >= self.inputs.len() as u32 { self.inputs.push(or(binding)); + } else { + // Binding is already claimed by another op in this fused kernel. + // Multiple ops can view the same underlying tensor through different + // "kernel ranks" (e.g. softmax treats a rank-4 tensor as rank-3 by + // factoring out the reduce axis). The info struct shared on the GPU + // must expose enough fields for the op with the highest rank so + // every op's shape_{i}/stride_{i} access is valid. + let new_input = or(binding); + let existing = &mut self.inputs[binding as usize]; + match (&mut existing.ty, &new_input.ty) { + ( + KernelInputType::TensorInfo(existing_info), + KernelInputType::TensorInfo(new_info), + ) if new_info.rank > existing_info.rank => { + existing_info.rank = new_info.rank; + } + (KernelInputType::QInfo(existing_info), KernelInputType::QInfo(new_info)) + if new_info.rank > existing_info.rank => + { + existing_info.rank = new_info.rank; + } + _ => {} + } } binding } diff --git a/fusor-ml/core/src/quantized/matmul/mod.rs b/fusor-ml/core/src/quantized/matmul/mod.rs index ab2a2f627..6afc28754 100644 --- a/fusor-ml/core/src/quantized/matmul/mod.rs +++ b/fusor-ml/core/src/quantized/matmul/mod.rs @@ -1,9 +1,15 @@ use crate::{ - DataType, DataTypeEnum, Device, Tensor, TensorData, + DataType, DataTypeEnum, Device, ElementWiseFunctions, Tensor, TensorData, compute_graph::NodeIndex, - mir::{inputs::MirValue, kernel::GenericKernel, operation::Operation}, + mir::{ + function::Function, + inputs::{MirValue, TensorInput}, + kernel::GenericKernel, + operation::Operation, + }, }; use fusor_gguf::GgmlType; +use std::fmt::Write; use super::QMatrix; @@ -16,9 +22,13 @@ pub use sgemm::{ChunkedSgemmConfig, GeneralSgemmConfig}; pub(crate) struct QMatMulOperation { pub(crate) input_datatype: DataTypeEnum, pub(crate) input: NodeIndex, + pub(crate) bias: Option, + pub(crate) bias_datatype: Option, pub(crate) matrix: QMatrix, pub(crate) in_shape: Box<[usize]>, pub(crate) out_shape: Box<[usize]>, + pub(crate) pre_element_wise: ElementWiseFunctions, + pub(crate) post_element_wise: ElementWiseFunctions, pub(crate) chunked_config: Option, pub(crate) general_config: Option, } @@ -38,14 +48,51 @@ impl QMatMulOperation { QMatMulOperation { input_datatype, input, + bias: None, + bias_datatype: None, matrix, in_shape: input_shape.into(), out_shape, + pre_element_wise: ElementWiseFunctions::empty(input_datatype), + post_element_wise: ElementWiseFunctions::empty(input_datatype), chunked_config: None, general_config: None, } } + pub(crate) fn with_bias(mut self, bias: &Tensor<1, T>) -> Self { + assert_eq!( + bias.shape()[0], + self.out_shape[self.out_shape.len() - 1], + "q_mat_mul bias shape must match output features" + ); + self.bias = Some(bias.key()); + self.bias_datatype = Some(bias.datatype()); + self + } + + pub(crate) fn apply_bias_and_post( + &self, + bias: Option<&TensorInput>, + n_index: &str, + value_expr: String, + post_fns: &[Function], + ) -> String { + let value_expr = if let Some(bias) = bias { + let mut bias_expr = String::new(); + write!(&mut bias_expr, "{bias}[").unwrap(); + bias.strided_index(&mut bias_expr, [n_index.to_string()]); + write!(&mut bias_expr, "]").unwrap(); + let bias_datatype = self.bias_datatype.unwrap_or(self.input_datatype); + format!("({value_expr} + {}({bias_expr}))", bias_datatype) + } else { + value_expr + }; + post_fns + .iter() + .fold(value_expr, |acc, function| function.call(vec![acc])) + } + fn elements_per_block(&self) -> u32 { self.matrix.datatype.block_size() as u32 } @@ -53,14 +100,13 @@ impl QMatMulOperation { fn sgemv(&self) -> bool { let m_dim_idx = self.in_shape.len() - 2; let m = self.in_shape[m_dim_idx]; - // Use SGEMV for tall and skinny matrices (small M, any K) - // SGEMV is more efficient when M is small because: - // - Each workgroup processes one M value independently - // - Less workgroup synchronization overhead - // - Better cache utilization for the K dimension - // SGEMM becomes more efficient for larger M where it can use - // tile-based processing with 16x16 workgroups - m <= 32 + // Use SGEMV for tall and skinny matrices (small M, any K). + // Decoder cross-attention cache init in encoder-decoder ASR models + // frequently lands in the 8..16-token range after audio subsampling. + // Routing those tiny widths through the SGEMM path has proven unstable + // on the current GPU backend, while SGEMV remains both stable and fast + // enough for these shapes. + m <= 16 } fn m_size(&self) -> u32 { @@ -71,39 +117,46 @@ impl QMatMulOperation { fn n_size(&self) -> u32 { self.matrix.shape[0] as u32 } + + pub(crate) fn matmul_datatype(&self) -> DataTypeEnum { + self.pre_element_wise.out_datatype() + } } impl Tensor { pub fn q_mat_mul(&self, other: &QMatrix) -> Self { + let in_shape = self.shape(); + // For F16/F32 matrices, dequantize and use regular mat_mul // because they don't have block structure like quantized types if matches!(other.datatype(), GgmlType::F16 | GgmlType::F32) { let dequantized: Tensor<2, T> = other.dequantize(); - // Broadcast the 2D matrix to match input tensor rank by unsqueezing batch dimensions - let mut broadcast_shape = Vec::from(self.shape().as_slice()); - broadcast_shape[R - 1] = other.shape()[0]; // Output dimension is first dim of weight - - // The weight matrix is [out_features, in_features], need to transpose for mat_mul - // self: [..., M, K] @ weight.T: [K, N] -> [..., M, N] - // Reshape weight to add batch dimensions: [1, 1, ..., K, N] + // Flatten all leading dimensions into a single rows dimension so we can use + // a plain 2D matmul instead of a broadcasted batched matmul. The broadcasted + // GPU path is particularly costly for encoder FFNs with tiny M and large N. + let rows = in_shape[..R - 1].iter().product::(); + let k = in_shape[R - 1]; + let n = other.shape()[0]; + + let input_2d: Tensor<2, T> = self.reshape([rows, k]); let weight_t = dequantized.transpose(0, 1); + let output_2d = input_2d.mat_mul(&weight_t); - // Create batch dimensions for the weight - let weight_shape: [usize; R] = std::array::from_fn(|i| { - if i < R - 2 { - 1 // Broadcast batch dimensions - } else if i == R - 2 { - other.shape()[1] // K dimension - } else { - other.shape()[0] // N dimension - } - }); - let weight_broadcast: Tensor = weight_t.reshape(weight_shape); - - return self.mat_mul(&weight_broadcast); + let out_shape: [usize; R] = + std::array::from_fn(|i| if i == R - 1 { n } else { in_shape[i] }); + return output_2d.reshape(out_shape); } + self.add_q_mat_mul(other) } + + pub fn q_mat_mul_bias(&self, other: &QMatrix, bias: &Tensor<1, T>) -> Self { + assert!( + !matches!(other.datatype(), GgmlType::F16 | GgmlType::F32), + "q_mat_mul_bias only supports quantized matrices" + ); + self.add_q_mat_mul_bias(other, bias) + } } #[cfg(test)] @@ -816,14 +869,26 @@ impl Operation for QMatMulOperation { fn visit_dependencies(&self, f: &mut dyn FnMut(NodeIndex)) { f(self.input); + if let Some(bias) = self.bias { + f(bias); + } } fn inputs(&self, nodes: &crate::compute_graph::ComputeGraphInner) -> Vec { let input = nodes.get_result(self.input).unwrap(); let q_matrix = self.matrix.clone(); let device = input.device(); - let output_tensor = TensorData::new_for_shape(device, &self.out_shape, input.datatype()); - vec![input.into(), q_matrix.into(), output_tensor.into()] + let output_tensor = TensorData::new_for_shape( + device, + &self.out_shape, + self.post_element_wise.out_datatype(), + ); + let mut inputs = vec![input.into(), q_matrix.into()]; + if let Some(bias) = self.bias { + inputs.push(nodes.get_result(bias).unwrap().into()); + } + inputs.push(output_tensor.into()); + inputs } // Related files/PRs in llama.cpp for reference: @@ -844,7 +909,15 @@ impl Operation for QMatMulOperation { let input_a = generic_kernel.add_tensor_input(rank, false, datatype); let input_b = generic_kernel.add_q_matrix_input(matrix_rank, self.matrix.datatype); - let output = generic_kernel.add_tensor_input(rank, true, datatype); + let bias = self.bias.map(|_| { + generic_kernel.add_tensor_input( + 1, + false, + self.bias_datatype.unwrap_or(self.input_datatype), + ) + }); + let output = + generic_kernel.add_tensor_input(rank, true, self.post_element_wise.out_datatype()); // For batched operations, we need to get the correct dimension indices let k_size = input_a.shape_binding(rank - 1).to_string(); // Last dimension is K @@ -864,6 +937,7 @@ impl Operation for QMatMulOperation { workgroup_shape, &input_a, &input_b, + bias.as_ref(), &output, &n_size, &m_size, @@ -873,7 +947,7 @@ impl Operation for QMatMulOperation { } fn output(&self, _: &crate::compute_graph::ComputeGraphInner, inputs: &[MirValue]) -> MirValue { - let output_tensor = inputs[2].as_tensor().unwrap(); + let output_tensor = inputs.last().unwrap().as_tensor().unwrap(); output_tensor.clone().into() } @@ -1021,3 +1095,512 @@ async fn test_fuzz_q_mat_mul_f16_sgemv() { println!("f16 sgemv test passed for shape {:?}", fusor_shape); } + +#[cfg(test)] +fn q8_matrix_from_rows(device: &Device, rows: &[Vec]) -> QMatrix { + assert!(!rows.is_empty()); + let cols = rows[0].len(); + assert_eq!(cols % 32, 0, "Q8_0 rows must be a multiple of 32 elements"); + for row in rows { + assert_eq!(row.len(), cols, "all rows must have the same width"); + } + + let block_size_bytes = 34; + let blocks_per_row = cols / 32; + let mut raw_bytes = vec![0u8; rows.len() * blocks_per_row * block_size_bytes]; + let scale = half::f16::from_f32(1.0).to_le_bytes(); + + for (row_idx, row) in rows.iter().enumerate() { + for block_idx in 0..blocks_per_row { + let offset = (row_idx * blocks_per_row + block_idx) * block_size_bytes; + raw_bytes[offset..offset + 2].copy_from_slice(&scale); + for element_idx in 0..32 { + raw_bytes[offset + 2 + element_idx] = row[block_idx * 32 + element_idx] as u8; + } + } + } + + QMatrix::from_parts( + device, + &raw_bytes, + vec![rows.len(), cols].into_boxed_slice(), + GgmlType::Q8_0, + ) + .unwrap() +} + +#[cfg(test)] +fn q8_test_matrix_with_rows(device: &Device, row_count: usize) -> QMatrix { + let rows: Vec> = (0..row_count) + .map(|row| { + (0..32) + .map(|value| ((row as i32 * 3 + value as i32) % 23 - 11) as i8) + .collect() + }) + .collect(); + q8_matrix_from_rows(device, &rows) +} + +#[cfg(test)] +async fn assert_close_3d(actual: &Tensor<3, f32>, expected: &Tensor<3, f32>) { + let actual = actual.as_slice().await.unwrap(); + let expected = expected.as_slice().await.unwrap(); + + assert_eq!(actual.shape(), expected.shape()); + for batch in 0..actual.shape()[0] { + for row in 0..actual.shape()[1] { + for col in 0..actual.shape()[2] { + let actual = actual[[batch, row, col]]; + let expected = expected[[batch, row, col]]; + assert!( + (actual - expected).abs() <= 1e-4, + "mismatch at [{batch}, {row}, {col}]: expected {expected}, got {actual}" + ); + } + } + } +} + +#[cfg(test)] +fn q8_test_matrix(device: &Device) -> QMatrix { + q8_matrix_from_rows( + device, + &[ + (1..=32).map(|value| value as i8).collect(), + (0..32).map(|value| (31 - value) as i8).collect(), + vec![1; 32], + vec![2; 32], + ], + ) +} + +#[cfg(test)] +fn q8_test_input_with_rows(device: &Device, rows: usize) -> Tensor<3, f32> { + let input_data: Vec>> = vec![{ + (0..rows) + .map(|row| { + (0..32) + .map(|value| (row as f32 + 1.0) * (value as f32 - 7.5) * 0.125) + .collect() + }) + .collect() + }]; + Tensor::<3, f32>::new(device, &input_data) +} + +#[cfg(test)] +fn q8_test_input(device: &Device) -> Tensor<3, f32> { + q8_test_input_with_rows(device, 17) +} + +#[cfg(test)] +fn q4_0_matrix_from_nibbles(device: &Device, rows: usize, cols: usize) -> (QMatrix, Vec) { + assert_eq!(cols % 32, 0, "Q4_0 rows must be a multiple of 32 elements"); + + let block_size_bytes = 18; + let blocks_per_row = cols / 32; + let mut raw_bytes = vec![0u8; rows * blocks_per_row * block_size_bytes]; + + for row in 0..rows { + for block in 0..blocks_per_row { + let offset = (row * blocks_per_row + block) * block_size_bytes; + let scale = half::f16::from_f32(0.015625 * (1 + (row + block) % 7) as f32); + raw_bytes[offset..offset + 2].copy_from_slice(&scale.to_le_bytes()); + for byte in 0..16 { + let low = ((row * 3 + block * 5 + byte) % 16) as u8; + let high = ((row * 7 + block * 11 + byte * 3 + 1) % 16) as u8; + raw_bytes[offset + 2 + byte] = low | (high << 4); + } + } + } + + let matrix = QMatrix::from_parts( + device, + &raw_bytes, + vec![rows, cols].into_boxed_slice(), + GgmlType::Q4_0, + ) + .unwrap(); + + (matrix, raw_bytes) +} + +#[cfg(test)] +fn q4_0_integer_reference( + input: &[Vec>], + raw_bytes: &[u8], + n: usize, +) -> Vec>> { + use fusor_gguf::{BlockQ4_0, BlockQ8_0, GgufBlock}; + + let batch = input.len(); + let m = input[0].len(); + let k = input[0][0].len(); + let blocks_per_row = k / BlockQ4_0::BLOCK_SIZE; + let blocks: &[BlockQ4_0] = bytemuck::cast_slice(raw_bytes); + let mut output = vec![vec![vec![0.0f32; n]; m]; batch]; + + for batch_idx in 0..batch { + for row in 0..m { + for col in 0..n { + let mut acc = 0.0f32; + for block_idx in 0..blocks_per_row { + let start = block_idx * BlockQ4_0::BLOCK_SIZE; + let activation: [f32; 32] = input[batch_idx][row] + [start..start + BlockQ4_0::BLOCK_SIZE] + .try_into() + .unwrap(); + let activation = BlockQ8_0::quantize(&activation); + let weight_block = &blocks[col * blocks_per_row + block_idx]; + acc += weight_block.vec_dot(&activation); + } + output[batch_idx][row][col] = acc; + } + } + } + + output +} + +#[cfg(test)] +fn q4_0_test_input(batch: usize, rows: usize, cols: usize) -> Vec>> { + (0..batch) + .map(|batch_idx| { + (0..rows) + .map(|row| { + (0..cols) + .map(|col| { + let value = ((batch_idx * 17 + row * 13 + col * 7) % 101) as f32; + (value - 50.0) * 0.017 + }) + .collect() + }) + .collect() + }) + .collect() +} + +#[cfg(test)] +fn q4k_matrix_from_bytes(device: &Device, rows: usize, cols: usize) -> (QMatrix, Vec) { + use fusor_gguf::BlockQ4K; + + assert_eq!( + cols % BlockQ4K::BLOCK_SIZE, + 0, + "Q4_K rows must be a multiple of 256 elements" + ); + + let block_size_bytes = std::mem::size_of::(); + let blocks_per_row = cols / BlockQ4K::BLOCK_SIZE; + let mut raw_bytes = vec![0u8; rows * blocks_per_row * block_size_bytes]; + + for row in 0..rows { + for block in 0..blocks_per_row { + let offset = (row * blocks_per_row + block) * block_size_bytes; + let scale = half::f16::from_f32(0.0025 * (1 + (row + block) % 9) as f32); + let min = half::f16::from_f32(0.0015 * (1 + (row * 3 + block) % 7) as f32); + raw_bytes[offset..offset + 2].copy_from_slice(&scale.to_le_bytes()); + raw_bytes[offset + 2..offset + 4].copy_from_slice(&min.to_le_bytes()); + for index in 0..12 { + raw_bytes[offset + 4 + index] = ((row * 19 + block * 23 + index * 29) % 251) as u8; + } + for index in 0..BlockQ4K::WEIGHTS_SIZE { + let low = ((row * 5 + block * 7 + index) % 16) as u8; + let high = ((row * 11 + block * 13 + index * 3 + 2) % 16) as u8; + raw_bytes[offset + 16 + index] = low | (high << 4); + } + } + } + + let matrix = QMatrix::from_parts( + device, + &raw_bytes, + vec![rows, cols].into_boxed_slice(), + GgmlType::Q4K, + ) + .unwrap(); + + (matrix, raw_bytes) +} + +#[cfg(test)] +fn q4k_integer_reference( + input: &[Vec>], + raw_bytes: &[u8], + n: usize, +) -> Vec>> { + use fusor_gguf::{BlockQ4K, BlockQ8K, GgufBlock}; + + let batch = input.len(); + let m = input[0].len(); + let k = input[0][0].len(); + let blocks_per_row = k / BlockQ4K::BLOCK_SIZE; + let blocks: &[BlockQ4K] = bytemuck::cast_slice(raw_bytes); + let mut output = vec![vec![vec![0.0f32; n]; m]; batch]; + + for batch_idx in 0..batch { + for row in 0..m { + for col in 0..n { + let mut acc = 0.0f32; + for block_idx in 0..blocks_per_row { + let start = block_idx * BlockQ4K::BLOCK_SIZE; + let activation: [f32; 256] = input[batch_idx][row] + [start..start + BlockQ4K::BLOCK_SIZE] + .try_into() + .unwrap(); + let activation = BlockQ8K::quantize(&activation); + let weight_block = &blocks[col * blocks_per_row + block_idx]; + acc += weight_block.vec_dot(&activation); + } + output[batch_idx][row][col] = acc; + } + } + } + + output +} + +#[cfg(test)] +#[tokio::test] +async fn test_q4_0_sgemm_matches_integer_activation_reference() { + let device = Device::test_instance(); + let rows = 32; + let cols = 64; + let m = 24; + let (q_matrix, raw_bytes) = q4_0_matrix_from_nibbles(&device, rows, cols); + let input_data = q4_0_test_input(1, m, cols); + let input = Tensor::<3, f32>::new(&device, &input_data); + + let result = input.q_mat_mul(&q_matrix); + assert_eq!( + device.compute_graph().node_variant_name(result.key()), + "QMatMul" + ); + let result = result.as_slice().await.unwrap(); + let expected = q4_0_integer_reference(&input_data, &raw_bytes, rows); + + for row in 0..m { + for col in 0..rows { + let actual = result[[0, row, col]]; + let expected = expected[0][row][col]; + assert!( + (actual - expected).abs() <= 1e-3, + "mismatch at [0, {row}, {col}]: expected {expected}, got {actual}" + ); + } + } +} + +#[cfg(test)] +#[tokio::test] +async fn test_q4_0_q_mat_mul_bias_matches_unfused_add() { + let device = Device::test_instance(); + let rows = 32; + let cols = 64; + let m = 7; + let (q_matrix, _) = q4_0_matrix_from_nibbles(&device, rows, cols); + let input_data = q4_0_test_input(1, m, cols); + let input = Tensor::<3, f32>::new(&device, &input_data); + let bias_data = (0..rows) + .map(|index| (index as f32 - 11.0) * 0.03125) + .collect::>(); + let bias = Tensor::<1, f32>::new(&device, &bias_data); + + let fused = input.q_mat_mul_bias(&q_matrix, &bias); + assert_eq!( + device.compute_graph().node_variant_name(fused.key()), + "QMatMul" + ); + let bias = bias.reshape([1, 1, rows]).broadcast_as([1, m, rows]); + let expected = input.q_mat_mul(&q_matrix) + bias; + assert_close_3d(&fused, &expected).await; +} + +#[cfg(test)] +#[tokio::test] +async fn test_q4k_sgemm_matches_integer_activation_reference() { + let device = Device::test_instance(); + let rows = 16; + let cols = 256; + let m = 20; + let (q_matrix, raw_bytes) = q4k_matrix_from_bytes(&device, rows, cols); + let input_data = q4_0_test_input(1, m, cols); + let input = Tensor::<3, f32>::new(&device, &input_data); + + let result = input.q_mat_mul(&q_matrix); + assert_eq!( + device.compute_graph().node_variant_name(result.key()), + "QMatMul" + ); + let result = result.as_slice().await.unwrap(); + let expected = q4k_integer_reference(&input_data, &raw_bytes, rows); + + for row in 0..m { + for col in 0..rows { + let actual = result[[0, row, col]]; + let expected = expected[0][row][col]; + assert!( + (actual - expected).abs() <= 1e-3, + "mismatch at [0, {row}, {col}]: expected {expected}, got {actual}" + ); + } + } +} + +#[cfg(test)] +#[tokio::test] +async fn test_q4_0_projection_matmul_matches_materialized_barrier() { + let device = Device::test_instance(); + let rows = 32; + let cols = 64; + let m = 24; + let (q_matrix, _) = q4_0_matrix_from_nibbles(&device, rows, cols); + let (k_matrix, _) = q4_0_matrix_from_nibbles(&device, rows, cols); + let input_data = q4_0_test_input(1, m, cols); + let input = Tensor::<3, f32>::new(&device, &input_data); + + let q = input.clone().q_mat_mul(&q_matrix); + let k = input.clone().q_mat_mul(&k_matrix); + let direct = q.mat_mul(&k.transpose(1, 2)); + + let q_barrier: Tensor<3, f32> = input.clone().q_mat_mul(&q_matrix).materialized().await; + let k_barrier: Tensor<3, f32> = input.q_mat_mul(&k_matrix).materialized().await; + let expected = q_barrier.mat_mul(&k_barrier.transpose(1, 2)); + + let direct = direct.as_slice().await.unwrap(); + let expected = expected.as_slice().await.unwrap(); + assert_eq!(direct.shape(), expected.shape()); + for row in 0..m { + for col in 0..m { + let actual = direct[[0, row, col]]; + let expected = expected[[0, row, col]]; + assert!( + (actual - expected).abs() <= 1e-3, + "mismatch at [0, {row}, {col}]: expected {expected}, got {actual}" + ); + } + } +} + +#[cfg(test)] +#[test] +fn test_q8_0_specialized_sgemv_enabled_on_metal() { + let device = Device::test_instance(); + if device.wgpu_adapter().get_info().backend != wgpu::Backend::Metal { + return; + } + if !device.subgroups_supported() || device.max_subgroup_size() < 2 * device.min_subgroup_size() + { + return; + } + + let op = QMatMulOperation::new( + DataTypeEnum::F32, + &[1, 8, 32], + NodeIndex::new(0), + q8_test_matrix_with_rows(&device, 3), + ); + + assert_eq!(sgemv::selected_sgemv_kernel_kind(&op, &device), "q_8_0"); +} + +#[cfg(test)] +#[tokio::test] +async fn test_q_mat_mul_metal_tiny_m_stays_quantized_and_correct() { + let device = Device::test_instance(); + if device.wgpu_adapter().get_info().backend != wgpu::Backend::Metal { + return; + } + + let q_matrix = q8_test_matrix_with_rows(&device, 3); + let input = q8_test_input_with_rows(&device, 8); + + let result: Tensor<3, f32> = input.q_mat_mul(&q_matrix); + assert_eq!( + device.compute_graph().node_variant_name(result.key()), + "QMatMul" + ); + + let dequantized: Tensor<2, f32> = q_matrix.dequantize(); + let expected_2d: Tensor<2, f32> = input.reshape([8, 32]).mat_mul(&dequantized.transpose(0, 1)); + let expected = expected_2d.reshape([1, 8, 3]); + + assert_close_3d(&result, &expected).await; +} + +#[cfg(test)] +#[tokio::test] +async fn test_q_mat_mul_materialized_many_matches_individual_on_metal() { + let device = Device::test_instance(); + if device.wgpu_adapter().get_info().backend != wgpu::Backend::Metal { + return; + } + + let input = q8_test_input_with_rows(&device, 13); + let matrix_a = q8_test_matrix_with_rows(&device, 3); + let matrix_b = q8_matrix_from_rows( + &device, + &[ + vec![3; 32], + (0..32).map(|value| (value as i8 % 7) - 3).collect(), + (0..32).map(|value| 5 - (value as i8 % 11)).collect(), + ], + ); + + let combined_a: Tensor<3, f32> = input.clone().q_mat_mul(&matrix_a); + let combined_b: Tensor<3, f32> = input.clone().q_mat_mul(&matrix_b); + + let expected_a: Tensor<3, f32> = input.clone().q_mat_mul(&matrix_a).materialized().await; + let expected_b: Tensor<3, f32> = input.q_mat_mul(&matrix_b).materialized().await; + + let combined = Tensor::materialized_many(&[&combined_a, &combined_b]).await; + + assert_close_3d(&combined[0], &expected_a).await; + assert_close_3d(&combined[1], &expected_b).await; +} + +#[cfg(test)] +#[tokio::test] +async fn test_q_mat_mul_post_elementwise_fuses() { + let device = Device::test_instance(); + let q_matrix = q8_test_matrix(&device); + let input = q8_test_input(&device); + + let fused: Tensor<3, f32> = input.q_mat_mul(&q_matrix) + 1.25; + assert_eq!(fused.count_kernels_to_resolve(), 1); + + let materialized = input.q_mat_mul(&q_matrix).materialized().await; + let expected: Tensor<3, f32> = materialized + 1.25; + + assert_close_3d(&fused, &expected).await; +} + +#[cfg(test)] +#[tokio::test] +async fn test_q_mat_mul_pre_elementwise_fuses() { + let device = Device::test_instance(); + let q_matrix = q8_test_matrix(&device); + let input = q8_test_input(&device); + + let fused: Tensor<3, f32> = (input.clone() + 0.5).q_mat_mul(&q_matrix); + assert_eq!(fused.count_kernels_to_resolve(), 1); + + let materialized_input = (input + 0.5).materialized().await; + let expected: Tensor<3, f32> = materialized_input.q_mat_mul(&q_matrix); + + assert_close_3d(&fused, &expected).await; +} + +#[cfg(test)] +#[tokio::test] +async fn test_q_mat_mul_view_chain_matches_materialized_barrier() { + let device = Device::test_instance(); + let q_matrix = q8_test_matrix(&device); + let input = q8_test_input(&device); + + let direct: Tensor<3, f32> = input.q_mat_mul(&q_matrix).transpose(1, 2) + 1.25; + let materialized = input.q_mat_mul(&q_matrix).materialized().await; + let staged: Tensor<3, f32> = materialized.transpose(1, 2) + 1.25; + + assert_close_3d(&direct, &staged).await; +} diff --git a/fusor-ml/core/src/quantized/matmul/sgemm/chunked.rs b/fusor-ml/core/src/quantized/matmul/sgemm/chunked.rs index 1c09cdbf4..d1b54343a 100644 --- a/fusor-ml/core/src/quantized/matmul/sgemm/chunked.rs +++ b/fusor-ml/core/src/quantized/matmul/sgemm/chunked.rs @@ -8,6 +8,7 @@ use crate::{ quantized::matmul::QMatMulOperation, }; use std::fmt::Write; +use std::sync::OnceLock; /// Size of the chunk we will dot at a time const MATRIX_ELEMENTS: u32 = 16; @@ -62,6 +63,7 @@ pub fn chunked_sgemm_with_config( kernel: &mut GenericKernel, input_a: &TensorInput, input_b: &QMatrixInput, + bias: Option<&TensorInput>, output: &TensorInput, k_size: &str, config: ChunkedSgemmConfig, @@ -84,6 +86,7 @@ pub fn chunked_sgemm_with_config( let sgemm_m_results_per_thread = config.m_results_per_thread; let cache_datatype = config.cache_datatype; let shuffled_loads = config.shuffled_loads; + let pre_element_wise_functions = OnceLock::new(); let cache_a_size = (sgemm_input_k_elements / 4) * (sgemm_input_m_elements / 4); let cache_a = kernel.add_global_array( @@ -238,12 +241,19 @@ pub fn chunked_sgemm_with_config( kernel, indices.iter().cloned(), |kernel: &mut GenericKernel| { + let pre_element_wise_functions = pre_element_wise_functions + .get_or_init(|| op.pre_element_wise.add_functions(kernel)); + let mut raw_input = String::new(); + write!(&mut raw_input, "{input_a}[")?; + input_a.strided_index(&mut raw_input, indices.iter().cloned()); + write!(&mut raw_input, "]")?; + let processed = pre_element_wise_functions + .iter() + .fold(raw_input, |acc, f| f.call(vec![acc])); write!( kernel, - "{cache_a}[chunk_index][col_index][row_index] = {cache_datatype}({input_a}[" + "{cache_a}[chunk_index][col_index][row_index] = {cache_datatype}({processed});" )?; - input_a.strided_index(kernel, indices.iter().cloned()); - write!(kernel, "]);")?; std::fmt::Result::Ok(()) }, ) @@ -376,17 +386,20 @@ pub fn chunked_sgemm_with_config( } writeln!(kernel, "}}").unwrap(); - write_acc_back(kernel, output, output_offset, &config).unwrap(); + write_acc_back(kernel, op, bias, output, output_offset, &config).unwrap(); } fn write_acc_back( kernel: &mut GenericKernel, + op: &QMatMulOperation, + bias: Option<&TensorInput>, output: &TensorInput, output_offset: &str, config: &ChunkedSgemmConfig, ) -> std::fmt::Result { let sgemm_n_results_per_thread = config.n_results_per_thread; let sgemm_m_results_per_thread = config.m_results_per_thread; + let post_element_wise_functions = op.post_element_wise.add_functions(kernel); for tile_m in 0..sgemm_m_results_per_thread { for tile_n in 0..sgemm_n_results_per_thread { writeln!( @@ -414,10 +427,16 @@ fn write_acc_back( write!(kernel, "let output_index = ")?; output.strided_index(kernel, output_indices.iter().cloned()); writeln!(kernel, ";")?; - writeln!( - kernel, - "{output}[output_index] = acc[{tile_m}][{tile_n}][y_offset][x_offset];" - )?; + let result = op.apply_bias_and_post( + bias, + &format!("{output_offset}.y + {tile_n} * 4u + y_offset"), + format!( + "{}(acc[{tile_m}][{tile_n}][y_offset][x_offset])", + op.input_datatype + ), + &post_element_wise_functions, + ); + writeln!(kernel, "{output}[output_index] = {result};")?; std::fmt::Result::Ok(()) }, )?; diff --git a/fusor-ml/core/src/quantized/matmul/sgemm/general.rs b/fusor-ml/core/src/quantized/matmul/sgemm/general.rs index 929fa74e3..437f9cb01 100644 --- a/fusor-ml/core/src/quantized/matmul/sgemm/general.rs +++ b/fusor-ml/core/src/quantized/matmul/sgemm/general.rs @@ -1,12 +1,15 @@ use crate::{ DataTypeEnum, dequantize_vec4_block, mir::{ + globals::KernelGlobalSpace, inputs::{QMatrixInput, TensorInput}, kernel::GenericKernel, }, quantized::matmul::QMatMulOperation, + util::maybe_vec_storage_type, }; use std::fmt::Write; +use std::sync::OnceLock; /// Configuration for general SGEMM algorithm #[derive(Debug, Clone, Copy)] @@ -36,6 +39,7 @@ pub fn general_sgemm_with_config( kernel: &mut GenericKernel, input_a: &TensorInput, input_b: &QMatrixInput, + bias: Option<&TensorInput>, output: &TensorInput, n_size: &str, m_size: &str, @@ -44,7 +48,10 @@ pub fn general_sgemm_with_config( ) { let global_id = kernel.global_id(); let elements_per_block = op.elements_per_block(); - let dtype = op.input_datatype; + let input_datatype = op.input_datatype; + let dtype = op.matmul_datatype(); + let pre_element_wise_functions = OnceLock::new(); + let post_element_wise_functions = OnceLock::new(); // Validate configuration config.validate().unwrap(); @@ -64,7 +71,12 @@ pub fn general_sgemm_with_config( writeln!(kernel, "block_batch = block_batch / {shape};").unwrap(); } - writeln!(kernel, "var acc = 0.0;").unwrap(); + let acc_storage_type = maybe_vec_storage_type(sgemm_vector_size, DataTypeEnum::F32); + writeln!( + kernel, + "var acc: {acc_storage_type} = {acc_storage_type}();" + ) + .unwrap(); writeln!( kernel, @@ -85,8 +97,10 @@ pub fn general_sgemm_with_config( kernel, op.matrix.datatype, "chunk".to_string(), - DataTypeEnum::F32, + dtype, |index, data, code| { + let pre_element_wise_functions = + pre_element_wise_functions.get_or_init(|| op.pre_element_wise.add_functions(code)); writeln!(code, "{{",).unwrap(); writeln!( code, @@ -98,7 +112,8 @@ pub fn general_sgemm_with_config( if local > 0 { write!(code, ", ").unwrap(); } - write!(code, "{input_a}[").unwrap(); + let mut raw_input = String::new(); + write!(&mut raw_input, "{input_a}[").unwrap(); let mut indices = vec![]; // Add batch indices first for dim in (0..input_a.rank()).rev().skip(2) { @@ -107,12 +122,16 @@ pub fn general_sgemm_with_config( // Then add M and K indices indices.push("y".to_string()); indices.push(format!("a_index_local_offset + {local}")); - input_a.strided_index(code, indices); - write!(code, "]").unwrap(); + input_a.strided_index(&mut raw_input, indices); + write!(&mut raw_input, "]").unwrap(); + let processed = pre_element_wise_functions + .iter() + .fold(raw_input, |acc, f| f.call(vec![acc])); + write!(code, "{processed}").unwrap(); } writeln!(code, ");").unwrap(); - writeln!(code, "acc += dot(a_values, {data});").unwrap(); + writeln!(code, "acc += {acc_storage_type}(a_values * {data});").unwrap(); writeln!(code, "}}").unwrap(); }, ); @@ -136,6 +155,358 @@ pub fn general_sgemm_with_config( output_indices.push("x".to_string()); output.strided_index(kernel, output_indices); writeln!(kernel, ";").unwrap(); - writeln!(kernel, "{output}[output_index] = acc;").unwrap(); + let acc_sum = match sgemm_vector_size { + 1 => "acc".to_string(), + 2..=4 => (0..sgemm_vector_size) + .map(|index| format!("acc[{index}]")) + .collect::>() + .join(" + "), + _ => (0..sgemm_vector_size) + .map(|index| format!("acc[{index}]")) + .collect::>() + .join(" + "), + }; + let post_fns = + post_element_wise_functions.get_or_init(|| op.post_element_wise.add_functions(kernel)); + let result = + op.apply_bias_and_post(bias, "x", format!("{input_datatype}({acc_sum})"), post_fns); + writeln!(kernel, "{output}[output_index] = {result};").unwrap(); + writeln!(kernel, "}}").unwrap(); +} + +#[allow(clippy::too_many_arguments)] +pub fn general_q4_0_sgemm( + op: &QMatMulOperation, + kernel: &mut GenericKernel, + input_a: &TensorInput, + input_b: &QMatrixInput, + bias: Option<&TensorInput>, + output: &TensorInput, + n_size: &str, + m_size: &str, + k_size: &str, + use_f16: bool, +) { + let global_id = kernel.global_id(); + let elements_per_block = op.elements_per_block(); + let input_datatype = op.input_datatype; + let pre_element_wise_functions = OnceLock::new(); + let post_element_wise_functions = OnceLock::new(); + + writeln!(kernel, "let x = {global_id}.x;").unwrap(); + writeln!(kernel, "let y = {global_id}.y;").unwrap(); + writeln!(kernel, "var block_batch = {global_id}.z;").unwrap(); + + for dim in (0..input_a.rank()).rev().skip(2) { + let shape = input_a.shape_binding(dim); + writeln!(kernel, "let block_batch_{dim} = block_batch % {shape};").unwrap(); + writeln!(kernel, "block_batch = block_batch / {shape};").unwrap(); + } + + writeln!( + kernel, + "let k_block_size = ({k_size} + {elements_per_block} - 1u) / {elements_per_block};" + ) + .unwrap(); + writeln!(kernel, "if x < {n_size} && y < {m_size} {{").unwrap(); + writeln!(kernel, "var acc = f32(0.0);").unwrap(); + writeln!(kernel, "for (var k = 0u; k < k_block_size; k += 1u) {{").unwrap(); + writeln!(kernel, "let chunk = &{input_b}[k + x * k_block_size];").unwrap(); + writeln!(kernel, "let a_index_base = k * {elements_per_block};").unwrap(); + writeln!(kernel, "var max_abs = f32(0.0);").unwrap(); + + for local in 0..32u32 { + let pre_element_wise_functions = + pre_element_wise_functions.get_or_init(|| op.pre_element_wise.add_functions(kernel)); + let mut raw_input = String::new(); + write!(&mut raw_input, "{input_a}[").unwrap(); + let mut indices = Vec::new(); + for dim in (0..input_a.rank()).rev().skip(2) { + indices.push(format!("block_batch_{dim}")); + } + indices.push("y".to_string()); + indices.push(format!("a_index_base + {local}u")); + input_a.strided_index(&mut raw_input, indices); + write!(&mut raw_input, "]").unwrap(); + let processed = pre_element_wise_functions + .iter() + .fold(raw_input, |acc, f| f.call(vec![acc])); + writeln!(kernel, "let a_{local} = f32({processed});").unwrap(); + writeln!(kernel, "max_abs = max(max_abs, abs(a_{local}));").unwrap(); + } + + writeln!(kernel, "var inv_scale = f32(0.0);").unwrap(); + writeln!(kernel, "if max_abs != f32(0.0) {{").unwrap(); + writeln!(kernel, "inv_scale = f32(127.0) / max_abs;").unwrap(); writeln!(kernel, "}}").unwrap(); + writeln!( + kernel, + "let q8_scale_raw = select(f32(0.0), max_abs / f32(127.0), max_abs != f32(0.0));" + ) + .unwrap(); + if use_f16 { + writeln!(kernel, "let q8_scale = f32(f16(q8_scale_raw));").unwrap(); + } else { + writeln!(kernel, "let q8_scale = q8_scale_raw;").unwrap(); + } + writeln!(kernel, "var block_sum = i32(0);").unwrap(); + + for local in 0..32u32 { + let byte = if local < 16 { local } else { local - 16 }; + let word = byte / 4; + let shift = (byte % 4) * 8 + if local < 16 { 0 } else { 4 }; + writeln!( + kernel, + "let q4_{local} = i32((chunk.data[{word}] >> {shift}u) & 0xFu) - 8;" + ) + .unwrap(); + writeln!( + kernel, + "let q8_{local} = i32(clamp(round(a_{local} * inv_scale), -127.0, 127.0));" + ) + .unwrap(); + writeln!(kernel, "block_sum += q4_{local} * q8_{local};").unwrap(); + } + + writeln!( + kernel, + "acc += f32(block_sum) * f32(chunk.scale) * q8_scale;" + ) + .unwrap(); + writeln!(kernel, "}}").unwrap(); + + write!(kernel, "let output_index = ").unwrap(); + let mut output_indices = Vec::new(); + for dim in (0..output.rank()).rev().skip(2) { + output_indices.push(format!("block_batch_{dim}")); + } + output_indices.push("y".to_string()); + output_indices.push("x".to_string()); + output.strided_index(kernel, output_indices); + writeln!(kernel, ";").unwrap(); + let post_fns = + post_element_wise_functions.get_or_init(|| op.post_element_wise.add_functions(kernel)); + let result = op.apply_bias_and_post(bias, "x", format!("{input_datatype}(acc)"), post_fns); + writeln!(kernel, "{output}[output_index] = {result};").unwrap(); + writeln!(kernel, "}}").unwrap(); +} + +#[allow(clippy::too_many_arguments)] +pub fn general_q4k_sgemm( + op: &QMatMulOperation, + kernel: &mut GenericKernel, + input_a: &TensorInput, + input_b: &QMatrixInput, + bias: Option<&TensorInput>, + output: &TensorInput, + n_size: &str, + m_size: &str, + k_size: &str, + use_f16: bool, +) { + let elements_per_block = op.elements_per_block(); + let input_datatype = op.input_datatype; + let pre_element_wise_functions = OnceLock::new(); + let post_element_wise_functions = OnceLock::new(); + let q8_cache = kernel.add_global_array( + KernelGlobalSpace::Workgroup, + DataTypeEnum::F32, + elements_per_block.to_string(), + ); + let max_cache = kernel.add_global_array( + KernelGlobalSpace::Workgroup, + DataTypeEnum::F32, + "64".to_string(), + ); + let scale_offset_fn = kernel.add_function( + "vec2", + r#" +var output = vec2(0.0, 0.0); +if group < 4u { + let shift = group * 8u; + output = vec2( + f32((scale_bytes_0 >> shift) & 0x3Fu), + f32((scale_bytes_1 >> shift) & 0x3Fu) + ); +} else { + let shift = (group - 4u) * 8u; + let high_shift = shift + 6u; + output = vec2( + f32(((scale_bytes_2 >> shift) & 0x0Fu) | (((scale_bytes_0 >> high_shift) & 0x03u) << 4u)), + f32(((scale_bytes_2 >> (shift + 4u)) & 0x0Fu) | (((scale_bytes_1 >> high_shift) & 0x03u) << 4u)) + ); +} +"#, + [ + ("scale_bytes_0".to_string(), "u32".to_string()), + ("scale_bytes_1".to_string(), "u32".to_string()), + ("scale_bytes_2".to_string(), "u32".to_string()), + ("group".to_string(), "u32".to_string()), + ], + ); + + let workgroup_id = kernel.workgroup_index(); + let local_id = kernel.workgroup_local_index(); + writeln!(kernel, "let local_id = {local_id};").unwrap(); + writeln!(kernel, "let x = {workgroup_id}.x * 64u + local_id;").unwrap(); + writeln!(kernel, "let y = {workgroup_id}.y;").unwrap(); + writeln!(kernel, "var block_batch = {workgroup_id}.z;").unwrap(); + + for dim in (0..input_a.rank()).rev().skip(2) { + let shape = input_a.shape_binding(dim); + writeln!(kernel, "let block_batch_{dim} = block_batch % {shape};").unwrap(); + writeln!(kernel, "block_batch = block_batch / {shape};").unwrap(); + } + + writeln!( + kernel, + "let k_block_size = ({k_size} + {elements_per_block} - 1u) / {elements_per_block};" + ) + .unwrap(); + writeln!(kernel, "var acc = f32(0.0);").unwrap(); + writeln!(kernel, "for (var k = 0u; k < k_block_size; k += 1u) {{").unwrap(); + writeln!(kernel, "let a_index_base = k * {elements_per_block};").unwrap(); + writeln!(kernel, "var thread_max_abs = f32(0.0);").unwrap(); + writeln!( + kernel, + "for (var offset = 0u; offset < {elements_per_block}; offset += 64u) {{" + ) + .unwrap(); + let pre_element_wise_functions = + pre_element_wise_functions.get_or_init(|| op.pre_element_wise.add_functions(kernel)); + let a_value = q4k_input_value_expr(input_a, "offset + local_id", pre_element_wise_functions); + writeln!(kernel, "let a_value = f32({a_value});").unwrap(); + writeln!( + kernel, + "thread_max_abs = max(thread_max_abs, abs(a_value));" + ) + .unwrap(); + writeln!(kernel, "}}").unwrap(); + writeln!(kernel, "{max_cache}[local_id] = thread_max_abs;").unwrap(); + writeln!(kernel, "workgroupBarrier();").unwrap(); + writeln!(kernel, "if local_id == 0u {{").unwrap(); + writeln!(kernel, "var block_max_abs = f32(0.0);").unwrap(); + writeln!(kernel, "for (var index = 0u; index < 64u; index += 1u) {{").unwrap(); + writeln!( + kernel, + "block_max_abs = max(block_max_abs, {max_cache}[index]);" + ) + .unwrap(); + writeln!(kernel, "}}").unwrap(); + writeln!(kernel, "{max_cache}[0] = block_max_abs;").unwrap(); + writeln!(kernel, "}}").unwrap(); + writeln!(kernel, "workgroupBarrier();").unwrap(); + writeln!(kernel, "let max_abs = {max_cache}[0];").unwrap(); + writeln!(kernel, "var inv_scale = f32(0.0);").unwrap(); + writeln!(kernel, "if max_abs != f32(0.0) {{").unwrap(); + writeln!(kernel, "inv_scale = f32(127.0) / max_abs;").unwrap(); + writeln!(kernel, "}}").unwrap(); + writeln!( + kernel, + "let q8_scale_raw = select(f32(0.0), max_abs / f32(127.0), max_abs != f32(0.0));" + ) + .unwrap(); + if use_f16 { + writeln!(kernel, "let q8_scale = f32(f16(q8_scale_raw));").unwrap(); + } else { + writeln!(kernel, "let q8_scale = q8_scale_raw;").unwrap(); + } + writeln!( + kernel, + "for (var offset = 0u; offset < {elements_per_block}; offset += 64u) {{" + ) + .unwrap(); + let a_value = q4k_input_value_expr(input_a, "offset + local_id", pre_element_wise_functions); + writeln!(kernel, "let a_value = f32({a_value});").unwrap(); + writeln!( + kernel, + "{q8_cache}[offset + local_id] = f32(i32(clamp(round(a_value * inv_scale), -127.0, 127.0)));" + ) + .unwrap(); + writeln!(kernel, "}}").unwrap(); + writeln!(kernel, "workgroupBarrier();").unwrap(); + + writeln!(kernel, "if x < {n_size} && y < {m_size} {{").unwrap(); + writeln!(kernel, "let chunk = &{input_b}[k + x * k_block_size];").unwrap(); + writeln!(kernel, "let scale_bytes_0 = chunk.scales[0];").unwrap(); + writeln!(kernel, "let scale_bytes_1 = chunk.scales[1];").unwrap(); + writeln!(kernel, "let scale_bytes_2 = chunk.scales[2];").unwrap(); + writeln!(kernel, "var scale_dot_sum = f32(0.0);").unwrap(); + writeln!(kernel, "var offset_y_sum = f32(0.0);").unwrap(); + writeln!( + kernel, + "for (var local_k = 0u; local_k < {elements_per_block}; local_k += 1u) {{" + ) + .unwrap(); + writeln!(kernel, "let pair = local_k / 64u;").unwrap(); + writeln!(kernel, "let within_pair = local_k % 64u;").unwrap(); + writeln!(kernel, "let high = within_pair >= 32u;").unwrap(); + writeln!(kernel, "let data_byte = pair * 32u + (within_pair % 32u);").unwrap(); + writeln!(kernel, "let data_word = data_byte / 4u;").unwrap(); + writeln!( + kernel, + "let data_shift = (data_byte % 4u) * 8u + select(0u, 4u, high);" + ) + .unwrap(); + writeln!( + kernel, + "let q4 = f32((chunk.data[data_word] >> data_shift) & 0xFu);" + ) + .unwrap(); + writeln!(kernel, "let group = pair * 2u + select(0u, 1u, high);").unwrap(); + let scale_offset = scale_offset_fn.call(vec![ + "scale_bytes_0".to_string(), + "scale_bytes_1".to_string(), + "scale_bytes_2".to_string(), + "group".to_string(), + ]); + writeln!(kernel, "let scale_offset = {scale_offset};").unwrap(); + writeln!(kernel, "let q8_f = {q8_cache}[local_k];").unwrap(); + writeln!(kernel, "scale_dot_sum += scale_offset.x * q4 * q8_f;").unwrap(); + writeln!(kernel, "offset_y_sum += scale_offset.y * q8_f;").unwrap(); + writeln!(kernel, "}}").unwrap(); + writeln!( + kernel, + "acc += q8_scale * (f32(chunk.scale) * scale_dot_sum - f32(chunk.min) * offset_y_sum);" + ) + .unwrap(); + writeln!(kernel, "}}").unwrap(); + writeln!(kernel, "workgroupBarrier();").unwrap(); + writeln!(kernel, "}}").unwrap(); + + writeln!(kernel, "if x < {n_size} && y < {m_size} {{").unwrap(); + write!(kernel, "let output_index = ").unwrap(); + let mut output_indices = Vec::new(); + for dim in (0..output.rank()).rev().skip(2) { + output_indices.push(format!("block_batch_{dim}")); + } + output_indices.push("y".to_string()); + output_indices.push("x".to_string()); + output.strided_index(kernel, output_indices); + writeln!(kernel, ";").unwrap(); + let post_fns = + post_element_wise_functions.get_or_init(|| op.post_element_wise.add_functions(kernel)); + let result = op.apply_bias_and_post(bias, "x", format!("{input_datatype}(acc)"), post_fns); + writeln!(kernel, "{output}[output_index] = {result};").unwrap(); + writeln!(kernel, "}}").unwrap(); +} + +fn q4k_input_value_expr( + input_a: &TensorInput, + local_expr: &str, + pre_element_wise_functions: &[crate::mir::function::Function], +) -> String { + let mut raw_input = String::new(); + write!(&mut raw_input, "{input_a}[").unwrap(); + let mut indices = Vec::new(); + for dim in (0..input_a.rank()).rev().skip(2) { + indices.push(format!("block_batch_{dim}")); + } + indices.push("y".to_string()); + indices.push(format!("a_index_base + {local_expr}")); + input_a.strided_index(&mut raw_input, indices); + write!(&mut raw_input, "]").unwrap(); + pre_element_wise_functions + .iter() + .fold(raw_input, |acc, f| f.call(vec![acc])) } diff --git a/fusor-ml/core/src/quantized/matmul/sgemm/mod.rs b/fusor-ml/core/src/quantized/matmul/sgemm/mod.rs index 5fb6a1950..a244ba7ba 100644 --- a/fusor-ml/core/src/quantized/matmul/sgemm/mod.rs +++ b/fusor-ml/core/src/quantized/matmul/sgemm/mod.rs @@ -7,12 +7,26 @@ use crate::{ }, quantized::matmul::QMatMulOperation, }; +use fusor_gguf::GgmlType; mod chunked; mod general; pub use chunked::{ChunkedSgemmConfig, chunked_sgemm_with_config}; -pub use general::{GeneralSgemmConfig, general_sgemm_with_config}; +pub use general::{ + GeneralSgemmConfig, general_q4_0_sgemm, general_q4k_sgemm, general_sgemm_with_config, +}; + +fn use_exact_q4k_sgemm(matrix: &QMatrix) -> bool { + matrix.datatype() == GgmlType::Q4K && matrix.shape()[0] <= 64 +} + +fn use_chunked_sgemm(matrix: &QMatrix) -> bool { + !matches!(matrix.datatype(), GgmlType::Q4_0) + && !use_exact_q4k_sgemm(matrix) + && dequantize_mat4x4_block_count(matrix.datatype()) > 0 + && matrix.device().subgroups_supported() +} #[allow(clippy::too_many_arguments)] pub(crate) fn sgemm( @@ -21,6 +35,7 @@ pub(crate) fn sgemm( _: &WorkgroupShape, input_a: &TensorInput, input_b: &QMatrixInput, + bias: Option<&TensorInput>, output: &TensorInput, _n_size: &str, // m size is always 1 for sgemv @@ -29,9 +44,38 @@ pub(crate) fn sgemm( graph: &crate::compute_graph::ComputeGraphInner, ) { // Use chunked sgemm for all types that support mat4x4 dequantization - if dequantize_mat4x4_block_count(input_b.datatype) > 0 && graph.device().subgroups_supported() { + if op.matrix.datatype() == GgmlType::Q4_0 || use_exact_q4k_sgemm(&op.matrix) { + let integer_sgemm = match op.matrix.datatype() { + GgmlType::Q4_0 => general_q4_0_sgemm, + GgmlType::Q4K => general_q4k_sgemm, + _ => unreachable!(), + }; + integer_sgemm( + op, + generic_kernel, + input_a, + input_b, + bias, + output, + _n_size, + _m_size, + k_size, + graph.device().f16_supported(), + ); + } else if dequantize_mat4x4_block_count(input_b.datatype) > 0 + && graph.device().subgroups_supported() + { let config = op.chunked_config.unwrap_or(ChunkedSgemmConfig::default()); - chunked_sgemm_with_config(op, generic_kernel, input_a, input_b, output, k_size, config); + chunked_sgemm_with_config( + op, + generic_kernel, + input_a, + input_b, + bias, + output, + k_size, + config, + ); } else { let config = op.general_config.unwrap_or(GeneralSgemmConfig::default()); general_sgemm_with_config( @@ -39,6 +83,7 @@ pub(crate) fn sgemm( generic_kernel, input_a, input_b, + bias, output, _n_size, _m_size, @@ -57,8 +102,7 @@ pub(crate) fn dispatch_size( batch_size: u32, ) -> [u32; 3] { // Use chunked dispatch size for all types that support mat4x4 dequantization - if dequantize_mat4x4_block_count(matrix.datatype()) > 0 && matrix.device().subgroups_supported() - { + if use_chunked_sgemm(matrix) { let config = op.chunked_config.unwrap_or(ChunkedSgemmConfig::default()); [ m.div_ceil(workgroup_shape.y() * config.m_results_per_thread * 4), @@ -79,13 +123,20 @@ pub(crate) fn workgroup_shape_constraints( _device: &Device, ) -> crate::mir::workgroup_shape::WorkgroupShapeConstraints { // Use chunked workgroup constraints for all types that support mat4x4 dequantization - if dequantize_mat4x4_block_count(matrix.datatype()) > 0 { + if use_chunked_sgemm(matrix) { let mut constraints = crate::mir::workgroup_shape::WorkgroupShapeConstraints::default(); constraints.add_constraint(0, Constraint::equals(16)); constraints.add_constraint(1, Constraint::equals(16)); constraints.add_constraint(2, Constraint::equals(1)); return constraints; } + if use_exact_q4k_sgemm(matrix) { + let mut constraints = crate::mir::workgroup_shape::WorkgroupShapeConstraints::default(); + constraints.add_constraint(0, Constraint::equals(64)); + constraints.add_constraint(1, Constraint::equals(1)); + constraints.add_constraint(2, Constraint::equals(1)); + return constraints; + } let mut constraints = crate::mir::workgroup_shape::WorkgroupShapeConstraints::default(); let second_dim = matrix.shape()[1]; let second_dim_block_width = second_dim.div_ceil(matrix.datatype().block_size()); diff --git a/fusor-ml/core/src/quantized/matmul/sgemv/general.rs b/fusor-ml/core/src/quantized/matmul/sgemv/general.rs index 057154bad..385efe4ed 100644 --- a/fusor-ml/core/src/quantized/matmul/sgemv/general.rs +++ b/fusor-ml/core/src/quantized/matmul/sgemv/general.rs @@ -8,7 +8,10 @@ use crate::{ }, quantized::matmul::{ QMatMulOperation, - sgemv::{SGEMV_CHUNK_SIZE, SGEMV_VECTOR_SIZE, decompose_workgroup_index}, + sgemv::{ + SGEMV_CHUNK_SIZE, SGEMV_VECTOR_SIZE, decompose_workgroup_index, + quantized_sgemv_subgroups_supported, + }, }, util::{ maybe_vec_storage_add, maybe_vec_storage_index, maybe_vec_storage_subgroup_add, @@ -16,6 +19,7 @@ use crate::{ }, }; use std::fmt::Write; +use std::sync::OnceLock; #[allow(clippy::too_many_arguments)] pub(crate) fn general_sgemv( @@ -24,6 +28,7 @@ pub(crate) fn general_sgemv( workgroup_size: &WorkgroupShape, input_a: &TensorInput, input_b: &QMatrixInput, + bias: Option<&TensorInput>, output: &TensorInput, n_size: &str, m_size: &str, @@ -31,10 +36,12 @@ pub(crate) fn general_sgemv( graph: &crate::compute_graph::ComputeGraphInner, ) { let blocksize = workgroup_size.x(); - let dtype = op.input_datatype; + let input_datatype = op.input_datatype; let workgroup_local_index = kernel.workgroup_local_index(); let elements_per_block = op.elements_per_block(); let device = graph.device(); + let pre_element_wise_functions = OnceLock::new(); + let post_element_wise_functions = OnceLock::new(); // Calculate n_workgroups for decomposing the linearized workgroup index let n_workgroups = format!("(({n_size} + {SGEMV_CHUNK_SIZE} - 1) / {SGEMV_CHUNK_SIZE})"); @@ -66,6 +73,13 @@ pub(crate) fn general_sgemv( ) .unwrap(); + // Fast path check: if all rows in this tile are valid, skip per-row bounds checks + writeln!( + kernel, + "let is_full_tile = workgroup_offset + {SGEMV_CHUNK_SIZE} <= {n_size};" + ) + .unwrap(); + // Always accumulate in f32 for precision, convert to output dtype at the end let acc_storage_type = maybe_vec_storage_type(SGEMV_CHUNK_SIZE, DataTypeEnum::F32); @@ -87,7 +101,7 @@ pub(crate) fn general_sgemv( debug_assert!(elements_per_block.is_multiple_of(SGEMV_VECTOR_SIZE)); writeln!( kernel, - "var a_cache = array, {chunk_blocks}>();" + "var a_cache = array, {chunk_blocks}>();" ) .unwrap(); @@ -98,9 +112,12 @@ pub(crate) fn general_sgemv( writeln!(kernel, "for (var i = 0u; i < {chunk_blocks}; i += 1u) {{").unwrap(); { // Get the values first + let pre_element_wise_functions = pre_element_wise_functions + .get_or_init(|| op.pre_element_wise.add_functions(kernel)); for i in 0..SGEMV_VECTOR_SIZE { writeln!(kernel, "let input_a_{i}_index = index * {elements_per_block} + i * {SGEMV_VECTOR_SIZE} + {i};").unwrap(); - write!(kernel, "let input_a_{i} = {input_a}[").unwrap(); + let mut raw_input = String::new(); + write!(&mut raw_input, "{input_a}[").unwrap(); let mut indices = Vec::new(); // Add batch indices first for dim in (0..input_a.rank()).rev().skip(2) { @@ -109,8 +126,12 @@ pub(crate) fn general_sgemv( // Then add M and K indices indices.push("m_idx".to_string()); indices.push(format!("input_a_{i}_index")); - input_a.strided_index(kernel, indices); - writeln!(kernel, "];").unwrap(); + input_a.strided_index(&mut raw_input, indices); + write!(&mut raw_input, "]").unwrap(); + let processed = pre_element_wise_functions + .iter() + .fold(raw_input, |acc, f| f.call(vec![acc])); + writeln!(kernel, "let input_a_{i} = f32({processed});").unwrap(); } // The pack them into a vector and write to the cache write!(kernel, "a_cache[i] = vec{SGEMV_VECTOR_SIZE}(").unwrap(); @@ -124,6 +145,49 @@ pub(crate) fn general_sgemv( } writeln!(kernel, "}}").unwrap(); + // Generate the computation body for one row + let generate_row_computation = + |kernel: &mut GenericKernel, op: &QMatMulOperation, input_b: &QMatrixInput| { + let index = if SGEMV_CHUNK_SIZE > 1 { + "(workgroup_offset + acc_offset)" + } else { + "workgroup_offset" + }; + writeln!( + kernel, + "let chunk = &{input_b}[{index} * k_block_size + index];" + ) + .unwrap(); + + let acc_indexed = maybe_vec_storage_index(SGEMV_CHUNK_SIZE, "acc", "acc_offset"); + // Always convert a_cache to f32 for the dot product since dequantize outputs f32 + // and we accumulate in f32 for precision + dequantize_vec4_block( + kernel, + op.matrix.datatype, + "chunk".to_string(), + DataTypeEnum::F32, + |index, data, code| { + writeln!(code, "{acc_indexed} += dot(a_cache[{index}], {data});").unwrap(); + }, + ); + }; + + // Fast path: all rows in tile are valid, no bounds checks needed + writeln!(kernel, "if is_full_tile {{").unwrap(); + if SGEMV_CHUNK_SIZE > 1 { + writeln!( + kernel, + "for (var acc_offset = 0u; acc_offset < {SGEMV_CHUNK_SIZE}; acc_offset += 1u) {{" + ) + .unwrap(); + } + generate_row_computation(kernel, op, input_b); + if SGEMV_CHUNK_SIZE > 1 { + writeln!(kernel, "}}").unwrap(); + } + writeln!(kernel, "}} else {{").unwrap(); + // Slow path: check bounds for each row if SGEMV_CHUNK_SIZE > 1 { writeln!( kernel, @@ -136,32 +200,14 @@ pub(crate) fn general_sgemv( } else { "workgroup_offset" }; - writeln!( - kernel, - "let chunk = &{input_b}[{index} * k_block_size + index];" - ) - .unwrap(); - - let acc_indexed = maybe_vec_storage_index(SGEMV_CHUNK_SIZE, "acc", "acc_offset"); - // Always convert a_cache to f32 for the dot product since dequantize outputs f32 - // and we accumulate in f32 for precision - dequantize_vec4_block( - kernel, - op.matrix.datatype, - "chunk".to_string(), - DataTypeEnum::F32, - |index, data, code| { - writeln!( - code, - "{acc_indexed} += dot(vec4(a_cache[{index}]), {data});" - ) - .unwrap(); - }, - ); - + writeln!(kernel, "let row_in_bounds = {index} < {n_size};").unwrap(); + writeln!(kernel, "if row_in_bounds {{").unwrap(); + generate_row_computation(kernel, op, input_b); + writeln!(kernel, "}}").unwrap(); if SGEMV_CHUNK_SIZE > 1 { writeln!(kernel, "}}").unwrap(); } + writeln!(kernel, "}}").unwrap(); writeln!(kernel, "index += {blocksize}u;").unwrap(); } @@ -169,7 +215,7 @@ pub(crate) fn general_sgemv( writeln!(kernel, "}}").unwrap(); // Reduce with subgroup operations if the device supports subgroups - if device.subgroups_supported() { + if quantized_sgemv_subgroups_supported(&device) { // Get the sum among all threads in the subgroup writeln!( kernel, @@ -264,7 +310,36 @@ pub(crate) fn general_sgemv( // If this is not the first simd thread in the workgroup, we can return early writeln!(kernel, "if {workgroup_local_index} != 0u {{ return; }}").unwrap(); + // Initialize post element-wise functions once before the loop + let post_fns = + post_element_wise_functions.get_or_init(|| op.post_element_wise.add_functions(kernel)); + + // Generate the output body for one row + let generate_row_output = |kernel: &mut GenericKernel| { + write!(kernel, "{output}[").unwrap(); + let mut output_indices = Vec::new(); + // Add batch indices first + for dim in (0..output.rank()).rev().skip(2) { + output_indices.push(format!("batch_idx_{dim}")); + } + // Then add M and N indices + output_indices.push("m_idx".to_string()); + output_indices.push("output_index".to_string()); + output.strided_index(kernel, output_indices); + // Convert from f32 accumulator to output dtype (single element per iteration) + let acc_val = maybe_vec_storage_index(SGEMV_CHUNK_SIZE, "acc", "acc_offset"); + let result = op.apply_bias_and_post( + bias, + "output_index", + format!("{input_datatype}({acc_val})"), + post_fns, + ); + writeln!(kernel, "] = {result};").unwrap(); + }; + // Write the output to the output tensor if this is the first thread in the workgroup + // Fast path: all rows in tile are valid, no bounds checks needed + writeln!(kernel, "if is_full_tile {{").unwrap(); if SGEMV_CHUNK_SIZE > 1 { writeln!( kernel, @@ -278,21 +353,31 @@ pub(crate) fn general_sgemv( } else { writeln!(kernel, "let output_index = workgroup_offset;").unwrap(); } - write!(kernel, "{output}[").unwrap(); - let mut output_indices = Vec::new(); - // Add batch indices first - for dim in (0..output.rank()).rev().skip(2) { - output_indices.push(format!("batch_idx_{dim}")); + generate_row_output(kernel); + } + if SGEMV_CHUNK_SIZE > 1 { + writeln!(kernel, "}}").unwrap(); + } + writeln!(kernel, "}} else {{").unwrap(); + // Slow path: check bounds for each row + if SGEMV_CHUNK_SIZE > 1 { + writeln!( + kernel, + "for (var acc_offset = 0u; acc_offset < {SGEMV_CHUNK_SIZE}; acc_offset += 1u) {{" + ) + .unwrap(); + } + { + if SGEMV_CHUNK_SIZE > 1 { + writeln!(kernel, "let output_index = workgroup_offset + acc_offset;").unwrap(); + } else { + writeln!(kernel, "let output_index = workgroup_offset;").unwrap(); } - // Then add M and N indices - output_indices.push("m_idx".to_string()); - output_indices.push("output_index".to_string()); - output.strided_index(kernel, output_indices); - // Convert from f32 accumulator to output dtype (single element per iteration) - let acc_val = maybe_vec_storage_index(SGEMV_CHUNK_SIZE, "acc", "acc_offset"); - writeln!(kernel, "] = {dtype}({acc_val});").unwrap(); + writeln!(kernel, "if output_index >= {n_size} {{ continue; }}").unwrap(); + generate_row_output(kernel); } if SGEMV_CHUNK_SIZE > 1 { writeln!(kernel, "}}").unwrap(); } + writeln!(kernel, "}}").unwrap(); } diff --git a/fusor-ml/core/src/quantized/matmul/sgemv/mod.rs b/fusor-ml/core/src/quantized/matmul/sgemv/mod.rs index abc2b4866..fec9f9e31 100644 --- a/fusor-ml/core/src/quantized/matmul/sgemv/mod.rs +++ b/fusor-ml/core/src/quantized/matmul/sgemv/mod.rs @@ -20,7 +20,7 @@ use crate::{ use crate::{ QMatrix, quantized::matmul::sgemv::{ - q_8_0::Q_8_0_SGEMV_CHUNK_SIZE, q_n::Q_N_SGEMV_CHUNK_SIZE, q4k::Q4K_SGEMV_CHUNK_SIZE, + q_8_0::Q_8_0_SGEMV_CHUNK_SIZE, q_n::q_n_sgemv_chunk_size, q4k::Q4K_SGEMV_CHUNK_SIZE, q5k::Q5K_SGEMV_CHUNK_SIZE, q6k::Q6K_SGEMV_CHUNK_SIZE, }, visit_tiled::distribute_workgroups, @@ -68,8 +68,12 @@ pub(crate) fn decompose_workgroup_index( /// Check if the device can support specialized SGEMV kernels that require 2 subgroups per workgroup. /// This requires the workgroup to be large enough to fit 2 subgroups, which means: /// max_subgroup_size >= 2 * min_subgroup_size +pub(crate) fn quantized_sgemv_subgroups_supported(device: &Device) -> bool { + device.subgroups_supported() +} + fn can_use_specialized_sgemv(device: &Device) -> bool { - if !device.subgroups_supported() { + if !quantized_sgemv_subgroups_supported(device) { return false; } // The workgroup is constrained to be <= max_subgroup_size, so we need @@ -77,6 +81,23 @@ fn can_use_specialized_sgemv(device: &Device) -> bool { device.max_subgroup_size() >= 2 * device.min_subgroup_size() } +fn can_use_specialized_q8_0_sgemv(device: &Device) -> bool { + can_use_specialized_sgemv(device) +} + +#[cfg(test)] +pub(crate) fn selected_sgemv_kernel_kind(op: &QMatMulOperation, device: &Device) -> &'static str { + let use_specialized = can_use_specialized_sgemv(device); + match op.matrix.datatype { + GgmlType::Q6K if use_specialized => "q6k", + GgmlType::Q4K if use_specialized => "q4k", + GgmlType::Q5K if use_specialized => "q5k", + GgmlType::Q4_0 | GgmlType::Q5_0 if use_specialized => "q_n", + GgmlType::Q8_0 if can_use_specialized_q8_0_sgemv(device) => "q_8_0", + _ => "general", + } +} + #[allow(clippy::too_many_arguments)] pub(crate) fn sgemv( op: &QMatMulOperation, @@ -84,6 +105,7 @@ pub(crate) fn sgemv( workgroup_size: &WorkgroupShape, input_a: &TensorInput, input_b: &QMatrixInput, + bias: Option<&TensorInput>, output: &TensorInput, _n_size: &str, _m_size: &str, @@ -100,6 +122,7 @@ pub(crate) fn sgemv( workgroup_size, input_a, input_b, + bias, output, _n_size, _m_size, @@ -111,6 +134,7 @@ pub(crate) fn sgemv( workgroup_size, input_a, input_b, + bias, output, _n_size, _m_size, @@ -122,6 +146,7 @@ pub(crate) fn sgemv( workgroup_size, input_a, input_b, + bias, output, _n_size, _m_size, @@ -133,17 +158,19 @@ pub(crate) fn sgemv( workgroup_size, input_a, input_b, + bias, output, _n_size, _m_size, k_size, ), - GgmlType::Q8_0 if use_specialized => q_8_0_sgemv( + GgmlType::Q8_0 if can_use_specialized_q8_0_sgemv(&device) => q_8_0_sgemv( op, generic_kernel, workgroup_size, input_a, input_b, + bias, output, _n_size, _m_size, @@ -155,6 +182,7 @@ pub(crate) fn sgemv( workgroup_size, input_a, input_b, + bias, output, _n_size, _m_size, @@ -175,7 +203,7 @@ pub(crate) fn n_workgroups(matrix: &QMatrix, n: u32) -> u32 { } else if matrix.datatype == GgmlType::Q5K { n.div_ceil(Q5K_SGEMV_CHUNK_SIZE * 2) } else if matches!(matrix.datatype, GgmlType::Q4_0 | GgmlType::Q5_0) { - n.div_ceil(Q_N_SGEMV_CHUNK_SIZE * 2) + n.div_ceil(q_n_sgemv_chunk_size(n) * 2) } else if matches!(matrix.datatype, GgmlType::Q8_0) { n.div_ceil(Q_8_0_SGEMV_CHUNK_SIZE * 2) } else { @@ -199,7 +227,7 @@ pub(crate) fn workgroup_shape_constraints( device: &Device, ) -> crate::mir::workgroup_shape::WorkgroupShapeConstraints { let mut constraints = crate::mir::workgroup_shape::WorkgroupShapeConstraints::default(); - if device.subgroups_supported() { + if quantized_sgemv_subgroups_supported(device) { constraints.add_constraint( 0, crate::mir::workgroup_shape::Constraint::more_than_or_equals( diff --git a/fusor-ml/core/src/quantized/matmul/sgemv/q4k.rs b/fusor-ml/core/src/quantized/matmul/sgemv/q4k.rs index ec3197eee..d8904725e 100644 --- a/fusor-ml/core/src/quantized/matmul/sgemv/q4k.rs +++ b/fusor-ml/core/src/quantized/matmul/sgemv/q4k.rs @@ -1,6 +1,7 @@ use crate::{ DataTypeEnum, mir::{ + function::Function, inputs::{QMatrixInput, TensorInput}, kernel::GenericKernel, workgroup_shape::WorkgroupShape, @@ -10,6 +11,7 @@ use crate::{ util::{maybe_vec_storage_index, maybe_vec_storage_subgroup_add, maybe_vec_storage_type}, }; use std::fmt::Write; +use std::sync::OnceLock; pub(crate) const Q4K_SGEMV_CHUNK_SIZE: u32 = 4; // This is the size of the chunk each thread will process at a time const SUBGROUP_COUNT: u32 = 2; @@ -18,6 +20,165 @@ const MASK1: u32 = 0b0011111100111111; const MASK2: u32 = 0b0000111100001111; const MASK3: u32 = 0b1100000011000000; +/// Generate WGSL for computing one row's contribution in the Q4K inner loop +fn generate_row_computation(kernel: &mut GenericKernel, offset: u32, input_b: &QMatrixInput) { + // Fetch and unpack the two sets of values from the cache + writeln!(kernel, "let first_values_offset = data_offset;").unwrap(); + writeln!(kernel, "let second_values_offset = data_offset + 16u;").unwrap(); + + // Keep track of the sum of each chunk + writeln!(kernel, "var first_sums = vec4();").unwrap(); + writeln!(kernel, "var second_sums = vec4();").unwrap(); + + // Perform the dot product of the values and scales + for j in 0..2 { + for (sum, cache, values) in [ + ("first_sums", "cached_a_low_values", "first_values_offset"), + ( + "second_sums", + "cached_a_high_values", + "second_values_offset", + ), + ] { + writeln!( + kernel, + "let value_u32_{values}_{j} = {input_b}[local_block_offset].data[{values} + {j}];" + ) + .unwrap(); + writeln!( + kernel, + "let first_four_values_{values}_{j} = vec4({cache}[{j}*4 + 0], {cache}[{j}*4 + 1], {cache}[{j}*4 + 2], {cache}[{j}*4 + 3]);" + ) + .unwrap(); + writeln!( + kernel, + "let second_four_values_{values}_{j} = vec4({cache}[{j}*4 + 8], {cache}[{j}*4 + 9], {cache}[{j}*4 + 10], {cache}[{j}*4 + 11]);" + ) + .unwrap(); + writeln!( + kernel, + "{sum} += vec4(first_four_values_{values}_{j}.x * f32(value_u32_{values}_{j} & 0x000F), first_four_values_{values}_{j}.y * f32(value_u32_{values}_{j} & 0x0F00), second_four_values_{values}_{j}.x * f32(value_u32_{values}_{j} & 0x00F0), second_four_values_{values}_{j}.y * f32(value_u32_{values}_{j} & 0xF000));" + ) + .unwrap(); + let shift_right_16 = shift_right_scale(16); + writeln!( + kernel, + "{sum} += vec4(first_four_values_{values}_{j}.z * f32(value_u32_{values}_{j} & 0x000F0000), first_four_values_{values}_{j}.w * f32(value_u32_{values}_{j} & 0x0F000000), second_four_values_{values}_{j}.z * f32(value_u32_{values}_{j} & 0x00F00000), second_four_values_{values}_{j}.w * f32(value_u32_{values}_{j} & 0xF0000000)) * f32({shift_right_16});" + ) + .unwrap(); + } + } + + // Load the block scale and min + writeln!( + kernel, + "let block_scale = f32({input_b}[local_block_offset].scale);" + ) + .unwrap(); + writeln!( + kernel, + "let block_min = f32({input_b}[local_block_offset].min);" + ) + .unwrap(); + // Load 8 scales into a cache + writeln!( + kernel, + "let first_32_scale_bits = {input_b}[local_block_offset].scales[0] >> (16 * scale_offset);" + ) + .unwrap(); + writeln!( + kernel, + "let second_32_scale_bits = {input_b}[local_block_offset].scales[1] >> (16 * scale_offset);" + ) + .unwrap(); + writeln!( + kernel, + "let third_32_scale_bits = {input_b}[local_block_offset].scales[2] >> (16 * scale_offset);" + ) + .unwrap(); + // Extract the scales from the bits into cached_scales + writeln!( + kernel, + "let first_two_scales = first_32_scale_bits & {MASK1};" + ) + .unwrap(); + writeln!( + kernel, + "let second_two_scales = second_32_scale_bits & {MASK1};" + ) + .unwrap(); + + writeln!(kernel, "let third_two_scales = ((third_32_scale_bits >> 0) & {MASK2}) | ((first_32_scale_bits & {MASK3}) >> 2);").unwrap(); + writeln!(kernel, "let fourth_two_scales = ((third_32_scale_bits >> 4) & {MASK2}) | ((second_32_scale_bits & {MASK3}) >> 2);").unwrap(); + + writeln!( + kernel, + "let odd_scales_unpacked = vec4(unpack4xU8(first_two_scales | (third_two_scales << 16)));" + ) + .unwrap(); + writeln!( + kernel, + "let even_scales_unpacked = vec4(unpack4xU8(second_two_scales | (fourth_two_scales << 16)));" + ) + .unwrap(); + + // Add the sums to the total sum + let indexed_sum = maybe_vec_storage_index(Q4K_SGEMV_CHUNK_SIZE, "sum", offset); + let shift_right_8 = shift_right_scale(8); + let shift_right_4 = shift_right_scale(4); + writeln!( + kernel, + "let small_shift_sums = vec4(first_sums[0], first_sums[2], second_sums[0], second_sums[2]);" + ) + .unwrap(); + writeln!( + kernel, + "let large_shift_sums = vec4(first_sums[1], first_sums[3], second_sums[1], second_sums[3]);" + ) + .unwrap(); + writeln!( + kernel, + "let shift_4 = vec4(1.0, {shift_right_4}, 1.0, {shift_right_4});" + ) + .unwrap(); + writeln!( + kernel, + r#"{indexed_sum} += block_scale * dot((small_shift_sums + f32({shift_right_8}) * large_shift_sums) * odd_scales_unpacked, shift_4) - + block_min * dot(vector_sum, even_scales_unpacked);"# + ) + .unwrap(); + // Move forward the block offset by one row + writeln!(kernel, "local_block_offset += k_block_size;").unwrap(); +} + +/// Generate WGSL for writing one row's output +fn generate_row_output( + op: &QMatMulOperation, + kernel: &mut GenericKernel, + offset: u32, + dtype: DataTypeEnum, + bias: Option<&TensorInput>, + output: &TensorInput, + post_element_wise_functions: &[Function], +) { + write!(kernel, "{output}[").unwrap(); + let mut output_indices = Vec::new(); + for dim in (0..output.rank()).rev().skip(2) { + output_indices.push(format!("batch_idx_{dim}")); + } + output_indices.push("m_idx".to_string()); + output_indices.push("row_index".to_string()); + output.strided_index(kernel, output_indices); + let indexed = maybe_vec_storage_index(Q4K_SGEMV_CHUNK_SIZE, "sum", offset); + let result = op.apply_bias_and_post( + bias, + "row_index", + format!("{dtype}({indexed})"), + post_element_wise_functions, + ); + writeln!(kernel, "] = {result};").unwrap(); +} + // https://github.com/ggml-org/llama.cpp/blob/6efcd65945a98cf6883cdd9de4c8ccd8c79d219a/ggml/src/ggml-metal/ggml-metal.metal#L5311 #[allow(clippy::too_many_arguments)] pub(crate) fn q4k_sgemv( @@ -26,6 +187,7 @@ pub(crate) fn q4k_sgemv( workgroup_shape: &WorkgroupShape, input_a: &TensorInput, input_b: &QMatrixInput, + bias: Option<&TensorInput>, output: &TensorInput, n_size: &str, m_size: &str, @@ -35,6 +197,8 @@ pub(crate) fn q4k_sgemv( let subgroup_index = kernel.subgroup_index(); let subgroup_local_index = kernel.subgroup_local_index(); let elements_per_block = op.elements_per_block(); + let pre_element_wise_functions = OnceLock::new(); + let post_element_wise_functions = OnceLock::new(); // Calculate n_workgroups for this kernel type (SUBGROUP_COUNT subgroups per workgroup, Q4K_SGEMV_CHUNK_SIZE per subgroup) let chunk_size = Q4K_SGEMV_CHUNK_SIZE * SUBGROUP_COUNT; @@ -74,6 +238,13 @@ pub(crate) fn q4k_sgemv( ) .unwrap(); + // Fast path check: if all rows in this tile are valid, skip per-row bounds checks + writeln!( + kernel, + "let is_full_tile = row + {Q4K_SGEMV_CHUNK_SIZE} <= {n_size};" + ) + .unwrap(); + writeln!(kernel, "let thread_id = {subgroup_local_index} >> 3;").unwrap(); writeln!(kernel, "let thread_local_id = {subgroup_local_index} & 7;").unwrap(); writeln!(kernel, "let half_subgroup_id = thread_local_id >> 2;").unwrap(); @@ -103,7 +274,10 @@ pub(crate) fn q4k_sgemv( for j in 0..8 { // Load all 4 values using strided indexing for (idx, offset) in [(0, 0), (1, 32), (2, 128), (3, 160)] { - write!(kernel, "let a_val_{j}_{idx} = {input_a}[").unwrap(); + let pre_element_wise_functions = pre_element_wise_functions + .get_or_init(|| op.pre_element_wise.add_functions(kernel)); + let mut raw_input = String::new(); + write!(&mut raw_input, "{input_a}[").unwrap(); let mut indices = Vec::new(); // Add batch indices first for dim in (0..input_a.rank()).rev().skip(2) { @@ -112,8 +286,12 @@ pub(crate) fn q4k_sgemv( // Then add M and K indices indices.push("m_idx".to_string()); indices.push(format!("vector_offset + {j} + {offset}")); - input_a.strided_index(kernel, indices); - writeln!(kernel, "];").unwrap(); + input_a.strided_index(&mut raw_input, indices); + write!(&mut raw_input, "]").unwrap(); + let processed = pre_element_wise_functions + .iter() + .fold(raw_input, |acc, f| f.call(vec![acc])); + writeln!(kernel, "let a_val_{j}_{idx} = f32({processed});").unwrap(); } writeln!( @@ -138,146 +316,24 @@ pub(crate) fn q4k_sgemv( writeln!(kernel, "var local_block_offset = block_offset + i;").unwrap(); + // Fast path: all rows in tile are valid, no bounds checks needed + writeln!(kernel, "if is_full_tile {{").unwrap(); for offset in 0..Q4K_SGEMV_CHUNK_SIZE { writeln!(kernel, "{{").unwrap(); - // Fetch and unpack the two sets of values from the cache - writeln!(kernel, "let first_values_offset = data_offset;").unwrap(); - writeln!(kernel, "let second_values_offset = data_offset + 16u;").unwrap(); - - // Keep track of the sum of each chunk - writeln!(kernel, "var first_sums = vec4();").unwrap(); - writeln!(kernel, "var second_sums = vec4();").unwrap(); - - // Perform the dot product of the values and scales - for j in 0..2 { - for (sum, cache, values) in [ - ("first_sums", "cached_a_low_values", "first_values_offset"), - ( - "second_sums", - "cached_a_high_values", - "second_values_offset", - ), - ] { - // Note: We add the values with a mask **without** shifting them - // this means the sums in the first_sums and second_sums - // will be scaled by different values. We correct this below - // by multiplying by the floating point values that correspond to the - // bit shifts. - writeln!( - kernel, - "let value_u32_{values}_{j} = {input_b}[local_block_offset].data[{values} + {j}];" - ) - .unwrap(); - writeln!( - kernel, - "let first_four_values_{values}_{j} = vec4({cache}[{j}*4 + 0], {cache}[{j}*4 + 1], {cache}[{j}*4 + 2], {cache}[{j}*4 + 3]);" - ) - .unwrap(); - writeln!( - kernel, - "let second_four_values_{values}_{j} = vec4({cache}[{j}*4 + 8], {cache}[{j}*4 + 9], {cache}[{j}*4 + 10], {cache}[{j}*4 + 11]);" - ) - .unwrap(); - writeln!( - kernel, - "{sum} += vec4(first_four_values_{values}_{j}.x * f32(value_u32_{values}_{j} & 0x000F), first_four_values_{values}_{j}.y * f32(value_u32_{values}_{j} & 0x0F00), second_four_values_{values}_{j}.x * f32(value_u32_{values}_{j} & 0x00F0), second_four_values_{values}_{j}.y * f32(value_u32_{values}_{j} & 0xF000));" - ) - .unwrap(); - let shift_right_16 = shift_right_scale(16); - writeln!( - kernel, - "{sum} += vec4(first_four_values_{values}_{j}.z * f32(value_u32_{values}_{j} & 0x000F0000), first_four_values_{values}_{j}.w * f32(value_u32_{values}_{j} & 0x0F000000), second_four_values_{values}_{j}.z * f32(value_u32_{values}_{j} & 0x00F00000), second_four_values_{values}_{j}.w * f32(value_u32_{values}_{j} & 0xF0000000)) * f32({shift_right_16});" - ) - .unwrap(); - } - } - - // Load the block scale and min - writeln!( - kernel, - "let block_scale = f32({input_b}[local_block_offset].scale);" - ) - .unwrap(); - writeln!( - kernel, - "let block_min = f32({input_b}[local_block_offset].min);" - ) - .unwrap(); - // Load 8 scales into a cache - writeln!( - kernel, - "let first_32_scale_bits = {input_b}[local_block_offset].scales[0] >> (16 * scale_offset);" - ) - .unwrap(); - writeln!( - kernel, - "let second_32_scale_bits = {input_b}[local_block_offset].scales[1] >> (16 * scale_offset);" - ) - .unwrap(); - writeln!( - kernel, - "let third_32_scale_bits = {input_b}[local_block_offset].scales[2] >> (16 * scale_offset);" - ) - .unwrap(); - // Extract the scales from the bits into cached_scales - writeln!( - kernel, - "let first_two_scales = first_32_scale_bits & {MASK1};" - ) - .unwrap(); - writeln!( - kernel, - "let second_two_scales = second_32_scale_bits & {MASK1};" - ) - .unwrap(); - - writeln!(kernel, "let third_two_scales = ((third_32_scale_bits >> 0) & {MASK2}) | ((first_32_scale_bits & {MASK3}) >> 2);").unwrap(); - writeln!(kernel, "let fourth_two_scales = ((third_32_scale_bits >> 4) & {MASK2}) | ((second_32_scale_bits & {MASK3}) >> 2);").unwrap(); - - writeln!( - kernel, - "let odd_scales_unpacked = vec4(unpack4xU8(first_two_scales | (third_two_scales << 16)));" - ) - .unwrap(); - writeln!( - kernel, - "let even_scales_unpacked = vec4(unpack4xU8(second_two_scales | (fourth_two_scales << 16)));" - ) - .unwrap(); - - // Add the sums to the total sum - let indexed_sum = maybe_vec_storage_index(Q4K_SGEMV_CHUNK_SIZE, "sum", offset); - // *_sums[0] needs to be shifted by 0 bits - // *_sums[1] needs to be shifted by 8 bits - // *_sums[2] needs to be shifted by 4 bits - // *_sums[3] needs to be shifted by 12 bits - let shift_right_8 = shift_right_scale(8); - let shift_right_4 = shift_right_scale(4); - writeln!( - kernel, - "let small_shift_sums = vec4(first_sums[0], first_sums[2], second_sums[0], second_sums[2]);" - ) - .unwrap(); - writeln!( - kernel, - "let large_shift_sums = vec4(first_sums[1], first_sums[3], second_sums[1], second_sums[3]);" - ) - .unwrap(); - writeln!( - kernel, - "let shift_4 = vec4(1.0, {shift_right_4}, 1.0, {shift_right_4});" - ) - .unwrap(); - writeln!( - kernel, - r#"{indexed_sum} += block_scale * dot((small_shift_sums + f32({shift_right_8}) * large_shift_sums) * odd_scales_unpacked, shift_4) - - block_min * dot(vector_sum, even_scales_unpacked);"# - ) - .unwrap(); - // Move forward the block offset by one row - writeln!(kernel, "local_block_offset += k_block_size;").unwrap(); + generate_row_computation(kernel, offset, input_b); + writeln!(kernel, "}}").unwrap(); + } + writeln!(kernel, "}} else {{").unwrap(); + // Slow path: check bounds for each row + for offset in 0..Q4K_SGEMV_CHUNK_SIZE { + writeln!(kernel, "{{").unwrap(); + writeln!(kernel, "let row_index = row + {offset};").unwrap(); + writeln!(kernel, "if row_index < {n_size} {{").unwrap(); + generate_row_computation(kernel, offset, input_b); + writeln!(kernel, "}}").unwrap(); writeln!(kernel, "}}").unwrap(); } + writeln!(kernel, "}}").unwrap(); // move forward the vector offset writeln!(kernel, "vector_offset += 4 * {elements_per_block};").unwrap(); @@ -292,26 +348,27 @@ pub(crate) fn q4k_sgemv( ) .unwrap(); + // Initialize post element-wise functions once before the loop + let post_fns = + post_element_wise_functions.get_or_init(|| op.post_element_wise.add_functions(kernel)); + + // Fast path: all rows in tile are valid, no bounds checks needed + writeln!(kernel, "if is_full_tile {{").unwrap(); for offset in 0..Q4K_SGEMV_CHUNK_SIZE { - // If this is not the first simd thread in the workgroup, we can return early writeln!(kernel, "if {subgroup_local_index} == 0u {{").unwrap(); - { - // Write the output to the output tensor if this is the first thread in the workgroup - // Convert from f32 accumulator to output dtype - write!(kernel, "{output}[").unwrap(); - let index = format!("row + {offset}"); - let mut output_indices = Vec::new(); - // Add batch indices first - for dim in (0..output.rank()).rev().skip(2) { - output_indices.push(format!("batch_idx_{dim}")); - } - // Then add M and N indices - output_indices.push("m_idx".to_string()); - output_indices.push(index); - output.strided_index(kernel, output_indices); - let indexed = maybe_vec_storage_index(Q4K_SGEMV_CHUNK_SIZE, "sum", offset); - writeln!(kernel, "] = {dtype}({indexed});").unwrap(); - } + writeln!(kernel, "let row_index = row + {offset}u;").unwrap(); + generate_row_output(op, kernel, offset, dtype, bias, output, post_fns); + writeln!(kernel, "}}").unwrap(); + } + writeln!(kernel, "}} else {{").unwrap(); + // Slow path: check bounds for each row + for offset in 0..Q4K_SGEMV_CHUNK_SIZE { + writeln!(kernel, "if {subgroup_local_index} == 0u {{").unwrap(); + writeln!(kernel, "let row_index = row + {offset}u;").unwrap(); + writeln!(kernel, "if row_index < {n_size} {{").unwrap(); + generate_row_output(op, kernel, offset, dtype, bias, output, post_fns); + writeln!(kernel, "}}").unwrap(); writeln!(kernel, "}}").unwrap(); } + writeln!(kernel, "}}").unwrap(); } diff --git a/fusor-ml/core/src/quantized/matmul/sgemv/q5k.rs b/fusor-ml/core/src/quantized/matmul/sgemv/q5k.rs index c4129553d..ed710a77c 100644 --- a/fusor-ml/core/src/quantized/matmul/sgemv/q5k.rs +++ b/fusor-ml/core/src/quantized/matmul/sgemv/q5k.rs @@ -1,6 +1,7 @@ use crate::{ DataTypeEnum, mir::{ + function::Function, inputs::{QMatrixInput, TensorInput}, kernel::GenericKernel, workgroup_shape::WorkgroupShape, @@ -10,6 +11,7 @@ use crate::{ util::{maybe_vec_storage_index, maybe_vec_storage_subgroup_add, maybe_vec_storage_type}, }; use std::fmt::Write; +use std::sync::OnceLock; pub(crate) const Q5K_SGEMV_CHUNK_SIZE: u32 = 4; // This is the size of the chunk each thread will process at a time const SUBGROUP_COUNT: u32 = 2; @@ -18,6 +20,209 @@ const MASK1: u32 = 0b0011111100111111; const MASK2: u32 = 0b0000111100001111; const MASK3: u32 = 0b1100000011000000; +/// Generate WGSL for computing one row's contribution in the Q5K inner loop +fn generate_row_computation(kernel: &mut GenericKernel, offset: u32, input_b: &QMatrixInput) { + // Fetch and unpack the two sets of values from the cache + writeln!(kernel, "let first_values_offset = data_offset;").unwrap(); + writeln!(kernel, "let second_values_offset = data_offset + 16u;").unwrap(); + + // Keep track of the sum of each chunk + writeln!(kernel, "var first_sums = vec4();").unwrap(); + writeln!(kernel, "var second_sums = vec4();").unwrap(); + + // Load the qh values - we need 2 u32s to cover 8 bytes + writeln!( + kernel, + "let qh_lo = {input_b}[local_block_offset].qh[qh_u32_base];" + ) + .unwrap(); + writeln!( + kernel, + "let qh_hi = {input_b}[local_block_offset].qh[qh_u32_base + 1u];" + ) + .unwrap(); + + // Perform the dot product of the values and scales + for j in 0..2 { + // Process qs values - same structure as Q4K but with qh bit addition + for (sum, cache, values, qh_bit_low, qh_bit_high) in [ + ( + "first_sums", + "cached_a_low_values", + "first_values_offset", + "qh_bit_low_first", + "qh_bit_high_first", + ), + ( + "second_sums", + "cached_a_high_values", + "second_values_offset", + "qh_bit_low_second", + "qh_bit_high_second", + ), + ] { + writeln!( + kernel, + "let value_u32_{values}_{j} = {input_b}[local_block_offset].qs[{values} + {j}];" + ) + .unwrap(); + + writeln!( + kernel, + "let first_four_values_{values}_{j} = vec4({cache}[{j}*4 + 0], {cache}[{j}*4 + 1], {cache}[{j}*4 + 2], {cache}[{j}*4 + 3]);" + ) + .unwrap(); + writeln!( + kernel, + "let second_four_values_{values}_{j} = vec4({cache}[{j}*4 + 8], {cache}[{j}*4 + 9], {cache}[{j}*4 + 10], {cache}[{j}*4 + 11]);" + ) + .unwrap(); + + let qh_source = if j == 0 { "qh_lo" } else { "qh_hi" }; + + writeln!( + kernel, + "let qh_bytes_{values}_{j} = unpack4xU8({qh_source});" + ) + .unwrap(); + + writeln!( + kernel, + "let qh_unpacked_lo_{values}_{j} = vec4((qh_bytes_{values}_{j}.x >> {qh_bit_low}) & 1u, (qh_bytes_{values}_{j}.y >> {qh_bit_low}) & 1u, (qh_bytes_{values}_{j}.z >> {qh_bit_low}) & 1u, (qh_bytes_{values}_{j}.w >> {qh_bit_low}) & 1u) * 16u;" + ) + .unwrap(); + + writeln!( + kernel, + "let qh_unpacked_hi_{values}_{j} = vec4((qh_bytes_{values}_{j}.x >> {qh_bit_high}) & 1u, (qh_bytes_{values}_{j}.y >> {qh_bit_high}) & 1u, (qh_bytes_{values}_{j}.z >> {qh_bit_high}) & 1u, (qh_bytes_{values}_{j}.w >> {qh_bit_high}) & 1u) * 16u;" + ) + .unwrap(); + + writeln!( + kernel, + "{sum} += vec4(first_four_values_{values}_{j}.x * f32((value_u32_{values}_{j} & 0x000F) + qh_unpacked_lo_{values}_{j}.x), first_four_values_{values}_{j}.y * f32((value_u32_{values}_{j} & 0x0F00) + (qh_unpacked_lo_{values}_{j}.y << 8)), second_four_values_{values}_{j}.x * f32((value_u32_{values}_{j} & 0x00F0) + (qh_unpacked_hi_{values}_{j}.x << 4)), second_four_values_{values}_{j}.y * f32((value_u32_{values}_{j} & 0xF000) + (qh_unpacked_hi_{values}_{j}.y << 12)));" + ) + .unwrap(); + + let shift_right_16 = shift_right_scale(16); + writeln!( + kernel, + "{sum} += vec4(first_four_values_{values}_{j}.z * f32((value_u32_{values}_{j} & 0x000F0000) + (qh_unpacked_lo_{values}_{j}.z << 16)), first_four_values_{values}_{j}.w * f32((value_u32_{values}_{j} & 0x0F000000) + (qh_unpacked_lo_{values}_{j}.w << 24)), second_four_values_{values}_{j}.z * f32((value_u32_{values}_{j} & 0x00F00000) + (qh_unpacked_hi_{values}_{j}.z << 20)), second_four_values_{values}_{j}.w * (f32(value_u32_{values}_{j} & 0xF0000000) + f32(qh_unpacked_hi_{values}_{j}.w) * 268435456.0)) * f32({shift_right_16});" + ) + .unwrap(); + } + } + + // Load the block scale and min + writeln!( + kernel, + "let block_scale = f32({input_b}[local_block_offset].scale);" + ) + .unwrap(); + writeln!( + kernel, + "let block_min = f32({input_b}[local_block_offset].min);" + ) + .unwrap(); + // Load 8 scales into a cache (same as Q4K) + writeln!( + kernel, + "let first_32_scale_bits = {input_b}[local_block_offset].scales[0] >> (16 * scale_offset);" + ) + .unwrap(); + writeln!( + kernel, + "let second_32_scale_bits = {input_b}[local_block_offset].scales[1] >> (16 * scale_offset);" + ) + .unwrap(); + writeln!( + kernel, + "let third_32_scale_bits = {input_b}[local_block_offset].scales[2] >> (16 * scale_offset);" + ) + .unwrap(); + // Extract the scales from the bits into cached_scales (same as Q4K) + writeln!( + kernel, + "let first_two_scales = first_32_scale_bits & {MASK1};" + ) + .unwrap(); + writeln!( + kernel, + "let second_two_scales = second_32_scale_bits & {MASK1};" + ) + .unwrap(); + + writeln!(kernel, "let third_two_scales = ((third_32_scale_bits >> 0) & {MASK2}) | ((first_32_scale_bits & {MASK3}) >> 2);").unwrap(); + writeln!(kernel, "let fourth_two_scales = ((third_32_scale_bits >> 4) & {MASK2}) | ((second_32_scale_bits & {MASK3}) >> 2);").unwrap(); + + writeln!( + kernel, + "let odd_scales_unpacked = vec4(unpack4xU8(first_two_scales | (third_two_scales << 16)));" + ) + .unwrap(); + writeln!( + kernel, + "let even_scales_unpacked = vec4(unpack4xU8(second_two_scales | (fourth_two_scales << 16)));" + ) + .unwrap(); + + // Add the sums to the total sum + let indexed_sum = maybe_vec_storage_index(Q5K_SGEMV_CHUNK_SIZE, "sum", offset); + let shift_right_8 = shift_right_scale(8); + let shift_right_4 = shift_right_scale(4); + writeln!( + kernel, + "let small_shift_sums = vec4(first_sums[0], first_sums[2], second_sums[0], second_sums[2]);" + ) + .unwrap(); + writeln!( + kernel, + "let large_shift_sums = vec4(first_sums[1], first_sums[3], second_sums[1], second_sums[3]);" + ) + .unwrap(); + writeln!( + kernel, + "let shift_4 = vec4(1.0, {shift_right_4}, 1.0, {shift_right_4});" + ) + .unwrap(); + writeln!( + kernel, + r#"{indexed_sum} += block_scale * dot((small_shift_sums + f32({shift_right_8}) * large_shift_sums) * odd_scales_unpacked, shift_4) - + block_min * dot(vector_sum, even_scales_unpacked);"# + ) + .unwrap(); + // Move forward the block offset by one row + writeln!(kernel, "local_block_offset += k_block_size;").unwrap(); +} + +/// Generate WGSL for writing one row's output +fn generate_row_output( + op: &QMatMulOperation, + kernel: &mut GenericKernel, + offset: u32, + dtype: DataTypeEnum, + bias: Option<&TensorInput>, + output: &TensorInput, + post_element_wise_functions: &[Function], +) { + write!(kernel, "{output}[").unwrap(); + let mut output_indices = Vec::new(); + for dim in (0..output.rank()).rev().skip(2) { + output_indices.push(format!("batch_idx_{dim}")); + } + output_indices.push("m_idx".to_string()); + output_indices.push("row_index".to_string()); + output.strided_index(kernel, output_indices); + let indexed = maybe_vec_storage_index(Q5K_SGEMV_CHUNK_SIZE, "sum", offset); + let result = op.apply_bias_and_post( + bias, + "row_index", + format!("{dtype}({indexed})"), + post_element_wise_functions, + ); + writeln!(kernel, "] = {result};").unwrap(); +} + // Q5K matmul kernel, similar to Q4K but with additional high bits (qh) // Q5K structure: scale (f16), min (f16), scales [12], qh [32], qs [128] // - Block has 256 elements in 4 chunks of 64 elements each @@ -39,6 +244,7 @@ pub(crate) fn q5k_sgemv( workgroup_shape: &WorkgroupShape, input_a: &TensorInput, input_b: &QMatrixInput, + bias: Option<&TensorInput>, output: &TensorInput, n_size: &str, m_size: &str, @@ -48,6 +254,8 @@ pub(crate) fn q5k_sgemv( let subgroup_index = kernel.subgroup_index(); let subgroup_local_index = kernel.subgroup_local_index(); let elements_per_block = op.elements_per_block(); + let pre_element_wise_functions = OnceLock::new(); + let post_element_wise_functions = OnceLock::new(); // Calculate n_workgroups for this kernel type (SUBGROUP_COUNT subgroups per workgroup, Q5K_SGEMV_CHUNK_SIZE per subgroup) let chunk_size = Q5K_SGEMV_CHUNK_SIZE * SUBGROUP_COUNT; @@ -87,6 +295,13 @@ pub(crate) fn q5k_sgemv( ) .unwrap(); + // Fast path check: if all rows in this tile are valid, skip per-row bounds checks + writeln!( + kernel, + "let is_full_tile = row + {Q5K_SGEMV_CHUNK_SIZE} <= {n_size};" + ) + .unwrap(); + writeln!(kernel, "let thread_id = {subgroup_local_index} >> 3;").unwrap(); writeln!(kernel, "let thread_local_id = {subgroup_local_index} & 7;").unwrap(); writeln!(kernel, "let half_subgroup_id = thread_local_id >> 2;").unwrap(); @@ -121,7 +336,10 @@ pub(crate) fn q5k_sgemv( // - 128: chunk 2/3 low nibbles, first 8 (second 128-element half) // - 160: chunk 2/3 high nibbles, first 8 (second 128-element half) for (idx, offset) in [(0, 0), (1, 32), (2, 128), (3, 160)] { - write!(kernel, "let a_val_{j}_{idx} = {input_a}[").unwrap(); + let pre_element_wise_functions = pre_element_wise_functions + .get_or_init(|| op.pre_element_wise.add_functions(kernel)); + let mut raw_input = String::new(); + write!(&mut raw_input, "{input_a}[").unwrap(); let mut indices = Vec::new(); // Add batch indices first for dim in (0..input_a.rank()).rev().skip(2) { @@ -130,8 +348,12 @@ pub(crate) fn q5k_sgemv( // Then add M and K indices indices.push("m_idx".to_string()); indices.push(format!("vector_offset + {j} + {offset}")); - input_a.strided_index(kernel, indices); - writeln!(kernel, "];").unwrap(); + input_a.strided_index(&mut raw_input, indices); + write!(&mut raw_input, "]").unwrap(); + let processed = pre_element_wise_functions + .iter() + .fold(raw_input, |acc, f| f.call(vec![acc])); + writeln!(kernel, "let a_val_{j}_{idx} = f32({processed});").unwrap(); } writeln!( @@ -174,203 +396,24 @@ pub(crate) fn q5k_sgemv( writeln!(kernel, "var local_block_offset = block_offset + i;").unwrap(); + // Fast path: all rows in tile are valid, no bounds checks needed + writeln!(kernel, "if is_full_tile {{").unwrap(); for offset in 0..Q5K_SGEMV_CHUNK_SIZE { writeln!(kernel, "{{").unwrap(); - // Fetch and unpack the two sets of values from the cache - writeln!(kernel, "let first_values_offset = data_offset;").unwrap(); - writeln!(kernel, "let second_values_offset = data_offset + 16u;").unwrap(); - - // Keep track of the sum of each chunk - writeln!(kernel, "var first_sums = vec4();").unwrap(); - writeln!(kernel, "var second_sums = vec4();").unwrap(); - - // Load the qh values - we need 2 u32s to cover 8 bytes - writeln!( - kernel, - "let qh_lo = {input_b}[local_block_offset].qh[qh_u32_base];" - ) - .unwrap(); - writeln!( - kernel, - "let qh_hi = {input_b}[local_block_offset].qh[qh_u32_base + 1u];" - ) - .unwrap(); - - // Perform the dot product of the values and scales - for j in 0..2 { - // Process qs values - same structure as Q4K but with qh bit addition - for (sum, cache, values, qh_bit_low, qh_bit_high) in [ - ( - "first_sums", - "cached_a_low_values", - "first_values_offset", - "qh_bit_low_first", - "qh_bit_high_first", - ), - ( - "second_sums", - "cached_a_high_values", - "second_values_offset", - "qh_bit_low_second", - "qh_bit_high_second", - ), - ] { - writeln!( - kernel, - "let value_u32_{values}_{j} = {input_b}[local_block_offset].qs[{values} + {j}];" - ) - .unwrap(); - - writeln!( - kernel, - "let first_four_values_{values}_{j} = vec4({cache}[{j}*4 + 0], {cache}[{j}*4 + 1], {cache}[{j}*4 + 2], {cache}[{j}*4 + 3]);" - ) - .unwrap(); - writeln!( - kernel, - "let second_four_values_{values}_{j} = vec4({cache}[{j}*4 + 8], {cache}[{j}*4 + 9], {cache}[{j}*4 + 10], {cache}[{j}*4 + 11]);" - ) - .unwrap(); - - // Extract qh bits for the 8 bytes we're processing - // qh_lo covers bytes 0-3 (elements 0-3 of this thread's 8-element group) - // qh_hi covers bytes 4-7 (elements 4-7) - // For j=0: we need bytes 0-3 (qh_lo) - // For j=1: we need bytes 4-7 (qh_hi) - let qh_source = if j == 0 { "qh_lo" } else { "qh_hi" }; - - // Extract 4 qh bits for low nibbles and high nibbles - // Each byte in qh contains 8 bits for 8 different elements at different chunks - // We need to unpack the bytes first, then extract the specific bit from each byte - // (Can't just shift the u32 because bits from higher bytes would spill into lower bytes) - writeln!( - kernel, - "let qh_bytes_{values}_{j} = unpack4xU8({qh_source});" - ) - .unwrap(); - - // Extract bit qh_bit_low from each of the 4 bytes for low nibbles - writeln!( - kernel, - "let qh_unpacked_lo_{values}_{j} = vec4((qh_bytes_{values}_{j}.x >> {qh_bit_low}) & 1u, (qh_bytes_{values}_{j}.y >> {qh_bit_low}) & 1u, (qh_bytes_{values}_{j}.z >> {qh_bit_low}) & 1u, (qh_bytes_{values}_{j}.w >> {qh_bit_low}) & 1u) * 16u;" - ) - .unwrap(); - - // Extract bit qh_bit_high from each of the 4 bytes for high nibbles - writeln!( - kernel, - "let qh_unpacked_hi_{values}_{j} = vec4((qh_bytes_{values}_{j}.x >> {qh_bit_high}) & 1u, (qh_bytes_{values}_{j}.y >> {qh_bit_high}) & 1u, (qh_bytes_{values}_{j}.z >> {qh_bit_high}) & 1u, (qh_bytes_{values}_{j}.w >> {qh_bit_high}) & 1u) * 16u;" - ) - .unwrap(); - - // Add qs values with qh high bits - // Low nibbles of qs: positions 0,8,16,24 in u32 (bytes 0-3) - // High nibbles of qs: positions 4,12,20,28 in u32 (bytes 0-3, upper nibble) - writeln!( - kernel, - "{sum} += vec4(first_four_values_{values}_{j}.x * f32((value_u32_{values}_{j} & 0x000F) + qh_unpacked_lo_{values}_{j}.x), first_four_values_{values}_{j}.y * f32((value_u32_{values}_{j} & 0x0F00) + (qh_unpacked_lo_{values}_{j}.y << 8)), second_four_values_{values}_{j}.x * f32((value_u32_{values}_{j} & 0x00F0) + (qh_unpacked_hi_{values}_{j}.x << 4)), second_four_values_{values}_{j}.y * f32((value_u32_{values}_{j} & 0xF000) + (qh_unpacked_hi_{values}_{j}.y << 12)));" - ) - .unwrap(); - - // Keep values in their bit positions and shift qh UP to match, then normalize with shift_right_16 - // Bit positions: 0x000F0000 (16), 0x0F000000 (24), 0x00F00000 (20), 0xF0000000 (28) - // NOTE: For the << 28 shift, we compute in f32 to avoid u32 overflow (16 << 28 = 2^32 overflows) - let shift_right_16 = shift_right_scale(16); - writeln!( - kernel, - "{sum} += vec4(first_four_values_{values}_{j}.z * f32((value_u32_{values}_{j} & 0x000F0000) + (qh_unpacked_lo_{values}_{j}.z << 16)), first_four_values_{values}_{j}.w * f32((value_u32_{values}_{j} & 0x0F000000) + (qh_unpacked_lo_{values}_{j}.w << 24)), second_four_values_{values}_{j}.z * f32((value_u32_{values}_{j} & 0x00F00000) + (qh_unpacked_hi_{values}_{j}.z << 20)), second_four_values_{values}_{j}.w * (f32(value_u32_{values}_{j} & 0xF0000000) + f32(qh_unpacked_hi_{values}_{j}.w) * 268435456.0)) * f32({shift_right_16});" - ) - .unwrap(); - } - } - - // Load the block scale and min - writeln!( - kernel, - "let block_scale = f32({input_b}[local_block_offset].scale);" - ) - .unwrap(); - writeln!( - kernel, - "let block_min = f32({input_b}[local_block_offset].min);" - ) - .unwrap(); - // Load 8 scales into a cache (same as Q4K) - writeln!( - kernel, - "let first_32_scale_bits = {input_b}[local_block_offset].scales[0] >> (16 * scale_offset);" - ) - .unwrap(); - writeln!( - kernel, - "let second_32_scale_bits = {input_b}[local_block_offset].scales[1] >> (16 * scale_offset);" - ) - .unwrap(); - writeln!( - kernel, - "let third_32_scale_bits = {input_b}[local_block_offset].scales[2] >> (16 * scale_offset);" - ) - .unwrap(); - // Extract the scales from the bits into cached_scales (same as Q4K) - writeln!( - kernel, - "let first_two_scales = first_32_scale_bits & {MASK1};" - ) - .unwrap(); - writeln!( - kernel, - "let second_two_scales = second_32_scale_bits & {MASK1};" - ) - .unwrap(); - - writeln!(kernel, "let third_two_scales = ((third_32_scale_bits >> 0) & {MASK2}) | ((first_32_scale_bits & {MASK3}) >> 2);").unwrap(); - writeln!(kernel, "let fourth_two_scales = ((third_32_scale_bits >> 4) & {MASK2}) | ((second_32_scale_bits & {MASK3}) >> 2);").unwrap(); - - writeln!( - kernel, - "let odd_scales_unpacked = vec4(unpack4xU8(first_two_scales | (third_two_scales << 16)));" - ) - .unwrap(); - writeln!( - kernel, - "let even_scales_unpacked = vec4(unpack4xU8(second_two_scales | (fourth_two_scales << 16)));" - ) - .unwrap(); - - // Add the sums to the total sum - let indexed_sum = maybe_vec_storage_index(Q5K_SGEMV_CHUNK_SIZE, "sum", offset); - // *_sums[0] needs to be shifted by 0 bits - // *_sums[1] needs to be shifted by 8 bits - // *_sums[2] needs to be shifted by 4 bits - // *_sums[3] needs to be shifted by 12 bits - let shift_right_8 = shift_right_scale(8); - let shift_right_4 = shift_right_scale(4); - writeln!( - kernel, - "let small_shift_sums = vec4(first_sums[0], first_sums[2], second_sums[0], second_sums[2]);" - ) - .unwrap(); - writeln!( - kernel, - "let large_shift_sums = vec4(first_sums[1], first_sums[3], second_sums[1], second_sums[3]);" - ) - .unwrap(); - writeln!( - kernel, - "let shift_4 = vec4(1.0, {shift_right_4}, 1.0, {shift_right_4});" - ) - .unwrap(); - // Add the final weighted sum (same structure as Q4K - without the Q5K -16 offset for now) - writeln!( - kernel, - r#"{indexed_sum} += block_scale * dot((small_shift_sums + f32({shift_right_8}) * large_shift_sums) * odd_scales_unpacked, shift_4) - - block_min * dot(vector_sum, even_scales_unpacked);"# - ) - .unwrap(); - // Move forward the block offset by one row - writeln!(kernel, "local_block_offset += k_block_size;").unwrap(); + generate_row_computation(kernel, offset, input_b); + writeln!(kernel, "}}").unwrap(); + } + writeln!(kernel, "}} else {{").unwrap(); + // Slow path: check bounds for each row + for offset in 0..Q5K_SGEMV_CHUNK_SIZE { + writeln!(kernel, "{{").unwrap(); + writeln!(kernel, "let row_index = row + {offset};").unwrap(); + writeln!(kernel, "if row_index < {n_size} {{").unwrap(); + generate_row_computation(kernel, offset, input_b); + writeln!(kernel, "}}").unwrap(); writeln!(kernel, "}}").unwrap(); } + writeln!(kernel, "}}").unwrap(); // move forward the vector offset writeln!(kernel, "vector_offset += 4 * {elements_per_block};").unwrap(); @@ -385,26 +428,31 @@ pub(crate) fn q5k_sgemv( ) .unwrap(); + // Initialize post element-wise functions once before the loop + let post_fns = + post_element_wise_functions.get_or_init(|| op.post_element_wise.add_functions(kernel)); + + // Fast path: all rows in tile are valid, no bounds checks needed + writeln!(kernel, "if is_full_tile {{").unwrap(); for offset in 0..Q5K_SGEMV_CHUNK_SIZE { - // If this is not the first simd thread in the workgroup, we can return early + writeln!(kernel, "{{").unwrap(); writeln!(kernel, "if {subgroup_local_index} == 0u {{").unwrap(); - { - // Write the output to the output tensor if this is the first thread in the workgroup - // Convert from f32 accumulator to output dtype - write!(kernel, "{output}[").unwrap(); - let index = format!("row + {offset}"); - let mut output_indices = Vec::new(); - // Add batch indices first - for dim in (0..output.rank()).rev().skip(2) { - output_indices.push(format!("batch_idx_{dim}")); - } - // Then add M and N indices - output_indices.push("m_idx".to_string()); - output_indices.push(index); - output.strided_index(kernel, output_indices); - let indexed = maybe_vec_storage_index(Q5K_SGEMV_CHUNK_SIZE, "sum", offset); - writeln!(kernel, "] = {dtype}({indexed});").unwrap(); - } + writeln!(kernel, "let row_index = row + {offset}u;").unwrap(); + generate_row_output(op, kernel, offset, dtype, bias, output, post_fns); + writeln!(kernel, "}}").unwrap(); writeln!(kernel, "}}").unwrap(); } + writeln!(kernel, "}} else {{").unwrap(); + // Slow path: check bounds for each row + for offset in 0..Q5K_SGEMV_CHUNK_SIZE { + writeln!(kernel, "{{").unwrap(); + writeln!(kernel, "if {subgroup_local_index} == 0u {{").unwrap(); + writeln!(kernel, "let row_index = row + {offset}u;").unwrap(); + writeln!(kernel, "if row_index < {n_size} {{").unwrap(); + generate_row_output(op, kernel, offset, dtype, bias, output, post_fns); + writeln!(kernel, "}}").unwrap(); + writeln!(kernel, "}}").unwrap(); + writeln!(kernel, "}}").unwrap(); + } + writeln!(kernel, "}}").unwrap(); } diff --git a/fusor-ml/core/src/quantized/matmul/sgemv/q6k.rs b/fusor-ml/core/src/quantized/matmul/sgemv/q6k.rs index 699596376..943e48b61 100644 --- a/fusor-ml/core/src/quantized/matmul/sgemv/q6k.rs +++ b/fusor-ml/core/src/quantized/matmul/sgemv/q6k.rs @@ -9,6 +9,7 @@ use crate::{ util::{maybe_vec_storage_index, maybe_vec_storage_subgroup_add, maybe_vec_storage_type}, }; use std::fmt::Write; +use std::sync::OnceLock; pub(crate) const Q6K_SGEMV_CHUNK_SIZE: u32 = 2; // This is the size of the chunk each thread will process at a time const PRELOAD: bool = false; @@ -21,6 +22,7 @@ pub(crate) fn q6k_sgemv( workgroup_shape: &WorkgroupShape, input_a: &TensorInput, input_b: &QMatrixInput, + bias: Option<&TensorInput>, output: &TensorInput, n_size: &str, m_size: &str, @@ -30,6 +32,8 @@ pub(crate) fn q6k_sgemv( let subgroup_index = kernel.subgroup_index(); let subgroup_local_index = kernel.subgroup_local_index(); let elements_per_block = op.elements_per_block(); + let pre_element_wise_functions = OnceLock::new(); + let post_element_wise_functions = OnceLock::new(); // Calculate n_workgroups for this kernel type (2 subgroups per workgroup, Q6K_SGEMV_CHUNK_SIZE per subgroup) let n_workgroups = format!( @@ -72,6 +76,13 @@ pub(crate) fn q6k_sgemv( ) .unwrap(); + // Fast path check: if all rows in this tile are valid, skip per-row bounds checks + writeln!( + kernel, + "let is_full_tile = row + {Q6K_SGEMV_CHUNK_SIZE} <= {n_size};" + ) + .unwrap(); + writeln!(kernel, "let block_offset = row * k_block_size;").unwrap(); writeln!(kernel, "let thread_id = {subgroup_local_index} / 2;").unwrap(); @@ -128,7 +139,10 @@ pub(crate) fn q6k_sgemv( ) .unwrap(); let load_value = |kernel: &mut GenericKernel, j: &str, offset: u32| { - write!(kernel, "f32({input_a}[").unwrap(); + let pre_element_wise_functions = pre_element_wise_functions + .get_or_init(|| op.pre_element_wise.add_functions(kernel)); + let mut raw_input = String::new(); + write!(&mut raw_input, "{input_a}[").unwrap(); let mut indices = Vec::new(); // Add batch indices first for dim in (0..input_a.rank()).rev().skip(2) { @@ -137,8 +151,12 @@ pub(crate) fn q6k_sgemv( // Then add M and K indices indices.push("m_idx".to_string()); indices.push(format!("{j} + vector_offset + {}", offset * 32)); - input_a.strided_index(kernel, indices); - write!(kernel, "])").unwrap(); + input_a.strided_index(&mut raw_input, indices); + write!(&mut raw_input, "]").unwrap(); + let processed = pre_element_wise_functions + .iter() + .fold(raw_input, |acc, f| f.call(vec![acc])); + write!(kernel, "f32({processed})").unwrap(); }; if PRELOAD { writeln!(kernel, "for (var j = 0u; j < 4; j += 1u) {{").unwrap(); @@ -150,6 +168,90 @@ pub(crate) fn q6k_sgemv( writeln!(kernel, "}}").unwrap(); } + // Generate the computation body for one row + let generate_row_computation = + |kernel: &mut GenericKernel, load_value: &dyn Fn(&mut GenericKernel, &str, u32)| { + if Q6K_SGEMV_CHUNK_SIZE > 1 { + writeln!( + kernel, + "let local_block_offset = i + block_offset + offset * k_block_size;" + ) + .unwrap(); + } else { + writeln!(kernel, "let local_block_offset = i + block_offset;").unwrap(); + } + writeln!(kernel, "let low_offset_1 = q_offset_l;").unwrap(); + writeln!( + kernel, + "let low_bytes_1 = unpack4xU8({input_b}[local_block_offset].data_low_bits[low_offset_1]);" + ) + .unwrap(); + writeln!(kernel, "let low_offset_2 = q_offset_l + 8;").unwrap(); + writeln!(kernel, "let low_bytes_2 = unpack4xU8({input_b}[local_block_offset].data_low_bits[low_offset_2]);").unwrap(); + writeln!(kernel, "let high_offset = q_offset_h;").unwrap(); + writeln!( + kernel, + "let high_bytes = unpack4xU8({input_b}[local_block_offset].data_high_bits[high_offset]);" + ) + .unwrap(); + writeln!(kernel, "let scale_offset = scale_index_offset;").unwrap(); + writeln!( + kernel, + "let scale_chunk_1 = unpack4xI8({input_b}[local_block_offset].scales[scale_offset]);" + ) + .unwrap(); + writeln!( + kernel, + "let scale_chunk_2 = unpack4xI8({input_b}[local_block_offset].scales[scale_offset + 1]);" + ) + .unwrap(); + writeln!( + kernel, + "let scales = vec4(f32(scale_chunk_1[scale_pair_offset]), f32(scale_chunk_1[2 + scale_pair_offset]), f32(scale_chunk_2[scale_pair_offset]), f32(scale_chunk_2[2 + scale_pair_offset]));" + ) + .unwrap(); + + writeln!( + kernel, + "let scale = f32({input_b}[local_block_offset].scale);" + ) + .unwrap(); + + writeln!(kernel, "var sums = vec4();").unwrap(); + writeln!(kernel, "for (var j = 0u; j < 4u; j += 1u) {{").unwrap(); + { + let first_four_bytes = 0b00001111u8; + let first_two_bytes = 0b00000011u8; + let second_two_bytes = 0b00001100u8; + let third_two_bytes = 0b00110000u8; + let fourth_two_bytes = 0b11000000u8; + let get_value = |kernel: &mut GenericKernel, j: &str, offset: u32| { + if PRELOAD { + write!(kernel, "cached_a_values[{j} + {offset}]").unwrap(); + } else { + load_value(kernel, j, offset); + } + }; + write!(kernel, "sums[0] += ").unwrap(); + get_value(kernel, "j", 0); + writeln!(kernel,"* f32(i32((low_bytes_1[j] & {first_four_bytes}) | ((high_bytes[j] & {first_two_bytes}) << 4)) - 32);").unwrap(); + write!(kernel, "sums[1] += ").unwrap(); + get_value(kernel, "j", 1); + writeln!(kernel,"* f32(i32((low_bytes_2[j] & {first_four_bytes}) | ((high_bytes[j] & {second_two_bytes}) << 2)) - 32);").unwrap(); + write!(kernel, "sums[2] += ").unwrap(); + get_value(kernel, "j", 2); + writeln!(kernel,"* f32(i32((low_bytes_1[j] >> 4) | ((high_bytes[j] & {third_two_bytes}) << 0)) - 32);").unwrap(); + write!(kernel, "sums[3] += ").unwrap(); + get_value(kernel, "j", 3); + writeln!(kernel,"* f32(i32((low_bytes_2[j] >> 4) | ((high_bytes[j] & {fourth_two_bytes}) >> 2)) - 32);").unwrap(); + } + writeln!(kernel, "}}").unwrap(); + let indexed = maybe_vec_storage_index(Q6K_SGEMV_CHUNK_SIZE, "sum", "offset"); + writeln!(kernel, "{indexed} += scale * dot(sums, scales);").unwrap(); + }; + + // Fast path: all rows in tile are valid, no bounds checks needed + writeln!(kernel, "if is_full_tile {{").unwrap(); if Q6K_SGEMV_CHUNK_SIZE > 1 { writeln!( kernel, @@ -158,87 +260,41 @@ pub(crate) fn q6k_sgemv( .unwrap(); } { - if Q6K_SGEMV_CHUNK_SIZE > 1 { - writeln!( - kernel, - "let local_block_offset = i + block_offset + offset * k_block_size;" - ) - .unwrap(); + let row_index = if Q6K_SGEMV_CHUNK_SIZE > 1 { + "row + offset" } else { - writeln!(kernel, "let local_block_offset = i + block_offset;").unwrap(); - } - writeln!(kernel, "let low_offset_1 = q_offset_l;").unwrap(); - writeln!( - kernel, - "let low_bytes_1 = unpack4xU8({input_b}[local_block_offset].data_low_bits[low_offset_1]);" - ) - .unwrap(); - writeln!(kernel, "let low_offset_2 = q_offset_l + 8;").unwrap(); - writeln!(kernel, "let low_bytes_2 = unpack4xU8({input_b}[local_block_offset].data_low_bits[low_offset_2]);").unwrap(); - writeln!(kernel, "let high_offset = q_offset_h;").unwrap(); - writeln!( - kernel, - "let high_bytes = unpack4xU8({input_b}[local_block_offset].data_high_bits[high_offset]);" - ) - .unwrap(); - writeln!(kernel, "let scale_offset = scale_index_offset;").unwrap(); - writeln!( - kernel, - "let scale_chunk_1 = unpack4xI8({input_b}[local_block_offset].scales[scale_offset]);" - ) - .unwrap(); - writeln!( - kernel, - "let scale_chunk_2 = unpack4xI8({input_b}[local_block_offset].scales[scale_offset + 1]);" - ) - .unwrap(); - writeln!( - kernel, - "let scales = vec4(f32(scale_chunk_1[scale_pair_offset]), f32(scale_chunk_1[2 + scale_pair_offset]), f32(scale_chunk_2[scale_pair_offset]), f32(scale_chunk_2[2 + scale_pair_offset]));" - ) - .unwrap(); - + "row" + }; + writeln!(kernel, "let row_index = {row_index};").unwrap(); + generate_row_computation(kernel, &load_value); + } + if Q6K_SGEMV_CHUNK_SIZE > 1 { + writeln!(kernel, "}}").unwrap(); + } + writeln!(kernel, "}} else {{").unwrap(); + // Slow path: check bounds for each row + if Q6K_SGEMV_CHUNK_SIZE > 1 { writeln!( kernel, - "let scale = f32({input_b}[local_block_offset].scale);" + "for (var offset = 0u; offset < {Q6K_SGEMV_CHUNK_SIZE}; offset += 1u) {{" ) .unwrap(); - - writeln!(kernel, "var sums = vec4();").unwrap(); - writeln!(kernel, "for (var j = 0u; j < 4u; j += 1u) {{").unwrap(); - { - let first_four_bytes = 0b00001111u8; - let first_two_bytes = 0b00000011u8; - let second_two_bytes = 0b00001100u8; - let third_two_bytes = 0b00110000u8; - let fourth_two_bytes = 0b11000000u8; - let get_value = |kernel: &mut GenericKernel, j: &str, offset: u32| { - if PRELOAD { - write!(kernel, "cached_a_values[{j} + {offset}]").unwrap(); - } else { - load_value(kernel, j, offset); - } - }; - write!(kernel, "sums[0] += ").unwrap(); - get_value(kernel, "j", 0); - writeln!(kernel,"* f32(i32((low_bytes_1[j] & {first_four_bytes}) | ((high_bytes[j] & {first_two_bytes}) << 4)) - 32);").unwrap(); - write!(kernel, "sums[1] += ").unwrap(); - get_value(kernel, "j", 1); - writeln!(kernel,"* f32(i32((low_bytes_2[j] & {first_four_bytes}) | ((high_bytes[j] & {second_two_bytes}) << 2)) - 32);").unwrap(); - write!(kernel, "sums[2] += ").unwrap(); - get_value(kernel, "j", 2); - writeln!(kernel,"* f32(i32((low_bytes_1[j] >> 4) | ((high_bytes[j] & {third_two_bytes}) << 0)) - 32);").unwrap(); - write!(kernel, "sums[3] += ").unwrap(); - get_value(kernel, "j", 3); - writeln!(kernel,"* f32(i32((low_bytes_2[j] >> 4) | ((high_bytes[j] & {fourth_two_bytes}) >> 2)) - 32);").unwrap(); - } + } + { + let row_index = if Q6K_SGEMV_CHUNK_SIZE > 1 { + "row + offset" + } else { + "row" + }; + writeln!(kernel, "let row_index = {row_index};").unwrap(); + writeln!(kernel, "if row_index < {n_size} {{").unwrap(); + generate_row_computation(kernel, &load_value); writeln!(kernel, "}}").unwrap(); - let indexed = maybe_vec_storage_index(Q6K_SGEMV_CHUNK_SIZE, "sum", "offset"); - writeln!(kernel, "{indexed} += scale * dot(sums, scales);").unwrap(); } if Q6K_SGEMV_CHUNK_SIZE > 1 { writeln!(kernel, "}}").unwrap(); } + writeln!(kernel, "}}").unwrap(); } writeln!(kernel, "}}").unwrap(); @@ -253,6 +309,30 @@ pub(crate) fn q6k_sgemv( // If this is not the first simd thread in the workgroup, we can return early writeln!(kernel, "if {subgroup_local_index} != 0u {{ return; }}").unwrap(); + // Initialize post element-wise functions once before the loop + let post_fns = + post_element_wise_functions.get_or_init(|| op.post_element_wise.add_functions(kernel)); + + // Generate the output body for one row + let generate_row_output = |kernel: &mut GenericKernel| { + write!(kernel, "{output}[").unwrap(); + let mut output_indices = Vec::new(); + // Add batch indices first + for dim in (0..output.rank()).rev().skip(2) { + output_indices.push(format!("batch_idx_{dim}")); + } + // Then add M and N indices + output_indices.push("m_idx".to_string()); + output_indices.push("row_index".to_string()); + output.strided_index(kernel, output_indices); + let indexed = maybe_vec_storage_index(Q6K_SGEMV_CHUNK_SIZE, "sum", "offset"); + let result = + op.apply_bias_and_post(bias, "row_index", format!("{dtype}({indexed})"), post_fns); + writeln!(kernel, "] = {result};").unwrap(); + }; + + // Fast path: all rows in tile are valid, no bounds checks needed + writeln!(kernel, "if is_full_tile {{").unwrap(); if Q6K_SGEMV_CHUNK_SIZE > 1 { writeln!( kernel, @@ -261,26 +341,39 @@ pub(crate) fn q6k_sgemv( .unwrap(); } { - // Write the output to the output tensor if this is the first thread in the workgroup - write!(kernel, "{output}[").unwrap(); let index = if Q6K_SGEMV_CHUNK_SIZE > 1 { "row + offset".to_string() } else { "row".to_string() }; - let mut output_indices = Vec::new(); - // Add batch indices first - for dim in (0..output.rank()).rev().skip(2) { - output_indices.push(format!("batch_idx_{dim}")); - } - // Then add M and N indices - output_indices.push("m_idx".to_string()); - output_indices.push(index); - output.strided_index(kernel, output_indices); - let indexed = maybe_vec_storage_index(Q6K_SGEMV_CHUNK_SIZE, "sum", "offset"); - writeln!(kernel, "] = {dtype}({indexed});").unwrap(); + writeln!(kernel, "let row_index = {index};").unwrap(); + generate_row_output(kernel); } if Q6K_SGEMV_CHUNK_SIZE > 1 { writeln!(kernel, "}}").unwrap(); } + writeln!(kernel, "}} else {{").unwrap(); + // Slow path: check bounds for each row + if Q6K_SGEMV_CHUNK_SIZE > 1 { + writeln!( + kernel, + "for (var offset = 0u; offset < {Q6K_SGEMV_CHUNK_SIZE}; offset += 1u) {{" + ) + .unwrap(); + } + { + let index = if Q6K_SGEMV_CHUNK_SIZE > 1 { + "row + offset".to_string() + } else { + "row".to_string() + }; + writeln!(kernel, "let row_index = {index};").unwrap(); + writeln!(kernel, "if row_index < {n_size} {{").unwrap(); + generate_row_output(kernel); + writeln!(kernel, "}}").unwrap(); + } + if Q6K_SGEMV_CHUNK_SIZE > 1 { + writeln!(kernel, "}}").unwrap(); + } + writeln!(kernel, "}}").unwrap(); } diff --git a/fusor-ml/core/src/quantized/matmul/sgemv/q_8_0.rs b/fusor-ml/core/src/quantized/matmul/sgemv/q_8_0.rs index d831fe875..c84a32028 100644 --- a/fusor-ml/core/src/quantized/matmul/sgemv/q_8_0.rs +++ b/fusor-ml/core/src/quantized/matmul/sgemv/q_8_0.rs @@ -1,4 +1,5 @@ use crate::{ + DataTypeEnum, mir::{ inputs::{QMatrixInput, TensorInput}, kernel::GenericKernel, @@ -8,12 +9,65 @@ use crate::{ util::{maybe_vec_storage_index, maybe_vec_storage_subgroup_add, maybe_vec_storage_type}, }; use std::fmt::Write; +use std::sync::OnceLock; pub(crate) const Q_8_0_SGEMV_CHUNK_SIZE: u32 = 4; // This is the size of the chunk each thread will process at a time const STEP_SIZE: u32 = 8; const SUBGROUP_COUNT: u32 = 2; const SUBGROUP_SIZE: u32 = 32; +/// Generate WGSL for computing one row's contribution in the inner loop +fn generate_row_computation( + kernel: &mut GenericKernel, + offset: u32, + dtype: DataTypeEnum, + input_b: &QMatrixInput, +) { + writeln!(kernel, "var local_sum = {dtype}();").unwrap(); + for data_offset in 0..(STEP_SIZE / 4) { + writeln!(kernel, "{{").unwrap(); + writeln!(kernel, "let block = vec4<{dtype}>(unpack4xI8({input_b}[block_offset].data[thread_local_id * 2u + {data_offset}]));").unwrap(); + writeln!(kernel, "let float_block = vec4<{dtype}>(cached_a_values[{data_offset} * 4u + 0], cached_a_values[{data_offset} * 4u + 1], cached_a_values[{data_offset} * 4u + 2], cached_a_values[{data_offset} * 4u + 3]);").unwrap(); + writeln!(kernel, "local_sum += dot(block, float_block);").unwrap(); + writeln!(kernel, "}}").unwrap(); + } + let indexed = maybe_vec_storage_index(Q_8_0_SGEMV_CHUNK_SIZE, "sum", offset); + writeln!( + kernel, + "{indexed} += local_sum * {dtype}({input_b}[block_offset].scale);" + ) + .unwrap(); + writeln!(kernel, "block_offset += k_block_size;").unwrap(); +} + +/// Generate WGSL for writing one row's output +fn generate_row_output( + op: &QMatMulOperation, + kernel: &mut GenericKernel, + offset: u32, + input_datatype: DataTypeEnum, + bias: Option<&TensorInput>, + output: &TensorInput, + post_element_wise_functions: &[crate::mir::function::Function], +) { + write!(kernel, "{output}[").unwrap(); + let mut output_indices = Vec::new(); + for dim in (0..output.rank()).rev().skip(2) { + output_indices.push(format!("batch_idx_{dim}")); + } + output_indices.push("m_idx".to_string()); + output_indices.push("row_index".to_string()); + output.strided_index(kernel, output_indices); + let indexed = maybe_vec_storage_index(Q_8_0_SGEMV_CHUNK_SIZE, "sum", offset); + let result = op.apply_bias_and_post( + bias, + "row_index", + format!("{input_datatype}({indexed})"), + post_element_wise_functions, + ); + writeln!(kernel, "] = {result};").unwrap(); +} + // https://github.com/ggml-org/llama.cpp/blob/6efcd65945a98cf6883cdd9de4c8ccd8c79d219a/ggml/src/ggml-metal/ggml-metal.metal#L2452 #[allow(clippy::too_many_arguments)] pub(crate) fn q_8_0_sgemv( @@ -22,15 +76,19 @@ pub(crate) fn q_8_0_sgemv( workgroup_shape: &WorkgroupShape, input_a: &TensorInput, input_b: &QMatrixInput, + bias: Option<&TensorInput>, output: &TensorInput, n_size: &str, m_size: &str, k_size: &str, ) { - let dtype = op.input_datatype; + let input_datatype = op.input_datatype; + let dtype = op.matmul_datatype(); let subgroup_index = kernel.subgroup_index(); let subgroup_local_index = kernel.subgroup_local_index(); let elements_per_block = op.elements_per_block(); + let pre_element_wise_functions = OnceLock::new(); + let post_element_wise_functions = OnceLock::new(); // Calculate n_workgroups for this kernel type (SUBGROUP_COUNT subgroups per workgroup, Q_8_0_SGEMV_CHUNK_SIZE per subgroup) let chunk_size = Q_8_0_SGEMV_CHUNK_SIZE * SUBGROUP_COUNT; @@ -70,6 +128,13 @@ pub(crate) fn q_8_0_sgemv( ) .unwrap(); + // Fast path check: if all rows in this tile are valid, skip per-row bounds checks + writeln!( + kernel, + "let is_full_tile = row + {Q_8_0_SGEMV_CHUNK_SIZE} <= {n_size};" + ) + .unwrap(); + writeln!(kernel, "let row_block_offset = row * k_block_size;").unwrap(); writeln!(kernel, "let thread_id = {subgroup_local_index} / 4;").unwrap(); @@ -102,8 +167,10 @@ pub(crate) fn q_8_0_sgemv( // First load the values of a into cached_a_values for j in 0..STEP_SIZE { writeln!(kernel, "{{").unwrap(); - - write!(kernel, "cached_a_values[{j}] = {input_a}[").unwrap(); + let pre_element_wise_functions = pre_element_wise_functions + .get_or_init(|| op.pre_element_wise.add_functions(kernel)); + let mut raw_input = String::new(); + write!(&mut raw_input, "{input_a}[").unwrap(); let mut indices = Vec::new(); // Add batch indices first for dim in (0..input_a.rank()).rev().skip(2) { @@ -112,35 +179,35 @@ pub(crate) fn q_8_0_sgemv( // Then add M and K indices indices.push("m_idx".to_string()); indices.push(format!("y_offset + {j}")); - input_a.strided_index(kernel, indices); - writeln!(kernel, "];").unwrap(); + input_a.strided_index(&mut raw_input, indices); + write!(&mut raw_input, "]").unwrap(); + let processed = pre_element_wise_functions + .iter() + .fold(raw_input, |acc, f| f.call(vec![acc])); + writeln!(kernel, "cached_a_values[{j}] = {processed};").unwrap(); writeln!(kernel, "}}").unwrap(); } writeln!(kernel, "var block_offset = row_block_offset + i;").unwrap(); + + // Fast path: all rows in tile are valid, no bounds checks needed + writeln!(kernel, "if is_full_tile {{").unwrap(); for offset in 0..Q_8_0_SGEMV_CHUNK_SIZE { writeln!(kernel, "{{").unwrap(); - - writeln!(kernel, "var local_sum = {dtype}();").unwrap(); - for data_offset in 0..(STEP_SIZE / 4) { - writeln!(kernel, "{{").unwrap(); - - writeln!(kernel, "let block = vec4<{dtype}>(unpack4xI8({input_b}[block_offset].data[thread_local_id * 2u + {data_offset}]));").unwrap(); - writeln!(kernel, "let float_block = vec4<{dtype}>(cached_a_values[{data_offset} * 4u + 0], cached_a_values[{data_offset} * 4u + 1], cached_a_values[{data_offset} * 4u + 2], cached_a_values[{data_offset} * 4u + 3]);").unwrap(); - writeln!(kernel, "local_sum += dot(block, float_block);").unwrap(); - - writeln!(kernel, "}}").unwrap(); - } - let indexed = maybe_vec_storage_index(Q_8_0_SGEMV_CHUNK_SIZE, "sum", offset); - writeln!( - kernel, - "{indexed} += local_sum * {dtype}({input_b}[block_offset].scale);" - ) - .unwrap(); - - writeln!(kernel, "block_offset += k_block_size;").unwrap(); + generate_row_computation(kernel, offset, dtype, input_b); writeln!(kernel, "}}").unwrap(); } + writeln!(kernel, "}} else {{").unwrap(); + // Slow path: check bounds for each row + for offset in 0..Q_8_0_SGEMV_CHUNK_SIZE { + writeln!(kernel, "{{").unwrap(); + writeln!(kernel, "let row_index = row + {offset};").unwrap(); + writeln!(kernel, "if row_index < {n_size} {{").unwrap(); + generate_row_computation(kernel, offset, dtype, input_b); + writeln!(kernel, "}}").unwrap(); + writeln!(kernel, "}}").unwrap(); + } + writeln!(kernel, "}}").unwrap(); writeln!(kernel, "y_offset += {elements_per_block} * {STEP_SIZE};").unwrap(); } @@ -154,29 +221,31 @@ pub(crate) fn q_8_0_sgemv( ) .unwrap(); + // Initialize post element-wise functions once before the loop + let post_fns = + post_element_wise_functions.get_or_init(|| op.post_element_wise.add_functions(kernel)); + + // Fast path: all rows in tile are valid, no bounds checks needed + writeln!(kernel, "if is_full_tile {{").unwrap(); for offset in 0..Q_8_0_SGEMV_CHUNK_SIZE { writeln!(kernel, "{{").unwrap(); - - // If this is not the first simd thread in the workgroup, we can return early writeln!(kernel, "if {subgroup_local_index} == 0u {{").unwrap(); - { - // Write the output to the output tensor if this is the first thread in the workgroup - write!(kernel, "{output}[").unwrap(); - let index = format!("row + {offset}"); - let mut output_indices = Vec::new(); - // Add batch indices first - for dim in (0..output.rank()).rev().skip(2) { - output_indices.push(format!("batch_idx_{dim}")); - } - // Then add M and N indices - output_indices.push("m_idx".to_string()); - output_indices.push(index); - output.strided_index(kernel, output_indices); - let indexed = maybe_vec_storage_index(Q_8_0_SGEMV_CHUNK_SIZE, "sum", offset); - writeln!(kernel, "] = {indexed};").unwrap(); - } + writeln!(kernel, "let row_index = row + {offset}u;").unwrap(); + generate_row_output(op, kernel, offset, input_datatype, bias, output, post_fns); writeln!(kernel, "}}").unwrap(); - writeln!(kernel, "}}").unwrap(); } + writeln!(kernel, "}} else {{").unwrap(); + // Slow path: check bounds for each row + for offset in 0..Q_8_0_SGEMV_CHUNK_SIZE { + writeln!(kernel, "{{").unwrap(); + writeln!(kernel, "if {subgroup_local_index} == 0u {{").unwrap(); + writeln!(kernel, "let row_index = row + {offset}u;").unwrap(); + writeln!(kernel, "if row_index < {n_size} {{").unwrap(); + generate_row_output(op, kernel, offset, input_datatype, bias, output, post_fns); + writeln!(kernel, "}}").unwrap(); + writeln!(kernel, "}}").unwrap(); + writeln!(kernel, "}}").unwrap(); + } + writeln!(kernel, "}}").unwrap(); } diff --git a/fusor-ml/core/src/quantized/matmul/sgemv/q_n.rs b/fusor-ml/core/src/quantized/matmul/sgemv/q_n.rs index ca928e1f1..e4dbafd02 100644 --- a/fusor-ml/core/src/quantized/matmul/sgemv/q_n.rs +++ b/fusor-ml/core/src/quantized/matmul/sgemv/q_n.rs @@ -10,8 +10,20 @@ use crate::{ util::{maybe_vec_storage_index, maybe_vec_storage_subgroup_add, maybe_vec_storage_type}, }; use std::fmt::Write; +use std::sync::OnceLock; + +const Q_N_SGEMV_SMALL_CHUNK_SIZE: u32 = 8; +const Q_N_SGEMV_LARGE_CHUNK_SIZE: u32 = 32; +const Q_N_SGEMV_LARGE_N_THRESHOLD: u32 = 1024; + +pub(crate) fn q_n_sgemv_chunk_size(n: u32) -> u32 { + if n >= Q_N_SGEMV_LARGE_N_THRESHOLD { + Q_N_SGEMV_LARGE_CHUNK_SIZE + } else { + Q_N_SGEMV_SMALL_CHUNK_SIZE + } +} -pub(crate) const Q_N_SGEMV_CHUNK_SIZE: u32 = 4; // This is the size of the chunk each thread will process at a time const SUBGROUP_COUNT: u32 = 2; const SUBGROUP_SIZE: u32 = 32; @@ -23,6 +35,7 @@ pub(crate) fn q_n_sgemv( workgroup_shape: &WorkgroupShape, input_a: &TensorInput, input_b: &QMatrixInput, + bias: Option<&TensorInput>, output: &TensorInput, n_size: &str, m_size: &str, @@ -32,9 +45,12 @@ pub(crate) fn q_n_sgemv( let subgroup_index = kernel.subgroup_index(); let subgroup_local_index = kernel.subgroup_local_index(); let elements_per_block = op.elements_per_block(); + let row_chunk_size = q_n_sgemv_chunk_size(op.n_size()); + let pre_element_wise_functions = OnceLock::new(); + let post_element_wise_functions = OnceLock::new(); - // Calculate n_workgroups for this kernel type (SUBGROUP_COUNT subgroups per workgroup, Q_N_SGEMV_CHUNK_SIZE per subgroup) - let chunk_size = Q_N_SGEMV_CHUNK_SIZE * SUBGROUP_COUNT; + // Calculate n_workgroups for this kernel type (SUBGROUP_COUNT subgroups per workgroup, row_chunk_size per subgroup) + let chunk_size = row_chunk_size * SUBGROUP_COUNT; let n_workgroups = format!("(({n_size} + {chunk_size} - 1) / {chunk_size})"); // Decompose linearized workgroup index into (n_workgroup_idx, m_idx, batch_idx) @@ -67,7 +83,14 @@ pub(crate) fn q_n_sgemv( writeln!(kernel, "let workgroup_offset = n_workgroup_idx;").unwrap(); writeln!( kernel, - "let row = ({SUBGROUP_COUNT} * workgroup_offset + {subgroup_index}) * {Q_N_SGEMV_CHUNK_SIZE};" + "let row = ({SUBGROUP_COUNT} * workgroup_offset + {subgroup_index}) * {row_chunk_size};" + ) + .unwrap(); + + // Fast path check: if all rows in this tile are valid, skip per-row bounds checks + writeln!( + kernel, + "let is_full_tile = row + {row_chunk_size} <= {n_size};" ) .unwrap(); @@ -85,7 +108,7 @@ pub(crate) fn q_n_sgemv( .unwrap(); // Always accumulate in f32 to avoid overflow, then convert to output dtype at the end - let sum_storage_type = maybe_vec_storage_type(Q_N_SGEMV_CHUNK_SIZE, DataTypeEnum::F32); + let sum_storage_type = maybe_vec_storage_type(row_chunk_size, DataTypeEnum::F32); writeln!(kernel, "var sum = {sum_storage_type}();",).unwrap(); writeln!(kernel, "var cached_a_values = array();",).unwrap(); @@ -109,7 +132,10 @@ pub(crate) fn q_n_sgemv( ("a_val_16", j + 16), ("a_val_17", j + 17), ] { - write!(kernel, "let {var_name} = f32({input_a}[").unwrap(); + let pre_element_wise_functions = pre_element_wise_functions + .get_or_init(|| op.pre_element_wise.add_functions(kernel)); + let mut raw_input = String::new(); + write!(&mut raw_input, "{input_a}[").unwrap(); let mut indices = Vec::new(); // Add batch indices first for dim in (0..input_a.rank()).rev().skip(2) { @@ -118,8 +144,12 @@ pub(crate) fn q_n_sgemv( // Then add M and K indices indices.push("m_idx".to_string()); indices.push(format!("y_offset + {offset}")); - input_a.strided_index(kernel, indices); - writeln!(kernel, "]);").unwrap(); + input_a.strided_index(&mut raw_input, indices); + write!(&mut raw_input, "]").unwrap(); + let processed = pre_element_wise_functions + .iter() + .fold(raw_input, |acc, f| f.call(vec![acc])); + writeln!(kernel, "let {var_name} = f32({processed});").unwrap(); } writeln!( @@ -156,22 +186,48 @@ pub(crate) fn q_n_sgemv( writeln!(kernel, "let vector_total = vector_sum;").unwrap(); writeln!(kernel, "var block_offset = row_block_offset + i;").unwrap(); - if Q_N_SGEMV_CHUNK_SIZE > 1 { + + // Fast path: all rows in tile are valid, no bounds checks needed + writeln!(kernel, "if is_full_tile {{").unwrap(); + if row_chunk_size > 1 { + writeln!( + kernel, + "for (var offset = 0u; offset < {row_chunk_size}; offset += 1u) {{" + ) + .unwrap(); + } + { + writeln!(kernel, "let row_index = row + offset;").unwrap(); + block_dot(kernel, op, input_b); + let indexed = maybe_vec_storage_index(row_chunk_size, "sum", "offset"); + writeln!(kernel, "{indexed} += product;").unwrap(); + } + if row_chunk_size > 1 { + writeln!(kernel, "block_offset += k_block_size;").unwrap(); + writeln!(kernel, "}}").unwrap(); + } + writeln!(kernel, "}} else {{").unwrap(); + // Slow path: check bounds for each row + if row_chunk_size > 1 { writeln!( kernel, - "for (var offset = 0u; offset < {Q_N_SGEMV_CHUNK_SIZE}; offset += 1u) {{" + "for (var offset = 0u; offset < {row_chunk_size}; offset += 1u) {{" ) .unwrap(); } { + writeln!(kernel, "let row_index = row + offset;").unwrap(); + writeln!(kernel, "if row_index < {n_size} {{").unwrap(); block_dot(kernel, op, input_b); - let indexed = maybe_vec_storage_index(Q_N_SGEMV_CHUNK_SIZE, "sum", "offset"); + let indexed = maybe_vec_storage_index(row_chunk_size, "sum", "offset"); writeln!(kernel, "{indexed} += product;").unwrap(); + writeln!(kernel, "}}").unwrap(); } - if Q_N_SGEMV_CHUNK_SIZE > 1 { + if row_chunk_size > 1 { writeln!(kernel, "block_offset += k_block_size;").unwrap(); writeln!(kernel, "}}").unwrap(); } + writeln!(kernel, "}}").unwrap(); writeln!(kernel, "y_offset += {elements_per_block} * 16;").unwrap(); } @@ -181,47 +237,81 @@ pub(crate) fn q_n_sgemv( writeln!( kernel, "sum = {};", - maybe_vec_storage_subgroup_add(Q_N_SGEMV_CHUNK_SIZE, "sum") + maybe_vec_storage_subgroup_add(row_chunk_size, "sum") ) .unwrap(); - // If this is not the first simd thread in the workgroup, we can return early - - if Q_N_SGEMV_CHUNK_SIZE > 1 { + // Initialize post element-wise functions once before the loop + let post_fns = + post_element_wise_functions.get_or_init(|| op.post_element_wise.add_functions(kernel)); + + // Generate the output body for one row + let generate_row_output = |kernel: &mut GenericKernel| { + write!(kernel, "{output}[").unwrap(); + let mut output_indices = Vec::new(); + // Add batch indices first + for dim in (0..output.rank()).rev().skip(2) { + output_indices.push(format!("batch_idx_{dim}")); + } + // Then add M and N indices + output_indices.push("m_idx".to_string()); + output_indices.push("row_index".to_string()); + output.strided_index(kernel, output_indices); + let indexed = maybe_vec_storage_index(row_chunk_size, "sum", "offset"); + let result = + op.apply_bias_and_post(bias, "row_index", format!("{dtype}({indexed})"), post_fns); + writeln!(kernel, "] = {result};").unwrap(); + }; + + // Fast path: all rows in tile are valid, no bounds checks needed + writeln!(kernel, "if is_full_tile {{").unwrap(); + if row_chunk_size > 1 { writeln!( kernel, - "for (var offset = 0u; offset < {Q_N_SGEMV_CHUNK_SIZE}; offset += 1u) {{" + "for (var offset = 0u; offset < {row_chunk_size}; offset += 1u) {{" ) .unwrap(); } { writeln!(kernel, "if {subgroup_local_index} == 0u {{").unwrap(); - { - // Write the output to the output tensor if this is the first thread in the workgroup - // Convert from f32 accumulator to output dtype - write!(kernel, "{output}[").unwrap(); - let index = if Q_N_SGEMV_CHUNK_SIZE > 1 { - "row + offset".to_string() - } else { - "row".to_string() - }; - let mut output_indices = Vec::new(); - // Add batch indices first - for dim in (0..output.rank()).rev().skip(2) { - output_indices.push(format!("batch_idx_{dim}")); - } - // Then add M and N indices - output_indices.push("m_idx".to_string()); - output_indices.push(index); - output.strided_index(kernel, output_indices); - let indexed = maybe_vec_storage_index(Q_N_SGEMV_CHUNK_SIZE, "sum", "offset"); - writeln!(kernel, "] = {dtype}({indexed});").unwrap(); - } + let index = if row_chunk_size > 1 { + "row + offset".to_string() + } else { + "row".to_string() + }; + writeln!(kernel, "let row_index = {index};").unwrap(); + generate_row_output(kernel); writeln!(kernel, "}}").unwrap(); } - if Q_N_SGEMV_CHUNK_SIZE > 1 { + if row_chunk_size > 1 { writeln!(kernel, "}}").unwrap(); } + writeln!(kernel, "}} else {{").unwrap(); + // Slow path: check bounds for each row + if row_chunk_size > 1 { + writeln!( + kernel, + "for (var offset = 0u; offset < {row_chunk_size}; offset += 1u) {{" + ) + .unwrap(); + } + { + writeln!(kernel, "if {subgroup_local_index} == 0u {{").unwrap(); + let index = if row_chunk_size > 1 { + "row + offset".to_string() + } else { + "row".to_string() + }; + writeln!(kernel, "let row_index = {index};").unwrap(); + writeln!(kernel, "if row_index < {n_size} {{").unwrap(); + generate_row_output(kernel); + writeln!(kernel, "}}").unwrap(); + writeln!(kernel, "}}").unwrap(); + } + if row_chunk_size > 1 { + writeln!(kernel, "}}").unwrap(); + } + writeln!(kernel, "}}").unwrap(); } fn block_dot(kernel: &mut GenericKernel, op: &QMatMulOperation, input_b: &QMatrixInput) { diff --git a/fusor-ml/core/src/resize.rs b/fusor-ml/core/src/resize.rs index f09bc2fad..a5554f5cd 100644 --- a/fusor-ml/core/src/resize.rs +++ b/fusor-ml/core/src/resize.rs @@ -20,6 +20,9 @@ pub(crate) struct ResizeOperation { pub(crate) current_shape: Box<[usize]>, pub(crate) new_shape: Box<[usize]>, pub(crate) fill_shape: Box<[usize]>, + /// Forbid layout-only lowering; always run the copy kernel. Used by + /// `contiguous()` to materialize a strided view into row-major memory. + pub(crate) force_copy: bool, } impl ResizeOperation { @@ -34,6 +37,22 @@ impl ResizeOperation { current_shape, new_shape, fill_shape, + force_copy: false, + } + } + + pub fn new_force_copy( + input: NodeIndex, + current_shape: Box<[usize]>, + new_shape: Box<[usize]>, + fill_shape: Box<[usize]>, + ) -> Self { + Self { + input, + current_shape, + new_shape, + fill_shape, + force_copy: true, } } } @@ -43,6 +62,9 @@ impl ResizeOperation { &self, graph: &crate::compute_graph::ComputeGraphInner, ) -> Option { + if self.force_copy { + return None; + } let full_fill = self.fill_shape == self.new_shape; let matching_size = self.current_shape.iter().product::() == self.new_shape.iter().product::(); @@ -239,6 +261,20 @@ impl Tensor { )) } + /// Materialize this tensor into a row-major contiguous buffer. Always + /// runs a copy kernel so subsequent ops (and `as_slice` reads) observe + /// data in logical shape order. + pub fn contiguous(&self) -> Tensor { + let shape: Box<[usize]> = (*self.shape()).as_slice().into(); + let input = self.key(); + self.add_resize(ResizeOperation::new_force_copy( + input, + shape.clone(), + shape.clone(), + shape, + )) + } + pub fn reshape(&self, new_shape: impl ShapeWithOneHole) -> Tensor { let new_shape = new_shape.resolve_shape(self.shape()); assert_eq!( diff --git a/fusor-ml/core/src/slice_assign.rs b/fusor-ml/core/src/slice_assign.rs index e9b49417e..0b2ca772d 100644 --- a/fusor-ml/core/src/slice_assign.rs +++ b/fusor-ml/core/src/slice_assign.rs @@ -244,3 +244,53 @@ async fn test_slice_assign_nonzero_offset() { assert_eq!(as_slice[[2, 3]], 30.); assert_eq!(as_slice[[2, 4]], 40.); } + +#[cfg(test)] +#[tokio::test] +async fn test_slice_assign_4d_preserves_prefix() { + use crate::Device; + + let device = Device::test_instance(); + + let data = vec![vec![ + vec![vec![0.1f32, 0.2, 0.3, 0.4]], + vec![vec![0.0f32, 0.0, 0.0, 0.0]], + ]]; + let tensor = Tensor::new(&device, &data); + let value_tensor = Tensor::new(&device, &vec![vec![vec![vec![0.5f32, 0.6, 0.7, 0.8]]]]); + let result = tensor.slice_assign([0..1, 1..2, 0..1, 0..4], &value_tensor); + let as_slice = result.as_slice().await.unwrap(); + + assert!((as_slice[[0, 0, 0, 0]] - 0.1).abs() < 1e-6); + assert!((as_slice[[0, 0, 0, 1]] - 0.2).abs() < 1e-6); + assert!((as_slice[[0, 0, 0, 2]] - 0.3).abs() < 1e-6); + assert!((as_slice[[0, 0, 0, 3]] - 0.4).abs() < 1e-6); + assert!((as_slice[[0, 1, 0, 0]] - 0.5).abs() < 1e-6); + assert!((as_slice[[0, 1, 0, 1]] - 0.6).abs() < 1e-6); + assert!((as_slice[[0, 1, 0, 2]] - 0.7).abs() < 1e-6); + assert!((as_slice[[0, 1, 0, 3]] - 0.8).abs() < 1e-6); +} + +#[cfg(test)] +#[tokio::test] +async fn test_cat_then_slice_assign_4d_preserves_prefix() { + use crate::Device; + + let device = Device::test_instance(); + + let data = vec![vec![vec![vec![0.1f32, 0.2, 0.3, 0.4]]]]; + let tensor = Tensor::new(&device, &data); + let grown = Tensor::cat([tensor, Tensor::zeros(&device, [1, 1, 1, 4])], 1); + let value_tensor = Tensor::new(&device, &vec![vec![vec![vec![0.5f32, 0.6, 0.7, 0.8]]]]); + let result = grown.slice_assign([0..1, 1..2, 0..1, 0..4], &value_tensor); + let as_slice = result.as_slice().await.unwrap(); + + assert!((as_slice[[0, 0, 0, 0]] - 0.1).abs() < 1e-6); + assert!((as_slice[[0, 0, 0, 1]] - 0.2).abs() < 1e-6); + assert!((as_slice[[0, 0, 0, 2]] - 0.3).abs() < 1e-6); + assert!((as_slice[[0, 0, 0, 3]] - 0.4).abs() < 1e-6); + assert!((as_slice[[0, 1, 0, 0]] - 0.5).abs() < 1e-6); + assert!((as_slice[[0, 1, 0, 1]] - 0.6).abs() < 1e-6); + assert!((as_slice[[0, 1, 0, 2]] - 0.7).abs() < 1e-6); + assert!((as_slice[[0, 1, 0, 3]] - 0.8).abs() < 1e-6); +} diff --git a/fusor-ml/core/src/tensor.rs b/fusor-ml/core/src/tensor.rs index 41e6b2286..27b86df53 100644 --- a/fusor-ml/core/src/tensor.rs +++ b/fusor-ml/core/src/tensor.rs @@ -173,6 +173,7 @@ impl TensorLayoutInfo { pub(crate) struct TensorInfo { shape: Box<[usize]>, datatype: DataTypeEnum, + is_zero_splat: bool, } impl Display for TensorInfo { @@ -183,7 +184,19 @@ impl Display for TensorInfo { impl TensorInfo { pub(crate) fn new(shape: Box<[usize]>, datatype: DataTypeEnum) -> Self { - Self { shape, datatype } + Self { + shape, + datatype, + is_zero_splat: false, + } + } + + pub(crate) fn new_zero_splat(shape: Box<[usize]>, datatype: DataTypeEnum) -> Self { + Self { + shape, + datatype, + is_zero_splat: true, + } } pub(crate) fn shape(&self) -> &[usize] { @@ -197,6 +210,14 @@ impl TensorInfo { pub(crate) fn datatype(&self) -> DataTypeEnum { self.datatype } + + pub(crate) fn is_zero_splat(&self) -> bool { + self.is_zero_splat + } + + pub(crate) fn mark_zero_splat(&mut self) { + self.is_zero_splat = true; + } } pub(crate) struct LazyTensorData { @@ -226,11 +247,16 @@ impl LazyTensorData { pub(crate) fn new(data: TensorData) -> Self { let device = data.device.clone(); let info = data.info.clone(); + let is_zero_splat = data.is_zero_splat(); let key = device.compute_graph().create_tensor(data); Self { device, - info: TensorInfo::new(info.shape().into(), info.datatype()), + info: if is_zero_splat { + TensorInfo::new_zero_splat(info.shape().into(), info.datatype()) + } else { + TensorInfo::new(info.shape().into(), info.datatype()) + }, key, } } @@ -315,7 +341,11 @@ impl LazyTensorData { // 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 info = if self.info.is_zero_splat() { + TensorInfo::new_zero_splat(new_layout.shape().into(), self.info.datatype()) + } else { + TensorInfo::new(new_layout.shape().into(), self.info.datatype()) + }; let key = device.compute_graph().create_map_layout(op); Self::from_parts(device, info, key) @@ -323,7 +353,11 @@ impl LazyTensorData { pub(crate) fn resize(&self, op: ResizeOperation) -> Self { let device = self.device.clone(); - let info = TensorInfo::new(op.new_shape.clone(), self.info.datatype()); + let info = if self.info.is_zero_splat() { + TensorInfo::new_zero_splat(op.new_shape.clone(), self.info.datatype()) + } else { + TensorInfo::new(op.new_shape.clone(), self.info.datatype()) + }; let key = device.compute_graph().create_resize(op); Self::from_parts(device, info, key) @@ -331,7 +365,9 @@ impl LazyTensorData { pub(crate) fn slice_assign(&self, op: SliceAssignOperation) -> Self { let device = self.device.clone(); - let info = self.info.clone(); + // slice_assign can introduce non-zero values even when the input tensor is a zero splat, + // so the result must not inherit the zero-splat marker from the input. + let info = TensorInfo::new(self.info.shape().into(), self.info.datatype()); let key = device.compute_graph().create_slice_assign(op); Self::from_parts(device, info, key) @@ -351,6 +387,10 @@ impl LazyTensorData { (result.data, result.total_kernels) } + pub(crate) fn is_zero_splat(&self) -> bool { + self.info.is_zero_splat() + } + pub fn graphvis(&self) -> Graph { self.device.compute_graph().graphvis(self.key) } @@ -361,11 +401,14 @@ pub(crate) struct TensorData { device: Device, buffer: Arc, info: TensorLayoutInfo, + is_zero_splat: bool, } impl PartialEq for TensorData { fn eq(&self, other: &Self) -> bool { - self.info == other.info && self.buffer == other.buffer + self.info == other.info + && self.buffer == other.buffer + && self.is_zero_splat == other.is_zero_splat } } @@ -402,6 +445,7 @@ impl TensorData { device: device.clone(), buffer, info: TensorLayoutInfo::new(layout, datatype), + is_zero_splat: false, } } @@ -435,7 +479,9 @@ impl TensorData { ); let strides = (0..shape.len()).map(|_| 0).collect(); let layout = Layout::from_parts(0, shape.into(), strides); - Self::new_from_parts(device, buffer, layout, datatype) + let mut tensor = Self::new_from_parts(device, buffer, layout, datatype); + tensor.is_zero_splat = raw_data.iter().all(|byte| *byte == 0); + tensor } fn new_inner<'a, D: DataType, I: Iterator>( @@ -484,6 +530,7 @@ impl TensorData { device: self.device.clone(), buffer: self.buffer.clone(), info: TensorLayoutInfo::new(layout, self.info.datatype), + is_zero_splat: self.is_zero_splat, } } @@ -495,6 +542,14 @@ impl TensorData { self.info.datatype } + pub(crate) fn is_zero_splat(&self) -> bool { + self.is_zero_splat + } + + pub(crate) fn mark_zero_splat(&mut self) { + self.is_zero_splat = true; + } + pub fn device(&self) -> &Device { &self.device } @@ -784,6 +839,10 @@ impl Tensor { } } + pub(crate) fn is_zero_splat(&self) -> bool { + self.data.is_zero_splat() + } + fn new_inner<'a, I: Iterator>( device: &Device, data: I, @@ -873,12 +932,94 @@ impl Tensor { } } + #[track_caller] + pub fn materialized(&self) -> impl Future + 'static { + let (data, _) = self.data.materialize(); + let device = self.device().clone(); + let (sender, receiver) = futures_channel::oneshot::channel(); + device.wgpu_queue().on_submitted_work_done(|| { + _ = sender.send(()); + }); + async move { + let _ = receiver.await; + Self::from_parts(LazyTensorData::new(data)) + } + } + + #[track_caller] + pub fn materialized_many( + tensors: &[&Self], + ) -> std::pin::Pin> + 'static>> { + if tensors.is_empty() { + return Box::pin(async { Vec::new() }); + } + + let device = tensors[0].device().clone(); + let keys: Vec<_> = tensors + .iter() + .map(|tensor| { + assert_eq!( + tensor.device().wgpu_device(), + device.wgpu_device(), + "all tensors in materialized_many must be on the same device" + ); + tensor.data.key + }) + .collect(); + let resolved = device.compute_graph().resolve_many(&keys, &device); + let data = resolved.data; + let (sender, receiver) = futures_channel::oneshot::channel(); + device.wgpu_queue().on_submitted_work_done(|| { + _ = sender.send(()); + }); + Box::pin(async move { + let _ = receiver.await; + data.into_iter() + .map(|data| Self::from_parts(LazyTensorData::new(data))) + .collect() + }) + } + + #[track_caller] + pub fn materialize_many( + tensors: &[&Self], + ) -> std::pin::Pin + 'static>> { + if tensors.is_empty() { + return Box::pin(async {}); + } + + let device = tensors[0].device().clone(); + let keys: Vec<_> = tensors + .iter() + .map(|tensor| { + assert_eq!( + tensor.device().wgpu_device(), + device.wgpu_device(), + "all tensors in materialize_many must be on the same device" + ); + tensor.data.key + }) + .collect(); + let _ = device.compute_graph().resolve_many(&keys, &device); + let (sender, receiver) = futures_channel::oneshot::channel(); + device.wgpu_queue().on_submitted_work_done(|| { + _ = sender.send(()); + }); + Box::pin(async move { + let _ = receiver.await; + }) + } + /// How many kernel calls are needed to fully resolve this tensor pub fn count_kernels_to_resolve(&self) -> usize { let (_, count) = self.data.materialize(); count } + pub(crate) fn is_materialized(&self) -> bool { + self.device().compute_graph().is_cached(self.data.key) + } + pub async fn as_slice( &self, ) -> Result, wgpu::BufferAsyncError> { @@ -959,6 +1100,14 @@ impl Tensor { Self::from_parts(self.data.q_mat_mul(operation)) } + pub(crate) fn add_q_mat_mul_bias(&self, other: &QMatrix, bias: &Tensor<1, D>) -> Self { + let operation = + QMatMulOperation::new(self.datatype(), self.shape(), self.data.key, other.clone()) + .with_bias(bias); + + Self::from_parts(self.data.q_mat_mul(operation)) + } + pub(crate) fn add_resize(&self, op: ResizeOperation) -> Tensor { Tensor { data: self.data.resize(op), diff --git a/fusor-ml/fusor/.claude/settings.local.json b/fusor-ml/fusor/.claude/settings.local.json new file mode 100644 index 000000000..b34331469 --- /dev/null +++ b/fusor-ml/fusor/.claude/settings.local.json @@ -0,0 +1,7 @@ +{ + "permissions": { + "allow": [ + "Bash(alias rt='cargo nextest run')" + ] + } +} diff --git a/fusor-ml/fusor/src/cache/attention_mask.rs b/fusor-ml/fusor/src/cache/attention_mask.rs index eda1b7fc2..d58b6f23d 100644 --- a/fusor-ml/fusor/src/cache/attention_mask.rs +++ b/fusor-ml/fusor/src/cache/attention_mask.rs @@ -155,4 +155,31 @@ mod tests { assert_eq!(output[[0, 1, 0]], 3.0); assert_eq!(output[[0, 1, 1]], 4.0); } + + #[tokio::test] + async fn test_attention_mask_causal_gpu_matches_cpu() { + let cpu = Device::cpu(); + let gpu = Device::new().await.expect("GPU required for this test"); + + let cpu_mask: AttentionMask = AttentionMask::causal(&cpu, 9); + let gpu_mask: AttentionMask = AttentionMask::causal(&gpu, 9); + + let cpu_data = cpu_mask.mask().clone().as_slice().await.unwrap(); + let gpu_data = gpu_mask.mask().clone().as_slice().await.unwrap(); + + for i in 0..9 { + for j in 0..9 { + let expected = cpu_data[[i, j]]; + let actual = gpu_data[[i, j]]; + if expected.is_infinite() { + assert!(actual.is_infinite() && actual.is_sign_negative()); + } else { + assert!( + (expected - actual).abs() < 1e-6, + "mismatch at [{i}, {j}]: expected {expected}, got {actual}" + ); + } + } + } + } } diff --git a/fusor-ml/fusor/src/cache/tensor_cache.rs b/fusor-ml/fusor/src/cache/tensor_cache.rs index adb86010e..170fb1e55 100644 --- a/fusor-ml/fusor/src/cache/tensor_cache.rs +++ b/fusor-ml/fusor/src/cache/tensor_cache.rs @@ -105,15 +105,17 @@ where *cached = cached.slice_assign(slice, v); self.current_seq_len = required_seq_len; // Return only the valid portion of the cache, not the full allocated tensor - cached + let current = cached .narrow(self.concat_dim, 0, self.current_seq_len) - .to_concrete() + .to_concrete(); + current } else { // First append - just store it - self.all_data = Some(v.clone()); + let current = v.to_concrete(); + self.all_data = Some(current.clone()); self.current_seq_len = seq_len; self.allocated_seq_len = seq_len; - v.clone() + current } } @@ -184,4 +186,90 @@ mod tests { assert_eq!(cache.current_seq_len(), 0); assert!(cache.current_data().is_none()); } + + #[tokio::test] + async fn test_tensor_cache_gpu_matches_cpu_multiple_appends() { + let cpu = Device::cpu(); + let gpu = Device::new().await.expect("GPU required for this test"); + + let mut cpu_cache: TensorCache<4, f32> = TensorCache::new(1, 16); + let mut gpu_cache: TensorCache<4, f32> = TensorCache::new(1, 16); + + let chunks = [ + vec![0.1f32, 0.2, 0.3, 0.4], + vec![0.5f32, 0.6, 0.7, 0.8], + vec![0.9f32, 1.0, 1.1, 1.2], + ]; + + for chunk in chunks { + let cpu_tensor: Tensor<4, f32> = Tensor::from_slice(&cpu, [1, 1, 1, 4], &chunk); + let gpu_tensor: Tensor<4, f32> = Tensor::from_slice(&gpu, [1, 1, 1, 4], &chunk); + + let gpu_direct = gpu_tensor.clone().as_slice().await.unwrap(); + for d in 0..4 { + assert!( + (gpu_direct[[0, 0, 0, d]] - chunk[d]).abs() < 1e-6, + "direct gpu mismatch at {d}: expected {}, got {}", + chunk[d], + gpu_direct[[0, 0, 0, d]] + ); + } + + let cpu_result = cpu_cache + .append(&cpu, &cpu_tensor) + .as_slice() + .await + .unwrap(); + let gpu_result = gpu_cache + .append(&gpu, &gpu_tensor) + .as_slice() + .await + .unwrap(); + + let shape = cpu_result.shape(); + for b in 0..shape[0] { + for s in 0..shape[1] { + for t in 0..shape[2] { + for d in 0..shape[3] { + assert!( + (cpu_result[[b, s, t, d]] - gpu_result[[b, s, t, d]]).abs() < 1e-6, + "mismatch at [{b}, {s}, {t}, {d}]: expected {}, got {}", + cpu_result[[b, s, t, d]], + gpu_result[[b, s, t, d]] + ); + } + } + } + } + } + } + + #[tokio::test] + async fn test_tensor_cache_gpu_lazy_trim_without_intermediate_reads() { + let gpu = Device::new().await.expect("GPU required for this test"); + let mut cache: TensorCache<4, f32> = TensorCache::new(1, 3); + + let chunks = [[1.0f32, 2.0], [3.0f32, 4.0], [5.0f32, 6.0], [7.0f32, 8.0]]; + + let mut result = None; + for chunk in chunks { + let tensor: Tensor<4, f32> = Tensor::from_slice(&gpu, [1, 1, 1, 2], &chunk); + result = Some(cache.append(&gpu, &tensor)); + } + + let result = result.unwrap(); + assert_eq!(result.shape(), [1, 3, 1, 2]); + + let output = result.as_slice().await.unwrap(); + let expected = [[3.0f32, 4.0], [5.0, 6.0], [7.0, 8.0]]; + for (seq, values) in expected.iter().enumerate() { + for (dim, expected) in values.iter().copied().enumerate() { + let actual = output[[0, seq, 0, dim]]; + assert!( + (actual - expected).abs() < 1e-6, + "mismatch at sequence {seq}, dim {dim}: expected {expected}, got {actual}" + ); + } + } + } } diff --git a/fusor-ml/fusor/src/composite/activations.rs b/fusor-ml/fusor/src/composite/activations.rs index d5b9c902b..9636edb5d 100644 --- a/fusor-ml/fusor/src/composite/activations.rs +++ b/fusor-ml/fusor/src/composite/activations.rs @@ -78,6 +78,28 @@ where } } +impl Tensor<3, f32> { + pub fn swiglu_split(&self, intermediate_size: usize) -> Tensor<3, f32> { + assert_eq!( + self.shape()[2], + intermediate_size * 2, + "swiglu_split expects the last dimension to be 2 * intermediate_size" + ); + match self { + Tensor::Cpu(_) => { + let states = self.narrow(2, 0, intermediate_size).to_concrete(); + let gate = self + .narrow(2, intermediate_size, intermediate_size) + .to_concrete() + .silu() + .to_concrete(); + states.mul_(&gate).to_concrete() + } + Tensor::Gpu(t) => Tensor::Gpu(t.swiglu_split(intermediate_size)), + } + } +} + #[cfg(test)] #[allow(clippy::identity_op, clippy::useless_conversion)] mod tests { @@ -120,6 +142,40 @@ mod tests { } } + #[tokio::test] + async fn test_swiglu_split_cpu_and_gpu() { + let data = [[ + [1.0, -2.0, 0.5, -0.25, 0.1, 1.5], + [3.0, 0.25, -1.0, 2.0, -0.5, 0.75], + ]]; + let cpu: Tensor<3, f32> = Tensor::Cpu(fusor_cpu::Tensor::from_slice( + [1, 2, 6], + &[ + 1.0, -2.0, 0.5, -0.25, 0.1, 1.5, 3.0, 0.25, -1.0, 2.0, -0.5, 0.75, + ], + )); + let expected = cpu + .narrow(2, 0, 3) + .to_concrete() + .mul_(&cpu.narrow(2, 3, 3).to_concrete().silu().to_concrete()) + .to_concrete(); + let cpu_output = cpu.swiglu_split(3); + let expected = expected.as_slice().await.unwrap(); + let cpu_output = cpu_output.as_slice().await.unwrap(); + for (actual, expected) in cpu_output.as_slice().iter().zip(expected.as_slice()) { + assert!((actual - expected).abs() < 1e-6); + } + + let device = crate::Device::new() + .await + .expect("GPU required for this test"); + let gpu: Tensor<3, f32> = Tensor::new(&device, &data); + let gpu_output = gpu.swiglu_split(3).as_slice().await.unwrap(); + for (actual, expected) in gpu_output.as_slice().iter().zip(expected.as_slice()) { + assert!((actual - expected).abs() < 1e-6); + } + } + fn gelu_ref(x: f32) -> f32 { 0.5 * x * (1.0 + ((2.0 / std::f32::consts::PI).sqrt() * (x + 0.044715 * x.powi(3))).tanh()) } diff --git a/fusor-ml/fusor/src/composite/conv.rs b/fusor-ml/fusor/src/composite/conv.rs index 4802fc806..5224aff94 100644 --- a/fusor-ml/fusor/src/composite/conv.rs +++ b/fusor-ml/fusor/src/composite/conv.rs @@ -111,15 +111,29 @@ where }); let windows_tensor: Tensor = padded.sliding_window_view(windows); - // Step 3: Prepare for matmul by reshaping and transposing + // Step 3: Prepare for matmul by reshaping and permuting. + // + // `sliding_window_view` produces: + // (batch, in_channels, out_spatial..., kernel...) + // We need: + // (batch, out_spatial..., in_channels, kernel...) + // before flattening to rows. A single transpose is only correct for 1D. let kernel_size: usize = weight_shape[spatial_start..].iter().product(); - // Transpose to move in_channels after spatial - let windows_transposed = windows_tensor.transpose(in_channels_axis, spatial_start); + let mut window_axes = [0usize; R2]; + window_axes[batch_axis] = batch_axis; + for i in 0..DIFF { + window_axes[1 + i] = spatial_start + i; + } + window_axes[1 + DIFF] = in_channels_axis; + for i in 0..DIFF { + window_axes[2 + DIFF + i] = R + i; + } + let windows_permuted = windows_tensor.permute(window_axes).to_concrete(); // Flatten to (batch * out_spatial_size, in_channels * kernel_size) let windows_flat: Tensor<2, D, _> = - windows_transposed.reshape([batch * out_spatial_size, in_channels * kernel_size]); + windows_permuted.reshape([batch * out_spatial_size, in_channels * kernel_size]); // Step 4: Reshape weight for matmul let weight_reshaped: Tensor<2, D, _> = @@ -130,27 +144,38 @@ where // Step 5: Matrix multiplication let output = windows_flat.mat_mul(&weight_t); - // Step 6: Reshape and transpose back to (batch, out_channels, ...out_spatial...) - let output_reshaped: Tensor<3, D, _> = - output.reshape([batch, out_spatial_size, out_channels]); - let output_transposed = output_reshaped.transpose(in_channels_axis, spatial_start); - - // Reshape to (batch, out_channels, ...out_spatial_dims...) - let mut output_shape = input_shape; - output_shape[in_channels_axis] = out_channels; + // Step 6: Reshape and permute back to (batch, out_channels, ...out_spatial...) + let mut output_spatial = [0usize; DIFF]; for i in 0..DIFF { let padded_len = input_shape[spatial_start + i] + 2 * padding[i]; let kernel_len = weight_shape[spatial_start + i]; - output_shape[spatial_start + i] = (padded_len - kernel_len) / strides[i] + 1; + output_spatial[i] = (padded_len - kernel_len) / strides[i] + 1; } - let output_final = output_transposed.reshape(output_shape); + + let mut output_reshape_shape = [0usize; R]; + output_reshape_shape[batch_axis] = batch; + for i in 0..DIFF { + output_reshape_shape[1 + i] = output_spatial[i]; + } + output_reshape_shape[R - 1] = out_channels; + + let output_reshaped: Tensor = output.reshape(output_reshape_shape); + let mut output_axes = [0usize; R]; + output_axes[batch_axis] = batch_axis; + output_axes[in_channels_axis] = R - 1; + for i in 0..DIFF { + output_axes[spatial_start + i] = 1 + i; + } + let output_final = output_reshaped.permute(output_axes).to_concrete(); // Step 7: Add bias if present if let Some(bias) = bias { // Bias shape: (out_channels,) - // Need to broadcast to (batch, out_channels, ...spatial...) - // Broadcast bias to the FULL output shape for correct addition - let bias_broadcast: Tensor = bias.broadcast_as(output_shape); + // Need to broadcast along the channel axis, not the trailing axis. + let mut bias_shape = [1; R]; + bias_shape[in_channels_axis] = out_channels; + let bias_unsqueezed = bias.unsqueeze(0); + let bias_broadcast: Tensor = bias_unsqueezed.reshape(bias_shape); output_final.add_(&bias_broadcast) } else { output_final.to_concrete() @@ -335,4 +360,45 @@ mod tests { assert!((result[[0, 2, 1]] - 18.0).abs() < 1e-5); assert!((result[[0, 2, 2]] - 22.0).abs() < 1e-5); } + + #[tokio::test] + async fn test_pad_axis_gpu_matches_cpu_4d() { + let gpu = crate::Device::new() + .await + .expect("GPU required for this test"); + + let input_data = vec![ + 1.0f32, 2.0, 3.0, 4.0, // + 5.0, 6.0, 7.0, 8.0, // + 9.0, 10.0, 11.0, 12.0, // + 13.0, 14.0, 15.0, 16.0, + ]; + let cpu_input: Tensor<4, f32> = + Tensor::Cpu(fusor_cpu::Tensor::from_slice([1, 1, 4, 4], &input_data)); + let gpu_input: Tensor<4, f32> = Tensor::from_slice(&gpu, [1, 1, 4, 4], &input_data); + + let cpu_padded = cpu_input + .pad_axis(2, 1) + .pad_axis(3, 1) + .as_slice() + .await + .unwrap(); + let gpu_padded = gpu_input + .pad_axis(2, 1) + .pad_axis(3, 1) + .as_slice() + .await + .unwrap(); + + for h in 0..cpu_padded.shape()[2] { + for w in 0..cpu_padded.shape()[3] { + let expected = cpu_padded[[0, 0, h, w]]; + let actual = gpu_padded[[0, 0, h, w]]; + assert!( + (expected - actual).abs() < 1e-6, + "mismatch at [{h}, {w}]: expected {expected}, got {actual}" + ); + } + } + } } diff --git a/fusor-ml/fusor/src/composite/flash_attention.rs b/fusor-ml/fusor/src/composite/flash_attention.rs index a34539645..5e2d08e9f 100644 --- a/fusor-ml/fusor/src/composite/flash_attention.rs +++ b/fusor-ml/fusor/src/composite/flash_attention.rs @@ -203,6 +203,7 @@ where #[cfg(test)] mod tests { use super::*; + use crate::Device; #[tokio::test] async fn test_flash_attention_cpu() { @@ -357,4 +358,79 @@ mod tests { } } } + + #[tokio::test] + async fn test_flash_attention_gpu_matches_cpu_with_causal_mask() { + let gpu_device = Device::new().await.expect("GPU required for this test"); + + let batch = 1; + let heads = 4; + let seq = 9; + let dim = 8; + + let q_data: Vec = (0..(batch * heads * seq * dim)) + .map(|i| ((i % 17) as f32 - 8.0) * 0.05) + .collect(); + let k_data: Vec = (0..(batch * heads * seq * dim)) + .map(|i| ((i % 19) as f32 - 9.0) * 0.04) + .collect(); + let v_data: Vec = (0..(batch * heads * seq * dim)) + .map(|i| ((i % 23) as f32 - 11.0) * 0.03) + .collect(); + + let mut mask_data = vec![0.0f32; seq * seq]; + for q in 0..seq { + for k in (q + 1)..seq { + mask_data[q * seq + k] = f32::NEG_INFINITY; + } + } + + let cpu_q: Tensor<4, f32> = Tensor::Cpu(fusor_cpu::Tensor::from_slice( + [batch, heads, seq, dim], + &q_data, + )); + let cpu_k: Tensor<4, f32> = Tensor::Cpu(fusor_cpu::Tensor::from_slice( + [batch, heads, seq, dim], + &k_data, + )); + let cpu_v: Tensor<4, f32> = Tensor::Cpu(fusor_cpu::Tensor::from_slice( + [batch, heads, seq, dim], + &v_data, + )); + let cpu_mask: Tensor<2, f32> = + Tensor::Cpu(fusor_cpu::Tensor::from_slice([seq, seq], &mask_data)); + + let gpu_q = Tensor::from_slice(&gpu_device, [batch, heads, seq, dim], &q_data); + let gpu_k = Tensor::from_slice(&gpu_device, [batch, heads, seq, dim], &k_data); + let gpu_v = Tensor::from_slice(&gpu_device, [batch, heads, seq, dim], &v_data); + let gpu_mask = Tensor::from_slice(&gpu_device, [seq, seq], &mask_data); + + let scale = 1.0 / (dim as f32).sqrt(); + + let cpu_out = cpu_q + .flash_attention(&cpu_k, &cpu_v, scale, Some((&cpu_mask, MaskKind::QKMask))) + .as_slice() + .await + .unwrap(); + let gpu_out = gpu_q + .flash_attention(&gpu_k, &gpu_v, scale, Some((&gpu_mask, MaskKind::QKMask))) + .as_slice() + .await + .unwrap(); + + for b in 0..batch { + for h in 0..heads { + for s in 0..seq { + for d in 0..dim { + let expected = cpu_out[[b, h, s, d]]; + let actual = gpu_out[[b, h, s, d]]; + assert!( + (expected - actual).abs() < 0.02, + "mismatch at [{b}, {h}, {s}, {d}]: expected {expected}, got {actual}" + ); + } + } + } + } + } } diff --git a/fusor-ml/fusor/src/composite/reductions.rs b/fusor-ml/fusor/src/composite/reductions.rs index 85b673744..7320fe6c6 100644 --- a/fusor-ml/fusor/src/composite/reductions.rs +++ b/fusor-ml/fusor/src/composite/reductions.rs @@ -256,6 +256,61 @@ where } } +impl Tensor +where + B: TensorBacking, +{ + pub fn argmax_last_dim(&self) -> Tensor + where + ConcreteTensor: TensorBacking, + { + assert_eq!(R, OUT_RANK + 1); + + match self { + Tensor::Cpu(t) => { + let concrete = t.as_ref().to_concrete(); + let shape = concrete.shape(); + let reduce_size = shape[R - 1]; + let output_shape: [usize; OUT_RANK] = std::array::from_fn(|i| shape[i]); + let output_elements = output_shape.iter().product(); + let mut output = ConcreteTensor::::zeros(output_shape); + + for linear_index in 0..output_elements { + let mut remaining = linear_index; + let output_index: [usize; OUT_RANK] = std::array::from_fn(|i| { + let stride = output_shape[i + 1..].iter().product::(); + let index = remaining / stride.max(1); + remaining %= stride.max(1); + index + }); + + let mut best_index = 0u32; + let mut best_value = f32::NEG_INFINITY; + for axis_index in 0..reduce_size { + let input_index: [usize; R] = std::array::from_fn(|i| { + if i < OUT_RANK { + output_index[i] + } else { + axis_index + } + }); + let value = concrete.get(input_index); + if value > best_value { + best_value = value; + best_index = axis_index as u32; + } + } + + output.set(output_index, best_index); + } + + Tensor::Cpu(fusor_cpu::Tensor::new(output)) + } + Tensor::Gpu(t) => Tensor::Gpu(t.argmax_last_dim::()), + } + } +} + /// Helper function to unsqueeze a reduced tensor back to original rank with size 1 at the axis. /// The reduced tensor has OUT_RANK dimensions (one less than original R). /// The result has R dimensions with size 1 at the reduced axis (standard keepdim semantics). @@ -436,4 +491,26 @@ mod tests { assert!((slice[[0, 0]] - expected).abs() < 0.001); assert!((slice[[1, 0]] - expected).abs() < 0.001); } + + #[tokio::test] + async fn test_argmax_last_dim_cpu_and_gpu() { + let cpu: Tensor<2, f32> = Tensor::Cpu(fusor_cpu::Tensor::from_slice( + [2, 4], + &[1.0, 5.0, 2.0, 5.0, 3.0, 0.0, 4.0, 1.0], + )); + let cpu_output = cpu.argmax_last_dim::<1>(); + let cpu_output = cpu_output.as_slice().await.unwrap(); + assert_eq!(cpu_output[[0]], 1); + assert_eq!(cpu_output[[1]], 2); + + let device = crate::Device::new() + .await + .expect("GPU required for this test"); + let gpu: Tensor<2, f32> = + Tensor::new(&device, &[[1.0, 5.0, 2.0, 5.0], [3.0, 0.0, 4.0, 1.0]]); + let gpu_output = gpu.argmax_last_dim::<1>(); + let gpu_output = gpu_output.as_slice().await.unwrap(); + assert_eq!(gpu_output[[0]], 1); + assert_eq!(gpu_output[[1]], 2); + } } diff --git a/fusor-ml/fusor/src/composite/rope.rs b/fusor-ml/fusor/src/composite/rope.rs index 989681779..d619780b0 100644 --- a/fusor-ml/fusor/src/composite/rope.rs +++ b/fusor-ml/fusor/src/composite/rope.rs @@ -153,6 +153,35 @@ where } } + /// Apply fused interleaved RoPE to the first `rotary_dim` elements and pass the rest through. + pub fn rope_partial_fused( + &self, + cos: &Tensor<2, D, ConcreteTensor>, + sin: &Tensor<2, D, ConcreteTensor>, + rotary_dim: usize, + ) -> Self { + match (self, cos, sin) { + (Tensor::Gpu(x), Tensor::Gpu(cos), Tensor::Gpu(sin)) => { + Tensor::Gpu(x.rope_partial_fused(cos, sin, rotary_dim)) + } + (Tensor::Cpu(_), Tensor::Cpu(_), Tensor::Cpu(_)) => { + let [_bz, _n_head, _sequence_length, head_dim] = self.shape(); + if rotary_dim == head_dim { + return self.rope_interleaved(cos, sin); + } + let rotated = self + .narrow(3, 0, rotary_dim) + .to_concrete() + .rope_interleaved(cos, sin); + let pass = self + .narrow(3, rotary_dim, head_dim - rotary_dim) + .to_concrete(); + crate::cat([rotated, pass], 3) + } + _ => panic!("All tensors must be on the same device"), + } + } + /// Apply fused normal RoPE (rotary position embedding). /// This pairs first half with second half: (0, head_dim/2), (1, head_dim/2+1), etc. /// @@ -264,6 +293,24 @@ impl RopeCache { (q, k) } + /// Apply interleaved RoPE to a prefix of query and key tensors. + pub fn forward_interleaved_partial( + &self, + q: &Tensor<4, f32>, + k: &Tensor<4, f32>, + start_pos: usize, + rotary_dim: usize, + ) -> (Tensor<4, f32>, Tensor<4, f32>) { + let [_b_sz, _n_head, seq_len, _n_embd] = q.shape(); + let cos = self.cos.narrow(0, start_pos, seq_len).to_concrete(); + let sin = self.sin.narrow(0, start_pos, seq_len).to_concrete(); + + let q = q.rope_partial_fused(&cos, &sin, rotary_dim); + let k = k.rope_partial_fused(&cos, &sin, rotary_dim); + + (q, k) + } + /// Access the raw sin tensor `[context_length, head_dim/2]`. pub fn sin(&self) -> &Tensor<2, f32> { &self.sin diff --git a/fusor-ml/fusor/src/composite/shape.rs b/fusor-ml/fusor/src/composite/shape.rs index bdb7c2960..f3d6fc899 100644 --- a/fusor-ml/fusor/src/composite/shape.rs +++ b/fusor-ml/fusor/src/composite/shape.rs @@ -479,6 +479,21 @@ mod tests { assert_eq!(slice[[1, 1]], 6.0); } + #[tokio::test] + async fn test_to_concrete_gpu_preserves_non_contiguous_view() { + let device = Device::new().await.expect("GPU required for this test"); + let data = [[1.0f32, 2.0, 3.0], [4.0, 5.0, 6.0]]; + let tensor = Tensor::from_slice(&device, [2, 3], data.as_flattened()); + + let concrete = tensor.narrow(1, 1, 2).to_concrete(); + let slice = concrete.as_slice().await.unwrap(); + + assert_eq!(slice[[0, 0]], 2.0); + assert_eq!(slice[[0, 1]], 3.0); + assert_eq!(slice[[1, 0]], 5.0); + assert_eq!(slice[[1, 1]], 6.0); + } + #[tokio::test] async fn test_chunk_cpu() { let data = [1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0]; diff --git a/fusor-ml/fusor/src/layers/batch_norm.rs b/fusor-ml/fusor/src/layers/batch_norm.rs new file mode 100644 index 000000000..922281936 --- /dev/null +++ b/fusor-ml/fusor/src/layers/batch_norm.rs @@ -0,0 +1,135 @@ +//! BatchNorm1d layer implementation for inference. + +use crate::{DataType, Device, SimdElement, Tensor, VarBuilder}; +use fusor_core::FloatDataType; +use fusor_cpu::{FloatOps, TensorBacking}; + +/// Inference-only BatchNorm1d. +/// +/// Input shape: (batch, channels, length) +pub struct BatchNorm1d { + weight: Option>, + bias: Option>, + running_mean: Tensor<1, D>, + running_var: Tensor<1, D>, + eps: D, +} + +impl BatchNorm1d +where + D: SimdElement + + DataType + + FloatDataType + + FloatOps + + Default + + std::ops::Add + + std::ops::Sub + + std::ops::Mul + + std::ops::Div, + crate::AddOp: fusor_cpu::SimdBinaryOp, + crate::SubOp: fusor_cpu::SimdBinaryOp, + crate::MulOp: fusor_cpu::SimdBinaryOp, + crate::DivOp: fusor_cpu::SimdBinaryOp, + crate::SqrtOp: fusor_cpu::SimdUnaryOp, +{ + /// Create a new BatchNorm1d layer. + pub fn new( + weight: Option>, + bias: Option>, + running_mean: Tensor<1, D>, + running_var: Tensor<1, D>, + eps: D, + ) -> Self { + Self { + weight, + bias, + running_mean, + running_var, + eps, + } + } + + /// Forward pass. + pub fn forward(&self, input: &Tensor<3, D, B>) -> Tensor<3, D> + where + B: TensorBacking<3, Elem = D>, + { + let shape = input.shape(); + let channels = shape[1]; + assert_eq!( + self.running_mean.shape()[0], + channels, + "running_mean channels ({}) must match input channels ({channels})", + self.running_mean.shape()[0] + ); + assert_eq!( + self.running_var.shape()[0], + channels, + "running_var channels ({}) must match input channels ({channels})", + self.running_var.shape()[0] + ); + + let mean_reshaped = self.running_mean.reshape([1, channels, 1]); + let mean = mean_reshaped.broadcast_as(shape); + let var_reshaped = self.running_var.reshape([1, channels, 1]); + let var = var_reshaped.broadcast_as(shape); + let normalized = (input.to_concrete() - mean).to_concrete().div_( + &(var + .to_concrete() + .add_scalar(self.eps) + .sqrt() + .broadcast_as(shape)), + ); + + let scaled = if let Some(weight) = &self.weight { + let weight_reshaped = weight.reshape([1, channels, 1]); + let weight = weight_reshaped.broadcast_as(shape); + normalized.mul_(&weight) + } else { + normalized + }; + + if let Some(bias) = &self.bias { + let bias_reshaped = bias.reshape([1, channels, 1]); + let bias = bias_reshaped.broadcast_as(shape); + scaled.add_(&bias) + } else { + scaled + } + } +} + +impl BatchNorm1d { + /// Load a BatchNorm1d layer from a GGUF var builder. + pub fn load(device: &Device, vb: &mut VarBuilder, eps: f32) -> crate::Result { + let weight = vb.get("weight", device).ok().map(|w| w.dequantize()); + let bias = vb.get("bias", device).ok().map(|b| b.dequantize()); + let running_mean = vb.get("running_mean", device)?.dequantize(); + let running_var = vb.get("running_var", device)?.dequantize(); + Ok(Self::new(weight, bias, running_mean, running_var, eps)) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn test_batch_norm_1d_inference() { + let device = Device::Cpu; + let input: Tensor<3, f32> = Tensor::from_slice(&device, [1, 2, 2], &[1.0, 3.0, 10.0, 14.0]); + let weight: Tensor<1, f32> = Tensor::from_slice(&device, [2], &[2.0, 0.5]); + let bias: Tensor<1, f32> = Tensor::from_slice(&device, [2], &[1.0, -1.0]); + let mean: Tensor<1, f32> = Tensor::from_slice(&device, [2], &[2.0, 12.0]); + let var: Tensor<1, f32> = Tensor::from_slice(&device, [2], &[4.0, 16.0]); + + let bn = BatchNorm1d::new(Some(weight), Some(bias), mean, var, 0.0); + let output = bn.forward(&input); + let result = output.as_slice().await.unwrap(); + + assert!((result[[0, 0, 0]] - 0.0).abs() < 1e-5); + assert!((result[[0, 0, 1]] - 2.0).abs() < 1e-5); + assert!((result[[0, 1, 0]] - -1.25).abs() < 1e-5); + assert!((result[[0, 1, 1]] - -0.75).abs() < 1e-5); + } +} diff --git a/fusor-ml/fusor/src/layers/conv1d.rs b/fusor-ml/fusor/src/layers/conv1d.rs index 11ff18d76..3834faff1 100644 --- a/fusor-ml/fusor/src/layers/conv1d.rs +++ b/fusor-ml/fusor/src/layers/conv1d.rs @@ -50,6 +50,16 @@ where + std::ops::Mul + std::ops::Add, { + /// Read-only access to the loaded weight tensor (out_channels, in_channels, kernel_size). + pub fn weight(&self) -> &Tensor<3, D, ConcreteTensor> { + &self.weight + } + + /// Read-only access to the loaded bias tensor. + pub fn bias(&self) -> Option<&Tensor<1, D, ConcreteTensor>> { + self.bias.as_ref() + } + /// Create a new Conv1d layer with given weights and configuration. /// /// Weight shape: (out_channels, in_channels, kernel_size) @@ -61,12 +71,23 @@ where ) -> Self { let shape = weight.shape(); let out_channels = shape[0]; - let in_channels = shape[1]; + let in_channels_per_group = shape[1]; let kernel_size = shape[2]; + let in_channels = in_channels_per_group * config.groups; // Validate configuration - assert_eq!(config.groups, 1, "Only groups=1 is currently supported"); assert_eq!(config.dilation, 1, "Only dilation=1 is currently supported"); + assert!( + config.groups > 0, + "groups must be greater than zero, got {}", + config.groups + ); + assert_eq!( + out_channels % config.groups, + 0, + "out_channels ({out_channels}) must be divisible by groups ({})", + config.groups + ); if let Some(ref b) = bias { assert_eq!( @@ -99,12 +120,61 @@ where crate::AddOp: fusor_cpu::SimdBinaryOp, fusor_cpu::SumOp: fusor_cpu::SimdReduceOp, { - input.conv( - &self.weight, - self.bias.as_ref(), - [self.config.padding], - [self.config.stride], - ) + if let Tensor::Gpu(input) = input { + let Tensor::Gpu(weight) = &self.weight else { + panic!("Conv1d input and weight must be on the same device"); + }; + let bias = self.bias.as_ref().map(|bias| { + let Tensor::Gpu(bias) = bias else { + panic!("Conv1d input and bias must be on the same device"); + }; + bias + }); + + return Tensor::Gpu(input.conv1d_grouped( + weight, + bias, + self.config.padding, + self.config.stride, + self.config.groups, + )); + } + + if self.config.groups == 1 { + return input.conv( + &self.weight, + self.bias.as_ref(), + [self.config.padding], + [self.config.stride], + ); + } + + let in_channels_per_group = self.in_channels / self.config.groups; + let out_channels_per_group = self.out_channels / self.config.groups; + let mut outputs = Vec::with_capacity(self.config.groups); + + for group in 0..self.config.groups { + let input_group = input + .narrow(1, group * in_channels_per_group, in_channels_per_group) + .to_concrete(); + let weight_group = self + .weight + .narrow(0, group * out_channels_per_group, out_channels_per_group) + .to_concrete(); + let bias_group = self.bias.as_ref().map(|bias| { + bias.narrow(0, group * out_channels_per_group, out_channels_per_group) + .to_concrete() + }); + + outputs.push(input_group.conv( + &weight_group, + bias_group.as_ref(), + [self.config.padding], + [self.config.stride], + )); + } + + Tensor::cat(outputs, 1) } /// Get the configuration. @@ -219,6 +289,100 @@ mod tests { assert_eq!(conv.config().stride, 3); } + #[tokio::test] + async fn test_conv1d_depthwise_groups() { + let config = Conv1dConfig { + padding: 1, + stride: 1, + groups: 2, + dilation: 1, + }; + + let input_data = [ + 1.0f32, 2.0, 3.0, 4.0, // + 10.0, 20.0, 30.0, 40.0, + ]; + let input: Tensor<3, f32> = + Tensor::Cpu(fusor_cpu::Tensor::from_slice([1, 2, 4], &input_data)); + + let weight_data = [ + 1.0f32, 0.0, -1.0, // + 0.5, 0.0, -0.5, + ]; + let weight: Tensor<3, f32> = + Tensor::Cpu(fusor_cpu::Tensor::from_slice([2, 1, 3], &weight_data)); + + let conv = Conv1d::new(weight, None, config); + let output = conv.forward(&input); + let result = output.as_slice().await.unwrap(); + + assert_eq!(result.shape(), &[1, 2, 4]); + assert!((result[[0, 0, 0]] - -2.0).abs() < 1e-5); + assert!((result[[0, 0, 1]] - -2.0).abs() < 1e-5); + assert!((result[[0, 0, 2]] - -2.0).abs() < 1e-5); + assert!((result[[0, 0, 3]] - 3.0).abs() < 1e-5); + assert!((result[[0, 1, 0]] - -10.0).abs() < 1e-5); + assert!((result[[0, 1, 1]] - -10.0).abs() < 1e-5); + assert!((result[[0, 1, 2]] - -10.0).abs() < 1e-5); + assert!((result[[0, 1, 3]] - 15.0).abs() < 1e-5); + } + + #[tokio::test] + async fn test_conv1d_depthwise_groups_cpu_vs_gpu_single_kernel() { + use crate::Device; + + let config = Conv1dConfig { + padding: 1, + stride: 1, + groups: 2, + dilation: 1, + }; + + let input_data = [ + 1.0f32, 2.0, 3.0, 4.0, // + 10.0, 20.0, 30.0, 40.0, + ]; + let weight_data = [ + 1.0f32, 0.0, -1.0, // + 0.5, 0.25, -0.5, + ]; + let bias_data = [0.5f32, -1.0]; + + let input_cpu: Tensor<3, f32> = + Tensor::Cpu(fusor_cpu::Tensor::from_slice([1, 2, 4], &input_data)); + let weight_cpu: Tensor<3, f32> = + Tensor::Cpu(fusor_cpu::Tensor::from_slice([2, 1, 3], &weight_data)); + let bias_cpu: Tensor<1, f32> = Tensor::Cpu(fusor_cpu::Tensor::from_slice([2], &bias_data)); + let conv_cpu = Conv1d::new(weight_cpu, Some(bias_cpu), config); + let output_cpu = conv_cpu.forward(&input_cpu); + let result_cpu = output_cpu.as_slice().await.unwrap(); + + let gpu_device = Device::new().await.expect("GPU required for this test"); + let input_gpu: Tensor<3, f32> = Tensor::from_slice(&gpu_device, [1, 2, 4], &input_data); + let weight_gpu: Tensor<3, f32> = Tensor::from_slice(&gpu_device, [2, 1, 3], &weight_data); + let bias_gpu: Tensor<1, f32> = Tensor::from_slice(&gpu_device, [2], &bias_data); + let conv_gpu = Conv1d::new(weight_gpu, Some(bias_gpu), config); + let output_gpu = conv_gpu.forward(&input_gpu); + if let Tensor::Gpu(output) = &output_gpu { + assert_eq!(output.count_kernels_to_resolve(), 1); + } + let result_gpu = output_gpu.as_slice().await.unwrap(); + + assert_eq!(result_cpu.shape(), result_gpu.shape()); + for batch in 0..1 { + for channel in 0..2 { + for position in 0..4 { + let cpu_val = result_cpu[[batch, channel, position]]; + let gpu_val = result_gpu[[batch, channel, position]]; + assert!( + (cpu_val - gpu_val).abs() < 1e-5, + "Mismatch at [{batch}, {channel}, {position}]: cpu={cpu_val}, gpu={gpu_val}" + ); + } + } + } + } + #[tokio::test] async fn test_conv1d_cpu_vs_gpu() { use crate::Device; diff --git a/fusor-ml/fusor/src/layers/conv2d.rs b/fusor-ml/fusor/src/layers/conv2d.rs new file mode 100644 index 000000000..6b50fcbee --- /dev/null +++ b/fusor-ml/fusor/src/layers/conv2d.rs @@ -0,0 +1,373 @@ +//! Conv2d layer implementation. + +use crate::{ConcreteTensor, MatmulImpl, SimdElement, Tensor}; +use fusor_core::{DataType, FloatDataType}; +use fusor_cpu::FloatOps; + +/// Configuration for Conv2d layers. +#[derive(Debug, Clone, Copy)] +pub struct Conv2dConfig { + pub padding: [usize; 2], + pub stride: [usize; 2], + pub groups: usize, + pub dilation: [usize; 2], +} + +impl Default for Conv2dConfig { + fn default() -> Self { + Self { + padding: [0, 0], + stride: [1, 1], + groups: 1, + dilation: [1, 1], + } + } +} + +/// 2D Convolution layer. +pub struct Conv2d { + weight: Tensor<4, D, ConcreteTensor>, // (out_channels, in_channels/groups, kernel_h, kernel_w) + bias: Option>>, // (out_channels,) + config: Conv2dConfig, + in_channels: usize, + out_channels: usize, + kernel_size: [usize; 2], +} + +impl Conv2d +where + D: SimdElement + + DataType + + FloatDataType + + FloatOps + + Default + + MatmulImpl + + std::ops::Mul + + std::ops::Add, +{ + /// Create a new Conv2d layer with given weights and configuration. + pub fn new( + weight: Tensor<4, D, ConcreteTensor>, + bias: Option>>, + config: Conv2dConfig, + ) -> Self { + let shape = weight.shape(); + let out_channels = shape[0]; + let in_channels_per_group = shape[1]; + let kernel_size = [shape[2], shape[3]]; + + assert!( + config.groups > 0, + "groups must be greater than zero, got {}", + config.groups + ); + assert_eq!( + config.dilation, + [1, 1], + "Only dilation=[1, 1] is currently supported" + ); + if let Some(ref b) = bias { + assert_eq!( + b.shape()[0], + out_channels, + "Bias shape must match out_channels" + ); + } + + Self { + weight, + bias, + config, + in_channels: in_channels_per_group * config.groups, + out_channels, + kernel_size, + } + } + + /// Forward pass. + pub fn forward( + &self, + input: &Tensor<4, D, ConcreteTensor>, + ) -> Tensor<4, D, ConcreteTensor> + where + crate::MulOp: fusor_cpu::SimdBinaryOp, + crate::AddOp: fusor_cpu::SimdBinaryOp, + fusor_cpu::SumOp: fusor_cpu::SimdReduceOp, + { + if self.config.groups == 1 { + return input.conv( + &self.weight, + self.bias.as_ref(), + self.config.padding, + self.config.stride, + ); + } + + let in_channels_per_group = self.in_channels / self.config.groups; + let out_channels_per_group = self.out_channels / self.config.groups; + let mut outputs = Vec::with_capacity(self.config.groups); + + for group in 0..self.config.groups { + let input_group = input + .narrow(1, group * in_channels_per_group, in_channels_per_group) + .to_concrete(); + let weight_group = self + .weight + .narrow(0, group * out_channels_per_group, out_channels_per_group) + .to_concrete(); + let bias_group = self.bias.as_ref().map(|bias| { + bias.narrow(0, group * out_channels_per_group, out_channels_per_group) + .to_concrete() + }); + + outputs.push(input_group.conv( + &weight_group, + bias_group.as_ref(), + self.config.padding, + self.config.stride, + )); + } + + Tensor::cat(outputs, 1) + } + + /// Get the configuration. + pub fn config(&self) -> &Conv2dConfig { + &self.config + } + + /// Get the number of input channels. + pub fn in_channels(&self) -> usize { + self.in_channels + } + + /// Get the number of output channels. + pub fn out_channels(&self) -> usize { + self.out_channels + } + + /// Get the kernel size. + pub fn kernel_size(&self) -> [usize; 2] { + self.kernel_size + } +} + +impl Conv2d { + /// Load a Conv2d layer from a GGUF var builder. + pub fn load( + device: &crate::Device, + vb: &mut crate::VarBuilder, + config: Conv2dConfig, + ) -> crate::Result { + let weight = vb.get("weight", device)?.dequantize(); + let bias = vb.get("bias", device).ok().map(|b| b.dequantize()); + Ok(Self::new(weight, bias, config)) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn test_conv2d_simple() { + let input: Tensor<4, f32> = Tensor::Cpu(fusor_cpu::Tensor::from_slice( + [1, 1, 3, 3], + &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0], + )); + let weight: Tensor<4, f32> = Tensor::Cpu(fusor_cpu::Tensor::from_slice( + [1, 1, 2, 2], + &[1.0, 0.0, 0.0, 1.0], + )); + + let conv = Conv2d::new(weight, None, Conv2dConfig::default()); + let output = conv.forward(&input); + let result = output.as_slice().await.unwrap(); + + assert_eq!(result.shape(), &[1, 1, 2, 2]); + assert!((result[[0, 0, 0, 0]] - 6.0).abs() < 1e-5); + assert!((result[[0, 0, 0, 1]] - 8.0).abs() < 1e-5); + assert!((result[[0, 0, 1, 0]] - 12.0).abs() < 1e-5); + assert!((result[[0, 0, 1, 1]] - 14.0).abs() < 1e-5); + } + + #[tokio::test] + async fn test_conv2d_depthwise_groups() { + let input: Tensor<4, f32> = Tensor::Cpu(fusor_cpu::Tensor::from_slice( + [1, 2, 3, 3], + &[ + 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, // + 10.0, 11.0, 12.0, 13.0, 14.0, 15.0, 16.0, 17.0, 18.0, + ], + )); + let weight: Tensor<4, f32> = Tensor::Cpu(fusor_cpu::Tensor::from_slice( + [2, 1, 2, 2], + &[ + 1.0, 0.0, 0.0, 1.0, // + 0.5, 0.0, 0.0, 0.5, + ], + )); + + let conv = Conv2d::new( + weight, + None, + Conv2dConfig { + groups: 2, + ..Default::default() + }, + ); + let output = conv.forward(&input); + let result = output.as_slice().await.unwrap(); + + assert_eq!(result.shape(), &[1, 2, 2, 2]); + assert!((result[[0, 0, 0, 0]] - 6.0).abs() < 1e-5); + assert!((result[[0, 0, 1, 1]] - 14.0).abs() < 1e-5); + assert!((result[[0, 1, 0, 0]] - 12.0).abs() < 1e-5); + assert!((result[[0, 1, 1, 1]] - 16.0).abs() < 1e-5); + } + + #[tokio::test] + async fn test_conv2d_asymmetric_kernel_order() { + let input: Tensor<4, f32> = Tensor::Cpu(fusor_cpu::Tensor::from_slice( + [1, 1, 3, 4], + &[ + 1.0, 2.0, 3.0, 4.0, // + 5.0, 6.0, 7.0, 8.0, // + 9.0, 10.0, 11.0, 12.0, + ], + )); + let weight: Tensor<4, f32> = Tensor::Cpu(fusor_cpu::Tensor::from_slice( + [1, 1, 2, 2], + &[ + 1.0, 2.0, // + 3.0, 4.0, + ], + )); + + let conv = Conv2d::new(weight, None, Conv2dConfig::default()); + let output = conv.forward(&input); + let result = output.as_slice().await.unwrap(); + + assert_eq!(result.shape(), &[1, 1, 2, 3]); + assert!((result[[0, 0, 0, 0]] - 44.0).abs() < 1e-5); + assert!((result[[0, 0, 0, 1]] - 54.0).abs() < 1e-5); + assert!((result[[0, 0, 0, 2]] - 64.0).abs() < 1e-5); + assert!((result[[0, 0, 1, 0]] - 84.0).abs() < 1e-5); + assert!((result[[0, 0, 1, 1]] - 94.0).abs() < 1e-5); + assert!((result[[0, 0, 1, 2]] - 104.0).abs() < 1e-5); + } + + #[tokio::test] + async fn test_conv2d_gpu_matches_cpu_stride2_padding1() { + let gpu = crate::Device::new() + .await + .expect("GPU required for this test"); + + let input_data = vec![ + 1.0f32, 2.0, 3.0, 4.0, // + 5.0, 6.0, 7.0, 8.0, // + 9.0, 10.0, 11.0, 12.0, // + 13.0, 14.0, 15.0, 16.0, + ]; + let weight_data = vec![1.0f32, 0.0, 0.0, 1.0]; + let bias_data = vec![0.25f32]; + + let cpu_input: Tensor<4, f32> = + Tensor::Cpu(fusor_cpu::Tensor::from_slice([1, 1, 4, 4], &input_data)); + let gpu_input: Tensor<4, f32> = Tensor::from_slice(&gpu, [1, 1, 4, 4], &input_data); + let cpu_weight: Tensor<4, f32> = + Tensor::Cpu(fusor_cpu::Tensor::from_slice([1, 1, 2, 2], &weight_data)); + let gpu_weight: Tensor<4, f32> = Tensor::from_slice(&gpu, [1, 1, 2, 2], &weight_data); + let cpu_bias: Tensor<1, f32> = Tensor::Cpu(fusor_cpu::Tensor::from_slice([1], &bias_data)); + let gpu_bias: Tensor<1, f32> = Tensor::from_slice(&gpu, [1], &bias_data); + + let config = Conv2dConfig { + padding: [1, 1], + stride: [2, 2], + groups: 1, + dilation: [1, 1], + }; + + let cpu_out = Conv2d::new(cpu_weight, Some(cpu_bias), config) + .forward(&cpu_input) + .as_slice() + .await + .unwrap(); + let gpu_out = Conv2d::new(gpu_weight, Some(gpu_bias), config) + .forward(&gpu_input) + .as_slice() + .await + .unwrap(); + + for h in 0..cpu_out.shape()[2] { + for w in 0..cpu_out.shape()[3] { + let expected = cpu_out[[0, 0, h, w]]; + let actual = gpu_out[[0, 0, h, w]]; + assert!( + (expected - actual).abs() < 1e-5, + "mismatch at [{h}, {w}]: expected {expected}, got {actual}" + ); + } + } + } + + #[tokio::test] + async fn test_conv2d_gpu_matches_cpu_depthwise_stride2_padding1() { + let gpu = crate::Device::new() + .await + .expect("GPU required for this test"); + + let input_data = vec![ + 1.0f32, 2.0, 3.0, 4.0, // + 5.0, 6.0, 7.0, 8.0, // + 9.0, 10.0, 11.0, 12.0, // + 13.0, 14.0, 15.0, 16.0, // + 17.0, 18.0, 19.0, 20.0, // + 21.0, 22.0, 23.0, 24.0, // + 25.0, 26.0, 27.0, 28.0, // + 29.0, 30.0, 31.0, 32.0, + ]; + let weight_data = vec![ + 1.0f32, 0.0, 0.0, 1.0, // + 0.5, 0.0, 0.0, 0.5, + ]; + + let cpu_input: Tensor<4, f32> = + Tensor::Cpu(fusor_cpu::Tensor::from_slice([1, 2, 4, 4], &input_data)); + let gpu_input: Tensor<4, f32> = Tensor::from_slice(&gpu, [1, 2, 4, 4], &input_data); + let cpu_weight: Tensor<4, f32> = + Tensor::Cpu(fusor_cpu::Tensor::from_slice([2, 1, 2, 2], &weight_data)); + let gpu_weight: Tensor<4, f32> = Tensor::from_slice(&gpu, [2, 1, 2, 2], &weight_data); + + let config = Conv2dConfig { + padding: [1, 1], + stride: [2, 2], + groups: 2, + dilation: [1, 1], + }; + + let cpu_out = Conv2d::new(cpu_weight, None, config) + .forward(&cpu_input) + .as_slice() + .await + .unwrap(); + let gpu_out = Conv2d::new(gpu_weight, None, config) + .forward(&gpu_input) + .as_slice() + .await + .unwrap(); + + for c in 0..cpu_out.shape()[1] { + for h in 0..cpu_out.shape()[2] { + for w in 0..cpu_out.shape()[3] { + let expected = cpu_out[[0, c, h, w]]; + let actual = gpu_out[[0, c, h, w]]; + assert!( + (expected - actual).abs() < 1e-5, + "mismatch at [{c}, {h}, {w}]: expected {expected}, got {actual}" + ); + } + } + } + } +} diff --git a/fusor-ml/fusor/src/layers/layer_norm.rs b/fusor-ml/fusor/src/layers/layer_norm.rs index a8b48aad2..1db2fe8d6 100644 --- a/fusor-ml/fusor/src/layers/layer_norm.rs +++ b/fusor-ml/fusor/src/layers/layer_norm.rs @@ -144,10 +144,16 @@ impl LayerNorm<1, f32> { ); Tensor::Cpu(fusor_cpu::Tensor::new(result)) } - Tensor::Gpu(_) => { - // Fall back to standard forward for GPU - self.forward(input) - } + Tensor::Gpu(input) => match &self.weight { + Tensor::Gpu(weight) => { + let gpu_bias = self.bias.as_ref().map(|bias| match bias { + Tensor::Gpu(bias) => bias, + _ => panic!("LayerNorm bias must be on GPU when input is on GPU"), + }); + Tensor::Gpu(input.layer_norm_fused(weight, gpu_bias, self.eps)) + } + _ => panic!("LayerNorm weight must be on GPU when input is on GPU"), + }, } } @@ -244,4 +250,194 @@ mod tests { assert!((result[[0, 0, 0]] - (-1.0)).abs() < 1e-4); assert!((result[[0, 0, 1]] - 1.0).abs() < 1e-4); } + + #[tokio::test] + async fn test_layer_norm_forward_fused_gpu_matches_cpu() { + let gpu_device = Device::new().await.expect("GPU required for this test"); + let input_data = [ + 1.0f32, 2.0, 3.0, 4.0, 5.0, 7.0, 11.0, 13.0, -4.0, -1.0, 0.5, 3.0, + ]; + let weight_data = [1.0f32, 0.5, 2.0, -1.0]; + let bias_data = [0.0f32, 0.25, -0.5, 1.0]; + + let cpu_input: Tensor<3, f32> = + Tensor::Cpu(fusor_cpu::Tensor::from_slice([1, 3, 4], &input_data)); + let gpu_input: Tensor<3, f32> = Tensor::from_slice(&gpu_device, [1, 3, 4], &input_data); + + let cpu_weight: Tensor<1, f32> = + Tensor::Cpu(fusor_cpu::Tensor::from_slice([4], &weight_data)); + let gpu_weight: Tensor<1, f32> = Tensor::from_slice(&gpu_device, [4], &weight_data); + let cpu_bias: Tensor<1, f32> = Tensor::Cpu(fusor_cpu::Tensor::from_slice([4], &bias_data)); + let gpu_bias: Tensor<1, f32> = Tensor::from_slice(&gpu_device, [4], &bias_data); + + let cpu_layer_norm = LayerNorm::new(cpu_weight, Some(cpu_bias), 1e-5); + let gpu_layer_norm = LayerNorm::new(gpu_weight, Some(gpu_bias), 1e-5); + + let cpu_output = cpu_layer_norm.forward_fused(&cpu_input); + let gpu_output = gpu_layer_norm.forward_fused(&gpu_input); + let cpu_output = cpu_output.as_slice().await.unwrap(); + let gpu_output = gpu_output.as_slice().await.unwrap(); + + for (cpu, gpu) in cpu_output.as_slice().iter().zip(gpu_output.as_slice()) { + let diff = (cpu - gpu).abs(); + assert!( + diff < 1e-4, + "CPU/GPU fused LayerNorm mismatch: {cpu} vs {gpu}" + ); + } + } + + #[tokio::test] + async fn test_layer_norm_forward_fused_gpu_matches_cpu_non_contiguous_no_bias() { + let gpu_device = Device::new().await.expect("GPU required for this test"); + let input_data = [ + 1.0f32, 2.0, 3.0, 4.0, 5.0, 7.0, 11.0, 13.0, -4.0, -1.0, 0.5, 3.0, 8.0, 6.0, 4.0, 2.0, + 0.25, 0.5, 1.0, 2.0, 3.0, 1.5, 0.75, 0.25, 9.0, 10.0, 12.0, 15.0, -8.0, -4.0, -2.0, + -1.0, + ]; + let weight_data = [1.0f32, 1.0, 1.0, 1.0]; + + let cpu_input = Tensor::Cpu(fusor_cpu::Tensor::from_slice([2, 4, 4], &input_data)); + let cpu_input = cpu_input.narrow(1, 1, 2); + let gpu_input_full: Tensor<3, f32> = + Tensor::from_slice(&gpu_device, [2, 4, 4], &input_data); + let gpu_input = gpu_input_full.narrow(1, 1, 2); + + let cpu_weight: Tensor<1, f32> = + Tensor::Cpu(fusor_cpu::Tensor::from_slice([4], &weight_data)); + let gpu_weight: Tensor<1, f32> = Tensor::from_slice(&gpu_device, [4], &weight_data); + + let cpu_layer_norm = LayerNorm::new(cpu_weight, None, 1e-5); + let gpu_layer_norm = LayerNorm::new(gpu_weight, None, 1e-5); + + let cpu_output = cpu_layer_norm.forward_fused(&cpu_input); + let gpu_output = gpu_layer_norm.forward_fused(&gpu_input); + let cpu_output = cpu_output.as_slice().await.unwrap(); + let gpu_output = gpu_output.as_slice().await.unwrap(); + + for (cpu, gpu) in cpu_output.as_slice().iter().zip(gpu_output.as_slice()) { + let diff = (cpu - gpu).abs(); + assert!( + diff < 1e-4, + "CPU/GPU fused LayerNorm mismatch: {cpu} vs {gpu}" + ); + } + } + + #[tokio::test] + async fn test_layer_norm_forward_fused_gpu_matches_cpu_hidden_320() { + let gpu_device = Device::new().await.expect("GPU required for this test"); + let hidden = 320; + let input_data: Vec = (0..2 * 3 * hidden) + .map(|i| (((i * 13) % 41) as f32 - 20.0) * 0.0625) + .collect(); + let weight_data: Vec = (0..hidden) + .map(|i| 0.85 + (i % 11) as f32 * 0.0125) + .collect(); + + let cpu_input: Tensor<3, f32> = + Tensor::Cpu(fusor_cpu::Tensor::from_slice([2, 3, hidden], &input_data)); + let gpu_input: Tensor<3, f32> = + Tensor::from_slice(&gpu_device, [2, 3, hidden], &input_data); + let cpu_weight: Tensor<1, f32> = + Tensor::Cpu(fusor_cpu::Tensor::from_slice([hidden], &weight_data)); + let gpu_weight: Tensor<1, f32> = Tensor::from_slice(&gpu_device, [hidden], &weight_data); + + let cpu_layer_norm = LayerNorm::new(cpu_weight, None, 1e-5); + let gpu_layer_norm = LayerNorm::new(gpu_weight, None, 1e-5); + + let cpu_output = cpu_layer_norm.forward_fused(&cpu_input); + let gpu_output = gpu_layer_norm.forward_fused(&gpu_input); + let cpu_output = cpu_output.as_slice().await.unwrap(); + let gpu_output = gpu_output.as_slice().await.unwrap(); + + for (cpu, gpu) in cpu_output.as_slice().iter().zip(gpu_output.as_slice()) { + let diff = (cpu - gpu).abs(); + assert!( + diff < 1e-4, + "CPU/GPU fused LayerNorm mismatch: {cpu} vs {gpu}" + ); + } + } + + #[tokio::test] + async fn test_layer_norm_forward_fused_gpu_unit_offset_pattern() { + let gpu_device = Device::new().await.expect("GPU required for this test"); + let hidden = 320; + let input_data: Vec = (0..2 * 3 * hidden) + .map(|i| (((i * 17) % 53) as f32 - 26.0) * 0.03125) + .collect(); + let ones = vec![1.0f32; hidden]; + let gamma_data: Vec = (0..hidden) + .map(|i| ((i * 7) % 29) as f32 * 0.01 - 0.14) + .collect(); + + let cpu_input: Tensor<3, f32> = + Tensor::Cpu(fusor_cpu::Tensor::from_slice([2, 3, hidden], &input_data)); + let gpu_input: Tensor<3, f32> = + Tensor::from_slice(&gpu_device, [2, 3, hidden], &input_data); + let cpu_weight: Tensor<1, f32> = + Tensor::Cpu(fusor_cpu::Tensor::from_slice([hidden], &ones)); + let gpu_weight: Tensor<1, f32> = Tensor::from_slice(&gpu_device, [hidden], &ones); + let cpu_gamma: Tensor<1, f32> = + Tensor::Cpu(fusor_cpu::Tensor::from_slice([hidden], &gamma_data)); + let gpu_gamma: Tensor<1, f32> = Tensor::from_slice(&gpu_device, [hidden], &gamma_data); + + let cpu_layer_norm = LayerNorm::new(cpu_weight, None, 1e-5); + let gpu_layer_norm = LayerNorm::new(gpu_weight, None, 1e-5); + + let cpu_normed = cpu_layer_norm.forward_fused(&cpu_input); + let cpu_gamma_offset = cpu_gamma.add_scalar(1.0); + let cpu_gamma_broadcast = cpu_gamma_offset.broadcast_as(cpu_normed.shape()); + let cpu_output = cpu_normed.mul_(&cpu_gamma_broadcast).to_concrete(); + let gpu_normed = gpu_layer_norm.forward_fused(&gpu_input); + let gpu_gamma_offset = gpu_gamma.add_scalar(1.0); + let gpu_gamma_broadcast = gpu_gamma_offset.broadcast_as(gpu_normed.shape()); + let gpu_output = gpu_normed.mul_(&gpu_gamma_broadcast).to_concrete(); + + let cpu_output = cpu_output.as_slice().await.unwrap(); + let gpu_output = gpu_output.as_slice().await.unwrap(); + + for (cpu, gpu) in cpu_output.as_slice().iter().zip(gpu_output.as_slice()) { + let diff = (cpu - gpu).abs(); + assert!( + diff < 1e-4, + "CPU/GPU unit-offset LayerNorm mismatch: {cpu} vs {gpu}" + ); + } + } + + #[tokio::test] + async fn test_layer_norm_forward_fused_gpu_low_variance_hidden_320() { + let gpu_device = Device::new().await.expect("GPU required for this test"); + let hidden = 320; + let input_data: Vec = (0..2 * 3 * hidden) + .map(|i| 1.0 + (((i * 13) % 17) as f32 - 8.0) * 1e-4) + .collect(); + let weight_data = vec![1.0f32; hidden]; + + let cpu_input: Tensor<3, f32> = + Tensor::Cpu(fusor_cpu::Tensor::from_slice([2, 3, hidden], &input_data)); + let gpu_input: Tensor<3, f32> = + Tensor::from_slice(&gpu_device, [2, 3, hidden], &input_data); + let cpu_weight: Tensor<1, f32> = + Tensor::Cpu(fusor_cpu::Tensor::from_slice([hidden], &weight_data)); + let gpu_weight: Tensor<1, f32> = Tensor::from_slice(&gpu_device, [hidden], &weight_data); + + let cpu_layer_norm = LayerNorm::new(cpu_weight, None, 1e-5); + let gpu_layer_norm = LayerNorm::new(gpu_weight, None, 1e-5); + + let cpu_output = cpu_layer_norm.forward_fused(&cpu_input); + let gpu_output = gpu_layer_norm.forward_fused(&gpu_input); + let cpu_output = cpu_output.as_slice().await.unwrap(); + let gpu_output = gpu_output.as_slice().await.unwrap(); + + for (cpu, gpu) in cpu_output.as_slice().iter().zip(gpu_output.as_slice()) { + let diff = (cpu - gpu).abs(); + assert!( + diff < 1e-4, + "CPU/GPU low-variance fused LayerNorm mismatch: {cpu} vs {gpu}" + ); + } + } } diff --git a/fusor-ml/fusor/src/layers/linear.rs b/fusor-ml/fusor/src/layers/linear.rs index 6bec403b5..9ae749761 100644 --- a/fusor-ml/fusor/src/layers/linear.rs +++ b/fusor-ml/fusor/src/layers/linear.rs @@ -30,6 +30,11 @@ impl Linear { self.bias.as_ref() } + /// Get the quantized weight matrix. + pub fn weight(&self) -> &QMatrix { + &self.weight + } + /// Get the input features size. pub fn in_features(&self) -> usize { self.weight.shape()[1] @@ -54,6 +59,14 @@ impl Linear { // f32-specific implementations for loading and forward impl Linear { + /// Forward pass without applying bias. + pub fn forward_no_bias(&self, input: &Tensor<3, f32, B>) -> Tensor<3, f32> + where + B: fusor_cpu::TensorBacking<3, Elem = f32>, + { + input.q_mat_mul(&self.weight) + } + /// Load a Linear layer from a VarBuilder. /// /// Expects: @@ -73,14 +86,13 @@ impl Linear { where B: fusor_cpu::TensorBacking<3, Elem = f32>, { - let output = input.q_mat_mul(&self.weight); - if let Some(bias) = &self.bias { - output.add_(bias) + input.q_mat_mul_bias(&self.weight, bias) } else { - output + self.forward_no_bias(input) } } + } // Generic forward implementations for Linear where T can be cast to/from f32 diff --git a/fusor-ml/fusor/src/layers/mod.rs b/fusor-ml/fusor/src/layers/mod.rs index 9477c92bf..b1ce04e79 100644 --- a/fusor-ml/fusor/src/layers/mod.rs +++ b/fusor-ml/fusor/src/layers/mod.rs @@ -4,13 +4,17 @@ //! //! All layers support loading from GGUF files via `VarBuilder` for f32 types. +mod batch_norm; mod conv1d; +mod conv2d; mod embedding; mod layer_norm; mod linear; mod rms_norm; +pub use batch_norm::BatchNorm1d; pub use conv1d::{Conv1d, Conv1dConfig}; +pub use conv2d::{Conv2d, Conv2dConfig}; pub use embedding::Embedding; pub use layer_norm::LayerNorm; pub use linear::Linear; diff --git a/fusor-ml/fusor/src/lib.rs b/fusor-ml/fusor/src/lib.rs index 5e326fdad..544efc948 100644 --- a/fusor-ml/fusor/src/lib.rs +++ b/fusor-ml/fusor/src/lib.rs @@ -28,7 +28,7 @@ pub use composite::{ }; pub use device::Device; pub use error::Error; -pub use fusor_types::FromArray; +pub use fusor_types::{FromArray, SlidingWindow}; /// Result type for fusor operations. pub type Result = std::result::Result; @@ -404,19 +404,113 @@ where /// Materialize the tensor to a concrete form. /// - /// For CPU tensors, this evaluates any lazy expressions. - /// For GPU tensors, this is a no-op as GPU tensors are already concrete. + /// For CPU tensors, this evaluates any lazy expressions. GPU tensors stay + /// lazy so graph resolution can keep the full operation chain in one + /// deferred GPU pass. pub fn to_concrete(&self) -> Tensor where B: TensorBacking, - D: SimdElement, + D: SimdElement + DataType, + { + match self { + Tensor::Cpu(t) => Tensor::Cpu(t.to_concrete()), + Tensor::Gpu(t) => Tensor::Gpu(t.clone()), + } + } + + /// Ensure any deferred GPU work for this tensor is submitted and resolved. + /// + /// This is primarily useful when a long-lived tensor would otherwise keep a + /// growing compute graph alive across many incremental decode steps. + pub fn materialize_blocking(&self) + where + B: TensorBacking, + D: SimdElement + DataType + 'static, + { + match self { + Tensor::Cpu(_) => {} + #[cfg(not(target_arch = "wasm32"))] + Tensor::Gpu(t) => pollster::block_on(t.materialize()), + #[cfg(target_arch = "wasm32")] + Tensor::Gpu(_) => {} + } + } + + pub fn to_materialized_blocking(&self) -> Tensor + where + B: TensorBacking, + D: SimdElement + DataType + 'static, { match self { Tensor::Cpu(t) => Tensor::Cpu(t.to_concrete()), + #[cfg(not(target_arch = "wasm32"))] + Tensor::Gpu(t) => Tensor::Gpu(pollster::block_on(t.materialized())), + #[cfg(target_arch = "wasm32")] Tensor::Gpu(t) => Tensor::Gpu(t.clone()), } } + pub fn materialize_many_blocking(tensors: &[&Self]) + where + B: TensorBacking, + D: SimdElement + DataType + 'static, + { + if tensors.is_empty() { + return; + } + + match tensors[0] { + Tensor::Cpu(_) => {} + #[cfg(not(target_arch = "wasm32"))] + Tensor::Gpu(_) => { + let gpu_tensors: Vec<_> = tensors + .iter() + .map(|tensor| match tensor { + Tensor::Gpu(t) => t, + Tensor::Cpu(_) => { + panic!("cannot mix CPU and GPU tensors in materialize_many_blocking") + } + }) + .collect(); + pollster::block_on(fusor_core::Tensor::materialize_many(&gpu_tensors)); + } + #[cfg(target_arch = "wasm32")] + Tensor::Gpu(_) => {} + } + } + + pub fn to_materialized_many_blocking(tensors: &[&Self]) -> Vec> + where + B: TensorBacking, + D: SimdElement + DataType + 'static, + { + if tensors.is_empty() { + return Vec::new(); + } + + match tensors[0] { + Tensor::Cpu(_) => tensors.iter().map(|tensor| tensor.to_concrete()).collect(), + #[cfg(not(target_arch = "wasm32"))] + Tensor::Gpu(_) => { + let gpu_tensors: Vec<_> = tensors + .iter() + .map(|tensor| match tensor { + Tensor::Gpu(t) => t, + Tensor::Cpu(_) => panic!( + "cannot mix CPU and GPU tensors in to_materialized_many_blocking" + ), + }) + .collect(); + pollster::block_on(fusor_core::Tensor::materialized_many(&gpu_tensors)) + .into_iter() + .map(Tensor::Gpu) + .collect() + } + #[cfg(target_arch = "wasm32")] + Tensor::Gpu(_) => tensors.iter().map(|tensor| tensor.to_concrete()).collect(), + } + } + /// Returns the shape of the tensor. pub fn shape(&self) -> [usize; R] where @@ -1158,6 +1252,34 @@ where _ => panic!("Cannot mix CPU and GPU tensors in q_mat_mul"), } } + + pub fn q_mat_mul_bias( + &self, + weights: &crate::QMatrix, + bias: &Tensor<1, f32, B2>, + ) -> Tensor + where + B2: TensorBacking<1, Elem = f32>, + (fusor_core::Tensor, fusor_core::Tensor<1, f32>): fusor_core::MaxRank, + (ConcreteTensor, ConcreteTensor): fusor_cpu::MaxRank, + AddOp: SimdBinaryOp, + { + use crate::QMatrix; + use fusor_gguf::GgmlType; + + if matches!(weights.ggml_type(), GgmlType::F16 | GgmlType::F32) { + return self.q_mat_mul(weights).add_(bias); + } + + match (self, weights, bias) { + (Tensor::Gpu(lhs), QMatrix::Gpu(rhs), Tensor::Gpu(bias)) => { + let bias = bias.clone(); + Tensor::Gpu(lhs.q_mat_mul_bias(rhs, &bias)) + } + (Tensor::Cpu(_), _, _) => self.q_mat_mul(weights).add_(bias), + _ => panic!("Cannot mix CPU and GPU tensors in q_mat_mul_bias"), + } + } } // Flatten operations diff --git a/fusor-ml/fusor/src/quantized.rs b/fusor-ml/fusor/src/quantized.rs index bfc8396e9..3eceff1aa 100644 --- a/fusor-ml/fusor/src/quantized.rs +++ b/fusor-ml/fusor/src/quantized.rs @@ -400,6 +400,188 @@ mod tests { ); } + #[tokio::test] + async fn test_gpu_qmatmul_tiny_m_q8_0_matches_cpu() { + let gpu_device = Device::new().await.expect("GPU required for this test"); + let cpu_device = Device::Cpu; + + let shape = [2, 32]; + let block_size_bytes = 34; + let mut raw_bytes = vec![0u8; 2 * block_size_bytes]; + + let scale_f16 = half::f16::from_f32(1.0); + raw_bytes[0..2].copy_from_slice(&scale_f16.to_le_bytes()); + for i in 0..32 { + raw_bytes[2 + i] = 1i8 as u8; + } + + raw_bytes[block_size_bytes..block_size_bytes + 2].copy_from_slice(&scale_f16.to_le_bytes()); + for i in 0..32 { + raw_bytes[block_size_bytes + 2 + i] = 2i8 as u8; + } + + let cpu_qmatrix = + QMatrix::from_raw_bytes(&cpu_device, shape, &raw_bytes, GgmlType::Q8_0).unwrap(); + let gpu_qmatrix = + QMatrix::from_raw_bytes(&gpu_device, shape, &raw_bytes, GgmlType::Q8_0).unwrap(); + + let input_data: Vec = (0..13) + .flat_map(|row| std::iter::repeat_n(0.25f32 * (row as f32 + 1.0), 32)) + .collect(); + let cpu_input: Tensor<3, f32> = Tensor::from_slice(&cpu_device, [1, 13, 32], &input_data); + let gpu_input: Tensor<3, f32> = Tensor::from_slice(&gpu_device, [1, 13, 32], &input_data); + + let cpu_output = cpu_input.q_mat_mul(&cpu_qmatrix).as_slice().await.unwrap(); + let gpu_output = gpu_input.q_mat_mul(&gpu_qmatrix).as_slice().await.unwrap(); + + for row in 0..13 { + for col in 0..2 { + let expected = cpu_output[[0, row, col]]; + let actual = gpu_output[[0, row, col]]; + assert!( + (expected - actual).abs() < 0.1, + "row {row} col {col}: expected {expected}, got {actual}" + ); + } + } + } + + #[tokio::test] + async fn test_gpu_qmatmul_bias_matches_unfused_add() { + let gpu_device = Device::new().await.expect("GPU required for this test"); + + let shape = [2, 32]; + let block_size_bytes = 34; + let mut raw_bytes = vec![0u8; 2 * block_size_bytes]; + + let scale_f16 = half::f16::from_f32(1.0); + raw_bytes[0..2].copy_from_slice(&scale_f16.to_le_bytes()); + for i in 0..32 { + raw_bytes[2 + i] = 1i8 as u8; + } + + raw_bytes[block_size_bytes..block_size_bytes + 2].copy_from_slice(&scale_f16.to_le_bytes()); + for i in 0..32 { + raw_bytes[block_size_bytes + 2 + i] = 2i8 as u8; + } + + let qmatrix = + QMatrix::from_raw_bytes(&gpu_device, shape, &raw_bytes, GgmlType::Q8_0).unwrap(); + let input_data: Vec>> = vec![ + (0..5) + .map(|row| vec![0.25f32 * (row as f32 + 1.0); 32]) + .collect(), + ]; + let input: Tensor<3, f32> = Tensor::new(&gpu_device, &input_data); + let bias: Tensor<1, f32> = Tensor::new(&gpu_device, &[0.5, -1.25]); + + let fused = input.q_mat_mul_bias(&qmatrix, &bias); + let fused = fused.as_slice().await.unwrap(); + let expected = input.q_mat_mul(&qmatrix).add_(&bias); + let expected = expected.as_slice().await.unwrap(); + + for row in 0..5 { + for col in 0..2 { + let expected = expected[[0, row, col]]; + let actual = fused[[0, row, col]]; + assert!( + (expected - actual).abs() < 0.1, + "row {row} col {col}: expected {expected}, got {actual}" + ); + } + } + } + + #[tokio::test] + async fn test_gpu_qmatmul_materialized_many_matches_individual() { + let gpu_device = Device::new().await.expect("GPU required for this test"); + let cpu_device = Device::Cpu; + + let shape = [2, 32]; + let block_size_bytes = 34; + let mut raw_bytes_a = vec![0u8; 2 * block_size_bytes]; + let mut raw_bytes_b = vec![0u8; 2 * block_size_bytes]; + + let scale_f16 = half::f16::from_f32(1.0); + raw_bytes_a[0..2].copy_from_slice(&scale_f16.to_le_bytes()); + raw_bytes_b[0..2].copy_from_slice(&scale_f16.to_le_bytes()); + for i in 0..32 { + raw_bytes_a[2 + i] = 1i8 as u8; + raw_bytes_b[2 + i] = 3i8 as u8; + } + + raw_bytes_a[block_size_bytes..block_size_bytes + 2] + .copy_from_slice(&scale_f16.to_le_bytes()); + raw_bytes_b[block_size_bytes..block_size_bytes + 2] + .copy_from_slice(&scale_f16.to_le_bytes()); + for i in 0..32 { + raw_bytes_a[block_size_bytes + 2 + i] = 2i8 as u8; + raw_bytes_b[block_size_bytes + 2 + i] = 4i8 as u8; + } + + let cpu_qmatrix_a = + QMatrix::from_raw_bytes(&cpu_device, shape, &raw_bytes_a, GgmlType::Q8_0).unwrap(); + let cpu_qmatrix_b = + QMatrix::from_raw_bytes(&cpu_device, shape, &raw_bytes_b, GgmlType::Q8_0).unwrap(); + let gpu_qmatrix_a = + QMatrix::from_raw_bytes(&gpu_device, shape, &raw_bytes_a, GgmlType::Q8_0).unwrap(); + let gpu_qmatrix_b = + QMatrix::from_raw_bytes(&gpu_device, shape, &raw_bytes_b, GgmlType::Q8_0).unwrap(); + + let input_data: Vec = (0..13) + .flat_map(|row| std::iter::repeat_n(0.25f32 * (row as f32 + 1.0), 32)) + .collect(); + let cpu_input: Tensor<3, f32> = Tensor::from_slice(&cpu_device, [1, 13, 32], &input_data); + let gpu_input: Tensor<3, f32> = Tensor::from_slice(&gpu_device, [1, 13, 32], &input_data); + + let cpu_output_a = cpu_input + .clone() + .q_mat_mul(&cpu_qmatrix_a) + .as_slice() + .await + .unwrap(); + let cpu_output_b = cpu_input + .q_mat_mul(&cpu_qmatrix_b) + .as_slice() + .await + .unwrap(); + + let gpu_output_a = gpu_input.clone().q_mat_mul(&gpu_qmatrix_a); + let gpu_output_b = gpu_input.q_mat_mul(&gpu_qmatrix_b); + let individual_a = gpu_output_a.clone().as_slice().await.unwrap(); + let individual_b = gpu_output_b.clone().as_slice().await.unwrap(); + let many = Tensor::to_materialized_many_blocking(&[&gpu_output_a, &gpu_output_b]); + let many_a = many[0].clone().as_slice().await.unwrap(); + let many_b = many[1].clone().as_slice().await.unwrap(); + + for row in 0..13 { + for col in 0..2 { + let cpu_a = cpu_output_a[[0, row, col]]; + let cpu_b = cpu_output_b[[0, row, col]]; + let gpu_a = individual_a[[0, row, col]]; + let gpu_b = individual_b[[0, row, col]]; + let shared_a = many_a[[0, row, col]]; + let shared_b = many_b[[0, row, col]]; + assert!( + (cpu_a - gpu_a).abs() < 0.1, + "individual A row {row} col {col}: expected {cpu_a}, got {gpu_a}" + ); + assert!( + (cpu_b - gpu_b).abs() < 0.1, + "individual B row {row} col {col}: expected {cpu_b}, got {gpu_b}" + ); + assert!( + (gpu_a - shared_a).abs() < 0.1, + "shared A row {row} col {col}: individual {gpu_a}, shared {shared_a}" + ); + assert!( + (gpu_b - shared_b).abs() < 0.1, + "shared B row {row} col {col}: individual {gpu_b}, shared {shared_b}" + ); + } + } + } + #[test] fn test_cpu_f32_qmatmul() { // Test F32 (non-quantized) qmatmul diff --git a/fusor-ml/nanochat-app/.claude/settings.local.json b/fusor-ml/nanochat-app/.claude/settings.local.json new file mode 100644 index 000000000..ef4a6e3a2 --- /dev/null +++ b/fusor-ml/nanochat-app/.claude/settings.local.json @@ -0,0 +1,10 @@ +{ + "permissions": { + "allow": [ + "WebFetch(domain:deepnight.net)", + "WebFetch(domain:rickyhan.com)", + "WebFetch(domain:kenlo.hashnode.dev)", + "WebFetch(domain:dev.aseprite.org)" + ] + } +} diff --git a/interfaces/kalosm-common/.claude/settings.local.json b/interfaces/kalosm-common/.claude/settings.local.json new file mode 100644 index 000000000..c8f17f933 --- /dev/null +++ b/interfaces/kalosm-common/.claude/settings.local.json @@ -0,0 +1,10 @@ +{ + "permissions": { + "allow": [ + "WebSearch", + "WebFetch(domain:github.com)", + "WebFetch(domain:raw.githubusercontent.com)", + "WebFetch(domain:docs.rs)" + ] + } +} diff --git a/interfaces/kalosm-sound/examples/transcribe.rs b/interfaces/kalosm-sound/examples/transcribe.rs index 028d45035..dd6d7fcb6 100644 --- a/interfaces/kalosm-sound/examples/transcribe.rs +++ b/interfaces/kalosm-sound/examples/transcribe.rs @@ -1,24 +1,118 @@ +use futures_util::Stream; use kalosm::sound::*; +use std::{path::PathBuf, pin::Pin}; use tokio::time::{Duration, Instant}; -#[tokio::main] -async fn main() -> Result<(), anyhow::Error> { - // Record audio from the microphone for 60 seconds. - let audio = MicInput::default() - .record_until(Instant::now() + Duration::from_secs(60)) - .await; +struct TimedAsyncSource { + source: S, + deadline: Instant, +} + +impl TimedAsyncSource { + fn new(source: S, deadline: Instant) -> Self { + Self { source, deadline } + } +} + +impl Stream for TimedAsyncSource { + type Item = f32; + + fn poll_next( + self: Pin<&mut Self>, + cx: &mut std::task::Context<'_>, + ) -> std::task::Poll> { + let myself = self.get_mut(); + if Instant::now() >= myself.deadline { + return std::task::Poll::Ready(None); + } + + let stream = myself.source.as_stream(); + let mut stream = std::pin::pin!(stream); + stream.as_mut().poll_next(cx) + } +} + +impl AsyncSource for TimedAsyncSource { + fn as_stream(&mut self) -> impl Stream + '_ { + self + } + + fn sample_rate(&self) -> u32 { + self.source.sample_rate() + } +} + +async fn print_segments(mut text: S) -> Result<(), anyhow::Error> +where + S: Stream + Unpin + Send, +{ + text.to_std_out().await?; + Ok(()) +} + +async fn run() -> Result<(), anyhow::Error> { + let source = if let Ok(dir) = std::env::var("RWHISPER_COHERE_DIR") { + WhisperSource::cohere_transcribe_03_2026_local(dir) + } else if let Ok(dir) = std::env::var("RWHISPER_MOONSHINE_DIR") { + WhisperSource::moonshine_streaming_local(dir) + } else if std::env::var("RWHISPER_COHERE").is_ok() { + WhisperSource::cohere_transcribe_03_2026() + } else if let Ok(dir) = std::env::var("RWHISPER_WHISPER_DIR") { + let dir = PathBuf::from(dir); + let model_path = dir.join("whisper-tiny-en.gguf.real"); + let model_path = if model_path.exists() { + model_path + } else { + dir.join("whisper-tiny-en.gguf") + }; + WhisperSource::new( + FileSource::local(model_path), + FileSource::local(dir.join("tokenizer-tiny-en.json")), + FileSource::local(dir.join("config-tiny-en.json")), + false, + Some(&[ + [1, 0], + [2, 0], + [2, 5], + [3, 0], + [3, 1], + [3, 2], + [3, 3], + [3, 4], + ]), + ) + } else if std::env::var("RWHISPER_MOONSHINE").is_ok() { + WhisperSource::moonshine_streaming_tiny() + } else { + WhisperSource::tiny_en() + }; + let duration = std::env::var("RWHISPER_MAX_SECONDS") + .ok() + .and_then(|value| value.parse::().ok()) + .unwrap_or(60); + let timestamped = std::env::var("RWHISPER_TIMESTAMPED").ok().as_deref() == Some("1"); - // Create a new small whisper model. let model = WhisperBuilder::default() - .with_source(WhisperSource::tiny_en()) + .with_source(source) .build() .await?; - // Transcribe the audio. - let mut text = model.transcribe(audio); - - // As the model transcribes the audio, print the text to the console. - text.to_std_out().await?; + let audio = TimedAsyncSource::new( + MicInput::default().stream(), + Instant::now() + Duration::from_secs(duration), + ); + let mut task = audio.transcribe(model); + if timestamped { + task = task.timestamped(); + } + print_segments(task).await?; Ok(()) } + +fn main() -> Result<(), anyhow::Error> { + tokio::runtime::Builder::new_current_thread() + .enable_all() + .build()? + .block_on(run()) +} diff --git a/models/kalosm-llama/.claude/settings.local.json b/models/kalosm-llama/.claude/settings.local.json new file mode 100644 index 000000000..59634938c --- /dev/null +++ b/models/kalosm-llama/.claude/settings.local.json @@ -0,0 +1,11 @@ +{ + "permissions": { + "allow": [ + "Bash(cargo test:*)", + "Bash(cargo check:*)", + "Bash(cargo doc:*)" + ], + "deny": [], + "ask": [] + } +} diff --git a/models/rbert/.claude/settings.local.json b/models/rbert/.claude/settings.local.json new file mode 100644 index 000000000..82398d962 --- /dev/null +++ b/models/rbert/.claude/settings.local.json @@ -0,0 +1,19 @@ +{ + "permissions": { + "allow": [ + "Bash(cargo check:*)", + "Read(//Users/evanalmloff/Desktop/Github/floneum/fusor-ml/core/src/**)", + "Read(//Users/evanalmloff/Desktop/Github/floneum/**)", + "Read(//Users/evanalmloff/Desktop/Github/floneum/fusor-ml/core/src/**)", + "Bash(cargo test)", + "Bash(cargo test:*)", + "Bash(RUST_BACKTRACE=1 cargo run --example embed)", + "Bash(cargo run:*)" + ], + "deny": [], + "ask": [], + "additionalDirectories": [ + "/Users/evanalmloff/Desktop/Github/floneum/fusor-ml/core" + ] + } +} \ No newline at end of file diff --git a/models/rwhisper/examples/transcribe.rs b/models/rwhisper/examples/transcribe.rs index 5a2cdfe85..7c6ec65e8 100644 --- a/models/rwhisper/examples/transcribe.rs +++ b/models/rwhisper/examples/transcribe.rs @@ -1,21 +1,72 @@ +use futures_util::Stream; use kalosm::sound::*; +use std::path::PathBuf; use tokio::time::{Duration, Instant}; +async fn print_segments(mut text: S) -> Result<(), anyhow::Error> +where + S: Stream + Unpin + Send, +{ + text.to_std_out().await?; + Ok(()) +} + #[tokio::main] async fn main() -> Result<(), anyhow::Error> { - // Record audio from the microphone for 60 seconds. - let audio = MicInput::default() - .record_until(Instant::now() + Duration::from_secs(60)) - .await; - - // Create a new small whisper model. - let model = WhisperBuilder::default().build().await?; + let source = if let Ok(path) = std::env::var("RWHISPER_COHERE_GGUF") { + WhisperSource::cohere_transcribe_03_2026_local(path) + } else if let Ok(path) = std::env::var("RWHISPER_MOONSHINE_GGUF") { + WhisperSource::moonshine_streaming_local(path) + } else if std::env::var("RWHISPER_COHERE").is_ok() { + WhisperSource::cohere_transcribe_03_2026() + } else if let Ok(dir) = std::env::var("RWHISPER_WHISPER_DIR") { + let dir = PathBuf::from(dir); + let model_path = dir.join("whisper-tiny-en.gguf.real"); + let model_path = if model_path.exists() { + model_path + } else { + dir.join("whisper-tiny-en.gguf") + }; + WhisperSource::new( + FileSource::local(model_path), + FileSource::local(dir.join("tokenizer-tiny-en.json")), + FileSource::local(dir.join("config-tiny-en.json")), + false, + Some(&[ + [1, 0], + [2, 0], + [2, 5], + [3, 0], + [3, 1], + [3, 2], + [3, 3], + [3, 4], + ]), + ) + } else if std::env::var("RWHISPER_MOONSHINE").is_ok() { + WhisperSource::moonshine_streaming_tiny() + } else { + WhisperSource::default() + }; + let duration = std::env::var("RWHISPER_MAX_SECONDS") + .ok() + .and_then(|value| value.parse::().ok()) + .unwrap_or(60); + let timestamped = std::env::var("RWHISPER_TIMESTAMPED").ok().as_deref() == Some("1"); - // Transcribe the audio. - let mut text = model.transcribe(audio); + let model = WhisperBuilder::default() + .with_source(source) + .build() + .await?; - // As the model transcribes the audio, print the text to the console. - text.to_std_out().await?; + let audio = MicInput::default() + .record_until(Instant::now() + Duration::from_secs(duration)) + .await; + let mut task = model.transcribe(audio); + if timestamped { + task = task.timestamped(); + } + print_segments(task).await?; Ok(()) } diff --git a/models/rwhisper/examples/transcribe_file.rs b/models/rwhisper/examples/transcribe_file.rs index dac40bd38..b121d2d94 100644 --- a/models/rwhisper/examples/transcribe_file.rs +++ b/models/rwhisper/examples/transcribe_file.rs @@ -1,29 +1,106 @@ +use futures_util::{stream, Stream}; use kalosm::sound::*; use rodio::Decoder; +use rodio::Source; +use std::path::PathBuf; +use std::pin::Pin; +use std::time::Duration; #[tokio::main] async fn main() -> Result<(), anyhow::Error> { eprintln!("Starting transcription..."); - // Create a new large whisper model + let source = if let Ok(path) = std::env::var("RWHISPER_COHERE_GGUF") { + WhisperSource::cohere_transcribe_03_2026_local(path) + } else if let Ok(_) = std::env::var("RWHISPER_MOONSHINE") { + WhisperSource::moonshine_streaming_tiny() + } else if std::env::var("RWHISPER_COHERE").is_ok() { + WhisperSource::cohere_transcribe_03_2026() + } else if let Ok(dir) = std::env::var("RWHISPER_WHISPER_DIR") { + let dir = PathBuf::from(dir); + let model_path = dir.join("whisper-tiny-en.gguf.real"); + let model_path = if model_path.exists() { + model_path + } else { + dir.join("whisper-tiny-en.gguf") + }; + WhisperSource::new( + FileSource::local(model_path), + FileSource::local(dir.join("tokenizer-tiny-en.json")), + FileSource::local(dir.join("config-tiny-en.json")), + false, + Some(&[ + [1, 0], + [2, 0], + [2, 5], + [3, 0], + [3, 1], + [3, 2], + [3, 3], + [3, 4], + ]), + ) + } else { + WhisperSource::large_v3_turbo() + }; + eprintln!("Building model..."); let model = WhisperBuilder::default() - .with_source(WhisperSource::tiny_en()) + .with_source(source) .build() .await?; eprintln!("Model built successfully"); // Load audio from a file let contents = std::fs::read("./models/rwhisper/examples/samples_jfk.wav").unwrap(); - let audio = Decoder::new(std::io::Cursor::new(contents.clone())).unwrap(); + let max_seconds = std::env::var("RWHISPER_MAX_SECONDS") + .ok() + .and_then(|value| value.parse::().ok()); + let streaming_chunk_ms = std::env::var("RWHISPER_STREAMING_CHUNK_MS") + .ok() + .and_then(|value| value.parse::().ok()); + let timestamped = std::env::var("RWHISPER_TIMESTAMPED").ok().as_deref() == Some("1"); - let (_stream, stream_handle) = rodio::OutputStream::try_default()?; - let sink = rodio::Sink::try_new(&stream_handle).unwrap(); - let rate = audio.sample_rate() as f32; + let audio = Decoder::new(std::io::Cursor::new(contents.clone())).unwrap(); + let rate = rodio::Source::sample_rate(&audio) as f32; // Transcribe the source audio into text eprintln!("Starting transcription..."); - let mut text = model.transcribe(audio); + let mut text: Pin>> = if let Some(chunk_ms) = streaming_chunk_ms + { + let channels = rodio::Source::channels(&audio); + let sample_rate = rodio::Source::sample_rate(&audio); + let samples = if let Some(max_seconds) = max_seconds { + audio + .take_duration(Duration::from_secs_f32(max_seconds)) + .collect::>() + } else { + audio.collect::>() + }; + let samples_per_chunk = + ((sample_rate as usize * chunk_ms) / 1000).max(1) * channels as usize; + let chunks = samples + .chunks(samples_per_chunk) + .map(|chunk| rodio::buffer::SamplesBuffer::new(channels, sample_rate, chunk.to_vec())) + .collect::>(); + let mut task = stream::iter(chunks).transcribe(model); + if timestamped { + task = task.timestamped(); + } + Box::pin(task) + } else if let Some(max_seconds) = max_seconds { + let mut task = model.transcribe(audio.take_duration(Duration::from_secs_f32(max_seconds))); + if timestamped { + task = task.timestamped(); + } + Box::pin(task) + } else { + let mut task = model.transcribe(audio); + if timestamped { + task = task.timestamped(); + } + Box::pin(task) + }; eprintln!("Waiting for segments..."); let mut segment_count = 0; @@ -37,10 +114,16 @@ async fn main() -> Result<(), anyhow::Error> { chunks.len() ); for chunk in chunks { - eprintln!("Chunk: {:?}", chunk); - print!("{chunk}"); - // Play the audio chunk if let Some(timestamp) = chunk.timestamp() { + println!( + "{:?}\t{:.3}..{:.3}", + chunk.text(), + timestamp.start, + timestamp.end + ); + // Play the audio chunk + let (_stream, stream_handle) = rodio::OutputStream::try_default()?; + let sink = rodio::Sink::try_new(&stream_handle).unwrap(); let start = timestamp.start; let end = timestamp.end; let start = (start * rate) as usize; @@ -50,6 +133,8 @@ async fn main() -> Result<(), anyhow::Error> { let audio_source = rodio::buffer::SamplesBuffer::new(1, rate as u32, audio_chunk); sink.append(audio_source); sink.sleep_until_end(); + } else { + print!("{chunk}"); } } } diff --git a/models/rwhisper/scripts/quantize_cohere_transcribe.py b/models/rwhisper/scripts/quantize_cohere_transcribe.py new file mode 100644 index 000000000..a79166ed6 --- /dev/null +++ b/models/rwhisper/scripts/quantize_cohere_transcribe.py @@ -0,0 +1,349 @@ +#!/usr/bin/env python3 + +import argparse +import json +import struct +from pathlib import Path + +import numpy as np + + +GGUF_VERSION = 3 +GGML_TYPE_F16 = 1 +GGML_TYPE_Q4_0 = 2 +GGML_TYPE_Q4K = 12 +METADATA_TYPE_U32 = 4 +METADATA_TYPE_STRING = 8 +ALIGNMENT = 32 +K_BLOCK_SIZE = 256 +Q4K_BLOCK_BYTES = 2 + 2 + 12 + (K_BLOCK_SIZE // 2) +Q4_0_BLOCK_SIZE = 32 +Q4_0_BLOCK_BYTES = 2 + (Q4_0_BLOCK_SIZE // 2) + + +def align_up(value: int, alignment: int) -> int: + return ((value + alignment - 1) // alignment) * alignment + + +def storage_shape(name: str, shape: tuple[int, ...]) -> tuple[int, ...]: + """Reshape a tensor for GGUF storage so it can fit a Q4K block. + + Only applied to `pos_bias_u`/`pos_bias_v` (rel-pos Conformer attention + biases) which the Rust runtime reshapes at use site anyway. Safe because + dequantization preserves the flat row-major buffer and `.reshape()` at + load recovers the logical `(num_heads, head_dim)` view. + """ + if len(shape) == 2 and (name.endswith(".pos_bias_u") or name.endswith(".pos_bias_v")): + numel = shape[0] * shape[1] + if numel % K_BLOCK_SIZE == 0: + return (numel // K_BLOCK_SIZE, K_BLOCK_SIZE) + return shape + + +def quantization_type(name: str, shape: tuple[int, ...]) -> int: + shape = storage_shape(name, shape) + if len(shape) != 2: + return GGML_TYPE_F16 + if shape[-1] % K_BLOCK_SIZE == 0: + return GGML_TYPE_Q4K + if shape[-1] % Q4_0_BLOCK_SIZE == 0: + return GGML_TYPE_Q4_0 + return GGML_TYPE_F16 + + +def load_safetensors_index(path: Path) -> tuple[int, list[tuple[str, str, tuple[int, ...], int, int]]]: + with path.open("rb") as handle: + header_len = struct.unpack(" int: + numel = 1 + for dim in shape: + numel *= dim + if ggml_type == GGML_TYPE_Q4K: + assert len(shape) == 2 + assert shape[-1] % K_BLOCK_SIZE == 0 + return (numel // K_BLOCK_SIZE) * Q4K_BLOCK_BYTES + if ggml_type == GGML_TYPE_Q4_0: + assert len(shape) == 2 + assert shape[-1] % Q4_0_BLOCK_SIZE == 0 + return (numel // Q4_0_BLOCK_SIZE) * Q4_0_BLOCK_BYTES + if ggml_type == GGML_TYPE_F16: + return numel * 2 + raise ValueError(f"unsupported ggml type: {ggml_type}") + + +def raw_tensor_bytes(mm: np.memmap, data_base: int, start: int, end: int) -> bytes: + return memoryview(mm[data_base + start : data_base + end]).tobytes() + + +def decode_tensor(mm: np.memmap, data_base: int, dtype: str, shape: tuple[int, ...], start: int, end: int) -> np.ndarray: + raw = memoryview(mm[data_base + start : data_base + end]) + numel = int(np.prod(shape)) + + if dtype == "F16": + return np.frombuffer(raw, dtype=" bytes: + first = np.zeros(4, dtype=np.uint8) + middle = np.zeros(4, dtype=np.uint8) + last = np.zeros(4, dtype=np.uint8) + + for i in range(4): + first[i] = (int(scales[i]) & 0x3F) | (((int(scales[i + 4]) >> 4) & 0x03) << 6) + middle[i] = (int(offsets[i]) & 0x3F) | (((int(offsets[i + 4]) >> 4) & 0x03) << 6) + last[i] = (int(scales[i + 4]) & 0x0F) | ((int(offsets[i + 4]) & 0x0F) << 4) + + return bytes(first.tolist() + middle.tolist() + last.tolist()) + + +def q4_0_bytes(array: np.ndarray) -> bytes: + assert array.ndim == 2 + rows, cols = array.shape + assert cols % Q4_0_BLOCK_SIZE == 0 + + blocks = np.asarray(array, dtype=np.float32).reshape(-1, Q4_0_BLOCK_SIZE) + + abs_blocks = np.abs(blocks) + max_idx = np.argmax(abs_blocks, axis=1) + max_signed = np.take_along_axis(blocks, max_idx[:, None], axis=1).squeeze(1) + scales = (max_signed / -8.0).astype(np.float32) + inv_scales = np.divide(1.0, scales, out=np.zeros_like(scales), where=scales != 0.0) + quantized = np.clip( + np.rint(blocks * inv_scales[:, None]) + 8.0, 0.0, 15.0 + ).astype(np.uint8) + + low = quantized[:, : Q4_0_BLOCK_SIZE // 2] + high = quantized[:, Q4_0_BLOCK_SIZE // 2 :] + packed_weights = (low | (high << 4)).astype(np.uint8) + + packed = np.empty( + blocks.shape[0], + dtype=np.dtype([("d", " bytes: + assert array.ndim == 2 + rows, cols = array.shape + assert cols % K_BLOCK_SIZE == 0 + + blocks = np.asarray(array, dtype=np.float32).reshape(-1, K_BLOCK_SIZE) + out = bytearray(blocks.shape[0] * Q4K_BLOCK_BYTES) + cursor = 0 + + for block in blocks: + local_scales = np.zeros(8, dtype=np.float32) + local_offsets = np.zeros(8, dtype=np.float32) + quantized_groups = [] + + groups = [block[i * 32 : (i + 1) * 32] for i in range(8)] + for group in groups: + group_min = float(group.min()) + group_max = float(group.max()) + offset = max(0.0, -group_min) + scale = max(0.0, (group_max + offset) / 15.0) + local_scales[len(quantized_groups)] = scale + local_offsets[len(quantized_groups)] = offset + quantized_groups.append(group) + + super_scale = float(local_scales.max() / 63.0) if local_scales.max() > 0 else 0.0 + super_min = float(local_offsets.max() / 63.0) if local_offsets.max() > 0 else 0.0 + + scale_codes = np.zeros(8, dtype=np.uint8) + offset_codes = np.zeros(8, dtype=np.uint8) + packed_weights = np.zeros(K_BLOCK_SIZE // 2, dtype=np.uint8) + + if super_scale > 0: + scale_codes = np.clip(np.rint(local_scales / super_scale), 0, 63).astype(np.uint8) + if super_min > 0: + offset_codes = np.clip(np.rint(local_offsets / super_min), 0, 63).astype(np.uint8) + + for pair in range(4): + low_idx = pair * 2 + high_idx = low_idx + 1 + + low_scale = float(scale_codes[low_idx]) * super_scale + high_scale = float(scale_codes[high_idx]) * super_scale + low_offset = float(offset_codes[low_idx]) * super_min + high_offset = float(offset_codes[high_idx]) * super_min + + low_group = groups[low_idx] + high_group = groups[high_idx] + + if low_scale > 0: + low_q = np.clip(np.rint((low_group + low_offset) / low_scale), 0, 15).astype(np.uint8) + else: + low_q = np.zeros(32, dtype=np.uint8) + if high_scale > 0: + high_q = np.clip(np.rint((high_group + high_offset) / high_scale), 0, 15).astype(np.uint8) + else: + high_q = np.zeros(32, dtype=np.uint8) + + packed_weights[pair * 32 : (pair + 1) * 32] = low_q | (high_q << 4) + + out[cursor : cursor + 2] = np.asarray(super_scale, dtype=" tuple[int, bytes]: + s_shape = storage_shape(name, shape) + ggml_type = quantization_type(name, shape) + if ggml_type == GGML_TYPE_Q4K: + array = decode_tensor(mm, data_base, dtype, shape, start, end).reshape(s_shape) + return GGML_TYPE_Q4K, q4k_bytes(array) + if ggml_type == GGML_TYPE_Q4_0: + array = decode_tensor(mm, data_base, dtype, shape, start, end).reshape(s_shape) + return GGML_TYPE_Q4_0, q4_0_bytes(array) + + if dtype == "F16": + return GGML_TYPE_F16, raw_tensor_bytes(mm, data_base, start, end) + + array = decode_tensor(mm, data_base, dtype, shape, start, end) + return GGML_TYPE_F16, np.asarray(array, dtype=np.float16).tobytes() + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--input", type=Path, required=True, help="Path to model.safetensors") + parser.add_argument("--output", type=Path, required=True, help="Path to output model.gguf") + parser.add_argument( + "--tokenizer", + type=Path, + help="Path to tokenizer.json (defaults to a tokenizer.json next to --input)", + ) + parser.add_argument( + "--config", + type=Path, + help="Path to config.json (defaults to a config.json next to --input)", + ) + args = parser.parse_args() + + tokenizer_path = args.tokenizer or args.input.with_name("tokenizer.json") + config_path = args.config or args.input.with_name("config.json") + if not tokenizer_path.exists(): + raise FileNotFoundError(f"missing tokenizer json: {tokenizer_path}") + if not config_path.exists(): + raise FileNotFoundError(f"missing config json: {config_path}") + tokenizer_json = tokenizer_path.read_text(encoding="utf-8") + config_json = config_path.read_text(encoding="utf-8") + + metadata = { + "general.architecture": ("string", "cohere_asr"), + "general.alignment": ("u32", ALIGNMENT), + "rwhisper.tokenizer.json": ("string", tokenizer_json), + "rwhisper.config.json": ("string", config_json), + } + + data_base, source_tensors = load_safetensors_index(args.input) + tensors = [] + for name, _, shape, _, _ in source_tensors: + s_shape = storage_shape(name, shape) + ggml_type = quantization_type(name, shape) + tensors.append((name, s_shape, ggml_type, tensor_size_bytes(s_shape, ggml_type))) + + header = bytearray() + header.extend(b"GGUF") + header.extend(struct.pack(" len(header): + handle.write(b"\0" * (tensor_data_offset - len(header))) + + cursor = 0 + for name, _, _, offset, _ in tensor_infos: + if offset > cursor: + handle.write(b"\0" * (offset - cursor)) + dtype, shape, start, end = source_map[name] + _, raw = tensor_bytes(mm, data_base, name, dtype, shape, start, end) + handle.write(raw) + cursor = offset + len(raw) + + q4k_count = sum( + 1 for name, shape, _, _ in tensors if quantization_type(name, shape) == GGML_TYPE_Q4K + ) + q4_0_count = sum( + 1 for name, shape, _, _ in tensors if quantization_type(name, shape) == GGML_TYPE_Q4_0 + ) + print(f"wrote {args.output} with {len(tensors)} tensors ({q4k_count} q4k, {q4_0_count} q4_0)") + + +if __name__ == "__main__": + main() diff --git a/models/rwhisper/scripts/quantize_moonshine_streaming.py b/models/rwhisper/scripts/quantize_moonshine_streaming.py new file mode 100644 index 000000000..555e210f3 --- /dev/null +++ b/models/rwhisper/scripts/quantize_moonshine_streaming.py @@ -0,0 +1,348 @@ +#!/usr/bin/env python3 + +import argparse +import json +import struct +from pathlib import Path + +import numpy as np + + +GGUF_VERSION = 3 +GGML_TYPE_F16 = 1 +GGML_TYPE_Q4_0 = 2 +GGML_TYPE_Q4K = 12 +METADATA_TYPE_U32 = 4 +METADATA_TYPE_STRING = 8 +ALIGNMENT = 32 +K_BLOCK_SIZE = 256 +Q4K_BLOCK_BYTES = 2 + 2 + 12 + (K_BLOCK_SIZE // 2) +Q4_0_BLOCK_SIZE = 32 +Q4_0_BLOCK_BYTES = 2 + (Q4_0_BLOCK_SIZE // 2) + + +def align_up(value: int, alignment: int) -> int: + return ((value + alignment - 1) // alignment) * alignment + + +def quantization_type(name: str, shape: tuple[int, ...]) -> int: + if len(shape) != 2: + return GGML_TYPE_F16 + if shape[-1] % K_BLOCK_SIZE == 0: + return GGML_TYPE_Q4K + if shape[-1] % Q4_0_BLOCK_SIZE == 0: + return GGML_TYPE_Q4_0 + return GGML_TYPE_F16 + + +def gguf_name(name: str) -> str: + if name == "proj_out.weight": + return "model.decoder.proj_out.weight" + return name + + +def load_safetensors_index(path: Path) -> tuple[int, list[tuple[str, str, tuple[int, ...], int, int]]]: + with path.open("rb") as handle: + header_len = struct.unpack(" int: + numel = 1 + for dim in shape: + numel *= dim + if ggml_type == GGML_TYPE_Q4K: + assert len(shape) == 2 + assert shape[-1] % K_BLOCK_SIZE == 0 + return (numel // K_BLOCK_SIZE) * Q4K_BLOCK_BYTES + if ggml_type == GGML_TYPE_Q4_0: + assert len(shape) == 2 + assert shape[-1] % Q4_0_BLOCK_SIZE == 0 + return (numel // Q4_0_BLOCK_SIZE) * Q4_0_BLOCK_BYTES + if ggml_type == GGML_TYPE_F16: + return numel * 2 + raise ValueError(f"unsupported ggml type: {ggml_type}") + + +def raw_tensor_bytes(mm: np.memmap, data_base: int, start: int, end: int) -> bytes: + return memoryview(mm[data_base + start : data_base + end]).tobytes() + + +def decode_tensor( + mm: np.memmap, + data_base: int, + dtype: str, + shape: tuple[int, ...], + start: int, + end: int, +) -> np.ndarray: + raw = memoryview(mm[data_base + start : data_base + end]) + numel = int(np.prod(shape)) + + if dtype == "F16": + return np.frombuffer(raw, dtype=" bytes: + first = np.zeros(4, dtype=np.uint8) + middle = np.zeros(4, dtype=np.uint8) + last = np.zeros(4, dtype=np.uint8) + + for i in range(4): + first[i] = (int(scales[i]) & 0x3F) | (((int(scales[i + 4]) >> 4) & 0x03) << 6) + middle[i] = (int(offsets[i]) & 0x3F) | (((int(offsets[i + 4]) >> 4) & 0x03) << 6) + last[i] = (int(scales[i + 4]) & 0x0F) | ((int(offsets[i + 4]) & 0x0F) << 4) + + return bytes(first.tolist() + middle.tolist() + last.tolist()) + + +def q4_0_bytes(array: np.ndarray) -> bytes: + assert array.ndim == 2 + rows, cols = array.shape + assert cols % Q4_0_BLOCK_SIZE == 0 + + blocks = np.asarray(array, dtype=np.float32).reshape(-1, Q4_0_BLOCK_SIZE) + + # ggml Q4_0: d = max_signed / -8, where max_signed is the value with largest magnitude + # (sign preserved). q = clamp(round(v / d) + 8, 0, 15). + abs_blocks = np.abs(blocks) + max_idx = np.argmax(abs_blocks, axis=1) + max_signed = np.take_along_axis(blocks, max_idx[:, None], axis=1).squeeze(1) + scales = (max_signed / -8.0).astype(np.float32) + inv_scales = np.divide(1.0, scales, out=np.zeros_like(scales), where=scales != 0.0) + quantized = np.clip( + np.rint(blocks * inv_scales[:, None]) + 8.0, 0.0, 15.0 + ).astype(np.uint8) + + # Pack: low nibble at position i (i < 16), high nibble at position i+16. + low = quantized[:, : Q4_0_BLOCK_SIZE // 2] + high = quantized[:, Q4_0_BLOCK_SIZE // 2 :] + packed_weights = (low | (high << 4)).astype(np.uint8) + + packed = np.empty( + blocks.shape[0], + dtype=np.dtype([("d", " bytes: + assert array.ndim == 2 + rows, cols = array.shape + assert cols % K_BLOCK_SIZE == 0 + + blocks = np.asarray(array, dtype=np.float32).reshape(-1, K_BLOCK_SIZE) + out = bytearray(blocks.shape[0] * Q4K_BLOCK_BYTES) + cursor = 0 + + for block in blocks: + local_scales = np.zeros(8, dtype=np.float32) + local_offsets = np.zeros(8, dtype=np.float32) + quantized_groups = [] + + groups = [block[i * 32 : (i + 1) * 32] for i in range(8)] + for group in groups: + group_min = float(group.min()) + group_max = float(group.max()) + offset = max(0.0, -group_min) + scale = max(0.0, (group_max + offset) / 15.0) + local_scales[len(quantized_groups)] = scale + local_offsets[len(quantized_groups)] = offset + quantized_groups.append(group) + + super_scale = float(local_scales.max() / 63.0) if local_scales.max() > 0 else 0.0 + super_min = float(local_offsets.max() / 63.0) if local_offsets.max() > 0 else 0.0 + + scale_codes = np.zeros(8, dtype=np.uint8) + offset_codes = np.zeros(8, dtype=np.uint8) + packed_weights = np.zeros(K_BLOCK_SIZE // 2, dtype=np.uint8) + + if super_scale > 0: + scale_codes = np.clip(np.rint(local_scales / super_scale), 0, 63).astype(np.uint8) + if super_min > 0: + offset_codes = np.clip(np.rint(local_offsets / super_min), 0, 63).astype(np.uint8) + + for pair in range(4): + low_idx = pair * 2 + high_idx = low_idx + 1 + + low_scale = float(scale_codes[low_idx]) * super_scale + high_scale = float(scale_codes[high_idx]) * super_scale + low_offset = float(offset_codes[low_idx]) * super_min + high_offset = float(offset_codes[high_idx]) * super_min + + low_group = groups[low_idx] + high_group = groups[high_idx] + + if low_scale > 0: + low_q = np.clip(np.rint((low_group + low_offset) / low_scale), 0, 15).astype(np.uint8) + else: + low_q = np.zeros(32, dtype=np.uint8) + if high_scale > 0: + high_q = np.clip(np.rint((high_group + high_offset) / high_scale), 0, 15).astype(np.uint8) + else: + high_q = np.zeros(32, dtype=np.uint8) + + packed_weights[pair * 32 : (pair + 1) * 32] = low_q | (high_q << 4) + + out[cursor : cursor + 2] = np.asarray(super_scale, dtype=" tuple[int, bytes]: + ggml_type = quantization_type(name, shape) + if ggml_type == GGML_TYPE_Q4K: + array = decode_tensor(mm, data_base, dtype, shape, start, end) + return GGML_TYPE_Q4K, q4k_bytes(array) + if ggml_type == GGML_TYPE_Q4_0: + array = decode_tensor(mm, data_base, dtype, shape, start, end) + return GGML_TYPE_Q4_0, q4_0_bytes(array) + + if dtype == "F16": + return GGML_TYPE_F16, raw_tensor_bytes(mm, data_base, start, end) + + array = decode_tensor(mm, data_base, dtype, shape, start, end) + return GGML_TYPE_F16, np.asarray(array, dtype=np.float16).tobytes() + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--input", type=Path, required=True, help="Path to model.safetensors") + parser.add_argument("--output", type=Path, required=True, help="Path to output model.gguf") + parser.add_argument( + "--tokenizer", + type=Path, + help="Path to tokenizer.json (defaults to a tokenizer.json next to --input)", + ) + parser.add_argument( + "--config", + type=Path, + help="Path to config.json (defaults to a config.json next to --input)", + ) + args = parser.parse_args() + + tokenizer_path = args.tokenizer or args.input.with_name("tokenizer.json") + config_path = args.config or args.input.with_name("config.json") + if not tokenizer_path.exists(): + raise FileNotFoundError(f"missing tokenizer json: {tokenizer_path}") + if not config_path.exists(): + raise FileNotFoundError(f"missing config json: {config_path}") + tokenizer_json = tokenizer_path.read_text(encoding="utf-8") + config_json = config_path.read_text(encoding="utf-8") + + metadata = { + "general.architecture": ("string", "moonshine_streaming_asr"), + "general.alignment": ("u32", ALIGNMENT), + "rwhisper.tokenizer.json": ("string", tokenizer_json), + "rwhisper.config.json": ("string", config_json), + } + + data_base, source_tensors = load_safetensors_index(args.input) + tensors = [] + for source_name, _, shape, _, _ in source_tensors: + name = gguf_name(source_name) + ggml_type = quantization_type(source_name, shape) + tensors.append((name, source_name, shape, ggml_type, tensor_size_bytes(shape, ggml_type))) + + header = bytearray() + header.extend(b"GGUF") + header.extend(struct.pack(" len(header): + handle.write(b"\0" * (tensor_data_offset - len(header))) + + cursor = 0 + for name, source_name, _, _, offset, _ in tensor_infos: + if offset > cursor: + handle.write(b"\0" * (offset - cursor)) + dtype, shape, start, end = source_map[source_name] + _, raw = tensor_bytes(mm, data_base, source_name, dtype, shape, start, end) + handle.write(raw) + cursor = offset + len(raw) + + q4k_count = sum( + 1 for _, source_name, shape, _, _ in tensors if quantization_type(source_name, shape) == GGML_TYPE_Q4K + ) + q4_0_count = sum( + 1 for _, source_name, shape, _, _ in tensors if quantization_type(source_name, shape) == GGML_TYPE_Q4_0 + ) + print(f"wrote {args.output} with {len(tensors)} tensors ({q4k_count} q4k, {q4_0_count} q4_0)") + + +if __name__ == "__main__": + main() diff --git a/models/rwhisper/src/cohere_audio.rs b/models/rwhisper/src/cohere_audio.rs new file mode 100644 index 000000000..59ba09e5d --- /dev/null +++ b/models/rwhisper/src/cohere_audio.rs @@ -0,0 +1,212 @@ +use half::bf16; +use rand::{Rng, SeedableRng}; +use rustfft::{num_complex::Complex, FftPlanner}; + +use crate::cohere_config::CohereConfig; + +const PREEMPH: f32 = 0.97; +const LOG_ZERO_GUARD: f32 = 5.960_464_5e-8_f32; + +fn quantize_bf16(value: f32) -> f32 { + bf16::from_f32(value).to_f32() +} + +fn sample_standard_normal(rng: &mut rand::rngs::StdRng) -> f32 { + let u1 = rng.random::().clamp(f32::MIN_POSITIVE, 1.0); + let u2 = rng.random::(); + (-2.0 * u1.ln()).sqrt() * (std::f32::consts::TAU * u2).cos() +} + +pub fn pcm_to_features( + cfg: &CohereConfig, + samples: &[f32], + filters: &[f32], +) -> (Vec, usize, usize) { + let sample_rate = cfg.preprocessor.sample_rate; + let n_fft = cfg.preprocessor.n_fft; + let win_length = (cfg.preprocessor.window_size * sample_rate as f32).round() as usize; + let hop_length = (cfg.preprocessor.window_stride * sample_rate as f32).round() as usize; + let n_mels = cfg.preprocessor.features; + let n_freqs = n_fft / 2 + 1; + let pad = n_fft / 2; + let total_frames = samples.len() / hop_length + 1; + let valid_frames = samples.len() / hop_length; + + let mut waveform = samples.to_vec(); + if cfg.preprocessor.dither > 0.0 { + let mut rng = rand::rngs::StdRng::seed_from_u64(samples.len() as u64); + for sample in &mut waveform { + *sample += cfg.preprocessor.dither * sample_standard_normal(&mut rng); + } + } + + if !waveform.is_empty() { + for i in (1..waveform.len()).rev() { + waveform[i] -= PREEMPH * waveform[i - 1]; + } + } + + let mut padded = vec![0.0f32; waveform.len() + 2 * pad]; + padded[pad..pad + waveform.len()].copy_from_slice(&waveform); + + let mut window = vec![0.0f32; n_fft]; + let pad_left = (n_fft - win_length) / 2; + for i in 0..win_length { + let phase = std::f32::consts::TAU * i as f32 / (win_length.saturating_sub(1)) as f32; + window[pad_left + i] = quantize_bf16(0.5 - 0.5 * phase.cos()); + } + + let mut planner = FftPlanner::new(); + let fft = planner.plan_fft_forward(n_fft); + let mut frame = vec![Complex::ZERO; n_fft]; + let mut fft_out = vec![0.0f32; n_freqs]; + let mut features = vec![cfg.preprocessor.pad_value; n_mels * total_frames]; + + for frame_idx in 0..total_frames { + let start = frame_idx * hop_length; + for i in 0..n_fft { + frame[i] = Complex::new(padded[start + i] * window[i], 0.0); + } + fft.process(&mut frame); + + for i in 0..n_freqs { + let re = frame[i].re; + let im = frame[i].im; + fft_out[i] = re * re + im * im; + } + + for mel in 0..n_mels { + let filter = &filters[mel * n_freqs..(mel + 1) * n_freqs]; + let mut sum = 0.0f32; + for i in 0..n_freqs { + sum += filter[i] * fft_out[i]; + } + features[mel * total_frames + frame_idx] = (sum + LOG_ZERO_GUARD).ln(); + } + } + + if cfg.preprocessor.normalize == "per_feature" && valid_frames > 1 { + for mel in 0..n_mels { + let row = &mut features[mel * total_frames..(mel + 1) * total_frames]; + let mean = row[..valid_frames].iter().copied().sum::() / valid_frames as f32; + let var = row[..valid_frames] + .iter() + .map(|value| { + let diff = *value - mean; + diff * diff + }) + .sum::() + / (valid_frames as f32 - 1.0); + let std = var.sqrt() + 1e-5; + for value in &mut row[..valid_frames] { + *value = (*value - mean) / std; + } + for value in &mut row[valid_frames..] { + *value = cfg.preprocessor.pad_value; + } + } + } + + (features, total_frames, valid_frames) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::cohere_config::{ + CohereConfig, CohereDecoderConfig, CohereDecoderWrapperConfig, CohereEncoderConfig, + CohereHeadConfig, CoherePreprocessorConfig, + }; + + fn test_config() -> CohereConfig { + CohereConfig { + batch_size: 64, + sample_rate: 16_000, + max_audio_clip_s: 35, + overlap_chunk_second: 5, + min_energy_window_samples: 1600, + supported_languages: vec!["en".to_owned()], + preprocessor: CoherePreprocessorConfig { + features: 128, + sample_rate: 16_000, + window_size: 0.025, + window_stride: 0.01, + n_fft: 512, + normalize: "per_feature".to_owned(), + dither: 1e-5, + pad_to: 0, + pad_value: 0.0, + window: "hann".to_owned(), + }, + encoder: CohereEncoderConfig { + att_context_size: [-1, -1], + conv_kernel_size: 9, + conv_norm_type: "batch_norm".to_owned(), + d_model: 1280, + dropout: 0.0, + feat_in: 128, + feat_out: 1280, + ff_expansion_factor: 4, + n_heads: 8, + n_layers: 48, + pos_emb_max_len: 5000, + self_attention_model: "rel_pos".to_owned(), + subsampling: "dw_striding".to_owned(), + subsampling_conv_channels: 256, + subsampling_factor: 8, + untie_biases: false, + xscaling: true, + }, + transf_decoder: CohereDecoderWrapperConfig { + config_dict: CohereDecoderConfig { + hidden_act: "relu".to_owned(), + hidden_size: 1024, + inner_size: 4096, + learn_positional_encodings: false, + max_sequence_length: 1024, + num_attention_heads: 8, + num_layers: 8, + }, + }, + head: CohereHeadConfig { + hidden_size: 1024, + log_softmax: true, + num_classes: 16_384, + }, + vocab_size: 16_384, + } + } + + #[test] + fn debug_feature_stats() { + let cfg = test_config(); + let path = concat!(env!("CARGO_MANIFEST_DIR"), "/examples/samples_jfk.wav"); + let reader = hound::WavReader::open(path).unwrap(); + let samples = reader + .into_samples::() + .map(|sample| sample.unwrap() as f32 / i16::MAX as f32) + .take(16_000) + .collect::>(); + let filter_bytes = include_bytes!("cohere_melfilters128.bytes").as_slice(); + let mut filterbank = vec![0.0f32; filter_bytes.len() / 4]; + ::read_f32_into( + filter_bytes, + &mut filterbank, + ); + + let (features, total_frames, valid_frames) = pcm_to_features(&cfg, &samples, &filterbank); + let sum: f32 = features.iter().sum(); + let mean = sum / features.len() as f32; + let var = features + .iter() + .map(|value| { + let diff = *value - mean; + diff * diff + }) + .sum::() + / features.len() as f32; + println!("frames total={total_frames} valid={valid_frames}"); + println!("sum={sum} mean={mean} std={}", var.sqrt()); + println!("first20={:?}", &features[..20]); + } +} diff --git a/models/rwhisper/src/cohere_config.rs b/models/rwhisper/src/cohere_config.rs new file mode 100644 index 000000000..913f4799b --- /dev/null +++ b/models/rwhisper/src/cohere_config.rs @@ -0,0 +1,79 @@ +use serde::Deserialize; + +#[derive(Debug, Clone, PartialEq, Deserialize)] +pub struct CohereConfig { + pub batch_size: usize, + pub sample_rate: usize, + pub max_audio_clip_s: usize, + pub overlap_chunk_second: usize, + #[serde(default = "default_min_energy_window_samples")] + pub min_energy_window_samples: usize, + pub supported_languages: Vec, + pub preprocessor: CoherePreprocessorConfig, + pub encoder: CohereEncoderConfig, + pub transf_decoder: CohereDecoderWrapperConfig, + pub head: CohereHeadConfig, + pub vocab_size: usize, +} + +fn default_min_energy_window_samples() -> usize { + 1600 +} + +#[derive(Debug, Clone, PartialEq, Deserialize)] +pub struct CoherePreprocessorConfig { + pub features: usize, + pub sample_rate: usize, + pub window_size: f32, + pub window_stride: f32, + pub n_fft: usize, + pub normalize: String, + pub dither: f32, + pub pad_to: usize, + pub pad_value: f32, + pub window: String, +} + +#[derive(Debug, Clone, PartialEq, Deserialize)] +pub struct CohereEncoderConfig { + pub att_context_size: [isize; 2], + pub conv_kernel_size: usize, + pub conv_norm_type: String, + pub d_model: usize, + pub dropout: f32, + pub feat_in: usize, + pub feat_out: isize, + pub ff_expansion_factor: usize, + pub n_heads: usize, + pub n_layers: usize, + pub pos_emb_max_len: usize, + pub self_attention_model: String, + pub subsampling: String, + pub subsampling_conv_channels: usize, + pub subsampling_factor: usize, + pub untie_biases: bool, + pub xscaling: bool, +} + +#[derive(Debug, Clone, PartialEq, Deserialize)] +pub struct CohereDecoderWrapperConfig { + pub config_dict: CohereDecoderConfig, +} + +#[derive(Debug, Clone, PartialEq, Deserialize)] +pub struct CohereDecoderConfig { + pub hidden_act: String, + pub hidden_size: usize, + pub inner_size: usize, + pub learn_positional_encodings: bool, + pub max_sequence_length: usize, + pub num_attention_heads: usize, + pub num_layers: usize, +} + +#[derive(Debug, Clone, PartialEq, Deserialize)] +pub struct CohereHeadConfig { + pub hidden_size: usize, + pub log_softmax: bool, + pub num_classes: usize, +} diff --git a/models/rwhisper/src/cohere_melfilters128.bytes b/models/rwhisper/src/cohere_melfilters128.bytes new file mode 100644 index 000000000..193a9b988 Binary files /dev/null and b/models/rwhisper/src/cohere_melfilters128.bytes differ diff --git a/models/rwhisper/src/cohere_runtime.rs b/models/rwhisper/src/cohere_runtime.rs new file mode 100644 index 000000000..123fe65ab --- /dev/null +++ b/models/rwhisper/src/cohere_runtime.rs @@ -0,0 +1,355 @@ +use futures_channel::mpsc::UnboundedSender; +use half::bf16; +use tokenizers::Tokenizer; + +use crate::{ + cohere_audio::pcm_to_features, + cohere_config::CohereConfig, + quantized::{cohere::Cohere, timestamps::extract_timestamps}, + DecodingResult, Segment, TokenChunk, WhisperLanguage, +}; +use fusor::{Device, Tensor, VarBuilder}; +use std::num::NonZeroUsize; +use std::ops::Range; +use std::time::Instant; + +pub(crate) struct CohereRuntime { + device: Device, + tokenizer: Tokenizer, + model: Cohere, + filterbank: Vec, + eos_token: u32, + max_new_tokens: usize, +} + +impl CohereRuntime { + pub(crate) fn new( + device: Device, + weights: &[u8], + tokenizer_bytes: &[u8], + config: CohereConfig, + ) -> Result { + let tokenizer = Tokenizer::from_bytes(tokenizer_bytes) + .map_err(crate::model::WhisperLoadingError::LoadTokenizer)?; + let eos_token = tokenizer + .token_to_id("<|endoftext|>") + .ok_or_else(|| fusor::Error::msg("missing <|endoftext|> token"))?; + let max_new_tokens = std::env::var("RWHISPER_COHERE_MAX_NEW_TOKENS") + .ok() + .and_then(|value| value.parse().ok()) + .unwrap_or(256); + let filter_bytes = include_bytes!("cohere_melfilters128.bytes").as_slice(); + let mut filterbank = vec![0.0f32; filter_bytes.len() / 4]; + ::read_f32_into( + filter_bytes, + &mut filterbank, + ); + for value in &mut filterbank { + *value = bf16::from_f32(*value).to_f32(); + } + + let mut reader = std::io::Cursor::new(weights); + let mut vb = VarBuilder::from_gguf(&mut reader).map_err(|err| { + crate::model::WhisperLoadingError::LoadModel(fusor::Error::msg(err.to_string())) + })?; + let model = Cohere::load(&device, &mut vb, config)?; + + Ok(Self { + device, + tokenizer, + model, + filterbank, + eos_token, + max_new_tokens, + }) + } + + fn prompt_ids(&self, language: WhisperLanguage) -> Result, fusor::Error> { + let language = language.to_string(); + let prompt = format!( + "<|startofcontext|><|startoftranscript|><|emo:undefined|><|{language}|><|{language}|><|pnc|><|noitn|><|notimestamp|><|nodiarize|>" + ); + let encoding = self + .tokenizer + .encode(prompt, false) + .map_err(|err| fusor::Error::msg(err.to_string()))?; + Ok(encoding.get_ids().to_vec()) + } + + async fn transcribe_clip( + &self, + samples: &[f32], + language: WhisperLanguage, + with_timestamps: bool, + ) -> Result { + let profile = std::env::var("RWHISPER_COHERE_PROFILE").ok().as_deref() == Some("1"); + let total_start = Instant::now(); + let prompt_ids = self.prompt_ids(language)?; + if profile { + eprintln!( + "cohere prompt ids: {:.3}s", + total_start.elapsed().as_secs_f32() + ); + } + let (features, total_frames, valid_frames) = + pcm_to_features(&self.model.config, samples, &self.filterbank); + if profile { + eprintln!( + "cohere features: {:.3}s total_frames={} valid_frames={}", + total_start.elapsed().as_secs_f32(), + total_frames, + valid_frames + ); + } + let input_features = Tensor::from_slice( + &self.device, + [1, self.model.config.preprocessor.features, total_frames], + &features, + ); + if profile { + eprintln!( + "cohere input tensor: {:.3}s", + total_start.elapsed().as_secs_f32() + ); + } + let (generated, token_timestamps) = if with_timestamps { + let (generated, cross_attentions, encoder_length) = self + .model + .generate_greedy_with_attention( + &input_features, + valid_frames, + &prompt_ids, + self.eos_token, + self.max_new_tokens, + ) + .await?; + let duration_s = samples.len() as f32 / self.model.config.sample_rate as f32; + let seconds_per_frame = if encoder_length == 0 { + 0.0 + } else { + duration_s / encoder_length as f32 + }; + let token_mask = vec![vec![true; generated.len()]]; + let token_timestamps = extract_timestamps( + None, + &cross_attentions, + const { NonZeroUsize::new(7).unwrap() }, + encoder_length, + seconds_per_frame, + token_mask, + ) + .await? + .into_iter() + .next(); + (generated, token_timestamps) + } else { + if profile { + eprintln!( + "cohere generate start: {:.3}s", + total_start.elapsed().as_secs_f32() + ); + } + ( + self.model + .generate_greedy( + &input_features, + valid_frames, + &prompt_ids, + self.eos_token, + self.max_new_tokens, + ) + .await?, + None, + ) + }; + if profile { + eprintln!( + "cohere generate done: {:.3}s generated_tokens={}", + total_start.elapsed().as_secs_f32(), + generated.len() + ); + } + + let mut remaining_tokens: Vec<_> = generated.iter().copied().enumerate().collect(); + remaining_tokens.reverse(); + let mut processed_tokens = Vec::new(); + let mut timestamp_start: Option = None; + let mut prev_text_len = 0; + let mut chunks = Vec::new(); + let mut current_text = String::new(); + let clip_end = samples.len() as f32 / self.model.config.sample_rate as f32; + + while let Some((index, token)) = remaining_tokens.pop() { + processed_tokens.push(token); + if let Some(timestamps) = &token_timestamps { + if timestamp_start.is_none() { + timestamp_start = timestamps.get(index).copied(); + } + } + let detokenized = self + .tokenizer + .decode(&processed_tokens, true) + .map_err(crate::model::WhisperError::Tokenizer)?; + if detokenized.len() > prev_text_len + && detokenized.chars().last().unwrap_or_default().is_ascii() + { + let timestamp = token_timestamps.as_ref().and_then(|timestamps| { + let start = timestamp_start?; + let end = timestamps.get(index).copied().unwrap_or(clip_end); + timestamp_start = Some(end); + Some(start..end) + }); + let text_range = current_text.len()..detokenized.len(); + current_text = detokenized; + prev_text_len = current_text.len(); + chunks.push(TokenChunk { + text_range, + timestamp, + }); + } else { + prev_text_len = detokenized.len(); + } + } + + if current_text.is_empty() && !processed_tokens.is_empty() { + current_text = self + .tokenizer + .decode(&processed_tokens, true) + .map_err(crate::model::WhisperError::Tokenizer)?; + if !current_text.is_empty() { + let timestamp = token_timestamps.as_ref().and_then(|timestamps| { + let start = timestamp_start.or_else(|| timestamps.first().copied())?; + let end = timestamps.last().copied().unwrap_or(clip_end); + Some(start..end) + }); + chunks.push(TokenChunk { + text_range: 0..current_text.len(), + timestamp, + }); + } + } + + Ok(DecodingResult { + text: current_text, + avg_logprob: 0.0, + no_speech_prob: 0.0, + compression_ratio: 0.0, + chunks, + }) + } + + pub(crate) async fn transcribe( + &self, + samples: Vec, + language: Option, + with_timestamps: bool, + result: UnboundedSender, + ) -> Result<(), crate::model::WhisperError> { + let language = language.unwrap_or(WhisperLanguage::English); + let mut result = result; + let chunk_ranges = split_audio_chunks_energy(&self.model.config, &samples); + let total_samples = samples.len(); + let start_time = cfg!(not(target_arch = "wasm32")).then(Instant::now); + + for range in chunk_ranges { + let mut decoding = self + .transcribe_clip(&samples[range.clone()], language, with_timestamps) + .await?; + let start_seconds = range.start as f32 / self.model.config.sample_rate as f32; + for chunk in &mut decoding.chunks { + if let Some(timestamp) = &mut chunk.timestamp { + timestamp.start += start_seconds; + timestamp.end += start_seconds; + } + } + let elapsed_time = start_time.map(|start| start.elapsed()); + let remaining_time = elapsed_time.map(|elapsed| { + let processed = range.end.max(1); + std::time::Duration::from_millis( + ((elapsed.as_millis() as usize / processed) + * (total_samples.saturating_sub(range.end))) as u64, + ) + }); + let segment = Segment { + sample_range: range.clone(), + start: start_seconds as f64, + duration: (range.end - range.start) as f64 / self.model.config.sample_rate as f64, + elapsed_time, + remaining_time, + progress: range.end as f32 / total_samples.max(1) as f32, + result: decoding, + }; + let _ = result.start_send(segment); + } + Ok(()) + } +} + +fn split_audio_chunks_energy(config: &CohereConfig, waveform: &[f32]) -> Vec> { + let sample_rate = config.sample_rate; + let chunk_size = (config.max_audio_clip_s * sample_rate).max(1); + let boundary_context_size = (config.overlap_chunk_second * sample_rate).max(1); + let fast_path_threshold = chunk_size.saturating_sub(boundary_context_size); + + if waveform.len() <= fast_path_threshold { + return vec![0..waveform.len()]; + } + + let mut chunks = Vec::new(); + let mut idx = 0usize; + while idx < waveform.len() { + if idx + chunk_size >= waveform.len() { + chunks.push(idx..waveform.len()); + break; + } + + let search_start = (idx + chunk_size) + .saturating_sub(boundary_context_size) + .max(idx); + let search_end = (idx + chunk_size).min(waveform.len()); + let split_point = if search_end <= search_start { + idx + chunk_size + } else { + find_split_point_energy( + waveform, + search_start, + search_end, + config.min_energy_window_samples, + ) + }; + let split_point = split_point.max(idx + 1).min(waveform.len()); + chunks.push(idx..split_point); + idx = split_point; + } + + chunks +} + +fn find_split_point_energy( + waveform: &[f32], + start_idx: usize, + end_idx: usize, + min_energy_window_samples: usize, +) -> usize { + let segment = &waveform[start_idx..end_idx]; + if segment.len() <= min_energy_window_samples { + return (start_idx + end_idx) / 2; + } + + let mut min_energy = f32::INFINITY; + let mut quietest_idx = start_idx; + let upper = segment.len().saturating_sub(min_energy_window_samples); + let step = min_energy_window_samples.max(1); + let mut i = 0usize; + while i < upper { + let window = &segment[i..(i + min_energy_window_samples)]; + let energy = + (window.iter().map(|sample| sample * sample).sum::() / window.len() as f32).sqrt(); + if energy < min_energy { + min_energy = energy; + quietest_idx = start_idx + i; + } + i += step; + } + quietest_idx +} diff --git a/models/rwhisper/src/lib.rs b/models/rwhisper/src/lib.rs index 79cfb8a4c..bde18fd9f 100644 --- a/models/rwhisper/src/lib.rs +++ b/models/rwhisper/src/lib.rs @@ -47,21 +47,80 @@ use std::{ ops::Range, pin::Pin, str::FromStr, - sync::{Arc, Mutex}, + sync::{ + atomic::{AtomicU64, Ordering}, + Arc, Mutex, + }, time::Duration, }; use futures_util::{FutureExt, Stream, StreamExt}; mod model; +mod moonshine_config; +mod moonshine_runtime; mod source; pub use source::*; -use crate::config::SAMPLE_RATE; +use crate::{config::SAMPLE_RATE, source::ModelFamily}; mod audio; +mod cohere_audio; +mod cohere_config; +mod cohere_runtime; mod config; mod quantized; +static NEXT_STREAM_SESSION_ID: AtomicU64 = AtomicU64::new(1); +const GGUF_TOKENIZER_JSON_METADATA_KEY: &str = "rwhisper.tokenizer.json"; +const GGUF_CONFIG_JSON_METADATA_KEY: &str = "rwhisper.config.json"; + +async fn load_source_bytes( + cache: &Cache, + source: &FileSource, + display_name: impl Display, + progress_handler: &mut impl FnMut(ModelLoadingProgress), +) -> Result, WhisperLoadingError> { + let mut create_progress = ModelLoadingProgress::downloading_progress(display_name.to_string()); + let bytes = match source { + FileSource::Local(path) => { + let size = std::fs::metadata(path) + .map_err(kalosm_common::CacheError::from)? + .len(); + progress_handler(create_progress(kalosm_model_types::FileLoadingProgress { + start_time: None, + cached_size: 0, + size, + progress: size, + })); + std::fs::read(path).map_err(kalosm_common::CacheError::from)? + } + _ => { + cache + .get_bytes(source, |progress| { + progress_handler(create_progress(progress)) + }) + .await? + } + }; + Ok(bytes) +} + +fn embedded_json_bytes( + model: &[u8], + metadata_key: &str, +) -> Result>, WhisperLoadingError> { + let mut reader = std::io::Cursor::new(model); + let vb = + fusor::VarBuilder::from_gguf(&mut reader).map_err(WhisperLoadingError::EmbeddedMetadata)?; + let Some(value) = vb.get_metadata(metadata_key) else { + return Ok(None); + }; + let value = value + .to_string() + .map_err(WhisperLoadingError::EmbeddedMetadata)?; + Ok(Some(value.as_bytes().to_vec())) +} + #[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)] struct DecodingResult { text: String, @@ -209,6 +268,10 @@ where stream: self, whisper: model, current_segment_task: None, + streaming_receiver: None, + streaming_session_id: None, + stream_finished: false, + sent_stream_finish: false, language: Some(WhisperLanguage::English), } } @@ -220,6 +283,10 @@ pub struct ChunkedTranscriptionTask { stream: S, whisper: Whisper, current_segment_task: Option, + streaming_receiver: Option>, + streaming_session_id: Option, + stream_finished: bool, + sent_stream_finish: bool, language: Option, } @@ -255,6 +322,81 @@ where ) -> std::task::Poll> { let myself = self.get_mut(); + if myself.whisper.inner.supports_streaming { + loop { + if myself.streaming_receiver.is_none() { + let (sender, receiver) = futures_channel::mpsc::unbounded(); + let session_id = NEXT_STREAM_SESSION_ID.fetch_add(1, Ordering::Relaxed); + myself + .whisper + .inner + .sender + .unbounded_send(WhisperMessage::StartStream { + session_id, + timestamps: myself.word_level_time_stamps, + lang: myself.language, + sender, + }) + .unwrap(); + myself.streaming_receiver = Some(receiver); + myself.streaming_session_id = Some(session_id); + } + + { + let mut task = myself.whisper.inner.task.lock().unwrap(); + let _ = task.poll_unpin(cx); + } + + if let Some(receiver) = &mut myself.streaming_receiver { + match receiver.poll_next_unpin(cx) { + std::task::Poll::Ready(Some(ready)) => { + return std::task::Poll::Ready(Some(ready)); + } + std::task::Poll::Ready(None) => return std::task::Poll::Ready(None), + std::task::Poll::Pending => {} + } + } + + if !myself.stream_finished { + match myself.stream.poll_next_unpin(cx) { + std::task::Poll::Ready(Some(source)) => { + let samples = + normalize_audio(source, myself.whisper.inner.apply_speech_filter); + myself + .whisper + .inner + .sender + .unbounded_send(WhisperMessage::AppendStream { + session_id: myself.streaming_session_id.unwrap(), + samples, + }) + .unwrap(); + continue; + } + std::task::Poll::Ready(None) => { + myself.stream_finished = true; + } + std::task::Poll::Pending => return std::task::Poll::Pending, + } + } + + if myself.stream_finished && !myself.sent_stream_finish { + myself + .whisper + .inner + .sender + .unbounded_send(WhisperMessage::FinishStream { + session_id: myself.streaming_session_id.unwrap(), + }) + .unwrap(); + myself.sent_stream_finish = true; + continue; + } + + return std::task::Poll::Pending; + } + } + loop { if let Some(task) = &mut myself.current_segment_task { match task.poll_next_unpin(cx) { @@ -353,9 +495,11 @@ impl ModelBuilder for WhisperBuilder { fn requires_download(&self) -> bool { let whisper = &self.model; let cache = Cache::default(); - !cache.exists(&whisper.model) - || !cache.exists(&whisper.tokenizer) - || !cache.exists(&whisper.config) + if !cache.exists(&whisper.model) { + return true; + } + !matches!(whisper.family, ModelFamily::MoonshineStreaming { .. }) + && (!cache.exists(&whisper.tokenizer) || !cache.exists(&whisper.config)) } } @@ -393,39 +537,42 @@ impl WhisperBuilder { self, mut progress_handler: impl FnMut(ModelLoadingProgress) + 'static, ) -> Result { - // Download section - let whisper = &self.model; - let tokenizer_source = &whisper.tokenizer; + let whisper = self.model.clone(); let model_source = &whisper.model; - let config_source = &whisper.config; - - let display_tokenizer_source = format!("Tokenizer ({tokenizer_source})"); - let mut create_progress = - ModelLoadingProgress::downloading_progress(display_tokenizer_source); - let tokenizer = self - .cache - .get_bytes(tokenizer_source, |progress| { - progress_handler(create_progress(progress)) - }) - .await?; - - let display_model_source = format!("Model ({model_source})"); - let mut create_progress = ModelLoadingProgress::downloading_progress(display_model_source); - let model = self - .cache - .get_bytes(model_source, |progress| { - progress_handler(create_progress(progress)) - }) - .await?; - - let display_config_source = format!("Config ({config_source})"); - let mut create_progress = ModelLoadingProgress::downloading_progress(display_config_source); - let config = self - .cache - .get_bytes(config_source, |progress| { - progress_handler(create_progress(progress)) - }) - .await?; + let model = load_source_bytes( + &self.cache, + model_source, + format!("Model ({model_source})"), + &mut progress_handler, + ) + .await?; + + let tokenizer = if let Some(tokenizer) = + embedded_json_bytes(&model, GGUF_TOKENIZER_JSON_METADATA_KEY)? + { + tokenizer + } else { + load_source_bytes( + &self.cache, + &whisper.tokenizer, + format!("Tokenizer ({})", whisper.tokenizer), + &mut progress_handler, + ) + .await? + }; + + let config = + if let Some(config) = embedded_json_bytes(&model, GGUF_CONFIG_JSON_METADATA_KEY)? { + config + } else { + load_source_bytes( + &self.cache, + &whisper.config, + format!("Config ({})", whisper.config), + &mut progress_handler, + ) + .await? + }; let (tx, rx) = futures_channel::mpsc::unbounded::(); let mut model = WhisperInner::new(self, &model, &tokenizer, &config).await?; @@ -433,14 +580,39 @@ impl WhisperBuilder { let task = Box::pin(async move { let mut rx = rx; while let Some(message) = rx.next().await { - model - .transcribe( - message.samples, - message.timestamps, - message.lang, - message.sender, - ) - .await; + match message { + WhisperMessage::Transcribe { + samples, + timestamps, + lang, + sender, + } => { + model.transcribe(samples, timestamps, lang, sender).await; + } + WhisperMessage::StartStream { + session_id, + timestamps, + lang, + sender, + } => { + if let Err(err) = model.start_stream(session_id, timestamps, lang, sender) { + tracing::error!("Error starting transcription stream: {err}"); + } + } + WhisperMessage::AppendStream { + session_id, + samples, + } => { + if let Err(err) = model.push_stream_audio(session_id, samples).await { + tracing::error!("Error updating transcription stream: {err}"); + } + } + WhisperMessage::FinishStream { session_id } => { + if let Err(err) = model.finish_stream(session_id).await { + tracing::error!("Error finishing transcription stream: {err}"); + } + } + } } }); @@ -448,6 +620,21 @@ impl WhisperBuilder { inner: Arc::new(WhisperTask { sender: tx, task: Mutex::new(task), + apply_speech_filter: match whisper.family { + ModelFamily::Whisper { + apply_speech_filter, + .. + } + | ModelFamily::MoonshineStreaming { + apply_speech_filter, + .. + } => apply_speech_filter, + ModelFamily::CohereTranscribe => false, + }, + supports_streaming: matches!( + whisper.family, + ModelFamily::MoonshineStreaming { .. } + ), }), }) } @@ -805,6 +992,8 @@ impl Display for WhisperLanguage { struct WhisperTask { sender: UnboundedSender, task: Mutex + 'static>>>, + apply_speech_filter: bool, + supports_streaming: bool, } #[derive(Clone)] @@ -833,7 +1022,7 @@ impl Whisper { ::Item: rodio::Sample, f32: FromSample<::Item>, { - let pcm_data: Vec<_> = normalize_audio(input); + let pcm_data: Vec<_> = normalize_audio(input, self.inner.apply_speech_filter); TranscriptionTask { word_level_time_stamps: false, audio: pcm_data, @@ -888,7 +1077,7 @@ impl Stream for TranscriptionTask { .whisper .inner .sender - .unbounded_send(WhisperMessage { + .unbounded_send(WhisperMessage::Transcribe { samples: pcm_data, timestamps: myself.word_level_time_stamps, lang: myself.language, @@ -907,20 +1096,41 @@ impl Stream for TranscriptionTask { } } -struct WhisperMessage { - samples: Vec, - timestamps: bool, - lang: Option, - sender: UnboundedSender, +enum WhisperMessage { + Transcribe { + samples: Vec, + timestamps: bool, + lang: Option, + sender: UnboundedSender, + }, + StartStream { + session_id: u64, + timestamps: bool, + lang: Option, + sender: UnboundedSender, + }, + AppendStream { + session_id: u64, + samples: Vec, + }, + FinishStream { + session_id: u64, + }, } -pub(crate) fn normalize_audio(input: S) -> Vec +pub(crate) fn normalize_audio(input: S, apply_speech_filter: bool) -> Vec where ::Item: rodio::Sample, f32: FromSample<::Item>, { let resample = UniformSourceIterator::new(input, 1, SAMPLE_RATE as u32); - let pass_filter = resample.low_pass(3000).high_pass(200).convert_samples(); - - pass_filter.collect::>() + if apply_speech_filter { + resample + .low_pass(3000) + .high_pass(200) + .convert_samples() + .collect::>() + } else { + resample.convert_samples().collect::>() + } } diff --git a/models/rwhisper/src/model.rs b/models/rwhisper/src/model.rs index 0a8b2b5c6..0762264e9 100644 --- a/models/rwhisper/src/model.rs +++ b/models/rwhisper/src/model.rs @@ -15,7 +15,9 @@ use tokenizers::Tokenizer; use super::{DecodingResult, Segment}; use crate::{ - audio, config::*, quantized::TextDecoderCache, Task, TaskType, TokenChunk, WhisperBuilder, + audio, cohere_config::CohereConfig, cohere_runtime::CohereRuntime, config::*, + moonshine_config::MoonshineStreamingConfig, moonshine_runtime::MoonshineRuntime, + quantized::TextDecoderCache, source::ModelFamily, Task, TaskType, TokenChunk, WhisperBuilder, WhisperLanguage, }; use kalosm_common::CacheError; @@ -55,6 +57,9 @@ pub enum WhisperLoadingError { /// An error that can occur when trying to load the whisper config. #[error("Failed to load config: {0}")] LoadConfig(serde_json::Error), + /// An error that can occur when trying to read embedded GGUF metadata. + #[error("Failed to read embedded GGUF metadata: {0}")] + EmbeddedMetadata(fusor::GgufReadError), /// Unsupported mel filter length #[error("Unsupported mel filter length: {0}; only 80 and 128 are supported")] UnsupportedMelFilterLength(usize), @@ -77,13 +82,19 @@ pub enum WhisperError { Compression(std::io::Error), } -pub(crate) struct WhisperInner { +pub(crate) struct WhisperRuntime { mel_filters: Vec, device: Device, decoder: Decoder, config: Config, } +pub(crate) enum WhisperInner { + Whisper(WhisperRuntime), + Cohere(CohereRuntime), + Moonshine(MoonshineRuntime), +} + impl WhisperInner { pub(crate) async fn new( settings: WhisperBuilder, @@ -91,7 +102,7 @@ impl WhisperInner { tokenizer: &[u8], config: &[u8], ) -> Result { - // Set FUSOR_USE_GPU=1 to use GPU, otherwise CPU + // Set FUSOR_USE_GPU=1 to use GPU, otherwise CPU. let use_gpu = std::env::var("FUSOR_USE_GPU") .map(|v| v == "1") .unwrap_or(false); @@ -101,48 +112,62 @@ impl WhisperInner { Device::cpu() }; - let tokenizer = - Tokenizer::from_bytes(tokenizer).map_err(WhisperLoadingError::LoadTokenizer)?; - let config: Config = - serde_json::from_slice(config).map_err(WhisperLoadingError::LoadConfig)?; - - let mel_bytes = match config.num_mel_bins { - 80 => include_bytes!("melfilters.bytes").as_slice(), - 128 => include_bytes!("melfilters128.bytes").as_slice(), - nmel => return Err(WhisperLoadingError::UnsupportedMelFilterLength(nmel)), - }; - let mut mel_filters = vec![0f32; mel_bytes.len() / 4]; - ::read_f32_into( - mel_bytes, - &mut mel_filters, - ); - let attention_heads = settings.model.heads; - - let model = ModelType::load(weights, &device, config.clone())?; - let language_token = if settings.model.multilingual { - let language = settings.language.unwrap_or(WhisperLanguage::English); - match token_id(&tokenizer, &format!("<|{language}|>")) { - Ok(token_id) => Some(token_id), - Err(_) => return Err(WhisperLoadingError::UnsupportedLanguage(language)), + match settings.model.family { + ModelFamily::Whisper { + multilingual, + heads, + .. + } => { + let tokenizer = + Tokenizer::from_bytes(tokenizer).map_err(WhisperLoadingError::LoadTokenizer)?; + let config: Config = + serde_json::from_slice(config).map_err(WhisperLoadingError::LoadConfig)?; + + let mel_bytes = match config.num_mel_bins { + 80 => include_bytes!("melfilters.bytes").as_slice(), + 128 => include_bytes!("melfilters128.bytes").as_slice(), + nmel => return Err(WhisperLoadingError::UnsupportedMelFilterLength(nmel)), + }; + let mut mel_filters = vec![0f32; mel_bytes.len() / 4]; + ::read_f32_into( + mel_bytes, + &mut mel_filters, + ); + + let model = ModelType::load(weights, &device, config.clone())?; + let language_token = if multilingual { + let language = settings.language.unwrap_or(WhisperLanguage::English); + match token_id(&tokenizer, &format!("<|{language}|>")) { + Ok(token_id) => Some(token_id), + Err(_) => return Err(WhisperLoadingError::UnsupportedLanguage(language)), + } + } else { + None + }; + let decoder = Decoder::new(model, tokenizer, 0, &device, language_token, heads)?; + + Ok(Self::Whisper(WhisperRuntime { + mel_filters, + device, + decoder, + config, + })) } - } else { - None - }; - let decoder = Decoder::new( - model, - tokenizer, - 0, - &device, - language_token, - attention_heads, - )?; - - Ok(Self { - mel_filters, - device, - decoder, - config, - }) + ModelFamily::CohereTranscribe => { + let config: CohereConfig = + serde_json::from_slice(config).map_err(WhisperLoadingError::LoadConfig)?; + Ok(Self::Cohere(CohereRuntime::new( + device, weights, tokenizer, config, + )?)) + } + ModelFamily::MoonshineStreaming { heads, .. } => { + let config: MoonshineStreamingConfig = + serde_json::from_slice(config).map_err(WhisperLoadingError::LoadConfig)?; + Ok(Self::Moonshine(MoonshineRuntime::new( + device, weights, tokenizer, config, heads, + )?)) + } + } } pub(crate) async fn transcribe( @@ -152,36 +177,96 @@ impl WhisperInner { language: Option, result: UnboundedSender, ) { - let mel = audio::pcm_to_mel(&self.config, &pcm_data, &self.mel_filters); - let mel_len = mel.len(); - let mel = Tensor::new(&self.device, &mel) - .reshape([self.config.num_mel_bins, mel_len / self.config.num_mel_bins]) - .to_concrete() - .cast(); - - if let Some(language) = language { - if let Err(err) = self.decoder.set_language_token(language) { - // Log error or send error message to result channel - // Continue with default language - tracing::error!("Error updating language token: {err}"); + match self { + Self::Whisper(runtime) => { + let mel = audio::pcm_to_mel(&runtime.config, &pcm_data, &runtime.mel_filters); + let mel_len = mel.len(); + let mel = Tensor::new(&runtime.device, &mel) + .reshape([ + runtime.config.num_mel_bins, + mel_len / runtime.config.num_mel_bins, + ]) + .to_concrete() + .cast(); + + if let Some(language) = language { + if let Err(err) = runtime.decoder.set_language_token(language) { + tracing::error!("Error updating language token: {err}"); + } + } + + if let Err(err) = runtime + .decoder + .run( + &mel, + pcm_data.len(), + Task { + task_type: TaskType::Unset, + word_level_time_stamps, + without_timestamps: true, + }, + result, + ) + .await + { + tracing::error!("Error transcribing audio: {err}"); + } + } + Self::Cohere(runtime) => { + if let Err(err) = runtime + .transcribe(pcm_data, language, word_level_time_stamps, result) + .await + { + tracing::error!("Error transcribing audio: {err}"); + } + } + Self::Moonshine(runtime) => { + if let Err(err) = runtime + .transcribe(pcm_data, language, word_level_time_stamps, result) + .await + { + tracing::error!("Error transcribing audio: {err}"); + } } } + } - if let Err(err) = self - .decoder - .run( - &mel, - pcm_data.len(), - Task { - task_type: TaskType::Unset, - word_level_time_stamps, - without_timestamps: true, - }, - result, - ) - .await - { - tracing::error!("Error transcribing audio: {err}"); + pub(crate) fn start_stream( + &mut self, + session_id: u64, + word_level_time_stamps: bool, + language: Option, + sender: UnboundedSender, + ) -> Result<(), WhisperLoadingError> { + match self { + Self::Moonshine(runtime) => { + runtime.start_stream(session_id, language, word_level_time_stamps, sender) + } + _ => Err(WhisperLoadingError::LoadModel(fusor::Error::msg( + "streaming is only supported for moonshine models", + ))), + } + } + + pub(crate) async fn push_stream_audio( + &mut self, + session_id: u64, + samples: Vec, + ) -> Result<(), WhisperError> { + match self { + Self::Moonshine(runtime) => runtime.push_stream_audio(session_id, samples).await, + _ => Err(WhisperError::Fusor(fusor::Error::msg( + "streaming is only supported for moonshine models", + ))), + } + } + + pub(crate) async fn finish_stream(&mut self, session_id: u64) -> Result<(), WhisperError> { + match self { + Self::Moonshine(runtime) => runtime.finish_stream(session_id).await, + _ => Err(WhisperError::Fusor(fusor::Error::msg( + "streaming is only supported for moonshine models", + ))), } } } @@ -661,7 +746,6 @@ impl Decoder { let start = range.start; let end = range.end; let start_time_offset = (start * HOP_LENGTH) as f64 / SAMPLE_RATE as f64; - let time_offset = (end * HOP_LENGTH) as f64 / SAMPLE_RATE as f64; let segment_duration = (segment_size * HOP_LENGTH) as f64 / SAMPLE_RATE as f64; @@ -716,7 +800,7 @@ impl Decoder { let segment = Segment { sample_range: (range.start * HOP_LENGTH) ..audio_frames.min(range.end * HOP_LENGTH), - start: time_offset, + start: start_time_offset, duration: segment_duration, remaining_time: remaining, elapsed_time: elapsed, diff --git a/models/rwhisper/src/moonshine_config.rs b/models/rwhisper/src/moonshine_config.rs new file mode 100644 index 000000000..4a3b0ea29 --- /dev/null +++ b/models/rwhisper/src/moonshine_config.rs @@ -0,0 +1,101 @@ +use serde::Deserialize; + +#[allow(dead_code)] +#[derive(Debug, Clone, Deserialize)] +pub(crate) struct MoonshineStreamingConfig { + pub(crate) encoder_config: MoonshineStreamingEncoderConfig, + pub(crate) hidden_size: usize, + pub(crate) intermediate_size: usize, + pub(crate) num_hidden_layers: usize, + pub(crate) num_attention_heads: usize, + pub(crate) num_key_value_heads: usize, + pub(crate) max_position_embeddings: usize, + pub(crate) use_cache: bool, + pub(crate) pad_token_id: Option, + pub(crate) bos_token_id: Option, + pub(crate) eos_token_id: u32, + pub(crate) decoder_start_token_id: Option, + pub(crate) attention_bias: bool, + pub(crate) attention_dropout: f32, + pub(crate) head_dim: Option, + pub(crate) pad_head_dim_to_multiple_of: Option, + pub(crate) tie_word_embeddings: bool, + pub(crate) vocab_size: usize, + #[serde(default)] + pub(crate) rope_parameters: Option, +} + +impl MoonshineStreamingConfig { + pub(crate) fn decoder_start_token(&self) -> u32 { + self.decoder_start_token_id + .or(self.bos_token_id) + .unwrap_or(1) + } + + pub(crate) fn decoder_head_dim(&self) -> usize { + self.head_dim + .unwrap_or(self.hidden_size / self.num_attention_heads.max(1)) + } + + pub(crate) fn encoder_hidden_size(&self) -> usize { + self.encoder_config.hidden_size + } + + pub(crate) fn rope_theta(&self) -> f32 { + self.rope_parameters + .as_ref() + .and_then(|params| params.rope_theta) + .unwrap_or(10_000.0) + } + + pub(crate) fn partial_rotary_factor(&self) -> f32 { + self.rope_parameters + .as_ref() + .and_then(|params| params.partial_rotary_factor) + .unwrap_or(0.5) + } + + pub(crate) fn decoder_rotary_dim(&self) -> usize { + let head_dim = self.decoder_head_dim(); + let factor = self.partial_rotary_factor().clamp(0.0, 1.0); + let rotary_dim = ((head_dim as f32 * factor).floor() as usize / 2) * 2; + rotary_dim.max(2).min(head_dim) + } +} + +#[allow(dead_code)] +#[derive(Debug, Clone, Deserialize)] +pub(crate) struct MoonshineStreamingEncoderConfig { + pub(crate) hidden_size: usize, + pub(crate) intermediate_size: usize, + pub(crate) hidden_act: String, + pub(crate) num_hidden_layers: usize, + pub(crate) num_attention_heads: usize, + pub(crate) num_key_value_heads: usize, + pub(crate) max_position_embeddings: usize, + pub(crate) attention_dropout: f32, + pub(crate) attention_bias: bool, + pub(crate) sample_rate: usize, + pub(crate) frame_ms: f32, + pub(crate) sliding_windows: Vec<[usize; 2]>, + pub(crate) head_dim: Option, +} + +impl MoonshineStreamingEncoderConfig { + pub(crate) fn frame_len(&self) -> usize { + ((self.sample_rate as f32 * self.frame_ms) / 1000.0).round() as usize + } + + pub(crate) fn head_dim(&self) -> usize { + self.head_dim + .unwrap_or(self.hidden_size / self.num_attention_heads.max(1)) + } +} + +#[derive(Debug, Clone, Deserialize)] +pub(crate) struct MoonshineRopeParameters { + #[serde(default)] + pub(crate) rope_theta: Option, + #[serde(default)] + pub(crate) partial_rotary_factor: Option, +} diff --git a/models/rwhisper/src/moonshine_runtime.rs b/models/rwhisper/src/moonshine_runtime.rs new file mode 100644 index 000000000..3740933dd --- /dev/null +++ b/models/rwhisper/src/moonshine_runtime.rs @@ -0,0 +1,863 @@ +use std::{collections::HashMap, num::NonZeroUsize, ops::Range}; + +use futures_channel::mpsc::UnboundedSender; +use tokenizers::Tokenizer; + +use crate::{ + moonshine_config::MoonshineStreamingConfig, + quantized::{ + moonshine::{Moonshine, MoonshineDecoderCache, MoonshineEncoderStreamState}, + timestamps::extract_timestamps, + }, + DecodingResult, Segment, TokenChunk, WhisperLanguage, +}; +use fusor::{Device, Tensor, VarBuilder}; + +fn moonshine_profile() -> bool { + std::env::var("RWHISPER_MOONSHINE_PROFILE").ok().as_deref() == Some("1") +} + +#[derive(Clone)] +struct CandidateDecode { + tokens: Vec, + token_timestamps: Vec, + clip_end_s: f32, + usable_samples: usize, +} + +struct MoonshineStreamState { + word_timestamps: bool, + sender: UnboundedSender, + last_candidate: Option, + encoder_state: MoonshineEncoderStreamState, + prepared_encoder_hidden_states: Option>, + decoder_cross_cache: MoonshineDecoderCache, + decoder_prefix_cache: MoonshineDecoderCache, + decoder_prefix_tokens: Vec, + last_decoded_finalized_frames: usize, + emitted_tokens: usize, + last_emitted_end_s: f32, +} + +pub(crate) struct MoonshineRuntime { + device: Device, + tokenizer: Tokenizer, + model: Moonshine, + start_token: u32, + eos_token: u32, + sample_rate: usize, + frame_len: usize, + max_tokens_per_second: f32, + stream_decode_interval_frames: usize, + stream_initial_holdback_tokens: usize, + stream_stable_holdback_tokens: usize, + alignment_heads: Option<&'static [[usize; 2]]>, + streams: HashMap, +} + +impl MoonshineRuntime { + pub(crate) fn new( + device: Device, + weights: &[u8], + tokenizer_bytes: &[u8], + config: MoonshineStreamingConfig, + alignment_heads: Option<&'static [[usize; 2]]>, + ) -> Result { + let tokenizer = Tokenizer::from_bytes(tokenizer_bytes) + .map_err(crate::model::WhisperLoadingError::LoadTokenizer)?; + let start_token = config.decoder_start_token(); + let eos_token = config.eos_token_id; + let sample_rate = config.encoder_config.sample_rate; + let frame_len = config.encoder_config.frame_len(); + let max_tokens_per_second = std::env::var("RWHISPER_MOONSHINE_MAX_TOKENS_PER_SECOND") + .ok() + .and_then(|value| value.parse().ok()) + .unwrap_or(6.5); + let stream_decode_interval_frames = std::env::var("RWHISPER_MOONSHINE_STREAM_DECODE_MS") + .ok() + .and_then(|value| value.parse::().ok()) + .map(|milliseconds| (milliseconds / 20).max(1)) + .unwrap_or(10); + let stream_stable_holdback_tokens = + std::env::var("RWHISPER_MOONSHINE_STREAM_HOLDBACK_TOKENS") + .ok() + .and_then(|value| value.parse::().ok()) + .unwrap_or(2); + let stream_initial_holdback_tokens = + std::env::var("RWHISPER_MOONSHINE_STREAM_INITIAL_HOLDBACK_TOKENS") + .ok() + .and_then(|value| value.parse::().ok()) + .unwrap_or(4); + + let mut reader = std::io::Cursor::new(weights); + let mut vb = VarBuilder::from_gguf(&mut reader).map_err(|err| { + crate::model::WhisperLoadingError::LoadModel(fusor::Error::msg(err.to_string())) + })?; + let model = Moonshine::load(&device, &mut vb, config) + .map_err(crate::model::WhisperLoadingError::LoadModel)?; + + Ok(Self { + device, + tokenizer, + model, + start_token, + eos_token, + sample_rate, + frame_len, + max_tokens_per_second, + stream_decode_interval_frames, + stream_initial_holdback_tokens, + stream_stable_holdback_tokens, + alignment_heads, + streams: HashMap::new(), + }) + } + + fn validate_language( + &self, + language: Option, + ) -> Result<(), crate::model::WhisperLoadingError> { + if let Some(language) = language { + if !matches!(language, WhisperLanguage::English) { + return Err(crate::model::WhisperLoadingError::UnsupportedLanguage( + language, + )); + } + } + Ok(()) + } + + fn usable_samples<'a>(&self, samples: &'a [f32]) -> &'a [f32] { + let usable = samples.len() / self.frame_len * self.frame_len; + &samples[..usable] + } + + fn max_new_tokens(&self, sample_count: usize) -> usize { + let seconds = sample_count as f32 / self.sample_rate as f32; + (seconds * self.max_tokens_per_second).ceil() as usize + 8 + } + + fn encode_prepared_full( + &mut self, + samples: &[f32], + ) -> Result<(Tensor<3, f32>, usize), crate::model::WhisperError> { + let encoder_hidden_states = self + .model + .encoder + .encode(&self.device, samples) + .map_err(crate::model::WhisperError::Fusor)?; + let total_frames = encoder_hidden_states.shape()[1]; + let adapted_encoder_hidden_states = self + .model + .decoder + .prepare_encoder_hidden_states(&encoder_hidden_states) + .map_err(crate::model::WhisperError::Fusor)?; + Ok((adapted_encoder_hidden_states, total_frames)) + } + + async fn greedy_generate( + &mut self, + encoder_hidden_states: &Tensor<3, f32>, + max_new_tokens: usize, + ) -> Result, crate::model::WhisperError> { + let profile = moonshine_profile(); + let started = profile.then(std::time::Instant::now); + let mut decoder_cache = MoonshineDecoderCache::default(); + let mut decoder_inputs = vec![self.start_token]; + let mut generated = Vec::new(); + let mut steps = 0usize; + + for _ in 0..max_new_tokens.max(1) { + steps += 1; + let logits = self + .model + .decoder + .decode_cached(&decoder_inputs, &encoder_hidden_states, &mut decoder_cache) + .map_err(crate::model::WhisperError::Fusor)?; + let last_index = logits.shape()[1].saturating_sub(1); + let last_logits = logits + .narrow(1, last_index, 1) + .squeeze(1) + .squeeze(0) + .to_concrete(); + let best_token = last_logits.argmax_last_dim::<0>().to_scalar().await?; + if best_token == self.eos_token { + break; + } + generated.push(best_token); + decoder_inputs.clear(); + decoder_inputs.push(best_token); + } + + if let Some(started) = started { + eprintln!( + "moonshine profile greedy_generate steps={steps} generated={} max_new_tokens={max_new_tokens} elapsed_ms={:.2}", + generated.len(), + started.elapsed().as_secs_f64() * 1000.0 + ); + } + + Ok(generated) + } + + async fn greedy_generate_from_prefix( + &mut self, + prefix_tokens: &[u32], + prefix_cache: &MoonshineDecoderCache, + encoder_hidden_states: &Tensor<3, f32>, + max_new_tokens: usize, + ) -> Result, crate::model::WhisperError> { + let profile = moonshine_profile(); + let started = profile.then(std::time::Instant::now); + let mut decoder_cache = prefix_cache.clone(); + let mut generated = prefix_tokens.to_vec(); + if prefix_tokens.len() >= max_new_tokens { + return Ok(generated); + } + + let mut decoder_inputs = vec![prefix_tokens.last().copied().unwrap_or(self.start_token)]; + let mut steps = 0usize; + for _ in 0..max_new_tokens.saturating_sub(prefix_tokens.len()).max(1) { + steps += 1; + let logits = self + .model + .decoder + .decode_cached(&decoder_inputs, &encoder_hidden_states, &mut decoder_cache) + .map_err(crate::model::WhisperError::Fusor)?; + let last_index = logits.shape()[1].saturating_sub(1); + let last_logits = logits + .narrow(1, last_index, 1) + .squeeze(1) + .squeeze(0) + .to_concrete(); + let best_token = last_logits.argmax_last_dim::<0>().to_scalar().await?; + if best_token == self.eos_token { + break; + } + generated.push(best_token); + decoder_inputs.clear(); + decoder_inputs.push(best_token); + } + + if let Some(started) = started { + eprintln!( + "moonshine profile greedy_generate_from_prefix prefix={} steps={steps} generated={} max_new_tokens={max_new_tokens} elapsed_ms={:.2}", + prefix_tokens.len(), + generated.len(), + started.elapsed().as_secs_f64() * 1000.0 + ); + } + + Ok(generated) + } + + async fn token_timestamps( + &mut self, + generated: &[u32], + encoder_hidden_states: &Tensor<3, f32>, + seconds_per_frame: f32, + ) -> Result, crate::model::WhisperError> { + if generated.is_empty() { + return Ok(Vec::new()); + } + let mut decoder_inputs = Vec::with_capacity(generated.len()); + decoder_inputs.push(self.start_token); + decoder_inputs.extend_from_slice(&generated[..generated.len().saturating_sub(1)]); + + let mut cross_attentions = Vec::new(); + let _ = self + .model + .decoder + .decode_prepared( + decoder_inputs.as_slice(), + encoder_hidden_states, + Some(&mut cross_attentions), + ) + .map_err(crate::model::WhisperError::Fusor)?; + + let mask = vec![vec![true; generated.len()]]; + let timestamps = extract_timestamps( + self.alignment_heads, + &cross_attentions, + const { NonZeroUsize::new(7).unwrap() }, + encoder_hidden_states.shape()[1], + seconds_per_frame, + mask, + ) + .await + .map_err(crate::model::WhisperError::Fusor)?; + + Ok(timestamps.into_iter().next().unwrap_or_default()) + } + + async fn decode_candidate_from_prepared( + &mut self, + adapted_encoder_hidden_states: &Tensor<3, f32>, + clip_end_s: f32, + usable_samples: usize, + compute_timestamps: bool, + reuse_prefix_tokens: &[u32], + reuse_prefix_cache: Option<&MoonshineDecoderCache>, + ) -> Result, crate::model::WhisperError> { + let total_frames = adapted_encoder_hidden_states.shape()[1]; + if usable_samples == 0 || total_frames == 0 { + return Ok(None); + } + + let profile = moonshine_profile(); + let started = profile.then(std::time::Instant::now); + let max_new_tokens = self.max_new_tokens(usable_samples); + let generated = if let Some(reuse_prefix_cache) = reuse_prefix_cache { + self.greedy_generate_from_prefix( + reuse_prefix_tokens, + reuse_prefix_cache, + adapted_encoder_hidden_states, + max_new_tokens, + ) + .await? + } else { + self.greedy_generate(adapted_encoder_hidden_states, max_new_tokens) + .await? + }; + let seconds_per_frame = if total_frames == 0 { + 0.0 + } else { + clip_end_s / total_frames as f32 + }; + let token_timestamps = if compute_timestamps { + self.token_timestamps(&generated, adapted_encoder_hidden_states, seconds_per_frame) + .await? + } else { + Vec::new() + }; + + if let Some(started) = started { + eprintln!( + "moonshine profile decode_candidate frames={total_frames} usable_s={:.3} reuse_prefix={} generated={} max_new_tokens={max_new_tokens} timestamps={} elapsed_ms={:.2}", + usable_samples as f32 / self.sample_rate as f32, + reuse_prefix_tokens.len(), + generated.len(), + compute_timestamps, + started.elapsed().as_secs_f64() * 1000.0 + ); + } + + Ok(Some(CandidateDecode { + tokens: generated, + token_timestamps, + clip_end_s, + usable_samples, + })) + } + + fn sync_stream_prefix_cache( + &mut self, + state: &mut MoonshineStreamState, + prefix_tokens: &[u32], + adapted_encoder_hidden_states: &Tensor<3, f32>, + ) -> Result<(), crate::model::WhisperError> { + if prefix_tokens.is_empty() { + if !state.decoder_prefix_tokens.is_empty() + || !state.decoder_prefix_cache.tokens.is_empty() + { + state.decoder_prefix_tokens.clear(); + state.decoder_prefix_cache = MoonshineDecoderCache::default(); + } + return Ok(()); + } + + let target_cache_tokens = &prefix_tokens[..prefix_tokens.len().saturating_sub(1)]; + let common_prefix = + self.common_prefix_len(&state.decoder_prefix_tokens, target_cache_tokens); + if common_prefix < state.decoder_prefix_tokens.len() { + state.decoder_prefix_tokens.clear(); + state.decoder_prefix_cache = state.decoder_cross_cache.clone(); + } + + let mut tokens_to_append = Vec::with_capacity( + target_cache_tokens + .len() + .saturating_sub(state.decoder_prefix_tokens.len()) + + 1, + ); + if state.decoder_prefix_cache.tokens.is_empty() { + state.decoder_prefix_cache = state.decoder_cross_cache.clone(); + tokens_to_append.push(self.start_token); + } + tokens_to_append + .extend_from_slice(&target_cache_tokens[state.decoder_prefix_tokens.len()..]); + + if !tokens_to_append.is_empty() { + let _ = self + .model + .decoder + .decode_cached( + &tokens_to_append, + adapted_encoder_hidden_states, + &mut state.decoder_prefix_cache, + ) + .map_err(crate::model::WhisperError::Fusor)?; + } + state.decoder_prefix_tokens = target_cache_tokens.to_vec(); + Ok(()) + } + + async fn populate_candidate_timestamps( + &mut self, + candidate: &mut CandidateDecode, + adapted_encoder_hidden_states: &Tensor<3, f32>, + ) -> Result<(), crate::model::WhisperError> { + if !candidate.token_timestamps.is_empty() || candidate.tokens.is_empty() { + return Ok(()); + } + let total_frames = adapted_encoder_hidden_states.shape()[1]; + let seconds_per_frame = if total_frames == 0 { + 0.0 + } else { + candidate.clip_end_s / total_frames as f32 + }; + candidate.token_timestamps = self + .token_timestamps( + &candidate.tokens, + adapted_encoder_hidden_states, + seconds_per_frame, + ) + .await?; + Ok(()) + } + + async fn update_stream_candidate( + &mut self, + state: &mut MoonshineStreamState, + samples: &[f32], + flush: bool, + ) -> Result, crate::model::WhisperError> { + let profile = moonshine_profile(); + let started = profile.then(std::time::Instant::now); + let encoder_append = self + .model + .encoder + .encode_stream(&self.device, &mut state.encoder_state, samples, flush) + .map_err(crate::model::WhisperError::Fusor)?; + if let Some(started) = started { + eprintln!( + "moonshine profile encode_stream samples={} flush={flush} hidden={} total_seen={} total_finalized={} elapsed_ms={:.2}", + samples.len(), + encoder_append.hidden_states.as_ref().map(|t| t.shape()[1]).unwrap_or(0), + encoder_append.total_seen_frames, + encoder_append.total_finalized_frames, + started.elapsed().as_secs_f64() * 1000.0 + ); + } + + let finalized_encoder_frames = state + .prepared_encoder_hidden_states + .as_ref() + .map(|tensor| tensor.shape()[1]) + .unwrap_or(0); + if encoder_append.hidden_states.is_none() + && encoder_append.total_finalized_frames == finalized_encoder_frames + { + return Ok(state.last_candidate.clone()); + } + + if let Some(hidden_states) = encoder_append.hidden_states { + let start_pos = state + .prepared_encoder_hidden_states + .as_ref() + .map(|tensor| tensor.shape()[1]) + .unwrap_or(0); + let prepared = self + .model + .decoder + .prepare_encoder_hidden_states_range(&hidden_states, start_pos) + .map_err(crate::model::WhisperError::Fusor)?; + state.prepared_encoder_hidden_states = + Some(match state.prepared_encoder_hidden_states.take() { + Some(existing) => Tensor::cat([existing, prepared], 1).to_concrete(), + None => prepared, + }); + } + + let Some(prepared_encoder_hidden_states) = state.prepared_encoder_hidden_states.as_ref() + else { + return Ok(None); + }; + let prepared_encoder_hidden_states = prepared_encoder_hidden_states.clone(); + let newly_finalized_frames = encoder_append + .total_finalized_frames + .saturating_sub(state.last_decoded_finalized_frames); + if !flush && newly_finalized_frames < self.stream_decode_interval_frames { + return Ok(state.last_candidate.clone()); + } + self.model + .decoder + .prepare_cross_attention_cache( + &prepared_encoder_hidden_states, + &mut state.decoder_cross_cache, + ) + .map_err(crate::model::WhisperError::Fusor)?; + + let clip_end_s = if encoder_append.total_seen_frames == 0 { + 0.0 + } else { + encoder_append.usable_input_samples as f32 / self.sample_rate as f32 + * encoder_append.total_finalized_frames as f32 + / encoder_append.total_seen_frames as f32 + }; + let reuse_prefix_tokens = state + .last_candidate + .as_ref() + .map(|candidate| { + candidate.tokens[..state.emitted_tokens.min(candidate.tokens.len())].to_vec() + }) + .unwrap_or_default(); + if !reuse_prefix_tokens.is_empty() { + self.sync_stream_prefix_cache( + state, + &reuse_prefix_tokens, + &prepared_encoder_hidden_states, + )?; + } + let reuse_prefix_cache = Some(if !reuse_prefix_tokens.is_empty() { + state.decoder_prefix_cache.clone() + } else { + state.decoder_cross_cache.clone() + }); + let candidate = self + .decode_candidate_from_prepared( + &prepared_encoder_hidden_states, + clip_end_s, + encoder_append.usable_input_samples, + false, + &reuse_prefix_tokens, + reuse_prefix_cache.as_ref(), + ) + .await?; + if candidate.is_some() { + state.last_decoded_finalized_frames = encoder_append.total_finalized_frames; + } + Ok(candidate) + } + + fn decoding_result_from_tokens( + &self, + tokens: &[u32], + token_timestamps: &[f32], + include_chunk_timestamps: bool, + clip_end_s: f32, + ) -> Result { + let mut processed_tokens = Vec::with_capacity(tokens.len()); + let mut timestamp_start: Option = None; + let mut prev_text_len = 0usize; + let mut chunks = Vec::new(); + let mut current_text = String::new(); + + for (index, token) in tokens.iter().copied().enumerate() { + processed_tokens.push(token); + if timestamp_start.is_none() { + timestamp_start = token_timestamps.get(index).copied(); + } + let detokenized = self + .tokenizer + .decode(&processed_tokens, true) + .map_err(crate::model::WhisperError::Tokenizer)?; + if detokenized.len() > prev_text_len + && detokenized.chars().last().unwrap_or_default().is_ascii() + { + let timestamp = if include_chunk_timestamps { + let start = timestamp_start.unwrap_or(0.0); + let end = token_timestamps.get(index).copied().unwrap_or(clip_end_s); + timestamp_start = Some(end); + Some(start..end) + } else { + None + }; + let text_range = current_text.len()..detokenized.len(); + current_text = detokenized; + prev_text_len = current_text.len(); + chunks.push(TokenChunk { + text_range, + timestamp, + }); + } else { + prev_text_len = detokenized.len(); + } + } + + if current_text.is_empty() && !tokens.is_empty() { + current_text = self + .tokenizer + .decode(tokens, true) + .map_err(crate::model::WhisperError::Tokenizer)?; + if !current_text.is_empty() { + chunks.push(TokenChunk { + text_range: 0..current_text.len(), + timestamp: include_chunk_timestamps.then_some(0.0..clip_end_s), + }); + } + } + + Ok(DecodingResult { + text: current_text, + avg_logprob: 0.0, + no_speech_prob: 0.0, + compression_ratio: 0.0, + chunks, + }) + } + + fn common_prefix_len(&self, left: &[u32], right: &[u32]) -> usize { + left.iter() + .zip(right.iter()) + .take_while(|(l, r)| l == r) + .count() + } + + fn stable_emit_len(&self, stable_len: usize, flush: bool) -> usize { + if flush { + stable_len + } else { + stable_len.saturating_sub(self.stream_stable_holdback_tokens) + } + } + + fn segment_from_token_range( + &self, + range: Range, + candidate: &CandidateDecode, + word_timestamps: bool, + fallback_start_s: f32, + ) -> Result, crate::model::WhisperError> { + if range.is_empty() || range.start >= candidate.tokens.len() { + return Ok(None); + } + + let prefix_text = if range.start == 0 { + String::new() + } else { + self.tokenizer + .decode(&candidate.tokens[..range.start], true) + .map_err(crate::model::WhisperError::Tokenizer)? + }; + let token_timestamps = if candidate.token_timestamps.is_empty() { + vec![] + } else { + candidate.token_timestamps[..range.end].to_vec() + }; + let mut result = self.decoding_result_from_tokens( + &candidate.tokens[..range.end], + &token_timestamps, + word_timestamps, + candidate.clip_end_s, + )?; + if result.text.len() >= prefix_text.len() && result.text.starts_with(&prefix_text) { + let prefix_len = prefix_text.len(); + result.text = result.text[prefix_len..].to_string(); + result.chunks = result + .chunks + .into_iter() + .filter_map(|mut chunk| { + if chunk.text_range.end <= prefix_len { + return None; + } + chunk.text_range.start = chunk.text_range.start.saturating_sub(prefix_len); + chunk.text_range.end -= prefix_len; + Some(chunk) + }) + .collect(); + } + if result.text.is_empty() { + return Ok(None); + } + + let total_tokens = candidate.tokens.len().max(1) as f32; + let mut start_s = token_timestamps + .get(range.start) + .copied() + .unwrap_or_else(|| candidate.clip_end_s * range.start as f32 / total_tokens); + let mut end_s = token_timestamps + .get(range.end.saturating_sub(1)) + .copied() + .unwrap_or_else(|| candidate.clip_end_s * range.end as f32 / total_tokens); + start_s = start_s.max(fallback_start_s); + if end_s <= start_s { + end_s = candidate.clip_end_s.max(start_s); + } + start_s = start_s.max(0.0); + let sample_start = (start_s * self.sample_rate as f32).round() as usize; + let sample_end = (end_s * self.sample_rate as f32).round() as usize; + let progress = if candidate.usable_samples == 0 { + 0.0 + } else { + sample_end as f32 / candidate.usable_samples as f32 + }; + + Ok(Some(Segment { + sample_range: sample_start..sample_end.max(sample_start), + start: start_s as f64, + duration: (end_s - start_s).max(0.0) as f64, + elapsed_time: None, + remaining_time: None, + progress: progress.clamp(0.0, 1.0), + result, + })) + } + + pub(crate) async fn transcribe( + &mut self, + samples: Vec, + language: Option, + word_timestamps: bool, + result: UnboundedSender, + ) -> Result<(), crate::model::WhisperError> { + if let Err(err) = self.validate_language(language) { + return Err(crate::model::WhisperError::Fusor(fusor::Error::msg( + err.to_string(), + ))); + } + let samples = self.usable_samples(&samples); + if samples.is_empty() { + return Ok(()); + } + let (adapted_encoder_hidden_states, total_frames) = self.encode_prepared_full(samples)?; + let clip_end_s = samples.len() as f32 / self.sample_rate as f32; + let Some(candidate) = self + .decode_candidate_from_prepared( + &adapted_encoder_hidden_states, + if total_frames == 0 { 0.0 } else { clip_end_s }, + samples.len(), + true, + &[], + None, + ) + .await? + else { + return Ok(()); + }; + if let Some(segment) = self.segment_from_token_range( + 0..candidate.tokens.len(), + &candidate, + word_timestamps, + 0.0, + )? { + let mut result = result; + let _ = result.start_send(segment); + } + Ok(()) + } + + pub(crate) fn start_stream( + &mut self, + session_id: u64, + language: Option, + word_timestamps: bool, + sender: UnboundedSender, + ) -> Result<(), crate::model::WhisperLoadingError> { + self.validate_language(language)?; + self.streams.insert( + session_id, + MoonshineStreamState { + word_timestamps, + sender, + last_candidate: None, + encoder_state: self.model.encoder.new_stream_state(), + prepared_encoder_hidden_states: None, + decoder_cross_cache: MoonshineDecoderCache::default(), + decoder_prefix_cache: MoonshineDecoderCache::default(), + decoder_prefix_tokens: Vec::new(), + last_decoded_finalized_frames: 0, + emitted_tokens: 0, + last_emitted_end_s: 0.0, + }, + ); + Ok(()) + } + + pub(crate) async fn push_stream_audio( + &mut self, + session_id: u64, + samples: Vec, + ) -> Result<(), crate::model::WhisperError> { + let Some(mut state) = self.streams.remove(&session_id) else { + return Ok(()); + }; + let Some(mut candidate) = self + .update_stream_candidate(&mut state, &samples, false) + .await? + else { + self.streams.insert(session_id, state); + return Ok(()); + }; + let stable = if let Some(previous) = &state.last_candidate { + self.stable_emit_len( + self.common_prefix_len(&previous.tokens, &candidate.tokens), + false, + ) + } else { + candidate + .tokens + .len() + .saturating_sub(self.stream_initial_holdback_tokens) + }; + if stable > state.emitted_tokens { + if state.word_timestamps { + if let Some(prepared_encoder_hidden_states) = + state.prepared_encoder_hidden_states.as_ref() + { + let prepared_encoder_hidden_states = prepared_encoder_hidden_states.clone(); + self.populate_candidate_timestamps( + &mut candidate, + &prepared_encoder_hidden_states, + ) + .await?; + } + } + if let Some(segment) = self.segment_from_token_range( + state.emitted_tokens..stable, + &candidate, + state.word_timestamps, + state.last_emitted_end_s, + )? { + state.last_emitted_end_s = segment.start as f32 + segment.duration as f32; + let mut sender = state.sender.clone(); + let _ = sender.start_send(segment); + } + state.emitted_tokens = stable; + } + state.last_candidate = Some(candidate); + self.streams.insert(session_id, state); + Ok(()) + } + + pub(crate) async fn finish_stream( + &mut self, + session_id: u64, + ) -> Result<(), crate::model::WhisperError> { + let Some(state) = self.streams.remove(&session_id) else { + return Ok(()); + }; + let mut state = state; + let Some(mut candidate) = self.update_stream_candidate(&mut state, &[], true).await? else { + return Ok(()); + }; + if state.word_timestamps { + if let Some(prepared_encoder_hidden_states) = + state.prepared_encoder_hidden_states.as_ref() + { + self.populate_candidate_timestamps(&mut candidate, prepared_encoder_hidden_states) + .await?; + } + } + if let Some(segment) = self.segment_from_token_range( + state.emitted_tokens..candidate.tokens.len(), + &candidate, + state.word_timestamps, + state.last_emitted_end_s, + )? { + let mut sender = state.sender; + let _ = sender.start_send(segment); + } + Ok(()) + } +} diff --git a/models/rwhisper/src/quantized/cohere.rs b/models/rwhisper/src/quantized/cohere.rs new file mode 100644 index 000000000..596af35c3 --- /dev/null +++ b/models/rwhisper/src/quantized/cohere.rs @@ -0,0 +1,1701 @@ +use fusor::{ + cache::{AttentionMask, KvCache, MaskCache, TensorCache}, + layers::{ + BatchNorm1d, Conv1d, Conv1dConfig, Conv2d, Conv2dConfig, Embedding, LayerNorm, Linear, + }, + Device, MaskKind, Result, Tensor, VarBuilder, +}; +use std::sync::Arc; + +use crate::cohere_config::{CohereConfig, CohereDecoderConfig, CohereEncoderConfig}; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::time::Instant; + +static COHERE_ENCODER_LAYER_CALLS: AtomicUsize = AtomicUsize::new(0); + +fn profile_enabled() -> bool { + std::env::var("RWHISPER_COHERE_PROFILE").ok().as_deref() == Some("1") +} + +fn cohere_layer_norm_3d(norm: &LayerNorm<1, f32>, input: &Tensor<3, f32>) -> Tensor<3, f32> { + if !input.is_gpu() { + return norm.forward(input); + } + + let profile = profile_enabled(); + let start = profile.then(Instant::now); + let [batch, time, hidden] = input.shape(); + let flat: Tensor<2, f32> = input.reshape([batch * time, hidden]).to_concrete(); + if let Some(start) = start { + eprintln!( + "cohere layer_norm reshape: {:.3}s", + start.elapsed().as_secs_f32() + ); + } + let mean = flat.mean_keepdim::<1>(1); + if let Some(start) = start { + eprintln!( + "cohere layer_norm mean: {:.3}s", + start.elapsed().as_secs_f32() + ); + } + let centered = (&flat - &mean.broadcast_as(flat.shape())).to_concrete(); + if let Some(start) = start { + eprintln!( + "cohere layer_norm center: {:.3}s", + start.elapsed().as_secs_f32() + ); + } + let normalized = centered.rms_norm_fused::<1, 1>(norm.weight(), norm.bias(), norm.eps()); + if let Some(start) = start { + eprintln!( + "cohere layer_norm rms: {:.3}s", + start.elapsed().as_secs_f32() + ); + } + normalized.reshape([batch, time, hidden]).to_concrete() +} + +fn conv_output_length(input: usize, kernel: usize, stride: usize, padding: usize) -> usize { + ((input + 2 * padding - kernel) / stride) + 1 +} + +fn rel_shift(x: &Tensor<4, f32>) -> Tensor<4, f32> { + let [batch, heads, q_len, pos_len] = x.shape(); + let zeros = Tensor::zeros(&x.device(), [batch, heads, q_len, 1]); + let padded = Tensor::cat([zeros, x.clone()], 3); + let reshaped = padded + .reshape([batch, heads, pos_len + 1, q_len]) + .to_concrete(); + reshaped + .narrow(2, 1, pos_len) + .reshape([batch, heads, q_len, pos_len]) + .to_concrete() +} + +fn valid_mask(device: &Device, batch: usize, seq_len: usize, valid_len: usize) -> Tensor<2, f32> { + let mut data = vec![0.0f32; batch * seq_len]; + for b in 0..batch { + for i in 0..valid_len.min(seq_len) { + data[b * seq_len + i] = 1.0; + } + } + Tensor::from_slice(device, [batch, seq_len], &data) +} + +fn encoder_attention_mask( + device: &Device, + batch: usize, + seq_len: usize, + valid_len: usize, +) -> Tensor<4, f32> { + let mut data = vec![0.0f32; batch * seq_len * seq_len]; + for b in 0..batch { + for q in 0..seq_len { + for k in 0..seq_len { + if q >= valid_len || k >= valid_len { + data[b * seq_len * seq_len + q * seq_len + k] = -1e9; + } + } + } + } + Tensor::from_slice(device, [batch, 1, seq_len, seq_len], &data) +} + +fn sigmoid(x: &Tensor<3, f32>) -> Tensor<3, f32> { + ((-x).exp().add_scalar(1.0)) + .to_concrete() + .div_scalar(1.0) + .recip() +} + +trait Reciprocal { + fn recip(&self) -> Self; +} + +impl Reciprocal for Tensor<3, f32> { + fn recip(&self) -> Self { + Tensor::splat(&self.device(), 1.0, self.shape()) + .div_(self) + .to_concrete() + } +} + +struct ConvSubsampling { + conv0: Conv2d, + conv1_dw: Conv2d, + conv1_pw: Conv2d, + conv2_dw: Conv2d, + conv2_pw: Conv2d, + out: Linear, +} + +impl ConvSubsampling { + fn load(device: &Device, vb: &mut VarBuilder, cfg: &CohereEncoderConfig) -> Result { + let conv_stride = [2, 2]; + let conv_padding = [1, 1]; + let depthwise_cfg = Conv2dConfig { + padding: conv_padding, + stride: conv_stride, + groups: cfg.subsampling_conv_channels, + dilation: [1, 1], + }; + let pointwise_cfg = Conv2dConfig { + padding: [0, 0], + stride: [1, 1], + groups: 1, + dilation: [1, 1], + }; + Ok(Self { + conv0: Conv2d::load( + device, + &mut vb.pp("conv.0"), + Conv2dConfig { + padding: conv_padding, + stride: conv_stride, + groups: 1, + dilation: [1, 1], + }, + )?, + conv1_dw: Conv2d::load(device, &mut vb.pp("conv.2"), depthwise_cfg)?, + conv1_pw: Conv2d::load(device, &mut vb.pp("conv.3"), pointwise_cfg)?, + conv2_dw: Conv2d::load(device, &mut vb.pp("conv.5"), depthwise_cfg)?, + conv2_pw: Conv2d::load(device, &mut vb.pp("conv.6"), pointwise_cfg)?, + out: Linear::load(device, &mut vb.pp("out"))?, + }) + } + + fn forward(&self, input: &Tensor<3, f32>, length: usize) -> (Tensor<3, f32>, usize) { + let mut x = input.transpose(1, 2).unsqueeze(1).to_concrete(); + x = self.conv0.forward(&x).relu().to_concrete(); + x = self.conv1_dw.forward(&x).to_concrete(); + x = self.conv1_pw.forward(&x).relu().to_concrete(); + x = self.conv2_dw.forward(&x).to_concrete(); + x = self.conv2_pw.forward(&x).relu().to_concrete(); + + let [batch, channels, time, freq] = x.shape(); + let x = x + .transpose(1, 2) + .reshape([batch, time, channels * freq]) + .to_concrete(); + let x = self.out.forward(&x); + + let mut reduced = length; + for _ in 0..3 { + reduced = conv_output_length(reduced, 3, 2, 1); + } + (x, reduced) + } +} + +struct RelPositionalEncoding { + pe: Tensor<2, f32>, +} + +impl RelPositionalEncoding { + fn new(device: &Device, d_model: usize, max_len: usize) -> Self { + let length = 2 * max_len - 1; + let mut data = vec![0.0f32; length * d_model]; + let center = max_len as isize - 1; + for pos_idx in 0..length { + let position = (center - pos_idx as isize) as f32; + for i in (0..d_model).step_by(2) { + let div = (-(10000.0f32).ln() * i as f32 / d_model as f32).exp(); + data[pos_idx * d_model + i] = (position * div).sin(); + if i + 1 < d_model { + data[pos_idx * d_model + i + 1] = (position * div).cos(); + } + } + } + Self { + pe: Tensor::from_slice(device, [length, d_model], &data), + } + } + + fn forward(&self, x: &Tensor<3, f32>) -> (Tensor<3, f32>, Tensor<3, f32>) { + let input_len = x.shape()[1]; + let center_pos = self.pe.shape()[0] / 2 + 1; + let start_pos = center_pos - input_len; + let pos_emb = self + .pe + .narrow(0, start_pos, 2 * input_len - 1) + .unsqueeze(0) + .to_concrete(); + (x.clone(), pos_emb) + } +} + +struct ConformerFeedForward { + linear1: Linear, + linear2: Linear, +} + +impl ConformerFeedForward { + fn load(device: &Device, vb: &mut VarBuilder) -> Result { + Ok(Self { + linear1: Linear::load(device, &mut vb.pp("linear1"))?, + linear2: Linear::load(device, &mut vb.pp("linear2"))?, + }) + } + + fn forward(&self, x: &Tensor<3, f32>) -> Tensor<3, f32> { + let profile = profile_enabled(); + let start = profile.then(Instant::now); + let x = self.linear1.forward(x); + if let Some(start) = start { + eprintln!("cohere ff linear1: {:.3}s", start.elapsed().as_secs_f32()); + } + let x = x.silu().to_concrete(); + if let Some(start) = start { + eprintln!("cohere ff silu: {:.3}s", start.elapsed().as_secs_f32()); + } + let x = self.linear2.forward(&x); + if let Some(start) = start { + eprintln!("cohere ff linear2: {:.3}s", start.elapsed().as_secs_f32()); + } + x + } +} + +struct ConformerConvolution { + pointwise_conv1: Conv1d, + depthwise_conv: Conv1d, + batch_norm: BatchNorm1d, + pointwise_conv2: Conv1d, +} + +impl ConformerConvolution { + fn load( + device: &Device, + vb: &mut VarBuilder, + d_model: usize, + kernel_size: usize, + ) -> Result { + Ok(Self { + pointwise_conv1: Conv1d::new( + vb.pp("pointwise_conv1").get("weight", device)?.dequantize(), + Some(vb.pp("pointwise_conv1").get("bias", device)?.dequantize()), + Conv1dConfig { + padding: 0, + stride: 1, + groups: 1, + dilation: 1, + }, + ), + depthwise_conv: Conv1d::new( + vb.pp("depthwise_conv").get("weight", device)?.dequantize(), + Some(vb.pp("depthwise_conv").get("bias", device)?.dequantize()), + Conv1dConfig { + padding: (kernel_size - 1) / 2, + stride: 1, + groups: d_model, + dilation: 1, + }, + ), + batch_norm: BatchNorm1d::load(device, &mut vb.pp("batch_norm"), 1e-5)?, + pointwise_conv2: Conv1d::new( + vb.pp("pointwise_conv2").get("weight", device)?.dequantize(), + Some(vb.pp("pointwise_conv2").get("bias", device)?.dequantize()), + Conv1dConfig { + padding: 0, + stride: 1, + groups: 1, + dilation: 1, + }, + ), + }) + } + + fn forward(&self, x: &Tensor<3, f32>, valid_mask: Option<&Tensor<2, f32>>) -> Tensor<3, f32> { + let x = x.transpose(1, 2).to_concrete(); + let x = self.pointwise_conv1.forward(&x); + let chunks = x.chunk(2, 1); + let x = (chunks[0].to_concrete() * sigmoid(&chunks[1].to_concrete())).to_concrete(); + let x = if let Some(mask) = valid_mask { + let mask_unsqueezed = mask.unsqueeze(1); + let mask = mask_unsqueezed.broadcast_as(x.shape()); + (x * mask).to_concrete() + } else { + x + }; + let x = self.depthwise_conv.forward(&x); + let x = self.batch_norm.forward(&x).silu().to_concrete(); + self.pointwise_conv2 + .forward(&x) + .transpose(1, 2) + .to_concrete() + } +} + +struct RelPositionMultiHeadAttention { + linear_q: Linear, + linear_k: Linear, + linear_v: Linear, + linear_pos: Linear, + linear_out: Linear, + pos_bias_u: Tensor<2, f32>, + pos_bias_v: Tensor<2, f32>, + num_heads: usize, + head_dim: usize, + scale: f32, +} + +impl RelPositionMultiHeadAttention { + fn load( + device: &Device, + vb: &mut VarBuilder, + d_model: usize, + num_heads: usize, + ) -> Result { + let head_dim = d_model / num_heads; + Ok(Self { + linear_q: Linear::load(device, &mut vb.pp("linear_q"))?, + linear_k: Linear::load(device, &mut vb.pp("linear_k"))?, + linear_v: Linear::load(device, &mut vb.pp("linear_v"))?, + linear_pos: Linear::load(device, &mut vb.pp("linear_pos"))?, + linear_out: Linear::load(device, &mut vb.pp("linear_out"))?, + pos_bias_u: vb.get("pos_bias_u", device)?.dequantize(), + pos_bias_v: vb.get("pos_bias_v", device)?.dequantize(), + num_heads, + head_dim, + scale: (head_dim as f32).powf(-0.5), + }) + } + + fn reshape_heads(&self, x: &Tensor<3, f32>) -> Tensor<4, f32> { + let [batch, time, _] = x.shape(); + x.reshape([batch, time, self.num_heads, self.head_dim]) + .transpose(1, 2) + .to_concrete() + } + + fn forward( + &self, + x: &Tensor<3, f32>, + pos_emb: &Tensor<3, f32>, + mask: Option<&Tensor<4, f32>>, + ) -> Tensor<3, f32> { + let [batch, time, _] = x.shape(); + let q = self.reshape_heads(&self.linear_q.forward(x)); + let k = self.reshape_heads(&self.linear_k.forward(x)); + let v = self.reshape_heads(&self.linear_v.forward(x)); + + let pos = if pos_emb.shape()[0] == 1 && batch > 1 { + pos_emb + .broadcast_as([batch, pos_emb.shape()[1], pos_emb.shape()[2]]) + .to_concrete() + } else { + pos_emb.to_concrete() + }; + let p = self + .linear_pos + .forward(&pos.to_concrete()) + .reshape([batch, pos.shape()[1], self.num_heads, self.head_dim]) + .transpose(1, 2) + .to_concrete(); + + let q_with_u = (q.clone() + + self + .pos_bias_u + .reshape([1, self.num_heads, 1, self.head_dim]) + .broadcast_as(q.shape())) + .to_concrete(); + let q_with_v = (q + self + .pos_bias_v + .reshape([1, self.num_heads, 1, self.head_dim]) + .broadcast_as([batch, self.num_heads, time, self.head_dim])) + .to_concrete(); + + let matrix_ac = q_with_u.mat_mul(&k.transpose(2, 3).to_concrete()); + let matrix_bd = rel_shift(&q_with_v.mat_mul(&p.transpose(2, 3).to_concrete())) + .narrow(3, 0, matrix_ac.shape()[3]) + .to_concrete(); + let mut scores = (matrix_ac + matrix_bd).to_concrete().mul_scalar(self.scale); + if let Some(mask) = mask { + scores = scores.add_(mask); + } + let attn = scores.softmax_last_dim_fused(); + let context = attn + .mat_mul(&v) + .transpose(1, 2) + .reshape([batch, time, self.num_heads * self.head_dim]) + .to_concrete(); + self.linear_out.forward(&context) + } +} + +struct ConformerLayer { + norm_feed_forward1: LayerNorm<1, f32>, + feed_forward1: ConformerFeedForward, + norm_self_att: LayerNorm<1, f32>, + self_attn: RelPositionMultiHeadAttention, + norm_conv: LayerNorm<1, f32>, + conv: ConformerConvolution, + norm_feed_forward2: LayerNorm<1, f32>, + feed_forward2: ConformerFeedForward, + norm_out: LayerNorm<1, f32>, +} + +impl ConformerLayer { + fn load(device: &Device, vb: &mut VarBuilder, cfg: &CohereEncoderConfig) -> Result { + Ok(Self { + norm_feed_forward1: LayerNorm::load(device, &mut vb.pp("norm_feed_forward1"), 1e-5)?, + feed_forward1: ConformerFeedForward::load(device, &mut vb.pp("feed_forward1"))?, + norm_self_att: LayerNorm::load(device, &mut vb.pp("norm_self_att"), 1e-5)?, + self_attn: RelPositionMultiHeadAttention::load( + device, + &mut vb.pp("self_attn"), + cfg.d_model, + cfg.n_heads, + )?, + norm_conv: LayerNorm::load(device, &mut vb.pp("norm_conv"), 1e-5)?, + conv: ConformerConvolution::load( + device, + &mut vb.pp("conv"), + cfg.d_model, + cfg.conv_kernel_size, + )?, + norm_feed_forward2: LayerNorm::load(device, &mut vb.pp("norm_feed_forward2"), 1e-5)?, + feed_forward2: ConformerFeedForward::load(device, &mut vb.pp("feed_forward2"))?, + norm_out: LayerNorm::load(device, &mut vb.pp("norm_out"), 1e-5)?, + }) + } + + fn forward( + &self, + x: &Tensor<3, f32>, + pos_emb: &Tensor<3, f32>, + mask: Option<&Tensor<4, f32>>, + valid_mask: Option<&Tensor<2, f32>>, + ) -> Tensor<3, f32> { + let profile = profile_enabled(); + let layer_idx = COHERE_ENCODER_LAYER_CALLS.fetch_add(1, Ordering::Relaxed) + 1; + let layer_start = profile.then(Instant::now); + let residual = x.clone(); + let normed_ff1 = cohere_layer_norm_3d(&self.norm_feed_forward1, x); + if let Some(layer_start) = layer_start { + eprintln!( + "cohere encoder layer {layer_idx} ff1 norm: {:.3}s", + layer_start.elapsed().as_secs_f32() + ); + } + let x = (residual + self.feed_forward1.forward(&normed_ff1).mul_scalar(0.5)).to_concrete(); + if let Some(layer_start) = layer_start { + eprintln!( + "cohere encoder layer {layer_idx} ff1: {:.3}s", + layer_start.elapsed().as_secs_f32() + ); + } + let residual = x.clone(); + let x = (residual + + self.self_attn.forward( + &cohere_layer_norm_3d(&self.norm_self_att, &x), + pos_emb, + mask, + )) + .to_concrete(); + if let Some(layer_start) = layer_start { + eprintln!( + "cohere encoder layer {layer_idx} attn: {:.3}s", + layer_start.elapsed().as_secs_f32() + ); + } + let residual = x.clone(); + let x = (residual + + self + .conv + .forward(&cohere_layer_norm_3d(&self.norm_conv, &x), valid_mask)) + .to_concrete(); + if let Some(layer_start) = layer_start { + eprintln!( + "cohere encoder layer {layer_idx} conv: {:.3}s", + layer_start.elapsed().as_secs_f32() + ); + } + let residual = x.clone(); + let out = cohere_layer_norm_3d( + &self.norm_out, + &(residual + + self + .feed_forward2 + .forward(&cohere_layer_norm_3d(&self.norm_feed_forward2, &x)) + .mul_scalar(0.5)) + .to_concrete(), + ); + if let Some(layer_start) = layer_start { + eprintln!( + "cohere encoder layer {layer_idx} ff2+out: {:.3}s", + layer_start.elapsed().as_secs_f32() + ); + } + out + } +} + +struct ConformerEncoder { + pre_encode: ConvSubsampling, + pos_enc: RelPositionalEncoding, + layers: Vec, +} + +impl ConformerEncoder { + fn load(device: &Device, vb: &mut VarBuilder, cfg: &CohereEncoderConfig) -> Result { + let mut layers = Vec::with_capacity(cfg.n_layers); + for i in 0..cfg.n_layers { + layers.push(ConformerLayer::load( + device, + &mut vb.pp(format!("layers.{i}")), + cfg, + )?); + } + Ok(Self { + pre_encode: ConvSubsampling::load(device, &mut vb.pp("pre_encode"), cfg)?, + pos_enc: RelPositionalEncoding::new(device, cfg.d_model, cfg.pos_emb_max_len), + layers, + }) + } + + fn forward(&self, input_features: &Tensor<3, f32>, length: usize) -> (Tensor<3, f32>, usize) { + let profile = profile_enabled(); + let encode_start = profile.then(Instant::now); + let (x, length) = self.pre_encode.forward(input_features, length); + if let Some(encode_start) = encode_start { + eprintln!( + "cohere pre_encode: {:.3}s time={} length={}", + encode_start.elapsed().as_secs_f32(), + x.shape()[1], + length + ); + } + let time = x.shape()[1]; + let (mut x, pos_emb) = self.pos_enc.forward(&x); + if let Some(encode_start) = encode_start { + eprintln!( + "cohere pos_enc: {:.3}s", + encode_start.elapsed().as_secs_f32() + ); + } + let valid = valid_mask(&x.device(), x.shape()[0], time, length); + if let Some(encode_start) = encode_start { + eprintln!( + "cohere valid mask: {:.3}s", + encode_start.elapsed().as_secs_f32() + ); + } + let att_mask = encoder_attention_mask(&x.device(), x.shape()[0], time, length); + if let Some(encode_start) = encode_start { + eprintln!( + "cohere att mask: {:.3}s", + encode_start.elapsed().as_secs_f32() + ); + } + for (i, layer) in self.layers.iter().enumerate() { + x = layer.forward(&x, &pos_emb, Some(&att_mask), Some(&valid)); + if profile && (i + 1) % 4 == 0 { + eprintln!( + "cohere encoder layer {}: {:.3}s", + i + 1, + encode_start.unwrap().elapsed().as_secs_f32() + ); + } + } + if let Some(encode_start) = encode_start { + eprintln!( + "cohere encode total: {:.3}s", + encode_start.elapsed().as_secs_f32() + ); + } + (x, length) + } +} + +struct FixedPositionalEncoding { + table: Embedding, +} + +impl FixedPositionalEncoding { + fn load(device: &Device, vb: &mut VarBuilder) -> Result { + Ok(Self { + table: Embedding::new_from_tensor(vb.get("pos_enc", device)?.dequantize()), + }) + } + + fn forward(&self, positions: &Tensor<2, u32>) -> Tensor<3, f32> { + self.table.forward(positions) + } +} + +struct DecoderAttention { + query_net: Linear, + key_net: Linear, + value_net: Linear, + out_projection: Linear, + num_heads: usize, + head_dim: usize, + scale: f32, +} + +struct DecoderAttentionCache { + kv_cache: KvCache, +} + +impl DecoderAttentionCache { + fn new(max_seq_len: usize) -> Self { + Self { + kv_cache: KvCache::new(1, max_seq_len), + } + } +} + +impl DecoderAttention { + fn load( + device: &Device, + vb: &mut VarBuilder, + hidden_size: usize, + num_heads: usize, + ) -> Result { + let head_dim = hidden_size / num_heads; + Ok(Self { + query_net: Linear::load(device, &mut vb.pp("query_net"))?, + key_net: Linear::load(device, &mut vb.pp("key_net"))?, + value_net: Linear::load(device, &mut vb.pp("value_net"))?, + out_projection: Linear::load(device, &mut vb.pp("out_projection"))?, + num_heads, + head_dim, + scale: (head_dim as f32).powf(-0.5), + }) + } + + fn reshape_heads(&self, x: &Tensor<3, f32>) -> Tensor<4, f32> { + let [batch, time, _] = x.shape(); + x.reshape([batch, time, self.num_heads, self.head_dim]) + .transpose(1, 2) + .to_concrete() + } + + fn forward_kv( + &self, + x: &Tensor<3, f32>, + cache: Option<&mut DecoderAttentionCache>, + ) -> (Tensor<3, f32>, Tensor<3, f32>) { + let device = x.device(); + let key_states = self.key_net.forward(x); + let value_states = self.value_net.forward(x); + match cache { + None => (key_states, value_states), + Some(cache) => { + let key_states_4d = key_states.unsqueeze(2).to_concrete(); + let value_states_4d = value_states.unsqueeze(2).to_concrete(); + let (k, v) = cache + .kv_cache + .append(&device, &key_states_4d, &value_states_4d); + (k.squeeze(2).to_concrete(), v.squeeze(2).to_concrete()) + } + } + } + + fn qkv_attention( + &self, + q: &Tensor<3, f32>, + k: &Tensor<3, f32>, + v: &Tensor<3, f32>, + attention_mask: Option<&AttentionMask>, + attention_output: Option<&mut TensorCache<4, f32>>, + ) -> Tensor<3, f32> { + let [batch, q_time, _] = q.shape(); + let q = self.reshape_heads(q); + let v = self.reshape_heads(v); + let k = self.reshape_heads(k); + + if attention_output.is_none() && attention_mask.is_some() { + let mask = attention_mask.map(|mask| (mask.mask(), MaskKind::QKMask)); + let context = q + .flash_attention(&k, &v, self.scale, mask) + .transpose(1, 2) + .reshape([batch, q_time, self.num_heads * self.head_dim]) + .to_concrete(); + return self.out_projection.forward(&context); + } + + let k = k.transpose(2, 3).to_concrete(); + let mut scores = q.mat_mul(&k).to_concrete().mul_scalar(self.scale); + if let Some(mask) = attention_mask { + mask.forward(&mut scores); + } + if let Some(output) = attention_output { + let last_query = scores.narrow(2, q_time - 1, 1).to_concrete(); + output.append(&q.device(), &last_query); + } + let attn = scores.softmax_last_dim_fused(); + let context = attn + .mat_mul(&v) + .transpose(1, 2) + .reshape([batch, q_time, self.num_heads * self.head_dim]) + .to_concrete(); + self.out_projection.forward(&context) + } + + fn forward( + &self, + hidden_states: &Tensor<3, f32>, + kv: (Tensor<3, f32>, Tensor<3, f32>), + attention_mask: Option<&AttentionMask>, + attention_output: Option<&mut TensorCache<4, f32>>, + ) -> Tensor<3, f32> { + let query_states = self.query_net.forward(hidden_states); + let (key_states, value_states) = &kv; + self.qkv_attention( + &query_states, + key_states, + value_states, + attention_mask, + attention_output, + ) + } +} + +struct DecoderFeedForward { + dense_in: Linear, + dense_out: Linear, + hidden_act: String, +} + +impl DecoderFeedForward { + fn load(device: &Device, vb: &mut VarBuilder, hidden_act: &str) -> Result { + Ok(Self { + dense_in: Linear::load(device, &mut vb.pp("dense_in"))?, + dense_out: Linear::load(device, &mut vb.pp("dense_out"))?, + hidden_act: hidden_act.to_owned(), + }) + } + + fn forward(&self, x: &Tensor<3, f32>) -> Tensor<3, f32> { + let x = self.dense_in.forward(x); + let x = match self.hidden_act.as_str() { + "relu" => x.relu().to_concrete(), + "silu" | "swish" => x.silu().to_concrete(), + _ => x.relu().to_concrete(), + }; + self.dense_out.forward(&x) + } +} + +struct TransformerDecoderLayer { + layer_norm_1: LayerNorm<1, f32>, + first_sub_layer: DecoderAttention, + layer_norm_2: LayerNorm<1, f32>, + second_sub_layer: DecoderAttention, + layer_norm_3: LayerNorm<1, f32>, + third_sub_layer: DecoderFeedForward, +} + +struct TransformerDecoderLayerCache { + self_attn: DecoderAttentionCache, + cross_attn_kv: (Tensor<3, f32>, Tensor<3, f32>), +} + +impl TransformerDecoderLayer { + fn load(device: &Device, vb: &mut VarBuilder, cfg: &CohereDecoderConfig) -> Result { + Ok(Self { + layer_norm_1: LayerNorm::load(device, &mut vb.pp("layer_norm_1"), 1e-5)?, + first_sub_layer: DecoderAttention::load( + device, + &mut vb.pp("first_sub_layer"), + cfg.hidden_size, + cfg.num_attention_heads, + )?, + layer_norm_2: LayerNorm::load(device, &mut vb.pp("layer_norm_2"), 1e-5)?, + second_sub_layer: DecoderAttention::load( + device, + &mut vb.pp("second_sub_layer"), + cfg.hidden_size, + cfg.num_attention_heads, + )?, + layer_norm_3: LayerNorm::load(device, &mut vb.pp("layer_norm_3"), 1e-5)?, + third_sub_layer: DecoderFeedForward::load( + device, + &mut vb.pp("third_sub_layer"), + &cfg.hidden_act, + )?, + }) + } + + fn forward( + &self, + hidden_states: &Tensor<3, f32>, + encoder_hidden_states: &Tensor<3, f32>, + self_attention_mask: &Tensor<4, f32>, + _cross_attention_mask: Option<&Tensor<4, f32>>, + ) -> Tensor<3, f32> { + let residual = hidden_states.clone(); + let self_kv = self.first_sub_layer.forward_kv( + &cohere_layer_norm_3d(&self.layer_norm_1, hidden_states), + None, + ); + let self_mask = AttentionMask::new( + self_attention_mask + .squeeze::<3>(0) + .squeeze::<2>(0) + .to_concrete(), + ); + let hidden_states = (residual + + self.first_sub_layer.forward( + &cohere_layer_norm_3d(&self.layer_norm_1, hidden_states), + self_kv, + Some(&self_mask), + None, + )) + .to_concrete(); + let residual = hidden_states.clone(); + let cross_kv = self + .second_sub_layer + .forward_kv(encoder_hidden_states, None); + let hidden_states = (residual + + self.second_sub_layer.forward( + &cohere_layer_norm_3d(&self.layer_norm_2, &hidden_states), + cross_kv, + None, + None, + )) + .to_concrete(); + let residual = hidden_states.clone(); + (residual + + self + .third_sub_layer + .forward(&cohere_layer_norm_3d(&self.layer_norm_3, &hidden_states))) + .to_concrete() + } + + fn forward_cached( + &self, + hidden_states: &Tensor<3, f32>, + self_attention_mask: &AttentionMask, + cache: &mut TransformerDecoderLayerCache, + attention_output: Option<&mut TensorCache<4, f32>>, + ) -> Tensor<3, f32> { + let profile = profile_enabled(); + let start = profile.then(Instant::now); + let ln1 = cohere_layer_norm_3d(&self.layer_norm_1, hidden_states); + if let Some(start) = start { + eprintln!( + "cohere decoder layer cached ln1: {:.3}s", + start.elapsed().as_secs_f32() + ); + } + let self_kv = self + .first_sub_layer + .forward_kv(&ln1, Some(&mut cache.self_attn)); + if let Some(start) = start { + eprintln!( + "cohere decoder layer cached self kv: {:.3}s", + start.elapsed().as_secs_f32() + ); + } + let residual = hidden_states.clone(); + let hidden_states = (residual + + self + .first_sub_layer + .forward(&ln1, self_kv, Some(self_attention_mask), None)) + .to_concrete(); + if let Some(start) = start { + eprintln!( + "cohere decoder layer cached self attn: {:.3}s", + start.elapsed().as_secs_f32() + ); + } + + let residual = hidden_states.clone(); + let hidden_states = (residual + + self.second_sub_layer.forward( + &cohere_layer_norm_3d(&self.layer_norm_2, &hidden_states), + cache.cross_attn_kv.clone(), + None, + attention_output, + )) + .to_concrete(); + if let Some(start) = start { + eprintln!( + "cohere decoder layer cached cross attn: {:.3}s", + start.elapsed().as_secs_f32() + ); + } + + let residual = hidden_states.clone(); + let hidden_states = (residual + + self + .third_sub_layer + .forward(&cohere_layer_norm_3d(&self.layer_norm_3, &hidden_states))) + .to_concrete(); + if let Some(start) = start { + eprintln!( + "cohere decoder layer cached mlp: {:.3}s", + start.elapsed().as_secs_f32() + ); + } + hidden_states + } +} + +struct TransformerDecoderEmbedding { + token_embedding: Embedding, + position_embedding: FixedPositionalEncoding, + layer_norm: LayerNorm<1, f32>, +} + +impl TransformerDecoderEmbedding { + fn load(device: &Device, vb: &mut VarBuilder) -> Result { + Ok(Self { + token_embedding: Embedding::load(device, &mut vb.pp("token_embedding"))?, + position_embedding: FixedPositionalEncoding::load( + device, + &mut vb.pp("position_embedding"), + )?, + layer_norm: LayerNorm::load(device, &mut vb.pp("layer_norm"), 1e-5)?, + }) + } + + fn forward(&self, input_ids: &Tensor<2, u32>, positions: &Tensor<2, u32>) -> Tensor<3, f32> { + cohere_layer_norm_3d( + &self.layer_norm, + &(self.token_embedding.forward(input_ids) + self.position_embedding.forward(positions)) + .to_concrete(), + ) + } +} + +struct TransformerDecoder { + embedding: TransformerDecoderEmbedding, + layers: Vec, + final_layer_norm: LayerNorm<1, f32>, + max_target_positions: usize, + mask_cache: Arc>, +} + +#[derive(Default)] +struct TransformerDecoderCache { + tokens: Vec, + layers: Vec, +} + +impl TransformerDecoder { + fn load(device: &Device, vb: &mut VarBuilder, cfg: &CohereDecoderConfig) -> Result { + let mut layers = Vec::with_capacity(cfg.num_layers); + for i in 0..cfg.num_layers { + layers.push(TransformerDecoderLayer::load( + device, + &mut vb.pp(format!("_decoder.layers.{i}")), + cfg, + )?); + } + Ok(Self { + embedding: TransformerDecoderEmbedding::load(device, &mut vb.pp("_embedding"))?, + layers, + final_layer_norm: LayerNorm::load( + device, + &mut vb.pp("_decoder.final_layer_norm"), + 1e-5, + )?, + max_target_positions: cfg.max_sequence_length, + mask_cache: Default::default(), + }) + } + + fn forward( + &self, + input_ids: &Tensor<2, u32>, + encoder_hidden_states: &Tensor<3, f32>, + encoder_length: usize, + ) -> Tensor<3, f32> { + let [batch, seq_len] = input_ids.shape(); + let positions: Vec = (0..seq_len as u32).collect(); + let positions = Tensor::from_slice(&input_ids.device(), [batch, seq_len], &positions); + let mut hidden_states = self.embedding.forward(input_ids, &positions); + let self_mask = AttentionMask::causal(&input_ids.device(), seq_len); + let cross_mask = if encoder_length < encoder_hidden_states.shape()[1] { + Some(encoder_attention_mask( + &input_ids.device(), + encoder_hidden_states.shape()[0], + encoder_hidden_states.shape()[1], + encoder_length, + )) + } else { + None + }; + for layer in &self.layers { + let self_mask_tensor = self_mask + .mask() + .clone() + .unsqueeze(0) + .unsqueeze(0) + .to_concrete(); + hidden_states = layer.forward( + &hidden_states, + encoder_hidden_states, + &self_mask_tensor, + cross_mask.as_ref(), + ); + } + cohere_layer_norm_3d(&self.final_layer_norm, &hidden_states) + } + + fn forward_cached( + &self, + tokens: &[u32], + encoder_hidden_states: &Tensor<3, f32>, + _encoder_length: usize, + cache: &mut TransformerDecoderCache, + mut attention_output: Option<&mut [TensorCache<4, f32>]>, + ) -> Tensor<3, f32> { + let profile = profile_enabled(); + let start = profile.then(Instant::now); + let index_pos = cache.tokens.len(); + cache.tokens.extend_from_slice(tokens); + let seq_len = tokens.len(); + assert!( + index_pos + seq_len <= self.max_target_positions, + "exceeded max sequence length" + ); + + let device = encoder_hidden_states.device(); + let self_mask = self.mask_cache.get_mask(seq_len, index_pos, None, &device); + let token_tensor: Tensor<1, u32> = Tensor::new(&device, tokens); + let token_tensor = token_tensor.unsqueeze(0).to_concrete(); + let positions: Vec = (index_pos as u32..(index_pos + seq_len) as u32).collect(); + let positions = Tensor::from_slice(&device, [1, seq_len], &positions); + let mut hidden_states = self.embedding.forward(&token_tensor, &positions); + if let Some(start) = start { + eprintln!( + "cohere decoder embedding: {:.3}s", + start.elapsed().as_secs_f32() + ); + } + + for (i, layer) in self.layers.iter().enumerate() { + if cache.layers.len() <= i { + if let Some(start) = start { + eprintln!( + "cohere decoder layer {i} cache init start: {:.3}s", + start.elapsed().as_secs_f32() + ); + } + let cross_attn_kv = layer + .second_sub_layer + .forward_kv(encoder_hidden_states, None); + cache.layers.push(TransformerDecoderLayerCache { + self_attn: DecoderAttentionCache::new(self.max_target_positions), + cross_attn_kv, + }); + if let Some(start) = start { + eprintln!( + "cohere decoder layer {i} cache init done: {:.3}s", + start.elapsed().as_secs_f32() + ); + } + } + let layer_cache = &mut cache.layers[i]; + let attention_output = attention_output.as_mut().map(|outputs| &mut outputs[i]); + hidden_states = + layer.forward_cached(&hidden_states, &self_mask, layer_cache, attention_output); + if let Some(start) = start { + eprintln!( + "cohere decoder layer {i} forward: {:.3}s", + start.elapsed().as_secs_f32() + ); + } + } + + let hidden_states = cohere_layer_norm_3d(&self.final_layer_norm, &hidden_states); + if let Some(start) = start { + eprintln!( + "cohere decoder final ln: {:.3}s", + start.elapsed().as_secs_f32() + ); + } + hidden_states + } +} + +pub struct Cohere { + pub config: CohereConfig, + encoder: ConformerEncoder, + decoder: TransformerDecoder, + encoder_decoder_proj: Option>, + lm_bias: Tensor<1, f32>, +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::cohere_audio::pcm_to_features; + use std::path::Path; + + #[tokio::test] + async fn debug_first_step_topk() { + let root = Path::new("/tmp/cohere-transcribe-03-2026"); + if !root.exists() { + return; + } + + let config: CohereConfig = + serde_json::from_slice(&std::fs::read(root.join("config.json")).unwrap()).unwrap(); + let weights = std::fs::read(root.join("model.gguf")).unwrap(); + let device = Device::cpu(); + let mut reader = std::io::Cursor::new(weights); + let mut vb = VarBuilder::from_gguf(&mut reader).unwrap(); + let model = Cohere::load(&device, &mut vb, config.clone()).unwrap(); + + let wav = hound::WavReader::open(concat!( + env!("CARGO_MANIFEST_DIR"), + "/examples/samples_jfk.wav" + )) + .unwrap() + .into_samples::() + .map(|sample| sample.unwrap() as f32 / 32768.0) + .take(16_000) + .collect::>(); + + let filter_bytes = include_bytes!("../cohere_melfilters128.bytes").as_slice(); + let mut filterbank = vec![0.0f32; filter_bytes.len() / 4]; + ::read_f32_into( + filter_bytes, + &mut filterbank, + ); + for value in &mut filterbank { + *value = half::bf16::from_f32(*value).to_f32(); + } + + let (features, total_frames, valid_frames) = pcm_to_features(&config, &wav, &filterbank); + assert_eq!(total_frames, 101); + assert_eq!(valid_frames, 100); + let input_features = Tensor::from_slice( + &device, + [1, config.preprocessor.features, total_frames], + &features, + ); + let prompt_ids = [7_u32, 4, 16, 62, 62, 5, 9, 11, 13]; + + let mut conv = input_features.transpose(1, 2).unsqueeze(1).to_concrete(); + conv = model + .encoder + .pre_encode + .conv0 + .forward(&conv) + .relu() + .to_concrete(); + conv = model + .encoder + .pre_encode + .conv1_dw + .forward(&conv) + .to_concrete(); + conv = model + .encoder + .pre_encode + .conv1_pw + .forward(&conv) + .relu() + .to_concrete(); + conv = model + .encoder + .pre_encode + .conv2_dw + .forward(&conv) + .to_concrete(); + conv = model + .encoder + .pre_encode + .conv2_pw + .forward(&conv) + .relu() + .to_concrete(); + let conv_slice = conv.clone().as_slice().await.unwrap(); + let conv_shape = conv_slice.shape(); + let (conv_b, conv_c, conv_t, conv_f) = + (conv_shape[0], conv_shape[1], conv_shape[2], conv_shape[3]); + let mut conv_data = Vec::with_capacity(conv_b * conv_c * conv_t * conv_f); + for b in 0..conv_b { + for c in 0..conv_c { + for t in 0..conv_t { + for f in 0..conv_f { + conv_data.push(conv_slice[[b, c, t, f]]); + } + } + } + } + let conv_sum: f32 = conv_data.iter().sum(); + let conv_mean = conv_sum / conv_data.len() as f32; + let conv_var = conv_data + .iter() + .map(|value| { + let diff = *value - conv_mean; + diff * diff + }) + .sum::() + / conv_data.len() as f32; + println!( + "conv shape={:?} sum={} mean={} std={} first20={:?}", + conv_slice.shape(), + conv_sum, + conv_mean, + conv_var.sqrt(), + &conv_data[..20] + ); + + let (pre_encode, pre_len) = model + .encoder + .pre_encode + .forward(&input_features, valid_frames); + let pre = pre_encode.clone().as_slice().await.unwrap(); + let pre_shape = pre.shape(); + let (pre_batch, pre_time, pre_hidden) = (pre_shape[0], pre_shape[1], pre_shape[2]); + let mut pre_data = Vec::with_capacity(pre_batch * pre_time * pre_hidden); + for b in 0..pre_batch { + for t in 0..pre_time { + for h in 0..pre_hidden { + pre_data.push(pre[[b, t, h]]); + } + } + } + let pre_sum: f32 = pre_data.iter().sum(); + let pre_mean = pre_sum / pre_data.len() as f32; + let pre_var = pre_data + .iter() + .map(|value| { + let diff = *value - pre_mean; + diff * diff + }) + .sum::() + / pre_data.len() as f32; + println!( + "pre shape={:?} len={} sum={} mean={} std={} first20={:?}", + pre.shape(), + pre_len, + pre_sum, + pre_mean, + pre_var.sqrt(), + &pre_data[..20] + ); + + let (encoder_hidden_states, encoder_length) = model.encode(&input_features, valid_frames); + let enc = encoder_hidden_states.clone().as_slice().await.unwrap(); + let shape = enc.shape(); + let (batch, time, hidden) = (shape[0], shape[1], shape[2]); + let mut enc_data = Vec::with_capacity(batch * time * hidden); + for b in 0..batch { + for t in 0..time { + for h in 0..hidden { + enc_data.push(enc[[b, t, h]]); + } + } + } + let enc_sum: f32 = enc_data.iter().sum(); + let enc_mean = enc_sum / enc_data.len() as f32; + let enc_var = enc_data + .iter() + .map(|value| { + let diff = *value - enc_mean; + diff * diff + }) + .sum::() + / enc_data.len() as f32; + println!( + "enc shape={:?} len={} sum={} mean={} std={} first20={:?}", + enc.shape(), + encoder_length, + enc_sum, + enc_mean, + enc_var.sqrt(), + &enc_data[..20] + ); + let input_ids = Tensor::from_slice(&device, [1, prompt_ids.len()], &prompt_ids); + let hidden_states = + model + .decoder + .forward(&input_ids, &encoder_hidden_states, encoder_length); + let last_hidden = hidden_states + .narrow(1, prompt_ids.len() - 1, 1) + .to_concrete(); + let logits = model.lm_head(&last_hidden).as_slice().await.unwrap(); + + let mut top = (0..config.vocab_size) + .map(|token_id| (token_id, logits[[0, 0, token_id]])) + .collect::>(); + top.sort_by(|a, b| b.1.total_cmp(&a.1)); + println!("top10={:?}", &top[..10]); + } + + #[tokio::test] + async fn compare_cpu_gpu_first_step() { + let root = Path::new("/tmp/cohere-transcribe-03-2026"); + if !root.exists() { + return; + } + + let gpu_device = match Device::new().await { + Ok(device) => device, + Err(_) => return, + }; + + let config: CohereConfig = + serde_json::from_slice(&std::fs::read(root.join("config.json")).unwrap()).unwrap(); + let weights = std::fs::read(root.join("model.gguf")).unwrap(); + + let cpu_device = Device::cpu(); + let mut cpu_reader = std::io::Cursor::new(weights.clone()); + let mut cpu_vb = VarBuilder::from_gguf(&mut cpu_reader).unwrap(); + let cpu_model = Cohere::load(&cpu_device, &mut cpu_vb, config.clone()).unwrap(); + + let mut gpu_reader = std::io::Cursor::new(weights); + let mut gpu_vb = VarBuilder::from_gguf(&mut gpu_reader).unwrap(); + let gpu_model = Cohere::load(&gpu_device, &mut gpu_vb, config.clone()).unwrap(); + + let wav = hound::WavReader::open(concat!( + env!("CARGO_MANIFEST_DIR"), + "/examples/samples_jfk.wav" + )) + .unwrap() + .into_samples::() + .map(|sample| sample.unwrap() as f32 / 32768.0) + .take(16_000) + .collect::>(); + + let filter_bytes = include_bytes!("../cohere_melfilters128.bytes").as_slice(); + let mut filterbank = vec![0.0f32; filter_bytes.len() / 4]; + ::read_f32_into( + filter_bytes, + &mut filterbank, + ); + for value in &mut filterbank { + *value = half::bf16::from_f32(*value).to_f32(); + } + + let (features, total_frames, valid_frames) = pcm_to_features(&config, &wav, &filterbank); + let cpu_input = Tensor::from_slice( + &cpu_device, + [1, config.preprocessor.features, total_frames], + &features, + ); + let gpu_input = Tensor::from_slice( + &gpu_device, + [1, config.preprocessor.features, total_frames], + &features, + ); + let prompt_ids = [7_u32, 4, 16, 62, 62, 5, 9, 11, 13]; + + let (cpu_pre, cpu_pre_len) = cpu_model + .encoder + .pre_encode + .forward(&cpu_input, valid_frames); + let (gpu_pre, gpu_pre_len) = gpu_model + .encoder + .pre_encode + .forward(&gpu_input, valid_frames); + assert_eq!(cpu_pre_len, gpu_pre_len); + eprintln!("compare: pre_encode slices"); + let cpu_pre = cpu_pre.as_slice().await.unwrap(); + let gpu_pre = gpu_pre.as_slice().await.unwrap(); + + let mut cpu_pre_max = 0.0f32; + let pre_shape = cpu_pre.shape(); + for i in 0..pre_shape[0] { + for j in 0..pre_shape[1] { + for k in 0..pre_shape[2] { + cpu_pre_max = cpu_pre_max.max((cpu_pre[[i, j, k]] - gpu_pre[[i, j, k]]).abs()); + } + } + } + println!("pre_encode max_abs_diff={cpu_pre_max}"); + + let (cpu_enc, cpu_enc_len) = cpu_model.encode(&cpu_input, valid_frames); + let (gpu_enc, gpu_enc_len) = gpu_model.encode(&gpu_input, valid_frames); + assert_eq!(cpu_enc_len, gpu_enc_len); + eprintln!("compare: encoder slices"); + let cpu_enc_slice = cpu_enc.clone().as_slice().await.unwrap(); + let gpu_enc_slice = gpu_enc.clone().as_slice().await.unwrap(); + + let mut cpu_enc_max = 0.0f32; + let enc_shape = cpu_enc_slice.shape(); + for i in 0..enc_shape[0] { + for j in 0..enc_shape[1] { + for k in 0..enc_shape[2] { + cpu_enc_max = cpu_enc_max + .max((cpu_enc_slice[[i, j, k]] - gpu_enc_slice[[i, j, k]]).abs()); + } + } + } + println!("encoder max_abs_diff={cpu_enc_max}"); + + eprintln!("compare: decoder cached init"); + let mut cpu_cache = TransformerDecoderCache::default(); + let mut gpu_cache = TransformerDecoderCache::default(); + let cpu_hidden = cpu_model + .decoder + .forward_cached(&prompt_ids, &cpu_enc, cpu_enc_len, &mut cpu_cache, None) + .narrow(1, prompt_ids.len() - 1, 1) + .to_concrete(); + let gpu_hidden = gpu_model + .decoder + .forward_cached(&prompt_ids, &gpu_enc, gpu_enc_len, &mut gpu_cache, None) + .narrow(1, prompt_ids.len() - 1, 1) + .to_concrete(); + eprintln!("compare: decoder hidden slices"); + let cpu_hidden_slice = cpu_hidden.clone().as_slice().await.unwrap(); + let gpu_hidden_slice = gpu_hidden.clone().as_slice().await.unwrap(); + + let mut cpu_hidden_max = 0.0f32; + let hidden_shape = cpu_hidden_slice.shape(); + for i in 0..hidden_shape[0] { + for j in 0..hidden_shape[1] { + for k in 0..hidden_shape[2] { + cpu_hidden_max = cpu_hidden_max + .max((cpu_hidden_slice[[i, j, k]] - gpu_hidden_slice[[i, j, k]]).abs()); + } + } + } + println!("decoder_hidden max_abs_diff={cpu_hidden_max}"); + + eprintln!("compare: logits slices"); + let cpu_logits = cpu_model.lm_head(&cpu_hidden).as_slice().await.unwrap(); + let gpu_logits = gpu_model.lm_head(&gpu_hidden).as_slice().await.unwrap(); + + let mut logits_max = 0.0f32; + let logits_shape = cpu_logits.shape(); + for i in 0..logits_shape[0] { + for j in 0..logits_shape[1] { + for k in 0..logits_shape[2] { + logits_max = + logits_max.max((cpu_logits[[i, j, k]] - gpu_logits[[i, j, k]]).abs()); + } + } + } + println!("logits max_abs_diff={logits_max}"); + + let mut cpu_top = (0..config.vocab_size) + .map(|token_id| (token_id, cpu_logits[[0, 0, token_id]])) + .collect::>(); + cpu_top.sort_by(|a, b| b.1.total_cmp(&a.1)); + let mut gpu_top = (0..config.vocab_size) + .map(|token_id| (token_id, gpu_logits[[0, 0, token_id]])) + .collect::>(); + gpu_top.sort_by(|a, b| b.1.total_cmp(&a.1)); + println!("cpu_top10={:?}", &cpu_top[..10]); + println!("gpu_top10={:?}", &gpu_top[..10]); + } +} + +impl Cohere { + pub fn load(device: &Device, vb: &mut VarBuilder, config: CohereConfig) -> Result { + let encoder = ConformerEncoder::load(device, &mut vb.pp("encoder"), &config.encoder)?; + let decoder = TransformerDecoder::load( + device, + &mut vb.pp("transf_decoder"), + &config.transf_decoder.config_dict, + )?; + let encoder_decoder_proj = + if config.encoder.d_model != config.transf_decoder.config_dict.hidden_size { + Some(Linear::load(device, &mut vb.pp("encoder_decoder_proj"))?) + } else { + None + }; + let lm_bias = vb + .pp("log_softmax") + .pp("mlp") + .pp("layer0") + .get("bias", device)? + .dequantize(); + Ok(Self { + config, + encoder, + decoder, + encoder_decoder_proj, + lm_bias, + }) + } + + fn encode(&self, input_features: &Tensor<3, f32>, length: usize) -> (Tensor<3, f32>, usize) { + let (encoder_hidden_states, encoder_length) = self.encoder.forward(input_features, length); + let encoder_hidden_states = if let Some(proj) = &self.encoder_decoder_proj { + proj.forward(&encoder_hidden_states) + } else { + encoder_hidden_states + }; + (encoder_hidden_states, encoder_length) + } + + fn lm_head(&self, hidden_states: &Tensor<3, f32>) -> Tensor<3, f32> { + hidden_states + .q_mat_mul( + self.decoder + .embedding + .token_embedding + .embeddings_quantized(), + ) + .add_(&self.lm_bias) + } + + pub async fn generate_greedy( + &self, + input_features: &Tensor<3, f32>, + length: usize, + prompt_ids: &[u32], + eos_token_id: u32, + max_new_tokens: usize, + ) -> Result> { + let profile = profile_enabled(); + let encode_start = Instant::now(); + let (encoder_hidden_states, encoder_length) = self.encode(input_features, length); + if profile { + eprintln!( + "cohere encode: {:.3}s", + encode_start.elapsed().as_secs_f32() + ); + } + let mut cache = TransformerDecoderCache::default(); + let mut tokens = prompt_ids.to_vec(); + let first_decode_start = Instant::now(); + let mut hidden_states = self.decoder.forward_cached( + prompt_ids, + &encoder_hidden_states, + encoder_length, + &mut cache, + None, + ); + if profile { + eprintln!( + "cohere decode init: {:.3}s", + first_decode_start.elapsed().as_secs_f32() + ); + } + + for step in 0..max_new_tokens { + let step_start = profile.then(Instant::now); + let last_hidden = hidden_states + .narrow(1, hidden_states.shape()[1] - 1, 1) + .to_concrete(); + let logits = self.lm_head(&last_hidden); + let logits = logits.as_slice().await?; + + let mut best_token = 0u32; + let mut best_value = f32::NEG_INFINITY; + for token_id in 0..self.config.vocab_size { + let value = logits[[0, 0, token_id]]; + if value > best_value { + best_value = value; + best_token = token_id as u32; + } + } + + if best_token == eos_token_id { + break; + } + tokens.push(best_token); + hidden_states = self.decoder.forward_cached( + &[best_token], + &encoder_hidden_states, + encoder_length, + &mut cache, + None, + ); + if let Some(step_start) = step_start { + eprintln!( + "cohere decode step {step}: {:.3}s token={best_token}", + step_start.elapsed().as_secs_f32() + ); + } + } + + Ok(tokens[prompt_ids.len()..].to_vec()) + } + + pub async fn generate_greedy_with_attention( + &self, + input_features: &Tensor<3, f32>, + length: usize, + prompt_ids: &[u32], + eos_token_id: u32, + max_new_tokens: usize, + ) -> Result<(Vec, Vec>, usize)> { + let profile = profile_enabled(); + let encode_start = Instant::now(); + let (encoder_hidden_states, encoder_length) = self.encode(input_features, length); + if profile { + eprintln!( + "cohere encode: {:.3}s", + encode_start.elapsed().as_secs_f32() + ); + } + let mut tokens = prompt_ids.to_vec(); + let mut cache = TransformerDecoderCache::default(); + let mut attention_output: Vec> = (0..self.decoder.layers.len()) + .map(|_| TensorCache::new(2, max_new_tokens)) + .collect(); + let first_decode_start = Instant::now(); + let mut hidden_states = self.decoder.forward_cached( + prompt_ids, + &encoder_hidden_states, + encoder_length, + &mut cache, + Some(&mut attention_output), + ); + if profile { + eprintln!( + "cohere decode init: {:.3}s", + first_decode_start.elapsed().as_secs_f32() + ); + } + + for step in 0..max_new_tokens { + let step_start = profile.then(Instant::now); + let last_hidden = hidden_states + .narrow(1, hidden_states.shape()[1] - 1, 1) + .to_concrete(); + let logits = self.lm_head(&last_hidden); + let logits = logits.as_slice().await?; + + let mut best_token = 0u32; + let mut best_value = f32::NEG_INFINITY; + for token_id in 0..self.config.vocab_size { + let value = logits[[0, 0, token_id]]; + if value > best_value { + best_value = value; + best_token = token_id as u32; + } + } + + if best_token == eos_token_id { + break; + } + tokens.push(best_token); + hidden_states = self.decoder.forward_cached( + &[best_token], + &encoder_hidden_states, + encoder_length, + &mut cache, + Some(&mut attention_output), + ); + if let Some(step_start) = step_start { + eprintln!( + "cohere decode step {step}: {:.3}s token={best_token}", + step_start.elapsed().as_secs_f32() + ); + } + } + + let collected_attentions = attention_output + .into_iter() + .filter_map(|attentions| { + if attentions.current_seq_len() == 0 { + None + } else { + attentions.current_data().cloned() + } + }) + .collect(); + + Ok(( + tokens[prompt_ids.len()..].to_vec(), + collected_attentions, + encoder_length, + )) + } +} diff --git a/models/rwhisper/src/quantized/mod.rs b/models/rwhisper/src/quantized/mod.rs index c1078329f..abddc5701 100644 --- a/models/rwhisper/src/quantized/mod.rs +++ b/models/rwhisper/src/quantized/mod.rs @@ -5,12 +5,14 @@ use std::{num::NonZeroUsize, sync::Arc}; use fusor::{ cache::{AttentionMask, KvCache, MaskCache, TensorCache}, layers::{Conv1d, Conv1dConfig, Embedding, LayerNorm, Linear}, - Device, Error, Result, Tensor, VarBuilder, + Device, Error, MaskKind, Result, Tensor, VarBuilder, }; use timestamps::extract_timestamps; -use crate::config::Config; +use crate::config::{Config, HOP_LENGTH, N_FRAMES, SAMPLE_RATE}; +pub(crate) mod cohere; +pub(crate) mod moonshine; pub(crate) mod timestamps; fn conv1d( @@ -155,23 +157,29 @@ impl MultiHeadAttention { let [n_batch, n_ctx_q, n_state] = q.shape(); let [_, n_ctx_kv, _] = k.shape(); let head_dim = n_state / self.n_head; - let scale = crate::WhisperDType::from((head_dim as f32).powf(-0.25)); + let scale = (head_dim as f32).powf(-0.5); // Reshape Q: [n_batch, n_ctx_q, n_state] -> [n_batch, n_head, n_ctx_q, head_dim] let q_target_dims = [n_batch, n_ctx_q, self.n_head, head_dim]; let q_reshaped = q.reshape(q_target_dims); - let q_transposed = q_reshaped.transpose(1, 2); - let q = q_transposed.mul_scalar(scale); + let q = q_reshaped.transpose(1, 2).to_concrete(); // Reshape K/V: [n_batch, n_ctx_kv, n_state] -> [n_batch, n_head, n_ctx_kv, head_dim] // For self-attention n_ctx_kv == n_ctx_q, for cross-attention they differ let kv_target_dims = [n_batch, n_ctx_kv, self.n_head, head_dim]; let k_reshaped = k.reshape(kv_target_dims); - let k_transposed = k_reshaped.transpose(1, 2); - let k_transposed2 = k_transposed.transpose(2, 3); - let k = k_transposed2.mul_scalar(scale); - let v_reshaped = v.reshape(kv_target_dims); - let v = v_reshaped.transpose(1, 2); + let k = k_reshaped.transpose(1, 2).to_concrete(); + let v = v.reshape(kv_target_dims).transpose(1, 2).to_concrete(); + + // When we do not need the raw attention scores for timestamp extraction, route through + // the fused attention kernel to avoid materializing the full QK matrix on GPU. + if attention_output.is_none() { + let mask = mask.map(|mask| (mask.mask(), MaskKind::QKMask)); + let wv = q.flash_attention(&k, &v, scale, mask); + return Ok(wv.transpose(1, 2).flatten_last_n::<1, _>()); + } + + let k = k.transpose(2, 3).to_concrete(); let mut qk = { let _enter = self.matmul_span.enter(); @@ -482,12 +490,13 @@ impl TextDecoder { for (i, block) in self.blocks.iter_mut().enumerate() { if cache.blocks.len() <= i { + let feature_attn_cache = block + .cross_attn + .as_ref() + .and_then(|(attn, _)| attn.forward_kv(&audio_features_batched, None).ok()); cache.blocks.push(ResidualAttentionBlockCache { attn: MultiHeadAttentionCache::new(self.max_target_positions), - feature_attn_cache: block - .cross_attn - .as_ref() - .and_then(|(attn, _)| attn.forward_kv(&audio_features_batched, None).ok()), + feature_attn_cache, }); } let block_cache = &mut cache.blocks[i]; @@ -557,10 +566,11 @@ impl Whisper { } extract_timestamps( - attention_heads, + Some(attention_heads), &attention_output_tensor, filter_width, - n_frames, + n_frames.min(N_FRAMES) / 2, + crate::WhisperDType::from((HOP_LENGTH * 2) as f32 / SAMPLE_RATE as f32), mask, ) .await diff --git a/models/rwhisper/src/quantized/moonshine.rs b/models/rwhisper/src/quantized/moonshine.rs new file mode 100644 index 000000000..147a4c519 --- /dev/null +++ b/models/rwhisper/src/quantized/moonshine.rs @@ -0,0 +1,1600 @@ +use std::sync::Arc; + +use fusor::{ + cache::{AttentionMask, KvCache, MaskCache, TensorCache}, + layers::{Conv1d, Conv1dConfig, Embedding, LayerNorm, Linear}, + Device, Error, MaskKind, Result, RopeCache, Tensor, VarBuilder, +}; + +use crate::moonshine_config::{MoonshineStreamingConfig, MoonshineStreamingEncoderConfig}; + +fn tensor1d(device: &Device, vb: &mut VarBuilder, name: &str) -> Result> { + let q = vb.get(name, device)?; + let shape = q.shape(); + Ok(if shape.len() == 1 { + q.dequantize() + } else { + let value: Tensor<2, f32> = q.dequantize(); + if value.shape()[0] == 1 { + value.squeeze(0).to_concrete() + } else { + value.squeeze(1).to_concrete() + } + }) +} + +fn conv1d(config: Conv1dConfig, device: &Device, vb: &mut VarBuilder) -> Result> { + let weight_q = vb.get("weight", device)?; + let weight_shape = weight_q.shape(); + let weight: Tensor<3, f32> = if weight_shape.len() == 3 { + weight_q.dequantize() + } else { + let out = weight_shape[0]; + let kernel = 5usize; + let in_channels = weight_shape[1] / kernel; + let weight_2d: Tensor<2, f32> = weight_q.dequantize(); + weight_2d.reshape([out, in_channels, kernel]).to_concrete() + }; + let bias = vb.get("bias", device).ok().map(|bias_q| { + let bias_shape = bias_q.shape(); + if bias_shape.len() == 1 { + bias_q.dequantize() + } else { + let bias_2d: Tensor<2, f32> = bias_q.dequantize(); + if bias_2d.shape()[0] == 1 { + bias_2d.squeeze(0).to_concrete() + } else { + bias_2d.squeeze(1).to_concrete() + } + } + }); + Ok(Conv1d::new(weight.to_concrete(), bias, config)) +} + +fn ones_1d(device: &Device, len: usize) -> Tensor<1, f32> { + Tensor::splat(device, 1.0, [len]) +} + +fn causal_pad_1d(x: &Tensor<3, f32>, channels: usize, pad: usize) -> Tensor<3, f32> { + let [batch, _, _] = x.shape(); + let zeros = Tensor::zeros(&x.device(), [batch, channels, pad]); + Tensor::cat([zeros, x.clone()], 2) +} + +fn flatten_frames(samples: &[f32], frame_len: usize, log_k: f32) -> Vec { + let usable_len = samples.len() / frame_len * frame_len; + let k = log_k.exp(); + let mut frames = Vec::with_capacity(usable_len); + for frame in samples[..usable_len].chunks_exact(frame_len) { + let mean = frame.iter().sum::() / frame_len as f32; + let mut centered = vec![0.0f32; frame_len]; + let mut rms = 0.0f32; + for (i, sample) in frame.iter().enumerate() { + let value = *sample - mean; + centered[i] = value; + rms += value * value; + } + rms = (rms / frame_len as f32 + 1e-6).sqrt(); + for value in centered { + let scaled = k * value / rms; + frames.push((scaled + (scaled * scaled + 1.0).sqrt()).ln()); + } + } + frames +} + +fn sliding_window_mask( + device: &Device, + seq_len: usize, + left: usize, + right: usize, +) -> Tensor<4, f32> { + let mut data = vec![f32::NEG_INFINITY; seq_len * seq_len]; + for q in 0..seq_len { + for k in 0..seq_len { + let dist = q as isize - k as isize; + let allowed = + (dist >= 0 && dist < left as isize) || (dist < 0 && -dist < right as isize); + if allowed { + data[q * seq_len + k] = 0.0; + } + } + } + Tensor::from_slice(device, [1, 1, seq_len, seq_len], &data) +} + +fn streaming_sliding_window_mask( + device: &Device, + query_start: usize, + query_len: usize, + key_len: usize, + left: usize, + right: usize, +) -> Tensor<4, f32> { + let mut data = vec![f32::NEG_INFINITY; query_len * key_len]; + for q in 0..query_len { + let query_index = query_start + q; + for k in 0..key_len { + let dist = query_index as isize - k as isize; + let allowed = + (dist >= 0 && dist < left as isize) || (dist < 0 && -dist < right as isize); + if allowed { + data[q * key_len + k] = 0.0; + } + } + } + Tensor::from_slice(device, [1, 1, query_len, key_len], &data) +} + +fn causal_mask(device: &Device, seq_len: usize) -> Tensor<4, f32> { + let mut data = vec![0.0f32; seq_len * seq_len]; + for q in 0..seq_len { + for k in (q + 1)..seq_len { + data[q * seq_len + k] = f32::NEG_INFINITY; + } + } + Tensor::from_slice(device, [1, 1, seq_len, seq_len], &data) +} + +fn stable_gelu_3d(x: &Tensor<3, f32, B>) -> Tensor<3, f32> +where + B: fusor::TensorBacking<3, Elem = f32>, +{ + let x = x.to_concrete(); + let x_squared = x.mul_(&x).to_concrete(); + let inner_factor = x_squared.mul_scalar(0.044715).add_scalar(1.0).to_concrete(); + let inner = x.mul_(&inner_factor).to_concrete(); + let tanh_input = inner + .mul_scalar((2.0 / std::f32::consts::PI).sqrt()) + .to_concrete(); + let tanh_input = tanh_input.clamp(-10.0, 10.0).to_concrete(); + let tanh_result = tanh_input.tanh().clamp(-1.0, 1.0).to_concrete(); + let one_plus_tanh = tanh_result.add_scalar(1.0).to_concrete(); + x.mul_(&one_plus_tanh).mul_scalar(0.5).to_concrete() +} + +struct UnitOffsetLayerNorm { + norm: LayerNorm<1, f32>, + gamma: Tensor<1, f32>, + unit_offset: f32, +} + +impl UnitOffsetLayerNorm { + fn load(device: &Device, vb: &mut VarBuilder, eps: f32) -> Result { + let gamma = tensor1d(device, vb, "gamma")?.to_concrete(); + let weight = ones_1d(device, gamma.shape()[0]).to_concrete(); + let norm = LayerNorm::new(weight, None, eps); + Ok(Self { + norm, + gamma, + unit_offset: 1.0, + }) + } + + fn forward(&self, x: &Tensor<3, f32, B>) -> Tensor<3, f32> + where + B: fusor::TensorBacking<3, Elem = f32>, + { + let normed = self.norm.forward(x); + let gamma = self.gamma.add_scalar(self.unit_offset); + let weight = gamma.broadcast_as(normed.shape()); + normed.mul_(&weight).to_concrete() + } +} + +#[derive(Clone)] +struct MoonshineAttentionCache { + kv_cache: KvCache, +} + +impl MoonshineAttentionCache { + fn new(max_seq_len: usize) -> Self { + Self { + kv_cache: KvCache::new(1, max_seq_len), + } + } +} + +struct MoonshineAttention { + q_proj: Linear, + k_proj: Linear, + v_proj: Linear, + o_proj: Linear, + num_heads: usize, + head_dim: usize, + scale: f32, + rope_cache: Option>, + rotary_dim: usize, +} + +impl MoonshineAttention { + fn load( + device: &Device, + vb: &mut VarBuilder, + num_heads: usize, + head_dim: usize, + rope_cache: Option>, + rotary_dim: usize, + ) -> Result { + Ok(Self { + q_proj: Linear::load(device, &mut vb.pp("q_proj"))?, + k_proj: Linear::load(device, &mut vb.pp("k_proj"))?, + v_proj: Linear::load(device, &mut vb.pp("v_proj"))?, + o_proj: Linear::load(device, &mut vb.pp("o_proj"))?, + num_heads, + head_dim, + scale: head_dim as f32, + rope_cache, + rotary_dim, + }) + } + + fn reshape_heads(&self, x: &Tensor<3, f32>) -> Tensor<4, f32> { + let [batch, seq_len, _] = x.shape(); + x.reshape([batch, seq_len, self.num_heads, self.head_dim]) + .transpose(1, 2) + .to_concrete() + } + + fn flatten_heads(&self, x: Tensor<4, f32>) -> Tensor<3, f32> { + let [batch, _, seq_len, _] = x.shape(); + x.transpose(1, 2) + .to_concrete() + .reshape([batch, seq_len, self.num_heads * self.head_dim]) + .to_concrete() + } + + fn apply_rope( + &self, + q: Tensor<4, f32>, + k: Tensor<4, f32>, + start_pos: usize, + ) -> (Tensor<4, f32>, Tensor<4, f32>) { + let Some(cache) = &self.rope_cache else { + return (q, k); + }; + if self.rotary_dim == 0 { + return (q, k); + } + if self.rotary_dim < self.head_dim { + return cache.forward_interleaved_partial(&q, &k, start_pos, self.rotary_dim); + } + let q_rot = q.narrow(3, 0, self.rotary_dim).to_concrete(); + let k_rot = k.narrow(3, 0, self.rotary_dim).to_concrete(); + let (q_rot, k_rot) = cache.forward_interleaved(&q_rot, &k_rot, start_pos); + if self.rotary_dim == self.head_dim { + return (q_rot, k_rot); + } + let q_pass = q + .narrow(3, self.rotary_dim, self.head_dim - self.rotary_dim) + .to_concrete(); + let k_pass = k + .narrow(3, self.rotary_dim, self.head_dim - self.rotary_dim) + .to_concrete(); + ( + Tensor::cat([q_rot, q_pass], 3), + Tensor::cat([k_rot, k_pass], 3), + ) + } + + fn apply_rope_3d( + &self, + q: Tensor<3, f32>, + k: Tensor<3, f32>, + start_pos: usize, + ) -> (Tensor<3, f32>, Tensor<3, f32>) { + if self.rope_cache.is_none() || self.rotary_dim == 0 { + return (q, k); + } + let q = self.reshape_heads(&q); + let k = self.reshape_heads(&k); + let (q, k) = self.apply_rope(q, k, start_pos); + (self.flatten_heads(q), self.flatten_heads(k)) + } + + fn project_kv(&self, hidden_states: &Tensor<3, f32>) -> (Tensor<3, f32>, Tensor<3, f32>) { + ( + self.k_proj.forward(hidden_states), + self.v_proj.forward(hidden_states), + ) + } + + fn append_kv( + &self, + key_states: Tensor<3, f32>, + value_states: Tensor<3, f32>, + cache: Option<&mut MoonshineAttentionCache>, + ) -> (Tensor<3, f32>, Tensor<3, f32>) { + match cache { + None => (key_states, value_states), + Some(cache) => { + let device = key_states.device(); + let key_states_4d = key_states.unsqueeze(2).to_concrete(); + let value_states_4d = value_states.unsqueeze(2).to_concrete(); + let (k, v) = cache + .kv_cache + .append(&device, &key_states_4d, &value_states_4d); + (k.squeeze(2).to_concrete(), v.squeeze(2).to_concrete()) + } + } + } + + fn qkv_attention_dense( + &self, + q: &Tensor<3, f32>, + k: &Tensor<3, f32>, + v: &Tensor<3, f32>, + attention_mask: Option<&Tensor<4, f32>>, + attention_output: Option<&mut Vec>>, + ) -> Result> { + let [batch, q_len, _] = q.shape(); + let q = self.reshape_heads(q); + let k = self.reshape_heads(k); + let v = self.reshape_heads(v); + let k_t = k.transpose(2, 3).to_concrete(); + let mut scores = q.mat_mul(&k_t).mul_scalar(self.scale.powf(-0.5)); + if let Some(mask) = attention_mask { + scores = scores.add_(mask).to_concrete(); + } + if let Some(outputs) = attention_output { + outputs.push(scores.clone()); + } + let scores = scores.to_concrete(); + let weights = scores.softmax_last_dim_fused().to_concrete(); + let context = weights + .mat_mul(&v) + .transpose(1, 2) + .to_concrete() + .reshape([batch, q_len, self.num_heads * self.head_dim]) + .to_concrete(); + Ok(self.o_proj.forward(&context)) + } + + fn qkv_attention_masked( + &self, + q: &Tensor<3, f32>, + k: &Tensor<3, f32>, + v: &Tensor<3, f32>, + attention_mask: Option<&AttentionMask>, + attention_output: Option<&mut TensorCache<4, f32>>, + ) -> Result> { + let [batch, q_len, _] = q.shape(); + let q = self.reshape_heads(q); + let k = self.reshape_heads(k); + let v = self.reshape_heads(v); + + if attention_output.is_none() { + let mask = attention_mask.map(|mask| (mask.mask(), MaskKind::QKMask)); + let context = q + .flash_attention(&k, &v, self.scale.powf(-0.5), mask) + .transpose(1, 2) + .reshape([batch, q_len, self.num_heads * self.head_dim]) + .to_concrete(); + return Ok(self.o_proj.forward(&context)); + } + + let k_t = k.transpose(2, 3).to_concrete(); + let mut scores = q.mat_mul(&k_t).mul_scalar(self.scale.powf(-0.5)); + if q.is_gpu() { + if let Some(mask) = attention_mask { + let mask = mask.mask().clone().unsqueeze(0).unsqueeze(0).to_concrete(); + scores = scores.add_(&mask).to_concrete(); + } + } else if let Some(mask) = attention_mask { + mask.forward(&mut scores); + } + if let Some(output) = attention_output { + let last_query = scores.narrow(2, q_len.saturating_sub(1), 1).to_concrete(); + output.append(&q.device(), &last_query); + } + let scores = scores.to_concrete(); + let weights = scores.softmax_last_dim_fused().to_concrete(); + let context = weights + .mat_mul(&v) + .transpose(1, 2) + .to_concrete() + .reshape([batch, q_len, self.num_heads * self.head_dim]) + .to_concrete(); + Ok(self.o_proj.forward(&context)) + } + + fn forward( + &self, + hidden_states: &Tensor<3, f32>, + key_value_states: &Tensor<3, f32>, + attention_mask: Option<&Tensor<4, f32>>, + attention_output: Option<&mut Vec>>, + ) -> Result> { + let query_states = self.q_proj.forward(hidden_states).to_concrete(); + let (key_states, value_states) = self.project_kv(key_value_states); + let key_states = key_states.to_concrete(); + let value_states = value_states.to_concrete(); + let (query_states, key_states) = self.apply_rope_3d(query_states, key_states, 0); + let query_states = query_states.to_concrete(); + let key_states = key_states.to_concrete(); + self.qkv_attention_dense( + &query_states, + &key_states, + &value_states, + attention_mask, + attention_output, + ) + } + + fn forward_cached( + &self, + hidden_states: &Tensor<3, f32>, + kv: (Tensor<3, f32>, Tensor<3, f32>), + attention_mask: Option<&AttentionMask>, + attention_output: Option<&mut TensorCache<4, f32>>, + ) -> Result> { + let query_states = self.q_proj.forward(hidden_states).to_concrete(); + let (key_states, value_states) = kv; + let key_states = key_states.to_concrete(); + let value_states = value_states.to_concrete(); + self.qkv_attention_masked( + &query_states, + &key_states, + &value_states, + attention_mask, + attention_output, + ) + } +} + +struct MoonshineEncoderMlp { + fc1: Linear, + fc2: Linear, +} + +impl MoonshineEncoderMlp { + fn load(device: &Device, vb: &mut VarBuilder) -> Result { + Ok(Self { + fc1: Linear::load(device, &mut vb.pp("fc1"))?, + fc2: Linear::load(device, &mut vb.pp("fc2"))?, + }) + } + + fn forward(&self, x: &Tensor<3, f32>) -> Tensor<3, f32> { + let hidden = self.fc1.forward(x).to_concrete(); + let hidden = stable_gelu_3d(&hidden); + self.fc2.forward(&hidden) + } +} + +struct MoonshineDecoderMlp { + fc1: Linear, + fc2: Linear, + intermediate_size: usize, +} + +impl MoonshineDecoderMlp { + fn load(device: &Device, vb: &mut VarBuilder, intermediate_size: usize) -> Result { + Ok(Self { + fc1: Linear::load(device, &mut vb.pp("fc1"))?, + fc2: Linear::load(device, &mut vb.pp("fc2"))?, + intermediate_size, + }) + } + + fn forward(&self, x: &Tensor<3, f32>) -> Tensor<3, f32> { + let hidden = self.fc1.forward(x).to_concrete(); + self.fc2 + .forward(&hidden.swiglu_split(self.intermediate_size)) + } +} + +struct MoonshineEncoderLayer { + self_attn: MoonshineAttention, + input_layernorm: UnitOffsetLayerNorm, + post_attention_layernorm: UnitOffsetLayerNorm, + mlp: MoonshineEncoderMlp, +} + +#[derive(Clone, Default)] +pub struct MoonshineEncoderLayerStreamState { + left_context_inputs: Option>, + pending_inputs: Option>, +} + +pub struct MoonshineEncoderStreamState { + sample_remainder: Vec, + linear_tail: Option>, + total_linear_frames: usize, + conv1_tail: Option>, + total_conv1_frames: usize, + layer_states: Vec, + total_seen_frames: usize, + total_finalized_frames: usize, +} + +pub struct MoonshineEncoderStreamAppend { + pub hidden_states: Option>, + pub total_seen_frames: usize, + pub total_finalized_frames: usize, + pub usable_input_samples: usize, +} + +impl MoonshineEncoderLayer { + fn load( + config: &MoonshineStreamingEncoderConfig, + device: &Device, + vb: &mut VarBuilder, + ) -> Result { + Ok(Self { + self_attn: MoonshineAttention::load( + device, + &mut vb.pp("self_attn"), + config.num_attention_heads, + config.head_dim(), + None, + 0, + )?, + input_layernorm: UnitOffsetLayerNorm::load( + device, + &mut vb.pp("input_layernorm"), + 1e-5, + )?, + post_attention_layernorm: UnitOffsetLayerNorm::load( + device, + &mut vb.pp("post_attention_layernorm"), + 1e-5, + )?, + mlp: MoonshineEncoderMlp::load(device, &mut vb.pp("mlp"))?, + }) + } + + fn forward( + &self, + hidden_states: &Tensor<3, f32>, + attention_mask: &Tensor<4, f32>, + ) -> Result> { + let attn_in = self.input_layernorm.forward(hidden_states); + let attn = self + .self_attn + .forward(&attn_in, &attn_in, Some(attention_mask), None)?; + let hidden_states = (hidden_states + &attn).to_concrete(); + let mlp_in = self.post_attention_layernorm.forward(&hidden_states); + let mlp = self.mlp.forward(&mlp_in); + Ok((hidden_states + mlp).to_concrete()) + } + + fn forward_stream( + &self, + device: &Device, + new_hidden_states: Option>, + state: &mut MoonshineEncoderLayerStreamState, + left: usize, + right: usize, + flush: bool, + ) -> Result>> { + let left_context_len = state + .left_context_inputs + .as_ref() + .map(|tensor| tensor.shape()[1]) + .unwrap_or(0); + let pending_len = state + .pending_inputs + .as_ref() + .map(|tensor| tensor.shape()[1]) + .unwrap_or(0); + let new_len = new_hidden_states + .as_ref() + .map(|tensor| tensor.shape()[1]) + .unwrap_or(0); + let working_len = pending_len + new_len; + if left_context_len + working_len == 0 { + return Ok(None); + } + + let mut parts = Vec::with_capacity(3); + if let Some(left_context_inputs) = &state.left_context_inputs { + parts.push(left_context_inputs.clone()); + } + if let Some(pending_inputs) = &state.pending_inputs { + parts.push(pending_inputs.clone()); + } + if let Some(new_hidden_states) = new_hidden_states { + parts.push(new_hidden_states); + } + let combined_inputs = Tensor::cat(parts, 1).to_concrete(); + + let emit_len = if flush { + working_len + } else { + working_len.saturating_sub(right) + }; + let pending_start = left_context_len + emit_len; + let pending_len = working_len.saturating_sub(emit_len); + state.pending_inputs = (pending_len > 0).then(|| { + combined_inputs + .narrow(1, pending_start, pending_len) + .to_concrete() + }); + + let finalized_input_end = left_context_len + emit_len; + let keep_left = left.min(finalized_input_end); + state.left_context_inputs = (keep_left > 0).then(|| { + combined_inputs + .narrow(1, finalized_input_end - keep_left, keep_left) + .to_concrete() + }); + + if emit_len == 0 { + return Ok(None); + } + + let key_value_inputs = self.input_layernorm.forward(&combined_inputs); + let query_inputs = key_value_inputs + .narrow(1, left_context_len, emit_len) + .to_concrete(); + let residual_inputs = combined_inputs + .narrow(1, left_context_len, emit_len) + .to_concrete(); + let attention_mask = streaming_sliding_window_mask( + device, + left_context_len, + emit_len, + combined_inputs.shape()[1], + left, + right, + ); + let attn = self.self_attn.forward( + &query_inputs, + &key_value_inputs, + Some(&attention_mask), + None, + )?; + let hidden_states = (residual_inputs + &attn).to_concrete(); + let mlp_in = self.post_attention_layernorm.forward(&hidden_states); + let mlp = self.mlp.forward(&mlp_in); + Ok(Some((hidden_states + mlp).to_concrete())) + } +} + +struct MoonshineDecoderLayer { + self_attn: MoonshineAttention, + encoder_attn: MoonshineAttention, + input_layernorm: LayerNorm<1, f32>, + post_attention_layernorm: LayerNorm<1, f32>, + final_layernorm: LayerNorm<1, f32>, + mlp: MoonshineDecoderMlp, +} + +#[derive(Clone)] +struct MoonshineDecoderLayerCache { + self_attn: MoonshineAttentionCache, + cross_attn_kv: (Tensor<3, f32>, Tensor<3, f32>), +} + +impl MoonshineDecoderLayer { + fn load( + config: &MoonshineStreamingConfig, + device: &Device, + vb: &mut VarBuilder, + rope_cache: Arc, + ) -> Result { + Ok(Self { + self_attn: MoonshineAttention::load( + device, + &mut vb.pp("self_attn"), + config.num_attention_heads, + config.decoder_head_dim(), + Some(rope_cache), + config.decoder_rotary_dim(), + )?, + encoder_attn: MoonshineAttention::load( + device, + &mut vb.pp("encoder_attn"), + config.num_attention_heads, + config.decoder_head_dim(), + None, + 0, + )?, + input_layernorm: LayerNorm::load(device, &mut vb.pp("input_layernorm"), 1e-5)?, + post_attention_layernorm: LayerNorm::load( + device, + &mut vb.pp("post_attention_layernorm"), + 1e-5, + )?, + final_layernorm: LayerNorm::load(device, &mut vb.pp("final_layernorm"), 1e-5)?, + mlp: MoonshineDecoderMlp::load(device, &mut vb.pp("mlp"), config.intermediate_size)?, + }) + } + + fn forward( + &mut self, + hidden_states: &Tensor<3, f32>, + encoder_hidden_states: &Tensor<3, f32>, + causal_mask: &Tensor<4, f32>, + attention_output: Option<&mut Vec>>, + ) -> Result> { + let self_attn_in = self.input_layernorm.forward_fused(hidden_states); + let self_attn = + self.self_attn + .forward(&self_attn_in, &self_attn_in, Some(causal_mask), None)?; + let hidden_states = (hidden_states + &self_attn).to_concrete(); + + let cross_attn_in = self.post_attention_layernorm.forward_fused(&hidden_states); + let cross_attn = self.encoder_attn.forward( + &cross_attn_in, + encoder_hidden_states, + None, + attention_output, + )?; + let hidden_states = (hidden_states + &cross_attn).to_concrete(); + + let mlp_in = self.final_layernorm.forward_fused(&hidden_states); + let mlp = self.mlp.forward(&mlp_in); + Ok((hidden_states + mlp).to_concrete()) + } + + fn forward_cached( + &mut self, + hidden_states: &Tensor<3, f32>, + self_attention_mask: &AttentionMask, + index_pos: usize, + cache: &mut MoonshineDecoderLayerCache, + attention_output: Option<&mut TensorCache<4, f32>>, + ) -> Result> { + let self_attn_in = self.input_layernorm.forward_fused(hidden_states); + let query_states = self.self_attn.q_proj.forward(&self_attn_in); + let (key_states, value_states) = self.self_attn.project_kv(&self_attn_in); + let (query_states, key_states) = + self.self_attn + .apply_rope_3d(query_states, key_states, index_pos); + let kv = self + .self_attn + .append_kv(key_states, value_states, Some(&mut cache.self_attn)); + let self_attn = self.self_attn.qkv_attention_masked( + &query_states, + &kv.0, + &kv.1, + Some(self_attention_mask), + None, + )?; + let hidden_states = (hidden_states + &self_attn).to_concrete(); + + let cross_attn_in = self.post_attention_layernorm.forward_fused(&hidden_states); + let cross_attn = self.encoder_attn.forward_cached( + &cross_attn_in, + cache.cross_attn_kv.clone(), + None, + attention_output, + )?; + let hidden_states = (hidden_states + &cross_attn).to_concrete(); + + let mlp_in = self.final_layernorm.forward_fused(&hidden_states); + let mlp = self.mlp.forward(&mlp_in); + Ok((hidden_states + mlp).to_concrete()) + } +} + +pub struct MoonshineEncoder { + linear: Linear, + conv1: Conv1d, + conv2: Conv1d, + compression_log_k: f32, + layers: Vec, + final_norm: UnitOffsetLayerNorm, + config: MoonshineStreamingEncoderConfig, +} + +impl MoonshineEncoderStreamState { + fn new(layer_count: usize) -> Self { + Self { + sample_remainder: Vec::new(), + linear_tail: None, + total_linear_frames: 0, + conv1_tail: None, + total_conv1_frames: 0, + layer_states: (0..layer_count) + .map(|_| MoonshineEncoderLayerStreamState::default()) + .collect(), + total_seen_frames: 0, + total_finalized_frames: 0, + } + } +} + +impl MoonshineEncoder { + fn load( + config: &MoonshineStreamingEncoderConfig, + device: &Device, + vb: &mut VarBuilder, + ) -> Result { + let conv_config = Conv1dConfig { + padding: 0, + stride: 2, + groups: 1, + dilation: 1, + }; + let layers = (0..config.num_hidden_layers) + .map(|idx| { + MoonshineEncoderLayer::load(config, device, &mut vb.pp(format!("layers.{idx}"))) + }) + .collect::>>()?; + let compression_log_k = { + let log_k_q = vb.pp("embedder.comp").get("log_k", device)?; + let log_k: Tensor<0, f32> = log_k_q.dequantize(); + pollster::block_on(log_k.to_scalar())? + }; + Ok(Self { + linear: Linear::load(device, &mut vb.pp("embedder.linear"))?, + conv1: conv1d(conv_config, device, &mut vb.pp("embedder.conv1"))?, + conv2: conv1d(conv_config, device, &mut vb.pp("embedder.conv2"))?, + compression_log_k, + layers, + final_norm: UnitOffsetLayerNorm::load(device, &mut vb.pp("final_norm"), 1e-5)?, + config: config.clone(), + }) + } + + fn preprocess(&self, samples: &[f32]) -> Vec { + flatten_frames(samples, self.config.frame_len(), self.compression_log_k) + } + + pub fn new_stream_state(&self) -> MoonshineEncoderStreamState { + MoonshineEncoderStreamState::new(self.layers.len()) + } + + fn ceil_div_2(value: usize) -> usize { + (value + 1) / 2 + } + + fn take_last_len( + tensor: &Tensor<3, f32>, + dim: usize, + max_len: usize, + ) -> Option> { + let len = tensor.shape()[dim].min(max_len); + (len > 0).then(|| { + tensor + .narrow(dim, tensor.shape()[dim] - len, len) + .to_concrete() + }) + } + + fn append_tail( + previous_tail: Option<&Tensor<3, f32>>, + new_values: &Tensor<3, f32>, + dim: usize, + max_len: usize, + ) -> Option> { + let combined = match previous_tail { + Some(previous_tail) => Tensor::cat([previous_tail.clone(), new_values.clone()], dim), + None => new_values.clone(), + }; + Self::take_last_len(&combined, dim, max_len) + } + + fn append_conv1_outputs( + &self, + new_linear: Option>, + state: &mut MoonshineEncoderStreamState, + ) -> Option> { + let Some(new_linear) = new_linear else { + return None; + }; + let old_linear_frames = state.total_linear_frames; + let old_conv1_frames = state.total_conv1_frames; + let new_linear_frames = new_linear.shape()[2]; + let total_linear_frames = old_linear_frames + new_linear_frames; + let total_conv1_frames = Self::ceil_div_2(total_linear_frames); + let new_conv1_frames = total_conv1_frames.saturating_sub(old_conv1_frames); + + let old_linear_tail = state.linear_tail.clone(); + state.linear_tail = Self::append_tail(state.linear_tail.as_ref(), &new_linear, 2, 4); + state.total_linear_frames = total_linear_frames; + state.total_conv1_frames = total_conv1_frames; + + if new_conv1_frames == 0 { + return None; + } + + let output = if old_linear_frames == 0 { + self.conv1 + .forward(&causal_pad_1d(&new_linear, self.config.hidden_size, 4)) + .silu() + .to_concrete() + } else { + let start_frame = old_conv1_frames.saturating_mul(2).saturating_sub(4); + let prefix_len = old_linear_frames.saturating_sub(start_frame); + let mut parts = Vec::with_capacity(2); + if prefix_len > 0 { + if let Some(tail) = &old_linear_tail { + parts.push( + tail.narrow(2, tail.shape()[2] - prefix_len, prefix_len) + .to_concrete(), + ); + } + } + parts.push(new_linear.clone()); + let window = Tensor::cat(parts, 2); + let pad = self + .conv1 + .kernel_size() + .saturating_sub(1) + .saturating_sub(old_conv1_frames.saturating_mul(self.conv1.config().stride)); + let window = if pad > 0 { + causal_pad_1d(&window, self.config.hidden_size, pad) + } else { + window + }; + self.conv1.forward(&window).silu().to_concrete() + }; + + Some( + output + .narrow(2, output.shape()[2] - new_conv1_frames, new_conv1_frames) + .to_concrete(), + ) + } + + fn append_encoder_outputs( + &self, + new_conv1: Option>, + state: &mut MoonshineEncoderStreamState, + ) -> Option> { + let Some(new_conv1) = new_conv1 else { + return None; + }; + let old_conv1_frames = state + .total_conv1_frames + .saturating_sub(new_conv1.shape()[2]); + let old_seen_frames = state.total_seen_frames; + let total_conv1_frames = state.total_conv1_frames; + let total_seen_frames = Self::ceil_div_2(total_conv1_frames); + let new_seen_frames = total_seen_frames.saturating_sub(old_seen_frames); + + let old_conv1_tail = state.conv1_tail.clone(); + state.conv1_tail = Self::append_tail(state.conv1_tail.as_ref(), &new_conv1, 2, 4); + state.total_seen_frames = total_seen_frames; + + if new_seen_frames == 0 { + return None; + } + + let output = if old_conv1_frames == 0 { + self.conv2 + .forward(&causal_pad_1d(&new_conv1, self.config.hidden_size * 2, 4)) + .transpose(1, 2) + .to_concrete() + } else { + let start_frame = old_seen_frames.saturating_mul(2).saturating_sub(4); + let prefix_len = old_conv1_frames.saturating_sub(start_frame); + let mut parts = Vec::with_capacity(2); + if prefix_len > 0 { + if let Some(tail) = &old_conv1_tail { + parts.push( + tail.narrow(2, tail.shape()[2] - prefix_len, prefix_len) + .to_concrete(), + ); + } + } + parts.push(new_conv1.clone()); + let window = Tensor::cat(parts, 2); + let pad = self + .conv2 + .kernel_size() + .saturating_sub(1) + .saturating_sub(old_seen_frames.saturating_mul(self.conv2.config().stride)); + let window = if pad > 0 { + causal_pad_1d(&window, self.config.hidden_size * 2, pad) + } else { + window + }; + self.conv2.forward(&window).transpose(1, 2).to_concrete() + }; + + Some( + output + .narrow(1, output.shape()[1] - new_seen_frames, new_seen_frames) + .to_concrete(), + ) + } + + pub fn encode_stream( + &self, + device: &Device, + state: &mut MoonshineEncoderStreamState, + samples: &[f32], + flush: bool, + ) -> Result { + if !samples.is_empty() { + state.sample_remainder.extend_from_slice(samples); + } + let frame_len = self.config.frame_len(); + let usable_sample_count = state.sample_remainder.len() / frame_len * frame_len; + let new_hidden_states = if usable_sample_count > 0 { + let frames = self.preprocess(&state.sample_remainder[..usable_sample_count]); + state.sample_remainder.drain(..usable_sample_count); + let frame_count = frames.len() / frame_len; + let hidden_states = Tensor::from_slice(device, [1, frame_count, frame_len], &frames); + let hidden_states = self.linear.forward(&hidden_states).silu().to_concrete(); + Some(hidden_states.transpose(1, 2).to_concrete()) + } else { + None + }; + + let mut hidden_states = self.append_conv1_outputs(new_hidden_states, state); + hidden_states = self.append_encoder_outputs(hidden_states, state); + + for (idx, layer) in self.layers.iter().enumerate() { + let window = self.config.sliding_windows[idx]; + hidden_states = layer.forward_stream( + device, + hidden_states, + &mut state.layer_states[idx], + window[0], + window[1], + flush, + )?; + } + + let finalized_hidden_states = hidden_states + .map(|hidden_states| self.final_norm.forward(&hidden_states).to_concrete()); + if let Some(finalized_hidden_states) = &finalized_hidden_states { + state.total_finalized_frames += finalized_hidden_states.shape()[1]; + } + + Ok(MoonshineEncoderStreamAppend { + hidden_states: finalized_hidden_states, + total_seen_frames: state.total_seen_frames, + total_finalized_frames: state.total_finalized_frames, + usable_input_samples: state.total_linear_frames * frame_len, + }) + } + + pub fn encode(&self, device: &Device, samples: &[f32]) -> Result> { + let frames = self.preprocess(samples); + let frame_len = self.config.frame_len(); + if frames.is_empty() { + return Err(Error::msg("moonshine input is too short")); + } + let frame_count = frames.len() / frame_len; + let mut hidden_states = Tensor::from_slice(device, [1, frame_count, frame_len], &frames); + hidden_states = self.linear.forward(&hidden_states).silu(); + let mut hidden_states = hidden_states.transpose(1, 2).to_concrete(); + hidden_states = self + .conv1 + .forward(&causal_pad_1d(&hidden_states, self.config.hidden_size, 4)) + .silu() + .to_concrete(); + hidden_states = self + .conv2 + .forward(&causal_pad_1d( + &hidden_states, + self.config.hidden_size * 2, + 4, + )) + .transpose(1, 2) + .to_concrete(); + + let seq_len = hidden_states.shape()[1]; + for (idx, layer) in self.layers.iter().enumerate() { + let window = self.config.sliding_windows[idx]; + let mask = sliding_window_mask(device, seq_len, window[0], window[1]); + hidden_states = layer.forward(&hidden_states, &mask)?; + } + + Ok(self.final_norm.forward(&hidden_states)) + } +} + +pub struct MoonshineDecoder { + token_embedding: Embedding, + pos_embedding: Embedding, + proj: Option>, + layers: Vec, + norm: LayerNorm<1, f32>, + output: Linear, + max_target_positions: usize, + mask_cache: Arc>, + config: MoonshineStreamingConfig, +} + +#[derive(Clone, Default)] +pub struct MoonshineDecoderCache { + pub(crate) tokens: Vec, + layers: Vec, + pub(crate) encoder_seq_len: usize, +} + +impl MoonshineDecoder { + fn load( + config: &MoonshineStreamingConfig, + device: &Device, + vb: &mut VarBuilder, + ) -> Result { + let rope_cache = Arc::new(RopeCache::new( + config.decoder_rotary_dim(), + config.max_position_embeddings, + config.rope_theta(), + device, + )?); + let layers = (0..config.num_hidden_layers) + .map(|idx| { + MoonshineDecoderLayer::load( + config, + device, + &mut vb.pp(format!("layers.{idx}")), + rope_cache.clone(), + ) + }) + .collect::>>()?; + let proj = (config.encoder_hidden_size() != config.hidden_size) + .then(|| Linear::load(device, &mut vb.pp("proj"))) + .transpose()?; + Ok(Self { + token_embedding: Embedding::load(device, &mut vb.pp("embed_tokens"))?, + pos_embedding: Embedding::load(device, &mut vb.pp("pos_emb"))?, + proj, + layers, + norm: LayerNorm::load(device, &mut vb.pp("norm"), 1e-5)?, + output: Linear::load(device, &mut vb.pp("proj_out"))?, + max_target_positions: config.max_position_embeddings, + mask_cache: Default::default(), + config: config.clone(), + }) + } + + fn adapt_encoder( + &self, + encoder_hidden_states: &Tensor<3, f32>, + start_pos: usize, + ) -> Result> { + let [_, seq_len, _] = encoder_hidden_states.shape(); + if start_pos + seq_len > self.config.max_position_embeddings { + return Err(Error::msg( + "moonshine encoder sequence exceeds max_position_embeddings", + )); + } + let positions: Vec = (start_pos as u32..(start_pos + seq_len) as u32).collect(); + let position_ids: Tensor<1, u32> = Tensor::new(&encoder_hidden_states.device(), &positions); + let position_embeddings: Tensor<2, f32> = self.pos_embedding.forward(&position_ids); + let position_embeddings = position_embeddings.unsqueeze(0).to_concrete(); + let hidden_states = (encoder_hidden_states + &position_embeddings).to_concrete(); + Ok(match &self.proj { + Some(proj) => proj.forward(&hidden_states), + None => hidden_states, + }) + } + + pub fn prepare_encoder_hidden_states( + &self, + encoder_hidden_states: &Tensor<3, f32>, + ) -> Result> { + self.prepare_encoder_hidden_states_range(encoder_hidden_states, 0) + } + + pub fn prepare_encoder_hidden_states_range( + &self, + encoder_hidden_states: &Tensor<3, f32>, + start_pos: usize, + ) -> Result> { + Ok(self + .adapt_encoder(encoder_hidden_states, start_pos)? + .to_concrete()) + } + + pub fn decode_prepared( + &mut self, + tokens: &[u32], + encoder_hidden_states: &Tensor<3, f32>, + attention_output: Option<&mut Vec>>, + ) -> Result> { + let device = encoder_hidden_states.device(); + let token_ids: Tensor<2, u32> = Tensor::new(&device, tokens) + .reshape([1, tokens.len()]) + .to_concrete(); + let mut hidden_states: Tensor<3, f32> = self.token_embedding.forward(&token_ids); + let mask = causal_mask(&device, tokens.len()); + + let mut attention_output = attention_output; + for layer in self.layers.iter_mut() { + hidden_states = layer.forward( + &hidden_states, + encoder_hidden_states, + &mask, + attention_output.as_mut().map(|outputs| &mut **outputs), + )?; + } + let hidden_states = self.norm.forward_fused(&hidden_states); + Ok(self.output.forward(&hidden_states)) + } + + pub fn decode_cached( + &mut self, + tokens: &[u32], + encoder_hidden_states: &Tensor<3, f32>, + cache: &mut MoonshineDecoderCache, + ) -> Result> { + let index_pos = cache.tokens.len(); + let seq_len = tokens.len(); + if index_pos + seq_len > self.max_target_positions { + return Err(Error::msg( + "moonshine decoder sequence exceeds max_position_embeddings", + )); + } + cache.tokens.extend_from_slice(tokens); + self.sync_cross_attention_kv(encoder_hidden_states, cache)?; + + let device = encoder_hidden_states.device(); + let self_mask = self.mask_cache.get_mask(seq_len, index_pos, None, &device); + let token_tensor: Tensor<1, u32> = Tensor::new(&device, tokens); + let token_tensor = token_tensor.unsqueeze(0).to_concrete(); + let mut hidden_states: Tensor<3, f32> = self.token_embedding.forward(&token_tensor); + + for (i, layer) in self.layers.iter_mut().enumerate() { + if cache.layers.len() <= i { + let cross_attn_kv = layer.encoder_attn.project_kv(encoder_hidden_states); + cache.layers.push(MoonshineDecoderLayerCache { + self_attn: MoonshineAttentionCache::new(self.max_target_positions), + cross_attn_kv, + }); + } + + hidden_states = layer.forward_cached( + &hidden_states, + &self_mask, + index_pos, + &mut cache.layers[i], + None, + )?; + } + + cache.encoder_seq_len = encoder_hidden_states.shape()[1]; + let hidden_states = self.norm.forward_fused(&hidden_states); + Ok(self.output.forward(&hidden_states)) + } + + pub fn prepare_cross_attention_cache( + &mut self, + encoder_hidden_states: &Tensor<3, f32>, + cache: &mut MoonshineDecoderCache, + ) -> Result<()> { + self.sync_cross_attention_kv(encoder_hidden_states, cache) + } + + pub fn decode( + &mut self, + tokens: &[u32], + encoder_hidden_states: &Tensor<3, f32>, + attention_output: Option<&mut Vec>>, + ) -> Result> { + let encoder_hidden_states = self.prepare_encoder_hidden_states(encoder_hidden_states)?; + self.decode_prepared(tokens, &encoder_hidden_states, attention_output) + } + + fn sync_cross_attention_kv( + &mut self, + encoder_hidden_states: &Tensor<3, f32>, + cache: &mut MoonshineDecoderCache, + ) -> Result<()> { + let encoder_seq_len = encoder_hidden_states.shape()[1]; + if encoder_seq_len == 0 { + cache.encoder_seq_len = 0; + return Ok(()); + } + if cache.layers.is_empty() { + for layer in self.layers.iter_mut() { + let cross_attn_kv = layer.encoder_attn.project_kv(encoder_hidden_states); + cache.layers.push(MoonshineDecoderLayerCache { + self_attn: MoonshineAttentionCache::new(self.max_target_positions), + cross_attn_kv, + }); + } + cache.encoder_seq_len = encoder_seq_len; + return Ok(()); + } + if cache.encoder_seq_len >= encoder_seq_len { + return Ok(()); + } + + let new_encoder_hidden_states = encoder_hidden_states + .narrow( + 1, + cache.encoder_seq_len, + encoder_seq_len - cache.encoder_seq_len, + ) + .to_concrete(); + for (layer, layer_cache) in self.layers.iter_mut().zip(cache.layers.iter_mut()) { + let new_cross_attn_kv = layer.encoder_attn.project_kv(&new_encoder_hidden_states); + layer_cache.cross_attn_kv = ( + Tensor::cat( + [layer_cache.cross_attn_kv.0.clone(), new_cross_attn_kv.0], + 1, + ) + .to_concrete(), + Tensor::cat( + [layer_cache.cross_attn_kv.1.clone(), new_cross_attn_kv.1], + 1, + ) + .to_concrete(), + ); + } + cache.encoder_seq_len = encoder_seq_len; + Ok(()) + } +} + +pub struct Moonshine { + pub encoder: MoonshineEncoder, + pub decoder: MoonshineDecoder, +} + +impl Moonshine { + pub fn load( + device: &Device, + vb: &mut VarBuilder, + config: MoonshineStreamingConfig, + ) -> Result { + let encoder = + MoonshineEncoder::load(&config.encoder_config, device, &mut vb.pp("model.encoder"))?; + let decoder = MoonshineDecoder::load(&config, device, &mut vb.pp("model.decoder"))?; + Ok(Self { encoder, decoder }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::moonshine_config::MoonshineStreamingConfig; + use std::{fs, io::Cursor, path::Path, path::PathBuf}; + + fn load_tiny_from_single_gguf() -> Option<(Moonshine, MoonshineStreamingConfig, Device)> { + load_tiny_with_device(Device::cpu()) + } + + fn load_tiny_with_device( + device: Device, + ) -> Option<(Moonshine, MoonshineStreamingConfig, Device)> { + let path = Path::new(env!("CARGO_MANIFEST_DIR")) + .join("artifacts") + .join("moonshine-streaming-tiny.gguf"); + let path = if path.exists() { + path + } else if let Ok(path) = std::env::var("RWHISPER_MOONSHINE_GGUF") { + PathBuf::from(path) + } else { + path + }; + if !path.exists() { + return None; + } + let weights = fs::read(&path).unwrap(); + let mut reader = Cursor::new(weights); + let mut vb = VarBuilder::from_gguf(&mut reader).unwrap(); + let config_json = vb + .get_metadata("rwhisper.config.json") + .expect("missing rwhisper.config.json metadata in GGUF") + .to_string() + .unwrap(); + let config: MoonshineStreamingConfig = serde_json::from_str(&config_json).unwrap(); + let model = Moonshine::load(&device, &mut vb, config.clone()).unwrap(); + Some((model, config, device)) + } + + fn read_jfk_pcm(target_rate: usize) -> Vec { + let path = Path::new(env!("CARGO_MANIFEST_DIR")).join("examples/samples_jfk.wav"); + let bytes = fs::read(&path).unwrap(); + let decoder = rodio::Decoder::new(std::io::Cursor::new(bytes)).unwrap(); + let resampled = rodio::source::UniformSourceIterator::new(decoder, 1, target_rate as u32); + resampled.map(|s: i16| s as f32 / i16::MAX as f32).collect() + } + + // Regression guard for the streaming encoder: chunked encode + flush must + // match the non-streaming encoder on real audio. Broken previously because + // `append_conv1_outputs` / `append_encoder_outputs` updated the conv tail + // before reading it for the next window, so later chunks padded their conv + // input with already-current frames instead of the prior chunk's tail. + #[tokio::test] + async fn streaming_encoder_matches_full_encoder_on_jfk() { + let Some((model, config, device)) = load_tiny_from_single_gguf() else { + return; + }; + + let frame_len = config.encoder_config.frame_len(); + let sample_rate = config.encoder_config.sample_rate; + let raw = read_jfk_pcm(sample_rate); + let usable_len = raw.len() / frame_len * frame_len; + let samples = raw[..usable_len].to_vec(); + + let full = model.encoder.encode(&device, &samples).unwrap(); + let full = full.as_slice().await.unwrap(); + + let chunk_samples = samples.len() / 2 / frame_len * frame_len; + let mut stream_state = model.encoder.new_stream_state(); + let mut streamed_parts = Vec::new(); + for chunk in samples.chunks(chunk_samples) { + let append = model + .encoder + .encode_stream(&device, &mut stream_state, chunk, false) + .unwrap(); + if let Some(hidden_states) = append.hidden_states { + streamed_parts.push(hidden_states); + } + } + let append = model + .encoder + .encode_stream(&device, &mut stream_state, &[], true) + .unwrap(); + if let Some(hidden_states) = append.hidden_states { + streamed_parts.push(hidden_states); + } + let streamed = Tensor::cat(streamed_parts, 1); + let streamed = streamed.as_slice().await.unwrap(); + + assert_eq!(full.shape(), streamed.shape()); + let max_diff = full + .as_slice() + .iter() + .zip(streamed.as_slice()) + .map(|(a, b)| (a - b).abs()) + .fold(0.0f32, f32::max); + eprintln!("streaming encoder max_diff vs full = {max_diff}"); + assert!( + max_diff < 0.1, + "streaming encoder diverged from full encoder on JFK: max_diff={max_diff}" + ); + } + + #[tokio::test] + async fn gpu_encoder_matches_cpu_on_zeros() { + let Some((cpu_model, config, cpu_device)) = load_tiny_with_device(Device::cpu()) else { + return; + }; + let Ok(gpu_device) = Device::new().await else { + return; + }; + let Some((gpu_model, _, gpu_device)) = load_tiny_with_device(gpu_device) else { + return; + }; + + let samples = vec![0.0f32; config.encoder_config.frame_len() * 80]; + let cpu = cpu_model.encoder.encode(&cpu_device, &samples).unwrap(); + let gpu = gpu_model.encoder.encode(&gpu_device, &samples).unwrap(); + + let cpu = cpu.as_slice().await.unwrap(); + let gpu = gpu.as_slice().await.unwrap(); + assert_eq!(cpu.shape(), gpu.shape()); + let max_diff = cpu + .as_slice() + .iter() + .zip(gpu.as_slice()) + .map(|(a, b)| (a - b).abs()) + .fold(0.0f32, f32::max); + eprintln!("moonshine GPU encoder max_diff vs CPU = {max_diff}"); + assert!( + max_diff < 0.1, + "GPU encoder diverged from CPU encoder: max_diff={max_diff}" + ); + } + + #[tokio::test] + async fn gpu_cached_decode_matches_cpu_on_zeros() { + let Some((mut cpu_model, config, cpu_device)) = load_tiny_with_device(Device::cpu()) else { + return; + }; + let Ok(gpu_device) = Device::new().await else { + return; + }; + let Some((mut gpu_model, _, gpu_device)) = load_tiny_with_device(gpu_device) else { + return; + }; + + let samples = vec![0.0f32; config.encoder_config.frame_len() * 80]; + let cpu_encoder_hidden_states = cpu_model.encoder.encode(&cpu_device, &samples).unwrap(); + let gpu_encoder_hidden_states = gpu_model.encoder.encode(&gpu_device, &samples).unwrap(); + let cpu_prepared = cpu_model + .decoder + .prepare_encoder_hidden_states(&cpu_encoder_hidden_states) + .unwrap(); + let gpu_prepared = gpu_model + .decoder + .prepare_encoder_hidden_states(&gpu_encoder_hidden_states) + .unwrap(); + + let tokens = [ + config.decoder_start_token(), + 123, + 456, + config.eos_token_id.saturating_sub(1), + ]; + let mut cpu_cache = MoonshineDecoderCache::default(); + let mut gpu_cache = MoonshineDecoderCache::default(); + let mut cpu_logits = cpu_model + .decoder + .decode_cached(&tokens[..1], &cpu_prepared, &mut cpu_cache) + .unwrap(); + let mut gpu_logits = gpu_model + .decoder + .decode_cached(&tokens[..1], &gpu_prepared, &mut gpu_cache) + .unwrap(); + for token in &tokens[1..] { + cpu_logits = cpu_model + .decoder + .decode_cached(&[*token], &cpu_prepared, &mut cpu_cache) + .unwrap(); + gpu_logits = gpu_model + .decoder + .decode_cached(&[*token], &gpu_prepared, &mut gpu_cache) + .unwrap(); + } + + let cpu_last = cpu_logits.squeeze(0).squeeze(0).to_concrete(); + let gpu_last = gpu_logits.squeeze(0).squeeze(0).to_concrete(); + let cpu_last = cpu_last.as_slice().await.unwrap(); + let gpu_last = gpu_last.as_slice().await.unwrap(); + assert_eq!(cpu_last.shape(), gpu_last.shape()); + let max_diff = cpu_last + .as_slice() + .iter() + .zip(gpu_last.as_slice()) + .map(|(a, b)| (a - b).abs()) + .fold(0.0f32, f32::max); + eprintln!("moonshine GPU cached decode max_diff vs CPU = {max_diff}"); + assert!( + max_diff < 0.1, + "GPU cached decode diverged from CPU: max_diff={max_diff}" + ); + } + + #[tokio::test] + async fn cached_decode_matches_full_decode() { + let Some((mut model, config, device)) = load_tiny_from_single_gguf() else { + return; + }; + + let samples = vec![0.0f32; config.encoder_config.frame_len() * 80]; + let encoder_hidden_states = model.encoder.encode(&device, &samples).unwrap(); + let adapted_encoder_hidden_states = model + .decoder + .prepare_encoder_hidden_states(&encoder_hidden_states) + .unwrap(); + + let tokens = [ + config.decoder_start_token(), + 123, + 456, + config.eos_token_id.saturating_sub(1), + ]; + let full_logits = model + .decoder + .decode_prepared(&tokens, &adapted_encoder_hidden_states, None) + .unwrap(); + let full_last = full_logits + .narrow(1, tokens.len() - 1, 1) + .squeeze(1) + .squeeze(0) + .to_concrete(); + let full_last = full_last.as_slice().await.unwrap(); + + let mut cache = MoonshineDecoderCache::default(); + let mut cached_logits = model + .decoder + .decode_cached(&tokens[..1], &adapted_encoder_hidden_states, &mut cache) + .unwrap(); + for token in &tokens[1..] { + cached_logits = model + .decoder + .decode_cached(&[*token], &adapted_encoder_hidden_states, &mut cache) + .unwrap(); + } + let cached_last = cached_logits.squeeze(0).squeeze(0).to_concrete(); + let cached_last = cached_last.as_slice().await.unwrap(); + + assert_eq!(full_last.shape(), cached_last.shape()); + let max_diff = full_last + .as_slice() + .iter() + .zip(cached_last.as_slice()) + .map(|(a, b)| (a - b).abs()) + .fold(0.0f32, f32::max); + assert!( + max_diff < 1e-3, + "cached decode diverged from full decode: max_diff={max_diff}" + ); + } +} diff --git a/models/rwhisper/src/quantized/timestamps.rs b/models/rwhisper/src/quantized/timestamps.rs index 3a5bf6014..d3ee95499 100644 --- a/models/rwhisper/src/quantized/timestamps.rs +++ b/models/rwhisper/src/quantized/timestamps.rs @@ -4,29 +4,43 @@ use fusor::Tensor; use std::num::NonZeroUsize; -use crate::config::{HOP_LENGTH, N_FRAMES, SAMPLE_RATE}; - /// Returns the token-level timestamps as a tensor of shape batch x timestamps -pub(super) async fn extract_timestamps( - // A list of (layer, head) pairs to use for timestamp determination - alignment_heads: &[[usize; 2]], +pub(crate) async fn extract_timestamps( + // A list of (layer, head) pairs to use for timestamp determination. + // If omitted, all available heads from all layers are used. + alignment_heads: Option<&[[usize; 2]]>, cross_attentions: &[Tensor<4, crate::WhisperDType>], filter_width: NonZeroUsize, n_frames: usize, + seconds_per_frame: crate::WhisperDType, mask: Vec>, ) -> fusor::Result>> { // Select relevant cross-attention heads let mut tensors_to_stack = Vec::new(); - for [layer, head] in alignment_heads.iter().copied() { - if let Some(attn) = cross_attentions.get(layer) { - let narrowed = attn.narrow(1, head, 1); - let squeezed = narrowed.squeeze(1); - tensors_to_stack.push(squeezed.to_concrete()); + if let Some(alignment_heads) = alignment_heads { + for [layer, head] in alignment_heads.iter().copied() { + if let Some(attn) = cross_attentions.get(layer) { + let narrowed = attn.narrow(1, head, 1); + let squeezed = narrowed.squeeze(1); + tensors_to_stack.push(squeezed.to_concrete()); + } + } + } else { + for attn in cross_attentions { + let heads = attn.shape()[1]; + for head in 0..heads { + let narrowed = attn.narrow(1, head, 1); + let squeezed = narrowed.squeeze(1); + tensors_to_stack.push(squeezed.to_concrete()); + } } } + if tensors_to_stack.is_empty() { + return Ok(Vec::new()); + } let stacked = Tensor::stack(tensors_to_stack.into_iter(), 0); let permuted = stacked.permute([1, 0, 2, 3]); - let weights = permuted.narrow(3, 0, n_frames.min(N_FRAMES) / 2); + let weights = permuted.narrow(3, 0, n_frames.min(permuted.shape()[3])); if weights.shape().contains(&0) { // No tokens to be aligned @@ -89,10 +103,7 @@ pub(super) async fn extract_timestamps( .zip(time_indices) .filter_map(|(is_jump, time_index)| { if is_jump { - Some( - time_index - / crate::WhisperDType::from((SAMPLE_RATE / (HOP_LENGTH * 2)) as f32), - ) + Some(time_index * seconds_per_frame) } else { None } diff --git a/models/rwhisper/src/source.rs b/models/rwhisper/src/source.rs index 63d15ad65..c274c9f73 100644 --- a/models/rwhisper/src/source.rs +++ b/models/rwhisper/src/source.rs @@ -1,4 +1,19 @@ use kalosm_model_types::FileSource; +use std::path::PathBuf; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum ModelFamily { + Whisper { + multilingual: bool, + heads: Option<&'static [[usize; 2]]>, + apply_speech_filter: bool, + }, + CohereTranscribe, + MoonshineStreaming { + heads: Option<&'static [[usize; 2]]>, + apply_speech_filter: bool, + }, +} /// Predefined Whisper model sources #[derive(Debug, Clone)] @@ -6,8 +21,7 @@ pub struct WhisperSource { pub(crate) model: FileSource, pub(crate) tokenizer: FileSource, pub(crate) config: FileSource, - pub(crate) multilingual: bool, - pub(crate) heads: Option<&'static [[usize; 2]]>, + pub(crate) family: ModelFamily, } impl Default for WhisperSource { @@ -24,16 +38,106 @@ impl WhisperSource { config: FileSource, multilingual: bool, heads: Option<&'static [[usize; 2]]>, + ) -> Self { + Self::new_with_family( + model, + tokenizer, + config, + ModelFamily::Whisper { + multilingual, + heads, + apply_speech_filter: true, + }, + ) + } + + pub(crate) fn new_with_family( + model: FileSource, + tokenizer: FileSource, + config: FileSource, + family: ModelFamily, ) -> Self { Self { model, tokenizer, config, - multilingual, - heads, + family, } } + /// Cohere Transcribe 03/2026 + pub fn cohere_transcribe_03_2026() -> Self { + let source = FileSource::huggingface( + "Demonthos/cohere-transcribe-03-2026-gguf".to_owned(), + "main".to_owned(), + "cohere-transcribe-03-2026.gguf".to_owned(), + ); + Self::new_with_family( + source.clone(), + source.clone(), + source, + ModelFamily::CohereTranscribe, + ) + } + + /// Cohere Transcribe 03/2026 from a local GGUF file. Tokenizer and config + /// metadata must be embedded in the GGUF (all files produced by + /// `scripts/quantize_cohere_transcribe.py` embed them). + pub fn cohere_transcribe_03_2026_local(path: impl Into) -> Self { + let source = FileSource::local(path.into()); + Self::new_with_family( + source.clone(), + source.clone(), + source, + ModelFamily::CohereTranscribe, + ) + } + + /// Moonshine streaming tiny English model. + pub fn moonshine_streaming_tiny() -> Self { + Self::moonshine_streaming_remote("moonshine-streaming-tiny.gguf") + } + + /// Moonshine streaming small English model. + pub fn moonshine_streaming_small() -> Self { + Self::moonshine_streaming_remote("moonshine-streaming-small.gguf") + } + + /// Moonshine streaming medium English model. + pub fn moonshine_streaming_medium() -> Self { + Self::moonshine_streaming_remote("moonshine-streaming-medium.gguf") + } + + fn moonshine_streaming_remote(file: &str) -> Self { + let repo = "Demonthos/moonshot-gguf".to_owned(); + let source = FileSource::huggingface(repo, "main".to_owned(), file.to_owned()); + Self::new_with_family( + source.clone(), + source.clone(), + source, + ModelFamily::MoonshineStreaming { + heads: None, + apply_speech_filter: false, + }, + ) + } + + /// Moonshine streaming English model from a local GGUF file. Tokenizer + /// and config metadata must be embedded in the GGUF (all files produced + /// by `scripts/quantize_moonshine_streaming.py` embed them). + pub fn moonshine_streaming_local(path: impl Into) -> Self { + let source = FileSource::local(path.into()); + Self::new_with_family( + source.clone(), + source.clone(), + source, + ModelFamily::MoonshineStreaming { + heads: None, + apply_speech_filter: false, + }, + ) + } + /// Tiny english model pub fn tiny_en() -> Self { let model = FileSource::huggingface( diff --git a/models/rwhisper/tests/jfk_regression.rs b/models/rwhisper/tests/jfk_regression.rs new file mode 100644 index 000000000..f69bf6c19 --- /dev/null +++ b/models/rwhisper/tests/jfk_regression.rs @@ -0,0 +1,276 @@ +use std::{fs, io::Cursor, path::Path}; + +use anyhow::{Context, Result}; +use futures_channel::mpsc; +use futures_util::{stream, StreamExt}; +use kalosm::sound::*; +use rodio::Decoder; +use rwhisper::TranscribeChunkedAudioStreamExt; + +// Reference text from JFK's January 20, 1961 inaugural address: +// "And so, my fellow Americans, ask not what your country can do for you, +// ask what you can do for your country." +const JFK_REFERENCE: &str = + "And so, my fellow Americans, ask not what your country can do for you, ask what you can do for your country."; + +fn normalize_text(text: &str) -> String { + let filtered = text + .chars() + .map(|c| { + if c.is_ascii_alphanumeric() || c.is_ascii_whitespace() { + c.to_ascii_lowercase() + } else { + ' ' + } + }) + .collect::(); + filtered.split_whitespace().collect::>().join(" ") +} + +fn word_error_rate(reference: &str, hypothesis: &str) -> f32 { + let reference = normalize_text(reference); + let hypothesis = normalize_text(hypothesis); + let reference = reference.split_whitespace().collect::>(); + let hypothesis = hypothesis.split_whitespace().collect::>(); + + let mut dp = vec![vec![0usize; hypothesis.len() + 1]; reference.len() + 1]; + for (i, row) in dp.iter_mut().enumerate() { + row[0] = i; + } + for (j, cell) in dp[0].iter_mut().enumerate() { + *cell = j; + } + + for i in 1..=reference.len() { + for j in 1..=hypothesis.len() { + let substitution_cost = usize::from(reference[i - 1] != hypothesis[j - 1]); + dp[i][j] = (dp[i - 1][j] + 1) + .min(dp[i][j - 1] + 1) + .min(dp[i - 1][j - 1] + substitution_cost); + } + } + + if reference.is_empty() { + 0.0 + } else { + dp[reference.len()][hypothesis.len()] as f32 / reference.len() as f32 + } +} + +fn moonshine_source(size: &str) -> WhisperSource { + let local = Path::new(env!("CARGO_MANIFEST_DIR")) + .join("artifacts") + .join(format!("moonshine-streaming-{size}.gguf")); + if local.exists() { + return WhisperSource::moonshine_streaming_local(local); + } + match size { + "tiny" => WhisperSource::moonshine_streaming_tiny(), + "small" => WhisperSource::moonshine_streaming_small(), + "medium" => WhisperSource::moonshine_streaming_medium(), + other => panic!("unknown Moonshine size: {other}"), + } +} + +async fn transcribe_jfk_sample(source: WhisperSource) -> Result { + let model = WhisperBuilder::default() + .with_source(source) + .build() + .await + .context("failed to build model")?; + + let contents = fs::read(Path::new(env!("CARGO_MANIFEST_DIR")).join("examples/samples_jfk.wav")) + .context("failed to read JFK sample")?; + let audio = Decoder::new(Cursor::new(contents)).context("failed to decode JFK sample")?; + + let mut stream = model.transcribe(audio); + let mut transcript = String::new(); + while let Some(segment) = stream.next().await { + if !transcript.is_empty() { + transcript.push(' '); + } + transcript.push_str(segment.text().trim()); + } + + Ok(transcript) +} + +async fn transcribe_jfk_sample_streaming( + source: WhisperSource, + chunk_ms: usize, + max_seconds: Option, +) -> Result { + let model = WhisperBuilder::default() + .with_source(source) + .build() + .await + .context("failed to build model")?; + + let contents = fs::read(Path::new(env!("CARGO_MANIFEST_DIR")).join("examples/samples_jfk.wav")) + .context("failed to read JFK sample")?; + let audio = Decoder::new(Cursor::new(contents)).context("failed to decode JFK sample")?; + let sample_rate = rodio::Source::sample_rate(&audio); + let channels = rodio::Source::channels(&audio); + let samples: Vec = if let Some(max_seconds) = max_seconds { + rodio::Source::take_duration(audio, std::time::Duration::from_secs_f32(max_seconds)) + .collect() + } else { + audio.collect() + }; + let samples_per_chunk = ((sample_rate as usize * chunk_ms) / 1000).max(1) * channels as usize; + let chunks = samples + .chunks(samples_per_chunk) + .map(|chunk| rodio::buffer::SamplesBuffer::new(channels, sample_rate, chunk.to_vec())) + .collect::>(); + + let mut stream = stream::iter(chunks).transcribe(model); + let mut transcript = String::new(); + while let Some(segment) = stream.next().await { + if !transcript.is_empty() { + transcript.push(' '); + } + transcript.push_str(segment.text().trim()); + } + + Ok(transcript) +} + +async fn first_streaming_emission_chunk( + source: WhisperSource, + chunk_ms: usize, +) -> Result> { + let model = WhisperBuilder::default() + .with_source(source) + .build() + .await + .context("failed to build model")?; + + let contents = fs::read(Path::new(env!("CARGO_MANIFEST_DIR")).join("examples/samples_jfk.wav")) + .context("failed to read JFK sample")?; + let audio = Decoder::new(Cursor::new(contents)).context("failed to decode JFK sample")?; + let sample_rate = rodio::Source::sample_rate(&audio); + let channels = rodio::Source::channels(&audio); + let samples: Vec = audio.collect(); + let samples_per_chunk = ((sample_rate as usize * chunk_ms) / 1000).max(1) * channels as usize; + let chunks = samples + .chunks(samples_per_chunk) + .map(|chunk| rodio::buffer::SamplesBuffer::new(channels, sample_rate, chunk.to_vec())) + .collect::>(); + let chunk_count = chunks.len(); + + let (tx, rx) = mpsc::unbounded(); + let mut stream = rx.transcribe(model); + + for (index, chunk) in chunks.into_iter().enumerate() { + tx.unbounded_send(chunk) + .map_err(|_| anyhow::anyhow!("failed to send chunk to transcription stream"))?; + if let Ok(Some(segment)) = + tokio::time::timeout(std::time::Duration::from_millis(250), stream.next()).await + { + if !segment.text().trim().is_empty() { + return Ok(Some(index)); + } + } + } + + drop(tx); + while let Some(segment) = stream.next().await { + if !segment.text().trim().is_empty() { + return Ok(Some(chunk_count)); + } + } + + Ok(None) +} + +#[test] +fn normalize_text_keeps_the_reference_stable() { + assert_eq!( + normalize_text(JFK_REFERENCE), + "and so my fellow americans ask not what your country can do for you ask what you can do for your country" + ); +} + +#[test] +fn word_error_rate_is_zero_for_exact_match() { + assert_eq!(word_error_rate(JFK_REFERENCE, JFK_REFERENCE), 0.0); +} + +#[tokio::test] +async fn moonshine_tiny_matches_the_jfk_reference() -> Result<()> { + let transcript = transcribe_jfk_sample(moonshine_source("tiny")).await?; + assert_eq!( + normalize_text(&transcript), + normalize_text(JFK_REFERENCE), + "unexpected Moonshine transcript: {transcript}" + ); + Ok(()) +} + +#[tokio::test] +async fn moonshine_small_matches_the_jfk_reference() -> Result<()> { + let transcript = transcribe_jfk_sample(moonshine_source("small")).await?; + assert_eq!( + normalize_text(&transcript), + normalize_text(JFK_REFERENCE), + "unexpected Moonshine transcript: {transcript}" + ); + Ok(()) +} + +#[tokio::test] +async fn moonshine_medium_matches_the_jfk_reference() -> Result<()> { + let transcript = transcribe_jfk_sample(moonshine_source("medium")).await?; + assert_eq!( + normalize_text(&transcript), + normalize_text(JFK_REFERENCE), + "unexpected Moonshine transcript: {transcript}" + ); + Ok(()) +} + +#[tokio::test] +async fn moonshine_tiny_streaming_matches_the_jfk_reference() -> Result<()> { + let transcript = transcribe_jfk_sample_streaming(moonshine_source("tiny"), 250, None).await?; + assert_eq!( + normalize_text(&transcript), + normalize_text(JFK_REFERENCE), + "unexpected Moonshine streaming transcript: {transcript}" + ); + Ok(()) +} + +#[tokio::test] +async fn moonshine_tiny_streaming_short_prefix_smoke() -> Result<()> { + let transcript = + transcribe_jfk_sample_streaming(moonshine_source("tiny"), 1_000, Some(2.0)).await?; + assert_eq!( + normalize_text(&transcript), + "and so my fellow americans", + "unexpected short Moonshine streaming transcript: {transcript}" + ); + Ok(()) +} + +#[tokio::test] +async fn moonshine_tiny_streaming_emits_before_the_clip_ends() -> Result<()> { + let first_chunk = first_streaming_emission_chunk(moonshine_source("tiny"), 250).await?; + let first_chunk = + first_chunk.context("streaming Moonshine never emitted a non-empty segment")?; + assert!( + first_chunk <= 8, + "expected streaming Moonshine to emit within the first 2 seconds; first emission chunk index was {first_chunk}" + ); + Ok(()) +} + +#[tokio::test] +async fn cohere_transcribe_matches_the_jfk_reference() -> Result<()> { + let transcript = transcribe_jfk_sample(WhisperSource::cohere_transcribe_03_2026()).await?; + assert_eq!( + normalize_text(&transcript), + normalize_text(JFK_REFERENCE), + "unexpected Cohere transcript: {transcript}" + ); + Ok(()) +}