Clio AI
gated_deltanet_recompute_w_u/submission/submission.py in the repo computes the Gated DeltaNet WY-transform forward outputs (w, u) from raw inputs (k, v, beta, A, g).
Core trick
Restructures Gated DeltaNet WY-transform into one kernel on [B,T,H,D] layout, computing beta*exp(g) inside the hot path instead of on host.
Tags
Description
gated_deltanet_recompute_w_u/submission/submission.py in the repo computes the Gated DeltaNet WY-transform forward outputs (w, u) from raw inputs (k, v, beta, A, g). Functionally, it applies the per-chunk WY matrix A over sequence tiles of size 64, producing u = A @ (v * beta) and w = A @ (k * beta * exp(g)). The optimized kernel runs directly on the original [B, T, H, D] layout, fuses both output paths in one kernel, and computes the gating terms inside the kernel instead of materializing chunked or pre-scaled temporaries on the host.
Performance-wise, the final promoted version is in the low single-digit microsecond range on the tuned benchmark shapes. On the official evaluator it measured about 7.07 us mean runtime with correctness passing. The main gains came from three structural improvements: removing chunk/unchunk materialization, removing pre-scaled 4D temporaries, and moving beta * exp(g) formation into the kernel hot path. The final round also showed that the best large-shape config uses block_sizes=[64, 64], which brought the tuned kernel down from about 7.97 us to about 7.10 us.