[Bugfix][Attention] Forward per-head FP8 descales through FA4 (#51363)

Signed-off-by: yiliu30 <yi4.liu@intel.com>
Signed-off-by: Yi Liu <yi4.liu@intel.com>
Co-authored-by: OpenAI Codex <noreply@openai.com>
Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com>

vllm/vllm_flash_attn/flash_attn_interface.py

@@ -397,6 +397,11 @@ def flash_attn_varlen_func(
                 "instead of passing None"
             )
 
+        # FA4 only accepts descales for FP8 inputs. The attention backends
+        # initialize these tensors even when the inputs are not quantized.
+        if v.dtype not in (torch.float8_e4m3fn, torch.float8_e5m2):
+            q_descale = k_descale = v_descale = None
+
         from vllm.vllm_flash_attn.cute.interface import _flash_attn_fwd
 
         out, softmax_lse, _, _ = _flash_attn_fwd(
@@ -423,6 +428,9 @@ def flash_attn_varlen_func(
             block_sparse_tensors=block_sparse_tensors,
             aux_tensors=aux_tensors,
             aux_tensor_leading_dims=aux_tensor_leading_dims,
+            q_descale=q_descale,
+            k_descale=k_descale,
+            v_descale=v_descale,
             output_scale=output_scale,
         )
     else: