From dc1332548294be616484648b558b5c5e4bbfb3f6 Mon Sep 17 00:00:00 2001 From: tpegolotti Date: Wed, 10 Jun 2026 13:22:58 +0200 Subject: [PATCH 01/56] add basic skeleton for gated delta update forward --- local_files/gated_basic.py | 0 mlx/backend/metal/CMakeLists.txt | 1 + mlx/backend/metal/gated_delta_update.cpp | 50 ++++++++++++++++++++++++ mlx/fast.cpp | 44 +++++++++++++++++++++ mlx/fast.h | 10 +++++ mlx/fast_primitives.h | 37 ++++++++++++++++++ python/src/fast.cpp | 26 ++++++++++++ 7 files changed, 168 insertions(+) create mode 100644 local_files/gated_basic.py create mode 100644 mlx/backend/metal/gated_delta_update.cpp diff --git a/local_files/gated_basic.py b/local_files/gated_basic.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/mlx/backend/metal/CMakeLists.txt b/mlx/backend/metal/CMakeLists.txt index e7a4d9d2af..0ba80646a8 100644 --- a/mlx/backend/metal/CMakeLists.txt +++ b/mlx/backend/metal/CMakeLists.txt @@ -135,6 +135,7 @@ target_sources( ${CMAKE_CURRENT_SOURCE_DIR}/logsumexp.cpp ${CMAKE_CURRENT_SOURCE_DIR}/matmul.cpp ${CMAKE_CURRENT_SOURCE_DIR}/scaled_dot_product_attention.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/gated_delta_update.cpp ${CMAKE_CURRENT_SOURCE_DIR}/metal.cpp ${CMAKE_CURRENT_SOURCE_DIR}/primitives.cpp ${CMAKE_CURRENT_SOURCE_DIR}/quantized.cpp diff --git a/mlx/backend/metal/gated_delta_update.cpp b/mlx/backend/metal/gated_delta_update.cpp new file mode 100644 index 0000000000..77e1b3d3ba --- /dev/null +++ b/mlx/backend/metal/gated_delta_update.cpp @@ -0,0 +1,50 @@ +// Copyright © 2024 Apple Inc. +#include + +#include "mlx/backend/common/compiled.h" +#include "mlx/backend/gpu/copy.h" +#include "mlx/backend/metal/device.h" +#include "mlx/backend/metal/kernels.h" +#include "mlx/backend/metal/kernels/defines.h" +#include "mlx/backend/metal/utils.h" +#include "mlx/fast_primitives.h" +#include "mlx/utils.h" + +namespace mlx::core::fast { + +bool GatedDeltaUpdate::use_fallback(Stream s) { + // always run on GPU for now + return false; +} + +void GatedDeltaUpdate::eval_gpu( + const std::vector& inputs, + std::vector& outputs) { + + auto& s = stream(); + auto& d = metal::device(s.device); + + auto& q = inputs[0]; + auto& k = inputs[1]; + auto& v = inputs[2]; + auto& g = inputs[3]; + auto& beta = inputs[4]; + auto& h0 = inputs[5]; + + auto& out = outputs[0]; + auto& hf = outputs[1]; + + // TODO: allocate outputs, dispatch Metal kernel + throw std::runtime_error("NYI"); +} + +bool GatedDeltaUpdate::is_equivalent(const Primitive& other) const { + const auto* p = dynamic_cast(&other); + if (p == nullptr) { + return false; + } + // TODO: compare chunk_size and other state fields once added + return true; +} + +} // namespace mlx::core::fast diff --git a/mlx/fast.cpp b/mlx/fast.cpp index a668fe9abd..4a0856e850 100644 --- a/mlx/fast.cpp +++ b/mlx/fast.cpp @@ -922,6 +922,50 @@ bool ScaledDotProductAttentionVJP::is_equivalent(const Primitive& other) const { has_sinks_ == a_other.has_sinks_; } +std::vector gated_delta_update_forward( + const array& queries, + const array& keys, + const array& values, + const array& gates, + const array& beta_, + const std::optional& initial_state, /* = std::nullopt */ + StreamOrDevice s /* = {} */) { + + // determine output dtype + auto promoted = promote_types(queries.dtype(), keys.dtype()); + auto out_dtype = issubdtype(promoted, float32) + ? promoted + : promote_types(promoted, float32); + + // cast all inputs + auto q = astype(queries, out_dtype, s); + auto k = astype(keys, out_dtype, s); + auto v = astype(values, out_dtype, s); + auto g = astype(gates, out_dtype, s); + auto beta = astype(beta_, out_dtype, s); + + int B = q.shape(0), H = q.shape(1); + int T = q.shape(2), Dk = q.shape(3); + int Dv = v.shape(3); + + auto h0 = initial_state.has_value() + ? astype(*initial_state, out_dtype, s) + : zeros({B, H, Dk, Dv}, out_dtype, s); + + auto fallback = [](std::vector inputs) -> std::vector { + throw std::runtime_error("NYI: GatedDeltaUpdate CPU fallback"); + return {}; + }; + + return array::make_arrays( + /* output shapes */ {{B, H, T, Dv}, {B, H, Dk, Dv}}, + /* dtypes */ {out_dtype, out_dtype}, + /* primitive */ std::make_shared(to_stream(s), fallback), + /* inputs */ {q, k, v, g, beta, h0} + ); +} + + bool Quantize::is_equivalent(const Primitive& other) const { const Quantize& p_other = static_cast(other); return ( diff --git a/mlx/fast.h b/mlx/fast.h index 934fadc2b7..796b159292 100644 --- a/mlx/fast.h +++ b/mlx/fast.h @@ -55,6 +55,16 @@ MLX_API array scaled_dot_product_attention( const std::optional& sinks = {}, StreamOrDevice s = {}); +MLX_API std::vector gated_delta_update_forward( + const array& queries, + const array& keys, + const array& values, + const array& gates, + const array& beta_, + const std::optional& initial_state = std::nullopt, + StreamOrDevice s = {} +); + using TemplateArg = std::variant; using ScalarArg = std::variant; diff --git a/mlx/fast_primitives.h b/mlx/fast_primitives.h index 0d2f861045..5a4609d44b 100644 --- a/mlx/fast_primitives.h +++ b/mlx/fast_primitives.h @@ -325,6 +325,43 @@ class ConvertFP8 : public Primitive { bool to_fp8_; }; +class GatedDeltaUpdate : public Custom { + public: + GatedDeltaUpdate( + Stream stream, + std::function(std::vector)> fallback + ) + : Custom(stream, std::move(fallback)) + {} + + static bool use_fallback( + /* TODO */ + Stream s); + static bool supports_bool_mask(); + + void eval_cpu(const std::vector& inputs, std::vector& outputs) + override { + throw std::runtime_error("NYI"); + } + + void eval_gpu(const std::vector& inputs, std::vector& outputs) + override; + + + bool is_equivalent(const Primitive& other) const override; + + DEFINE_NAME(GatedDeltaUpdate); + DEFINE_INPUT_OUTPUT_SHAPE() + auto state() const { + return std::make_tuple( + nullptr); /* TODO */ + } + + private: + /* TODO */ +}; + + class Quantize : public Custom { public: explicit Quantize( diff --git a/python/src/fast.cpp b/python/src/fast.cpp index cd30b0bacd..94d48b0f5d 100644 --- a/python/src/fast.cpp +++ b/python/src/fast.cpp @@ -332,6 +332,32 @@ void init_fast(nb::module_& parent_module) { scale = D ** -0.5 out = mx.fast.scaled_dot_product_attention(q, k, v, scale=scale, mask="causal") )pbdoc"); + + m.def( + "gated_delta_update_forward", + &mlx::core::fast::gated_delta_update_forward, + "q"_a, + "k"_a, + "v"_a, + "gates"_a, + "beta"_a, + "initial_state"_a = nb::none(), // optional, defaults to None + "stream"_a = nb::none(), // optional, defaults to None + R"( + Chunked gated delta network forward pass. + + Args: + q: Queries [B, H, T, Dk] + k: Keys [B, H, T, Dk] + v: Values [B, H, T, Dv] + gates: Log-decay gates [B, H, T] + beta: Delta update rates [B, H, T] + initial_state: Optional initial hidden state [B, H, Dk, Dv] + + Returns: + Tuple of (output [B, H, T, Dv], final_state [B, H, Dk, Dv]) + )" + ); m.def( "metal_kernel", From 6f35421e9a738a9d95aea0899cc8ab2c83b38e88 Mon Sep 17 00:00:00 2001 From: tpegolotti Date: Wed, 10 Jun 2026 16:18:07 +0200 Subject: [PATCH 02/56] add skeleton for metal kernel - compile and run ok --- mlx/backend/metal/gated_delta_update.cpp | 32 ++++++++++++++++--- mlx/backend/metal/kernels.h | 8 +++++ mlx/backend/metal/kernels/CMakeLists.txt | 1 + .../metal/kernels/gated_delta_update.metal | 19 +++++++++++ .../metal/kernels/gated_delta_update_impl.h | 24 ++++++++++++++ mlx/backend/metal/nojit_kernels.cpp | 10 ++++++ 6 files changed, 90 insertions(+), 4 deletions(-) create mode 100644 mlx/backend/metal/kernels/gated_delta_update.metal create mode 100644 mlx/backend/metal/kernels/gated_delta_update_impl.h diff --git a/mlx/backend/metal/gated_delta_update.cpp b/mlx/backend/metal/gated_delta_update.cpp index 77e1b3d3ba..0c95b2ffdd 100644 --- a/mlx/backend/metal/gated_delta_update.cpp +++ b/mlx/backend/metal/gated_delta_update.cpp @@ -12,8 +12,29 @@ namespace mlx::core::fast { +namespace { + +void gated_delta_update_forward_metal( + const Stream& s, + metal::Device& d){ + + std::string base_name = "gated_delta_update_fwd_float_64_64"; + std::string hash_name = base_name; + metal::MTLFCList func_consts = {}; + + auto kernel = get_steel_gated_delta_forward_kernel( + d, + base_name, + hash_name, + func_consts); +} + + + +} + bool GatedDeltaUpdate::use_fallback(Stream s) { - // always run on GPU for now + // TODO: finish implementation return false; } @@ -34,8 +55,11 @@ void GatedDeltaUpdate::eval_gpu( auto& out = outputs[0]; auto& hf = outputs[1]; - // TODO: allocate outputs, dispatch Metal kernel - throw std::runtime_error("NYI"); + out.set_data(allocator::malloc(out.nbytes())); + hf.set_data(allocator::malloc(hf.nbytes())); + + gated_delta_update_forward_metal(s,d); + // throw std::runtime_error("NYI"); } bool GatedDeltaUpdate::is_equivalent(const Primitive& other) const { @@ -43,7 +67,7 @@ bool GatedDeltaUpdate::is_equivalent(const Primitive& other) const { if (p == nullptr) { return false; } - // TODO: compare chunk_size and other state fields once added + // TODO: finish implementation return true; } diff --git a/mlx/backend/metal/kernels.h b/mlx/backend/metal/kernels.h index 973041932a..7fb8e3be80 100644 --- a/mlx/backend/metal/kernels.h +++ b/mlx/backend/metal/kernels.h @@ -417,6 +417,14 @@ MTL::ComputePipelineState* get_steel_attention_nax_kernel( int wn, const array& m); +MTL::ComputePipelineState* get_steel_gated_delta_forward_kernel( + metal::Device& d, + const std::string& kernel_name, + const std::string& hash_name, + const metal::MTLFCList& func_consts + // TODO: More parameters? + ); + // Create a GPU kernel template definition for JIT compilation template std::string get_template_definition( diff --git a/mlx/backend/metal/kernels/CMakeLists.txt b/mlx/backend/metal/kernels/CMakeLists.txt index 6d9a0883f0..5f70ae9b0d 100644 --- a/mlx/backend/metal/kernels/CMakeLists.txt +++ b/mlx/backend/metal/kernels/CMakeLists.txt @@ -54,6 +54,7 @@ build_kernel(random) build_kernel(rms_norm) build_kernel(rope) build_kernel(scaled_dot_product_attention sdpa_vector.h) +build_kernel(gated_delta_update gated_delta_update_impl.h) if(MLX_METAL_VERSION GREATER_EQUAL 320) build_kernel(fence) endif() diff --git a/mlx/backend/metal/kernels/gated_delta_update.metal b/mlx/backend/metal/kernels/gated_delta_update.metal new file mode 100644 index 0000000000..909b561ce7 --- /dev/null +++ b/mlx/backend/metal/kernels/gated_delta_update.metal @@ -0,0 +1,19 @@ +// mlx/backend/metal/kernels/gated_delta_update.metal + +#include "mlx/backend/metal/kernels/utils.h" +#include "mlx/backend/metal/kernels/gated_delta_update_impl.h" + +using namespace metal; + +#define instantiate_gated_delta_update(type, dk, dv) \ + instantiate_kernel( \ + "gated_delta_update_fwd_" #type "_" #dk "_" #dv, \ + gated_delta_update_fwd, \ + type, \ + dk, \ + dv) + +#define instantiate_gated_delta_update_dims(type) \ + instantiate_gated_delta_update(type, 64, 64) + +instantiate_gated_delta_update_dims(float) diff --git a/mlx/backend/metal/kernels/gated_delta_update_impl.h b/mlx/backend/metal/kernels/gated_delta_update_impl.h new file mode 100644 index 0000000000..b9459a9550 --- /dev/null +++ b/mlx/backend/metal/kernels/gated_delta_update_impl.h @@ -0,0 +1,24 @@ +#pragma once + +#include +#include "mlx/backend/metal/kernels/utils.h" + +using namespace metal; + +template +[[kernel]] void gated_delta_update_fwd( + const device T* q [[buffer(0)]], + const device T* k [[buffer(1)]], + const device T* v [[buffer(2)]], + const device T* g [[buffer(3)]], + const device T* beta [[buffer(4)]], + const device T* h0 [[buffer(5)]], + device T* out [[buffer(6)]], + device T* hf [[buffer(7)]], + constant int& B [[buffer(8)]], + constant int& H [[buffer(9)]], + constant int& T_len [[buffer(10)]], + uint3 tid [[thread_position_in_grid]] +) { + // kernel implementation +} diff --git a/mlx/backend/metal/nojit_kernels.cpp b/mlx/backend/metal/nojit_kernels.cpp index 9f6f8782f5..04c948f681 100644 --- a/mlx/backend/metal/nojit_kernels.cpp +++ b/mlx/backend/metal/nojit_kernels.cpp @@ -492,4 +492,14 @@ MTL::ComputePipelineState* get_steel_attention_nax_kernel( return d.get_kernel(kernel_name, hash_name, func_consts); } +MTL::ComputePipelineState* get_steel_gated_delta_forward_kernel( + metal::Device& d, + const std::string& kernel_name, + const std::string& hash_name, + const metal::MTLFCList& func_consts + // TODO: need more parameters? + ) { + return d.get_kernel(kernel_name, hash_name, func_consts); +} + } // namespace mlx::core From 8e513973f722d43063fa6280d0208c2a687b3144 Mon Sep 17 00:00:00 2001 From: tpegolotti Date: Thu, 11 Jun 2026 12:10:46 +0200 Subject: [PATCH 03/56] copied implementation over from mlx_lm as baseline --- local_files/references.py | 0 mlx/backend/metal/gated_delta_update.cpp | 69 +++++++++----- .../metal/kernels/gated_delta_update.metal | 25 ++--- .../metal/kernels/gated_delta_update_impl.h | 92 ++++++++++++++++--- mlx/fast.cpp | 11 ++- 5 files changed, 143 insertions(+), 54 deletions(-) create mode 100644 local_files/references.py diff --git a/local_files/references.py b/local_files/references.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/mlx/backend/metal/gated_delta_update.cpp b/mlx/backend/metal/gated_delta_update.cpp index 0c95b2ffdd..2c9797696b 100644 --- a/mlx/backend/metal/gated_delta_update.cpp +++ b/mlx/backend/metal/gated_delta_update.cpp @@ -12,27 +12,6 @@ namespace mlx::core::fast { -namespace { - -void gated_delta_update_forward_metal( - const Stream& s, - metal::Device& d){ - - std::string base_name = "gated_delta_update_fwd_float_64_64"; - std::string hash_name = base_name; - metal::MTLFCList func_consts = {}; - - auto kernel = get_steel_gated_delta_forward_kernel( - d, - base_name, - hash_name, - func_consts); -} - - - -} - bool GatedDeltaUpdate::use_fallback(Stream s) { // TODO: finish implementation return false; @@ -54,12 +33,54 @@ void GatedDeltaUpdate::eval_gpu( auto& out = outputs[0]; auto& hf = outputs[1]; + + int B = q.shape(0); + int Hk = q.shape(1); // key heads + int T = q.shape(2); + int Dk = q.shape(3); + int Hv = v.shape(1); // value heads + int Dv = v.shape(3); out.set_data(allocator::malloc(out.nbytes())); hf.set_data(allocator::malloc(hf.nbytes())); - - gated_delta_update_forward_metal(s,d); - // throw std::runtime_error("NYI"); + + std::string base_name = "gated_delta_step_" + + get_type_string(q.dtype()) // "float" + + "_" + get_type_string(h0.dtype()) // "float" + + "_" + std::to_string(Dk) + + "_" + std::to_string(Dv) + + "_" + std::to_string(Hk) + + "_" + std::to_string(Hv); + // e.g. "gated_delta_step_float_float_64_64_4_4" + std::string hash_name = base_name; + metal::MTLFCList func_consts = {}; + + auto kernel = get_steel_gated_delta_forward_kernel( + d, + base_name, + hash_name, + func_consts); + + auto& compute_encoder = metal::get_command_encoder(s); + + compute_encoder.set_compute_pipeline_state(kernel); + + compute_encoder.set_input_array(q, 0); + compute_encoder.set_input_array(k, 1); + compute_encoder.set_input_array(v, 2); + compute_encoder.set_input_array(g, 3); + compute_encoder.set_input_array(beta, 4); + compute_encoder.set_input_array(h0, 5); + compute_encoder.set_bytes(T, 6); + compute_encoder.set_output_array(out, 7); + compute_encoder.set_output_array(hf, 8); + + // auto grid = MTL::Size(1, 1, 1); + // auto threads = MTL::Size(1, 1, 1); + auto grid = MTL::Size(32, Dv, B * Hv); + auto threads = MTL::Size(32, 4, 1); + compute_encoder.dispatch_threads(grid, threads); + // throw std::runtime_error("NYI"); } bool GatedDeltaUpdate::is_equivalent(const Primitive& other) const { diff --git a/mlx/backend/metal/kernels/gated_delta_update.metal b/mlx/backend/metal/kernels/gated_delta_update.metal index 909b561ce7..a76576ebd7 100644 --- a/mlx/backend/metal/kernels/gated_delta_update.metal +++ b/mlx/backend/metal/kernels/gated_delta_update.metal @@ -1,19 +1,20 @@ -// mlx/backend/metal/kernels/gated_delta_update.metal - #include "mlx/backend/metal/kernels/utils.h" #include "mlx/backend/metal/kernels/gated_delta_update_impl.h" using namespace metal; -#define instantiate_gated_delta_update(type, dk, dv) \ - instantiate_kernel( \ - "gated_delta_update_fwd_" #type "_" #dk "_" #dv, \ - gated_delta_update_fwd, \ - type, \ - dk, \ - dv) +#define instantiate_gated_delta_update(in_type, st_type, dk, dv, hk, hv) \ + instantiate_kernel( \ + "gated_delta_step_" #in_type "_" #st_type "_" #dk "_" #dv "_" #hk "_" #hv, \ + gated_delta_step, \ + in_type, \ + st_type, \ + dk, \ + dv, \ + hk, \ + hv) -#define instantiate_gated_delta_update_dims(type) \ - instantiate_gated_delta_update(type, 64, 64) +#define instantiate_gated_delta_update_dims(in_type, st_type) \ + instantiate_gated_delta_update(in_type, st_type, 64, 64, 4, 4) -instantiate_gated_delta_update_dims(float) +instantiate_gated_delta_update_dims(float, float) diff --git a/mlx/backend/metal/kernels/gated_delta_update_impl.h b/mlx/backend/metal/kernels/gated_delta_update_impl.h index b9459a9550..67caa8302e 100644 --- a/mlx/backend/metal/kernels/gated_delta_update_impl.h +++ b/mlx/backend/metal/kernels/gated_delta_update_impl.h @@ -5,20 +5,84 @@ using namespace metal; -template -[[kernel]] void gated_delta_update_fwd( - const device T* q [[buffer(0)]], - const device T* k [[buffer(1)]], - const device T* v [[buffer(2)]], - const device T* g [[buffer(3)]], - const device T* beta [[buffer(4)]], - const device T* h0 [[buffer(5)]], - device T* out [[buffer(6)]], - device T* hf [[buffer(7)]], - constant int& B [[buffer(8)]], - constant int& H [[buffer(9)]], - constant int& T_len [[buffer(10)]], - uint3 tid [[thread_position_in_grid]] +template +[[kernel]] void gated_delta_step( + const device InT* q [[buffer(0)]], + const device InT* k [[buffer(1)]], + const device InT* v [[buffer(2)]], + const device InT* g [[buffer(3)]], // [B, T, Hv] or [B, T, Hv, Dk] + const device InT* beta [[buffer(4)]], // [B, T, Hv] + const device StT* state_in [[buffer(5)]], // [B, Hv, Dv, Dk] + constant int& T [[buffer(6)]], + device InT* y [[buffer(7)]], // [B, T, Hv, Dv] + device StT* state_out [[buffer(8)]], // [B, Hv, Dv, Dk] + uint3 thread_position_in_grid [[thread_position_in_grid]], + uint3 thread_position_in_threadgroup [[thread_position_in_threadgroup]], + uint thread_index_in_simdgroup [[thread_index_in_simdgroup]] ) { // kernel implementation + auto n = thread_position_in_grid.z; + auto b_idx = n / Hv; + auto hv_idx = n % Hv; + auto hk_idx = hv_idx / (Hv / Hk); + constexpr int n_per_t = Dk / 32; + + // q, k: [B, T, Hk, Dk] + auto q_ = q + b_idx * T * Hk * Dk + hk_idx * Dk; + auto k_ = k + b_idx * T * Hk * Dk + hk_idx * Dk; + + // v, y: [B, T, Hv, Dv] + auto v_ = v + b_idx * T * Hv * Dv + hv_idx * Dv; + y += b_idx * T * Hv * Dv + hv_idx * Dv; + + auto dk_idx = thread_position_in_threadgroup.x; + auto dv_idx = thread_position_in_grid.y; + + // state_in, state_out: [B, Hv, Dv, Dk] + auto i_state = state_in + (n * Dv + dv_idx) * Dk; + auto o_state = state_out + (n * Dv + dv_idx) * Dk; + + float state[n_per_t]; + for (int i = 0; i < n_per_t; ++i) { + auto s_idx = n_per_t * dk_idx + i; + state[i] = static_cast(i_state[s_idx]); + } + + // g: [B, T, Hv] + auto g_ = g + b_idx * T * Hv; + auto beta_ = beta + b_idx * T * Hv; + + for (int t = 0; t < T; ++t) { + float kv_mem = 0.0f; + for (int i = 0; i < n_per_t; ++i) { + auto s_idx = n_per_t * dk_idx + i; + state[i] = state[i] * g_[hv_idx]; + kv_mem += state[i] * k_[s_idx]; + } + kv_mem = simd_sum(kv_mem); + + auto delta = (v_[dv_idx] - kv_mem) * beta_[hv_idx]; + + float out = 0.0f; + for (int i = 0; i < n_per_t; ++i) { + auto s_idx = n_per_t * dk_idx + i; + state[i] = state[i] + k_[s_idx] * delta; + out += state[i] * q_[s_idx]; + } + out = simd_sum(out); + if (thread_index_in_simdgroup == 0) { + y[dv_idx] = static_cast(out); + } + // Increment data pointers to next time step + q_ += Hk * Dk; + k_ += Hk * Dk; + v_ += Hv * Dv; + y += Hv * Dv; + g_ += Hv; + beta_ += Hv; + } + for (int i = 0; i < n_per_t; ++i) { + auto s_idx = n_per_t * dk_idx + i; + o_state[s_idx] = static_cast(state[i]); + } } diff --git a/mlx/fast.cpp b/mlx/fast.cpp index 4a0856e850..c5c1344a58 100644 --- a/mlx/fast.cpp +++ b/mlx/fast.cpp @@ -944,13 +944,16 @@ std::vector gated_delta_update_forward( auto g = astype(gates, out_dtype, s); auto beta = astype(beta_, out_dtype, s); - int B = q.shape(0), H = q.shape(1); - int T = q.shape(2), Dk = q.shape(3); + int B = q.shape(0); + int Hk = q.shape(1); // key heads + int T = q.shape(2); + int Dk = q.shape(3); + int Hv = v.shape(1); // value heads int Dv = v.shape(3); auto h0 = initial_state.has_value() ? astype(*initial_state, out_dtype, s) - : zeros({B, H, Dk, Dv}, out_dtype, s); + : zeros({B, Hv, Dk, Dv}, out_dtype, s); auto fallback = [](std::vector inputs) -> std::vector { throw std::runtime_error("NYI: GatedDeltaUpdate CPU fallback"); @@ -958,7 +961,7 @@ std::vector gated_delta_update_forward( }; return array::make_arrays( - /* output shapes */ {{B, H, T, Dv}, {B, H, Dk, Dv}}, + /* output shapes */ {{B, T, Hv, Dv}, {B, Hv, Dk, Dv}}, /* dtypes */ {out_dtype, out_dtype}, /* primitive */ std::make_shared(to_stream(s), fallback), /* inputs */ {q, k, v, g, beta, h0} From 3b81c8b4af9fc7d7177c9be5eb3ca715715a72e8 Mon Sep 17 00:00:00 2001 From: tpegolotti Date: Thu, 11 Jun 2026 15:43:53 +0200 Subject: [PATCH 04/56] make shapes consistent with reference --- mlx/backend/metal/gated_delta_update.cpp | 10 +++++----- mlx/fast.cpp | 10 +++++----- 2 files changed, 10 insertions(+), 10 deletions(-) diff --git a/mlx/backend/metal/gated_delta_update.cpp b/mlx/backend/metal/gated_delta_update.cpp index 2c9797696b..41df7a3ef7 100644 --- a/mlx/backend/metal/gated_delta_update.cpp +++ b/mlx/backend/metal/gated_delta_update.cpp @@ -35,11 +35,11 @@ void GatedDeltaUpdate::eval_gpu( auto& hf = outputs[1]; int B = q.shape(0); - int Hk = q.shape(1); // key heads - int T = q.shape(2); - int Dk = q.shape(3); - int Hv = v.shape(1); // value heads - int Dv = v.shape(3); + int T = q.shape(1); + int Hk = q.shape(2); + int Dk = q.shape(3); + int Hv = v.shape(2); + int Dv = v.shape(3); out.set_data(allocator::malloc(out.nbytes())); hf.set_data(allocator::malloc(hf.nbytes())); diff --git a/mlx/fast.cpp b/mlx/fast.cpp index c5c1344a58..99aa06df0f 100644 --- a/mlx/fast.cpp +++ b/mlx/fast.cpp @@ -945,11 +945,11 @@ std::vector gated_delta_update_forward( auto beta = astype(beta_, out_dtype, s); int B = q.shape(0); - int Hk = q.shape(1); // key heads - int T = q.shape(2); - int Dk = q.shape(3); - int Hv = v.shape(1); // value heads - int Dv = v.shape(3); + int T = q.shape(1); + int Hk = q.shape(2); + int Dk = q.shape(3); + int Hv = v.shape(2); + int Dv = v.shape(3); auto h0 = initial_state.has_value() ? astype(*initial_state, out_dtype, s) From 6a7f080c23847c1a784573ffe17b7c06f585fd87 Mon Sep 17 00:00:00 2001 From: tpegolotti Date: Fri, 12 Jun 2026 17:22:36 +0200 Subject: [PATCH 05/56] Add skeleton for chunkwise implementation --- mlx/backend/metal/gated_delta_update.cpp | 147 +++++++++++++----- .../metal/kernels/gated_delta_update.metal | 44 +++++- .../metal/kernels/gated_delta_update_impl.h | 141 +++++++++++------ mlx/fast_primitives.h | 2 +- 4 files changed, 235 insertions(+), 99 deletions(-) diff --git a/mlx/backend/metal/gated_delta_update.cpp b/mlx/backend/metal/gated_delta_update.cpp index 41df7a3ef7..1a3e5bc845 100644 --- a/mlx/backend/metal/gated_delta_update.cpp +++ b/mlx/backend/metal/gated_delta_update.cpp @@ -35,52 +35,117 @@ void GatedDeltaUpdate::eval_gpu( auto& hf = outputs[1]; int B = q.shape(0); - int T = q.shape(1); - int Hk = q.shape(2); - int Dk = q.shape(3); - int Hv = v.shape(2); - int Dv = v.shape(3); + int T = q.shape(1); + int Hk = q.shape(2); + int Dk = q.shape(3); + int Hv = v.shape(2); + int Dv = v.shape(3); + + const int C = 16; + int n_chunks = T / C; // TODO: make general + + std::string kernel_name = C == 1 ? "seq_gated_delta_" : "chunk_gated_delta_"; + std::string suffix = get_type_string(q.dtype()) // "float" + + "_" + get_type_string(h0.dtype()) // "float" + + "_" + std::to_string(Dk) + + "_" + std::to_string(Dv) + + "_" + std::to_string(Hk) + + "_" + std::to_string(Hv); - out.set_data(allocator::malloc(out.nbytes())); - hf.set_data(allocator::malloc(hf.nbytes())); - std::string base_name = "gated_delta_step_" - + get_type_string(q.dtype()) // "float" - + "_" + get_type_string(h0.dtype()) // "float" - + "_" + std::to_string(Dk) - + "_" + std::to_string(Dv) - + "_" + std::to_string(Hk) - + "_" + std::to_string(Hv); - // e.g. "gated_delta_step_float_float_64_64_4_4" - std::string hash_name = base_name; - metal::MTLFCList func_consts = {}; - - auto kernel = get_steel_gated_delta_forward_kernel( - d, - base_name, - hash_name, - func_consts); + std::string base_name = kernel_name + suffix; + + base_name += C >= 1 ? "_" + std::to_string(C) : ""; + + std::string hash_name = base_name; + + metal::MTLFCList func_consts = {}; + + auto delta_kernel = get_steel_gated_delta_forward_kernel( + d, + base_name, + hash_name, + func_consts); auto& compute_encoder = metal::get_command_encoder(s); - compute_encoder.set_compute_pipeline_state(kernel); - - compute_encoder.set_input_array(q, 0); - compute_encoder.set_input_array(k, 1); - compute_encoder.set_input_array(v, 2); - compute_encoder.set_input_array(g, 3); - compute_encoder.set_input_array(beta, 4); - compute_encoder.set_input_array(h0, 5); - compute_encoder.set_bytes(T, 6); - compute_encoder.set_output_array(out, 7); - compute_encoder.set_output_array(hf, 8); - - // auto grid = MTL::Size(1, 1, 1); - // auto threads = MTL::Size(1, 1, 1); - auto grid = MTL::Size(32, Dv, B * Hv); - auto threads = MTL::Size(32, 4, 1); - compute_encoder.dispatch_threads(grid, threads); - // throw std::runtime_error("NYI"); + + out.set_data(allocator::malloc(out.nbytes())); + hf.set_data(allocator::malloc(hf.nbytes())); + + + if (C > 1){ + // allocate full W and U -- [B, T, H, D] + array W({B, T, Hk, Dk}, q.dtype(), nullptr, {}); + array U({B, T, Hv, Dv}, q.dtype(), nullptr, {}); + W.set_data(allocator::malloc(W.nbytes())); + U.set_data(allocator::malloc(U.nbytes())); + + compute_encoder.add_temporary(W); + compute_encoder.add_temporary(U); + + // --- kernel 1: compute full W and U --- + // all chunks dispatched simultaneously + std::string make_wy_name = "make_wy_" + + get_type_string(q.dtype()) // "float" + + "_" + std::to_string(Dk) + + "_" + std::to_string(Dv) + + "_" + std::to_string(Hk) + + "_" + std::to_string(Hv) + + "_" + std::to_string(C); + + auto wy_kernel = d.get_kernel(make_wy_name); + + compute_encoder.set_compute_pipeline_state(wy_kernel); + compute_encoder.set_input_array(k, 0); + compute_encoder.set_input_array(v, 1); + compute_encoder.set_input_array(g, 2); + compute_encoder.set_input_array(beta, 3); + compute_encoder.set_output_array(W, 4); + compute_encoder.set_output_array(U, 5); + compute_encoder.set_bytes(T, 6); + + auto grid_wy = MTL::Size(Dk, Dv, B * Hv * n_chunks); + auto threads_wy = MTL::Size(32, 4, 1); + compute_encoder.dispatch_threads(grid_wy, threads_wy); + + // --- kernel 2: gated delta -- one dispatch, loop inside kernel --- + compute_encoder.set_compute_pipeline_state(delta_kernel); + compute_encoder.set_input_array(q, 0); + compute_encoder.set_input_array(k, 1); + compute_encoder.set_input_array(W, 2); + compute_encoder.set_input_array(U, 3); + compute_encoder.set_input_array(h0, 4); // initial state in + compute_encoder.set_input_array(g, 5); + compute_encoder.set_output_array(out, 6); + compute_encoder.set_output_array(hf, 7); // final state out + compute_encoder.set_bytes(T, 8); + compute_encoder.set_bytes(n_chunks, 9); + + auto grid = MTL::Size(32, Dv, B * Hv); + auto threads = MTL::Size(32, 4, 1); + compute_encoder.dispatch_threads(grid, threads); + + } else { + compute_encoder.set_compute_pipeline_state(delta_kernel); + + compute_encoder.set_input_array(q, 0); + compute_encoder.set_input_array(k, 1); + compute_encoder.set_input_array(v, 2); + compute_encoder.set_input_array(g, 3); + compute_encoder.set_input_array(beta, 4); + compute_encoder.set_input_array(h0, 5); + compute_encoder.set_bytes(T, 6); + compute_encoder.set_output_array(out, 7); + compute_encoder.set_output_array(hf, 8); + + // auto grid = MTL::Size(1, 1, 1); + // auto threads = MTL::Size(1, 1, 1); + auto grid = MTL::Size(32, Dv, B * Hv); + auto threads = MTL::Size(32, 4, 1); + compute_encoder.dispatch_threads(grid, threads); + } + // throw std::runtime_error("NYI"); } bool GatedDeltaUpdate::is_equivalent(const Primitive& other) const { diff --git a/mlx/backend/metal/kernels/gated_delta_update.metal b/mlx/backend/metal/kernels/gated_delta_update.metal index a76576ebd7..e41e4b40f7 100644 --- a/mlx/backend/metal/kernels/gated_delta_update.metal +++ b/mlx/backend/metal/kernels/gated_delta_update.metal @@ -3,10 +3,10 @@ using namespace metal; -#define instantiate_gated_delta_update(in_type, st_type, dk, dv, hk, hv) \ +#define instantiate_gated_delta_update_seq(in_type, st_type, dk, dv, hk, hv) \ instantiate_kernel( \ - "gated_delta_step_" #in_type "_" #st_type "_" #dk "_" #dv "_" #hk "_" #hv, \ - gated_delta_step, \ + "seq_gated_delta_" #in_type "_" #st_type "_" #dk "_" #dv "_" #hk "_" #hv, \ + gated_delta_seq, \ in_type, \ st_type, \ dk, \ @@ -14,7 +14,39 @@ using namespace metal; hk, \ hv) -#define instantiate_gated_delta_update_dims(in_type, st_type) \ - instantiate_gated_delta_update(in_type, st_type, 64, 64, 4, 4) +#define instantiate_gated_delta_update_seq_dims(in_type, st_type) \ + instantiate_gated_delta_update_seq(in_type, st_type, 64, 64, 4, 4) -instantiate_gated_delta_update_dims(float, float) + +#define instantiate_gated_delta_update_chunk(in_type, st_type, dk, dv, hk, hv, c) \ + instantiate_kernel( \ + "chunk_gated_delta_" #in_type "_" #st_type "_" #dk "_" #dv "_" #hk "_" #hv "_" #c, \ + gated_delta_chunk, \ + in_type, \ + st_type, \ + dk, \ + dv, \ + hk, \ + hv, \ + c) + +#define instantiate_gated_delta_update_chunk_dims(in_type, st_type) \ + instantiate_gated_delta_update_chunk(in_type, st_type, 64, 64, 4, 4, 16) + +#define instantiate_make_wy(in_type, dk, dv, hk, hv, c) \ + instantiate_kernel( \ + "make_wy_" #in_type "_" #dk "_" #dv "_" #hk "_" #hv "_" #c, \ + make_wy, \ + in_type, \ + dk, \ + dv, \ + hk, \ + hv, \ + c) + +#define instantiate_make_wy_dims(in_type) \ + instantiate_make_wy(in_type, 64, 64, 4, 4, 16) + +instantiate_gated_delta_update_seq_dims(float, float) +instantiate_gated_delta_update_chunk_dims(float, float) +instantiate_make_wy_dims(float) \ No newline at end of file diff --git a/mlx/backend/metal/kernels/gated_delta_update_impl.h b/mlx/backend/metal/kernels/gated_delta_update_impl.h index 67caa8302e..6a83979d0c 100644 --- a/mlx/backend/metal/kernels/gated_delta_update_impl.h +++ b/mlx/backend/metal/kernels/gated_delta_update_impl.h @@ -5,8 +5,47 @@ using namespace metal; +template +[[kernel]] void make_wy( + const device InT* k [[buffer(0)]], + const device InT* v [[buffer(1)]], + const device InT* g [[buffer(2)]], // [B, T, Hv] or [B, T, Hv, Dk] + const device InT* beta [[buffer(3)]], // [B, Hv, Dv, Dk] + device InT* W [[buffer(4)]], // [B, T, Hv] + device InT* U [[buffer(5)]], // [B, T, Hv] + constant int& T [[buffer(6)]], + uint3 thread_position_in_grid [[thread_position_in_grid]], + uint3 thread_position_in_threadgroup [[thread_position_in_threadgroup]], + uint thread_index_in_simdgroup [[thread_index_in_simdgroup]] +) { + +} + +template +[[kernel]] void gated_delta_chunk( + const device InT* q [[buffer(0)]], + const device InT* k [[buffer(1)]], + const device InT* W [[buffer(2)]], + const device InT* U [[buffer(3)]], // [B, T, Hv] or [B, T, Hv, Dk] + const device StT* state_in [[buffer(4)]], // [B, Hv, Dv, Dk] + const device InT* g [[buffer(5)]], // [B, T, Hv] + device InT* y [[buffer(6)]], // [B, T, Hv, Dv] + device StT* state_out [[buffer(7)]], // [B, Hv, Dv, Dk] + constant int& T [[buffer(8)]], + uint3 thread_position_in_grid [[thread_position_in_grid]], + uint3 thread_position_in_threadgroup [[thread_position_in_threadgroup]], + uint thread_index_in_simdgroup [[thread_index_in_simdgroup]] +) { + +} + + +/* + auto grid = MTL::Size(32, Dv, B * Hv); + auto threads = MTL::Size(32, 4, 1); + */ template -[[kernel]] void gated_delta_step( +[[kernel]] void gated_delta_seq( const device InT* q [[buffer(0)]], const device InT* k [[buffer(1)]], const device InT* v [[buffer(2)]], @@ -22,67 +61,67 @@ template ) { // kernel implementation auto n = thread_position_in_grid.z; - auto b_idx = n / Hv; - auto hv_idx = n % Hv; - auto hk_idx = hv_idx / (Hv / Hk); - constexpr int n_per_t = Dk / 32; + auto b_idx = n / Hv; + auto hv_idx = n % Hv; + auto hk_idx = hv_idx / (Hv / Hk); + constexpr int n_per_t = Dk / 32; + + // q, k: [B, T, Hk, Dk] + auto q_ = q + b_idx * T * Hk * Dk + hk_idx * Dk; + auto k_ = k + b_idx * T * Hk * Dk + hk_idx * Dk; - // q, k: [B, T, Hk, Dk] - auto q_ = q + b_idx * T * Hk * Dk + hk_idx * Dk; - auto k_ = k + b_idx * T * Hk * Dk + hk_idx * Dk; + // v, y: [B, T, Hv, Dv] + auto v_ = v + b_idx * T * Hv * Dv + hv_idx * Dv; + y += b_idx * T * Hv * Dv + hv_idx * Dv; - // v, y: [B, T, Hv, Dv] - auto v_ = v + b_idx * T * Hv * Dv + hv_idx * Dv; - y += b_idx * T * Hv * Dv + hv_idx * Dv; + auto dk_idx = thread_position_in_threadgroup.x; + auto dv_idx = thread_position_in_grid.y; - auto dk_idx = thread_position_in_threadgroup.x; - auto dv_idx = thread_position_in_grid.y; + // state_in, state_out: [B, Hv, Dv, Dk] + auto i_state = state_in + (n * Dv + dv_idx) * Dk; + auto o_state = state_out + (n * Dv + dv_idx) * Dk; - // state_in, state_out: [B, Hv, Dv, Dk] - auto i_state = state_in + (n * Dv + dv_idx) * Dk; - auto o_state = state_out + (n * Dv + dv_idx) * Dk; + float state[n_per_t]; + for (int i = 0; i < n_per_t; ++i) { + auto s_idx = n_per_t * dk_idx + i; + state[i] = static_cast(i_state[s_idx]); + } - float state[n_per_t]; + // g: [B, T, Hv] + auto g_ = g + b_idx * T * Hv; + auto beta_ = beta + b_idx * T * Hv; + + for (int t = 0; t < T; ++t) { + float kv_mem = 0.0f; for (int i = 0; i < n_per_t; ++i) { auto s_idx = n_per_t * dk_idx + i; - state[i] = static_cast(i_state[s_idx]); + state[i] = state[i] * g_[hv_idx]; + kv_mem += state[i] * k_[s_idx]; } + kv_mem = simd_sum(kv_mem); - // g: [B, T, Hv] - auto g_ = g + b_idx * T * Hv; - auto beta_ = beta + b_idx * T * Hv; - - for (int t = 0; t < T; ++t) { - float kv_mem = 0.0f; - for (int i = 0; i < n_per_t; ++i) { - auto s_idx = n_per_t * dk_idx + i; - state[i] = state[i] * g_[hv_idx]; - kv_mem += state[i] * k_[s_idx]; - } - kv_mem = simd_sum(kv_mem); - - auto delta = (v_[dv_idx] - kv_mem) * beta_[hv_idx]; + auto delta = (v_[dv_idx] - kv_mem) * beta_[hv_idx]; - float out = 0.0f; - for (int i = 0; i < n_per_t; ++i) { - auto s_idx = n_per_t * dk_idx + i; - state[i] = state[i] + k_[s_idx] * delta; - out += state[i] * q_[s_idx]; - } - out = simd_sum(out); - if (thread_index_in_simdgroup == 0) { - y[dv_idx] = static_cast(out); - } - // Increment data pointers to next time step - q_ += Hk * Dk; - k_ += Hk * Dk; - v_ += Hv * Dv; - y += Hv * Dv; - g_ += Hv; - beta_ += Hv; - } + float out = 0.0f; for (int i = 0; i < n_per_t; ++i) { auto s_idx = n_per_t * dk_idx + i; - o_state[s_idx] = static_cast(state[i]); + state[i] = state[i] + k_[s_idx] * delta; + out += state[i] * q_[s_idx]; + } + out = simd_sum(out); + if (thread_index_in_simdgroup == 0) { + y[dv_idx] = static_cast(out); } + // Increment data pointers to next time step + q_ += Hk * Dk; + k_ += Hk * Dk; + v_ += Hv * Dv; + y += Hv * Dv; + g_ += Hv; + beta_ += Hv; + } + for (int i = 0; i < n_per_t; ++i) { + auto s_idx = n_per_t * dk_idx + i; + o_state[s_idx] = static_cast(state[i]); + } } diff --git a/mlx/fast_primitives.h b/mlx/fast_primitives.h index 5a4609d44b..9cc8537ffd 100644 --- a/mlx/fast_primitives.h +++ b/mlx/fast_primitives.h @@ -358,7 +358,7 @@ class GatedDeltaUpdate : public Custom { } private: - /* TODO */ + /* TODO */ }; From bb53c0fe598e331ab235de8c00e94b8535a83237 Mon Sep 17 00:00:00 2001 From: tpegolotti Date: Tue, 16 Jun 2026 15:54:24 +0200 Subject: [PATCH 06/56] base chunkwise implementation --- mlx/backend/metal/gated_delta_update.cpp | 9 +- .../metal/kernels/gated_delta_update_impl.h | 224 +++++++++++++++++- 2 files changed, 224 insertions(+), 9 deletions(-) diff --git a/mlx/backend/metal/gated_delta_update.cpp b/mlx/backend/metal/gated_delta_update.cpp index 1a3e5bc845..46278624f9 100644 --- a/mlx/backend/metal/gated_delta_update.cpp +++ b/mlx/backend/metal/gated_delta_update.cpp @@ -84,8 +84,7 @@ void GatedDeltaUpdate::eval_gpu( compute_encoder.add_temporary(W); compute_encoder.add_temporary(U); - // --- kernel 1: compute full W and U --- - // all chunks dispatched simultaneously + // kernel 1: compute full W and U std::string make_wy_name = "make_wy_" + get_type_string(q.dtype()) // "float" + "_" + std::to_string(Dk) @@ -105,11 +104,11 @@ void GatedDeltaUpdate::eval_gpu( compute_encoder.set_output_array(U, 5); compute_encoder.set_bytes(T, 6); - auto grid_wy = MTL::Size(Dk, Dv, B * Hv * n_chunks); + auto grid_wy = MTL::Size(32, 4, B * Hv * n_chunks); auto threads_wy = MTL::Size(32, 4, 1); compute_encoder.dispatch_threads(grid_wy, threads_wy); - // --- kernel 2: gated delta -- one dispatch, loop inside kernel --- + // kernel 2: gated delta compute_encoder.set_compute_pipeline_state(delta_kernel); compute_encoder.set_input_array(q, 0); compute_encoder.set_input_array(k, 1); @@ -124,6 +123,8 @@ void GatedDeltaUpdate::eval_gpu( auto grid = MTL::Size(32, Dv, B * Hv); auto threads = MTL::Size(32, 4, 1); + // auto grid = MTL::Size(1, 1, 1); + // auto threads = MTL::Size(1, 1, 1); compute_encoder.dispatch_threads(grid, threads); } else { diff --git a/mlx/backend/metal/kernels/gated_delta_update_impl.h b/mlx/backend/metal/kernels/gated_delta_update_impl.h index 6a83979d0c..95ddcb234d 100644 --- a/mlx/backend/metal/kernels/gated_delta_update_impl.h +++ b/mlx/backend/metal/kernels/gated_delta_update_impl.h @@ -18,15 +18,116 @@ template uint3 thread_position_in_threadgroup [[thread_position_in_threadgroup]], uint thread_index_in_simdgroup [[thread_index_in_simdgroup]] ) { - + auto n = thread_position_in_grid.z; + auto n_chunks = T / C; + auto chunk = n % n_chunks; + auto bh_idx = n / n_chunks; + auto b_idx = bh_idx / Hv; + auto hv_idx = bh_idx % Hv; + auto hk_idx = hv_idx / (Hv / Hk); + + constexpr int n_per_dk = Dk / 32; + constexpr int n_per_dv = Dv / 32; + + auto dk_idx = thread_position_in_threadgroup.x; + auto simd_group_id = thread_position_in_threadgroup.y; + const int num_simdgroups = 4; + + int offset_t = chunk * C; + auto g_ = g + b_idx * T * Hv + offset_t * Hv; + auto k_ = k + b_idx * T * Hk * Dk + offset_t * Hk * Dk + hk_idx * Dk; + auto v_ = v + b_idx * T * Hv * Dv + offset_t * Hv * Dv + hv_idx * Dv; + auto beta_ = beta + b_idx * T * Hv + offset_t * Hv; + auto W_ = W + b_idx * T * Hv * Dk + offset_t * Hv * Dk + hv_idx * Dk; + auto U_ = U + b_idx * T * Hv * Dv + offset_t * Hv * Dv + hv_idx * Dv; + + // threadgroup memory + threadgroup float K_tg[C][Dk]; + threadgroup float KKt[C][C]; + + for (int i = simd_group_id; i < C; i += num_simdgroups) { + for (int d = dk_idx; d < Dk; d += 32) { + K_tg[i][d] = k_[i * Hk * Dk + d]; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + // compute gamma (cumprod) + float gamma[C]; + gamma[0] = g_[hv_idx]; + for (int i = 1; i < C; i++) { + gamma[i] = gamma[i-1] * g_[i * Hv + hv_idx]; + } + + // compute KKt = K @ K.T + unsigned int counter = 0; + for (int i = 0; i < C; i++) { + for (int j = 0; j <= i; j++) { + if (counter % num_simdgroups == simd_group_id) { + float sum = 0; + for (int d = dk_idx; d < Dk; d += 32) { + sum += K_tg[i][d] * K_tg[j][d]; + } + sum = simd_sum(sum); + if (thread_index_in_simdgroup == 0) { + KKt[i][j] = sum; + } + } + counter++; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + float u[C][n_per_dv]; + float w[C][n_per_dk]; + + // initialize + for (int i = 0; i < C; i++) { + float beta_i = beta_[i * Hv + hv_idx]; + for (int p = 0; p < n_per_dv; p++) { + int dv = dk_idx * n_per_dv + p; + u[i][p] = beta_i * v_[i * Hv * Dv + dv]; + } + for (int p = 0; p < n_per_dk; p++) { + int dk = dk_idx * n_per_dk + p; + w[i][p] = beta_i * K_tg[i][dk]; + } + } + + + // forward substitution + for (int i = 1; i < C; i++) { + float beta_i = beta_[i * Hv + hv_idx]; + for (int j = 0; j < i; j++) { + float a_U = beta_i * (gamma[i] / gamma[j]) * KKt[i][j]; + float a_W = beta_i * KKt[i][j]; + for (int p = 0; p < n_per_dv; p++) { + u[i][p] -= a_U * u[j][p]; + } + for (int p = 0; p < n_per_dk; p++) { + w[i][p] -= a_W * w[j][p]; + } + } + } + + for (int i = 0; i < C; i++) { + for (int p = 0; p < n_per_dv; p++) { + int dv = dk_idx * n_per_dv + p; + U_[i * Hv * Dv + dv] = u[i][p]; + } + for (int p = 0; p < n_per_dk; p++) { + int dk = dk_idx * n_per_dk + p; + W_[i * Hk * Dk + dk] = w[i][p]; + } + } } template [[kernel]] void gated_delta_chunk( const device InT* q [[buffer(0)]], const device InT* k [[buffer(1)]], - const device InT* W [[buffer(2)]], - const device InT* U [[buffer(3)]], // [B, T, Hv] or [B, T, Hv, Dk] + const device InT* W [[buffer(2)]], // [B, T, Hv, Dk] + const device InT* U [[buffer(3)]], // [B, T, Hv, Dv] const device StT* state_in [[buffer(4)]], // [B, Hv, Dv, Dk] const device InT* g [[buffer(5)]], // [B, T, Hv] device InT* y [[buffer(6)]], // [B, T, Hv, Dv] @@ -36,9 +137,122 @@ template uint3 thread_position_in_threadgroup [[thread_position_in_threadgroup]], uint thread_index_in_simdgroup [[thread_index_in_simdgroup]] ) { - -} + auto n = thread_position_in_grid.z; + auto b_idx = n / Hv; + auto hv_idx = n % Hv; + auto hk_idx = hv_idx / (Hv / Hk); + constexpr int n_per_dk = Dk / 32; + + auto dk_idx = thread_position_in_threadgroup.x; // simd group id + auto dv_idx = thread_position_in_grid.y; + + // set up pointers + // g: [B, T, Hv] + auto g_ = g + b_idx * T * Hv; + + // q, k: [B, T, Hk, Dk] + auto q_ = q + b_idx * T * Hk * Dk + hk_idx * Dk; + auto k_ = k + b_idx * T * Hk * Dk + hk_idx * Dk; + + // W: [B, T, Hk, Dk], U: [B, T, Hv, Dv] + auto W_ = W + b_idx * T * Hv * Dk + hv_idx * Dk; + auto U_ = U + b_idx * T * Hv * Dv + hv_idx * Dv; + + // v, y: [B, T, Hv, Dv] + y += b_idx * T * Hv * Dv + hv_idx * Dv; + + // state_in, state_out: [B, Hv, Dv, Dk] + auto i_state = state_in + (n * Dv + dv_idx) * Dk; + auto o_state = state_out + (n * Dv + dv_idx) * Dk; + + float state[n_per_dk]; + for (int i = 0; i < n_per_dk; ++i) { + auto s_idx = n_per_dk * dk_idx + i; + state[i] = static_cast(i_state[s_idx]); + } + + float gamma[C]; + float delta[C]; + + for (int t = 0; t < T; t+=C) { + gamma[0] = g_[hv_idx]; + for (int i = 1; i < C; i++) { + gamma[i] = gamma[i-1] * g_[i * Hv + hv_idx]; + } + float gamma_C = gamma[C-1]; + + for (int c = 0; c < C; c++) { + // W_left[c] @ S^T + float ws = 0; + float g_c = gamma[c]; + for (int i = 0; i < n_per_dk; i++) { + int s_idx = dk_idx * n_per_dk + i; + float w_left = g_c * W_[c * Hk * Dk + s_idx]; + ws += w_left * state[i]; + } + ws = simd_sum(ws); + + // delta = U - ws + delta[c] = U_[c * Hv * Dv + dv_idx] - ws; + } + + + // Compute Output: O = Q_left S^T + Q K^T delta + for (int c = 0; c < C; c++) { + // Q_left[c] S^T + float out_c = 0; + float g_c = gamma[c]; + for (int i = 0; i < n_per_dk; i++) { + int s_idx = dk_idx * n_per_dk + i; + float q_left = g_c * q_[c * Hk * Dk + s_idx]; + out_c += q_left * state[i]; + } + out_c = simd_sum(out_c); + + // Q K^T delta + for (int j = 0; j <= c; j++) { + float qkt = 0; + for (int i = 0; i < n_per_dk; i++) { + int s_idx = dk_idx * n_per_dk + i; + qkt += q_[c * Hk * Dk + s_idx] * k_[j * Hk * Dk + s_idx]; + } + qkt = simd_sum(qkt); + out_c += (gamma[c] / gamma[j]) * qkt * delta[j]; + } + + if (thread_index_in_simdgroup == 0) { + y[c * Hv * Dv + dv_idx] = out_c; + } + } + // Update state: S = gamma_C * S + delta.T @ K + // gamma_C * S + for (int i = 0; i < n_per_dk; i++) { + state[i] *= gamma_C; + } + + // delta.T @ K + for (int c = 0; c < C; c++) { + float k_right = (gamma_C / gamma[c]); + for (int i = 0; i < n_per_dk; i++) { + int s_idx = dk_idx * n_per_dk + i; + state[i] += k_right * k_[c * Hk * Dk + s_idx] * delta[c]; + } + } + + // update pointers + q_ += C * Hk * Dk; + k_ += C * Hk * Dk; + U_ += C * Hv * Dv; + W_ += C * Hv * Dk; + y += C * Hv * Dv; + g_ += C * Hv; + } + for (int i = 0; i < n_per_dk; ++i) { + auto s_idx = n_per_dk * dk_idx + i; + o_state[s_idx] = static_cast(state[i]); + } +} /* auto grid = MTL::Size(32, Dv, B * Hv); From a1835d83e28adc1e3f11d4385a03342e2d548128 Mon Sep 17 00:00:00 2001 From: Tommaso Pegolotti <31819389+tpegolotti@users.noreply.github.com> Date: Tue, 16 Jun 2026 23:06:13 +0200 Subject: [PATCH 07/56] Delete local_files directory --- local_files/gated_basic.py | 0 local_files/references.py | 0 2 files changed, 0 insertions(+), 0 deletions(-) delete mode 100644 local_files/gated_basic.py delete mode 100644 local_files/references.py diff --git a/local_files/gated_basic.py b/local_files/gated_basic.py deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/local_files/references.py b/local_files/references.py deleted file mode 100644 index e69de29bb2..0000000000 From 747182c87bee7117d5eb8265e78f977f5a827114 Mon Sep 17 00:00:00 2001 From: tpegolotti Date: Thu, 18 Jun 2026 17:32:52 +0200 Subject: [PATCH 08/56] simdgroup matrices work for C=8 --- mlx/backend/metal/gated_delta_update.cpp | 10 +- .../metal/kernels/gated_delta_update.metal | 10 +- .../metal/kernels/gated_delta_update_impl.h | 217 +++++++++++------- mlx/fast.cpp | 3 +- mlx/fast.h | 1 + mlx/fast_primitives.h | 7 +- python/src/fast.cpp | 1 + 7 files changed, 150 insertions(+), 99 deletions(-) diff --git a/mlx/backend/metal/gated_delta_update.cpp b/mlx/backend/metal/gated_delta_update.cpp index 46278624f9..a64d567a96 100644 --- a/mlx/backend/metal/gated_delta_update.cpp +++ b/mlx/backend/metal/gated_delta_update.cpp @@ -41,10 +41,10 @@ void GatedDeltaUpdate::eval_gpu( int Hv = v.shape(2); int Dv = v.shape(3); - const int C = 16; + int C = chunk_size; int n_chunks = T / C; // TODO: make general - std::string kernel_name = C == 1 ? "seq_gated_delta_" : "chunk_gated_delta_"; + std::string kernel_name = C < 1 ? "seq_gated_delta_" : "chunk_gated_delta_"; std::string suffix = get_type_string(q.dtype()) // "float" + "_" + get_type_string(h0.dtype()) // "float" + "_" + std::to_string(Dk) @@ -52,7 +52,7 @@ void GatedDeltaUpdate::eval_gpu( + "_" + std::to_string(Hk) + "_" + std::to_string(Hv); - + // printf("C: %d\nname: %s\n",C,kernel_name.c_str()); std::string base_name = kernel_name + suffix; base_name += C >= 1 ? "_" + std::to_string(C) : ""; @@ -121,8 +121,8 @@ void GatedDeltaUpdate::eval_gpu( compute_encoder.set_bytes(T, 8); compute_encoder.set_bytes(n_chunks, 9); - auto grid = MTL::Size(32, Dv, B * Hv); - auto threads = MTL::Size(32, 4, 1); + auto grid = MTL::Size(32, Dv / 8, B * Hv); + auto threads = MTL::Size(32, 1, 1); // get one simdgroup to work // auto grid = MTL::Size(1, 1, 1); // auto threads = MTL::Size(1, 1, 1); compute_encoder.dispatch_threads(grid, threads); diff --git a/mlx/backend/metal/kernels/gated_delta_update.metal b/mlx/backend/metal/kernels/gated_delta_update.metal index e41e4b40f7..8edba33b0e 100644 --- a/mlx/backend/metal/kernels/gated_delta_update.metal +++ b/mlx/backend/metal/kernels/gated_delta_update.metal @@ -31,7 +31,10 @@ using namespace metal; c) #define instantiate_gated_delta_update_chunk_dims(in_type, st_type) \ - instantiate_gated_delta_update_chunk(in_type, st_type, 64, 64, 4, 4, 16) + instantiate_gated_delta_update_chunk(in_type, st_type, 64, 64, 4, 4, 32) \ + instantiate_gated_delta_update_chunk(in_type, st_type, 64, 64, 4, 4, 16) \ + instantiate_gated_delta_update_chunk(in_type, st_type, 64, 64, 4, 4, 8) \ + instantiate_gated_delta_update_chunk(in_type, st_type, 64, 64, 1, 1, 8) \ #define instantiate_make_wy(in_type, dk, dv, hk, hv, c) \ instantiate_kernel( \ @@ -45,7 +48,10 @@ using namespace metal; c) #define instantiate_make_wy_dims(in_type) \ - instantiate_make_wy(in_type, 64, 64, 4, 4, 16) + instantiate_make_wy(in_type, 64, 64, 4, 4, 32) \ + instantiate_make_wy(in_type, 64, 64, 4, 4, 16) \ + instantiate_make_wy(in_type, 64, 64, 4, 4, 8) \ + instantiate_make_wy(in_type, 64, 64, 1, 1, 8) \ instantiate_gated_delta_update_seq_dims(float, float) instantiate_gated_delta_update_chunk_dims(float, float) diff --git a/mlx/backend/metal/kernels/gated_delta_update_impl.h b/mlx/backend/metal/kernels/gated_delta_update_impl.h index 95ddcb234d..5865445146 100644 --- a/mlx/backend/metal/kernels/gated_delta_update_impl.h +++ b/mlx/backend/metal/kernels/gated_delta_update_impl.h @@ -3,6 +3,14 @@ #include #include "mlx/backend/metal/kernels/utils.h" +// importing steel +// #include "mlx/backend/metal/kernels/steel/attn/loader.h" +// #include "mlx/backend/metal/kernels/steel/attn/mma.h" +// #include "mlx/backend/metal/kernels/steel/attn/params.h" +// #include "mlx/backend/metal/kernels/steel/attn/transforms.h" +// #include "mlx/backend/metal/kernels/steel/gemm/params.h" +// #include "mlx/backend/metal/kernels/steel/utils.h" + using namespace metal; template @@ -124,27 +132,31 @@ template template [[kernel]] void gated_delta_chunk( - const device InT* q [[buffer(0)]], - const device InT* k [[buffer(1)]], - const device InT* W [[buffer(2)]], // [B, T, Hv, Dk] - const device InT* U [[buffer(3)]], // [B, T, Hv, Dv] - const device StT* state_in [[buffer(4)]], // [B, Hv, Dv, Dk] - const device InT* g [[buffer(5)]], // [B, T, Hv] - device InT* y [[buffer(6)]], // [B, T, Hv, Dv] - device StT* state_out [[buffer(7)]], // [B, Hv, Dv, Dk] - constant int& T [[buffer(8)]], - uint3 thread_position_in_grid [[thread_position_in_grid]], - uint3 thread_position_in_threadgroup [[thread_position_in_threadgroup]], - uint thread_index_in_simdgroup [[thread_index_in_simdgroup]] + const device InT* q [[buffer(0)]], + const device InT* k [[buffer(1)]], + const device InT* W [[buffer(2)]], // [B, T, Hv, Dk] + const device InT* U [[buffer(3)]], // [B, T, Hv, Dv] + const device StT* state_in [[buffer(4)]], // [B, Hv, Dv, Dk] + const device InT* g [[buffer(5)]], // [B, T, Hv] + device InT* y [[buffer(6)]], // [B, T, Hv, Dv] + device StT* state_out [[buffer(7)]], // [B, Hv, Dv, Dk] + constant int& T [[buffer(8)]], + uint3 thread_position_in_grid [[thread_position_in_grid]], + uint3 thread_position_in_threadgroup [[thread_position_in_threadgroup]], + uint thread_index_in_simdgroup [[thread_index_in_simdgroup]] ) { auto n = thread_position_in_grid.z; auto b_idx = n / Hv; auto hv_idx = n % Hv; auto hk_idx = hv_idx / (Hv / Hk); - constexpr int n_per_dk = Dk / 32; + + // get coord from steel + const short qid = thread_index_in_simdgroup / 4; + const short fm = (qid & 4) + ((thread_index_in_simdgroup / 2) % 4); + const short fn = (qid & 2) * 2 + (thread_index_in_simdgroup % 2) * 2; auto dk_idx = thread_position_in_threadgroup.x; // simd group id - auto dv_idx = thread_position_in_grid.y; + auto dv_idx = thread_position_in_grid.y * 8; // set up pointers // g: [B, T, Hv] @@ -165,79 +177,107 @@ template auto i_state = state_in + (n * Dv + dv_idx) * Dk; auto o_state = state_out + (n * Dv + dv_idx) * Dk; - float state[n_per_dk]; - for (int i = 0; i < n_per_dk; ++i) { - auto s_idx = n_per_dk * dk_idx + i; - state[i] = static_cast(i_state[s_idx]); + // threadgroup tiles + threadgroup float Q_left_tg[C][Dk]; // gamma-scaled Q + threadgroup float K_tg[C][Dk]; // raw K + threadgroup float W_tg[C][Dk]; // gamma-scaled W (W_left) + threadgroup float U_tg[C][8]; // U chunk, 8 dv values + threadgroup float K_right_tg[C][Dk]; + + simdgroup_float8x8 S_tile[8]; + simdgroup_float8x8 W_tile, K_tile, Q_tile; + simdgroup_float8x8 WS_tile; + simdgroup_float8x8 U_tile; + simdgroup_float8x8 delta_tile; + simdgroup_float8x8 tmp_tile; + simdgroup_float8x8 QKt_tile; + simdgroup_float8x8 out_tile; + simdgroup_float8x8 KD_tile; + + thread auto& kd = KD_tile.thread_elements(); + thread auto& d_e = delta_tile.thread_elements(); + thread auto& u_e = U_tile.thread_elements(); + thread auto& ws_e = WS_tile.thread_elements(); + thread auto& qkt_e = QKt_tile.thread_elements(); + thread auto& o_e = out_tile.thread_elements(); + + // load initial state into threadgroup + for (int kk = 0; kk < Dk; kk += 8) { + simdgroup_load(S_tile[kk/8], i_state + kk, Dk, ulong2(0, 0), true); } + threadgroup_barrier(mem_flags::mem_threadgroup); - float gamma[C]; - float delta[C]; - - for (int t = 0; t < T; t+=C) { - gamma[0] = g_[hv_idx]; - for (int i = 1; i < C; i++) { - gamma[i] = gamma[i-1] * g_[i * Hv + hv_idx]; - } - float gamma_C = gamma[C-1]; - - for (int c = 0; c < C; c++) { - // W_left[c] @ S^T - float ws = 0; - float g_c = gamma[c]; - for (int i = 0; i < n_per_dk; i++) { - int s_idx = dk_idx * n_per_dk + i; - float w_left = g_c * W_[c * Hk * Dk + s_idx]; - ws += w_left * state[i]; - } - ws = simd_sum(ws); - - // delta = U - ws - delta[c] = U_[c * Hv * Dv + dv_idx] - ws; - } - - - // Compute Output: O = Q_left S^T + Q K^T delta - for (int c = 0; c < C; c++) { - // Q_left[c] S^T - float out_c = 0; - float g_c = gamma[c]; - for (int i = 0; i < n_per_dk; i++) { - int s_idx = dk_idx * n_per_dk + i; - float q_left = g_c * q_[c * Hk * Dk + s_idx]; - out_c += q_left * state[i]; - } - out_c = simd_sum(out_c); - - // Q K^T delta - for (int j = 0; j <= c; j++) { - float qkt = 0; - for (int i = 0; i < n_per_dk; i++) { - int s_idx = dk_idx * n_per_dk + i; - qkt += q_[c * Hk * Dk + s_idx] * k_[j * Hk * Dk + s_idx]; - } - qkt = simd_sum(qkt); - out_c += (gamma[c] / gamma[j]) * qkt * delta[j]; - } - - if (thread_index_in_simdgroup == 0) { - y[c * Hv * Dv + dv_idx] = out_c; - } - } - - // Update state: S = gamma_C * S + delta.T @ K - // gamma_C * S - for (int i = 0; i < n_per_dk; i++) { - state[i] *= gamma_C; - } - // delta.T @ K - for (int c = 0; c < C; c++) { - float k_right = (gamma_C / gamma[c]); - for (int i = 0; i < n_per_dk; i++) { - int s_idx = dk_idx * n_per_dk + i; - state[i] += k_right * k_[c * Hk * Dk + s_idx] * delta[c]; - } + float gamma[C]; + + for (int t = 0; t < T; t += C) { + gamma[0] = g_[hv_idx]; + for (int i = 1; i < C; i++) { + gamma[i] = gamma[i-1] * g_[i * Hv + hv_idx]; + } + float gamma_C = gamma[C-1]; + + for (int c = 0; c < C; c++) { + float kr = gamma_C / gamma[c]; + for (int d = dk_idx; d < Dk; d += 32) { + Q_left_tg[c][d] = gamma[c] * q_[c * Hk * Dk + d]; + K_tg[c][d] = k_[c * Hk * Dk + d]; + W_tg[c][d] = gamma[c] * W_[c * Hv * Dk + d]; + K_right_tg[c][d] = kr * K_tg[c][d]; + } + } + + for (int idx = dk_idx; idx < C * 8; idx += 32) { + int c = idx / 8; + int j = idx % 8; + U_tg[c][j] = U_[c * Hv * Dv + dv_idx + j]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + // WS = W_left @ S^T + WS_tile = make_filled_simdgroup_matrix(0.f); + for (int kk = 0; kk < Dk; kk += 8) { + simdgroup_load(W_tile, &W_tg[0][kk], Dk); + simdgroup_multiply_accumulate(WS_tile, W_tile, S_tile[kk/8], WS_tile); + } + + // delta = U - WS + simdgroup_load(U_tile, &U_tg[0][0], 8); + + d_e[0] = u_e[0] - ws_e[0]; + d_e[1] = u_e[1] - ws_e[1]; + + + // Q_left @ S^T + tmp_tile = make_filled_simdgroup_matrix(0.f); + QKt_tile = make_filled_simdgroup_matrix(0.f); + + for (int kk = 0; kk < Dk; kk += 8) { + simdgroup_load(Q_tile, &Q_left_tg[0][kk], Dk); + simdgroup_load(K_tile, &K_tg[0][kk], Dk, ulong2(0, 0), true); + + simdgroup_multiply_accumulate(tmp_tile, Q_tile, S_tile[kk/8], tmp_tile); + simdgroup_multiply_accumulate(QKt_tile, Q_tile, K_tile, QKt_tile); + } + + + // element 0 at (fm, fn), element 1 at (fm, fn+1) + qkt_e[0] *= fn > fm ? 0.f : (1.0f / gamma[fn]); + qkt_e[1] *= fn + 1 > fm ? 0.f : (1.0f / gamma[fn+1]); + + simdgroup_multiply_accumulate(out_tile, QKt_tile, delta_tile, tmp_tile); + + y[fm * Hv * Dv + dv_idx + fn] = static_cast(o_e[0]); + y[fm * Hv * Dv + dv_idx + fn + 1] = static_cast(o_e[1]); + + + for (int kk = 0; kk < Dk; kk += 8) { + simdgroup_load(K_tile, &K_right_tg[0][kk], Dk, ulong2(0, 0), true); + simdgroup_multiply(KD_tile, K_tile, delta_tile); + + thread auto& s_e = S_tile[kk/8].thread_elements(); + s_e[0] = gamma_C * s_e[0] + kd[0]; + s_e[1] = gamma_C * s_e[1] + kd[1]; } // update pointers @@ -247,10 +287,11 @@ template W_ += C * Hv * Dk; y += C * Hv * Dv; g_ += C * Hv; - } - for (int i = 0; i < n_per_dk; ++i) { - auto s_idx = n_per_dk * dk_idx + i; - o_state[s_idx] = static_cast(state[i]); + } + + // o_state is [Dv, Dk]: o_state[dv][dk] = st_out[dk][dv] + for (int kk = 0; kk < Dk; kk += 8) { + simdgroup_store(S_tile[kk/8], o_state + kk, Dk, ulong2(0,0), true); // writes [Dk_block, Dv] layout } } diff --git a/mlx/fast.cpp b/mlx/fast.cpp index 99aa06df0f..67ce47cc31 100644 --- a/mlx/fast.cpp +++ b/mlx/fast.cpp @@ -929,6 +929,7 @@ std::vector gated_delta_update_forward( const array& gates, const array& beta_, const std::optional& initial_state, /* = std::nullopt */ + const int C, StreamOrDevice s /* = {} */) { // determine output dtype @@ -963,7 +964,7 @@ std::vector gated_delta_update_forward( return array::make_arrays( /* output shapes */ {{B, T, Hv, Dv}, {B, Hv, Dk, Dv}}, /* dtypes */ {out_dtype, out_dtype}, - /* primitive */ std::make_shared(to_stream(s), fallback), + /* primitive */ std::make_shared(to_stream(s), fallback, C), /* inputs */ {q, k, v, g, beta, h0} ); } diff --git a/mlx/fast.h b/mlx/fast.h index 796b159292..d4cd810463 100644 --- a/mlx/fast.h +++ b/mlx/fast.h @@ -62,6 +62,7 @@ MLX_API std::vector gated_delta_update_forward( const array& gates, const array& beta_, const std::optional& initial_state = std::nullopt, + const int C = 16, StreamOrDevice s = {} ); diff --git a/mlx/fast_primitives.h b/mlx/fast_primitives.h index 9cc8537ffd..305cfaa601 100644 --- a/mlx/fast_primitives.h +++ b/mlx/fast_primitives.h @@ -329,9 +329,10 @@ class GatedDeltaUpdate : public Custom { public: GatedDeltaUpdate( Stream stream, - std::function(std::vector)> fallback + std::function(std::vector)> fallback, + int C ) - : Custom(stream, std::move(fallback)) + : Custom(stream, std::move(fallback)), chunk_size(C) {} static bool use_fallback( @@ -358,7 +359,7 @@ class GatedDeltaUpdate : public Custom { } private: - /* TODO */ + int chunk_size; }; diff --git a/python/src/fast.cpp b/python/src/fast.cpp index 94d48b0f5d..86c971359a 100644 --- a/python/src/fast.cpp +++ b/python/src/fast.cpp @@ -342,6 +342,7 @@ void init_fast(nb::module_& parent_module) { "gates"_a, "beta"_a, "initial_state"_a = nb::none(), // optional, defaults to None + "C"_a = 16, // optional, defaults to None "stream"_a = nb::none(), // optional, defaults to None R"( Chunked gated delta network forward pass. From 6f0ceea28c43bd95e3ab47065a58c70c3ebd3991 Mon Sep 17 00:00:00 2001 From: tpegolotti Date: Fri, 19 Jun 2026 00:08:37 +0200 Subject: [PATCH 09/56] ran pre-commit --- mlx/backend/metal/gated_delta_update.cpp | 259 ++++--- mlx/backend/metal/kernels.h | 2 +- .../metal/kernels/gated_delta_update.metal | 78 +- .../metal/kernels/gated_delta_update_impl.h | 673 +++++++++--------- mlx/backend/metal/nojit_kernels.cpp | 4 +- mlx/fast.cpp | 85 ++- mlx/fast.h | 3 +- mlx/fast_primitives.h | 15 +- python/src/fast.cpp | 27 +- 9 files changed, 563 insertions(+), 583 deletions(-) diff --git a/mlx/backend/metal/gated_delta_update.cpp b/mlx/backend/metal/gated_delta_update.cpp index a64d567a96..6250b4047e 100644 --- a/mlx/backend/metal/gated_delta_update.cpp +++ b/mlx/backend/metal/gated_delta_update.cpp @@ -13,149 +13,138 @@ namespace mlx::core::fast { bool GatedDeltaUpdate::use_fallback(Stream s) { - // TODO: finish implementation - return false; + // TODO: finish implementation + return false; } void GatedDeltaUpdate::eval_gpu( const std::vector& inputs, std::vector& outputs) { - - auto& s = stream(); - auto& d = metal::device(s.device); - - auto& q = inputs[0]; - auto& k = inputs[1]; - auto& v = inputs[2]; - auto& g = inputs[3]; - auto& beta = inputs[4]; - auto& h0 = inputs[5]; - - auto& out = outputs[0]; - auto& hf = outputs[1]; - - int B = q.shape(0); - int T = q.shape(1); - int Hk = q.shape(2); - int Dk = q.shape(3); - int Hv = v.shape(2); - int Dv = v.shape(3); - - int C = chunk_size; - int n_chunks = T / C; // TODO: make general - - std::string kernel_name = C < 1 ? "seq_gated_delta_" : "chunk_gated_delta_"; - std::string suffix = get_type_string(q.dtype()) // "float" - + "_" + get_type_string(h0.dtype()) // "float" - + "_" + std::to_string(Dk) - + "_" + std::to_string(Dv) - + "_" + std::to_string(Hk) - + "_" + std::to_string(Hv); - - // printf("C: %d\nname: %s\n",C,kernel_name.c_str()); - std::string base_name = kernel_name + suffix; - - base_name += C >= 1 ? "_" + std::to_string(C) : ""; - - std::string hash_name = base_name; - - metal::MTLFCList func_consts = {}; - - auto delta_kernel = get_steel_gated_delta_forward_kernel( - d, - base_name, - hash_name, - func_consts); - - auto& compute_encoder = metal::get_command_encoder(s); - - - out.set_data(allocator::malloc(out.nbytes())); - hf.set_data(allocator::malloc(hf.nbytes())); - - - if (C > 1){ - // allocate full W and U -- [B, T, H, D] - array W({B, T, Hk, Dk}, q.dtype(), nullptr, {}); - array U({B, T, Hv, Dv}, q.dtype(), nullptr, {}); - W.set_data(allocator::malloc(W.nbytes())); - U.set_data(allocator::malloc(U.nbytes())); - - compute_encoder.add_temporary(W); - compute_encoder.add_temporary(U); - - // kernel 1: compute full W and U - std::string make_wy_name = "make_wy_" - + get_type_string(q.dtype()) // "float" - + "_" + std::to_string(Dk) - + "_" + std::to_string(Dv) - + "_" + std::to_string(Hk) - + "_" + std::to_string(Hv) - + "_" + std::to_string(C); - - auto wy_kernel = d.get_kernel(make_wy_name); - - compute_encoder.set_compute_pipeline_state(wy_kernel); - compute_encoder.set_input_array(k, 0); - compute_encoder.set_input_array(v, 1); - compute_encoder.set_input_array(g, 2); - compute_encoder.set_input_array(beta, 3); - compute_encoder.set_output_array(W, 4); - compute_encoder.set_output_array(U, 5); - compute_encoder.set_bytes(T, 6); - - auto grid_wy = MTL::Size(32, 4, B * Hv * n_chunks); - auto threads_wy = MTL::Size(32, 4, 1); - compute_encoder.dispatch_threads(grid_wy, threads_wy); - - // kernel 2: gated delta - compute_encoder.set_compute_pipeline_state(delta_kernel); - compute_encoder.set_input_array(q, 0); - compute_encoder.set_input_array(k, 1); - compute_encoder.set_input_array(W, 2); - compute_encoder.set_input_array(U, 3); - compute_encoder.set_input_array(h0, 4); // initial state in - compute_encoder.set_input_array(g, 5); - compute_encoder.set_output_array(out, 6); - compute_encoder.set_output_array(hf, 7); // final state out - compute_encoder.set_bytes(T, 8); - compute_encoder.set_bytes(n_chunks, 9); - - auto grid = MTL::Size(32, Dv / 8, B * Hv); - auto threads = MTL::Size(32, 1, 1); // get one simdgroup to work - // auto grid = MTL::Size(1, 1, 1); - // auto threads = MTL::Size(1, 1, 1); - compute_encoder.dispatch_threads(grid, threads); - - } else { - compute_encoder.set_compute_pipeline_state(delta_kernel); - - compute_encoder.set_input_array(q, 0); - compute_encoder.set_input_array(k, 1); - compute_encoder.set_input_array(v, 2); - compute_encoder.set_input_array(g, 3); - compute_encoder.set_input_array(beta, 4); - compute_encoder.set_input_array(h0, 5); - compute_encoder.set_bytes(T, 6); - compute_encoder.set_output_array(out, 7); - compute_encoder.set_output_array(hf, 8); - - // auto grid = MTL::Size(1, 1, 1); - // auto threads = MTL::Size(1, 1, 1); - auto grid = MTL::Size(32, Dv, B * Hv); - auto threads = MTL::Size(32, 4, 1); - compute_encoder.dispatch_threads(grid, threads); - } - // throw std::runtime_error("NYI"); + auto& s = stream(); + auto& d = metal::device(s.device); + + auto& q = inputs[0]; + auto& k = inputs[1]; + auto& v = inputs[2]; + auto& g = inputs[3]; + auto& beta = inputs[4]; + auto& h0 = inputs[5]; + + auto& out = outputs[0]; + auto& hf = outputs[1]; + + int B = q.shape(0); + int T = q.shape(1); + int Hk = q.shape(2); + int Dk = q.shape(3); + int Hv = v.shape(2); + int Dv = v.shape(3); + + int C = chunk_size; + int n_chunks = T / C; // TODO: make general + + std::string kernel_name = C < 1 ? "seq_gated_delta_" : "chunk_gated_delta_"; + std::string suffix = get_type_string(q.dtype()) // "float" + + "_" + get_type_string(h0.dtype()) // "float" + + "_" + std::to_string(Dk) + "_" + std::to_string(Dv) + "_" + + std::to_string(Hk) + "_" + std::to_string(Hv); + + // printf("C: %d\nname: %s\n",C,kernel_name.c_str()); + std::string base_name = kernel_name + suffix; + + base_name += C >= 1 ? "_" + std::to_string(C) : ""; + + std::string hash_name = base_name; + + metal::MTLFCList func_consts = {}; + + auto delta_kernel = get_steel_gated_delta_forward_kernel( + d, base_name, hash_name, func_consts); + + auto& compute_encoder = metal::get_command_encoder(s); + + out.set_data(allocator::malloc(out.nbytes())); + hf.set_data(allocator::malloc(hf.nbytes())); + + if (C > 1) { + // allocate full W and U -- [B, T, H, D] + array W({B, T, Hk, Dk}, q.dtype(), nullptr, {}); + array U({B, T, Hv, Dv}, q.dtype(), nullptr, {}); + W.set_data(allocator::malloc(W.nbytes())); + U.set_data(allocator::malloc(U.nbytes())); + + compute_encoder.add_temporary(W); + compute_encoder.add_temporary(U); + + // kernel 1: compute full W and U + std::string make_wy_name = "make_wy_" + + get_type_string(q.dtype()) // "float" + + "_" + std::to_string(Dk) + "_" + std::to_string(Dv) + "_" + + std::to_string(Hk) + "_" + std::to_string(Hv) + "_" + std::to_string(C); + + auto wy_kernel = d.get_kernel(make_wy_name); + + compute_encoder.set_compute_pipeline_state(wy_kernel); + compute_encoder.set_input_array(k, 0); + compute_encoder.set_input_array(v, 1); + compute_encoder.set_input_array(g, 2); + compute_encoder.set_input_array(beta, 3); + compute_encoder.set_output_array(W, 4); + compute_encoder.set_output_array(U, 5); + compute_encoder.set_bytes(T, 6); + + auto grid_wy = MTL::Size(32, 4, B * Hv * n_chunks); + auto threads_wy = MTL::Size(32, 4, 1); + compute_encoder.dispatch_threads(grid_wy, threads_wy); + + // kernel 2: gated delta + compute_encoder.set_compute_pipeline_state(delta_kernel); + compute_encoder.set_input_array(q, 0); + compute_encoder.set_input_array(k, 1); + compute_encoder.set_input_array(W, 2); + compute_encoder.set_input_array(U, 3); + compute_encoder.set_input_array(h0, 4); // initial state in + compute_encoder.set_input_array(g, 5); + compute_encoder.set_output_array(out, 6); + compute_encoder.set_output_array(hf, 7); // final state out + compute_encoder.set_bytes(T, 8); + compute_encoder.set_bytes(n_chunks, 9); + + auto grid = MTL::Size(32, Dv / 8, B * Hv); + auto threads = MTL::Size(32, 1, 1); // get one simdgroup to work + // auto grid = MTL::Size(1, 1, 1); + // auto threads = MTL::Size(1, 1, 1); + compute_encoder.dispatch_threads(grid, threads); + + } else { + compute_encoder.set_compute_pipeline_state(delta_kernel); + + compute_encoder.set_input_array(q, 0); + compute_encoder.set_input_array(k, 1); + compute_encoder.set_input_array(v, 2); + compute_encoder.set_input_array(g, 3); + compute_encoder.set_input_array(beta, 4); + compute_encoder.set_input_array(h0, 5); + compute_encoder.set_bytes(T, 6); + compute_encoder.set_output_array(out, 7); + compute_encoder.set_output_array(hf, 8); + + // auto grid = MTL::Size(1, 1, 1); + // auto threads = MTL::Size(1, 1, 1); + auto grid = MTL::Size(32, Dv, B * Hv); + auto threads = MTL::Size(32, 4, 1); + compute_encoder.dispatch_threads(grid, threads); + } + // throw std::runtime_error("NYI"); } bool GatedDeltaUpdate::is_equivalent(const Primitive& other) const { - const auto* p = dynamic_cast(&other); - if (p == nullptr) { - return false; - } - // TODO: finish implementation - return true; + const auto* p = dynamic_cast(&other); + if (p == nullptr) { + return false; + } + // TODO: finish implementation + return true; } } // namespace mlx::core::fast diff --git a/mlx/backend/metal/kernels.h b/mlx/backend/metal/kernels.h index 7fb8e3be80..e09bffda3d 100644 --- a/mlx/backend/metal/kernels.h +++ b/mlx/backend/metal/kernels.h @@ -423,7 +423,7 @@ MTL::ComputePipelineState* get_steel_gated_delta_forward_kernel( const std::string& hash_name, const metal::MTLFCList& func_consts // TODO: More parameters? - ); +); // Create a GPU kernel template definition for JIT compilation template diff --git a/mlx/backend/metal/kernels/gated_delta_update.metal b/mlx/backend/metal/kernels/gated_delta_update.metal index 8edba33b0e..8a87d9fbb1 100644 --- a/mlx/backend/metal/kernels/gated_delta_update.metal +++ b/mlx/backend/metal/kernels/gated_delta_update.metal @@ -1,12 +1,13 @@ -#include "mlx/backend/metal/kernels/utils.h" #include "mlx/backend/metal/kernels/gated_delta_update_impl.h" +#include "mlx/backend/metal/kernels/utils.h" using namespace metal; -#define instantiate_gated_delta_update_seq(in_type, st_type, dk, dv, hk, hv) \ +#define instantiate_gated_delta_update_seq(in_type, st_type, dk, dv, hk, hv) \ instantiate_kernel( \ - "seq_gated_delta_" #in_type "_" #st_type "_" #dk "_" #dv "_" #hk "_" #hv, \ - gated_delta_seq, \ + "seq_gated_delta_" #in_type "_" #st_type "_" #dk "_" #dv "_" #hk \ + "_" #hv, \ + gated_delta_seq, \ in_type, \ st_type, \ dk, \ @@ -17,42 +18,45 @@ using namespace metal; #define instantiate_gated_delta_update_seq_dims(in_type, st_type) \ instantiate_gated_delta_update_seq(in_type, st_type, 64, 64, 4, 4) - -#define instantiate_gated_delta_update_chunk(in_type, st_type, dk, dv, hk, hv, c) \ - instantiate_kernel( \ - "chunk_gated_delta_" #in_type "_" #st_type "_" #dk "_" #dv "_" #hk "_" #hv "_" #c, \ - gated_delta_chunk, \ - in_type, \ - st_type, \ - dk, \ - dv, \ - hk, \ - hv, \ +#define instantiate_gated_delta_update_chunk( \ + in_type, st_type, dk, dv, hk, hv, c) \ + instantiate_kernel( \ + "chunk_gated_delta_" #in_type "_" #st_type "_" #dk "_" #dv "_" #hk \ + "_" #hv "_" #c, \ + gated_delta_chunk, \ + in_type, \ + st_type, \ + dk, \ + dv, \ + hk, \ + hv, \ c) -#define instantiate_gated_delta_update_chunk_dims(in_type, st_type) \ - instantiate_gated_delta_update_chunk(in_type, st_type, 64, 64, 4, 4, 32) \ - instantiate_gated_delta_update_chunk(in_type, st_type, 64, 64, 4, 4, 16) \ - instantiate_gated_delta_update_chunk(in_type, st_type, 64, 64, 4, 4, 8) \ - instantiate_gated_delta_update_chunk(in_type, st_type, 64, 64, 1, 1, 8) \ - -#define instantiate_make_wy(in_type, dk, dv, hk, hv, c) \ - instantiate_kernel( \ - "make_wy_" #in_type "_" #dk "_" #dv "_" #hk "_" #hv "_" #c, \ - make_wy, \ - in_type, \ - dk, \ - dv, \ - hk, \ - hv, \ +#define instantiate_gated_delta_update_chunk_dims(in_type, st_type) \ + instantiate_gated_delta_update_chunk(in_type, st_type, 64, 64, 4, 4, 32) \ + instantiate_gated_delta_update_chunk(in_type, st_type, 64, 64, 4, 4, 16) \ + instantiate_gated_delta_update_chunk( \ + in_type, st_type, 64, 64, 4, 4, 8) \ + instantiate_gated_delta_update_chunk( \ + in_type, st_type, 64, 64, 1, 1, 8) + +#define instantiate_make_wy(in_type, dk, dv, hk, hv, c) \ + instantiate_kernel( \ + "make_wy_" #in_type "_" #dk "_" #dv "_" #hk "_" #hv "_" #c, \ + make_wy, \ + in_type, \ + dk, \ + dv, \ + hk, \ + hv, \ c) -#define instantiate_make_wy_dims(in_type) \ - instantiate_make_wy(in_type, 64, 64, 4, 4, 32) \ - instantiate_make_wy(in_type, 64, 64, 4, 4, 16) \ - instantiate_make_wy(in_type, 64, 64, 4, 4, 8) \ - instantiate_make_wy(in_type, 64, 64, 1, 1, 8) \ +#define instantiate_make_wy_dims(in_type) \ + instantiate_make_wy(in_type, 64, 64, 4, 4, 32) \ + instantiate_make_wy(in_type, 64, 64, 4, 4, 16) \ + instantiate_make_wy(in_type, 64, 64, 4, 4, 8) \ + instantiate_make_wy(in_type, 64, 64, 1, 1, 8) instantiate_gated_delta_update_seq_dims(float, float) -instantiate_gated_delta_update_chunk_dims(float, float) -instantiate_make_wy_dims(float) \ No newline at end of file + instantiate_gated_delta_update_chunk_dims(float, float) + instantiate_make_wy_dims(float) \ No newline at end of file diff --git a/mlx/backend/metal/kernels/gated_delta_update_impl.h b/mlx/backend/metal/kernels/gated_delta_update_impl.h index 5865445146..2d87be29c3 100644 --- a/mlx/backend/metal/kernels/gated_delta_update_impl.h +++ b/mlx/backend/metal/kernels/gated_delta_update_impl.h @@ -15,368 +15,365 @@ using namespace metal; template [[kernel]] void make_wy( - const device InT* k [[buffer(0)]], - const device InT* v [[buffer(1)]], - const device InT* g [[buffer(2)]], // [B, T, Hv] or [B, T, Hv, Dk] - const device InT* beta [[buffer(3)]], // [B, Hv, Dv, Dk] - device InT* W [[buffer(4)]], // [B, T, Hv] - device InT* U [[buffer(5)]], // [B, T, Hv] - constant int& T [[buffer(6)]], - uint3 thread_position_in_grid [[thread_position_in_grid]], - uint3 thread_position_in_threadgroup [[thread_position_in_threadgroup]], - uint thread_index_in_simdgroup [[thread_index_in_simdgroup]] -) { - auto n = thread_position_in_grid.z; - auto n_chunks = T / C; - auto chunk = n % n_chunks; - auto bh_idx = n / n_chunks; - auto b_idx = bh_idx / Hv; - auto hv_idx = bh_idx % Hv; - auto hk_idx = hv_idx / (Hv / Hk); - - constexpr int n_per_dk = Dk / 32; - constexpr int n_per_dv = Dv / 32; - - auto dk_idx = thread_position_in_threadgroup.x; - auto simd_group_id = thread_position_in_threadgroup.y; - const int num_simdgroups = 4; - - int offset_t = chunk * C; - auto g_ = g + b_idx * T * Hv + offset_t * Hv; - auto k_ = k + b_idx * T * Hk * Dk + offset_t * Hk * Dk + hk_idx * Dk; - auto v_ = v + b_idx * T * Hv * Dv + offset_t * Hv * Dv + hv_idx * Dv; - auto beta_ = beta + b_idx * T * Hv + offset_t * Hv; - auto W_ = W + b_idx * T * Hv * Dk + offset_t * Hv * Dk + hv_idx * Dk; - auto U_ = U + b_idx * T * Hv * Dv + offset_t * Hv * Dv + hv_idx * Dv; - - // threadgroup memory - threadgroup float K_tg[C][Dk]; - threadgroup float KKt[C][C]; - - for (int i = simd_group_id; i < C; i += num_simdgroups) { - for (int d = dk_idx; d < Dk; d += 32) { - K_tg[i][d] = k_[i * Hk * Dk + d]; - } - } - threadgroup_barrier(mem_flags::mem_threadgroup); - - // compute gamma (cumprod) - float gamma[C]; - gamma[0] = g_[hv_idx]; - for (int i = 1; i < C; i++) { - gamma[i] = gamma[i-1] * g_[i * Hv + hv_idx]; - } - - // compute KKt = K @ K.T - unsigned int counter = 0; - for (int i = 0; i < C; i++) { - for (int j = 0; j <= i; j++) { - if (counter % num_simdgroups == simd_group_id) { - float sum = 0; - for (int d = dk_idx; d < Dk; d += 32) { - sum += K_tg[i][d] * K_tg[j][d]; - } - sum = simd_sum(sum); - if (thread_index_in_simdgroup == 0) { - KKt[i][j] = sum; - } - } - counter++; - } - } - threadgroup_barrier(mem_flags::mem_threadgroup); - - float u[C][n_per_dv]; - float w[C][n_per_dk]; - - // initialize - for (int i = 0; i < C; i++) { - float beta_i = beta_[i * Hv + hv_idx]; - for (int p = 0; p < n_per_dv; p++) { - int dv = dk_idx * n_per_dv + p; - u[i][p] = beta_i * v_[i * Hv * Dv + dv]; - } - for (int p = 0; p < n_per_dk; p++) { - int dk = dk_idx * n_per_dk + p; - w[i][p] = beta_i * K_tg[i][dk]; - } - } - - - // forward substitution - for (int i = 1; i < C; i++) { - float beta_i = beta_[i * Hv + hv_idx]; - for (int j = 0; j < i; j++) { - float a_U = beta_i * (gamma[i] / gamma[j]) * KKt[i][j]; - float a_W = beta_i * KKt[i][j]; - for (int p = 0; p < n_per_dv; p++) { - u[i][p] -= a_U * u[j][p]; - } - for (int p = 0; p < n_per_dk; p++) { - w[i][p] -= a_W * w[j][p]; - } - } - } - - for (int i = 0; i < C; i++) { - for (int p = 0; p < n_per_dv; p++) { - int dv = dk_idx * n_per_dv + p; - U_[i * Hv * Dv + dv] = u[i][p]; - } - for (int p = 0; p < n_per_dk; p++) { - int dk = dk_idx * n_per_dk + p; - W_[i * Hk * Dk + dk] = w[i][p]; - } - } -} - -template -[[kernel]] void gated_delta_chunk( - const device InT* q [[buffer(0)]], - const device InT* k [[buffer(1)]], - const device InT* W [[buffer(2)]], // [B, T, Hv, Dk] - const device InT* U [[buffer(3)]], // [B, T, Hv, Dv] - const device StT* state_in [[buffer(4)]], // [B, Hv, Dv, Dk] - const device InT* g [[buffer(5)]], // [B, T, Hv] - device InT* y [[buffer(6)]], // [B, T, Hv, Dv] - device StT* state_out [[buffer(7)]], // [B, Hv, Dv, Dk] - constant int& T [[buffer(8)]], - uint3 thread_position_in_grid [[thread_position_in_grid]], + const device InT* k [[buffer(0)]], + const device InT* v [[buffer(1)]], + const device InT* g [[buffer(2)]], // [B, T, Hv] or [B, T, Hv, Dk] + const device InT* beta [[buffer(3)]], // [B, Hv, Dv, Dk] + device InT* W [[buffer(4)]], // [B, T, Hv] + device InT* U [[buffer(5)]], // [B, T, Hv] + constant int& T [[buffer(6)]], + uint3 thread_position_in_grid [[thread_position_in_grid]], uint3 thread_position_in_threadgroup [[thread_position_in_threadgroup]], - uint thread_index_in_simdgroup [[thread_index_in_simdgroup]] -) { - auto n = thread_position_in_grid.z; - auto b_idx = n / Hv; - auto hv_idx = n % Hv; - auto hk_idx = hv_idx / (Hv / Hk); - - // get coord from steel - const short qid = thread_index_in_simdgroup / 4; - const short fm = (qid & 4) + ((thread_index_in_simdgroup / 2) % 4); - const short fn = (qid & 2) * 2 + (thread_index_in_simdgroup % 2) * 2; - - auto dk_idx = thread_position_in_threadgroup.x; // simd group id - auto dv_idx = thread_position_in_grid.y * 8; - - // set up pointers - // g: [B, T, Hv] - auto g_ = g + b_idx * T * Hv; - - // q, k: [B, T, Hk, Dk] - auto q_ = q + b_idx * T * Hk * Dk + hk_idx * Dk; - auto k_ = k + b_idx * T * Hk * Dk + hk_idx * Dk; - - // W: [B, T, Hk, Dk], U: [B, T, Hv, Dv] - auto W_ = W + b_idx * T * Hv * Dk + hv_idx * Dk; - auto U_ = U + b_idx * T * Hv * Dv + hv_idx * Dv; - - // v, y: [B, T, Hv, Dv] - y += b_idx * T * Hv * Dv + hv_idx * Dv; - - // state_in, state_out: [B, Hv, Dv, Dk] - auto i_state = state_in + (n * Dv + dv_idx) * Dk; - auto o_state = state_out + (n * Dv + dv_idx) * Dk; - - // threadgroup tiles - threadgroup float Q_left_tg[C][Dk]; // gamma-scaled Q - threadgroup float K_tg[C][Dk]; // raw K - threadgroup float W_tg[C][Dk]; // gamma-scaled W (W_left) - threadgroup float U_tg[C][8]; // U chunk, 8 dv values - threadgroup float K_right_tg[C][Dk]; - - simdgroup_float8x8 S_tile[8]; - simdgroup_float8x8 W_tile, K_tile, Q_tile; - simdgroup_float8x8 WS_tile; - simdgroup_float8x8 U_tile; - simdgroup_float8x8 delta_tile; - simdgroup_float8x8 tmp_tile; - simdgroup_float8x8 QKt_tile; - simdgroup_float8x8 out_tile; - simdgroup_float8x8 KD_tile; - - thread auto& kd = KD_tile.thread_elements(); - thread auto& d_e = delta_tile.thread_elements(); - thread auto& u_e = U_tile.thread_elements(); - thread auto& ws_e = WS_tile.thread_elements(); - thread auto& qkt_e = QKt_tile.thread_elements(); - thread auto& o_e = out_tile.thread_elements(); - - // load initial state into threadgroup - for (int kk = 0; kk < Dk; kk += 8) { - simdgroup_load(S_tile[kk/8], i_state + kk, Dk, ulong2(0, 0), true); - } - threadgroup_barrier(mem_flags::mem_threadgroup); - - - float gamma[C]; - - for (int t = 0; t < T; t += C) { - gamma[0] = g_[hv_idx]; - for (int i = 1; i < C; i++) { - gamma[i] = gamma[i-1] * g_[i * Hv + hv_idx]; + uint thread_index_in_simdgroup [[thread_index_in_simdgroup]]) { + auto n = thread_position_in_grid.z; + auto n_chunks = T / C; + auto chunk = n % n_chunks; + auto bh_idx = n / n_chunks; + auto b_idx = bh_idx / Hv; + auto hv_idx = bh_idx % Hv; + auto hk_idx = hv_idx / (Hv / Hk); + + constexpr int n_per_dk = Dk / 32; + constexpr int n_per_dv = Dv / 32; + + auto dk_idx = thread_position_in_threadgroup.x; + auto simd_group_id = thread_position_in_threadgroup.y; + const int num_simdgroups = 4; + + int offset_t = chunk * C; + auto g_ = g + b_idx * T * Hv + offset_t * Hv; + auto k_ = k + b_idx * T * Hk * Dk + offset_t * Hk * Dk + hk_idx * Dk; + auto v_ = v + b_idx * T * Hv * Dv + offset_t * Hv * Dv + hv_idx * Dv; + auto beta_ = beta + b_idx * T * Hv + offset_t * Hv; + auto W_ = W + b_idx * T * Hv * Dk + offset_t * Hv * Dk + hv_idx * Dk; + auto U_ = U + b_idx * T * Hv * Dv + offset_t * Hv * Dv + hv_idx * Dv; + + // threadgroup memory + threadgroup float K_tg[C][Dk]; + threadgroup float KKt[C][C]; + + for (int i = simd_group_id; i < C; i += num_simdgroups) { + for (int d = dk_idx; d < Dk; d += 32) { + K_tg[i][d] = k_[i * Hk * Dk + d]; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + // compute gamma (cumprod) + float gamma[C]; + gamma[0] = g_[hv_idx]; + for (int i = 1; i < C; i++) { + gamma[i] = gamma[i - 1] * g_[i * Hv + hv_idx]; + } + + // compute KKt = K @ K.T + unsigned int counter = 0; + for (int i = 0; i < C; i++) { + for (int j = 0; j <= i; j++) { + if (counter % num_simdgroups == simd_group_id) { + float sum = 0; + for (int d = dk_idx; d < Dk; d += 32) { + sum += K_tg[i][d] * K_tg[j][d]; } - float gamma_C = gamma[C-1]; - - for (int c = 0; c < C; c++) { - float kr = gamma_C / gamma[c]; - for (int d = dk_idx; d < Dk; d += 32) { - Q_left_tg[c][d] = gamma[c] * q_[c * Hk * Dk + d]; - K_tg[c][d] = k_[c * Hk * Dk + d]; - W_tg[c][d] = gamma[c] * W_[c * Hv * Dk + d]; - K_right_tg[c][d] = kr * K_tg[c][d]; - } + sum = simd_sum(sum); + if (thread_index_in_simdgroup == 0) { + KKt[i][j] = sum; } + } + counter++; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + float u[C][n_per_dv]; + float w[C][n_per_dk]; + + // initialize + for (int i = 0; i < C; i++) { + float beta_i = beta_[i * Hv + hv_idx]; + for (int p = 0; p < n_per_dv; p++) { + int dv = dk_idx * n_per_dv + p; + u[i][p] = beta_i * v_[i * Hv * Dv + dv]; + } + for (int p = 0; p < n_per_dk; p++) { + int dk = dk_idx * n_per_dk + p; + w[i][p] = beta_i * K_tg[i][dk]; + } + } + + // forward substitution + for (int i = 1; i < C; i++) { + float beta_i = beta_[i * Hv + hv_idx]; + for (int j = 0; j < i; j++) { + float a_U = beta_i * (gamma[i] / gamma[j]) * KKt[i][j]; + float a_W = beta_i * KKt[i][j]; + for (int p = 0; p < n_per_dv; p++) { + u[i][p] -= a_U * u[j][p]; + } + for (int p = 0; p < n_per_dk; p++) { + w[i][p] -= a_W * w[j][p]; + } + } + } - for (int idx = dk_idx; idx < C * 8; idx += 32) { - int c = idx / 8; - int j = idx % 8; - U_tg[c][j] = U_[c * Hv * Dv + dv_idx + j]; - } - threadgroup_barrier(mem_flags::mem_threadgroup); + for (int i = 0; i < C; i++) { + for (int p = 0; p < n_per_dv; p++) { + int dv = dk_idx * n_per_dv + p; + U_[i * Hv * Dv + dv] = u[i][p]; + } + for (int p = 0; p < n_per_dk; p++) { + int dk = dk_idx * n_per_dk + p; + W_[i * Hk * Dk + dk] = w[i][p]; + } + } +} - // WS = W_left @ S^T - WS_tile = make_filled_simdgroup_matrix(0.f); - for (int kk = 0; kk < Dk; kk += 8) { - simdgroup_load(W_tile, &W_tg[0][kk], Dk); - simdgroup_multiply_accumulate(WS_tile, W_tile, S_tile[kk/8], WS_tile); - } +template +[[kernel]] void gated_delta_chunk( + const device InT* q [[buffer(0)]], + const device InT* k [[buffer(1)]], + const device InT* W [[buffer(2)]], // [B, T, Hv, Dk] + const device InT* U [[buffer(3)]], // [B, T, Hv, Dv] + const device StT* state_in [[buffer(4)]], // [B, Hv, Dv, Dk] + const device InT* g [[buffer(5)]], // [B, T, Hv] + device InT* y [[buffer(6)]], // [B, T, Hv, Dv] + device StT* state_out [[buffer(7)]], // [B, Hv, Dv, Dk] + constant int& T [[buffer(8)]], + uint3 thread_position_in_grid [[thread_position_in_grid]], + uint3 thread_position_in_threadgroup [[thread_position_in_threadgroup]], + uint thread_index_in_simdgroup [[thread_index_in_simdgroup]]) { + auto n = thread_position_in_grid.z; + auto b_idx = n / Hv; + auto hv_idx = n % Hv; + auto hk_idx = hv_idx / (Hv / Hk); + + // get coord from steel + const short qid = thread_index_in_simdgroup / 4; + const short fm = (qid & 4) + ((thread_index_in_simdgroup / 2) % 4); + const short fn = (qid & 2) * 2 + (thread_index_in_simdgroup % 2) * 2; + + auto dk_idx = thread_position_in_threadgroup.x; // simd group id + auto dv_idx = thread_position_in_grid.y * 8; + + // set up pointers + // g: [B, T, Hv] + auto g_ = g + b_idx * T * Hv; + + // q, k: [B, T, Hk, Dk] + auto q_ = q + b_idx * T * Hk * Dk + hk_idx * Dk; + auto k_ = k + b_idx * T * Hk * Dk + hk_idx * Dk; + + // W: [B, T, Hk, Dk], U: [B, T, Hv, Dv] + auto W_ = W + b_idx * T * Hv * Dk + hv_idx * Dk; + auto U_ = U + b_idx * T * Hv * Dv + hv_idx * Dv; + + // v, y: [B, T, Hv, Dv] + y += b_idx * T * Hv * Dv + hv_idx * Dv; + + // state_in, state_out: [B, Hv, Dv, Dk] + auto i_state = state_in + (n * Dv + dv_idx) * Dk; + auto o_state = state_out + (n * Dv + dv_idx) * Dk; + + // threadgroup tiles + threadgroup float Q_left_tg[C][Dk]; // gamma-scaled Q + threadgroup float K_tg[C][Dk]; // raw K + threadgroup float W_tg[C][Dk]; // gamma-scaled W (W_left) + threadgroup float U_tg[C][8]; // U chunk, 8 dv values + threadgroup float K_right_tg[C][Dk]; + + simdgroup_float8x8 S_tile[8]; + simdgroup_float8x8 W_tile, K_tile, Q_tile; + simdgroup_float8x8 WS_tile; + simdgroup_float8x8 U_tile; + simdgroup_float8x8 delta_tile; + simdgroup_float8x8 tmp_tile; + simdgroup_float8x8 QKt_tile; + simdgroup_float8x8 out_tile; + simdgroup_float8x8 KD_tile; + + thread auto& kd = KD_tile.thread_elements(); + thread auto& d_e = delta_tile.thread_elements(); + thread auto& u_e = U_tile.thread_elements(); + thread auto& ws_e = WS_tile.thread_elements(); + thread auto& qkt_e = QKt_tile.thread_elements(); + thread auto& o_e = out_tile.thread_elements(); + + // load initial state into threadgroup + for (int kk = 0; kk < Dk; kk += 8) { + simdgroup_load(S_tile[kk / 8], i_state + kk, Dk, ulong2(0, 0), true); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + float gamma[C]; + + for (int t = 0; t < T; t += C) { + gamma[0] = g_[hv_idx]; + for (int i = 1; i < C; i++) { + gamma[i] = gamma[i - 1] * g_[i * Hv + hv_idx]; + } + float gamma_C = gamma[C - 1]; + + for (int c = 0; c < C; c++) { + float kr = gamma_C / gamma[c]; + for (int d = dk_idx; d < Dk; d += 32) { + Q_left_tg[c][d] = gamma[c] * q_[c * Hk * Dk + d]; + K_tg[c][d] = k_[c * Hk * Dk + d]; + W_tg[c][d] = gamma[c] * W_[c * Hv * Dk + d]; + K_right_tg[c][d] = kr * K_tg[c][d]; + } + } - // delta = U - WS - simdgroup_load(U_tile, &U_tg[0][0], 8); - - d_e[0] = u_e[0] - ws_e[0]; - d_e[1] = u_e[1] - ws_e[1]; - + for (int idx = dk_idx; idx < C * 8; idx += 32) { + int c = idx / 8; + int j = idx % 8; + U_tg[c][j] = U_[c * Hv * Dv + dv_idx + j]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); - // Q_left @ S^T - tmp_tile = make_filled_simdgroup_matrix(0.f); - QKt_tile = make_filled_simdgroup_matrix(0.f); + // WS = W_left @ S^T + WS_tile = make_filled_simdgroup_matrix(0.f); + for (int kk = 0; kk < Dk; kk += 8) { + simdgroup_load(W_tile, &W_tg[0][kk], Dk); + simdgroup_multiply_accumulate(WS_tile, W_tile, S_tile[kk / 8], WS_tile); + } - for (int kk = 0; kk < Dk; kk += 8) { - simdgroup_load(Q_tile, &Q_left_tg[0][kk], Dk); - simdgroup_load(K_tile, &K_tg[0][kk], Dk, ulong2(0, 0), true); + // delta = U - WS + simdgroup_load(U_tile, &U_tg[0][0], 8); - simdgroup_multiply_accumulate(tmp_tile, Q_tile, S_tile[kk/8], tmp_tile); - simdgroup_multiply_accumulate(QKt_tile, Q_tile, K_tile, QKt_tile); - } + d_e[0] = u_e[0] - ws_e[0]; + d_e[1] = u_e[1] - ws_e[1]; + // Q_left @ S^T + tmp_tile = make_filled_simdgroup_matrix(0.f); + QKt_tile = make_filled_simdgroup_matrix(0.f); - // element 0 at (fm, fn), element 1 at (fm, fn+1) - qkt_e[0] *= fn > fm ? 0.f : (1.0f / gamma[fn]); - qkt_e[1] *= fn + 1 > fm ? 0.f : (1.0f / gamma[fn+1]); + for (int kk = 0; kk < Dk; kk += 8) { + simdgroup_load(Q_tile, &Q_left_tg[0][kk], Dk); + simdgroup_load(K_tile, &K_tg[0][kk], Dk, ulong2(0, 0), true); - simdgroup_multiply_accumulate(out_tile, QKt_tile, delta_tile, tmp_tile); + simdgroup_multiply_accumulate(tmp_tile, Q_tile, S_tile[kk / 8], tmp_tile); + simdgroup_multiply_accumulate(QKt_tile, Q_tile, K_tile, QKt_tile); + } - y[fm * Hv * Dv + dv_idx + fn] = static_cast(o_e[0]); - y[fm * Hv * Dv + dv_idx + fn + 1] = static_cast(o_e[1]); + // element 0 at (fm, fn), element 1 at (fm, fn+1) + qkt_e[0] *= fn > fm ? 0.f : (1.0f / gamma[fn]); + qkt_e[1] *= fn + 1 > fm ? 0.f : (1.0f / gamma[fn + 1]); + simdgroup_multiply_accumulate(out_tile, QKt_tile, delta_tile, tmp_tile); - for (int kk = 0; kk < Dk; kk += 8) { - simdgroup_load(K_tile, &K_right_tg[0][kk], Dk, ulong2(0, 0), true); - simdgroup_multiply(KD_tile, K_tile, delta_tile); + y[fm * Hv * Dv + dv_idx + fn] = static_cast(o_e[0]); + y[fm * Hv * Dv + dv_idx + fn + 1] = static_cast(o_e[1]); - thread auto& s_e = S_tile[kk/8].thread_elements(); - s_e[0] = gamma_C * s_e[0] + kd[0]; - s_e[1] = gamma_C * s_e[1] + kd[1]; - } + for (int kk = 0; kk < Dk; kk += 8) { + simdgroup_load(K_tile, &K_right_tg[0][kk], Dk, ulong2(0, 0), true); + simdgroup_multiply(KD_tile, K_tile, delta_tile); - // update pointers - q_ += C * Hk * Dk; - k_ += C * Hk * Dk; - U_ += C * Hv * Dv; - W_ += C * Hv * Dk; - y += C * Hv * Dv; - g_ += C * Hv; + thread auto& s_e = S_tile[kk / 8].thread_elements(); + s_e[0] = gamma_C * s_e[0] + kd[0]; + s_e[1] = gamma_C * s_e[1] + kd[1]; } - // o_state is [Dv, Dk]: o_state[dv][dk] = st_out[dk][dv] - for (int kk = 0; kk < Dk; kk += 8) { - simdgroup_store(S_tile[kk/8], o_state + kk, Dk, ulong2(0,0), true); // writes [Dk_block, Dv] layout - } + // update pointers + q_ += C * Hk * Dk; + k_ += C * Hk * Dk; + U_ += C * Hv * Dv; + W_ += C * Hv * Dk; + y += C * Hv * Dv; + g_ += C * Hv; + } + + // o_state is [Dv, Dk]: o_state[dv][dk] = st_out[dk][dv] + for (int kk = 0; kk < Dk; kk += 8) { + simdgroup_store( + S_tile[kk / 8], + o_state + kk, + Dk, + ulong2(0, 0), + true); // writes [Dk_block, Dv] layout + } } /* - auto grid = MTL::Size(32, Dv, B * Hv); + auto grid = MTL::Size(32, Dv, B * Hv); auto threads = MTL::Size(32, 4, 1); */ template [[kernel]] void gated_delta_seq( - const device InT* q [[buffer(0)]], - const device InT* k [[buffer(1)]], - const device InT* v [[buffer(2)]], - const device InT* g [[buffer(3)]], // [B, T, Hv] or [B, T, Hv, Dk] - const device InT* beta [[buffer(4)]], // [B, T, Hv] - const device StT* state_in [[buffer(5)]], // [B, Hv, Dv, Dk] - constant int& T [[buffer(6)]], - device InT* y [[buffer(7)]], // [B, T, Hv, Dv] - device StT* state_out [[buffer(8)]], // [B, Hv, Dv, Dk] - uint3 thread_position_in_grid [[thread_position_in_grid]], - uint3 thread_position_in_threadgroup [[thread_position_in_threadgroup]], - uint thread_index_in_simdgroup [[thread_index_in_simdgroup]] -) { - // kernel implementation - auto n = thread_position_in_grid.z; - auto b_idx = n / Hv; - auto hv_idx = n % Hv; - auto hk_idx = hv_idx / (Hv / Hk); - constexpr int n_per_t = Dk / 32; - - // q, k: [B, T, Hk, Dk] - auto q_ = q + b_idx * T * Hk * Dk + hk_idx * Dk; - auto k_ = k + b_idx * T * Hk * Dk + hk_idx * Dk; - - // v, y: [B, T, Hv, Dv] - auto v_ = v + b_idx * T * Hv * Dv + hv_idx * Dv; - y += b_idx * T * Hv * Dv + hv_idx * Dv; - - auto dk_idx = thread_position_in_threadgroup.x; - auto dv_idx = thread_position_in_grid.y; - - // state_in, state_out: [B, Hv, Dv, Dk] - auto i_state = state_in + (n * Dv + dv_idx) * Dk; - auto o_state = state_out + (n * Dv + dv_idx) * Dk; - - float state[n_per_t]; - for (int i = 0; i < n_per_t; ++i) { - auto s_idx = n_per_t * dk_idx + i; - state[i] = static_cast(i_state[s_idx]); - } - - // g: [B, T, Hv] - auto g_ = g + b_idx * T * Hv; - auto beta_ = beta + b_idx * T * Hv; - - for (int t = 0; t < T; ++t) { - float kv_mem = 0.0f; - for (int i = 0; i < n_per_t; ++i) { - auto s_idx = n_per_t * dk_idx + i; - state[i] = state[i] * g_[hv_idx]; - kv_mem += state[i] * k_[s_idx]; - } - kv_mem = simd_sum(kv_mem); - - auto delta = (v_[dv_idx] - kv_mem) * beta_[hv_idx]; - - float out = 0.0f; - for (int i = 0; i < n_per_t; ++i) { - auto s_idx = n_per_t * dk_idx + i; - state[i] = state[i] + k_[s_idx] * delta; - out += state[i] * q_[s_idx]; - } - out = simd_sum(out); - if (thread_index_in_simdgroup == 0) { - y[dv_idx] = static_cast(out); - } - // Increment data pointers to next time step - q_ += Hk * Dk; - k_ += Hk * Dk; - v_ += Hv * Dv; - y += Hv * Dv; - g_ += Hv; - beta_ += Hv; - } - for (int i = 0; i < n_per_t; ++i) { - auto s_idx = n_per_t * dk_idx + i; - o_state[s_idx] = static_cast(state[i]); - } + const device InT* q [[buffer(0)]], + const device InT* k [[buffer(1)]], + const device InT* v [[buffer(2)]], + const device InT* g [[buffer(3)]], // [B, T, Hv] or [B, T, Hv, Dk] + const device InT* beta [[buffer(4)]], // [B, T, Hv] + const device StT* state_in [[buffer(5)]], // [B, Hv, Dv, Dk] + constant int& T [[buffer(6)]], + device InT* y [[buffer(7)]], // [B, T, Hv, Dv] + device StT* state_out [[buffer(8)]], // [B, Hv, Dv, Dk] + uint3 thread_position_in_grid [[thread_position_in_grid]], + uint3 thread_position_in_threadgroup [[thread_position_in_threadgroup]], + uint thread_index_in_simdgroup [[thread_index_in_simdgroup]]) { + // kernel implementation + auto n = thread_position_in_grid.z; + auto b_idx = n / Hv; + auto hv_idx = n % Hv; + auto hk_idx = hv_idx / (Hv / Hk); + constexpr int n_per_t = Dk / 32; + + // q, k: [B, T, Hk, Dk] + auto q_ = q + b_idx * T * Hk * Dk + hk_idx * Dk; + auto k_ = k + b_idx * T * Hk * Dk + hk_idx * Dk; + + // v, y: [B, T, Hv, Dv] + auto v_ = v + b_idx * T * Hv * Dv + hv_idx * Dv; + y += b_idx * T * Hv * Dv + hv_idx * Dv; + + auto dk_idx = thread_position_in_threadgroup.x; + auto dv_idx = thread_position_in_grid.y; + + // state_in, state_out: [B, Hv, Dv, Dk] + auto i_state = state_in + (n * Dv + dv_idx) * Dk; + auto o_state = state_out + (n * Dv + dv_idx) * Dk; + + float state[n_per_t]; + for (int i = 0; i < n_per_t; ++i) { + auto s_idx = n_per_t * dk_idx + i; + state[i] = static_cast(i_state[s_idx]); + } + + // g: [B, T, Hv] + auto g_ = g + b_idx * T * Hv; + auto beta_ = beta + b_idx * T * Hv; + + for (int t = 0; t < T; ++t) { + float kv_mem = 0.0f; + for (int i = 0; i < n_per_t; ++i) { + auto s_idx = n_per_t * dk_idx + i; + state[i] = state[i] * g_[hv_idx]; + kv_mem += state[i] * k_[s_idx]; + } + kv_mem = simd_sum(kv_mem); + + auto delta = (v_[dv_idx] - kv_mem) * beta_[hv_idx]; + + float out = 0.0f; + for (int i = 0; i < n_per_t; ++i) { + auto s_idx = n_per_t * dk_idx + i; + state[i] = state[i] + k_[s_idx] * delta; + out += state[i] * q_[s_idx]; + } + out = simd_sum(out); + if (thread_index_in_simdgroup == 0) { + y[dv_idx] = static_cast(out); + } + // Increment data pointers to next time step + q_ += Hk * Dk; + k_ += Hk * Dk; + v_ += Hv * Dv; + y += Hv * Dv; + g_ += Hv; + beta_ += Hv; + } + for (int i = 0; i < n_per_t; ++i) { + auto s_idx = n_per_t * dk_idx + i; + o_state[s_idx] = static_cast(state[i]); + } } diff --git a/mlx/backend/metal/nojit_kernels.cpp b/mlx/backend/metal/nojit_kernels.cpp index 04c948f681..fa662afac0 100644 --- a/mlx/backend/metal/nojit_kernels.cpp +++ b/mlx/backend/metal/nojit_kernels.cpp @@ -498,8 +498,8 @@ MTL::ComputePipelineState* get_steel_gated_delta_forward_kernel( const std::string& hash_name, const metal::MTLFCList& func_consts // TODO: need more parameters? - ) { - return d.get_kernel(kernel_name, hash_name, func_consts); +) { + return d.get_kernel(kernel_name, hash_name, func_consts); } } // namespace mlx::core diff --git a/mlx/fast.cpp b/mlx/fast.cpp index 67ce47cc31..785d46e288 100644 --- a/mlx/fast.cpp +++ b/mlx/fast.cpp @@ -923,53 +923,50 @@ bool ScaledDotProductAttentionVJP::is_equivalent(const Primitive& other) const { } std::vector gated_delta_update_forward( - const array& queries, - const array& keys, - const array& values, - const array& gates, - const array& beta_, - const std::optional& initial_state, /* = std::nullopt */ - const int C, - StreamOrDevice s /* = {} */) { - - // determine output dtype - auto promoted = promote_types(queries.dtype(), keys.dtype()); - auto out_dtype = issubdtype(promoted, float32) - ? promoted - : promote_types(promoted, float32); - - // cast all inputs - auto q = astype(queries, out_dtype, s); - auto k = astype(keys, out_dtype, s); - auto v = astype(values, out_dtype, s); - auto g = astype(gates, out_dtype, s); - auto beta = astype(beta_, out_dtype, s); - - int B = q.shape(0); - int T = q.shape(1); - int Hk = q.shape(2); - int Dk = q.shape(3); - int Hv = v.shape(2); - int Dv = v.shape(3); - - auto h0 = initial_state.has_value() - ? astype(*initial_state, out_dtype, s) - : zeros({B, Hv, Dk, Dv}, out_dtype, s); - - auto fallback = [](std::vector inputs) -> std::vector { - throw std::runtime_error("NYI: GatedDeltaUpdate CPU fallback"); - return {}; - }; + const array& queries, + const array& keys, + const array& values, + const array& gates, + const array& beta_, + const std::optional& initial_state, /* = std::nullopt */ + const int C, + StreamOrDevice s /* = {} */) { + // determine output dtype + auto promoted = promote_types(queries.dtype(), keys.dtype()); + auto out_dtype = issubdtype(promoted, float32) + ? promoted + : promote_types(promoted, float32); + + // cast all inputs + auto q = astype(queries, out_dtype, s); + auto k = astype(keys, out_dtype, s); + auto v = astype(values, out_dtype, s); + auto g = astype(gates, out_dtype, s); + auto beta = astype(beta_, out_dtype, s); + + int B = q.shape(0); + int T = q.shape(1); + int Hk = q.shape(2); + int Dk = q.shape(3); + int Hv = v.shape(2); + int Dv = v.shape(3); + + auto h0 = initial_state.has_value() ? astype(*initial_state, out_dtype, s) + : zeros({B, Hv, Dk, Dv}, out_dtype, s); + + auto fallback = [](std::vector inputs) -> std::vector { + throw std::runtime_error("NYI: GatedDeltaUpdate CPU fallback"); + return {}; + }; - return array::make_arrays( - /* output shapes */ {{B, T, Hv, Dv}, {B, Hv, Dk, Dv}}, - /* dtypes */ {out_dtype, out_dtype}, - /* primitive */ std::make_shared(to_stream(s), fallback, C), - /* inputs */ {q, k, v, g, beta, h0} - ); + return array::make_arrays( + /* output shapes */ {{B, T, Hv, Dv}, {B, Hv, Dk, Dv}}, + /* dtypes */ {out_dtype, out_dtype}, + /* primitive */ + std::make_shared(to_stream(s), fallback, C), + /* inputs */ {q, k, v, g, beta, h0}); } - bool Quantize::is_equivalent(const Primitive& other) const { const Quantize& p_other = static_cast(other); return ( diff --git a/mlx/fast.h b/mlx/fast.h index d4cd810463..2dbd03e1d0 100644 --- a/mlx/fast.h +++ b/mlx/fast.h @@ -63,8 +63,7 @@ MLX_API std::vector gated_delta_update_forward( const array& beta_, const std::optional& initial_state = std::nullopt, const int C = 16, - StreamOrDevice s = {} -); + StreamOrDevice s = {}); using TemplateArg = std::variant; using ScalarArg = std::variant; diff --git a/mlx/fast_primitives.h b/mlx/fast_primitives.h index 305cfaa601..79ce867d82 100644 --- a/mlx/fast_primitives.h +++ b/mlx/fast_primitives.h @@ -327,13 +327,11 @@ class ConvertFP8 : public Primitive { class GatedDeltaUpdate : public Custom { public: - GatedDeltaUpdate( + GatedDeltaUpdate( Stream stream, std::function(std::vector)> fallback, - int C - ) - : Custom(stream, std::move(fallback)), chunk_size(C) - {} + int C) + : Custom(stream, std::move(fallback)), chunk_size(C) {} static bool use_fallback( /* TODO */ @@ -348,21 +346,18 @@ class GatedDeltaUpdate : public Custom { void eval_gpu(const std::vector& inputs, std::vector& outputs) override; - bool is_equivalent(const Primitive& other) const override; - + DEFINE_NAME(GatedDeltaUpdate); DEFINE_INPUT_OUTPUT_SHAPE() auto state() const { - return std::make_tuple( - nullptr); /* TODO */ + return std::make_tuple(nullptr); /* TODO */ } private: int chunk_size; }; - class Quantize : public Custom { public: explicit Quantize( diff --git a/python/src/fast.cpp b/python/src/fast.cpp index 86c971359a..5cedf8273f 100644 --- a/python/src/fast.cpp +++ b/python/src/fast.cpp @@ -332,19 +332,19 @@ void init_fast(nb::module_& parent_module) { scale = D ** -0.5 out = mx.fast.scaled_dot_product_attention(q, k, v, scale=scale, mask="causal") )pbdoc"); - + m.def( - "gated_delta_update_forward", - &mlx::core::fast::gated_delta_update_forward, - "q"_a, - "k"_a, - "v"_a, - "gates"_a, - "beta"_a, - "initial_state"_a = nb::none(), // optional, defaults to None - "C"_a = 16, // optional, defaults to None - "stream"_a = nb::none(), // optional, defaults to None - R"( + "gated_delta_update_forward", + &mlx::core::fast::gated_delta_update_forward, + "q"_a, + "k"_a, + "v"_a, + "gates"_a, + "beta"_a, + "initial_state"_a = nb::none(), // optional, defaults to None + "C"_a = 16, // optional, defaults to None + "stream"_a = nb::none(), // optional, defaults to None + R"( Chunked gated delta network forward pass. Args: @@ -357,8 +357,7 @@ void init_fast(nb::module_& parent_module) { Returns: Tuple of (output [B, H, T, Dv], final_state [B, H, Dk, Dv]) - )" - ); + )"); m.def( "metal_kernel", From 9f11eab8f5879cf245243d5affc8a982d2f42277 Mon Sep 17 00:00:00 2001 From: tpegolotti Date: Fri, 19 Jun 2026 16:36:06 +0200 Subject: [PATCH 10/56] full simdgroup main loop --- .../metal/kernels/gated_delta_update.metal | 9 +-- .../metal/kernels/gated_delta_update_impl.h | 65 +++++++++---------- 2 files changed, 35 insertions(+), 39 deletions(-) diff --git a/mlx/backend/metal/kernels/gated_delta_update.metal b/mlx/backend/metal/kernels/gated_delta_update.metal index 8a87d9fbb1..43de9e871e 100644 --- a/mlx/backend/metal/kernels/gated_delta_update.metal +++ b/mlx/backend/metal/kernels/gated_delta_update.metal @@ -15,8 +15,9 @@ using namespace metal; hk, \ hv) -#define instantiate_gated_delta_update_seq_dims(in_type, st_type) \ - instantiate_gated_delta_update_seq(in_type, st_type, 64, 64, 4, 4) +#define instantiate_gated_delta_update_seq_dims(in_type, st_type) \ + instantiate_gated_delta_update_seq(in_type, st_type, 64, 64, 4, 4) \ + instantiate_gated_delta_update_seq(in_type, st_type, 64, 64, 8, 8) #define instantiate_gated_delta_update_chunk( \ in_type, st_type, dk, dv, hk, hv, c) \ @@ -38,7 +39,7 @@ using namespace metal; instantiate_gated_delta_update_chunk( \ in_type, st_type, 64, 64, 4, 4, 8) \ instantiate_gated_delta_update_chunk( \ - in_type, st_type, 64, 64, 1, 1, 8) + in_type, st_type, 64, 64, 8, 8, 8) #define instantiate_make_wy(in_type, dk, dv, hk, hv, c) \ instantiate_kernel( \ @@ -55,7 +56,7 @@ using namespace metal; instantiate_make_wy(in_type, 64, 64, 4, 4, 32) \ instantiate_make_wy(in_type, 64, 64, 4, 4, 16) \ instantiate_make_wy(in_type, 64, 64, 4, 4, 8) \ - instantiate_make_wy(in_type, 64, 64, 1, 1, 8) + instantiate_make_wy(in_type, 64, 64, 8, 8, 8) instantiate_gated_delta_update_seq_dims(float, float) instantiate_gated_delta_update_chunk_dims(float, float) diff --git a/mlx/backend/metal/kernels/gated_delta_update_impl.h b/mlx/backend/metal/kernels/gated_delta_update_impl.h index 2d87be29c3..87d16348c0 100644 --- a/mlx/backend/metal/kernels/gated_delta_update_impl.h +++ b/mlx/backend/metal/kernels/gated_delta_update_impl.h @@ -174,13 +174,6 @@ template auto i_state = state_in + (n * Dv + dv_idx) * Dk; auto o_state = state_out + (n * Dv + dv_idx) * Dk; - // threadgroup tiles - threadgroup float Q_left_tg[C][Dk]; // gamma-scaled Q - threadgroup float K_tg[C][Dk]; // raw K - threadgroup float W_tg[C][Dk]; // gamma-scaled W (W_left) - threadgroup float U_tg[C][8]; // U chunk, 8 dv values - threadgroup float K_right_tg[C][Dk]; - simdgroup_float8x8 S_tile[8]; simdgroup_float8x8 W_tile, K_tile, Q_tile; simdgroup_float8x8 WS_tile; @@ -197,6 +190,9 @@ template thread auto& ws_e = WS_tile.thread_elements(); thread auto& qkt_e = QKt_tile.thread_elements(); thread auto& o_e = out_tile.thread_elements(); + thread auto& w_e = W_tile.thread_elements(); + thread auto& q_e = Q_tile.thread_elements(); + thread auto& k_e = K_tile.thread_elements(); // load initial state into threadgroup for (int kk = 0; kk < Dk; kk += 8) { @@ -204,41 +200,33 @@ template } threadgroup_barrier(mem_flags::mem_threadgroup); - float gamma[C]; + threadgroup float gamma[C]; for (int t = 0; t < T; t += C) { - gamma[0] = g_[hv_idx]; - for (int i = 1; i < C; i++) { - gamma[i] = gamma[i - 1] * g_[i * Hv + hv_idx]; - } - float gamma_C = gamma[C - 1]; - - for (int c = 0; c < C; c++) { - float kr = gamma_C / gamma[c]; - for (int d = dk_idx; d < Dk; d += 32) { - Q_left_tg[c][d] = gamma[c] * q_[c * Hk * Dk + d]; - K_tg[c][d] = k_[c * Hk * Dk + d]; - W_tg[c][d] = gamma[c] * W_[c * Hv * Dk + d]; - K_right_tg[c][d] = kr * K_tg[c][d]; - } - } + float g_val = (thread_index_in_simdgroup < C) + ? g_[thread_index_in_simdgroup * Hv + hv_idx] + : 1.0f; + + float gamma_val = simd_prefix_inclusive_product(g_val); - for (int idx = dk_idx; idx < C * 8; idx += 32) { - int c = idx / 8; - int j = idx % 8; - U_tg[c][j] = U_[c * Hv * Dv + dv_idx + j]; + // Only write within bounds + if (thread_index_in_simdgroup < C) { + gamma[thread_index_in_simdgroup] = gamma_val; } - threadgroup_barrier(mem_flags::mem_threadgroup); // WS = W_left @ S^T WS_tile = make_filled_simdgroup_matrix(0.f); for (int kk = 0; kk < Dk; kk += 8) { - simdgroup_load(W_tile, &W_tg[0][kk], Dk); + simdgroup_load(W_tile, W_ + kk, Hv * Dk); + + w_e[0] *= gamma[fm]; + w_e[1] *= gamma[fm]; + simdgroup_multiply_accumulate(WS_tile, W_tile, S_tile[kk / 8], WS_tile); } // delta = U - WS - simdgroup_load(U_tile, &U_tg[0][0], 8); + simdgroup_load(U_tile, U_ + dv_idx, Hv * Dv); d_e[0] = u_e[0] - ws_e[0]; d_e[1] = u_e[1] - ws_e[1]; @@ -248,8 +236,11 @@ template QKt_tile = make_filled_simdgroup_matrix(0.f); for (int kk = 0; kk < Dk; kk += 8) { - simdgroup_load(Q_tile, &Q_left_tg[0][kk], Dk); - simdgroup_load(K_tile, &K_tg[0][kk], Dk, ulong2(0, 0), true); + simdgroup_load(Q_tile, q_ + kk, Hk * Dk); + simdgroup_load(K_tile, k_ + kk, Hk * Dk, ulong2(0, 0), true); + + q_e[0] *= gamma[fm]; + q_e[1] *= gamma[fm]; simdgroup_multiply_accumulate(tmp_tile, Q_tile, S_tile[kk / 8], tmp_tile); simdgroup_multiply_accumulate(QKt_tile, Q_tile, K_tile, QKt_tile); @@ -265,12 +256,16 @@ template y[fm * Hv * Dv + dv_idx + fn + 1] = static_cast(o_e[1]); for (int kk = 0; kk < Dk; kk += 8) { - simdgroup_load(K_tile, &K_right_tg[0][kk], Dk, ulong2(0, 0), true); + simdgroup_load(K_tile, k_ + kk, Hk * Dk, ulong2(0, 0), true); + + k_e[0] *= (gamma[C - 1] / gamma[fn]); + k_e[1] *= (gamma[C - 1] / gamma[fn + 1]); + simdgroup_multiply(KD_tile, K_tile, delta_tile); thread auto& s_e = S_tile[kk / 8].thread_elements(); - s_e[0] = gamma_C * s_e[0] + kd[0]; - s_e[1] = gamma_C * s_e[1] + kd[1]; + s_e[0] = gamma[C - 1] * s_e[0] + kd[0]; + s_e[1] = gamma[C - 1] * s_e[1] + kd[1]; } // update pointers From 7933e8accf9f7428bf6a1a8594fc3cb10514a9e1 Mon Sep 17 00:00:00 2001 From: tpegolotti Date: Mon, 22 Jun 2026 17:23:35 +0200 Subject: [PATCH 11/56] fused wy into gated computation --- mlx/backend/metal/gated_delta_update.cpp | 90 +++++---- .../metal/kernels/gated_delta_update.metal | 27 ++- .../metal/kernels/gated_delta_update_impl.h | 190 +++++++++++++++++- 3 files changed, 269 insertions(+), 38 deletions(-) diff --git a/mlx/backend/metal/gated_delta_update.cpp b/mlx/backend/metal/gated_delta_update.cpp index 6250b4047e..2ac73faddf 100644 --- a/mlx/backend/metal/gated_delta_update.cpp +++ b/mlx/backend/metal/gated_delta_update.cpp @@ -43,7 +43,10 @@ void GatedDeltaUpdate::eval_gpu( int C = chunk_size; int n_chunks = T / C; // TODO: make general - std::string kernel_name = C < 1 ? "seq_gated_delta_" : "chunk_gated_delta_"; + // std::string kernel_name = C < 1 ? "seq_gated_delta_" : + // "chunk_gated_delta_"; + std::string kernel_name = + C < 1 ? "seq_gated_delta_" : "gated_delta_fused_chunk_"; std::string suffix = get_type_string(q.dtype()) // "float" + "_" + get_type_string(h0.dtype()) // "float" + "_" + std::to_string(Dk) + "_" + std::to_string(Dv) + "_" + @@ -68,47 +71,66 @@ void GatedDeltaUpdate::eval_gpu( if (C > 1) { // allocate full W and U -- [B, T, H, D] - array W({B, T, Hk, Dk}, q.dtype(), nullptr, {}); - array U({B, T, Hv, Dv}, q.dtype(), nullptr, {}); - W.set_data(allocator::malloc(W.nbytes())); - U.set_data(allocator::malloc(U.nbytes())); - - compute_encoder.add_temporary(W); - compute_encoder.add_temporary(U); - - // kernel 1: compute full W and U - std::string make_wy_name = "make_wy_" + - get_type_string(q.dtype()) // "float" - + "_" + std::to_string(Dk) + "_" + std::to_string(Dv) + "_" + - std::to_string(Hk) + "_" + std::to_string(Hv) + "_" + std::to_string(C); - - auto wy_kernel = d.get_kernel(make_wy_name); - - compute_encoder.set_compute_pipeline_state(wy_kernel); - compute_encoder.set_input_array(k, 0); - compute_encoder.set_input_array(v, 1); - compute_encoder.set_input_array(g, 2); - compute_encoder.set_input_array(beta, 3); - compute_encoder.set_output_array(W, 4); - compute_encoder.set_output_array(U, 5); - compute_encoder.set_bytes(T, 6); - - auto grid_wy = MTL::Size(32, 4, B * Hv * n_chunks); - auto threads_wy = MTL::Size(32, 4, 1); - compute_encoder.dispatch_threads(grid_wy, threads_wy); + // array W({B, T, Hk, Dk}, q.dtype(), nullptr, {}); + // array U({B, T, Hv, Dv}, q.dtype(), nullptr, {}); + // W.set_data(allocator::malloc(W.nbytes())); + // U.set_data(allocator::malloc(U.nbytes())); + + // compute_encoder.add_temporary(W); + // compute_encoder.add_temporary(U); + + // // kernel 1: compute full W and U + // std::string make_wy_name = "make_wy_" + + // get_type_string(q.dtype()) // "float" + // + "_" + std::to_string(Dk) + "_" + std::to_string(Dv) + "_" + + // std::to_string(Hk) + "_" + std::to_string(Hv) + "_" + + // std::to_string(C); + + // auto wy_kernel = d.get_kernel(make_wy_name); + + // compute_encoder.set_compute_pipeline_state(wy_kernel); + // compute_encoder.set_input_array(k, 0); + // compute_encoder.set_input_array(v, 1); + // compute_encoder.set_input_array(g, 2); + // compute_encoder.set_input_array(beta, 3); + // compute_encoder.set_output_array(W, 4); + // compute_encoder.set_output_array(U, 5); + // compute_encoder.set_bytes(T, 6); + + // auto grid_wy = MTL::Size(32, 4, B * Hv * n_chunks); + // auto threads_wy = MTL::Size(32, 4, 1); + // compute_encoder.dispatch_threads(grid_wy, threads_wy); + + // // kernel 2: gated delta + // compute_encoder.set_compute_pipeline_state(delta_kernel); + // compute_encoder.set_input_array(q, 0); + // compute_encoder.set_input_array(k, 1); + // compute_encoder.set_input_array(W, 2); + // compute_encoder.set_input_array(U, 3); + // compute_encoder.set_input_array(h0, 4); // initial state in + // compute_encoder.set_input_array(g, 5); + // compute_encoder.set_output_array(out, 6); + // compute_encoder.set_output_array(hf, 7); // final state out + // compute_encoder.set_bytes(T, 8); + // compute_encoder.set_bytes(n_chunks, 9); + + // auto grid = MTL::Size(32, Dv / 8, B * Hv); + // auto threads = MTL::Size(32, 1, 1); // get one simdgroup to work + // // auto grid = MTL::Size(1, 1, 1); + // // auto threads = MTL::Size(1, 1, 1); + // compute_encoder.dispatch_threads(grid, threads); // kernel 2: gated delta compute_encoder.set_compute_pipeline_state(delta_kernel); compute_encoder.set_input_array(q, 0); compute_encoder.set_input_array(k, 1); - compute_encoder.set_input_array(W, 2); - compute_encoder.set_input_array(U, 3); - compute_encoder.set_input_array(h0, 4); // initial state in - compute_encoder.set_input_array(g, 5); + compute_encoder.set_input_array(v, 2); + compute_encoder.set_input_array(h0, 3); // initial state in + compute_encoder.set_input_array(g, 4); + compute_encoder.set_input_array(beta, 5); compute_encoder.set_output_array(out, 6); compute_encoder.set_output_array(hf, 7); // final state out compute_encoder.set_bytes(T, 8); - compute_encoder.set_bytes(n_chunks, 9); auto grid = MTL::Size(32, Dv / 8, B * Hv); auto threads = MTL::Size(32, 1, 1); // get one simdgroup to work diff --git a/mlx/backend/metal/kernels/gated_delta_update.metal b/mlx/backend/metal/kernels/gated_delta_update.metal index 43de9e871e..0ed621431b 100644 --- a/mlx/backend/metal/kernels/gated_delta_update.metal +++ b/mlx/backend/metal/kernels/gated_delta_update.metal @@ -58,6 +58,31 @@ using namespace metal; instantiate_make_wy(in_type, 64, 64, 4, 4, 8) \ instantiate_make_wy(in_type, 64, 64, 8, 8, 8) +#define instantiate_gated_delta_update_fused_chunk( \ + in_type, st_type, dk, dv, hk, hv, c) \ + instantiate_kernel( \ + "gated_delta_fused_chunk_" #in_type "_" #st_type "_" #dk "_" #dv "_" #hk \ + "_" #hv "_" #c, \ + gated_delta_fused_chunk, \ + in_type, \ + st_type, \ + dk, \ + dv, \ + hk, \ + hv, \ + c) + +#define instantiate_gated_delta_update_fused_chunk_dims(in_type, st_type) \ + instantiate_gated_delta_update_fused_chunk( \ + in_type, st_type, 64, 64, 4, 4, 32) \ + instantiate_gated_delta_update_fused_chunk( \ + in_type, st_type, 64, 64, 4, 4, 16) \ + instantiate_gated_delta_update_fused_chunk( \ + in_type, st_type, 64, 64, 4, 4, 8) \ + instantiate_gated_delta_update_fused_chunk( \ + in_type, st_type, 64, 64, 8, 8, 8) + instantiate_gated_delta_update_seq_dims(float, float) instantiate_gated_delta_update_chunk_dims(float, float) - instantiate_make_wy_dims(float) \ No newline at end of file + instantiate_make_wy_dims(float) + instantiate_gated_delta_update_fused_chunk_dims(float, float) \ No newline at end of file diff --git a/mlx/backend/metal/kernels/gated_delta_update_impl.h b/mlx/backend/metal/kernels/gated_delta_update_impl.h index 87d16348c0..706232555f 100644 --- a/mlx/backend/metal/kernels/gated_delta_update_impl.h +++ b/mlx/backend/metal/kernels/gated_delta_update_impl.h @@ -13,6 +13,193 @@ using namespace metal; +#define AT(TILE, IDX) TILE.thread_elements()[IDX] + +template +[[kernel]] void gated_delta_fused_chunk( + const device InT* q [[buffer(0)]], + const device InT* k [[buffer(1)]], + const device InT* v [[buffer(2)]], + const device StT* state_in [[buffer(3)]], + const device InT* g [[buffer(4)]], + const device InT* beta [[buffer(5)]], + device InT* y [[buffer(6)]], + device StT* state_out [[buffer(7)]], + constant int& T [[buffer(8)]], + uint3 thread_position_in_grid [[thread_position_in_grid]], + uint thread_index_in_simdgroup [[thread_index_in_simdgroup]]) { + auto n = thread_position_in_grid.z; + auto b_idx = n / Hv; + auto hv_idx = n % Hv; + auto hk_idx = hv_idx / (Hv / Hk); + + const short qid = thread_index_in_simdgroup / 4; + const short fm = (qid & 4) + ((thread_index_in_simdgroup / 2) % 4); + const short fn = (qid & 2) * 2 + (thread_index_in_simdgroup % 2) * 2; + + auto dv_idx = thread_position_in_grid.y * 8; + + // set up pointers + // g: [B, T, Hv] + auto g_ = g + b_idx * T * Hv; + + // q, k: [B, T, Hk, Dk] + auto q_ = q + b_idx * T * Hk * Dk + hk_idx * Dk; + auto k_ = k + b_idx * T * Hk * Dk + hk_idx * Dk; + + // v, y: [B, T, Hv, Dv] + y += b_idx * T * Hv * Dv + hv_idx * Dv; + auto v_ = v + b_idx * T * Hv * Dv + hv_idx * Dv; + auto beta_ = beta + b_idx * T * Hv; + + // state_in, state_out: [B, Hv, Dv, Dk] + auto i_state = state_in + (n * Dv + dv_idx) * Dk; + auto o_state = state_out + (n * Dv + dv_idx) * Dk; + + simdgroup_float8x8 S_tile[Dk / 8]; + + // simdgroup matrices + simdgroup_float8x8 V_tile, K_tile, KT_tile, Q_tile; + simdgroup_float8x8 W_tile, U_tile; + simdgroup_float8x8 WS_tile; + simdgroup_float8x8 delta_tile; + simdgroup_float8x8 tmp_tile; + simdgroup_float8x8 QKt_tile; + simdgroup_float8x8 out_tile; + simdgroup_float8x8 KD_tile; + + // tiles for WY form computation + simdgroup_float8x8 KKtK_tile, KKtV_tile, X_tile, KKt_tile; + + threadgroup float gamma[C]; + + // load initial state into threadgroup + for (int kk = 0; kk < Dk; kk += 8) { + simdgroup_load(S_tile[kk / 8], i_state + kk, Dk, ulong2(0, 0), true); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + for (int t = 0; t < T; t += C) { + float g_val = (thread_index_in_simdgroup < C) + ? g_[thread_index_in_simdgroup * Hv + hv_idx] + : 1.0f; + + float gamma_val = simd_prefix_inclusive_product(g_val); + + if (thread_index_in_simdgroup < C) { + gamma[thread_index_in_simdgroup] = gamma_val; + } + + float beta_fm = beta_[fm * Hv + hv_idx]; + + KKt_tile = make_filled_simdgroup_matrix(0.f); + for (int kk = 0; kk < Dk; kk += 8) { + simdgroup_load(K_tile, k_ + kk, Dk * Hk); + simdgroup_load(KT_tile, k_ + kk, Dk * Hk, ulong2(0, 0), true); + simdgroup_multiply_accumulate(KKt_tile, K_tile, KT_tile, KKt_tile); + } + + KKtK_tile = KKt_tile; + KKtV_tile = KKt_tile; + + AT(KKtK_tile, 0) = fn >= fm ? 0.f : AT(KKtK_tile, 0) * beta_fm; + AT(KKtV_tile, 0) = + fn >= fm ? 0.f : AT(KKtV_tile, 0) * beta_fm * (gamma[fm] / gamma[fn]); + + AT(KKtK_tile, 1) = fn + 1 >= fm ? 0.f : AT(KKtK_tile, 1) * beta_fm; + AT(KKtV_tile, 1) = fn + 1 >= fm + ? 0.f + : AT(KKtV_tile, 1) * beta_fm * (gamma[fm] / gamma[fn + 1]); + + WS_tile = make_filled_simdgroup_matrix(0.f); + // Use the Neumann series: (I - T)^-1 = sum T^k, instead of doing forward + // substitution to compute W (and U). For C=8 this is 8 matmuls, for generic + // C this probably does not work well? Also fuse the WS in the same loop -> + // less memory movements. + for (int kk = 0; kk < Dk; kk += 8) { + simdgroup_load(K_tile, k_ + kk, Dk * Hk); + AT(K_tile, 0) *= beta_fm; + AT(K_tile, 1) *= beta_fm; + W_tile = K_tile; + for (int iter = 0; iter < C - 1; iter++) { + simdgroup_multiply(X_tile, KKtK_tile, W_tile); + AT(W_tile, 0) = AT(K_tile, 0) - AT(X_tile, 0); + AT(W_tile, 1) = AT(K_tile, 1) - AT(X_tile, 1); + } + + AT(W_tile, 0) *= gamma[fm]; + AT(W_tile, 1) *= gamma[fm]; + + // WS = W_left @ S^T + simdgroup_multiply_accumulate(WS_tile, W_tile, S_tile[kk / 8], WS_tile); + } + + simdgroup_load(V_tile, v_ + dv_idx, Dv * Hv); + AT(V_tile, 0) *= beta_fm; + AT(V_tile, 1) *= beta_fm; + U_tile = V_tile; + for (int iter = 0; iter < C - 1; iter++) { + simdgroup_multiply(X_tile, KKtV_tile, U_tile); + AT(U_tile, 0) = AT(V_tile, 0) - AT(X_tile, 0); + AT(U_tile, 1) = AT(V_tile, 1) - AT(X_tile, 1); + } + + // delta = U - WS + AT(delta_tile, 0) = AT(U_tile, 0) - AT(WS_tile, 0); + AT(delta_tile, 1) = AT(U_tile, 1) - AT(WS_tile, 1); + + // Q_left @ S^T + tmp_tile = make_filled_simdgroup_matrix(0.f); + QKt_tile = make_filled_simdgroup_matrix(0.f); + + for (int kk = 0; kk < Dk; kk += 8) { + simdgroup_load(Q_tile, q_ + kk, Hk * Dk); + simdgroup_load(K_tile, k_ + kk, Hk * Dk, ulong2(0, 0), true); + + AT(Q_tile, 0) *= gamma[fm]; + AT(Q_tile, 1) *= gamma[fm]; + + simdgroup_multiply_accumulate(tmp_tile, Q_tile, S_tile[kk / 8], tmp_tile); + simdgroup_multiply_accumulate(QKt_tile, Q_tile, K_tile, QKt_tile); + } + + // element 0 at (fm, fn), element 1 at (fm, fn+1) + AT(QKt_tile, 0) *= fn > fm ? 0.f : (1.0f / gamma[fn]); + AT(QKt_tile, 1) *= fn + 1 > fm ? 0.f : (1.0f / gamma[fn + 1]); + + simdgroup_multiply_accumulate(out_tile, QKt_tile, delta_tile, tmp_tile); + + y[fm * Hv * Dv + dv_idx + fn] = static_cast(AT(out_tile, 0)); + y[fm * Hv * Dv + dv_idx + fn + 1] = static_cast(AT(out_tile, 1)); + + for (int kk = 0; kk < Dk; kk += 8) { + simdgroup_load(K_tile, k_ + kk, Hk * Dk, ulong2(0, 0), true); + + AT(K_tile, 0) *= (gamma[C - 1] / gamma[fn]); + AT(K_tile, 1) *= (gamma[C - 1] / gamma[fn + 1]); + + simdgroup_multiply(KD_tile, K_tile, delta_tile); + + AT(S_tile[kk / 8], 0) = + gamma[C - 1] * AT(S_tile[kk / 8], 0) + AT(KD_tile, 0); + AT(S_tile[kk / 8], 1) = + gamma[C - 1] * AT(S_tile[kk / 8], 1) + AT(KD_tile, 1); + } + + // advance pointers + q_ += C * Hk * Dk; + k_ += C * Hk * Dk; + v_ += C * Hv * Dv; + beta_ += C * Hv; + y += C * Hv * Dv; + g_ += C * Hv; + } + + for (int kk = 0; kk < Dk; kk += 8) { + simdgroup_store(S_tile[kk / 8], o_state + kk, Dk, ulong2(0, 0), true); + } +} + template [[kernel]] void make_wy( const device InT* k [[buffer(0)]], @@ -140,7 +327,6 @@ template device StT* state_out [[buffer(7)]], // [B, Hv, Dv, Dk] constant int& T [[buffer(8)]], uint3 thread_position_in_grid [[thread_position_in_grid]], - uint3 thread_position_in_threadgroup [[thread_position_in_threadgroup]], uint thread_index_in_simdgroup [[thread_index_in_simdgroup]]) { auto n = thread_position_in_grid.z; auto b_idx = n / Hv; @@ -152,7 +338,6 @@ template const short fm = (qid & 4) + ((thread_index_in_simdgroup / 2) % 4); const short fn = (qid & 2) * 2 + (thread_index_in_simdgroup % 2) * 2; - auto dk_idx = thread_position_in_threadgroup.x; // simd group id auto dv_idx = thread_position_in_grid.y * 8; // set up pointers @@ -209,7 +394,6 @@ template float gamma_val = simd_prefix_inclusive_product(g_val); - // Only write within bounds if (thread_index_in_simdgroup < C) { gamma[thread_index_in_simdgroup] = gamma_val; } From bd05b9297bfff37fa97ffe6d8c2047d9b52cc74d Mon Sep 17 00:00:00 2001 From: tpegolotti Date: Tue, 23 Jun 2026 11:06:00 +0200 Subject: [PATCH 12/56] add elementwise macros --- .../metal/kernels/gated_delta_update_impl.h | 77 +++++++++++-------- 1 file changed, 46 insertions(+), 31 deletions(-) diff --git a/mlx/backend/metal/kernels/gated_delta_update_impl.h b/mlx/backend/metal/kernels/gated_delta_update_impl.h index 706232555f..ad208b4497 100644 --- a/mlx/backend/metal/kernels/gated_delta_update_impl.h +++ b/mlx/backend/metal/kernels/gated_delta_update_impl.h @@ -14,6 +14,37 @@ using namespace metal; #define AT(TILE, IDX) TILE.thread_elements()[IDX] +#define SUB(TILE0, TILE1, TILE2) \ + { \ + AT(TILE0, 0) = AT(TILE1, 0) - AT(TILE2, 0); \ + AT(TILE0, 1) = AT(TILE1, 1) - AT(TILE2, 1); \ + } +#define FMA(TILE0, S, TILE1, TILE2) \ + { \ + AT(TILE0, 0) = S * AT(TILE1, 0) + AT(TILE2, 0); \ + AT(TILE0, 1) = S * AT(TILE1, 1) + AT(TILE2, 1); \ + } + +#define SCALE(TILE0, S) \ + { \ + AT(TILE0, 0) *= S; \ + AT(TILE0, 1) *= S; \ + } +#define SCALE2(TILE0, S0, S1) \ + { \ + AT(TILE0, 0) *= S0; \ + AT(TILE0, 1) *= S1; \ + } +#define SCALE_TRI(TILE0, S0, S1) \ + { \ + AT(TILE0, 0) *= fn > fm ? 0.f : S0; \ + AT(TILE0, 1) *= fn + 1 > fm ? 0.f : S1; \ + } +#define SCALE_TRIEQ(TILE0, S0, S1) \ + { \ + AT(TILE0, 0) *= fn >= fm ? 0.f : S0; \ + AT(TILE0, 1) *= fn + 1 >= fm ? 0.f : S1; \ + } template [[kernel]] void gated_delta_fused_chunk( @@ -77,7 +108,6 @@ template for (int kk = 0; kk < Dk; kk += 8) { simdgroup_load(S_tile[kk / 8], i_state + kk, Dk, ulong2(0, 0), true); } - threadgroup_barrier(mem_flags::mem_threadgroup); for (int t = 0; t < T; t += C) { float g_val = (thread_index_in_simdgroup < C) @@ -102,14 +132,12 @@ template KKtK_tile = KKt_tile; KKtV_tile = KKt_tile; - AT(KKtK_tile, 0) = fn >= fm ? 0.f : AT(KKtK_tile, 0) * beta_fm; - AT(KKtV_tile, 0) = - fn >= fm ? 0.f : AT(KKtV_tile, 0) * beta_fm * (gamma[fm] / gamma[fn]); - - AT(KKtK_tile, 1) = fn + 1 >= fm ? 0.f : AT(KKtK_tile, 1) * beta_fm; - AT(KKtV_tile, 1) = fn + 1 >= fm - ? 0.f - : AT(KKtV_tile, 1) * beta_fm * (gamma[fm] / gamma[fn + 1]); + // elementwise multiplication by Gamma and beta (for V) and beta (for K) + SCALE_TRIEQ(KKtK_tile, beta_fm, beta_fm) + SCALE_TRIEQ( + KKtV_tile, + beta_fm * (gamma[fm] / gamma[fn]), + beta_fm * (gamma[fm] / gamma[fn + 1])) WS_tile = make_filled_simdgroup_matrix(0.f); // Use the Neumann series: (I - T)^-1 = sum T^k, instead of doing forward @@ -118,17 +146,14 @@ template // less memory movements. for (int kk = 0; kk < Dk; kk += 8) { simdgroup_load(K_tile, k_ + kk, Dk * Hk); - AT(K_tile, 0) *= beta_fm; - AT(K_tile, 1) *= beta_fm; + SCALE(K_tile, beta_fm) W_tile = K_tile; for (int iter = 0; iter < C - 1; iter++) { simdgroup_multiply(X_tile, KKtK_tile, W_tile); - AT(W_tile, 0) = AT(K_tile, 0) - AT(X_tile, 0); - AT(W_tile, 1) = AT(K_tile, 1) - AT(X_tile, 1); + SUB(W_tile, K_tile, X_tile) } - AT(W_tile, 0) *= gamma[fm]; - AT(W_tile, 1) *= gamma[fm]; + SCALE(W_tile, gamma[fm]) // WS = W_left @ S^T simdgroup_multiply_accumulate(WS_tile, W_tile, S_tile[kk / 8], WS_tile); @@ -140,32 +165,26 @@ template U_tile = V_tile; for (int iter = 0; iter < C - 1; iter++) { simdgroup_multiply(X_tile, KKtV_tile, U_tile); - AT(U_tile, 0) = AT(V_tile, 0) - AT(X_tile, 0); - AT(U_tile, 1) = AT(V_tile, 1) - AT(X_tile, 1); + SUB(U_tile, V_tile, X_tile) } // delta = U - WS - AT(delta_tile, 0) = AT(U_tile, 0) - AT(WS_tile, 0); - AT(delta_tile, 1) = AT(U_tile, 1) - AT(WS_tile, 1); + SUB(delta_tile, U_tile, WS_tile) // Q_left @ S^T tmp_tile = make_filled_simdgroup_matrix(0.f); QKt_tile = make_filled_simdgroup_matrix(0.f); - for (int kk = 0; kk < Dk; kk += 8) { simdgroup_load(Q_tile, q_ + kk, Hk * Dk); simdgroup_load(K_tile, k_ + kk, Hk * Dk, ulong2(0, 0), true); - AT(Q_tile, 0) *= gamma[fm]; - AT(Q_tile, 1) *= gamma[fm]; + SCALE(Q_tile, gamma[fm]) simdgroup_multiply_accumulate(tmp_tile, Q_tile, S_tile[kk / 8], tmp_tile); simdgroup_multiply_accumulate(QKt_tile, Q_tile, K_tile, QKt_tile); } - // element 0 at (fm, fn), element 1 at (fm, fn+1) - AT(QKt_tile, 0) *= fn > fm ? 0.f : (1.0f / gamma[fn]); - AT(QKt_tile, 1) *= fn + 1 > fm ? 0.f : (1.0f / gamma[fn + 1]); + SCALE_TRI(QKt_tile, (1.0f / gamma[fn]), (1.0f / gamma[fn + 1])) simdgroup_multiply_accumulate(out_tile, QKt_tile, delta_tile, tmp_tile); @@ -175,15 +194,11 @@ template for (int kk = 0; kk < Dk; kk += 8) { simdgroup_load(K_tile, k_ + kk, Hk * Dk, ulong2(0, 0), true); - AT(K_tile, 0) *= (gamma[C - 1] / gamma[fn]); - AT(K_tile, 1) *= (gamma[C - 1] / gamma[fn + 1]); + SCALE2(K_tile, (gamma[C - 1] / gamma[fn]), (gamma[C - 1] / gamma[fn + 1])) simdgroup_multiply(KD_tile, K_tile, delta_tile); - AT(S_tile[kk / 8], 0) = - gamma[C - 1] * AT(S_tile[kk / 8], 0) + AT(KD_tile, 0); - AT(S_tile[kk / 8], 1) = - gamma[C - 1] * AT(S_tile[kk / 8], 1) + AT(KD_tile, 1); + FMA(S_tile[kk / 8], gamma[C - 1], S_tile[kk / 8], KD_tile) } // advance pointers From 73d10aae5f4158f74eb8042f8c886a00edce1293 Mon Sep 17 00:00:00 2001 From: tpegolotti Date: Wed, 24 Jun 2026 15:21:21 +0200 Subject: [PATCH 13/56] add gated delta benchmark --- benchmarks/python/gated_delta_bench.py | 286 ++++++++++++++++++ mlx/backend/metal/gated_delta_update.cpp | 4 +- .../metal/kernels/gated_delta_update.metal | 19 +- .../metal/kernels/gated_delta_update_impl.h | 12 +- 4 files changed, 310 insertions(+), 11 deletions(-) create mode 100644 benchmarks/python/gated_delta_bench.py diff --git a/benchmarks/python/gated_delta_bench.py b/benchmarks/python/gated_delta_bench.py new file mode 100644 index 0000000000..9fb4ff0a21 --- /dev/null +++ b/benchmarks/python/gated_delta_bench.py @@ -0,0 +1,286 @@ +import argparse +import itertools +import os +import time +from datetime import datetime +from typing import Optional, Tuple + +import mlx.core as mx + +N_warmup = 5 +N_iter_bench = 40 +N_iter_func = 8 + +os.environ["MTL_CAPTURE_ENABLED"] = "1" + + +def bench(f, *args): + for i in range(N_warmup): + f(*args) + + s = time.perf_counter_ns() + for i in range(N_iter_bench): + f(*args) + e = time.perf_counter_ns() + return (e - s) * 1e-9 + + +def do_kernel_bench(f, *args): + q_out = args[0] + + for i in range(N_iter_func): + out, hf = f(*args) + + mx.eval(out, hf) + return q_out + + +def profile(f, *args): + C = args[-1] + now = datetime.now() + timestamp = now.strftime("%Y_%m_%d_%H_%M") + trace_file = f"traces/mlx_trace_{C}_{timestamp}.gputrace" + + mx.metal.start_capture(trace_file) + + for i in range(N_iter_func): + f(*args) + + mx.metal.stop_capture() + print(f"Writing trace: ") + + +def gated_delta_ref( + q: mx.array, # [B, T, H, Dk] + k: mx.array, # [B, T, H, Dk] + v: mx.array, # [B, T, H, Dv] + g: mx.array, # [B, T, H] or [B, T, H, Dk] + beta: mx.array, # [B, T, H] + state: Optional[mx.array] = None, # [B, H, Dv, Dk] +) -> Tuple[mx.array, mx.array]: + """ + Implements: + S_t = a_t S_{t-1} + b_t (v_t - a_t S_{t-1} k_t) k_t^T + o_t = S_t q_t + """ + B, T, H, Dk = q.shape + Dv = v.shape[-1] + + if state is None: + state = mx.zeros((B, H, Dv, Dk), dtype=mx.float32) + + outputs = [] + for t in range(T): + q_t = q[:, t] # [B, H, Dk] + k_t = k[:, t] # [B, H, Dk] + v_t = v[:, t] # [B, H, Dv] + g_t = g[:, t] # [B, H] or [B, H, Dk] + beta_t = beta[:, t] # [B, H] + + # decay + if g_t.ndim == 2: + decay = g_t[..., None, None] # [B, H, 1, 1] + else: + decay = g_t[..., None, :] # [B, H, 1, Dk] + + # S = a S + state = state * decay + + # kv = S * k_t, [B, H, Dv, Dk] * [B, H, 1, Dk] = [B, H, Dv, Dk] -> reduction on Dk + kv_mem = (state * k_t[..., None, :]).sum(axis=-1) # [B, H, Dv] + + # delta = b_t * (v_t - kv) + delta = (v_t - kv_mem) * beta_t[..., None] # [B, H, Dv] + + # S = S + delta * k_t^T, [B, H, Dv, 1] * [B, H, 1, Dk] = [B, H, Dv, Dk] + state = state + delta[..., None] * k_t[..., None, :] # [B, H, Dv, Dk] + + # o_t = S_t * q_t, [B, H, Dv, Dk] * [B, H, 1, Dk] = [B, H, Dv, Dk] -> reduction on Dk + o_t = (state * q_t[..., None, :]).sum(axis=-1) # [B, H, Dv] + outputs.append(o_t) + + return mx.stack(outputs, axis=1), state # [B, T, H, Dv], [B, H, Dv, Dk] + + +def benchmark_shape(B, T, H, Dk, Dv, chunk_sizes): + mx.random.seed(42) + + q = mx.random.normal(shape=(B, T, H, Dk)) + k = mx.random.normal(shape=(B, T, H, Dk)) + k = k / (mx.linalg.norm(k, axis=-1, keepdims=True) + 1e-6) + v = mx.random.normal(shape=(B, T, H, Dv)) + g = mx.random.normal(shape=(B, T, H)) * 0.1 - 1.0 + b = mx.sigmoid(mx.random.normal(shape=(B, T, H))) + h0 = mx.zeros((B, H, Dv, Dk), dtype=mx.float32) + + shape_str = f"B={B} T={T} H={H} Dk={Dk} Dv={Dv}" + + out_ref, hf_ref = gated_delta_ref(q, k, v, g, b, state=h0) + mx.eval(out_ref, hf_ref) + + out, hf = mx.fast.gated_delta_update_forward(q, k, v, g, b, initial_state=h0, C=0) + mx.eval(out, hf) + + atol = 1e-2 + out_close = mx.all(mx.abs(out - out_ref) < atol).item() + hf_close = mx.all(mx.abs(hf - hf_ref) < atol).item() + + out_max_diff = mx.abs(out - out_ref).max().item() + hf_max_diff = mx.abs(hf - hf_ref).max().item() + + assert out_close, f"output mismatch! max diff: {out_max_diff:.2e}" + assert hf_close, f"state mismatch! max diff: {hf_max_diff:.2e}" + + ms_seq = ( + bench(do_kernel_bench, mx.fast.gated_delta_update_forward, q, k, v, g, b, h0, 0) + * 1000 + ) + + speedups = [] + + non_zero_Cs = [C for C in chunk_sizes if C != 0] + for C in non_zero_Cs: + try: + h0 = mx.zeros((B, H, Dv, Dk), dtype=mx.float32) + out, hf = mx.fast.gated_delta_update_forward( + q, k, v, g, b, initial_state=h0, C=C + ) + mx.eval(out, hf) + + err_out = mx.abs(out - out_ref).max().item() + err_hf = mx.abs(hf - hf_ref).max().item() + + ms_c = ( + bench( + do_kernel_bench, + mx.fast.gated_delta_update_forward, + q, + k, + v, + g, + b, + h0, + C, + ) + * 1000 + ) + + speedup = ms_seq / ms_c if ms_c > 0 else float("nan") + + speedups.append(speedup) + + out_close = mx.all(mx.abs(out - out_ref) < atol).item() + hf_close = mx.all(mx.abs(hf - hf_ref) < atol).item() + + out_max_diff = mx.abs(out - out_ref).max().item() + hf_max_diff = mx.abs(hf - hf_ref).max().item() + + if not out_close or not hf_close: + raise Exception + + except Exception as e: + speedups.append("err") + + return shape_str, f"{ms_seq:.3f}", speedups + + +def run_benchmark(run_full): + if run_full: + Bs = [1, 2, 4, 8, 16] + Ts = [2048, 4096, 8192, 16384] + Hs = [16, 24, 32] + Dks = [64] + Dvs = [64] + else: + Bs = [1, 8, 16] + Ts = [4096, 8192] + Hs = [16] + Dks = [64] + Dvs = [64] + + rows = [] + + chunk_sizes = [0, 8] + non_zero_Cs = [C for C in chunk_sizes if C != 0] + + headers = ["B", "T", "H", "Dk", "Dv", "time_seq (ms)"] + [ + f"C={C} speedup" for C in non_zero_Cs + ] + + col_widths = [5, 5, 5, 5, 5, 20] + [20] * (len(non_zero_Cs)) + fmt = " ".join(f"{{:<{w}}}" for w in col_widths) + print(fmt.format(*headers)) + print("-" * (sum(col_widths) + 10)) + + for B, T, H, Dk, Dv in itertools.product(Bs, Ts, Hs, Dks, Dvs): + + shapes_s, base_time, speedups = benchmark_shape(B, T, H, Dk, Dv, chunk_sizes) + + row = [f"{B}", f"{T}", f"{H}", f"{Dk}", f"{Dv}", base_time] + for speed in speedups: + row.append(f"{speed:.2f}x") + + print(fmt.format(*row)) + + +def run_profile(): + B = 8 + H = 8 + T = 4096 * 2 + Dk = 64 + Dv = 64 + CS = [8] # , 16]#, 32] + + mx.random.seed(42) + + q = mx.random.normal(shape=(B, T, H, Dk)) + k = mx.random.normal(shape=(B, T, H, Dk)) + k = k / (mx.linalg.norm(k, axis=-1, keepdims=True) + 1e-6) + v = mx.random.normal(shape=(B, T, H, Dv)) + g = mx.random.normal(shape=(B, T, H)) * 0.1 - 1.0 + b = mx.sigmoid(mx.random.normal(shape=(B, T, H))) + h0 = mx.zeros((B, H, Dv, Dk), dtype=mx.float32) + + mx.eval(q, k, v, g, b, h0) + + out_ref, hf_ref = gated_delta_ref(q, k, v, g, b, state=h0) + mx.eval(out_ref, hf_ref) + + for C in CS: + out, hf = mx.fast.gated_delta_update_forward( + q, k, v, g, b, initial_state=h0, C=C + ) + mx.eval(out, hf) + + assert list(out.shape) == [B, T, H, Dv] + assert list(hf.shape) == [B, H, Dv, Dk] + assert not mx.any(mx.isnan(out)).item(), "NaNs in output!" + assert not mx.any(mx.isnan(hf)).item(), "NaNs in final state!" + + # correctness check + atol = 1e-3 + out_close = mx.all(mx.abs(out - out_ref) < atol).item() + hf_close = mx.all(mx.abs(hf - hf_ref) < atol).item() + + out_max_diff = mx.abs(out - out_ref).max().item() + hf_max_diff = mx.abs(hf - hf_ref).max().item() + + assert out_close, f"output mismatch! max diff: {out_max_diff:.2e}" + assert hf_close, f"state mismatch! max diff: {hf_max_diff:.2e}" + + profile( + do_kernel_bench, mx.fast.gated_delta_update_forward, q, k, v, g, b, h0, C + ) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Gated delta chunk kernel runner") + parser.add_argument("--profile", "-p", action="store_true") + parser.add_argument("--full", "-f", action="store_true") + args = parser.parse_args() + + if args.profile: + run_profile() + exit() + + run_benchmark(args.full) diff --git a/mlx/backend/metal/gated_delta_update.cpp b/mlx/backend/metal/gated_delta_update.cpp index 2ac73faddf..3f564cb826 100644 --- a/mlx/backend/metal/gated_delta_update.cpp +++ b/mlx/backend/metal/gated_delta_update.cpp @@ -115,7 +115,7 @@ void GatedDeltaUpdate::eval_gpu( // compute_encoder.set_bytes(n_chunks, 9); // auto grid = MTL::Size(32, Dv / 8, B * Hv); - // auto threads = MTL::Size(32, 1, 1); // get one simdgroup to work + // auto threads = MTL::Size(32, 1, 1); // // auto grid = MTL::Size(1, 1, 1); // // auto threads = MTL::Size(1, 1, 1); // compute_encoder.dispatch_threads(grid, threads); @@ -133,7 +133,7 @@ void GatedDeltaUpdate::eval_gpu( compute_encoder.set_bytes(T, 8); auto grid = MTL::Size(32, Dv / 8, B * Hv); - auto threads = MTL::Size(32, 1, 1); // get one simdgroup to work + auto threads = MTL::Size(32, 1, 1); // auto grid = MTL::Size(1, 1, 1); // auto threads = MTL::Size(1, 1, 1); compute_encoder.dispatch_threads(grid, threads); diff --git a/mlx/backend/metal/kernels/gated_delta_update.metal b/mlx/backend/metal/kernels/gated_delta_update.metal index 0ed621431b..828573a61f 100644 --- a/mlx/backend/metal/kernels/gated_delta_update.metal +++ b/mlx/backend/metal/kernels/gated_delta_update.metal @@ -15,9 +15,14 @@ using namespace metal; hk, \ hv) -#define instantiate_gated_delta_update_seq_dims(in_type, st_type) \ - instantiate_gated_delta_update_seq(in_type, st_type, 64, 64, 4, 4) \ - instantiate_gated_delta_update_seq(in_type, st_type, 64, 64, 8, 8) +#define instantiate_gated_delta_update_seq_dims(in_type, st_type) \ + instantiate_gated_delta_update_seq(in_type, st_type, 64, 64, 4, 4) \ + instantiate_gated_delta_update_seq(in_type, st_type, 64, 64, 8, 8) \ + instantiate_gated_delta_update_seq(in_type, st_type, 64, 64, 16, 16) \ + instantiate_gated_delta_update_seq( \ + in_type, st_type, 64, 64, 24, 24) \ + instantiate_gated_delta_update_seq( \ + in_type, st_type, 64, 64, 32, 32) #define instantiate_gated_delta_update_chunk( \ in_type, st_type, dk, dv, hk, hv, c) \ @@ -80,7 +85,13 @@ using namespace metal; instantiate_gated_delta_update_fused_chunk( \ in_type, st_type, 64, 64, 4, 4, 8) \ instantiate_gated_delta_update_fused_chunk( \ - in_type, st_type, 64, 64, 8, 8, 8) + in_type, st_type, 64, 64, 8, 8, 8) \ + instantiate_gated_delta_update_fused_chunk( \ + in_type, st_type, 64, 64, 16, 16, 8) \ + instantiate_gated_delta_update_fused_chunk( \ + in_type, st_type, 64, 64, 24, 24, 8) \ + instantiate_gated_delta_update_fused_chunk( \ + in_type, st_type, 64, 64, 32, 32, 8) instantiate_gated_delta_update_seq_dims(float, float) instantiate_gated_delta_update_chunk_dims(float, float) diff --git a/mlx/backend/metal/kernels/gated_delta_update_impl.h b/mlx/backend/metal/kernels/gated_delta_update_impl.h index ad208b4497..5e06a751c3 100644 --- a/mlx/backend/metal/kernels/gated_delta_update_impl.h +++ b/mlx/backend/metal/kernels/gated_delta_update_impl.h @@ -65,8 +65,10 @@ template auto hk_idx = hv_idx / (Hv / Hk); const short qid = thread_index_in_simdgroup / 4; - const short fm = (qid & 4) + ((thread_index_in_simdgroup / 2) % 4); - const short fn = (qid & 2) * 2 + (thread_index_in_simdgroup % 2) * 2; + const short fm = (qid & 4) + + ((thread_index_in_simdgroup / 2) % 4); // row coordinate of the held tile + const short fn = (qid & 2) * 2 + + (thread_index_in_simdgroup % 2) * 2; // column coordinate of the held tile auto dv_idx = thread_position_in_grid.y * 8; @@ -160,8 +162,7 @@ template } simdgroup_load(V_tile, v_ + dv_idx, Dv * Hv); - AT(V_tile, 0) *= beta_fm; - AT(V_tile, 1) *= beta_fm; + SCALE(V_tile, beta_fm) U_tile = V_tile; for (int iter = 0; iter < C - 1; iter++) { simdgroup_multiply(X_tile, KKtV_tile, U_tile); @@ -171,7 +172,6 @@ template // delta = U - WS SUB(delta_tile, U_tile, WS_tile) - // Q_left @ S^T tmp_tile = make_filled_simdgroup_matrix(0.f); QKt_tile = make_filled_simdgroup_matrix(0.f); for (int kk = 0; kk < Dk; kk += 8) { @@ -180,10 +180,12 @@ template SCALE(Q_tile, gamma[fm]) + // Q_left @ S^T simdgroup_multiply_accumulate(tmp_tile, Q_tile, S_tile[kk / 8], tmp_tile); simdgroup_multiply_accumulate(QKt_tile, Q_tile, K_tile, QKt_tile); } + // (Q @ K) * Gamma, in the paper they use M here but it's probably a typo. SCALE_TRI(QKt_tile, (1.0f / gamma[fn]), (1.0f / gamma[fn + 1])) simdgroup_multiply_accumulate(out_tile, QKt_tile, delta_tile, tmp_tile); From a2e0c60fb6457c9526e4b140b028770531cb7dc7 Mon Sep 17 00:00:00 2001 From: tpegolotti Date: Wed, 24 Jun 2026 15:39:35 +0200 Subject: [PATCH 14/56] fix name --- benchmarks/python/gated_delta_bench.py | 19 ++++++------------- mlx/fast.cpp | 2 +- mlx/fast.h | 2 +- python/src/fast.cpp | 4 ++-- 4 files changed, 10 insertions(+), 17 deletions(-) diff --git a/benchmarks/python/gated_delta_bench.py b/benchmarks/python/gated_delta_bench.py index 9fb4ff0a21..251b125d18 100644 --- a/benchmarks/python/gated_delta_bench.py +++ b/benchmarks/python/gated_delta_bench.py @@ -118,7 +118,7 @@ def benchmark_shape(B, T, H, Dk, Dv, chunk_sizes): out_ref, hf_ref = gated_delta_ref(q, k, v, g, b, state=h0) mx.eval(out_ref, hf_ref) - out, hf = mx.fast.gated_delta_update_forward(q, k, v, g, b, initial_state=h0, C=0) + out, hf = mx.fast.gated_delta_update(q, k, v, g, b, initial_state=h0, C=0) mx.eval(out, hf) atol = 1e-2 @@ -132,8 +132,7 @@ def benchmark_shape(B, T, H, Dk, Dv, chunk_sizes): assert hf_close, f"state mismatch! max diff: {hf_max_diff:.2e}" ms_seq = ( - bench(do_kernel_bench, mx.fast.gated_delta_update_forward, q, k, v, g, b, h0, 0) - * 1000 + bench(do_kernel_bench, mx.fast.gated_delta_update, q, k, v, g, b, h0, 0) * 1000 ) speedups = [] @@ -142,9 +141,7 @@ def benchmark_shape(B, T, H, Dk, Dv, chunk_sizes): for C in non_zero_Cs: try: h0 = mx.zeros((B, H, Dv, Dk), dtype=mx.float32) - out, hf = mx.fast.gated_delta_update_forward( - q, k, v, g, b, initial_state=h0, C=C - ) + out, hf = mx.fast.gated_delta_update(q, k, v, g, b, initial_state=h0, C=C) mx.eval(out, hf) err_out = mx.abs(out - out_ref).max().item() @@ -153,7 +150,7 @@ def benchmark_shape(B, T, H, Dk, Dv, chunk_sizes): ms_c = ( bench( do_kernel_bench, - mx.fast.gated_delta_update_forward, + mx.fast.gated_delta_update, q, k, v, @@ -247,9 +244,7 @@ def run_profile(): mx.eval(out_ref, hf_ref) for C in CS: - out, hf = mx.fast.gated_delta_update_forward( - q, k, v, g, b, initial_state=h0, C=C - ) + out, hf = mx.fast.gated_delta_update(q, k, v, g, b, initial_state=h0, C=C) mx.eval(out, hf) assert list(out.shape) == [B, T, H, Dv] @@ -268,9 +263,7 @@ def run_profile(): assert out_close, f"output mismatch! max diff: {out_max_diff:.2e}" assert hf_close, f"state mismatch! max diff: {hf_max_diff:.2e}" - profile( - do_kernel_bench, mx.fast.gated_delta_update_forward, q, k, v, g, b, h0, C - ) + profile(do_kernel_bench, mx.fast.gated_delta_update, q, k, v, g, b, h0, C) if __name__ == "__main__": diff --git a/mlx/fast.cpp b/mlx/fast.cpp index 785d46e288..905c83b4ed 100644 --- a/mlx/fast.cpp +++ b/mlx/fast.cpp @@ -922,7 +922,7 @@ bool ScaledDotProductAttentionVJP::is_equivalent(const Primitive& other) const { has_sinks_ == a_other.has_sinks_; } -std::vector gated_delta_update_forward( +std::vector gated_delta_update( const array& queries, const array& keys, const array& values, diff --git a/mlx/fast.h b/mlx/fast.h index 2dbd03e1d0..18d9079a39 100644 --- a/mlx/fast.h +++ b/mlx/fast.h @@ -55,7 +55,7 @@ MLX_API array scaled_dot_product_attention( const std::optional& sinks = {}, StreamOrDevice s = {}); -MLX_API std::vector gated_delta_update_forward( +MLX_API std::vector gated_delta_update( const array& queries, const array& keys, const array& values, diff --git a/python/src/fast.cpp b/python/src/fast.cpp index 5cedf8273f..2a926979ad 100644 --- a/python/src/fast.cpp +++ b/python/src/fast.cpp @@ -334,8 +334,8 @@ void init_fast(nb::module_& parent_module) { )pbdoc"); m.def( - "gated_delta_update_forward", - &mlx::core::fast::gated_delta_update_forward, + "gated_delta_update", + &mlx::core::fast::gated_delta_update, "q"_a, "k"_a, "v"_a, From 7237966dfd3ea0ccdd45e65e72f826ba09f73efc Mon Sep 17 00:00:00 2001 From: tpegolotti Date: Wed, 24 Jun 2026 15:41:20 +0200 Subject: [PATCH 15/56] fix default C value --- mlx/fast.h | 2 +- python/src/fast.cpp | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/mlx/fast.h b/mlx/fast.h index 18d9079a39..24d17bdae1 100644 --- a/mlx/fast.h +++ b/mlx/fast.h @@ -62,7 +62,7 @@ MLX_API std::vector gated_delta_update( const array& gates, const array& beta_, const std::optional& initial_state = std::nullopt, - const int C = 16, + const int C = 8, StreamOrDevice s = {}); using TemplateArg = std::variant; diff --git a/python/src/fast.cpp b/python/src/fast.cpp index 2a926979ad..773419fc24 100644 --- a/python/src/fast.cpp +++ b/python/src/fast.cpp @@ -342,7 +342,7 @@ void init_fast(nb::module_& parent_module) { "gates"_a, "beta"_a, "initial_state"_a = nb::none(), // optional, defaults to None - "C"_a = 16, // optional, defaults to None + "C"_a = 8, // optional, defaults to None "stream"_a = nb::none(), // optional, defaults to None R"( Chunked gated delta network forward pass. From aa44a664fbb6550f0ed1680b48424e9e0168ae9e Mon Sep 17 00:00:00 2001 From: tpegolotti Date: Wed, 24 Jun 2026 16:41:40 +0200 Subject: [PATCH 16/56] update bench script --- benchmarks/python/gated_delta_bench.py | 16 +++++++--------- 1 file changed, 7 insertions(+), 9 deletions(-) diff --git a/benchmarks/python/gated_delta_bench.py b/benchmarks/python/gated_delta_bench.py index 251b125d18..c13a42b7ae 100644 --- a/benchmarks/python/gated_delta_bench.py +++ b/benchmarks/python/gated_delta_bench.py @@ -47,7 +47,7 @@ def profile(f, *args): f(*args) mx.metal.stop_capture() - print(f"Writing trace: ") + print(f"Writing trace: {trace_file}") def gated_delta_ref( @@ -204,15 +204,13 @@ def run_benchmark(run_full): f"C={C} speedup" for C in non_zero_Cs ] - col_widths = [5, 5, 5, 5, 5, 20] + [20] * (len(non_zero_Cs)) - fmt = " ".join(f"{{:<{w}}}" for w in col_widths) + col_widths = [6, 6, 6, 6, 6, 15] + [15] * (len(non_zero_Cs)) + fmt = "".join(f"{{:<{w}}}" for w in col_widths) print(fmt.format(*headers)) - print("-" * (sum(col_widths) + 10)) + print("-" * (sum(col_widths))) for B, T, H, Dk, Dv in itertools.product(Bs, Ts, Hs, Dks, Dvs): - shapes_s, base_time, speedups = benchmark_shape(B, T, H, Dk, Dv, chunk_sizes) - row = [f"{B}", f"{T}", f"{H}", f"{Dk}", f"{Dv}", base_time] for speed in speedups: row.append(f"{speed:.2f}x") @@ -221,8 +219,8 @@ def run_benchmark(run_full): def run_profile(): - B = 8 - H = 8 + B = 1 + H = 16 T = 4096 * 2 Dk = 64 Dv = 64 @@ -253,7 +251,7 @@ def run_profile(): assert not mx.any(mx.isnan(hf)).item(), "NaNs in final state!" # correctness check - atol = 1e-3 + atol = 1e-2 out_close = mx.all(mx.abs(out - out_ref) < atol).item() hf_close = mx.all(mx.abs(hf - hf_ref) < atol).item() From 8635620fac0aa8c9e978d44d3fbfe8f872c438c6 Mon Sep 17 00:00:00 2001 From: tpegolotti Date: Thu, 25 Jun 2026 13:18:31 +0200 Subject: [PATCH 17/56] Remove old chunkwise implementation --- mlx/backend/metal/gated_delta_update.cpp | 168 ++++------- .../metal/kernels/gated_delta_update.metal | 43 +-- .../metal/kernels/gated_delta_update_impl.h | 280 ------------------ 3 files changed, 65 insertions(+), 426 deletions(-) diff --git a/mlx/backend/metal/gated_delta_update.cpp b/mlx/backend/metal/gated_delta_update.cpp index 3f564cb826..41eeadc8e1 100644 --- a/mlx/backend/metal/gated_delta_update.cpp +++ b/mlx/backend/metal/gated_delta_update.cpp @@ -43,121 +43,81 @@ void GatedDeltaUpdate::eval_gpu( int C = chunk_size; int n_chunks = T / C; // TODO: make general - // std::string kernel_name = C < 1 ? "seq_gated_delta_" : - // "chunk_gated_delta_"; - std::string kernel_name = - C < 1 ? "seq_gated_delta_" : "gated_delta_fused_chunk_"; std::string suffix = get_type_string(q.dtype()) // "float" + "_" + get_type_string(h0.dtype()) // "float" + "_" + std::to_string(Dk) + "_" + std::to_string(Dv) + "_" + std::to_string(Hk) + "_" + std::to_string(Hv); - // printf("C: %d\nname: %s\n",C,kernel_name.c_str()); - std::string base_name = kernel_name + suffix; - - base_name += C >= 1 ? "_" + std::to_string(C) : ""; - - std::string hash_name = base_name; - - metal::MTLFCList func_consts = {}; - - auto delta_kernel = get_steel_gated_delta_forward_kernel( - d, base_name, hash_name, func_consts); - auto& compute_encoder = metal::get_command_encoder(s); out.set_data(allocator::malloc(out.nbytes())); hf.set_data(allocator::malloc(hf.nbytes())); - if (C > 1) { - // allocate full W and U -- [B, T, H, D] - // array W({B, T, Hk, Dk}, q.dtype(), nullptr, {}); - // array U({B, T, Hv, Dv}, q.dtype(), nullptr, {}); - // W.set_data(allocator::malloc(W.nbytes())); - // U.set_data(allocator::malloc(U.nbytes())); - - // compute_encoder.add_temporary(W); - // compute_encoder.add_temporary(U); - - // // kernel 1: compute full W and U - // std::string make_wy_name = "make_wy_" + - // get_type_string(q.dtype()) // "float" - // + "_" + std::to_string(Dk) + "_" + std::to_string(Dv) + "_" + - // std::to_string(Hk) + "_" + std::to_string(Hv) + "_" + - // std::to_string(C); - - // auto wy_kernel = d.get_kernel(make_wy_name); - - // compute_encoder.set_compute_pipeline_state(wy_kernel); - // compute_encoder.set_input_array(k, 0); - // compute_encoder.set_input_array(v, 1); - // compute_encoder.set_input_array(g, 2); - // compute_encoder.set_input_array(beta, 3); - // compute_encoder.set_output_array(W, 4); - // compute_encoder.set_output_array(U, 5); - // compute_encoder.set_bytes(T, 6); - - // auto grid_wy = MTL::Size(32, 4, B * Hv * n_chunks); - // auto threads_wy = MTL::Size(32, 4, 1); - // compute_encoder.dispatch_threads(grid_wy, threads_wy); - - // // kernel 2: gated delta - // compute_encoder.set_compute_pipeline_state(delta_kernel); - // compute_encoder.set_input_array(q, 0); - // compute_encoder.set_input_array(k, 1); - // compute_encoder.set_input_array(W, 2); - // compute_encoder.set_input_array(U, 3); - // compute_encoder.set_input_array(h0, 4); // initial state in - // compute_encoder.set_input_array(g, 5); - // compute_encoder.set_output_array(out, 6); - // compute_encoder.set_output_array(hf, 7); // final state out - // compute_encoder.set_bytes(T, 8); - // compute_encoder.set_bytes(n_chunks, 9); - - // auto grid = MTL::Size(32, Dv / 8, B * Hv); - // auto threads = MTL::Size(32, 1, 1); - // // auto grid = MTL::Size(1, 1, 1); - // // auto threads = MTL::Size(1, 1, 1); - // compute_encoder.dispatch_threads(grid, threads); - - // kernel 2: gated delta - compute_encoder.set_compute_pipeline_state(delta_kernel); - compute_encoder.set_input_array(q, 0); - compute_encoder.set_input_array(k, 1); - compute_encoder.set_input_array(v, 2); - compute_encoder.set_input_array(h0, 3); // initial state in - compute_encoder.set_input_array(g, 4); - compute_encoder.set_input_array(beta, 5); - compute_encoder.set_output_array(out, 6); - compute_encoder.set_output_array(hf, 7); // final state out - compute_encoder.set_bytes(T, 8); - - auto grid = MTL::Size(32, Dv / 8, B * Hv); - auto threads = MTL::Size(32, 1, 1); - // auto grid = MTL::Size(1, 1, 1); - // auto threads = MTL::Size(1, 1, 1); - compute_encoder.dispatch_threads(grid, threads); - - } else { - compute_encoder.set_compute_pipeline_state(delta_kernel); - - compute_encoder.set_input_array(q, 0); - compute_encoder.set_input_array(k, 1); - compute_encoder.set_input_array(v, 2); - compute_encoder.set_input_array(g, 3); - compute_encoder.set_input_array(beta, 4); - compute_encoder.set_input_array(h0, 5); - compute_encoder.set_bytes(T, 6); - compute_encoder.set_output_array(out, 7); - compute_encoder.set_output_array(hf, 8); - - // auto grid = MTL::Size(1, 1, 1); - // auto threads = MTL::Size(1, 1, 1); - auto grid = MTL::Size(32, Dv, B * Hv); - auto threads = MTL::Size(32, 4, 1); - compute_encoder.dispatch_threads(grid, threads); + switch (C) { + case 8: { + std::string kernel_name = "gated_delta_fused_chunk_"; + + // printf("C: %d\nname: %s\n",C,kernel_name.c_str()); + std::string base_name = kernel_name + suffix; + + base_name += "_" + std::to_string(C); + + std::string hash_name = base_name; + + metal::MTLFCList func_consts = {}; + + auto delta_kernel = get_steel_gated_delta_forward_kernel( + d, base_name, hash_name, func_consts); + + compute_encoder.set_compute_pipeline_state(delta_kernel); + compute_encoder.set_input_array(q, 0); + compute_encoder.set_input_array(k, 1); + compute_encoder.set_input_array(v, 2); + compute_encoder.set_input_array(h0, 3); // initial state in + compute_encoder.set_input_array(g, 4); + compute_encoder.set_input_array(beta, 5); + compute_encoder.set_output_array(out, 6); + compute_encoder.set_output_array(hf, 7); // final state out + compute_encoder.set_bytes(T, 8); + + auto grid = MTL::Size(32, Dv / 8, B * Hv); + auto threads = MTL::Size(32, 1, 1); + compute_encoder.dispatch_threads(grid, threads); + break; + } + case 1: + case 0: { + std::string kernel_name = "seq_gated_delta_"; + std::string base_name = kernel_name + suffix; + std::string hash_name = base_name; + + metal::MTLFCList func_consts = {}; + + auto delta_kernel = get_steel_gated_delta_forward_kernel( + d, base_name, hash_name, func_consts); + + compute_encoder.set_compute_pipeline_state(delta_kernel); + + compute_encoder.set_input_array(q, 0); + compute_encoder.set_input_array(k, 1); + compute_encoder.set_input_array(v, 2); + compute_encoder.set_input_array(g, 3); + compute_encoder.set_input_array(beta, 4); + compute_encoder.set_input_array(h0, 5); + compute_encoder.set_bytes(T, 6); + compute_encoder.set_output_array(out, 7); + compute_encoder.set_output_array(hf, 8); + + auto grid = MTL::Size(32, Dv, B * Hv); + auto threads = MTL::Size(32, 4, 1); + compute_encoder.dispatch_threads(grid, threads); + break; + } + default: { + throw std::runtime_error( + "NYI: Only sequential and chunk size 8 are supported"); + } } - // throw std::runtime_error("NYI"); } bool GatedDeltaUpdate::is_equivalent(const Primitive& other) const { diff --git a/mlx/backend/metal/kernels/gated_delta_update.metal b/mlx/backend/metal/kernels/gated_delta_update.metal index 828573a61f..8fb6e5dd54 100644 --- a/mlx/backend/metal/kernels/gated_delta_update.metal +++ b/mlx/backend/metal/kernels/gated_delta_update.metal @@ -24,45 +24,6 @@ using namespace metal; instantiate_gated_delta_update_seq( \ in_type, st_type, 64, 64, 32, 32) -#define instantiate_gated_delta_update_chunk( \ - in_type, st_type, dk, dv, hk, hv, c) \ - instantiate_kernel( \ - "chunk_gated_delta_" #in_type "_" #st_type "_" #dk "_" #dv "_" #hk \ - "_" #hv "_" #c, \ - gated_delta_chunk, \ - in_type, \ - st_type, \ - dk, \ - dv, \ - hk, \ - hv, \ - c) - -#define instantiate_gated_delta_update_chunk_dims(in_type, st_type) \ - instantiate_gated_delta_update_chunk(in_type, st_type, 64, 64, 4, 4, 32) \ - instantiate_gated_delta_update_chunk(in_type, st_type, 64, 64, 4, 4, 16) \ - instantiate_gated_delta_update_chunk( \ - in_type, st_type, 64, 64, 4, 4, 8) \ - instantiate_gated_delta_update_chunk( \ - in_type, st_type, 64, 64, 8, 8, 8) - -#define instantiate_make_wy(in_type, dk, dv, hk, hv, c) \ - instantiate_kernel( \ - "make_wy_" #in_type "_" #dk "_" #dv "_" #hk "_" #hv "_" #c, \ - make_wy, \ - in_type, \ - dk, \ - dv, \ - hk, \ - hv, \ - c) - -#define instantiate_make_wy_dims(in_type) \ - instantiate_make_wy(in_type, 64, 64, 4, 4, 32) \ - instantiate_make_wy(in_type, 64, 64, 4, 4, 16) \ - instantiate_make_wy(in_type, 64, 64, 4, 4, 8) \ - instantiate_make_wy(in_type, 64, 64, 8, 8, 8) - #define instantiate_gated_delta_update_fused_chunk( \ in_type, st_type, dk, dv, hk, hv, c) \ instantiate_kernel( \ @@ -94,6 +55,4 @@ using namespace metal; in_type, st_type, 64, 64, 32, 32, 8) instantiate_gated_delta_update_seq_dims(float, float) - instantiate_gated_delta_update_chunk_dims(float, float) - instantiate_make_wy_dims(float) - instantiate_gated_delta_update_fused_chunk_dims(float, float) \ No newline at end of file + instantiate_gated_delta_update_fused_chunk_dims(float, float) \ No newline at end of file diff --git a/mlx/backend/metal/kernels/gated_delta_update_impl.h b/mlx/backend/metal/kernels/gated_delta_update_impl.h index 5e06a751c3..9f057cfe80 100644 --- a/mlx/backend/metal/kernels/gated_delta_update_impl.h +++ b/mlx/backend/metal/kernels/gated_delta_update_impl.h @@ -3,14 +3,6 @@ #include #include "mlx/backend/metal/kernels/utils.h" -// importing steel -// #include "mlx/backend/metal/kernels/steel/attn/loader.h" -// #include "mlx/backend/metal/kernels/steel/attn/mma.h" -// #include "mlx/backend/metal/kernels/steel/attn/params.h" -// #include "mlx/backend/metal/kernels/steel/attn/transforms.h" -// #include "mlx/backend/metal/kernels/steel/gemm/params.h" -// #include "mlx/backend/metal/kernels/steel/utils.h" - using namespace metal; #define AT(TILE, IDX) TILE.thread_elements()[IDX] @@ -217,278 +209,6 @@ template } } -template -[[kernel]] void make_wy( - const device InT* k [[buffer(0)]], - const device InT* v [[buffer(1)]], - const device InT* g [[buffer(2)]], // [B, T, Hv] or [B, T, Hv, Dk] - const device InT* beta [[buffer(3)]], // [B, Hv, Dv, Dk] - device InT* W [[buffer(4)]], // [B, T, Hv] - device InT* U [[buffer(5)]], // [B, T, Hv] - constant int& T [[buffer(6)]], - uint3 thread_position_in_grid [[thread_position_in_grid]], - uint3 thread_position_in_threadgroup [[thread_position_in_threadgroup]], - uint thread_index_in_simdgroup [[thread_index_in_simdgroup]]) { - auto n = thread_position_in_grid.z; - auto n_chunks = T / C; - auto chunk = n % n_chunks; - auto bh_idx = n / n_chunks; - auto b_idx = bh_idx / Hv; - auto hv_idx = bh_idx % Hv; - auto hk_idx = hv_idx / (Hv / Hk); - - constexpr int n_per_dk = Dk / 32; - constexpr int n_per_dv = Dv / 32; - - auto dk_idx = thread_position_in_threadgroup.x; - auto simd_group_id = thread_position_in_threadgroup.y; - const int num_simdgroups = 4; - - int offset_t = chunk * C; - auto g_ = g + b_idx * T * Hv + offset_t * Hv; - auto k_ = k + b_idx * T * Hk * Dk + offset_t * Hk * Dk + hk_idx * Dk; - auto v_ = v + b_idx * T * Hv * Dv + offset_t * Hv * Dv + hv_idx * Dv; - auto beta_ = beta + b_idx * T * Hv + offset_t * Hv; - auto W_ = W + b_idx * T * Hv * Dk + offset_t * Hv * Dk + hv_idx * Dk; - auto U_ = U + b_idx * T * Hv * Dv + offset_t * Hv * Dv + hv_idx * Dv; - - // threadgroup memory - threadgroup float K_tg[C][Dk]; - threadgroup float KKt[C][C]; - - for (int i = simd_group_id; i < C; i += num_simdgroups) { - for (int d = dk_idx; d < Dk; d += 32) { - K_tg[i][d] = k_[i * Hk * Dk + d]; - } - } - threadgroup_barrier(mem_flags::mem_threadgroup); - - // compute gamma (cumprod) - float gamma[C]; - gamma[0] = g_[hv_idx]; - for (int i = 1; i < C; i++) { - gamma[i] = gamma[i - 1] * g_[i * Hv + hv_idx]; - } - - // compute KKt = K @ K.T - unsigned int counter = 0; - for (int i = 0; i < C; i++) { - for (int j = 0; j <= i; j++) { - if (counter % num_simdgroups == simd_group_id) { - float sum = 0; - for (int d = dk_idx; d < Dk; d += 32) { - sum += K_tg[i][d] * K_tg[j][d]; - } - sum = simd_sum(sum); - if (thread_index_in_simdgroup == 0) { - KKt[i][j] = sum; - } - } - counter++; - } - } - threadgroup_barrier(mem_flags::mem_threadgroup); - - float u[C][n_per_dv]; - float w[C][n_per_dk]; - - // initialize - for (int i = 0; i < C; i++) { - float beta_i = beta_[i * Hv + hv_idx]; - for (int p = 0; p < n_per_dv; p++) { - int dv = dk_idx * n_per_dv + p; - u[i][p] = beta_i * v_[i * Hv * Dv + dv]; - } - for (int p = 0; p < n_per_dk; p++) { - int dk = dk_idx * n_per_dk + p; - w[i][p] = beta_i * K_tg[i][dk]; - } - } - - // forward substitution - for (int i = 1; i < C; i++) { - float beta_i = beta_[i * Hv + hv_idx]; - for (int j = 0; j < i; j++) { - float a_U = beta_i * (gamma[i] / gamma[j]) * KKt[i][j]; - float a_W = beta_i * KKt[i][j]; - for (int p = 0; p < n_per_dv; p++) { - u[i][p] -= a_U * u[j][p]; - } - for (int p = 0; p < n_per_dk; p++) { - w[i][p] -= a_W * w[j][p]; - } - } - } - - for (int i = 0; i < C; i++) { - for (int p = 0; p < n_per_dv; p++) { - int dv = dk_idx * n_per_dv + p; - U_[i * Hv * Dv + dv] = u[i][p]; - } - for (int p = 0; p < n_per_dk; p++) { - int dk = dk_idx * n_per_dk + p; - W_[i * Hk * Dk + dk] = w[i][p]; - } - } -} - -template -[[kernel]] void gated_delta_chunk( - const device InT* q [[buffer(0)]], - const device InT* k [[buffer(1)]], - const device InT* W [[buffer(2)]], // [B, T, Hv, Dk] - const device InT* U [[buffer(3)]], // [B, T, Hv, Dv] - const device StT* state_in [[buffer(4)]], // [B, Hv, Dv, Dk] - const device InT* g [[buffer(5)]], // [B, T, Hv] - device InT* y [[buffer(6)]], // [B, T, Hv, Dv] - device StT* state_out [[buffer(7)]], // [B, Hv, Dv, Dk] - constant int& T [[buffer(8)]], - uint3 thread_position_in_grid [[thread_position_in_grid]], - uint thread_index_in_simdgroup [[thread_index_in_simdgroup]]) { - auto n = thread_position_in_grid.z; - auto b_idx = n / Hv; - auto hv_idx = n % Hv; - auto hk_idx = hv_idx / (Hv / Hk); - - // get coord from steel - const short qid = thread_index_in_simdgroup / 4; - const short fm = (qid & 4) + ((thread_index_in_simdgroup / 2) % 4); - const short fn = (qid & 2) * 2 + (thread_index_in_simdgroup % 2) * 2; - - auto dv_idx = thread_position_in_grid.y * 8; - - // set up pointers - // g: [B, T, Hv] - auto g_ = g + b_idx * T * Hv; - - // q, k: [B, T, Hk, Dk] - auto q_ = q + b_idx * T * Hk * Dk + hk_idx * Dk; - auto k_ = k + b_idx * T * Hk * Dk + hk_idx * Dk; - - // W: [B, T, Hk, Dk], U: [B, T, Hv, Dv] - auto W_ = W + b_idx * T * Hv * Dk + hv_idx * Dk; - auto U_ = U + b_idx * T * Hv * Dv + hv_idx * Dv; - - // v, y: [B, T, Hv, Dv] - y += b_idx * T * Hv * Dv + hv_idx * Dv; - - // state_in, state_out: [B, Hv, Dv, Dk] - auto i_state = state_in + (n * Dv + dv_idx) * Dk; - auto o_state = state_out + (n * Dv + dv_idx) * Dk; - - simdgroup_float8x8 S_tile[8]; - simdgroup_float8x8 W_tile, K_tile, Q_tile; - simdgroup_float8x8 WS_tile; - simdgroup_float8x8 U_tile; - simdgroup_float8x8 delta_tile; - simdgroup_float8x8 tmp_tile; - simdgroup_float8x8 QKt_tile; - simdgroup_float8x8 out_tile; - simdgroup_float8x8 KD_tile; - - thread auto& kd = KD_tile.thread_elements(); - thread auto& d_e = delta_tile.thread_elements(); - thread auto& u_e = U_tile.thread_elements(); - thread auto& ws_e = WS_tile.thread_elements(); - thread auto& qkt_e = QKt_tile.thread_elements(); - thread auto& o_e = out_tile.thread_elements(); - thread auto& w_e = W_tile.thread_elements(); - thread auto& q_e = Q_tile.thread_elements(); - thread auto& k_e = K_tile.thread_elements(); - - // load initial state into threadgroup - for (int kk = 0; kk < Dk; kk += 8) { - simdgroup_load(S_tile[kk / 8], i_state + kk, Dk, ulong2(0, 0), true); - } - threadgroup_barrier(mem_flags::mem_threadgroup); - - threadgroup float gamma[C]; - - for (int t = 0; t < T; t += C) { - float g_val = (thread_index_in_simdgroup < C) - ? g_[thread_index_in_simdgroup * Hv + hv_idx] - : 1.0f; - - float gamma_val = simd_prefix_inclusive_product(g_val); - - if (thread_index_in_simdgroup < C) { - gamma[thread_index_in_simdgroup] = gamma_val; - } - - // WS = W_left @ S^T - WS_tile = make_filled_simdgroup_matrix(0.f); - for (int kk = 0; kk < Dk; kk += 8) { - simdgroup_load(W_tile, W_ + kk, Hv * Dk); - - w_e[0] *= gamma[fm]; - w_e[1] *= gamma[fm]; - - simdgroup_multiply_accumulate(WS_tile, W_tile, S_tile[kk / 8], WS_tile); - } - - // delta = U - WS - simdgroup_load(U_tile, U_ + dv_idx, Hv * Dv); - - d_e[0] = u_e[0] - ws_e[0]; - d_e[1] = u_e[1] - ws_e[1]; - - // Q_left @ S^T - tmp_tile = make_filled_simdgroup_matrix(0.f); - QKt_tile = make_filled_simdgroup_matrix(0.f); - - for (int kk = 0; kk < Dk; kk += 8) { - simdgroup_load(Q_tile, q_ + kk, Hk * Dk); - simdgroup_load(K_tile, k_ + kk, Hk * Dk, ulong2(0, 0), true); - - q_e[0] *= gamma[fm]; - q_e[1] *= gamma[fm]; - - simdgroup_multiply_accumulate(tmp_tile, Q_tile, S_tile[kk / 8], tmp_tile); - simdgroup_multiply_accumulate(QKt_tile, Q_tile, K_tile, QKt_tile); - } - - // element 0 at (fm, fn), element 1 at (fm, fn+1) - qkt_e[0] *= fn > fm ? 0.f : (1.0f / gamma[fn]); - qkt_e[1] *= fn + 1 > fm ? 0.f : (1.0f / gamma[fn + 1]); - - simdgroup_multiply_accumulate(out_tile, QKt_tile, delta_tile, tmp_tile); - - y[fm * Hv * Dv + dv_idx + fn] = static_cast(o_e[0]); - y[fm * Hv * Dv + dv_idx + fn + 1] = static_cast(o_e[1]); - - for (int kk = 0; kk < Dk; kk += 8) { - simdgroup_load(K_tile, k_ + kk, Hk * Dk, ulong2(0, 0), true); - - k_e[0] *= (gamma[C - 1] / gamma[fn]); - k_e[1] *= (gamma[C - 1] / gamma[fn + 1]); - - simdgroup_multiply(KD_tile, K_tile, delta_tile); - - thread auto& s_e = S_tile[kk / 8].thread_elements(); - s_e[0] = gamma[C - 1] * s_e[0] + kd[0]; - s_e[1] = gamma[C - 1] * s_e[1] + kd[1]; - } - - // update pointers - q_ += C * Hk * Dk; - k_ += C * Hk * Dk; - U_ += C * Hv * Dv; - W_ += C * Hv * Dk; - y += C * Hv * Dv; - g_ += C * Hv; - } - - // o_state is [Dv, Dk]: o_state[dv][dk] = st_out[dk][dv] - for (int kk = 0; kk < Dk; kk += 8) { - simdgroup_store( - S_tile[kk / 8], - o_state + kk, - Dk, - ulong2(0, 0), - true); // writes [Dk_block, Dv] layout - } -} - /* auto grid = MTL::Size(32, Dv, B * Hv); auto threads = MTL::Size(32, 4, 1); From 05713abd5ba09a6410e190518fe6a1135a502afa Mon Sep 17 00:00:00 2001 From: tpegolotti Date: Thu, 25 Jun 2026 13:18:58 +0200 Subject: [PATCH 18/56] Start fallback implementation --- mlx/fast.cpp | 16 ++++++++++++++-- 1 file changed, 14 insertions(+), 2 deletions(-) diff --git a/mlx/fast.cpp b/mlx/fast.cpp index 905c83b4ed..a1c99e5c94 100644 --- a/mlx/fast.cpp +++ b/mlx/fast.cpp @@ -954,9 +954,21 @@ std::vector gated_delta_update( auto h0 = initial_state.has_value() ? astype(*initial_state, out_dtype, s) : zeros({B, Hv, Dk, Dv}, out_dtype, s); - auto fallback = [](std::vector inputs) -> std::vector { + auto fallback = [B, T, Hk, Dk, Hv, Dv, s](std::vector inputs) { + auto q = astype(inputs[0], float32, s); + auto k = astype(inputs[1], float32, s); + auto v = astype(inputs[2], float32, s); + auto g = astype(inputs[3], float32, s); + auto beta = astype(inputs[4], float32, s); + auto state = astype(inputs[5], float32, s); + + std::vector outputs; + for (int t = 0; t < T; t++) { + // TODO. Implement + } throw std::runtime_error("NYI: GatedDeltaUpdate CPU fallback"); - return {}; + auto out = concatenate(outputs, 1, s); + return std::vector{out, state}; }; return array::make_arrays( From c45cd0f141526ca19fef78a96bba8ebfbe45da4c Mon Sep 17 00:00:00 2001 From: tpegolotti Date: Fri, 26 Jun 2026 13:31:31 +0200 Subject: [PATCH 19/56] Add contiguous memory copy and padding to handle generic T to get it to work on Qwen --- mlx/backend/metal/gated_delta_update.cpp | 2 + .../metal/kernels/gated_delta_update.metal | 44 ++++++++++++------- mlx/fast.cpp | 31 +++++++++++-- python/src/fast.cpp | 6 +-- 4 files changed, 61 insertions(+), 22 deletions(-) diff --git a/mlx/backend/metal/gated_delta_update.cpp b/mlx/backend/metal/gated_delta_update.cpp index 41eeadc8e1..59bfe3d76f 100644 --- a/mlx/backend/metal/gated_delta_update.cpp +++ b/mlx/backend/metal/gated_delta_update.cpp @@ -40,6 +40,8 @@ void GatedDeltaUpdate::eval_gpu( int Hv = v.shape(2); int Dv = v.shape(3); + // printf("%d %d %d %d %d %d\n",B,T,Hk,Hv,Dk,Dv); + int C = chunk_size; int n_chunks = T / C; // TODO: make general diff --git a/mlx/backend/metal/kernels/gated_delta_update.metal b/mlx/backend/metal/kernels/gated_delta_update.metal index 8fb6e5dd54..86b0788aa9 100644 --- a/mlx/backend/metal/kernels/gated_delta_update.metal +++ b/mlx/backend/metal/kernels/gated_delta_update.metal @@ -22,7 +22,15 @@ using namespace metal; instantiate_gated_delta_update_seq( \ in_type, st_type, 64, 64, 24, 24) \ instantiate_gated_delta_update_seq( \ - in_type, st_type, 64, 64, 32, 32) + in_type, st_type, 64, 64, 32, 32) \ + instantiate_gated_delta_update_seq( \ + in_type, st_type, 128, 128, 16, 16) \ + instantiate_gated_delta_update_seq( \ + in_type, st_type, 128, 128, 24, 24) \ + instantiate_gated_delta_update_seq( \ + in_type, st_type, 128, 128, 32, 32) \ + instantiate_gated_delta_update_seq( \ + in_type, st_type, 128, 128, 16, 32) #define instantiate_gated_delta_update_fused_chunk( \ in_type, st_type, dk, dv, hk, hv, c) \ @@ -38,21 +46,25 @@ using namespace metal; hv, \ c) -#define instantiate_gated_delta_update_fused_chunk_dims(in_type, st_type) \ - instantiate_gated_delta_update_fused_chunk( \ - in_type, st_type, 64, 64, 4, 4, 32) \ - instantiate_gated_delta_update_fused_chunk( \ - in_type, st_type, 64, 64, 4, 4, 16) \ - instantiate_gated_delta_update_fused_chunk( \ - in_type, st_type, 64, 64, 4, 4, 8) \ - instantiate_gated_delta_update_fused_chunk( \ - in_type, st_type, 64, 64, 8, 8, 8) \ - instantiate_gated_delta_update_fused_chunk( \ - in_type, st_type, 64, 64, 16, 16, 8) \ - instantiate_gated_delta_update_fused_chunk( \ - in_type, st_type, 64, 64, 24, 24, 8) \ - instantiate_gated_delta_update_fused_chunk( \ - in_type, st_type, 64, 64, 32, 32, 8) +#define instantiate_gated_delta_update_fused_chunk_dims(in_type, st_type) \ + instantiate_gated_delta_update_fused_chunk( \ + in_type, st_type, 128, 128, 16, 16, 8) \ + instantiate_gated_delta_update_fused_chunk( \ + in_type, st_type, 128, 128, 24, 24, 8) \ + instantiate_gated_delta_update_fused_chunk( \ + in_type, st_type, 128, 128, 32, 32, 8) \ + instantiate_gated_delta_update_fused_chunk( \ + in_type, st_type, 64, 64, 4, 4, 8) \ + instantiate_gated_delta_update_fused_chunk( \ + in_type, st_type, 64, 64, 8, 8, 8) \ + instantiate_gated_delta_update_fused_chunk( \ + in_type, st_type, 64, 64, 16, 16, 8) \ + instantiate_gated_delta_update_fused_chunk( \ + in_type, st_type, 64, 64, 24, 24, 8) \ + instantiate_gated_delta_update_fused_chunk( \ + in_type, st_type, 64, 64, 32, 32, 8) \ + instantiate_gated_delta_update_fused_chunk( \ + in_type, st_type, 128, 128, 16, 32, 8) instantiate_gated_delta_update_seq_dims(float, float) instantiate_gated_delta_update_fused_chunk_dims(float, float) \ No newline at end of file diff --git a/mlx/fast.cpp b/mlx/fast.cpp index a1c99e5c94..c73d55f554 100644 --- a/mlx/fast.cpp +++ b/mlx/fast.cpp @@ -944,6 +944,12 @@ std::vector gated_delta_update( auto g = astype(gates, out_dtype, s); auto beta = astype(beta_, out_dtype, s); + q = contiguous(q, false, s); + k = contiguous(k, false, s); + v = contiguous(v, false, s); + g = contiguous(g, false, s); + beta = contiguous(beta, false, s); + int B = q.shape(0); int T = q.shape(1); int Hk = q.shape(2); @@ -951,8 +957,20 @@ std::vector gated_delta_update( int Hv = v.shape(2); int Dv = v.shape(3); + int pad_size = (C != 0) ? (C - T % C) % C : 0; + int T_padded = T > 1 ? T + pad_size : T; + + if (pad_size > 1) { + q = pad(q, {1}, {0}, {pad_size}, array(0.0f, out_dtype), "constant", s); + k = pad(k, {1}, {0}, {pad_size}, array(0.0f, out_dtype), "constant", s); + v = pad(v, {1}, {0}, {pad_size}, array(0.0f, out_dtype), "constant", s); + g = pad(g, {1}, {0}, {pad_size}, array(1.0f, out_dtype), "constant", s); + beta = + pad(beta, {1}, {0}, {pad_size}, array(0.0f, out_dtype), "constant", s); + } + auto h0 = initial_state.has_value() ? astype(*initial_state, out_dtype, s) - : zeros({B, Hv, Dk, Dv}, out_dtype, s); + : zeros({B, Hv, Dv, Dk}, out_dtype, s); auto fallback = [B, T, Hk, Dk, Hv, Dv, s](std::vector inputs) { auto q = astype(inputs[0], float32, s); @@ -971,12 +989,19 @@ std::vector gated_delta_update( return std::vector{out, state}; }; - return array::make_arrays( - /* output shapes */ {{B, T, Hv, Dv}, {B, Hv, Dk, Dv}}, + auto result = array::make_arrays( + /* output shapes */ {{B, T_padded, Hv, Dv}, {B, Hv, Dv, Dk}}, /* dtypes */ {out_dtype, out_dtype}, /* primitive */ std::make_shared(to_stream(s), fallback, C), /* inputs */ {q, k, v, g, beta, h0}); + + // slice output back to original T + if (pad_size > 0) { + result[0] = slice(result[0], {0, 0, 0, 0}, {B, T, Hv, Dv}, s); + } + + return result; } bool Quantize::is_equivalent(const Primitive& other) const { diff --git a/python/src/fast.cpp b/python/src/fast.cpp index 773419fc24..26db40cff2 100644 --- a/python/src/fast.cpp +++ b/python/src/fast.cpp @@ -339,10 +339,10 @@ void init_fast(nb::module_& parent_module) { "q"_a, "k"_a, "v"_a, - "gates"_a, + "gamma"_a, "beta"_a, "initial_state"_a = nb::none(), // optional, defaults to None - "C"_a = 8, // optional, defaults to None + "C"_a = 0, // optional, defaults to 0 "stream"_a = nb::none(), // optional, defaults to None R"( Chunked gated delta network forward pass. @@ -351,7 +351,7 @@ void init_fast(nb::module_& parent_module) { q: Queries [B, H, T, Dk] k: Keys [B, H, T, Dk] v: Values [B, H, T, Dv] - gates: Log-decay gates [B, H, T] + gamma: beta: Delta update rates [B, H, T] initial_state: Optional initial hidden state [B, H, Dk, Dv] From 2ccb06636f52067e12dde8231aadc1dd79fbdd1c Mon Sep 17 00:00:00 2001 From: tpegolotti Date: Fri, 26 Jun 2026 13:49:48 +0200 Subject: [PATCH 20/56] Update gated delta benchmarking. Added Qwen3.5 dimensions --- benchmarks/python/gated_delta_bench.py | 114 +++++++++++++++---------- 1 file changed, 68 insertions(+), 46 deletions(-) diff --git a/benchmarks/python/gated_delta_bench.py b/benchmarks/python/gated_delta_bench.py index c13a42b7ae..518198d3b8 100644 --- a/benchmarks/python/gated_delta_bench.py +++ b/benchmarks/python/gated_delta_bench.py @@ -13,6 +13,10 @@ os.environ["MTL_CAPTURE_ENABLED"] = "1" +RED_BOLD = "\033[1;31m" +GREEN = "\033[0;32m" +RESET = "\033[0m" + def bench(f, *args): for i in range(N_warmup): @@ -51,24 +55,29 @@ def profile(f, *args): def gated_delta_ref( - q: mx.array, # [B, T, H, Dk] - k: mx.array, # [B, T, H, Dk] - v: mx.array, # [B, T, H, Dv] - g: mx.array, # [B, T, H] or [B, T, H, Dk] - beta: mx.array, # [B, T, H] - state: Optional[mx.array] = None, # [B, H, Dv, Dk] + q: mx.array, # [B, T, Hk, Dk] + k: mx.array, # [B, T, Hk, Dk] + v: mx.array, # [B, T, Hv, Dv] + g: mx.array, # [B, T, Hv] or [B, T, Hv, Dk] + beta: mx.array, # [B, T, Hv] + state: Optional[mx.array] = None, # [B, Hv, Dv, Dk] ) -> Tuple[mx.array, mx.array]: """ Implements: S_t = a_t S_{t-1} + b_t (v_t - a_t S_{t-1} k_t) k_t^T o_t = S_t q_t """ - B, T, H, Dk = q.shape + B, T, Hk, Dk = q.shape Dv = v.shape[-1] + Hv = v.shape[-2] if state is None: state = mx.zeros((B, H, Dv, Dk), dtype=mx.float32) + if (repeat_factor := Hv // Hk) > 1: + q = mx.repeat(q, repeat_factor, -2) + k = mx.repeat(k, repeat_factor, -2) + outputs = [] for t in range(T): q_t = q[:, t] # [B, H, Dk] @@ -102,18 +111,18 @@ def gated_delta_ref( return mx.stack(outputs, axis=1), state # [B, T, H, Dv], [B, H, Dv, Dk] -def benchmark_shape(B, T, H, Dk, Dv, chunk_sizes): +def benchmark_shape(B, T, Hk, Hv, Dk, Dv, chunk_sizes): mx.random.seed(42) - q = mx.random.normal(shape=(B, T, H, Dk)) - k = mx.random.normal(shape=(B, T, H, Dk)) + q = mx.random.normal(shape=(B, T, Hk, Dk)) + k = mx.random.normal(shape=(B, T, Hk, Dk)) k = k / (mx.linalg.norm(k, axis=-1, keepdims=True) + 1e-6) - v = mx.random.normal(shape=(B, T, H, Dv)) - g = mx.random.normal(shape=(B, T, H)) * 0.1 - 1.0 - b = mx.sigmoid(mx.random.normal(shape=(B, T, H))) - h0 = mx.zeros((B, H, Dv, Dk), dtype=mx.float32) + v = mx.random.normal(shape=(B, T, Hv, Dv)) + g = mx.random.normal(shape=(B, T, Hv)) * 0.1 - 1.0 + b = mx.sigmoid(mx.random.normal(shape=(B, T, Hv))) + h0 = mx.zeros((B, Hv, Dv, Dk), dtype=mx.float32) - shape_str = f"B={B} T={T} H={H} Dk={Dk} Dv={Dv}" + shape_str = f"B={B} T={T} Hk={Hk} Hv={Hv} Dk={Dk} Dv={Dv}" out_ref, hf_ref = gated_delta_ref(q, k, v, g, b, state=h0) mx.eval(out_ref, hf_ref) @@ -121,15 +130,22 @@ def benchmark_shape(B, T, H, Dk, Dv, chunk_sizes): out, hf = mx.fast.gated_delta_update(q, k, v, g, b, initial_state=h0, C=0) mx.eval(out, hf) - atol = 1e-2 + atol = 1e-1 out_close = mx.all(mx.abs(out - out_ref) < atol).item() hf_close = mx.all(mx.abs(hf - hf_ref) < atol).item() out_max_diff = mx.abs(out - out_ref).max().item() hf_max_diff = mx.abs(hf - hf_ref).max().item() - assert out_close, f"output mismatch! max diff: {out_max_diff:.2e}" - assert hf_close, f"state mismatch! max diff: {hf_max_diff:.2e}" + # assert out_close, f"output mismatch! max diff: {out_max_diff:.2e}" + # assert hf_close, f"state mismatch! max diff: {hf_max_diff:.2e}" + + # for b in range(B): + # diff = mx.abs(out[b] - out_ref[b]).max().item() + # print(f"batch {b}: {diff:.2e}") + # exit() + if not out_close or not hf_close: + print(f"{RED_BOLD}", end="") ms_seq = ( bench(do_kernel_bench, mx.fast.gated_delta_update, q, k, v, g, b, h0, 0) * 1000 @@ -140,7 +156,7 @@ def benchmark_shape(B, T, H, Dk, Dv, chunk_sizes): non_zero_Cs = [C for C in chunk_sizes if C != 0] for C in non_zero_Cs: try: - h0 = mx.zeros((B, H, Dv, Dk), dtype=mx.float32) + h0 = mx.zeros((B, Hv, Dv, Dk), dtype=mx.float32) out, hf = mx.fast.gated_delta_update(q, k, v, g, b, initial_state=h0, C=C) mx.eval(out, hf) @@ -164,8 +180,6 @@ def benchmark_shape(B, T, H, Dk, Dv, chunk_sizes): speedup = ms_seq / ms_c if ms_c > 0 else float("nan") - speedups.append(speedup) - out_close = mx.all(mx.abs(out - out_ref) < atol).item() hf_close = mx.all(mx.abs(hf - hf_ref) < atol).item() @@ -175,8 +189,10 @@ def benchmark_shape(B, T, H, Dk, Dv, chunk_sizes): if not out_close or not hf_close: raise Exception + speedups.append(speedup) + except Exception as e: - speedups.append("err") + speedups.append(-1) return shape_str, f"{ms_seq:.3f}", speedups @@ -186,55 +202,61 @@ def run_benchmark(run_full): Bs = [1, 2, 4, 8, 16] Ts = [2048, 4096, 8192, 16384] Hs = [16, 24, 32] - Dks = [64] - Dvs = [64] + Dks = [128] + Dvs = [128] else: Bs = [1, 8, 16] - Ts = [4096, 8192] - Hs = [16] - Dks = [64] - Dvs = [64] + Ts = [10, 512, 1024, 2048] + # Hs = [16] + Hks = [16] + Hvs = [32] + Dks = [128] + Dvs = [128] rows = [] chunk_sizes = [0, 8] non_zero_Cs = [C for C in chunk_sizes if C != 0] - headers = ["B", "T", "H", "Dk", "Dv", "time_seq (ms)"] + [ + headers = ["B", "T", "Hk", "Hv", "Dk", "Dv", "time_seq (ms)"] + [ f"C={C} speedup" for C in non_zero_Cs ] - col_widths = [6, 6, 6, 6, 6, 15] + [15] * (len(non_zero_Cs)) + col_widths = [6, 6, 6, 6, 6, 6, 15] + [15] * (len(non_zero_Cs)) fmt = "".join(f"{{:<{w}}}" for w in col_widths) print(fmt.format(*headers)) print("-" * (sum(col_widths))) - for B, T, H, Dk, Dv in itertools.product(Bs, Ts, Hs, Dks, Dvs): - shapes_s, base_time, speedups = benchmark_shape(B, T, H, Dk, Dv, chunk_sizes) - row = [f"{B}", f"{T}", f"{H}", f"{Dk}", f"{Dv}", base_time] + for B, T, Hk, Hv, Dk, Dv in itertools.product(Bs, Ts, Hks, Hvs, Dks, Dvs): + shapes_s, base_time, speedups = benchmark_shape( + B, T, Hk, Hv, Dk, Dv, chunk_sizes + ) + row = [f"{B}", f"{T}", f"{Hk}", f"{Hv}", f"{Dk}", f"{Dv}", base_time] for speed in speedups: row.append(f"{speed:.2f}x") - print(fmt.format(*row)) + print(fmt.format(*row), end="") + print(f"{RESET}") def run_profile(): B = 1 - H = 16 - T = 4096 * 2 - Dk = 64 - Dv = 64 + Hk = 16 + Hv = 32 + T = 512 + Dk = 128 + Dv = 128 CS = [8] # , 16]#, 32] mx.random.seed(42) - q = mx.random.normal(shape=(B, T, H, Dk)) - k = mx.random.normal(shape=(B, T, H, Dk)) + q = mx.random.normal(shape=(B, T, Hk, Dk)) + k = mx.random.normal(shape=(B, T, Hk, Dk)) k = k / (mx.linalg.norm(k, axis=-1, keepdims=True) + 1e-6) - v = mx.random.normal(shape=(B, T, H, Dv)) - g = mx.random.normal(shape=(B, T, H)) * 0.1 - 1.0 - b = mx.sigmoid(mx.random.normal(shape=(B, T, H))) - h0 = mx.zeros((B, H, Dv, Dk), dtype=mx.float32) + v = mx.random.normal(shape=(B, T, Hv, Dv)) + g = mx.random.normal(shape=(B, T, Hv)) * 0.1 - 1.0 + b = mx.sigmoid(mx.random.normal(shape=(B, T, Hv))) + h0 = mx.zeros((B, Hv, Dv, Dk), dtype=mx.float32) mx.eval(q, k, v, g, b, h0) @@ -245,8 +267,8 @@ def run_profile(): out, hf = mx.fast.gated_delta_update(q, k, v, g, b, initial_state=h0, C=C) mx.eval(out, hf) - assert list(out.shape) == [B, T, H, Dv] - assert list(hf.shape) == [B, H, Dv, Dk] + assert list(out.shape) == [B, T, Hv, Dv] + assert list(hf.shape) == [B, Hv, Dv, Dk] assert not mx.any(mx.isnan(out)).item(), "NaNs in output!" assert not mx.any(mx.isnan(hf)).item(), "NaNs in final state!" From 6a159af6fb9b3eba50a5422f9a513a7b2109b525 Mon Sep 17 00:00:00 2001 From: tpegolotti Date: Thu, 9 Jul 2026 01:48:52 -0700 Subject: [PATCH 21/56] Added first nax version --- mlx/backend/metal/gated_delta_update.cpp | 34 +- .../metal/kernels/gated_delta_update.metal | 117 +++- .../metal/kernels/gated_delta_update_impl.h | 586 +++++++++++++++++- mlx/fast.cpp | 3 + 4 files changed, 708 insertions(+), 32 deletions(-) diff --git a/mlx/backend/metal/gated_delta_update.cpp b/mlx/backend/metal/gated_delta_update.cpp index 59bfe3d76f..a0fd48092f 100644 --- a/mlx/backend/metal/gated_delta_update.cpp +++ b/mlx/backend/metal/gated_delta_update.cpp @@ -43,7 +43,6 @@ void GatedDeltaUpdate::eval_gpu( // printf("%d %d %d %d %d %d\n",B,T,Hk,Hv,Dk,Dv); int C = chunk_size; - int n_chunks = T / C; // TODO: make general std::string suffix = get_type_string(q.dtype()) // "float" + "_" + get_type_string(h0.dtype()) // "float" @@ -55,7 +54,40 @@ void GatedDeltaUpdate::eval_gpu( out.set_data(allocator::malloc(out.nbytes())); hf.set_data(allocator::malloc(hf.nbytes())); + fill_gpu(array(0, out.dtype()), out, s); + switch (C) { + case 16: { + std::string kernel_name = "gated_delta_fused_nax_"; + + // printf("C: %d\nname: %s\n",C,kernel_name.c_str()); + std::string base_name = kernel_name + suffix; + + base_name += "_" + std::to_string(C); + + std::string hash_name = base_name; + + metal::MTLFCList func_consts = {}; + + auto delta_kernel = get_steel_gated_delta_forward_kernel( + d, base_name, hash_name, func_consts); + + compute_encoder.set_compute_pipeline_state(delta_kernel); + compute_encoder.set_input_array(q, 0); + compute_encoder.set_input_array(k, 1); + compute_encoder.set_input_array(v, 2); + compute_encoder.set_input_array(h0, 3); // initial state in + compute_encoder.set_input_array(g, 4); + compute_encoder.set_input_array(beta, 5); + compute_encoder.set_output_array(out, 6); + compute_encoder.set_output_array(hf, 7); // final state out + compute_encoder.set_bytes(T, 8); + + auto grid = MTL::Size(32, Dv / 16, B * Hv); + auto threads = MTL::Size(32, 1, 1); + compute_encoder.dispatch_threads(grid, threads); + break; + } case 8: { std::string kernel_name = "gated_delta_fused_chunk_"; diff --git a/mlx/backend/metal/kernels/gated_delta_update.metal b/mlx/backend/metal/kernels/gated_delta_update.metal index 86b0788aa9..1e9adba9e2 100644 --- a/mlx/backend/metal/kernels/gated_delta_update.metal +++ b/mlx/backend/metal/kernels/gated_delta_update.metal @@ -46,25 +46,102 @@ using namespace metal; hv, \ c) -#define instantiate_gated_delta_update_fused_chunk_dims(in_type, st_type) \ - instantiate_gated_delta_update_fused_chunk( \ - in_type, st_type, 128, 128, 16, 16, 8) \ - instantiate_gated_delta_update_fused_chunk( \ - in_type, st_type, 128, 128, 24, 24, 8) \ - instantiate_gated_delta_update_fused_chunk( \ - in_type, st_type, 128, 128, 32, 32, 8) \ - instantiate_gated_delta_update_fused_chunk( \ - in_type, st_type, 64, 64, 4, 4, 8) \ - instantiate_gated_delta_update_fused_chunk( \ - in_type, st_type, 64, 64, 8, 8, 8) \ - instantiate_gated_delta_update_fused_chunk( \ - in_type, st_type, 64, 64, 16, 16, 8) \ - instantiate_gated_delta_update_fused_chunk( \ - in_type, st_type, 64, 64, 24, 24, 8) \ - instantiate_gated_delta_update_fused_chunk( \ - in_type, st_type, 64, 64, 32, 32, 8) \ - instantiate_gated_delta_update_fused_chunk( \ - in_type, st_type, 128, 128, 16, 32, 8) +#define instantiate_gated_delta_update_fused_chunk_dims(in_type, st_type) \ + instantiate_gated_delta_update_fused_chunk( \ + in_type, st_type, 128, 128, 16, 16, 8) \ + instantiate_gated_delta_update_fused_chunk( \ + in_type, st_type, 128, 128, 24, 24, 8) \ + instantiate_gated_delta_update_fused_chunk( \ + in_type, st_type, 128, 128, 32, 32, 8) \ + instantiate_gated_delta_update_fused_chunk( \ + in_type, st_type, 64, 64, 4, 4, 8) \ + instantiate_gated_delta_update_fused_chunk( \ + in_type, st_type, 64, 64, 8, 8, 8) \ + instantiate_gated_delta_update_fused_chunk( \ + in_type, st_type, 64, 64, 16, 16, 8) \ + instantiate_gated_delta_update_fused_chunk( \ + in_type, st_type, 64, 64, 24, 24, 8) \ + instantiate_gated_delta_update_fused_chunk( \ + in_type, st_type, 64, 64, 32, 32, 8) \ + instantiate_gated_delta_update_fused_chunk( \ + in_type, st_type, 128, 128, 16, 32, 8) \ + instantiate_gated_delta_update_fused_chunk( \ + in_type, st_type, 32, 32, 16, 32, 8) \ + instantiate_gated_delta_update_fused_chunk( \ + in_type, \ + st_type, \ + 32, \ + 32, \ + 32, \ + 32, \ + 8) \ + instantiate_gated_delta_update_fused_chunk( \ + in_type, \ + st_type, \ + 16, \ + 16, \ + 16, \ + 32, \ + 8) + +#define instantiate_gated_delta_update_fused_nax( \ + in_type, st_type, dk, dv, hk, hv, c) \ + instantiate_kernel( \ + "gated_delta_fused_nax_" #in_type "_" #st_type "_" #dk "_" #dv "_" #hk \ + "_" #hv "_" #c, \ + gated_delta_fused_nax, \ + in_type, \ + st_type, \ + dk, \ + dv, \ + hk, \ + hv, \ + c) + +#define instantiate_gated_delta_update_fused_nax_dims(in_type, st_type) \ + instantiate_gated_delta_update_fused_nax( \ + in_type, st_type, 128, 128, 16, 16, 16) \ + instantiate_gated_delta_update_fused_nax( \ + in_type, st_type, 128, 128, 24, 24, 16) \ + instantiate_gated_delta_update_fused_nax( \ + in_type, st_type, 128, 128, 32, 32, 16) \ + instantiate_gated_delta_update_fused_nax( \ + in_type, st_type, 64, 64, 4, 4, 16) \ + instantiate_gated_delta_update_fused_nax( \ + in_type, st_type, 64, 64, 8, 8, 16) \ + instantiate_gated_delta_update_fused_nax( \ + in_type, st_type, 64, 64, 16, 16, 16) \ + instantiate_gated_delta_update_fused_nax( \ + in_type, st_type, 64, 64, 24, 24, 16) \ + instantiate_gated_delta_update_fused_nax( \ + in_type, st_type, 64, 64, 32, 32, 16) \ + instantiate_gated_delta_update_fused_nax( \ + in_type, st_type, 128, 128, 16, 32, 16) \ + instantiate_gated_delta_update_fused_nax( \ + in_type, \ + st_type, \ + 32, \ + 32, \ + 16, \ + 32, \ + 16) \ + instantiate_gated_delta_update_fused_nax( \ + in_type, \ + st_type, \ + 32, \ + 32, \ + 32, \ + 32, \ + 16) \ + instantiate_gated_delta_update_fused_nax( \ + in_type, \ + st_type, \ + 16, \ + 16, \ + 16, \ + 32, \ + 16) instantiate_gated_delta_update_seq_dims(float, float) - instantiate_gated_delta_update_fused_chunk_dims(float, float) \ No newline at end of file + instantiate_gated_delta_update_fused_chunk_dims(float, float) + instantiate_gated_delta_update_fused_nax_dims(float, float) \ No newline at end of file diff --git a/mlx/backend/metal/kernels/gated_delta_update_impl.h b/mlx/backend/metal/kernels/gated_delta_update_impl.h index 9f057cfe80..601242b07c 100644 --- a/mlx/backend/metal/kernels/gated_delta_update_impl.h +++ b/mlx/backend/metal/kernels/gated_delta_update_impl.h @@ -3,7 +3,17 @@ #include #include "mlx/backend/metal/kernels/utils.h" +#include +#include + +#include "mlx/backend/metal/kernels/steel/gemm/nax.h" +#include "mlx/backend/metal/kernels/steel/gemm/params.h" +#include "mlx/backend/metal/kernels/steel/gemm/transforms.h" +#include "mlx/backend/metal/kernels/steel/utils.h" + using namespace metal; +using namespace mpp; +using namespace mpp::tensor_ops; #define AT(TILE, IDX) TILE.thread_elements()[IDX] #define SUB(TILE0, TILE1, TILE2) \ @@ -38,6 +48,544 @@ using namespace metal; AT(TILE0, 1) *= fn + 1 >= fm ? 0.f : S1; \ } +// NAX MACROS I can probably do a nice template instead of doing this + +// fm = base_fm + (idx >> 2) * 8; // idx>>2 = idx/4 -> 0 for idx 0-3, 1 for +// idx 4-7 fn = base_fn + (idx % 4); // 4 consecutive columns +#define AT_NAX(TILE, IDX) TILE.elems()[IDX] + +#define SUB_NAX(TILE0, TILE1, TILE2) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + AT_NAX(TILE0, _i) = AT_NAX(TILE1, _i) - AT_NAX(TILE2, _i); \ + } \ + } + +#define FMA_NAX(TILE0, S, TILE1, TILE2) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < mlx::steel::BaseNAXFrag::kElemsPerFrag; _i++) { \ + (TILE0)[_i] = (S) * (TILE1)[_i] + (TILE2)[_i]; \ + } \ + } + +#define SCALE_NAX(TILE0, S) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + AT_NAX(TILE0, _i) *= (S); \ + } \ + } + +#define SCALE_ROW_NAX(TILE0, S) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; \ + AT_NAX(TILE0, _i) *= (S)[mlx::steel::BaseNAXFrag::get_coord(_w).y]; \ + } \ + } + +#define SCALE_BETA_NAX(TILE0, BETA2) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; \ + AT_NAX(TILE0, _i) *= (BETA2)[_w >> 2]; \ + } \ + } + +// gamma is indexed by absolute column; get_coord(idx).x gives that column. +#define SCALE2_NAX(TILE0, GAMMA) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerFrag; _i++) { \ + const short _fm = mlx::steel::BaseNAXFrag::get_coord(_i).y; \ + AT_NAX(TILE0, _i) *= (GAMMA)[(C) - 1] / (GAMMA)[_fm]; \ + } \ + } + +#define SCALE_TRI_NAX(TILE0, GAMMA) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerFrag; _i++) { \ + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); /* {fn, fm} */ \ + AT_NAX(TILE0, _i) *= (_c.x > _c.y) ? 0.f : (1.0f / (GAMMA)[_c.x]); \ + } \ + } + +#define SCALE_TRIEQ_NAX1(TILE0, BETA) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerFrag; _i++) { \ + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); /* {fn, fm} */ \ + AT_NAX(TILE0, _i) *= (_c.x >= _c.y) ? 0.f : (BETA)[_i >> 2]; \ + } \ + } + +#define SCALE_TRIEQ_NAX(TILE0, BETA, GAMMA) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerFrag; _i++) { \ + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); /* {fn, fm} */ \ + const short _fn = _c.x; \ + const short _fm = _c.y; \ + const float _s = (BETA)[_i >> 2] * ((GAMMA)[_fm] / (GAMMA)[_fn]); \ + AT_NAX(TILE0, _i) *= (_fn >= _fm) ? 0.f : _s; \ + } \ + } + +namespace mlx { +namespace steel { +template < + typename CType, + typename AType, + typename BType, + bool transpose_a = false, + bool transpose_b = false, + mpp::tensor_ops::matmul2d_descriptor::mode Mode = + mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate> +METAL_FUNC static constexpr void mmak( + thread BaseNAXFrag::dtype_frag_t& C, + const thread BaseNAXFrag::dtype_frag_t& A0, + const thread BaseNAXFrag::dtype_frag_t& A1, + metal::bool_constant, + const thread BaseNAXFrag::dtype_frag_t& B0, + const thread BaseNAXFrag::dtype_frag_t& B1, + metal::bool_constant) { + // K = 32: two K-fragments per operand, single 16x16 output. + constexpr auto desc = mpp::tensor_ops::matmul2d_descriptor( + 16, 16, 32, transpose_a, transpose_b, true, Mode); + + mpp::tensor_ops::matmul2d gemm_op; + + // Create matmul operands in registers + auto ct_a = + gemm_op.template get_left_input_cooperative_tensor(); + auto ct_b = + gemm_op + .template get_right_input_cooperative_tensor(); + + // Create matmul output in register + auto ct_c = gemm_op.template get_destination_cooperative_tensor< + decltype(ct_a), + decltype(ct_b), + CType>(); + + STEEL_PRAGMA_UNROLL + for (short i = 0; i < BaseNAXFrag::kElemsPerFrag; i++) { + ct_a[i] = A0[i]; + ct_a[BaseNAXFrag::kElemsPerFrag + i] = A1[i]; + ct_b[i] = B0[i]; + ct_b[BaseNAXFrag::kElemsPerFrag + i] = B1[i]; + ct_c[i] = C[i]; + } + + gemm_op.run(ct_a, ct_b, ct_c); + + STEEL_PRAGMA_UNROLL + for (short i = 0; i < BaseNAXFrag::kElemsPerFrag; i++) { + C[i] = ct_c[i]; + } +} + +template < + typename CType, + typename AType, + typename BType, + bool transpose_a = false, + bool transpose_b = false, + mpp::tensor_ops::matmul2d_descriptor::mode Mode = + mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate> +METAL_FUNC static constexpr void mman( + thread BaseNAXFrag::dtype_frag_t& C0, + thread BaseNAXFrag::dtype_frag_t& C1, + const thread BaseNAXFrag::dtype_frag_t& A, + metal::bool_constant, + const thread BaseNAXFrag::dtype_frag_t& B0, + const thread BaseNAXFrag::dtype_frag_t& B1, + metal::bool_constant) { + // N = 32: single A fragment, two N-fragments for B and C output. + // template parameters are M, N, K where + // Tensor dimensions where M x K tensor A,K x N tensor B, and M x N tensor C. + constexpr auto desc = mpp::tensor_ops::matmul2d_descriptor( + 16, 32, 16, transpose_a, transpose_b, true, Mode); + + mpp::tensor_ops::matmul2d gemm_op; + + // Create matmul operands in registers + auto ct_a = + gemm_op.template get_left_input_cooperative_tensor(); + auto ct_b = + gemm_op + .template get_right_input_cooperative_tensor(); + + // Create matmul output in register + auto ct_c = gemm_op.template get_destination_cooperative_tensor< + decltype(ct_a), + decltype(ct_b), + CType>(); + + STEEL_PRAGMA_UNROLL + for (short i = 0; i < BaseNAXFrag::kElemsPerFrag; i++) { + ct_a[i] = A[i]; + } + + STEEL_PRAGMA_UNROLL + for (short i = 0; i < BaseNAXFrag::kElemsPerFrag; i++) { + ct_b[i] = B0[i]; + ct_b[BaseNAXFrag::kElemsPerFrag + i] = B1[i]; + } + + STEEL_PRAGMA_UNROLL + for (short i = 0; i < BaseNAXFrag::kElemsPerFrag; i++) { + ct_c[i] = C0[i]; + ct_c[BaseNAXFrag::kElemsPerFrag + i] = C1[i]; + } + + gemm_op.run(ct_a, ct_b, ct_c); + + // Copy out both N-fragments + STEEL_PRAGMA_UNROLL + for (short i = 0; i < BaseNAXFrag::kElemsPerFrag; i++) { + C0[i] = ct_c[i]; + C1[i] = ct_c[BaseNAXFrag::kElemsPerFrag + i]; + } +} +} // namespace steel +} // namespace mlx + +template +[[kernel]] void gated_delta_fused_nax( + const device InT* q [[buffer(0)]], + const device InT* k [[buffer(1)]], + const device InT* v [[buffer(2)]], + const device StT* state_in [[buffer(3)]], + const device InT* g [[buffer(4)]], + const device InT* beta [[buffer(5)]], + device InT* y [[buffer(6)]], + device StT* state_out [[buffer(7)]], + constant int& T [[buffer(8)]], + uint3 thread_position_in_grid [[thread_position_in_grid]], + uint thread_index_in_simdgroup [[thread_index_in_simdgroup]]) { + auto n = thread_position_in_grid.z; + auto b_idx = n / Hv; + auto hv_idx = n % Hv; + auto hk_idx = hv_idx / (Hv / Hk); + + const ushort simd_lane_id = __metal_get_thread_index_in_simdgroup(ushort()); + const short qid = simd_lane_id >> 2; + const short fm = ((qid & 4) | ((simd_lane_id >> 1) & 3)); + const short fn = ((qid & 2) | (simd_lane_id & 1)) * 4; + + auto dv_idx = thread_position_in_grid.y * 16; + + // set up pointers + // g: [B, T, Hv] + auto g_ = g + b_idx * T * Hv; + + // q, k: [B, T, Hk, Dk] + auto q_ = q + b_idx * T * Hk * Dk + hk_idx * Dk; + auto k_ = k + b_idx * T * Hk * Dk + hk_idx * Dk; + + // v, y: [B, T, Hv, Dv] + y += b_idx * T * Hv * Dv + hv_idx * Dv; + auto v_ = v + b_idx * T * Hv * Dv + hv_idx * Dv; + auto beta_ = beta + b_idx * T * Hv; + + // state_in, state_out: [B, Hv, Dv, Dk] + auto i_state = state_in + (n * Dv + dv_idx) * Dk; + auto o_state = state_out + (n * Dv + dv_idx) * Dk; + + threadgroup float gamma[C]; + float beta_fm[2]; + + mlx::steel::NAXTile S_tile; + S_tile.load(i_state, Dk); + + mlx::steel::NAXTile K_tile, Q_tile; + mlx::steel::NAXTile KP_tile; // panel + mlx::steel::NAXTile W_tile; // panel + mlx::steel::NAXTile XP_tile; // panel + + mlx::steel::NAXTile X_tile; + + mlx::steel::NAXTile V_tile; + mlx::steel::NAXTile U_tile; + mlx::steel::NAXTile WS_tile; + mlx::steel::NAXTile delta_tile; + mlx::steel::NAXTile tmp_tile; + mlx::steel::NAXTile QKt_tile; + mlx::steel::NAXTile out_tile; + mlx::steel::NAXTile K1_tile, Q1_tile; + mlx::steel::NAXTile KD_tile; + mlx::steel::NAXTile Tinv_tile; + + mlx::steel::NAXTile KKtK_tile, KKtV_tile, KKt_tile; + + // Annoyingly one dimension needs to be 32 so for now let's see if this works + + mlx::steel::NAXTile I_tile; // Declare and init eye + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(I_tile)::kElemsPerFrag; _i++) { + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); /* {fn, fm} */ + const short _fn = _c.x; + const short _fm = _c.y; + AT_NAX(I_tile, _i) = (_fn == _fm) ? 1.0f : 0.0f; + } + mlx::steel::NAXTile TMP_tile; + + mlx::steel::NAXTile Z; + mlx::steel::NAXTile junk; + + junk.clear(); + Z.clear(); + + for (int t = 0; t < T; t += C) { + float g_val = (thread_index_in_simdgroup < C) + ? g_[thread_index_in_simdgroup * Hv + hv_idx] + : 1.0f; + + float gamma_val = simd_prefix_inclusive_product(g_val); + + if (thread_index_in_simdgroup < C) { + gamma[thread_index_in_simdgroup] = gamma_val; + } + + beta_fm[0] = beta_[fm * Hv + hv_idx]; + beta_fm[1] = + beta_[(fm + mlx::steel::BaseNAXFrag::kElemRowsJump) * Hv + hv_idx]; + + // I can only do matmuls with a dimension 32. This one is _easy_ + // i just do two tiles for the reduction + KKt_tile.clear(); + for (int kk = 0; kk < Dk; kk += 32) { // two 16-tiles per iter + K_tile.load(k_ + kk, Dk * Hk); + mlx::steel::mmak( + KKt_tile.frag_at(0, 0), + K_tile.frag_at(0, 0), + K_tile.frag_at(0, 1), + metal::bool_constant{}, + K_tile.frag_at(0, 0), + K_tile.frag_at(0, 1), + metal::bool_constant{}); + } + +#define OUTPUT(T) \ + T.store(y, 16); \ + return; + + // KKt_tile.store(y,16); + // return; + + KKtK_tile = KKt_tile; + KKtV_tile = KKt_tile; + + SCALE_TRIEQ_NAX1(KKtK_tile, beta_fm) + SCALE_TRIEQ_NAX(KKtV_tile, beta_fm, gamma) + + // W = (I - diag(b)(KK.T))^-1 diag(b)K + // diag(b)K + KP_tile.load(k_, Dk * Hk); + SCALE_BETA_NAX(KP_tile, beta_fm) + + // (I - diag(b)(KK.T))^-1 = sum T^k + // Tinv = (I + L_W)^{-1} = sum -L_W_k + // T0 - T + T2 - T3 + T4 - T5 + T6 - T7 + T8 + // I - T(I - T(I - T(I - T(I - T(I - T(I - T(I - T(I)))))))) + SUB_NAX(Tinv_tile, I_tile, KKtK_tile) + for (int iter = 0; iter < C - 2; iter++) { + mlx::steel::mmak< + float, + float, + float, + false, + false, + mpp::tensor_ops::matmul2d_descriptor::mode::multiply>( + TMP_tile.frag_at(0, 0), + KKtK_tile.frag_at(0, 0), + Z.frag_at(0, 0), + metal::bool_constant{}, + Tinv_tile.frag_at(0, 0), + Z.frag_at(0, 0), + metal::bool_constant{}); + // W = K - X + SUB_NAX(Tinv_tile, I_tile, TMP_tile) + } + // OUTPUT(Tinv_tile) + + // W = Tinv @ (diag(beta) K) + STEEL_PRAGMA_UNROLL + for (short nn = 0; nn < Dk / 16; nn += 2) { + mlx::steel::mman< + float, + float, + float, + false, + false, + mpp::tensor_ops::matmul2d_descriptor::mode::multiply>( + W_tile.frag_at(0, nn), + W_tile.frag_at(0, nn + 1), + Tinv_tile.frag_at(0, 0), // A = Tinv (reused across N) + metal::bool_constant{}, + KP_tile.frag_at(0, nn), // B = beta-scaled K + KP_tile.frag_at(0, nn + 1), + metal::bool_constant{}); + } + SCALE_ROW_NAX(W_tile, gamma) + + V_tile.load(v_ + dv_idx, Dv * Hv); + SCALE_BETA_NAX(V_tile, beta_fm) + Z.clear(); + SUB_NAX(Tinv_tile, I_tile, KKtV_tile) + U_tile = V_tile; + for (int iter = 0; iter < C - 2; iter++) { + mlx::steel::mmak< + float, + float, + float, + false, + false, + mpp::tensor_ops::matmul2d_descriptor::mode::multiply>( + TMP_tile.frag_at(0, 0), + KKtV_tile.frag_at(0, 0), + Z.frag_at(0, 0), + metal::bool_constant{}, + Tinv_tile.frag_at(0, 0), + Z.frag_at(0, 0), + metal::bool_constant{}); + + SUB_NAX(Tinv_tile, I_tile, TMP_tile) + } + + // U = Tinv @ diag(b)V + mlx::steel::mmak< + float, + float, + float, + false, + false, + mpp::tensor_ops::matmul2d_descriptor::mode::multiply>( + U_tile.frag_at(0, 0), + Tinv_tile.frag_at(0, 0), + Z.frag_at(0, 0), + metal::bool_constant{}, + V_tile.frag_at(0, 0), + Z.frag_at(0, 0), + metal::bool_constant{}); + + // WS = W @ S.T, easy reduction with stride 2 over K + WS_tile.clear(); + STEEL_PRAGMA_UNROLL + for (short kk = 0; kk < Dk / 16; kk += 2) { + mlx::steel::mmak( + WS_tile.frag_at(0, 0), + W_tile.frag_at(0, kk), + W_tile.frag_at(0, kk + 1), + metal::bool_constant{}, + S_tile.frag_at(0, kk), + S_tile.frag_at(0, kk + 1), + metal::bool_constant{}); + } + // OUTPUT(WS_tile) + + SUB_NAX(delta_tile, U_tile, WS_tile) + // STEEL_PRAGMA_UNROLL + // for (short _i = 0; _i < decltype(delta_tile)::kElemsPerFrag; _i++) { + // const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); + // const short _fn = _c.x; + // const short _fm = _c.y; + // AT_NAX(delta_tile, _i) = _fn >= 8 ? 0 : AT_NAX(delta_tile, _i); + // } + // OUTPUT(delta_tile) + + tmp_tile.clear(); + QKt_tile.clear(); + for (int kk = 0; kk < Dk; kk += 32) { + Q_tile.load(q_ + kk, Hk * Dk); + K_tile.load(k_ + kk, Hk * Dk); // this one is transposed + + SCALE_ROW_NAX(Q_tile, gamma) + + // Q_left @ S^T + mlx::steel::mmak( + tmp_tile.frag_at(0, 0), + Q_tile.frag_at(0, 0), + Q_tile.frag_at(0, 1), + metal::bool_constant{}, + S_tile.frag_at(0, kk / 16), + S_tile.frag_at(0, kk / 16 + 1), + metal::bool_constant{}); + + // Q @ K^T + mlx::steel::mmak( + QKt_tile.frag_at(0, 0), + Q_tile.frag_at(0, 0), + Q_tile.frag_at(0, 1), + metal::bool_constant{}, + K_tile.frag_at(0, 0), + K_tile.frag_at(0, 1), + metal::bool_constant{}); + } + // OUTPUT(tmp_tile) + + SCALE_TRI_NAX(QKt_tile, gamma) + // OUTPUT(QKt_tile) + + out_tile = tmp_tile; + mlx::steel::mmak( + out_tile.frag_at(0, 0), + QKt_tile.frag_at(0, 0), + Z.frag_at(0, 0), + metal::bool_constant{}, + delta_tile.frag_at(0, 0), + Z.frag_at(0, 0), + metal::bool_constant{}); + + out_tile.store(y + dv_idx, Hv * Dv); + // OUTPUT(out_tile); + + for (int kk = 0; kk < Dk; kk += 16) { + K1_tile.load(k_ + kk, Hk * Dk); + SCALE2_NAX(K1_tile, gamma) + + // KD = delta @ K.T so that I can sum to S directly + KD_tile.clear(); + mlx::steel::mmak< + float, + InT, + float, + true, + false, + mpp::tensor_ops::matmul2d_descriptor::mode::multiply>( + KD_tile.frag_at(0, 0), + delta_tile.frag_at(0, 0), + Z.frag_at(0, 0), + metal::bool_constant{}, + K1_tile.frag_at(0, 0), + Z.frag_at(0, 0), + metal::bool_constant{}); + + FMA_NAX( + S_tile.frag_at(0, kk / 16), + gamma[C - 1], + S_tile.frag_at(0, kk / 16), + KD_tile.frag_at(0, 0)) + } + // OUTPUT(out_tile); + + // advance pointers + q_ += C * Hk * Dk; + k_ += C * Hk * Dk; + v_ += C * Hv * Dv; + beta_ += C * Hv; + y += C * Hv * Dv; + g_ += C * Hv; + } + S_tile.store(o_state, Dk); +} + template [[kernel]] void gated_delta_fused_chunk( const device InT* q [[buffer(0)]], @@ -98,6 +646,10 @@ template threadgroup float gamma[C]; + simdgroup_float8x8 I_tile = make_filled_simdgroup_matrix(0.f); + AT(I_tile, 0) = (fm == fn) ? 1.0f : 0.0f; + AT(I_tile, 1) = (fm == fn + 1) ? 1.0f : 0.0f; + // load initial state into threadgroup for (int kk = 0; kk < Dk; kk += 8) { simdgroup_load(S_tile[kk / 8], i_state + kk, Dk, ulong2(0, 0), true); @@ -122,6 +674,9 @@ template simdgroup_load(KT_tile, k_ + kk, Dk * Hk, ulong2(0, 0), true); simdgroup_multiply_accumulate(KKt_tile, K_tile, KT_tile, KKt_tile); } +#define OUTPUT(T) \ + simdgroup_store(T, y); \ + return; KKtK_tile = KKt_tile; KKtV_tile = KKt_tile; @@ -133,25 +688,30 @@ template beta_fm * (gamma[fm] / gamma[fn]), beta_fm * (gamma[fm] / gamma[fn + 1])) + // Tinv = (I + L_W)^{-1} = sum -L_W_k + // T0 - T + T2 - T3 + T4 - T5 + T6 - T7 + T8 + // I - T(I - T(I - T(I - T(I - T(I - T(I - T(I - T(I)))))))) + simdgroup_float8x8 Tinv; + SUB(Tinv, I_tile, KKtK_tile) + for (int iter = 0; iter < C - 2; iter++) { + simdgroup_float8x8 TMP; + simdgroup_multiply(TMP, KKtK_tile, Tinv); // L @ Tinv + SUB(Tinv, I_tile, TMP) // Tinv = I - L @ Tinv + } + // OUTPUT(Tinv) + WS_tile = make_filled_simdgroup_matrix(0.f); - // Use the Neumann series: (I - T)^-1 = sum T^k, instead of doing forward - // substitution to compute W (and U). For C=8 this is 8 matmuls, for generic - // C this probably does not work well? Also fuse the WS in the same loop -> - // less memory movements. for (int kk = 0; kk < Dk; kk += 8) { simdgroup_load(K_tile, k_ + kk, Dk * Hk); SCALE(K_tile, beta_fm) - W_tile = K_tile; - for (int iter = 0; iter < C - 1; iter++) { - simdgroup_multiply(X_tile, KKtK_tile, W_tile); - SUB(W_tile, K_tile, X_tile) - } - SCALE(W_tile, gamma[fm]) + // W = Tinv @ (beta * K) + simdgroup_multiply(W_tile, Tinv, K_tile); - // WS = W_left @ S^T + SCALE(W_tile, gamma[fm]) simdgroup_multiply_accumulate(WS_tile, W_tile, S_tile[kk / 8], WS_tile); } + // OUTPUT(WS_tile) simdgroup_load(V_tile, v_ + dv_idx, Dv * Hv); SCALE(V_tile, beta_fm) @@ -163,6 +723,7 @@ template // delta = U - WS SUB(delta_tile, U_tile, WS_tile) + // OUTPUT(delta_tile) tmp_tile = make_filled_simdgroup_matrix(0.f); QKt_tile = make_filled_simdgroup_matrix(0.f); @@ -176,14 +737,17 @@ template simdgroup_multiply_accumulate(tmp_tile, Q_tile, S_tile[kk / 8], tmp_tile); simdgroup_multiply_accumulate(QKt_tile, Q_tile, K_tile, QKt_tile); } + // OUTPUT(tmp_tile) // (Q @ K) * Gamma, in the paper they use M here but it's probably a typo. SCALE_TRI(QKt_tile, (1.0f / gamma[fn]), (1.0f / gamma[fn + 1])) + // OUTPUT(QKt_tile) simdgroup_multiply_accumulate(out_tile, QKt_tile, delta_tile, tmp_tile); y[fm * Hv * Dv + dv_idx + fn] = static_cast(AT(out_tile, 0)); y[fm * Hv * Dv + dv_idx + fn + 1] = static_cast(AT(out_tile, 1)); + // OUTPUT(out_tile) for (int kk = 0; kk < Dk; kk += 8) { simdgroup_load(K_tile, k_ + kk, Hk * Dk, ulong2(0, 0), true); diff --git a/mlx/fast.cpp b/mlx/fast.cpp index c73d55f554..885319f08c 100644 --- a/mlx/fast.cpp +++ b/mlx/fast.cpp @@ -944,6 +944,9 @@ std::vector gated_delta_update( auto g = astype(gates, out_dtype, s); auto beta = astype(beta_, out_dtype, s); + // wrong, not as fast. This kind of operations need to be done in + // the eval gpu inside a function that checks for shapes and strides + // TODO: mode this q = contiguous(q, false, s); k = contiguous(k, false, s); v = contiguous(v, false, s); From 03e36763a6fb5e8b7bd5a92ed85d9cf4166e7923 Mon Sep 17 00:00:00 2001 From: tpegolotti Date: Mon, 13 Jul 2026 01:44:27 -0700 Subject: [PATCH 22/56] improve inverse computation --- .../metal/kernels/gated_delta_update_impl.h | 169 ++++++++++++------ 1 file changed, 111 insertions(+), 58 deletions(-) diff --git a/mlx/backend/metal/kernels/gated_delta_update_impl.h b/mlx/backend/metal/kernels/gated_delta_update_impl.h index 601242b07c..f4d9b01af9 100644 --- a/mlx/backend/metal/kernels/gated_delta_update_impl.h +++ b/mlx/backend/metal/kernels/gated_delta_update_impl.h @@ -21,6 +21,11 @@ using namespace mpp::tensor_ops; AT(TILE0, 0) = AT(TILE1, 0) - AT(TILE2, 0); \ AT(TILE0, 1) = AT(TILE1, 1) - AT(TILE2, 1); \ } +#define ADD(TILE0, TILE1, TILE2) \ + { \ + AT(TILE0, 0) = AT(TILE1, 0) + AT(TILE2, 0); \ + AT(TILE0, 1) = AT(TILE1, 1) + AT(TILE2, 1); \ + } #define FMA(TILE0, S, TILE1, TILE2) \ { \ AT(TILE0, 0) = S * AT(TILE1, 0) + AT(TILE2, 0); \ @@ -96,14 +101,15 @@ using namespace mpp::tensor_ops; } \ } -// gamma is indexed by absolute column; get_coord(idx).x gives that column. -#define SCALE2_NAX(TILE0, GAMMA) \ - { \ - STEEL_PRAGMA_UNROLL \ - for (short _i = 0; _i < decltype(TILE0)::kElemsPerFrag; _i++) { \ - const short _fm = mlx::steel::BaseNAXFrag::get_coord(_i).y; \ - AT_NAX(TILE0, _i) *= (GAMMA)[(C) - 1] / (GAMMA)[_fm]; \ - } \ +#define SCALE2_NAX(TILE0, GAMMA) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + const short _fm = mlx::steel::BaseNAXFrag::get_coord( \ + _i % mlx::steel::BaseNAXFrag::kElemsPerFrag) \ + .y; \ + AT_NAX(TILE0, _i) *= (GAMMA)[(C) - 1] / (GAMMA)[_fm]; \ + } \ } #define SCALE_TRI_NAX(TILE0, GAMMA) \ @@ -274,6 +280,18 @@ template auto hv_idx = n % Hv; auto hk_idx = hv_idx / (Hv / Hk); +#define OUTPUT_16(T) \ + T.store(y, 16); \ + return; + +#define OUTPUT_S(T, S) \ + T.store(y, S); \ + return; + +// pick the 3rd argument as the macro name +#define OUTPUT_GET(_1, _2, NAME, ...) NAME +#define OUTPUT(...) OUTPUT_GET(__VA_ARGS__, OUTPUT_S, OUTPUT_16)(__VA_ARGS__) + const ushort simd_lane_id = __metal_get_thread_index_in_simdgroup(ushort()); const short qid = simd_lane_id >> 2; const short fm = ((qid & 4) | ((simd_lane_id >> 1) & 3)); @@ -319,12 +337,12 @@ template mlx::steel::NAXTile QKt_tile; mlx::steel::NAXTile out_tile; mlx::steel::NAXTile K1_tile, Q1_tile; - mlx::steel::NAXTile KD_tile; + mlx::steel::NAXTile KD_tile; mlx::steel::NAXTile Tinv_tile; mlx::steel::NAXTile KKtK_tile, KKtV_tile, KKt_tile; - // Annoyingly one dimension needs to be 32 so for now let's see if this works + mlx::steel::NAXTile P_tile; mlx::steel::NAXTile I_tile; // Declare and init eye STEEL_PRAGMA_UNROLL @@ -372,10 +390,6 @@ template metal::bool_constant{}); } -#define OUTPUT(T) \ - T.store(y, 16); \ - return; - // KKt_tile.store(y,16); // return; @@ -394,8 +408,20 @@ template // Tinv = (I + L_W)^{-1} = sum -L_W_k // T0 - T + T2 - T3 + T4 - T5 + T6 - T7 + T8 // I - T(I - T(I - T(I - T(I - T(I - T(I - T(I - T(I)))))))) - SUB_NAX(Tinv_tile, I_tile, KKtK_tile) - for (int iter = 0; iter < C - 2; iter++) { + // Update 10.7: compute the above as (1 - T)(1 + T^2)(1 + T^4)(1 + T^8) + // => 6 mults instead of 14 + + Z.clear(); + STEEL_PRAGMA_UNROLL + for (short e = 0; e < decltype(P_tile)::kElemsPerFrag; e++) { + AT_NAX(P_tile, e) = AT_NAX(KKtK_tile, e); + } + + // S = I + x = I - KKtV + SUB_NAX(Tinv_tile, I_tile, KKtK_tile) // Tinv = I - T (2 terms) + STEEL_PRAGMA_UNROLL + for (int step = 1; (1 << step) < C; step++) { // + // P2 = P · P mlx::steel::mmak< float, float, @@ -403,15 +429,23 @@ template false, false, mpp::tensor_ops::matmul2d_descriptor::mode::multiply>( - TMP_tile.frag_at(0, 0), - KKtK_tile.frag_at(0, 0), + P_tile.frag_at(0, 0), + P_tile.frag_at(0, 0), Z.frag_at(0, 0), metal::bool_constant{}, + P_tile.frag_at(0, 0), + Z.frag_at(0, 0), + metal::bool_constant{}); + + // Tinv = S · P2 + mlx::steel::mmak( Tinv_tile.frag_at(0, 0), + Tinv_tile.frag_at(0, 0), + Z.frag_at(0, 0), + metal::bool_constant{}, + P_tile.frag_at(0, 0), Z.frag_at(0, 0), metal::bool_constant{}); - // W = K - X - SUB_NAX(Tinv_tile, I_tile, TMP_tile) } // OUTPUT(Tinv_tile) @@ -435,12 +469,18 @@ template } SCALE_ROW_NAX(W_tile, gamma) - V_tile.load(v_ + dv_idx, Dv * Hv); - SCALE_BETA_NAX(V_tile, beta_fm) Z.clear(); + STEEL_PRAGMA_UNROLL + for (short e = 0; e < decltype(P_tile)::kElemsPerFrag; e++) { + AT_NAX(P_tile, e) = AT_NAX(KKtV_tile, e); + } + + // S = I + x = I - KKtV + // Tinv = (1 - x)(1 + x^2)(1 + x^4)(1 + x^8) SUB_NAX(Tinv_tile, I_tile, KKtV_tile) - U_tile = V_tile; - for (int iter = 0; iter < C - 2; iter++) { + STEEL_PRAGMA_UNROLL + for (int step = 1; (1 << step) < C; step++) { + // P^2 mlx::steel::mmak< float, float, @@ -448,17 +488,26 @@ template false, false, mpp::tensor_ops::matmul2d_descriptor::mode::multiply>( - TMP_tile.frag_at(0, 0), - KKtV_tile.frag_at(0, 0), + P_tile.frag_at(0, 0), + P_tile.frag_at(0, 0), Z.frag_at(0, 0), metal::bool_constant{}, - Tinv_tile.frag_at(0, 0), + P_tile.frag_at(0, 0), Z.frag_at(0, 0), metal::bool_constant{}); - SUB_NAX(Tinv_tile, I_tile, TMP_tile) + // Tinv = S · P2 + mlx::steel::mmak( + Tinv_tile.frag_at(0, 0), + Tinv_tile.frag_at(0, 0), + Z.frag_at(0, 0), + metal::bool_constant{}, + P_tile.frag_at(0, 0), + Z.frag_at(0, 0), + metal::bool_constant{}); } - + V_tile.load(v_ + dv_idx, Dv * Hv); + SCALE_BETA_NAX(V_tile, beta_fm) // U = Tinv @ diag(b)V mlx::steel::mmak< float, @@ -491,14 +540,6 @@ template // OUTPUT(WS_tile) SUB_NAX(delta_tile, U_tile, WS_tile) - // STEEL_PRAGMA_UNROLL - // for (short _i = 0; _i < decltype(delta_tile)::kElemsPerFrag; _i++) { - // const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); - // const short _fn = _c.x; - // const short _fm = _c.y; - // AT_NAX(delta_tile, _i) = _fn >= 8 ? 0 : AT_NAX(delta_tile, _i); - // } - // OUTPUT(delta_tile) tmp_tile.clear(); QKt_tile.clear(); @@ -546,13 +587,12 @@ template out_tile.store(y + dv_idx, Hv * Dv); // OUTPUT(out_tile); - for (int kk = 0; kk < Dk; kk += 16) { - K1_tile.load(k_ + kk, Hk * Dk); - SCALE2_NAX(K1_tile, gamma) + for (int nn = 0; nn < Dk / 16; nn += 2) { + K_tile.load(k_ + nn * 16, Hk * Dk); + SCALE2_NAX(K_tile, gamma) - // KD = delta @ K.T so that I can sum to S directly - KD_tile.clear(); - mlx::steel::mmak< + // KD = delta.T @ K so that I can sum to S directly + mlx::steel::mman< float, InT, float, @@ -560,18 +600,23 @@ template false, mpp::tensor_ops::matmul2d_descriptor::mode::multiply>( KD_tile.frag_at(0, 0), + KD_tile.frag_at(0, 1), delta_tile.frag_at(0, 0), - Z.frag_at(0, 0), metal::bool_constant{}, - K1_tile.frag_at(0, 0), - Z.frag_at(0, 0), + K_tile.frag_at(0, 0), + K_tile.frag_at(0, 1), metal::bool_constant{}); FMA_NAX( - S_tile.frag_at(0, kk / 16), + S_tile.frag_at(0, nn), gamma[C - 1], - S_tile.frag_at(0, kk / 16), + S_tile.frag_at(0, nn), KD_tile.frag_at(0, 0)) + FMA_NAX( + S_tile.frag_at(0, nn + 1), + gamma[C - 1], + S_tile.frag_at(0, nn + 1), + KD_tile.frag_at(0, 1)) } // OUTPUT(out_tile); @@ -691,12 +736,15 @@ template // Tinv = (I + L_W)^{-1} = sum -L_W_k // T0 - T + T2 - T3 + T4 - T5 + T6 - T7 + T8 // I - T(I - T(I - T(I - T(I - T(I - T(I - T(I - T(I)))))))) - simdgroup_float8x8 Tinv; + simdgroup_float8x8 Tinv, P; // + // S = I + x = I - KKtK + AT(P, 0) = AT(KKtK_tile, 0); + AT(P, 1) = AT(KKtK_tile, 1); SUB(Tinv, I_tile, KKtK_tile) - for (int iter = 0; iter < C - 2; iter++) { - simdgroup_float8x8 TMP; - simdgroup_multiply(TMP, KKtK_tile, Tinv); // L @ Tinv - SUB(Tinv, I_tile, TMP) // Tinv = I - L @ Tinv + + for (int step = 1; (1 << step) < C; step++) { + simdgroup_multiply(P, P, P); + simdgroup_multiply_accumulate(Tinv, Tinv, P, Tinv); } // OUTPUT(Tinv) @@ -713,13 +761,18 @@ template } // OUTPUT(WS_tile) + AT(P, 0) = AT(KKtV_tile, 0); + AT(P, 1) = AT(KKtV_tile, 1); + SUB(Tinv, I_tile, KKtV_tile) + for (int step = 1; (1 << step) < C; step++) { + simdgroup_multiply(P, P, P); + simdgroup_multiply_accumulate(Tinv, Tinv, P, Tinv); + } + + // U = Tinv @ (beta * V) simdgroup_load(V_tile, v_ + dv_idx, Dv * Hv); SCALE(V_tile, beta_fm) - U_tile = V_tile; - for (int iter = 0; iter < C - 1; iter++) { - simdgroup_multiply(X_tile, KKtV_tile, U_tile); - SUB(U_tile, V_tile, X_tile) - } + simdgroup_multiply(U_tile, Tinv, V_tile); // delta = U - WS SUB(delta_tile, U_tile, WS_tile) From ed76b70a2e30b8f287e8a1ccba055934470791ea Mon Sep 17 00:00:00 2001 From: tpegolotti Date: Mon, 13 Jul 2026 02:14:48 -0700 Subject: [PATCH 23/56] Fuse status update with output computation --- .../metal/kernels/gated_delta_update_impl.h | 77 +++++++++++++------ 1 file changed, 53 insertions(+), 24 deletions(-) diff --git a/mlx/backend/metal/kernels/gated_delta_update_impl.h b/mlx/backend/metal/kernels/gated_delta_update_impl.h index f4d9b01af9..363dda962a 100644 --- a/mlx/backend/metal/kernels/gated_delta_update_impl.h +++ b/mlx/backend/metal/kernels/gated_delta_update_impl.h @@ -568,27 +568,7 @@ template K_tile.frag_at(0, 0), K_tile.frag_at(0, 1), metal::bool_constant{}); - } - // OUTPUT(tmp_tile) - - SCALE_TRI_NAX(QKt_tile, gamma) - // OUTPUT(QKt_tile) - - out_tile = tmp_tile; - mlx::steel::mmak( - out_tile.frag_at(0, 0), - QKt_tile.frag_at(0, 0), - Z.frag_at(0, 0), - metal::bool_constant{}, - delta_tile.frag_at(0, 0), - Z.frag_at(0, 0), - metal::bool_constant{}); - out_tile.store(y + dv_idx, Hv * Dv); - // OUTPUT(out_tile); - - for (int nn = 0; nn < Dk / 16; nn += 2) { - K_tile.load(k_ + nn * 16, Hk * Dk); SCALE2_NAX(K_tile, gamma) // KD = delta.T @ K so that I can sum to S directly @@ -608,16 +588,65 @@ template metal::bool_constant{}); FMA_NAX( - S_tile.frag_at(0, nn), + S_tile.frag_at(0, kk / 16), gamma[C - 1], - S_tile.frag_at(0, nn), + S_tile.frag_at(0, kk / 16), KD_tile.frag_at(0, 0)) FMA_NAX( - S_tile.frag_at(0, nn + 1), + S_tile.frag_at(0, kk / 16 + 1), gamma[C - 1], - S_tile.frag_at(0, nn + 1), + S_tile.frag_at(0, kk / 16 + 1), KD_tile.frag_at(0, 1)) } + // OUTPUT(tmp_tile) + + SCALE_TRI_NAX(QKt_tile, gamma) + // OUTPUT(QKt_tile) + + out_tile = tmp_tile; + mlx::steel::mmak( + out_tile.frag_at(0, 0), + QKt_tile.frag_at(0, 0), + Z.frag_at(0, 0), + metal::bool_constant{}, + delta_tile.frag_at(0, 0), + Z.frag_at(0, 0), + metal::bool_constant{}); + + out_tile.store(y + dv_idx, Hv * Dv); + // OUTPUT(out_tile); + + // for (int nn = 0; nn < Dk / 16; nn += 2) { + // K_tile.load(k_ + nn * 16, Hk * Dk); + // SCALE2_NAX(K_tile, gamma) + + // // KD = delta.T @ K so that I can sum to S directly + // mlx::steel::mman< + // float, + // InT, + // float, + // true, + // false, + // mpp::tensor_ops::matmul2d_descriptor::mode::multiply>( + // KD_tile.frag_at(0, 0), + // KD_tile.frag_at(0, 1), + // delta_tile.frag_at(0, 0), + // metal::bool_constant{}, + // K_tile.frag_at(0, 0), + // K_tile.frag_at(0, 1), + // metal::bool_constant{}); + + // FMA_NAX( + // S_tile.frag_at(0, nn), + // gamma[C - 1], + // S_tile.frag_at(0, nn), + // KD_tile.frag_at(0, 0)) + // FMA_NAX( + // S_tile.frag_at(0, nn + 1), + // gamma[C - 1], + // S_tile.frag_at(0, nn + 1), + // KD_tile.frag_at(0, 1)) + // } // OUTPUT(out_tile); // advance pointers From 6d6133e1c0f9b7f50560a0c0e926a77b3d2401cf Mon Sep 17 00:00:00 2001 From: tpegolotti Date: Mon, 13 Jul 2026 04:29:24 -0700 Subject: [PATCH 24/56] Half matmuls in invers by fusion --- .../metal/kernels/gated_delta_update_impl.h | 178 ++++++++++-------- 1 file changed, 104 insertions(+), 74 deletions(-) diff --git a/mlx/backend/metal/kernels/gated_delta_update_impl.h b/mlx/backend/metal/kernels/gated_delta_update_impl.h index 363dda962a..33ffe713da 100644 --- a/mlx/backend/metal/kernels/gated_delta_update_impl.h +++ b/mlx/backend/metal/kernels/gated_delta_update_impl.h @@ -259,6 +259,70 @@ METAL_FUNC static constexpr void mman( C1[i] = ct_c[BaseNAXFrag::kElemsPerFrag + i]; } } + +template < + typename CType, + typename AType, + typename BType, + bool transpose_a = false, + bool transpose_b = false, + mpp::tensor_ops::matmul2d_descriptor::mode Mode = + mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate> +METAL_FUNC static constexpr void mmam( + thread BaseNAXFrag::dtype_frag_t& C0, + thread BaseNAXFrag::dtype_frag_t& C1, + const thread BaseNAXFrag::dtype_frag_t& A0, + const thread BaseNAXFrag::dtype_frag_t& A1, + metal::bool_constant, + const thread BaseNAXFrag::dtype_frag_t& B, + metal::bool_constant) { + // N = 32: single A fragment, two N-fragments for B and C output. + // template parameters are M, N, K where + // Tensor dimensions where M x K tensor A,K x N tensor B, and M x N tensor C. + constexpr auto desc = mpp::tensor_ops::matmul2d_descriptor( + 32, 16, 16, transpose_a, transpose_b, true, Mode); + + mpp::tensor_ops::matmul2d gemm_op; + + // Create matmul operands in registers + auto ct_a = + gemm_op.template get_left_input_cooperative_tensor(); + auto ct_b = + gemm_op + .template get_right_input_cooperative_tensor(); + + // Create matmul output in register + auto ct_c = gemm_op.template get_destination_cooperative_tensor< + decltype(ct_a), + decltype(ct_b), + CType>(); + + STEEL_PRAGMA_UNROLL + for (short i = 0; i < BaseNAXFrag::kElemsPerFrag; i++) { + ct_a[i] = A0[i]; + ct_a[BaseNAXFrag::kElemsPerFrag + i] = A1[i]; + } + + STEEL_PRAGMA_UNROLL + for (short i = 0; i < BaseNAXFrag::kElemsPerFrag; i++) { + ct_b[i] = B[i]; + } + + STEEL_PRAGMA_UNROLL + for (short i = 0; i < BaseNAXFrag::kElemsPerFrag; i++) { + ct_c[i] = C0[i]; + ct_c[BaseNAXFrag::kElemsPerFrag + i] = 0; + } + + gemm_op.run(ct_a, ct_b, ct_c); + + // Copy out both N-fragments + STEEL_PRAGMA_UNROLL + for (short i = 0; i < BaseNAXFrag::kElemsPerFrag; i++) { + C0[i] = ct_c[i]; + C1[i] = ct_c[BaseNAXFrag::kElemsPerFrag + i]; + } +} } // namespace steel } // namespace mlx @@ -355,9 +419,7 @@ template mlx::steel::NAXTile TMP_tile; mlx::steel::NAXTile Z; - mlx::steel::NAXTile junk; - junk.clear(); Z.clear(); for (int t = 0; t < T; t += C) { @@ -419,32 +481,32 @@ template // S = I + x = I - KKtV SUB_NAX(Tinv_tile, I_tile, KKtK_tile) // Tinv = I - T (2 terms) - STEEL_PRAGMA_UNROLL - for (int step = 1; (1 << step) < C; step++) { // - // P2 = P · P - mlx::steel::mmak< - float, - float, - float, - false, - false, - mpp::tensor_ops::matmul2d_descriptor::mode::multiply>( - P_tile.frag_at(0, 0), - P_tile.frag_at(0, 0), - Z.frag_at(0, 0), - metal::bool_constant{}, - P_tile.frag_at(0, 0), - Z.frag_at(0, 0), - metal::bool_constant{}); + // P^2 + mlx::steel::mmak< + float, + float, + float, + false, + false, + mpp::tensor_ops::matmul2d_descriptor::mode::multiply>( + P_tile.frag_at(0, 0), + P_tile.frag_at(0, 0), + Z.frag_at(0, 0), + metal::bool_constant{}, + P_tile.frag_at(0, 0), + Z.frag_at(0, 0), + metal::bool_constant{}); + STEEL_PRAGMA_UNROLL + for (int step = 1; (1 << step) < C; step++) { // Tinv = S · P2 - mlx::steel::mmak( + mlx::steel::mmam( Tinv_tile.frag_at(0, 0), + P_tile.frag_at(0, 0), Tinv_tile.frag_at(0, 0), - Z.frag_at(0, 0), + P_tile.frag_at(0, 0), metal::bool_constant{}, P_tile.frag_at(0, 0), - Z.frag_at(0, 0), metal::bool_constant{}); } // OUTPUT(Tinv_tile) @@ -478,32 +540,33 @@ template // S = I + x = I - KKtV // Tinv = (1 - x)(1 + x^2)(1 + x^4)(1 + x^8) SUB_NAX(Tinv_tile, I_tile, KKtV_tile) + + // P^2 + mlx::steel::mmak< + float, + float, + float, + false, + false, + mpp::tensor_ops::matmul2d_descriptor::mode::multiply>( + P_tile.frag_at(0, 0), + P_tile.frag_at(0, 0), + Z.frag_at(0, 0), + metal::bool_constant{}, + P_tile.frag_at(0, 0), + Z.frag_at(0, 0), + metal::bool_constant{}); + STEEL_PRAGMA_UNROLL for (int step = 1; (1 << step) < C; step++) { - // P^2 - mlx::steel::mmak< - float, - float, - float, - false, - false, - mpp::tensor_ops::matmul2d_descriptor::mode::multiply>( - P_tile.frag_at(0, 0), - P_tile.frag_at(0, 0), - Z.frag_at(0, 0), - metal::bool_constant{}, - P_tile.frag_at(0, 0), - Z.frag_at(0, 0), - metal::bool_constant{}); - // Tinv = S · P2 - mlx::steel::mmak( + mlx::steel::mmam( Tinv_tile.frag_at(0, 0), + P_tile.frag_at(0, 0), Tinv_tile.frag_at(0, 0), - Z.frag_at(0, 0), + P_tile.frag_at(0, 0), metal::bool_constant{}, P_tile.frag_at(0, 0), - Z.frag_at(0, 0), metal::bool_constant{}); } V_tile.load(v_ + dv_idx, Dv * Hv); @@ -616,39 +679,6 @@ template out_tile.store(y + dv_idx, Hv * Dv); // OUTPUT(out_tile); - // for (int nn = 0; nn < Dk / 16; nn += 2) { - // K_tile.load(k_ + nn * 16, Hk * Dk); - // SCALE2_NAX(K_tile, gamma) - - // // KD = delta.T @ K so that I can sum to S directly - // mlx::steel::mman< - // float, - // InT, - // float, - // true, - // false, - // mpp::tensor_ops::matmul2d_descriptor::mode::multiply>( - // KD_tile.frag_at(0, 0), - // KD_tile.frag_at(0, 1), - // delta_tile.frag_at(0, 0), - // metal::bool_constant{}, - // K_tile.frag_at(0, 0), - // K_tile.frag_at(0, 1), - // metal::bool_constant{}); - - // FMA_NAX( - // S_tile.frag_at(0, nn), - // gamma[C - 1], - // S_tile.frag_at(0, nn), - // KD_tile.frag_at(0, 0)) - // FMA_NAX( - // S_tile.frag_at(0, nn + 1), - // gamma[C - 1], - // S_tile.frag_at(0, nn + 1), - // KD_tile.frag_at(0, 1)) - // } - // OUTPUT(out_tile); - // advance pointers q_ += C * Hk * Dk; k_ += C * Hk * Dk; From c898f44bb4c33f58de308a2d5b7d09f4847ce604 Mon Sep 17 00:00:00 2001 From: tpegolotti Date: Mon, 13 Jul 2026 08:48:36 -0700 Subject: [PATCH 25/56] Remove KP --- .../metal/kernels/gated_delta_update_impl.h | 44 +++++++------------ 1 file changed, 17 insertions(+), 27 deletions(-) diff --git a/mlx/backend/metal/kernels/gated_delta_update_impl.h b/mlx/backend/metal/kernels/gated_delta_update_impl.h index 33ffe713da..37e3e3e0cc 100644 --- a/mlx/backend/metal/kernels/gated_delta_update_impl.h +++ b/mlx/backend/metal/kernels/gated_delta_update_impl.h @@ -387,9 +387,7 @@ template S_tile.load(i_state, Dk); mlx::steel::NAXTile K_tile, Q_tile; - mlx::steel::NAXTile KP_tile; // panel mlx::steel::NAXTile W_tile; // panel - mlx::steel::NAXTile XP_tile; // panel mlx::steel::NAXTile X_tile; @@ -400,7 +398,6 @@ template mlx::steel::NAXTile tmp_tile; mlx::steel::NAXTile QKt_tile; mlx::steel::NAXTile out_tile; - mlx::steel::NAXTile K1_tile, Q1_tile; mlx::steel::NAXTile KD_tile; mlx::steel::NAXTile Tinv_tile; @@ -440,6 +437,7 @@ template // I can only do matmuls with a dimension 32. This one is _easy_ // i just do two tiles for the reduction KKt_tile.clear(); + // KP_tile.load(k_, Dk * Hk); for (int kk = 0; kk < Dk; kk += 32) { // two 16-tiles per iter K_tile.load(k_ + kk, Dk * Hk); mlx::steel::mmak( @@ -461,11 +459,6 @@ template SCALE_TRIEQ_NAX1(KKtK_tile, beta_fm) SCALE_TRIEQ_NAX(KKtV_tile, beta_fm, gamma) - // W = (I - diag(b)(KK.T))^-1 diag(b)K - // diag(b)K - KP_tile.load(k_, Dk * Hk); - SCALE_BETA_NAX(KP_tile, beta_fm) - // (I - diag(b)(KK.T))^-1 = sum T^k // Tinv = (I + L_W)^{-1} = sum -L_W_k // T0 - T + T2 - T3 + T4 - T5 + T6 - T7 + T8 @@ -473,12 +466,6 @@ template // Update 10.7: compute the above as (1 - T)(1 + T^2)(1 + T^4)(1 + T^8) // => 6 mults instead of 14 - Z.clear(); - STEEL_PRAGMA_UNROLL - for (short e = 0; e < decltype(P_tile)::kElemsPerFrag; e++) { - AT_NAX(P_tile, e) = AT_NAX(KKtK_tile, e); - } - // S = I + x = I - KKtV SUB_NAX(Tinv_tile, I_tile, KKtK_tile) // Tinv = I - T (2 terms) // P^2 @@ -490,10 +477,10 @@ template false, mpp::tensor_ops::matmul2d_descriptor::mode::multiply>( P_tile.frag_at(0, 0), - P_tile.frag_at(0, 0), + KKtK_tile.frag_at(0, 0), Z.frag_at(0, 0), metal::bool_constant{}, - P_tile.frag_at(0, 0), + KKtK_tile.frag_at(0, 0), Z.frag_at(0, 0), metal::bool_constant{}); @@ -511,9 +498,19 @@ template } // OUTPUT(Tinv_tile) + // W = (I - diag(b)(KK.T))^-1 diag(b)K + // diag(b)K + // SCALE_BETA_NAX(KP_tile, beta_fm) // W = Tinv @ (diag(beta) K) STEEL_PRAGMA_UNROLL for (short nn = 0; nn < Dk / 16; nn += 2) { + K_tile.load(k_ + nn * 16, Dk * Hk); + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(K_tile)::kElemsPerTile; _i++) { + const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; + AT_NAX(K_tile, _i) *= beta_fm[_w >> 2]; + } + mlx::steel::mman< float, float, @@ -525,18 +522,11 @@ template W_tile.frag_at(0, nn + 1), Tinv_tile.frag_at(0, 0), // A = Tinv (reused across N) metal::bool_constant{}, - KP_tile.frag_at(0, nn), // B = beta-scaled K - KP_tile.frag_at(0, nn + 1), + K_tile.frag_at(0, 0), // B = beta-scaled K + K_tile.frag_at(0, 1), metal::bool_constant{}); } SCALE_ROW_NAX(W_tile, gamma) - - Z.clear(); - STEEL_PRAGMA_UNROLL - for (short e = 0; e < decltype(P_tile)::kElemsPerFrag; e++) { - AT_NAX(P_tile, e) = AT_NAX(KKtV_tile, e); - } - // S = I + x = I - KKtV // Tinv = (1 - x)(1 + x^2)(1 + x^4)(1 + x^8) SUB_NAX(Tinv_tile, I_tile, KKtV_tile) @@ -550,10 +540,10 @@ template false, mpp::tensor_ops::matmul2d_descriptor::mode::multiply>( P_tile.frag_at(0, 0), - P_tile.frag_at(0, 0), + KKtV_tile.frag_at(0, 0), Z.frag_at(0, 0), metal::bool_constant{}, - P_tile.frag_at(0, 0), + KKtV_tile.frag_at(0, 0), Z.frag_at(0, 0), metal::bool_constant{}); From dd1124a61ab0d06fe1f0e7c29acbb28e17018508 Mon Sep 17 00:00:00 2001 From: tpegolotti Date: Tue, 14 Jul 2026 04:25:05 -0700 Subject: [PATCH 26/56] Add log decay for stability --- .../metal/kernels/gated_delta_update_impl.h | 98 +++++++++++-------- 1 file changed, 57 insertions(+), 41 deletions(-) diff --git a/mlx/backend/metal/kernels/gated_delta_update_impl.h b/mlx/backend/metal/kernels/gated_delta_update_impl.h index 37e3e3e0cc..13bff8c805 100644 --- a/mlx/backend/metal/kernels/gated_delta_update_impl.h +++ b/mlx/backend/metal/kernels/gated_delta_update_impl.h @@ -83,13 +83,14 @@ using namespace mpp::tensor_ops; } \ } -#define SCALE_ROW_NAX(TILE0, S) \ - { \ - STEEL_PRAGMA_UNROLL \ - for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ - const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; \ - AT_NAX(TILE0, _i) *= (S)[mlx::steel::BaseNAXFrag::get_coord(_w).y]; \ - } \ +#define SCALE_ROW_NAX(TILE0, S) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; \ + AT_NAX(TILE0, _i) *= \ + metal::exp((S)[mlx::steel::BaseNAXFrag::get_coord(_w).y]); \ + } \ } #define SCALE_BETA_NAX(TILE0, BETA2) \ @@ -108,7 +109,7 @@ using namespace mpp::tensor_ops; const short _fm = mlx::steel::BaseNAXFrag::get_coord( \ _i % mlx::steel::BaseNAXFrag::kElemsPerFrag) \ .y; \ - AT_NAX(TILE0, _i) *= (GAMMA)[(C) - 1] / (GAMMA)[_fm]; \ + AT_NAX(TILE0, _i) *= metal::exp((GAMMA)[(C) - 1] - (GAMMA)[_fm]); \ } \ } @@ -117,7 +118,8 @@ using namespace mpp::tensor_ops; STEEL_PRAGMA_UNROLL \ for (short _i = 0; _i < decltype(TILE0)::kElemsPerFrag; _i++) { \ const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); /* {fn, fm} */ \ - AT_NAX(TILE0, _i) *= (_c.x > _c.y) ? 0.f : (1.0f / (GAMMA)[_c.x]); \ + AT_NAX(TILE0, _i) *= \ + (_c.x > _c.y) ? 0.f : metal::exp((GAMMA)[_c.y] - (GAMMA)[_c.x]); \ } \ } @@ -137,7 +139,8 @@ using namespace mpp::tensor_ops; const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); /* {fn, fm} */ \ const short _fn = _c.x; \ const short _fm = _c.y; \ - const float _s = (BETA)[_i >> 2] * ((GAMMA)[_fm] / (GAMMA)[_fn]); \ + const float _s = \ + (BETA)[_i >> 2] * metal::exp((GAMMA)[_fm] - (GAMMA)[_fn]); \ AT_NAX(TILE0, _i) *= (_fn >= _fm) ? 0.f : _s; \ } \ } @@ -421,10 +424,11 @@ template for (int t = 0; t < T; t += C) { float g_val = (thread_index_in_simdgroup < C) - ? g_[thread_index_in_simdgroup * Hv + hv_idx] - : 1.0f; + ? metal::log( + metal::max(g_[thread_index_in_simdgroup * Hv + hv_idx], 1e-6f)) + : 0.0f; - float gamma_val = simd_prefix_inclusive_product(g_val); + float gamma_val = simd_prefix_inclusive_sum(g_val); if (thread_index_in_simdgroup < C) { gamma[thread_index_in_simdgroup] = gamma_val; @@ -449,7 +453,7 @@ template K_tile.frag_at(0, 1), metal::bool_constant{}); } - + // OUTPUT(KKt_tile) // KKt_tile.store(y,16); // return; @@ -459,6 +463,7 @@ template SCALE_TRIEQ_NAX1(KKtK_tile, beta_fm) SCALE_TRIEQ_NAX(KKtV_tile, beta_fm, gamma) + // OUTPUT(KKtV_tile) // (I - diag(b)(KK.T))^-1 = sum T^k // Tinv = (I + L_W)^{-1} = sum -L_W_k // T0 - T + T2 - T3 + T4 - T5 + T6 - T7 + T8 @@ -559,6 +564,7 @@ template P_tile.frag_at(0, 0), metal::bool_constant{}); } + // OUTPUT(Tinv_tile) V_tile.load(v_ + dv_idx, Dv * Hv); SCALE_BETA_NAX(V_tile, beta_fm) // U = Tinv @ diag(b)V @@ -577,6 +583,8 @@ template Z.frag_at(0, 0), metal::bool_constant{}); + // OUTPUT(U_tile) + // WS = W @ S.T, easy reduction with stride 2 over K WS_tile.clear(); STEEL_PRAGMA_UNROLL @@ -600,8 +608,17 @@ template Q_tile.load(q_ + kk, Hk * Dk); K_tile.load(k_ + kk, Hk * Dk); // this one is transposed - SCALE_ROW_NAX(Q_tile, gamma) + // Q @ K^T + mlx::steel::mmak( + QKt_tile.frag_at(0, 0), + Q_tile.frag_at(0, 0), + Q_tile.frag_at(0, 1), + metal::bool_constant{}, + K_tile.frag_at(0, 0), + K_tile.frag_at(0, 1), + metal::bool_constant{}); + SCALE_ROW_NAX(Q_tile, gamma) // Q_left @ S^T mlx::steel::mmak( tmp_tile.frag_at(0, 0), @@ -612,16 +629,6 @@ template S_tile.frag_at(0, kk / 16 + 1), metal::bool_constant{}); - // Q @ K^T - mlx::steel::mmak( - QKt_tile.frag_at(0, 0), - Q_tile.frag_at(0, 0), - Q_tile.frag_at(0, 1), - metal::bool_constant{}, - K_tile.frag_at(0, 0), - K_tile.frag_at(0, 1), - metal::bool_constant{}); - SCALE2_NAX(K_tile, gamma) // KD = delta.T @ K so that I can sum to S directly @@ -642,12 +649,12 @@ template FMA_NAX( S_tile.frag_at(0, kk / 16), - gamma[C - 1], + metal::exp(gamma[C - 1]), S_tile.frag_at(0, kk / 16), KD_tile.frag_at(0, 0)) FMA_NAX( S_tile.frag_at(0, kk / 16 + 1), - gamma[C - 1], + metal::exp(gamma[C - 1]), S_tile.frag_at(0, kk / 16 + 1), KD_tile.frag_at(0, 1)) } @@ -751,10 +758,11 @@ template for (int t = 0; t < T; t += C) { float g_val = (thread_index_in_simdgroup < C) - ? g_[thread_index_in_simdgroup * Hv + hv_idx] - : 1.0f; + ? metal::log( + metal::max(g_[thread_index_in_simdgroup * Hv + hv_idx], 1e-6f)) + : 0.0f; - float gamma_val = simd_prefix_inclusive_product(g_val); + float gamma_val = simd_prefix_inclusive_sum(g_val); if (thread_index_in_simdgroup < C) { gamma[thread_index_in_simdgroup] = gamma_val; @@ -771,7 +779,7 @@ template #define OUTPUT(T) \ simdgroup_store(T, y); \ return; - + // OUTPUT(KKt_tile) KKtK_tile = KKt_tile; KKtV_tile = KKt_tile; @@ -779,9 +787,10 @@ template SCALE_TRIEQ(KKtK_tile, beta_fm, beta_fm) SCALE_TRIEQ( KKtV_tile, - beta_fm * (gamma[fm] / gamma[fn]), - beta_fm * (gamma[fm] / gamma[fn + 1])) + beta_fm * (metal::exp(gamma[fm] - gamma[fn])), + beta_fm * (metal::exp(gamma[fm] - gamma[fn + 1]))) + // OUTPUT(KKtV_tile) // Tinv = (I + L_W)^{-1} = sum -L_W_k // T0 - T + T2 - T3 + T4 - T5 + T6 - T7 + T8 // I - T(I - T(I - T(I - T(I - T(I - T(I - T(I - T(I)))))))) @@ -805,7 +814,7 @@ template // W = Tinv @ (beta * K) simdgroup_multiply(W_tile, Tinv, K_tile); - SCALE(W_tile, gamma[fm]) + SCALE(W_tile, metal::exp(gamma[fm])) simdgroup_multiply_accumulate(WS_tile, W_tile, S_tile[kk / 8], WS_tile); } // OUTPUT(WS_tile) @@ -817,12 +826,13 @@ template simdgroup_multiply(P, P, P); simdgroup_multiply_accumulate(Tinv, Tinv, P, Tinv); } + // OUTPUT(Tinv) // U = Tinv @ (beta * V) simdgroup_load(V_tile, v_ + dv_idx, Dv * Hv); SCALE(V_tile, beta_fm) simdgroup_multiply(U_tile, Tinv, V_tile); - + // OUTPUT(U_tile) // delta = U - WS SUB(delta_tile, U_tile, WS_tile) // OUTPUT(delta_tile) @@ -833,16 +843,19 @@ template simdgroup_load(Q_tile, q_ + kk, Hk * Dk); simdgroup_load(K_tile, k_ + kk, Hk * Dk, ulong2(0, 0), true); - SCALE(Q_tile, gamma[fm]) - + // SCALE(Q_tile, gamma[fm]) + simdgroup_multiply_accumulate(QKt_tile, Q_tile, K_tile, QKt_tile); // Q_left @ S^T + SCALE(Q_tile, metal::exp(gamma[fm])) simdgroup_multiply_accumulate(tmp_tile, Q_tile, S_tile[kk / 8], tmp_tile); - simdgroup_multiply_accumulate(QKt_tile, Q_tile, K_tile, QKt_tile); } // OUTPUT(tmp_tile) // (Q @ K) * Gamma, in the paper they use M here but it's probably a typo. - SCALE_TRI(QKt_tile, (1.0f / gamma[fn]), (1.0f / gamma[fn + 1])) + SCALE_TRI( + QKt_tile, + metal::exp(gamma[fm] - gamma[fn]), + metal::exp(gamma[fm] - gamma[fn + 1])) // OUTPUT(QKt_tile) simdgroup_multiply_accumulate(out_tile, QKt_tile, delta_tile, tmp_tile); @@ -854,11 +867,14 @@ template for (int kk = 0; kk < Dk; kk += 8) { simdgroup_load(K_tile, k_ + kk, Hk * Dk, ulong2(0, 0), true); - SCALE2(K_tile, (gamma[C - 1] / gamma[fn]), (gamma[C - 1] / gamma[fn + 1])) + SCALE2( + K_tile, + metal::exp(gamma[C - 1] - gamma[fn]), + metal::exp(gamma[C - 1] - gamma[fn + 1])) simdgroup_multiply(KD_tile, K_tile, delta_tile); - FMA(S_tile[kk / 8], gamma[C - 1], S_tile[kk / 8], KD_tile) + FMA(S_tile[kk / 8], metal::exp(gamma[C - 1]), S_tile[kk / 8], KD_tile) } // advance pointers From d558dbb77f0466375a1b6b7fd330eb82abaf2648 Mon Sep 17 00:00:00 2001 From: tpegolotti Date: Tue, 14 Jul 2026 04:39:32 -0700 Subject: [PATCH 27/56] Make ensure row contiguous in eval_gpu --- mlx/backend/metal/gated_delta_update.cpp | 63 +++++++++++++++++++++--- mlx/fast.cpp | 10 ++-- 2 files changed, 62 insertions(+), 11 deletions(-) diff --git a/mlx/backend/metal/gated_delta_update.cpp b/mlx/backend/metal/gated_delta_update.cpp index a0fd48092f..5a4c10387b 100644 --- a/mlx/backend/metal/gated_delta_update.cpp +++ b/mlx/backend/metal/gated_delta_update.cpp @@ -17,18 +17,61 @@ bool GatedDeltaUpdate::use_fallback(Stream s) { return false; } +inline array +ensure_row_contiguous(const array& x, metal::Device& d, const Stream& s) { + if (!x.flags().row_contiguous) { + array x_copy = contiguous_copy_gpu(x, s); + metal::get_command_encoder(s).add_temporary(x_copy); + return x_copy; + } else { + return x; + } +} + +#define PRINT_STRIDES(arr) \ + printf( \ + "%s strides: %lld %lld %lld %lld\n", \ + #arr, \ + arr.strides()[0], \ + arr.strides()[1], \ + arr.strides()[2], \ + arr.strides()[3]) + +#define PRINT_SHAPES(arr) \ + printf( \ + "%s shapes: %lld %lld %lld %lld\n", \ + #arr, \ + arr.shape()[0], \ + arr.shape()[1], \ + arr.shape()[2], \ + arr.shape()[3]) + +#define PRINT_ARR(arr) \ + if (arr.flags().row_contiguous) \ + printf("%s is row contiguous\n", #arr); \ + PRINT_SHAPES(arr); \ + PRINT_STRIDES(arr); \ + printf("\n"); + void GatedDeltaUpdate::eval_gpu( const std::vector& inputs, std::vector& outputs) { auto& s = stream(); auto& d = metal::device(s.device); - auto& q = inputs[0]; - auto& k = inputs[1]; - auto& v = inputs[2]; - auto& g = inputs[3]; - auto& beta = inputs[4]; - auto& h0 = inputs[5]; + // auto& q = inputs[0]; + // auto& k = inputs[1]; + // auto& v = inputs[2]; + // auto& g = inputs[3]; + // auto& beta = inputs[4]; + // auto& h0 = inputs[5]; + + auto q = ensure_row_contiguous(inputs[0], d, s); + auto k = ensure_row_contiguous(inputs[1], d, s); + auto v = ensure_row_contiguous(inputs[2], d, s); + auto g = ensure_row_contiguous(inputs[3], d, s); + auto beta = ensure_row_contiguous(inputs[4], d, s); + auto h0 = ensure_row_contiguous(inputs[5], d, s); auto& out = outputs[0]; auto& hf = outputs[1]; @@ -40,6 +83,12 @@ void GatedDeltaUpdate::eval_gpu( int Hv = v.shape(2); int Dv = v.shape(3); + // PRINT_ARR(q); + // PRINT_ARR(k); + // PRINT_ARR(v); + // PRINT_ARR(g); + // PRINT_ARR(beta); + // PRINT_ARR(h0); // printf("%d %d %d %d %d %d\n",B,T,Hk,Hv,Dk,Dv); int C = chunk_size; @@ -84,6 +133,7 @@ void GatedDeltaUpdate::eval_gpu( compute_encoder.set_bytes(T, 8); auto grid = MTL::Size(32, Dv / 16, B * Hv); + // auto grid = MTL::Size(32, 1, 1); auto threads = MTL::Size(32, 1, 1); compute_encoder.dispatch_threads(grid, threads); break; @@ -115,6 +165,7 @@ void GatedDeltaUpdate::eval_gpu( compute_encoder.set_bytes(T, 8); auto grid = MTL::Size(32, Dv / 8, B * Hv); + // auto grid = MTL::Size(32, 1, 1); auto threads = MTL::Size(32, 1, 1); compute_encoder.dispatch_threads(grid, threads); break; diff --git a/mlx/fast.cpp b/mlx/fast.cpp index 885319f08c..161004744f 100644 --- a/mlx/fast.cpp +++ b/mlx/fast.cpp @@ -947,11 +947,11 @@ std::vector gated_delta_update( // wrong, not as fast. This kind of operations need to be done in // the eval gpu inside a function that checks for shapes and strides // TODO: mode this - q = contiguous(q, false, s); - k = contiguous(k, false, s); - v = contiguous(v, false, s); - g = contiguous(g, false, s); - beta = contiguous(beta, false, s); + // q = contiguous(q, false, s); + // k = contiguous(k, false, s); + // v = contiguous(v, false, s); + // g = contiguous(g, false, s); + // beta = contiguous(beta, false, s); int B = q.shape(0); int T = q.shape(1); From b71225269e289af6631e3e7bb806fe9bfd8331f0 Mon Sep 17 00:00:00 2001 From: tpegolotti Date: Tue, 14 Jul 2026 06:16:07 -0700 Subject: [PATCH 28/56] Add fallback --- mlx/backend/metal/gated_delta_update.cpp | 7 +- mlx/fast.cpp | 81 ++++++++++++++++++++---- 2 files changed, 74 insertions(+), 14 deletions(-) diff --git a/mlx/backend/metal/gated_delta_update.cpp b/mlx/backend/metal/gated_delta_update.cpp index 5a4c10387b..0f1fefccfc 100644 --- a/mlx/backend/metal/gated_delta_update.cpp +++ b/mlx/backend/metal/gated_delta_update.cpp @@ -13,7 +13,10 @@ namespace mlx::core::fast { bool GatedDeltaUpdate::use_fallback(Stream s) { - // TODO: finish implementation + // TODO: finish implementation. What else is needed? + if (s.device == Device::cpu) { + return true; + } return false; } @@ -200,7 +203,7 @@ void GatedDeltaUpdate::eval_gpu( } default: { throw std::runtime_error( - "NYI: Only sequential and chunk size 8 are supported"); + "NYI: Only sequential and chunk size 8,16 are supported"); } } } diff --git a/mlx/fast.cpp b/mlx/fast.cpp index 161004744f..d63f424590 100644 --- a/mlx/fast.cpp +++ b/mlx/fast.cpp @@ -930,8 +930,10 @@ std::vector gated_delta_update( const array& beta_, const std::optional& initial_state, /* = std::nullopt */ const int C, - StreamOrDevice s /* = {} */) { + StreamOrDevice s_ /* = {} */) { // determine output dtype + auto s = to_stream(s_); + auto promoted = promote_types(queries.dtype(), keys.dtype()); auto out_dtype = issubdtype(promoted, float32) ? promoted @@ -985,25 +987,80 @@ std::vector gated_delta_update( std::vector outputs; for (int t = 0; t < T; t++) { - // TODO. Implement + auto get_t = [&](const array& a, int t) { + Shape start(a.ndim(), 0), stop = a.shape(); + start[1] = t; + stop[1] = t + 1; // [:, t:t+1, ...] + return squeeze(slice(a, start, stop, s), 1, s); // drop the time axis + }; + // q_t = q[:, t] # [B, H, Dk] + // k_t = k[:, t] # [B, H, Dk] + // v_t = v[:, t] # [B, H, Dv] + // g_t = g[:, t] # [B, H] or [B, H, Dk] + // beta_t = beta[:, t] # [B, H] + auto q_t = get_t(q, t); + auto k_t = get_t(k, t); + auto v_t = get_t(v, t); + auto g_t = get_t(g, t); + auto beta_t = get_t(beta, t); + + // if g_t.ndim == 2: + // decay = g_t[..., None, None] # [B, H, 1, 1] + // else: + // decay = g_t[..., None, :] # [B, H, 1, Dk] + auto decay = (g_t.ndim() == 2) + ? expand_dims(g_t, {-1, -2}, s) // [B,H,1,1] + : expand_dims(g_t, -2, s); // [B,H,1,Dk] + + // state = state * decay + state = multiply(state, decay, s); + + // kv_mem = (state * k_t[..., None, :]).sum(axis=-1) + auto kv = + sum(multiply(state, expand_dims(k_t, -2, s), s), + -1, + false, + s); // [B,H,Dv] + + // delta = (v_t - kv_mem) * beta_t[..., None] # [B, H, Dv] + auto delta = subtract(v_t, kv, s); + delta = multiply(delta, expand_dims(beta_t, -1, s), s); + + // state = state + delta[..., None] * k_t[..., None, :] # [B, H, Dv, Dk] + state = + add(state, + multiply(expand_dims(delta, -1, s), expand_dims(k_t, -2, s), s), + s); + + // o_t = (state * q_t[..., None, :]).sum(axis=-1) + auto o_t = sum(multiply(state, expand_dims(q_t, -2, s), s), -1, false, s); + outputs.push_back(o_t); } - throw std::runtime_error("NYI: GatedDeltaUpdate CPU fallback"); - auto out = concatenate(outputs, 1, s); + // mx.stack(outputs, axis=1) + auto out = stack(outputs, 1, s); return std::vector{out, state}; }; - auto result = array::make_arrays( - /* output shapes */ {{B, T_padded, Hv, Dv}, {B, Hv, Dv, Dk}}, - /* dtypes */ {out_dtype, out_dtype}, - /* primitive */ - std::make_shared(to_stream(s), fallback, C), - /* inputs */ {q, k, v, g, beta, h0}); + if (!GatedDeltaUpdate::use_fallback(s)) { + auto result = array::make_arrays( + /* output shapes */ {{B, T_padded, Hv, Dv}, {B, Hv, Dv, Dk}}, + /* dtypes */ {out_dtype, out_dtype}, + /* primitive */ + std::make_shared(s, fallback, C), + /* inputs */ {q, k, v, g, beta, h0}); + + // slice output back to original T + if (pad_size > 0) { + result[0] = slice(result[0], {0, 0, 0, 0}, {B, T, Hv, Dv}, s); + } + + return result; + } - // slice output back to original T + auto result = fallback({q, k, v, g, beta, h0}); if (pad_size > 0) { result[0] = slice(result[0], {0, 0, 0, 0}, {B, T, Hv, Dv}, s); } - return result; } From fb25ec5d946a86bd05ba7444f65149725715c879 Mon Sep 17 00:00:00 2001 From: tpegolotti Date: Wed, 29 Jul 2026 05:31:10 -0700 Subject: [PATCH 29/56] Reverted invert --- .../metal/kernels/gated_delta_update_impl.h | 370 ++++++------------ 1 file changed, 130 insertions(+), 240 deletions(-) diff --git a/mlx/backend/metal/kernels/gated_delta_update_impl.h b/mlx/backend/metal/kernels/gated_delta_update_impl.h index 13bff8c805..09c27184ce 100644 --- a/mlx/backend/metal/kernels/gated_delta_update_impl.h +++ b/mlx/backend/metal/kernels/gated_delta_update_impl.h @@ -62,7 +62,7 @@ using namespace mpp::tensor_ops; #define SUB_NAX(TILE0, TILE1, TILE2) \ { \ STEEL_PRAGMA_UNROLL \ - for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerFrag; _i++) { \ AT_NAX(TILE0, _i) = AT_NAX(TILE1, _i) - AT_NAX(TILE2, _i); \ } \ } @@ -102,15 +102,14 @@ using namespace mpp::tensor_ops; } \ } -#define SCALE2_NAX(TILE0, GAMMA) \ - { \ - STEEL_PRAGMA_UNROLL \ - for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ - const short _fm = mlx::steel::BaseNAXFrag::get_coord( \ - _i % mlx::steel::BaseNAXFrag::kElemsPerFrag) \ - .y; \ - AT_NAX(TILE0, _i) *= metal::exp((GAMMA)[(C) - 1] - (GAMMA)[_fm]); \ - } \ +#define SCALE2_NAX(TILE0, GAMMA) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; \ + const short _fm = mlx::steel::BaseNAXFrag::get_coord(_w).y; \ + AT_NAX(TILE0, _i) *= metal::exp((GAMMA)[(C) - 1] - (GAMMA)[_fm]); \ + } \ } #define SCALE_TRI_NAX(TILE0, GAMMA) \ @@ -155,7 +154,7 @@ template < bool transpose_b = false, mpp::tensor_ops::matmul2d_descriptor::mode Mode = mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate> -METAL_FUNC static constexpr void mmak( +METAL_FUNC static constexpr void mma( thread BaseNAXFrag::dtype_frag_t& C, const thread BaseNAXFrag::dtype_frag_t& A0, const thread BaseNAXFrag::dtype_frag_t& A1, @@ -163,20 +162,17 @@ METAL_FUNC static constexpr void mmak( const thread BaseNAXFrag::dtype_frag_t& B0, const thread BaseNAXFrag::dtype_frag_t& B1, metal::bool_constant) { - // K = 32: two K-fragments per operand, single 16x16 output. + // M=16, N=16, K=32: A and B each two K-fragments, single 16x16 C. constexpr auto desc = mpp::tensor_ops::matmul2d_descriptor( 16, 16, 32, transpose_a, transpose_b, true, Mode); mpp::tensor_ops::matmul2d gemm_op; - // Create matmul operands in registers auto ct_a = gemm_op.template get_left_input_cooperative_tensor(); auto ct_b = gemm_op .template get_right_input_cooperative_tensor(); - - // Create matmul output in register auto ct_c = gemm_op.template get_destination_cooperative_tensor< decltype(ct_a), decltype(ct_b), @@ -203,34 +199,26 @@ template < typename CType, typename AType, typename BType, - bool transpose_a = false, - bool transpose_b = false, + bool transpose_a, + bool transpose_b, mpp::tensor_ops::matmul2d_descriptor::mode Mode = mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate> -METAL_FUNC static constexpr void mman( - thread BaseNAXFrag::dtype_frag_t& C0, - thread BaseNAXFrag::dtype_frag_t& C1, +METAL_FUNC static constexpr void mma( + thread BaseNAXFrag::dtype_frag_t& C, const thread BaseNAXFrag::dtype_frag_t& A, metal::bool_constant, - const thread BaseNAXFrag::dtype_frag_t& B0, - const thread BaseNAXFrag::dtype_frag_t& B1, + const thread BaseNAXFrag::dtype_frag_t& B, metal::bool_constant) { - // N = 32: single A fragment, two N-fragments for B and C output. - // template parameters are M, N, K where - // Tensor dimensions where M x K tensor A,K x N tensor B, and M x N tensor C. constexpr auto desc = mpp::tensor_ops::matmul2d_descriptor( 16, 32, 16, transpose_a, transpose_b, true, Mode); mpp::tensor_ops::matmul2d gemm_op; - // Create matmul operands in registers auto ct_a = gemm_op.template get_left_input_cooperative_tensor(); auto ct_b = gemm_op .template get_right_input_cooperative_tensor(); - - // Create matmul output in register auto ct_c = gemm_op.template get_destination_cooperative_tensor< decltype(ct_a), decltype(ct_b), @@ -239,27 +227,17 @@ METAL_FUNC static constexpr void mman( STEEL_PRAGMA_UNROLL for (short i = 0; i < BaseNAXFrag::kElemsPerFrag; i++) { ct_a[i] = A[i]; - } - - STEEL_PRAGMA_UNROLL - for (short i = 0; i < BaseNAXFrag::kElemsPerFrag; i++) { - ct_b[i] = B0[i]; - ct_b[BaseNAXFrag::kElemsPerFrag + i] = B1[i]; - } - - STEEL_PRAGMA_UNROLL - for (short i = 0; i < BaseNAXFrag::kElemsPerFrag; i++) { - ct_c[i] = C0[i]; - ct_c[BaseNAXFrag::kElemsPerFrag + i] = C1[i]; + ct_b[i] = B[i]; + ct_b[BaseNAXFrag::kElemsPerFrag + i] = 0.0; + ct_c[i] = C[i]; + ct_c[BaseNAXFrag::kElemsPerFrag + i] = 0.0; } gemm_op.run(ct_a, ct_b, ct_c); - // Copy out both N-fragments STEEL_PRAGMA_UNROLL for (short i = 0; i < BaseNAXFrag::kElemsPerFrag; i++) { - C0[i] = ct_c[i]; - C1[i] = ct_c[BaseNAXFrag::kElemsPerFrag + i]; + C[i] = ct_c[i]; } } @@ -271,20 +249,19 @@ template < bool transpose_b = false, mpp::tensor_ops::matmul2d_descriptor::mode Mode = mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate> -METAL_FUNC static constexpr void mmam( - thread BaseNAXFrag::dtype_frag_t& C0, - thread BaseNAXFrag::dtype_frag_t& C1, - const thread BaseNAXFrag::dtype_frag_t& A0, - const thread BaseNAXFrag::dtype_frag_t& A1, +METAL_FUNC static constexpr void mman( + thread BaseNAXFrag::dtype_frag_t& Cn0, + thread BaseNAXFrag::dtype_frag_t& Cn1, + const thread BaseNAXFrag::dtype_frag_t& A, metal::bool_constant, - const thread BaseNAXFrag::dtype_frag_t& B, + const thread BaseNAXFrag::dtype_frag_t& Bn0, + const thread BaseNAXFrag::dtype_frag_t& Bn1, metal::bool_constant) { - // N = 32: single A fragment, two N-fragments for B and C output. - // template parameters are M, N, K where - // Tensor dimensions where M x K tensor A,K x N tensor B, and M x N tensor C. + // M=16, N=32, K=16: single A (K=16), B and C two N-fragments each. constexpr auto desc = mpp::tensor_ops::matmul2d_descriptor( - 32, 16, 16, transpose_a, transpose_b, true, Mode); + 16, 32, 16, transpose_a, transpose_b, true, Mode); + // Create matmul op mpp::tensor_ops::matmul2d gemm_op; // Create matmul operands in registers @@ -300,32 +277,37 @@ METAL_FUNC static constexpr void mmam( decltype(ct_b), CType>(); + // Load A in to left operand registers STEEL_PRAGMA_UNROLL for (short i = 0; i < BaseNAXFrag::kElemsPerFrag; i++) { - ct_a[i] = A0[i]; - ct_a[BaseNAXFrag::kElemsPerFrag + i] = A1[i]; + ct_a[i] = A[i]; } + // Load B into right operand registers STEEL_PRAGMA_UNROLL for (short i = 0; i < BaseNAXFrag::kElemsPerFrag; i++) { - ct_b[i] = B[i]; + ct_b[i] = Bn0[i]; + ct_b[BaseNAXFrag::kElemsPerFrag + i] = Bn1[i]; } + // Load C into output registers (op handles accumulation) STEEL_PRAGMA_UNROLL for (short i = 0; i < BaseNAXFrag::kElemsPerFrag; i++) { - ct_c[i] = C0[i]; - ct_c[BaseNAXFrag::kElemsPerFrag + i] = 0; + ct_c[i] = Cn0[i]; + ct_c[BaseNAXFrag::kElemsPerFrag + i] = Cn1[i]; } + // Do matmul gemm_op.run(ct_a, ct_b, ct_c); - // Copy out both N-fragments + // Copy out results STEEL_PRAGMA_UNROLL for (short i = 0; i < BaseNAXFrag::kElemsPerFrag; i++) { - C0[i] = ct_c[i]; - C1[i] = ct_c[BaseNAXFrag::kElemsPerFrag + i]; + Cn0[i] = ct_c[i]; + Cn1[i] = ct_c[BaseNAXFrag::kElemsPerFrag + i]; } } + } // namespace steel } // namespace mlx @@ -342,29 +324,16 @@ template constant int& T [[buffer(8)]], uint3 thread_position_in_grid [[thread_position_in_grid]], uint thread_index_in_simdgroup [[thread_index_in_simdgroup]]) { - auto n = thread_position_in_grid.z; + auto n = thread_position_in_grid.z; // 7 auto b_idx = n / Hv; auto hv_idx = n % Hv; auto hk_idx = hv_idx / (Hv / Hk); -#define OUTPUT_16(T) \ - T.store(y, 16); \ - return; - -#define OUTPUT_S(T, S) \ - T.store(y, S); \ - return; - -// pick the 3rd argument as the macro name -#define OUTPUT_GET(_1, _2, NAME, ...) NAME -#define OUTPUT(...) OUTPUT_GET(__VA_ARGS__, OUTPUT_S, OUTPUT_16)(__VA_ARGS__) + auto dv_idx = thread_position_in_grid.y * 16; const ushort simd_lane_id = __metal_get_thread_index_in_simdgroup(ushort()); const short qid = simd_lane_id >> 2; const short fm = ((qid & 4) | ((simd_lane_id >> 1) & 3)); - const short fn = ((qid & 2) | (simd_lane_id & 1)) * 4; - - auto dv_idx = thread_position_in_grid.y * 16; // set up pointers // g: [B, T, Hv] @@ -386,29 +355,24 @@ template threadgroup float gamma[C]; float beta_fm[2]; - mlx::steel::NAXTile S_tile; + mlx::steel::NAXTile S_tile; S_tile.load(i_state, Dk); - mlx::steel::NAXTile K_tile, Q_tile; - mlx::steel::NAXTile W_tile; // panel - - mlx::steel::NAXTile X_tile; + mlx::steel::NAXTile K_tile, Q_tile; + mlx::steel::NAXTile W_tile; // panel - mlx::steel::NAXTile V_tile; - mlx::steel::NAXTile U_tile; - mlx::steel::NAXTile WS_tile; - mlx::steel::NAXTile delta_tile; - mlx::steel::NAXTile tmp_tile; - mlx::steel::NAXTile QKt_tile; - mlx::steel::NAXTile out_tile; - mlx::steel::NAXTile KD_tile; - mlx::steel::NAXTile Tinv_tile; + mlx::steel::NAXTile V_tile; + mlx::steel::NAXTile U_tile; + mlx::steel::NAXTile WS_tile; + mlx::steel::NAXTile delta_tile; + mlx::steel::NAXTile tmp_tile; + mlx::steel::NAXTile QKt_tile; + mlx::steel::NAXTile out_tile; + mlx::steel::NAXTile Tinv_tile; - mlx::steel::NAXTile KKtK_tile, KKtV_tile, KKt_tile; + mlx::steel::NAXTile KKtK_tile, KKtV_tile, KKt_tile; - mlx::steel::NAXTile P_tile; - - mlx::steel::NAXTile I_tile; // Declare and init eye + mlx::steel::NAXTile I_tile; STEEL_PRAGMA_UNROLL for (short _i = 0; _i < decltype(I_tile)::kElemsPerFrag; _i++) { const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); /* {fn, fm} */ @@ -416,16 +380,13 @@ template const short _fm = _c.y; AT_NAX(I_tile, _i) = (_fn == _fm) ? 1.0f : 0.0f; } - mlx::steel::NAXTile TMP_tile; - - mlx::steel::NAXTile Z; - - Z.clear(); + mlx::steel::NAXTile TMP_tile; for (int t = 0; t < T; t += C) { float g_val = (thread_index_in_simdgroup < C) ? metal::log( - metal::max(g_[thread_index_in_simdgroup * Hv + hv_idx], 1e-6f)) + metal::clamp( + g_[thread_index_in_simdgroup * Hv + hv_idx], 1e-6f, 1.0f)) : 0.0f; float gamma_val = simd_prefix_inclusive_sum(g_val); @@ -438,13 +399,12 @@ template beta_fm[1] = beta_[(fm + mlx::steel::BaseNAXFrag::kElemRowsJump) * Hv + hv_idx]; - // I can only do matmuls with a dimension 32. This one is _easy_ - // i just do two tiles for the reduction KKt_tile.clear(); // KP_tile.load(k_, Dk * Hk); for (int kk = 0; kk < Dk; kk += 32) { // two 16-tiles per iter K_tile.load(k_ + kk, Dk * Hk); - mlx::steel::mmak( + + mlx::steel::mma( KKt_tile.frag_at(0, 0), K_tile.frag_at(0, 0), K_tile.frag_at(0, 1), @@ -453,9 +413,6 @@ template K_tile.frag_at(0, 1), metal::bool_constant{}); } - // OUTPUT(KKt_tile) - // KKt_tile.store(y,16); - // return; KKtK_tile = KKt_tile; KKtV_tile = KKt_tile; @@ -463,58 +420,30 @@ template SCALE_TRIEQ_NAX1(KKtK_tile, beta_fm) SCALE_TRIEQ_NAX(KKtV_tile, beta_fm, gamma) - // OUTPUT(KKtV_tile) - // (I - diag(b)(KK.T))^-1 = sum T^k - // Tinv = (I + L_W)^{-1} = sum -L_W_k - // T0 - T + T2 - T3 + T4 - T5 + T6 - T7 + T8 - // I - T(I - T(I - T(I - T(I - T(I - T(I - T(I - T(I)))))))) - // Update 10.7: compute the above as (1 - T)(1 + T^2)(1 + T^4)(1 + T^8) - // => 6 mults instead of 14 - - // S = I + x = I - KKtV - SUB_NAX(Tinv_tile, I_tile, KKtK_tile) // Tinv = I - T (2 terms) - // P^2 - mlx::steel::mmak< - float, - float, - float, - false, - false, - mpp::tensor_ops::matmul2d_descriptor::mode::multiply>( - P_tile.frag_at(0, 0), - KKtK_tile.frag_at(0, 0), - Z.frag_at(0, 0), - metal::bool_constant{}, - KKtK_tile.frag_at(0, 0), - Z.frag_at(0, 0), - metal::bool_constant{}); - + Tinv_tile = I_tile; STEEL_PRAGMA_UNROLL - for (int step = 1; (1 << step) < C; step++) { - // Tinv = S · P2 - mlx::steel::mmam( - Tinv_tile.frag_at(0, 0), - P_tile.frag_at(0, 0), - Tinv_tile.frag_at(0, 0), - P_tile.frag_at(0, 0), + for (int step = 0; step < C - 1; step++) { + // TMP = T_K · Tinv + mlx::steel::mma< + float, + float, + float, + false, + false, + mpp::tensor_ops::matmul2d_descriptor::mode::multiply>( + TMP_tile.frag_at(0, 0), + KKtK_tile.frag_at(0, 0), metal::bool_constant{}, - P_tile.frag_at(0, 0), + Tinv_tile.frag_at(0, 0), metal::bool_constant{}); + // Tinv = I - TMP + SUB_NAX(Tinv_tile, I_tile, TMP_tile) } - // OUTPUT(Tinv_tile) - // W = (I - diag(b)(KK.T))^-1 diag(b)K - // diag(b)K - // SCALE_BETA_NAX(KP_tile, beta_fm) - // W = Tinv @ (diag(beta) K) STEEL_PRAGMA_UNROLL for (short nn = 0; nn < Dk / 16; nn += 2) { K_tile.load(k_ + nn * 16, Dk * Hk); - STEEL_PRAGMA_UNROLL - for (short _i = 0; _i < decltype(K_tile)::kElemsPerTile; _i++) { - const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; - AT_NAX(K_tile, _i) *= beta_fm[_w >> 2]; - } + SCALE_BETA_NAX(K_tile, beta_fm) mlx::steel::mman< float, @@ -532,43 +461,29 @@ template metal::bool_constant{}); } SCALE_ROW_NAX(W_tile, gamma) - // S = I + x = I - KKtV - // Tinv = (1 - x)(1 + x^2)(1 + x^4)(1 + x^8) - SUB_NAX(Tinv_tile, I_tile, KKtV_tile) - - // P^2 - mlx::steel::mmak< - float, - float, - float, - false, - false, - mpp::tensor_ops::matmul2d_descriptor::mode::multiply>( - P_tile.frag_at(0, 0), - KKtV_tile.frag_at(0, 0), - Z.frag_at(0, 0), - metal::bool_constant{}, - KKtV_tile.frag_at(0, 0), - Z.frag_at(0, 0), - metal::bool_constant{}); + Tinv_tile = I_tile; STEEL_PRAGMA_UNROLL - for (int step = 1; (1 << step) < C; step++) { - // Tinv = S · P2 - mlx::steel::mmam( - Tinv_tile.frag_at(0, 0), - P_tile.frag_at(0, 0), - Tinv_tile.frag_at(0, 0), - P_tile.frag_at(0, 0), + for (int step = 0; step < C - 1; step++) { + mlx::steel::mma< + float, + float, + float, + false, + false, + mpp::tensor_ops::matmul2d_descriptor::mode::multiply>( + TMP_tile.frag_at(0, 0), + KKtV_tile.frag_at(0, 0), metal::bool_constant{}, - P_tile.frag_at(0, 0), + Tinv_tile.frag_at(0, 0), metal::bool_constant{}); + SUB_NAX(Tinv_tile, I_tile, TMP_tile) } - // OUTPUT(Tinv_tile) + V_tile.load(v_ + dv_idx, Dv * Hv); SCALE_BETA_NAX(V_tile, beta_fm) // U = Tinv @ diag(b)V - mlx::steel::mmak< + mlx::steel::mma< float, float, float, @@ -577,19 +492,14 @@ template mpp::tensor_ops::matmul2d_descriptor::mode::multiply>( U_tile.frag_at(0, 0), Tinv_tile.frag_at(0, 0), - Z.frag_at(0, 0), metal::bool_constant{}, V_tile.frag_at(0, 0), - Z.frag_at(0, 0), metal::bool_constant{}); - // OUTPUT(U_tile) - - // WS = W @ S.T, easy reduction with stride 2 over K WS_tile.clear(); STEEL_PRAGMA_UNROLL for (short kk = 0; kk < Dk / 16; kk += 2) { - mlx::steel::mmak( + mlx::steel::mma( WS_tile.frag_at(0, 0), W_tile.frag_at(0, kk), W_tile.frag_at(0, kk + 1), @@ -598,7 +508,6 @@ template S_tile.frag_at(0, kk + 1), metal::bool_constant{}); } - // OUTPUT(WS_tile) SUB_NAX(delta_tile, U_tile, WS_tile) @@ -609,7 +518,7 @@ template K_tile.load(k_ + kk, Hk * Dk); // this one is transposed // Q @ K^T - mlx::steel::mmak( + mlx::steel::mma( QKt_tile.frag_at(0, 0), Q_tile.frag_at(0, 0), Q_tile.frag_at(0, 1), @@ -620,7 +529,7 @@ template SCALE_ROW_NAX(Q_tile, gamma) // Q_left @ S^T - mlx::steel::mmak( + mlx::steel::mma( tmp_tile.frag_at(0, 0), Q_tile.frag_at(0, 0), Q_tile.frag_at(0, 1), @@ -628,53 +537,40 @@ template S_tile.frag_at(0, kk / 16), S_tile.frag_at(0, kk / 16 + 1), metal::bool_constant{}); + } - SCALE2_NAX(K_tile, gamma) + SCALE_TRI_NAX(QKt_tile, gamma) + + out_tile = tmp_tile; + mlx::steel::mma( + out_tile.frag_at(0, 0), + QKt_tile.frag_at(0, 0), + metal::bool_constant{}, + delta_tile.frag_at(0, 0), + metal::bool_constant{}); + + out_tile.store(y + dv_idx, Hv * Dv); + + SCALE_NAX(S_tile, metal::exp(gamma[C - 1])) - // KD = delta.T @ K so that I can sum to S directly + for (int kk = 0; kk < Dk; kk += 32) { + K_tile.load(k_ + kk, Hk * Dk); + SCALE2_NAX(K_tile, gamma) mlx::steel::mman< float, - InT, + float, float, true, false, - mpp::tensor_ops::matmul2d_descriptor::mode::multiply>( - KD_tile.frag_at(0, 0), - KD_tile.frag_at(0, 1), + mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate>( + S_tile.frag_at(0, kk / 16), + S_tile.frag_at(0, kk / 16 + 1), delta_tile.frag_at(0, 0), metal::bool_constant{}, K_tile.frag_at(0, 0), K_tile.frag_at(0, 1), metal::bool_constant{}); - - FMA_NAX( - S_tile.frag_at(0, kk / 16), - metal::exp(gamma[C - 1]), - S_tile.frag_at(0, kk / 16), - KD_tile.frag_at(0, 0)) - FMA_NAX( - S_tile.frag_at(0, kk / 16 + 1), - metal::exp(gamma[C - 1]), - S_tile.frag_at(0, kk / 16 + 1), - KD_tile.frag_at(0, 1)) } - // OUTPUT(tmp_tile) - - SCALE_TRI_NAX(QKt_tile, gamma) - // OUTPUT(QKt_tile) - - out_tile = tmp_tile; - mlx::steel::mmak( - out_tile.frag_at(0, 0), - QKt_tile.frag_at(0, 0), - Z.frag_at(0, 0), - metal::bool_constant{}, - delta_tile.frag_at(0, 0), - Z.frag_at(0, 0), - metal::bool_constant{}); - - out_tile.store(y + dv_idx, Hv * Dv); - // OUTPUT(out_tile); // advance pointers q_ += C * Hk * Dk; @@ -700,7 +596,7 @@ template constant int& T [[buffer(8)]], uint3 thread_position_in_grid [[thread_position_in_grid]], uint thread_index_in_simdgroup [[thread_index_in_simdgroup]]) { - auto n = thread_position_in_grid.z; + auto n = thread_position_in_grid.z; // 7; auto b_idx = n / Hv; auto hv_idx = n % Hv; auto hk_idx = hv_idx / (Hv / Hk); @@ -713,6 +609,11 @@ template auto dv_idx = thread_position_in_grid.y * 8; +#define OUTPUT(T) \ + if (true) { \ + simdgroup_store(T, y); \ + return; \ + } // set up pointers // g: [B, T, Hv] auto g_ = g + b_idx * T * Hv; @@ -743,7 +644,7 @@ template simdgroup_float8x8 KD_tile; // tiles for WY form computation - simdgroup_float8x8 KKtK_tile, KKtV_tile, X_tile, KKt_tile; + simdgroup_float8x8 KKtK_tile, KKtV_tile, KKt_tile; threadgroup float gamma[C]; @@ -755,7 +656,6 @@ template for (int kk = 0; kk < Dk; kk += 8) { simdgroup_load(S_tile[kk / 8], i_state + kk, Dk, ulong2(0, 0), true); } - for (int t = 0; t < T; t += C) { float g_val = (thread_index_in_simdgroup < C) ? metal::log( @@ -773,13 +673,11 @@ template KKt_tile = make_filled_simdgroup_matrix(0.f); for (int kk = 0; kk < Dk; kk += 8) { simdgroup_load(K_tile, k_ + kk, Dk * Hk); + simdgroup_load(KT_tile, k_ + kk, Dk * Hk, ulong2(0, 0), true); simdgroup_multiply_accumulate(KKt_tile, K_tile, KT_tile, KKt_tile); } -#define OUTPUT(T) \ - simdgroup_store(T, y); \ - return; - // OUTPUT(KKt_tile) + KKtK_tile = KKt_tile; KKtV_tile = KKt_tile; @@ -790,7 +688,6 @@ template beta_fm * (metal::exp(gamma[fm] - gamma[fn])), beta_fm * (metal::exp(gamma[fm] - gamma[fn + 1]))) - // OUTPUT(KKtV_tile) // Tinv = (I + L_W)^{-1} = sum -L_W_k // T0 - T + T2 - T3 + T4 - T5 + T6 - T7 + T8 // I - T(I - T(I - T(I - T(I - T(I - T(I - T(I - T(I)))))))) @@ -804,7 +701,6 @@ template simdgroup_multiply(P, P, P); simdgroup_multiply_accumulate(Tinv, Tinv, P, Tinv); } - // OUTPUT(Tinv) WS_tile = make_filled_simdgroup_matrix(0.f); for (int kk = 0; kk < Dk; kk += 8) { @@ -817,7 +713,6 @@ template SCALE(W_tile, metal::exp(gamma[fm])) simdgroup_multiply_accumulate(WS_tile, W_tile, S_tile[kk / 8], WS_tile); } - // OUTPUT(WS_tile) AT(P, 0) = AT(KKtV_tile, 0); AT(P, 1) = AT(KKtV_tile, 1); @@ -826,16 +721,13 @@ template simdgroup_multiply(P, P, P); simdgroup_multiply_accumulate(Tinv, Tinv, P, Tinv); } - // OUTPUT(Tinv) // U = Tinv @ (beta * V) simdgroup_load(V_tile, v_ + dv_idx, Dv * Hv); SCALE(V_tile, beta_fm) simdgroup_multiply(U_tile, Tinv, V_tile); - // OUTPUT(U_tile) // delta = U - WS SUB(delta_tile, U_tile, WS_tile) - // OUTPUT(delta_tile) tmp_tile = make_filled_simdgroup_matrix(0.f); QKt_tile = make_filled_simdgroup_matrix(0.f); @@ -849,31 +741,29 @@ template SCALE(Q_tile, metal::exp(gamma[fm])) simdgroup_multiply_accumulate(tmp_tile, Q_tile, S_tile[kk / 8], tmp_tile); } - // OUTPUT(tmp_tile) // (Q @ K) * Gamma, in the paper they use M here but it's probably a typo. SCALE_TRI( QKt_tile, metal::exp(gamma[fm] - gamma[fn]), metal::exp(gamma[fm] - gamma[fn + 1])) - // OUTPUT(QKt_tile) simdgroup_multiply_accumulate(out_tile, QKt_tile, delta_tile, tmp_tile); y[fm * Hv * Dv + dv_idx + fn] = static_cast(AT(out_tile, 0)); y[fm * Hv * Dv + dv_idx + fn + 1] = static_cast(AT(out_tile, 1)); - // OUTPUT(out_tile) for (int kk = 0; kk < Dk; kk += 8) { simdgroup_load(K_tile, k_ + kk, Hk * Dk, ulong2(0, 0), true); - SCALE2( K_tile, metal::exp(gamma[C - 1] - gamma[fn]), metal::exp(gamma[C - 1] - gamma[fn + 1])) simdgroup_multiply(KD_tile, K_tile, delta_tile); + // SCALE(K_tile, 100000) + // FMA(S_tile[kk / 8], 1, S_tile[kk / 8], KD_tile) FMA(S_tile[kk / 8], metal::exp(gamma[C - 1]), S_tile[kk / 8], KD_tile) } @@ -974,4 +864,4 @@ template auto s_idx = n_per_t * dk_idx + i; o_state[s_idx] = static_cast(state[i]); } -} +} \ No newline at end of file From 2c49bfdfc69a2c4c4e4e8f8ba8f043fc13daf21d Mon Sep 17 00:00:00 2001 From: tpegolotti Date: Thu, 30 Jul 2026 06:26:19 -0700 Subject: [PATCH 30/56] Change threadgroup grid --- mlx/backend/metal/gated_delta_update.cpp | 6 ++---- mlx/backend/metal/kernels/gated_delta_update_impl.h | 11 +++++++++-- 2 files changed, 11 insertions(+), 6 deletions(-) diff --git a/mlx/backend/metal/gated_delta_update.cpp b/mlx/backend/metal/gated_delta_update.cpp index 0f1fefccfc..ed29053f97 100644 --- a/mlx/backend/metal/gated_delta_update.cpp +++ b/mlx/backend/metal/gated_delta_update.cpp @@ -136,8 +136,7 @@ void GatedDeltaUpdate::eval_gpu( compute_encoder.set_bytes(T, 8); auto grid = MTL::Size(32, Dv / 16, B * Hv); - // auto grid = MTL::Size(32, 1, 1); - auto threads = MTL::Size(32, 1, 1); + auto threads = MTL::Size(32, 4, 1); compute_encoder.dispatch_threads(grid, threads); break; } @@ -168,8 +167,7 @@ void GatedDeltaUpdate::eval_gpu( compute_encoder.set_bytes(T, 8); auto grid = MTL::Size(32, Dv / 8, B * Hv); - // auto grid = MTL::Size(32, 1, 1); - auto threads = MTL::Size(32, 1, 1); + auto threads = MTL::Size(32, 4, 1); compute_encoder.dispatch_threads(grid, threads); break; } diff --git a/mlx/backend/metal/kernels/gated_delta_update_impl.h b/mlx/backend/metal/kernels/gated_delta_update_impl.h index 09c27184ce..0da3bbbb43 100644 --- a/mlx/backend/metal/kernels/gated_delta_update_impl.h +++ b/mlx/backend/metal/kernels/gated_delta_update_impl.h @@ -323,6 +323,7 @@ template device StT* state_out [[buffer(7)]], constant int& T [[buffer(8)]], uint3 thread_position_in_grid [[thread_position_in_grid]], + uint3 thread_position_in_threadgroup [[thread_position_in_threadgroup]], uint thread_index_in_simdgroup [[thread_index_in_simdgroup]]) { auto n = thread_position_in_grid.z; // 7 auto b_idx = n / Hv; @@ -330,6 +331,7 @@ template auto hk_idx = hv_idx / (Hv / Hk); auto dv_idx = thread_position_in_grid.y * 16; + const short sg_id = thread_position_in_threadgroup.y; // 0..3 const ushort simd_lane_id = __metal_get_thread_index_in_simdgroup(ushort()); const short qid = simd_lane_id >> 2; @@ -352,7 +354,9 @@ template auto i_state = state_in + (n * Dv + dv_idx) * Dk; auto o_state = state_out + (n * Dv + dv_idx) * Dk; - threadgroup float gamma[C]; + threadgroup float gamma_all[C * 4]; + threadgroup float* gamma = gamma_all + sg_id * C; + float beta_fm[2]; mlx::steel::NAXTile S_tile; @@ -595,6 +599,7 @@ template device StT* state_out [[buffer(7)]], constant int& T [[buffer(8)]], uint3 thread_position_in_grid [[thread_position_in_grid]], + uint3 thread_position_in_threadgroup [[thread_position_in_threadgroup]], uint thread_index_in_simdgroup [[thread_index_in_simdgroup]]) { auto n = thread_position_in_grid.z; // 7; auto b_idx = n / Hv; @@ -608,6 +613,7 @@ template (thread_index_in_simdgroup % 2) * 2; // column coordinate of the held tile auto dv_idx = thread_position_in_grid.y * 8; + const short sg_id = thread_position_in_threadgroup.y; // 0..3 #define OUTPUT(T) \ if (true) { \ @@ -646,7 +652,8 @@ template // tiles for WY form computation simdgroup_float8x8 KKtK_tile, KKtV_tile, KKt_tile; - threadgroup float gamma[C]; + threadgroup float gamma_all[C * 4]; + threadgroup float* gamma = gamma_all + sg_id * C; simdgroup_float8x8 I_tile = make_filled_simdgroup_matrix(0.f); AT(I_tile, 0) = (fm == fn) ? 1.0f : 0.0f; From 8344147cf3226f4ee029465c5d51ccb8a047b263 Mon Sep 17 00:00:00 2001 From: tpegolotti Date: Thu, 30 Jul 2026 06:34:21 -0700 Subject: [PATCH 31/56] Remove one inverse --- .../metal/kernels/gated_delta_update_impl.h | 20 +------------------ 1 file changed, 1 insertion(+), 19 deletions(-) diff --git a/mlx/backend/metal/kernels/gated_delta_update_impl.h b/mlx/backend/metal/kernels/gated_delta_update_impl.h index 0da3bbbb43..d645f15c96 100644 --- a/mlx/backend/metal/kernels/gated_delta_update_impl.h +++ b/mlx/backend/metal/kernels/gated_delta_update_impl.h @@ -422,7 +422,6 @@ template KKtV_tile = KKt_tile; SCALE_TRIEQ_NAX1(KKtK_tile, beta_fm) - SCALE_TRIEQ_NAX(KKtV_tile, beta_fm, gamma) Tinv_tile = I_tile; STEEL_PRAGMA_UNROLL @@ -466,24 +465,7 @@ template } SCALE_ROW_NAX(W_tile, gamma) - Tinv_tile = I_tile; - STEEL_PRAGMA_UNROLL - for (int step = 0; step < C - 1; step++) { - mlx::steel::mma< - float, - float, - float, - false, - false, - mpp::tensor_ops::matmul2d_descriptor::mode::multiply>( - TMP_tile.frag_at(0, 0), - KKtV_tile.frag_at(0, 0), - metal::bool_constant{}, - Tinv_tile.frag_at(0, 0), - metal::bool_constant{}); - SUB_NAX(Tinv_tile, I_tile, TMP_tile) - } - + SCALE_TRI_NAX(Tinv_tile, gamma) V_tile.load(v_ + dv_idx, Dv * Hv); SCALE_BETA_NAX(V_tile, beta_fm) // U = Tinv @ diag(b)V From f187201a60f4f683e2679ae2a5debfd859f34142 Mon Sep 17 00:00:00 2001 From: tpegolotti Date: Fri, 31 Jul 2026 06:28:11 -0700 Subject: [PATCH 32/56] Improve inverse --- .../metal/kernels/gated_delta_update_impl.h | 240 +++++++++--------- 1 file changed, 113 insertions(+), 127 deletions(-) diff --git a/mlx/backend/metal/kernels/gated_delta_update_impl.h b/mlx/backend/metal/kernels/gated_delta_update_impl.h index d645f15c96..3a3ed683de 100644 --- a/mlx/backend/metal/kernels/gated_delta_update_impl.h +++ b/mlx/backend/metal/kernels/gated_delta_update_impl.h @@ -67,6 +67,14 @@ using namespace mpp::tensor_ops; } \ } +#define ADD_NAX(TILE0, TILE1, TILE2) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerFrag; _i++) { \ + AT_NAX(TILE0, _i) = AT_NAX(TILE1, _i) + AT_NAX(TILE2, _i); \ + } \ + } + #define FMA_NAX(TILE0, S, TILE1, TILE2) \ { \ STEEL_PRAGMA_UNROLL \ @@ -281,18 +289,8 @@ METAL_FUNC static constexpr void mman( STEEL_PRAGMA_UNROLL for (short i = 0; i < BaseNAXFrag::kElemsPerFrag; i++) { ct_a[i] = A[i]; - } - - // Load B into right operand registers - STEEL_PRAGMA_UNROLL - for (short i = 0; i < BaseNAXFrag::kElemsPerFrag; i++) { ct_b[i] = Bn0[i]; ct_b[BaseNAXFrag::kElemsPerFrag + i] = Bn1[i]; - } - - // Load C into output registers (op handles accumulation) - STEEL_PRAGMA_UNROLL - for (short i = 0; i < BaseNAXFrag::kElemsPerFrag; i++) { ct_c[i] = Cn0[i]; ct_c[BaseNAXFrag::kElemsPerFrag + i] = Cn1[i]; } @@ -311,6 +309,82 @@ METAL_FUNC static constexpr void mman( } // namespace steel } // namespace mlx +#define MM16x16x16(C, CO, A, TA, AO, B, TB, BO) \ + mlx::steel::mma< \ + float, \ + float, \ + float, \ + TA, \ + TB, \ + mpp::tensor_ops::matmul2d_descriptor::mode::multiply>( \ + C.frag_at(0, (CO)), \ + A.frag_at(0, (AO)), \ + metal::bool_constant{}, \ + B.frag_at(0, (BO)), \ + metal::bool_constant{}); + +#define MMA16x16x16(C, CO, A, TA, AO, B, TB, BO) \ + mlx::steel::mma< \ + float, \ + float, \ + float, \ + TA, \ + TB, \ + mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate>( \ + C.frag_at(0, (CO)), \ + A.frag_at(0, (AO)), \ + metal::bool_constant{}, \ + B.frag_at(0, (BO)), \ + metal::bool_constant{}); + +#define MMA16x16x32(C, CO, A, TA, AO, B, TB, BO) \ + mlx::steel::mma< \ + float, \ + float, \ + float, \ + TA, \ + TB, \ + mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate>( \ + C.frag_at(0, (CO)), \ + A.frag_at(0, (AO)), \ + A.frag_at(0, (AO) + 1), \ + metal::bool_constant{}, \ + B.frag_at(0, (BO)), \ + B.frag_at(0, (BO) + 1), \ + metal::bool_constant{}); + +#define MM16x32x16(C, CO, A, TA, AO, B, TB, BO) \ + mlx::steel::mman< \ + float, \ + float, \ + float, \ + TA, \ + TB, \ + mpp::tensor_ops::matmul2d_descriptor::mode::multiply>( \ + C.frag_at(0, (CO)), \ + C.frag_at(0, (CO) + 1), \ + A.frag_at(0, (AO)), \ + metal::bool_constant{}, \ + B.frag_at(0, (BO)), \ + B.frag_at(0, (BO) + 1), \ + metal::bool_constant{}); + +#define MMA16x32x16(C, CO, A, TA, AO, B, TB, BO) \ + mlx::steel::mman< \ + float, \ + float, \ + float, \ + TA, \ + TB, \ + mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate>( \ + C.frag_at(0, (CO)), \ + C.frag_at(0, (CO) + 1), \ + A.frag_at(0, (AO)), \ + metal::bool_constant{}, \ + B.frag_at(0, (BO)), \ + B.frag_at(0, (BO) + 1), \ + metal::bool_constant{}); + template [[kernel]] void gated_delta_fused_nax( const device InT* q [[buffer(0)]], @@ -372,7 +446,7 @@ template mlx::steel::NAXTile tmp_tile; mlx::steel::NAXTile QKt_tile; mlx::steel::NAXTile out_tile; - mlx::steel::NAXTile Tinv_tile; + mlx::steel::NAXTile Tinv_tile, P; mlx::steel::NAXTile KKtK_tile, KKtV_tile, KKt_tile; @@ -407,61 +481,27 @@ template // KP_tile.load(k_, Dk * Hk); for (int kk = 0; kk < Dk; kk += 32) { // two 16-tiles per iter K_tile.load(k_ + kk, Dk * Hk); - - mlx::steel::mma( - KKt_tile.frag_at(0, 0), - K_tile.frag_at(0, 0), - K_tile.frag_at(0, 1), - metal::bool_constant{}, - K_tile.frag_at(0, 0), - K_tile.frag_at(0, 1), - metal::bool_constant{}); + MMA16x16x32(KKt_tile, 0, K_tile, false, 0, K_tile, true, 0) } KKtK_tile = KKt_tile; KKtV_tile = KKt_tile; SCALE_TRIEQ_NAX1(KKtK_tile, beta_fm) + MM16x16x16(P, 0, KKtK_tile, false, 0, KKtK_tile, false, 0) - Tinv_tile = I_tile; - STEEL_PRAGMA_UNROLL - for (int step = 0; step < C - 1; step++) { - // TMP = T_K · Tinv - mlx::steel::mma< - float, - float, - float, - false, - false, - mpp::tensor_ops::matmul2d_descriptor::mode::multiply>( - TMP_tile.frag_at(0, 0), - KKtK_tile.frag_at(0, 0), - metal::bool_constant{}, - Tinv_tile.frag_at(0, 0), - metal::bool_constant{}); - // Tinv = I - TMP - SUB_NAX(Tinv_tile, I_tile, TMP_tile) - } + ADD_NAX(Tinv_tile, I_tile, P) + STEEL_PRAGMA_UNROLL for (int step = 0; step < 6; step++){ + MM16x16x16(TMP_tile, 0, P, false, 0, Tinv_tile, false, 0) + ADD_NAX(Tinv_tile, I_tile, TMP_tile)} - STEEL_PRAGMA_UNROLL - for (short nn = 0; nn < Dk / 16; nn += 2) { + MM16x16x16(TMP_tile, 0, KKtK_tile, false, 0, Tinv_tile, false, 0) + SUB_NAX(Tinv_tile, Tinv_tile, TMP_tile) + + STEEL_PRAGMA_UNROLL for (short nn = 0; nn < Dk / 16; nn += 2) { K_tile.load(k_ + nn * 16, Dk * Hk); SCALE_BETA_NAX(K_tile, beta_fm) - - mlx::steel::mman< - float, - float, - float, - false, - false, - mpp::tensor_ops::matmul2d_descriptor::mode::multiply>( - W_tile.frag_at(0, nn), - W_tile.frag_at(0, nn + 1), - Tinv_tile.frag_at(0, 0), // A = Tinv (reused across N) - metal::bool_constant{}, - K_tile.frag_at(0, 0), // B = beta-scaled K - K_tile.frag_at(0, 1), - metal::bool_constant{}); + MM16x32x16(W_tile, nn, Tinv_tile, false, 0, K_tile, false, 0) } SCALE_ROW_NAX(W_tile, gamma) @@ -469,30 +509,12 @@ template V_tile.load(v_ + dv_idx, Dv * Hv); SCALE_BETA_NAX(V_tile, beta_fm) // U = Tinv @ diag(b)V - mlx::steel::mma< - float, - float, - float, - false, - false, - mpp::tensor_ops::matmul2d_descriptor::mode::multiply>( - U_tile.frag_at(0, 0), - Tinv_tile.frag_at(0, 0), - metal::bool_constant{}, - V_tile.frag_at(0, 0), - metal::bool_constant{}); - - WS_tile.clear(); + MM16x16x16(U_tile, 0, Tinv_tile, false, 0, V_tile, false, 0) + + WS_tile.clear(); STEEL_PRAGMA_UNROLL for (short kk = 0; kk < Dk / 16; kk += 2) { - mlx::steel::mma( - WS_tile.frag_at(0, 0), - W_tile.frag_at(0, kk), - W_tile.frag_at(0, kk + 1), - metal::bool_constant{}, - S_tile.frag_at(0, kk), - S_tile.frag_at(0, kk + 1), - metal::bool_constant{}); + MMA16x16x32(WS_tile, 0, W_tile, false, kk, S_tile, true, kk) } SUB_NAX(delta_tile, U_tile, WS_tile) @@ -501,61 +523,28 @@ template QKt_tile.clear(); for (int kk = 0; kk < Dk; kk += 32) { Q_tile.load(q_ + kk, Hk * Dk); - K_tile.load(k_ + kk, Hk * Dk); // this one is transposed + K_tile.load(k_ + kk, Hk * Dk); // Q @ K^T - mlx::steel::mma( - QKt_tile.frag_at(0, 0), - Q_tile.frag_at(0, 0), - Q_tile.frag_at(0, 1), - metal::bool_constant{}, - K_tile.frag_at(0, 0), - K_tile.frag_at(0, 1), - metal::bool_constant{}); - - SCALE_ROW_NAX(Q_tile, gamma) - // Q_left @ S^T - mlx::steel::mma( - tmp_tile.frag_at(0, 0), - Q_tile.frag_at(0, 0), - Q_tile.frag_at(0, 1), - metal::bool_constant{}, - S_tile.frag_at(0, kk / 16), - S_tile.frag_at(0, kk / 16 + 1), - metal::bool_constant{}); + MMA16x16x32(QKt_tile, 0, Q_tile, false, 0, K_tile, true, 0) + + // Q_left @ S^T + SCALE_ROW_NAX(Q_tile, gamma) + MMA16x16x32(tmp_tile, 0, Q_tile, false, 0, S_tile, true, kk / 16) } SCALE_TRI_NAX(QKt_tile, gamma) out_tile = tmp_tile; - mlx::steel::mma( - out_tile.frag_at(0, 0), - QKt_tile.frag_at(0, 0), - metal::bool_constant{}, - delta_tile.frag_at(0, 0), - metal::bool_constant{}); - - out_tile.store(y + dv_idx, Hv * Dv); + MMA16x16x16(out_tile, 0, QKt_tile, false, 0, delta_tile, false, 0) + out_tile.store(y + dv_idx, Hv * Dv); SCALE_NAX(S_tile, metal::exp(gamma[C - 1])) for (int kk = 0; kk < Dk; kk += 32) { K_tile.load(k_ + kk, Hk * Dk); SCALE2_NAX(K_tile, gamma) - mlx::steel::mman< - float, - float, - float, - true, - false, - mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate>( - S_tile.frag_at(0, kk / 16), - S_tile.frag_at(0, kk / 16 + 1), - delta_tile.frag_at(0, 0), - metal::bool_constant{}, - K_tile.frag_at(0, 0), - K_tile.frag_at(0, 1), - metal::bool_constant{}); + MMA16x32x16(S_tile, kk / 16, delta_tile, true, 0, K_tile, false, 0) } // advance pointers @@ -703,13 +692,10 @@ template simdgroup_multiply_accumulate(WS_tile, W_tile, S_tile[kk / 8], WS_tile); } - AT(P, 0) = AT(KKtV_tile, 0); - AT(P, 1) = AT(KKtV_tile, 1); - SUB(Tinv, I_tile, KKtV_tile) - for (int step = 1; (1 << step) < C; step++) { - simdgroup_multiply(P, P, P); - simdgroup_multiply_accumulate(Tinv, Tinv, P, Tinv); - } + SCALE_TRI( + Tinv, + metal::exp(gamma[fm] - gamma[fn]), + metal::exp(gamma[fm] - gamma[fn + 1])) // U = Tinv @ (beta * V) simdgroup_load(V_tile, v_ + dv_idx, Dv * Hv); From 5be2b4830bf23cd2f3bac839997d2180e8b30103 Mon Sep 17 00:00:00 2001 From: tpegolotti Date: Mon, 3 Aug 2026 08:18:28 -0700 Subject: [PATCH 33/56] Pre PR changes: remove C as parameter, remove explicit padding, fix typing --- mlx/backend/metal/gated_delta_update.cpp | 31 +-- .../metal/kernels/gated_delta_update.metal | 131 +++--------- .../metal/kernels/gated_delta_update_impl.h | 194 ++++++++++-------- mlx/fast.cpp | 34 +-- mlx/fast.h | 1 - mlx/fast_primitives.h | 6 +- python/src/fast.cpp | 1 - 7 files changed, 166 insertions(+), 232 deletions(-) diff --git a/mlx/backend/metal/gated_delta_update.cpp b/mlx/backend/metal/gated_delta_update.cpp index ed29053f97..264446ac9b 100644 --- a/mlx/backend/metal/gated_delta_update.cpp +++ b/mlx/backend/metal/gated_delta_update.cpp @@ -62,13 +62,6 @@ void GatedDeltaUpdate::eval_gpu( auto& s = stream(); auto& d = metal::device(s.device); - // auto& q = inputs[0]; - // auto& k = inputs[1]; - // auto& v = inputs[2]; - // auto& g = inputs[3]; - // auto& beta = inputs[4]; - // auto& h0 = inputs[5]; - auto q = ensure_row_contiguous(inputs[0], d, s); auto k = ensure_row_contiguous(inputs[1], d, s); auto v = ensure_row_contiguous(inputs[2], d, s); @@ -86,15 +79,17 @@ void GatedDeltaUpdate::eval_gpu( int Hv = v.shape(2); int Dv = v.shape(3); - // PRINT_ARR(q); - // PRINT_ARR(k); - // PRINT_ARR(v); - // PRINT_ARR(g); - // PRINT_ARR(beta); - // PRINT_ARR(h0); - // printf("%d %d %d %d %d %d\n",B,T,Hk,Hv,Dk,Dv); - - int C = chunk_size; + int C = 1; + const char* threashold_env = std::getenv("GATED_DELTA_THRESH"); + int threshold = threashold_env ? std::stoi(threashold_env) : 16; + if (T > threshold) { + if (metal::is_nax_available()) + C = 16; + else + C = 8; + } + const char* chunk_env = std::getenv("GATED_DELTA_CHUNK"); + C = chunk_env ? std::stoi(chunk_env) : C; std::string suffix = get_type_string(q.dtype()) // "float" + "_" + get_type_string(h0.dtype()) // "float" @@ -111,8 +106,6 @@ void GatedDeltaUpdate::eval_gpu( switch (C) { case 16: { std::string kernel_name = "gated_delta_fused_nax_"; - - // printf("C: %d\nname: %s\n",C,kernel_name.c_str()); std::string base_name = kernel_name + suffix; base_name += "_" + std::to_string(C); @@ -142,8 +135,6 @@ void GatedDeltaUpdate::eval_gpu( } case 8: { std::string kernel_name = "gated_delta_fused_chunk_"; - - // printf("C: %d\nname: %s\n",C,kernel_name.c_str()); std::string base_name = kernel_name + suffix; base_name += "_" + std::to_string(C); diff --git a/mlx/backend/metal/kernels/gated_delta_update.metal b/mlx/backend/metal/kernels/gated_delta_update.metal index 1e9adba9e2..0d41a100a3 100644 --- a/mlx/backend/metal/kernels/gated_delta_update.metal +++ b/mlx/backend/metal/kernels/gated_delta_update.metal @@ -15,22 +15,13 @@ using namespace metal; hk, \ hv) -#define instantiate_gated_delta_update_seq_dims(in_type, st_type) \ - instantiate_gated_delta_update_seq(in_type, st_type, 64, 64, 4, 4) \ - instantiate_gated_delta_update_seq(in_type, st_type, 64, 64, 8, 8) \ - instantiate_gated_delta_update_seq(in_type, st_type, 64, 64, 16, 16) \ - instantiate_gated_delta_update_seq( \ - in_type, st_type, 64, 64, 24, 24) \ - instantiate_gated_delta_update_seq( \ - in_type, st_type, 64, 64, 32, 32) \ - instantiate_gated_delta_update_seq( \ - in_type, st_type, 128, 128, 16, 16) \ - instantiate_gated_delta_update_seq( \ - in_type, st_type, 128, 128, 24, 24) \ - instantiate_gated_delta_update_seq( \ - in_type, st_type, 128, 128, 32, 32) \ - instantiate_gated_delta_update_seq( \ - in_type, st_type, 128, 128, 16, 32) +#define instantiate_gated_delta_update_seq_dims(in_type, st_type) \ + instantiate_gated_delta_update_seq(in_type, st_type, 128, 128, 24, 24) \ + instantiate_gated_delta_update_seq(in_type, st_type, 128, 128, 32, 32) \ + instantiate_gated_delta_update_seq( \ + in_type, st_type, 128, 128, 16, 32) \ + instantiate_gated_delta_update_seq( \ + in_type, st_type, 128, 128, 16, 48) #define instantiate_gated_delta_update_fused_chunk( \ in_type, st_type, dk, dv, hk, hv, c) \ @@ -46,43 +37,15 @@ using namespace metal; hv, \ c) -#define instantiate_gated_delta_update_fused_chunk_dims(in_type, st_type) \ - instantiate_gated_delta_update_fused_chunk( \ - in_type, st_type, 128, 128, 16, 16, 8) \ - instantiate_gated_delta_update_fused_chunk( \ - in_type, st_type, 128, 128, 24, 24, 8) \ - instantiate_gated_delta_update_fused_chunk( \ - in_type, st_type, 128, 128, 32, 32, 8) \ - instantiate_gated_delta_update_fused_chunk( \ - in_type, st_type, 64, 64, 4, 4, 8) \ - instantiate_gated_delta_update_fused_chunk( \ - in_type, st_type, 64, 64, 8, 8, 8) \ - instantiate_gated_delta_update_fused_chunk( \ - in_type, st_type, 64, 64, 16, 16, 8) \ - instantiate_gated_delta_update_fused_chunk( \ - in_type, st_type, 64, 64, 24, 24, 8) \ - instantiate_gated_delta_update_fused_chunk( \ - in_type, st_type, 64, 64, 32, 32, 8) \ - instantiate_gated_delta_update_fused_chunk( \ - in_type, st_type, 128, 128, 16, 32, 8) \ - instantiate_gated_delta_update_fused_chunk( \ - in_type, st_type, 32, 32, 16, 32, 8) \ - instantiate_gated_delta_update_fused_chunk( \ - in_type, \ - st_type, \ - 32, \ - 32, \ - 32, \ - 32, \ - 8) \ - instantiate_gated_delta_update_fused_chunk( \ - in_type, \ - st_type, \ - 16, \ - 16, \ - 16, \ - 32, \ - 8) +#define instantiate_gated_delta_update_fused_chunk_dims(in_type, st_type) \ + instantiate_gated_delta_update_fused_chunk( \ + in_type, st_type, 128, 128, 16, 32, 8) \ + instantiate_gated_delta_update_fused_chunk( \ + in_type, st_type, 128, 128, 16, 48, 8) \ + instantiate_gated_delta_update_fused_chunk( \ + in_type, st_type, 128, 128, 24, 24, 8) \ + instantiate_gated_delta_update_fused_chunk( \ + in_type, st_type, 128, 128, 32, 32, 8) #define instantiate_gated_delta_update_fused_nax( \ in_type, st_type, dk, dv, hk, hv, c) \ @@ -98,50 +61,20 @@ using namespace metal; hv, \ c) -#define instantiate_gated_delta_update_fused_nax_dims(in_type, st_type) \ - instantiate_gated_delta_update_fused_nax( \ - in_type, st_type, 128, 128, 16, 16, 16) \ - instantiate_gated_delta_update_fused_nax( \ - in_type, st_type, 128, 128, 24, 24, 16) \ - instantiate_gated_delta_update_fused_nax( \ - in_type, st_type, 128, 128, 32, 32, 16) \ - instantiate_gated_delta_update_fused_nax( \ - in_type, st_type, 64, 64, 4, 4, 16) \ - instantiate_gated_delta_update_fused_nax( \ - in_type, st_type, 64, 64, 8, 8, 16) \ - instantiate_gated_delta_update_fused_nax( \ - in_type, st_type, 64, 64, 16, 16, 16) \ - instantiate_gated_delta_update_fused_nax( \ - in_type, st_type, 64, 64, 24, 24, 16) \ - instantiate_gated_delta_update_fused_nax( \ - in_type, st_type, 64, 64, 32, 32, 16) \ - instantiate_gated_delta_update_fused_nax( \ - in_type, st_type, 128, 128, 16, 32, 16) \ - instantiate_gated_delta_update_fused_nax( \ - in_type, \ - st_type, \ - 32, \ - 32, \ - 16, \ - 32, \ - 16) \ - instantiate_gated_delta_update_fused_nax( \ - in_type, \ - st_type, \ - 32, \ - 32, \ - 32, \ - 32, \ - 16) \ - instantiate_gated_delta_update_fused_nax( \ - in_type, \ - st_type, \ - 16, \ - 16, \ - 16, \ - 32, \ - 16) +#define instantiate_gated_delta_update_fused_nax_dims(in_type, st_type) \ + instantiate_gated_delta_update_fused_nax( \ + in_type, st_type, 128, 128, 16, 32, 16) \ + instantiate_gated_delta_update_fused_nax( \ + in_type, st_type, 128, 128, 16, 48, 16) \ + instantiate_gated_delta_update_fused_nax( \ + in_type, st_type, 128, 128, 24, 24, 16) \ + instantiate_gated_delta_update_fused_nax( \ + in_type, st_type, 128, 128, 32, 32, 16) -instantiate_gated_delta_update_seq_dims(float, float) - instantiate_gated_delta_update_fused_chunk_dims(float, float) - instantiate_gated_delta_update_fused_nax_dims(float, float) \ No newline at end of file +instantiate_gated_delta_update_seq_dims(float, float); +instantiate_gated_delta_update_fused_chunk_dims(float, float); +instantiate_gated_delta_update_fused_nax_dims(float, float); + +instantiate_gated_delta_update_seq_dims(bfloat16_t, float); +// instantiate_gated_delta_update_fused_chunk_dims(bfloat16_t, float); +instantiate_gated_delta_update_fused_nax_dims(bfloat16_t, float); \ No newline at end of file diff --git a/mlx/backend/metal/kernels/gated_delta_update_impl.h b/mlx/backend/metal/kernels/gated_delta_update_impl.h index 3a3ed683de..a7f4627c2b 100644 --- a/mlx/backend/metal/kernels/gated_delta_update_impl.h +++ b/mlx/backend/metal/kernels/gated_delta_update_impl.h @@ -399,7 +399,7 @@ template uint3 thread_position_in_grid [[thread_position_in_grid]], uint3 thread_position_in_threadgroup [[thread_position_in_threadgroup]], uint thread_index_in_simdgroup [[thread_index_in_simdgroup]]) { - auto n = thread_position_in_grid.z; // 7 + auto n = thread_position_in_grid.z; auto b_idx = n / Hv; auto hv_idx = n % Hv; auto hk_idx = hv_idx / (Hv / Hk); @@ -460,55 +460,64 @@ template } mlx::steel::NAXTile TMP_tile; - for (int t = 0; t < T; t += C) { - float g_val = (thread_index_in_simdgroup < C) + auto process_chunk = [&](const short valid_rows, auto bounded_tag) { + constexpr bool B = decltype(bounded_tag)::value; + + auto load_seq = [&](thread auto& tile, auto src, int ld) { + if constexpr (B) { + tile.load_rows(src, ld, valid_rows); + } else { + tile.load(src, ld); + } + }; + + auto g_val = (thread_index_in_simdgroup < (uint)valid_rows) ? metal::log( metal::clamp( g_[thread_index_in_simdgroup * Hv + hv_idx], 1e-6f, 1.0f)) : 0.0f; - - float gamma_val = simd_prefix_inclusive_sum(g_val); - + auto gamma_val = simd_prefix_inclusive_sum(g_val); if (thread_index_in_simdgroup < C) { - gamma[thread_index_in_simdgroup] = gamma_val; + gamma[thread_index_in_simdgroup] = static_cast(gamma_val); } - beta_fm[0] = beta_[fm * Hv + hv_idx]; - beta_fm[1] = - beta_[(fm + mlx::steel::BaseNAXFrag::kElemRowsJump) * Hv + hv_idx]; + beta_fm[0] = (fm < valid_rows) ? beta_[fm * Hv + hv_idx] : 0.0f; + const short fm1 = fm + mlx::steel::BaseNAXFrag::kElemRowsJump; + beta_fm[1] = (fm1 < valid_rows) ? beta_[fm1 * Hv + hv_idx] : 0.0f; KKt_tile.clear(); - // KP_tile.load(k_, Dk * Hk); - for (int kk = 0; kk < Dk; kk += 32) { // two 16-tiles per iter - K_tile.load(k_ + kk, Dk * Hk); - MMA16x16x32(KKt_tile, 0, K_tile, false, 0, K_tile, true, 0) + for (int kk = 0; kk < Dk; kk += 32) { + load_seq(K_tile, k_ + kk, Dk * Hk); + MMA16x16x32(KKt_tile, 0, K_tile, false, 0, K_tile, true, 0); } KKtK_tile = KKt_tile; KKtV_tile = KKt_tile; - SCALE_TRIEQ_NAX1(KKtK_tile, beta_fm) - MM16x16x16(P, 0, KKtK_tile, false, 0, KKtK_tile, false, 0) + SCALE_TRIEQ_NAX1(KKtK_tile, beta_fm); + MM16x16x16(P, 0, KKtK_tile, false, 0, KKtK_tile, false, 0); - ADD_NAX(Tinv_tile, I_tile, P) - STEEL_PRAGMA_UNROLL for (int step = 0; step < 6; step++){ - MM16x16x16(TMP_tile, 0, P, false, 0, Tinv_tile, false, 0) - ADD_NAX(Tinv_tile, I_tile, TMP_tile)} + ADD_NAX(Tinv_tile, I_tile, P); + STEEL_PRAGMA_UNROLL + for (int step = 0; step < 6; step++) { + MM16x16x16(TMP_tile, 0, P, false, 0, Tinv_tile, false, 0); + ADD_NAX(Tinv_tile, I_tile, TMP_tile); + } - MM16x16x16(TMP_tile, 0, KKtK_tile, false, 0, Tinv_tile, false, 0) - SUB_NAX(Tinv_tile, Tinv_tile, TMP_tile) + MM16x16x16(TMP_tile, 0, KKtK_tile, false, 0, Tinv_tile, false, 0); + SUB_NAX(Tinv_tile, Tinv_tile, TMP_tile); - STEEL_PRAGMA_UNROLL for (short nn = 0; nn < Dk / 16; nn += 2) { - K_tile.load(k_ + nn * 16, Dk * Hk); - SCALE_BETA_NAX(K_tile, beta_fm) - MM16x32x16(W_tile, nn, Tinv_tile, false, 0, K_tile, false, 0) + STEEL_PRAGMA_UNROLL + for (short nn = 0; nn < Dk / 16; nn += 2) { + load_seq(K_tile, k_ + nn * 16, Dk * Hk); + SCALE_BETA_NAX(K_tile, beta_fm); + MM16x32x16(W_tile, nn, Tinv_tile, false, 0, K_tile, false, 0); } SCALE_ROW_NAX(W_tile, gamma) SCALE_TRI_NAX(Tinv_tile, gamma) - V_tile.load(v_ + dv_idx, Dv * Hv); - SCALE_BETA_NAX(V_tile, beta_fm) - // U = Tinv @ diag(b)V + load_seq(V_tile, v_ + dv_idx, Dv * Hv); + SCALE_BETA_NAX(V_tile, beta_fm); MM16x16x16(U_tile, 0, Tinv_tile, false, 0, V_tile, false, 0) WS_tile.clear(); @@ -522,32 +531,43 @@ template tmp_tile.clear(); QKt_tile.clear(); for (int kk = 0; kk < Dk; kk += 32) { - Q_tile.load(q_ + kk, Hk * Dk); - K_tile.load(k_ + kk, Hk * Dk); + load_seq(Q_tile, q_ + kk, Hk * Dk); + load_seq(K_tile, k_ + kk, Hk * Dk); - // Q @ K^T - MMA16x16x32(QKt_tile, 0, Q_tile, false, 0, K_tile, true, 0) + MMA16x16x32(QKt_tile, 0, Q_tile, false, 0, K_tile, true, 0); - // Q_left @ S^T - SCALE_ROW_NAX(Q_tile, gamma) - MMA16x16x32(tmp_tile, 0, Q_tile, false, 0, S_tile, true, kk / 16) + SCALE_ROW_NAX(Q_tile, gamma); + MMA16x16x32(tmp_tile, 0, Q_tile, false, 0, S_tile, true, kk / 16); } SCALE_TRI_NAX(QKt_tile, gamma) out_tile = tmp_tile; - MMA16x16x16(out_tile, 0, QKt_tile, false, 0, delta_tile, false, 0) - out_tile.store(y + dv_idx, Hv * Dv); + MMA16x16x16(out_tile, 0, QKt_tile, false, 0, delta_tile, false, 0); + + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(out_tile)::kElemsPerFrag; _i++) { + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); // {fn, fm} + const short _fn = _c.x; + const short _fm = _c.y; + if (_fm < valid_rows) { + y[_fm * Hv * Dv + dv_idx + _fn] = + static_cast(AT_NAX(out_tile, _i)); + } + } - SCALE_NAX(S_tile, metal::exp(gamma[C - 1])) + SCALE_NAX(S_tile, metal::exp(gamma[C - 1])); for (int kk = 0; kk < Dk; kk += 32) { - K_tile.load(k_ + kk, Hk * Dk); - SCALE2_NAX(K_tile, gamma) - MMA16x32x16(S_tile, kk / 16, delta_tile, true, 0, K_tile, false, 0) + load_seq(K_tile, k_ + kk, Hk * Dk); + SCALE2_NAX(K_tile, gamma); + MMA16x32x16(S_tile, kk / 16, delta_tile, true, 0, K_tile, false, 0); } + }; - // advance pointers + int t = 0; + for (; t + C <= T; t += C) { + process_chunk(C, metal::false_type{}); q_ += C * Hk * Dk; k_ += C * Hk * Dk; v_ += C * Hv * Dv; @@ -555,6 +575,10 @@ template y += C * Hv * Dv; g_ += C * Hv; } + if (t < T) { + process_chunk(short(T - t), metal::true_type{}); + } + S_tile.store(o_state, Dk); } @@ -572,7 +596,7 @@ template uint3 thread_position_in_grid [[thread_position_in_grid]], uint3 thread_position_in_threadgroup [[thread_position_in_threadgroup]], uint thread_index_in_simdgroup [[thread_index_in_simdgroup]]) { - auto n = thread_position_in_grid.z; // 7; + auto n = thread_position_in_grid.z; auto b_idx = n / Hv; auto hv_idx = n % Hv; auto hk_idx = hv_idx / (Hv / Hk); @@ -630,47 +654,62 @@ template AT(I_tile, 0) = (fm == fn) ? 1.0f : 0.0f; AT(I_tile, 1) = (fm == fn + 1) ? 1.0f : 0.0f; - // load initial state into threadgroup + // load initial state into registers for (int kk = 0; kk < Dk; kk += 8) { simdgroup_load(S_tile[kk / 8], i_state + kk, Dk, ulong2(0, 0), true); } - for (int t = 0; t < T; t += C) { - float g_val = (thread_index_in_simdgroup < C) + + auto process_chunk = [&](const short valid_rows, auto bounded_tag) { + constexpr bool B = decltype(bounded_tag)::value; + + // non-transposed load + auto load_M = [&](thread simdgroup_float8x8& M, auto src, int ld) { + if constexpr (B) { + AT(M, 0) = (fm < valid_rows) ? float(src[fm * ld + fn]) : 0.f; + AT(M, 1) = (fm < valid_rows) ? float(src[fm * ld + fn + 1]) : 0.f; + } else { + simdgroup_load(M, src, ld); + } + }; + + // transposed load + auto load_MT = [&](thread simdgroup_float8x8& M, auto src, int ld) { + if constexpr (B) { + AT(M, 0) = (fn < valid_rows) ? float(src[fn * ld + fm]) : 0.f; + AT(M, 1) = (fn + 1 < valid_rows) ? float(src[(fn + 1) * ld + fm]) : 0.f; + } else { + simdgroup_load(M, src, ld, ulong2(0, 0), true); + } + }; + + float g_val = (thread_index_in_simdgroup < (uint)valid_rows) ? metal::log( metal::max(g_[thread_index_in_simdgroup * Hv + hv_idx], 1e-6f)) : 0.0f; - float gamma_val = simd_prefix_inclusive_sum(g_val); - if (thread_index_in_simdgroup < C) { gamma[thread_index_in_simdgroup] = gamma_val; } - float beta_fm = beta_[fm * Hv + hv_idx]; + float beta_fm = (fm < valid_rows) ? beta_[fm * Hv + hv_idx] : 0.0f; KKt_tile = make_filled_simdgroup_matrix(0.f); for (int kk = 0; kk < Dk; kk += 8) { - simdgroup_load(K_tile, k_ + kk, Dk * Hk); - - simdgroup_load(KT_tile, k_ + kk, Dk * Hk, ulong2(0, 0), true); + load_M(K_tile, k_ + kk, Dk * Hk); + load_MT(KT_tile, k_ + kk, Dk * Hk); simdgroup_multiply_accumulate(KKt_tile, K_tile, KT_tile, KKt_tile); } KKtK_tile = KKt_tile; KKtV_tile = KKt_tile; - // elementwise multiplication by Gamma and beta (for V) and beta (for K) SCALE_TRIEQ(KKtK_tile, beta_fm, beta_fm) SCALE_TRIEQ( KKtV_tile, beta_fm * (metal::exp(gamma[fm] - gamma[fn])), beta_fm * (metal::exp(gamma[fm] - gamma[fn + 1]))) - // Tinv = (I + L_W)^{-1} = sum -L_W_k - // T0 - T + T2 - T3 + T4 - T5 + T6 - T7 + T8 - // I - T(I - T(I - T(I - T(I - T(I - T(I - T(I - T(I)))))))) - simdgroup_float8x8 Tinv, P; // - // S = I + x = I - KKtK + simdgroup_float8x8 Tinv, P; AT(P, 0) = AT(KKtK_tile, 0); AT(P, 1) = AT(KKtK_tile, 1); SUB(Tinv, I_tile, KKtK_tile) @@ -682,12 +721,9 @@ template WS_tile = make_filled_simdgroup_matrix(0.f); for (int kk = 0; kk < Dk; kk += 8) { - simdgroup_load(K_tile, k_ + kk, Dk * Hk); + load_M(K_tile, k_ + kk, Dk * Hk); SCALE(K_tile, beta_fm) - - // W = Tinv @ (beta * K) simdgroup_multiply(W_tile, Tinv, K_tile); - SCALE(W_tile, metal::exp(gamma[fm])) simdgroup_multiply_accumulate(WS_tile, W_tile, S_tile[kk / 8], WS_tile); } @@ -697,27 +733,21 @@ template metal::exp(gamma[fm] - gamma[fn]), metal::exp(gamma[fm] - gamma[fn + 1])) - // U = Tinv @ (beta * V) - simdgroup_load(V_tile, v_ + dv_idx, Dv * Hv); + load_M(V_tile, v_ + dv_idx, Dv * Hv); SCALE(V_tile, beta_fm) simdgroup_multiply(U_tile, Tinv, V_tile); - // delta = U - WS SUB(delta_tile, U_tile, WS_tile) tmp_tile = make_filled_simdgroup_matrix(0.f); QKt_tile = make_filled_simdgroup_matrix(0.f); for (int kk = 0; kk < Dk; kk += 8) { - simdgroup_load(Q_tile, q_ + kk, Hk * Dk); - simdgroup_load(K_tile, k_ + kk, Hk * Dk, ulong2(0, 0), true); - - // SCALE(Q_tile, gamma[fm]) + load_M(Q_tile, q_ + kk, Hk * Dk); + load_MT(K_tile, k_ + kk, Hk * Dk); simdgroup_multiply_accumulate(QKt_tile, Q_tile, K_tile, QKt_tile); - // Q_left @ S^T SCALE(Q_tile, metal::exp(gamma[fm])) simdgroup_multiply_accumulate(tmp_tile, Q_tile, S_tile[kk / 8], tmp_tile); } - // (Q @ K) * Gamma, in the paper they use M here but it's probably a typo. SCALE_TRI( QKt_tile, metal::exp(gamma[fm] - gamma[fn]), @@ -725,24 +755,25 @@ template simdgroup_multiply_accumulate(out_tile, QKt_tile, delta_tile, tmp_tile); - y[fm * Hv * Dv + dv_idx + fn] = static_cast(AT(out_tile, 0)); - y[fm * Hv * Dv + dv_idx + fn + 1] = static_cast(AT(out_tile, 1)); + if (fm < valid_rows) { + y[fm * Hv * Dv + dv_idx + fn] = static_cast(AT(out_tile, 0)); + y[fm * Hv * Dv + dv_idx + fn + 1] = static_cast(AT(out_tile, 1)); + } for (int kk = 0; kk < Dk; kk += 8) { - simdgroup_load(K_tile, k_ + kk, Hk * Dk, ulong2(0, 0), true); + load_MT(K_tile, k_ + kk, Hk * Dk); SCALE2( K_tile, metal::exp(gamma[C - 1] - gamma[fn]), metal::exp(gamma[C - 1] - gamma[fn + 1])) - simdgroup_multiply(KD_tile, K_tile, delta_tile); - // SCALE(K_tile, 100000) - - // FMA(S_tile[kk / 8], 1, S_tile[kk / 8], KD_tile) FMA(S_tile[kk / 8], metal::exp(gamma[C - 1]), S_tile[kk / 8], KD_tile) } + }; - // advance pointers + int t = 0; + for (; t + C <= T; t += C) { + process_chunk(C, metal::false_type{}); q_ += C * Hk * Dk; k_ += C * Hk * Dk; v_ += C * Hv * Dv; @@ -750,6 +781,9 @@ template y += C * Hv * Dv; g_ += C * Hv; } + if (t < T) { + process_chunk(short(T - t), metal::true_type{}); + } for (int kk = 0; kk < Dk; kk += 8) { simdgroup_store(S_tile[kk / 8], o_state + kk, Dk, ulong2(0, 0), true); diff --git a/mlx/fast.cpp b/mlx/fast.cpp index d63f424590..36a79d2131 100644 --- a/mlx/fast.cpp +++ b/mlx/fast.cpp @@ -929,13 +929,13 @@ std::vector gated_delta_update( const array& gates, const array& beta_, const std::optional& initial_state, /* = std::nullopt */ - const int C, StreamOrDevice s_ /* = {} */) { // determine output dtype auto s = to_stream(s_); + // TODO fix this auto promoted = promote_types(queries.dtype(), keys.dtype()); - auto out_dtype = issubdtype(promoted, float32) + auto out_dtype = issubdtype(promoted, floating) ? promoted : promote_types(promoted, float32); @@ -962,20 +962,8 @@ std::vector gated_delta_update( int Hv = v.shape(2); int Dv = v.shape(3); - int pad_size = (C != 0) ? (C - T % C) % C : 0; - int T_padded = T > 1 ? T + pad_size : T; - - if (pad_size > 1) { - q = pad(q, {1}, {0}, {pad_size}, array(0.0f, out_dtype), "constant", s); - k = pad(k, {1}, {0}, {pad_size}, array(0.0f, out_dtype), "constant", s); - v = pad(v, {1}, {0}, {pad_size}, array(0.0f, out_dtype), "constant", s); - g = pad(g, {1}, {0}, {pad_size}, array(1.0f, out_dtype), "constant", s); - beta = - pad(beta, {1}, {0}, {pad_size}, array(0.0f, out_dtype), "constant", s); - } - - auto h0 = initial_state.has_value() ? astype(*initial_state, out_dtype, s) - : zeros({B, Hv, Dv, Dk}, out_dtype, s); + auto h0 = initial_state.has_value() ? astype(*initial_state, float32, s) + : zeros({B, Hv, Dv, Dk}, float32, s); auto fallback = [B, T, Hk, Dk, Hv, Dv, s](std::vector inputs) { auto q = astype(inputs[0], float32, s); @@ -1043,24 +1031,16 @@ std::vector gated_delta_update( if (!GatedDeltaUpdate::use_fallback(s)) { auto result = array::make_arrays( - /* output shapes */ {{B, T_padded, Hv, Dv}, {B, Hv, Dv, Dk}}, - /* dtypes */ {out_dtype, out_dtype}, + /* output shapes */ {{B, T, Hv, Dv}, {B, Hv, Dv, Dk}}, + /* dtypes */ {out_dtype, float32}, /* primitive */ - std::make_shared(s, fallback, C), + std::make_shared(s, fallback), /* inputs */ {q, k, v, g, beta, h0}); - // slice output back to original T - if (pad_size > 0) { - result[0] = slice(result[0], {0, 0, 0, 0}, {B, T, Hv, Dv}, s); - } - return result; } auto result = fallback({q, k, v, g, beta, h0}); - if (pad_size > 0) { - result[0] = slice(result[0], {0, 0, 0, 0}, {B, T, Hv, Dv}, s); - } return result; } diff --git a/mlx/fast.h b/mlx/fast.h index 24d17bdae1..81ebf6d953 100644 --- a/mlx/fast.h +++ b/mlx/fast.h @@ -62,7 +62,6 @@ MLX_API std::vector gated_delta_update( const array& gates, const array& beta_, const std::optional& initial_state = std::nullopt, - const int C = 8, StreamOrDevice s = {}); using TemplateArg = std::variant; diff --git a/mlx/fast_primitives.h b/mlx/fast_primitives.h index 79ce867d82..62eaa7c51e 100644 --- a/mlx/fast_primitives.h +++ b/mlx/fast_primitives.h @@ -329,9 +329,8 @@ class GatedDeltaUpdate : public Custom { public: GatedDeltaUpdate( Stream stream, - std::function(std::vector)> fallback, - int C) - : Custom(stream, std::move(fallback)), chunk_size(C) {} + std::function(std::vector)> fallback) + : Custom(stream, std::move(fallback)) {} static bool use_fallback( /* TODO */ @@ -355,7 +354,6 @@ class GatedDeltaUpdate : public Custom { } private: - int chunk_size; }; class Quantize : public Custom { diff --git a/python/src/fast.cpp b/python/src/fast.cpp index 26db40cff2..925f555a0b 100644 --- a/python/src/fast.cpp +++ b/python/src/fast.cpp @@ -342,7 +342,6 @@ void init_fast(nb::module_& parent_module) { "gamma"_a, "beta"_a, "initial_state"_a = nb::none(), // optional, defaults to None - "C"_a = 0, // optional, defaults to 0 "stream"_a = nb::none(), // optional, defaults to None R"( Chunked gated delta network forward pass. From 4a292527c99dc4b6ff48301d11d09e78fd924b5e Mon Sep 17 00:00:00 2001 From: tpegolotti Date: Tue, 4 Aug 2026 15:26:51 +0200 Subject: [PATCH 34/56] remove log space for C=8 and force cleanup inline --- .../metal/kernels/gated_delta_update_impl.h | 59 +++++++++---------- 1 file changed, 29 insertions(+), 30 deletions(-) diff --git a/mlx/backend/metal/kernels/gated_delta_update_impl.h b/mlx/backend/metal/kernels/gated_delta_update_impl.h index a7f4627c2b..1484df7db0 100644 --- a/mlx/backend/metal/kernels/gated_delta_update_impl.h +++ b/mlx/backend/metal/kernels/gated_delta_update_impl.h @@ -460,7 +460,8 @@ template } mlx::steel::NAXTile TMP_tile; - auto process_chunk = [&](const short valid_rows, auto bounded_tag) { + auto process_chunk = [&](const short valid_rows, + auto bounded_tag) __attribute__((always_inline)) { constexpr bool B = decltype(bounded_tag)::value; auto load_seq = [&](thread auto& tile, auto src, int ld) { @@ -659,7 +660,9 @@ template simdgroup_load(S_tile[kk / 8], i_state + kk, Dk, ulong2(0, 0), true); } - auto process_chunk = [&](const short valid_rows, auto bounded_tag) { + auto process_chunk = [&](thread simdgroup_float8x8* S_tile, + const short valid_rows, + auto bounded_tag) __attribute__((always_inline)) { constexpr bool B = decltype(bounded_tag)::value; // non-transposed load @@ -682,18 +685,22 @@ template } }; - float g_val = (thread_index_in_simdgroup < (uint)valid_rows) - ? metal::log( - metal::max(g_[thread_index_in_simdgroup * Hv + hv_idx], 1e-6f)) - : 0.0f; - float gamma_val = simd_prefix_inclusive_sum(g_val); + float g_val = (thread_index_in_simdgroup < C) + ? g_[thread_index_in_simdgroup * Hv + hv_idx] + : 1.0f; + + float gamma_val = simd_prefix_inclusive_product(g_val); + if (thread_index_in_simdgroup < C) { gamma[thread_index_in_simdgroup] = gamma_val; } + gamma[C - 1] = metal::max(gamma[C - 1], 1e-9f); + float beta_fm = (fm < valid_rows) ? beta_[fm * Hv + hv_idx] : 0.0f; KKt_tile = make_filled_simdgroup_matrix(0.f); + STEEL_PRAGMA_UNROLL for (int kk = 0; kk < Dk; kk += 8) { load_M(K_tile, k_ + kk, Dk * Hk); load_MT(KT_tile, k_ + kk, Dk * Hk); @@ -701,37 +708,30 @@ template } KKtK_tile = KKt_tile; - KKtV_tile = KKt_tile; - SCALE_TRIEQ(KKtK_tile, beta_fm, beta_fm) - SCALE_TRIEQ( - KKtV_tile, - beta_fm * (metal::exp(gamma[fm] - gamma[fn])), - beta_fm * (metal::exp(gamma[fm] - gamma[fn + 1]))) simdgroup_float8x8 Tinv, P; AT(P, 0) = AT(KKtK_tile, 0); AT(P, 1) = AT(KKtK_tile, 1); SUB(Tinv, I_tile, KKtK_tile) + STEEL_PRAGMA_UNROLL for (int step = 1; (1 << step) < C; step++) { simdgroup_multiply(P, P, P); simdgroup_multiply_accumulate(Tinv, Tinv, P, Tinv); } WS_tile = make_filled_simdgroup_matrix(0.f); + STEEL_PRAGMA_UNROLL for (int kk = 0; kk < Dk; kk += 8) { load_M(K_tile, k_ + kk, Dk * Hk); SCALE(K_tile, beta_fm) simdgroup_multiply(W_tile, Tinv, K_tile); - SCALE(W_tile, metal::exp(gamma[fm])) + SCALE(W_tile, gamma[fm]) simdgroup_multiply_accumulate(WS_tile, W_tile, S_tile[kk / 8], WS_tile); } - SCALE_TRI( - Tinv, - metal::exp(gamma[fm] - gamma[fn]), - metal::exp(gamma[fm] - gamma[fn + 1])) + SCALE_TRI(Tinv, (gamma[fm] / gamma[fn]), (gamma[fm] / gamma[fn + 1])) load_M(V_tile, v_ + dv_idx, Dv * Hv); SCALE(V_tile, beta_fm) @@ -740,18 +740,16 @@ template tmp_tile = make_filled_simdgroup_matrix(0.f); QKt_tile = make_filled_simdgroup_matrix(0.f); + STEEL_PRAGMA_UNROLL for (int kk = 0; kk < Dk; kk += 8) { load_M(Q_tile, q_ + kk, Hk * Dk); load_MT(K_tile, k_ + kk, Hk * Dk); + SCALE(Q_tile, gamma[fm]) simdgroup_multiply_accumulate(QKt_tile, Q_tile, K_tile, QKt_tile); - SCALE(Q_tile, metal::exp(gamma[fm])) simdgroup_multiply_accumulate(tmp_tile, Q_tile, S_tile[kk / 8], tmp_tile); } - SCALE_TRI( - QKt_tile, - metal::exp(gamma[fm] - gamma[fn]), - metal::exp(gamma[fm] - gamma[fn + 1])) + SCALE_TRI(QKt_tile, (1.0f / gamma[fn]), (1.0f / gamma[fn + 1])) simdgroup_multiply_accumulate(out_tile, QKt_tile, delta_tile, tmp_tile); @@ -760,20 +758,20 @@ template y[fm * Hv * Dv + dv_idx + fn + 1] = static_cast(AT(out_tile, 1)); } + STEEL_PRAGMA_UNROLL for (int kk = 0; kk < Dk; kk += 8) { load_MT(K_tile, k_ + kk, Hk * Dk); - SCALE2( - K_tile, - metal::exp(gamma[C - 1] - gamma[fn]), - metal::exp(gamma[C - 1] - gamma[fn + 1])) + SCALE2(K_tile, (gamma[C - 1] / gamma[fn]), (gamma[C - 1] / gamma[fn + 1])) + simdgroup_multiply(KD_tile, K_tile, delta_tile); - FMA(S_tile[kk / 8], metal::exp(gamma[C - 1]), S_tile[kk / 8], KD_tile) + FMA(S_tile[kk / 8], gamma[C - 1], S_tile[kk / 8], KD_tile) } }; int t = 0; + STEEL_PRAGMA_UNROLL for (; t + C <= T; t += C) { - process_chunk(C, metal::false_type{}); + process_chunk(S_tile, C, metal::false_type{}); q_ += C * Hk * Dk; k_ += C * Hk * Dk; v_ += C * Hv * Dv; @@ -782,9 +780,10 @@ template g_ += C * Hv; } if (t < T) { - process_chunk(short(T - t), metal::true_type{}); + process_chunk(S_tile, short(T - t), metal::true_type{}); } + STEEL_PRAGMA_UNROLL for (int kk = 0; kk < Dk; kk += 8) { simdgroup_store(S_tile[kk / 8], o_state + kk, Dk, ulong2(0, 0), true); } From d86cd3e34adc1338aaba27bda05703126a85c7cd Mon Sep 17 00:00:00 2001 From: tpegolotti Date: Tue, 4 Aug 2026 06:57:42 -0700 Subject: [PATCH 35/56] update typing --- mlx/backend/metal/gated_delta_update.cpp | 7 +- .../metal/kernels/gated_delta_update.metal | 90 ++++++++----------- .../metal/kernels/gated_delta_update_impl.h | 51 ++++++----- 3 files changed, 70 insertions(+), 78 deletions(-) diff --git a/mlx/backend/metal/gated_delta_update.cpp b/mlx/backend/metal/gated_delta_update.cpp index 264446ac9b..1e6cf14e71 100644 --- a/mlx/backend/metal/gated_delta_update.cpp +++ b/mlx/backend/metal/gated_delta_update.cpp @@ -91,10 +91,9 @@ void GatedDeltaUpdate::eval_gpu( const char* chunk_env = std::getenv("GATED_DELTA_CHUNK"); C = chunk_env ? std::stoi(chunk_env) : C; - std::string suffix = get_type_string(q.dtype()) // "float" - + "_" + get_type_string(h0.dtype()) // "float" - + "_" + std::to_string(Dk) + "_" + std::to_string(Dv) + "_" + - std::to_string(Hk) + "_" + std::to_string(Hv); + std::string suffix = get_type_string(q.dtype()) + "_" + std::to_string(Dk) + + "_" + std::to_string(Dv) + "_" + std::to_string(Hk) + "_" + + std::to_string(Hv); auto& compute_encoder = metal::get_command_encoder(s); diff --git a/mlx/backend/metal/kernels/gated_delta_update.metal b/mlx/backend/metal/kernels/gated_delta_update.metal index 0d41a100a3..ac02b7fadb 100644 --- a/mlx/backend/metal/kernels/gated_delta_update.metal +++ b/mlx/backend/metal/kernels/gated_delta_update.metal @@ -3,78 +3,66 @@ using namespace metal; -#define instantiate_gated_delta_update_seq(in_type, st_type, dk, dv, hk, hv) \ - instantiate_kernel( \ - "seq_gated_delta_" #in_type "_" #st_type "_" #dk "_" #dv "_" #hk \ - "_" #hv, \ - gated_delta_seq, \ - in_type, \ - st_type, \ - dk, \ - dv, \ - hk, \ +#define instantiate_gated_delta_update_seq(in_type, dk, dv, hk, hv) \ + instantiate_kernel( \ + "seq_gated_delta_" #in_type "_" #dk "_" #dv "_" #hk "_" #hv, \ + gated_delta_seq, \ + in_type, \ + dk, \ + dv, \ + hk, \ hv) -#define instantiate_gated_delta_update_seq_dims(in_type, st_type) \ - instantiate_gated_delta_update_seq(in_type, st_type, 128, 128, 24, 24) \ - instantiate_gated_delta_update_seq(in_type, st_type, 128, 128, 32, 32) \ - instantiate_gated_delta_update_seq( \ - in_type, st_type, 128, 128, 16, 32) \ - instantiate_gated_delta_update_seq( \ - in_type, st_type, 128, 128, 16, 48) +#define instantiate_gated_delta_update_seq_dims(in_type) \ + instantiate_gated_delta_update_seq(in_type, 128, 128, 24, 24) \ + instantiate_gated_delta_update_seq(in_type, 128, 128, 32, 32) \ + instantiate_gated_delta_update_seq(in_type, 128, 128, 16, 32) \ + instantiate_gated_delta_update_seq(in_type, 128, 128, 16, 48) -#define instantiate_gated_delta_update_fused_chunk( \ - in_type, st_type, dk, dv, hk, hv, c) \ +#define instantiate_gated_delta_update_fused_chunk(in_type, dk, dv, hk, hv, c) \ instantiate_kernel( \ - "gated_delta_fused_chunk_" #in_type "_" #st_type "_" #dk "_" #dv "_" #hk \ - "_" #hv "_" #c, \ + "gated_delta_fused_chunk_" #in_type "_" #dk "_" #dv "_" #hk "_" #hv \ + "_" #c, \ gated_delta_fused_chunk, \ in_type, \ - st_type, \ dk, \ dv, \ hk, \ hv, \ c) -#define instantiate_gated_delta_update_fused_chunk_dims(in_type, st_type) \ - instantiate_gated_delta_update_fused_chunk( \ - in_type, st_type, 128, 128, 16, 32, 8) \ - instantiate_gated_delta_update_fused_chunk( \ - in_type, st_type, 128, 128, 16, 48, 8) \ - instantiate_gated_delta_update_fused_chunk( \ - in_type, st_type, 128, 128, 24, 24, 8) \ - instantiate_gated_delta_update_fused_chunk( \ - in_type, st_type, 128, 128, 32, 32, 8) +#define instantiate_gated_delta_update_fused_chunk_dims(in_type) \ + instantiate_gated_delta_update_fused_chunk(in_type, 128, 128, 16, 32, 8) \ + instantiate_gated_delta_update_fused_chunk(in_type, 128, 128, 16, 48, 8) \ + instantiate_gated_delta_update_fused_chunk( \ + in_type, 128, 128, 24, 24, 8) \ + instantiate_gated_delta_update_fused_chunk( \ + in_type, 128, 128, 32, 32, 8) -#define instantiate_gated_delta_update_fused_nax( \ - in_type, st_type, dk, dv, hk, hv, c) \ +#define instantiate_gated_delta_update_fused_nax(in_type, dk, dv, hk, hv, c) \ instantiate_kernel( \ - "gated_delta_fused_nax_" #in_type "_" #st_type "_" #dk "_" #dv "_" #hk \ - "_" #hv "_" #c, \ + "gated_delta_fused_nax_" #in_type "_" #dk "_" #dv "_" #hk "_" #hv \ + "_" #c, \ gated_delta_fused_nax, \ in_type, \ - st_type, \ dk, \ dv, \ hk, \ hv, \ c) -#define instantiate_gated_delta_update_fused_nax_dims(in_type, st_type) \ - instantiate_gated_delta_update_fused_nax( \ - in_type, st_type, 128, 128, 16, 32, 16) \ - instantiate_gated_delta_update_fused_nax( \ - in_type, st_type, 128, 128, 16, 48, 16) \ - instantiate_gated_delta_update_fused_nax( \ - in_type, st_type, 128, 128, 24, 24, 16) \ - instantiate_gated_delta_update_fused_nax( \ - in_type, st_type, 128, 128, 32, 32, 16) +#define instantiate_gated_delta_update_fused_nax_dims(in_type) \ + instantiate_gated_delta_update_fused_nax(in_type, 128, 128, 16, 32, 16) \ + instantiate_gated_delta_update_fused_nax(in_type, 128, 128, 16, 48, 16) \ + instantiate_gated_delta_update_fused_nax( \ + in_type, 128, 128, 24, 24, 16) \ + instantiate_gated_delta_update_fused_nax( \ + in_type, 128, 128, 32, 32, 16) -instantiate_gated_delta_update_seq_dims(float, float); -instantiate_gated_delta_update_fused_chunk_dims(float, float); -instantiate_gated_delta_update_fused_nax_dims(float, float); +instantiate_gated_delta_update_seq_dims(float); +instantiate_gated_delta_update_fused_chunk_dims(float); +instantiate_gated_delta_update_fused_nax_dims(float); -instantiate_gated_delta_update_seq_dims(bfloat16_t, float); -// instantiate_gated_delta_update_fused_chunk_dims(bfloat16_t, float); -instantiate_gated_delta_update_fused_nax_dims(bfloat16_t, float); \ No newline at end of file +instantiate_gated_delta_update_seq_dims(bfloat16_t); +instantiate_gated_delta_update_fused_chunk_dims(bfloat16_t); +instantiate_gated_delta_update_fused_nax_dims(bfloat16_t); \ No newline at end of file diff --git a/mlx/backend/metal/kernels/gated_delta_update_impl.h b/mlx/backend/metal/kernels/gated_delta_update_impl.h index 1484df7db0..9a920194ed 100644 --- a/mlx/backend/metal/kernels/gated_delta_update_impl.h +++ b/mlx/backend/metal/kernels/gated_delta_update_impl.h @@ -385,16 +385,16 @@ METAL_FUNC static constexpr void mman( B.frag_at(0, (BO) + 1), \ metal::bool_constant{}); -template +template [[kernel]] void gated_delta_fused_nax( const device InT* q [[buffer(0)]], const device InT* k [[buffer(1)]], const device InT* v [[buffer(2)]], - const device StT* state_in [[buffer(3)]], + const device float* state_in [[buffer(3)]], const device InT* g [[buffer(4)]], const device InT* beta [[buffer(5)]], device InT* y [[buffer(6)]], - device StT* state_out [[buffer(7)]], + device float* state_out [[buffer(7)]], constant int& T [[buffer(8)]], uint3 thread_position_in_grid [[thread_position_in_grid]], uint3 thread_position_in_threadgroup [[thread_position_in_threadgroup]], @@ -448,7 +448,7 @@ template mlx::steel::NAXTile out_tile; mlx::steel::NAXTile Tinv_tile, P; - mlx::steel::NAXTile KKtK_tile, KKtV_tile, KKt_tile; + mlx::steel::NAXTile KKtK_tile, KKt_tile; mlx::steel::NAXTile I_tile; STEEL_PRAGMA_UNROLL @@ -493,7 +493,6 @@ template } KKtK_tile = KKt_tile; - KKtV_tile = KKt_tile; SCALE_TRIEQ_NAX1(KKtK_tile, beta_fm); MM16x16x16(P, 0, KKtK_tile, false, 0, KKtK_tile, false, 0); @@ -583,16 +582,16 @@ template S_tile.store(o_state, Dk); } -template +template [[kernel]] void gated_delta_fused_chunk( const device InT* q [[buffer(0)]], const device InT* k [[buffer(1)]], const device InT* v [[buffer(2)]], - const device StT* state_in [[buffer(3)]], + const device float* state_in [[buffer(3)]], const device InT* g [[buffer(4)]], const device InT* beta [[buffer(5)]], device InT* y [[buffer(6)]], - device StT* state_out [[buffer(7)]], + device float* state_out [[buffer(7)]], constant int& T [[buffer(8)]], uint3 thread_position_in_grid [[thread_position_in_grid]], uint3 thread_position_in_threadgroup [[thread_position_in_threadgroup]], @@ -646,7 +645,7 @@ template simdgroup_float8x8 KD_tile; // tiles for WY form computation - simdgroup_float8x8 KKtK_tile, KKtV_tile, KKt_tile; + simdgroup_float8x8 KKtK_tile, KKt_tile; threadgroup float gamma_all[C * 4]; threadgroup float* gamma = gamma_all + sg_id * C; @@ -668,35 +667,41 @@ template // non-transposed load auto load_M = [&](thread simdgroup_float8x8& M, auto src, int ld) { if constexpr (B) { - AT(M, 0) = (fm < valid_rows) ? float(src[fm * ld + fn]) : 0.f; - AT(M, 1) = (fm < valid_rows) ? float(src[fm * ld + fn + 1]) : 0.f; + AT(M, 0) = + static_cast((fm < valid_rows) ? (src[fm * ld + fn]) : 0.f); + AT(M, 1) = static_cast( + (fm < valid_rows) ? (src[fm * ld + fn + 1]) : 0.f); } else { - simdgroup_load(M, src, ld); + // simdgroup_load(M, src, ld); + AT(M, 0) = static_cast(src[fm * ld + fn]); + AT(M, 1) = static_cast(src[fm * ld + fn + 1]); } }; // transposed load auto load_MT = [&](thread simdgroup_float8x8& M, auto src, int ld) { if constexpr (B) { - AT(M, 0) = (fn < valid_rows) ? float(src[fn * ld + fm]) : 0.f; - AT(M, 1) = (fn + 1 < valid_rows) ? float(src[(fn + 1) * ld + fm]) : 0.f; + AT(M, 0) = + static_cast((fn < valid_rows) ? (src[fn * ld + fm]) : 0.f); + AT(M, 1) = static_cast( + (fn + 1 < valid_rows) ? (src[(fn + 1) * ld + fm]) : 0.f); } else { - simdgroup_load(M, src, ld, ulong2(0, 0), true); + // simdgroup_load(M, src, ld, ulong2(0, 0), true); + AT(M, 0) = static_cast(src[fn * ld + fm]); + AT(M, 1) = static_cast(src[(fn + 1) * ld + fm]); } }; - float g_val = (thread_index_in_simdgroup < C) + float g_val = (thread_index_in_simdgroup < (uint)valid_rows) ? g_[thread_index_in_simdgroup * Hv + hv_idx] : 1.0f; float gamma_val = simd_prefix_inclusive_product(g_val); - if (thread_index_in_simdgroup < C) { + if (thread_index_in_simdgroup < (uint)valid_rows) { gamma[thread_index_in_simdgroup] = gamma_val; } - gamma[C - 1] = metal::max(gamma[C - 1], 1e-9f); - float beta_fm = (fm < valid_rows) ? beta_[fm * Hv + hv_idx] : 0.0f; KKt_tile = make_filled_simdgroup_matrix(0.f); @@ -793,17 +798,17 @@ template auto grid = MTL::Size(32, Dv, B * Hv); auto threads = MTL::Size(32, 4, 1); */ -template +template [[kernel]] void gated_delta_seq( const device InT* q [[buffer(0)]], const device InT* k [[buffer(1)]], const device InT* v [[buffer(2)]], const device InT* g [[buffer(3)]], // [B, T, Hv] or [B, T, Hv, Dk] const device InT* beta [[buffer(4)]], // [B, T, Hv] - const device StT* state_in [[buffer(5)]], // [B, Hv, Dv, Dk] + const device float* state_in [[buffer(5)]], // [B, Hv, Dv, Dk] constant int& T [[buffer(6)]], device InT* y [[buffer(7)]], // [B, T, Hv, Dv] - device StT* state_out [[buffer(8)]], // [B, Hv, Dv, Dk] + device float* state_out [[buffer(8)]], // [B, Hv, Dv, Dk] uint3 thread_position_in_grid [[thread_position_in_grid]], uint3 thread_position_in_threadgroup [[thread_position_in_threadgroup]], uint thread_index_in_simdgroup [[thread_index_in_simdgroup]]) { @@ -870,6 +875,6 @@ template } for (int i = 0; i < n_per_t; ++i) { auto s_idx = n_per_t * dk_idx + i; - o_state[s_idx] = static_cast(state[i]); + o_state[s_idx] = static_cast(state[i]); } } \ No newline at end of file From 3471ad169c0aaf98f307ed92aebb5df360533463 Mon Sep 17 00:00:00 2001 From: tpegolotti Date: Tue, 4 Aug 2026 07:52:54 -0700 Subject: [PATCH 36/56] added back log space --- .../metal/kernels/gated_delta_update_impl.h | 86 +++++++++---------- 1 file changed, 40 insertions(+), 46 deletions(-) diff --git a/mlx/backend/metal/kernels/gated_delta_update_impl.h b/mlx/backend/metal/kernels/gated_delta_update_impl.h index 9a920194ed..2363866e54 100644 --- a/mlx/backend/metal/kernels/gated_delta_update_impl.h +++ b/mlx/backend/metal/kernels/gated_delta_update_impl.h @@ -91,14 +91,14 @@ using namespace mpp::tensor_ops; } \ } -#define SCALE_ROW_NAX(TILE0, S) \ - { \ - STEEL_PRAGMA_UNROLL \ - for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ - const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; \ - AT_NAX(TILE0, _i) *= \ - metal::exp((S)[mlx::steel::BaseNAXFrag::get_coord(_w).y]); \ - } \ +#define SCALE_ROW_NAX(TILE0, S) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; \ + AT_NAX(TILE0, _i) *= \ + metal::fast::exp((S)[mlx::steel::BaseNAXFrag::get_coord(_w).y]); \ + } \ } #define SCALE_BETA_NAX(TILE0, BETA2) \ @@ -110,14 +110,14 @@ using namespace mpp::tensor_ops; } \ } -#define SCALE2_NAX(TILE0, GAMMA) \ - { \ - STEEL_PRAGMA_UNROLL \ - for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ - const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; \ - const short _fm = mlx::steel::BaseNAXFrag::get_coord(_w).y; \ - AT_NAX(TILE0, _i) *= metal::exp((GAMMA)[(C) - 1] - (GAMMA)[_fm]); \ - } \ +#define SCALE2_NAX(TILE0, GAMMA) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; \ + const short _fm = mlx::steel::BaseNAXFrag::get_coord(_w).y; \ + AT_NAX(TILE0, _i) *= metal::fast::exp((GAMMA)[(C) - 1] - (GAMMA)[_fm]); \ + } \ } #define SCALE_TRI_NAX(TILE0, GAMMA) \ @@ -125,8 +125,9 @@ using namespace mpp::tensor_ops; STEEL_PRAGMA_UNROLL \ for (short _i = 0; _i < decltype(TILE0)::kElemsPerFrag; _i++) { \ const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); /* {fn, fm} */ \ - AT_NAX(TILE0, _i) *= \ - (_c.x > _c.y) ? 0.f : metal::exp((GAMMA)[_c.y] - (GAMMA)[_c.x]); \ + AT_NAX(TILE0, _i) *= (_c.x > _c.y) \ + ? 0.f \ + : metal::fast::exp((GAMMA)[_c.y] - (GAMMA)[_c.x]); \ } \ } @@ -139,19 +140,6 @@ using namespace mpp::tensor_ops; } \ } -#define SCALE_TRIEQ_NAX(TILE0, BETA, GAMMA) \ - { \ - STEEL_PRAGMA_UNROLL \ - for (short _i = 0; _i < decltype(TILE0)::kElemsPerFrag; _i++) { \ - const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); /* {fn, fm} */ \ - const short _fn = _c.x; \ - const short _fm = _c.y; \ - const float _s = \ - (BETA)[_i >> 2] * metal::exp((GAMMA)[_fm] - (GAMMA)[_fn]); \ - AT_NAX(TILE0, _i) *= (_fn >= _fm) ? 0.f : _s; \ - } \ - } - namespace mlx { namespace steel { template < @@ -472,11 +460,10 @@ template } }; - auto g_val = (thread_index_in_simdgroup < (uint)valid_rows) - ? metal::log( - metal::clamp( - g_[thread_index_in_simdgroup * Hv + hv_idx], 1e-6f, 1.0f)) + float g_val = (thread_index_in_simdgroup < (uint)valid_rows) + ? metal::fast::log(g_[thread_index_in_simdgroup * Hv + hv_idx]) : 0.0f; + auto gamma_val = simd_prefix_inclusive_sum(g_val); if (thread_index_in_simdgroup < C) { gamma[thread_index_in_simdgroup] = static_cast(gamma_val); @@ -556,7 +543,7 @@ template } } - SCALE_NAX(S_tile, metal::exp(gamma[C - 1])); + SCALE_NAX(S_tile, metal::fast::exp(gamma[C - 1])); for (int kk = 0; kk < Dk; kk += 32) { load_seq(K_tile, k_ + kk, Hk * Dk); @@ -693,15 +680,22 @@ template }; float g_val = (thread_index_in_simdgroup < (uint)valid_rows) - ? g_[thread_index_in_simdgroup * Hv + hv_idx] - : 1.0f; + ? metal::fast::log(g_[thread_index_in_simdgroup * Hv + hv_idx]) + : 0.0f; - float gamma_val = simd_prefix_inclusive_product(g_val); + float gamma_val = simd_prefix_inclusive_sum(g_val); - if (thread_index_in_simdgroup < (uint)valid_rows) { + if (thread_index_in_simdgroup < C) { gamma[thread_index_in_simdgroup] = gamma_val; } + float gamma_fm = metal::fast::exp(gamma[fm]); + float gamma_fmdfn = metal::fast::exp(gamma[fm] - gamma[fn]); + float gamma_fmdfn1 = metal::fast::exp(gamma[fm] - gamma[fn + 1]); + float gamma_Cdfn = metal::fast::exp(gamma[C - 1] - gamma[fn]); + float gamma_Cdfn1 = metal::fast::exp(gamma[C - 1] - gamma[fn + 1]); + float gamma_C = metal::fast::exp(gamma[C - 1]); + float beta_fm = (fm < valid_rows) ? beta_[fm * Hv + hv_idx] : 0.0f; KKt_tile = make_filled_simdgroup_matrix(0.f); @@ -732,11 +726,11 @@ template load_M(K_tile, k_ + kk, Dk * Hk); SCALE(K_tile, beta_fm) simdgroup_multiply(W_tile, Tinv, K_tile); - SCALE(W_tile, gamma[fm]) + SCALE(W_tile, gamma_fm) simdgroup_multiply_accumulate(WS_tile, W_tile, S_tile[kk / 8], WS_tile); } - SCALE_TRI(Tinv, (gamma[fm] / gamma[fn]), (gamma[fm] / gamma[fn + 1])) + SCALE_TRI(Tinv, gamma_fmdfn, gamma_fmdfn1) load_M(V_tile, v_ + dv_idx, Dv * Hv); SCALE(V_tile, beta_fm) @@ -749,12 +743,12 @@ template for (int kk = 0; kk < Dk; kk += 8) { load_M(Q_tile, q_ + kk, Hk * Dk); load_MT(K_tile, k_ + kk, Hk * Dk); - SCALE(Q_tile, gamma[fm]) simdgroup_multiply_accumulate(QKt_tile, Q_tile, K_tile, QKt_tile); + SCALE(Q_tile, gamma_fm) simdgroup_multiply_accumulate(tmp_tile, Q_tile, S_tile[kk / 8], tmp_tile); } - SCALE_TRI(QKt_tile, (1.0f / gamma[fn]), (1.0f / gamma[fn + 1])) + SCALE_TRI(QKt_tile, gamma_fmdfn, gamma_fmdfn1) simdgroup_multiply_accumulate(out_tile, QKt_tile, delta_tile, tmp_tile); @@ -766,10 +760,10 @@ template STEEL_PRAGMA_UNROLL for (int kk = 0; kk < Dk; kk += 8) { load_MT(K_tile, k_ + kk, Hk * Dk); - SCALE2(K_tile, (gamma[C - 1] / gamma[fn]), (gamma[C - 1] / gamma[fn + 1])) + SCALE2(K_tile, gamma_Cdfn, gamma_Cdfn1) simdgroup_multiply(KD_tile, K_tile, delta_tile); - FMA(S_tile[kk / 8], gamma[C - 1], S_tile[kk / 8], KD_tile) + FMA(S_tile[kk / 8], gamma_C, S_tile[kk / 8], KD_tile) } }; From ddeebd65312755f8da20bce48fd242aa617a9ca0 Mon Sep 17 00:00:00 2001 From: tpegolotti Date: Tue, 4 Aug 2026 08:27:20 -0700 Subject: [PATCH 37/56] cleanup .metal file --- .../metal/kernels/gated_delta_update.metal | 39 ++++++------------- 1 file changed, 12 insertions(+), 27 deletions(-) diff --git a/mlx/backend/metal/kernels/gated_delta_update.metal b/mlx/backend/metal/kernels/gated_delta_update.metal index ac02b7fadb..a5bd7e5442 100644 --- a/mlx/backend/metal/kernels/gated_delta_update.metal +++ b/mlx/backend/metal/kernels/gated_delta_update.metal @@ -13,12 +13,6 @@ using namespace metal; hk, \ hv) -#define instantiate_gated_delta_update_seq_dims(in_type) \ - instantiate_gated_delta_update_seq(in_type, 128, 128, 24, 24) \ - instantiate_gated_delta_update_seq(in_type, 128, 128, 32, 32) \ - instantiate_gated_delta_update_seq(in_type, 128, 128, 16, 32) \ - instantiate_gated_delta_update_seq(in_type, 128, 128, 16, 48) - #define instantiate_gated_delta_update_fused_chunk(in_type, dk, dv, hk, hv, c) \ instantiate_kernel( \ "gated_delta_fused_chunk_" #in_type "_" #dk "_" #dv "_" #hk "_" #hv \ @@ -31,14 +25,6 @@ using namespace metal; hv, \ c) -#define instantiate_gated_delta_update_fused_chunk_dims(in_type) \ - instantiate_gated_delta_update_fused_chunk(in_type, 128, 128, 16, 32, 8) \ - instantiate_gated_delta_update_fused_chunk(in_type, 128, 128, 16, 48, 8) \ - instantiate_gated_delta_update_fused_chunk( \ - in_type, 128, 128, 24, 24, 8) \ - instantiate_gated_delta_update_fused_chunk( \ - in_type, 128, 128, 32, 32, 8) - #define instantiate_gated_delta_update_fused_nax(in_type, dk, dv, hk, hv, c) \ instantiate_kernel( \ "gated_delta_fused_nax_" #in_type "_" #dk "_" #dv "_" #hk "_" #hv \ @@ -51,18 +37,17 @@ using namespace metal; hv, \ c) -#define instantiate_gated_delta_update_fused_nax_dims(in_type) \ - instantiate_gated_delta_update_fused_nax(in_type, 128, 128, 16, 32, 16) \ - instantiate_gated_delta_update_fused_nax(in_type, 128, 128, 16, 48, 16) \ - instantiate_gated_delta_update_fused_nax( \ - in_type, 128, 128, 24, 24, 16) \ - instantiate_gated_delta_update_fused_nax( \ - in_type, 128, 128, 32, 32, 16) +#define instantiate_gated_delta_dims(in_type, dk, dv, hk, hv) \ + instantiate_gated_delta_update_seq(in_type, dk, dv, hk, hv) \ + instantiate_gated_delta_update_fused_chunk(in_type, dk, dv, hk, hv, 8) \ + instantiate_gated_delta_update_fused_nax( \ + in_type, dk, dv, hk, hv, 16) -instantiate_gated_delta_update_seq_dims(float); -instantiate_gated_delta_update_fused_chunk_dims(float); -instantiate_gated_delta_update_fused_nax_dims(float); +#define instantiate_gated_delta(in_type) \ + instantiate_gated_delta_dims(in_type, 128, 128, 24, 24) \ + instantiate_gated_delta_dims(in_type, 128, 128, 32, 32) \ + instantiate_gated_delta_dims(in_type, 128, 128, 16, 32) \ + instantiate_gated_delta_dims(in_type, 128, 128, 16, 48) -instantiate_gated_delta_update_seq_dims(bfloat16_t); -instantiate_gated_delta_update_fused_chunk_dims(bfloat16_t); -instantiate_gated_delta_update_fused_nax_dims(bfloat16_t); \ No newline at end of file +instantiate_gated_delta(float); +instantiate_gated_delta(bfloat16_t); \ No newline at end of file From d76b0e95dd2417e40a4286752c86347687be2416 Mon Sep 17 00:00:00 2001 From: tpegolotti Date: Wed, 5 Aug 2026 00:25:34 -0700 Subject: [PATCH 38/56] Add clamping back --- mlx/backend/metal/kernels/gated_delta_update_impl.h | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/mlx/backend/metal/kernels/gated_delta_update_impl.h b/mlx/backend/metal/kernels/gated_delta_update_impl.h index 2363866e54..742d221f9e 100644 --- a/mlx/backend/metal/kernels/gated_delta_update_impl.h +++ b/mlx/backend/metal/kernels/gated_delta_update_impl.h @@ -461,7 +461,8 @@ template }; float g_val = (thread_index_in_simdgroup < (uint)valid_rows) - ? metal::fast::log(g_[thread_index_in_simdgroup * Hv + hv_idx]) + ? metal::fast::log( + metal::max(g_[thread_index_in_simdgroup * Hv + hv_idx], 1e-6)) : 0.0f; auto gamma_val = simd_prefix_inclusive_sum(g_val); @@ -680,7 +681,8 @@ template }; float g_val = (thread_index_in_simdgroup < (uint)valid_rows) - ? metal::fast::log(g_[thread_index_in_simdgroup * Hv + hv_idx]) + ? metal::fast::log( + metal::max(g_[thread_index_in_simdgroup * Hv + hv_idx], 1e-6)) : 0.0f; float gamma_val = simd_prefix_inclusive_sum(g_val); From 5913f773af2277aa1f7a131541675aea93ff2dcd Mon Sep 17 00:00:00 2001 From: tpegolotti Date: Wed, 5 Aug 2026 04:20:56 -0700 Subject: [PATCH 39/56] Back to Horner for NAX --- mlx/backend/metal/kernels/gated_delta_update_impl.h | 13 ++++--------- 1 file changed, 4 insertions(+), 9 deletions(-) diff --git a/mlx/backend/metal/kernels/gated_delta_update_impl.h b/mlx/backend/metal/kernels/gated_delta_update_impl.h index 742d221f9e..35ac09210a 100644 --- a/mlx/backend/metal/kernels/gated_delta_update_impl.h +++ b/mlx/backend/metal/kernels/gated_delta_update_impl.h @@ -483,18 +483,13 @@ template KKtK_tile = KKt_tile; SCALE_TRIEQ_NAX1(KKtK_tile, beta_fm); - MM16x16x16(P, 0, KKtK_tile, false, 0, KKtK_tile, false, 0); - - ADD_NAX(Tinv_tile, I_tile, P); + ADD_NAX(Tinv_tile, I_tile, KKtK_tile); STEEL_PRAGMA_UNROLL - for (int step = 0; step < 6; step++) { - MM16x16x16(TMP_tile, 0, P, false, 0, Tinv_tile, false, 0); - ADD_NAX(Tinv_tile, I_tile, TMP_tile); + for (int step = 0; step < 15; step++) { + MM16x16x16(TMP_tile, 0, KKtK_tile, false, 0, Tinv_tile, false, 0); + SUB_NAX(Tinv_tile, I_tile, TMP_tile); } - MM16x16x16(TMP_tile, 0, KKtK_tile, false, 0, Tinv_tile, false, 0); - SUB_NAX(Tinv_tile, Tinv_tile, TMP_tile); - STEEL_PRAGMA_UNROLL for (short nn = 0; nn < Dk / 16; nn += 2) { load_seq(K_tile, k_ + nn * 16, Dk * Hk); From 486889d78574cae742c73b8ae834ff220cac5954 Mon Sep 17 00:00:00 2001 From: tpegolotti Date: Wed, 5 Aug 2026 15:51:38 +0200 Subject: [PATCH 40/56] Add test and benchmark --- benchmarks/python/gated_delta_bench.py | 280 +++++------------------ mlx/backend/metal/gated_delta_update.cpp | 3 + python/tests/test_fast_gated_delta.py | 180 +++++++++++++++ 3 files changed, 246 insertions(+), 217 deletions(-) create mode 100644 python/tests/test_fast_gated_delta.py diff --git a/benchmarks/python/gated_delta_bench.py b/benchmarks/python/gated_delta_bench.py index 518198d3b8..128c86c748 100644 --- a/benchmarks/python/gated_delta_bench.py +++ b/benchmarks/python/gated_delta_bench.py @@ -1,4 +1,5 @@ import argparse +import csv import itertools import os import time @@ -6,294 +7,139 @@ from typing import Optional, Tuple import mlx.core as mx - -N_warmup = 5 -N_iter_bench = 40 -N_iter_func = 8 - -os.environ["MTL_CAPTURE_ENABLED"] = "1" +import numpy as np RED_BOLD = "\033[1;31m" GREEN = "\033[0;32m" RESET = "\033[0m" +N_warmup = 8 +N_iter_bench = 80 +N_iter_func = 5 + + +# similar to ./blas/bench_gemm.py def bench(f, *args): - for i in range(N_warmup): + for _ in range(N_warmup): f(*args) + mx.synchronize() s = time.perf_counter_ns() - for i in range(N_iter_bench): + for _ in range(N_iter_bench): f(*args) + mx.synchronize() e = time.perf_counter_ns() - return (e - s) * 1e-9 + return (e - s) * 1e-9 # total seconds for N_iter_bench * N_iter_func calls def do_kernel_bench(f, *args): - q_out = args[0] - - for i in range(N_iter_func): + ys = [] + for _ in range(N_iter_func): out, hf = f(*args) - - mx.eval(out, hf) - return q_out - - -def profile(f, *args): - C = args[-1] - now = datetime.now() - timestamp = now.strftime("%Y_%m_%d_%H_%M") - trace_file = f"traces/mlx_trace_{C}_{timestamp}.gputrace" - - mx.metal.start_capture(trace_file) - - for i in range(N_iter_func): - f(*args) - - mx.metal.stop_capture() - print(f"Writing trace: {trace_file}") - - -def gated_delta_ref( - q: mx.array, # [B, T, Hk, Dk] - k: mx.array, # [B, T, Hk, Dk] - v: mx.array, # [B, T, Hv, Dv] - g: mx.array, # [B, T, Hv] or [B, T, Hv, Dk] - beta: mx.array, # [B, T, Hv] - state: Optional[mx.array] = None, # [B, Hv, Dv, Dk] -) -> Tuple[mx.array, mx.array]: - """ - Implements: - S_t = a_t S_{t-1} + b_t (v_t - a_t S_{t-1} k_t) k_t^T - o_t = S_t q_t - """ - B, T, Hk, Dk = q.shape - Dv = v.shape[-1] - Hv = v.shape[-2] - - if state is None: - state = mx.zeros((B, H, Dv, Dk), dtype=mx.float32) - - if (repeat_factor := Hv // Hk) > 1: - q = mx.repeat(q, repeat_factor, -2) - k = mx.repeat(k, repeat_factor, -2) - - outputs = [] - for t in range(T): - q_t = q[:, t] # [B, H, Dk] - k_t = k[:, t] # [B, H, Dk] - v_t = v[:, t] # [B, H, Dv] - g_t = g[:, t] # [B, H] or [B, H, Dk] - beta_t = beta[:, t] # [B, H] - - # decay - if g_t.ndim == 2: - decay = g_t[..., None, None] # [B, H, 1, 1] - else: - decay = g_t[..., None, :] # [B, H, 1, Dk] - - # S = a S - state = state * decay - - # kv = S * k_t, [B, H, Dv, Dk] * [B, H, 1, Dk] = [B, H, Dv, Dk] -> reduction on Dk - kv_mem = (state * k_t[..., None, :]).sum(axis=-1) # [B, H, Dv] - - # delta = b_t * (v_t - kv) - delta = (v_t - kv_mem) * beta_t[..., None] # [B, H, Dv] - - # S = S + delta * k_t^T, [B, H, Dv, 1] * [B, H, 1, Dk] = [B, H, Dv, Dk] - state = state + delta[..., None] * k_t[..., None, :] # [B, H, Dv, Dk] - - # o_t = S_t * q_t, [B, H, Dv, Dk] * [B, H, 1, Dk] = [B, H, Dv, Dk] -> reduction on Dk - o_t = (state * q_t[..., None, :]).sum(axis=-1) # [B, H, Dv] - outputs.append(o_t) - - return mx.stack(outputs, axis=1), state # [B, T, H, Dv], [B, H, Dv, Dk] + ys.append(out) + ys.append(hf) + mx.eval(ys) + return ys def benchmark_shape(B, T, Hk, Hv, Dk, Dv, chunk_sizes): mx.random.seed(42) - q = mx.random.normal(shape=(B, T, Hk, Dk)) k = mx.random.normal(shape=(B, T, Hk, Dk)) k = k / (mx.linalg.norm(k, axis=-1, keepdims=True) + 1e-6) v = mx.random.normal(shape=(B, T, Hv, Dv)) g = mx.random.normal(shape=(B, T, Hv)) * 0.1 - 1.0 b = mx.sigmoid(mx.random.normal(shape=(B, T, Hv))) - h0 = mx.zeros((B, Hv, Dv, Dk), dtype=mx.float32) shape_str = f"B={B} T={T} Hk={Hk} Hv={Hv} Dk={Dk} Dv={Dv}" + denom = N_iter_bench * N_iter_func - out_ref, hf_ref = gated_delta_ref(q, k, v, g, b, state=h0) - mx.eval(out_ref, hf_ref) - - out, hf = mx.fast.gated_delta_update(q, k, v, g, b, initial_state=h0, C=0) - mx.eval(out, hf) - - atol = 1e-1 - out_close = mx.all(mx.abs(out - out_ref) < atol).item() - hf_close = mx.all(mx.abs(hf - hf_ref) < atol).item() - - out_max_diff = mx.abs(out - out_ref).max().item() - hf_max_diff = mx.abs(hf - hf_ref).max().item() - - # assert out_close, f"output mismatch! max diff: {out_max_diff:.2e}" - # assert hf_close, f"state mismatch! max diff: {hf_max_diff:.2e}" - - # for b in range(B): - # diff = mx.abs(out[b] - out_ref[b]).max().item() - # print(f"batch {b}: {diff:.2e}") - # exit() - if not out_close or not hf_close: - print(f"{RED_BOLD}", end="") - + os.environ["GATED_DELTA_CHUNK"] = "0" + h0 = mx.zeros((B, Hv, Dv, Dk), dtype=mx.float32) + mx.eval(*mx.fast.gated_delta_update(q, k, v, g, b, initial_state=h0)) ms_seq = ( - bench(do_kernel_bench, mx.fast.gated_delta_update, q, k, v, g, b, h0, 0) * 1000 + bench(do_kernel_bench, mx.fast.gated_delta_update, q, k, v, g, b, h0) + / denom + * 1e3 ) speedups = [] - - non_zero_Cs = [C for C in chunk_sizes if C != 0] - for C in non_zero_Cs: + for C in (c for c in chunk_sizes if c != 0): try: + os.environ["GATED_DELTA_CHUNK"] = str(C) h0 = mx.zeros((B, Hv, Dv, Dk), dtype=mx.float32) - out, hf = mx.fast.gated_delta_update(q, k, v, g, b, initial_state=h0, C=C) - mx.eval(out, hf) - - err_out = mx.abs(out - out_ref).max().item() - err_hf = mx.abs(hf - hf_ref).max().item() - + mx.eval(*mx.fast.gated_delta_update(q, k, v, g, b, initial_state=h0)) ms_c = ( - bench( - do_kernel_bench, - mx.fast.gated_delta_update, - q, - k, - v, - g, - b, - h0, - C, - ) - * 1000 + bench(do_kernel_bench, mx.fast.gated_delta_update, q, k, v, g, b, h0) + / denom + * 1e3 ) + speedups.append(ms_seq / ms_c if ms_c > 0 else float("nan")) + except Exception as ex: + print(f" chunk {C} failed: {ex}") + speedups.append(float("nan")) - speedup = ms_seq / ms_c if ms_c > 0 else float("nan") - - out_close = mx.all(mx.abs(out - out_ref) < atol).item() - hf_close = mx.all(mx.abs(hf - hf_ref) < atol).item() - - out_max_diff = mx.abs(out - out_ref).max().item() - hf_max_diff = mx.abs(hf - hf_ref).max().item() - - if not out_close or not hf_close: - raise Exception + return shape_str, f"{ms_seq:.3f}", speedups, ms_seq - speedups.append(speedup) - except Exception as e: - speedups.append(-1) - - return shape_str, f"{ms_seq:.3f}", speedups - - -def run_benchmark(run_full): +def run_benchmark(run_full, to_csv=False, csv_path="benchmark_results.csv"): if run_full: - Bs = [1, 2, 4, 8, 16] - Ts = [2048, 4096, 8192, 16384] - Hs = [16, 24, 32] + Bs = [1, 4, 8, 16] + Ts = [1, 8, 16, 512, 2048, 4096] + Hks = [16] + Hvs = [32] Dks = [128] Dvs = [128] else: Bs = [1, 8, 16] - Ts = [10, 512, 1024, 2048] - # Hs = [16] + Ts = [8, 512, 1024, 2048] Hks = [16] Hvs = [32] Dks = [128] Dvs = [128] - rows = [] - - chunk_sizes = [0, 8] + chunk_sizes = [0, 8, 16] non_zero_Cs = [C for C in chunk_sizes if C != 0] headers = ["B", "T", "Hk", "Hv", "Dk", "Dv", "time_seq (ms)"] + [ - f"C={C} speedup" for C in non_zero_Cs + f"C={C} (speedup)" for C in non_zero_Cs ] - col_widths = [6, 6, 6, 6, 6, 6, 15] + [15] * (len(non_zero_Cs)) + col_widths = [6, 6, 6, 6, 6, 6, 15] + [25] * (len(non_zero_Cs)) fmt = "".join(f"{{:<{w}}}" for w in col_widths) + + rows = [] + print(fmt.format(*headers)) print("-" * (sum(col_widths))) for B, T, Hk, Hv, Dk, Dv in itertools.product(Bs, Ts, Hks, Hvs, Dks, Dvs): - shapes_s, base_time, speedups = benchmark_shape( + shapes_s, base_time_s, speedups, base_time = benchmark_shape( B, T, Hk, Hv, Dk, Dv, chunk_sizes ) - row = [f"{B}", f"{T}", f"{Hk}", f"{Hv}", f"{Dk}", f"{Dv}", base_time] + row = [f"{B}", f"{T}", f"{Hk}", f"{Hv}", f"{Dk}", f"{Dv}", base_time_s] for speed in speedups: - row.append(f"{speed:.2f}x") + row.append(f"{(base_time / speed):<8.2f} ({speed:<5.2f}x)") print(fmt.format(*row), end="") print(f"{RESET}") + rows.append(row) -def run_profile(): - B = 1 - Hk = 16 - Hv = 32 - T = 512 - Dk = 128 - Dv = 128 - CS = [8] # , 16]#, 32] - - mx.random.seed(42) - - q = mx.random.normal(shape=(B, T, Hk, Dk)) - k = mx.random.normal(shape=(B, T, Hk, Dk)) - k = k / (mx.linalg.norm(k, axis=-1, keepdims=True) + 1e-6) - v = mx.random.normal(shape=(B, T, Hv, Dv)) - g = mx.random.normal(shape=(B, T, Hv)) * 0.1 - 1.0 - b = mx.sigmoid(mx.random.normal(shape=(B, T, Hv))) - h0 = mx.zeros((B, Hv, Dv, Dk), dtype=mx.float32) - - mx.eval(q, k, v, g, b, h0) - - out_ref, hf_ref = gated_delta_ref(q, k, v, g, b, state=h0) - mx.eval(out_ref, hf_ref) - - for C in CS: - out, hf = mx.fast.gated_delta_update(q, k, v, g, b, initial_state=h0, C=C) - mx.eval(out, hf) - - assert list(out.shape) == [B, T, Hv, Dv] - assert list(hf.shape) == [B, Hv, Dv, Dk] - assert not mx.any(mx.isnan(out)).item(), "NaNs in output!" - assert not mx.any(mx.isnan(hf)).item(), "NaNs in final state!" - - # correctness check - atol = 1e-2 - out_close = mx.all(mx.abs(out - out_ref) < atol).item() - hf_close = mx.all(mx.abs(hf - hf_ref) < atol).item() - - out_max_diff = mx.abs(out - out_ref).max().item() - hf_max_diff = mx.abs(hf - hf_ref).max().item() - - assert out_close, f"output mismatch! max diff: {out_max_diff:.2e}" - assert hf_close, f"state mismatch! max diff: {hf_max_diff:.2e}" - - profile(do_kernel_bench, mx.fast.gated_delta_update, q, k, v, g, b, h0, C) + if to_csv: + with open(csv_path, "w", newline="") as f: + writer = csv.writer(f) + writer.writerow(headers) + writer.writerows(rows) + print(f"\nResults also written to {csv_path}") if __name__ == "__main__": - parser = argparse.ArgumentParser(description="Gated delta chunk kernel runner") - parser.add_argument("--profile", "-p", action="store_true") + parser = argparse.ArgumentParser(description="Gated delta benchmark") parser.add_argument("--full", "-f", action="store_true") + parser.add_argument("--csv", "-c", action="store_true") + parser.add_argument("--csv_out", "-co", default="benchmark_results.csv") args = parser.parse_args() - if args.profile: - run_profile() - exit() - - run_benchmark(args.full) + run_benchmark(args.full, to_csv=args.csv, csv_path=args.csv_out) diff --git a/mlx/backend/metal/gated_delta_update.cpp b/mlx/backend/metal/gated_delta_update.cpp index 1e6cf14e71..2fd7a19469 100644 --- a/mlx/backend/metal/gated_delta_update.cpp +++ b/mlx/backend/metal/gated_delta_update.cpp @@ -91,6 +91,9 @@ void GatedDeltaUpdate::eval_gpu( const char* chunk_env = std::getenv("GATED_DELTA_CHUNK"); C = chunk_env ? std::stoi(chunk_env) : C; + if (!metal::is_nax_available()) + C = std::min(C, 8); // override in case nax is not available. + std::string suffix = get_type_string(q.dtype()) + "_" + std::to_string(Dk) + "_" + std::to_string(Dv) + "_" + std::to_string(Hk) + "_" + std::to_string(Hv); diff --git a/python/tests/test_fast_gated_delta.py b/python/tests/test_fast_gated_delta.py new file mode 100644 index 0000000000..7fbba0333b --- /dev/null +++ b/python/tests/test_fast_gated_delta.py @@ -0,0 +1,180 @@ +import os +import unittest + +import mlx.core as mx +import mlx_tests +import numpy as np +import torch + + +def gated_delta_oracle( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + beta: torch.Tensor, + g: torch.Tensor, + scale: float = None, + initial_state: torch.Tensor = None, + output_final_state: bool = False, +): + """ + Reference PyTorch implementation of recurrent gated delta rule. + Taken from: https://github.com/fla-org/flash-linear-attention/blob/main/fla/ops/gated_delta_rule/naive.py + + Args: + q: [B, T, H, K] + k: [B, T, H, K] + v: [B, T, H, V] + beta: [B, T, H] + g: [B, T, H] <--- Difference: with out kernel: this is expected as a log. + scale: float, optional <--- Difference: This is done by the qwen3_5 model. + initial_state: [B, H, K, V], optional <- Difference: last two dimensions are transposed. + output_final_state: bool + + Returns: + o: [B, T, H, V] + final_state: [B, H, K, V] if output_final_state else None + """ + q, k, v, beta, g = map( + lambda x: x.transpose(1, 2).contiguous().to(torch.float32), [q, k, v, beta, g] + ) + B, H, T, K, V = *k.shape, v.shape[-1] + o = torch.zeros(B, H, T, V).to(v) + h = torch.zeros(B, H, K, V).to(v) + if initial_state is not None: + h = initial_state.to(torch.float32) + if scale is None: + scale = 1 / (q.shape[-1] ** 0.5) + q = q * scale + + for i in range(T): + b_q = q[:, :, i] + b_k = k[:, :, i] + b_v = v[:, :, i].clone() + h = h.clone() * g[:, :, i].exp()[..., None, None] + b_beta = beta[:, :, i] + b_v = b_v - (h.clone() * b_k[..., None]).sum(-2) + b_v = b_v * b_beta[..., None] + h = h.clone() + b_k.unsqueeze(-1) * b_v.unsqueeze(-2) + o[:, :, i] = torch.einsum("bhd,bhdm->bhm", b_q, h) + + if not output_final_state: + h = None + o = o.transpose(1, 2).contiguous() + return o, h + + +def runner(dims, stream=mx.gpu, reference=True): + B, Hk, Hv, T, Dk, Dv = dims + + q = mx.random.normal(shape=(B, T, Hk, Dk)) * 10 + k = mx.random.normal(shape=(B, T, Hk, Dk)) * 10 + k = k / (mx.linalg.norm(k, axis=-1, keepdims=True) + 1e-6) + v = mx.random.normal(shape=(B, T, Hv, Dv)) + g = mx.random.uniform(shape=(B, T, Hv)) + b = mx.sigmoid(mx.random.normal(shape=(B, T, Hv))) + h0 = mx.zeros((B, Hv, Dv, Dk), dtype=mx.float32) + + if reference: + # Prepare reference inputs + qpt = torch.from_numpy(np.array(q)) + kpt = torch.from_numpy(np.array(k)) + vpt = torch.from_numpy(np.array(v)) + bpt = torch.from_numpy(np.array(b)) + gpt = torch.from_numpy(np.array(g)) + h0pt = torch.from_numpy(np.array(h0)).transpose(-1, -2).contiguous() + + out_on_py, hf_on_py = gated_delta_oracle( + qpt, + kpt, + vpt, + bpt, + torch.log(gpt), + scale=1.0, + initial_state=h0pt, + output_final_state=True, + ) + + out_on = mx.array(out_on_py.detach().cpu().numpy()) # [B, T, Hv, Dv] + hf_on = mx.swapaxes( + mx.array(hf_on_py.detach().cpu().numpy()), -1, -2 + ) # -> [B, Hv, Dv, Dk] + out_ref = mx.array(out_on) + hf_ref = mx.array(hf_on) + else: + # use fallback for tests once fallback is validated + out_ref, hf_ref = mx.fast.gated_delta_update( + q, k, v, g, b, initial_state=h0, stream=mx.cpu + ) + + mx.eval(out_ref, hf_ref) + + out, hf = mx.fast.gated_delta_update(q, k, v, g, b, initial_state=h0, stream=stream) + + mx.eval(out, hf) + return (out, hf), (out_ref, hf_ref) + + +class TestGatedDelta(mlx_tests.MLXTestCase): + base_dims = (1, 32, 32, 1, 128, 128) + unaligned_dims = (1, 32, 32, 33, 128, 128) + big_batch_dims = (128, 32, 32, 16, 128, 128) + large_t_dims = (2, 32, 32, 1111, 128, 128) + + fallback_dims = [base_dims, unaligned_dims, big_batch_dims] + gpu_dims = fallback_dims + [large_t_dims] + + # base + def test_gated_delta_fallback(self): + for dims in self.fallback_dims: + (out, hf), (out_ref, hf_ref) = runner(dims, mx.cpu) + msg = f"Failed on Dimensions: {dims}" + self.assertTrue( + mx.allclose(out_ref, out, atol=1e-4, rtol=1e-4), msg="Out " + msg + ) + self.assertTrue( + mx.allclose(hf_ref, hf, atol=1e-4, rtol=1e-4), msg="State " + msg + ) + + @unittest.skipIf(not mx.is_available(mx.gpu), "No GPU available") + def test_gated_delta_sequential(self): + os.environ["GATED_DELTA_CHUNK"] = "0" + for dims in self.gpu_dims: + (out, hf), (out_ref, hf_ref) = runner(dims, reference=False) + msg = f"Failed on Dimensions: {dims}" + self.assertTrue( + mx.allclose(out_ref, out, atol=1e-4, rtol=1e-4), msg="Out " + msg + ) + self.assertTrue( + mx.allclose(hf_ref, hf, atol=1e-4, rtol=1e-4), msg="State " + msg + ) + + @unittest.skipIf(not mx.is_available(mx.gpu), "No GPU available") + def test_gated_delta_simdgroup(self): + os.environ["GATED_DELTA_CHUNK"] = "8" + for dims in self.gpu_dims: + (out, hf), (out_ref, hf_ref) = runner(dims, reference=False) + msg = f"Failed on Dimensions: {dims}" + self.assertTrue( + mx.allclose(out_ref, out, atol=1e-4, rtol=1e-4), msg="Out " + msg + ) + self.assertTrue( + mx.allclose(hf_ref, hf, atol=1e-4, rtol=1e-4), msg="State " + msg + ) + + @unittest.skipIf(not mx.is_available(mx.gpu), "No GPU available") + def test_gated_delta_nax(self): + os.environ["GATED_DELTA_CHUNK"] = "16" + for dims in self.gpu_dims: + (out, hf), (out_ref, hf_ref) = runner(dims, reference=False) + msg = f"Failed on Dimensions: {dims}" + self.assertTrue( + mx.allclose(out_ref, out, atol=1e-4, rtol=1e-4), msg="Out " + msg + ) + self.assertTrue( + mx.allclose(hf_ref, hf, atol=1e-4, rtol=1e-4), msg="State " + msg + ) + + +if __name__ == "__main__": + mlx_tests.MLXTestRunner(failfast=True) From 7c112e4acc16e195b821813b55f5f606bcb54e93 Mon Sep 17 00:00:00 2001 From: tpegolotti Date: Wed, 5 Aug 2026 16:17:53 +0200 Subject: [PATCH 41/56] Fix sign --- mlx/backend/metal/kernels/gated_delta_update_impl.h | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mlx/backend/metal/kernels/gated_delta_update_impl.h b/mlx/backend/metal/kernels/gated_delta_update_impl.h index 35ac09210a..b6271a9971 100644 --- a/mlx/backend/metal/kernels/gated_delta_update_impl.h +++ b/mlx/backend/metal/kernels/gated_delta_update_impl.h @@ -483,7 +483,7 @@ template KKtK_tile = KKt_tile; SCALE_TRIEQ_NAX1(KKtK_tile, beta_fm); - ADD_NAX(Tinv_tile, I_tile, KKtK_tile); + SUB_NAX(Tinv_tile, I_tile, KKtK_tile); STEEL_PRAGMA_UNROLL for (int step = 0; step < 15; step++) { MM16x16x16(TMP_tile, 0, KKtK_tile, false, 0, Tinv_tile, false, 0); From df345688291d6fdfe9eddc9808976414d4b4668b Mon Sep 17 00:00:00 2001 From: tpegolotti Date: Wed, 5 Aug 2026 16:18:58 +0200 Subject: [PATCH 42/56] Increase nax atol --- python/tests/test_fast_gated_delta.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/python/tests/test_fast_gated_delta.py b/python/tests/test_fast_gated_delta.py index 7fbba0333b..ee9b6d9734 100644 --- a/python/tests/test_fast_gated_delta.py +++ b/python/tests/test_fast_gated_delta.py @@ -169,10 +169,10 @@ def test_gated_delta_nax(self): (out, hf), (out_ref, hf_ref) = runner(dims, reference=False) msg = f"Failed on Dimensions: {dims}" self.assertTrue( - mx.allclose(out_ref, out, atol=1e-4, rtol=1e-4), msg="Out " + msg + mx.allclose(out_ref, out, atol=1e-1, rtol=1e-4), msg="Out " + msg ) self.assertTrue( - mx.allclose(hf_ref, hf, atol=1e-4, rtol=1e-4), msg="State " + msg + mx.allclose(hf_ref, hf, atol=1e-1, rtol=1e-4), msg="State " + msg ) From e4a6d524bdca74d42e7c9a2eaaa4e7d5ee273dae Mon Sep 17 00:00:00 2001 From: tpegolotti Date: Wed, 5 Aug 2026 16:30:04 +0200 Subject: [PATCH 43/56] Removed scaling from input --- python/tests/test_fast_gated_delta.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/python/tests/test_fast_gated_delta.py b/python/tests/test_fast_gated_delta.py index ee9b6d9734..707b90f113 100644 --- a/python/tests/test_fast_gated_delta.py +++ b/python/tests/test_fast_gated_delta.py @@ -67,8 +67,8 @@ def gated_delta_oracle( def runner(dims, stream=mx.gpu, reference=True): B, Hk, Hv, T, Dk, Dv = dims - q = mx.random.normal(shape=(B, T, Hk, Dk)) * 10 - k = mx.random.normal(shape=(B, T, Hk, Dk)) * 10 + q = mx.random.normal(shape=(B, T, Hk, Dk)) + k = mx.random.normal(shape=(B, T, Hk, Dk)) k = k / (mx.linalg.norm(k, axis=-1, keepdims=True) + 1e-6) v = mx.random.normal(shape=(B, T, Hv, Dv)) g = mx.random.uniform(shape=(B, T, Hv)) From 19c1cb56bf9529d711f99067f14cfbeb67e694ae Mon Sep 17 00:00:00 2001 From: tpegolotti Date: Wed, 5 Aug 2026 13:42:02 -0700 Subject: [PATCH 44/56] Add mask support in fallback --- benchmarks/python/gated_delta_bench.py | 2 +- mlx/backend/metal/gated_delta_update.cpp | 27 +++++++++- mlx/fast.cpp | 66 ++++++++++++------------ mlx/fast.h | 1 + mlx/fast_primitives.h | 7 ++- python/src/fast.cpp | 3 +- python/tests/test_fast_gated_delta.py | 45 +++++++++++++++- 7 files changed, 109 insertions(+), 42 deletions(-) diff --git a/benchmarks/python/gated_delta_bench.py b/benchmarks/python/gated_delta_bench.py index 128c86c748..277897699c 100644 --- a/benchmarks/python/gated_delta_bench.py +++ b/benchmarks/python/gated_delta_bench.py @@ -86,7 +86,7 @@ def benchmark_shape(B, T, Hk, Hv, Dk, Dv, chunk_sizes): def run_benchmark(run_full, to_csv=False, csv_path="benchmark_results.csv"): if run_full: Bs = [1, 4, 8, 16] - Ts = [1, 8, 16, 512, 2048, 4096] + Ts = [8, 64, 256, 512, 1024, 2048, 4096] Hks = [16] Hvs = [32] Dks = [128] diff --git a/mlx/backend/metal/gated_delta_update.cpp b/mlx/backend/metal/gated_delta_update.cpp index 2fd7a19469..b0b57f95e2 100644 --- a/mlx/backend/metal/gated_delta_update.cpp +++ b/mlx/backend/metal/gated_delta_update.cpp @@ -12,11 +12,34 @@ namespace mlx::core::fast { -bool GatedDeltaUpdate::use_fallback(Stream s) { - // TODO: finish implementation. What else is needed? +bool GatedDeltaUpdate::use_fallback( + const int Hk, + const int Dk, + const int Hv, + const int Dv, + const bool has_mask, + Stream s) { if (s.device == Device::cpu) { return true; } + + if (has_mask) { + return true; + } + + if (Dk != 128 || Dv != 128) { + return true; + } + + const bool supported_heads = (Hk == 24 && Hv == 24) || + (Hk == 32 && Hv == 32) || (Hk == 16 && Hv == 32) || + (Hk == 16 && Hv == 48); + if (!supported_heads) { + return true; + } + + return false; + return false; } diff --git a/mlx/fast.cpp b/mlx/fast.cpp index 36a79d2131..4ec16abe97 100644 --- a/mlx/fast.cpp +++ b/mlx/fast.cpp @@ -929,11 +929,11 @@ std::vector gated_delta_update( const array& gates, const array& beta_, const std::optional& initial_state, /* = std::nullopt */ + const std::optional& mask_, /* = std::nullopt */ StreamOrDevice s_ /* = {} */) { // determine output dtype auto s = to_stream(s_); - // TODO fix this auto promoted = promote_types(queries.dtype(), keys.dtype()); auto out_dtype = issubdtype(promoted, floating) ? promoted @@ -946,15 +946,6 @@ std::vector gated_delta_update( auto g = astype(gates, out_dtype, s); auto beta = astype(beta_, out_dtype, s); - // wrong, not as fast. This kind of operations need to be done in - // the eval gpu inside a function that checks for shapes and strides - // TODO: mode this - // q = contiguous(q, false, s); - // k = contiguous(k, false, s); - // v = contiguous(v, false, s); - // g = contiguous(g, false, s); - // beta = contiguous(beta, false, s); - int B = q.shape(0); int T = q.shape(1); int Hk = q.shape(2); @@ -965,7 +956,11 @@ std::vector gated_delta_update( auto h0 = initial_state.has_value() ? astype(*initial_state, float32, s) : zeros({B, Hv, Dv, Dk}, float32, s); - auto fallback = [B, T, Hk, Dk, Hv, Dv, s](std::vector inputs) { + bool has_mask = mask_.has_value(); + auto mask = has_mask ? astype(*mask_, bool_, s) : array(false); + + auto fallback = [B, T, Hk, Dk, Hv, Dv, has_mask, s]( + std::vector inputs) { auto q = astype(inputs[0], float32, s); auto k = astype(inputs[1], float32, s); auto v = astype(inputs[2], float32, s); @@ -973,63 +968,66 @@ std::vector gated_delta_update( auto beta = astype(inputs[4], float32, s); auto state = astype(inputs[5], float32, s); + array mask = has_mask ? astype(inputs[6], bool_, s) : array(false); + const array zero = array(0.0f, float32); + std::vector outputs; for (int t = 0; t < T; t++) { auto get_t = [&](const array& a, int t) { Shape start(a.ndim(), 0), stop = a.shape(); start[1] = t; - stop[1] = t + 1; // [:, t:t+1, ...] - return squeeze(slice(a, start, stop, s), 1, s); // drop the time axis + stop[1] = t + 1; + return squeeze(slice(a, start, stop, s), 1, s); }; - // q_t = q[:, t] # [B, H, Dk] - // k_t = k[:, t] # [B, H, Dk] - // v_t = v[:, t] # [B, H, Dv] - // g_t = g[:, t] # [B, H] or [B, H, Dk] - // beta_t = beta[:, t] # [B, H] auto q_t = get_t(q, t); auto k_t = get_t(k, t); auto v_t = get_t(v, t); auto g_t = get_t(g, t); auto beta_t = get_t(beta, t); + array mask_t = has_mask ? get_t(mask, t) : array(true); + + auto state_prev = state; - // if g_t.ndim == 2: - // decay = g_t[..., None, None] # [B, H, 1, 1] - // else: - // decay = g_t[..., None, :] # [B, H, 1, Dk] auto decay = (g_t.ndim() == 2) ? expand_dims(g_t, {-1, -2}, s) // [B,H,1,1] : expand_dims(g_t, -2, s); // [B,H,1,Dk] - // state = state * decay - state = multiply(state, decay, s); + auto state_next = multiply(state_prev, decay, s); - // kv_mem = (state * k_t[..., None, :]).sum(axis=-1) auto kv = - sum(multiply(state, expand_dims(k_t, -2, s), s), + sum(multiply(state_next, expand_dims(k_t, -2, s), s), -1, false, s); // [B,H,Dv] - // delta = (v_t - kv_mem) * beta_t[..., None] # [B, H, Dv] auto delta = subtract(v_t, kv, s); delta = multiply(delta, expand_dims(beta_t, -1, s), s); - // state = state + delta[..., None] * k_t[..., None, :] # [B, H, Dv, Dk] - state = - add(state, + state_next = + add(state_next, multiply(expand_dims(delta, -1, s), expand_dims(k_t, -2, s), s), s); - // o_t = (state * q_t[..., None, :]).sum(axis=-1) + if (has_mask) { + auto state_mask = expand_dims(mask_t, {-1, -2, -3}, s); + state = where(state_mask, state_next, state_prev, s); + } else { + state = state_next; + } + auto o_t = sum(multiply(state, expand_dims(q_t, -2, s), s), -1, false, s); + + if (has_mask) { + auto out_mask = expand_dims(mask_t, {-1, -2}, s); + o_t = where(out_mask, o_t, zero, s); + } outputs.push_back(o_t); } - // mx.stack(outputs, axis=1) auto out = stack(outputs, 1, s); return std::vector{out, state}; }; - if (!GatedDeltaUpdate::use_fallback(s)) { + if (!GatedDeltaUpdate::use_fallback(Hk, Dk, Hv, Dv, has_mask, s)) { auto result = array::make_arrays( /* output shapes */ {{B, T, Hv, Dv}, {B, Hv, Dv, Dk}}, /* dtypes */ {out_dtype, float32}, @@ -1040,7 +1038,7 @@ std::vector gated_delta_update( return result; } - auto result = fallback({q, k, v, g, beta, h0}); + auto result = fallback({q, k, v, g, beta, h0, mask}); return result; } diff --git a/mlx/fast.h b/mlx/fast.h index 81ebf6d953..052a23811d 100644 --- a/mlx/fast.h +++ b/mlx/fast.h @@ -62,6 +62,7 @@ MLX_API std::vector gated_delta_update( const array& gates, const array& beta_, const std::optional& initial_state = std::nullopt, + const std::optional& mask = std::nullopt, StreamOrDevice s = {}); using TemplateArg = std::variant; diff --git a/mlx/fast_primitives.h b/mlx/fast_primitives.h index 62eaa7c51e..712de05b40 100644 --- a/mlx/fast_primitives.h +++ b/mlx/fast_primitives.h @@ -333,9 +333,12 @@ class GatedDeltaUpdate : public Custom { : Custom(stream, std::move(fallback)) {} static bool use_fallback( - /* TODO */ + const int Hk, + const int Dk, + const int Hv, + const int Dv, + const bool has_mask, Stream s); - static bool supports_bool_mask(); void eval_cpu(const std::vector& inputs, std::vector& outputs) override { diff --git a/python/src/fast.cpp b/python/src/fast.cpp index 925f555a0b..7021e361a9 100644 --- a/python/src/fast.cpp +++ b/python/src/fast.cpp @@ -342,6 +342,7 @@ void init_fast(nb::module_& parent_module) { "gamma"_a, "beta"_a, "initial_state"_a = nb::none(), // optional, defaults to None + "mask"_a = nb::none(), // optional, defaults to None "stream"_a = nb::none(), // optional, defaults to None R"( Chunked gated delta network forward pass. @@ -353,7 +354,7 @@ void init_fast(nb::module_& parent_module) { gamma: beta: Delta update rates [B, H, T] initial_state: Optional initial hidden state [B, H, Dk, Dv] - + mask: Optional Returns: Tuple of (output [B, H, T, Dv], final_state [B, H, Dk, Dv]) )"); diff --git a/python/tests/test_fast_gated_delta.py b/python/tests/test_fast_gated_delta.py index 707b90f113..d596972ff2 100644 --- a/python/tests/test_fast_gated_delta.py +++ b/python/tests/test_fast_gated_delta.py @@ -73,7 +73,7 @@ def runner(dims, stream=mx.gpu, reference=True): v = mx.random.normal(shape=(B, T, Hv, Dv)) g = mx.random.uniform(shape=(B, T, Hv)) b = mx.sigmoid(mx.random.normal(shape=(B, T, Hv))) - h0 = mx.zeros((B, Hv, Dv, Dk), dtype=mx.float32) + h0 = mx.random.normal((B, Hv, Dv, Dk), dtype=mx.float32) if reference: # Prepare reference inputs @@ -124,7 +124,6 @@ class TestGatedDelta(mlx_tests.MLXTestCase): fallback_dims = [base_dims, unaligned_dims, big_batch_dims] gpu_dims = fallback_dims + [large_t_dims] - # base def test_gated_delta_fallback(self): for dims in self.fallback_dims: (out, hf), (out_ref, hf_ref) = runner(dims, mx.cpu) @@ -136,6 +135,48 @@ def test_gated_delta_fallback(self): mx.allclose(hf_ref, hf, atol=1e-4, rtol=1e-4), msg="State " + msg ) + def test_gated_delta_fallback_masked(self): + for dims in self.fallback_dims: + + B, Hk, Hv, T, Dk, Dv = dims + + q = mx.random.normal(shape=(B, T, Hk, Dk)) + k = mx.random.normal(shape=(B, T, Hk, Dk)) + k = k / (mx.linalg.norm(k, axis=-1, keepdims=True) + 1e-6) + v = mx.random.normal(shape=(B, T, Hv, Dv)) + g = mx.random.uniform(shape=(B, T, Hv)) + b = mx.sigmoid(mx.random.normal(shape=(B, T, Hv))) + h0 = mx.random.normal((B, Hv, Dv, Dk), dtype=mx.float32) + + # make a mask + lengths = mx.random.randint(1, T + 1, shape=(B,)) + mask = mx.arange(T)[None, :] < lengths[:, None] + # mask one input in python + mask_float = mask.astype(q.dtype) + km = k * mask_float[..., None, None] + vm = v * mask_float[..., None, None] + qm = q * mask_float[..., None, None] + bm = b * mask_float[..., None] + gm = mx.where(mask[..., None], g, 1.0) + + out_ref, hf_ref = mx.fast.gated_delta_update( + qm, km, vm, gm, bm, initial_state=h0, stream=mx.cpu + ) + + mx.eval(out_ref, hf_ref) + out, hf = mx.fast.gated_delta_update( + q, k, v, g, b, initial_state=h0, mask=mask, stream=mx.cpu + ) + mx.eval(out, hf) + + msg = f"Failed on Dimensions: {dims}" + self.assertTrue( + mx.allclose(out_ref, out, atol=1e-4, rtol=1e-4), msg="Out " + msg + ) + self.assertTrue( + mx.allclose(hf_ref, hf, atol=1e-4, rtol=1e-4), msg="State " + msg + ) + @unittest.skipIf(not mx.is_available(mx.gpu), "No GPU available") def test_gated_delta_sequential(self): os.environ["GATED_DELTA_CHUNK"] = "0" From 21973cfe48c107ebad58673db4c25e308597c39e Mon Sep 17 00:00:00 2001 From: tpegolotti Date: Thu, 6 Aug 2026 03:27:16 -0700 Subject: [PATCH 45/56] Fix non metal build link errors --- mlx/backend/no_gpu/primitives.cpp | 11 +++++++++++ mlx/fast_primitives.h | 2 -- 2 files changed, 11 insertions(+), 2 deletions(-) diff --git a/mlx/backend/no_gpu/primitives.cpp b/mlx/backend/no_gpu/primitives.cpp index 0e05e9d19f..c4067bcd8f 100644 --- a/mlx/backend/no_gpu/primitives.cpp +++ b/mlx/backend/no_gpu/primitives.cpp @@ -36,6 +36,16 @@ bool fast::ScaledDotProductAttention::use_fallback( return true; } +bool fast::GatedDeltaUpdate::use_fallback( + const int Hk, + const int Dk, + const int Hv, + const int Dv, + const bool has_mask, + Stream s) { + return true; +} + bool fast::ScaledDotProductAttention::supports_bool_mask() { return false; } @@ -170,6 +180,7 @@ NO_GPU_MULTI(RMSNormVJP) NO_GPU_USE_FALLBACK(RoPE) NO_GPU_MULTI(ScaledDotProductAttention) NO_GPU_MULTI(ScaledDotProductAttentionVJP) +NO_GPU_MULTI(GatedDeltaUpdate) NO_GPU_MULTI(ConvertFP8) NO_GPU_MULTI(Quantize) NO_GPU_MULTI(CustomKernel) diff --git a/mlx/fast_primitives.h b/mlx/fast_primitives.h index 712de05b40..9bd3954b2a 100644 --- a/mlx/fast_primitives.h +++ b/mlx/fast_primitives.h @@ -348,8 +348,6 @@ class GatedDeltaUpdate : public Custom { void eval_gpu(const std::vector& inputs, std::vector& outputs) override; - bool is_equivalent(const Primitive& other) const override; - DEFINE_NAME(GatedDeltaUpdate); DEFINE_INPUT_OUTPUT_SHAPE() auto state() const { From eba6ea70c053f2e4f401aadb5e8df58fb9eafe81 Mon Sep 17 00:00:00 2001 From: tpegolotti Date: Thu, 6 Aug 2026 05:03:16 -0700 Subject: [PATCH 46/56] Fixing more linker errors --- mlx/backend/cuda/primitives.cpp | 14 + mlx/backend/metal/CMakeLists.txt | 1 + mlx/backend/metal/gated_delta_update.cpp | 13 +- mlx/backend/metal/kernels/CMakeLists.txt | 2 + .../metal/kernels/gated_delta_update.metal | 20 +- .../metal/kernels/gated_delta_update_impl.h | 535 +----------------- 6 files changed, 32 insertions(+), 553 deletions(-) diff --git a/mlx/backend/cuda/primitives.cpp b/mlx/backend/cuda/primitives.cpp index 94f260767f..ad3991a7e5 100644 --- a/mlx/backend/cuda/primitives.cpp +++ b/mlx/backend/cuda/primitives.cpp @@ -24,6 +24,16 @@ namespace mlx::core { throw std::runtime_error(#func " has no CUDA implementation."); \ } +bool fast::GatedDeltaUpdate::use_fallback( + const int Hk, + const int Dk, + const int Hv, + const int Dv, + const bool has_mask, + Stream s) { + return true; +} + NO_GPU_MULTI(LUF) NO_GPU_MULTI(QRF) NO_GPU_MULTI(SVD) @@ -32,6 +42,10 @@ NO_GPU(Cholesky) NO_GPU_MULTI(Eig) NO_GPU_MULTI(Eigh) +namespace fast { +NO_GPU_MULTI(GatedDeltaUpdate) +} + namespace distributed { NO_GPU_MULTI(Send) NO_GPU_MULTI(Recv) diff --git a/mlx/backend/metal/CMakeLists.txt b/mlx/backend/metal/CMakeLists.txt index 0ba80646a8..8b1266f36f 100644 --- a/mlx/backend/metal/CMakeLists.txt +++ b/mlx/backend/metal/CMakeLists.txt @@ -101,6 +101,7 @@ if(MLX_METAL_JIT) kernels/fp4.h) make_jit_source(steel/attn/kernels/steel_attention_nax) + make_jit_source(gated_delta_update_nax) else() message( diff --git a/mlx/backend/metal/gated_delta_update.cpp b/mlx/backend/metal/gated_delta_update.cpp index b0b57f95e2..1572cf7676 100644 --- a/mlx/backend/metal/gated_delta_update.cpp +++ b/mlx/backend/metal/gated_delta_update.cpp @@ -102,6 +102,10 @@ void GatedDeltaUpdate::eval_gpu( int Hv = v.shape(2); int Dv = v.shape(3); +#if __METAL_VERSION__ >= 400 +#error "NAX branch IS compiled, __METAL_VERSION__ >= 400" +#endif + int C = 1; const char* threashold_env = std::getenv("GATED_DELTA_THRESH"); int threshold = threashold_env ? std::stoi(threashold_env) : 16; @@ -222,13 +226,4 @@ void GatedDeltaUpdate::eval_gpu( } } -bool GatedDeltaUpdate::is_equivalent(const Primitive& other) const { - const auto* p = dynamic_cast(&other); - if (p == nullptr) { - return false; - } - // TODO: finish implementation - return true; -} - } // namespace mlx::core::fast diff --git a/mlx/backend/metal/kernels/CMakeLists.txt b/mlx/backend/metal/kernels/CMakeLists.txt index 5f70ae9b0d..c899458f64 100644 --- a/mlx/backend/metal/kernels/CMakeLists.txt +++ b/mlx/backend/metal/kernels/CMakeLists.txt @@ -168,6 +168,8 @@ if(NOT MLX_METAL_JIT) ${STEEL_NAX_HEADERS}) build_kernel(quantized_nax quantized_nax.h ${STEEL_NAX_HEADERS}) + build_kernel(gated_delta_update_nax gated_delta_update_nax.h + ${STEEL_NAX_HEADERS}) build_kernel(fp_quantized_nax fp4.h fp8.h fp_quantized_nax.h ${STEEL_NAX_HEADERS}) diff --git a/mlx/backend/metal/kernels/gated_delta_update.metal b/mlx/backend/metal/kernels/gated_delta_update.metal index a5bd7e5442..fba60ed104 100644 --- a/mlx/backend/metal/kernels/gated_delta_update.metal +++ b/mlx/backend/metal/kernels/gated_delta_update.metal @@ -25,23 +25,9 @@ using namespace metal; hv, \ c) -#define instantiate_gated_delta_update_fused_nax(in_type, dk, dv, hk, hv, c) \ - instantiate_kernel( \ - "gated_delta_fused_nax_" #in_type "_" #dk "_" #dv "_" #hk "_" #hv \ - "_" #c, \ - gated_delta_fused_nax, \ - in_type, \ - dk, \ - dv, \ - hk, \ - hv, \ - c) - -#define instantiate_gated_delta_dims(in_type, dk, dv, hk, hv) \ - instantiate_gated_delta_update_seq(in_type, dk, dv, hk, hv) \ - instantiate_gated_delta_update_fused_chunk(in_type, dk, dv, hk, hv, 8) \ - instantiate_gated_delta_update_fused_nax( \ - in_type, dk, dv, hk, hv, 16) +#define instantiate_gated_delta_dims(in_type, dk, dv, hk, hv) \ + instantiate_gated_delta_update_seq(in_type, dk, dv, hk, hv) \ + instantiate_gated_delta_update_fused_chunk(in_type, dk, dv, hk, hv, 8) #define instantiate_gated_delta(in_type) \ instantiate_gated_delta_dims(in_type, 128, 128, 24, 24) \ diff --git a/mlx/backend/metal/kernels/gated_delta_update_impl.h b/mlx/backend/metal/kernels/gated_delta_update_impl.h index b6271a9971..095a744a71 100644 --- a/mlx/backend/metal/kernels/gated_delta_update_impl.h +++ b/mlx/backend/metal/kernels/gated_delta_update_impl.h @@ -6,14 +6,7 @@ #include #include -#include "mlx/backend/metal/kernels/steel/gemm/nax.h" -#include "mlx/backend/metal/kernels/steel/gemm/params.h" -#include "mlx/backend/metal/kernels/steel/gemm/transforms.h" -#include "mlx/backend/metal/kernels/steel/utils.h" - -using namespace metal; -using namespace mpp; -using namespace mpp::tensor_ops; +#define FULL_UNROLL _Pragma("clang loop unroll(full)") #define AT(TILE, IDX) TILE.thread_elements()[IDX] #define SUB(TILE0, TILE1, TILE2) \ @@ -53,518 +46,6 @@ using namespace mpp::tensor_ops; AT(TILE0, 1) *= fn + 1 >= fm ? 0.f : S1; \ } -// NAX MACROS I can probably do a nice template instead of doing this - -// fm = base_fm + (idx >> 2) * 8; // idx>>2 = idx/4 -> 0 for idx 0-3, 1 for -// idx 4-7 fn = base_fn + (idx % 4); // 4 consecutive columns -#define AT_NAX(TILE, IDX) TILE.elems()[IDX] - -#define SUB_NAX(TILE0, TILE1, TILE2) \ - { \ - STEEL_PRAGMA_UNROLL \ - for (short _i = 0; _i < decltype(TILE0)::kElemsPerFrag; _i++) { \ - AT_NAX(TILE0, _i) = AT_NAX(TILE1, _i) - AT_NAX(TILE2, _i); \ - } \ - } - -#define ADD_NAX(TILE0, TILE1, TILE2) \ - { \ - STEEL_PRAGMA_UNROLL \ - for (short _i = 0; _i < decltype(TILE0)::kElemsPerFrag; _i++) { \ - AT_NAX(TILE0, _i) = AT_NAX(TILE1, _i) + AT_NAX(TILE2, _i); \ - } \ - } - -#define FMA_NAX(TILE0, S, TILE1, TILE2) \ - { \ - STEEL_PRAGMA_UNROLL \ - for (short _i = 0; _i < mlx::steel::BaseNAXFrag::kElemsPerFrag; _i++) { \ - (TILE0)[_i] = (S) * (TILE1)[_i] + (TILE2)[_i]; \ - } \ - } - -#define SCALE_NAX(TILE0, S) \ - { \ - STEEL_PRAGMA_UNROLL \ - for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ - AT_NAX(TILE0, _i) *= (S); \ - } \ - } - -#define SCALE_ROW_NAX(TILE0, S) \ - { \ - STEEL_PRAGMA_UNROLL \ - for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ - const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; \ - AT_NAX(TILE0, _i) *= \ - metal::fast::exp((S)[mlx::steel::BaseNAXFrag::get_coord(_w).y]); \ - } \ - } - -#define SCALE_BETA_NAX(TILE0, BETA2) \ - { \ - STEEL_PRAGMA_UNROLL \ - for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ - const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; \ - AT_NAX(TILE0, _i) *= (BETA2)[_w >> 2]; \ - } \ - } - -#define SCALE2_NAX(TILE0, GAMMA) \ - { \ - STEEL_PRAGMA_UNROLL \ - for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ - const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; \ - const short _fm = mlx::steel::BaseNAXFrag::get_coord(_w).y; \ - AT_NAX(TILE0, _i) *= metal::fast::exp((GAMMA)[(C) - 1] - (GAMMA)[_fm]); \ - } \ - } - -#define SCALE_TRI_NAX(TILE0, GAMMA) \ - { \ - STEEL_PRAGMA_UNROLL \ - for (short _i = 0; _i < decltype(TILE0)::kElemsPerFrag; _i++) { \ - const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); /* {fn, fm} */ \ - AT_NAX(TILE0, _i) *= (_c.x > _c.y) \ - ? 0.f \ - : metal::fast::exp((GAMMA)[_c.y] - (GAMMA)[_c.x]); \ - } \ - } - -#define SCALE_TRIEQ_NAX1(TILE0, BETA) \ - { \ - STEEL_PRAGMA_UNROLL \ - for (short _i = 0; _i < decltype(TILE0)::kElemsPerFrag; _i++) { \ - const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); /* {fn, fm} */ \ - AT_NAX(TILE0, _i) *= (_c.x >= _c.y) ? 0.f : (BETA)[_i >> 2]; \ - } \ - } - -namespace mlx { -namespace steel { -template < - typename CType, - typename AType, - typename BType, - bool transpose_a = false, - bool transpose_b = false, - mpp::tensor_ops::matmul2d_descriptor::mode Mode = - mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate> -METAL_FUNC static constexpr void mma( - thread BaseNAXFrag::dtype_frag_t& C, - const thread BaseNAXFrag::dtype_frag_t& A0, - const thread BaseNAXFrag::dtype_frag_t& A1, - metal::bool_constant, - const thread BaseNAXFrag::dtype_frag_t& B0, - const thread BaseNAXFrag::dtype_frag_t& B1, - metal::bool_constant) { - // M=16, N=16, K=32: A and B each two K-fragments, single 16x16 C. - constexpr auto desc = mpp::tensor_ops::matmul2d_descriptor( - 16, 16, 32, transpose_a, transpose_b, true, Mode); - - mpp::tensor_ops::matmul2d gemm_op; - - auto ct_a = - gemm_op.template get_left_input_cooperative_tensor(); - auto ct_b = - gemm_op - .template get_right_input_cooperative_tensor(); - auto ct_c = gemm_op.template get_destination_cooperative_tensor< - decltype(ct_a), - decltype(ct_b), - CType>(); - - STEEL_PRAGMA_UNROLL - for (short i = 0; i < BaseNAXFrag::kElemsPerFrag; i++) { - ct_a[i] = A0[i]; - ct_a[BaseNAXFrag::kElemsPerFrag + i] = A1[i]; - ct_b[i] = B0[i]; - ct_b[BaseNAXFrag::kElemsPerFrag + i] = B1[i]; - ct_c[i] = C[i]; - } - - gemm_op.run(ct_a, ct_b, ct_c); - - STEEL_PRAGMA_UNROLL - for (short i = 0; i < BaseNAXFrag::kElemsPerFrag; i++) { - C[i] = ct_c[i]; - } -} - -template < - typename CType, - typename AType, - typename BType, - bool transpose_a, - bool transpose_b, - mpp::tensor_ops::matmul2d_descriptor::mode Mode = - mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate> -METAL_FUNC static constexpr void mma( - thread BaseNAXFrag::dtype_frag_t& C, - const thread BaseNAXFrag::dtype_frag_t& A, - metal::bool_constant, - const thread BaseNAXFrag::dtype_frag_t& B, - metal::bool_constant) { - constexpr auto desc = mpp::tensor_ops::matmul2d_descriptor( - 16, 32, 16, transpose_a, transpose_b, true, Mode); - - mpp::tensor_ops::matmul2d gemm_op; - - auto ct_a = - gemm_op.template get_left_input_cooperative_tensor(); - auto ct_b = - gemm_op - .template get_right_input_cooperative_tensor(); - auto ct_c = gemm_op.template get_destination_cooperative_tensor< - decltype(ct_a), - decltype(ct_b), - CType>(); - - STEEL_PRAGMA_UNROLL - for (short i = 0; i < BaseNAXFrag::kElemsPerFrag; i++) { - ct_a[i] = A[i]; - ct_b[i] = B[i]; - ct_b[BaseNAXFrag::kElemsPerFrag + i] = 0.0; - ct_c[i] = C[i]; - ct_c[BaseNAXFrag::kElemsPerFrag + i] = 0.0; - } - - gemm_op.run(ct_a, ct_b, ct_c); - - STEEL_PRAGMA_UNROLL - for (short i = 0; i < BaseNAXFrag::kElemsPerFrag; i++) { - C[i] = ct_c[i]; - } -} - -template < - typename CType, - typename AType, - typename BType, - bool transpose_a = false, - bool transpose_b = false, - mpp::tensor_ops::matmul2d_descriptor::mode Mode = - mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate> -METAL_FUNC static constexpr void mman( - thread BaseNAXFrag::dtype_frag_t& Cn0, - thread BaseNAXFrag::dtype_frag_t& Cn1, - const thread BaseNAXFrag::dtype_frag_t& A, - metal::bool_constant, - const thread BaseNAXFrag::dtype_frag_t& Bn0, - const thread BaseNAXFrag::dtype_frag_t& Bn1, - metal::bool_constant) { - // M=16, N=32, K=16: single A (K=16), B and C two N-fragments each. - constexpr auto desc = mpp::tensor_ops::matmul2d_descriptor( - 16, 32, 16, transpose_a, transpose_b, true, Mode); - - // Create matmul op - mpp::tensor_ops::matmul2d gemm_op; - - // Create matmul operands in registers - auto ct_a = - gemm_op.template get_left_input_cooperative_tensor(); - auto ct_b = - gemm_op - .template get_right_input_cooperative_tensor(); - - // Create matmul output in register - auto ct_c = gemm_op.template get_destination_cooperative_tensor< - decltype(ct_a), - decltype(ct_b), - CType>(); - - // Load A in to left operand registers - STEEL_PRAGMA_UNROLL - for (short i = 0; i < BaseNAXFrag::kElemsPerFrag; i++) { - ct_a[i] = A[i]; - ct_b[i] = Bn0[i]; - ct_b[BaseNAXFrag::kElemsPerFrag + i] = Bn1[i]; - ct_c[i] = Cn0[i]; - ct_c[BaseNAXFrag::kElemsPerFrag + i] = Cn1[i]; - } - - // Do matmul - gemm_op.run(ct_a, ct_b, ct_c); - - // Copy out results - STEEL_PRAGMA_UNROLL - for (short i = 0; i < BaseNAXFrag::kElemsPerFrag; i++) { - Cn0[i] = ct_c[i]; - Cn1[i] = ct_c[BaseNAXFrag::kElemsPerFrag + i]; - } -} - -} // namespace steel -} // namespace mlx - -#define MM16x16x16(C, CO, A, TA, AO, B, TB, BO) \ - mlx::steel::mma< \ - float, \ - float, \ - float, \ - TA, \ - TB, \ - mpp::tensor_ops::matmul2d_descriptor::mode::multiply>( \ - C.frag_at(0, (CO)), \ - A.frag_at(0, (AO)), \ - metal::bool_constant{}, \ - B.frag_at(0, (BO)), \ - metal::bool_constant{}); - -#define MMA16x16x16(C, CO, A, TA, AO, B, TB, BO) \ - mlx::steel::mma< \ - float, \ - float, \ - float, \ - TA, \ - TB, \ - mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate>( \ - C.frag_at(0, (CO)), \ - A.frag_at(0, (AO)), \ - metal::bool_constant{}, \ - B.frag_at(0, (BO)), \ - metal::bool_constant{}); - -#define MMA16x16x32(C, CO, A, TA, AO, B, TB, BO) \ - mlx::steel::mma< \ - float, \ - float, \ - float, \ - TA, \ - TB, \ - mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate>( \ - C.frag_at(0, (CO)), \ - A.frag_at(0, (AO)), \ - A.frag_at(0, (AO) + 1), \ - metal::bool_constant{}, \ - B.frag_at(0, (BO)), \ - B.frag_at(0, (BO) + 1), \ - metal::bool_constant{}); - -#define MM16x32x16(C, CO, A, TA, AO, B, TB, BO) \ - mlx::steel::mman< \ - float, \ - float, \ - float, \ - TA, \ - TB, \ - mpp::tensor_ops::matmul2d_descriptor::mode::multiply>( \ - C.frag_at(0, (CO)), \ - C.frag_at(0, (CO) + 1), \ - A.frag_at(0, (AO)), \ - metal::bool_constant{}, \ - B.frag_at(0, (BO)), \ - B.frag_at(0, (BO) + 1), \ - metal::bool_constant{}); - -#define MMA16x32x16(C, CO, A, TA, AO, B, TB, BO) \ - mlx::steel::mman< \ - float, \ - float, \ - float, \ - TA, \ - TB, \ - mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate>( \ - C.frag_at(0, (CO)), \ - C.frag_at(0, (CO) + 1), \ - A.frag_at(0, (AO)), \ - metal::bool_constant{}, \ - B.frag_at(0, (BO)), \ - B.frag_at(0, (BO) + 1), \ - metal::bool_constant{}); - -template -[[kernel]] void gated_delta_fused_nax( - const device InT* q [[buffer(0)]], - const device InT* k [[buffer(1)]], - const device InT* v [[buffer(2)]], - const device float* state_in [[buffer(3)]], - const device InT* g [[buffer(4)]], - const device InT* beta [[buffer(5)]], - device InT* y [[buffer(6)]], - device float* state_out [[buffer(7)]], - constant int& T [[buffer(8)]], - uint3 thread_position_in_grid [[thread_position_in_grid]], - uint3 thread_position_in_threadgroup [[thread_position_in_threadgroup]], - uint thread_index_in_simdgroup [[thread_index_in_simdgroup]]) { - auto n = thread_position_in_grid.z; - auto b_idx = n / Hv; - auto hv_idx = n % Hv; - auto hk_idx = hv_idx / (Hv / Hk); - - auto dv_idx = thread_position_in_grid.y * 16; - const short sg_id = thread_position_in_threadgroup.y; // 0..3 - - const ushort simd_lane_id = __metal_get_thread_index_in_simdgroup(ushort()); - const short qid = simd_lane_id >> 2; - const short fm = ((qid & 4) | ((simd_lane_id >> 1) & 3)); - - // set up pointers - // g: [B, T, Hv] - auto g_ = g + b_idx * T * Hv; - - // q, k: [B, T, Hk, Dk] - auto q_ = q + b_idx * T * Hk * Dk + hk_idx * Dk; - auto k_ = k + b_idx * T * Hk * Dk + hk_idx * Dk; - - // v, y: [B, T, Hv, Dv] - y += b_idx * T * Hv * Dv + hv_idx * Dv; - auto v_ = v + b_idx * T * Hv * Dv + hv_idx * Dv; - auto beta_ = beta + b_idx * T * Hv; - - // state_in, state_out: [B, Hv, Dv, Dk] - auto i_state = state_in + (n * Dv + dv_idx) * Dk; - auto o_state = state_out + (n * Dv + dv_idx) * Dk; - - threadgroup float gamma_all[C * 4]; - threadgroup float* gamma = gamma_all + sg_id * C; - - float beta_fm[2]; - - mlx::steel::NAXTile S_tile; - S_tile.load(i_state, Dk); - - mlx::steel::NAXTile K_tile, Q_tile; - mlx::steel::NAXTile W_tile; // panel - - mlx::steel::NAXTile V_tile; - mlx::steel::NAXTile U_tile; - mlx::steel::NAXTile WS_tile; - mlx::steel::NAXTile delta_tile; - mlx::steel::NAXTile tmp_tile; - mlx::steel::NAXTile QKt_tile; - mlx::steel::NAXTile out_tile; - mlx::steel::NAXTile Tinv_tile, P; - - mlx::steel::NAXTile KKtK_tile, KKt_tile; - - mlx::steel::NAXTile I_tile; - STEEL_PRAGMA_UNROLL - for (short _i = 0; _i < decltype(I_tile)::kElemsPerFrag; _i++) { - const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); /* {fn, fm} */ - const short _fn = _c.x; - const short _fm = _c.y; - AT_NAX(I_tile, _i) = (_fn == _fm) ? 1.0f : 0.0f; - } - mlx::steel::NAXTile TMP_tile; - - auto process_chunk = [&](const short valid_rows, - auto bounded_tag) __attribute__((always_inline)) { - constexpr bool B = decltype(bounded_tag)::value; - - auto load_seq = [&](thread auto& tile, auto src, int ld) { - if constexpr (B) { - tile.load_rows(src, ld, valid_rows); - } else { - tile.load(src, ld); - } - }; - - float g_val = (thread_index_in_simdgroup < (uint)valid_rows) - ? metal::fast::log( - metal::max(g_[thread_index_in_simdgroup * Hv + hv_idx], 1e-6)) - : 0.0f; - - auto gamma_val = simd_prefix_inclusive_sum(g_val); - if (thread_index_in_simdgroup < C) { - gamma[thread_index_in_simdgroup] = static_cast(gamma_val); - } - - beta_fm[0] = (fm < valid_rows) ? beta_[fm * Hv + hv_idx] : 0.0f; - const short fm1 = fm + mlx::steel::BaseNAXFrag::kElemRowsJump; - beta_fm[1] = (fm1 < valid_rows) ? beta_[fm1 * Hv + hv_idx] : 0.0f; - - KKt_tile.clear(); - for (int kk = 0; kk < Dk; kk += 32) { - load_seq(K_tile, k_ + kk, Dk * Hk); - MMA16x16x32(KKt_tile, 0, K_tile, false, 0, K_tile, true, 0); - } - - KKtK_tile = KKt_tile; - - SCALE_TRIEQ_NAX1(KKtK_tile, beta_fm); - SUB_NAX(Tinv_tile, I_tile, KKtK_tile); - STEEL_PRAGMA_UNROLL - for (int step = 0; step < 15; step++) { - MM16x16x16(TMP_tile, 0, KKtK_tile, false, 0, Tinv_tile, false, 0); - SUB_NAX(Tinv_tile, I_tile, TMP_tile); - } - - STEEL_PRAGMA_UNROLL - for (short nn = 0; nn < Dk / 16; nn += 2) { - load_seq(K_tile, k_ + nn * 16, Dk * Hk); - SCALE_BETA_NAX(K_tile, beta_fm); - MM16x32x16(W_tile, nn, Tinv_tile, false, 0, K_tile, false, 0); - } - SCALE_ROW_NAX(W_tile, gamma) - - SCALE_TRI_NAX(Tinv_tile, gamma) - load_seq(V_tile, v_ + dv_idx, Dv * Hv); - SCALE_BETA_NAX(V_tile, beta_fm); - MM16x16x16(U_tile, 0, Tinv_tile, false, 0, V_tile, false, 0) - - WS_tile.clear(); - STEEL_PRAGMA_UNROLL - for (short kk = 0; kk < Dk / 16; kk += 2) { - MMA16x16x32(WS_tile, 0, W_tile, false, kk, S_tile, true, kk) - } - - SUB_NAX(delta_tile, U_tile, WS_tile) - - tmp_tile.clear(); - QKt_tile.clear(); - for (int kk = 0; kk < Dk; kk += 32) { - load_seq(Q_tile, q_ + kk, Hk * Dk); - load_seq(K_tile, k_ + kk, Hk * Dk); - - MMA16x16x32(QKt_tile, 0, Q_tile, false, 0, K_tile, true, 0); - - SCALE_ROW_NAX(Q_tile, gamma); - MMA16x16x32(tmp_tile, 0, Q_tile, false, 0, S_tile, true, kk / 16); - } - - SCALE_TRI_NAX(QKt_tile, gamma) - - out_tile = tmp_tile; - MMA16x16x16(out_tile, 0, QKt_tile, false, 0, delta_tile, false, 0); - - STEEL_PRAGMA_UNROLL - for (short _i = 0; _i < decltype(out_tile)::kElemsPerFrag; _i++) { - const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); // {fn, fm} - const short _fn = _c.x; - const short _fm = _c.y; - if (_fm < valid_rows) { - y[_fm * Hv * Dv + dv_idx + _fn] = - static_cast(AT_NAX(out_tile, _i)); - } - } - - SCALE_NAX(S_tile, metal::fast::exp(gamma[C - 1])); - - for (int kk = 0; kk < Dk; kk += 32) { - load_seq(K_tile, k_ + kk, Hk * Dk); - SCALE2_NAX(K_tile, gamma); - MMA16x32x16(S_tile, kk / 16, delta_tile, true, 0, K_tile, false, 0); - } - }; - - int t = 0; - for (; t + C <= T; t += C) { - process_chunk(C, metal::false_type{}); - q_ += C * Hk * Dk; - k_ += C * Hk * Dk; - v_ += C * Hv * Dv; - beta_ += C * Hv; - y += C * Hv * Dv; - g_ += C * Hv; - } - if (t < T) { - process_chunk(short(T - t), metal::true_type{}); - } - - S_tile.store(o_state, Dk); -} - template [[kernel]] void gated_delta_fused_chunk( const device InT* q [[buffer(0)]], @@ -696,7 +177,7 @@ template float beta_fm = (fm < valid_rows) ? beta_[fm * Hv + hv_idx] : 0.0f; KKt_tile = make_filled_simdgroup_matrix(0.f); - STEEL_PRAGMA_UNROLL + FULL_UNROLL for (int kk = 0; kk < Dk; kk += 8) { load_M(K_tile, k_ + kk, Dk * Hk); load_MT(KT_tile, k_ + kk, Dk * Hk); @@ -711,14 +192,14 @@ template AT(P, 1) = AT(KKtK_tile, 1); SUB(Tinv, I_tile, KKtK_tile) - STEEL_PRAGMA_UNROLL + FULL_UNROLL for (int step = 1; (1 << step) < C; step++) { simdgroup_multiply(P, P, P); simdgroup_multiply_accumulate(Tinv, Tinv, P, Tinv); } WS_tile = make_filled_simdgroup_matrix(0.f); - STEEL_PRAGMA_UNROLL + FULL_UNROLL for (int kk = 0; kk < Dk; kk += 8) { load_M(K_tile, k_ + kk, Dk * Hk); SCALE(K_tile, beta_fm) @@ -736,7 +217,7 @@ template tmp_tile = make_filled_simdgroup_matrix(0.f); QKt_tile = make_filled_simdgroup_matrix(0.f); - STEEL_PRAGMA_UNROLL + FULL_UNROLL for (int kk = 0; kk < Dk; kk += 8) { load_M(Q_tile, q_ + kk, Hk * Dk); load_MT(K_tile, k_ + kk, Hk * Dk); @@ -754,7 +235,7 @@ template y[fm * Hv * Dv + dv_idx + fn + 1] = static_cast(AT(out_tile, 1)); } - STEEL_PRAGMA_UNROLL + FULL_UNROLL for (int kk = 0; kk < Dk; kk += 8) { load_MT(K_tile, k_ + kk, Hk * Dk); SCALE2(K_tile, gamma_Cdfn, gamma_Cdfn1) @@ -765,7 +246,7 @@ template }; int t = 0; - STEEL_PRAGMA_UNROLL + FULL_UNROLL for (; t + C <= T; t += C) { process_chunk(S_tile, C, metal::false_type{}); q_ += C * Hk * Dk; @@ -779,7 +260,7 @@ template process_chunk(S_tile, short(T - t), metal::true_type{}); } - STEEL_PRAGMA_UNROLL + FULL_UNROLL for (int kk = 0; kk < Dk; kk += 8) { simdgroup_store(S_tile[kk / 8], o_state + kk, Dk, ulong2(0, 0), true); } From b2bd2f45f5d96947f47e7baf274ba25c8a063e16 Mon Sep 17 00:00:00 2001 From: tpegolotti Date: Thu, 6 Aug 2026 05:22:46 -0700 Subject: [PATCH 47/56] Changed lambda function to macro --- .../metal/kernels/gated_delta_update_impl.h | 269 ++++++++++-------- 1 file changed, 145 insertions(+), 124 deletions(-) diff --git a/mlx/backend/metal/kernels/gated_delta_update_impl.h b/mlx/backend/metal/kernels/gated_delta_update_impl.h index 095a744a71..27558cdf48 100644 --- a/mlx/backend/metal/kernels/gated_delta_update_impl.h +++ b/mlx/backend/metal/kernels/gated_delta_update_impl.h @@ -46,6 +46,149 @@ AT(TILE0, 1) *= fn + 1 >= fm ? 0.f : S1; \ } +// lambdas are not supported in metal 14 so porting to macros. + +// non transposed +#define LOAD_M(M, SRC, LD, B) \ + if constexpr (B) { \ + AT(M, 0) = \ + static_cast((fm < valid_rows) ? ((SRC)[fm * (LD) + fn]) : 0.f); \ + AT(M, 1) = static_cast( \ + (fm < valid_rows) ? ((SRC)[fm * (LD) + fn + 1]) : 0.f); \ + } else { \ + AT(M, 0) = static_cast((SRC)[fm * (LD) + fn]); \ + AT(M, 1) = static_cast((SRC)[fm * (LD) + fn + 1]); \ + } + +// transposed load: sequence is the column -> mask fn / fn+1 +#define LOAD_MT(M, SRC, LD, B) \ + if constexpr (B) { \ + AT(M, 0) = \ + static_cast((fn < valid_rows) ? ((SRC)[fn * (LD) + fm]) : 0.f); \ + AT(M, 1) = static_cast( \ + (fn + 1 < valid_rows) ? ((SRC)[(fn + 1) * (LD) + fm]) : 0.f); \ + } else { \ + AT(M, 0) = static_cast((SRC)[fn * (LD) + fm]); \ + AT(M, 1) = static_cast((SRC)[(fn + 1) * (LD) + fm]); \ + } + +// non-transposed load +#define LOAD_M(M, SRC, LD, B) \ + if constexpr (B) { \ + AT(M, 0) = \ + static_cast((fm < valid_rows) ? ((SRC)[fm * (LD) + fn]) : 0.f); \ + AT(M, 1) = static_cast( \ + (fm < valid_rows) ? ((SRC)[fm * (LD) + fn + 1]) : 0.f); \ + } else { \ + AT(M, 0) = static_cast((SRC)[fm * (LD) + fn]); \ + AT(M, 1) = static_cast((SRC)[fm * (LD) + fn + 1]); \ + } + +// transposed load +#define LOAD_MT(M, SRC, LD, B) \ + if constexpr (B) { \ + AT(M, 0) = \ + static_cast((fn < valid_rows) ? ((SRC)[fn * (LD) + fm]) : 0.f); \ + AT(M, 1) = static_cast( \ + (fn + 1 < valid_rows) ? ((SRC)[(fn + 1) * (LD) + fm]) : 0.f); \ + } else { \ + AT(M, 0) = static_cast((SRC)[fn * (LD) + fm]); \ + AT(M, 1) = static_cast((SRC)[(fn + 1) * (LD) + fm]); \ + } + +#define PROCESS_CHUNK_SG(B, S_tile, VALID) \ + { \ + const short valid_rows = (VALID); \ + \ + float g_val = (thread_index_in_simdgroup < (uint)valid_rows) \ + ? metal::fast::log( \ + metal::max(g_[thread_index_in_simdgroup * Hv + hv_idx], 1e-6)) \ + : 0.0f; \ + \ + float gamma_val = simd_prefix_inclusive_sum(g_val); \ + \ + if (thread_index_in_simdgroup < C) { \ + gamma[thread_index_in_simdgroup] = gamma_val; \ + } \ + \ + float gamma_fm = metal::fast::exp(gamma[fm]); \ + float gamma_fmdfn = metal::fast::exp(gamma[fm] - gamma[fn]); \ + float gamma_fmdfn1 = metal::fast::exp(gamma[fm] - gamma[fn + 1]); \ + float gamma_Cdfn = metal::fast::exp(gamma[C - 1] - gamma[fn]); \ + float gamma_Cdfn1 = metal::fast::exp(gamma[C - 1] - gamma[fn + 1]); \ + float gamma_C = metal::fast::exp(gamma[C - 1]); \ + \ + float beta_fm = (fm < valid_rows) ? beta_[fm * Hv + hv_idx] : 0.0f; \ + \ + KKt_tile = make_filled_simdgroup_matrix(0.f); \ + FULL_UNROLL \ + for (int kk = 0; kk < Dk; kk += 8) { \ + LOAD_M(K_tile, k_ + kk, Dk * Hk, B) \ + LOAD_MT(KT_tile, k_ + kk, Dk * Hk, B) \ + simdgroup_multiply_accumulate(KKt_tile, K_tile, KT_tile, KKt_tile); \ + } \ + \ + KKtK_tile = KKt_tile; \ + SCALE_TRIEQ(KKtK_tile, beta_fm, beta_fm) \ + \ + simdgroup_float8x8 Tinv, P; \ + AT(P, 0) = AT(KKtK_tile, 0); \ + AT(P, 1) = AT(KKtK_tile, 1); \ + SUB(Tinv, I_tile, KKtK_tile) \ + \ + FULL_UNROLL \ + for (int step = 1; (1 << step) < C; step++) { \ + simdgroup_multiply(P, P, P); \ + simdgroup_multiply_accumulate(Tinv, Tinv, P, Tinv); \ + } \ + \ + WS_tile = make_filled_simdgroup_matrix(0.f); \ + FULL_UNROLL \ + for (int kk = 0; kk < Dk; kk += 8) { \ + LOAD_M(K_tile, k_ + kk, Dk * Hk, B) \ + SCALE(K_tile, beta_fm) \ + simdgroup_multiply(W_tile, Tinv, K_tile); \ + SCALE(W_tile, gamma_fm) \ + simdgroup_multiply_accumulate(WS_tile, W_tile, S_tile[kk / 8], WS_tile); \ + } \ + \ + SCALE_TRI(Tinv, gamma_fmdfn, gamma_fmdfn1) \ + \ + LOAD_M(V_tile, v_ + dv_idx, Dv * Hv, B) \ + SCALE(V_tile, beta_fm) \ + simdgroup_multiply(U_tile, Tinv, V_tile); \ + SUB(delta_tile, U_tile, WS_tile) \ + \ + tmp_tile = make_filled_simdgroup_matrix(0.f); \ + QKt_tile = make_filled_simdgroup_matrix(0.f); \ + FULL_UNROLL \ + for (int kk = 0; kk < Dk; kk += 8) { \ + LOAD_M(Q_tile, q_ + kk, Hk * Dk, B) \ + LOAD_MT(K_tile, k_ + kk, Hk * Dk, B) \ + simdgroup_multiply_accumulate(QKt_tile, Q_tile, K_tile, QKt_tile); \ + SCALE(Q_tile, gamma_fm) \ + simdgroup_multiply_accumulate( \ + tmp_tile, Q_tile, S_tile[kk / 8], tmp_tile); \ + } \ + \ + SCALE_TRI(QKt_tile, gamma_fmdfn, gamma_fmdfn1) \ + \ + simdgroup_multiply_accumulate(out_tile, QKt_tile, delta_tile, tmp_tile); \ + \ + if (fm < valid_rows) { \ + y[fm * Hv * Dv + dv_idx + fn] = static_cast(AT(out_tile, 0)); \ + y[fm * Hv * Dv + dv_idx + fn + 1] = static_cast(AT(out_tile, 1)); \ + } \ + \ + FULL_UNROLL \ + for (int kk = 0; kk < Dk; kk += 8) { \ + LOAD_MT(K_tile, k_ + kk, Hk * Dk, B) \ + SCALE2(K_tile, gamma_Cdfn, gamma_Cdfn1) \ + simdgroup_multiply(KD_tile, K_tile, delta_tile); \ + FMA(S_tile[kk / 8], gamma_C, S_tile[kk / 8], KD_tile) \ + } \ + } + template [[kernel]] void gated_delta_fused_chunk( const device InT* q [[buffer(0)]], @@ -123,132 +266,10 @@ template simdgroup_load(S_tile[kk / 8], i_state + kk, Dk, ulong2(0, 0), true); } - auto process_chunk = [&](thread simdgroup_float8x8* S_tile, - const short valid_rows, - auto bounded_tag) __attribute__((always_inline)) { - constexpr bool B = decltype(bounded_tag)::value; - - // non-transposed load - auto load_M = [&](thread simdgroup_float8x8& M, auto src, int ld) { - if constexpr (B) { - AT(M, 0) = - static_cast((fm < valid_rows) ? (src[fm * ld + fn]) : 0.f); - AT(M, 1) = static_cast( - (fm < valid_rows) ? (src[fm * ld + fn + 1]) : 0.f); - } else { - // simdgroup_load(M, src, ld); - AT(M, 0) = static_cast(src[fm * ld + fn]); - AT(M, 1) = static_cast(src[fm * ld + fn + 1]); - } - }; - - // transposed load - auto load_MT = [&](thread simdgroup_float8x8& M, auto src, int ld) { - if constexpr (B) { - AT(M, 0) = - static_cast((fn < valid_rows) ? (src[fn * ld + fm]) : 0.f); - AT(M, 1) = static_cast( - (fn + 1 < valid_rows) ? (src[(fn + 1) * ld + fm]) : 0.f); - } else { - // simdgroup_load(M, src, ld, ulong2(0, 0), true); - AT(M, 0) = static_cast(src[fn * ld + fm]); - AT(M, 1) = static_cast(src[(fn + 1) * ld + fm]); - } - }; - - float g_val = (thread_index_in_simdgroup < (uint)valid_rows) - ? metal::fast::log( - metal::max(g_[thread_index_in_simdgroup * Hv + hv_idx], 1e-6)) - : 0.0f; - - float gamma_val = simd_prefix_inclusive_sum(g_val); - - if (thread_index_in_simdgroup < C) { - gamma[thread_index_in_simdgroup] = gamma_val; - } - - float gamma_fm = metal::fast::exp(gamma[fm]); - float gamma_fmdfn = metal::fast::exp(gamma[fm] - gamma[fn]); - float gamma_fmdfn1 = metal::fast::exp(gamma[fm] - gamma[fn + 1]); - float gamma_Cdfn = metal::fast::exp(gamma[C - 1] - gamma[fn]); - float gamma_Cdfn1 = metal::fast::exp(gamma[C - 1] - gamma[fn + 1]); - float gamma_C = metal::fast::exp(gamma[C - 1]); - - float beta_fm = (fm < valid_rows) ? beta_[fm * Hv + hv_idx] : 0.0f; - - KKt_tile = make_filled_simdgroup_matrix(0.f); - FULL_UNROLL - for (int kk = 0; kk < Dk; kk += 8) { - load_M(K_tile, k_ + kk, Dk * Hk); - load_MT(KT_tile, k_ + kk, Dk * Hk); - simdgroup_multiply_accumulate(KKt_tile, K_tile, KT_tile, KKt_tile); - } - - KKtK_tile = KKt_tile; - SCALE_TRIEQ(KKtK_tile, beta_fm, beta_fm) - - simdgroup_float8x8 Tinv, P; - AT(P, 0) = AT(KKtK_tile, 0); - AT(P, 1) = AT(KKtK_tile, 1); - SUB(Tinv, I_tile, KKtK_tile) - - FULL_UNROLL - for (int step = 1; (1 << step) < C; step++) { - simdgroup_multiply(P, P, P); - simdgroup_multiply_accumulate(Tinv, Tinv, P, Tinv); - } - - WS_tile = make_filled_simdgroup_matrix(0.f); - FULL_UNROLL - for (int kk = 0; kk < Dk; kk += 8) { - load_M(K_tile, k_ + kk, Dk * Hk); - SCALE(K_tile, beta_fm) - simdgroup_multiply(W_tile, Tinv, K_tile); - SCALE(W_tile, gamma_fm) - simdgroup_multiply_accumulate(WS_tile, W_tile, S_tile[kk / 8], WS_tile); - } - - SCALE_TRI(Tinv, gamma_fmdfn, gamma_fmdfn1) - - load_M(V_tile, v_ + dv_idx, Dv * Hv); - SCALE(V_tile, beta_fm) - simdgroup_multiply(U_tile, Tinv, V_tile); - SUB(delta_tile, U_tile, WS_tile) - - tmp_tile = make_filled_simdgroup_matrix(0.f); - QKt_tile = make_filled_simdgroup_matrix(0.f); - FULL_UNROLL - for (int kk = 0; kk < Dk; kk += 8) { - load_M(Q_tile, q_ + kk, Hk * Dk); - load_MT(K_tile, k_ + kk, Hk * Dk); - simdgroup_multiply_accumulate(QKt_tile, Q_tile, K_tile, QKt_tile); - SCALE(Q_tile, gamma_fm) - simdgroup_multiply_accumulate(tmp_tile, Q_tile, S_tile[kk / 8], tmp_tile); - } - - SCALE_TRI(QKt_tile, gamma_fmdfn, gamma_fmdfn1) - - simdgroup_multiply_accumulate(out_tile, QKt_tile, delta_tile, tmp_tile); - - if (fm < valid_rows) { - y[fm * Hv * Dv + dv_idx + fn] = static_cast(AT(out_tile, 0)); - y[fm * Hv * Dv + dv_idx + fn + 1] = static_cast(AT(out_tile, 1)); - } - - FULL_UNROLL - for (int kk = 0; kk < Dk; kk += 8) { - load_MT(K_tile, k_ + kk, Hk * Dk); - SCALE2(K_tile, gamma_Cdfn, gamma_Cdfn1) - - simdgroup_multiply(KD_tile, K_tile, delta_tile); - FMA(S_tile[kk / 8], gamma_C, S_tile[kk / 8], KD_tile) - } - }; - int t = 0; FULL_UNROLL for (; t + C <= T; t += C) { - process_chunk(S_tile, C, metal::false_type{}); + PROCESS_CHUNK_SG(false, S_tile, C); q_ += C * Hk * Dk; k_ += C * Hk * Dk; v_ += C * Hv * Dv; @@ -257,7 +278,7 @@ template g_ += C * Hv; } if (t < T) { - process_chunk(S_tile, short(T - t), metal::true_type{}); + PROCESS_CHUNK_SG(true, S_tile, short(T - t)); } FULL_UNROLL From 5f5065e070161b1bd469cc5d44819cf9d5cdc860 Mon Sep 17 00:00:00 2001 From: tpegolotti Date: Thu, 6 Aug 2026 05:27:53 -0700 Subject: [PATCH 48/56] adding nax files --- .../metal/kernels/gated_delta_update_nax.h | 524 ++++++++++++++++++ .../kernels/gated_delta_update_nax.metal | 28 + 2 files changed, 552 insertions(+) create mode 100644 mlx/backend/metal/kernels/gated_delta_update_nax.h create mode 100644 mlx/backend/metal/kernels/gated_delta_update_nax.metal diff --git a/mlx/backend/metal/kernels/gated_delta_update_nax.h b/mlx/backend/metal/kernels/gated_delta_update_nax.h new file mode 100644 index 0000000000..54b0abc4a1 --- /dev/null +++ b/mlx/backend/metal/kernels/gated_delta_update_nax.h @@ -0,0 +1,524 @@ +#pragma once + +#include +#include "mlx/backend/metal/kernels/utils.h" + +#include +#include + +#include "mlx/backend/metal/kernels/steel/gemm/nax.h" + +using namespace metal; +using namespace mpp; +using namespace mpp::tensor_ops; + +// NAX MACROS I can probably do a nice template instead of doing this +// fm = base_fm + (idx >> 2) * 8; // idx>>2 = idx/4 -> 0 for idx 0-3, 1 for +// idx 4-7 fn = base_fn + (idx % 4); // 4 consecutive columns +#define AT_NAX(TILE, IDX) TILE.elems()[IDX] + +#define SUB_NAX(TILE0, TILE1, TILE2) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerFrag; _i++) { \ + AT_NAX(TILE0, _i) = AT_NAX(TILE1, _i) - AT_NAX(TILE2, _i); \ + } \ + } + +#define ADD_NAX(TILE0, TILE1, TILE2) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerFrag; _i++) { \ + AT_NAX(TILE0, _i) = AT_NAX(TILE1, _i) + AT_NAX(TILE2, _i); \ + } \ + } + +#define FMA_NAX(TILE0, S, TILE1, TILE2) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < mlx::steel::BaseNAXFrag::kElemsPerFrag; _i++) { \ + (TILE0)[_i] = (S) * (TILE1)[_i] + (TILE2)[_i]; \ + } \ + } + +#define SCALE_NAX(TILE0, S) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + AT_NAX(TILE0, _i) *= (S); \ + } \ + } + +#define SCALE_ROW_NAX(TILE0, S) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; \ + AT_NAX(TILE0, _i) *= \ + metal::fast::exp((S)[mlx::steel::BaseNAXFrag::get_coord(_w).y]); \ + } \ + } + +#define SCALE_BETA_NAX(TILE0, BETA2) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; \ + AT_NAX(TILE0, _i) *= (BETA2)[_w >> 2]; \ + } \ + } + +#define SCALE2_NAX(TILE0, GAMMA) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; \ + const short _fm = mlx::steel::BaseNAXFrag::get_coord(_w).y; \ + AT_NAX(TILE0, _i) *= metal::fast::exp((GAMMA)[(C) - 1] - (GAMMA)[_fm]); \ + } \ + } + +#define SCALE_TRI_NAX(TILE0, GAMMA) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerFrag; _i++) { \ + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); /* {fn, fm} */ \ + AT_NAX(TILE0, _i) *= (_c.x > _c.y) \ + ? 0.f \ + : metal::fast::exp((GAMMA)[_c.y] - (GAMMA)[_c.x]); \ + } \ + } + +#define SCALE_TRIEQ_NAX1(TILE0, BETA) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerFrag; _i++) { \ + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); /* {fn, fm} */ \ + AT_NAX(TILE0, _i) *= (_c.x >= _c.y) ? 0.f : (BETA)[_i >> 2]; \ + } \ + } + +namespace mlx { +namespace steel { +template < + typename CType, + typename AType, + typename BType, + bool transpose_a = false, + bool transpose_b = false, + mpp::tensor_ops::matmul2d_descriptor::mode Mode = + mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate> +METAL_FUNC static constexpr void mma( + thread BaseNAXFrag::dtype_frag_t& C, + const thread BaseNAXFrag::dtype_frag_t& A0, + const thread BaseNAXFrag::dtype_frag_t& A1, + metal::bool_constant, + const thread BaseNAXFrag::dtype_frag_t& B0, + const thread BaseNAXFrag::dtype_frag_t& B1, + metal::bool_constant) { + // M=16, N=16, K=32: A and B each two K-fragments, single 16x16 C. + constexpr auto desc = mpp::tensor_ops::matmul2d_descriptor( + 16, 16, 32, transpose_a, transpose_b, true, Mode); + + mpp::tensor_ops::matmul2d gemm_op; + + auto ct_a = + gemm_op.template get_left_input_cooperative_tensor(); + auto ct_b = + gemm_op + .template get_right_input_cooperative_tensor(); + auto ct_c = gemm_op.template get_destination_cooperative_tensor< + decltype(ct_a), + decltype(ct_b), + CType>(); + + STEEL_PRAGMA_UNROLL + for (short i = 0; i < BaseNAXFrag::kElemsPerFrag; i++) { + ct_a[i] = A0[i]; + ct_a[BaseNAXFrag::kElemsPerFrag + i] = A1[i]; + ct_b[i] = B0[i]; + ct_b[BaseNAXFrag::kElemsPerFrag + i] = B1[i]; + ct_c[i] = C[i]; + } + + gemm_op.run(ct_a, ct_b, ct_c); + + STEEL_PRAGMA_UNROLL + for (short i = 0; i < BaseNAXFrag::kElemsPerFrag; i++) { + C[i] = ct_c[i]; + } +} + +template < + typename CType, + typename AType, + typename BType, + bool transpose_a, + bool transpose_b, + mpp::tensor_ops::matmul2d_descriptor::mode Mode = + mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate> +METAL_FUNC static constexpr void mma( + thread BaseNAXFrag::dtype_frag_t& C, + const thread BaseNAXFrag::dtype_frag_t& A, + metal::bool_constant, + const thread BaseNAXFrag::dtype_frag_t& B, + metal::bool_constant) { + constexpr auto desc = mpp::tensor_ops::matmul2d_descriptor( + 16, 32, 16, transpose_a, transpose_b, true, Mode); + + mpp::tensor_ops::matmul2d gemm_op; + + auto ct_a = + gemm_op.template get_left_input_cooperative_tensor(); + auto ct_b = + gemm_op + .template get_right_input_cooperative_tensor(); + auto ct_c = gemm_op.template get_destination_cooperative_tensor< + decltype(ct_a), + decltype(ct_b), + CType>(); + + STEEL_PRAGMA_UNROLL + for (short i = 0; i < BaseNAXFrag::kElemsPerFrag; i++) { + ct_a[i] = A[i]; + ct_b[i] = B[i]; + ct_b[BaseNAXFrag::kElemsPerFrag + i] = 0.0; + ct_c[i] = C[i]; + ct_c[BaseNAXFrag::kElemsPerFrag + i] = 0.0; + } + + gemm_op.run(ct_a, ct_b, ct_c); + + STEEL_PRAGMA_UNROLL + for (short i = 0; i < BaseNAXFrag::kElemsPerFrag; i++) { + C[i] = ct_c[i]; + } +} + +template < + typename CType, + typename AType, + typename BType, + bool transpose_a = false, + bool transpose_b = false, + mpp::tensor_ops::matmul2d_descriptor::mode Mode = + mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate> +METAL_FUNC static constexpr void mman( + thread BaseNAXFrag::dtype_frag_t& Cn0, + thread BaseNAXFrag::dtype_frag_t& Cn1, + const thread BaseNAXFrag::dtype_frag_t& A, + metal::bool_constant, + const thread BaseNAXFrag::dtype_frag_t& Bn0, + const thread BaseNAXFrag::dtype_frag_t& Bn1, + metal::bool_constant) { + // M=16, N=32, K=16: single A (K=16), B and C two N-fragments each. + constexpr auto desc = mpp::tensor_ops::matmul2d_descriptor( + 16, 32, 16, transpose_a, transpose_b, true, Mode); + + // Create matmul op + mpp::tensor_ops::matmul2d gemm_op; + + // Create matmul operands in registers + auto ct_a = + gemm_op.template get_left_input_cooperative_tensor(); + auto ct_b = + gemm_op + .template get_right_input_cooperative_tensor(); + + // Create matmul output in register + auto ct_c = gemm_op.template get_destination_cooperative_tensor< + decltype(ct_a), + decltype(ct_b), + CType>(); + + // Load A in to left operand registers + STEEL_PRAGMA_UNROLL + for (short i = 0; i < BaseNAXFrag::kElemsPerFrag; i++) { + ct_a[i] = A[i]; + ct_b[i] = Bn0[i]; + ct_b[BaseNAXFrag::kElemsPerFrag + i] = Bn1[i]; + ct_c[i] = Cn0[i]; + ct_c[BaseNAXFrag::kElemsPerFrag + i] = Cn1[i]; + } + + // Do matmul + gemm_op.run(ct_a, ct_b, ct_c); + + // Copy out results + STEEL_PRAGMA_UNROLL + for (short i = 0; i < BaseNAXFrag::kElemsPerFrag; i++) { + Cn0[i] = ct_c[i]; + Cn1[i] = ct_c[BaseNAXFrag::kElemsPerFrag + i]; + } +} + +} // namespace steel +} // namespace mlx + +#define MM16x16x16(C, CO, A, TA, AO, B, TB, BO) \ + mlx::steel::mma< \ + float, \ + float, \ + float, \ + TA, \ + TB, \ + mpp::tensor_ops::matmul2d_descriptor::mode::multiply>( \ + C.frag_at(0, (CO)), \ + A.frag_at(0, (AO)), \ + metal::bool_constant{}, \ + B.frag_at(0, (BO)), \ + metal::bool_constant{}); + +#define MMA16x16x16(C, CO, A, TA, AO, B, TB, BO) \ + mlx::steel::mma< \ + float, \ + float, \ + float, \ + TA, \ + TB, \ + mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate>( \ + C.frag_at(0, (CO)), \ + A.frag_at(0, (AO)), \ + metal::bool_constant{}, \ + B.frag_at(0, (BO)), \ + metal::bool_constant{}); + +#define MMA16x16x32(C, CO, A, TA, AO, B, TB, BO) \ + mlx::steel::mma< \ + float, \ + float, \ + float, \ + TA, \ + TB, \ + mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate>( \ + C.frag_at(0, (CO)), \ + A.frag_at(0, (AO)), \ + A.frag_at(0, (AO) + 1), \ + metal::bool_constant{}, \ + B.frag_at(0, (BO)), \ + B.frag_at(0, (BO) + 1), \ + metal::bool_constant{}); + +#define MM16x32x16(C, CO, A, TA, AO, B, TB, BO) \ + mlx::steel::mman< \ + float, \ + float, \ + float, \ + TA, \ + TB, \ + mpp::tensor_ops::matmul2d_descriptor::mode::multiply>( \ + C.frag_at(0, (CO)), \ + C.frag_at(0, (CO) + 1), \ + A.frag_at(0, (AO)), \ + metal::bool_constant{}, \ + B.frag_at(0, (BO)), \ + B.frag_at(0, (BO) + 1), \ + metal::bool_constant{}); + +#define MMA16x32x16(C, CO, A, TA, AO, B, TB, BO) \ + mlx::steel::mman< \ + float, \ + float, \ + float, \ + TA, \ + TB, \ + mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate>( \ + C.frag_at(0, (CO)), \ + C.frag_at(0, (CO) + 1), \ + A.frag_at(0, (AO)), \ + metal::bool_constant{}, \ + B.frag_at(0, (BO)), \ + B.frag_at(0, (BO) + 1), \ + metal::bool_constant{}); + +template +[[kernel]] void gated_delta_fused_nax( + const device InT* q [[buffer(0)]], + const device InT* k [[buffer(1)]], + const device InT* v [[buffer(2)]], + const device float* state_in [[buffer(3)]], + const device InT* g [[buffer(4)]], + const device InT* beta [[buffer(5)]], + device InT* y [[buffer(6)]], + device float* state_out [[buffer(7)]], + constant int& T [[buffer(8)]], + uint3 thread_position_in_grid [[thread_position_in_grid]], + uint3 thread_position_in_threadgroup [[thread_position_in_threadgroup]], + uint thread_index_in_simdgroup [[thread_index_in_simdgroup]]) { + auto n = thread_position_in_grid.z; + auto b_idx = n / Hv; + auto hv_idx = n % Hv; + auto hk_idx = hv_idx / (Hv / Hk); + + auto dv_idx = thread_position_in_grid.y * 16; + const short sg_id = thread_position_in_threadgroup.y; // 0..3 + + const ushort simd_lane_id = __metal_get_thread_index_in_simdgroup(ushort()); + const short qid = simd_lane_id >> 2; + const short fm = ((qid & 4) | ((simd_lane_id >> 1) & 3)); + + // set up pointers + // g: [B, T, Hv] + auto g_ = g + b_idx * T * Hv; + + // q, k: [B, T, Hk, Dk] + auto q_ = q + b_idx * T * Hk * Dk + hk_idx * Dk; + auto k_ = k + b_idx * T * Hk * Dk + hk_idx * Dk; + + // v, y: [B, T, Hv, Dv] + y += b_idx * T * Hv * Dv + hv_idx * Dv; + auto v_ = v + b_idx * T * Hv * Dv + hv_idx * Dv; + auto beta_ = beta + b_idx * T * Hv; + + // state_in, state_out: [B, Hv, Dv, Dk] + auto i_state = state_in + (n * Dv + dv_idx) * Dk; + auto o_state = state_out + (n * Dv + dv_idx) * Dk; + + threadgroup float gamma_all[C * 4]; + threadgroup float* gamma = gamma_all + sg_id * C; + + float beta_fm[2]; + + mlx::steel::NAXTile S_tile; + S_tile.load(i_state, Dk); + + mlx::steel::NAXTile K_tile, Q_tile; + mlx::steel::NAXTile W_tile; // panel + + mlx::steel::NAXTile V_tile; + mlx::steel::NAXTile U_tile; + mlx::steel::NAXTile WS_tile; + mlx::steel::NAXTile delta_tile; + mlx::steel::NAXTile tmp_tile; + mlx::steel::NAXTile QKt_tile; + mlx::steel::NAXTile out_tile; + mlx::steel::NAXTile Tinv_tile, P; + + mlx::steel::NAXTile KKtK_tile, KKt_tile; + + mlx::steel::NAXTile I_tile; + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(I_tile)::kElemsPerFrag; _i++) { + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); /* {fn, fm} */ + const short _fn = _c.x; + const short _fm = _c.y; + AT_NAX(I_tile, _i) = (_fn == _fm) ? 1.0f : 0.0f; + } + mlx::steel::NAXTile TMP_tile; + + auto process_chunk = [&](const short valid_rows, + auto bounded_tag) __attribute__((always_inline)) { + constexpr bool B = decltype(bounded_tag)::value; + + auto load_seq = [&](thread auto& tile, auto src, int ld) { + if constexpr (B) { + tile.load_rows(src, ld, valid_rows); + } else { + tile.load(src, ld); + } + }; + + float g_val = (thread_index_in_simdgroup < (uint)valid_rows) + ? metal::fast::log( + metal::max(g_[thread_index_in_simdgroup * Hv + hv_idx], 1e-6)) + : 0.0f; + + auto gamma_val = simd_prefix_inclusive_sum(g_val); + if (thread_index_in_simdgroup < C) { + gamma[thread_index_in_simdgroup] = static_cast(gamma_val); + } + + beta_fm[0] = (fm < valid_rows) ? beta_[fm * Hv + hv_idx] : 0.0f; + const short fm1 = fm + mlx::steel::BaseNAXFrag::kElemRowsJump; + beta_fm[1] = (fm1 < valid_rows) ? beta_[fm1 * Hv + hv_idx] : 0.0f; + + KKt_tile.clear(); + for (int kk = 0; kk < Dk; kk += 32) { + load_seq(K_tile, k_ + kk, Dk * Hk); + MMA16x16x32(KKt_tile, 0, K_tile, false, 0, K_tile, true, 0); + } + + KKtK_tile = KKt_tile; + + SCALE_TRIEQ_NAX1(KKtK_tile, beta_fm); + SUB_NAX(Tinv_tile, I_tile, KKtK_tile); + STEEL_PRAGMA_UNROLL + for (int step = 0; step < 15; step++) { + MM16x16x16(TMP_tile, 0, KKtK_tile, false, 0, Tinv_tile, false, 0); + SUB_NAX(Tinv_tile, I_tile, TMP_tile); + } + + STEEL_PRAGMA_UNROLL + for (short nn = 0; nn < Dk / 16; nn += 2) { + load_seq(K_tile, k_ + nn * 16, Dk * Hk); + SCALE_BETA_NAX(K_tile, beta_fm); + MM16x32x16(W_tile, nn, Tinv_tile, false, 0, K_tile, false, 0); + } + SCALE_ROW_NAX(W_tile, gamma) + + SCALE_TRI_NAX(Tinv_tile, gamma) + load_seq(V_tile, v_ + dv_idx, Dv * Hv); + SCALE_BETA_NAX(V_tile, beta_fm); + MM16x16x16(U_tile, 0, Tinv_tile, false, 0, V_tile, false, 0) + + WS_tile.clear(); + STEEL_PRAGMA_UNROLL + for (short kk = 0; kk < Dk / 16; kk += 2) { + MMA16x16x32(WS_tile, 0, W_tile, false, kk, S_tile, true, kk) + } + + SUB_NAX(delta_tile, U_tile, WS_tile) + + tmp_tile.clear(); + QKt_tile.clear(); + for (int kk = 0; kk < Dk; kk += 32) { + load_seq(Q_tile, q_ + kk, Hk * Dk); + load_seq(K_tile, k_ + kk, Hk * Dk); + + MMA16x16x32(QKt_tile, 0, Q_tile, false, 0, K_tile, true, 0); + + SCALE_ROW_NAX(Q_tile, gamma); + MMA16x16x32(tmp_tile, 0, Q_tile, false, 0, S_tile, true, kk / 16); + } + + SCALE_TRI_NAX(QKt_tile, gamma) + + out_tile = tmp_tile; + MMA16x16x16(out_tile, 0, QKt_tile, false, 0, delta_tile, false, 0); + + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(out_tile)::kElemsPerFrag; _i++) { + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); // {fn, fm} + const short _fn = _c.x; + const short _fm = _c.y; + if (_fm < valid_rows) { + y[_fm * Hv * Dv + dv_idx + _fn] = + static_cast(AT_NAX(out_tile, _i)); + } + } + + SCALE_NAX(S_tile, metal::fast::exp(gamma[C - 1])); + + for (int kk = 0; kk < Dk; kk += 32) { + load_seq(K_tile, k_ + kk, Hk * Dk); + SCALE2_NAX(K_tile, gamma); + MMA16x32x16(S_tile, kk / 16, delta_tile, true, 0, K_tile, false, 0); + } + }; + + int t = 0; + for (; t + C <= T; t += C) { + process_chunk(C, metal::false_type{}); + q_ += C * Hk * Dk; + k_ += C * Hk * Dk; + v_ += C * Hv * Dv; + beta_ += C * Hv; + y += C * Hv * Dv; + g_ += C * Hv; + } + if (t < T) { + process_chunk(short(T - t), metal::true_type{}); + } + + S_tile.store(o_state, Dk); +} \ No newline at end of file diff --git a/mlx/backend/metal/kernels/gated_delta_update_nax.metal b/mlx/backend/metal/kernels/gated_delta_update_nax.metal new file mode 100644 index 0000000000..b17707fb40 --- /dev/null +++ b/mlx/backend/metal/kernels/gated_delta_update_nax.metal @@ -0,0 +1,28 @@ +#include "mlx/backend/metal/kernels/gated_delta_update_nax.h" +#include "mlx/backend/metal/kernels/utils.h" + +using namespace metal; + +#define instantiate_gated_delta_update_fused_nax(in_type, dk, dv, hk, hv, c) \ + instantiate_kernel( \ + "gated_delta_fused_nax_" #in_type "_" #dk "_" #dv "_" #hk "_" #hv \ + "_" #c, \ + gated_delta_fused_nax, \ + in_type, \ + dk, \ + dv, \ + hk, \ + hv, \ + c) + +#define instantiate_gated_delta_dims(in_type, dk, dv, hk, hv) \ + instantiate_gated_delta_update_fused_nax(in_type, dk, dv, hk, hv, 16) + +#define instantiate_gated_delta(in_type) \ + instantiate_gated_delta_dims(in_type, 128, 128, 24, 24) \ + instantiate_gated_delta_dims(in_type, 128, 128, 32, 32) \ + instantiate_gated_delta_dims(in_type, 128, 128, 16, 32) \ + instantiate_gated_delta_dims(in_type, 128, 128, 16, 48) + +instantiate_gated_delta(float); +instantiate_gated_delta(bfloat16_t); \ No newline at end of file From b9000984afdde9b22ca873beb15ca5a4531e034f Mon Sep 17 00:00:00 2001 From: tpegolotti Date: Thu, 6 Aug 2026 05:45:41 -0700 Subject: [PATCH 49/56] Updating kernel getters --- mlx/backend/metal/gated_delta_update.cpp | 12 ++++++------ mlx/backend/metal/jit_kernels.cpp | 14 ++++++++++++++ mlx/backend/metal/kernels.h | 2 +- mlx/backend/metal/nojit_kernels.cpp | 2 +- 4 files changed, 22 insertions(+), 8 deletions(-) diff --git a/mlx/backend/metal/gated_delta_update.cpp b/mlx/backend/metal/gated_delta_update.cpp index 1572cf7676..23c3f1e553 100644 --- a/mlx/backend/metal/gated_delta_update.cpp +++ b/mlx/backend/metal/gated_delta_update.cpp @@ -143,8 +143,8 @@ void GatedDeltaUpdate::eval_gpu( metal::MTLFCList func_consts = {}; - auto delta_kernel = get_steel_gated_delta_forward_kernel( - d, base_name, hash_name, func_consts); + auto delta_kernel = + get_gated_delta_kernel(d, base_name, hash_name, func_consts); compute_encoder.set_compute_pipeline_state(delta_kernel); compute_encoder.set_input_array(q, 0); @@ -172,8 +172,8 @@ void GatedDeltaUpdate::eval_gpu( metal::MTLFCList func_consts = {}; - auto delta_kernel = get_steel_gated_delta_forward_kernel( - d, base_name, hash_name, func_consts); + auto delta_kernel = + get_gated_delta_kernel(d, base_name, hash_name, func_consts); compute_encoder.set_compute_pipeline_state(delta_kernel); compute_encoder.set_input_array(q, 0); @@ -199,8 +199,8 @@ void GatedDeltaUpdate::eval_gpu( metal::MTLFCList func_consts = {}; - auto delta_kernel = get_steel_gated_delta_forward_kernel( - d, base_name, hash_name, func_consts); + auto delta_kernel = + get_gated_delta_kernel(d, base_name, hash_name, func_consts); compute_encoder.set_compute_pipeline_state(delta_kernel); diff --git a/mlx/backend/metal/jit_kernels.cpp b/mlx/backend/metal/jit_kernels.cpp index 1384e06c50..05fb1b6815 100644 --- a/mlx/backend/metal/jit_kernels.cpp +++ b/mlx/backend/metal/jit_kernels.cpp @@ -39,6 +39,10 @@ const char* fp_quantized_nax() { const char* steel_attention_nax() { return ""; } +const char* gated_delta_update_nax() { + return ""; +} +} } // namespace metal #endif // MLX_METAL_NO_NAX @@ -1322,4 +1326,14 @@ MTL::ComputePipelineState* get_steel_attention_nax_kernel( return d.get_kernel(kernel_name, lib, hash_name, func_consts); } +MTL::ComputePipelineState* get_gated_delta_kernel( + metal::Device& d, + const std::string& kernel_name, + const std::string& hash_name, + const metal::MTLFCList& func_consts + // TODO: need more parameters? +) { + return d.get_kernel(kernel_name, hash_name, func_consts); +} + } // namespace mlx::core diff --git a/mlx/backend/metal/kernels.h b/mlx/backend/metal/kernels.h index e09bffda3d..ef78039886 100644 --- a/mlx/backend/metal/kernels.h +++ b/mlx/backend/metal/kernels.h @@ -417,7 +417,7 @@ MTL::ComputePipelineState* get_steel_attention_nax_kernel( int wn, const array& m); -MTL::ComputePipelineState* get_steel_gated_delta_forward_kernel( +MTL::ComputePipelineState* get_gated_delta_kernel( metal::Device& d, const std::string& kernel_name, const std::string& hash_name, diff --git a/mlx/backend/metal/nojit_kernels.cpp b/mlx/backend/metal/nojit_kernels.cpp index fa662afac0..83ad4b3e6b 100644 --- a/mlx/backend/metal/nojit_kernels.cpp +++ b/mlx/backend/metal/nojit_kernels.cpp @@ -492,7 +492,7 @@ MTL::ComputePipelineState* get_steel_attention_nax_kernel( return d.get_kernel(kernel_name, hash_name, func_consts); } -MTL::ComputePipelineState* get_steel_gated_delta_forward_kernel( +MTL::ComputePipelineState* get_gated_delta_kernel( metal::Device& d, const std::string& kernel_name, const std::string& hash_name, From f87efb39005efda8c60fac357bc52dd75a145025 Mon Sep 17 00:00:00 2001 From: tpegolotti Date: Thu, 6 Aug 2026 05:54:35 -0700 Subject: [PATCH 50/56] Fix format error --- mlx/backend/metal/jit_kernels.cpp | 11 ----------- 1 file changed, 11 deletions(-) diff --git a/mlx/backend/metal/jit_kernels.cpp b/mlx/backend/metal/jit_kernels.cpp index 05fb1b6815..47b9cccc5c 100644 --- a/mlx/backend/metal/jit_kernels.cpp +++ b/mlx/backend/metal/jit_kernels.cpp @@ -42,7 +42,6 @@ const char* steel_attention_nax() { const char* gated_delta_update_nax() { return ""; } -} } // namespace metal #endif // MLX_METAL_NO_NAX @@ -1326,14 +1325,4 @@ MTL::ComputePipelineState* get_steel_attention_nax_kernel( return d.get_kernel(kernel_name, lib, hash_name, func_consts); } -MTL::ComputePipelineState* get_gated_delta_kernel( - metal::Device& d, - const std::string& kernel_name, - const std::string& hash_name, - const metal::MTLFCList& func_consts - // TODO: need more parameters? -) { - return d.get_kernel(kernel_name, hash_name, func_consts); -} - } // namespace mlx::core From 92906e0e5e364b313dfcf31f532f56e56032b67b Mon Sep 17 00:00:00 2001 From: tpegolotti Date: Thu, 6 Aug 2026 06:48:06 -0700 Subject: [PATCH 51/56] add jit getter --- mlx/backend/metal/jit_kernels.cpp | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/mlx/backend/metal/jit_kernels.cpp b/mlx/backend/metal/jit_kernels.cpp index 47b9cccc5c..0907e300a4 100644 --- a/mlx/backend/metal/jit_kernels.cpp +++ b/mlx/backend/metal/jit_kernels.cpp @@ -1325,4 +1325,12 @@ MTL::ComputePipelineState* get_steel_attention_nax_kernel( return d.get_kernel(kernel_name, lib, hash_name, func_consts); } +MTL::ComputePipelineState* get_gated_delta_kernel( + metal::Device& d, + const std::string& kernel_name, + const std::string& hash_name, + const metal::MTLFCList& func_consts) { + return d.get_kernel(kernel_name, hash_name, func_consts); +} + } // namespace mlx::core From 79d9a712b97d5b511ee2ce52467604830c9e97f6 Mon Sep 17 00:00:00 2001 From: tpegolotti Date: Thu, 6 Aug 2026 07:12:50 -0700 Subject: [PATCH 52/56] adding torch check on test --- python/tests/test_fast_gated_delta.py | 1 + 1 file changed, 1 insertion(+) diff --git a/python/tests/test_fast_gated_delta.py b/python/tests/test_fast_gated_delta.py index d596972ff2..3171e7bbdd 100644 --- a/python/tests/test_fast_gated_delta.py +++ b/python/tests/test_fast_gated_delta.py @@ -124,6 +124,7 @@ class TestGatedDelta(mlx_tests.MLXTestCase): fallback_dims = [base_dims, unaligned_dims, big_batch_dims] gpu_dims = fallback_dims + [large_t_dims] + @unittest.skipIf(not has_torch, "requires Torch") def test_gated_delta_fallback(self): for dims in self.fallback_dims: (out, hf), (out_ref, hf_ref) = runner(dims, mx.cpu) From df5e2da012ca4d0c676b21854d872ce3ab3e0b24 Mon Sep 17 00:00:00 2001 From: tpegolotti Date: Thu, 6 Aug 2026 07:26:27 -0700 Subject: [PATCH 53/56] adding torch check on test --- python/tests/test_fast_gated_delta.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/python/tests/test_fast_gated_delta.py b/python/tests/test_fast_gated_delta.py index 3171e7bbdd..f56bfb663c 100644 --- a/python/tests/test_fast_gated_delta.py +++ b/python/tests/test_fast_gated_delta.py @@ -4,7 +4,13 @@ import mlx.core as mx import mlx_tests import numpy as np -import torch + +try: + import torch + + has_torch = True +except ImportError as e: + has_torch = False def gated_delta_oracle( From 8861d8527129fc457d0952cf05bdad0e89bc4c07 Mon Sep 17 00:00:00 2001 From: tpegolotti Date: Thu, 6 Aug 2026 08:35:26 -0700 Subject: [PATCH 54/56] Fixing jit linkin error --- mlx/backend/metal/gated_delta_update.cpp | 2 +- mlx/backend/metal/jit/includes.h | 3 +++ mlx/backend/metal/jit_kernels.cpp | 14 ++++++++++++++ mlx/backend/metal/kernels.h | 10 +++++++--- .../metal/kernels/gated_delta_update_nax.h | 1 - mlx/backend/metal/nojit_kernels.cpp | 12 +++++++++--- python/tests/test_fast_gated_delta.py | 16 ++++++++-------- 7 files changed, 42 insertions(+), 16 deletions(-) diff --git a/mlx/backend/metal/gated_delta_update.cpp b/mlx/backend/metal/gated_delta_update.cpp index 23c3f1e553..5cd46cb462 100644 --- a/mlx/backend/metal/gated_delta_update.cpp +++ b/mlx/backend/metal/gated_delta_update.cpp @@ -144,7 +144,7 @@ void GatedDeltaUpdate::eval_gpu( metal::MTLFCList func_consts = {}; auto delta_kernel = - get_gated_delta_kernel(d, base_name, hash_name, func_consts); + get_gated_delta_nax_kernel(d, base_name, hash_name, func_consts); compute_encoder.set_compute_pipeline_state(delta_kernel); compute_encoder.set_input_array(q, 0); diff --git a/mlx/backend/metal/jit/includes.h b/mlx/backend/metal/jit/includes.h index ac9fb81e26..3001a31f1d 100644 --- a/mlx/backend/metal/jit/includes.h +++ b/mlx/backend/metal/jit/includes.h @@ -58,4 +58,7 @@ const char* fp_quantized_nax(); const char* steel_attention_nax(); +const char* gated_delta_update(); +const char* gated_delta_update_nax(); + } // namespace mlx::core::metal diff --git a/mlx/backend/metal/jit_kernels.cpp b/mlx/backend/metal/jit_kernels.cpp index 0907e300a4..b75ab7a874 100644 --- a/mlx/backend/metal/jit_kernels.cpp +++ b/mlx/backend/metal/jit_kernels.cpp @@ -1333,4 +1333,18 @@ MTL::ComputePipelineState* get_gated_delta_kernel( return d.get_kernel(kernel_name, hash_name, func_consts); } +MTL::ComputePipelineState* get_gated_delta_nax_kernel( + metal::Device& d, + const std::string& kernel_name, + const std::string& hash_name, + const metal::MTLFCList& func_consts) { + const auto& lib_name = kernel_name; + auto lib = d.get_library(lib_name, [&]() { + std::string kernel_source; + concatenate(kernel_source, metal::utils(), metal::gated_delta_update_nax()); + return kernel_source; + }); + return d.get_kernel(kernel_name, lib, hash_name, func_consts); +} + } // namespace mlx::core diff --git a/mlx/backend/metal/kernels.h b/mlx/backend/metal/kernels.h index ef78039886..a3cfc2e6ba 100644 --- a/mlx/backend/metal/kernels.h +++ b/mlx/backend/metal/kernels.h @@ -421,9 +421,13 @@ MTL::ComputePipelineState* get_gated_delta_kernel( metal::Device& d, const std::string& kernel_name, const std::string& hash_name, - const metal::MTLFCList& func_consts - // TODO: More parameters? -); + const metal::MTLFCList& func_consts); + +MTL::ComputePipelineState* get_gated_delta_nax_kernel( + metal::Device& d, + const std::string& kernel_name, + const std::string& hash_name, + const metal::MTLFCList& func_consts); // Create a GPU kernel template definition for JIT compilation template diff --git a/mlx/backend/metal/kernels/gated_delta_update_nax.h b/mlx/backend/metal/kernels/gated_delta_update_nax.h index 54b0abc4a1..72392bcba6 100644 --- a/mlx/backend/metal/kernels/gated_delta_update_nax.h +++ b/mlx/backend/metal/kernels/gated_delta_update_nax.h @@ -1,7 +1,6 @@ #pragma once #include -#include "mlx/backend/metal/kernels/utils.h" #include #include diff --git a/mlx/backend/metal/nojit_kernels.cpp b/mlx/backend/metal/nojit_kernels.cpp index 83ad4b3e6b..2ea601863c 100644 --- a/mlx/backend/metal/nojit_kernels.cpp +++ b/mlx/backend/metal/nojit_kernels.cpp @@ -496,9 +496,15 @@ MTL::ComputePipelineState* get_gated_delta_kernel( metal::Device& d, const std::string& kernel_name, const std::string& hash_name, - const metal::MTLFCList& func_consts - // TODO: need more parameters? -) { + const metal::MTLFCList& func_consts) { + return d.get_kernel(kernel_name, hash_name, func_consts); +} + +MTL::ComputePipelineState* get_gated_delta_nax_kernel( + metal::Device& d, + const std::string& kernel_name, + const std::string& hash_name, + const metal::MTLFCList& func_consts) { return d.get_kernel(kernel_name, hash_name, func_consts); } diff --git a/python/tests/test_fast_gated_delta.py b/python/tests/test_fast_gated_delta.py index f56bfb663c..9b9c347d6b 100644 --- a/python/tests/test_fast_gated_delta.py +++ b/python/tests/test_fast_gated_delta.py @@ -14,14 +14,14 @@ def gated_delta_oracle( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - beta: torch.Tensor, - g: torch.Tensor, - scale: float = None, - initial_state: torch.Tensor = None, - output_final_state: bool = False, + q, + k, + v, + beta, + g, + scale=None, + initial_state=None, + output_final_state=False, ): """ Reference PyTorch implementation of recurrent gated delta rule. From 05a89f2859d252aa716b29a309b1ce425b5a1760 Mon Sep 17 00:00:00 2001 From: tpegolotti Date: Fri, 7 Aug 2026 04:53:43 -0700 Subject: [PATCH 55/56] add support for different head sizes in fallback --- mlx/fast.cpp | 6 ++++++ python/tests/test_fast_gated_delta.py | 11 ++++++++++- 2 files changed, 16 insertions(+), 1 deletion(-) diff --git a/mlx/fast.cpp b/mlx/fast.cpp index 4ec16abe97..e552868799 100644 --- a/mlx/fast.cpp +++ b/mlx/fast.cpp @@ -968,6 +968,12 @@ std::vector gated_delta_update( auto beta = astype(inputs[4], float32, s); auto state = astype(inputs[5], float32, s); + if (Hv != Hk) { + int repeat_factor = Hv / Hk; + q = repeat(q, repeat_factor, 2, s); + k = repeat(k, repeat_factor, 2, s); + } + array mask = has_mask ? astype(inputs[6], bool_, s) : array(false); const array zero = array(0.0f, float32); diff --git a/python/tests/test_fast_gated_delta.py b/python/tests/test_fast_gated_delta.py index 9b9c347d6b..b2334dcdbb 100644 --- a/python/tests/test_fast_gated_delta.py +++ b/python/tests/test_fast_gated_delta.py @@ -73,6 +73,9 @@ def gated_delta_oracle( def runner(dims, stream=mx.gpu, reference=True): B, Hk, Hv, T, Dk, Dv = dims + assert Hv % Hk == 0 + repeat_factor = Hv // Hk + q = mx.random.normal(shape=(B, T, Hk, Dk)) k = mx.random.normal(shape=(B, T, Hk, Dk)) k = k / (mx.linalg.norm(k, axis=-1, keepdims=True) + 1e-6) @@ -90,6 +93,10 @@ def runner(dims, stream=mx.gpu, reference=True): gpt = torch.from_numpy(np.array(g)) h0pt = torch.from_numpy(np.array(h0)).transpose(-1, -2).contiguous() + if repeat_factor > 1: + qpt = qpt.repeat_interleave(repeat_factor, dim=2) + kpt = kpt.repeat_interleave(repeat_factor, dim=2) + out_on_py, hf_on_py = gated_delta_oracle( qpt, kpt, @@ -126,8 +133,10 @@ class TestGatedDelta(mlx_tests.MLXTestCase): unaligned_dims = (1, 32, 32, 33, 128, 128) big_batch_dims = (128, 32, 32, 16, 128, 128) large_t_dims = (2, 32, 32, 1111, 128, 128) + diff_heads = (1, 16, 32, 33, 128, 128) + diff_heads2 = (1, 16, 48, 33, 128, 128) - fallback_dims = [base_dims, unaligned_dims, big_batch_dims] + fallback_dims = [base_dims, unaligned_dims, big_batch_dims, diff_heads, diff_heads2] gpu_dims = fallback_dims + [large_t_dims] @unittest.skipIf(not has_torch, "requires Torch") From 4ad8724334f942dc95bdc0c962fb7bda9b59b77c Mon Sep 17 00:00:00 2001 From: tpegolotti Date: Fri, 7 Aug 2026 15:21:24 +0200 Subject: [PATCH 56/56] Remove debug macros --- mlx/backend/metal/gated_delta_update.cpp | 29 ------------------------ 1 file changed, 29 deletions(-) diff --git a/mlx/backend/metal/gated_delta_update.cpp b/mlx/backend/metal/gated_delta_update.cpp index 5cd46cb462..ae223c0379 100644 --- a/mlx/backend/metal/gated_delta_update.cpp +++ b/mlx/backend/metal/gated_delta_update.cpp @@ -54,31 +54,6 @@ ensure_row_contiguous(const array& x, metal::Device& d, const Stream& s) { } } -#define PRINT_STRIDES(arr) \ - printf( \ - "%s strides: %lld %lld %lld %lld\n", \ - #arr, \ - arr.strides()[0], \ - arr.strides()[1], \ - arr.strides()[2], \ - arr.strides()[3]) - -#define PRINT_SHAPES(arr) \ - printf( \ - "%s shapes: %lld %lld %lld %lld\n", \ - #arr, \ - arr.shape()[0], \ - arr.shape()[1], \ - arr.shape()[2], \ - arr.shape()[3]) - -#define PRINT_ARR(arr) \ - if (arr.flags().row_contiguous) \ - printf("%s is row contiguous\n", #arr); \ - PRINT_SHAPES(arr); \ - PRINT_STRIDES(arr); \ - printf("\n"); - void GatedDeltaUpdate::eval_gpu( const std::vector& inputs, std::vector& outputs) { @@ -102,10 +77,6 @@ void GatedDeltaUpdate::eval_gpu( int Hv = v.shape(2); int Dv = v.shape(3); -#if __METAL_VERSION__ >= 400 -#error "NAX branch IS compiled, __METAL_VERSION__ >= 400" -#endif - int C = 1; const char* threashold_env = std::getenv("GATED_DELTA_THRESH"); int threshold = threashold_env ? std::stoi(threashold_env) : 16;