This is an automated email from the ASF dual-hosted git repository.
tlopex 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 efac9ec823 [Fix][Relax][Metal] Constrain wide-head prefill tiling
(#20235)
efac9ec823 is described below
commit efac9ec823c8c550e60d08229d497d546ffa513f
Author: Akaash Parthasarathy <[email protected]>
AuthorDate: Tue Sep 1 22:03:37 2026 -0700
[Fix][Relax][Metal] Constrain wide-head prefill tiling (#20235)
Apply the existing low-storage WebGPU prefill configuration to Metal for
wide attention heads. This keeps generated threadgroup allocations below
Metal device limits, including Gemma 4 global attention heads.
---
python/tvm/relax/frontend/nn/llm/_kernel_common.py | 259 ++++++++++++++++++---
.../tvm/relax/frontend/nn/llm/_prefill_kernels.py | 4 +-
python/tvm/relax/frontend/nn/llm/tree_attn.py | 56 +----
.../relax/test_frontend_nn_llm_kernel_config.py | 244 +++++++++++++++++++
4 files changed, 479 insertions(+), 84 deletions(-)
diff --git a/python/tvm/relax/frontend/nn/llm/_kernel_common.py
b/python/tvm/relax/frontend/nn/llm/_kernel_common.py
index 2de72360ee..1ff2db07fd 100644
--- a/python/tvm/relax/frontend/nn/llm/_kernel_common.py
+++ b/python/tvm/relax/frontend/nn/llm/_kernel_common.py
@@ -444,32 +444,235 @@ def _make_prefill_macros(tile_x, tile_y, tile_z, tile_o,
bdx, num_warps, group_s
return init_states, compute_s_gemm, softmax_update_causal, compute_o_gemm,
softmax_update_valid_length, advance_tile_batch, paged_store_output_lse,
softmax_update_causal_padded_left
-def _get_prefill_kernel_config(h_kv, h_q, d, dtype, target: Target):
+def _get_prefill_shared_memory_usage(
+ tile_x, tile_z, d, dtype, *, d_v=None, merged_kv=False
+):
+ """Return shared bytes, where ``d`` is Q/K width and ``d_v`` is V/output
width.
+
+ ``merged_kv`` denotes MLA's single shared KV buffer. Otherwise K and V
occupy
+ separate buffers, and ``d_v`` defaults to ``d`` for standard attention.
+ """
+ if d_v is None:
+ d_v = d
+ dtype_bytes = (DataType(dtype).bits + 7) // 8
+ kv_elements = tile_z * d if merged_kv else tile_z * (d + d_v)
+ qkv_bytes = (tile_x * d + kv_elements) * dtype_bytes
+ softmax_bytes = (tile_x * tile_z + 3 * tile_x) * 4
+ return qkv_bytes + softmax_bytes
+
+
+def _get_prefill_vector_size(extent, load_vec):
+ """Return the scheduler's vector width for a contiguous extent."""
+ return min(load_vec, extent & ~(extent - 1))
+
+
+def _get_prefill_tile_size(x, y, num_threads):
+ """Return the scheduler's per-thread 2D tile, or ``None`` if none is
legal."""
+ if (x * y) % num_threads != 0:
+ return None
+ elements_per_thread = (x * y) // num_threads
+ inner_y = math.ceil(math.sqrt(elements_per_thread))
+ while inner_y <= elements_per_thread:
+ if elements_per_thread % inner_y == 0:
+ inner_x = elements_per_thread // inner_y
+ if y % inner_y == 0 and x % inner_x == 0:
+ return inner_x, inner_y
+ inner_y += 1
+ return None
+
+
+def _get_prefill_load_config(x, y, num_threads, load_vec):
+ """Return ``(vector width, tile x, tile y)`` for a scheduled load, if
legal."""
+ if (x * y) % num_threads != 0:
+ return None
+ elements_per_thread = (x * y) // num_threads
+ vec_size = min(
+ _get_prefill_vector_size(y, load_vec),
+ _get_prefill_vector_size(elements_per_thread, load_vec),
+ )
+ tile = _get_prefill_tile_size(x, y // vec_size, num_threads)
+ if tile is None:
+ return None
+ return vec_size, *tile
+
+
+def _is_prefill_kernel_config_legal(
+ tile_x, tile_y, tile_z, d_v, load_vec, bdx, num_warps, merged_kv
+):
+ """Check the factorization assumptions made by the prefill schedulers."""
+ num_threads = bdx * num_warps
+ return all(
+ (
+ _get_prefill_tile_size(tile_x, tile_z, num_threads) is not None,
+ _get_prefill_tile_size(tile_x, d_v, num_threads) is not None,
+ _get_prefill_load_config(tile_x, tile_y, num_threads, load_vec) is
not None,
+ _get_prefill_load_config(tile_z, tile_y, num_threads, load_vec) is
not None,
+ merged_kv
+ or _get_prefill_load_config(tile_z, d_v, num_threads, load_vec) is
not None,
+ )
+ )
+
+
+def _fit_prefill_config_to_shared_memory(
+ tile_x,
+ tile_y,
+ tile_z,
+ preferred_tile_z,
+ max_tile_x,
+ d,
+ d_v,
+ dtype,
+ load_vec,
+ bdx,
+ num_warps,
+ merged_kv,
+ max_shared_memory_per_block,
+):
+ """Reduce the query and key tiles until the prefill kernel fits shared
memory."""
+ if (
+ _get_prefill_shared_memory_usage(
+ tile_x, tile_z, d, dtype, d_v=d_v, merged_kv=merged_kv
+ )
+ <= max_shared_memory_per_block
+ and _is_prefill_kernel_config_legal(
+ tile_x, tile_y, tile_z, d_v, load_vec, bdx, num_warps, merged_kv
+ )
+ ):
+ return tile_x, num_warps, tile_z
+
+ candidate_num_warps = sorted({num_warps, min(num_warps, 2), 1},
reverse=True)
+ for warps in candidate_num_warps:
+ for candidate_tile_z in range(tile_z, 0, -1):
+ if not _is_prefill_kernel_config_legal(
+ tile_x,
+ tile_y,
+ candidate_tile_z,
+ d_v,
+ load_vec,
+ bdx,
+ warps,
+ merged_kv,
+ ):
+ continue
+ if (
+ _get_prefill_shared_memory_usage(
+ tile_x,
+ candidate_tile_z,
+ d,
+ dtype,
+ d_v=d_v,
+ merged_kv=merged_kv,
+ )
+ <= max_shared_memory_per_block
+ ):
+ return tile_x, warps, candidate_tile_z
+
+ # If the heuristic query tile cannot be scheduled, prefer the established
+ # two-warp low-storage configuration, then try nearby query tiles in either
+ # direction, preferring the larger tile on ties. For each query tile, avoid
+ # enlarging the original key tile unless thread factorization requires it.
+ fallback_num_warps = sorted(
+ candidate_num_warps,
+ key=lambda warps: (warps != min(num_warps, 2), -warps),
+ )
+ alternative_tile_x = sorted(
+ (candidate for candidate in range(1, max_tile_x + 1) if candidate !=
tile_x),
+ key=lambda candidate: (abs(candidate - tile_x), -candidate),
+ )
+ fallback_tile_z_ceiling = max(tile_z, bdx * num_warps)
+ candidate_tile_z_ranges = (
+ range(preferred_tile_z, 0, -1),
+ range(preferred_tile_z + 1, fallback_tile_z_ceiling + 1),
+ )
+ for warps in fallback_num_warps:
+ for candidate_tile_x in alternative_tile_x:
+ for candidate_tile_z_range in candidate_tile_z_ranges:
+ for candidate_tile_z in candidate_tile_z_range:
+ if not _is_prefill_kernel_config_legal(
+ candidate_tile_x,
+ tile_y,
+ candidate_tile_z,
+ d_v,
+ load_vec,
+ bdx,
+ warps,
+ merged_kv,
+ ):
+ continue
+ if (
+ _get_prefill_shared_memory_usage(
+ candidate_tile_x,
+ candidate_tile_z,
+ d,
+ dtype,
+ d_v=d_v,
+ merged_kv=merged_kv,
+ )
+ <= max_shared_memory_per_block
+ ):
+ return candidate_tile_x, warps, candidate_tile_z
+
+ required = _get_prefill_shared_memory_usage(
+ tile_x, tile_z, d, dtype, d_v=d_v, merged_kv=merged_kv
+ )
+ raise ValueError(
+ "Unable to find a legal prefill tile within the target's shared-memory
limit: "
+ f"initial tile requires {required} bytes, target allows "
+ f"{max_shared_memory_per_block} bytes"
+ )
+
+
+def _get_prefill_kernel_config(
+ h_kv, h_q, d, dtype, target: Target, *, d_v=None, merged_kv=False
+):
+ if d_v is None:
+ d_v = d
NUM_BLKS = 16
- LOAD_VEC = 8 // ((DataType(dtype).bits + 7) // 8) # 8 bytes
+ dtype_bytes = (DataType(dtype).bits + 7) // 8
+ LOAD_VEC = 8 // dtype_bytes # 8 bytes
group_size = h_q // h_kv
bdx = 32
num_warps = 4
+ # Preserve the largest query tile considered by the existing heuristic.
+ max_tile_x = 64 // dtype_bytes
tile_x, tile_y, tile_z = (
- 64 // ((DataType(dtype).bits + 7) // 8) // max(d // 128, 1),
+ max_tile_x // max(d // 128, 1),
d,
- 64 // ((DataType(dtype).bits + 7) // 8) // max(d // 128, 1),
+ max_tile_x // max(d // 128, 1),
)
original_tile_y = tile_y
original_tile_z = tile_z
while (tile_x * tile_z) % (bdx * num_warps) != 0:
tile_z += original_tile_z
- while (tile_x * tile_y) % (bdx * num_warps) != 0:
- tile_y += original_tile_y
+ if target.kind.name != "metal":
+ while (tile_x * tile_y) % (bdx * num_warps) != 0:
+ tile_y += original_tile_y
- # Otherwise we would exceed maxComputeWorkgroupStorageSize
+ # Preserve the established WebGPU config, which targets WebGPU's portable
limits.
if (
target.kind.name == "webgpu"
and ((d + 127) // 128) * ((DataType(dtype).bits + 15) // 16) >= 4
):
tile_z = 8
num_warps = 2
+ 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(
+ tile_x,
+ tile_y,
+ tile_z,
+ original_tile_z,
+ max_tile_x,
+ d,
+ d_v,
+ dtype,
+ LOAD_VEC,
+ bdx,
+ num_warps,
+ merged_kv,
+ max_shared_memory_per_block,
+ )
if target.kind.name == "opencl" and (
("android" in str(target.host)) or ("adreno" in str(target.attrs))
):
@@ -477,6 +680,14 @@ def _get_prefill_kernel_config(h_kv, h_q, d, dtype,
target: Target):
NUM_BLKS = group_size * 8
check_thread_limits(target, bdx=bdx, bdy=num_warps, bdz=1, gdz=1)
+ if not _is_prefill_kernel_config_legal(
+ tile_x, tile_y, tile_z, d_v, LOAD_VEC, bdx, num_warps, merged_kv
+ ):
+ raise ValueError(
+ "Prefill tile is incompatible with the scheduler's thread
factorization: "
+ f"tile=({tile_x}, {tile_y}, {tile_z}, {d_v}), "
+ f"threads={bdx * num_warps}"
+ )
return NUM_BLKS, LOAD_VEC, group_size, bdx, num_warps, tile_x, tile_y,
tile_z
@@ -484,30 +695,15 @@ def _get_prefill_kernel_config(h_kv, h_q, d, dtype,
target: Target):
def _schedule_prefill_kernel(sch: s_tir.Schedule, load_vec, bdx, num_warps,
tile_x, tile_y, tile_z, transform_k_load: bool, merged_qk_load: bool) ->
tvm.s_tir.Schedule:
get_extent = lambda *lps: [int(sch.get(lp).extent) for lp in lps]
- def get_vecsize(extent):
- return min(load_vec, (extent & ~(extent - 1)))
-
- def getxy_vecsize(x, y, t):
- assert (x * y) % t == 0
- return min(get_vecsize(y), get_vecsize(x * y // t))
-
- def get_tile_size(x, y, t):
- cnt = (x * y) // t
- assert (x * y) % t == 0
- tile_y = math.ceil(math.sqrt(cnt))
- while (cnt % tile_y != 0 or y % tile_y != 0 or x % (cnt // tile_y) !=
0) and tile_y <= cnt:
- tile_y += 1
- assert tile_y <= cnt
- tile_x = cnt // tile_y
- return tile_x, tile_y
-
def apply_to_qkv_load(sch: s_tir.Schedule, block):
loop_x, loop_y = sch.get_loops(block)[-2:]
x_extent, y_extent = get_extent(loop_x, loop_y)
- vec_size = getxy_vecsize(x_extent, y_extent, bdx * num_warps)
+ load_config = _get_prefill_load_config(
+ x_extent, y_extent, bdx * num_warps, load_vec
+ )
+ assert load_config is not None
+ vec_size, tile_x, tile_y = load_config
yo, yv = sch.split(loop_y, [None, vec_size])
- yo_extent = y_extent // vec_size
- tile_x, tile_y = get_tile_size(x_extent, yo_extent, (bdx * num_warps))
xo, xi = sch.split(loop_x, [tile_x, None])
yo, yi = sch.split(yo, [tile_y, None])
sch.reorder(xi, yi, xo, yo)
@@ -522,7 +718,7 @@ def _schedule_prefill_kernel(sch: s_tir.Schedule, load_vec,
bdx, num_warps, tile
xo, xi = sch.split(loop_x, factors=[None, tile[0]])
yo, yi = sch.split(loop_y, factors=[None, tile[1]])
sch.reorder(xo, yo, xi, yi)
- yiv_extent = get_vecsize(tile[1])
+ yiv_extent = _get_prefill_vector_size(tile[1], load_vec)
yio, yiv = sch.split(yi, [None, yiv_extent])
sch.unroll(yio)
sch.vectorize(yiv)
@@ -546,7 +742,7 @@ def _schedule_prefill_kernel(sch: s_tir.Schedule, load_vec,
bdx, num_warps, tile
sch.reorder(ko, xi, yi, ki)
else:
sch.reorder(ko, ki, xi, yi)
- yiv_extent = get_vecsize(tile[1])
+ yiv_extent = _get_prefill_vector_size(tile[1], load_vec)
yio, yiv = sch.split(yi, [None, yiv_extent])
sch.unroll(yio)
sch.vectorize(yiv)
@@ -561,8 +757,9 @@ def _schedule_prefill_kernel(sch: s_tir.Schedule, load_vec,
bdx, num_warps, tile
if transform_k_load and not merged_qk_load:
sch.transform_layout("K_load", ("write", 0), lambda i, j: (j, i))
- tile_s = get_tile_size(tile_x, tile_z, bdx * num_warps)
- tile_o = get_tile_size(tile_x, tile_y, bdx * num_warps)
+ tile_s = _get_prefill_tile_size(tile_x, tile_z, bdx * num_warps)
+ tile_o = _get_prefill_tile_size(tile_x, tile_y, bdx * num_warps)
+ assert tile_s is not None and tile_o is not None
apply_to_gemm(sch, sch.get_sblock("S_gemm"), tile_s, k_major=True)
apply_to_gemm(sch, sch.get_sblock("O_gemm"), tile_o, k_major=False)
apply_to_so_ewise(sch, sch.get_sblock("S_store"), tile_s)
diff --git a/python/tvm/relax/frontend/nn/llm/_prefill_kernels.py
b/python/tvm/relax/frontend/nn/llm/_prefill_kernels.py
index 9eeb18f8cf..c60f11642b 100644
--- a/python/tvm/relax/frontend/nn/llm/_prefill_kernels.py
+++ b/python/tvm/relax/frontend/nn/llm/_prefill_kernels.py
@@ -793,7 +793,7 @@ def _attention_prefill_ragged_cpu(h_kv, h_q, d_qk, d_v,
dtype, rope_scaling: dic
def _attention_prefill_ragged(h_kv, h_q, d_qk, d_v, dtype, rope_scaling:
dict[str, Any], target: Target):
- NUM_BLKS, LOAD_VEC, group_size, bdx, num_warps, tile_x, tile_y, tile_z =
_get_prefill_kernel_config(h_kv, h_q, d_qk, dtype, target)
+ NUM_BLKS, LOAD_VEC, group_size, bdx, num_warps, tile_x, tile_y, tile_z =
_get_prefill_kernel_config(h_kv, h_q, d_qk, dtype, target, d_v=d_v)
init_states, compute_s_gemm, softmax_update_causal, compute_o_gemm, _,
advance_tile_batch, paged_store_output_lse, *_ = _make_prefill_macros(tile_x,
tile_y, tile_z, d_v, bdx, num_warps, group_size)
@T.prim_func(s_tir=True)
@@ -926,7 +926,7 @@ def _attention_prefill_ragged(h_kv, h_q, d_qk, d_v, dtype,
rope_scaling: dict[st
def _attention_prefill_mla(h_q, d_latent, d_rope, dtype, sliding_window: bool,
target: Target, page_size: int = 16):
d_qk = d_latent + d_rope
- NUM_BLKS, LOAD_VEC, group_size, bdx, num_warps, tile_x, tile_y, tile_z =
_get_prefill_kernel_config(1, h_q, d_qk, dtype, target)
+ NUM_BLKS, LOAD_VEC, group_size, bdx, num_warps, tile_x, tile_y, tile_z =
_get_prefill_kernel_config(1, h_q, d_qk, dtype, target, d_v=d_latent,
merged_kv=True)
init_states, compute_s_gemm, softmax_update_causal, compute_o_gemm, _,
advance_tile_batch, paged_store_output_lse, *_ = _make_prefill_macros(tile_x,
tile_y, tile_z, d_latent, bdx, num_warps, group_size)
global_symbol = "batch_prefill_paged_kv_mla"
diff --git a/python/tvm/relax/frontend/nn/llm/tree_attn.py
b/python/tvm/relax/frontend/nn/llm/tree_attn.py
index 9c427ecd08..9cf0ccb678 100644
--- a/python/tvm/relax/frontend/nn/llm/tree_attn.py
+++ b/python/tvm/relax/frontend/nn/llm/tree_attn.py
@@ -23,7 +23,6 @@ import math
from typing import Any
from tvm import s_tir, tirx
-from tvm.runtime import DataType
from tvm.script import tirx as T
from tvm.target import Target
@@ -36,9 +35,9 @@ from ._kernel_common import (
_alloc_tile_walk_state,
_declare_length_info,
_get_kv_chunk_len,
+ _get_prefill_kernel_config,
_get_seq_offset,
_rope,
- check_thread_limits,
)
# mypy: disable-error-code="attr-defined,valid-type,no-redef"
@@ -283,31 +282,9 @@ def tree_attn(h_kv, h_q, d, dtype, rope_scaling: dict[str,
Any], target: Target)
The generated IR module.
"""
# pylint: disable=invalid-name,line-too-long
- NUM_BLKS = 16
- LOAD_VEC = 8 // ((DataType(dtype).bits + 7) // 8) # 8 bytes
- group_size = h_q // h_kv
-
- bdx = 32
- num_warps = 4
- tile_x, tile_y, tile_z = (
- 64 // ((DataType(dtype).bits + 7) // 8) // max(d // 128, 1),
- d,
- 64 // ((DataType(dtype).bits + 7) // 8) // max(d // 128, 1),
+ NUM_BLKS, LOAD_VEC, group_size, bdx, num_warps, tile_x, tile_y, tile_z = (
+ _get_prefill_kernel_config(h_kv, h_q, d, dtype, target)
)
- original_tile_y = tile_y
- original_tile_z = tile_z
- while (tile_x * tile_z) % (bdx * num_warps) != 0:
- tile_z += original_tile_z
- while (tile_x * tile_y) % (bdx * num_warps) != 0:
- tile_y += original_tile_y
-
- # Otherwise we would exceed maxComputeWorkgroupStorageSize
- if (
- target.kind.name == "webgpu"
- and ((d + 127) // 128) * ((DataType(dtype).bits + 15) // 16) >= 4
- ):
- tile_z = 8
- num_warps = 2
# fmt: off
@T.prim_func(s_tir=True)
@@ -819,32 +796,9 @@ def tree_attn_with_paged_kv_cache(
The generated IR module.
"""
# pylint: disable=invalid-name, line-too-long
- NUM_BLKS = 16
- LOAD_VEC = 8 // ((DataType(dtype).bits + 7) // 8) # 8 bytes
- group_size = h_q // h_kv
-
- bdx = 32
- num_warps = 4
- tile_x, tile_y, tile_z = (
- 64 // ((DataType(dtype).bits + 7) // 8) // max(d // 128, 1),
- d,
- 64 // ((DataType(dtype).bits + 7) // 8) // max(d // 128, 1),
+ NUM_BLKS, LOAD_VEC, group_size, bdx, num_warps, tile_x, tile_y, tile_z = (
+ _get_prefill_kernel_config(h_kv, h_q, d, dtype, target)
)
- original_tile_y = tile_y
- original_tile_z = tile_z
- while (tile_x * tile_z) % (bdx * num_warps) != 0:
- tile_z += original_tile_z
- while (tile_x * tile_y) % (bdx * num_warps) != 0:
- tile_y += original_tile_y
-
- # Otherwise we would exceed maxComputeWorkgroupStorageSize
- if (
- target.kind.name == "webgpu"
- and ((d + 127) // 128) * ((DataType(dtype).bits + 15) // 16) >= 4
- ):
- tile_z = 8
- num_warps = 2
- check_thread_limits(target, bdx=bdx, bdy=num_warps, bdz=1, gdz=1)
global_symbol = "tree_attn_paged_kv"
sliding_window = False # Sliding window is not supported in this kernel.
diff --git a/tests/python/relax/test_frontend_nn_llm_kernel_config.py
b/tests/python/relax/test_frontend_nn_llm_kernel_config.py
new file mode 100644
index 0000000000..b0561dffc6
--- /dev/null
+++ b/tests/python/relax/test_frontend_nn_llm_kernel_config.py
@@ -0,0 +1,244 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import pytest
+
+import tvm
+import tvm.testing
+from tvm.relax.frontend.nn.llm._kernel_common import (
+ _get_prefill_kernel_config,
+ _get_prefill_shared_memory_usage,
+)
+from tvm.relax.frontend.nn.llm._prefill_kernels import (
+ _attention_prefill,
+ _attention_prefill_mla,
+ _attention_prefill_ragged,
+)
+from tvm.relax.frontend.nn.llm.tree_attn import tree_attn,
tree_attn_with_paged_kv_cache
+
+
+def _get_allocated_shared_memory(func):
+ mod = tvm.IRModule.from_expr(func)
+ mod = tvm.s_tir.transform.ConvertBlocksToOpaque()(mod)
+ mod = tvm.s_tir.transform.LowerOpaqueBlock()(mod)
+ allocated_bytes = tvm.s_tir.analysis.calculate_allocated_bytes(mod)
+ (function_allocations,) = allocated_bytes.values()
+ return function_allocations["shared"]
+
+
+def test_wide_head_prefill_fits_metal_shared_memory():
+ target = tvm.target.Target("metal")
+ config = _get_prefill_kernel_config(
+ h_kv=1,
+ h_q=8,
+ d=512,
+ dtype="float16",
+ target=target,
+ )
+
+ _, _, _, _, num_warps, tile_x, _, tile_z = config
+ assert num_warps == 2
+ assert tile_z == 8
+ assert _get_prefill_shared_memory_usage(tile_x, tile_z, 512, "float16") ==
24_928
+ assert 24_928 <= int(target.attrs["max_shared_memory_per_block"])
+
+
+def test_non_power_of_two_head_prefill_reduces_query_tile_for_metal():
+ target = tvm.target.Target("metal")
+ config = _get_prefill_kernel_config(
+ h_kv=1,
+ h_q=8,
+ d=384,
+ dtype="float16",
+ target=target,
+ )
+ func = _attention_prefill(1, 8, 384, "float16", False, {}, target)
+
+ assert config == (16, 4, 8, 32, 2, 8, 384, 8)
+ assert func.attrs["tirx.is_scheduled"]
+ assert _get_allocated_shared_memory(func) == 18_784
+ assert 18_784 <= int(target.attrs["max_shared_memory_per_block"])
+ assert tvm.tirx.build(func, target=target).imports[0].inspect_source()
+
+
+def test_reduced_query_tile_keeps_the_real_head_dimension():
+ target = tvm.target.Target("metal")
+ config = _get_prefill_kernel_config(
+ h_kv=1,
+ h_q=8,
+ d=416,
+ dtype="float16",
+ target=target,
+ )
+ func = _attention_prefill(1, 8, 416, "float16", False, {}, target)
+
+ assert config == (16, 4, 8, 32, 2, 8, 416, 8)
+ assert func.attrs["tirx.is_scheduled"]
+ assert _get_allocated_shared_memory(func) == 20_320
+ assert 20_320 <= int(target.attrs["max_shared_memory_per_block"])
+ assert tvm.tirx.build(func, target=target).imports[0].inspect_source()
+
+
+def test_fallback_can_select_larger_nearby_query_tile():
+ target = tvm.target.Target("metal")
+ config = _get_prefill_kernel_config(
+ h_kv=1,
+ h_q=8,
+ d=640,
+ dtype="float16",
+ target=target,
+ )
+ func = _attention_prefill(1, 8, 640, "float16", False, {}, target)
+
+ assert config == (16, 4, 8, 32, 2, 8, 640, 8)
+ assert func.attrs["tirx.is_scheduled"]
+ assert _get_allocated_shared_memory(func) == 31_072
+ assert 31_072 <= int(target.attrs["max_shared_memory_per_block"])
+ assert tvm.tirx.build(func, target=target).imports[0].inspect_source()
+
+
[email protected](
+ ("d", "dtype", "expected_config", "expected_shared_memory"),
+ [
+ (672, "float16", (16, 4, 8, 32, 2, 8, 672, 8), 32_608),
+ (768, "float16", (16, 4, 8, 32, 1, 4, 768, 8), 30_896),
+ (384, "float32", (16, 2, 8, 32, 1, 4, 384, 8), 30_896),
+ ],
+)
+def test_fallback_can_expand_key_tile_for_factorization(
+ d, dtype, expected_config, expected_shared_memory
+):
+ target = tvm.target.Target("metal")
+ config = _get_prefill_kernel_config(
+ h_kv=1,
+ h_q=8,
+ d=d,
+ dtype=dtype,
+ target=target,
+ )
+ func = _attention_prefill(1, 8, d, dtype, False, {}, target)
+
+ assert config == expected_config
+ assert func.attrs["tirx.is_scheduled"]
+ assert _get_allocated_shared_memory(func) == expected_shared_memory
+ assert expected_shared_memory <=
int(target.attrs["max_shared_memory_per_block"])
+ assert tvm.tirx.build(func, target=target).imports[0].inspect_source()
+
+
+def test_normal_head_prefill_keeps_existing_metal_config():
+ config = _get_prefill_kernel_config(
+ h_kv=1,
+ h_q=8,
+ d=256,
+ dtype="float16",
+ target=tvm.target.Target("metal"),
+ )
+
+ assert config == (16, 4, 8, 32, 4, 16, 256, 16)
+
+
+def test_wide_head_prefill_keeps_existing_webgpu_config():
+ config = _get_prefill_kernel_config(
+ h_kv=1,
+ h_q=8,
+ d=512,
+ dtype="float16",
+ target=tvm.target.Target("webgpu"),
+ )
+
+ assert config == (16, 4, 8, 32, 2, 8, 512, 8)
+
+
+def test_wide_head_prefill_uses_target_shared_memory_limit():
+ target = tvm.target.Target({"kind": "metal",
"max_shared_memory_per_block": 65_536})
+ config = _get_prefill_kernel_config(
+ h_kv=1,
+ h_q=8,
+ d=512,
+ dtype="float16",
+ target=target,
+ )
+
+ assert config == (16, 4, 8, 32, 4, 8, 512, 16)
+ assert _get_prefill_shared_memory_usage(8, 16, 512, "float16") == 41_568
+ assert 41_568 <= int(target.attrs["max_shared_memory_per_block"])
+
+
+def test_wide_head_prefill_rejects_unachievable_shared_memory_limit():
+ target = tvm.target.Target({"kind": "metal",
"max_shared_memory_per_block": 1_024})
+
+ with pytest.raises(ValueError, match="target allows 1024 bytes"):
+ _get_prefill_kernel_config(
+ h_kv=1,
+ h_q=8,
+ d=512,
+ dtype="float16",
+ target=target,
+ )
+
+
+def test_ragged_prefill_accounts_for_wider_value_dimension():
+ target = tvm.target.Target("metal")
+ config = _get_prefill_kernel_config(
+ h_kv=1,
+ h_q=8,
+ d=256,
+ dtype="float16",
+ target=target,
+ d_v=512,
+ )
+ func = _attention_prefill_ragged(1, 8, 256, 512, "float16", {}, target)
+
+ assert config == (16, 4, 8, 32, 4, 16, 256, 8)
+ assert _get_prefill_shared_memory_usage(16, 8, 256, "float16", d_v=512) ==
21_184
+ assert _get_allocated_shared_memory(func) == 21_184
+ assert 21_184 <= int(target.attrs["max_shared_memory_per_block"])
+
+
+def test_mla_prefill_accounts_for_merged_kv_buffer():
+ target = tvm.target.Target("metal")
+ config = _get_prefill_kernel_config(
+ h_kv=1,
+ h_q=8,
+ d=576,
+ dtype="float16",
+ target=target,
+ d_v=512,
+ merged_kv=True,
+ )
+ func = _attention_prefill_mla(8, 512, 64, "float16", False, target)
+
+ assert config == (16, 4, 8, 32, 4, 8, 576, 16)
+ assert (
+ _get_prefill_shared_memory_usage(8, 16, 576, "float16", d_v=512,
merged_kv=True) == 28_256
+ )
+ assert _get_allocated_shared_memory(func) == 28_256
+ assert 28_256 <= int(target.attrs["max_shared_memory_per_block"])
+
+
[email protected]("kernel", [tree_attn, tree_attn_with_paged_kv_cache])
+def test_wide_head_tree_attention_has_legal_metal_schedule(kernel):
+ target = tvm.target.Target("metal")
+ func = kernel(1, 8, 512, "float16", {}, target)
+
+ assert func.attrs["tirx.is_scheduled"]
+ assert _get_allocated_shared_memory(func) == 24_928
+ assert 24_928 <= int(target.attrs["max_shared_memory_per_block"])
+
+
+if __name__ == "__main__":
+ tvm.testing.main()
