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

Fix Mamba2-based models' default attention backend (#15117)

Browse Source 
Signed-off-by: Roi Koren <roik@nvidia.com>
This commit is contained in:
roikoren755
2025-12-14 21:08:48 +02:00
committed by GitHub
parent b11af135eb
commit 3f0482174a
1 changed files with 36 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -1286,6 +1286,29 @@ class ServerArgs:
else:
self.quantization = model_config.quantization
self.moe_runner_backend = "flashinfer_cutlass"
if 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
assert self.attention_backend != "triton", (
"NemotronHForCausalLM does not support triton attention backend,"
"as the first layer might not be an attention layer"
)
elif model_arch in [
"Qwen3MoeForCausalLM",
"Qwen3VLMoeForConditionalGeneration",
@@ -1350,7 +1373,6 @@ class ServerArgs:
self.disable_radix_cache = True
self.disable_overlap_schedule = False
elif model_arch in [
"NemotronHForCausalLM",
"FalconH1ForCausalLM",
"JetNemotronForCausalLM",
"JetVLMForConditionalGeneration",
@@ -1361,6 +1383,19 @@ class ServerArgs:
"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 = "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
def _handle_sampling_backend(self):
if self.sampling_backend is None: