This is an automated email from the ASF dual-hosted git repository.
akaashrp pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new 3e523be8ce [FIX][WebGPU] Fit wide FP32 prefill heads in shared memory
(#20658)
3e523be8ce is described below
commit 3e523be8ce96a509d61bef37d721d53139139907
Author: Akaash Parthasarathy <[email protected]>
AuthorDate: Wed Oct 14 11:27:21 2026 -0700
[FIX][WebGPU] Fit wide FP32 prefill heads in shared memory (#20658)
WebGPU prefill with 512-wide FP32 attention heads exceeds the 32 KiB
shared-memory limit with the current tile. Use a 4-by-4 tile and 16
threads for this case, bringing the allocation to 24,688 bytes.
The regression test covers causal and sliding-window prefill and
verifies that both kernels compile to WebGPU.
---
python/tvm/relax/frontend/nn/llm/_kernel_common.py | 7 +++++++
tests/python/relax/test_frontend_nn_llm_kernel_config.py | 12 ++++++++++++
2 files changed, 19 insertions(+)
diff --git a/python/tvm/relax/frontend/nn/llm/_kernel_common.py
b/python/tvm/relax/frontend/nn/llm/_kernel_common.py
index 9722a85c52..ec60ede494 100644
--- a/python/tvm/relax/frontend/nn/llm/_kernel_common.py
+++ b/python/tvm/relax/frontend/nn/llm/_kernel_common.py
@@ -634,6 +634,13 @@ def _get_prefill_kernel_config(
):
tile_z = 8
num_warps = 2
+ if dtype == "float32" and d == 512 and d_v == 512:
+ # A 4x4 tile needs 16 threads and 24,688 shared bytes. The usual
+ # 64-thread configuration cannot fit these wide FP32 heads.
+ bdx = 16
+ num_warps = 1
+ tile_x = 4
+ tile_z = 4
if target.kind.name == "metal":
max_shared_memory_per_block =
int(target.attrs["max_shared_memory_per_block"])
tile_x, num_warps, tile_z = _fit_prefill_config_to_shared_memory(
diff --git a/tests/python/relax/test_frontend_nn_llm_kernel_config.py
b/tests/python/relax/test_frontend_nn_llm_kernel_config.py
index 85fb823f90..9056dee2bb 100644
--- a/tests/python/relax/test_frontend_nn_llm_kernel_config.py
+++ b/tests/python/relax/test_frontend_nn_llm_kernel_config.py
@@ -58,6 +58,18 @@ def test_wide_head_prefill_fits_metal_shared_memory():
assert 24_928 <= int(target.attrs["max_shared_memory_per_block"])
[email protected]("sliding_window", [False, True])
+def test_wide_fp32_head_prefill_fits_webgpu_shared_memory(sliding_window):
+ target = tvm.target.Target("webgpu")
+ config = _get_prefill_kernel_config(1, 8, 512, "float32", target)
+ func = _attention_prefill(1, 8, 512, "float32", sliding_window, {}, target)
+
+ assert config == (16, 2, 8, 16, 1, 4, 512, 4)
+ assert _get_allocated_shared_memory(func) <=
int(target.attrs["max_shared_memory_per_block"])
+ mod = tvm.tirx.compile(func.with_attr("global_symbol", "main"),
target=target).mod
+ assert mod.imports[0].inspect_source()
+
+
def test_non_power_of_two_head_prefill_reduces_query_tile_for_metal():
target = tvm.target.Target("metal")
config = _get_prefill_kernel_config(
