From 4f2f31b82e03917ef33c9008337909ae418dc323 Mon Sep 17 00:00:00 2001 From: Yi Liu Date: Wed, 12 Aug 2026 02:05:20 +0800 Subject: [PATCH] [Bugfix][Attention] Forward per-head FP8 descales through FA4 (#51363) Signed-off-by: yiliu30 Signed-off-by: Yi Liu Co-authored-by: OpenAI Codex Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com> --- vllm/vllm_flash_attn/flash_attn_interface.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/vllm/vllm_flash_attn/flash_attn_interface.py b/vllm/vllm_flash_attn/flash_attn_interface.py index d3f8a36e70..81768e1209 100644 --- a/vllm/vllm_flash_attn/flash_attn_interface.py +++ b/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: