From 6c5f4660238dd9d81df603c9e6bfc6122ad97f94 Mon Sep 17 00:00:00 2001 From: Liangsheng Yin Date: Tue, 23 Jun 2026 14:21:59 -0700 Subject: [PATCH] [server_args] Reland FA4 page_size auto-force for combined --attention-backend fa4 (#28976) --- python/sglang/srt/server_args.py | 10 ++++- .../unit/server_args/test_server_args.py | 42 +++++++++++++++++++ 2 files changed, 51 insertions(+), 1 deletion(-) diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index f92d625c0..e46ea92be 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -4667,9 +4667,17 @@ class ServerArgs: self.attention_backend = "triton" if ( - self.prefill_attention_backend == "fa4" + ( + self.attention_backend == "fa4" + or self.decode_attention_backend == "fa4" + or self.prefill_attention_backend == "fa4" + ) and not self.use_mla_backend() and is_sm100_supported() + # EAGLE topk>1 spec runs the two-pass page-tree cascade, which the FA4 + # CUTLASS kernel aborts on at page_size>1. That path only works at + # page_size==1, so skip the 128 auto-force for it and keep the default. + and (self.speculative_eagle_topk or 0) <= 1 ): logger.warning( f"FA4 backend only supports page size 128 for non-MLA model architectures, changing page_size from {self.page_size} to 128." diff --git a/test/registered/unit/server_args/test_server_args.py b/test/registered/unit/server_args/test_server_args.py index f1e7643d3..5ed5b9687 100644 --- a/test/registered/unit/server_args/test_server_args.py +++ b/test/registered/unit/server_args/test_server_args.py @@ -256,6 +256,48 @@ class TestHiSparseDsaBackendPolicy(unittest.TestCase): server_args._validate_hisparse_kv_cache_dtype() +class TestFa4PageSizeAutoForce(CustomTestCase): + """FA4 requires page_size 128 for non-MLA models on SM100. The auto-force + must trigger for `--attention-backend fa4` (combined) too, not only for the + explicit `--prefill-attention-backend fa4` path.""" + + def _make_args(self, attention_backend, prefill=None, decode=None, page_size=1): + args = ServerArgs(model_path="dummy") + args.attention_backend = attention_backend + args.prefill_attention_backend = prefill + args.decode_attention_backend = decode + args.page_size = page_size + # Short-circuit get_model_config(): the fa4 page_size branch only needs + # use_mla_backend() (mocked) and is_sm100_supported() (mocked), not a + # real model_config. Pre-set the attribute so get_model_config returns + # early without touching ModelConfig.from_server_args. + args.model_config = MagicMock() + args.model_config.hf_config.dual_chunk_attention_config = None + return args + + @patch("sglang.srt.server_args.is_sm100_supported", return_value=True) + @patch("sglang.srt.server_args.ServerArgs.use_mla_backend", return_value=False) + def test_combined_attention_backend_fa4_forces_page_size_128( + self, _mock_mla, _mock_sm100 + ): + # `--attention-backend fa4` (combined): prefill/decode fields stay None. + args = self._make_args(attention_backend="fa4") + + args._handle_attention_backend_compatibility() + + self.assertEqual(args.page_size, 128) + + @patch("sglang.srt.server_args.is_sm100_supported", return_value=True) + @patch("sglang.srt.server_args.ServerArgs.use_mla_backend", return_value=False) + def test_explicit_prefill_fa4_forces_page_size_128(self, _mock_mla, _mock_sm100): + # `--prefill-attention-backend fa4`: the previously-covered path. + args = self._make_args(attention_backend=None, prefill="fa4", page_size=1) + + args._handle_attention_backend_compatibility() + + self.assertEqual(args.page_size, 128) + + class TestContextParallelServerArgs(CustomTestCase): def setUp(self): self.parser = server_args_module.argparse.ArgumentParser()