tlopex commented on code in PR #20235:
URL: https://github.com/apache/tvm/pull/20235#discussion_r3901482673
##########
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):
Review Comment:
This search only varies `num_warps` and `tile_z`, so it can incorrectly
report that no legal Metal configuration exists when reducing `tile_x` would
fit comfortably.
For example:
_get_prefill_kernel_config(
h_kv=1,
h_q=8,
d=384,
dtype="float16",
target=tvm.target.Target("metal"),
)
currently raises `ValueError`, reporting that the initial tile requires
512120 bytes. However, `(tile_x=8, tile_z=8, num_warps=2)` satisfies the
scheduler's factorization constraints and uses only 18784 bytes of shared
memory. I verified that this configuration produces the scheduled PrimFunc and
passes Metal codegen.
Could we also search smaller `tile_x` candidates, recomputing the
corresponding factorization constraints, before concluding that the limit is
unachievable? Please add a `d=384` regression test as well.
--
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]
