[ExecuTorch][WebGPU] Packed-word-dequant f16 steel q4gsw prefill GEMM#20798
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/20798
Note: Links to docs will display an error until the docs builds have been completed. ❌ 3 New Failures, 12 Pending, 1 Unrelated FailureAs of commit 06fe79f with merge base f4b01a8 ( 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. |
This PR needs a
|
SS-JIA
left a comment
There was a problem hiding this comment.
Review automatically exported from Phabricator review in Meta.
Pull Request resolved: #20798 **+24-28% microbench GFLOP/s (+11-16% end-to-end prefill tok/s) on the f16 `steel` q4gsw prefill GEMM, Apple M4 Pro / Chrome Canary — bit-exact.** **Problem:** the f16 `steel` prefill GEMM (the `half` variant) dequantizes weights per-nibble — each of the 256 threads extracts ONE 4-bit weight from a full `u32` word, so every packed word is re-loaded ~8x and the per-column scale is re-read ~16x across the `BK=16` tile. That redundant global traffic dominates the kernel on Apple. **Solution:** a packed-word-dequant staging variant, bit-exact to the `half` kernel — identical `As` staging, compute loop, epilogue, and `64x64` tile / 256-thread / `BK=16` geometry; only the weight-dequant staging differs. Before (`half`): each thread extracts 1 nibble from a re-loaded word; scale re-read per nibble. After (`pwdq`): threads `[0,BN)` each stage one full `BK`-column — load the column's two `u32` words ONCE, unpack all 16 nibbles, read the per-column scale ONCE. **Implementation:** - New `PWDQ` fork in the shared `q4gsw_linear_gemm_steel.wgsl` template (a `shader_variants` entry in `q4gsw_linear_gemm_steel.yaml` generates `q4gsw_linear_gemm_steel_half_pwdq_wgsl.h`) — no standalone shader file. Selected at runtime when the device negotiated shader-f16 (`ctx->shader_f16_supported`) AND `group_size % BK == 0` (the hoisted scale must be constant across the `BK` tile), else the per-nibble `half` kernel. - No new build option, dispatch, or bindings — same tile/count as `steel`; only the generated shader variant differs. **Constraints:** requires `K % BK == 0` (the steel route already guarantees it, making `K_packed = K/2` a multiple of 8 so every column is `u32`-aligned) and `group_size % BK == 0` (all real q4 group sizes: 32/64/128). f16-multiply / f32-accumulate, so numerically identical to the `half` kernel. Co-authored-with: Claude Code. ghstack-source-id: 401564878 @exported-using-ghexport Differential Revision: [D111163510](https://our.internmc.facebook.com/intern/diff/D111163510/)
Pull Request resolved: #20798 **+24-28% microbench GFLOP/s (+11-16% end-to-end prefill tok/s) on the f16 `steel` q4gsw prefill GEMM, Apple M4 Pro / Chrome Canary — bit-exact.** **Problem:** the f16 `steel` prefill GEMM (the `half` variant) dequantizes weights per-nibble — each of the 256 threads extracts ONE 4-bit weight from a full `u32` word, so every packed word is re-loaded ~8x and the per-column scale is re-read ~16x across the `BK=16` tile. That redundant global traffic dominates the kernel on Apple. **Solution:** a packed-word-dequant staging variant, bit-exact to the `half` kernel — identical `As` staging, compute loop, epilogue, and `64x64` tile / 256-thread / `BK=16` geometry; only the weight-dequant staging differs. Before (`half`): each thread extracts 1 nibble from a re-loaded word; scale re-read per nibble. After (`pwdq`): threads `[0,BN)` each stage one full `BK`-column — load the column's two `u32` words ONCE, unpack all 16 nibbles, read the per-column scale ONCE. **Implementation:** - New `PWDQ` fork in the shared `q4gsw_linear_gemm_steel.wgsl` template (a `shader_variants` entry in `q4gsw_linear_gemm_steel.yaml` generates `q4gsw_linear_gemm_steel_half_pwdq_wgsl.h`) — no standalone shader file. Selected at runtime when the device negotiated shader-f16 (`ctx->shader_f16_supported`) AND `group_size % BK == 0` (the hoisted scale must be constant across the `BK` tile), else the per-nibble `half` kernel. - No new build option, dispatch, or bindings — same tile/count as `steel`; only the generated shader variant differs. **Constraints:** requires `K % BK == 0` (the steel route already guarantees it, making `K_packed = K/2` a multiple of 8 so every column is `u32`-aligned) and `group_size % BK == 0` (all real q4 group sizes: 32/64/128). f16-multiply / f32-accumulate, so numerically identical to the `half` kernel. Co-authored-with: Claude Code. ghstack-source-id: 401564878 @exported-using-ghexport Differential Revision: [D111163510](https://our.internmc.facebook.com/intern/diff/D111163510/)
Pull Request resolved: #20798 **+24-28% microbench GFLOP/s (+11-16% end-to-end prefill tok/s) on the f16 `steel` q4gsw prefill GEMM, Apple M4 Pro / Chrome Canary — bit-exact.** **Problem:** the f16 `steel` prefill GEMM (the `half` variant) dequantizes weights per-nibble — each of the 256 threads extracts ONE 4-bit weight from a full `u32` word, so every packed word is re-loaded ~8x and the per-column scale is re-read ~16x across the `BK=16` tile. That redundant global traffic dominates the kernel on Apple. **Solution:** a packed-word-dequant staging variant, bit-exact to the `half` kernel — identical `As` staging, compute loop, epilogue, and `64x64` tile / 256-thread / `BK=16` geometry; only the weight-dequant staging differs. Before (`half`): each thread extracts 1 nibble from a re-loaded word; scale re-read per nibble. After (`pwdq`): threads `[0,BN)` each stage one full `BK`-column — load the column's two `u32` words ONCE, unpack all 16 nibbles, read the per-column scale ONCE. **Implementation:** - New `PWDQ` fork in the shared `q4gsw_linear_gemm_steel.wgsl` template (a `shader_variants` entry in `q4gsw_linear_gemm_steel.yaml` generates `q4gsw_linear_gemm_steel_half_pwdq_wgsl.h`) — no standalone shader file. Selected at runtime when the device negotiated shader-f16 (`ctx->shader_f16_supported`) AND `group_size % BK == 0` (the hoisted scale must be constant across the `BK` tile), else the per-nibble `half` kernel. - No new build option, dispatch, or bindings — same tile/count as `steel`; only the generated shader variant differs. **Constraints:** requires `K % BK == 0` (the steel route already guarantees it, making `K_packed = K/2` a multiple of 8 so every column is `u32`-aligned) and `group_size % BK == 0` (all real q4 group sizes: 32/64/128). f16-multiply / f32-accumulate, so numerically identical to the `half` kernel. Co-authored-with: Claude Code. ghstack-source-id: 401564878 @exported-using-ghexport Differential Revision: [D111163510](https://our.internmc.facebook.com/intern/diff/D111163510/)
Pull Request resolved: #20798 **+24-28% microbench GFLOP/s (+11-16% end-to-end prefill tok/s) on the f16 `steel` q4gsw prefill GEMM, Apple M4 Pro / Chrome Canary — bit-exact.** **Problem:** the f16 `steel` prefill GEMM (the `half` variant) dequantizes weights per-nibble — each of the 256 threads extracts ONE 4-bit weight from a full `u32` word, so every packed word is re-loaded ~8x and the per-column scale is re-read ~16x across the `BK=16` tile. That redundant global traffic dominates the kernel on Apple. **Solution:** a packed-word-dequant staging variant, bit-exact to the `half` kernel — identical `As` staging, compute loop, epilogue, and `64x64` tile / 256-thread / `BK=16` geometry; only the weight-dequant staging differs. Before (`half`): each thread extracts 1 nibble from a re-loaded word; scale re-read per nibble. After (`pwdq`): threads `[0,BN)` each stage one full `BK`-column — load the column's two `u32` words ONCE, unpack all 16 nibbles, read the per-column scale ONCE. **Implementation:** - New `PWDQ` fork in the shared `q4gsw_linear_gemm_steel.wgsl` template (a `shader_variants` entry in `q4gsw_linear_gemm_steel.yaml` generates `q4gsw_linear_gemm_steel_half_pwdq_wgsl.h`) — no standalone shader file. Selected at runtime when the device negotiated shader-f16 (`ctx->shader_f16_supported`) AND `group_size % BK == 0` (the hoisted scale must be constant across the `BK` tile), else the per-nibble `half` kernel. - No new build option, dispatch, or bindings — same tile/count as `steel`; only the generated shader variant differs. **Constraints:** requires `K % BK == 0` (the steel route already guarantees it, making `K_packed = K/2` a multiple of 8 so every column is `u32`-aligned) and `group_size % BK == 0` (all real q4 group sizes: 32/64/128). f16-multiply / f32-accumulate, so numerically identical to the `half` kernel. Co-authored-with: Claude Code. ghstack-source-id: 401564878 @exported-using-ghexport Differential Revision: [D111163510](https://our.internmc.facebook.com/intern/diff/D111163510/)
Pull Request resolved: #20798 **+24-28% microbench GFLOP/s (+11-16% end-to-end prefill tok/s) on the f16 `steel` q4gsw prefill GEMM, Apple M4 Pro / Chrome Canary — bit-exact.** **Problem:** the f16 `steel` prefill GEMM (the `half` variant) dequantizes weights per-nibble — each of the 256 threads extracts ONE 4-bit weight from a full `u32` word, so every packed word is re-loaded ~8x and the per-column scale is re-read ~16x across the `BK=16` tile. That redundant global traffic dominates the kernel on Apple. **Solution:** a packed-word-dequant staging variant, bit-exact to the `half` kernel — identical `As` staging, compute loop, epilogue, and `64x64` tile / 256-thread / `BK=16` geometry; only the weight-dequant staging differs. Before (`half`): each thread extracts 1 nibble from a re-loaded word; scale re-read per nibble. After (`pwdq`): threads `[0,BN)` each stage one full `BK`-column — load the column's two `u32` words ONCE, unpack all 16 nibbles, read the per-column scale ONCE. **Implementation:** - New `PWDQ` fork in the shared `q4gsw_linear_gemm_steel.wgsl` template (a `shader_variants` entry in `q4gsw_linear_gemm_steel.yaml` generates `q4gsw_linear_gemm_steel_half_pwdq_wgsl.h`) — no standalone shader file. Selected at runtime when the device negotiated shader-f16 (`ctx->shader_f16_supported`) AND `group_size % BK == 0` (the hoisted scale must be constant across the `BK` tile), else the per-nibble `half` kernel. - No new build option, dispatch, or bindings — same tile/count as `steel`; only the generated shader variant differs. **Constraints:** requires `K % BK == 0` (the steel route already guarantees it, making `K_packed = K/2` a multiple of 8 so every column is `u32`-aligned) and `group_size % BK == 0` (all real q4 group sizes: 32/64/128). f16-multiply / f32-accumulate, so numerically identical to the `half` kernel. Co-authored-with: Claude Code. ghstack-source-id: 401564878 @exported-using-ghexport Differential Revision: [D111163510](https://our.internmc.facebook.com/intern/diff/D111163510/)
Pull Request resolved: #20798 **+24-28% microbench GFLOP/s (+11-16% end-to-end prefill tok/s) on the f16 `steel` q4gsw prefill GEMM, Apple M4 Pro / Chrome Canary — bit-exact.** **Problem:** the f16 `steel` prefill GEMM (the `half` variant) dequantizes weights per-nibble — each of the 256 threads extracts ONE 4-bit weight from a full `u32` word, so every packed word is re-loaded ~8x and the per-column scale is re-read ~16x across the `BK=16` tile. That redundant global traffic dominates the kernel on Apple. **Solution:** a packed-word-dequant staging variant, bit-exact to the `half` kernel — identical `As` staging, compute loop, epilogue, and `64x64` tile / 256-thread / `BK=16` geometry; only the weight-dequant staging differs. Before (`half`): each thread extracts 1 nibble from a re-loaded word; scale re-read per nibble. After (`pwdq`): threads `[0,BN)` each stage one full `BK`-column — load the column's two `u32` words ONCE, unpack all 16 nibbles, read the per-column scale ONCE. **Implementation:** - New `PWDQ` fork in the shared `q4gsw_linear_gemm_steel.wgsl` template (a `shader_variants` entry in `q4gsw_linear_gemm_steel.yaml` generates `q4gsw_linear_gemm_steel_half_pwdq_wgsl.h`) — no standalone shader file. Selected at runtime when the device negotiated shader-f16 (`ctx->shader_f16_supported`) AND `group_size % BK == 0` (the hoisted scale must be constant across the `BK` tile), else the per-nibble `half` kernel. - No new build option, dispatch, or bindings — same tile/count as `steel`; only the generated shader variant differs. **Constraints:** requires `K % BK == 0` (the steel route already guarantees it, making `K_packed = K/2` a multiple of 8 so every column is `u32`-aligned) and `group_size % BK == 0` (all real q4 group sizes: 32/64/128). f16-multiply / f32-accumulate, so numerically identical to the `half` kernel. Co-authored-with: Claude Code. ghstack-source-id: 401564878 @exported-using-ghexport Differential Revision: [D111163510](https://our.internmc.facebook.com/intern/diff/D111163510/)
Pull Request resolved: #20798 **+24-28% microbench GFLOP/s (+11-16% end-to-end prefill tok/s) on the f16 `steel` q4gsw prefill GEMM, Apple M4 Pro / Chrome Canary — bit-exact.** **Problem:** the f16 `steel` prefill GEMM (the `half` variant) dequantizes weights per-nibble — each of the 256 threads extracts ONE 4-bit weight from a full `u32` word, so every packed word is re-loaded ~8x and the per-column scale is re-read ~16x across the `BK=16` tile. That redundant global traffic dominates the kernel on Apple. **Solution:** a packed-word-dequant staging variant, bit-exact to the `half` kernel — identical `As` staging, compute loop, epilogue, and `64x64` tile / 256-thread / `BK=16` geometry; only the weight-dequant staging differs. Before (`half`): each thread extracts 1 nibble from a re-loaded word; scale re-read per nibble. After (`pwdq`): threads `[0,BN)` each stage one full `BK`-column — load the column's two `u32` words ONCE, unpack all 16 nibbles, read the per-column scale ONCE. **Implementation:** - New `PWDQ` fork in the shared `q4gsw_linear_gemm_steel.wgsl` template (a `shader_variants` entry in `q4gsw_linear_gemm_steel.yaml` generates `q4gsw_linear_gemm_steel_half_pwdq_wgsl.h`) — no standalone shader file. Selected at runtime when the device negotiated shader-f16 (`ctx->shader_f16_supported`) AND `group_size % BK == 0` (the hoisted scale must be constant across the `BK` tile), else the per-nibble `half` kernel. - No new build option, dispatch, or bindings — same tile/count as `steel`; only the generated shader variant differs. **Constraints:** requires `K % BK == 0` (the steel route already guarantees it, making `K_packed = K/2` a multiple of 8 so every column is `u32`-aligned) and `group_size % BK == 0` (all real q4 group sizes: 32/64/128). f16-multiply / f32-accumulate, so numerically identical to the `half` kernel. Co-authored-with: Claude Code. ghstack-source-id: 401564878 @exported-using-ghexport Differential Revision: [D111163510](https://our.internmc.facebook.com/intern/diff/D111163510/)
Pull Request resolved: #20798 **+24-28% microbench GFLOP/s (+11-16% end-to-end prefill tok/s) on the f16 `steel` q4gsw prefill GEMM, Apple M4 Pro / Chrome Canary — bit-exact.** **Problem:** the f16 `steel` prefill GEMM (the `half` variant) dequantizes weights per-nibble — each of the 256 threads extracts ONE 4-bit weight from a full `u32` word, so every packed word is re-loaded ~8x and the per-column scale is re-read ~16x across the `BK=16` tile. That redundant global traffic dominates the kernel on Apple. **Solution:** a packed-word-dequant staging variant, bit-exact to the `half` kernel — identical `As` staging, compute loop, epilogue, and `64x64` tile / 256-thread / `BK=16` geometry; only the weight-dequant staging differs. Before (`half`): each thread extracts 1 nibble from a re-loaded word; scale re-read per nibble. After (`pwdq`): threads `[0,BN)` each stage one full `BK`-column — load the column's two `u32` words ONCE, unpack all 16 nibbles, read the per-column scale ONCE. **Implementation:** - New `PWDQ` fork in the shared `q4gsw_linear_gemm_steel.wgsl` template (a `shader_variants` entry in `q4gsw_linear_gemm_steel.yaml` generates `q4gsw_linear_gemm_steel_half_pwdq_wgsl.h`) — no standalone shader file. Selected at runtime when the device negotiated shader-f16 (`ctx->shader_f16_supported`) AND `group_size % BK == 0` (the hoisted scale must be constant across the `BK` tile), else the per-nibble `half` kernel. - No new build option, dispatch, or bindings — same tile/count as `steel`; only the generated shader variant differs. **Constraints:** requires `K % BK == 0` (the steel route already guarantees it, making `K_packed = K/2` a multiple of 8 so every column is `u32`-aligned) and `group_size % BK == 0` (all real q4 group sizes: 32/64/128). f16-multiply / f32-accumulate, so numerically identical to the `half` kernel. Co-authored-with: Claude Code. ghstack-source-id: 401564878 @exported-using-ghexport Differential Revision: [D111163510](https://our.internmc.facebook.com/intern/diff/D111163510/)
Pull Request resolved: #20798 **+24-28% microbench GFLOP/s (+11-16% end-to-end prefill tok/s) on the f16 `steel` q4gsw prefill GEMM, Apple M4 Pro / Chrome Canary — bit-exact.** **Problem:** the f16 `steel` prefill GEMM (the `half` variant) dequantizes weights per-nibble — each of the 256 threads extracts ONE 4-bit weight from a full `u32` word, so every packed word is re-loaded ~8x and the per-column scale is re-read ~16x across the `BK=16` tile. That redundant global traffic dominates the kernel on Apple. **Solution:** a packed-word-dequant staging variant, bit-exact to the `half` kernel — identical `As` staging, compute loop, epilogue, and `64x64` tile / 256-thread / `BK=16` geometry; only the weight-dequant staging differs. Before (`half`): each thread extracts 1 nibble from a re-loaded word; scale re-read per nibble. After (`pwdq`): threads `[0,BN)` each stage one full `BK`-column — load the column's two `u32` words ONCE, unpack all 16 nibbles, read the per-column scale ONCE. **Implementation:** - New `PWDQ` fork in the shared `q4gsw_linear_gemm_steel.wgsl` template (a `shader_variants` entry in `q4gsw_linear_gemm_steel.yaml` generates `q4gsw_linear_gemm_steel_half_pwdq_wgsl.h`) — no standalone shader file. Selected at runtime when the device negotiated shader-f16 (`ctx->shader_f16_supported`) AND `group_size % BK == 0` (the hoisted scale must be constant across the `BK` tile), else the per-nibble `half` kernel. - No new build option, dispatch, or bindings — same tile/count as `steel`; only the generated shader variant differs. **Constraints:** requires `K % BK == 0` (the steel route already guarantees it, making `K_packed = K/2` a multiple of 8 so every column is `u32`-aligned) and `group_size % BK == 0` (all real q4 group sizes: 32/64/128). f16-multiply / f32-accumulate, so numerically identical to the `half` kernel. Co-authored-with: Claude Code. ghstack-source-id: 401564878 @exported-using-ghexport Differential Revision: [D111163510](https://our.internmc.facebook.com/intern/diff/D111163510/)
Pull Request resolved: #20798 **+24-28% microbench GFLOP/s (+11-16% end-to-end prefill tok/s) on the f16 `steel` q4gsw prefill GEMM, Apple M4 Pro / Chrome Canary — bit-exact.** **Problem:** the f16 `steel` prefill GEMM (the `half` variant) dequantizes weights per-nibble — each of the 256 threads extracts ONE 4-bit weight from a full `u32` word, so every packed word is re-loaded ~8x and the per-column scale is re-read ~16x across the `BK=16` tile. That redundant global traffic dominates the kernel on Apple. **Solution:** a packed-word-dequant staging variant, bit-exact to the `half` kernel — identical `As` staging, compute loop, epilogue, and `64x64` tile / 256-thread / `BK=16` geometry; only the weight-dequant staging differs. Before (`half`): each thread extracts 1 nibble from a re-loaded word; scale re-read per nibble. After (`pwdq`): threads `[0,BN)` each stage one full `BK`-column — load the column's two `u32` words ONCE, unpack all 16 nibbles, read the per-column scale ONCE. **Implementation:** - New `PWDQ` fork in the shared `q4gsw_linear_gemm_steel.wgsl` template (a `shader_variants` entry in `q4gsw_linear_gemm_steel.yaml` generates `q4gsw_linear_gemm_steel_half_pwdq_wgsl.h`) — no standalone shader file. Selected at runtime when the device negotiated shader-f16 (`ctx->shader_f16_supported`) AND `group_size % BK == 0` (the hoisted scale must be constant across the `BK` tile), else the per-nibble `half` kernel. - No new build option, dispatch, or bindings — same tile/count as `steel`; only the generated shader variant differs. **Constraints:** requires `K % BK == 0` (the steel route already guarantees it, making `K_packed = K/2` a multiple of 8 so every column is `u32`-aligned) and `group_size % BK == 0` (all real q4 group sizes: 32/64/128). f16-multiply / f32-accumulate, so numerically identical to the `half` kernel. Co-authored-with: Claude Code. ghstack-source-id: 401564878 @exported-using-ghexport Differential Revision: [D111163510](https://our.internmc.facebook.com/intern/diff/D111163510/)
Stack from ghstack (oldest at bottom):
+24-28% microbench GFLOP/s (+11-16% end-to-end prefill tok/s) on the f16
steelq4gsw prefill GEMM, Apple M4 Pro / Chrome Canary — bit-exact.Problem: the f16
steelprefill GEMM (thehalfvariant) dequantizes weights per-nibble — each of the 256 threads extracts ONE 4-bit weight from a fullu32word, so every packed word is re-loaded ~8x and the per-column scale is re-read ~16x across theBK=16tile. That redundant global traffic dominates the kernel on Apple.Solution: a packed-word-dequant staging variant, bit-exact to the
halfkernel — identicalAsstaging, compute loop, epilogue, and64x64tile / 256-thread /BK=16geometry; only the weight-dequant staging differs.Before (
half): each thread extracts 1 nibble from a re-loaded word; scale re-read per nibble.After (
pwdq): threads[0,BN)each stage one fullBK-column — load the column's twou32words ONCE, unpack all 16 nibbles, read the per-column scale ONCE.Implementation:
PWDQfork in the sharedq4gsw_linear_gemm_steel.wgsltemplate (ashader_variantsentry inq4gsw_linear_gemm_steel.yamlgeneratesq4gsw_linear_gemm_steel_half_pwdq_wgsl.h) — no standalone shader file. Selected at runtime when the device negotiated shader-f16 (ctx->shader_f16_supported) ANDgroup_size % BK == 0(the hoisted scale must be constant across theBKtile), else the per-nibblehalfkernel.steel; only the generated shader variant differs.Constraints: requires
K % BK == 0(the steel route already guarantees it, makingK_packed = K/2a multiple of 8 so every column isu32-aligned) andgroup_size % BK == 0(all real q4 group sizes: 32/64/128). f16-multiply / f32-accumulate, so numerically identical to thehalfkernel.Co-authored-with: Claude Code.
@exported-using-ghexport
Differential Revision: D111163510
Differential Revision: D111163510