flashinfer.attn_scores.precompile_paged_mqa_logits

flashinfer.attn_scores.precompile_paged_mqa_logits(device: device | None = None, variants: Tuple[str, ...] = ('fp8', 'fp4'), output_dtypes: Tuple[dtype, ...] | None = None, batch_sizes: Sequence[int] | None = None) → None

Pre-compile paged MQA logits kernels for common static configs.

Populates the on-disk CuTe-DSL kernel cache so subsequent calls to fp8_paged_mqa_logits() and fp4_paged_mqa_logits() skip compilation on first use. Call once during deployment setup or as part of a package-build step.

Only the configs listed below are covered; anything else (fp16 epilogue, other head_dim / num_heads, num_epi_subtiles != 1) still compiles on first use. Measured on sm_100a: ~9s for the 8 fp8 kernels, ~3s per output dtype for the 9 fp4.

Parameters:

  • device – CUDA device to target. Defaults to the current CUDA device, so a worker that has set its per-rank device gets kernels for that device without passing anything.

  • variants – Which precisions to build. A deployment normally runs one indexer precision, so pass e.g. ("fp8",) to avoid compiling kernels that will never be called.

  • output_dtypes – Which output dtypes to warm. output_dtype is part of the compile cache key, so a dtype not warmed here still compiles on first use. Defaults to each variant’s common set: float32 for FP8, and both bfloat16 and float32 for FP4 – the API default plus the dtype consumers with a float logits ABI require. Pass an explicit tuple to build only what you run.

  • batch_sizes – Batch sizes whose GPU schedule kernel should be warmed. The schedule kernel specializes on ceil(batch_size / 32) * 32, so it is compiled per 32-row bucket and is NOT covered by the shape sweep above. Sizes in the same bucket collapse to one build. Defaults to None, which warms no schedule buckets – a deployment that captures CUDA graphs for a known set of batch sizes should pass them, or the first capture of each bucket pays compilation.