[Nikhil Jain (nikhilJain17)]

Motivation

We want to reduce the TTFT of running Qwen3.8-27B on Apple devices via WebGPU.

Kernel-level profiling shows that the wide tiling matmul kernel has the best opportunity for a speedup on prefill:

Solution

We introduce two new kernels:

1. Register tiling
   The existing MatMulNBitsWideTile kernel uses 32×128 output tiles and 128-thread workgroups, with each thread computing a 32×1 output column. Threads cooperatively load a 32×32 activation tile into shared memory and individually load their weight columns. They read activation chunks from shared memory and use subgroupShuffle to distribute them across the subgroup, accumulating partial dot products along K.

We switch to 64×64 output tiles and 64-thread workgroups. Each 32-thread subgroup handles a 64×32 tile, with each thread computing 2×32 outputs. Activations are stored in each thread's local memory, and weights are loaded and broadcasted using subgroupShuffle.

So essentially we change the tile dimensions, and also replace the shared activation buffer (and associated workgroup barriers) with duplicate activation loads.

2. Subgroup Matrix Multiply
   We have a second kernel which uses the experimental subgroupMatrixMulitplyAccumulate operation. Each 128-thread workgroup is comprised of four subgroups of 32 threads. The output tile size is 64x64. Threads cooperatively load activations and weights into shared memory, and then each subgroup does an individual 8x8 matrix multiply and accumulates the results. This kernel was faster in testing, but requires an experimental instruction which may not be supported by all browsers.

Testing

Tested on Mac M4 Pro running Qwen3.8-27B.

Prompt tokens	Baseline TTFT	Register TTFT	Experimental matrix TTFT
256	2.903 s / 2.913 s¹	2.764 s (−4.8%)	2.695 s (−7.5%)
512	5.777 s / 5.783 s¹	5.413 s (−6.3%)	5.304 s (−8.3%)
1,024	12.027 s	11.321 s (−5.9%)	10.276 s (−10.0%)
2,048	24.144 s	22.706 s (−6.0%)	20.994 s (−10.4%)
4,096²	45.75 s	41.69 s (−8.9%)	41.40 s (−9.5%)
8,192²	94.63 s	87.54 s (−7.5%)	88.13 s (−6.9%)

[azure-pipelines]

Azure Pipelines:
There may be pipelines that require an authorized user to comment /azp run to run.

[Nikhil Jain (nikhilJain17)]

@microsoft-github-policy-service agree