Use Accelerate for quantized SDPA prefill on Apple (#22760) - #22760
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/22760
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: ✅ No FailuresAs of commit 73be6b1 with merge base aeb4224 ( 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 D119391820. |
This PR needs a
|
Summary: Dequantize multi-row int8 Q and K into per-thread scratch on Apple and use the platform BLAS for QK as the Apple library is very performant, more so than ET's INT8 path. Decode and non-Apple platforms keep the existing int8 dot-product kernel as it beats the Eigen path. Differential Revision: D119391820
7551800 to
1341216
Compare
digantdesai
left a comment
There was a problem hiding this comment.
Review automatically exported from Phabricator review in Meta.
Summary: Replace the iterator-returning `std::min_element` / `std::max_element` implementations of `torch::executor::vec_minf` and `vec_maxf` with four independent value-reduction lanes. A shared compile-time implementation computes only the requested extrema, with a narrowly scoped Clang vectorization hint for NEON/SSE2 targets and no fast-math requirement. Add `torch::executor::vec_minmaxf(const float* x, size_t size, float* min_out, float* max_out)` to compute both extrema together. Wire the per-tensor and both serial/parallel per-token `choose_qparams` paths to the fused helper, eliminating their separate minimum and maximum scans without changing scale/zero-point calculations. Differential Revision: D119391738
Summary: Dequantize multi-row int8 Q and K into per-thread scratch on Apple and use the platform BLAS for QK as the Apple library is very performant, more so than ET's INT8 path. Decode and non-Apple platforms keep the existing int8 dot-product kernel as it beats the Eigen path. Reviewed By: digantdesai Differential Revision: D119391820
1341216 to
73be6b1
Compare
Summary:
Dequantize multi-row int8 Q and K into per-thread scratch on Apple and use the platform BLAS for QK as the Apple library is very performant, more so than ET's INT8 path. Decode and non-Apple platforms keep the existing int8 dot-product kernel as it beats the Eigen path.
Reviewed By: digantdesai
Differential Revision: D119391820