From 439f759497682895b1a13571f00fa95c1461a4cc Mon Sep 17 00:00:00 2001 From: Akaash Parthasarathy Date: Thu, 13 Aug 2026 00:18:24 -0400 Subject: [PATCH 1/3] [FIX][RELAX][METAL] Constrain wide-head prefill tiling 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. --- .../relax/frontend/nn/llm/_kernel_common.py | 4 +- .../test_frontend_nn_llm_kernel_config.py | 38 +++++++++++++++++++ 2 files changed, 40 insertions(+), 2 deletions(-) create mode 100644 tests/python/relax/test_frontend_nn_llm_kernel_config.py diff --git a/python/tvm/relax/frontend/nn/llm/_kernel_common.py b/python/tvm/relax/frontend/nn/llm/_kernel_common.py index 2de72360ee29..a4ab693f2dd6 100644 --- a/python/tvm/relax/frontend/nn/llm/_kernel_common.py +++ b/python/tvm/relax/frontend/nn/llm/_kernel_common.py @@ -463,9 +463,9 @@ def _get_prefill_kernel_config(h_kv, h_q, d, dtype, target: Target): while (tile_x * tile_y) % (bdx * num_warps) != 0: tile_y += original_tile_y - # Otherwise we would exceed maxComputeWorkgroupStorageSize + # Otherwise we would exceed the per-workgroup storage limit on WebGPU and Metal. if ( - target.kind.name == "webgpu" + target.kind.name in ("webgpu", "metal") and ((d + 127) // 128) * ((DataType(dtype).bits + 15) // 16) >= 4 ): tile_z = 8 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 000000000000..6dfbef9d9880 --- /dev/null +++ b/tests/python/relax/test_frontend_nn_llm_kernel_config.py @@ -0,0 +1,38 @@ +# 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 tvm +import tvm.testing +from tvm.relax.frontend.nn.llm._kernel_common import _get_prefill_kernel_config + + +def test_wide_head_prefill_uses_constrained_metal_tile(): + config = _get_prefill_kernel_config( + h_kv=1, + h_q=8, + d=512, + dtype="float16", + target=tvm.target.Target("metal"), + ) + + _, _, _, _, num_warps, _, _, tile_z = config + assert num_warps == 2 + assert tile_z == 8 + + +if __name__ == "__main__": + tvm.testing.main() From f28be7d2004acae02f0e719c2c5a6888ae648b97 Mon Sep 17 00:00:00 2001 From: Akaash Parthasarathy Date: Sun, 30 Aug 2026 12:34:49 -0400 Subject: [PATCH 2/3] [FIX][RELAX][METAL] Select legal prefill tiles by resource usage Choose Metal prefill tiling from the target shared-memory limit and the kernel's exact Q/K/V layout. Account for unequal ragged-attention dimensions and MLA's merged KV buffer, and reject configurations that do not satisfy the scheduler's factorization constraints.\n\nShare configuration and factorization helpers with both tree-attention paths while preserving existing WebGPU behavior. Add allocation-level coverage for standard, ragged, MLA, and tree-attention kernels, custom Metal limits, and impossible limits. --- .../relax/frontend/nn/llm/_kernel_common.py | 198 +++++++++++++++--- .../relax/frontend/nn/llm/_prefill_kernels.py | 4 +- python/tvm/relax/frontend/nn/llm/tree_attn.py | 56 +---- .../test_frontend_nn_llm_kernel_config.py | 131 +++++++++++- 4 files changed, 305 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 a4ab693f2dd6..5fe8f584a512 100644 --- a/python/tvm/relax/frontend/nn/llm/_kernel_common.py +++ b/python/tvm/relax/frontend/nn/llm/_kernel_common.py @@ -444,7 +444,142 @@ def softmax_update_causal_padded_left( 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, + d, + d_v, + dtype, + load_vec, + bdx, + num_warps, + merged_kv, + max_shared_memory_per_block, +): + """Reduce the key tile 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 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 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 group_size = h_q // h_kv @@ -463,13 +598,28 @@ def _get_prefill_kernel_config(h_kv, h_q, d, dtype, target: Target): while (tile_x * tile_y) % (bdx * num_warps) != 0: tile_y += original_tile_y - # Otherwise we would exceed the per-workgroup storage limit on WebGPU and Metal. + # Preserve the established WebGPU config, which targets WebGPU's portable limits. if ( - target.kind.name in ("webgpu", "metal") + 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"]) + num_warps, tile_z = _fit_prefill_config_to_shared_memory( + tile_x, + tile_y, + tile_z, + 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 +627,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 +642,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 +665,7 @@ def apply_to_so_ewise(sch: s_tir.Schedule, block, 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 +689,7 @@ def apply_to_gemm(sch: s_tir.Schedule, block, tile, r_len=16, k_major=False): 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 +704,9 @@ def apply_to_md(sch, block): 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 9eeb18f8cf93..c60f11642bc8 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 batch_prefill_ragged_kv( # pylint: disable=too-many-branches 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 batch_prefill_ragged_kv( # pylint: disable=too-many-branches 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 9c427ecd0825..9cf0ccb6788d 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 @@ 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 @@ _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 index 6dfbef9d9880..cac0281d01da 100644 --- a/tests/python/relax/test_frontend_nn_llm_kernel_config.py +++ b/tests/python/relax/test_frontend_nn_llm_kernel_config.py @@ -15,23 +15,146 @@ # 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 +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_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_uses_constrained_metal_tile(): + +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=tvm.target.Target("metal"), + target=target, ) - _, _, _, _, num_warps, _, _, tile_z = config + _, _, _, _, 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_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"]) + + +@pytest.mark.parametrize("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__": From c10c00f51d97d1e3a88006ff4e85c394abbb5537 Mon Sep 17 00:00:00 2001 From: Akaash Parthasarathy Date: Tue, 1 Sep 2026 19:43:05 -0400 Subject: [PATCH 3/3] [FIX][RELAX][METAL] Search nearby legal prefill tiles Search query tiles around the heuristic when adjusting warp and key tile sizes cannot satisfy Metal shared-memory and scheduler constraints. Recheck factorization for every candidate and keep the reduction tile equal to the real head dimension. Add Metal codegen regressions for non-power-of-two head dimensions, nearby larger query tiles, and key-tile expansion. --- .../relax/frontend/nn/llm/_kernel_common.py | 71 ++++++++++++++-- .../test_frontend_nn_llm_kernel_config.py | 83 +++++++++++++++++++ 2 files changed, 145 insertions(+), 9 deletions(-) diff --git a/python/tvm/relax/frontend/nn/llm/_kernel_common.py b/python/tvm/relax/frontend/nn/llm/_kernel_common.py index 5fe8f584a512..1ff2db07fd11 100644 --- a/python/tvm/relax/frontend/nn/llm/_kernel_common.py +++ b/python/tvm/relax/frontend/nn/llm/_kernel_common.py @@ -517,6 +517,8 @@ def _fit_prefill_config_to_shared_memory( tile_x, tile_y, tile_z, + preferred_tile_z, + max_tile_x, d, d_v, dtype, @@ -526,7 +528,7 @@ def _fit_prefill_config_to_shared_memory( merged_kv, max_shared_memory_per_block, ): - """Reduce the key tile until the prefill kernel fits shared memory.""" + """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 @@ -536,7 +538,7 @@ def _fit_prefill_config_to_shared_memory( tile_x, tile_y, tile_z, d_v, load_vec, bdx, num_warps, merged_kv ) ): - return num_warps, tile_z + 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: @@ -563,7 +565,52 @@ def _fit_prefill_config_to_shared_memory( ) <= max_shared_memory_per_block ): - return warps, candidate_tile_z + 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 @@ -581,22 +628,26 @@ def _get_prefill_kernel_config( 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 # Preserve the established WebGPU config, which targets WebGPU's portable limits. if ( @@ -607,10 +658,12 @@ def _get_prefill_kernel_config( num_warps = 2 if target.kind.name == "metal": max_shared_memory_per_block = int(target.attrs["max_shared_memory_per_block"]) - num_warps, tile_z = _fit_prefill_config_to_shared_memory( + 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, 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 cac0281d01da..b0561dffc674 100644 --- a/tests/python/relax/test_frontend_nn_llm_kernel_config.py +++ b/tests/python/relax/test_frontend_nn_llm_kernel_config.py @@ -24,6 +24,7 @@ _get_prefill_shared_memory_usage, ) from tvm.relax.frontend.nn.llm._prefill_kernels import ( + _attention_prefill, _attention_prefill_mla, _attention_prefill_ragged, ) @@ -56,6 +57,88 @@ def test_wide_head_prefill_fits_metal_shared_memory(): 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() + + +@pytest.mark.parametrize( + ("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,