[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>
Yi Liu committed
Aug 12, 2026 at 02:05 UTC
4f2f31b82e03917ef33c9008337909ae418dc323
1 file changed
+8
vllm/vllm_flash_attn/flash_attn_interface.py
+8
@@ -397,6 +397,11 @@ def flash_attn_varlen_func(
397
"instead of passing None"
398
)
399
400
+ # FA4 only accepts descales for FP8 inputs. The attention backends
401
+ # initialize these tensors even when the inputs are not quantized.
402
+ if v.dtype not in (torch.float8_e4m3fn, torch.float8_e5m2):
403
+ q_descale = k_descale = v_descale = None
404
+
405
from vllm.vllm_flash_attn.cute.interface import _flash_attn_fwd
406
407
out, softmax_lse, _, _ = _flash_attn_fwd(
@@ -423,6 +428,9 @@ def flash_attn_varlen_func(
428
block_sparse_tensors=block_sparse_tensors,
429
aux_tensors=aux_tensors,
430
aux_tensor_leading_dims=aux_tensor_leading_dims,
431
+ q_descale=q_descale,
432
+ k_descale=k_descale,
433
+ v_descale=v_descale,
434
output_scale=output_scale,
435
)
436
else: