For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /max/get-started.md).
Mojo function
select_mla_prefill_config
def select_mla_prefill_config[qkv_dtype: DType, *, rope_gmem_dtype: DType, rope_mma_dtype: DType, scale_dtype_: Optional[DType] = None](*, num_q_heads: Int, group: Int, depth: Int, page_size: Int, v_depth: Int) -> MLAConfig[qkv_dtype, rope_gmem_dtype=rope_gmem_dtype, rope_mma_dtype=rope_mma_dtype, scale_dtype_=scale_dtype_]
Selects the supported SM100 MLA-prefill config for these dims.
Shared by the generic / blockscale / per-token-scale prefill kernels, which differ only in the config's dtype parameters, so the single-O fallback policy lives in ONE place instead of three hand-synced copies.
v_depth is the per-head V / output width (v_head_dim); -1 means "V
width == nope width" (the DeepSeek shape). The standard 2-O config is tried
first, so when v_head_dim == qk_nope_head_dim the result is byte-identical
to the pre-decoupling path. A wide V (e.g. v_head_dim=256) overflows the
2-O TMEM layout (standard BN=0), so fall back to a single-O (num_q=1)
config at the TMEM-max BN, then to a BN capped at the supported() floor so
= 2 KV stages still fit shared memory. If none is supported, return the standard config so the caller's
supported()assert reports the real dims.
Returns:
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!