flashinfer.mla.cake_sparse_mla_sm120_dsv4_nvfp4_prefill

flashinfer.mla.cake_sparse_mla_sm120_dsv4_nvfp4_prefill(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, lse_scale: float = 1.0, head_tiles: int | None = None) → Dict[str, int]

Run the allocation-free Cake SM120 NVFP4 sparse-MLA prefill (one launch, no scratch).

Same tensor contract as cake_sparse_mla_sm120_dsv4_nvfp4_decode() (q [T, H, 512] BF16, HND / NHD / 3-D packed caches with runtime page size and stride, [T, topk] int32 indices with -1 masks, optional lengths / sink / second cache, caller-owned output and out_lse), but one CTA runs every chunk of a token: the candidate list is limited to 16 chunks (topk + extra_topk <= 1024 slots) and the head count needs a prefill instance (16 .. 128, multiples of 16). Writes output and out_lse in place and returns the resolved plan {"head_tiles", "num_ctas"} (one CTA per (token, head block) item); head_tiles (see cake_sparse_mla_sm120_dsv4_nvfp4_prefill_head_tiles()) overrides cake_sparse_mla_sm120_dsv4_nvfp4_plan_prefill().

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.

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

  • head_tiles (int, optional) – Overrides cake_sparse_mla_sm120_dsv4_nvfp4_plan_prefill() with one of the head count’s instances (see cake_sparse_mla_sm120_dsv4_nvfp4_prefill_head_tiles()).

Returns:

The resolved plan {"head_tiles", "num_ctas"}.

Return type:

dict[str, int]