Shared surface for the gated-delta-net (GDN) linear-attention recurrence.
Two export paths reach the SAME tract operator and must agree on its contract, so the contract lives here once:
t2n_extra::gated_delta_scan(seeop/extras/scan_ops.py), a torch custom op whose portable lowering is atract_core_scan, and- a reified GDN module captured by a
ModuleInfoExtractor, which has no portable lowering and only ever emits the native operator.
What the two paths do NOT share is how they obtain the six operands (a
traced custom-op node against a module boundary) and the layout those
operands arrive in, so emit_native_gdn_recurrent takes a head_major
flag instead of assuming one.
Functions:
| Name | Description |
|---|---|
emit_native_gdn_recurrent |
Emit |
gated_delta_fake |
Meta kernel for |
gated_delta_reference |
Pure-torch gated-delta recurrence over the time axis (axis 2). |
l2norm |
L2-normalize the last axis, with HF's |
native_gdn_min_version |
First tract release carrying the fused operator. |
native_gdn_reject_reason |
Why these operands can NOT become tract's fused operator. |
emit_native_gdn_recurrent(g, op_helper, operands: Sequence['NTensor'], outputs: Sequence['NTensor'], head_major: bool) -> T.List[str]
Emit tract_transformers_gdn_recurrent(q, k, v, log_decay, beta, s0).
tract's operand layout puts the sequence axis at 1 and the head axis at 2
([B, S, H, W] for q/k/v, [B, S, H] for log-decay and beta), while the
state keeps [B, H, W, W]. With head_major the per-step operands and
the first output arrive as [B, H, S, ...] instead, so transpose them on
the way in and take one back on the way out. At S == 1 those transposes
are pure shape changes, but they are what makes the emitted graph match
the operator's documented layout rather than its flat-index reading of a
single step.
outputs is (output, final_state), already created by the caller in ITS
own layout, since the two paths build output tensors differently.
gated_delta_fake(q: Tensor, k: Tensor, v: Tensor, g: Tensor, beta: Tensor, s0: Tensor) -> T.Tuple[torch.Tensor, torch.Tensor]
Meta kernel for t2n_extra::gated_delta_scan.
y takes the PROMOTED dtype, not q's: the recurrence accumulates against
the state, which stays f32 even for an f16 model, so the eager op returns
f32 there. A fake that claimed f16 would disagree with the real op under
torch.export / opcheck / torch.compile.
gated_delta_reference(q: Tensor, k: Tensor, v: Tensor, g: Tensor, beta: Tensor, s0: Tensor) -> T.Tuple[torch.Tensor, torch.Tensor]
Pure-torch gated-delta recurrence over the time axis (axis 2).
The single reference for the whole repo: the t2n_extra custom op's eager
body, the reified module's eager forward and the tests all call this, so a
divergence cannot hide in one of the copies. Matches HF's
torch_recurrent_gated_delta_rule for pre-normalized q/k.
L2-normalize the last axis, with HF's use_qk_l2norm_in_kernel eps.
Shared so the pre-normalization applied OUTSIDE the op cannot drift from the eps tract's fused operator uses internally.
First tract release carrying the fused operator.
native_gdn_reject_reason(inference_target, operands: Sequence['NTensor'], head_major: bool) -> T.Optional[str]
Why these operands can NOT become tract's fused operator.
tract_transformers_gdn_recurrent is a fused SINGLE decode step (it also
folds the q/k l2-norm and the 1 / sqrt(head_dim) output scale), with
hard constraints checked by tract at load time. Returns None when the
traced tensors satisfy all of them, else a short reason for logging
before falling back.
operands is (query, key, value, log_decay, beta, state); head_major
says whether the per-step tensors carry (B, H, S, ...) rather than
tract's own (B, S, H, ...).