akaashrp commented on code in PR #20235:
URL: https://github.com/apache/tvm/pull/20235#discussion_r3909455856
##########
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:
Thanks for pointing this out. Updated in c10c00f51d: the Metal fallback now
searches nearby `tile_x` candidates, rechecks the scheduler factorization
constraints for every tuple, and keeps `tile_y` equal to the real head
dimension. The d=384 case selects `(tile_x=8, tile_z=8, num_warps=2)`. I also
added d=416/d=640 and expanded-`tile_z` regressions with Metal codegen.
--
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]
