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

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

Browse Source 
Signed-off-by: Shahar Mor <smor@nvidia.com>
This commit is contained in:
shaharmor98
2026-02-19 10:56:06 +02:00
committed by GitHub
parent 0be30d4b0d
commit 82a0bafc1c
8 changed files with 366 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
python/sglang/srt/layers/attention/mamba/ops/__init__.py
View File

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

277
python/sglang/srt/layers/attention/mamba/ops/ssu_dispatch.py Normal file
View File

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

9
python/sglang/srt/managers/scheduler.py
View File

@@ -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(

3
python/sglang/srt/model_executor/model_runner.py
View File

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

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

@@ -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(