Skip to content

Commit 4daf06e

Browse files
committed
[ExecuTorch][WebGPU] Steel q4gsw prefill GEMM — test coverage
Pull Request resolved: #20731 **Add steel-GEMM coverage to the `et_vk.linear_q4gsw` golden sweep (stacked on the steel op diff).** The op diff routes M>1 q4gsw prefill to the new steel GEMM on a >=256-invocation device (`K % 16 == 0`), falling back to shmem/register-tiled otherwise. The existing M>1 CONFIGS (`q_proj_4k`, `gate_proj_pf`, `down_proj_pf`, `shmem_edge`) already exercise steel on such a device via the shape-discovering native sweep; this adds one small config that isolates the steel branch specifically and documents the routing. **Changes:** - `test_quantized_linear.py` / `test_webgpu_native.cpp`: add the `steel` config (M=96, K=2048, N=256) — below the shmem thresholds (K<4096, N<2048) so pre-steel it was register-tiled, which uniquely pins the steel branch; M=96 exercises the partial 64-row tile (edge masking). - Document that M>1 `K % 16 == 0` shapes prefer steel on a >=256-invocation device (lvp) and fall back on a <256 device (SwiftShader) — the same fp64 golden validates whichever kernel runs. Co-authored-with: Claude Code. ghstack-source-id: 401515169 @exported-using-ghexport Differential Revision: [D110660966](https://our.internmc.facebook.com/intern/diff/D110660966/)
1 parent bf12277 commit 4daf06e

2 files changed

Lines changed: 6 additions & 1 deletion

File tree

backends/webgpu/test/ops/test_quantized_linear.py

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -59,7 +59,10 @@ class Q4gswConfig:
5959
# requires N % 8 == 0 (torchao pads N for the scale layout), so odd-N / N=1 are
6060
# not exportable -- bicol's has1 odd-N guard is defensive (mirrors coop4's
6161
# general-N robustness) and unreachable through this op.
62-
# Prefill shapes routing to the shmem GEMM (K>=4096 or N>=2048); M=128.
62+
# M>1 prefill: prefer the steel GEMM (K%16==0) on a >=256-invocation device
63+
# (e.g. lvp); else shmem (K>=4096 or N>=2048) or register-tiled (SwiftShader
64+
# caps at 128). Same fp64 golden regardless of which kernel runs.
65+
Q4gswConfig("steel", 96, 2048, 256), # steel-isolating (K<4096, N<2048)
6366
Q4gswConfig("gate_proj_pf", 128, 2048, 8192), # gate/up prefill (shmem via N)
6467
Q4gswConfig("down_proj_pf", 128, 8192, 2048), # down prefill (shmem via K)
6568
Q4gswConfig("shmem_edge", 130, 4096, 2056), # partial 32-tile bounds

backends/webgpu/test/test_webgpu_native.cpp

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -275,6 +275,8 @@ const Q4gswConfig kQ4gswConfigs[] = {
275275
// scale over 64-256 K-groups). q4gsw requires N % 8 == 0, so odd-N is not
276276
// exportable; bicol's has1 odd-N guard is defensive (mirrors coop4
277277
// general-N robustness).
278+
// M>1: steel GEMM on a >=256-invocation device (K%16==0), else shmem/tiled.
279+
{"steel", 96, 2048, 256, 1e-4f, 1e-3f, true, false}, // steel-isolating
278280
{"gate_proj_pf", 128, 2048, 8192, 1e-4f, 1e-3f, true, false}, // shmem via N
279281
{"down_proj_pf", 128, 8192, 2048, 1e-3f, 1e-2f, true, false}, // shmem via K
280282
{"shmem_edge", 130, 4096, 2056, 1e-4f, 1e-3f, true, false}, // partial tiles

0 commit comments

Comments
 (0)