flashinfer.nvfp4_attention_sm120.nvfp4_attention_sm120_fwd

flashinfer.nvfp4_attention_sm120.nvfp4_attention_sm120_fwd(q_fp4: Tensor, k_fp4: Tensor, v_fp4_t: Tensor, q_scale: Tensor, k_scale: Tensor, v_scale_t: Tensor, qk_correction: Tensor, sm_scale: float | None = None, causal: bool = False, per_block_mean: bool = True, out: Tensor | None = None, lse: Tensor | None = None, out_dtype: dtype = torch.bfloat16, softmax_scale: float | None = None, return_lse: bool = True, unpadded_k_len: int | None = None) → Tensor | Tuple[Tensor, Tensor]

Run SM120 NVFP4 attention on pre-quantized Q/K/V tensors.

The packed tensors should be produced by nvfp4_attention_sm120_quantize_qkv(). q_fp4 uses layout [batch, num_qo_heads, M_pad, head_dim / 2] and k_fp4 uses [batch, num_kv_heads, N_pad, head_dim / 2]. v_fp4_t and v_scale_t are stored transposed with num_kv_heads.

Parameters:

  • q_fp4 (torch.Tensor) – Packed NVFP4 Q/K/V tensors.

  • k_fp4 (torch.Tensor) – Packed NVFP4 Q/K/V tensors.

  • v_fp4_t (torch.Tensor) – Packed NVFP4 Q/K/V tensors.

  • q_scale (torch.Tensor) – Per-vector FP8 scale factors for Q/K/V.

  • k_scale (torch.Tensor) – Per-vector FP8 scale factors for Q/K/V.

  • v_scale_t (torch.Tensor) – Per-vector FP8 scale factors for Q/K/V.

  • qk_correction (torch.Tensor) – Compact FP32 correction term returned by nvfp4_attention_sm120_quantize_qkv(), one row per 128-token Q block.

  • sm_scale (Optional[float], optional) – Scale applied to QK scores before softmax. Defaults to 1 / sqrt(head_dim) when omitted.

  • causal (bool, optional) – Whether to apply a causal mask.

  • per_block_mean (bool, optional) – Must match the value used by nvfp4_attention_sm120_quantize_qkv.

  • out (Optional[torch.Tensor], optional) – Optional output and log-sum-exp buffers.

  • lse (Optional[torch.Tensor], optional) – Optional output and log-sum-exp buffers.

  • out_dtype (torch.dtype, optional) – Output dtype used when out is not provided.

  • softmax_scale (Optional[float], optional) – Deprecated alias for sm_scale.

  • return_lse (bool, optional) – Whether to compute and return the log-sum-exp tensor. Defaults to True for compatibility with the legacy (out, lse) return contract. Set to False to return only the attention output.

  • unpadded_k_len (Optional[int], optional) – Logical K/V sequence length. Values at or beyond this position are masked before softmax. Defaults to the physical N_pad extent for backward compatibility.

Returns:

The attention output when return_lse is False; otherwise, the attention output and log-sum-exp tensor.

Return type:

Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]