From 4f6f5d25c801d8bf8cea49c33f95866939fe1a2c Mon Sep 17 00:00:00 2001 From: Qiaolin Yu Date: Wed, 21 Jan 2026 06:54:02 -0800 Subject: [PATCH] Support fa4 decoding (#16034) --- .../srt/layers/attention/flashattention_backend.py | 10 ++++++++-- python/sglang/srt/server_args.py | 4 ---- 2 files changed, 8 insertions(+), 6 deletions(-) diff --git a/python/sglang/srt/layers/attention/flashattention_backend.py b/python/sglang/srt/layers/attention/flashattention_backend.py index 33d4253b2..367dff9b7 100644 --- a/python/sglang/srt/layers/attention/flashattention_backend.py +++ b/python/sglang/srt/layers/attention/flashattention_backend.py @@ -375,8 +375,15 @@ class FlashAttentionBackend(AttentionBackend): # If num_splits == 0, we use a heuristic to automatically determine the number of splits. # We set nums splits to 1 if deterministic inference is enabled. # See https://thinkingmachines.ai/blog/defeating-nondeterminism-in-llm-inference/ for more details. + # Furthermore, FA4 does not support num_splits=0 with CUDA Graph, so we set num_splits to 1 if CUDA Graph is enabled. self.num_splits = ( - 1 if model_runner.server_args.enable_deterministic_inference else 0 + 1 + if model_runner.server_args.enable_deterministic_inference + or ( + self.fa_impl_ver == 4 + and not model_runner.server_args.disable_cuda_graph + ) + else 0 ) def init_forward_metadata(self, forward_batch: ForwardBatch): @@ -1056,7 +1063,6 @@ class FlashAttentionBackend(AttentionBackend): k_rope: Optional[torch.Tensor] = None, sinks: Optional[torch.Tensor] = None, ) -> torch.Tensor: - assert self.fa_impl_ver in [3], "Only FA3 support decoding" if k is not None: assert v is not None if save_kv_cache: diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 478712244..2797e0670 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -1769,10 +1769,6 @@ class ServerArgs: ) self.attention_backend = "triton" - if self.attention_backend == "fa4" or self.decode_attention_backend == "fa4": - raise ValueError( - "FA4 backend is only supported for prefill. Please use `--prefill-attention-backend fa4` instead." - ) if self.prefill_attention_backend == "fa4" and not self.use_mla_backend(): logger.warning( f"FA4 backend only supports page size 128 for non-MLA model architectures, changing page_size from {self.page_size} to 128."