flashinfer.nvfp4_attention_sm120.nvfp4_attention_sm120_quantize_qkv

flashinfer.nvfp4_attention_sm120.nvfp4_attention_sm120_quantize_qkv(q: Tensor, k: Tensor, v: Tensor, per_block_mean: bool = True) → Tuple[Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor]

Preprocess and quantize dense Q/K/V tensors for SM120 NVFP4 attention.

The input layout is [batch, num_qo_heads, M, head_dim] for Q and [batch, num_kv_heads, N, head_dim] for K/V. Inputs must be contiguous CUDA tensors with the same batch size, head dimension, dtype, and device. num_qo_heads must be greater than or equal to and divisible by num_kv_heads. Q is padded independently to M_pad = round_up(M, 128); K/V are padded to N_pad = round_up(N, 128).

Parameters:

  • q (torch.Tensor) – Dense Q/K/V tensors with dtype torch.float16 or torch.bfloat16.

  • k (torch.Tensor) – Dense Q/K/V tensors with dtype torch.float16 or torch.bfloat16.

  • v (torch.Tensor) – Dense Q/K/V tensors with dtype torch.float16 or torch.bfloat16.

  • per_block_mean (bool, optional) – Whether to center Q per 128-token block. When False, Q is centered once across the full sequence.

Returns:

q_fp4, k_fp4, transposed v_fp4_t, scale tensors q_scale, k_scale, v_scale_t, and the compact FP32 QK correction with shape [batch, num_qo_heads, M_pad / 128, N_pad] ([batch, num_qo_heads, 1, N_pad] when per_block_mean=False).

Return type:

Tuple[torch.Tensor, …]