[bugfix]fix fa4 decoding (#19388)
Co-authored-by: caoyizhong.cyz <caoyizhong.cyz@alibaba-inc.com>
This commit is contained in:
@@ -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"
|
||||
|
||||
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user