flashinfer.cake_vsa.plan_cake_vsa

flashinfer.cake_vsa.plan_cake_vsa(indptr: Tensor | None, indices: Tensor | None, block_mask: Tensor | None, kv_block_lens: Tensor | None, q2k_indices: Tensor | None, q2k_num: Tensor | None, *, M: int, N: int, R: int, C: int, num_qo_heads: int, num_kv_heads: int, head_dim: int, q_data_type: dtype, sm_scale: float | None, device: device) → dict[str, Any]

Create stable metadata and workspaces for the source-level backend.

Parameters:

  • indptr (Optional[torch.Tensor]) – CSR-style row pointers for a block pattern shared by all heads, with shape (M // R + 1,). Required with indices when neither block_mask nor q2k_indices is supplied.

  • indices (Optional[torch.Tensor]) – Column indices corresponding to indptr. Duplicate or non-packed rows are canonicalized before a fixed-stride source kernel can use them.

  • block_mask (Optional[torch.Tensor]) – Boolean block mask with shape (num_qo_heads, M // R, N // C) or (num_kv_heads, M // R, N // C).

  • kv_block_lens (Optional[torch.Tensor]) – Valid-token count for each KV block, with shape (N // C,). This is supported only for 64-token blocks.

  • q2k_indices (Optional[torch.Tensor]) – Direct block-64 selections as contiguous int32 metadata with shape (num_qo_heads, M // R, topk).

  • q2k_num (Optional[torch.Tensor]) – Number of active q2k_indices entries per row, as contiguous int32 metadata with shape (num_qo_heads, M // R).

  • M (int) – Query sequence length.

  • N (int) – Key/value sequence length.

  • R (int) – Query block size. Cake supports 64 and 128.

  • C (int) – Key/value block size, which must equal R.

  • num_qo_heads (int) – Number of query/output heads.

  • num_kv_heads (int) – Number of key/value heads.

  • head_dim (int) – Per-head dimension. Cake supports 64, 96, and 128.

  • q_data_type (torch.dtype) – Planned Q/K/V dtype, either torch.float16 or torch.bfloat16.

  • sm_scale (Optional[float]) – Softmax scale. None selects 1 / sqrt(head_dim).

  • device (torch.device) – SM100 or SM103 CUDA device that will execute the plan.

Returns:

Validated metadata and reusable workspaces consumed by run_cake_vsa().

Return type:

dict[str, Any]