flashinfer.attn_scores.compute_paged_mqa_logits_schedule

flashinfer.attn_scores.compute_paged_mqa_logits_schedule(context_lens: Tensor, device: device | None = None, *, use_gpu_kernel: bool = True, out: Tensor | None = None) → Tensor

Compute the CTA schedule tensor for paged MQA logits kernels.

Returns [num_sms+1, 2] int32 on CUDA, ready to pass as schedule_meta to fp8_paged_mqa_logits() or fp4_paged_mqa_logits().

Parameters:

  • context_lens – [B] int32, on CPU or CUDA.

  • device – target CUDA device. Defaults to context_lens.device (or cuda:0 if CPU).

  • use_gpu_kernel – if True (default), compute entirely on-GPU via PagedMQALogitsScheduleKernel — no D2H copy, CUDA-graph-capturable. Falls back to CPU numpy if CuTe DSL is unavailable.

  • out – optional pre-allocated [num_sms+1, 2] int32 on CUDA. Required for CUDA-graph capture (static buffer). If None, a new tensor is allocated each call. This provides static storage, not static contents: the address stays stable, but the values must be recomputed into it whenever ceil(context_lens / 256) changes – including on every graph replay where the lengths may have moved across a 256-token boundary.

Returns:

[num_sms+1, 2] int32 on CUDA (out if provided).

Return type:

schedule_meta