[server_args] Reland FA4 page_size auto-force for combined --attention-backend fa4 (#28976)

This commit is contained in:
Liangsheng Yin
2026-06-23 14:21:59 -07:00
committed by GitHub
parent 93015a9e6b
commit 6c5f466023
2 changed files with 51 additions and 1 deletions
+9 -1
View File
@@ -4667,9 +4667,17 @@ class ServerArgs:
self.attention_backend = "triton" self.attention_backend = "triton"
if ( 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 not self.use_mla_backend()
and is_sm100_supported() 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( logger.warning(
f"FA4 backend only supports page size 128 for non-MLA model architectures, changing page_size from {self.page_size} to 128." f"FA4 backend only supports page size 128 for non-MLA model architectures, changing page_size from {self.page_size} to 128."
@@ -256,6 +256,48 @@ class TestHiSparseDsaBackendPolicy(unittest.TestCase):
server_args._validate_hisparse_kv_cache_dtype() 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): class TestContextParallelServerArgs(CustomTestCase):
def setUp(self): def setUp(self):
self.parser = server_args_module.argparse.ArgumentParser() self.parser = server_args_module.argparse.ArgumentParser()