Skip to content

Commit 8134bb2

Browse files
authored
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
1 parent 179c4ee commit 8134bb2

2 files changed

Lines changed: 481 additions & 52 deletions

File tree

0 commit comments

Comments
 (0)