Add parallel chunked prefill for channelwise gated delta rule (#21061)#21061
Add parallel chunked prefill for channelwise gated delta rule (#21061)#21061JakeStevens wants to merge 1 commit into
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/21061
Note: Links to docs will display an error until the docs builds have been completed. ❗ 1 Active SEVsThere are 1 currently active SEVs. If your PR is affected, please view them below: ❌ 2 New Failures, 1 Unrelated FailureAs of commit 85d15b3 with merge base 667c91b ( NEW FAILURES - The following jobs have failed:
BROKEN TRUNK - The following job failed but were present on the merge base:👉 Rebase onto the `viable/strict` branch to avoid these failures
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
|
@JakeStevens has exported this pull request. If you are a Meta employee, you can view the originating Diff in D112597348. |
This PR needs a
|
…h#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
1509cce to
85d15b3
Compare
…h#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
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