From d9e1a26dac60d22726e8710ad00675ed22e8f247 Mon Sep 17 00:00:00 2001 From: Cheng Date: Thu, 29 Oct 2026 02:42:01 -0700 Subject: [PATCH] Respect MLX_ENABLE_TF32 in gated_delta_update --- mlx/backend/metal/gated_delta_update.cpp | 22 ++++++++++++++++------ 1 file changed, 16 insertions(+), 6 deletions(-) diff --git a/mlx/backend/metal/gated_delta_update.cpp b/mlx/backend/metal/gated_delta_update.cpp index 846f4ae9f0..bb9f7b567b 100644 --- a/mlx/backend/metal/gated_delta_update.cpp +++ b/mlx/backend/metal/gated_delta_update.cpp @@ -15,14 +15,24 @@ namespace mlx::core::fast { namespace { -inline int gated_delta_chunk_size(int T) { - bool nax = metal::is_nax_available(); +inline int gated_delta_chunk_size(Dtype dtype, bool vectorized_g, int T) { + if (vectorized_g) { + // Only sequential kernel supports vectorized g. + return 1; + } + bool nax = + metal::is_nax_available() && (env::enable_tf32() || dtype != float32); int C = config::get("MLX_GATED_DELTA_CHUNK", T > 8 ? (nax ? 16 : 8) : 1); return (C == 16 && !nax) ? 8 : C; } -inline int gated_delta_chunk_size_vjp() { - bool nax = metal::is_nax_available(); +inline int gated_delta_chunk_size_vjp(Dtype dtype, bool vectorized_g) { + if (vectorized_g) { + // Only sequential kernel supports vectorized g. + return 1; + } + bool nax = + metal::is_nax_available() && (env::enable_tf32() || dtype != float32); return nax ? config::get("MLX_GATED_DELTA_CHUNK_VJP", 16) : 1; } @@ -106,7 +116,7 @@ void GatedDeltaUpdate::eval_gpu( int Dv = v.shape(3); bool vectorized_g = (g.ndim() == 4); - int C = vectorized_g ? 1 : gated_delta_chunk_size(T); + int C = gated_delta_chunk_size(q.dtype(), vectorized_g, T); const int ckpt = gated_delta_ckpt(C); std::string suffix = @@ -262,7 +272,7 @@ void GatedDeltaUpdateVJP::eval_gpu( // 16 = chunked NAX, anything else = sequential. bool vectorized_g = (g.ndim() == 4); - int C = vectorized_g ? 1 : gated_delta_chunk_size_vjp(); + int C = gated_delta_chunk_size_vjp(q.dtype(), vectorized_g); const bool chunked = (C == 16); const int n_chunks = (T + 15) / 16;