Add parallel chunked prefill for channelwise gated delta rule (#21061) (#21061)
Summary: Route the channelwise gated delta rule by sequence length: T == 1 keeps the two-pass token recurrence for autoregressive decode, while T != 1 uses a chunkwise WY/UT formulation for prefill. The chunked path computes per-channel log-decay prefixes, causal query/key terms, the beta-folded triangular transform, WY pseudo-keys and pseudo-values, and inter-chunk state carry. It handles a ragged final chunk without a separate tail implementation. Parallelize independent (batch, head) work across the ExecuTorch threadpool. Each worker receives a disjoint slice of one temporary scratch arena, avoiding shared mutable buffers while amortizing allocation across chunks. Reviewed By: billmguo Differential Revision: D112597348
J
Jacob Stevens committed
8134bb23df6effa1a4976760a9a62121bb77a1c8
Parent: 179c4ee
Committed by GitHub <noreply@github.com>
on 7/23/2026, 8:08:19 PM