flashinfer.mla.cake_sparse_mla_sm120_dsv4_nvfp4_plan_splits

flashinfer.mla.cake_sparse_mla_sm120_dsv4_nvfp4_plan_splits(*, num_tokens: int, num_heads: int, topk: int, extra_topk: int = 0, num_sms: int, max_splits: int = 16, head_tiles: int = 1) → Tuple[int, int]

Return (num_splits, chunks_per_block) for one decode call.

Measured split rules (RTX PRO 6000 Blackwell / RTX 5090; GB10 below), with the grid counted in CTAs of 16 * head_tiles heads:

  • two chunks never split: one CTA pipelining both chunks beats two CTAs plus a merge;

  • while the unsplit grid covers at most 40 % of the SMs, take the largest power-of-two split that keeps the grid within 80 % of the SMs and leaves at least two chunks per CTA (one chunk per CTA only pays off for a lone CTA);

  • once the unsplit grid fills the SMs, split in two only for 16-chunk work whose doubled grid ends in a wave at least half full;

  • everything else runs unsplit;

  • independently, chunks_per_block <= max_chunks_per_block always holds (the CTA index table), so more than 1024 candidates force num_splits >= ceil(chunks / 16); max_splits below that raises.

GB10 (SM121, 48 SMs; num_sms < 64): the LPDDR gather saturates with a handful of CTAs, so “fill 80 % of the SMs” over-splits every small grid there. The same loop runs with tighter caps, measured on a 108-row sweep:

  • the split grid stays within SMs / 3 CTAs (16): one token with 128 heads (8 CTAs) is best at two splits, with <= 64 heads at four;

  • the distinct gather streams tokens * splits stay within SMs / 6 (8): eight distinct tokens never gain from a split (+10-28 %), only the shared single-token gather does;

  • at least two chunks per CTA also for a lone CTA (one chunk per CTA is 4-10 % slower than two);

  • no wave-quantization split on a full grid (16-chunk rows are flat).

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).

  • max_splits (int, optional) – Upper bound on num_splits.

  • head_tiles (int, optional) – Sixteen-head tiles per CTA the grid is counted in (1 or 2).

Returns:

(num_splits, chunks_per_block).

Return type:

tuple[int, int]