Skip to content

Commit 928bffa

Browse files
JacobSzwejbkaJacobSzwejbka
committed
Reland recurrent gated delta rule custom op
1 parent 4dc1b5b commit 928bffa

16 files changed

Lines changed: 1243 additions & 102 deletions

‎.ci/scripts/test_llama.sh‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -124,9 +124,9 @@ if [[ "${MODE}" =~ .*qnn.* ]]; then
124124
source "$(dirname "${BASH_SOURCE[0]}")/../../backends/qualcomm/scripts/install_qnn_sdk.sh"
125125
install_qnn
126126

127-
export EXECUTORCH_ROOT="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")/.." && pwd)"
127+
export EXECUTORCH_ROOT="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")/../.." && pwd)"
128128
export LD_LIBRARY_PATH="${QNN_SDK_ROOT}/lib/x86_64-linux-clang"
129-
export PYTHONPATH=".."
129+
export PYTHONPATH="${EXECUTORCH_ROOT}/..${PYTHONPATH:+:${PYTHONPATH}}"
130130
cp schema/program.fbs exir/_serialize/program.fbs
131131
cp schema/scalar_type.fbs exir/_serialize/scalar_type.fbs
132132
cp -f build-x86/backends/qualcomm/PyQnnManagerAdaptor.cpython-310-x86_64-linux-gnu.so backends/qualcomm/python

‎.ci/scripts/test_qnn_static_llama_eval.sh‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -15,9 +15,9 @@ source "$(dirname "${BASH_SOURCE[0]}")/utils.sh"
1515
source "$(dirname "${BASH_SOURCE[0]}")/../../backends/qualcomm/scripts/install_qnn_sdk.sh"
1616
install_qnn
1717

18-
export EXECUTORCH_ROOT="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")/.." && pwd)"
18+
export EXECUTORCH_ROOT="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")/../.." && pwd)"
1919
export LD_LIBRARY_PATH="${QNN_SDK_ROOT}/lib/x86_64-linux-clang"
20-
export PYTHONPATH=".."
20+
export PYTHONPATH="${EXECUTORCH_ROOT}/..${PYTHONPATH:+:${PYTHONPATH}}"
2121
cp schema/program.fbs exir/_serialize/program.fbs
2222
cp schema/scalar_type.fbs exir/_serialize/scalar_type.fbs
2323
cp -f build-x86/backends/qualcomm/PyQnnManagerAdaptor.cpython-310-x86_64-linux-gnu.so backends/qualcomm/python

‎.ci/scripts/test_qnn_static_llm.sh‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -20,9 +20,9 @@ fi
2020
source "$(dirname "${BASH_SOURCE[0]}")/../../backends/qualcomm/scripts/install_qnn_sdk.sh"
2121
install_qnn
2222

23-
export EXECUTORCH_ROOT="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")/.." && pwd)"
23+
export EXECUTORCH_ROOT="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")/../.." && pwd)"
2424
export LD_LIBRARY_PATH="${QNN_SDK_ROOT}/lib/x86_64-linux-clang"
25-
export PYTHONPATH=".."
25+
export PYTHONPATH="${EXECUTORCH_ROOT}/..${PYTHONPATH:+:${PYTHONPATH}}"
2626
cp schema/program.fbs exir/_serialize/program.fbs
2727
cp schema/scalar_type.fbs exir/_serialize/scalar_type.fbs
2828
cp -f build-x86/backends/qualcomm/PyQnnManagerAdaptor.cpython-310-x86_64-linux-gnu.so backends/qualcomm/python

‎.ci/scripts/unittest-linux-cmake.sh‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,14 @@ if ! python -c "import tosa_serializer" >/dev/null 2>&1; then
2020
fi
2121

2222
# NOTE: Will be removed when tosa-tools is installed via pypi
23+
if grep -q "^cmake\\.verbose = true$" "${TOSA_SERIALIZATION_DIR}/pyproject.toml"; then
24+
sed -i "s/^cmake\\.verbose = true$/build.verbose = true/" \
25+
"${TOSA_SERIALIZATION_DIR}/pyproject.toml"
26+
fi
27+
if grep -q "^dynamic = \\[\"version\"\\]$" "${TOSA_SERIALIZATION_DIR}/pyproject.toml"; then
28+
sed -i "s/^dynamic = \\[\"version\"\\]$/version = \"0.0.0\"/" \
29+
"${TOSA_SERIALIZATION_DIR}/pyproject.toml"
30+
fi
2331
python -m pip install pybind11==2.10.4
2432
CMAKE_POLICY_VERSION_MINIMUM=3.5 BUILD_PYBIND=1 \
2533
python -m pip install --no-dependencies \

‎.github/workflows/pull.yml‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -534,8 +534,8 @@ jobs:
534534
output=$(ls -la cmake-out/test/size_test)
535535
arr=($output)
536536
size=${arr[4]}
537-
# Current CI size: 48008 (gcc9-nopytorch, 2026-03-06)
538-
threshold="48500"
537+
# Current CI size: 52168 (gcc9-nopytorch, 2026-07-06)
538+
threshold="53000"
539539
if [[ "$size" -le "$threshold" ]]; then
540540
echo "Success $size <= $threshold"
541541
else

