Gated delta net recurrence export to tract_transformers_gdn_recurrent.
HF's Qwen3.5 linear attention computes its core with a Python loop over
the sequence axis (torch_recurrent_gated_delta_rule /
torch_chunk_gated_delta_rule). Traced, that loop unrolls at the traced
sequence length and the exported graph is frozen to it (S=1 for the
decode-path trace). Reifying the rule as ONE custom op keeps the graph
S-generic: tract's GatedDeltaNetRecurrent runs the recurrence
sequentially for any S.
The exporter swaps the module-bound rule functions with
:class:GatedDeltaNetRecurrentReified (see the Qwen3.5 architecture
handler's prepare_model_for_export); the extractor below then emits
the op instead of tracing the shim's internals.
Op contract (must match tract transformers/src/ops/gdn_recurrent.rs):
query/key [b, S, hk, w] (key width == value width)
value [b, S, hv, w] with hv = G * hk (GQA groups resolved
from the shapes; value head h reads
query/key head h // G)
log_decay/beta [b, S, hv]
initial_state [b, hv, w, w]
-> output [b, S, hv, w] (query dtype), final_state [b, hv, w, w]
The op L2-normalizes q/k (eps 1e-6) and scales q by 1/sqrt(w) itself,
i.e. HF use_qk_l2norm_in_kernel=True semantics. HF repeat-interleaves
q/k to hv heads BEFORE calling the rule; the Qwen3.5 handler neutralizes
that repeat so the boundary receives hk-head q/k and the exported graph
never materializes the broadcast (the shim repeats internally, so its
eager math stays HF-exact).
Classes:
| Name | Description |
|---|---|
CausalConvUpdateExtractor |
Emit |
CausalConvUpdateReified |
Drop-in core for HF's |
GatedDeltaNetRecurrentExtractor |
Emit |
GatedDeltaNetRecurrentReified |
Drop-in for HF's (chunk|recurrent)_gated_delta_rule functions. |
Bases: Module
Drop-in core for HF's torch_causal_conv1d_update.
Depthwise causal conv1d + SiLU with a k-sample carry state. Eager
forward matches HF exactly but avoids F.conv1d (whose fake-tensor
meta kernel breaks under the offload tracing strategy): the conv is
an unfold + mul + sum. At export the call is reified as one
tract_transformers_causal_conv1d_update op.
Layout: hidden_states [b, C, S], conv_state [b, C, k],
weight [C, k] -> (out [b, C, S], final_state [b, C, k]).
The caller keeps responsibility for writing final_state back into
the HF cache (a plain copy_, handled by the aten copy handler).
Bases: Module
Drop-in for HF's (chunk|recurrent)_gated_delta_rule functions.
Eager forward reproduces torch_recurrent_gated_delta_rule
semantics exactly (the chunked variant computes the same recurrence
with different blocking, so substituting it is numerically
equivalent up to fp reordering). At export the whole call is
reified as one tract_transformers_gdn_recurrent op.