Skip to content
8 changes: 8 additions & 0 deletions vllm/vllm_flash_attn/flash_attn_interface.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand All @@ -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:
Expand Down
Loading