‎examples/models/llama/attention.py‎

Lines changed: 93 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
import logging
12
from abc import ABC, abstractmethod
23
from enum import Enum
34
from typing import Any, Dict, Optional, Tuple, Type, TypedDict
@@ -56,6 +57,8 @@ def forward(
5657

5758

5859
ATTENTION_REGISTRY: Dict[str, Type[Attention]] = {}
60+
_RECURRENT_GATED_DELTA_RULE_OP = None
61+
_TRIED_LOADING_RECURRENT_GATED_DELTA_RULE_OP = False
5962

6063

6164
def register_attention(name: str):
@@ -68,6 +71,38 @@ def decorator(cls: Type[Attention]):
6871
return decorator
6972

7073

74+
def _get_recurrent_gated_delta_rule_op():
75+
global _RECURRENT_GATED_DELTA_RULE_OP
76+
global _TRIED_LOADING_RECURRENT_GATED_DELTA_RULE_OP
77+
78+
if _TRIED_LOADING_RECURRENT_GATED_DELTA_RULE_OP:
79+
return _RECURRENT_GATED_DELTA_RULE_OP
80+
81+
_TRIED_LOADING_RECURRENT_GATED_DELTA_RULE_OP = True
82+
try:
83+
_RECURRENT_GATED_DELTA_RULE_OP = (
84+
torch.ops.llama.recurrent_gated_delta_rule.default
85+
)
86+
return _RECURRENT_GATED_DELTA_RULE_OP
87+
except (AttributeError, RuntimeError):
88+
pass
89+
90+
try:
91+
from executorch.extension.llm.custom_ops import custom_ops # noqa: F401
92+
except (ImportError, OSError, RuntimeError):
93+
logging.debug("Failed to import custom ops library", exc_info=True)
94+
return None
95+
96+
try:
97+
_RECURRENT_GATED_DELTA_RULE_OP = (
98+
torch.ops.llama.recurrent_gated_delta_rule.default
99+
)
100+
except (AttributeError, RuntimeError):
101+
_RECURRENT_GATED_DELTA_RULE_OP = None
102+
103+
return _RECURRENT_GATED_DELTA_RULE_OP
104+
105+
71106
class KVCache(nn.Module):
72107
def __init__(
73108
self,
@@ -762,28 +797,43 @@ def _apply_causal_conv(self, mixed_qkv: torch.Tensor) -> torch.Tensor:
762797
out = F.silu(out[:, :, -seq_len:]).to(mixed_qkv.dtype)
763798
return out.transpose(1, 2).contiguous()
764799

765-
def _recurrent_gated_delta_rule(
800+
def _gated_delta_rule_op(
766801
self,
767802
query: torch.Tensor,
768803
key: torch.Tensor,
769804
value: torch.Tensor,
770805
g: torch.Tensor,
771806
beta: torch.Tensor,
772807
) -> torch.Tensor:
773-
# query/key/value: (batch, seq_len, num_heads, head_dim)
774-
# g/beta: (batch, seq_len, num_heads)
775-
initial_dtype = query.dtype
776-
query = _l2norm(query, dim=-1, eps=1e-6)
777-
key = _l2norm(key, dim=-1, eps=1e-6)
778-
query, key, value, beta, g = [
779-
x.transpose(1, 2).contiguous().to(torch.float32)
780-
for x in (query, key, value, beta, g)
781-
]
808+
batch_size = query.shape[0]
809+
recurrent_gated_delta_rule_op = _get_recurrent_gated_delta_rule_op()
810+
if recurrent_gated_delta_rule_op is not None:
811+
return recurrent_gated_delta_rule_op(
812+
query,
813+
key,
814+
value,
815+
g,
816+
beta,
817+
self.recurrent_state[:batch_size],
818+
)
819+
return self._naive_gated_delta_rule_op(
820+
query,
821+
key,
822+
value,
823+
g,
824+
beta,
825+
)
782826

783-
batch_size, num_heads, sequence_length, k_head_dim = key.shape
827+
def _naive_gated_delta_rule_op(
828+
self,
829+
query: torch.Tensor,
830+
key: torch.Tensor,
831+
value: torch.Tensor,
832+
g: torch.Tensor,
833+
beta: torch.Tensor,
834+
) -> torch.Tensor:
835+
batch_size, num_heads, sequence_length, _ = key.shape
784836
v_head_dim = value.shape[-1]
785-
scale = 1.0 / (query.shape[-1] ** 0.5)
786-
query = query * scale
787837

788838
core_attn_out = torch.zeros(
789839
batch_size,
@@ -817,6 +867,36 @@ def _recurrent_gated_delta_rule(
817867
last_recurrent_state.to(self.recurrent_state.dtype)
818868
)
819869

870+
return core_attn_out
871+
872+
def _recurrent_gated_delta_rule(
873+
self,
874+
query: torch.Tensor,
875+
key: torch.Tensor,
876+
value: torch.Tensor,
877+
g: torch.Tensor,
878+
beta: torch.Tensor,
879+
) -> torch.Tensor:
880+
# query/key/value: (batch, seq_len, num_heads, head_dim)
881+
# g/beta: (batch, seq_len, num_heads)
882+
initial_dtype = query.dtype
883+
query = _l2norm(query, dim=-1, eps=1e-6)
884+
key = _l2norm(key, dim=-1, eps=1e-6)
885+
query, key, value, beta, g = [
886+
x.transpose(1, 2).contiguous().to(torch.float32)
887+
for x in (query, key, value, beta, g)
888+
]
889+
890+
scale = 1.0 / (query.shape[-1] ** 0.5)
891+
query = query * scale
892+
893+
core_attn_out = self._gated_delta_rule_op(
894+
query,
895+
key,
896+
value,
897+
g,
898+
beta,
899+
)
820900
return core_attn_out.transpose(1, 2).contiguous().to(initial_dtype)
821901

822902
def forward(

‎examples/models/llama/tests/test_export_llama_lib.py‎

Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,10 @@
55
# This source code is licensed under the BSD-style license found in the
66
# LICENSE file in the root directory of this source tree.
77

8+
import json
9+
import tempfile
810
import unittest
11+
from pathlib import Path
912

1013
from executorch.devtools.backend_debug import get_delegation_info
1114

@@ -25,6 +28,7 @@
2528

2629
from executorch.examples.models.llama.export_llama_lib import (
2730
_export_llama,
31+
_prepare_for_llama_export,
2832
build_args_parser,
2933
get_quantizer_and_quant_params,
3034
)
@@ -41,6 +45,39 @@
4145

4246

4347
class ExportLlamaLibTest(unittest.TestCase):
48+
def _make_tiny_qwen35_params(self) -> dict:
49+
return {
50+
"dim": 64,
51+
"hidden_dim": 128,
52+
"n_heads": 4,
53+
"head_dim": 16,
54+
"n_kv_heads": 2,
55+
"n_layers": 4,
56+
"norm_eps": 1e-6,
57+
"rope_theta": 10000000.0,
58+
"use_scaled_rope": False,
59+
"vocab_size": 256,
60+
"use_hf_rope": True,
61+
"partial_rotary_factor": 0.25,
62+
"attention_qkv_bias": False,
63+
"use_qk_norm": True,
64+
"qk_norm_before_rope": True,
65+
"attention_type": "mha",
66+
"use_q_gate": True,
67+
"rms_norm_add_unit_offset": True,
68+
"linear_conv_kernel_dim": 4,
69+
"linear_key_head_dim": 8,
70+
"linear_value_head_dim": 8,
71+
"linear_num_key_heads": 4,
72+
"linear_num_value_heads": 4,
73+
"layer_types": [
74+
"linear_attention",
75+
"full_attention",
76+
"linear_attention",
77+
"full_attention",
78+
],
79+
}
80+
4481
def test_has_expected_ops_and_op_counts(self):
4582
"""
4683
Checks the presence of unwanted expensive ops.
@@ -70,6 +107,41 @@ def test_has_expected_ops_and_op_counts(self):
70107
for op, _op_info in delegation_info.delegation_by_operator.items():
71108
self.assertTrue(op not in UNWANTED_OPS)
72109

110+
def test_tiny_qwen35_export_uses_recurrent_gated_delta_rule(self):
111+
with tempfile.TemporaryDirectory() as temp_dir:
112+
params_path = Path(temp_dir) / "tiny_qwen35.json"
113+
params_path.write_text(json.dumps(self._make_tiny_qwen35_params()))
114+
115+
parser = build_args_parser()
116+
args = parser.parse_args(
117+
[
118+
"--model",
119+
"qwen3_5_0_8b",
120+
"--params",
121+
str(params_path),
122+
"--use_kv_cache",
123+
"--disable_dynamic_shape",
124+
"--max_seq_length",
125+
"8",
126+
"--max_context_length",
127+
"8",
128+
]
129+
)
130+
131+
llm_config = LlmConfig.from_args(args)
132+
builder = _prepare_for_llama_export(llm_config).export()
133+
assert builder.pre_autograd_graph_module is not None
134+
135+
recurrent_nodes = [
136+
node
137+
for node in builder.pre_autograd_graph_module.graph.nodes
138+
if "auto_functionalized_v2" in str(node.target)
139+
and node.args
140+
and "llama.recurrent_gated_delta_rule" in str(node.args[0])
141+
]
142+
143+
self.assertEqual(len(recurrent_nodes), 2)
144+
73145
@unittest.skipUnless(HAS_ARM_BACKEND, "ARM backend not available")
74146
def test_get_quantizer_and_quant_params_returns_tosa_quantizer(self):
75147
llm_config = LlmConfig()

0 commit comments

Comments
 (0)