From 2980c22280218b815c4c7b478dfd1c42a1199cc5 Mon Sep 17 00:00:00 2001 From: Jake Stevens Date: Fri, 17 Jul 2026 13:48:55 -0700 Subject: [PATCH 1/3] 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; } } } From 00d6ca7806cf57de8c2a9ec5c1b30a39d21bf5cd Mon Sep 17 00:00:00 2001 From: Jake Stevens Date: Fri, 17 Jul 2026 13:48:55 -0700 Subject: [PATCH 2/3] Add parallel chunked prefill for channelwise gated delta rule Summary: Route the channelwise gated delta rule by sequence length: T == 1 keeps the two-pass token recurrence for autoregressive decode, while T != 1 uses a chunkwise WY/UT formulation for prefill. The chunked path computes per-channel log-decay prefixes, causal query/key terms, the beta-folded triangular transform, WY pseudo-keys and pseudo-values, and inter-chunk state carry. It handles a ragged final chunk without a separate tail implementation. Parallelize independent (batch, head) work across the ExecuTorch threadpool. Each worker receives a disjoint slice of one temporary scratch arena, avoiding shared mutable buffers while amortizing allocation across chunks. Differential Revision: D112597348 --- extension/llm/custom_ops/op_sdpa.cpp | 445 ++++++++++++++++-- extension/llm/custom_ops/test_update_cache.py | 45 +- 2 files changed, 439 insertions(+), 51 deletions(-) diff --git a/extension/llm/custom_ops/op_sdpa.cpp b/extension/llm/custom_ops/op_sdpa.cpp index 9ab0f99356a..61e71b738c1 100644 --- a/extension/llm/custom_ops/op_sdpa.cpp +++ b/extension/llm/custom_ops/op_sdpa.cpp @@ -14,16 +14,18 @@ #include #include // @lint-ignore CLANGTIDY facebook-unused-include-check +#include #include #include #include #include -#include +#include #include +// parallel_for + get_thread_num have serial fallbacks without a threadpool. +#include #ifdef ET_USE_THREADPOOL #include -#include #endif #include @@ -691,7 +693,11 @@ Tensor& sdpa_with_kv_cache_out( return output; } -std::tuple channelwise_gated_delta_rule_out( +// Shared full-sequence recurrence: streams the K x V state twice per token (see +// the two-pass comments inline). The T==1 and T!=1 entry points both delegate +// here for now; they exist so the decode and prefill paths can later diverge +// (e.g. a chunkwise-parallel prefill kernel) without touching callers. +void channelwise_gated_delta_rule_recurrence( RuntimeContext& ctx, const Tensor& query, const Tensor& key, @@ -701,50 +707,6 @@ std::tuple channelwise_gated_delta_rule_out( const Tensor& initial_state, Tensor& out, Tensor& final_state_out) { - std::tuple ret(out, final_state_out); - ET_KERNEL_CHECK( - ctx, - validate_channelwise_gated_delta_rule_args( - query, key, value, decay, beta, initial_state), - InvalidArgument, - ret); - ET_KERNEL_CHECK_MSG( - ctx, - !tensor_memory_ranges_overlap(initial_state, final_state_out), - InvalidArgument, - ret, - "channelwise_gated_delta_rule final_state_out must not alias initial_state."); - ET_KERNEL_CHECK_MSG( - ctx, - resize_tensor(out, value.sizes()) == Error::Ok, - InvalidArgument, - ret, - "Failed to resize channelwise_gated_delta_rule output tensor."); - ET_KERNEL_CHECK_MSG( - ctx, - resize_tensor(final_state_out, initial_state.sizes()) == Error::Ok, - InvalidArgument, - ret, - "Failed to resize channelwise_gated_delta_rule final_state tensor."); - ET_KERNEL_CHECK( - ctx, out.scalar_type() == ScalarType::Float, InvalidArgument, ret); - ET_KERNEL_CHECK( - ctx, - final_state_out.scalar_type() == ScalarType::Float, - InvalidArgument, - ret); - ET_KERNEL_CHECK( - ctx, - is_contiguous_dim_order(out.dim_order().data(), out.dim()), - InvalidArgument, - ret); - ET_KERNEL_CHECK( - ctx, - is_contiguous_dim_order( - final_state_out.dim_order().data(), final_state_out.dim()), - InvalidArgument, - ret); - const auto batch_size = query.size(0); const auto num_heads = query.size(1); const auto sequence_length = query.size(2); @@ -864,6 +826,395 @@ std::tuple channelwise_gated_delta_rule_out( } } } +} + +// T == 1 (autoregressive decode) route. TODO: specialize with a single-token +// fast path; delegates to the shared recurrence for now. +void channelwise_gated_delta_rule_decode( + RuntimeContext& ctx, + const Tensor& query, + const Tensor& key, + const Tensor& value, + const Tensor& decay, + const Tensor& beta, + const Tensor& initial_state, + Tensor& out, + Tensor& final_state_out) { + channelwise_gated_delta_rule_recurrence( + ctx, query, key, value, decay, beta, initial_state, out, final_state_out); +} + +// Chunk length for the chunkwise recurrence. Kept small so that +// exp(-cumsum(log decay)) over a chunk stays finite (see eg_inv in step 1). +constexpr int64_t CHUNK_SIZE = 32; + +// Round a float count up to a 64-byte (16-float) boundary so each per-worker +// scratch slab is cache-line aligned (avoids false sharing between workers). +inline int64_t chunk_scratch_align(int64_t n) { + constexpr int64_t kAlign = 16; // floats; 16 * 4B = 64B + return (n + kAlign - 1) / kAlign * kAlign; +} + +// One worker's scratch for the chunked kernel, carved from a single arena so +// there is one allocate_temp for the whole op and disjoint per-worker slabs. +struct ChunkScratch { + float* gc; // [CHUNK_SIZE, K] cumulative log-decay + float* eg; // [CHUNK_SIZE, K] exp(gc) + float* eg_inv; // [CHUNK_SIZE, K] exp(-gc) + float* w; // [CHUNK_SIZE, K] WY pseudo-keys + float* u; // [CHUNK_SIZE, V] WY pseudo-values + float* pv; // [CHUNK_SIZE, V] u - w @ S + float* Aqk; // [CHUNK_SIZE, CHUNK_SIZE] + float* A; // [CHUNK_SIZE, CHUNK_SIZE] k.k -> UT inverse (beta-folded) + float* de; // [CHUNK_SIZE] decay from row r to chunk end + + // Floats needed by one worker. Single source of truth for sizing + carving. + static int64_t elems(int64_t k_head_dim, int64_t v_head_dim) { + return 4 * chunk_scratch_align(CHUNK_SIZE * k_head_dim) + // gc,eg,eg_inv,w + 2 * chunk_scratch_align(CHUNK_SIZE * v_head_dim) + // u, pv + 2 * chunk_scratch_align(CHUNK_SIZE * CHUNK_SIZE) + // Aqk, A + chunk_scratch_align(CHUNK_SIZE); // de + } + + static ChunkScratch view(float* p, int64_t k_head_dim, int64_t v_head_dim) { + const int64_t nk = chunk_scratch_align(CHUNK_SIZE * k_head_dim); + const int64_t nv = chunk_scratch_align(CHUNK_SIZE * v_head_dim); + const int64_t nbb = chunk_scratch_align(CHUNK_SIZE * CHUNK_SIZE); + ChunkScratch s; + s.gc = p; + p += nk; + s.eg = p; + p += nk; + s.eg_inv = p; + p += nk; + s.w = p; + p += nk; + s.u = p; + p += nv; + s.pv = p; + p += nv; + s.Aqk = p; + p += nbb; + s.A = p; + p += nbb; + s.de = p; + return s; + } +}; + +// T != 1 (prefill) route. +void channelwise_gated_delta_rule_chunked( + RuntimeContext& ctx, + const Tensor& query, + const Tensor& key, + const Tensor& value, + const Tensor& decay, + const Tensor& beta, + const Tensor& initial_state, + Tensor& out, + Tensor& final_state_out) { + const auto batch_size = query.size(0); + const auto num_heads = query.size(1); + const auto sequence_length = query.size(2); + const auto k_head_dim = query.size(3); + const auto v_head_dim = value.size(3); + const auto num_chunks = (sequence_length + CHUNK_SIZE - 1) / CHUNK_SIZE; + + const auto* query_data = query.const_data_ptr(); + const auto* key_data = key.const_data_ptr(); + const auto* value_data = value.const_data_ptr(); + const auto* decay_data = decay.const_data_ptr(); + const auto* beta_data = beta.const_data_ptr(); + 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(); + + // Per (b, h) strides. q/k/decay are [.., T, K]; v/out are [.., T, V]. + const auto qk_head_stride = sequence_length * k_head_dim; + const auto v_head_stride = sequence_length * v_head_dim; + const auto beta_head_stride = sequence_length; + const auto state_head_stride = k_head_dim * v_head_dim; + + // One scratch arena for the whole op: a single temp allocation carved into a + // disjoint slab per worker (indexed by thread id), reused across chunks. The + // arena is NOT zeroed; every scratch element is assigned before it is read + // (only the causal j <= r triangles of Aqk/A and rows < cur are ever used). + const int64_t size_per_thread = ChunkScratch::elems(k_head_dim, v_head_dim); +#ifdef ET_USE_THREADPOOL + const int64_t num_thread = + ::executorch::extension::threadpool::get_threadpool()->get_thread_count(); +#else + const int64_t num_thread = 1; +#endif + const int64_t size_bytes = size_per_thread * num_thread * sizeof(float); + constexpr std::size_t kArenaAlign = 64; + std::unique_ptr fallback_buf; + float* arena = nullptr; + Result scratch = ctx.allocate_temp(size_bytes, kArenaAlign); + if (scratch.ok()) { + arena = reinterpret_cast(scratch.get()); + } else { + // allocate_temp guarantees 64B alignment; the heap fallback must match it + // so per-worker slabs stay cache-line separated (else false sharing + // returns). Over-allocate and align the base into the buffer. + fallback_buf = std::make_unique(size_bytes + kArenaAlign - 1); + void* base = fallback_buf.get(); + std::size_t space = static_cast(size_bytes) + kArenaAlign - 1; + arena = reinterpret_cast(std::align( + kArenaAlign, static_cast(size_bytes), base, space)); + } + + // Heads are independent, so parallelize over (b, h); each worker gets its own + // scratch slab (via get_thread_num()) and writes disjoint output/state. + torch::executor::parallel_for( + 0, + batch_size * num_heads, + /*grain_size=*/1, + [&](int64_t bh_begin, int64_t bh_end) { + for (int64_t bh = bh_begin; bh < bh_end; ++bh) { + const int64_t tid = torch::executor::get_thread_num(); + const ChunkScratch sc = ChunkScratch::view( + arena + tid * size_per_thread, k_head_dim, v_head_dim); + const auto* __restrict__ q_head = query_data + bh * qk_head_stride; + const auto* __restrict__ k_head = key_data + bh * qk_head_stride; + const auto* __restrict__ v_head = value_data + bh * v_head_stride; + const auto* __restrict__ d_head = decay_data + bh * qk_head_stride; + const auto* __restrict__ beta_head = + beta_data + bh * beta_head_stride; + const auto* __restrict__ init_head = + initial_state_data + bh * state_head_stride; + auto* __restrict__ S = state_data + bh * state_head_stride; + auto* __restrict__ o_head = output_data + bh * v_head_stride; + + // Seed the running state from the (read-only) initial_state. + for (const auto idx : c10::irange(state_head_stride)) { + S[idx] = init_head[idx]; + } + + for (const auto chunk : c10::irange(num_chunks)) { + const int64_t base = chunk * CHUNK_SIZE; + // Ragged tail: the last chunk may hold fewer than CHUNK_SIZE + // tokens. + const int64_t cur = + std::min(CHUNK_SIZE, sequence_length - base); + const auto* __restrict__ q_c = q_head + base * k_head_dim; + const auto* __restrict__ k_c = k_head + base * k_head_dim; + const auto* __restrict__ v_c = v_head + base * v_head_dim; + const auto* __restrict__ d_c = d_head + base * k_head_dim; + const auto* __restrict__ beta_c = beta_head + base; + auto* __restrict__ o_c = o_head + base * v_head_dim; + + // 1. gc = cumsum_r(log decay) (per channel; resets each chunk) + for (const auto k_idx : c10::irange(k_head_dim)) { + float acc = 0.0f; + for (const auto r : c10::irange(cur)) { + acc += std::log(d_c[r * k_head_dim + k_idx]); + sc.gc[r * k_head_dim + k_idx] = acc; + sc.eg[r * k_head_dim + k_idx] = std::exp(acc); + sc.eg_inv[r * k_head_dim + k_idx] = std::exp(-acc); + } + } + + // 2. Aqk (causal, incl diag, no beta) and A = -beta_r * + // strictlower(Akk). + // ratio = exp(gc[r,k] - gc[j,k]) = eg[r,k] * eg_inv[j,k], + // factored so the exp is hoisted out of the O(BT^2) pair loop. + // Safe because BT is small enough that exp(-gc) does not + // overflow (see eg_inv above). + for (const auto r : c10::irange(cur)) { + const float beta_r = beta_c[r]; + for (const auto j : c10::irange(r + 1)) { + float aqk = 0.0f, akk = 0.0f; + for (const auto k_idx : c10::irange(k_head_dim)) { + const float ratio = sc.eg[r * k_head_dim + k_idx] * + sc.eg_inv[j * k_head_dim + k_idx]; + aqk += q_c[r * k_head_dim + k_idx] * + k_c[j * k_head_dim + k_idx] * ratio; + if (j < r) { + akk += k_c[r * k_head_dim + k_idx] * + k_c[j * k_head_dim + k_idx] * ratio; + } + } + // Causal: only j <= r is written, and step 5 reads only j <= r. + // The upper triangle is left uninitialized (arena is not + // zeroed). + sc.Aqk[r * CHUNK_SIZE + j] = aqk; + sc.A[r * CHUNK_SIZE + j] = (j < r) ? -beta_r * akk : 0.0f; + } + } + + // 3. UT inverse by forward substitution, then + I and right-scale + // by + // diag(beta) (column j scaled by beta_j) -> A is now the WY + // transform. + for (const auto r : c10::irange(int64_t{1}, cur)) { + for (const auto j : c10::irange(r)) { + float s = sc.A[r * CHUNK_SIZE + j]; + for (const auto m : c10::irange(j + 1, r)) { + s += sc.A[r * CHUNK_SIZE + m] * sc.A[m * CHUNK_SIZE + j]; + } + sc.A[r * CHUNK_SIZE + j] = s; + } + } + for (const auto r : c10::irange(cur)) { + sc.A[r * CHUNK_SIZE + r] = 1.0f; // + I + for (const auto j : c10::irange(r + 1)) { + sc.A[r * CHUNK_SIZE + j] *= beta_c[j]; // right diag(beta) + } + } + + // 4. WY packing: w = A @ (exp(gc) * k), u = A @ v (A lower-tri: j + // <= r) + for (const auto r : c10::irange(cur)) { + for (const auto k_idx : c10::irange(k_head_dim)) { + float acc = 0.0f; + for (const auto j : c10::irange(r + 1)) { + acc += sc.A[r * CHUNK_SIZE + j] * + sc.eg[j * k_head_dim + k_idx] * + k_c[j * k_head_dim + k_idx]; + } + sc.w[r * k_head_dim + k_idx] = acc; + } + for (const auto v_idx : c10::irange(v_head_dim)) { + float acc = 0.0f; + for (const auto j : c10::irange(r + 1)) { + acc += sc.A[r * CHUNK_SIZE + j] * v_c[j * v_head_dim + v_idx]; + } + sc.u[r * v_head_dim + v_idx] = acc; + } + } + + // 5. pv = u - w @ S_old ; o = (q ⊙ exp gc) @ S_old + Aqk @ pv + for (const auto r : c10::irange(cur)) { + for (const auto v_idx : c10::irange(v_head_dim)) { + float wS = 0.0f, qS = 0.0f; + for (const auto k_idx : c10::irange(k_head_dim)) { + const float s = S[k_idx * v_head_dim + v_idx]; + wS += sc.w[r * k_head_dim + k_idx] * s; + qS += q_c[r * k_head_dim + k_idx] * + sc.eg[r * k_head_dim + k_idx] * s; + } + sc.pv[r * v_head_dim + v_idx] = + sc.u[r * v_head_dim + v_idx] - wS; + o_c[r * v_head_dim + v_idx] = qS; + } + } + for (const auto r : c10::irange(cur)) { + for (const auto j : c10::irange(r + 1)) { + const float a = sc.Aqk[r * CHUNK_SIZE + j]; + for (const auto v_idx : c10::irange(v_head_dim)) { + o_c[r * v_head_dim + v_idx] += + a * sc.pv[j * v_head_dim + v_idx]; + } + } + } + + // 6. S = diag(exp gc_last) S + (k ⊙ decay_to_end)^T @ pv. + // decay_to_end must stay exp(gc_last - gc[r]) (a small, safe + // diff); eg[r] can underflow to 0 deep in a chunk, so + // eg_last/eg[r] would be 0/0. Precompute it per channel to keep + // it out of the v loop. + for (const auto k_idx : c10::irange(k_head_dim)) { + const float gc_last = sc.gc[(cur - 1) * k_head_dim + k_idx]; + const float eg_last = sc.eg[(cur - 1) * k_head_dim + k_idx]; + for (const auto r : c10::irange(cur)) { + sc.de[r] = std::exp(gc_last - sc.gc[r * k_head_dim + k_idx]); + } + for (const auto v_idx : c10::irange(v_head_dim)) { + float acc = S[k_idx * v_head_dim + v_idx] * eg_last; + for (const auto r : c10::irange(cur)) { + acc += k_c[r * k_head_dim + k_idx] * sc.de[r] * + sc.pv[r * v_head_dim + v_idx]; + } + S[k_idx * v_head_dim + v_idx] = acc; + } + } + } // chunk + } // bh + }); // parallel_for +} + +std::tuple channelwise_gated_delta_rule_out( + RuntimeContext& ctx, + const Tensor& query, + const Tensor& key, + const Tensor& value, + const Tensor& decay, + const Tensor& beta, + const Tensor& initial_state, + Tensor& out, + Tensor& final_state_out) { + std::tuple ret(out, final_state_out); + // Validate everything before resizing any output, so a bad call leaves the + // caller's out/final_state_out tensors untouched. + ET_KERNEL_CHECK( + ctx, + validate_channelwise_gated_delta_rule_args( + query, key, value, decay, beta, initial_state), + InvalidArgument, + ret); + ET_KERNEL_CHECK_MSG( + ctx, + !tensor_memory_ranges_overlap(initial_state, final_state_out), + InvalidArgument, + ret, + "channelwise_gated_delta_rule final_state_out must not alias initial_state."); + ET_KERNEL_CHECK( + ctx, out.scalar_type() == ScalarType::Float, InvalidArgument, ret); + ET_KERNEL_CHECK( + ctx, + final_state_out.scalar_type() == ScalarType::Float, + InvalidArgument, + ret); + ET_KERNEL_CHECK( + ctx, + is_contiguous_dim_order(out.dim_order().data(), out.dim()), + InvalidArgument, + ret); + ET_KERNEL_CHECK( + ctx, + is_contiguous_dim_order( + final_state_out.dim_order().data(), final_state_out.dim()), + InvalidArgument, + ret); + ET_KERNEL_CHECK_MSG( + ctx, + resize_tensor(out, value.sizes()) == Error::Ok, + InvalidArgument, + ret, + "Failed to resize channelwise_gated_delta_rule output tensor."); + ET_KERNEL_CHECK_MSG( + ctx, + resize_tensor(final_state_out, initial_state.sizes()) == Error::Ok, + InvalidArgument, + ret, + "Failed to resize channelwise_gated_delta_rule final_state tensor."); + + // Route on sequence length: T == 1 is autoregressive decode (token-by-token + // recurrence), T != 1 is prefill (chunkwise WY kernel). + if (query.size(2) == 1) { + channelwise_gated_delta_rule_decode( + ctx, + query, + key, + value, + decay, + beta, + initial_state, + out, + final_state_out); + } else { + channelwise_gated_delta_rule_chunked( + ctx, + query, + key, + value, + decay, + beta, + initial_state, + out, + final_state_out); + } return ret; } diff --git a/extension/llm/custom_ops/test_update_cache.py b/extension/llm/custom_ops/test_update_cache.py index 2db3c637883..682ddb1b93d 100644 --- a/extension/llm/custom_ops/test_update_cache.py +++ b/extension/llm/custom_ops/test_update_cache.py @@ -646,7 +646,16 @@ def test_channelwise_gated_delta_rule_out_rejects_initial_state_overlap( def test_channelwise_gated_delta_rule_chunked_matches_full_sequence(self): torch.manual_seed(0) - query, key, value, decay, beta, initial_state = self._make_inputs(seq_len=6) + # seq_len > CHUNK_SIZE (128 in the C++ kernel): the single full call runs + # the internal chunk loop across a chunk boundary (128 + 72), while the + # per-segment calls each stay within a single internal chunk. Equality + # validates the inter-chunk state carry independent of chunk boundaries. + # fp32 accumulation over a 128-token chunk needs a looser tol than the + # single-chunk tests above. + seq_len = 200 + query, key, value, decay, beta, initial_state = self._make_inputs( + seq_len=seq_len + ) full_output, full_state = torch.ops.llama.channelwise_gated_delta_rule( query, @@ -659,7 +668,7 @@ def test_channelwise_gated_delta_rule_chunked_matches_full_sequence(self): chunk_state = initial_state chunk_outputs = [] - for start, end in ((0, 2), (2, 5), (5, 6)): + for start, end in ((0, 64), (64, 192), (192, 200)): chunk_output, chunk_state = torch.ops.llama.channelwise_gated_delta_rule( query[:, :, start:end, :], key[:, :, start:end, :], @@ -671,8 +680,36 @@ def test_channelwise_gated_delta_rule_chunked_matches_full_sequence(self): chunk_outputs.append(chunk_output) chunked_output = torch.cat(chunk_outputs, dim=2) - self.assertTrue(torch.allclose(chunked_output, full_output, atol=1e-5)) - self.assertTrue(torch.allclose(chunk_state, full_state, atol=1e-5)) + self.assertTrue( + torch.allclose(chunked_output, full_output, atol=1e-3, rtol=1e-3) + ) + self.assertTrue(torch.allclose(chunk_state, full_state, atol=1e-3, rtol=1e-3)) + + def test_channelwise_gated_delta_rule_multichunk_matches_reference(self): + torch.manual_seed(0) + + # Drive the prefill route past a single CHUNK_SIZE (32) so the kernel's + # internal chunk loop runs multiple times against the token-by-token + # reference: 130 exercises a ragged final chunk (4 * 32 + 2), 256 + # exercises an exact multiple (8 * 32). fp32 accumulation over a full + # chunk needs a looser tol than the single-chunk tests. + for seq_len in (130, 256): + with self.subTest(seq_len=seq_len): + inputs = self._make_inputs(seq_len=seq_len) + expected_output, expected_state = ( + self._reference_channelwise_gated_delta_rule(*inputs) + ) + + actual_output, actual_state = ( + torch.ops.llama.channelwise_gated_delta_rule(*inputs) + ) + + self.assertTrue( + torch.allclose(actual_output, expected_output, atol=1e-3, rtol=1e-3) + ) + self.assertTrue( + torch.allclose(actual_state, expected_state, atol=1e-3, rtol=1e-3) + ) def test_channelwise_gated_delta_rule_exports(self): class Module(torch.nn.Module): From dac45054a6a958876df3cace1de8ea984d545bc1 Mon Sep 17 00:00:00 2001 From: Jake Stevens Date: Fri, 17 Jul 2026 13:48:55 -0700 Subject: [PATCH 3/3] Reorder channelwise gated delta rule chunked hot loops for autovectorization (#21021) Summary: Reorder the chunked prefill inner loops (steps 1, 4, 5, 6) so the innermost loop runs contiguously over the head dimension (k or v) instead of striding down a column of the state / pv. This lets the compiler autovectorize the now-unit-stride AXPYs; hand-written at::vec was tried and was slower than the compiler output, so the loops stay scalar. Differential Revision: D112598714 --- extension/llm/custom_ops/op_sdpa.cpp | 102 +++++++++++++++++---------- 1 file changed, 65 insertions(+), 37 deletions(-) diff --git a/extension/llm/custom_ops/op_sdpa.cpp b/extension/llm/custom_ops/op_sdpa.cpp index 61e71b738c1..f030b098c93 100644 --- a/extension/llm/custom_ops/op_sdpa.cpp +++ b/extension/llm/custom_ops/op_sdpa.cpp @@ -1004,14 +1004,22 @@ void channelwise_gated_delta_rule_chunked( const auto* __restrict__ beta_c = beta_head + base; auto* __restrict__ o_c = o_head + base * v_head_dim; - // 1. gc = cumsum_r(log decay) (per channel; resets each chunk) - for (const auto k_idx : c10::irange(k_head_dim)) { - float acc = 0.0f; - for (const auto r : c10::irange(cur)) { - acc += std::log(d_c[r * k_head_dim + k_idx]); - sc.gc[r * k_head_dim + k_idx] = acc; - sc.eg[r * k_head_dim + k_idx] = std::exp(acc); - sc.eg_inv[r * k_head_dim + k_idx] = std::exp(-acc); + // 1. gc = cumsum_r(log decay) (per channel; resets each chunk). + // Reordered so the inner loop runs contiguously over k; the running + // cumsum is read back from the previous gc row. + for (const auto r : c10::irange(cur)) { + const float* __restrict__ d_row = d_c + r * k_head_dim; + const float* __restrict__ gc_prev = + r == 0 ? nullptr : sc.gc + (r - 1) * k_head_dim; + float* __restrict__ gc_row = sc.gc + r * k_head_dim; + float* __restrict__ eg_row = sc.eg + r * k_head_dim; + float* __restrict__ eg_inv_row = sc.eg_inv + r * k_head_dim; + for (const auto k_idx : c10::irange(k_head_dim)) { + const float prev = gc_prev == nullptr ? 0.0f : gc_prev[k_idx]; + const float acc = prev + std::log(d_row[k_idx]); + gc_row[k_idx] = acc; + eg_row[k_idx] = std::exp(acc); + eg_inv_row[k_idx] = std::exp(-acc); } } @@ -1064,47 +1072,64 @@ void channelwise_gated_delta_rule_chunked( } // 4. WY packing: w = A @ (exp(gc) * k), u = A @ v (A lower-tri: j - // <= r) + // <= r). Reordered so the inner loop runs contiguously over k / v. for (const auto r : c10::irange(cur)) { + float* __restrict__ w_row = sc.w + r * k_head_dim; + float* __restrict__ u_row = sc.u + r * v_head_dim; for (const auto k_idx : c10::irange(k_head_dim)) { - float acc = 0.0f; - for (const auto j : c10::irange(r + 1)) { - acc += sc.A[r * CHUNK_SIZE + j] * - sc.eg[j * k_head_dim + k_idx] * - k_c[j * k_head_dim + k_idx]; - } - sc.w[r * k_head_dim + k_idx] = acc; + w_row[k_idx] = 0.0f; } for (const auto v_idx : c10::irange(v_head_dim)) { - float acc = 0.0f; - for (const auto j : c10::irange(r + 1)) { - acc += sc.A[r * CHUNK_SIZE + j] * v_c[j * v_head_dim + v_idx]; + u_row[v_idx] = 0.0f; + } + for (const auto j : c10::irange(r + 1)) { + const float arj = sc.A[r * CHUNK_SIZE + j]; + const float* __restrict__ eg_row = sc.eg + j * k_head_dim; + const float* __restrict__ k_row = k_c + j * k_head_dim; + for (const auto k_idx : c10::irange(k_head_dim)) { + w_row[k_idx] += arj * eg_row[k_idx] * k_row[k_idx]; + } + const float* __restrict__ v_row = v_c + j * v_head_dim; + for (const auto v_idx : c10::irange(v_head_dim)) { + u_row[v_idx] += arj * v_row[v_idx]; } - sc.u[r * v_head_dim + v_idx] = acc; } } - // 5. pv = u - w @ S_old ; o = (q ⊙ exp gc) @ S_old + Aqk @ pv + // 5. pv = u - w @ S_old ; o = (q ⊙ exp gc) @ S_old. Reordered so + // the inner loop runs contiguously over v; pv/o_c accumulate the + // k-sum, then pv is finalized as u - (w@S) with a single + // subtraction (bit-identical to the k-then-subtract form). for (const auto r : c10::irange(cur)) { + float* __restrict__ pv_row = sc.pv + r * v_head_dim; + float* __restrict__ o_row = o_c + r * v_head_dim; for (const auto v_idx : c10::irange(v_head_dim)) { - float wS = 0.0f, qS = 0.0f; - for (const auto k_idx : c10::irange(k_head_dim)) { - const float s = S[k_idx * v_head_dim + v_idx]; - wS += sc.w[r * k_head_dim + k_idx] * s; - qS += q_c[r * k_head_dim + k_idx] * - sc.eg[r * k_head_dim + k_idx] * s; + pv_row[v_idx] = 0.0f; + o_row[v_idx] = 0.0f; + } + for (const auto k_idx : c10::irange(k_head_dim)) { + const float wrk = sc.w[r * k_head_dim + k_idx]; + const float qrk = + q_c[r * k_head_dim + k_idx] * sc.eg[r * k_head_dim + k_idx]; + const float* __restrict__ s_row = S + k_idx * v_head_dim; + for (const auto v_idx : c10::irange(v_head_dim)) { + pv_row[v_idx] += wrk * s_row[v_idx]; + o_row[v_idx] += qrk * s_row[v_idx]; } - sc.pv[r * v_head_dim + v_idx] = - sc.u[r * v_head_dim + v_idx] - wS; - o_c[r * v_head_dim + v_idx] = qS; + } + const float* __restrict__ u_row = sc.u + r * v_head_dim; + for (const auto v_idx : c10::irange(v_head_dim)) { + pv_row[v_idx] = u_row[v_idx] - pv_row[v_idx]; } } + // 5b. o += Aqk @ pv (AXPY over the contiguous v). for (const auto r : c10::irange(cur)) { + float* __restrict__ o_row = o_c + r * v_head_dim; for (const auto j : c10::irange(r + 1)) { const float a = sc.Aqk[r * CHUNK_SIZE + j]; + const float* __restrict__ pv_row = sc.pv + j * v_head_dim; for (const auto v_idx : c10::irange(v_head_dim)) { - o_c[r * v_head_dim + v_idx] += - a * sc.pv[j * v_head_dim + v_idx]; + o_row[v_idx] += a * pv_row[v_idx]; } } } @@ -1120,13 +1145,16 @@ void channelwise_gated_delta_rule_chunked( for (const auto r : c10::irange(cur)) { sc.de[r] = std::exp(gc_last - sc.gc[r * k_head_dim + k_idx]); } + float* __restrict__ S_row = S + k_idx * v_head_dim; for (const auto v_idx : c10::irange(v_head_dim)) { - float acc = S[k_idx * v_head_dim + v_idx] * eg_last; - for (const auto r : c10::irange(cur)) { - acc += k_c[r * k_head_dim + k_idx] * sc.de[r] * - sc.pv[r * v_head_dim + v_idx]; + S_row[v_idx] *= eg_last; + } + for (const auto r : c10::irange(cur)) { + const float c = k_c[r * k_head_dim + k_idx] * sc.de[r]; + const float* __restrict__ pv_row = sc.pv + r * v_head_dim; + for (const auto v_idx : c10::irange(v_head_dim)) { + S_row[v_idx] += c * pv_row[v_idx]; } - S[k_idx * v_head_dim + v_idx] = acc; } } } // chunk