[bugfix]fix fa4 decoding (#19388)

Co-authored-by: caoyizhong.cyz <caoyizhong.cyz@alibaba-inc.com>
This commit is contained in:
Yizhong Cao
2026-02-27 09:48:21 +08:00
committed by GitHub
parent 054bd71086
commit e567215e44
2 changed files with 10 additions and 1 deletions

View File

@@ -209,7 +209,8 @@ def attn_backend_wrapper(runner: "ModelRunner", full_attn_backend: "AttentionBac
assert (
runner.server_args.attention_backend == "triton"
or runner.server_args.attention_backend == "trtllm_mha"
), "triton or trtllm_mha backend are the only supported backends on Blackwell GPUs for hybrid GDN models, use --attention-backend triton or --attention-backend trtllm_mha to specify the backend."
or runner.server_args.attention_backend == "fa4"
), "triton or trtllm_mha or fa4 backend are the only supported backends on Blackwell GPUs for hybrid GDN models, use --attention-backend triton or --attention-backend trtllm_mha to specify the backend."
if is_npu():
assert (
runner.server_args.attention_backend == "ascend"

View File

@@ -1136,6 +1136,14 @@ class FlashAttentionBackend(AttentionBackend):
if sinks is not None:
kwargs["sinks"] = sinks
flash_attn_with_kvcache_base = flash_attn_with_kvcache_fa3
flash_attn_with_kvcache = (
flash_attn_with_kvcache_fa4
if self.fa_impl_ver == 4
else flash_attn_with_kvcache_base
)
k_descale, v_descale = None, None
# only use kv scaling if: 1) fp8 kv is explicitly enabled, 2) RadixAttention
# has corresponding quantization method so that layer.k_scale is not None,