Pack gated_delta_seq: 8 value rows per SIMD-group

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;
     }

mlx/backend/metal/kernels/gated_delta_update_impl.h

@@ -267,10 +267,6 @@ template <typename InT, int Dk, int Dv, int Hk, int Hv, int C>
   }
 }
 
-/*
-        auto grid   = MTL::Size(32, Dv, B * Hv);
-    auto threads = MTL::Size(32, 4, 1);
- */
 template <typename InT, int Dk, int Dv, int Hk, int Hv>
 [[kernel]] void gated_delta_seq(
     const device InT* q [[buffer(0)]],
@@ -283,59 +279,87 @@ template <typename InT, int Dk, int Dv, int Hk, int Hv>
     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<float>(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<float>(i_state[i]);
   }
 
   // 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];
+    float gt = static_cast<float>(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<float>(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<float>(v_[dv_idx]) - kv_mem) *
+        static_cast<float>(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<float>(k_[e]) * delta;
+        acc += state[e] * static_cast<float>(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<InT>(out);
     }
     // Increment data pointers to next time step
@@ -346,8 +370,7 @@ template <typename InT, int Dk, int Dv, int Hk, int Hv>
     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<float>(state[i]);
+  for (int i = 0; i < values_per_lane; ++i) {
+    o_state[i] = state[i];
   }
-}
+}