Feat/add fi selective state update kernel call (#18070)

Signed-off-by: Shahar Mor <smor@nvidia.com>

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",
+]

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,
+    )

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(

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",

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(

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

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."""
 

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