flashinfer.get_dcp_spec_workspace_size_bytes

flashinfer.get_dcp_spec_workspace_size_bytes(batch_size: int, q_len_per_req: int, num_qo_heads: int, num_split: int = 16, *, head_dim: int = 128) → int

Bytes for Cake FMHA Split-KV BF16 partial-O and FP32 partial-LSE scratch.

Pass head_dim=256 for the FP8/page64 D256 ratio-16 production profile; the default remains the D128 profile.

Parameters:

  • batch_size (int) – Number of requests in the batch.

  • q_len_per_req (int) – Number of speculative query rows per request.

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

  • num_split (int) – Maximum Split-KV fanout to reserve, in the inclusive range [2, 16].

  • head_dim (int) – Keyword-only query/output head dimension; supported values are 128 and 256.

Returns:

Required caller-owned workspace size in bytes.

Return type:

int