flashinfer.cake_vsa.run_cake_vsa

flashinfer.cake_vsa.run_cake_vsa(plan: dict[str, Any], q: Tensor, k: Tensor, v: Tensor, *, out: Tensor | None, lse: Tensor | None, return_lse: bool, backend: str)

Run one explicit source-level route; no external fallback is available.

Parameters:

  • plan (dict[str, Any]) – Metadata returned by plan_cake_vsa().

  • q (torch.Tensor) – Contiguous query tensor with shape (M, num_qo_heads, head_dim).

  • k (torch.Tensor) – Contiguous key tensor with shape (N, num_kv_heads, head_dim).

  • v (torch.Tensor) – Contiguous value tensor with the same shape and dtype as k.

  • out (Optional[torch.Tensor]) – Optional output buffer matching q.

  • lse (Optional[torch.Tensor]) – Optional float32 log-sum-exp buffer with shape (M, num_qo_heads). It is accepted only when return_lse is true.

  • return_lse (bool) – Return log-sum-exp values with the output. D64/D96, BF16 GQA, ultrasparse, and long-sequence routes do not support this option.

  • backend (str) – Must be "cake".

Returns:

Attention output, or (output, lse) when return_lse is true.

Return type:

torch.Tensor or tuple[torch.Tensor, torch.Tensor]