Transaction-Science/open-standards
Watch
1
0
Fork
You've already forked open-standards
 0
Code Issues Pull requests Projects Releases 15 Packages Wiki Activity Actions 1

jouleclaw-backend-apple: 84 clippy warnings -> 0, and -D warnings now covers both configurations #437

Merged
dcharlot merged 2 commits from fix/backend-apple-lints into main 2026-11-01 07:42:04 -05:00
Conversation 0 Commits 2 Files changed 5 +100 -92
5 changed files with 100 additions and 92 deletions
Download patch file Download diff file Expand all files Collapse all files

16
.github/workflows/standards-build.yml vendored
View file

@ -597,12 +597,14 @@ jobs:
# And the crate omni loads its weights through, which was linted
# by nothing at all: 161 warnings when first measured, now 0.
#
# Default features only. Under `rayon fabric-meter` this crate is
# still clean, but clippy also reports jouleclaw-backend-apple
# through the path dependency (`--no-deps` does not exclude
# workspace members), and that crate has 43 warnings of its own --
# mostly `unused_unsafe` and `unused_imports` that are unused in
# THAT configuration and needed in others, so they want cfg work
# rather than deletion.
# Both configurations. `--no-deps` does not exclude workspace path
# dependencies, so the fabric-meter line also covers
# jouleclaw-backend-apple, which is why that crate is linted on its
# own below rather than only incidentally.
lint jouleclaw/jouleclaw-rs -p jouleclaw-loader-gguf
lint jouleclaw/jouleclaw-rs -p jouleclaw-loader-gguf --features "rayon fabric-meter"
# The Apple backend: 84 warnings when first measured on its own, now
# 0. It has no features, only a macOS-conditional dependency block,
# so there is one configuration to keep clean.
lint jouleclaw/jouleclaw-rs -p jouleclaw-backend-apple
lint sandbox/sandbox-rs -p sandbox-meter -p sandbox-proof

2
jouleclaw/jouleclaw-rs/crates/jouleclaw-backend-apple/examples/mpp_diag.rs
View file

@ -13,7 +13,7 @@ fn main() {
MTLComputeCommandEncoder, MTLComputePipelineState, MTLCreateSystemDefaultDevice,
MTLDevice, MTLLibrary, MTLResourceOptions, MTLSize,
};
fn f32_to_bf16(v: f32) -> u16 { let b = v.to_bits(); (((b + (((b >> 16) & 1) + 0x7fff)) >> 16)) as u16 }
fn f32_to_bf16(v: f32) -> u16 { let b = v.to_bits(); ((b + (((b >> 16) & 1) + 0x7fff)) >> 16) as u16 }
const KERNEL: &str = r#"
#include <metal_stdlib>

2
jouleclaw/jouleclaw-rs/crates/jouleclaw-backend-apple/examples/mpp_gemm_bench.rs
View file

@ -25,7 +25,7 @@ fn main() {
MTLComputeCommandEncoder, MTLComputePipelineState, MTLCreateSystemDefaultDevice,
MTLDevice, MTLLibrary, MTLResourceOptions, MTLSize,
};
fn f32_to_bf16(v: f32) -> u16 { let b = v.to_bits(); (((b + (((b >> 16) & 1) + 0x7fff)) >> 16)) as u16 }
fn f32_to_bf16(v: f32) -> u16 { let b = v.to_bits(); ((b + (((b >> 16) & 1) + 0x7fff)) >> 16) as u16 }
fn bf16_to_f32(h: u16) -> f32 { f32::from_bits((h as u32) << 16) }
const KERNEL: &str = r#"

8
jouleclaw/jouleclaw-rs/crates/jouleclaw-backend-apple/src/lib.rs
View file

@ -20,6 +20,14 @@
//! returns empty and the kernel constructors return `None`. The runtime
//! falls back to the reference backend.
// ── lints this crate has decided against ────────────────────────────────────
//
// The same two decisions jouleclaw-omni and jouleclaw-loader-gguf record: a
// GPU dispatch signature is its kernel's parameter list, and these kernels
// walk several buffers at one index.
#![allow(clippy::too_many_arguments)]
#![allow(clippy::needless_range_loop)]
mod accelerate;
mod gpu_load;
mod metal_compute;

164
jouleclaw/jouleclaw-rs/crates/jouleclaw-backend-apple/src/mps_matmul.rs
View file

