Docs / Fused and custom Metal kernels

Fused and custom Metal kernels

nuvec's fast module: fused RMSNorm, LayerNorm, RoPE, attention, cross entropy and Gated DeltaNet kernels, plus custom Metal kernels written from source.

Fused kernels (nuvec::fast)🔗

The fast module exposes MLX’s fused operations and a few of nuvec’s own. Each one is a single Metal kernel on the GPU and an equivalent composition of ordinary operations elsewhere; the composition also defines its derivatives and batching, so they work under every transform.

The Gated DeltaNet kernels pick mlx-lm’s packed kernel where it applies: key head size 128, value head size a multiple of 8, float32 gates and state. MLX_GDN_PACKED=0 keeps every shape on the generic kernel, as in mlx-lm, so the two paths can be compared in one build.

Attention picks MLX’s kernel for the shape: vector kernels for short queries (decode), a two-pass split-K kernel for long contexts, and the tiled flash-style kernel or the composition for prefill; at head size 512, long causal prefill in half precision runs MLX’s blocked path (256 queries at a time over only the keys they see). A sliding window skips the keys outside it in every kernel. plan_attention reports the choice, and the scratch memory it holds, from the shapes alone; max_score_bytes bounds every score-sized intermediate (the tiled kernel, or the composition in blocks of queries). With 2 to 8 causal queries (speculative verification), each row gives the bits of the same row decoded alone. softmax and logsumexp over the last axis use fused kernels too.

use nuvec::fast::{AttentionOptions, plan_attention};
use nuvec::DType;

# fn main() -> nuvec::Result<()> {
let opts = AttentionOptions::causal().max_score_bytes(512 << 20);
let plan = plan_attention(&[1, 16, 1024, 256], &[1, 2, 33_792, 256], &[1, 2, 33_792, 256], DType::BF16, &opts)?;
assert!(plan.max_intermediate_bytes <= 512 << 20);
# Ok(())
# }

Custom kernels🔗

MetalKernel runs your own kernel body, written in Metal Shading Language, the way MLX’s mx.fast.metal_kernel does. You give a name, the input and output names and the body; nuvec generates the function signature:

use nuvec::prelude::*;
use nuvec::fast::{KernelConfig, MetalKernel, TemplateArg};

# fn main() -> Result<()> {
# if !nuvec::metal::is_available() { return Ok(()); } // needs a Metal GPU
# let n = 1024;
# let x = nuvec::ops::ones([n], DType::F16);

let scale = MetalKernel::new(
    "scale",
    &["inp"],
    &["out"],
    "uint i = thread_position_in_grid.x;\nout[i] = static_cast<T>(inp[i] * N);",
);
let config = KernelConfig {
    grid: [n as usize, 1, 1],
    threadgroup: [256, 1, 1],
    template: vec![("T".into(), TemplateArg::DType(DType::F16)), ("N".into(), TemplateArg::Int(3))],
    ..Default::default()
};
let y = scale.call(&[x], &[(&[n], DType::F16)], &config).remove(0);
# assert_eq!(y.astype(DType::F32).sum(..).item::<f32>()?, 3072.0);
# Ok(())
# }

Signature rules, as in MLX:

• an input is const device T* name; a 0-d input is const constant T& name;
• name_shape, name_strides and name_ndim are added when the body uses them;
• an output is device T* name, or device atomic<T>* name with atomic outputs (optionally initialized);
• thread attributes the body mentions (thread_position_in_grid, thread_index_in_simdgroup, …) become parameters;
• template arguments become constexpr values and type aliases; MLX’s bfloat16_t, float16_t and elem_to_loc are defined.

Kernels compile once per generated source and run with one thread per grid position. They have no derivative or batching rule of their own: wrap them in custom_function to give them one. mlx-lm’s own Gated DeltaNet kernel source runs unchanged through MetalKernel.

✏️ Edit this page on GitHub