Uh oh!
There was an error while loading. Please reload this page.
[ExecuTorch][WebGPU] Add Llama K16 online causal attention - #21132
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/21132
Note: Links to docs will display an error until the docs builds have been completed. ❌ 48 New Failures, 28 PendingAs of commit 2a2acde with merge base 28a7fac ( NEW FAILURES - The following jobs have failed:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
This PR needs a |
fbsource master [ghstack-poisoned]
Uh oh!
There was an error while loading. Please reload this page.
Pull Request resolved: #21132 Materializing the full attention-weight matrix for Llama prefill is memory- and bandwidth-heavy and does not scale to longer sequences. This adds a single-pass online-softmax K16 causal-attention kernel for the exact Llama Hq32/Hkv8/G4/D64 geometry with fp16 KV, guarded by explicit adapter limits, so accepted prefill shapes through S512 never materialize attention weights. S1 keeps the existing FlashDecoding path, and unsupported geometry, storage, or capabilities fall back to the materialized implementation. No Vulkan analogue (WebGPU-specific): the Vulkan backend has only a materialized attention (compute the weights, then a separate multi-pass softmax), with no online-softmax kernel. Key changes: - runtime/ops/sdpa/streaming_attention_k16_causal_bound.wgsl (+ generated header): the online-softmax single-pass causal kernel. - Sdpa.cpp: exact Llama geometry and capability guards, S512 transition handling, and fallback to the materialized route. ghstack-source-id: 411961446 @exported-using-ghexport Differential Revision: [D113171738](https://our.internmc.facebook.com/intern/diff/D113171738/)
Pull Request resolved: #21132 Materializing the full attention-weight matrix for Llama prefill is memory- and bandwidth-heavy and does not scale to longer sequences. This adds a single-pass online-softmax K16 causal-attention kernel for the exact Llama Hq32/Hkv8/G4/D64 geometry with fp16 KV, guarded by explicit adapter limits, so accepted prefill shapes through S512 never materialize attention weights. S1 keeps the existing FlashDecoding path, and unsupported geometry, storage, or capabilities fall back to the materialized implementation. No Vulkan analogue (WebGPU-specific): the Vulkan backend has only a materialized attention (compute the weights, then a separate multi-pass softmax), with no online-softmax kernel. Key changes: - runtime/ops/sdpa/streaming_attention_k16_causal_bound.wgsl (+ generated header): the online-softmax single-pass causal kernel. - Sdpa.cpp: exact Llama geometry and capability guards, S512 transition handling, and fallback to the materialized route. ghstack-source-id: 411961446 @exported-using-ghexport Differential Revision: [D113171738](https://our.internmc.facebook.com/intern/diff/D113171738/)
Pull Request resolved: #21132 Materializing the full attention-weight matrix for Llama prefill is memory- and bandwidth-heavy and does not scale to longer sequences. This adds a single-pass online-softmax K16 causal-attention kernel for the exact Llama Hq32/Hkv8/G4/D64 geometry with fp16 KV, guarded by explicit adapter limits, so accepted prefill shapes through S512 never materialize attention weights. S1 keeps the existing FlashDecoding path, and unsupported geometry, storage, or capabilities fall back to the materialized implementation. No Vulkan analogue (WebGPU-specific): the Vulkan backend has only a materialized attention (compute the weights, then a separate multi-pass softmax), with no online-softmax kernel. Key changes: - runtime/ops/sdpa/streaming_attention_k16_causal_bound.wgsl (+ generated header): the online-softmax single-pass causal kernel. - Sdpa.cpp: exact Llama geometry and capability guards, S512 transition handling, and fallback to the materialized route. ghstack-source-id: 411961446 @exported-using-ghexport Differential Revision: [D113171738](https://our.internmc.facebook.com/intern/diff/D113171738/)
Pull Request resolved: #21132 Materializing the full attention-weight matrix for Llama prefill is memory- and bandwidth-heavy and does not scale to longer sequences. This adds a single-pass online-softmax K16 causal-attention kernel for the exact Llama Hq32/Hkv8/G4/D64 geometry with fp16 KV, guarded by explicit adapter limits, so accepted prefill shapes through S512 never materialize attention weights. S1 keeps the existing FlashDecoding path, and unsupported geometry, storage, or capabilities fall back to the materialized implementation. No Vulkan analogue (WebGPU-specific): the Vulkan backend has only a materialized attention (compute the weights, then a separate multi-pass softmax), with no online-softmax kernel. Key changes: - runtime/ops/sdpa/streaming_attention_k16_causal_bound.wgsl (+ generated header): the online-softmax single-pass causal kernel. - Sdpa.cpp: exact Llama geometry and capability guards, S512 transition handling, and fallback to the materialized route. ghstack-source-id: 411961446 @exported-using-ghexport Differential Revision: [D113171738](https://our.internmc.facebook.com/intern/diff/D113171738/)
Pull Request resolved: #21132 Materializing the full attention-weight matrix for Llama prefill is memory- and bandwidth-heavy and does not scale to longer sequences. This adds a single-pass online-softmax K16 causal-attention kernel for the exact Llama Hq32/Hkv8/G4/D64 geometry with fp16 KV, guarded by explicit adapter limits, so accepted prefill shapes through S512 never materialize attention weights. S1 keeps the existing FlashDecoding path, and unsupported geometry, storage, or capabilities fall back to the materialized implementation. No Vulkan analogue (WebGPU-specific): the Vulkan backend has only a materialized attention (compute the weights, then a separate multi-pass softmax), with no online-softmax kernel. Key changes: - runtime/ops/sdpa/streaming_attention_k16_causal_bound.wgsl (+ generated header): the online-softmax single-pass causal kernel. - Sdpa.cpp: exact Llama geometry and capability guards, S512 transition handling, and fallback to the materialized route. ghstack-source-id: 411961446 @exported-using-ghexport Differential Revision: [D113171738](https://our.internmc.facebook.com/intern/diff/D113171738/)
Stack from ghstack (oldest at bottom):
Materializing the full attention-weight matrix for Llama prefill is memory- and
bandwidth-heavy and does not scale to longer sequences. This adds a single-pass
online-softmax K16 causal-attention kernel for the exact Llama Hq32/Hkv8/G4/D64
geometry with fp16 KV, guarded by explicit adapter limits, so accepted prefill
shapes through S512 never materialize attention weights. S1 keeps the existing
FlashDecoding path, and unsupported geometry, storage, or capabilities fall back
to the materialized implementation. No Vulkan analogue (WebGPU-specific): the
Vulkan backend has only a materialized attention (compute the weights, then a
separate multi-pass softmax), with no online-softmax kernel.
Key changes:
header): the online-softmax single-pass causal kernel.
handling, and fallback to the materialized route.
@exported-using-ghexport
Differential Revision: D113171738
Differential Revision: D113171738