Support fa4 decoding (#16034)

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:

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."