minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[refactor] Migrate the page_size resolution chain (stack 12/15) (#30074)

Browse Source 
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
Cheng Wan
2026-07-04 02:22:28 -07:00
committed by GitHub
co-authored by Claude Fable 5
parent abbb41a214
commit 276fbfe880
5 changed files with 410 additions and 135 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/arg_groups/overrides.py
+193
View File
@@ -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
python/sglang/srt/runtime_context.py
+1
View File
@@ -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):
python/sglang/srt/server_args.py
+29 -132
View File
@@ -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."
# 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,
)
elif not is_musa():
self.page_size = 1
else:
self.page_size = 64
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"
test/registered/unit/server_args/test_server_args.py
+2 -2
View File
@@ -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.
test/registered/unit/test_model_overrides.py
+184
View File
@@ -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,