flashinfer.prefill.fmha_v2_prefill_sm120

flashinfer.prefill.fmha_v2_prefill_sm120(query: Tensor, key: Tensor, value: Tensor, out: Tensor, num_heads: int, head_dim: int, seq_len: int, scale_softmax: float, scale_bmm1: float | None = None, scale_bmm2: float | None = None, causal: bool = True, return_lse: bool = False, lse: Tensor | None = None, scale_bmm1_d: Tensor | None = None, scale_bmm2_d: Tensor | None = None) → Tensor | Tuple[Tensor, Tensor]

Run SM120 FMHA v2 with contiguous, separate Q/K/V tensors.

Supports FP8 E4M3 MHA, GQA, and MQA self-attention with head dimensions of 64, 128, or 256 and BF16 output. The query head count must be divisible by the K/V head count. scale_bmm1 combines the Q and K dequantization scales with the attention scale; scale_bmm2 is the V dequantization scale. Set causal=False for dense / bidirectional attention. Q/K/V/O use fixed-shape [batch, seq_len, heads, head_dim] storage, so every batch item has the same sequence length. Use trtllm_fmha_v2_prefill() for ragged inputs. FP8 calls may pass scale_bmm1_d and scale_bmm2_d as persistent one-element FP32 CUDA model weights. scale_bmm1_d replaces scale_bmm1 and must contain the full fused q_scale * k_scale / sqrt(head_dim) value; scale_bmm2_d replaces scale_bmm2 and contains the V dequantization scale. If either is omitted, the kernel uses the corresponding host-encoded scale.

This entry point is validated for SM120 and SM121.

Parameters:

  • query (torch.Tensor) – Contiguous FP8 E4M3 tensors in BSHD layout. Q has num_heads heads; K and V have the same, possibly smaller, head count.

  • key (torch.Tensor) – Contiguous FP8 E4M3 tensors in BSHD layout. Q has num_heads heads; K and V have the same, possibly smaller, head count.

  • value (torch.Tensor) – Contiguous FP8 E4M3 tensors in BSHD layout. Q has num_heads heads; K and V have the same, possibly smaller, head count.

  • out (torch.Tensor) – Caller-owned contiguous BF16 output tensor in BSHD layout.

  • num_heads (int) – Fixed attention geometry shared by every batch item.

  • head_dim (int) – Fixed attention geometry shared by every batch item.

  • seq_len (int) – Fixed attention geometry shared by every batch item.

  • scale_softmax (float) – Softmax scale passed to the FMHA kernel.

  • scale_bmm1 (float, optional) – Host-encoded QK and V dequantization scales.

  • scale_bmm2 (float, optional) – Host-encoded QK and V dequantization scales.

  • causal (bool) – Apply bottom-right causal masking when true.

  • return_lse (bool) – Return the log-sum-exp tensor together with out.

  • lse (torch.Tensor, optional) – Caller-owned contiguous FP32 [B, S, H, 2] LSE tensor.

  • scale_bmm1_d (torch.Tensor, optional) – Persistent one-element FP32 CUDA scale tensors overriding host scales.

  • scale_bmm2_d (torch.Tensor, optional) – Persistent one-element FP32 CUDA scale tensors overriding host scales.