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
+30 -133
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."
)
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"