[ExecuTorch][WebGPU] Route exact prefill shapes to the BK64 GEMM - #21130
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/21130
Note: Links to docs will display an error until the docs builds have been completed. ❌ 37 New Failures, 30 PendingAs of commit 1b7fdc1 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
|
SS-JIA
left a comment
There was a problem hiding this comment.
Review automatically exported from Phabricator review in Meta.
fbsource master [ghstack-poisoned]
Pull Request resolved: #21130 Llama prefill drives the ordinary quantized-linear projections at a small set of fixed batch-row counts, and the generic Steel schedule leaves throughput on the table for them. This adds a frozen BK64 Steel quantized-GEMM route for the exact accepted Llama ordinary-projection shapes at live M128, M508, and M512, selected only when the capability and dynamic-route guards all pass. M511 and other prefill sizes stay on the generic Steel schedule and M1 stays on bicol decode, so the route fails closed outside its accepted shapes. Mirrors Vulkan xplat/executorch/backends/vulkan/runtime/graph/ops/glsl/q4gsw_linear_gemm__w_4x8.glsl (the 4-bit grouped-symmetric-weight GEMM), specialized here for the BK64 tile. Key changes: - runtime/ops/quantized_linear/q4gsw_steel_bk64.wgsl (+ generated header): the BK64-tiled q4gsw prefill GEMM kernel. - QuantizedLinear.cpp, WebGPUUtils.h: exact-shape predicate, capability guard, and dynamic re-entry into and out of the BK64 route. ghstack-source-id: 411961445 @exported-using-ghexport Differential Revision: [D113171739](https://our.internmc.facebook.com/intern/diff/D113171739/)
Pull Request resolved: #21130 Llama prefill drives the ordinary quantized-linear projections at a small set of fixed batch-row counts, and the generic Steel schedule leaves throughput on the table for them. This adds a frozen BK64 Steel quantized-GEMM route for the exact accepted Llama ordinary-projection shapes at live M128, M508, and M512, selected only when the capability and dynamic-route guards all pass. M511 and other prefill sizes stay on the generic Steel schedule and M1 stays on bicol decode, so the route fails closed outside its accepted shapes. Mirrors Vulkan xplat/executorch/backends/vulkan/runtime/graph/ops/glsl/q4gsw_linear_gemm__w_4x8.glsl (the 4-bit grouped-symmetric-weight GEMM), specialized here for the BK64 tile. Key changes: - runtime/ops/quantized_linear/q4gsw_steel_bk64.wgsl (+ generated header): the BK64-tiled q4gsw prefill GEMM kernel. - QuantizedLinear.cpp, WebGPUUtils.h: exact-shape predicate, capability guard, and dynamic re-entry into and out of the BK64 route. ghstack-source-id: 411961445 @exported-using-ghexport Differential Revision: [D113171739](https://our.internmc.facebook.com/intern/diff/D113171739/)
Pull Request resolved: #21130 Llama prefill drives the ordinary quantized-linear projections at a small set of fixed batch-row counts, and the generic Steel schedule leaves throughput on the table for them. This adds a frozen BK64 Steel quantized-GEMM route for the exact accepted Llama ordinary-projection shapes at live M128, M508, and M512, selected only when the capability and dynamic-route guards all pass. M511 and other prefill sizes stay on the generic Steel schedule and M1 stays on bicol decode, so the route fails closed outside its accepted shapes. Mirrors Vulkan xplat/executorch/backends/vulkan/runtime/graph/ops/glsl/q4gsw_linear_gemm__w_4x8.glsl (the 4-bit grouped-symmetric-weight GEMM), specialized here for the BK64 tile. Key changes: - runtime/ops/quantized_linear/q4gsw_steel_bk64.wgsl (+ generated header): the BK64-tiled q4gsw prefill GEMM kernel. - QuantizedLinear.cpp, WebGPUUtils.h: exact-shape predicate, capability guard, and dynamic re-entry into and out of the BK64 route. ghstack-source-id: 411961445 @exported-using-ghexport Differential Revision: [D113171739](https://our.internmc.facebook.com/intern/diff/D113171739/)
Stack from ghstack (oldest at bottom):
Llama prefill drives the ordinary quantized-linear projections at a small set
of fixed batch-row counts, and the generic Steel schedule leaves throughput on
the table for them. This adds a frozen BK64 Steel quantized-GEMM route for the
exact accepted Llama ordinary-projection shapes at live M128, M508, and M512,
selected only when the capability and dynamic-route guards all pass. M511 and
other prefill sizes stay on the generic Steel schedule and M1 stays on bicol
decode, so the route fails closed outside its accepted shapes. Mirrors Vulkan
xplat/executorch/backends/vulkan/runtime/graph/ops/glsl/q4gsw_linear_gemm__w_4x8.glsl
(the 4-bit grouped-symmetric-weight GEMM), specialized here for the BK64 tile.
Key changes:
BK64-tiled q4gsw prefill GEMM kernel.
and dynamic re-entry into and out of the BK64 route.
@exported-using-ghexport
Differential Revision: D113171739
Differential Revision: D113171739