Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 6 additions & 1 deletion python/sglang/srt/arg_groups/overrides.py
Original file line number Diff line number Diff line change
Expand Up @@ -416,7 +416,12 @@ def supports_mamba_cache_extra_buffer(view: Any, hf_config: Any) -> bool:
the configured linear-attention backend (pure read)."""
from sglang.srt.configs.linear_attn_model_registry import get_linear_attn_spec

if get_platform().is_xpu:
# Inkling has no no_buffer path (inkling.py asserts enable_mamba_extra_buffer),
# so rejecting it on XPU makes the model unlaunchable there.
if get_platform().is_xpu and hf_config.architectures[0] not in (
"InklingForConditionalGeneration",
"InklingForConditionalGenerationMTP",
):
return False
spec = get_linear_attn_spec(hf_config)
if hf_config.architectures[0] in _MAMBA_EXTRA_BUFFER_ARCHS or (
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
from unittest.mock import patch

from sglang.srt.arg_groups.mamba_hook import validate_mamba_extra_buffer
from sglang.srt.arg_groups.overrides import supports_mamba_cache_extra_buffer
from sglang.srt.runtime_context import override_platform
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase
Expand Down Expand Up @@ -63,5 +64,32 @@ def test_out_of_tree_validation_uses_platform_capability(
_validate()


def _supports(arch: str) -> bool:
return supports_mamba_cache_extra_buffer(
SimpleNamespace(linear_attn_backend="triton"),
SimpleNamespace(architectures=[arch]),
)


class TestMambaExtraBufferXpuArchs(CustomTestCase):
def test_inkling_is_admitted_on_xpu(self):
"""Inkling has no no_buffer path, so rejecting it on XPU fails every launch."""
with override_platform(is_xpu=True):
for arch in (
"InklingForConditionalGeneration",
"InklingForConditionalGenerationMTP",
):
with self.subTest(arch=arch):
self.assertTrue(_supports(arch))

def test_other_archs_stay_rejected_on_xpu(self):
for arch in ("KimiK3ForConditionalGeneration", "NemotronHForCausalLM"):
with self.subTest(arch=arch):
with override_platform(is_xpu=False):
self.assertTrue(_supports(arch))
with override_platform(is_xpu=True):
self.assertFalse(_supports(arch))


if __name__ == "__main__":
unittest.main()
Loading