From baee84404940d3ff95b23c9defb138c7166294bb Mon Sep 17 00:00:00 2001 From: Jake Stevens Date: Fri, 17 Jul 2026 13:38:41 -0700 Subject: [PATCH] Two-pass optimization for channelwise gated delta rule kernel (#21020) Summary: Reduce the channelwise gated delta rule kernel from four state traversals to two per token. The first pass computes the predicted value while folding in decay; the second applies decay and the rank-one update while accumulating the output. This preserves operation grouping and numerical behavior while reducing state-memory traffic. Add a standalone microbenchmark covering decode and representative prefill lengths. Differential Revision: D112596724 --- extension/llm/custom_ops/BUCK | 11 +++ .../bench_channelwise_gated_delta_rule.py | 78 +++++++++++++++++++ extension/llm/custom_ops/op_sdpa.cpp | 63 ++++++++------- 3 files changed, 125 insertions(+), 27 deletions(-) create mode 100644 extension/llm/custom_ops/bench_channelwise_gated_delta_rule.py diff --git a/extension/llm/custom_ops/BUCK b/extension/llm/custom_ops/BUCK index d70b985136d..16a25b419ac 100644 --- a/extension/llm/custom_ops/BUCK +++ b/extension/llm/custom_ops/BUCK @@ -91,3 +91,14 @@ fbcode_target(_kind = runtime.python_test, "//caffe2:torch", ], ) + +fbcode_target(_kind = runtime.python_binary, + name = "bench_channelwise_gated_delta_rule", + srcs = ["bench_channelwise_gated_delta_rule.py"], + 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 = ["//caffe2:torch"], +) diff --git a/extension/llm/custom_ops/bench_channelwise_gated_delta_rule.py b/extension/llm/custom_ops/bench_channelwise_gated_delta_rule.py new file mode 100644 index 00000000000..9fd43066d72 --- /dev/null +++ b/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() diff --git a/extension/llm/custom_ops/op_sdpa.cpp b/extension/llm/custom_ops/op_sdpa.cpp index 497715956fa..9ab0f99356a 100644 --- a/extension/llm/custom_ops/op_sdpa.cpp +++ b/extension/llm/custom_ops/op_sdpa.cpp @@ -773,8 +773,20 @@ std::tuple channelwise_gated_delta_rule_out( const auto* initial_state_data = initial_state.const_data_ptr(); auto* state_data = final_state_out.mutable_data_ptr(); auto* output_data = out.mutable_data_ptr(); - std::vector v_pred(v_head_dim); - std::vector delta(v_head_dim); + + const int64_t scratch_numel = 2 * v_head_dim; + std::unique_ptr fallback_scratch; + float* scratch_data = nullptr; + Result scratch = ctx.allocate_temp( + scratch_numel * sizeof(float), /*alignment=*/alignof(float)); + if (scratch.ok()) { + scratch_data = reinterpret_cast(scratch.get()); + } else { + fallback_scratch = std::make_unique(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 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; } } }