flashinfer.msa_ops.msa_proxy_score_fp4

flashinfer.msa_ops.msa_proxy_score_fp4(q_fp4: Tensor, k_fp4: Tensor, q_scale: Tensor, k_scale: Tensor, q_global_scale: float, k_global_scale: float, 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

NVFP4 MSA dense proxy pass for SM120/SM121 (the FP4 counterpart of msa_proxy_score()).

Same contract and output as msa_proxy_score() (per-KV-block max of the unscaled, causally-masked Q K^T logits), but Q/K arrive pre-quantized as packed NVFP4 (e2m1 + per-16 e4m3 block scales + per-tensor global scales), so the index K is read from HBM at ~4 bits/elem. The full-KV index read is the dominant decode-step cost, so this is the bandwidth-saving path; msa_proxy_score() stays as the bf16 precision reference.

Numerics equal a torch dequant of the same packed inputs (not the bf16 reference, which differs by fp4 rounding). The two global scales are folded into the logits as q_global_scale * k_global_scale before the block-max. Short q (decode, including multi-token verify) uses a head-fused packed schedule on the fp4 tensor cores (MmaMXF4NVF4Op) that scores all group_size heads of a kv_head from one shared index-K read; longer q uses the general tensor-core schedule. Both flat and paged K are supported; see the dispatch below for the exact regime bounds.

Parameters:

  • q_fp4 (torch.Tensor) – (total_q, num_qo_heads, 64) uint8, packed e2m1 (2 nibbles/byte).

  • k_fp4 (torch.Tensor) – Flat (total_k, num_kv_heads, 64) with cu_seqlens_k, or paged (num_pages, num_kv_heads, 128, 64) with page_table + seqused_k. num_qo_heads must be a multiple of num_kv_heads.

  • q_scale (torch.Tensor) – Flat uint8 e4m3 block scales in the cuBLAS 128x4 tiled layout, indexed by logical row token*num_heads + head (paged K: (page*num_kv_heads + kv_head)*128 + token_in_page). Produced by nvfp4_quantize() (128x4 layout, sf_vec_size=16); shared with the attention + decode kernels.

  • k_scale (torch.Tensor) – Flat uint8 e4m3 block scales in the cuBLAS 128x4 tiled layout, indexed by logical row token*num_heads + head (paged K: (page*num_kv_heads + kv_head)*128 + token_in_page). Produced by nvfp4_quantize() (128x4 layout, sf_vec_size=16); shared with the attention + decode kernels.

  • q_global_scale (float) – Per-tensor inverse global scales (1 / global_scale).

  • k_global_scale (float) – Per-tensor inverse global scales (1 / global_scale).

  • cu_seqlens_q (torch.Tensor) – As in msa_proxy_score().

  • cu_seqlens_k (torch.Tensor) – As in msa_proxy_score().

  • page_table (torch.Tensor, optional) – As in msa_proxy_score().

  • seqused_k (torch.Tensor, optional) – As in msa_proxy_score().

  • causal (bool) – As in msa_proxy_score().

  • max_seqlen_q (int, optional) – As in msa_proxy_score().

  • max_k_tiles (int, optional) – As in msa_proxy_score().

  • output (torch.Tensor, optional) – As in msa_proxy_score().

  • reduce_heads (bool) – As in msa_proxy_score().

  • q_offset (int or torch.Tensor, optional) – As in msa_proxy_score().

Returns:

Float32 max_score 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