From 276fbfe88025fcb3eb821db10eab894205d2c270 Mon Sep 17 00:00:00 2001 From: Cheng Wan <54331508+ch-wan@users.noreply.github.com> Date: Sat, 4 Jul 2026 02:22:28 -0700 Subject: [PATCH] [refactor] Migrate the page_size resolution chain (stack 12/15) (#30074) Co-authored-by: Claude Fable 5 --- python/sglang/srt/arg_groups/overrides.py | 193 ++++++++++++++++++ python/sglang/srt/runtime_context.py | 1 + python/sglang/srt/server_args.py | 163 +++------------ .../unit/server_args/test_server_args.py | 4 +- test/registered/unit/test_model_overrides.py | 184 +++++++++++++++++ 5 files changed, 410 insertions(+), 135 deletions(-) diff --git a/python/sglang/srt/arg_groups/overrides.py b/python/sglang/srt/arg_groups/overrides.py index c6f053715..2e675ba4c 100644 --- a/python/sglang/srt/arg_groups/overrides.py +++ b/python/sglang/srt/arg_groups/overrides.py @@ -43,6 +43,7 @@ from sglang.srt.utils.common import ( is_cuda, is_flashinfer_available, is_hip, + is_musa, is_npu, is_sm90_supported, is_sm100_supported, @@ -380,6 +381,53 @@ def _lfm2_overrides(server_args: Any, hf_config: Any) -> dict: return {} +@_register_for( + "Qwen3NextForCausalLM", + "Qwen3_5MoeForConditionalGeneration", + "InternS2PreviewForConditionalGeneration", + "Qwen3_5ForConditionalGeneration", +) +def _qwen3_5_hybrid_overrides(server_args: Any, hf_config: Any) -> dict: + if not is_sm100_supported() or server_args.attention_backend is not None: + return {} + sm100_default_attn_backend = "triton" + # trtllm_mha requires speculative_eagle_topk == 1 and page_size > 1. + # _get_default_attn_backend handles the eagle_topk check. + # There is only one case where page_size=1 is required, + # which is when radix cache is enabled and both extra_buffer + # and spec decoding are disabled. + default_attn_backend = server_args._get_default_attn_backend( + use_mla_backend=server_args.use_mla_backend(), + model_config=server_args.get_model_config(), + ) + if default_attn_backend == "trtllm_mha" and not ( + not server_args.enable_mamba_extra_buffer() + and not server_args.disable_radix_cache + and server_args.speculative_algorithm is None + ): + sm100_default_attn_backend = "trtllm_mha" + return { + "attention_backend": sm100_default_attn_backend, + "page_size": 64 if sm100_default_attn_backend == "trtllm_mha" else 1, + } + + +@_register_for("Qwen3VLForConditionalGeneration") +def _qwen3vl_overrides(server_args: Any, hf_config: Any) -> dict: + from sglang.srt.environ import envs + + if ( + is_hip() + and envs.SGLANG_USE_AITER_UNIFIED_ATTN.get() + and server_args.page_size is None + ): + logger.info( + "Setting page_size=16 for aiter unified attention on Qwen3VLForConditionalGeneration." + ) + return {"page_size": 16} + return {} + + @_register_for("Glm4MoeForCausalLM") def _glm4_moe_overrides(server_args: Any, hf_config: Any) -> dict: logger.info( @@ -537,6 +585,72 @@ def _attention_backend_default(view: Any) -> dict: return {} +@register_post_process +def _mla_backend_page_constraints(view: Any) -> dict: + """Page-size constraints of the MLA/TRTLLM backend family (the raises and + the cutedsl prefill fallback stay in the handler; only the page snaps are + declared). The snaps chain on a local value exactly as the legacy blocks + chained on self.page_size.""" + page_size = view.page_size + if ( + view.attention_backend == "flashmla" + or view.decode_attention_backend == "flashmla" + ): + logger.warning( + "FlashMLA only supports a page_size of 64, change page_size to 64." + ) + page_size = 64 + if ( + view.attention_backend == "cutlass_mla" + or view.decode_attention_backend == "cutlass_mla" + ): + logger.warning( + "Cutlass MLA only supports a page_size of 128, change page_size to 128." + ) + page_size = 128 + if ( + view.attention_backend == "trtllm_mla" + or view.decode_attention_backend == "trtllm_mla" + ): + if page_size not in [32, 64]: + logger.warning( + f"TensorRT-LLM MLA only supports page_size of 32 or 64, changing page_size from {page_size} to 64." + ) + page_size = 64 + if ( + view.attention_backend == "tokenspeed_mla" + or view.decode_attention_backend == "tokenspeed_mla" + ): + if page_size not in [32, 64]: + logger.warning( + f"tokenspeed_mla only supports page_size of 32 or 64, changing page_size from {page_size} to 64." + ) + page_size = 64 + if ( + view.attention_backend == "cutedsl_mla" + or view.decode_attention_backend == "cutedsl_mla" + or view.prefill_attention_backend == "cutedsl_mla" + ): + if page_size not in [32, 64]: + logger.warning( + f"CuteDSL MLA only supports page_size of 32 or 64, changing page_size from {page_size} to 64." + ) + page_size = 64 + if ( + view.attention_backend == "trtllm_mha" + or view.decode_attention_backend == "trtllm_mha" + or view.prefill_attention_backend == "trtllm_mha" + ): + if page_size not in [16, 32, 64]: + logger.warning( + f"TensorRT-LLM MHA only supports page_size of 16, 32 or 64, changing page_size from {page_size} to 64." + ) + page_size = 64 + if page_size != view.page_size: + return {"page_size": page_size} + return {} + + @register_post_process def _attention_backend_fa3_fp8_fallback(view: Any) -> dict: if view.attention_backend == "fa3" and view.kv_cache_dtype == "fp8_e5m2": @@ -548,6 +662,28 @@ def _attention_backend_fa3_fp8_fallback(view: Any) -> dict: return {} +@register_post_process +def _fa4_page_constraint(view: Any) -> dict: + if ( + ( + view.attention_backend == "fa4" + or view.decode_attention_backend == "fa4" + or view.prefill_attention_backend == "fa4" + ) + and not view.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 (view.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 {view.page_size} to 128." + ) + return {"page_size": 128} + return {} + + @register_post_process def _attention_backend_platform_fallbacks(view: Any) -> dict: if ( @@ -571,6 +707,24 @@ def _attention_backend_platform_fallbacks(view: Any) -> dict: return {} +@register_post_process +def _intel_xpu_page_constraint(view: Any) -> dict: + _, decode_backend = view.get_attention_backends() + if decode_backend == "intel_xpu": + if view.use_mla_backend(): + supported_page_sizes = [16, 32, 64, 128] + msg = "Intel XPU attention backend for MLA Decode" + else: + supported_page_sizes = [64, 128] + msg = "Intel XPU attention backend" + if view.page_size not in supported_page_sizes: + logger.warning( + f"{msg} only supports page_sizes of {supported_page_sizes}, changing page_size from {view.page_size} to 128." + ) + return {"page_size": 128} + return {} + + @register_post_process def _attention_backend_dual_chunk(view: Any) -> dict: if ( @@ -588,6 +742,30 @@ def _attention_backend_dual_chunk(view: Any) -> dict: return {} +@register_post_process +def _page_size_default(view: Any) -> dict: + if view.page_size is not None: + return {} + from sglang.srt.environ import envs + + # SHUFFLE 5D vectorized KV layout (aiter backend + pa_decode_gluon) + # is tuned for and prefers page_size=64 — making it the default + # when the layout flag is set avoids users having to pass + # --page-size 64 explicitly. The env var is only consumed by the + # ROCm AITER backend, so the auto-bump is gated on HIP; on other + # platforms the SHUFFLE 5D pool has no consumer kernels and the + # env var is silently ignored (see MHATokenToKVPool). + if is_hip() and envs.SGLANG_AITER_KV_CACHE_LAYOUT.get().lower() == "vectorized_5d": + logger.info( + "Setting page_size=64 as default for " + "SGLANG_AITER_KV_CACHE_LAYOUT=vectorized_5d." + ) + return {"page_size": 64} + if not is_musa(): + return {"page_size": 1} + return {"page_size": 64} + + @register_post_process def _dllm_attention_backend(view: Any) -> dict: if view.dllm_algorithm is None: @@ -613,6 +791,21 @@ def _dllm_attention_backend(view: Any) -> dict: return {} +@register_post_process +def _dllm_page_size(view: Any) -> dict: + if view.dllm_algorithm is None or view.disable_radix_cache: + return {} + from sglang.srt.dllm.config import DllmConfig + + config = DllmConfig.from_server_args(view) + if view.page_size % config.block_size != 0: + logger.warning( + f"Setting page size to {config.block_size} for diffusion LLM inference" + ) + return {"page_size": config.block_size} + return {} + + @dataclasses.dataclass(frozen=True) class OverrideRecord: """Provenance of one resolved write: ``base`` is the value before this diff --git a/python/sglang/srt/runtime_context.py b/python/sglang/srt/runtime_context.py index 063c93158..4dcc6297a 100644 --- a/python/sglang/srt/runtime_context.py +++ b/python/sglang/srt/runtime_context.py @@ -319,6 +319,7 @@ class Flags(_StaticFlags): swa_full_tokens_ratio: float = 0.8 disable_hybrid_swa_memory: bool = False sampling_backend: str | None = None + page_size: int | None = None def freeze(self) -> None: for field in dataclasses.fields(self): diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 1e642739c..822207431 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -750,7 +750,10 @@ class ServerArgs: float, "How conservative the schedule policy is. A larger value means more conservative scheduling. Use a larger value if you see requests being retracted frequently.", ] = 1.0 - page_size: A[Optional[int], "The number of tokens in a page."] = None + page_size: A[ + Optional[int], + Arg(help="The number of tokens in a page.", model_overridable=True), + ] = None swa_full_tokens_ratio: A[ float, Arg( @@ -4334,30 +4337,8 @@ class ServerArgs: "InternS2PreviewForConditionalGeneration", "Qwen3_5ForConditionalGeneration", ]: - sm100_default_attn_backend = "triton" - if is_sm100_supported(): - # trtllm_mha requires speculative_eagle_topk == 1 and page_size > 1. - # _get_default_attn_backend handles the eagle_topk check. - # There is only one case where page_size=1 is required, - # which is when radix cache is enabled and both extra_buffer - # and spec decoding are disabled. - default_attn_backend = self._get_default_attn_backend( - use_mla_backend=self.use_mla_backend(), - model_config=self.get_model_config(), - ) - if default_attn_backend == "trtllm_mha" and not ( - not self.enable_mamba_extra_buffer() - and not self.disable_radix_cache - and self.speculative_algorithm is None - ): - sm100_default_attn_backend = "trtllm_mha" - - if self.attention_backend is None: - self.attention_backend = sm100_default_attn_backend - self.page_size = ( - 64 if sm100_default_attn_backend == "trtllm_mha" else 1 - ) - + # Attention backend + page size defaults moved to the override + # registry (arg_groups/overrides.py: _qwen3_5_hybrid_overrides). self._handle_mamba_radix_cache(model_arch=model_arch) elif model_arch == "MiniCPMV4_6ForConditionalGeneration": @@ -4428,16 +4409,8 @@ class ServerArgs: # MiniMaxM2ForCausalLM (enable_tf32_matmul) moved to the override registry # (arg_groups/overrides.py: _minimax_m2_overrides). - if ( - model_arch in ["Qwen3VLForConditionalGeneration"] - and is_hip() - and envs.SGLANG_USE_AITER_UNIFIED_ATTN.get() - and self.page_size is None - ): - self.page_size = 16 - logger.info( - "Setting page_size=16 for aiter unified attention on Qwen3VLForConditionalGeneration." - ) + # Qwen3VL aiter unified-attention page_size moved to the override registry + # (arg_groups/overrides.py: _qwen3vl_overrides). if envs.SGLANG_EMBEDDINGS_SPARSE_HEAD.is_set(): self.disable_overlap_schedule = True @@ -4640,6 +4613,9 @@ class ServerArgs: _attention_backend_dual_chunk, _attention_backend_fa3_fp8_fallback, _attention_backend_platform_fallbacks, + _fa4_page_constraint, + _intel_xpu_page_constraint, + _mla_backend_page_constraints, run_post_process_pass, ) @@ -4675,24 +4651,11 @@ class ServerArgs: logger.info("Radix cache is disabled for Whisper") self.disable_radix_cache = True - # Major NVIDIA platforms backends - if ( - self.attention_backend == "flashmla" - or self.decode_attention_backend == "flashmla" - ): - logger.warning( - "FlashMLA only supports a page_size of 64, change page_size to 64." - ) - self.page_size = 64 - - if ( - self.attention_backend == "cutlass_mla" - or self.decode_attention_backend == "cutlass_mla" - ): - logger.warning( - "Cutlass MLA only supports a page_size of 128, change page_size to 128." - ) - self.page_size = 128 + # Major NVIDIA platforms backends: the page-size snaps of this family + # moved to the resolution pipeline (arg_groups/overrides.py: + # _mla_backend_page_constraints); the raises and the cutedsl prefill + # fallback stay below. + run_post_process_pass(self, _mla_backend_page_constraints) if ( self.attention_backend == "trtllm_mla" @@ -4703,12 +4666,6 @@ class ServerArgs: "TRTLLM MLA backend is only supported on Blackwell GPUs (SM100/SM12x). Please use a different backend." ) - if self.page_size not in [32, 64]: - logger.warning( - f"TensorRT-LLM MLA only supports page_size of 32 or 64, changing page_size from {self.page_size} to 64." - ) - self.page_size = 64 - if self.kv_cache_dtype not in ["fp8_e4m3", "fp4_e2m1", "bf16", "auto"]: raise ValueError( "TensorRT-LLM MLA backend only supports kv-cache-dtype of fp8_e4m3, fp4_e2m1, bf16, or auto." @@ -4722,11 +4679,6 @@ class ServerArgs: raise ValueError( "tokenspeed_mla backend is only supported on Blackwell GPUs (SM100/SM12x)." ) - if self.page_size not in [32, 64]: - logger.warning( - f"tokenspeed_mla only supports page_size of 32 or 64, changing page_size from {self.page_size} to 64." - ) - self.page_size = 64 if self.kv_cache_dtype not in ["fp8_e4m3"]: raise ValueError( "tokenspeed_mla backend requires kv-cache-dtype=fp8_e4m3, " @@ -4745,11 +4697,6 @@ class ServerArgs: raise ValueError( "CuteDSL MLA backend is only supported on Blackwell GPUs (SM100). Please use a different backend." ) - if self.page_size not in [32, 64]: - logger.warning( - f"CuteDSL MLA only supports page_size of 32 or 64, changing page_size from {self.page_size} to 64." - ) - self.page_size = 64 if self.kv_cache_dtype not in [ "fp8_e4m3", "bf16", @@ -4791,31 +4738,9 @@ class ServerArgs: "TRTLLM MHA backend for decode is only supported on Hopper (SM90), Blackwell (SM100) and (SM120) GPUs. Please use a different decode backend." ) - if self.page_size not in [16, 32, 64]: - logger.warning( - f"TensorRT-LLM MHA only supports page_size of 16, 32 or 64, changing page_size from {self.page_size} to 64." - ) - self.page_size = 64 - run_post_process_pass(self, _attention_backend_fa3_fp8_fallback) - if ( - ( - 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." - ) - self.page_size = 128 + run_post_process_pass(self, _fa4_page_constraint) # AMD platforms backends if self.attention_backend == "aiter": @@ -4831,19 +4756,7 @@ class ServerArgs: "intel_xpu backend is only supported on decode for MLA models, please set --decode-attention-backend to intel_xpu and do not set --attention-backend or --prefill-attention-backend to intel_xpu for prefill instead use triton." ) - if decode_backend == "intel_xpu": - if self.use_mla_backend(): - supported_page_sizes = [16, 32, 64, 128] - msg = "Intel XPU attention backend for MLA Decode" - else: - supported_page_sizes = [64, 128] - msg = "Intel XPU attention backend" - - if self.page_size not in supported_page_sizes: - logger.warning( - f"{msg} only supports page_sizes of {supported_page_sizes}, changing page_size from {self.page_size} to 128." - ) - self.page_size = 128 + run_post_process_pass(self, _intel_xpu_page_constraint) # Dual chunk flash attention backend run_post_process_pass(self, _attention_backend_dual_chunk) @@ -4934,27 +4847,14 @@ class ServerArgs: raise RuntimeError("KV4 is not tested on non-CUDA platforms.") def _handle_page_size(self): - if self.page_size is None: - # SHUFFLE 5D vectorized KV layout (aiter backend + pa_decode_gluon) - # is tuned for and prefers page_size=64 — making it the default - # when the layout flag is set avoids users having to pass - # --page-size 64 explicitly. The env var is only consumed by the - # ROCm AITER backend, so the auto-bump is gated on HIP; on other - # platforms the SHUFFLE 5D pool has no consumer kernels and the - # env var is silently ignored (see MHATokenToKVPool). - if ( - is_hip() - and envs.SGLANG_AITER_KV_CACHE_LAYOUT.get().lower() == "vectorized_5d" - ): - self.page_size = 64 - logger.info( - "Setting page_size=64 as default for " - "SGLANG_AITER_KV_CACHE_LAYOUT=vectorized_5d." - ) - elif not is_musa(): - self.page_size = 1 - else: - self.page_size = 64 + # Moved to the resolution pipeline (arg_groups/overrides.py: + # _page_size_default), invoked here at its legacy slot. + from sglang.srt.arg_groups.overrides import ( + _page_size_default, + run_post_process_pass, + ) + + run_post_process_pass(self, _page_size_default) def _handle_amd_specifics(self): if is_hip(): @@ -6383,14 +6283,11 @@ class ServerArgs: self.disable_overlap_schedule = True if not self.disable_radix_cache: - from sglang.srt.dllm.config import DllmConfig + # The page_size adjustment moved to the resolution pipeline + # (arg_groups/overrides.py: _dllm_page_size). + from sglang.srt.arg_groups.overrides import _dllm_page_size - config = DllmConfig.from_server_args(self) - if self.page_size % config.block_size != 0: - logger.warning( - f"Setting page size to {config.block_size} for diffusion LLM inference" - ) - self.page_size = config.block_size + run_post_process_pass(self, _dllm_page_size) if self.enable_hierarchical_cache: logger.warning( "Hierarchical cache is disabled because of using diffusion LLM inference" diff --git a/test/registered/unit/server_args/test_server_args.py b/test/registered/unit/server_args/test_server_args.py index a6dd00a5f..eac1cb4b6 100644 --- a/test/registered/unit/server_args/test_server_args.py +++ b/test/registered/unit/server_args/test_server_args.py @@ -311,7 +311,7 @@ class TestFa4PageSizeAutoForce(CustomTestCase): 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.arg_groups.overrides.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 @@ -323,7 +323,7 @@ class TestFa4PageSizeAutoForce(CustomTestCase): self.assertEqual(args.page_size, 128) - @patch("sglang.srt.server_args.is_sm100_supported", return_value=True) + @patch("sglang.srt.arg_groups.overrides.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. diff --git a/test/registered/unit/test_model_overrides.py b/test/registered/unit/test_model_overrides.py index 7c1870bee..f50465889 100644 --- a/test/registered/unit/test_model_overrides.py +++ b/test/registered/unit/test_model_overrides.py @@ -69,6 +69,7 @@ class TestModelOverridableWhitelist(CustomTestCase): "disable_hybrid_swa_memory", "sampling_backend", "attention_backend", + "page_size", } ), ) @@ -742,6 +743,189 @@ class TestGoldenModelOverrides(_IsolatedPublish): _dllm_attention_backend(_view(dllm_algorithm=None)), {} ) + def test_page_size_default_pass(self): + from sglang.srt.arg_groups.overrides import ResolvedView, _page_size_default + + # user-set page_size: nothing to declare + self.assertEqual( + _page_size_default(ResolvedView(SimpleNamespace(page_size=64))), {} + ) + # default fill on non-HIP/non-MUSA platforms is 1 + with patch.object(overrides_module, "is_hip", return_value=False): + with patch.object(overrides_module, "is_musa", return_value=False): + self.assertEqual( + _page_size_default(ResolvedView(SimpleNamespace(page_size=None))), + {"page_size": 1}, + ) + with patch.object(overrides_module, "is_musa", return_value=True): + self.assertEqual( + _page_size_default(ResolvedView(SimpleNamespace(page_size=None))), + {"page_size": 64}, + ) + + def test_dllm_page_size_pass(self): + from sglang.srt.arg_groups.overrides import ResolvedView, _dllm_page_size + + def _view(**kw): + defaults = dict( + dllm_algorithm="LowConfidence", disable_radix_cache=False, page_size=1 + ) + defaults.update(kw) + return ResolvedView(SimpleNamespace(**defaults)) + + with patch( + "sglang.srt.dllm.config.DllmConfig.from_server_args", + return_value=SimpleNamespace(block_size=32), + ): + self.assertEqual(_view() and _dllm_page_size(_view()), {"page_size": 32}) + self.assertEqual(_dllm_page_size(_view(page_size=64)), {}) # aligned + self.assertEqual(_dllm_page_size(_view(dllm_algorithm=None)), {}) + self.assertEqual(_dllm_page_size(_view(disable_radix_cache=True)), {}) + + def test_page_size_leaf_materializes_end_state(self): + sa = self._construct("LlamaForCausalLM", "llama") + declared = {f for _s, d in sa._resolved_overrides for f in d} + self.assertIn("page_size", declared) # default fill declared + self.assertEqual(self._publish(sa).page_size, sa.page_size) + + def test_qwen3_5_hybrid_coupled_declaration(self): + from sglang.srt.arg_groups.overrides import _qwen3_5_hybrid_overrides + + def _args(default_backend, **kw): + defaults = dict( + attention_backend=None, + _get_default_attn_backend=lambda **_: default_backend, + use_mla_backend=lambda: False, + get_model_config=lambda: None, + enable_mamba_extra_buffer=lambda: False, + disable_radix_cache=False, + speculative_algorithm=None, + ) + defaults.update(kw) + return SimpleNamespace(**defaults) + + with patch.object(overrides_module, "is_sm100_supported", return_value=True): + # radix on + no extra buffer + no spec -> page_size=1 path + self.assertEqual( + _qwen3_5_hybrid_overrides(_args("trtllm_mha"), None), + {"attention_backend": "triton", "page_size": 1}, + ) + # spec decoding present -> trtllm_mha + page 64 (coupled) + self.assertEqual( + _qwen3_5_hybrid_overrides( + _args("trtllm_mha", speculative_algorithm="EAGLE"), None + ), + {"attention_backend": "trtllm_mha", "page_size": 64}, + ) + # user-set backend: nothing declared + self.assertEqual( + _qwen3_5_hybrid_overrides( + _args("trtllm_mha", attention_backend="fa3"), None + ), + {}, + ) + with patch.object(overrides_module, "is_sm100_supported", return_value=False): + self.assertEqual(_qwen3_5_hybrid_overrides(_args("fa3"), None), {}) + + def test_qwen3vl_page_size(self): + from sglang.srt.arg_groups.overrides import _qwen3vl_overrides + + with patch.object(overrides_module, "is_hip", return_value=True): + with patch("sglang.srt.environ.envs.SGLANG_USE_AITER_UNIFIED_ATTN") as e: + e.get.return_value = True + self.assertEqual( + _qwen3vl_overrides(SimpleNamespace(page_size=None), None), + {"page_size": 16}, + ) + self.assertEqual( + _qwen3vl_overrides(SimpleNamespace(page_size=64), None), {} + ) + + def test_page_constraint_passes_at_callable_level(self): + from sglang.srt.arg_groups.overrides import ( + ResolvedView, + _fa4_page_constraint, + _intel_xpu_page_constraint, + _mla_backend_page_constraints, + ) + + def _view(**kw): + defaults = dict( + attention_backend=None, + decode_attention_backend=None, + prefill_attention_backend=None, + page_size=1, + ) + defaults.update(kw) + return ResolvedView(SimpleNamespace(**defaults)) + + # flashmla snaps to 64 (unconditional within the backend match) + self.assertEqual( + _mla_backend_page_constraints(_view(attention_backend="flashmla")), + {"page_size": 64}, + ) + # trtllm_mla with already-valid page: no declaration + self.assertEqual( + _mla_backend_page_constraints( + _view(attention_backend="trtllm_mla", page_size=32) + ), + {}, + ) + # chained: flashmla via decode -> 64, then trtllm_mha accepts 64 + self.assertEqual( + _mla_backend_page_constraints( + _view( + decode_attention_backend="flashmla", + prefill_attention_backend="trtllm_mha", + ) + ), + {"page_size": 64}, + ) + # no matching backend: nothing declared + self.assertEqual(_mla_backend_page_constraints(_view()), {}) + + with patch.object(overrides_module, "is_sm100_supported", return_value=True): + self.assertEqual( + _fa4_page_constraint( + _view( + attention_backend="fa4", + use_mla_backend=lambda: False, + speculative_eagle_topk=None, + ) + ), + {"page_size": 128}, + ) + self.assertEqual( + _fa4_page_constraint( + _view( + attention_backend="fa4", + use_mla_backend=lambda: False, + speculative_eagle_topk=2, # EAGLE topk>1 keeps default + ) + ), + {}, + ) + + self.assertEqual( + _intel_xpu_page_constraint( + _view( + get_attention_backends=lambda: (None, "intel_xpu"), + use_mla_backend=lambda: False, + ) + ), + {"page_size": 128}, + ) + self.assertEqual( + _intel_xpu_page_constraint( + _view( + get_attention_backends=lambda: (None, "intel_xpu"), + use_mla_backend=lambda: True, + page_size=16, # MLA decode accepts 16 + ) + ), + {}, + ) + def test_monolith_attention_families_at_callable_level(self): from sglang.srt.arg_groups.overrides import ( _falcon_h1_jet_overrides,