flashinfer.mla.cake_sparse_mla_sm120_dsv4_nvfp4_plan_prefill

flashinfer.mla.cake_sparse_mla_sm120_dsv4_nvfp4_plan_prefill(*, num_tokens: int, num_heads: int, topk: int, extra_topk: int = 0, num_sms: int) → int

Return head_tiles for one prefill call (measured crossover, see _PREFILL_POLICY).

For num_heads >= 64 with a two-tile instance (head count divisible by 32): two tiles (32 heads per CTA) up to the two-tile token limit, the largest instance above it (4 tiles for head counts divisible by 64, 2 for 96 heads; 80 / 112 heads only have the one-tile instance). The limit is 96 tokens at every candidate count on both SKUs, except 128 heads with 512 to 639 candidates on the large die (RTX PRO 6000), where it is 128 tokens. Narrower head counts take their largest instance at every token count: two tiles for 32 heads, one tile for 16 and 48 heads.

Parameters:

  • num_tokens (int) – Query tokens in the call (q.shape[0]).

  • num_heads (int) – Query heads per token (q.shape[1]).

  • topk (int) – Candidate slots per token in the main cache (indices.shape[1]).

  • extra_topk (int, optional) – Candidate slots per token in the second (compressed) cache; 0 without one.

  • num_sms (int) – Streaming multiprocessors of the target device; keys the thresholds on which the two sm_120a SKUs disagree (torch.cuda.get_device_properties(d).multi_processor_count).

Returns:

head_tiles for the prefill launch: 1, 2 or 4 sixteen-head tiles per CTA.

Return type:

int