Skip to content

[ExecuTorch][WebGPU] Packed-word-dequant f16 steel q4gsw prefill GEMM#20798

Merged
meta-codesync[bot] merged 6 commits into
gh/JCNTH/13/basefrom
gh/JCNTH/13/head
Jul 9, 2026
Merged

[ExecuTorch][WebGPU] Packed-word-dequant f16 steel q4gsw prefill GEMM#20798
meta-codesync[bot] merged 6 commits into
gh/JCNTH/13/basefrom
gh/JCNTH/13/head

Conversation

@JCNTH

@JCNTH JCNTH commented Jul 9, 2026

Copy link
Copy Markdown

Stack from ghstack (oldest at bottom):

+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.
@exported-using-ghexport

Differential Revision: D111163510

Differential Revision: D111163510

[ghstack-poisoned]
@pytorch-bot

pytorch-bot Bot commented Jul 9, 2026

Copy link
Copy Markdown

🔗 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 Failure

As of commit 06fe79f with merge base f4b01a8 (image):

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

  • pull / android / build-android (gh) (trunk failure)
    ImportError: /opt/hostedtoolcache/Python/3.11.15/x64/lib/python3.11/site-packages/executorch/extension/training/pybindings/_training_lib.cpython-311-x86_64-linux-gnu.so: undefined symbol: _ZN3c104impl3cow23materialize_cow_storageERNS_11StorageImplE

This comment was automatically generated by Dr. CI and updates every 15 minutes.

@github-actions

github-actions Bot commented Jul 9, 2026

Copy link
Copy Markdown

This PR needs a release notes: label

If your change should be included in the release notes (i.e. would users of this library care about this change?), please use a label starting with release notes:. This helps us keep track and include your important work in the next release notes.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "release notes: none"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.

[ghstack-poisoned]
[ghstack-poisoned]

@SS-JIA SS-JIA left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Review automatically exported from Phabricator review in Meta.

[ghstack-poisoned]
[ghstack-poisoned]
[ghstack-poisoned]
@meta-codesync meta-codesync Bot merged commit 6676cfa into gh/JCNTH/13/base Jul 9, 2026
177 of 183 checks passed
@meta-codesync meta-codesync Bot deleted the gh/JCNTH/13/head branch July 9, 2026 21:40
@meta-codesync meta-codesync Bot temporarily deployed to cherry-pick-bot July 9, 2026 21:40 Inactive
JCNTH added a commit that referenced this pull request Jul 9, 2026
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/)
JCNTH added a commit that referenced this pull request Jul 9, 2026
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/)
JCNTH added a commit that referenced this pull request Jul 9, 2026
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/)
JCNTH added a commit that referenced this pull request Jul 9, 2026
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/)
JCNTH added a commit that referenced this pull request Jul 9, 2026
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/)
JCNTH added a commit that referenced this pull request Jul 9, 2026
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/)
JCNTH added a commit that referenced this pull request Jul 9, 2026
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/)
JCNTH added a commit that referenced this pull request Jul 9, 2026
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/)
JCNTH added a commit that referenced this pull request Jul 9, 2026
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/)
JCNTH added a commit that referenced this pull request Jul 9, 2026
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/)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. meta-exported

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants