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

refactor linear attention backend (#18622)

Browse Source 
Co-authored-by: yizhang2077 <1109276519@qq.com>
This commit is contained in:
Yuhao Yang
2026-02-25 23:02:44 +08:00
committed by GitHub
co-authored by yizhang2077
parent 471acd98b9
commit c7c4a1cbbd
14 changed files with 1453 additions and 843 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/environ.py
-2
View File
@@ -185,8 +185,6 @@ class Envs:
SGLANG_GRAMMAR_MAX_POLL_ITERATIONS = EnvInt(10000) SGLANG_GRAMMAR_MAX_POLL_ITERATIONS = EnvInt(10000)
SGLANG_DISABLE_OUTLINES_DISK_CACHE = EnvBool(False) SGLANG_DISABLE_OUTLINES_DISK_CACHE = EnvBool(False)
# CuTe DSL GDN Decode
SGLANG_USE_CUTEDSL_GDN_DECODE = EnvBool(False)
# Test & Debug # Test & Debug
SGLANG_DETECT_SLOW_RANK = EnvBool(False) SGLANG_DETECT_SLOW_RANK = EnvBool(False)
python/sglang/srt/layers/attention/attention_registry.py
+10 -4
View File
@@ -189,15 +189,21 @@ def attn_backend_wrapper(runner: "ModelRunner", full_attn_backend: "AttentionBac
if cfg := runner.mambaish_config: if cfg := runner.mambaish_config:
from sglang.srt.layers.attention.fla.utils import check_environments from sglang.srt.layers.attention.fla.utils import check_environments
from sglang.srt.layers.attention.hybrid_linear_attn_backend import ( from sglang.srt.layers.attention.hybrid_linear_attn_backend import (
GDNAttnBackend,
HybridLinearAttnBackend, HybridLinearAttnBackend,
KimiLinearAttnBackend,
LightningAttentionBackend,
Mamba2AttnBackend, Mamba2AttnBackend,
) )
from sglang.srt.layers.attention.linear.gdn_backend import GDNAttnBackend
from sglang.srt.layers.attention.linear.kda_backend import KDAAttnBackend
from sglang.srt.layers.attention.linear.lightning_backend import (
LightningAttentionBackend,
)
from sglang.srt.layers.attention.linear.utils import (
initialize_linear_attn_config,
)
from sglang.srt.utils import is_blackwell, is_npu from sglang.srt.utils import is_blackwell, is_npu
check_environments() check_environments()
initialize_linear_attn_config(runner.server_args)
if runner.hybrid_gdn_config is not None: if runner.hybrid_gdn_config is not None:
if is_blackwell(): if is_blackwell():
assert ( assert (
@@ -213,7 +219,7 @@ def attn_backend_wrapper(runner: "ModelRunner", full_attn_backend: "AttentionBac
elif runner.mamba2_config is not None: elif runner.mamba2_config is not None:
linear_attn_backend = Mamba2AttnBackend(runner) linear_attn_backend = Mamba2AttnBackend(runner)
elif runner.kimi_linear_config is not None: elif runner.kimi_linear_config is not None:
linear_attn_backend = KimiLinearAttnBackend(runner) linear_attn_backend = KDAAttnBackend(runner)
elif runner.hybrid_lightning_config is not None: elif runner.hybrid_lightning_config is not None:
linear_attn_backend = LightningAttentionBackend(runner) linear_attn_backend = LightningAttentionBackend(runner)
else: else:
python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py
+3 -837
View File
@@ -1,32 +1,12 @@
import logging import logging
import math
from typing import Optional, Union from typing import Optional, Union
import torch import torch
import triton import triton
import triton.language as tl import triton.language as tl
from einops import rearrange
from sglang.srt.environ import Envs
from sglang.srt.layers.attention.base_attn_backend import AttentionBackend from sglang.srt.layers.attention.base_attn_backend import AttentionBackend
from sglang.srt.layers.attention.fla.fused_gdn_gating import fused_gdn_gating from sglang.srt.layers.attention.mamba.causal_conv1d_triton import PAD_SLOT_ID
from sglang.srt.layers.attention.fla.fused_recurrent import (
fused_recurrent_gated_delta_rule_update,
)
from sglang.srt.layers.attention.fla.fused_sigmoid_gating_recurrent import (
fused_sigmoid_gating_delta_rule_update,
)
from sglang.srt.layers.attention.linear.lightning_attn import (
BailingLinearKernel,
linear_decode_forward_triton,
)
from sglang.srt.layers.attention.linear.linear_metadata import BailingLinearMetadata
from sglang.srt.layers.attention.linear.seg_la import SegLaMeta, seg_la_fwd
from sglang.srt.layers.attention.mamba.causal_conv1d_triton import (
PAD_SLOT_ID,
causal_conv1d_fn,
causal_conv1d_update,
)
from sglang.srt.layers.attention.mamba.mamba import MambaMixer2 from sglang.srt.layers.attention.mamba.mamba import MambaMixer2
from sglang.srt.layers.attention.mamba.mamba2_metadata import ( from sglang.srt.layers.attention.mamba.mamba2_metadata import (
ForwardMetadata, ForwardMetadata,
@@ -36,69 +16,18 @@ from sglang.srt.layers.attention.mamba.mamba_state_scatter_triton import (
fused_mamba_state_scatter_with_mask, fused_mamba_state_scatter_with_mask,
) )
from sglang.srt.layers.radix_attention import RadixAttention from sglang.srt.layers.radix_attention import RadixAttention
from sglang.srt.layers.radix_linear_attention import RadixLinearAttention from sglang.srt.mem_cache.memory_pool import HybridReqToTokenPool
from sglang.srt.mem_cache.memory_pool import HybridReqToTokenPool, MambaPool
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
from sglang.srt.model_executor.model_runner import ModelRunner from sglang.srt.model_executor.model_runner import ModelRunner
from sglang.srt.server_args import get_global_server_args from sglang.srt.server_args import get_global_server_args
from sglang.srt.speculative.eagle_info import EagleDraftInput, EagleVerifyInput from sglang.srt.speculative.eagle_info import EagleDraftInput, EagleVerifyInput
from sglang.srt.speculative.spec_info import SpecInput from sglang.srt.speculative.spec_info import SpecInput
from sglang.srt.utils import cpu_has_amx_support, is_cpu, is_cuda, is_npu from sglang.srt.utils import is_cpu
from sglang.srt.utils.common import rank0_log
if not is_cpu(): if not is_cpu():
# fix import error on CPU device, no impacts when non-CPU path
try:
from sglang.jit_kernel.cutedsl_gdn import (
cutedsl_fused_sigmoid_gating_delta_rule_update,
)
except ModuleNotFoundError:
# CuTe DSL path requires cuda-python (cuda.bindings.*). Keep runtime usable
# by falling back to non-CuTe kernels when it's unavailable.
cutedsl_fused_sigmoid_gating_delta_rule_update = None
from sglang.srt.layers.attention.fla.chunk import chunk_gated_delta_rule
from sglang.srt.layers.attention.fla.chunk_delta_h import ( from sglang.srt.layers.attention.fla.chunk_delta_h import (
CHUNK_SIZE as FLA_CHUNK_SIZE, CHUNK_SIZE as FLA_CHUNK_SIZE,
) )
from sglang.srt.layers.attention.fla.kda import chunk_kda
if is_cuda():
from sglang.srt.layers.attention.mamba.causal_conv1d import (
causal_conv1d_fn as causal_conv1d_fn_cuda,
)
causal_conv1d_fn = causal_conv1d_fn_cuda
elif is_npu():
from sgl_kernel_npu.fla.chunk import chunk_gated_delta_rule_npu
from sgl_kernel_npu.fla.fused_sigmoid_gating_recurrent import (
fused_sigmoid_gating_delta_rule_update_npu,
)
from sgl_kernel_npu.mamba.causal_conv1d import (
causal_conv1d_fn_npu,
causal_conv1d_update_npu,
)
chunk_gated_delta_rule = chunk_gated_delta_rule_npu
fused_sigmoid_gating_delta_rule_update = fused_sigmoid_gating_delta_rule_update_npu
causal_conv1d_fn = causal_conv1d_fn_npu
causal_conv1d_update = causal_conv1d_update_npu
elif is_cpu():
assert (
cpu_has_amx_support()
), "CPU requires AMX support for hybrid linear attn backend"
from sgl_kernel.mamba import (
causal_conv1d_fn_cpu,
causal_conv1d_update_cpu,
chunk_gated_delta_rule_cpu,
)
chunk_gated_delta_rule = chunk_gated_delta_rule_cpu
causal_conv1d_fn = causal_conv1d_fn_cpu
causal_conv1d_update = causal_conv1d_update_cpu
fused_sigmoid_gating_delta_rule_update = (
torch.ops.sgl_kernel.fused_sigmoid_gating_delta_rule_update_cpu
)
fused_gdn_gating = torch.ops.sgl_kernel.fused_gdn_gating_cpu
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@@ -661,410 +590,6 @@ class MambaAttnBackendBase(AttentionBackend):
] ]
class KimiLinearAttnBackend(MambaAttnBackendBase):
"""Attention backend using Mamba kernel."""
def forward_decode(
self,
layer: RadixLinearAttention,
mixed_qkv: torch.Tensor,
a: torch.Tensor,
b: torch.Tensor,
**kwargs,
):
q_proj_states, k_proj_states, v_proj_states = torch.split(
mixed_qkv,
[layer.q_dim, layer.k_dim, layer.v_dim],
dim=-1,
)
q_conv_weights, k_conv_weights, v_conv_weights = layer.conv_weights
q_conv_bias, k_conv_bias, v_conv_bias = layer.bias
layer_cache = self.req_to_token_pool.mamba2_layer_cache(layer.layer_id)
q_conv_state, k_conv_state, v_conv_state = layer_cache.conv
ssm_states = layer_cache.temporal
query_start_loc = self.forward_metadata.query_start_loc
cache_indices = self.forward_metadata.mamba_cache_indices
q_conv_state = q_conv_state.transpose(-1, -2)
k_conv_state = k_conv_state.transpose(-1, -2)
v_conv_state = v_conv_state.transpose(-1, -2)
q = causal_conv1d_update(
q_proj_states,
q_conv_state,
q_conv_weights,
q_conv_bias,
activation="silu",
conv_state_indices=cache_indices,
)
k = causal_conv1d_update(
k_proj_states,
k_conv_state,
k_conv_weights,
k_conv_bias,
activation="silu",
conv_state_indices=cache_indices,
)
v = causal_conv1d_update(
v_proj_states,
v_conv_state,
v_conv_weights,
v_conv_bias,
activation="silu",
conv_state_indices=cache_indices,
)
q = rearrange(q, "n (h d) -> 1 n h d", d=layer.head_q_dim)
k = rearrange(k, "n (h d) -> 1 n h d", d=layer.head_k_dim)
v = rearrange(v, "n (h d) -> 1 n h d", d=layer.head_v_dim)
return fused_sigmoid_gating_delta_rule_update(
A_log=layer.A_log,
dt_bias=layer.dt_bias,
q=q,
k=k,
v=v,
a=a,
b=b,
initial_state_source=ssm_states,
initial_state_indices=cache_indices,
cu_seqlens=query_start_loc,
use_qk_l2norm_in_kernel=True,
softplus_beta=1.0,
softplus_threshold=20.0,
is_kda=True,
)
def forward_extend(
self,
layer: RadixLinearAttention,
forward_batch: ForwardBatch,
mixed_qkv: torch.Tensor,
a: torch.Tensor,
b: torch.Tensor,
**kwargs, # Unused, for compatibility with HybridLinearAttnBackend
):
from sglang.srt.layers.attention.mamba.causal_conv1d_triton import (
causal_conv1d_fn,
)
q_proj_states, k_proj_states, v_proj_states = torch.split(
mixed_qkv,
[layer.q_dim, layer.k_dim, layer.v_dim],
dim=-1,
)
q_conv_weights, k_conv_weights, v_conv_weights = layer.conv_weights
q_conv_bias, k_conv_bias, v_conv_bias = layer.bias
query_start_loc = self.forward_metadata.query_start_loc
cache_indices = self.forward_metadata.mamba_cache_indices
mamba_cache_params = self.req_to_token_pool.mamba2_layer_cache(layer.layer_id)
conv_state_q, conv_state_k, conv_state_v = mamba_cache_params.conv
# deal with strides
conv_state_q = conv_state_q.transpose(-1, -2)
conv_state_k = conv_state_k.transpose(-1, -2)
conv_state_v = conv_state_v.transpose(-1, -2)
ssm_states = mamba_cache_params.temporal
has_initial_state = forward_batch.extend_prefix_lens > 0
q_proj_states = q_proj_states.transpose(0, 1)
k_proj_states = k_proj_states.transpose(0, 1)
v_proj_states = v_proj_states.transpose(0, 1)
q = causal_conv1d_fn(
q_proj_states,
q_conv_weights,
q_conv_bias,
activation="silu",
conv_states=conv_state_q,
has_initial_state=has_initial_state,
cache_indices=cache_indices,
query_start_loc=query_start_loc,
seq_lens_cpu=forward_batch.extend_seq_lens_cpu,
).transpose(0, 1)
k = causal_conv1d_fn(
k_proj_states,
k_conv_weights,
k_conv_bias,
activation="silu",
conv_states=conv_state_k,
has_initial_state=has_initial_state,
cache_indices=cache_indices,
query_start_loc=query_start_loc,
seq_lens_cpu=forward_batch.extend_seq_lens_cpu,
).transpose(0, 1)
v = causal_conv1d_fn(
v_proj_states,
v_conv_weights,
v_conv_bias,
activation="silu",
conv_states=conv_state_v,
has_initial_state=has_initial_state,
cache_indices=cache_indices,
query_start_loc=query_start_loc,
seq_lens_cpu=forward_batch.extend_seq_lens_cpu,
).transpose(0, 1)
q = rearrange(q, "n (h d) -> 1 n h d", d=layer.head_q_dim)
k = rearrange(k, "n (h d) -> 1 n h d", d=layer.head_k_dim)
v = rearrange(v, "n (h d) -> 1 n h d", d=layer.head_v_dim)
core_attn_out = chunk_kda(
q=q,
k=k,
v=v,
g=a,
beta=b,
initial_state=ssm_states,
initial_state_indices=cache_indices,
use_qk_l2norm_in_kernel=True,
cu_seqlens=query_start_loc,
)
return core_attn_out
class GDNAttnBackend(MambaAttnBackendBase):
"""Attention backend using Mamba kernel."""
def __init__(self, model_runner: ModelRunner):
super().__init__(model_runner)
self.conv_states_shape = (
model_runner.req_to_token_pool.mamba_pool.mamba_cache.conv[0].shape
)
if not is_cpu() and not is_npu():
assert (
self.conv_states_shape[-1] < FLA_CHUNK_SIZE
), f"{self.conv_states_shape[-1]=} should be less than {FLA_CHUNK_SIZE}"
use_cutedsl = Envs.SGLANG_USE_CUTEDSL_GDN_DECODE.get()
if use_cutedsl and cutedsl_fused_sigmoid_gating_delta_rule_update is None:
rank0_log(
"CuTe DSL GDN decode requested but unavailable "
"(missing cuda.bindings). Falling back to FLA decode kernel."
)
use_cutedsl = False
rank0_log(f"CuTe DSL GDN decode enabled: {use_cutedsl}")
self._kernel_func = (
cutedsl_fused_sigmoid_gating_delta_rule_update
if use_cutedsl
else fused_sigmoid_gating_delta_rule_update
)
def forward_decode(
self,
layer: RadixLinearAttention,
forward_batch: ForwardBatch,
mixed_qkv: torch.Tensor,
a: torch.Tensor,
b: torch.Tensor,
**kwargs, # Unused, for compatibility with HybridLinearAttnBackend
):
layer_cache = self.req_to_token_pool.mamba2_layer_cache(layer.layer_id)
conv_states = layer_cache.conv[0]
ssm_states = layer_cache.temporal
query_start_loc = self.forward_metadata.query_start_loc
cache_indices = self.forward_metadata.mamba_cache_indices
mixed_qkv = causal_conv1d_update(
mixed_qkv,
conv_states,
layer.conv_weights,
layer.bias,
layer.activation,
conv_state_indices=cache_indices,
)
query, key, value = torch.split(
mixed_qkv,
[layer.q_dim, layer.k_dim, layer.v_dim],
dim=-1,
)
# Reshape from [bs, h*d] to [1, bs, h, d]
bs = forward_batch.batch_size
query = query.view(1, bs, layer.num_q_heads, layer.head_q_dim)
key = key.view(1, bs, layer.num_k_heads, layer.head_k_dim)
value = value.view(1, bs, layer.num_v_heads, layer.head_v_dim)
core_attn_out = self._kernel_func(
A_log=layer.A_log,
dt_bias=layer.dt_bias,
q=query,
k=key,
v=value,
a=a,
b=b,
initial_state_source=ssm_states,
initial_state_indices=cache_indices,
cu_seqlens=query_start_loc,
use_qk_l2norm_in_kernel=True,
softplus_beta=1.0,
softplus_threshold=20.0,
)
self._track_mamba_state_decode(
forward_batch, conv_states, ssm_states, cache_indices
)
return core_attn_out
def forward_extend(
self,
layer: RadixLinearAttention,
forward_batch: ForwardBatch,
mixed_qkv: torch.Tensor,
a: torch.Tensor,
b: torch.Tensor,
**kwargs, # Unused, for compatibility with HybridLinearAttnBackend
):
seq_len = mixed_qkv.shape[0]
is_target_verify = forward_batch.forward_mode.is_target_verify()
forward_metadata = self.forward_metadata
query_start_loc = forward_metadata.query_start_loc
cache_indices = forward_metadata.mamba_cache_indices
retrieve_next_token = forward_metadata.retrieve_next_token
retrieve_next_sibling = forward_metadata.retrieve_next_sibling
retrieve_parent_token = forward_metadata.retrieve_parent_token
mamba_cache_params = self.req_to_token_pool.mamba2_layer_cache(layer.layer_id)
conv_states = mamba_cache_params.conv[0]
ssm_states = mamba_cache_params.temporal
if is_target_verify:
assert isinstance(mamba_cache_params, MambaPool.SpeculativeState)
intermediate_state_cache = mamba_cache_params.intermediate_ssm
intermediate_conv_window_cache = (
mamba_cache_params.intermediate_conv_window[0]
)
has_initial_states = torch.ones(
seq_len // forward_batch.spec_info.draft_token_num,
dtype=torch.bool,
device=forward_batch.input_ids.device,
)
intermediate_state_indices = torch.arange(
cache_indices.shape[0], dtype=torch.int32, device=cache_indices.device
)
else:
has_initial_states = forward_batch.extend_prefix_lens > 0
if is_target_verify:
batch_size = seq_len // forward_batch.spec_info.draft_token_num
draft_token_num = forward_batch.spec_info.draft_token_num
mixed_qkv_reshaped = mixed_qkv.view(
batch_size, draft_token_num, -1
).transpose(1, 2)
mixed_qkv_processed = causal_conv1d_update(
mixed_qkv_reshaped,
conv_states,
layer.conv_weights,
layer.bias,
layer.activation,
conv_state_indices=cache_indices[:batch_size],
intermediate_conv_window=intermediate_conv_window_cache,
intermediate_state_indices=intermediate_state_indices[:batch_size],
retrieve_next_token=retrieve_next_token,
retrieve_next_sibling=retrieve_next_sibling,
retrieve_parent_token=retrieve_parent_token,
)
mixed_qkv = mixed_qkv_processed.transpose(1, 2).view(seq_len, -1)
else:
mixed_qkv = mixed_qkv.transpose(0, 1)
if (
forward_batch.mamba_track_mask is not None
and forward_batch.mamba_track_mask.any()
):
conv_dst = forward_batch.mamba_track_indices
# Gather all slices at once: [:, track_conv_indices] -> [d, num_masked, slice_len]
# track_conv_indices is already filtered and clamped in _init_track_conv_indices
mixed_qkv_to_track = mixed_qkv[
:, forward_metadata.track_conv_indices
].transpose(0, 1)
# Apply mask and assign to destinations
mask_indices = forward_batch.mamba_track_mask.nonzero(as_tuple=True)[0]
conv_states[conv_dst[mask_indices]] = mixed_qkv_to_track
mixed_qkv = causal_conv1d_fn(
mixed_qkv,
layer.conv_weights,
layer.bias,
activation=layer.activation,
conv_states=conv_states,
has_initial_state=has_initial_states,
cache_indices=cache_indices,
query_start_loc=query_start_loc,
seq_lens_cpu=forward_batch.extend_seq_lens_cpu,
).transpose(0, 1)[:seq_len]
query, key, value = torch.split(
mixed_qkv,
[layer.q_dim, layer.k_dim, layer.v_dim],
dim=-1,
)
actual_seq_len = query.shape[0]
query = query.view(1, actual_seq_len, layer.num_q_heads, layer.head_q_dim)
key = key.view(1, actual_seq_len, layer.num_k_heads, layer.head_k_dim)
value = value.view(1, actual_seq_len, layer.num_v_heads, layer.head_v_dim)
g, beta = fused_gdn_gating(layer.A_log, a, b, layer.dt_bias)
if is_target_verify:
core_attn_out = fused_recurrent_gated_delta_rule_update(
q=query,
k=key,
v=value,
g=g,
beta=beta,
initial_state_source=ssm_states,
initial_state_indices=cache_indices,
cu_seqlens=query_start_loc,
use_qk_l2norm_in_kernel=True,
disable_state_update=True,
intermediate_states_buffer=intermediate_state_cache,
intermediate_state_indices=intermediate_state_indices,
cache_steps=forward_batch.spec_info.draft_token_num,
retrieve_parent_token=retrieve_parent_token,
)
else:
# Only cuda env uses fuse ssm_states update
recurrent_state = ssm_states
recurrent_state_indices_args = {"initial_state_indices": cache_indices}
if is_npu() or is_cpu():
recurrent_state = ssm_states[cache_indices]
recurrent_state_indices_args = {}
core_attn_out, last_recurrent_state, h = chunk_gated_delta_rule(
q=query,
k=key,
v=value,
g=g,
beta=beta,
initial_state=recurrent_state,
cu_seqlens=query_start_loc,
head_first=False,
use_qk_l2norm_in_kernel=True,
**recurrent_state_indices_args,
)
if is_npu() or is_cpu():
last_recurrent_state = last_recurrent_state.to(
ssm_states.dtype, copy=False
)
ssm_states[cache_indices] = last_recurrent_state
self._track_mamba_state_extend(
forward_batch, h, ssm_states, forward_metadata
)
return core_attn_out
class Mamba2AttnBackend(MambaAttnBackendBase): class Mamba2AttnBackend(MambaAttnBackendBase):
"""Attention backend wrapper for Mamba2Mixer kernels.""" """Attention backend wrapper for Mamba2Mixer kernels."""
@@ -1154,365 +679,6 @@ class Mamba2AttnBackend(MambaAttnBackendBase):
) )
class LightningAttentionBackend(MambaAttnBackendBase):
"""
Note about the init:
- If no spec decoding
- FlashAttentionBackend will be init once when the server starts.
- If spec decoding
- FlashAttentionBackend will be init once for the target worker
- FlashAttentionMultiStepBackend will be once for the draft worker
- It will spawn num_steps FlashAttentionBackend for the draft worker
Note about CUDA Graph:
- We only support CUDA Graph for Decode (Normal Decode and Draft Decode) and Target Verify.
- We don't support CUDA Graph for Extend and Draft Extend.
- When server init, init_cuda_graph_state will be called first and then init_cuda_graph_capture will be called.
- For each forward batch, init_replay_cuda_graph will be called first and then replay the graph.
"""
def __init__(self, model_runner: ModelRunner):
super().__init__(model_runner)
assert not (
model_runner.sliding_window_size is not None
and model_runner.model_config.is_encoder_decoder
), "Sliding window and cross attention are not supported together"
# extra metadata for handling speculative decoding topk > 1, extended draft decode and verify
self.max_context_len = model_runner.model_config.context_len
self.device = model_runner.device
self.decode_cuda_graph_metadata = {}
self.kv_cache_dtype = model_runner.kv_cache_dtype
self.kv_cache_dtype_str = model_runner.server_args.kv_cache_dtype
self.BLOCK = (
model_runner.model_config.block
if hasattr(model_runner.model_config, "block")
else 256
)
total_num_heads = model_runner.model_config.hf_config.num_attention_heads
num_hidden_layers = model_runner.model_config.hf_config.num_hidden_layers
self.tp_slope = LightningAttentionBackend._build_slope_tensor(
total_num_heads, num_hidden_layers, self.device
)
self.linear_backend = getattr(
model_runner.model_config.hf_config, "linear_backend", "seg_la"
)
logger.info(
f"linear_backend for linear attention in hybrid_linear_backend: {self.linear_backend}"
)
def init_forward_metadata(self, forward_batch: ForwardBatch):
metadata = self._forward_metadata(forward_batch)
self.forward_metadata = BailingLinearMetadata.prepare_mixed(
metadata.query_start_loc,
metadata.mamba_cache_indices,
forward_batch,
)
def init_forward_metadata_capture_cuda_graph(
self,
bs: int,
num_tokens: int,
req_pool_indices: torch.Tensor,
seq_lens: torch.Tensor,
encoder_lens: Optional[torch.Tensor],
forward_mode: ForwardMode,
spec_info: Optional[Union[EagleDraftInput, EagleVerifyInput]],
):
metadata = self._capture_metadata(bs, req_pool_indices, forward_mode, spec_info)
self.forward_metadata = BailingLinearMetadata.prepare_decode(
metadata.query_start_loc, metadata.mamba_cache_indices, bs, seq_lens
)
def init_forward_metadata_replay_cuda_graph(
self,
bs: int,
req_pool_indices: torch.Tensor,
seq_lens: torch.Tensor,
seq_lens_sum: int,
encoder_lens: Optional[torch.Tensor],
forward_mode: ForwardMode,
spec_info: Optional[Union[EagleDraftInput, EagleVerifyInput]],
seq_lens_cpu: Optional[torch.Tensor],
):
metadata = self._replay_metadata(
bs, req_pool_indices, forward_mode, spec_info, seq_lens_cpu
)
self.forward_metadata = BailingLinearMetadata.prepare_decode(
metadata.query_start_loc, metadata.mamba_cache_indices, bs, seq_lens
)
@staticmethod
def _build_slope_tensor(
n_attention_heads: int, num_hidden_layers: int, device="cuda"
):
def get_slopes(n):
def get_slopes_power_of_2(n):
start = 2 ** (-(2 ** -(math.log2(n) - 3)))
ratio = start
return [start * ratio**i for i in range(n)]
if math.log2(n).is_integer():
return get_slopes_power_of_2(n)
else:
closest_power_of_2 = 2 ** math.floor(math.log2(n))
return (
get_slopes_power_of_2(closest_power_of_2)
+ get_slopes(2 * closest_power_of_2)[0::2][: n - closest_power_of_2]
)
slopes = torch.tensor(
get_slopes(n_attention_heads), dtype=torch.float32
).reshape(n_attention_heads, 1, 1)
from sglang.srt.layers.dp_attention import (
get_attention_tp_rank,
get_attention_tp_size,
)
tp_heads = n_attention_heads // get_attention_tp_size()
tp_rank = get_attention_tp_rank()
if num_hidden_layers <= 1:
slope_rate_list = [slopes * (1 + 1e-5)]
else:
slope_rate_list = [
slopes * (1 - layer_id / (num_hidden_layers - 1) + 1e-5)
for layer_id in range(num_hidden_layers)
]
tp_slope = [
slope_rate_list[layer_id][tp_rank * tp_heads : (tp_rank + 1) * tp_heads]
.contiguous()
.to(device)
for layer_id in range(num_hidden_layers)
]
return tp_slope
def _prefill_and_mix_infer(
self,
q,
k,
v,
kv_cache,
state_indices_tensor,
forward_batch,
layer,
metadata,
):
hidden = []
for _prefill_idx in range(metadata.num_prefills):
if _prefill_idx >= forward_batch.extend_start_loc.shape[0]:
break
if _prefill_idx >= state_indices_tensor.shape[0]:
break
_start = forward_batch.extend_start_loc[_prefill_idx]
if _prefill_idx + 1 < forward_batch.extend_start_loc.shape[0]:
_end = forward_batch.extend_start_loc[_prefill_idx + 1]
else:
if (
forward_batch.extend_seq_lens is not None
and _prefill_idx < forward_batch.extend_seq_lens.shape[0]
and metadata.num_decodes > 0
):
seq_len = forward_batch.extend_seq_lens[_prefill_idx]
_end = _start + seq_len
else:
_end = q.shape[0]
slot_id = state_indices_tensor[_prefill_idx]
qs = q[_start:_end].transpose(0, 1).contiguous()
ks = k[_start:_end].transpose(0, 1).contiguous()
vs = v[_start:_end].transpose(0, 1).contiguous()
slice_layer_cache = kv_cache[slot_id, ...]
out_slice = BailingLinearKernel.jit_linear_forward_prefix(
qs,
ks,
vs,
slice_layer_cache,
self.tp_slope[layer.layer_id],
self.BLOCK,
layer_idx=layer.layer_id,
)
hidden.append(out_slice.contiguous())
if metadata.num_decodes > 0:
hidden.append(
self._decode_infer(
q, k, v, kv_cache, state_indices_tensor, metadata, layer
)
)
if not hidden:
return torch.empty((0, q.size(-1)), device=q.device, dtype=q.dtype)
hidden = torch.concat(hidden, dim=0).contiguous()
return hidden
def _decode_infer(self, q, k, v, kv_cache, state_indices_tensor, metadata, layer):
num_prefill_tokens = metadata.num_prefill_tokens
num_prefills = metadata.num_prefills
q = q[num_prefill_tokens:].unsqueeze(2).contiguous()
k = k[num_prefill_tokens:].unsqueeze(2).contiguous()
v = v[num_prefill_tokens:].unsqueeze(2).contiguous()
slot_id = state_indices_tensor[num_prefills:]
assert slot_id.shape[0] == q.shape[0], (
f"slot_id length {slot_id.shape[0]} does not match decode batch size {q.shape[0]}. "
"This indicates a bug in the upstream logic that should be investigated."
)
hidden = linear_decode_forward_triton(
q, k, v, kv_cache, self.tp_slope[layer.layer_id], slot_id, 32
)
return hidden
def _linear_attention_entry(
self,
q,
k,
v,
kv_cache,
state_indices_tensor,
metadata,
layer,
mask=None,
temp_cache=None,
intermediate_state_indices=None,
):
q_offsets = metadata.query_start_loc
seg_meta = SegLaMeta(
batch_size=metadata.batch_size,
q_offsets=metadata.query_start_loc,
s_offsets=state_indices_tensor,
q_lengths=q_offsets.diff(),
s_scales=metadata.has_initial_states,
max_q_length=None,
mask=mask,
)
hidden = seg_la_fwd(
q=q,
k=k,
v=v,
s=kv_cache,
decay_scales=self.tp_slope[layer.layer_id],
meta=seg_meta,
caches=temp_cache,
cache_indices=intermediate_state_indices,
decouple=True,
)
return hidden
def forward_extend(
self,
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
layer: RadixAttention,
forward_batch: ForwardBatch,
save_kv_cache=True,
**kwargs,
):
q_rope = kwargs["q_rope"] if "q_rope" in kwargs else None
k_rope = kwargs["k_rope"] if "k_rope" in kwargs else None
layer_id = layer.layer_id if layer else kwargs["layer_id"]
metadata = self.forward_metadata
if self.kv_cache_dtype_str != "auto" and layer.k_scale is not None:
q = q.to(self.kv_cache_dtype)
query_start_loc = self.forward_metadata.query_start_loc
cache_indices = self.forward_metadata.mamba_cache_indices
mamba_cache_params = self.req_to_token_pool.mamba2_layer_cache(layer_id)
ssm_states = mamba_cache_params.temporal
# logger.warning(
# f"---mix {layer.layer_id=}, {query_start_loc=}, {cache_indices=}, {ssm_states.shape=}"
# )
if self.linear_backend == "minimax":
o = self._prefill_and_mix_infer(
q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim),
k,
v,
ssm_states,
cache_indices,
forward_batch,
layer,
metadata,
)
elif self.linear_backend == "seg_la":
intermediate_state_indices = (
torch.arange(
cache_indices.shape[0],
dtype=torch.int32,
device=cache_indices.device,
)
if forward_batch.forward_mode.is_target_verify()
else None
)
o = self._linear_attention_entry(
q,
k,
v,
ssm_states,
cache_indices,
metadata,
layer,
temp_cache=(
mamba_cache_params.intermediate_ssm
if forward_batch.forward_mode.is_target_verify()
else None
),
intermediate_state_indices=intermediate_state_indices,
)
else:
raise ValueError(
f"linear backend: {self.linear_backend} is not support for now"
)
return o.view(-1, layer.tp_q_head_num * layer.v_head_dim)
def forward_decode(
self,
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
layer: RadixAttention,
forward_batch: ForwardBatch,
save_kv_cache=True,
**kwargs,
) -> torch.Tensor:
q_rope = kwargs["q_rope"] if "q_rope" in kwargs else None
k_rope = kwargs["k_rope"] if "k_rope" in kwargs else None
layer_id = layer.layer_id if layer else kwargs["layer_id"]
# Use precomputed metadata across all layers
metadata = self.forward_metadata
if self.kv_cache_dtype_str != "auto":
q = q.to(self.kv_cache_dtype)
# Do linear attention
query_start_loc = self.forward_metadata.query_start_loc
cache_indices = self.forward_metadata.mamba_cache_indices
mamba_cache_params = self.req_to_token_pool.mamba2_layer_cache(layer_id)
ssm_states = mamba_cache_params.temporal
# logger.warning(
# f"---mix {layer.layer_id=}, {query_start_loc.shape=}, {cache_indices.shape=}, {ssm_states.shape=}"
# )
if self.linear_backend == "minimax":
o = self._decode_infer(q, k, v, ssm_states, cache_indices, metadata, layer)
elif self.linear_backend == "seg_la":
o = self._linear_attention_entry(
q, k, v, ssm_states, cache_indices, metadata, layer
)
else:
raise ValueError(
f"linear backend: {self.linear_backend} is not support for now"
)
return o.view(-1, layer.tp_q_head_num * layer.v_head_dim)
class HybridLinearAttnBackend(AttentionBackend): class HybridLinearAttnBackend(AttentionBackend):
"""Manages a full and linear attention backend""" """Manages a full and linear attention backend"""
python/sglang/srt/layers/attention/linear/__init__.py
View File
python/sglang/srt/layers/attention/linear/gdn_backend.py
+379
View File
@@ -0,0 +1,379 @@
from typing import Tuple, Union
import torch
from sglang.srt.layers.attention.fla.fused_gdn_gating import fused_gdn_gating
from sglang.srt.layers.attention.hybrid_linear_attn_backend import MambaAttnBackendBase
from sglang.srt.layers.attention.linear.kernels.gdn_triton import TritonGDNKernel
from sglang.srt.layers.attention.linear.utils import (
LinearAttnKernelBackend,
get_linear_attn_decode_backend,
get_linear_attn_prefill_backend,
)
from sglang.srt.layers.attention.mamba.causal_conv1d_triton import (
causal_conv1d_fn,
causal_conv1d_update,
)
from sglang.srt.layers.radix_linear_attention import RadixLinearAttention
from sglang.srt.mem_cache.memory_pool import MambaPool
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
from sglang.srt.model_executor.model_runner import ModelRunner
from sglang.srt.utils import is_cpu, is_cuda, is_npu
from sglang.srt.utils.common import rank0_log
if not is_cpu():
from sglang.srt.layers.attention.fla.chunk_delta_h import (
CHUNK_SIZE as FLA_CHUNK_SIZE,
)
if is_cuda():
from sglang.srt.layers.attention.mamba.causal_conv1d import (
causal_conv1d_fn as causal_conv1d_fn_cuda,
)
causal_conv1d_fn = causal_conv1d_fn_cuda
elif is_npu():
from sgl_kernel_npu.mamba.causal_conv1d import (
causal_conv1d_fn_npu,
causal_conv1d_update_npu,
)
causal_conv1d_fn = causal_conv1d_fn_npu
causal_conv1d_update = causal_conv1d_update_npu
elif is_cpu():
from sgl_kernel.mamba import causal_conv1d_fn_cpu, causal_conv1d_update_cpu
causal_conv1d_fn = causal_conv1d_fn_cpu
causal_conv1d_update = causal_conv1d_update_cpu
fused_gdn_gating = torch.ops.sgl_kernel.fused_gdn_gating_cpu
class GDNKernelDispatcher:
"""Dispatches GDN kernel calls to the appropriate backend per mode."""
def __init__(
self,
decode_backend: LinearAttnKernelBackend,
prefill_backend: LinearAttnKernelBackend,
):
triton_kernel = TritonGDNKernel()
if decode_backend.is_triton():
self.decode_kernel = triton_kernel
elif decode_backend.is_cutedsl():
if not is_cuda():
raise ValueError("CuTe DSL backend requires CUDA")
from sglang.srt.layers.attention.linear.kernels.gdn_cutedsl import (
CuteDSLGDNKernel,
)
self.decode_kernel = CuteDSLGDNKernel()
else:
raise ValueError(f"Unsupported GDN decode backend: {decode_backend}")
if prefill_backend.is_triton():
self.extend_kernel = triton_kernel
elif prefill_backend.is_cutedsl():
raise ValueError(
"CuTe DSL backend only supports decode, not prefill. "
"Use --linear-attn-prefill-backend triton instead."
)
else:
raise ValueError(f"Unsupported GDN prefill backend: {prefill_backend}")
self.verify_kernel = triton_kernel
rank0_log(
f"GDN kernel dispatcher: decode={self.decode_kernel.__class__.__name__}, "
f"extend={self.extend_kernel.__class__.__name__}, "
f"verify={self.verify_kernel.__class__.__name__}"
)
def decode(
self,
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
a: torch.Tensor,
b: torch.Tensor,
*,
A_log: torch.Tensor,
dt_bias: torch.Tensor,
ssm_states: torch.Tensor,
cache_indices: torch.Tensor,
query_start_loc: torch.Tensor,
**kwargs,
) -> torch.Tensor:
return self.decode_kernel.decode(
q,
k,
v,
a,
b,
A_log=A_log,
dt_bias=dt_bias,
ssm_states=ssm_states,
cache_indices=cache_indices,
query_start_loc=query_start_loc,
**kwargs,
)
def extend(
self,
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
g: torch.Tensor,
beta: torch.Tensor,
*,
ssm_states: torch.Tensor,
cache_indices: torch.Tensor,
query_start_loc: torch.Tensor,
**kwargs,
) -> tuple:
return self.extend_kernel.extend(
q,
k,
v,
g,
beta,
ssm_states=ssm_states,
cache_indices=cache_indices,
query_start_loc=query_start_loc,
**kwargs,
)
def target_verify(
self,
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
g: torch.Tensor,
beta: torch.Tensor,
*,
ssm_states: torch.Tensor,
cache_indices: torch.Tensor,
query_start_loc: torch.Tensor,
**kwargs,
) -> torch.Tensor:
return self.verify_kernel.target_verify(
q,
k,
v,
g,
beta,
ssm_states=ssm_states,
cache_indices=cache_indices,
query_start_loc=query_start_loc,
**kwargs,
)
class GDNAttnBackend(MambaAttnBackendBase):
"""Attention backend for GDN (Gated Delta Network) linear attention."""
def __init__(self, model_runner: ModelRunner):
super().__init__(model_runner)
self.conv_states_shape = (
model_runner.req_to_token_pool.mamba_pool.mamba_cache.conv[0].shape
)
if not is_cpu() and not is_npu():
assert (
self.conv_states_shape[-1] < FLA_CHUNK_SIZE
), f"{self.conv_states_shape[-1]=} should be less than {FLA_CHUNK_SIZE}"
decode_backend = get_linear_attn_decode_backend()
prefill_backend = get_linear_attn_prefill_backend()
self.kernel_dispatcher = GDNKernelDispatcher(decode_backend, prefill_backend)
def forward_decode(
self,
layer: RadixLinearAttention,
forward_batch: ForwardBatch,
mixed_qkv: Union[torch.Tensor, Tuple[torch.Tensor, ...]],
a: torch.Tensor,
b: torch.Tensor,
**kwargs,
):
layer_cache = self.req_to_token_pool.mamba2_layer_cache(layer.layer_id)
conv_states = layer_cache.conv[0]
ssm_states = layer_cache.temporal
query_start_loc = self.forward_metadata.query_start_loc
cache_indices = self.forward_metadata.mamba_cache_indices
assert isinstance(mixed_qkv, torch.Tensor)
mixed_qkv = causal_conv1d_update(
mixed_qkv,
conv_states,
layer.conv_weights,
layer.bias,
layer.activation,
conv_state_indices=cache_indices,
)
query, key, value = torch.split(
mixed_qkv,
[layer.q_dim, layer.k_dim, layer.v_dim],
dim=-1,
)
# Reshape from [bs, h*d] to [1, bs, h, d]
bs = forward_batch.batch_size
query = query.view(1, bs, layer.num_q_heads, layer.head_q_dim)
key = key.view(1, bs, layer.num_k_heads, layer.head_k_dim)
value = value.view(1, bs, layer.num_v_heads, layer.head_v_dim)
core_attn_out = self.kernel_dispatcher.decode(
q=query,
k=key,
v=value,
a=a,
b=b,
A_log=layer.A_log,
dt_bias=layer.dt_bias,
ssm_states=ssm_states,
cache_indices=cache_indices,
query_start_loc=query_start_loc,
)
self._track_mamba_state_decode(
forward_batch, conv_states, ssm_states, cache_indices
)
return core_attn_out
def forward_extend(
self,
layer: RadixLinearAttention,
forward_batch: ForwardBatch,
mixed_qkv: Union[torch.Tensor, Tuple[torch.Tensor, ...]],
a: torch.Tensor,
b: torch.Tensor,
**kwargs,
):
assert isinstance(mixed_qkv, torch.Tensor)
seq_len = mixed_qkv.shape[0]
is_target_verify = forward_batch.forward_mode.is_target_verify()
forward_metadata = self.forward_metadata
query_start_loc = forward_metadata.query_start_loc
cache_indices = forward_metadata.mamba_cache_indices
retrieve_next_token = forward_metadata.retrieve_next_token
retrieve_next_sibling = forward_metadata.retrieve_next_sibling
retrieve_parent_token = forward_metadata.retrieve_parent_token
mamba_cache_params = self.req_to_token_pool.mamba2_layer_cache(layer.layer_id)
conv_states = mamba_cache_params.conv[0]
ssm_states = mamba_cache_params.temporal
if is_target_verify:
assert isinstance(mamba_cache_params, MambaPool.SpeculativeState)
intermediate_state_cache = mamba_cache_params.intermediate_ssm
intermediate_conv_window_cache = (
mamba_cache_params.intermediate_conv_window[0]
)
has_initial_states = torch.ones(
seq_len // forward_batch.spec_info.draft_token_num,
dtype=torch.bool,
device=forward_batch.input_ids.device,
)
intermediate_state_indices = torch.arange(
cache_indices.shape[0], dtype=torch.int32, device=cache_indices.device
)
else:
has_initial_states = forward_batch.extend_prefix_lens > 0
if is_target_verify:
batch_size = seq_len // forward_batch.spec_info.draft_token_num
draft_token_num = forward_batch.spec_info.draft_token_num
mixed_qkv_reshaped = mixed_qkv.view(
batch_size, draft_token_num, -1
).transpose(1, 2)
mixed_qkv_processed = causal_conv1d_update(
mixed_qkv_reshaped,
conv_states,
layer.conv_weights,
layer.bias,
layer.activation,
conv_state_indices=cache_indices[:batch_size],
intermediate_conv_window=intermediate_conv_window_cache,
intermediate_state_indices=intermediate_state_indices[:batch_size],
retrieve_next_token=retrieve_next_token,
retrieve_next_sibling=retrieve_next_sibling,
retrieve_parent_token=retrieve_parent_token,
)
mixed_qkv = mixed_qkv_processed.transpose(1, 2).view(seq_len, -1)
else:
mixed_qkv = mixed_qkv.transpose(0, 1)
if (
forward_batch.mamba_track_mask is not None
and forward_batch.mamba_track_mask.any()
):
conv_dst = forward_batch.mamba_track_indices
mixed_qkv_to_track = mixed_qkv[
:, forward_metadata.track_conv_indices
].transpose(0, 1)
mask_indices = forward_batch.mamba_track_mask.nonzero(as_tuple=True)[0]
conv_states[conv_dst[mask_indices]] = mixed_qkv_to_track
mixed_qkv = causal_conv1d_fn(
mixed_qkv,
layer.conv_weights,
layer.bias,
activation=layer.activation,
conv_states=conv_states,
has_initial_state=has_initial_states,
cache_indices=cache_indices,
query_start_loc=query_start_loc,
seq_lens_cpu=forward_batch.extend_seq_lens_cpu,
).transpose(0, 1)[:seq_len]
query, key, value = torch.split(
mixed_qkv,
[layer.q_dim, layer.k_dim, layer.v_dim],
dim=-1,
)
actual_seq_len = query.shape[0]
query = query.view(1, actual_seq_len, layer.num_q_heads, layer.head_q_dim)
key = key.view(1, actual_seq_len, layer.num_k_heads, layer.head_k_dim)
value = value.view(1, actual_seq_len, layer.num_v_heads, layer.head_v_dim)
g, beta = fused_gdn_gating(layer.A_log, a, b, layer.dt_bias)
if is_target_verify:
core_attn_out = self.kernel_dispatcher.target_verify(
q=query,
k=key,
v=value,
g=g,
beta=beta,
ssm_states=ssm_states,
cache_indices=cache_indices,
query_start_loc=query_start_loc,
intermediate_states_buffer=intermediate_state_cache,
intermediate_state_indices=intermediate_state_indices,
cache_steps=forward_batch.spec_info.draft_token_num,
retrieve_parent_token=retrieve_parent_token,
)
else:
core_attn_out, last_recurrent_state, h = self.kernel_dispatcher.extend(
q=query,
k=key,
v=value,
g=g,
beta=beta,
ssm_states=ssm_states,
cache_indices=cache_indices,
query_start_loc=query_start_loc,
)
if is_npu() or is_cpu():
last_recurrent_state = last_recurrent_state.to(
ssm_states.dtype, copy=False
)
ssm_states[cache_indices] = last_recurrent_state
self._track_mamba_state_extend(
forward_batch, h, ssm_states, forward_metadata
)
return core_attn_out
python/sglang/srt/layers/attention/linear/kda_backend.py
+285
View File
@@ -0,0 +1,285 @@
from typing import Tuple, Union
import torch
from einops import rearrange
from sglang.srt.layers.attention.hybrid_linear_attn_backend import MambaAttnBackendBase
from sglang.srt.layers.attention.linear.kernels.kda_triton import TritonKDAKernel
from sglang.srt.layers.attention.linear.utils import (
LinearAttnKernelBackend,
get_linear_attn_decode_backend,
get_linear_attn_prefill_backend,
)
from sglang.srt.layers.attention.mamba.causal_conv1d_triton import (
causal_conv1d_fn,
causal_conv1d_update,
)
from sglang.srt.layers.radix_linear_attention import RadixLinearAttention
from sglang.srt.utils import is_cpu, is_npu
from sglang.srt.utils.common import rank0_log
# KDA always uses the triton causal_conv1d_fn (no CUDA override).
# Only causal_conv1d_update needs platform-specific overrides for decode.
if is_npu():
from sgl_kernel_npu.mamba.causal_conv1d import causal_conv1d_update_npu
causal_conv1d_update = causal_conv1d_update_npu
elif is_cpu():
from sgl_kernel.mamba import causal_conv1d_update_cpu
causal_conv1d_update = causal_conv1d_update_cpu
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
from sglang.srt.model_executor.model_runner import ModelRunner
class KDAKernelDispatcher:
"""Dispatches KDA kernel calls to the appropriate backend per mode."""
def __init__(
self,
decode_backend: LinearAttnKernelBackend,
prefill_backend: LinearAttnKernelBackend,
):
triton_kernel = TritonKDAKernel()
if decode_backend.is_triton():
self.decode_kernel = triton_kernel
else:
raise ValueError(
f"Unsupported KDA decode backend: {decode_backend}. "
"KDA currently only supports 'triton'."
)
if prefill_backend.is_triton():
self.extend_kernel = triton_kernel
else:
raise ValueError(
f"Unsupported KDA prefill backend: {prefill_backend}. "
"KDA currently only supports 'triton'."
)
rank0_log(
f"KDA kernel dispatcher: decode={self.decode_kernel.__class__.__name__}, "
f"extend={self.extend_kernel.__class__.__name__}"
)
def decode(
self,
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
a: torch.Tensor,
b: torch.Tensor,
*,
A_log: torch.Tensor,
dt_bias: torch.Tensor,
ssm_states: torch.Tensor,
cache_indices: torch.Tensor,
query_start_loc: torch.Tensor,
**kwargs,
) -> torch.Tensor:
return self.decode_kernel.decode(
q,
k,
v,
a,
b,
A_log=A_log,
dt_bias=dt_bias,
ssm_states=ssm_states,
cache_indices=cache_indices,
query_start_loc=query_start_loc,
**kwargs,
)
def extend(
self,
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
g: torch.Tensor,
beta: torch.Tensor,
*,
ssm_states: torch.Tensor,
cache_indices: torch.Tensor,
query_start_loc: torch.Tensor,
**kwargs,
) -> torch.Tensor:
return self.extend_kernel.extend(
q,
k,
v,
g,
beta,
ssm_states=ssm_states,
cache_indices=cache_indices,
query_start_loc=query_start_loc,
**kwargs,
)
class KDAAttnBackend(MambaAttnBackendBase):
"""Attention backend for KDA (Kimi Delta Attention) linear attention."""
def __init__(self, model_runner: ModelRunner):
super().__init__(model_runner)
decode_backend = get_linear_attn_decode_backend()
prefill_backend = get_linear_attn_prefill_backend()
self.kernel_dispatcher = KDAKernelDispatcher(decode_backend, prefill_backend)
def forward_decode(
self,
layer: RadixLinearAttention,
mixed_qkv: Union[torch.Tensor, Tuple[torch.Tensor, ...]],
a: torch.Tensor,
b: torch.Tensor,
**kwargs,
):
q_proj_states, k_proj_states, v_proj_states = torch.split(
mixed_qkv,
[layer.q_dim, layer.k_dim, layer.v_dim],
dim=-1,
)
q_conv_weights, k_conv_weights, v_conv_weights = layer.conv_weights
q_conv_bias, k_conv_bias, v_conv_bias = layer.bias
layer_cache = self.req_to_token_pool.mamba2_layer_cache(layer.layer_id)
q_conv_state, k_conv_state, v_conv_state = layer_cache.conv
ssm_states = layer_cache.temporal
query_start_loc = self.forward_metadata.query_start_loc
cache_indices = self.forward_metadata.mamba_cache_indices
q_conv_state = q_conv_state.transpose(-1, -2)
k_conv_state = k_conv_state.transpose(-1, -2)
v_conv_state = v_conv_state.transpose(-1, -2)
q = causal_conv1d_update(
q_proj_states,
q_conv_state,
q_conv_weights,
q_conv_bias,
activation="silu",
conv_state_indices=cache_indices,
)
k = causal_conv1d_update(
k_proj_states,
k_conv_state,
k_conv_weights,
k_conv_bias,
activation="silu",
conv_state_indices=cache_indices,
)
v = causal_conv1d_update(
v_proj_states,
v_conv_state,
v_conv_weights,
v_conv_bias,
activation="silu",
conv_state_indices=cache_indices,
)
q = rearrange(q, "n (h d) -> 1 n h d", d=layer.head_q_dim)
k = rearrange(k, "n (h d) -> 1 n h d", d=layer.head_k_dim)
v = rearrange(v, "n (h d) -> 1 n h d", d=layer.head_v_dim)
return self.kernel_dispatcher.decode(
q=q,
k=k,
v=v,
a=a,
b=b,
A_log=layer.A_log,
dt_bias=layer.dt_bias,
ssm_states=ssm_states,
cache_indices=cache_indices,
query_start_loc=query_start_loc,
)
def forward_extend(
self,
layer: RadixLinearAttention,
forward_batch: ForwardBatch,
mixed_qkv: Union[torch.Tensor, Tuple[torch.Tensor, ...]],
a: torch.Tensor,
b: torch.Tensor,
**kwargs,
):
q_proj_states, k_proj_states, v_proj_states = torch.split(
mixed_qkv,
[layer.q_dim, layer.k_dim, layer.v_dim],
dim=-1,
)
q_conv_weights, k_conv_weights, v_conv_weights = layer.conv_weights
q_conv_bias, k_conv_bias, v_conv_bias = layer.bias
query_start_loc = self.forward_metadata.query_start_loc
cache_indices = self.forward_metadata.mamba_cache_indices
mamba_cache_params = self.req_to_token_pool.mamba2_layer_cache(layer.layer_id)
conv_state_q, conv_state_k, conv_state_v = mamba_cache_params.conv
# deal with strides
conv_state_q = conv_state_q.transpose(-1, -2)
conv_state_k = conv_state_k.transpose(-1, -2)
conv_state_v = conv_state_v.transpose(-1, -2)
ssm_states = mamba_cache_params.temporal
has_initial_state = forward_batch.extend_prefix_lens > 0
q_proj_states = q_proj_states.transpose(0, 1)
k_proj_states = k_proj_states.transpose(0, 1)
v_proj_states = v_proj_states.transpose(0, 1)
q = causal_conv1d_fn(
q_proj_states,
q_conv_weights,
q_conv_bias,
activation="silu",
conv_states=conv_state_q,
has_initial_state=has_initial_state,
cache_indices=cache_indices,
query_start_loc=query_start_loc,
seq_lens_cpu=forward_batch.extend_seq_lens_cpu,
).transpose(0, 1)
k = causal_conv1d_fn(
k_proj_states,
k_conv_weights,
k_conv_bias,
activation="silu",
conv_states=conv_state_k,
has_initial_state=has_initial_state,
cache_indices=cache_indices,
query_start_loc=query_start_loc,
seq_lens_cpu=forward_batch.extend_seq_lens_cpu,
).transpose(0, 1)
v = causal_conv1d_fn(
v_proj_states,
v_conv_weights,
v_conv_bias,
activation="silu",
conv_states=conv_state_v,
has_initial_state=has_initial_state,
cache_indices=cache_indices,
query_start_loc=query_start_loc,
seq_lens_cpu=forward_batch.extend_seq_lens_cpu,
).transpose(0, 1)
q = rearrange(q, "n (h d) -> 1 n h d", d=layer.head_q_dim)
k = rearrange(k, "n (h d) -> 1 n h d", d=layer.head_k_dim)
v = rearrange(v, "n (h d) -> 1 n h d", d=layer.head_v_dim)
core_attn_out = self.kernel_dispatcher.extend(
q=q,
k=k,
v=v,
g=a,
beta=b,
ssm_states=ssm_states,
cache_indices=cache_indices,
query_start_loc=query_start_loc,
)
return core_attn_out
python/sglang/srt/layers/attention/linear/kernels/__init__.py
View File
python/sglang/srt/layers/attention/linear/kernels/gdn_cutedsl.py
+47
View File
@@ -0,0 +1,47 @@
import torch
from sglang.jit_kernel.cutedsl_gdn import cutedsl_fused_sigmoid_gating_delta_rule_update
from sglang.srt.layers.attention.linear.kernels.kernel_backend import (
LinearAttnKernelBase,
)
class CuteDSLGDNKernel(LinearAttnKernelBase):
"""CuTe DSL kernel for GDN decode (CUDA only)."""
def decode(
self,
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
a: torch.Tensor,
b: torch.Tensor,
*,
A_log: torch.Tensor,
dt_bias: torch.Tensor,
ssm_states: torch.Tensor,
cache_indices: torch.Tensor,
query_start_loc: torch.Tensor,
**kwargs,
) -> torch.Tensor:
return cutedsl_fused_sigmoid_gating_delta_rule_update(
A_log=A_log,
dt_bias=dt_bias,
q=q,
k=k,
v=v,
a=a,
b=b,
initial_state_source=ssm_states,
initial_state_indices=cache_indices,
cu_seqlens=query_start_loc,
use_qk_l2norm_in_kernel=True,
softplus_beta=1.0,
softplus_threshold=20.0,
)
def extend(self, *args, **kwargs):
raise NotImplementedError("CuteDSLGDNKernel only supports decode")
def target_verify(self, *args, **kwargs):
raise NotImplementedError("CuteDSLGDNKernel only supports decode")
python/sglang/srt/layers/attention/linear/kernels/gdn_triton.py
+131
View File
@@ -0,0 +1,131 @@
import torch
from sglang.srt.layers.attention.linear.kernels.kernel_backend import (
LinearAttnKernelBase,
)
from sglang.srt.utils import is_cpu, is_npu
if not is_cpu():
from sglang.srt.layers.attention.fla.chunk import chunk_gated_delta_rule
from sglang.srt.layers.attention.fla.fused_recurrent import (
fused_recurrent_gated_delta_rule_update,
)
from sglang.srt.layers.attention.fla.fused_sigmoid_gating_recurrent import (
fused_sigmoid_gating_delta_rule_update,
)
if is_npu():
from sgl_kernel_npu.fla.chunk import chunk_gated_delta_rule_npu
from sgl_kernel_npu.fla.fused_sigmoid_gating_recurrent import (
fused_sigmoid_gating_delta_rule_update_npu,
)
chunk_gated_delta_rule = chunk_gated_delta_rule_npu
fused_sigmoid_gating_delta_rule_update = fused_sigmoid_gating_delta_rule_update_npu
elif is_cpu():
from sgl_kernel.mamba import chunk_gated_delta_rule_cpu
chunk_gated_delta_rule = chunk_gated_delta_rule_cpu
fused_sigmoid_gating_delta_rule_update = (
torch.ops.sgl_kernel.fused_sigmoid_gating_delta_rule_update_cpu
)
class TritonGDNKernel(LinearAttnKernelBase):
"""Triton-based kernel for GDN (Gated Delta Network) linear attention."""
def decode(
self,
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
a: torch.Tensor,
b: torch.Tensor,
*,
A_log: torch.Tensor,
dt_bias: torch.Tensor,
ssm_states: torch.Tensor,
cache_indices: torch.Tensor,
query_start_loc: torch.Tensor,
**kwargs,
) -> torch.Tensor:
return fused_sigmoid_gating_delta_rule_update(
A_log=A_log,
dt_bias=dt_bias,
q=q,
k=k,
v=v,
a=a,
b=b,
initial_state_source=ssm_states,
initial_state_indices=cache_indices,
cu_seqlens=query_start_loc,
use_qk_l2norm_in_kernel=True,
softplus_beta=1.0,
softplus_threshold=20.0,
)
def extend(
self,
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
g: torch.Tensor,
beta: torch.Tensor,
*,
ssm_states: torch.Tensor,
cache_indices: torch.Tensor,
query_start_loc: torch.Tensor,
**kwargs,
) -> tuple:
recurrent_state = ssm_states
recurrent_state_indices_args = {"initial_state_indices": cache_indices}
if is_npu() or is_cpu():
recurrent_state = ssm_states[cache_indices]
recurrent_state_indices_args = {}
return chunk_gated_delta_rule(
q=q,
k=k,
v=v,
g=g,
beta=beta,
initial_state=recurrent_state,
cu_seqlens=query_start_loc,
head_first=False,
use_qk_l2norm_in_kernel=True,
**recurrent_state_indices_args,
)
def target_verify(
self,
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
g: torch.Tensor,
beta: torch.Tensor,
*,
ssm_states: torch.Tensor,
cache_indices: torch.Tensor,
query_start_loc: torch.Tensor,
intermediate_states_buffer: torch.Tensor,
intermediate_state_indices: torch.Tensor,
cache_steps: int,
retrieve_parent_token: torch.Tensor,
**kwargs,
) -> torch.Tensor:
return fused_recurrent_gated_delta_rule_update(
q=q,
k=k,
v=v,
g=g,
beta=beta,
initial_state_source=ssm_states,
initial_state_indices=cache_indices,
cu_seqlens=query_start_loc,
use_qk_l2norm_in_kernel=True,
disable_state_update=True,
intermediate_states_buffer=intermediate_states_buffer,
intermediate_state_indices=intermediate_state_indices,
cache_steps=cache_steps,
retrieve_parent_token=retrieve_parent_token,
)
python/sglang/srt/layers/attention/linear/kernels/kda_triton.py
+73
View File
@@ -0,0 +1,73 @@
import torch
from sglang.srt.layers.attention.linear.kernels.kernel_backend import (
LinearAttnKernelBase,
)
from sglang.srt.utils import is_cpu
if not is_cpu():
from sglang.srt.layers.attention.fla.fused_sigmoid_gating_recurrent import (
fused_sigmoid_gating_delta_rule_update,
)
from sglang.srt.layers.attention.fla.kda import chunk_kda
class TritonKDAKernel(LinearAttnKernelBase):
"""Triton-based kernel for KDA (Kimi Delta Attention) linear attention."""
def decode(
self,
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
a: torch.Tensor,
b: torch.Tensor,
*,
A_log: torch.Tensor,
dt_bias: torch.Tensor,
ssm_states: torch.Tensor,
cache_indices: torch.Tensor,
query_start_loc: torch.Tensor,
**kwargs,
) -> torch.Tensor:
return fused_sigmoid_gating_delta_rule_update(
A_log=A_log,
dt_bias=dt_bias,
q=q,
k=k,
v=v,
a=a,
b=b,
initial_state_source=ssm_states,
initial_state_indices=cache_indices,
cu_seqlens=query_start_loc,
use_qk_l2norm_in_kernel=True,
softplus_beta=1.0,
softplus_threshold=20.0,
is_kda=True,
)
def extend(
self,
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
g: torch.Tensor,
beta: torch.Tensor,
*,
ssm_states: torch.Tensor,
cache_indices: torch.Tensor,
query_start_loc: torch.Tensor,
**kwargs,
) -> torch.Tensor:
return chunk_kda(
q=q,
k=k,
v=v,
g=g,
beta=beta,
initial_state=ssm_states,
initial_state_indices=cache_indices,
use_qk_l2norm_in_kernel=True,
cu_seqlens=query_start_loc,
)
python/sglang/srt/layers/attention/linear/kernels/kernel_backend.py
+60
View File
@@ -0,0 +1,60 @@
from abc import ABC, abstractmethod
import torch
class LinearAttnKernelBase(ABC):
"""Abstract base class for linear attention kernel implementations.
Each concrete implementation wraps a specific kernel (Triton, CuTe DSL, etc.)
and provides decode/extend/target_verify methods with a unified interface.
"""
@abstractmethod
def decode(
self,
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
a: torch.Tensor,
b: torch.Tensor,
*,
A_log: torch.Tensor,
dt_bias: torch.Tensor,
ssm_states: torch.Tensor,
cache_indices: torch.Tensor,
query_start_loc: torch.Tensor,
**kwargs,
) -> torch.Tensor: ...
@abstractmethod
def extend(
self,
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
g: torch.Tensor,
beta: torch.Tensor,
*,
ssm_states: torch.Tensor,
cache_indices: torch.Tensor,
query_start_loc: torch.Tensor,
**kwargs,
) -> tuple: ...
def target_verify(
self,
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
g: torch.Tensor,
beta: torch.Tensor,
*,
ssm_states: torch.Tensor,
cache_indices: torch.Tensor,
query_start_loc: torch.Tensor,
**kwargs,
) -> torch.Tensor:
raise NotImplementedError(
f"{self.__class__.__name__} does not support target_verify"
)
python/sglang/srt/layers/attention/linear/lightning_backend.py
+372
View File
@@ -0,0 +1,372 @@
import logging
import math
from typing import Optional, Union
import torch
from sglang.srt.layers.attention.hybrid_linear_attn_backend import MambaAttnBackendBase
from sglang.srt.layers.attention.linear.lightning_attn import (
BailingLinearKernel,
linear_decode_forward_triton,
)
from sglang.srt.layers.attention.linear.linear_metadata import BailingLinearMetadata
from sglang.srt.layers.attention.linear.seg_la import SegLaMeta, seg_la_fwd
from sglang.srt.layers.radix_attention import RadixAttention
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
from sglang.srt.model_executor.model_runner import ModelRunner
from sglang.srt.speculative.eagle_info import EagleDraftInput, EagleVerifyInput
logger = logging.getLogger(__name__)
class LightningAttentionBackend(MambaAttnBackendBase):
"""
Note about the init:
- If no spec decoding
- FlashAttentionBackend will be init once when the server starts.
- If spec decoding
- FlashAttentionBackend will be init once for the target worker
- FlashAttentionMultiStepBackend will be once for the draft worker
- It will spawn num_steps FlashAttentionBackend for the draft worker
Note about CUDA Graph:
- We only support CUDA Graph for Decode (Normal Decode and Draft Decode) and Target Verify.
- We don't support CUDA Graph for Extend and Draft Extend.
- When server init, init_cuda_graph_state will be called first and then init_cuda_graph_capture will be called.
- For each forward batch, init_replay_cuda_graph will be called first and then replay the graph.
"""
def __init__(self, model_runner: ModelRunner):
super().__init__(model_runner)
assert not (
model_runner.sliding_window_size is not None
and model_runner.model_config.is_encoder_decoder
), "Sliding window and cross attention are not supported together"
# extra metadata for handling speculative decoding topk > 1, extended draft decode and verify
self.max_context_len = model_runner.model_config.context_len
self.device = model_runner.device
self.decode_cuda_graph_metadata = {}
self.kv_cache_dtype = model_runner.kv_cache_dtype
self.kv_cache_dtype_str = model_runner.server_args.kv_cache_dtype
self.BLOCK = (
model_runner.model_config.block
if hasattr(model_runner.model_config, "block")
else 256
)
total_num_heads = model_runner.model_config.hf_config.num_attention_heads
num_hidden_layers = model_runner.model_config.hf_config.num_hidden_layers
self.tp_slope = LightningAttentionBackend._build_slope_tensor(
total_num_heads, num_hidden_layers, self.device
)
self.linear_backend = getattr(
model_runner.model_config.hf_config, "linear_backend", "seg_la"
)
logger.info(
f"linear_backend for linear attention in hybrid_linear_backend: {self.linear_backend}"
)
def init_forward_metadata(self, forward_batch: ForwardBatch):
metadata = self._forward_metadata(forward_batch)
self.forward_metadata = BailingLinearMetadata.prepare_mixed(
metadata.query_start_loc,
metadata.mamba_cache_indices,
forward_batch,
)
def init_forward_metadata_capture_cuda_graph(
self,
bs: int,
num_tokens: int,
req_pool_indices: torch.Tensor,
seq_lens: torch.Tensor,
encoder_lens: Optional[torch.Tensor],
forward_mode: ForwardMode,
spec_info: Optional[Union[EagleDraftInput, EagleVerifyInput]],
):
metadata = self._capture_metadata(bs, req_pool_indices, forward_mode, spec_info)
self.forward_metadata = BailingLinearMetadata.prepare_decode(
metadata.query_start_loc, metadata.mamba_cache_indices, bs, seq_lens
)
def init_forward_metadata_replay_cuda_graph(
self,
bs: int,
req_pool_indices: torch.Tensor,
seq_lens: torch.Tensor,
seq_lens_sum: int,
encoder_lens: Optional[torch.Tensor],
forward_mode: ForwardMode,
spec_info: Optional[Union[EagleDraftInput, EagleVerifyInput]],
seq_lens_cpu: Optional[torch.Tensor],
):
metadata = self._replay_metadata(
bs, req_pool_indices, forward_mode, spec_info, seq_lens_cpu
)
self.forward_metadata = BailingLinearMetadata.prepare_decode(
metadata.query_start_loc, metadata.mamba_cache_indices, bs, seq_lens
)
@staticmethod
def _build_slope_tensor(
n_attention_heads: int, num_hidden_layers: int, device="cuda"
):
def get_slopes(n):
def get_slopes_power_of_2(n):
start = 2 ** (-(2 ** -(math.log2(n) - 3)))
ratio = start
return [start * ratio**i for i in range(n)]
if math.log2(n).is_integer():
return get_slopes_power_of_2(n)
else:
closest_power_of_2 = 2 ** math.floor(math.log2(n))
return (
get_slopes_power_of_2(closest_power_of_2)
+ get_slopes(2 * closest_power_of_2)[0::2][: n - closest_power_of_2]
)
slopes = torch.tensor(
get_slopes(n_attention_heads), dtype=torch.float32
).reshape(n_attention_heads, 1, 1)
from sglang.srt.layers.dp_attention import (
get_attention_tp_rank,
get_attention_tp_size,
)
tp_heads = n_attention_heads // get_attention_tp_size()
tp_rank = get_attention_tp_rank()
if num_hidden_layers <= 1:
slope_rate_list = [slopes * (1 + 1e-5)]
else:
slope_rate_list = [
slopes * (1 - layer_id / (num_hidden_layers - 1) + 1e-5)
for layer_id in range(num_hidden_layers)
]
tp_slope = [
slope_rate_list[layer_id][tp_rank * tp_heads : (tp_rank + 1) * tp_heads]
.contiguous()
.to(device)
for layer_id in range(num_hidden_layers)
]
return tp_slope
def _prefill_and_mix_infer(
self,
q,
k,
v,
kv_cache,
state_indices_tensor,
forward_batch,
layer,
metadata,
):
hidden = []
for _prefill_idx in range(metadata.num_prefills):
if _prefill_idx >= forward_batch.extend_start_loc.shape[0]:
break
if _prefill_idx >= state_indices_tensor.shape[0]:
break
_start = forward_batch.extend_start_loc[_prefill_idx]
if _prefill_idx + 1 < forward_batch.extend_start_loc.shape[0]:
_end = forward_batch.extend_start_loc[_prefill_idx + 1]
else:
if (
forward_batch.extend_seq_lens is not None
and _prefill_idx < forward_batch.extend_seq_lens.shape[0]
and metadata.num_decodes > 0
):
seq_len = forward_batch.extend_seq_lens[_prefill_idx]
_end = _start + seq_len
else:
_end = q.shape[0]
slot_id = state_indices_tensor[_prefill_idx]
qs = q[_start:_end].transpose(0, 1).contiguous()
ks = k[_start:_end].transpose(0, 1).contiguous()
vs = v[_start:_end].transpose(0, 1).contiguous()
slice_layer_cache = kv_cache[slot_id, ...]
out_slice = BailingLinearKernel.jit_linear_forward_prefix(
qs,
ks,
vs,
slice_layer_cache,
self.tp_slope[layer.layer_id],
self.BLOCK,
layer_idx=layer.layer_id,
)
hidden.append(out_slice.contiguous())
if metadata.num_decodes > 0:
hidden.append(
self._decode_infer(
q, k, v, kv_cache, state_indices_tensor, metadata, layer
)
)
if not hidden:
return torch.empty((0, q.size(-1)), device=q.device, dtype=q.dtype)
hidden = torch.concat(hidden, dim=0).contiguous()
return hidden
def _decode_infer(self, q, k, v, kv_cache, state_indices_tensor, metadata, layer):
num_prefill_tokens = metadata.num_prefill_tokens
num_prefills = metadata.num_prefills
q = q[num_prefill_tokens:].unsqueeze(2).contiguous()
k = k[num_prefill_tokens:].unsqueeze(2).contiguous()
v = v[num_prefill_tokens:].unsqueeze(2).contiguous()
slot_id = state_indices_tensor[num_prefills:]
assert slot_id.shape[0] == q.shape[0], (
f"slot_id length {slot_id.shape[0]} does not match decode batch size {q.shape[0]}. "
"This indicates a bug in the upstream logic that should be investigated."
)
hidden = linear_decode_forward_triton(
q, k, v, kv_cache, self.tp_slope[layer.layer_id], slot_id, 32
)
return hidden
def _linear_attention_entry(
self,
q,
k,
v,
kv_cache,
state_indices_tensor,
metadata,
layer,
mask=None,
temp_cache=None,
intermediate_state_indices=None,
):
q_offsets = metadata.query_start_loc
seg_meta = SegLaMeta(
batch_size=metadata.batch_size,
q_offsets=metadata.query_start_loc,
s_offsets=state_indices_tensor,
q_lengths=q_offsets.diff(),
s_scales=metadata.has_initial_states,
max_q_length=None,
mask=mask,
)
hidden = seg_la_fwd(
q=q,
k=k,
v=v,
s=kv_cache,
decay_scales=self.tp_slope[layer.layer_id],
meta=seg_meta,
caches=temp_cache,
cache_indices=intermediate_state_indices,
decouple=True,
)
return hidden
def forward_extend(
self,
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
layer: RadixAttention,
forward_batch: ForwardBatch,
save_kv_cache=True,
**kwargs,
):
q_rope = kwargs["q_rope"] if "q_rope" in kwargs else None
k_rope = kwargs["k_rope"] if "k_rope" in kwargs else None
layer_id = layer.layer_id if layer else kwargs["layer_id"]
metadata = self.forward_metadata
if self.kv_cache_dtype_str != "auto" and layer.k_scale is not None:
q = q.to(self.kv_cache_dtype)
query_start_loc = self.forward_metadata.query_start_loc
cache_indices = self.forward_metadata.mamba_cache_indices
mamba_cache_params = self.req_to_token_pool.mamba2_layer_cache(layer_id)
ssm_states = mamba_cache_params.temporal
if self.linear_backend == "minimax":
o = self._prefill_and_mix_infer(
q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim),
k,
v,
ssm_states,
cache_indices,
forward_batch,
layer,
metadata,
)
elif self.linear_backend == "seg_la":
intermediate_state_indices = (
torch.arange(
cache_indices.shape[0],
dtype=torch.int32,
device=cache_indices.device,
)
if forward_batch.forward_mode.is_target_verify()
else None
)
o = self._linear_attention_entry(
q,
k,
v,
ssm_states,
cache_indices,
metadata,
layer,
temp_cache=(
mamba_cache_params.intermediate_ssm
if forward_batch.forward_mode.is_target_verify()
else None
),
intermediate_state_indices=intermediate_state_indices,
)
else:
raise ValueError(
f"linear backend: {self.linear_backend} is not support for now"
)
return o.view(-1, layer.tp_q_head_num * layer.v_head_dim)
def forward_decode(
self,
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
layer: RadixAttention,
forward_batch: ForwardBatch,
save_kv_cache=True,
**kwargs,
) -> torch.Tensor:
q_rope = kwargs["q_rope"] if "q_rope" in kwargs else None
k_rope = kwargs["k_rope"] if "k_rope" in kwargs else None
layer_id = layer.layer_id if layer else kwargs["layer_id"]
# Use precomputed metadata across all layers
metadata = self.forward_metadata
if self.kv_cache_dtype_str != "auto":
q = q.to(self.kv_cache_dtype)
# Do linear attention
query_start_loc = self.forward_metadata.query_start_loc
cache_indices = self.forward_metadata.mamba_cache_indices
mamba_cache_params = self.req_to_token_pool.mamba2_layer_cache(layer_id)
ssm_states = mamba_cache_params.temporal
if self.linear_backend == "minimax":
o = self._decode_infer(q, k, v, ssm_states, cache_indices, metadata, layer)
elif self.linear_backend == "seg_la":
o = self._linear_attention_entry(
q, k, v, ssm_states, cache_indices, metadata, layer
)
else:
raise ValueError(
f"linear backend: {self.linear_backend} is not support for now"
)
return o.view(-1, layer.tp_q_head_num * layer.v_head_dim)
python/sglang/srt/layers/attention/linear/utils.py
+64
View File
@@ -0,0 +1,64 @@
from __future__ import annotations
import logging
from enum import Enum
from typing import TYPE_CHECKING, Optional
from sglang.srt.utils.common import rank0_log
if TYPE_CHECKING:
from sglang.srt.server_args import ServerArgs
logger = logging.getLogger(__name__)
class LinearAttnKernelBackend(Enum):
TRITON = "triton"
CUTEDSL = "cutedsl"
def is_triton(self):
return self == LinearAttnKernelBackend.TRITON
def is_cutedsl(self):
return self == LinearAttnKernelBackend.CUTEDSL
LINEAR_ATTN_DECODE_BACKEND: Optional[LinearAttnKernelBackend] = None
LINEAR_ATTN_PREFILL_BACKEND: Optional[LinearAttnKernelBackend] = None
def initialize_linear_attn_config(server_args: ServerArgs):
global LINEAR_ATTN_DECODE_BACKEND
global LINEAR_ATTN_PREFILL_BACKEND
base = server_args.linear_attn_backend
decode = server_args.linear_attn_decode_backend or base
prefill = server_args.linear_attn_prefill_backend or base
LINEAR_ATTN_DECODE_BACKEND = LinearAttnKernelBackend(decode)
LINEAR_ATTN_PREFILL_BACKEND = LinearAttnKernelBackend(prefill)
rank0_log(
f"Linear attention kernel backend: "
f"decode={LINEAR_ATTN_DECODE_BACKEND.value}, "
f"prefill={LINEAR_ATTN_PREFILL_BACKEND.value}"
)
def get_linear_attn_decode_backend() -> LinearAttnKernelBackend:
global LINEAR_ATTN_DECODE_BACKEND
if LINEAR_ATTN_DECODE_BACKEND is None:
logger.warning(
"LINEAR_ATTN_DECODE_BACKEND is not initialized, using triton backend"
)
LINEAR_ATTN_DECODE_BACKEND = LinearAttnKernelBackend.TRITON
return LINEAR_ATTN_DECODE_BACKEND
def get_linear_attn_prefill_backend() -> LinearAttnKernelBackend:
global LINEAR_ATTN_PREFILL_BACKEND
if LINEAR_ATTN_PREFILL_BACKEND is None:
logger.warning(
"LINEAR_ATTN_PREFILL_BACKEND is not initialized, using triton backend"
)
LINEAR_ATTN_PREFILL_BACKEND = LinearAttnKernelBackend.TRITON
return LINEAR_ATTN_PREFILL_BACKEND
python/sglang/srt/server_args.py
+29
View File
@@ -216,6 +216,7 @@ MAMBA_SSM_DTYPE_CHOICES = ["float32", "bfloat16", "float16"]
MAMBA_SCHEDULER_STRATEGY_CHOICES = ["auto", "no_buffer", "extra_buffer"] MAMBA_SCHEDULER_STRATEGY_CHOICES = ["auto", "no_buffer", "extra_buffer"]
MAMBA_BACKEND_CHOICES = ["triton", "flashinfer"] MAMBA_BACKEND_CHOICES = ["triton", "flashinfer"]
LINEAR_ATTN_KERNEL_BACKEND_CHOICES = ["triton", "cutedsl"]
# Allow external code to add more choices # Allow external code to add more choices
@@ -524,6 +525,9 @@ class ServerArgs:
mamba_full_memory_ratio: float = 0.9 mamba_full_memory_ratio: float = 0.9
mamba_scheduler_strategy: str = "auto" mamba_scheduler_strategy: str = "auto"
mamba_track_interval: int = 256 mamba_track_interval: int = 256
linear_attn_backend: str = "triton"
linear_attn_decode_backend: Optional[str] = None
linear_attn_prefill_backend: Optional[str] = None
# Hierarchical cache # Hierarchical cache
enable_hierarchical_cache: bool = False enable_hierarchical_cache: bool = False
@@ -4305,6 +4309,31 @@ class ServerArgs:
help="Choose the kernel backend for Mamba SSM operations. Default is 'triton'. " help="Choose the kernel backend for Mamba SSM operations. Default is 'triton'. "
"Options: 'triton' (default), 'flashinfer' (requires FlashInfer with Mamba support).", "Options: 'triton' (default), 'flashinfer' (requires FlashInfer with Mamba support).",
) )
parser.add_argument(
"--linear-attn-backend",
type=str,
choices=LINEAR_ATTN_KERNEL_BACKEND_CHOICES,
default=ServerArgs.linear_attn_backend,
help="The default kernel backend for linear attention (GDN/KDA). "
"Can be overridden per-mode by --linear-attn-decode-backend "
"and --linear-attn-prefill-backend.",
)
parser.add_argument(
"--linear-attn-decode-backend",
type=str,
choices=LINEAR_ATTN_KERNEL_BACKEND_CHOICES,
default=ServerArgs.linear_attn_decode_backend,
help="Override the kernel backend for linear attention decode. "
"If not set, uses --linear-attn-backend.",
)
parser.add_argument(
"--linear-attn-prefill-backend",
type=str,
choices=LINEAR_ATTN_KERNEL_BACKEND_CHOICES,
default=ServerArgs.linear_attn_prefill_backend,
help="Override the kernel backend for linear attention prefill/extend. "
"If not set, uses --linear-attn-backend.",
)
# Hierarchical cache # Hierarchical cache
parser.add_argument( parser.add_argument(