[refactor] Migrate the page_size resolution chain (stack 12/15) (#30074)
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
co-authored by
Claude Fable 5
parent
abbb41a214
commit
276fbfe880
@@ -43,6 +43,7 @@ from sglang.srt.utils.common import (
|
|||||||
is_cuda,
|
is_cuda,
|
||||||
is_flashinfer_available,
|
is_flashinfer_available,
|
||||||
is_hip,
|
is_hip,
|
||||||
|
is_musa,
|
||||||
is_npu,
|
is_npu,
|
||||||
is_sm90_supported,
|
is_sm90_supported,
|
||||||
is_sm100_supported,
|
is_sm100_supported,
|
||||||
@@ -380,6 +381,53 @@ def _lfm2_overrides(server_args: Any, hf_config: Any) -> dict:
|
|||||||
return {}
|
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")
|
@_register_for("Glm4MoeForCausalLM")
|
||||||
def _glm4_moe_overrides(server_args: Any, hf_config: Any) -> dict:
|
def _glm4_moe_overrides(server_args: Any, hf_config: Any) -> dict:
|
||||||
logger.info(
|
logger.info(
|
||||||
@@ -537,6 +585,72 @@ def _attention_backend_default(view: Any) -> dict:
|
|||||||
return {}
|
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
|
@register_post_process
|
||||||
def _attention_backend_fa3_fp8_fallback(view: Any) -> dict:
|
def _attention_backend_fa3_fp8_fallback(view: Any) -> dict:
|
||||||
if view.attention_backend == "fa3" and view.kv_cache_dtype == "fp8_e5m2":
|
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 {}
|
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
|
@register_post_process
|
||||||
def _attention_backend_platform_fallbacks(view: Any) -> dict:
|
def _attention_backend_platform_fallbacks(view: Any) -> dict:
|
||||||
if (
|
if (
|
||||||
@@ -571,6 +707,24 @@ def _attention_backend_platform_fallbacks(view: Any) -> dict:
|
|||||||
return {}
|
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
|
@register_post_process
|
||||||
def _attention_backend_dual_chunk(view: Any) -> dict:
|
def _attention_backend_dual_chunk(view: Any) -> dict:
|
||||||
if (
|
if (
|
||||||
@@ -588,6 +742,30 @@ def _attention_backend_dual_chunk(view: Any) -> dict:
|
|||||||
return {}
|
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
|
@register_post_process
|
||||||
def _dllm_attention_backend(view: Any) -> dict:
|
def _dllm_attention_backend(view: Any) -> dict:
|
||||||
if view.dllm_algorithm is None:
|
if view.dllm_algorithm is None:
|
||||||
@@ -613,6 +791,21 @@ def _dllm_attention_backend(view: Any) -> dict:
|
|||||||
return {}
|
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)
|
@dataclasses.dataclass(frozen=True)
|
||||||
class OverrideRecord:
|
class OverrideRecord:
|
||||||
"""Provenance of one resolved write: ``base`` is the value before this
|
"""Provenance of one resolved write: ``base`` is the value before this
|
||||||
|
|||||||
@@ -319,6 +319,7 @@ class Flags(_StaticFlags):
|
|||||||
swa_full_tokens_ratio: float = 0.8
|
swa_full_tokens_ratio: float = 0.8
|
||||||
disable_hybrid_swa_memory: bool = False
|
disable_hybrid_swa_memory: bool = False
|
||||||
sampling_backend: str | None = None
|
sampling_backend: str | None = None
|
||||||
|
page_size: int | None = None
|
||||||
|
|
||||||
def freeze(self) -> None:
|
def freeze(self) -> None:
|
||||||
for field in dataclasses.fields(self):
|
for field in dataclasses.fields(self):
|
||||||
|
|||||||
@@ -750,7 +750,10 @@ class ServerArgs:
|
|||||||
float,
|
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.",
|
"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
|
] = 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[
|
swa_full_tokens_ratio: A[
|
||||||
float,
|
float,
|
||||||
Arg(
|
Arg(
|
||||||
@@ -4334,30 +4337,8 @@ class ServerArgs:
|
|||||||
"InternS2PreviewForConditionalGeneration",
|
"InternS2PreviewForConditionalGeneration",
|
||||||
"Qwen3_5ForConditionalGeneration",
|
"Qwen3_5ForConditionalGeneration",
|
||||||
]:
|
]:
|
||||||
sm100_default_attn_backend = "triton"
|
# Attention backend + page size defaults moved to the override
|
||||||
if is_sm100_supported():
|
# registry (arg_groups/overrides.py: _qwen3_5_hybrid_overrides).
|
||||||
# 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
|
|
||||||
)
|
|
||||||
|
|
||||||
self._handle_mamba_radix_cache(model_arch=model_arch)
|
self._handle_mamba_radix_cache(model_arch=model_arch)
|
||||||
|
|
||||||
elif model_arch == "MiniCPMV4_6ForConditionalGeneration":
|
elif model_arch == "MiniCPMV4_6ForConditionalGeneration":
|
||||||
@@ -4428,16 +4409,8 @@ class ServerArgs:
|
|||||||
# MiniMaxM2ForCausalLM (enable_tf32_matmul) moved to the override registry
|
# MiniMaxM2ForCausalLM (enable_tf32_matmul) moved to the override registry
|
||||||
# (arg_groups/overrides.py: _minimax_m2_overrides).
|
# (arg_groups/overrides.py: _minimax_m2_overrides).
|
||||||
|
|
||||||
if (
|
# Qwen3VL aiter unified-attention page_size moved to the override registry
|
||||||
model_arch in ["Qwen3VLForConditionalGeneration"]
|
# (arg_groups/overrides.py: _qwen3vl_overrides).
|
||||||
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."
|
|
||||||
)
|
|
||||||
|
|
||||||
if envs.SGLANG_EMBEDDINGS_SPARSE_HEAD.is_set():
|
if envs.SGLANG_EMBEDDINGS_SPARSE_HEAD.is_set():
|
||||||
self.disable_overlap_schedule = True
|
self.disable_overlap_schedule = True
|
||||||
@@ -4640,6 +4613,9 @@ class ServerArgs:
|
|||||||
_attention_backend_dual_chunk,
|
_attention_backend_dual_chunk,
|
||||||
_attention_backend_fa3_fp8_fallback,
|
_attention_backend_fa3_fp8_fallback,
|
||||||
_attention_backend_platform_fallbacks,
|
_attention_backend_platform_fallbacks,
|
||||||
|
_fa4_page_constraint,
|
||||||
|
_intel_xpu_page_constraint,
|
||||||
|
_mla_backend_page_constraints,
|
||||||
run_post_process_pass,
|
run_post_process_pass,
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -4675,24 +4651,11 @@ class ServerArgs:
|
|||||||
logger.info("Radix cache is disabled for Whisper")
|
logger.info("Radix cache is disabled for Whisper")
|
||||||
self.disable_radix_cache = True
|
self.disable_radix_cache = True
|
||||||
|
|
||||||
# Major NVIDIA platforms backends
|
# Major NVIDIA platforms backends: the page-size snaps of this family
|
||||||
if (
|
# moved to the resolution pipeline (arg_groups/overrides.py:
|
||||||
self.attention_backend == "flashmla"
|
# _mla_backend_page_constraints); the raises and the cutedsl prefill
|
||||||
or self.decode_attention_backend == "flashmla"
|
# fallback stay below.
|
||||||
):
|
run_post_process_pass(self, _mla_backend_page_constraints)
|
||||||
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
|
|
||||||
|
|
||||||
if (
|
if (
|
||||||
self.attention_backend == "trtllm_mla"
|
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."
|
"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"]:
|
if self.kv_cache_dtype not in ["fp8_e4m3", "fp4_e2m1", "bf16", "auto"]:
|
||||||
raise ValueError(
|
raise ValueError(
|
||||||
"TensorRT-LLM MLA backend only supports kv-cache-dtype of fp8_e4m3, fp4_e2m1, bf16, or auto."
|
"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(
|
raise ValueError(
|
||||||
"tokenspeed_mla backend is only supported on Blackwell GPUs (SM100/SM12x)."
|
"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"]:
|
if self.kv_cache_dtype not in ["fp8_e4m3"]:
|
||||||
raise ValueError(
|
raise ValueError(
|
||||||
"tokenspeed_mla backend requires kv-cache-dtype=fp8_e4m3, "
|
"tokenspeed_mla backend requires kv-cache-dtype=fp8_e4m3, "
|
||||||
@@ -4745,11 +4697,6 @@ class ServerArgs:
|
|||||||
raise ValueError(
|
raise ValueError(
|
||||||
"CuteDSL MLA backend is only supported on Blackwell GPUs (SM100). Please use a different backend."
|
"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 [
|
if self.kv_cache_dtype not in [
|
||||||
"fp8_e4m3",
|
"fp8_e4m3",
|
||||||
"bf16",
|
"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."
|
"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)
|
run_post_process_pass(self, _attention_backend_fa3_fp8_fallback)
|
||||||
|
|
||||||
if (
|
run_post_process_pass(self, _fa4_page_constraint)
|
||||||
(
|
|
||||||
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
|
|
||||||
|
|
||||||
# AMD platforms backends
|
# AMD platforms backends
|
||||||
if self.attention_backend == "aiter":
|
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."
|
"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":
|
run_post_process_pass(self, _intel_xpu_page_constraint)
|
||||||
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
|
|
||||||
|
|
||||||
# Dual chunk flash attention backend
|
# Dual chunk flash attention backend
|
||||||
run_post_process_pass(self, _attention_backend_dual_chunk)
|
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.")
|
raise RuntimeError("KV4 is not tested on non-CUDA platforms.")
|
||||||
|
|
||||||
def _handle_page_size(self):
|
def _handle_page_size(self):
|
||||||
if self.page_size is None:
|
# Moved to the resolution pipeline (arg_groups/overrides.py:
|
||||||
# SHUFFLE 5D vectorized KV layout (aiter backend + pa_decode_gluon)
|
# _page_size_default), invoked here at its legacy slot.
|
||||||
# is tuned for and prefers page_size=64 — making it the default
|
from sglang.srt.arg_groups.overrides import (
|
||||||
# when the layout flag is set avoids users having to pass
|
_page_size_default,
|
||||||
# --page-size 64 explicitly. The env var is only consumed by the
|
run_post_process_pass,
|
||||||
# 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).
|
run_post_process_pass(self, _page_size_default)
|
||||||
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
|
|
||||||
|
|
||||||
def _handle_amd_specifics(self):
|
def _handle_amd_specifics(self):
|
||||||
if is_hip():
|
if is_hip():
|
||||||
@@ -6383,14 +6283,11 @@ class ServerArgs:
|
|||||||
self.disable_overlap_schedule = True
|
self.disable_overlap_schedule = True
|
||||||
|
|
||||||
if not self.disable_radix_cache:
|
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)
|
run_post_process_pass(self, _dllm_page_size)
|
||||||
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
|
|
||||||
if self.enable_hierarchical_cache:
|
if self.enable_hierarchical_cache:
|
||||||
logger.warning(
|
logger.warning(
|
||||||
"Hierarchical cache is disabled because of using diffusion LLM inference"
|
"Hierarchical cache is disabled because of using diffusion LLM inference"
|
||||||
|
|||||||
@@ -311,7 +311,7 @@ class TestFa4PageSizeAutoForce(CustomTestCase):
|
|||||||
args.model_config.hf_config.dual_chunk_attention_config = None
|
args.model_config.hf_config.dual_chunk_attention_config = None
|
||||||
return args
|
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)
|
@patch("sglang.srt.server_args.ServerArgs.use_mla_backend", return_value=False)
|
||||||
def test_combined_attention_backend_fa4_forces_page_size_128(
|
def test_combined_attention_backend_fa4_forces_page_size_128(
|
||||||
self, _mock_mla, _mock_sm100
|
self, _mock_mla, _mock_sm100
|
||||||
@@ -323,7 +323,7 @@ class TestFa4PageSizeAutoForce(CustomTestCase):
|
|||||||
|
|
||||||
self.assertEqual(args.page_size, 128)
|
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)
|
@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):
|
def test_explicit_prefill_fa4_forces_page_size_128(self, _mock_mla, _mock_sm100):
|
||||||
# `--prefill-attention-backend fa4`: the previously-covered path.
|
# `--prefill-attention-backend fa4`: the previously-covered path.
|
||||||
|
|||||||
@@ -69,6 +69,7 @@ class TestModelOverridableWhitelist(CustomTestCase):
|
|||||||
"disable_hybrid_swa_memory",
|
"disable_hybrid_swa_memory",
|
||||||
"sampling_backend",
|
"sampling_backend",
|
||||||
"attention_backend",
|
"attention_backend",
|
||||||
|
"page_size",
|
||||||
}
|
}
|
||||||
),
|
),
|
||||||
)
|
)
|
||||||
@@ -742,6 +743,189 @@ class TestGoldenModelOverrides(_IsolatedPublish):
|
|||||||
_dllm_attention_backend(_view(dllm_algorithm=None)), {}
|
_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):
|
def test_monolith_attention_families_at_callable_level(self):
|
||||||
from sglang.srt.arg_groups.overrides import (
|
from sglang.srt.arg_groups.overrides import (
|
||||||
_falcon_h1_jet_overrides,
|
_falcon_h1_jet_overrides,
|
||||||
|
|||||||
Reference in New Issue
Block a user