flashinfer.msa_ops.msa_proxy_score

flashinfer.msa_ops.msa_proxy_score(q: Tensor, k: Tensor, cu_seqlens_q: Tensor, cu_seqlens_k: Tensor | None = None, *, page_table: Tensor | None = None, seqused_k: Tensor | None = None, causal: bool = True, max_seqlen_q: int | None = None, max_k_tiles: int | None = None, output: Tensor | None = None, reduce_heads: bool = False, q_offset=None) → Tensor

MSA dense proxy pass for SM120/SM121: per-KV-block max attention logits.

Computes max_score[h, t, q], the maximum of the unscaled, causally-masked Q K^T logits over the 128 tokens of KV block t, for every query token and query head. The output feeds directly into msa_topk_select(). KV blocks beyond a sequence’s valid range, or entirely above the causal limit, yield -inf.

There is no softmax and no V: this is MSA pipeline stage 1.

Single-token decode uses a dim-parallel scalar schedule that streams index-K straight to registers. Short multi-token decode (MTP) uses a head-fused packed schedule that scores all group_size heads of a kv_head from one shared index-K read. fp8 K and prefill use the general schedule; see the dispatch below for the exact regime bounds.

Parameters:

  • q (torch.Tensor) – (total_q, num_qo_heads, 128), bf16 or fp16 (the cheap proxy Q).

  • k (torch.Tensor) – Flat (total_k, num_kv_heads, 128) with cu_seqlens_k, or paged (num_pages, num_kv_heads, 128, 128) with page_table + seqused_k. May be fp8 E4M3 (upconverted in-kernel). num_qo_heads must be a multiple of num_kv_heads.

  • cu_seqlens_q (torch.Tensor) – (batch_size + 1,) int32 cumulative query lengths.

  • cu_seqlens_k (torch.Tensor, optional) – (batch_size + 1,) int32 cumulative KV lengths. Required for flat KV layout and unused for paged KV layout.

  • page_table (torch.Tensor, optional) – Page-table mapping for paged KV layout. Required together with seqused_k when k is paged.

  • seqused_k (torch.Tensor, optional) – Per-sequence valid KV-token counts. Required together with page_table for paged KV layout.

  • causal (bool) – Right-aligned causal masking, applied before the block max.

  • max_seqlen_q (int, optional) – Maximum query length. Required by some CUDA Graph / compilation paths.

  • max_k_tiles (int, optional) – Number of KV-block columns in the output; defaults to the maximum ceil(seqlen_k / 128) over the batch.

  • output (torch.Tensor, optional) – Pre-allocated float32 output. Shape is (num_qo_heads, max_k_tiles, total_q) normally, or (1, max_k_tiles, total_q) when reduce_heads=True.

  • reduce_heads (bool) – If True, max-reduce the per-head max_score over the query-head axis and return a single (1, max_k_tiles, total_q) score. Use this when the query heads are an indexer’s proxy heads that collapse to one shared block selection per query (MiniMax-M3 indexer semantics). Defaults to False: per-head max_score, where each head selects its own blocks. The reduction is currently a post-kernel amax over the per-head buffer (the kernel is one CTA per head, so a cross-head epilogue would need cross-CTA float atomics); folding it into the kernel is a possible future optimization (saves materializing the per-head buffer).

  • q_offset (int or torch.Tensor, optional) – Optional query-position offset used by the causal alignment logic.

Returns:

Float32 max_score ready for msa_topk_select(): (num_qo_heads, max_k_tiles, total_q), or (1, max_k_tiles, total_q) when reduce_heads=True.

Return type:

torch.Tensor