SIGN IN SIGN UP

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