diff --git a/python/sglang/srt/layers/attention/attention_registry.py b/python/sglang/srt/layers/attention/attention_registry.py index 6eabb10b5..9020e5260 100644 --- a/python/sglang/srt/layers/attention/attention_registry.py +++ b/python/sglang/srt/layers/attention/attention_registry.py @@ -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" diff --git a/python/sglang/srt/layers/attention/flashattention_backend.py b/python/sglang/srt/layers/attention/flashattention_backend.py index c0adc510e..6c5d67220 100644 --- a/python/sglang/srt/layers/attention/flashattention_backend.py +++ b/python/sglang/srt/layers/attention/flashattention_backend.py @@ -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,