@ -192,7 +192,7 @@ pub mod inner {
fn context() -> Option<&'static MpsContext> {
static CTX: OnceLock<Option<MpsContext>> = OnceLock::new();
CTX.get_or_init(|| unsafe {
CTX.get_or_init(|| {
let device = MTLCreateSystemDefaultDevice()?;
let queue = device.newCommandQueue()?;
Some(MpsContext {
@ -335,7 +335,7 @@ pub mod inner {
assert_eq!(w_bf16.len(), n * k * 2);
let ctx = context().ok_or("no Metal device")?;
let mut half: Vec<u16> = Vec::with_capacity(n * k);
for c in w_bf16.chunks_exact(2) {
for c in w_bf16.as_chunks::<2>().0 {
let v = f32::from_bits((u16::from_le_bytes([c[0], c[1]]) as u32) << 16);
if v.is_finite() && v.abs() > 65504.0 {
return Err(format!("weight {v} overflows f16"));
@ -451,11 +451,11 @@ kernel void silu_mul_f32_to_f16(
fn silu_mul_pipeline()
-> Option<&'static Retained<ProtocolObject<dyn objc2_metal::MTLComputePipelineState>>> {
use objc2_foundation::NSString;
use objc2_metal::{MTLComputePipelineState as _, MTLLibrary as _};
use objc2_metal::MTLLibrary as _;
static P: OnceLock<
Option<Retained<ProtocolObject<dyn objc2_metal::MTLComputePipelineState>>>,
> = OnceLock::new();
P.get_or_init(|| unsafe {
P.get_or_init(|| {
let ctx = context()?;
let src = NSString::from_str(SILU_MUL_MSL);
let lib = ctx
@ -494,11 +494,11 @@ kernel void silu_mul_f32_to_f32(
fn silu_mul_f32_pipeline()
-> Option<&'static Retained<ProtocolObject<dyn objc2_metal::MTLComputePipelineState>>> {
use objc2_foundation::NSString;
use objc2_metal::{MTLComputePipelineState as _, MTLLibrary as _};
use objc2_metal::MTLLibrary as _;
static P: OnceLock<
Option<Retained<ProtocolObject<dyn objc2_metal::MTLComputePipelineState>>>,
> = OnceLock::new();
P.get_or_init(|| unsafe {
P.get_or_init(|| {
let ctx = context()?;
let lib = ctx.device
.newLibraryWithSource_options_error(&NSString::from_str(SILU_MUL_F32_MSL), None)
@ -528,11 +528,11 @@ kernel void residual_add_f32(
fn residual_add_pipeline()
-> Option<&'static Retained<ProtocolObject<dyn objc2_metal::MTLComputePipelineState>>> {
use objc2_foundation::NSString;
use objc2_metal::{MTLComputePipelineState as _, MTLLibrary as _};
use objc2_metal::MTLLibrary as _;
static P: OnceLock<
Option<Retained<ProtocolObject<dyn objc2_metal::MTLComputePipelineState>>>,
> = OnceLock::new();
P.get_or_init(|| unsafe {
P.get_or_init(|| {
let ctx = context()?;
let lib = ctx.device
.newLibraryWithSource_options_error(&NSString::from_str(RESIDUAL_ADD_MSL), None)
@ -636,7 +636,7 @@ kernel void silu_mul_rows(
use objc2_metal::MTLLibrary as _;
static P: OnceLock<Option<Retained<ProtocolObject<dyn objc2_metal::MTLComputePipelineState>>>> =
OnceLock::new();
P.get_or_init(|| unsafe {
P.get_or_init(|| {
let ctx = context()?;
let lib = ctx.device
.newLibraryWithSource_options_error(&NSString::from_str(RMSNORM_ROWS_MSL), None).ok()?;
@ -651,7 +651,7 @@ kernel void silu_mul_rows(
use objc2_metal::MTLLibrary as _;
static P: OnceLock<Option<Retained<ProtocolObject<dyn objc2_metal::MTLComputePipelineState>>>> =
OnceLock::new();
P.get_or_init(|| unsafe {
P.get_or_init(|| {
let ctx = context()?;
let lib = ctx.device
.newLibraryWithSource_options_error(&NSString::from_str(ROPE_ROWS_MSL), None).ok()?;
@ -666,7 +666,7 @@ kernel void silu_mul_rows(
use objc2_metal::MTLLibrary as _;
static P: OnceLock<Option<Retained<ProtocolObject<dyn objc2_metal::MTLComputePipelineState>>>> =
OnceLock::new();
P.get_or_init(|| unsafe {
P.get_or_init(|| {
let ctx = context()?;
let lib = ctx.device
.newLibraryWithSource_options_error(&NSString::from_str(SILU_MUL_ROWS_MSL), None).ok()?;
@ -676,12 +676,10 @@ kernel void silu_mul_rows(
.as_ref()
}
// SAFETY: pipeline state objects are immutable + thread-safe.
struct PipelineCell(
Option<Retained<ProtocolObject<dyn objc2_metal::MTLComputePipelineState>>>,
);
unsafe impl Send for PipelineCell {}
unsafe impl Sync for PipelineCell {}
/// The `(buffer, offset, length)` triples a dispatch keeps alive until
/// its command buffer completes. Named because the inline form is three
/// generic layers deep and appears in two places.
type KeepAlive = Vec<(Retained<ProtocolObject<dyn MTLBuffer>>, usize, usize)>;
/// One MoE layer in ONE command buffer: per pick
/// `(gate, up, down)` resident F16 weights; gate/up GEMMs
@ -867,7 +865,7 @@ kernel void relu2_f32_to_f16(
static P: OnceLock<
Option<Retained<ProtocolObject<dyn objc2_metal::MTLComputePipelineState>>>,
> = OnceLock::new();
P.get_or_init(|| unsafe {
P.get_or_init(|| {
let ctx = context()?;
let src = NSString::from_str(RELU2_MSL);
let lib = ctx.device.newLibraryWithSource_options_error(&src, None).ok()?;
@ -1054,7 +1052,7 @@ kernel void rmsnorm_f32_to_f16(
static P: OnceLock<
Option<Retained<ProtocolObject<dyn objc2_metal::MTLComputePipelineState>>>,
> = OnceLock::new();
P.get_or_init(|| unsafe {
P.get_or_init(|| {
let ctx = context()?;
let src = NSString::from_str(RMSNORM_MSL);
let lib = ctx
@ -1105,7 +1103,7 @@ kernel void gemm_bf16_mpp(device bfloat* A_buf [[buffer(0)]],
static P: OnceLock<
Option<Retained<ProtocolObject<dyn objc2_metal::MTLComputePipelineState>>>,
> = OnceLock::new();
P.get_or_init(|| unsafe {
P.get_or_init(|| {
let ctx = context()?;
let src = NSString::from_str(MATMUL2D_MSL);
// Compile FAILS on pre-Metal-4 (the MPP header is
@ -1149,7 +1147,7 @@ kernel void gemm_bf16_mpp(device bfloat* A_buf [[buffer(0)]],
// C declared [N,M] (transposed) makes the example
// grid/slice correct for NON-SQUARE too. Tile dims:
// M%64, N%32, K%32 (≥ one tile each).
if m % 64 != 0 || m < 64 || n % 32 != 0 || n < 32 || k % 32 != 0 {
if !m.is_multiple_of(64) || m < 64 || !n.is_multiple_of(32) || n < 32 || !k.is_multiple_of(32) {
return Err(format!("matmul2d shape unsupported: M={m} N={n} K={k}"));
}
assert_eq!(a.len(), m * k);
@ -1407,7 +1405,7 @@ kernel void qknorm_rope(
static P: OnceLock<
Option<Retained<ProtocolObject<dyn objc2_metal::MTLComputePipelineState>>>,
> = OnceLock::new();
P.get_or_init(|| unsafe {
P.get_or_init(|| {
let ctx = context()?;
let src = NSString::from_str(QKNORM_ROPE_MSL);
let lib = ctx
@ -1539,7 +1537,7 @@ kernel void rope_rotate_half(
static P: OnceLock<
Option<Retained<ProtocolObject<dyn objc2_metal::MTLComputePipelineState>>>,
> = OnceLock::new();
P.get_or_init(|| unsafe {
P.get_or_init(|| {
let ctx = context()?;
let src = NSString::from_str(ROPE_ROTATE_HALF_MSL);
let lib = ctx
@ -1584,7 +1582,7 @@ kernel void scale_qk(
static P: OnceLock<
Option<Retained<ProtocolObject<dyn objc2_metal::MTLComputePipelineState>>>,
> = OnceLock::new();
P.get_or_init(|| unsafe {
P.get_or_init(|| {
let ctx = context()?;
let src = NSString::from_str(SCALE_QK_MSL);
let lib = ctx
@ -1606,7 +1604,7 @@ kernel void scale_qk(
static P: OnceLock<
Option<Retained<ProtocolObject<dyn objc2_metal::MTLComputePipelineState>>>,
> = OnceLock::new();
P.get_or_init(|| unsafe {
P.get_or_init(|| {
let ctx = context()?;
let src = NSString::from_str(ROPE_INTERLEAVED_MSL);
let lib = ctx
@ -1898,7 +1896,7 @@ kernel void attn_decode(
static P: OnceLock<
Option<Retained<ProtocolObject<dyn objc2_metal::MTLComputePipelineState>>>,
> = OnceLock::new();
P.get_or_init(|| unsafe {
P.get_or_init(|| {
let ctx = context()?;
let src = NSString::from_str(KV_APPEND_MSL);
let lib = ctx.device.newLibraryWithSource_options_error(&src, None).ok()?;
@ -1915,7 +1913,7 @@ kernel void attn_decode(
static P: OnceLock<
Option<Retained<ProtocolObject<dyn objc2_metal::MTLComputePipelineState>>>,
> = OnceLock::new();
P.get_or_init(|| unsafe {
P.get_or_init(|| {
let ctx = context()?;
let src = NSString::from_str(ATTN_DECODE_MSL);
let lib = ctx.device.newLibraryWithSource_options_error(&src, None).ok()?;
@ -1957,7 +1955,7 @@ kernel void attn_decode(
let cap = cap.min(4096); // threadgroup score scratch bound
let ctx = context()?;
let bytes = n_kv * cap * hd * 2;
unsafe {
{
let k = ctx
.device
.newBufferWithLength_options(bytes, MTLResourceOptions::StorageModeShared)?;
@ -1970,6 +1968,9 @@ kernel void attn_decode(
pub fn len(&self) -> usize {
self.len
}
pub fn is_empty(&self) -> bool {
self.len == 0
}
pub fn hd(&self) -> usize {
self.hd
}
@ -2062,7 +2063,6 @@ kernel void attn_decode(
Some(scale_qk_pipeline().ok_or("scale_qk pipeline unavailable")?)
}
};
let qk_pipe = qk_pipe.as_deref();
let append_pipe = kv_append_pipeline().ok_or("kv_append pipeline unavailable")?;
let attn_pipe = attn_decode_pipeline().ok_or("attn pipeline unavailable")?;
// Slot bookkeeping (CPU side, before encode).
@ -2612,7 +2612,7 @@ kernel void gemv_q8(device const uchar* W [[buffer(0)]],
static P: OnceLock<
Option<Retained<ProtocolObject<dyn objc2_metal::MTLComputePipelineState>>>,
> = OnceLock::new();
P.get_or_init(|| unsafe {
P.get_or_init(|| {
let ctx = context()?;
let lib = ctx
.device
@ -2642,7 +2642,7 @@ kernel void gemv_q8(device const uchar* W [[buffer(0)]],
) -> Result<(), String> {
assert_eq!(a.len(), m * k);
assert_eq!(c.len(), m * out_dim);
if k % 32 != 0 {
if !k.is_multiple_of(32) {
return Err(format!("k ({k}) not a multiple of 32"));
}
let nb = k / 32;
@ -2671,7 +2671,7 @@ kernel void gemv_q8(device const uchar* W [[buffer(0)]],
};
unsafe {
use objc2_metal::{
MTLCommandEncoder, MTLComputeCommandEncoder, MTLComputePipelineState as _, MTLSize,
MTLCommandEncoder, MTLComputeCommandEncoder, MTLSize,
};
let _g = ctx.encode_lock.lock().unwrap();
for i in 0..m {
@ -2819,7 +2819,7 @@ kernel void gemv_q4k(device const uchar* W [[buffer(0)]],
static P: OnceLock<
Option<Retained<ProtocolObject<dyn objc2_metal::MTLComputePipelineState>>>,
> = OnceLock::new();
P.get_or_init(|| unsafe {
P.get_or_init(|| {
let ctx = context()?;
let lib = ctx
.device
@ -2939,7 +2939,7 @@ kernel void gemm_q4k_m4(device const uchar* W [[buffer(0)]],
static P: OnceLock<
Option<Retained<ProtocolObject<dyn objc2_metal::MTLComputePipelineState>>>,
> = OnceLock::new();
P.get_or_init(|| unsafe {
P.get_or_init(|| {
let ctx = context()?;
let lib = ctx
.device
@ -3040,7 +3040,7 @@ kernel void sgemm_q4k(device const float* A [[buffer(0)]],
static P: OnceLock<
Option<Retained<ProtocolObject<dyn objc2_metal::MTLComputePipelineState>>>,
> = OnceLock::new();
P.get_or_init(|| unsafe {
P.get_or_init(|| {
let ctx = context()?;
let lib = ctx
.device
@ -3059,10 +3059,10 @@ kernel void sgemm_q4k(device const float* A [[buffer(0)]],
) -> Result<(), String> {
assert_eq!(a.len(), m * k);
assert_eq!(c.len(), m * out_dim);
if k % 256 != 0 {
if !k.is_multiple_of(256) {
return Err(format!("k ({k}) not a multiple of 256"));
}
if out_dim % 8 != 0 {
if !out_dim.is_multiple_of(8) {
return Err(format!("out_dim ({out_dim}) not a multiple of 8"));
}
let nb = k / 256;
@ -3095,7 +3095,7 @@ kernel void sgemm_q4k(device const float* A [[buffer(0)]],
let mp = m.div_ceil(8) * 8;
unsafe {
use objc2_metal::{
MTLCommandEncoder, MTLComputeCommandEncoder, MTLComputePipelineState as _, MTLSize,
MTLCommandEncoder, MTLComputeCommandEncoder, MTLSize,
};
// Zero-padded activation [mp, k].
let a_buf = ctx
@ -3212,7 +3212,7 @@ kernel void sgemm_q6k(device const float* A [[buffer(0)]],
static P: OnceLock<
Option<Retained<ProtocolObject<dyn objc2_metal::MTLComputePipelineState>>>,
> = OnceLock::new();
P.get_or_init(|| unsafe {
P.get_or_init(|| {
let ctx = context()?;
let lib = ctx
.device
@ -3230,10 +3230,10 @@ kernel void sgemm_q6k(device const float* A [[buffer(0)]],
) -> Result<(), String> {
assert_eq!(a.len(), m * k);
assert_eq!(c.len(), m * out_dim);
if k % 256 != 0 {
if !k.is_multiple_of(256) {
return Err(format!("k ({k}) not a multiple of 256"));
}
if out_dim % 8 != 0 {
if !out_dim.is_multiple_of(8) {
return Err(format!("out_dim ({out_dim}) not a multiple of 8"));
}
let nb = k / 256;
@ -3266,7 +3266,7 @@ kernel void sgemm_q6k(device const float* A [[buffer(0)]],
let mp = m.div_ceil(8) * 8;
unsafe {
use objc2_metal::{
MTLCommandEncoder, MTLComputeCommandEncoder, MTLComputePipelineState as _, MTLSize,
MTLCommandEncoder, MTLComputeCommandEncoder, MTLSize,
};
let a_buf = ctx
.device
@ -3314,7 +3314,7 @@ kernel void sgemm_q6k(device const float* A [[buffer(0)]],
) -> Result<(), String> {
assert_eq!(a.len(), m * k);
assert_eq!(c.len(), m * out_dim);
if k % 256 != 0 {
if !k.is_multiple_of(256) {
return Err(format!("k ({k}) not a multiple of 256"));
}
let nb = k / 256;
@ -3343,7 +3343,7 @@ kernel void sgemm_q6k(device const float* A [[buffer(0)]],
};
unsafe {
use objc2_metal::{
MTLCommandEncoder, MTLComputeCommandEncoder, MTLComputePipelineState as _, MTLSize,
MTLCommandEncoder, MTLComputeCommandEncoder, MTLSize,
};
let _g = ctx.encode_lock.lock().unwrap();
for i in 0..m {
@ -3400,7 +3400,7 @@ kernel void sgemm_q6k(device const float* A [[buffer(0)]],
) -> Result<(), String> {
assert_eq!(a.len(), m * k);
assert_eq!(c.len(), m * out_dim);
if k % 256 != 0 {
if !k.is_multiple_of(256) {
return Err(format!("k ({k}) not a multiple of 256"));
}
let nb = k / 256;
@ -3436,7 +3436,7 @@ kernel void sgemm_q6k(device const float* A [[buffer(0)]],
let n_tiles = m.div_ceil(TILE);
unsafe {
use objc2_metal::{
MTLCommandEncoder, MTLComputeCommandEncoder, MTLComputePipelineState as _, MTLSize,
MTLCommandEncoder, MTLComputeCommandEncoder, MTLSize,
};
let set_u32 = |enc: &ProtocolObject<dyn MTLComputeCommandEncoder>, v: u32, idx: usize| {
enc.setBytes_length_atIndex(
@ -3445,7 +3445,7 @@ kernel void sgemm_q6k(device const float* A [[buffer(0)]],
let _g = ctx.encode_lock.lock().unwrap();
let cb = ctx.queue.commandBuffer().ok_or("commandBuffer failed")?;
// Keep tile buffers alive until after the single wait.
let mut keep: Vec<(Retained<ProtocolObject<dyn MTLBuffer>>, usize, usize)> = Vec::new();
let mut keep: KeepAlive = Vec::new();
for t in 0..n_tiles {
let row0 = t * TILE;
let rows = (m - row0).min(TILE);
@ -3511,7 +3511,7 @@ kernel void sgemm_q6k(device const float* A [[buffer(0)]],
}
};
unsafe {
use objc2_metal::{MTLCommandEncoder, MTLComputeCommandEncoder, MTLComputePipelineState as _, MTLSize};
use objc2_metal::{MTLCommandEncoder, MTLComputeCommandEncoder, MTLSize};
let buf_a = ctx.device.newBufferWithBytes_length_options(
std::ptr::NonNull::new(a.as_ptr() as *mut _).unwrap(), k * 4,
MTLResourceOptions::StorageModeShared).ok_or("a")?;
@ -3631,7 +3631,7 @@ kernel void gemv_q6k(device const uchar* W [[buffer(0)]],
static P: OnceLock<
Option<Retained<ProtocolObject<dyn objc2_metal::MTLComputePipelineState>>>,
> = OnceLock::new();
P.get_or_init(|| unsafe {
P.get_or_init(|| {
let ctx = context()?;
let lib = ctx
.device
@ -3740,7 +3740,7 @@ kernel void gemm_q6k_m4(device const uchar* W [[buffer(0)]],
static P: OnceLock<
Option<Retained<ProtocolObject<dyn objc2_metal::MTLComputePipelineState>>>,
> = OnceLock::new();
P.get_or_init(|| unsafe {
P.get_or_init(|| {
let ctx = context()?;
let lib = ctx
.device
@ -3760,7 +3760,7 @@ kernel void gemm_q6k_m4(device const uchar* W [[buffer(0)]],
) -> Result<(), String> {
assert_eq!(a.len(), m * k);
assert_eq!(c.len(), m * out_dim);
if k % 256 != 0 {
if !k.is_multiple_of(256) {
return Err(format!("k ({k}) not a multiple of 256"));
}
let nb = k / 256;
@ -3796,7 +3796,7 @@ kernel void gemm_q6k_m4(device const uchar* W [[buffer(0)]],
let n_tiles = m.div_ceil(TILE);
unsafe {
use objc2_metal::{
MTLCommandEncoder, MTLComputeCommandEncoder, MTLComputePipelineState as _, MTLSize,
MTLCommandEncoder, MTLComputeCommandEncoder, MTLSize,
};
let set_u32 = |enc: &ProtocolObject<dyn MTLComputeCommandEncoder>, v: u32, idx: usize| {
enc.setBytes_length_atIndex(
@ -3804,7 +3804,7 @@ kernel void gemm_q6k_m4(device const uchar* W [[buffer(0)]],
};
let _g = ctx.encode_lock.lock().unwrap();
let cb = ctx.queue.commandBuffer().ok_or("commandBuffer failed")?;
let mut keep: Vec<(Retained<ProtocolObject<dyn MTLBuffer>>, usize, usize)> = Vec::new();
let mut keep: KeepAlive = Vec::new();
for t in 0..n_tiles {
let row0 = t * TILE;
let rows = (m - row0).min(TILE);
@ -3853,7 +3853,7 @@ kernel void gemm_q6k_m4(device const uchar* W [[buffer(0)]],
) -> Result<(), String> {
assert_eq!(a.len(), m * k);
assert_eq!(c.len(), m * out_dim);
if k % 256 != 0 {
if !k.is_multiple_of(256) {
return Err(format!("k ({k}) not a multiple of 256"));
}
let nb = k / 256;
@ -3882,7 +3882,7 @@ kernel void gemm_q6k_m4(device const uchar* W [[buffer(0)]],
};
unsafe {
use objc2_metal::{
MTLCommandEncoder, MTLComputeCommandEncoder, MTLComputePipelineState as _, MTLSize,
MTLCommandEncoder, MTLComputeCommandEncoder, MTLSize,
};
let _g = ctx.encode_lock.lock().unwrap();
for i in 0..m {
@ -3935,15 +3935,15 @@ kernel void gemm_q6k_m4(device const uchar* W [[buffer(0)]],
/// (gate/up Q4_K, down Q6_K). The gate/up/act intermediates stay
/// in GPU buffers across stages (no CPU round-trip); the 3 GEMVs
/// + activation collapse from 3 command buffers (3 commit+waits)
/// into one. `a` is post-norm `[h]`; gate/up Q4_K `[f·(h/256)·144]`,
/// down Q6_K `[h·(f/256)·210]`; `out` `[h]`. h,f multiples of 256.
/// into one. `a` is post-norm `[h]`; gate/up Q4_K `[f·(h/256)·144]`,
/// down Q6_K `[h·(f/256)·210]`; `out` `[h]`. h,f multiples of 256.
pub fn mlp_block_q4k_q6k(
a: &[f32], gate_w: &[u8], up_w: &[u8], down_w: &[u8],
h: usize, f: usize, out: &mut [f32],
) -> Result<(), String> {
assert_eq!(a.len(), h);
assert_eq!(out.len(), h);
if h % 256 != 0 || f % 256 != 0 {
if !h.is_multiple_of(256) || !f.is_multiple_of(256) {
return Err(format!("h ({h}) and f ({f}) must be multiples of 256"));
}
let nb_g = h / 256; // gate/up contraction blocks
@ -4063,7 +4063,7 @@ kernel void gemm_q6k_m4(device const uchar* W [[buffer(0)]],
) -> Result<(), String> {
assert_eq!(a.len(), h);
assert_eq!(out.len(), q_out + 2 * kv_out);
if h % 256 != 0 {
if !h.is_multiple_of(256) {
return Err(format!("h ({h}) must be a multiple of 256"));
}
let nb = h / 256;
@ -4091,7 +4091,7 @@ kernel void gemm_q6k_m4(device const uchar* W [[buffer(0)]],
let v_buf = upload(v_w)?;
unsafe {
use objc2_metal::{
MTLCommandEncoder, MTLComputeCommandEncoder, MTLComputePipelineState as _, MTLSize,
MTLCommandEncoder, MTLComputeCommandEncoder, MTLSize,
};
let buf_a = ctx.device.newBufferWithBytes_length_options(
std::ptr::NonNull::new(a.as_ptr() as *mut _).unwrap(), h * 4,
@ -4255,7 +4255,7 @@ kernel void attn_decode_f32(
static P: OnceLock<
Option<Retained<ProtocolObject<dyn objc2_metal::MTLComputePipelineState>>>,
> = OnceLock::new();
P.get_or_init(|| unsafe {
P.get_or_init(|| {
let ctx = context()?;
let src = NSString::from_str(RMSNORM_F32F32_MSL);
let lib = ctx.device.newLibraryWithSource_options_error(&src, None).ok()?;
@ -4272,7 +4272,7 @@ kernel void attn_decode_f32(
static P: OnceLock<
Option<Retained<ProtocolObject<dyn objc2_metal::MTLComputePipelineState>>>,
> = OnceLock::new();
P.get_or_init(|| unsafe {
P.get_or_init(|| {
let ctx = context()?;
let src = NSString::from_str(ATTN_DECODE_F32_MSL);
let lib = ctx.device.newLibraryWithSource_options_error(&src, None).ok()?;
@ -4361,7 +4361,7 @@ kernel void batched_causal_attn_f32(
static P: OnceLock<
Option<Retained<ProtocolObject<dyn objc2_metal::MTLComputePipelineState>>>,
> = OnceLock::new();
P.get_or_init(|| unsafe {
P.get_or_init(|| {
let ctx = context()?;
let src = NSString::from_str(BATCHED_CAUSAL_ATTN_MSL);
let lib = ctx.device.newLibraryWithSource_options_error(&src, None).ok()?;
@ -4391,7 +4391,7 @@ kernel void batched_causal_attn_f32(
let pipe = batched_causal_attn_pipeline().ok_or("batched_causal_attn pipeline")?;
unsafe {
use objc2_metal::{
MTLCommandEncoder, MTLComputeCommandEncoder, MTLComputePipelineState as _, MTLSize,
MTLCommandEncoder, MTLComputeCommandEncoder, MTLSize,
};
let up = |data: &[f32]| {
mctx.device.newBufferWithBytes_length_options(
@ -4466,11 +4466,11 @@ kernel void batched_causal_attn_f32(
assert_eq!(k_out.len(), n_kv * hd);
assert_eq!(v_out.len(), n_kv * hd);
assert_eq!(o_out.len(), k_dim);
if k_dim % 256 != 0 {
if !k_dim.is_multiple_of(256) {
return Err(format!("k_dim ({k_dim}) must be a multiple of 256"));
}
let o_in = n_q * hd;
if o_in % 256 != 0 {
if !o_in.is_multiple_of(256) {
return Err(format!("n_q*hd ({o_in}) must be a multiple of 256"));
}
let nb_in = k_dim / 256; // 256-blocks per input row for q/k/v
@ -4519,7 +4519,6 @@ kernel void batched_causal_attn_f32(
Some(scale_qk_pipeline().ok_or("scale_qk pipeline unavailable")?)
}
};
let qk_pipe = qk_pipe.as_deref();
let append_pipe = kv_append_pipeline().ok_or("kv_append pipeline unavailable")?;
let attn_pipe = attn_decode_f32_pipeline().ok_or("attn pipeline unavailable")?;
@ -4562,7 +4561,7 @@ kernel void batched_causal_attn_f32(
unsafe {
use objc2_metal::{
MTLCommandEncoder, MTLComputeCommandEncoder, MTLComputePipelineState as _, MTLSize,
MTLCommandEncoder, MTLComputeCommandEncoder, MTLSize,
};
let upload = |data: *const u8, len: usize, what: &str| {
ctx.device
@ -4843,11 +4842,11 @@ kernel void batched_causal_attn_f32(
assert_eq!(ffn_norm.len(), k_dim);
assert_eq!(hd, kv.hd);
assert_eq!(n_kv, kv.n_kv);
if k_dim % 256 != 0 || ffn % 256 != 0 {
if !k_dim.is_multiple_of(256) || !ffn.is_multiple_of(256) {
return Err(format!("k_dim ({k_dim}) and ffn ({ffn}) must be multiples of 256"));
}
let o_in = n_q * hd;
if o_in % 256 != 0 {
if !o_in.is_multiple_of(256) {
return Err(format!("n_q*hd ({o_in}) must be a multiple of 256"));
}
let nb_in = k_dim / 256; // input blocks for q/k/v + gate/up
@ -4895,7 +4894,6 @@ kernel void batched_causal_attn_f32(
}
QkStage::ScaleQk { .. } => Some(scale_qk_pipeline().ok_or("scale_qk pipeline")?),
};
let qk_pipe = qk_pipe.as_deref();
let append_pipe = kv_append_pipeline().ok_or("kv_append pipeline")?;
let attn_pipe = attn_decode_f32_pipeline().ok_or("attn pipeline")?;
@ -5238,7 +5236,7 @@ kernel void batched_causal_attn_f32(
let kd = n_kv * hd;
assert_eq!(k_out.len(), m * kd);
assert_eq!(v_out.len(), m * kd);
if hidden % 256 != 0 || ffn % 256 != 0 || qd % 256 != 0 {
if !hidden.is_multiple_of(256) || !ffn.is_multiple_of(256) || !qd.is_multiple_of(256) {
return Err("hidden, ffn, n_q*hd must be multiples of 256".into());
}
if m == 0 || m > 4096 {
@ -5293,7 +5291,7 @@ kernel void batched_causal_attn_f32(
let mp = m.div_ceil(8) * 8; // padded row count for the 8×8 tiled GEMM
unsafe {
use objc2_metal::{
MTLCommandEncoder, MTLComputeCommandEncoder, MTLComputePipelineState as _, MTLSize,
MTLCommandEncoder, MTLComputeCommandEncoder, MTLSize,
};
let alloc = |elems: usize| {
ctx.device.newBufferWithLength_options(elems * 4, MTLResourceOptions::StorageModeShared)
@ -6457,7 +6455,7 @@ mod f16_tests {
};
let mut got = vec![0f32; m * n];
super::gemm_bt_resident_f16(&a_f, &weight, m, &mut got).unwrap();
let bf = |v: f32| f32::from_bits(((v.to_bits() >> 16) as u32) << 16);
let bf = |v: f32| f32::from_bits((v.to_bits() >> 16) << 16);
for i in 0..m {
for j in 0..n {
let mut acc = 0f32;
@ -6534,8 +6532,8 @@ mod f16_tests {
(-2.0, 0xc000),
(65504.0, 0x7bff),
// Subnormals: 2⁻¹⁵ → man 512; 2⁻²⁴ → man 1.
(3.0517578125e-5, 0x0200),
(5.9604644775390625e-8, 0x0001),
(3.051_757_8e-5, 0x0200),
(5.960_464_5e-8, 0x0001),
// Mid-subnormal: 1e-6 ≈ 16.78 ULP → 17.
(1e-6, 0x0011),
// Below half the smallest subnormal → zero.
@ -6553,7 +6551,7 @@ mod f16_tests {
#[test]
fn attn_block_matches_cpu_reference() {
let to_bf16 = |v: f32| -> u16 { (v.to_bits() >> 16) as u16 };
let bf = |v: f32| f32::from_bits(((v.to_bits() >> 16) as u32) << 16);
let bf = |v: f32| f32::from_bits((v.to_bits() >> 16) << 16);
let (kd, hd, n_qh, n_kvh, cap, steps) = (32usize, 8usize, 4usize, 2usize, 3usize, 5usize);
let (nq, nkv) = (n_qh * hd, n_kvh * hd);
let nw: Vec<f32> = (0..kd).map(|i| 0.8 + ((i % 5) as f32) * 0.1).collect();
@ -6688,7 +6686,7 @@ mod f16_tests {
#[test]
fn attn_block_interleaved_matches_cpu_reference() {
let to_bf16 = |v: f32| -> u16 { (v.to_bits() >> 16) as u16 };
let bf = |v: f32| f32::from_bits(((v.to_bits() >> 16) as u32) << 16);
let bf = |v: f32| f32::from_bits((v.to_bits() >> 16) << 16);
let (kd, hd, n_qh, n_kvh, cap, steps) = (32usize, 8usize, 4usize, 2usize, 3usize, 5usize);
let (nq, nkv) = (n_qh * hd, n_kvh * hd);
let nw: Vec<f32> = (0..kd).map(|i| 0.8 + ((i % 5) as f32) * 0.1).collect();
@ -6814,7 +6812,7 @@ mod f16_tests {
#[test]
fn attn_block_rotate_half_matches_cpu_reference() {
let to_bf16 = |v: f32| -> u16 { (v.to_bits() >> 16) as u16 };
let bf = |v: f32| f32::from_bits(((v.to_bits() >> 16) as u32) << 16);
let bf = |v: f32| f32::from_bits((v.to_bits() >> 16) << 16);
let (kd, hd, n_qh, n_kvh, cap, steps) = (32usize, 8usize, 4usize, 2usize, 3usize, 5usize);
let (nq, nkv) = (n_qh * hd, n_kvh * hd);
let nw: Vec<f32> = (0..kd).map(|i| 0.8 + ((i % 5) as f32) * 0.1).collect();
@ -6981,7 +6979,7 @@ mod f16_tests {
#[test]
fn attn_block_scale_qk_matches_cpu_reference() {
let to_bf16 = |v: f32| -> u16 { (v.to_bits() >> 16) as u16 };
let bf = |v: f32| f32::from_bits(((v.to_bits() >> 16) as u32) << 16);
let bf = |v: f32| f32::from_bits((v.to_bits() >> 16) << 16);
let (kd, hd, n_qh, n_kvh, cap, steps) = (32usize, 8usize, 4usize, 2usize, 3usize, 5usize);
let (nq, nkv) = (n_qh * hd, n_kvh * hd);
let nw: Vec<f32> = (0..kd).map(|i| 0.8 + ((i % 5) as f32) * 0.1).collect();
@ -7083,7 +7081,7 @@ mod f16_tests {
#[test]
fn norm_qkv_rope_window_matches_cpu_reference() {
let to_bf16 = |v: f32| -> u16 { (v.to_bits() >> 16) as u16 };
let bf = |v: f32| f32::from_bits(((v.to_bits() >> 16) as u32) << 16);
let bf = |v: f32| f32::from_bits((v.to_bits() >> 16) << 16);
let (kd, hd, n_qh, n_kvh) = (64usize, 8usize, 4usize, 2usize);
let (nq, nkv) = (n_qh * hd, n_kvh * hd);
let h: Vec<f32> = (0..kd).map(|i| ((i % 13) as f32 - 6.0) * 0.21).collect();
@ -7171,7 +7169,7 @@ mod f16_tests {
#[test]
fn norm_gemm_window_matches_cpu_reference() {
let to_bf16 = |v: f32| -> u16 { (v.to_bits() >> 16) as u16 };
let bf = |v: f32| f32::from_bits(((v.to_bits() >> 16) as u32) << 16);
let bf = |v: f32| f32::from_bits((v.to_bits() >> 16) << 16);
let (k, n1, n2) = (64usize, 24usize, 16usize);
let h: Vec<f32> = (0..k).map(|i| ((i % 13) as f32 - 6.0) * 0.21).collect();
let nw: Vec<f32> = (0..k).map(|i| 0.8 + ((i % 5) as f32) * 0.1).collect();
@ -7221,7 +7219,7 @@ mod f16_tests {
#[test]
fn moe_window_matches_cpu_reference() {
let to_bf16 = |v: f32| -> u16 { (v.to_bits() >> 16) as u16 };
let bf = |v: f32| f32::from_bits(((v.to_bits() >> 16) as u32) << 16);
let bf = |v: f32| f32::from_bits((v.to_bits() >> 16) << 16);
let (k, f, h, n_p) = (32usize, 24usize, 16usize, 3usize);
let a: Vec<f32> = (0..k).map(|i| ((i % 11) as f32 - 5.0) * 0.09).collect();
let mk_w = |seed: usize, rows: usize, cols: usize| -> (Vec<f32>, Vec<u8>) {