diff --git a/examples/models/llama/attention.py b/examples/models/llama/attention.py index d43533b5a70..4b56744c088 100644 --- a/examples/models/llama/attention.py +++ b/examples/models/llama/attention.py @@ -1,5 +1,7 @@ +import logging from abc import ABC, abstractmethod from enum import Enum +from functools import cache from typing import Any, Dict, Optional, Tuple, Type, TypedDict import torch @@ -68,6 +70,32 @@ def decorator(cls: Type[Attention]): return decorator +@cache +def _get_gated_delta_rule_op() -> Optional[Any]: + """Return the fused gated delta rule op, or None if it is unavailable. + + ``channelwise_gated_delta_rule`` accepts a per-head scalar decay as a 3D + ``[B, H, T]`` tensor, which is the layout GatedDeltaNet produces. Resolve it + lazily (importing the custom ops library on first use) and memoize so + environments without the custom op fall back to the Python recurrence. + """ + try: + return torch.ops.llama.channelwise_gated_delta_rule.default + except (AttributeError, RuntimeError): + pass + + try: + from executorch.extension.llm.custom_ops import custom_ops # noqa: F401 + except (AssertionError, ImportError, OSError, RuntimeError): + logging.debug("Failed to import the ExecuTorch custom ops library") + return None + + try: + return torch.ops.llama.channelwise_gated_delta_rule.default + except (AttributeError, RuntimeError): + return None + + class KVCache(nn.Module): def __init__( self, @@ -762,28 +790,17 @@ def _apply_causal_conv(self, mixed_qkv: torch.Tensor) -> torch.Tensor: out = F.silu(out[:, :, -seq_len:]).to(mixed_qkv.dtype) return out.transpose(1, 2).contiguous() - def _recurrent_gated_delta_rule( + def _naive_gated_delta_rule( self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, - g: torch.Tensor, + decay: torch.Tensor, beta: torch.Tensor, - ) -> torch.Tensor: - # query/key/value: (batch, seq_len, num_heads, head_dim) - # g/beta: (batch, seq_len, num_heads) - initial_dtype = query.dtype - query = _l2norm(query, dim=-1, eps=1e-6) - key = _l2norm(key, dim=-1, eps=1e-6) - query, key, value, beta, g = [ - x.transpose(1, 2).contiguous().to(torch.float32) - for x in (query, key, value, beta, g) - ] - - batch_size, num_heads, sequence_length, k_head_dim = key.shape + initial_state: torch.Tensor, + ) -> Tuple[torch.Tensor, torch.Tensor]: + batch_size, num_heads, sequence_length, _ = key.shape v_head_dim = value.shape[-1] - scale = 1.0 / (query.shape[-1] ** 0.5) - query = query * scale core_attn_out = torch.zeros( batch_size, @@ -793,16 +810,16 @@ def _recurrent_gated_delta_rule( device=value.device, dtype=value.dtype, ) - last_recurrent_state = self.recurrent_state[:batch_size].to(value.dtype) + last_recurrent_state = initial_state for i in range(sequence_length): q_t = query[:, :, i] k_t = key[:, :, i] v_t = value[:, :, i] - g_t = g[:, :, i].exp().unsqueeze(-1).unsqueeze(-1) + decay_t = decay[:, :, i].unsqueeze(-1).unsqueeze(-1) beta_t = beta[:, :, i].unsqueeze(-1) - last_recurrent_state = last_recurrent_state * g_t + last_recurrent_state = last_recurrent_state * decay_t kv_mem = (last_recurrent_state * k_t.unsqueeze(-1)).sum(dim=-2) delta = (v_t - kv_mem) * beta_t last_recurrent_state = last_recurrent_state + k_t.unsqueeze( @@ -812,6 +829,49 @@ def _recurrent_gated_delta_rule( dim=-2 ) + return core_attn_out, last_recurrent_state + + def _recurrent_gated_delta_rule( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + ) -> torch.Tensor: + # query/key/value: (batch, seq_len, num_heads, head_dim) + # g/beta: (batch, seq_len, num_heads) + initial_dtype = query.dtype + query = _l2norm(query, dim=-1, eps=1e-6) + key = _l2norm(key, dim=-1, eps=1e-6) + query, key, value, beta, g = [ + x.transpose(1, 2).contiguous().to(torch.float32) + for x in (query, key, value, beta, g) + ] + + batch_size = key.shape[0] + scale = 1.0 / (query.shape[-1] ** 0.5) + query = query * scale + # The op consumes the decay itself, not its log. + decay = g.exp() + initial_state = self.recurrent_state[:batch_size].to(torch.float32) + + op = _get_gated_delta_rule_op() + # The fused op is an ExecuTorch portable kernel; it cannot run in eager + # mode on CUDA (there is no kernel runtime context, so it aborts with + # "No temp allocator provided"). Restrict it to export/compile tracing + # and fall back to the Python recurrence for eager CUDA. + if op is not None and query.is_cuda and not torch.compiler.is_compiling(): + op = None + if op is not None: + core_attn_out, last_recurrent_state = op( + query, key, value, decay, beta, initial_state + ) + else: + core_attn_out, last_recurrent_state = self._naive_gated_delta_rule( + query, key, value, decay, beta, initial_state + ) + with torch.no_grad(): self.recurrent_state[:batch_size].copy_( last_recurrent_state.to(self.recurrent_state.dtype) diff --git a/examples/models/llama/tests/BUCK b/examples/models/llama/tests/BUCK index c01fa9f2151..563c2f95435 100644 --- a/examples/models/llama/tests/BUCK +++ b/examples/models/llama/tests/BUCK @@ -20,6 +20,10 @@ fbcode_target(_kind = python_unittest, srcs = [ "test_qwen3_5_attention.py", ], + preload_deps = [ + "//executorch/extension/llm/custom_ops:custom_ops_aot_lib", + "//executorch/extension/llm/custom_ops:custom_ops_aot_py", + ], deps = [ "//caffe2:torch", "//executorch/examples/models/llama:llama_transformer", diff --git a/examples/models/llama/tests/test_qwen3_5_attention.py b/examples/models/llama/tests/test_qwen3_5_attention.py index 5a9f67d57cf..d2ead9cd0bb 100644 --- a/examples/models/llama/tests/test_qwen3_5_attention.py +++ b/examples/models/llama/tests/test_qwen3_5_attention.py @@ -5,8 +5,10 @@ # LICENSE file in the root directory of this source tree. import unittest +from unittest import mock import torch +from executorch.examples.models.llama import attention as attention_module from executorch.examples.models.llama.attention import ATTENTION_REGISTRY from executorch.examples.models.llama.model_args import ModelArgs from executorch.examples.models.llama.norm import RMSNorm @@ -123,6 +125,64 @@ def test_gated_deltanet_no_input_pos_does_not_leak_state(self): torch.allclose(state_after_first, state_after_second, atol=1e-5) ) + def test_gated_deltanet_fused_op_matches_python_recurrence(self): + # The custom op takes the per-head decay as a 3D [B, H, T] tensor and + # must reproduce the Python token-by-token recurrence exactly, for both + # the decode (T == 1) and prefill (T > 1) routes. + from executorch.extension.llm.custom_ops import custom_ops # noqa: F401 + + op = torch.ops.llama.channelwise_gated_delta_rule.default + + for seq_len in (1, 5, 40): + with self.subTest(seq_len=seq_len): + torch.manual_seed(seq_len) + args = self._make_args( + max_seq_len=64, + max_context_len=64, + use_kv_cache=True, + use_q_gate=True, + linear_conv_kernel_dim=4, + linear_key_head_dim=4, + linear_value_head_dim=4, + linear_num_key_heads=2, + linear_num_value_heads=4, + ) + rope = Rope(args) + attn = ATTENTION_REGISTRY["gated_deltanet"](args, 0, rope).eval() + + x = torch.randn(1, seq_len, args.dim) + dummy_freq = torch.zeros(1, 1) + # input_pos != 0 keeps the seeded state instead of resetting it, + # so the non-trivial initial_state actually reaches the kernel. + input_pos = torch.tensor([1], dtype=torch.long) + initial_conv = torch.randn_like(attn.conv_state) + initial_recurrent = torch.randn_like(attn.recurrent_state) + + def run(resolved_op): + with torch.no_grad(): + # Both buffers are mutated in place by forward; restore + # them so each path starts from identical state. + attn.conv_state.copy_(initial_conv) + attn.recurrent_state.copy_(initial_recurrent) + with mock.patch.object( + attention_module, + "_get_gated_delta_rule_op", + return_value=resolved_op, + ): + out, _ = attn( + x, dummy_freq, dummy_freq, input_pos=input_pos + ) + return out, attn.recurrent_state.clone() + + op_out, op_state = run(op) + naive_out, naive_state = run(None) + + self.assertFalse(torch.allclose(op_state, initial_recurrent)) + self.assertTrue(torch.allclose(op_out, naive_out, atol=1e-4, rtol=1e-4)) + self.assertTrue( + torch.allclose(op_state, naive_state, atol=1e-4, rtol=1e-4) + ) + if __name__ == "__main__": unittest.main() diff --git a/extension/llm/custom_ops/bench_channelwise_gated_delta_rule.py b/extension/llm/custom_ops/bench_channelwise_gated_delta_rule.py index 9fd43066d72..e47baddbdea 100644 --- a/extension/llm/custom_ops/bench_channelwise_gated_delta_rule.py +++ b/extension/llm/custom_ops/bench_channelwise_gated_delta_rule.py @@ -17,6 +17,7 @@ python -m executorch.extension.llm.custom_ops.bench_channelwise_gated_delta_rule """ +import itertools import time import torch @@ -26,21 +27,29 @@ _OP = torch.ops.llama.channelwise_gated_delta_rule.default -def _make_inputs(b: int, h: int, t: int, k: int, v: int): +def _make_inputs(b: int, h: int, t: int, k: int, v: int, scalar_decay: bool = False): + decay_shape = [b, h, t] if scalar_decay else [b, h, t, k] 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(decay_shape), 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 + b: int, + h: int, + t: int, + k: int, + v: int, + iters: int, + warmup: int, + scalar_decay: bool = False, ) -> tuple[float, float]: - q, key, val, decay, beta, state = _make_inputs(b, h, t, k, v) + q, key, val, decay, beta, state = _make_inputs(b, h, t, k, v, scalar_decay) for _ in range(warmup): _OP(q, key, val, decay, beta, state) start = time.perf_counter() @@ -68,10 +77,27 @@ def main() -> None: 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}") + print( + f"{'config':<15}{'decay':<12}{'B':>3}{'H':>4}{'T':>6}" + f"{'ms/call':>12}{'us/token':>12}" + ) + for (label, b, h, t, iters), scalar_decay in itertools.product( + configs, (False, True) + ): + ms, us = _bench_one( + b, + h, + t, + k, + v, + iters=iters, + warmup=max(10, iters // 10), + scalar_decay=scalar_decay, + ) + print( + f"{label:<15}{'scalar' if scalar_decay else 'channelwise':<12}" + f"{b:>3}{h:>4}{t:>6}{ms:>12.4f}{us:>12.3f}" + ) if __name__ == "__main__": diff --git a/extension/llm/custom_ops/custom_ops.py b/extension/llm/custom_ops/custom_ops.py index bd4bd7864ba..ffefc4d05f5 100644 --- a/extension/llm/custom_ops/custom_ops.py +++ b/extension/llm/custom_ops/custom_ops.py @@ -396,9 +396,10 @@ def _validate_channelwise_gated_delta_rule_params( assert ( value.dim() == 4 ), f"Expected value to be 4 dimensional but got {value.dim()} dimensions." - assert ( - decay.dim() == 4 - ), f"Expected decay to be 4 dimensional but got {decay.dim()} dimensions." + assert decay.dim() in (3, 4), ( + "Expected decay to be 4 dimensional (channelwise) or 3 dimensional " + f"(scalar) but got {decay.dim()} dimensions." + ) assert ( beta.dim() == 3 ), f"Expected beta to be 3 dimensional but got {beta.dim()} dimensions." @@ -419,7 +420,7 @@ def _validate_channelwise_gated_delta_rule_params( ), f"Expected {name} to be float32 but got {tensor.dtype}" assert query.size(0) == key.size(0) and query.shape[2:] == key.shape[2:] - assert key.shape == decay.shape + assert decay.shape == (key.shape if decay.dim() == 4 else key.shape[:3]) assert key.shape[:3] == value.shape[:3] assert beta.shape == key.shape[:3] assert query.size(1) % key.size(1) == 0 diff --git a/extension/llm/custom_ops/op_sdpa.cpp b/extension/llm/custom_ops/op_sdpa.cpp index 79ad3388e02..0be786d7c42 100644 --- a/extension/llm/custom_ops/op_sdpa.cpp +++ b/extension/llm/custom_ops/op_sdpa.cpp @@ -203,7 +203,9 @@ bool validate_channelwise_gated_delta_rule_args( ET_CHECK_OR_RETURN_FALSE(query.dim() == 4, "query must be a 4D tensor"); ET_CHECK_OR_RETURN_FALSE(key.dim() == 4, "key must be a 4D tensor"); ET_CHECK_OR_RETURN_FALSE(value.dim() == 4, "value must be a 4D tensor"); - ET_CHECK_OR_RETURN_FALSE(decay.dim() == 4, "decay must be a 4D tensor"); + ET_CHECK_OR_RETURN_FALSE( + decay.dim() == 4 || decay.dim() == 3, + "decay must be a 4D tensor (channelwise) or a 3D tensor (scalar)"); ET_CHECK_OR_RETURN_FALSE(beta.dim() == 3, "beta must be a 3D tensor"); ET_CHECK_OR_RETURN_FALSE( initial_state.dim() == 4, "initial_state must be a 4D tensor"); @@ -228,8 +230,9 @@ bool validate_channelwise_gated_delta_rule_args( "query and key must match in batch/sequence/head-dim"); ET_CHECK_OR_RETURN_FALSE( key.size(0) == decay.size(0) && key.size(1) == decay.size(1) && - key.size(2) == decay.size(2) && key.size(3) == decay.size(3), - "key and decay must have matching shapes"); + key.size(2) == decay.size(2) && + (decay.dim() == 3 || key.size(3) == decay.size(3)), + "decay must be [B, Hkv, T, K] (channelwise) or [B, Hkv, T] (scalar)"); ET_CHECK_OR_RETURN_FALSE( key.size(0) == value.size(0) && key.size(1) == value.size(1) && key.size(2) == value.size(2), @@ -252,7 +255,7 @@ bool validate_channelwise_gated_delta_rule_args( {&query, &key, &value, &decay, &beta, &initial_state}) { ET_CHECK_OR_RETURN_FALSE( is_contiguous_dim_order((*tensor).dim_order().data(), (*tensor).dim()), - "channelwise gated delta rule expects contiguous inputs"); + "gated delta rule expects contiguous inputs"); } return true; @@ -795,7 +798,14 @@ ET_INLINE void vec_state_update_and_output( // the two-pass comments inline). Used by the T==1 decode path and as the exact // prefill fallback when decay contains non-positive entries -- the chunked // kernel takes log(decay) and cannot represent decay <= 0. -void channelwise_gated_delta_rule_recurrence( +// +// kScalarDecay selects the decay layout: channelwise decay is [B, Hkv, T, K] +// (one value per state row), scalar decay is [B, Hkv, T] and is broadcast +// across the K rows. Templating rather than passing a runtime stride keeps the +// channelwise codegen unchanged and turns the scalar read into a +// loop-invariant load. +template +void gated_delta_rule_recurrence( RuntimeContext& ctx, const Tensor& query, const Tensor& key, @@ -805,6 +815,7 @@ void channelwise_gated_delta_rule_recurrence( const Tensor& initial_state, Tensor& out, Tensor& final_state_out) { + constexpr int64_t decay_channel_stride = kScalarDecay ? 0 : 1; const auto batch_size = query.size(0); const auto num_query_heads = query.size(1); const auto num_kv_heads = key.size(1); @@ -828,6 +839,10 @@ void channelwise_gated_delta_rule_recurrence( const auto beta_batch_stride = num_kv_heads * sequence_length; const auto beta_head_stride = sequence_length; + const auto decay_seq_stride = kScalarDecay ? int64_t{1} : k_head_dim; + const auto decay_head_stride = sequence_length * decay_seq_stride; + const auto decay_batch_stride = num_kv_heads * decay_head_stride; + const auto state_batch_stride = num_kv_heads * k_head_dim * v_head_dim; const auto state_head_stride = k_head_dim * v_head_dim; @@ -868,7 +883,8 @@ void channelwise_gated_delta_rule_recurrence( batch * state_batch_stride + kv_head * state_head_stride; const auto* k_head = key_data + kv_offset; - const auto* decay_head = decay_data + kv_offset; + const auto* decay_head = + decay_data + batch * decay_batch_stride + kv_head * decay_head_stride; const auto* value_head = value_data + value_offset; const auto* beta_head = beta_data + beta_offset; const auto* initial_state_head = initial_state_data + state_offset; @@ -884,7 +900,7 @@ void channelwise_gated_delta_rule_recurrence( for (int64_t token = 0; token < sequence_length; ++token) { const auto* k_t = k_head + token * qk_seq_stride; - const auto* decay_t = decay_head + token * qk_seq_stride; + const auto* decay_t = decay_head + token * decay_seq_stride; const auto* v_t = value_head + token * value_seq_stride; const float beta_t = beta_head[token]; @@ -900,7 +916,7 @@ void channelwise_gated_delta_rule_recurrence( // 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]; + const float decay_value = decay_t[k_idx * decay_channel_stride]; const float key_value = k_t[k_idx]; const auto* state_row = state_head + k_idx * v_head_dim; vec_axpy(v_pred, state_row, decay_value * key_value, v_head_dim); @@ -917,7 +933,7 @@ void channelwise_gated_delta_rule_recurrence( auto* output_t = output_group + token * value_seq_stride; 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 decay_value = decay_t[k_idx * decay_channel_stride]; const float key_value = k_t[k_idx]; const float query_value = q_t[k_idx]; auto* state_row = state_head + k_idx * v_head_dim; @@ -934,7 +950,7 @@ void channelwise_gated_delta_rule_recurrence( } for (int64_t k_idx = 0; k_idx < k_head_dim; ++k_idx) { - const float decay_value = decay_t[k_idx]; + const float decay_value = decay_t[k_idx * decay_channel_stride]; const float key_value = k_t[k_idx]; auto* state_row = state_head + k_idx * v_head_dim; int64_t idx = 0; @@ -999,10 +1015,11 @@ inline int64_t chunk_scratch_align(int64_t n) { // 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. +// D below is K for channelwise decay and 1 for scalar decay. 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* gc; // [CHUNK_SIZE, D] cumulative log-decay + float* eg; // [CHUNK_SIZE, D] exp(gc) + float* eg_inv; // [CHUNK_SIZE, D] 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 @@ -1011,24 +1028,32 @@ struct ChunkScratch { 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 + static int64_t + elems(int64_t k_head_dim, int64_t v_head_dim, int64_t decay_channels) { + return 3 * + chunk_scratch_align(CHUNK_SIZE * decay_channels) + // gc,eg,eg_inv + chunk_scratch_align(CHUNK_SIZE * k_head_dim) + // 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) { + static ChunkScratch view( + float* p, + int64_t k_head_dim, + int64_t v_head_dim, + int64_t decay_channels) { + const int64_t nd = chunk_scratch_align(CHUNK_SIZE * decay_channels); 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; + p += nd; s.eg = p; - p += nk; + p += nd; s.eg_inv = p; - p += nk; + p += nd; s.w = p; p += nk; s.u = p; @@ -1044,8 +1069,9 @@ struct ChunkScratch { } }; -// T != 1 (prefill) route. -void channelwise_gated_delta_rule_chunked( +// T != 1 (prefill) route. See gated_delta_rule_recurrence for kScalarDecay. +template +void gated_delta_rule_chunked( RuntimeContext& ctx, const Tensor& query, const Tensor& key, @@ -1055,6 +1081,7 @@ void channelwise_gated_delta_rule_chunked( const Tensor& initial_state, Tensor& out, Tensor& final_state_out) { + constexpr int64_t decay_channel_stride = kScalarDecay ? 0 : 1; const auto batch_size = query.size(0); const auto num_query_heads = query.size(1); const auto num_kv_heads = key.size(1); @@ -1073,17 +1100,21 @@ void channelwise_gated_delta_rule_chunked( auto* state_data = final_state_out.mutable_data_ptr(); auto* output_data = out.mutable_data_ptr(); - // Per-head strides. q/k/decay are [.., T, K]; v/out are [.., T, V]. + // Per-head strides. q/k are [.., T, K]; v/out are [.., T, V]; decay is + // [.., T, K] channelwise and [.., T] scalar. 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; + const auto decay_channels = kScalarDecay ? int64_t{1} : k_head_dim; + const auto decay_head_stride = sequence_length * decay_channels; // 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); + const int64_t size_per_thread = + ChunkScratch::elems(k_head_dim, v_head_dim, decay_channels); #ifdef ET_USE_THREADPOOL const int64_t num_thread = ::executorch::extension::threadpool::get_threadpool()->get_thread_count(); @@ -1127,10 +1158,13 @@ void channelwise_gated_delta_rule_chunked( const int64_t query_head_begin = kv_head * query_heads_per_kv; 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); + arena + tid * size_per_thread, + k_head_dim, + v_head_dim, + decay_channels); 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__ d_head = decay_data + bh * decay_head_stride; const auto* __restrict__ beta_head = beta_data + bh * beta_head_stride; const auto* __restrict__ init_head = @@ -1150,7 +1184,7 @@ void channelwise_gated_delta_rule_chunked( std::min(CHUNK_SIZE, sequence_length - base); 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__ d_c = d_head + base * decay_channels; const auto* __restrict__ beta_c = beta_head + base; const auto* __restrict__ single_q_c = query_heads_per_kv == 1 ? query_data + @@ -1167,15 +1201,17 @@ void channelwise_gated_delta_rule_chunked( // 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. + // cumsum is read back from the previous gc row. Scalar decay has a + // single column, so this sweep costs one log/exp pair per token + // instead of k_head_dim of them. for (const auto r : c10::irange(cur)) { - const float* __restrict__ d_row = d_c + r * k_head_dim; + const float* __restrict__ d_row = d_c + r * decay_channels; 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)) { + r == 0 ? nullptr : sc.gc + (r - 1) * decay_channels; + float* __restrict__ gc_row = sc.gc + r * decay_channels; + float* __restrict__ eg_row = sc.eg + r * decay_channels; + float* __restrict__ eg_inv_row = sc.eg_inv + r * decay_channels; + for (const auto k_idx : c10::irange(decay_channels)) { 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; @@ -1194,8 +1230,9 @@ void channelwise_gated_delta_rule_chunked( 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]; + const int64_t d_idx = k_idx * decay_channel_stride; + const float ratio = sc.eg[r * decay_channels + d_idx] * + sc.eg_inv[j * decay_channels + d_idx]; if (query_heads_per_kv == 1) { aqk += single_q_c[r * k_head_dim + k_idx] * k_c[j * k_head_dim + k_idx] * ratio; @@ -1245,10 +1282,11 @@ void channelwise_gated_delta_rule_chunked( } 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__ eg_row = sc.eg + j * decay_channels; 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]; + w_row[k_idx] += + arj * eg_row[k_idx * decay_channel_stride] * k_row[k_idx]; } const float* __restrict__ v_row = v_c + j * v_head_dim; vec_axpy(u_row, v_row, arj, v_head_dim); @@ -1272,7 +1310,7 @@ void channelwise_gated_delta_rule_chunked( const float* __restrict__ s_row = S + k_idx * v_head_dim; const float qrk = query_heads_per_kv == 1 ? single_q_c[r * k_head_dim + k_idx] * - sc.eg[r * k_head_dim + k_idx] + sc.eg[r * decay_channels + k_idx * decay_channel_stride] : 0.0f; for (const auto v_idx : c10::irange(v_head_dim)) { pv_row[v_idx] += wrk * s_row[v_idx]; @@ -1315,8 +1353,9 @@ void channelwise_gated_delta_rule_chunked( for (const auto j : c10::irange(r + 1)) { float aqk = 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]; + const int64_t d_idx = k_idx * decay_channel_stride; + const float ratio = sc.eg[r * decay_channels + d_idx] * + sc.eg_inv[j * decay_channels + d_idx]; aqk += q_c[r * k_head_dim + k_idx] * k_c[j * k_head_dim + k_idx] * ratio; } @@ -1331,7 +1370,7 @@ void channelwise_gated_delta_rule_chunked( } for (const auto k_idx : c10::irange(k_head_dim)) { const float qrk = q_c[r * k_head_dim + k_idx] * - sc.eg[r * k_head_dim + k_idx]; + sc.eg[r * decay_channels + k_idx * decay_channel_stride]; const float* __restrict__ s_row = S + k_idx * v_head_dim; vec_axpy(o_row, s_row, qrk, v_head_dim); } @@ -1349,10 +1388,16 @@ void channelwise_gated_delta_rule_chunked( // 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]); + const int64_t d_idx = k_idx * decay_channel_stride; + const float gc_last = sc.gc[(cur - 1) * decay_channels + d_idx]; + const float eg_last = sc.eg[(cur - 1) * decay_channels + d_idx]; + // Every state row shares one decay column under scalar decay, so + // the exp() sweep only has to run for the first row. + if (!kScalarDecay || k_idx == 0) { + for (const auto r : c10::irange(cur)) { + sc.de[r] = + std::exp(gc_last - sc.gc[r * decay_channels + d_idx]); + } } float* __restrict__ S_row = S + k_idx * v_head_dim; vec_scale(S_row, eg_last, v_head_dim); @@ -1379,14 +1424,15 @@ void channelwise_gated_delta_rule_chunked( // whole head's output. A non-positive decay is likewise undefined (std::log). // Detect both so callers can route to the exact recurrence, which is stable for // tiny decay and handles decay <= 0. NaN decay is treated as unsafe as well. -bool channelwise_gated_delta_rule_chunked_is_safe(const Tensor& decay) { +bool gated_delta_rule_chunked_is_safe(const Tensor& decay) { // Margin below the float32 expf overflow threshold (std::exp overflows around // 88.7); 80 leaves headroom for the downstream eg * eg_inv product. constexpr float kMaxNegLogSum = 80.0f; const auto* decay_data = decay.const_data_ptr(); const auto batch_heads = decay.size(0) * decay.size(1); const auto sequence_length = decay.size(2); - const auto k_head_dim = decay.size(3); + // Scalar decay ([B, Hkv, T]) carries a single column per token. + const auto k_head_dim = decay.dim() == 4 ? decay.size(3) : 1; const auto head_stride = sequence_length * k_head_dim; for (int64_t bh = 0; bh < batch_heads; ++bh) { const float* decay_head = decay_data + bh * head_stride; @@ -1412,6 +1458,79 @@ bool channelwise_gated_delta_rule_chunked_is_safe(const Tensor& decay) { return true; } +// Bridge the runtime decay layout to the compile-time kernel specialization. +void run_gated_delta_rule_recurrence( + RuntimeContext& ctx, + const Tensor& query, + const Tensor& key, + const Tensor& value, + const Tensor& decay, + const Tensor& beta, + const Tensor& initial_state, + bool scalar_decay, + Tensor& out, + Tensor& final_state_out) { + if (scalar_decay) { + gated_delta_rule_recurrence( + ctx, + query, + key, + value, + decay, + beta, + initial_state, + out, + final_state_out); + } else { + gated_delta_rule_recurrence( + ctx, + query, + key, + value, + decay, + beta, + initial_state, + out, + final_state_out); + } +} + +void run_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, + bool scalar_decay, + Tensor& out, + Tensor& final_state_out) { + if (scalar_decay) { + gated_delta_rule_chunked( + ctx, + query, + key, + value, + decay, + beta, + initial_state, + out, + final_state_out); + } else { + gated_delta_rule_chunked( + ctx, + query, + key, + value, + decay, + beta, + initial_state, + out, + final_state_out); + } +} + std::tuple channelwise_gated_delta_rule_out( RuntimeContext& ctx, const Tensor& query, @@ -1482,19 +1601,9 @@ std::tuple channelwise_gated_delta_rule_out( // unsafe // -- non-positive decay (log undefined) or tiny positive decay whose // within-chunk cumulative -log(decay) would overflow exp(-gc) to inf -> NaN. - if (query.size(2) == 1) { - channelwise_gated_delta_rule_recurrence( - ctx, - query, - key, - value, - decay, - beta, - initial_state, - out, - final_state_out); - } else if (channelwise_gated_delta_rule_chunked_is_safe(decay)) { - channelwise_gated_delta_rule_chunked( + const bool scalar_decay = decay.dim() == 3; + if (query.size(2) != 1 && gated_delta_rule_chunked_is_safe(decay)) { + run_gated_delta_rule_chunked( ctx, query, key, @@ -1502,10 +1611,11 @@ std::tuple channelwise_gated_delta_rule_out( decay, beta, initial_state, + scalar_decay, out, final_state_out); } else { - channelwise_gated_delta_rule_recurrence( + run_gated_delta_rule_recurrence( ctx, query, key, @@ -1513,6 +1623,7 @@ std::tuple channelwise_gated_delta_rule_out( decay, beta, initial_state, + scalar_decay, out, final_state_out); } diff --git a/extension/llm/custom_ops/op_sdpa_test.cpp b/extension/llm/custom_ops/op_sdpa_test.cpp index 1e07800be9f..f14233c089a 100644 --- a/extension/llm/custom_ops/op_sdpa_test.cpp +++ b/extension/llm/custom_ops/op_sdpa_test.cpp @@ -60,6 +60,167 @@ op_channelwise_gated_delta_rule( Most tests are generated by FACTO */ +namespace { + +// A scalar decay of [B, Hkv, T] must produce exactly what the channelwise +// kernel produces when the same value is broadcast across the K state rows. +void test_scalar_decay_matches_broadcast_channelwise( + int32_t sequence_length, + int32_t query_heads_per_kv) { + TensorFactory tfFloat; + constexpr int32_t batch_size = 2; + constexpr int32_t num_kv_heads = 2; + constexpr int32_t k_head_dim = 6; + constexpr int32_t v_head_dim = 4; + const int32_t num_query_heads = num_kv_heads * query_heads_per_kv; + + auto fill = [](std::vector& values, float scale, size_t period) { + for (size_t i = 0; i < values.size(); ++i) { + values[i] = + static_cast(i % period) - static_cast(period / 2); + values[i] *= scale; + } + }; + + const size_t query_numel = static_cast(batch_size) * num_query_heads * + sequence_length * k_head_dim; + const size_t key_numel = static_cast(batch_size) * num_kv_heads * + sequence_length * k_head_dim; + const size_t value_numel = static_cast(batch_size) * num_kv_heads * + sequence_length * v_head_dim; + const size_t beta_numel = + static_cast(batch_size) * num_kv_heads * sequence_length; + const size_t state_numel = + static_cast(batch_size) * num_kv_heads * k_head_dim * v_head_dim; + + std::vector query_values(query_numel); + std::vector key_values(key_numel); + std::vector value_values(value_numel); + std::vector beta_values(beta_numel); + std::vector scalar_decay_values(beta_numel); + std::vector initial_state_values(state_numel); + fill(query_values, 0.11f, 7); + fill(key_values, 0.13f, 5); + fill(value_values, 0.17f, 9); + fill(beta_values, 0.07f, 6); + fill(initial_state_values, 0.05f, 11); + for (size_t i = 0; i < beta_numel; ++i) { + scalar_decay_values[i] = 0.5f + static_cast(i % 5) * 0.1f; + } + + // Same gate value repeated across every channel of the head. + std::vector channelwise_decay_values(beta_numel * k_head_dim); + for (size_t i = 0; i < beta_numel; ++i) { + std::fill( + channelwise_decay_values.begin() + i * k_head_dim, + channelwise_decay_values.begin() + (i + 1) * k_head_dim, + scalar_decay_values[i]); + } + + auto query = tfFloat.make( + {batch_size, num_query_heads, sequence_length, k_head_dim}, query_values); + auto key = tfFloat.make( + {batch_size, num_kv_heads, sequence_length, k_head_dim}, key_values); + auto value = tfFloat.make( + {batch_size, num_kv_heads, sequence_length, v_head_dim}, value_values); + auto beta = + tfFloat.make({batch_size, num_kv_heads, sequence_length}, beta_values); + auto initial_state = tfFloat.make( + {batch_size, num_kv_heads, k_head_dim, v_head_dim}, initial_state_values); + auto scalar_decay = tfFloat.make( + {batch_size, num_kv_heads, sequence_length}, scalar_decay_values); + auto channelwise_decay = tfFloat.make( + {batch_size, num_kv_heads, sequence_length, k_head_dim}, + channelwise_decay_values); + + auto scalar_out = + tfFloat.zeros({batch_size, num_query_heads, sequence_length, v_head_dim}); + auto scalar_state = + tfFloat.zeros({batch_size, num_kv_heads, k_head_dim, v_head_dim}); + auto channelwise_out = + tfFloat.zeros({batch_size, num_query_heads, sequence_length, v_head_dim}); + auto channelwise_state = + tfFloat.zeros({batch_size, num_kv_heads, k_head_dim, v_head_dim}); + + executorch::runtime::KernelRuntimeContext scalar_context{}; + op_channelwise_gated_delta_rule( + scalar_context, + query, + key, + value, + scalar_decay, + beta, + initial_state, + scalar_out, + scalar_state); + ASSERT_EQ(scalar_context.failure_state(), torch::executor::Error::Ok); + + executorch::runtime::KernelRuntimeContext channelwise_context{}; + op_channelwise_gated_delta_rule( + channelwise_context, + query, + key, + value, + channelwise_decay, + beta, + initial_state, + channelwise_out, + channelwise_state); + ASSERT_EQ(channelwise_context.failure_state(), torch::executor::Error::Ok); + + EXPECT_TENSOR_CLOSE_WITH_TOL(scalar_out, channelwise_out, 1e-5, 1e-5); + EXPECT_TENSOR_CLOSE_WITH_TOL(scalar_state, channelwise_state, 1e-5, 1e-5); +} + +} // namespace + +TEST(ChannelwiseGatedDeltaRuleTest, ScalarDecayDecodeMatchesBroadcast) { + test_scalar_decay_matches_broadcast_channelwise( + /*sequence_length=*/1, /*query_heads_per_kv=*/1); +} + +TEST(ChannelwiseGatedDeltaRuleTest, ScalarDecayPrefillMatchesBroadcast) { + test_scalar_decay_matches_broadcast_channelwise( + /*sequence_length=*/40, /*query_heads_per_kv=*/1); +} + +TEST(ChannelwiseGatedDeltaRuleTest, ScalarDecayGroupedMatchesBroadcast) { + test_scalar_decay_matches_broadcast_channelwise( + /*sequence_length=*/40, /*query_heads_per_kv=*/3); +} + +TEST(ChannelwiseGatedDeltaRuleTest, ScalarDecayGroupedDecodeMatchesBroadcast) { + test_scalar_decay_matches_broadcast_channelwise( + /*sequence_length=*/1, /*query_heads_per_kv=*/3); +} + +TEST(ChannelwiseGatedDeltaRuleTest, RejectsMismatchedScalarDecay) { + TensorFactory tfFloat; + + executorch::aten::Tensor query = tfFloat.ones({1, 1, 3, 2}); + executorch::aten::Tensor key = tfFloat.ones({1, 1, 3, 2}); + executorch::aten::Tensor value = tfFloat.ones({1, 1, 3, 2}); + executorch::aten::Tensor decay = tfFloat.ones({1, 1, 2}); + executorch::aten::Tensor beta = tfFloat.ones({1, 1, 3}); + executorch::aten::Tensor initial_state = tfFloat.ones({1, 1, 2, 2}); + executorch::aten::Tensor out = tfFloat.zeros({1, 1, 3, 2}); + executorch::aten::Tensor final_state_out = tfFloat.zeros({1, 1, 2, 2}); + + executorch::runtime::KernelRuntimeContext context{}; + op_channelwise_gated_delta_rule( + context, + query, + key, + value, + decay, + beta, + initial_state, + out, + final_state_out); + + EXPECT_NE(context.failure_state(), torch::executor::Error::Ok); +} + TEST(ChannelwiseGatedDeltaRuleTest, RejectsPartiallyOverlappingFinalStateOut) { TensorFactory tfFloat; diff --git a/extension/llm/custom_ops/test_gated_delta.py b/extension/llm/custom_ops/test_gated_delta.py index edd6f6e7305..b11ec3735f7 100644 --- a/extension/llm/custom_ops/test_gated_delta.py +++ b/extension/llm/custom_ops/test_gated_delta.py @@ -6,6 +6,7 @@ # pyre-unsafe +import itertools import sys import unittest @@ -23,14 +24,19 @@ def _make_inputs( k_head_dim: int = 5, v_head_dim: int = 6, num_state_heads: int | None = None, + scalar_decay: bool = False, ): if num_state_heads is None: num_state_heads = num_heads query = torch.randn(batch_size, num_heads, seq_len, k_head_dim) key = torch.randn(batch_size, num_state_heads, seq_len, k_head_dim) value = torch.randn(batch_size, num_state_heads, seq_len, v_head_dim) - # Per-key-channel decay, passed already exponentiated (in (0, 1)). - decay = torch.rand(batch_size, num_state_heads, seq_len, k_head_dim) + # Decay is passed already exponentiated (in (0, 1)), per key channel or + # — with scalar_decay — one gate per (batch, state head, token). + decay_shape = [batch_size, num_state_heads, seq_len] + if not scalar_decay: + decay_shape.append(k_head_dim) + decay = torch.rand(decay_shape) beta = torch.sigmoid(torch.randn(batch_size, num_state_heads, seq_len)) initial_state = torch.randn(batch_size, num_state_heads, k_head_dim, v_head_dim) return query, key, value, decay, beta, initial_state @@ -49,6 +55,10 @@ def _reference_channelwise_gated_delta_rule( query.size(0), query.size(1), query.size(2), value.size(3) ) heads_per_state = query.size(1) // key.size(1) + # A scalar gate applies to every state row, so broadcast it out and + # share the channelwise body below. + if decay.dim() == 3: + decay = decay.unsqueeze(-1).expand_as(key) for token in range(query.size(2)): # Per-key-channel decay: [B, H, K, 1], already exponentiated. @@ -92,6 +102,49 @@ def test_channelwise_gated_delta_rule_grouped_matches_reference(self): torch.allclose(actual_state, expected_state, atol=1e-3, rtol=1e-3) ) + def test_scalar_decay_matches_reference(self): + torch.manual_seed(0) + + # T == 1 hits the decode recurrence; 35 and 130 hit the chunked prefill + # route with a ragged final chunk (CHUNK_SIZE is 32 in the C++ kernel). + for seq_len, num_heads, num_state_heads in ( + (1, 4, 2), + (35, 4, 2), + (130, 3, 3), + ): + with self.subTest(seq_len=seq_len, num_state_heads=num_state_heads): + inputs = self._make_inputs( + batch_size=2, + num_heads=num_heads, + num_state_heads=num_state_heads, + seq_len=seq_len, + scalar_decay=True, + ) + self.assertEqual(inputs[3].dim(), 3) + + 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_scalar_decay_rejects_mismatched_sequence_length(self): + query, key, value, _, beta, initial_state = self._make_inputs(seq_len=4) + short_decay = torch.rand(query.size(0), key.size(1), 3) + + with self.assertRaises(RuntimeError): + torch.ops.llama.channelwise_gated_delta_rule( + query, key, value, short_decay, beta, initial_state + ) + def test_channelwise_gated_delta_rule_rejects_uneven_groups(self): inputs = self._make_inputs(num_heads=3, num_state_heads=2) with self.assertRaises(RuntimeError): @@ -380,32 +433,39 @@ def forward(self, query, key, value, decay, beta, initial_state): query, key, value, decay, beta, initial_state ) - inputs = self._make_inputs() - - # Static export: the op must survive as a single graph node (the Meta - # impl lets it trace without running the real kernel). - ep = torch.export.export(Module(), inputs) - targets = [str(n.target) for n in ep.graph.nodes if n.op == "call_function"] - self.assertIn("llama.channelwise_gated_delta_rule.default", targets) - - # Dynamic sequence length: one graph shared across prefill/decode. - seq = torch.export.Dim("seq", min=1, max=128) - dynamic_shapes = ( - {2: seq}, # query [B, H, T, K] - {2: seq}, # key [B, H, T, K] - {2: seq}, # value [B, H, T, V] - {2: seq}, # decay [B, H, T, K] - {2: seq}, # beta [B, H, T] - {}, # initial_state [B, H, K, V] (no T dim) - ) - ep_dyn = torch.export.export(Module(), inputs, dynamic_shapes=dynamic_shapes) - self.assertTrue( - any( - "channelwise_gated_delta_rule" in str(n.target) - for n in ep_dyn.graph.nodes - if n.op == "call_function" - ) - ) + for scalar_decay in (False, True): + with self.subTest(scalar_decay=scalar_decay): + inputs = self._make_inputs(scalar_decay=scalar_decay) + + # Static export: the op must survive as a single graph node (the + # Meta impl lets it trace without running the real kernel). + ep = torch.export.export(Module(), inputs) + targets = [ + str(n.target) for n in ep.graph.nodes if n.op == "call_function" + ] + self.assertIn("llama.channelwise_gated_delta_rule.default", targets) + + # Dynamic sequence length: one graph shared across + # prefill/decode. T is dim 2 for both decay layouts. + seq = torch.export.Dim("seq", min=1, max=128) + dynamic_shapes = ( + {2: seq}, # query [B, H, T, K] + {2: seq}, # key [B, H, T, K] + {2: seq}, # value [B, H, T, V] + {2: seq}, # decay [B, H, T, K] or [B, H, T] + {2: seq}, # beta [B, H, T] + {}, # initial_state [B, H, K, V] (no T dim) + ) + ep_dyn = torch.export.export( + Module(), inputs, dynamic_shapes=dynamic_shapes + ) + self.assertTrue( + any( + "channelwise_gated_delta_rule" in str(n.target) + for n in ep_dyn.graph.nodes + if n.op == "call_function" + ) + ) @unittest.skipUnless( sys.platform == "linux", @@ -437,10 +497,11 @@ def forward(self, query, key, value, decay, beta, initial_state): ) runtime = Runtime.get() - for seq_len in (4, 1): # T != 1 (chunked route) and T == 1 (decode route) - with self.subTest(seq_len=seq_len): + # T != 1 (chunked route) and T == 1 (decode route), each decay layout. + for seq_len, scalar_decay in itertools.product((4, 1), (False, True)): + with self.subTest(seq_len=seq_len, scalar_decay=scalar_decay): torch.manual_seed(seq_len) - inputs = self._make_inputs(seq_len=seq_len) + inputs = self._make_inputs(seq_len=seq_len, scalar_decay=scalar_decay) expected_output, expected_state = ( self._reference_channelwise_gated_delta_rule(*inputs) )