flashinfer.mla.cake_sparse_mla_sm120_dsv4_nvfp4_plan_head_tiles

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

Return the number of 16-head tiles per decode CTA (1 or 2).

Two tiles (32 heads per CTA sharing one candidate gather) follow the measured one-tile vs two-tile sweeps on RTX PRO 6000, with ctas the one-tile grid num_tokens * num_heads / 16:

  • ctas >= SMs (a full wave or more): two tiles for every head count;

  • below one wave the halved grid costs SM coverage and the split planner’s extra split costs a merge, so the shared gather has to pay for it: H = 32 pairs at ctas >= 2/3 SMs, or at ctas >= SMs / 3 with at least 8 chunks; H = 64 pairs only for SMs / 3 <= ctas < 2/3 SMs with at least 8 chunks; H >= 96 pairs only with at least 16 chunks at ctas >= SMs / 3.

GB10 (SM121, 48 SMs; num_sms < 64) differs at a full wave with two chunks: a two-chunk CTA gathers too little to pay for the doubled serial MMA of two tiles once the one-tile grid is more than two waves (H = 64 and H = 128 at 32 tokens: one tile 2-3 % faster), while up to two waves the pair still folds the partial second wave into one resident wave (H = 128 at 8 tokens: 1.04-1.07x) and H = 32 always pairs because one CTA then covers the token’s whole head set (1.05-1.07x); beyond eight one-tile waves the pair wins again (H = 64 at 128 tokens 1.04-1.05x, H = 128 flat), so the one-tile band is 2 waves < ctas <= 8 waves. Rows with >= 4 chunks pair at a full wave on both. Sub-wave, the H = 64 band is inclusive at ctas == 2/3 SMs on GB10 (32 CTAs of 8 chunks: two tiles 1.03-1.04x; GB202 never lands exactly on 2/3 of its SM count); every other sub-wave rule is the same on both.

Head counts not divisible by 32 have no two-tile instance.

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:

1 or 2 sixteen-head tiles per decode CTA.

Return type:

int