chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

refactor mamba radix cache logic in server_args (#17645)

Browse Source
This commit is contained in:
Yi Zhang
2026-01-26 17:02:49 +08:00
committed by GitHub
parent f6f1b6d000
commit 5844cb2fd8
1 changed files with 91 additions and 116 deletions
Download Patch File Download Diff File Expand all files Collapse all files

207
python/sglang/srt/server_args.py
View File

@@ -1449,14 +1449,11 @@ class ServerArgs:
f"Using {self.attention_backend} as attention backend for {model_arch}."
)
elif model_arch in ["KimiLinearForCausalLM"]:
logger.warning(
f"Disabling Radix Cache for {model_arch} as it is not yet supported."
self._handle_mamba_radix_cache(
model_arch=model_arch,
support_mamba_cache=False,
)
self.disable_radix_cache = True
elif model_arch in ["NemotronHForCausalLM"]:
assert (
not self.enable_mamba_extra_buffer()
), f"mamba extra_buffer is not supported for {model_arch} model"
model_config = self.get_model_config()
if model_config.quantization in [
"modelopt",
@@ -1475,30 +1472,11 @@ class ServerArgs:
self.quantization = model_config.quantization
self.moe_runner_backend = "flashinfer_cutlass"
if not self.disable_radix_cache and self.speculative_algorithm is not None:
logger.warning(
"Disabling radix cache since speculative decoding for NemotronHForCausalLM is not supported with radix cache yet."
)
self.disable_radix_cache = True
elif not self.disable_radix_cache:
logger.warning(
"Disabling overlap schedule since MambaRadixCache is not compatible with "
"overlap schedule currently, try to use --disable-radix-cache if overlap schedule is necessary"
)
self.disable_overlap_schedule = True
if is_sm100_supported():
if self.attention_backend is None:
self.attention_backend = "flashinfer"
logger.info(
"Use flashinfer as attention backend on sm100 for NemotronHForCausalLM"
)
if self.attention_backend == "trtllm_mha":
logger.warning(
"Disabling radix cache since trtllm_mha does not support page_size = 1, which is required by MambaRadixCache. "
"Try to use --attention-backend triton if radix cache is necessary."
)
self.disable_radix_cache = True
self.disable_overlap_schedule = False
self._handle_mamba_radix_cache(
model_arch=model_arch,
support_mamba_cache_extra_buffer=False,
sm100_default_attention_backend="flashinfer",
)
assert self.attention_backend != "triton", (
"NemotronHForCausalLM does not support triton attention backend,"
"as the first layer might not be an attention layer"
@@ -1541,21 +1519,12 @@ class ServerArgs:
logger.info(
"Use flashinfer_trtllm as MoE runner backend on sm100 for Qwen3NextForCausalLM"
)
if self.attention_backend is None:
self.attention_backend = "triton"
logger.info(
"Use triton as attention backend on sm100 for Qwen3NextForCausalLM"
)
if (
not self.disable_radix_cache
and self.attention_backend == "trtllm_mha"
):
logger.warning(
"Disabling radix cache since trtllm_mha does not support page_size = 1, which is required by MambaRadixCache. "
"Try to use --attention-backend triton if radix cache is necessary."
)
self.disable_radix_cache = True
self.disable_overlap_schedule = False
self._handle_mamba_radix_cache(
model_arch=model_arch,
support_mamba_cache_extra_buffer=True,
sm100_default_attention_backend="triton",
)
elif model_arch in ["Glm4MoeForCausalLM"]:
if is_sm100_supported():
quantization_config = getattr(hf_config, "quantization_config", None)
@@ -1578,83 +1547,23 @@ class ServerArgs:
"Use flashinfer_trtllm as MoE runner backend on sm100 for Glm4MoeForCausalLM"
)
# Mamba radix cache v2
if self.enable_mamba_extra_buffer():
assert (
is_cuda()
), "Mamba extra_buffer is only supported on CUDA devices with FLA backend"
if self.speculative_num_draft_tokens is not None:
assert (
self.mamba_track_interval >= self.speculative_num_draft_tokens
), f"mamba_track_interval {self.mamba_track_interval} must be greater than or equal to speculative_num_draft_tokens {self.speculative_num_draft_tokens}"
if self.page_size is not None:
assert (
self.mamba_track_interval % self.page_size == 0
), f"mamba_track_interval {self.mamba_track_interval} must be divisible by page_size {self.page_size}"
assert (
max(FLA_CHUNK_SIZE, self.page_size)
% min(FLA_CHUNK_SIZE, self.page_size)
== 0
), f"For SSM models with extra buffer, either FLA_CHUNK_SIZE or page_size must be divisible by the other, got {FLA_CHUNK_SIZE=}, {self.page_size=}"
elif not self.disable_radix_cache:
logger.warning(
"Disabling overlap schedule since MambaRadixCache no_buffer is not compatible with "
"overlap schedule currently, try to use --mamba-scheduler-strategy extra_buffer to enable overlap schedule"
)
self.disable_overlap_schedule = True
elif model_arch in [
"FalconH1ForCausalLM",
"JetNemotronForCausalLM",
"JetVLMForConditionalGeneration",
]:
assert (
not self.enable_mamba_extra_buffer()
), f"mamba extra_buffer is not supported for {model_arch} model"
if not self.disable_radix_cache:
logger.warning(
"Disabling overlap schedule since mamba no_buffer is not compatible with "
"overlap schedule, try to use --disable-radix-cache if overlap schedule is necessary"
)
self.disable_overlap_schedule = True
if is_sm100_supported():
if self.attention_backend is None:
self.attention_backend = "triton"
logger.info(
f"Use triton as attention backend on sm100 for {model_arch}"
)
if self.attention_backend == "trtllm_mha":
logger.warning(
"Disabling radix cache since trtllm_mha does not support page_size = 1, which is required by MambaRadixCache. "
"Try to use --attention-backend triton if radix cache is necessary."
)
self.disable_radix_cache = True
self.disable_overlap_schedule = False
self._handle_mamba_radix_cache(
model_arch=model_arch,
support_mamba_cache_extra_buffer=False,
sm100_default_attention_backend="triton",
)
elif model_arch in ["Lfm2ForCausalLM"]:
assert (
not self.enable_mamba_extra_buffer()
), f"mamba extra_buffer is not supported for {model_arch} model"
if not self.disable_radix_cache:
logger.warning(
"Disabling overlap schedule since mamba no_buffer is not compatible with "
"overlap schedule, try to use --disable-radix-cache if overlap schedule is necessary"
)
self.disable_overlap_schedule = True
if is_sm100_supported():
if self.attention_backend is None:
self.attention_backend = "flashinfer"
logger.info(
f"Use flashinfer as attention backend on sm100 for {model_arch}"
)
if self.attention_backend == "trtllm_mha":
logger.warning(
"Disabling radix cache since trtllm_mha does not support page_size = 1, which is required by MambaRadixCache. "
"Try to use --attention-backend flashinfer if radix cache is necessary."
)
self.disable_radix_cache = True
self.disable_overlap_schedule = False
self._handle_mamba_radix_cache(
model_arch=model_arch,
support_mamba_cache_extra_buffer=False,
sm100_default_attention_backend="flashinfer",
)
assert self.attention_backend != "triton", (
f"{model_arch} does not support triton attention backend, "
"as the first layer might not be an attention layer"
@@ -1692,6 +1601,72 @@ class ServerArgs:
):
self.enable_flashinfer_allreduce_fusion = True
def _handle_mamba_radix_cache(
self,
model_arch: str,
support_mamba_cache: bool = True,
support_mamba_cache_extra_buffer: bool = True,
sm100_default_attention_backend: str = None,
):
if (
is_sm100_supported()
and self.attention_backend is None
and sm100_default_attention_backend is not None
):
self.attention_backend = sm100_default_attention_backend
logger.info(
f"Use {sm100_default_attention_backend} as attention backend on sm100 for {model_arch}"
)
if not support_mamba_cache:
logger.warning(
f"Disabling Radix Cache for {model_arch} as it is not yet supported."
)
self.disable_radix_cache = True
return
if not support_mamba_cache_extra_buffer:
assert (
not self.enable_mamba_extra_buffer()
), f"mamba extra_buffer is not supported for {model_arch} model"
elif self.enable_mamba_extra_buffer(): # extra_buffer
assert (
is_cuda()
), "Mamba extra_buffer is only supported on CUDA devices with FLA backend"
if self.speculative_num_draft_tokens is not None:
assert (
self.mamba_track_interval >= self.speculative_num_draft_tokens
), f"mamba_track_interval {self.mamba_track_interval} must be greater than or equal to speculative_num_draft_tokens {self.speculative_num_draft_tokens}"
if self.page_size is not None:
assert (
self.mamba_track_interval % self.page_size == 0
), f"mamba_track_interval {self.mamba_track_interval} must be divisible by page_size {self.page_size}"
assert (
max(FLA_CHUNK_SIZE, self.page_size)
% min(FLA_CHUNK_SIZE, self.page_size)
== 0
), f"For SSM models with extra buffer, either FLA_CHUNK_SIZE or page_size must be divisible by the other, got {FLA_CHUNK_SIZE=}, {self.page_size=}"
elif not self.disable_radix_cache: # no_buffer
if self.speculative_algorithm is None:
logger.warning(
"Disabling overlap schedule since mamba no_buffer is not compatible with "
"overlap schedule, try to use --disable-radix-cache if overlap schedule is necessary"
)
self.disable_overlap_schedule = True
if self.attention_backend == "trtllm_mha":
logger.warning(
"Disabling radix cache since trtllm_mha does not support page_size = 1, which is required by MambaRadixCache. "
"Try to use --attention-backend triton if radix cache is necessary."
)
self.disable_radix_cache = True
self.disable_overlap_schedule = False
else:
logger.warning(
f"Disabling radix cache since speculative decoding for {model_arch} is not supported with radix cache yet."
)
self.disable_radix_cache = True
def _handle_sampling_backend(self):
if self.sampling_backend is None:
self.sampling_backend = (