flashinfer.attn_scores.padded_context_len

flashinfer.attn_scores.padded_context_len(max_context_len: int) → int

Return the minimum allocated context dimension for paged MQA logits output.

The kernel may write unconditionally into SPLIT_KV-padded trailing positions, so the output tensor must be allocated with at least this many columns.

Use this to pre-allocate the out parameter:

out = torch.empty((B * next_n, padded_context_len(max_ctx)), dtype=…, device=”cuda”) logits = fp8_paged_mqa_logits(…, out=out)