flashinfer.mla.cake_sparse_mla_sm120_dsv4_nvfp4_decode

flashinfer.mla.cake_sparse_mla_sm120_dsv4_nvfp4_decode(q: Tensor, kv_cache: Tensor, indices: Tensor, output: Tensor, out_lse: Tensor, sm_scale: float, *, topk_length: Tensor | None = None, attn_sink: Tensor | None = None, extra_kv_cache: Tensor | None = None, extra_indices: Tensor | None = None, extra_topk_length: Tensor | None = None, mid_out: Tensor | None = None, mid_lse: Tensor | None = None, lse_scale: float = 1.0, num_splits: int | None = None, max_splits: int = 16, head_tiles: int | None = None) → Dict[str, int]

Run the allocation-free Cake SM120 NVFP4 sparse-MLA decode.

Writes output and out_lse in place and returns the resolved plan {"head_tiles", "num_splits", "chunks_per_block"}. mid_out / mid_lse are required when the plan splits (num_splits > 1); size them with cake_sparse_mla_sm120_dsv4_nvfp4_num_chunks splits to cover every plan. head_tiles (1 or 2, head counts divisible by 32) and num_splits override the planners.

Parameters:

  • q (torch.Tensor) – [T, H, 512] contiguous CUDA BF16 queries (448 NoPE dims, then 64 RoPE dims).

  • kv_cache (torch.Tensor) – Paged NVFP4 cache, uint8 with 384 bytes per token: [num_pages, page_size, 384], HND [num_pages, 1, page_size, 384] or NHD [num_pages, page_size, 1, 384]; page size and page stride are read from the view at runtime.

  • indices (torch.Tensor) – [T, topk] (or [T, 1, topk]) int32 candidate slots into kv_cache; -1 masks a slot.

  • output (torch.Tensor) – [T, H, 512] contiguous BF16 attention output, written in place.

  • out_lse (torch.Tensor) – [T, H] contiguous float32, written in place with the base-2 log-sum-exp times lse_scale.

  • sm_scale (float) – Softmax scale applied to the QK logits.

  • topk_length (torch.Tensor, optional) – 1-D int32 with at least T entries: valid candidates per token in indices (every slot when omitted).

  • attn_sink (torch.Tensor, optional) – 1-D float32 with at least H entries: per-head attention-sink value that joins the softmax normalisation.

  • extra_kv_cache (torch.Tensor, optional) – Second (compressed) cache in the same format; used together with extra_indices.

  • extra_indices (torch.Tensor, optional) – [T, extra_topk] int32 slots into extra_kv_cache; required with it.

  • extra_topk_length (torch.Tensor, optional) – 1-D int32 with at least T entries: valid candidates per token in extra_indices.

  • mid_out (torch.Tensor, optional) – [T, H, >= num_splits, 512] contiguous BF16 split partials; required when the plan splits.

  • mid_lse (torch.Tensor, optional) – [T, H, >= num_splits] contiguous float32 split log-sum-exps; required with mid_out.

  • lse_scale (float, optional) – Multiplier applied to the base-2 log-sum-exp written to out_lse.

  • num_splits (int, optional) – Overrides cake_sparse_mla_sm120_dsv4_nvfp4_plan_splits().

  • max_splits (int, optional) – Upper bound for the split planner.

  • head_tiles (int, optional) – Overrides cake_sparse_mla_sm120_dsv4_nvfp4_plan_head_tiles(): 1 or 2 (head counts divisible by 32).

Returns:

The resolved plan {"head_tiles", "num_splits", "chunks_per_block"}.

Return type:

dict[str, int]