Implementation:Sgl project Sglang Mamba Ops

Knowledge Sources	Sgl_project_Sglang
Domains	Kernel, State Space Models, Convolution
Last Updated	2026-02-10 00:00 GMT

Overview

Python interface for Mamba state-space model causal convolution and gated delta rule kernels, supporting both GPU and CPU execution.

Description

The mamba.py module provides GPU and CPU variants of causal 1D convolution operations for Mamba/SSM model inference. causal_conv1d_fwd performs the forward pass of causal 1D convolution on GPU with support for variable-length sequences via query_start_loc, conv state management via conv_states and cache_indices, initial state handling, optional SiLU activation, and pad slot configuration. causal_conv1d_update handles single-step state updates during decode on GPU, updating the convolution state and producing a single output step. causal_conv1d_fn_cpu and causal_conv1d_update_cpu are CPU equivalents that accept similar parameters with slight differences (e.g., seq_lens_cpu for the CPU forward, string-based activation specification). chunk_gated_delta_rule_cpu computes chunked gated delta rule attention on CPU, taking query, key, value, gate, beta, initial state, and cumulative sequence lengths, returning the core attention output, last recurrent state, and a placeholder for future hidden state support.

Usage

Use the GPU variants (causal_conv1d_fwd, causal_conv1d_update) for Mamba model inference on CUDA devices during prefill and decode phases respectively. Use the CPU variants for CPU-based inference. Use chunk_gated_delta_rule_cpu for gated delta rule linear attention on CPU.

Code Reference

Source Location

• Repository: Sgl_project_Sglang
• File: sgl-kernel/python/sgl_kernel/mamba.py
• Lines: 1-121

Signature

def causal_conv1d_fwd(
    x: torch.Tensor,
    weight: torch.Tensor,
    bias_: Optional[torch.Tensor],
    conv_states: Optional[torch.Tensor],
    query_start_loc: Optional[torch.Tensor],
    cache_indices: Optional[torch.Tensor],
    has_initial_state: Optional[torch.Tensor],
    silu_activation: bool,
    pad_slot_id: int,
): ...

def causal_conv1d_update(
    x: torch.Tensor,
    conv_state: torch.Tensor,
    weight: torch.Tensor,
    bias_: Optional[torch.Tensor],
    silu_activation: bool,
    cache_seqlens: Optional[torch.Tensor],
    conv_state_indices: Optional[torch.Tensor],
    pad_slot_id: int,
): ...

def causal_conv1d_fn_cpu(
    mixed_qkv_transposed, conv_weights, bias, activation,
    conv_states, has_initial_state, cache_indices,
    query_start_loc, seq_lens_cpu,
): ...

def causal_conv1d_update_cpu(
    mixed_qkv, conv_states, conv_weights, bias,
    activation, conv_state_indices,
): ...

def chunk_gated_delta_rule_cpu(
    q, k, v, g, beta, initial_state,
    cu_seqlens, head_first, use_qk_l2norm_in_kernel,
): ...

Import

from sgl_kernel.mamba import (
    causal_conv1d_fwd,
    causal_conv1d_update,
    causal_conv1d_fn_cpu,
    causal_conv1d_update_cpu,
    chunk_gated_delta_rule_cpu,
)

I/O Contract

Inputs

Name	Type	Required	Description
x	torch.Tensor	Yes	Input tensor for convolution
weight	torch.Tensor	Yes	Convolution weight tensor
bias_	Optional[torch.Tensor]	No	Optional bias tensor
conv_states	Optional[torch.Tensor]	No	Convolution state buffer for caching
query_start_loc	Optional[torch.Tensor]	No	Start locations for variable-length sequences
cache_indices	Optional[torch.Tensor]	No	Indices into conv state cache
has_initial_state	Optional[torch.Tensor]	No	Boolean flags indicating initial state presence
silu_activation	bool	Yes	Whether to apply SiLU activation
pad_slot_id	int	Yes	Padding slot identifier
cache_seqlens	Optional[torch.Tensor]	No	Cached sequence lengths (update only)
conv_state_indices	Optional[torch.Tensor]	No	State indices for update
q, k, v	torch.Tensor	Yes (delta rule)	Query, key, value tensors
g	torch.Tensor	Yes (delta rule)	Gate tensor
beta	torch.Tensor	Yes (delta rule)	Beta scaling tensor
initial_state	torch.Tensor	Yes (delta rule)	Initial recurrent state
cu_seqlens	torch.Tensor	Yes (delta rule)	Cumulative sequence lengths

Outputs

Name	Type	Description
(in-place)	-	causal_conv1d_fwd and causal_conv1d_update modify x and conv_states in-place
core_attn_out	torch.Tensor	Attention output from chunk_gated_delta_rule_cpu
last_recurrent_state	torch.Tensor	Final recurrent state from chunk_gated_delta_rule_cpu
h	None	Placeholder for future hidden state support

Usage Examples

from sgl_kernel.mamba import causal_conv1d_fwd, causal_conv1d_update

# Prefill: forward pass with variable-length sequences
causal_conv1d_fwd(
    x, weight, bias, conv_states,
    query_start_loc, cache_indices,
    has_initial_state, silu_activation=True,
    pad_slot_id=-1
)

# Decode: single-step state update
causal_conv1d_update(
    x, conv_state, weight, bias,
    silu_activation=True,
    cache_seqlens=seqlens,
    conv_state_indices=indices,
    pad_slot_id=-1
)

Related Pages

• Environment:Sgl_project_Sglang_CUDA_GPU_Runtime