Reland two-pass channelwise gated delta rule kernel + OSS-safe bench BUCK (#21105)

Differential Revision: D113076546

Pull Request resolved: #21105

Author: JakeStevens

extension/llm/custom_ops/BUCK

@@ -91,3 +91,20 @@ fbcode_target(_kind = runtime.python_test,
         "//caffe2:torch",
     ],
 )
+
+fbcode_target(_kind = runtime.python_library,
+    name = "bench_channelwise_gated_delta_rule_lib",
+    srcs = ["bench_channelwise_gated_delta_rule.py"],
+    base_module = "executorch.extension.llm.custom_ops",
+    deps = ["//caffe2:torch"],
+)
+
+fbcode_target(_kind = runtime.python_binary,
+    name = "bench_channelwise_gated_delta_rule",
+    main_module = "executorch.extension.llm.custom_ops.bench_channelwise_gated_delta_rule",
+    preload_deps = [
+        ":custom_ops_aot_lib_mkl_noomp",
+        ":custom_ops_aot_py",
+    ],
+    deps = [":bench_channelwise_gated_delta_rule_lib"],
+)

extension/llm/custom_ops/bench_channelwise_gated_delta_rule.py

@@ -0,0 +1,78 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+# pyre-unsafe
+
+"""Microbenchmark for the ``channelwise_gated_delta_rule`` custom op.
+
+Times the fused recurrent kernel in isolation (no model) at representative
+GatedDeltaNet sizes, for both the decode (T=1) and prefill (T>1) regimes. Use it
+to compare kernel variants (e.g. naive vs. fused) — rebuild the custom op, run
+this, and diff the numbers.
+
+Run (in an env where the custom op is built):
+    python -m executorch.extension.llm.custom_ops.bench_channelwise_gated_delta_rule
+"""
+
+import time
+
+import torch
+
+from executorch.extension.llm.custom_ops import custom_ops  # noqa: F401
+
+_OP = torch.ops.llama.channelwise_gated_delta_rule.default
+
+
+def _make_inputs(b: int, h: int, t: int, k: int, v: int):
+    return (
+        torch.randn(b, h, t, k),
+        torch.randn(b, h, t, k),
+        torch.randn(b, h, t, v),
+        torch.rand(b, h, t, k),
+        torch.rand(b, h, t),
+        torch.randn(b, h, k, v),
+    )
+
+
+def _bench_one(
+    b: int, h: int, t: int, k: int, v: int, iters: int, warmup: int
+) -> tuple[float, float]:
+    q, key, val, decay, beta, state = _make_inputs(b, h, t, k, v)
+    for _ in range(warmup):
+        _OP(q, key, val, decay, beta, state)
+    start = time.perf_counter()
+    for _ in range(iters):
+        _OP(q, key, val, decay, beta, state)
+    elapsed = time.perf_counter() - start
+    ms_per_call = elapsed / iters * 1e3
+    us_per_token = ms_per_call * 1e3 / t
+    return ms_per_call, us_per_token
+
+
+def main() -> None:
+    torch.manual_seed(0)
+    torch.set_num_threads(1)  # single-thread: the kernel has no internal threading
+
+    k = v = 128  # head dim (KDA / Qwen3-Next GatedDeltaNet)
+    # (label, batch, heads, seq_len, iters)
+    configs = [
+        ("decode  T=1", 1, 32, 1, 2000),
+        ("prefill T=128", 1, 32, 128, 200),
+        ("prefill T=512", 1, 32, 512, 50),
+    ]
+
+    print(
+        f"channelwise_gated_delta_rule microbenchmark "
+        f"(K=V={k}, fp32, 1 thread, state={32 * k * v * 4 // 1024}KB)"
+    )
+    print(f"{'config':<15}{'B':>3}{'H':>4}{'T':>6}{'ms/call':>12}{'us/token':>12}")
+    for label, b, h, t, iters in configs:
+        ms, us = _bench_one(b, h, t, k, v, iters=iters, warmup=max(10, iters // 10))
+        print(f"{label:<15}{b:>3}{h:>4}{t:>6}{ms:>12.4f}{us:>12.3f}")
+
+
+if __name__ == "__main__":
+    main()

extension/llm/custom_ops/op_sdpa.cpp

@@ -773,8 +773,20 @@ std::tuple<Tensor&, Tensor&> channelwise_gated_delta_rule_out(
   const auto* initial_state_data = initial_state.const_data_ptr<float>();
   auto* state_data = final_state_out.mutable_data_ptr<float>();
   auto* output_data = out.mutable_data_ptr<float>();
-  std::vector<float> v_pred(v_head_dim);
-  std::vector<float> delta(v_head_dim);
+
+  const int64_t scratch_numel = 2 * v_head_dim;
+  std::unique_ptr<float[]> fallback_scratch;
+  float* scratch_data = nullptr;
+  Result<void*> scratch = ctx.allocate_temp(
+      scratch_numel * sizeof(float), /*alignment=*/alignof(float));
+  if (scratch.ok()) {
+    scratch_data = reinterpret_cast<float*>(scratch.get());
+  } else {
+    fallback_scratch = std::make_unique<float[]>(scratch_numel);
+    scratch_data = fallback_scratch.get();
+  }
+  float* v_pred = scratch_data;
+  float* delta = scratch_data + v_head_dim;
 
   for (int64_t batch = 0; batch < batch_size; ++batch) {
     for (int64_t head = 0; head < num_heads; ++head) {
@@ -809,47 +821,44 @@ std::tuple<Tensor&, Tensor&> channelwise_gated_delta_rule_out(
         const float beta_t = beta_head[token];
         auto* output_t = output_head + token * value_seq_stride;
 
-        // 1. Per-key-channel decay: scale each state row by decay_t[k].
+        // The recurrence needs only two passes over the K x V state S:
+        //   pass 1 (read-only): v_pred = (Diag(decay) S)^T k
+        //   pass 2 (read+write): S = Diag(decay) S + k (x) delta;  o = S^T q
+        // Decay is folded into both passes (never materialized separately), and
+        // the rank-1 write and output readout share pass 2, so S is streamed
+        // twice per token instead of four times. Multiply groupings match the
+        // naive form, so results are identical up to floating-point contraction
+        // (e.g. compiler FMA).
+
+        // Pass 1: predicted value off the decayed state (S left untouched).
+        std::fill(v_pred, v_pred + v_head_dim, 0.0f);
         for (int64_t k_idx = 0; k_idx < k_head_dim; ++k_idx) {
           const float decay_value = decay_t[k_idx];
-          auto* state_row = state_head + k_idx * v_head_dim;
-          for (int64_t v_idx = 0; v_idx < v_head_dim; ++v_idx) {
-            state_row[v_idx] *= decay_value;
-          }
-        }
-
-        // 2. v_pred = S^T k: the memory's predicted value for this key, read
-        // off the decayed state.
-        std::fill(v_pred.begin(), v_pred.end(), 0.0f);
-        for (int64_t k_idx = 0; k_idx < k_head_dim; ++k_idx) {
           const float key_value = k_t[k_idx];
           const auto* state_row = state_head + k_idx * v_head_dim;
           for (int64_t v_idx = 0; v_idx < v_head_dim; ++v_idx) {
-            v_pred[v_idx] += state_row[v_idx] * key_value;
+            v_pred[v_idx] += (state_row[v_idx] * decay_value) * key_value;
           }
         }
 
-        // 3. delta = beta * (v - v_pred).
+        // delta = beta * (v - v_pred).
         for (int64_t v_idx = 0; v_idx < v_head_dim; ++v_idx) {
           delta[v_idx] = (v_t[v_idx] - v_pred[v_idx]) * beta_t;
         }
 
-        // 4. Rank-1 write: S += k (x) delta.
-        for (int64_t k_idx = 0; k_idx < k_head_dim; ++k_idx) {
-          const float key_value = k_t[k_idx];
-          auto* state_row = state_head + k_idx * v_head_dim;
-          for (int64_t v_idx = 0; v_idx < v_head_dim; ++v_idx) {
-            state_row[v_idx] += key_value * delta[v_idx];
-          }
-        }
-
-        // 5. Output readout: o_t = S^T q.
+        // Pass 2: apply decay + rank-1 write in place, and read back the
+        // updated row for the output projection in the same sweep.
         std::fill(output_t, output_t + v_head_dim, 0.0f);
         for (int64_t k_idx = 0; k_idx < k_head_dim; ++k_idx) {
+          const float decay_value = decay_t[k_idx];
+          const float key_value = k_t[k_idx];
           const float query_value = q_t[k_idx];
-          const auto* state_row = state_head + k_idx * v_head_dim;
+          auto* state_row = state_head + k_idx * v_head_dim;
           for (int64_t v_idx = 0; v_idx < v_head_dim; ++v_idx) {
-            output_t[v_idx] += state_row[v_idx] * query_value;
+            const float updated =
+                state_row[v_idx] * decay_value + key_value * delta[v_idx];
+            state_row[v_idx] = updated;
+            output_t[v_idx] += updated * query_value;
           }
         }
       }