Handlers for t2n_extra::* scan-shaped ops.
Currently provides ssm_scan for Mamba's selective state-space scan.
The handler emits a mamba_ssm_scan NNEF fragment call which wraps a
tract_core_scan over a per-step mamba_ssm_step body. Tract's pulse
declutter compiles the scan into a streaming graph, so the prefill
cost is one tract call instead of one per token.
Functions:
| Name | Description |
|---|---|
gated_delta_scan |
Emit a |
ssm_scan |
Emit a |
ssm_scan_y |
Pulse-friendly variant of |
Emit a gated_delta_scan fragment call (Qwen3.5 gated-delta-net).
Torch side
t2n_extra::gated_delta_scan(q, k, v, g, beta, s0) -> (y, s_final)
with (time axis at 2):
q, k (B, H, T, hk) v (B, H, T, hv)
g, beta (B, H, T) s0 (B, H, hk, hv)
The scan iterates axis 0, so pre-transpose the per-step inputs to put
time first (s0 is the state, left as-is), then the fragment wraps
tract_core_scan. Outputs: y (B, H, T, hv), s_final (B, H, hk, hv).
From tract 0.23.5 a single-step graph instead emits tract's fused
tract_transformers_gdn_recurrent (CPU/CUDA/Metal kernels) when it fits
the operator's constraints, see _native_gdn_reject_reason. That operator
additionally folds the q/k l2-norm and a 1 / sqrt(head_dim) output
scale, so it assumes the Qwen3.5 convention this op is written for: q
passed as l2norm(q) / sqrt(head_k_dim) and k as l2norm(k).
Re-normalizing an already-normalized q/k is a no-op, and the internal
scale (tract reads head_dim off the same axis, so it always matches
ours) then restores the one normalizing q stripped, so both lowerings agree
(to ~6e-5 relative, from the operator's 1e-6 norm epsilon; under f16
resolution, and covered by check_io). Feeding raw, un-normalized q/k
would NOT be equivalent.
Emit a mamba_ssm_scan fragment call.
Signature on the torch side
t2n_extra::ssm_scan(discrete_A, deltaB_u, C, h_init) -> (scan_outputs, h_final)
The fragment scans along axis 0 of its inputs (matches the GRU/LSTM convention). The handler pre-transposes the SSM tensors so the time axis lands at position 0 before the scan:
discrete_A (B, D, T, N) -> (T, B, D, N)
deltaB_u (B, D, T, N) -> (T, B, D, N)
C (B, T, N) -> (T, B, N)
h_init (B, D, N) -- unchanged (state)
After the scan
scan_y (T, B, D) -> (B, D, T) to match scan_outputs's
PyTorch shape (stack on last axis).
h_final (B, D, N)