From 82a0bafc1c3a0680b6389f5ae9c48def7800dcbd Mon Sep 17 00:00:00 2001 From: shaharmor98 <17088876+shaharmor98@users.noreply.github.com> Date: Thu, 19 Feb 2026 10:56:06 +0200 Subject: [PATCH] Feat/add fi selective state update kernel call (#18070) Signed-off-by: Shahar Mor --- .../layers/attention/mamba/ops/__init__.py | 13 +- .../attention/mamba/ops/ssu_dispatch.py | 277 ++++++++++++++++++ python/sglang/srt/managers/scheduler.py | 9 + .../sglang/srt/model_executor/model_runner.py | 3 + python/sglang/srt/server_args.py | 28 ++ test/registered/layers/mamba/conftest.py | 19 ++ .../models/test_nvidia_nemotron_3_nano.py | 13 + test/run_suite.py | 6 +- 8 files changed, 366 insertions(+), 2 deletions(-) create mode 100644 python/sglang/srt/layers/attention/mamba/ops/ssu_dispatch.py create mode 100644 test/registered/layers/mamba/conftest.py diff --git a/python/sglang/srt/layers/attention/mamba/ops/__init__.py b/python/sglang/srt/layers/attention/mamba/ops/__init__.py index 809ff36fb..6496105ae 100644 --- a/python/sglang/srt/layers/attention/mamba/ops/__init__.py +++ b/python/sglang/srt/layers/attention/mamba/ops/__init__.py @@ -1,2 +1,13 @@ -from .mamba_ssm import selective_state_update +from .mamba_ssm import PAD_SLOT_ID from .ssd_combined import mamba_chunk_scan_combined +from .ssu_dispatch import ( + initialize_mamba_selective_state_update_backend, + selective_state_update, +) + +__all__ = [ + "PAD_SLOT_ID", + "selective_state_update", + "mamba_chunk_scan_combined", + "initialize_mamba_selective_state_update_backend", +] diff --git a/python/sglang/srt/layers/attention/mamba/ops/ssu_dispatch.py b/python/sglang/srt/layers/attention/mamba/ops/ssu_dispatch.py new file mode 100644 index 000000000..77d586eee --- /dev/null +++ b/python/sglang/srt/layers/attention/mamba/ops/ssu_dispatch.py @@ -0,0 +1,277 @@ +from __future__ import annotations + +import logging +from abc import ABC, abstractmethod +from typing import TYPE_CHECKING + +import torch + +if TYPE_CHECKING: + from sglang.srt.server_args import ServerArgs + +logger = logging.getLogger(__name__) + + +class MambaSSUBackend(ABC): + @property + @abstractmethod + def name(self) -> str: + """Human-readable name used for logging.""" + + @abstractmethod + def __call__( + self, + state: torch.Tensor, + x: torch.Tensor, + dt: torch.Tensor, + A: torch.Tensor, + B: torch.Tensor, + C: torch.Tensor, + D: torch.Tensor | None = None, + z: torch.Tensor | None = None, + dt_bias: torch.Tensor | None = None, + dt_softplus: bool = False, + state_batch_indices: torch.Tensor | None = None, + pad_slot_id: int = -1, + out: torch.Tensor | None = None, + disable_state_update: bool = False, + intermediate_states_buffer: torch.Tensor | None = None, + cache_steps: int | None = None, + retrieve_parent_token: torch.Tensor | None = None, + intermediate_state_indices: torch.Tensor | None = None, + ) -> None: ... + + +class TritonSSUBackend(MambaSSUBackend): + """Triton-based selective-state-update backend.""" + + def __init__(self) -> None: + from sglang.srt.layers.attention.mamba.ops.mamba_ssm import ( + selective_state_update, + ) + + self._kernel = selective_state_update + + @property + def name(self) -> str: + return "triton" + + def __call__( + self, + state: torch.Tensor, + x: torch.Tensor, + dt: torch.Tensor, + A: torch.Tensor, + B: torch.Tensor, + C: torch.Tensor, + D: torch.Tensor | None = None, + z: torch.Tensor | None = None, + dt_bias: torch.Tensor | None = None, + dt_softplus: bool = False, + state_batch_indices: torch.Tensor | None = None, + pad_slot_id: int = -1, + out: torch.Tensor | None = None, + disable_state_update: bool = False, + intermediate_states_buffer: torch.Tensor | None = None, + cache_steps: int | None = None, + retrieve_parent_token: torch.Tensor | None = None, + intermediate_state_indices: torch.Tensor | None = None, + ) -> None: + self._kernel( + state, + x, + dt, + A, + B, + C, + D=D, + z=z, + dt_bias=dt_bias, + dt_softplus=dt_softplus, + state_batch_indices=state_batch_indices, + pad_slot_id=pad_slot_id, + out=out, + disable_state_update=disable_state_update, + intermediate_states_buffer=intermediate_states_buffer, + cache_steps=cache_steps, + retrieve_parent_token=retrieve_parent_token, + intermediate_state_indices=intermediate_state_indices, + ) + + +class FlashInferSSUBackend(MambaSSUBackend): + """FlashInfer-based selective-state-update backend.""" + + def __init__(self) -> None: + from flashinfer.mamba import selective_state_update + + self._kernel = selective_state_update + + @property + def name(self) -> str: + return "flashinfer" + + def __call__( + self, + state: torch.Tensor, + x: torch.Tensor, + dt: torch.Tensor, + A: torch.Tensor, + B: torch.Tensor, + C: torch.Tensor, + D: torch.Tensor | None = None, + z: torch.Tensor | None = None, + dt_bias: torch.Tensor | None = None, + dt_softplus: bool = False, + state_batch_indices: torch.Tensor | None = None, + pad_slot_id: int = -1, + out: torch.Tensor | None = None, + disable_state_update: bool = False, + intermediate_states_buffer: torch.Tensor | None = None, + cache_steps: int | None = None, + retrieve_parent_token: torch.Tensor | None = None, + intermediate_state_indices: torch.Tensor | None = None, + ) -> None: + if retrieve_parent_token is not None: + raise ValueError( + "FlashInfer backend does not support retrieve_parent_token. " + "Use --mamba-backend triton for EAGLE tree attention." + ) + # FlashInfer expects cache_steps as an int (0 when unused). + self._kernel( + state, + x, + dt, + A, + B, + C, + D=D, + z=z, + dt_bias=dt_bias, + dt_softplus=dt_softplus, + state_batch_indices=state_batch_indices, + pad_slot_id=pad_slot_id, + out=out, + disable_state_update=disable_state_update, + intermediate_states_buffer=intermediate_states_buffer, + cache_steps=0 if cache_steps is None else cache_steps, + intermediate_state_indices=intermediate_state_indices, + ) + + +_BACKEND_REGISTRY: dict[str, type[MambaSSUBackend]] = { + "triton": TritonSSUBackend, + "flashinfer": FlashInferSSUBackend, +} + +_mamba_ssu_backend: MambaSSUBackend | None = None + + +def initialize_mamba_selective_state_update_backend(server_args: ServerArgs) -> None: + """Instantiate the selective-state-update backend from server config. + + This should be called once during scheduler initialization. + + Args: + server_args: Server arguments containing ``mamba_backend`` setting. + + Raises: + ValueError: If the requested backend is unavailable or cannot be imported. + """ + global _mamba_ssu_backend + + requested = server_args.mamba_backend or "triton" + + backend_cls = _BACKEND_REGISTRY.get(requested) + if backend_cls is None: + raise ValueError( + f"Unknown mamba backend '{requested}'. " + f"Available backends: {list(_BACKEND_REGISTRY.keys())}" + ) + + try: + _mamba_ssu_backend = backend_cls() + except ImportError: + raise ValueError( + f"Mamba backend '{requested}' requested but its dependencies are not " + f"available. Install the required package or use a different " + f"--mamba-backend value." + ) + + logger.info( + "Mamba selective_state_update backend initialized: %s", + _mamba_ssu_backend.name, + ) + + +def selective_state_update( + state: torch.Tensor, + x: torch.Tensor, + dt: torch.Tensor, + A: torch.Tensor, + B: torch.Tensor, + C: torch.Tensor, + D: torch.Tensor | None = None, + z: torch.Tensor | None = None, + dt_bias: torch.Tensor | None = None, + dt_softplus: bool = False, + state_batch_indices: torch.Tensor | None = None, + pad_slot_id: int = -1, + out: torch.Tensor | None = None, + disable_state_update: bool = False, + intermediate_states_buffer: torch.Tensor | None = None, + cache_steps: int | None = None, + retrieve_parent_token: torch.Tensor | None = None, + intermediate_state_indices: torch.Tensor | None = None, +) -> None: + """Dispatch selective-state-update to the configured backend. + + This function provides a unified interface regardless of the underlying + backend. Backend-specific argument adaptation is handled inside each + :class:`MambaSSUBackend` subclass. + + Args: + state: SSM state tensor (batch, nheads, dim, dstate) + x: Input tensor + dt: Delta time tensor + A: A matrix + B: B matrix + C: C matrix + D: Optional D vector + z: Optional z tensor for gating + dt_bias: Optional dt bias + dt_softplus: Whether to apply softplus to dt + state_batch_indices: Optional batch indices for state + out: Preallocated output tensor (in-place updated) + disable_state_update: If True, don't write back to state (for speculative verify) + intermediate_states_buffer: Buffer to cache intermediate states + cache_steps: Total number of steps in the buffer + retrieve_parent_token: (batch, T) tensor of parent token indices for EAGLE tree attention + intermediate_state_indices: (batch,) tensor of indices for intermediate_states_buffer operations. + If provided, uses these indices instead of state_batch_indices for the buffer. + """ + assert _mamba_ssu_backend is not None, ( + "Mamba selective_state_update backend not initialized. " + "Call initialize_mamba_selective_state_update_backend() first." + ) + + _mamba_ssu_backend( + state, + x, + dt, + A, + B, + C, + D=D, + z=z, + dt_bias=dt_bias, + dt_softplus=dt_softplus, + state_batch_indices=state_batch_indices, + pad_slot_id=pad_slot_id, + out=out, + disable_state_update=disable_state_update, + intermediate_states_buffer=intermediate_states_buffer, + cache_steps=cache_steps, + retrieve_parent_token=retrieve_parent_token, + intermediate_state_indices=intermediate_state_indices, + ) diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 3b7153143..b00f7765f 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -61,6 +61,9 @@ from sglang.srt.distributed.parallel_state import get_tp_group from sglang.srt.dllm.mixin.scheduler import SchedulerDllmMixin from sglang.srt.environ import envs from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder +from sglang.srt.layers.attention.mamba.ops import ( + initialize_mamba_selective_state_update_backend, +) from sglang.srt.layers.dp_attention import ( compute_dp_attention_world_info, get_attention_cp_group, @@ -356,6 +359,9 @@ class Scheduler( # Init moe config and GEMM config (FP8 GEMM, etc.) self.init_moe_gemm_config() + # Init mamba backend + self.init_mamba_backend() + # Launch a model worker and draft model worker if using speculative decoding self.init_model_worker() @@ -489,6 +495,9 @@ class Scheduler( reasoning_parser.detector.think_end_token, add_special_tokens=False )[0] + def init_mamba_backend(self) -> None: + initialize_mamba_selective_state_update_backend(self.server_args) + def init_moe_gemm_config(self): # For the MM models, check the text_config for MoE settings config_to_check = getattr( diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index eb4a4c18a..094b1d317 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -1831,6 +1831,9 @@ class ModelRunner(ModelRunnerKVCacheMixin): return False backend_str = self.server_args.moe_runner_backend + + # TODO smor- support other cases for flashinfer autotune, such as, mamba backend + if backend_str not in [ "flashinfer_trtllm", "flashinfer_mxfp4", diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 89688d34d..88cbcf1d5 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -215,6 +215,8 @@ MAMBA_SSM_DTYPE_CHOICES = ["float32", "bfloat16", "float16"] MAMBA_SCHEDULER_STRATEGY_CHOICES = ["auto", "no_buffer", "extra_buffer"] +MAMBA_BACKEND_CHOICES = ["triton", "flashinfer"] + # Allow external code to add more choices def add_load_format_choices(choices): @@ -459,6 +461,7 @@ class ServerArgs: None # auto-detect based on hardware/kv_cache_dtype ) disable_flashinfer_autotune: bool = False + mamba_backend: str = "triton" # Speculative decoding speculative_algorithm: Optional[str] = None @@ -735,6 +738,7 @@ class ServerArgs: # Set kernel backends. self._handle_sampling_backend() self._handle_attention_backend_compatibility() + self._handle_mamba_backend() self._handle_kv4_compatibility() self._handle_page_size() self._handle_amd_specifics() @@ -2059,6 +2063,22 @@ class ServerArgs: if self.grammar_backend is None: self.grammar_backend = "xgrammar" + def _handle_mamba_backend(self): + if self.mamba_backend == "flashinfer": + if is_flashinfer_available(): + try: + import flashinfer.mamba # noqa: F401 + + logger.info("Successfully imported FlashInfer mamba module") + except (ImportError, AttributeError): + raise ValueError( + "FlashInfer mamba module not available, please check flashinfer installation." + ) + else: + raise ValueError( + "FlashInfer mamba module not available, please check flashinfer installation." + ) + def _handle_context_parallelism(self): if self.attn_cp_size > 1: # The tp_size is the world size, not the real tensor parallel size @@ -4220,6 +4240,14 @@ class ServerArgs: default=ServerArgs.mamba_track_interval, help="The interval to track the mamba state during decode.", ) + parser.add_argument( + "--mamba-backend", + type=str, + choices=MAMBA_BACKEND_CHOICES, + default=ServerArgs.mamba_backend, + help="Choose the kernel backend for Mamba SSM operations. Default is 'triton'. " + "Options: 'triton' (default), 'flashinfer' (requires FlashInfer with Mamba support).", + ) # Hierarchical cache parser.add_argument( diff --git a/test/registered/layers/mamba/conftest.py b/test/registered/layers/mamba/conftest.py new file mode 100644 index 000000000..606a5ee1a --- /dev/null +++ b/test/registered/layers/mamba/conftest.py @@ -0,0 +1,19 @@ +import pytest + +from sglang.srt.layers.attention.mamba.ops import ssu_dispatch +from sglang.srt.layers.attention.mamba.ops.ssu_dispatch import ( + initialize_mamba_selective_state_update_backend, +) +from sglang.srt.server_args import ServerArgs + + +@pytest.fixture(scope="session", autouse=True) +def _init_mamba_ssu_backend(): + """Initialize the Mamba SSU dispatch backend for the test session. + + In production this happens in Scheduler.init_mamba_backend(). Tests have no + scheduler, so we do it here via the same public API. + """ + initialize_mamba_selective_state_update_backend(ServerArgs(model_path="dummy")) + yield + ssu_dispatch._mamba_ssu_backend = None diff --git a/test/registered/models/test_nvidia_nemotron_3_nano.py b/test/registered/models/test_nvidia_nemotron_3_nano.py index abd22e3f2..2129abca3 100644 --- a/test/registered/models/test_nvidia_nemotron_3_nano.py +++ b/test/registered/models/test_nvidia_nemotron_3_nano.py @@ -26,6 +26,19 @@ class TestNvidiaNemotron3Nano30BBF16(LMEvalMixin, DefaultServerBase): ] + NEMOTRON_3_NANO_THINKING_ARGS +class TestNvidiaNemotron3Nano30BBF16FlashInfer(LMEvalMixin, DefaultServerBase): + """Test Nemotron-3-Nano-30B BF16 model with lm-eval GSM8K evaluation using flashinfer mamba backend.""" + + model = "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16" + model_config_name = "lm_eval_configs/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16.yaml" + other_args = [ + "--tp-size", + "2", + "--mamba-backend", + "flashinfer", + ] + NEMOTRON_3_NANO_THINKING_ARGS + + class TestNvidiaNemotron3Nano30BFP8(LMEvalMixin, DefaultServerBase): """Test Nemotron-3-Nano-30B FP8 model with lm-eval GSM8K evaluation.""" diff --git a/test/run_suite.py b/test/run_suite.py index 4d6a56471..93c1150b4 100644 --- a/test/run_suite.py +++ b/test/run_suite.py @@ -188,7 +188,11 @@ def run_a_suite(args): auto_partition_size = args.auto_partition_size # All tests (per-commit and nightly) are now in registered/ - files = glob.glob("registered/**/*.py", recursive=True) + files = [ + f + for f in glob.glob("registered/**/*.py", recursive=True) + if not f.endswith("/conftest.py") + ] # Strict: all registered files must have proper registration sanity_check = True