From 10c6a841e56a6a072da98f27b374c799c57bdad8 Mon Sep 17 00:00:00 2001 From: Yanzhao Wang Date: Tue, 25 Aug 2026 23:46:45 -0700 Subject: [PATCH] Pack gated_delta_seq to 8 value rows per SIMD-group --- mlx/backend/metal/gated_delta_update.cpp | 4 +- .../metal/kernels/gated_delta_update_impl.h | 101 +++++++++++------- 2 files changed, 64 insertions(+), 41 deletions(-) diff --git a/mlx/backend/metal/gated_delta_update.cpp b/mlx/backend/metal/gated_delta_update.cpp index 662e574788..76f235dc6d 100644 --- a/mlx/backend/metal/gated_delta_update.cpp +++ b/mlx/backend/metal/gated_delta_update.cpp @@ -182,8 +182,8 @@ void GatedDeltaUpdate::eval_gpu( 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); + auto grid = MTL::Size(32, Dv / 8, B * Hv); + auto threads = MTL::Size(32, 2, 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 78d2441498..dd1fdd3b88 100644 --- a/mlx/backend/metal/kernels/gated_delta_update_impl.h +++ b/mlx/backend/metal/kernels/gated_delta_update_impl.h @@ -267,10 +267,6 @@ template } } -/* - 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)]], @@ -283,34 +279,44 @@ template device InT* y [[buffer(7)]], // [B, T, Hv, Dv] 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]]) { - // kernel implementation + constexpr int lanes_per_row = 4; + constexpr int rows_per_simdgroup = 32 / lanes_per_row; + constexpr int values_per_lane = Dk / lanes_per_row; + constexpr int partials_per_lane = values_per_lane / 4; + static_assert(Dk == 128, "Packed sequential gated delta requires Dk == 128"); + static_assert( + Dv % rows_per_simdgroup == 0, + "Packed sequential gated delta requires Dv divisible by 8"); + 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 lane = thread_index_in_simdgroup; + auto row_in_simdgroup = lane / lanes_per_row; + auto lane_in_row = lane & (lanes_per_row - 1); + auto row_group = thread_position_in_grid.y; + auto dv_idx = row_group * rows_per_simdgroup + row_in_simdgroup; // 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; + auto q_ = q + (b_idx * T * Hk + hk_idx) * Dk + lane_in_row * values_per_lane; + auto k_ = k + (b_idx * T * Hk + hk_idx) * Dk + lane_in_row * values_per_lane; // 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]); + auto i_state = + state_in + (n * Dv + dv_idx) * Dk + lane_in_row * values_per_lane; + auto o_state = + state_out + (n * Dv + dv_idx) * Dk + lane_in_row * values_per_lane; + + float state[values_per_lane]; + for (int i = 0; i < values_per_lane; ++i) { + state[i] = static_cast(i_state[i]); } // g: [B, T, Hv] @@ -318,24 +324,42 @@ template 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]; + float gt = static_cast(g_[hv_idx]); + + // Preserve the unpacked four-element partials. Reduce the first three + // tree levels locally and the last two across the four-lane row. + float part[partials_per_lane]; + for (int pb = 0; pb < partials_per_lane; ++pb) { + float acc = 0.0f; + for (int i = 0; i < 4; ++i) { + int e = pb * 4 + i; + state[e] = state[e] * gt; + acc += state[e] * static_cast(k_[e]); + } + part[pb] = acc; } - 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]; + float kv_mem = ((part[0] + part[1]) + (part[2] + part[3])) + + ((part[4] + part[5]) + (part[6] + part[7])); + kv_mem += simd_shuffle_xor(kv_mem, 1); + kv_mem += simd_shuffle_xor(kv_mem, 2); + + auto delta = (static_cast(v_[dv_idx]) - kv_mem) * + static_cast(beta_[hv_idx]); + + for (int pb = 0; pb < partials_per_lane; ++pb) { + float acc = 0.0f; + for (int i = 0; i < 4; ++i) { + int e = pb * 4 + i; + state[e] = state[e] + static_cast(k_[e]) * delta; + acc += state[e] * static_cast(q_[e]); + } + part[pb] = acc; } - out = simd_sum(out); - if (thread_index_in_simdgroup == 0) { + float out = ((part[0] + part[1]) + (part[2] + part[3])) + + ((part[4] + part[5]) + (part[6] + part[7])); + out += simd_shuffle_xor(out, 1); + out += simd_shuffle_xor(out, 2); + if (lane_in_row == 0) { y[dv_idx] = static_cast(out); } // Increment data pointers to next time step @@ -346,8 +370,7 @@ template 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]); + for (int i = 0; i < values_per_lane; ++i) { + o_state[i] = state[i]; } -} \ No newline at end of file +}