flashinfer.mla.cake_sparse_mla_sm120_dsv4_nvfp4_scratch_bytes

flashinfer.mla.cake_sparse_mla_sm120_dsv4_nvfp4_scratch_bytes(num_tokens: int, num_heads: int, topk: int, extra_topk: int = 0) → int

Workspace bytes that cover every split plan of this shape (partials + LSE + alignment slack).

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.

Returns:

Bytes of caller-owned scratch that cover every split plan of the shape.

Return type:

int