Support scalar (non-channelwise) decay in the GDN kernels

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)

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",

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 @@
             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()

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__":

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