flashinfer.mla.cake_sparse_mla_sm120_dsv4_nvfp4_select_kernel

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

"decode" or "prefill" for one backend="cake" call (measured crossover).

Decode always serves head counts without a prefill instance (8 heads) and candidate lists beyond the prefill’s 16-chunk table (more than 1024 candidates); partial 64-wide chunks are served by the prefill’s masked tail. Otherwise, with candidates = topk + extra_topk:

• num_heads >= 64: decode for num_tokens <= 8 at 512 and more candidates and for num_tokens <= 16 below; prefill above;

• num_heads < 64: at 512 and more candidates decode for num_tokens <= 8 on devices with >= 180 SMs (RTX PRO 6000) and for num_tokens <= 32 on smaller devices (RTX 5090); below 512 candidates prefill at every token count.

num_sms keys the thresholds on which the two sm_120a SKUs disagree (the decode split planner’s wave structure follows the SM count).

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:

"decode" or "prefill".

Return type:

str