JCNTH opened a new pull request, #20425:
URL: https://github.com/apache/tvm/pull/20425
`dlight.gpu.GEMV` has target configs for cuda, metal, rocm, vulkan and two
OpenCL variants. WebGPU falls through to the generic `else` (`TS, TR = 1, 64`),
so a q4f16_1 decode GEMV runs one output row per 64-thread workgroup. This is
what ships today: the decode matmuls in the prebuilt WebLLM Llama-3.2-1B lib
(`binary-mlc-llm-libs`, `v0_2_84/base`) are all `@workgroup_size(64, 1, 1)`.
This adds a `webgpu` branch:
```python
VEC_C = 1; LOAD_V_SHARED = False; LOAD_V_VEC = -1; UNROLL = 64
SUPPORT_WARP_SHUFFLE = bool(target.attrs.get("supports_subgroups", False))
TS, TR = (8, 16) if isinstance(len_S, int) else (1, 64)
```
## End to end (WebLLM 0.2.85, Llama-3.2-1B-Instruct-q4f16_1)
Two model libs built from source with MLC-LLM, identical except for this
change (4 of 89 WGSL kernels differ: the four static decode GEMVs), loaded side
by side, generations interleaved, 512 tokens × 7 rounds, Apple M4 Max, Chrome
Canary:
| lib | decode tok/s |
|---|---|
| prebuilt `v0_2_84/base` (what ships) | 201.3 |
| built from source, upstream | 199.9 |
| built from source, this PR | **222.7** (1.11×) |
Temperature-0 output is identical across all three.
## Kernel level
GPU timestamp queries, every q4f16_1 GEMV of one decode step, kernels
produced the way MLC-LLM produces them (fusion passes →
`LowBatchGemvSpecialize` → production dlight rule order, `lm_head` with its
symbolic `vocab_size`), one weight buffer per layer so each token streams the
whole model:
| model | fallback | this PR | |
|---|---|---|---|
| Llama-3.1-8B | 11.67 ms GEMV / token | 9.28 ms | 1.26× |
| Llama-3.2-1B | 2.28 ms | 1.83 ms | 1.25× |
### Why (8, 16)
8B, ms per token over the (TS, TR) grid; `(1, 64)` is the fallback:
| | TR=8 | TR=16 | TR=32 | TR=64 |
|---|---|---|---|---|
| TS=1 | 24.18 | 12.71 | – | 11.69 |
| TS=2 | – | 9.38 | – | 11.73 |
| TS=4 | 9.68 | 9.44 | 11.95 | 11.86 |
| TS=8 | 9.65 | **9.29** | 11.82 | – |
| TS=16 | 9.85 | 9.64 | – | – |
The gain needs both `TR <= 16` and `TS >= 2`; neither helps alone. `(2,
16)`, `(4, 16)` and `(8, 16)` are within run-to-run spread; `TS = 8` also keeps
`N / TS` under the 65535 workgroups-per-dimension limit for vocabulary-sized
outputs. The 1B model gives the same ranking.
### Other constants
- `VEC_C = 1`: `VEC_C = 4` is within spread, and `CodeGenWebGPU` rejects
vectors wider than 4 lanes.
- `UNROLL = 64`, not the 256 of cuda/metal/vulkan: at 256 the WGSL for
22016×4096 grows from 63 to 248 lines and is slower than the fallback.
- `LOAD_V_SHARED = False`: `True` fails an iter-affine-map divisibility
check against the `// 8` nibble packing.
- `SUPPORT_WARP_SHUFFLE` only sizes the shared-memory budget; the shuffle
itself is gated by `thread_warp_size`, which `UpdateWebGPUAttrs` derives from
`supports_subgroups`. Subgroups add 1–2%.
- **Symbolic `len_S` keeps `(1, 64)`, and this is required.** MLC's
`lm_head` takes this path (symbolic `vocab_size`). With `TS > 1` on a symbolic
extent the final store of the cross-thread reduction is emitted without the
split predicate: for `vocab_size = 128257` the kernel writes 7 rows past the
end. (That unguarded store is pre-existing and reachable today through the
OpenCL Adreno config, which sets `TS, TR = 2, 32` unconditionally; not
addressed here.)
## Correctness
Every kernel (all five shapes, both models, every config in the grid)
checked against an f64 CPU reference with the output NaN-poisoned and guard
slots after the last row: no NaNs, no writes past the end. Relative L2 vs f64
on the worst shape (14336 → 4096) goes 0.24% → 0.60% from the longer per-thread
f16 accumulation chains (4-bit group quantization itself is ~9.5%). Through the
WebLLM API, top-5 logprobs shift by a mean 0.016 nat at short context with
identical greedy text.
Measured on one adapter (Apple M4 Max). `low_batch_gemv.py` was checked
separately: its fallback is already the best config at batch 1–2 on this
adapter, so it is left alone.
## Tests
`test_decode_gemv_webgpu` pins the scheduled TIR.
`tests/python/s_tir/dlight/test_gpu_gemv.py`: 15 passed. pre-commit clean.
🤖 Generated with [Claude Code](https://claude.com/claude-code)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
