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

Support LingV2_5 model (#18598)

Browse Source 
Co-authored-by: zhangkaihong.zkh <zhangkaihong.zkh@antgroup.com>
Co-authored-by: 有禾 <zhangdonghao.zdh@antgroup.com>
Co-authored-by: yudian0504 <138860534+yudian0504@users.noreply.github.com>
Co-authored-by: 悠扬 <youyang.zmy@antgroup.com>
Co-authored-by: xinxingyang <xinxing.yangxx@antgroup.com>
Co-authored-by: zmy460290 <zmy460290@antgroup.com>
This commit is contained in:
ant-yy
2026-02-13 16:09:15 +08:00
committed by GitHub
parent 013a199bc6
commit d97eb111a3
16 changed files with 4042 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
python/sglang/srt/configs/__init__.py
View File

@@ -1,4 +1,5 @@
from sglang.srt.configs.afmoe import AfmoeConfig
from sglang.srt.configs.bailing_hybrid import BailingHybridConfig
from sglang.srt.configs.chatglm import ChatGLMConfig
from sglang.srt.configs.dbrx import DbrxConfig
from sglang.srt.configs.deepseekvl2 import DeepseekVL2Config
@@ -30,6 +31,7 @@ from sglang.srt.configs.step3p5 import Step3p5Config
__all__ = [
"AfmoeConfig",
"BailingHybridConfig",
"ExaoneConfig",
"ChatGLMConfig",
"DbrxConfig",

188
python/sglang/srt/configs/bailing_hybrid.py Normal file
View File

@@ -0,0 +1,188 @@
# coding=utf-8
# Copyright 2024 The Qwen team, Alibaba Group and the HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""BailingHybrid model configuration"""
import enum
from transformers.configuration_utils import PretrainedConfig
from transformers.utils import logging
from sglang.srt.configs.mamba_utils import Mamba2CacheParams, Mamba2StateShape
logger = logging.get_logger(__name__)
class HybridLayerType(enum.Enum):
full_attention = "attention"
linear_attention = "linear_attention"
class BailingHybridConfig(PretrainedConfig):
model_type = "bailing_hybrid"
keys_to_ignore_at_inference = ["past_key_values"]
def __init__(
self,
vocab_size=157184,
hidden_size=2048,
intermediate_size=5120,
num_hidden_layers=20,
num_attention_heads=16,
num_key_value_heads=4,
hidden_act="silu",
use_qkv_bias=False, # bailing only
use_bias=False, # bailing only
rms_norm_eps=1e-06,
tie_word_embeddings=False, # PretrainedConfig key, here change default value.
embedding_dropout=0.0,
attention_dropout=0.0,
output_dropout=0.0,
initializer_range=0.02,
max_position_embeddings=32768,
rope_theta=600000.0,
use_cache=True,
max_window_layers=20,
rope_scaling=None,
pad_token_id=156892,
eos_token_id=156892,
num_experts=256,
num_shared_experts=1,
num_experts_per_tok=8,
n_group=8,
topk_group=4,
moe_intermediate_size=512,
first_k_dense_replace=1,
head_dim=128,
output_router_logits=False,
use_qk_norm=True,
num_nextn_predict_layers=0,
mtp_loss_scaling_factor=0,
moe_router_enable_expert_bias=True,
routed_scaling_factor=1.0,
layer_group_size=1,
group_norm_size=1,
linear_silu=False,
kv_lora_rank=512,
q_lora_rank=None,
qk_rope_head_dim=64,
v_head_dim=128,
qk_nope_head_dim=128,
rope_interleave=True,
**kwargs,
):
self.num_hidden_layers = num_hidden_layers
self.vocab_size = vocab_size
self.hidden_size = hidden_size
self.intermediate_size = intermediate_size
self.num_attention_heads = num_attention_heads
self.num_key_value_heads = num_key_value_heads
self.hidden_act = hidden_act
self.use_qkv_bias = use_qkv_bias
self.use_bias = use_bias
self.rms_norm_eps = rms_norm_eps
self.embedding_dropout = embedding_dropout
self.attention_dropout = attention_dropout
self.output_dropout = output_dropout
self.num_nextn_predict_layers = num_nextn_predict_layers
self.mtp_loss_scaling_factor = mtp_loss_scaling_factor
self.initializer_range = initializer_range
self.max_position_embeddings = max_position_embeddings
self.rope_theta = rope_theta
self.use_cache = use_cache
self.max_window_layers = max_window_layers
self.head_dim = head_dim or self.hidden_size // self.num_attention_heads
self.rope_scaling = rope_scaling
self.use_qk_norm = use_qk_norm
self.moe_router_enable_expert_bias = moe_router_enable_expert_bias
self.routed_scaling_factor = routed_scaling_factor
# MoE configs
self.num_experts = num_experts
self.num_shared_experts = num_shared_experts
self.num_experts_per_tok = num_experts_per_tok
self.n_group = n_group
self.topk_group = topk_group
self.moe_intermediate_size = moe_intermediate_size
self.first_k_dense_replace = first_k_dense_replace
self.output_router_logits = output_router_logits
# Linear configs
self.layer_group_size = layer_group_size
self.group_norm_size = group_norm_size
self.linear_silu = linear_silu
self.num_linear_key_value_heads = num_attention_heads
# mla
self.kv_lora_rank = kv_lora_rank
self.q_lora_rank = q_lora_rank
self.qk_rope_head_dim = qk_rope_head_dim
self.v_head_dim = v_head_dim
self.qk_nope_head_dim = qk_nope_head_dim
self.qk_head_dim = qk_nope_head_dim + qk_rope_head_dim
self.rope_interleave = rope_interleave
self.for_nextn_model = False
super().__init__(
pad_token_id=pad_token_id,
eos_token_id=eos_token_id,
tie_word_embeddings=tie_word_embeddings,
**kwargs,
)
@property
def layers_block_type(self):
if self.for_nextn_model:
return [HybridLayerType.full_attention.value]
layer_type_list = []
for l in range(self.num_hidden_layers):
if (l + 1) % self.layer_group_size == 0:
layer_type_list.append(HybridLayerType.full_attention.value)
else:
layer_type_list.append(HybridLayerType.linear_attention.value)
return layer_type_list
@property
def linear_layer_ids(self):
return [
i
for i, type_value in enumerate(self.layers_block_type)
if type_value == HybridLayerType.linear_attention.value
]
@property
def full_attention_layer_ids(self):
return [
i
for i, type_value in enumerate(self.layers_block_type)
if type_value == HybridLayerType.full_attention.value
]
@property
def mamba2_cache_params(self) -> Mamba2CacheParams:
from sglang.srt.layers.dp_attention import get_attention_tp_size
shape = Mamba2StateShape.create(
tp_world_size=get_attention_tp_size(),
intermediate_size=0,
n_groups=0,
num_heads=self.num_linear_key_value_heads,
head_dim=self.head_dim,
state_size=self.head_dim,
conv_kernel=1,
)
return Mamba2CacheParams(shape=shape, layers=self.linear_layer_ids)

20
python/sglang/srt/configs/model_config.py
View File

@@ -308,6 +308,7 @@ class ModelConfig:
if is_draft_model and self.hf_config.architectures[0] in [
"BailingMoeV2ForCausalLM",
"BailingMoeForCausalLM",
"BailingMoeV2_5ForCausalLM",
]:
self.hf_config.architectures[0] = "BailingMoeForCausalLMNextN"
if (
@@ -482,6 +483,25 @@ class ModelConfig:
self.qk_rope_head_dim = self.hf_config.qk_rope_head_dim
self.v_head_dim = self.hf_config.v_head_dim
self.qk_nope_head_dim = self.hf_config.qk_nope_head_dim
elif (
"BailingMoeV2_5ForCausalLM" in self.hf_config.architectures
or "BailingMoeForCausalLMNextN" in self.hf_config.architectures
):
self.head_dim = self.hf_text_config.head_dim
self.attention_arch = AttentionArch.MLA
self.kv_lora_rank = self.hf_text_config.kv_lora_rank
self.qk_nope_head_dim = self.hf_text_config.qk_nope_head_dim
self.qk_rope_head_dim = self.hf_text_config.qk_rope_head_dim
self.v_head_dim = self.hf_config.v_head_dim
# Handle rope scaling with yarn
self.scaling = 1 / math.sqrt(self.qk_nope_head_dim + self.qk_rope_head_dim)
if self.hf_config.rope_scaling:
mscale_all_dim = self.hf_config.rope_scaling.get(
"mscale_all_dim", False
)
scaling_factor = self.hf_config.rope_scaling["factor"]
mscale = yarn_get_mscale(scaling_factor, float(mscale_all_dim))
self.scaling = self.scaling * mscale * mscale
else:
if (
"MistralModel" in self.hf_config.architectures

3
python/sglang/srt/layers/attention/attention_registry.py
View File

@@ -192,6 +192,7 @@ def attn_backend_wrapper(runner: "ModelRunner", full_attn_backend: "AttentionBac
GDNAttnBackend,
HybridLinearAttnBackend,
KimiLinearAttnBackend,
LightningAttentionBackend,
Mamba2AttnBackend,
)
from sglang.srt.utils import is_blackwell, is_npu
@@ -213,6 +214,8 @@ def attn_backend_wrapper(runner: "ModelRunner", full_attn_backend: "AttentionBac
linear_attn_backend = Mamba2AttnBackend(runner)
elif runner.kimi_linear_config is not None:
linear_attn_backend = KimiLinearAttnBackend(runner)
elif runner.hybrid_lightning_config is not None:
linear_attn_backend = LightningAttentionBackend(runner)
else:
raise ValueError(
"Expected hybrid GDN or NemotronH models, but got unknown model."

31
python/sglang/srt/layers/attention/fla/layernorm_gated.py
View File

@@ -81,6 +81,7 @@ def _layer_norm_fwd_1pass_kernel(
HAS_Z: tl.constexpr,
NORM_BEFORE_GATE: tl.constexpr,
IS_RMS_NORM: tl.constexpr,
ACTIVATION: tl.constexpr,
):
# Map the program id to the starting row of X and Y it should compute.
row_start = tl.program_id(0) * ROWS_PER_BLOCK
@@ -109,7 +110,10 @@ def _layer_norm_fwd_1pass_kernel(
if HAS_Z and not NORM_BEFORE_GATE:
Z_base = Z + rows[:, None] * stride_z_row + col_offsets
z = tl.load(Z_base, mask=mask, other=0.0).to(tl.float32)
x *= z * tl.sigmoid(z)
if ACTIVATION == "swish" or ACTIVATION == "silu":
x *= z * tl.sigmoid(z)
elif ACTIVATION == "sigmoid":
x *= tl.sigmoid(z)
# Compute mean and variance per row (reduce along axis 1)
if not IS_RMS_NORM:
@@ -152,7 +156,10 @@ def _layer_norm_fwd_1pass_kernel(
if HAS_Z and NORM_BEFORE_GATE:
Z_base = Z + rows[:, None] * stride_z_row + col_offsets
z = tl.load(Z_base, mask=mask, other=0.0).to(tl.float32)
y *= z * tl.sigmoid(z)
if ACTIVATION == "swish" or ACTIVATION == "silu":
y *= z * tl.sigmoid(z)
elif ACTIVATION == "sigmoid":
y *= tl.sigmoid(z)
# Write output
tl.store(Y_base, y, mask=mask)
@@ -182,6 +189,7 @@ def _layer_norm_fwd(
group_size=None,
norm_before_gate=True,
is_rms_norm=False,
activation: str = "swish",
):
M, N = x.shape
if group_size is None:
@@ -242,6 +250,7 @@ def _layer_norm_fwd(
NORM_BEFORE_GATE=norm_before_gate,
IS_RMS_NORM=is_rms_norm,
num_warps=num_warps,
ACTIVATION=activation,
)
return out, mean, rstd
@@ -260,6 +269,7 @@ def rms_norm_gated(
group_size=None,
norm_before_gate=True,
is_rms_norm=False,
activation: str = "swish",
):
"""If z is not None, we do norm(x) * silu(z) if norm_before_gate, else norm(x * silu(z))"""
@@ -276,6 +286,8 @@ def rms_norm_gated(
weight = weight.contiguous()
if bias is not None:
bias = bias.contiguous()
if _is_npu:
assert activation == "swish", "NPU only supports swish activation"
y, mean, rstd = _layer_norm_fwd(
x,
weight,
@@ -285,6 +297,7 @@ def rms_norm_gated(
group_size=group_size,
norm_before_gate=norm_before_gate,
is_rms_norm=is_rms_norm,
activation=activation,
)
return y.reshape(x_shape_og)
@@ -302,6 +315,7 @@ class LayerNormFn(torch.autograd.Function):
group_size=None,
norm_before_gate=True,
is_rms_norm=False,
activation: str = "swish",
):
return rms_norm_gated(
x=x,
@@ -312,6 +326,7 @@ class LayerNormFn(torch.autograd.Function):
group_size=group_size,
norm_before_gate=norm_before_gate,
is_rms_norm=is_rms_norm,
activation=activation,
)
@@ -324,9 +339,10 @@ def layernorm_fn(
group_size=None,
norm_before_gate=True,
is_rms_norm=False,
activation: str = "swish",
):
return LayerNormFn.apply(
x, weight, bias, z, eps, group_size, norm_before_gate, is_rms_norm
x, weight, bias, z, eps, group_size, norm_before_gate, is_rms_norm, activation
)
@@ -382,6 +398,7 @@ class RMSNorm(torch.nn.Module):
norm_before_gate=True,
device=None,
dtype=None,
activation: str = "swish",
):
"""If group_size is not None, we do GroupNorm with each group having group_size elements.
group_size=None is equivalent to group_size=hidden_size (i.e. there's only 1 group).
@@ -389,6 +406,7 @@ class RMSNorm(torch.nn.Module):
factory_kwargs = {"device": device, "dtype": dtype}
super().__init__()
self.eps = eps
self.activation = activation
self.weight = torch.nn.Parameter(torch.empty(hidden_size, **factory_kwargs))
self.register_parameter("bias", None)
self.group_size = group_size
@@ -402,8 +420,10 @@ class RMSNorm(torch.nn.Module):
"""If z is not None, we do norm(x) * silu(z) if norm_before_gate, else norm(x * silu(z))"""
if _use_cpu:
assert (
self.norm_before_gate and self.group_size is None
), "CPU rmsnorm_gated currently only supports norm before gate without group size"
self.norm_before_gate
and self.group_size is None
and self.activation == "swish"
), "CPU rmsnorm_gated currently only supports norm before gate without group size or activation other than swish"
return torch.ops.sgl_kernel.fused_rmsnorm_gated_cpu(
x, self.weight, z, self.eps
)
@@ -417,4 +437,5 @@ class RMSNorm(torch.nn.Module):
group_size=self.group_size,
norm_before_gate=self.norm_before_gate,
is_rms_norm=True,
activation=self.activation,
)

369
python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py
View File

@@ -1,3 +1,5 @@
import logging
import math
from typing import Optional, Tuple, Union
import torch
@@ -14,6 +16,12 @@ from sglang.srt.layers.attention.fla.fused_recurrent import (
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,
@@ -84,6 +92,8 @@ elif is_cpu():
)
fused_gdn_gating = torch.ops.sgl_kernel.fused_gdn_gating_cpu
logger = logging.getLogger(__name__)
# Kernel to track mamba states if needed based on track mask
@triton.jit
@@ -1125,6 +1135,365 @@ 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):
"""Manages a full and linear attention backend"""

767
python/sglang/srt/layers/attention/linear/lightning_attn.py Normal file
View File

@@ -0,0 +1,767 @@
# Adapted from https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/layers/mamba/linear_attn.py
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import torch
import triton
import triton.language as tl
from einops import rearrange
@triton.jit
def _fwd_diag_kernel(
Q,
K,
V,
Out,
S,
b: tl.constexpr,
h: tl.constexpr,
n,
d: tl.constexpr,
e: tl.constexpr,
BLOCK: tl.constexpr,
NUM_BLOCK,
CBLOCK: tl.constexpr,
):
# This kernel computes the diagonal blocks of the attention matrix
# Each diagonal block represents attention
# where queries attend to keys in the same block
off = tl.program_id(0)
off_bh = off // NUM_BLOCK # batch-head index
off_block = off % NUM_BLOCK # block index within the sequence
off_cblock = tl.program_id(1) # sub-block index within a block
off_h = off_bh % h # head index
# Calculate base offsets for the current batch and head
qk_offset = off_bh * n * d
v_offset = off_bh * n * e
o_offset = off_bh * n * e
# Calculate offsets for the current block
block_offset = off_block * BLOCK
qk_block_offset = block_offset * d
v_block_offset = block_offset * e
o_block_offset = block_offset * e
# Calculate offsets for the current sub-block
cblock_offset = off_cblock * CBLOCK
q_cblock_offset = cblock_offset * d
o_cblock_offset = cblock_offset * e
# Calculate pointers to the query, key, value, and output tensors
Q_block_ptr = (
Q
+ qk_offset
+ qk_block_offset
+ q_cblock_offset
+ tl.arange(0, CBLOCK)[:, None] * d
+ tl.arange(0, d)[None, :]
)
K_trans_block_ptr = (
K
+ qk_offset
+ qk_block_offset
+ tl.arange(0, CBLOCK)[None, :] * d
+ tl.arange(0, d)[:, None]
)
V_block_ptr = (
V
+ v_offset
+ v_block_offset
+ tl.arange(0, CBLOCK)[:, None] * e
+ tl.arange(0, e)[None, :]
)
O_block_ptr = (
Out
+ o_offset
+ o_block_offset
+ o_cblock_offset
+ tl.arange(0, CBLOCK)[:, None] * e
+ tl.arange(0, e)[None, :]
)
# Load the decay rate for the current head
S_block_ptr = S + off_h
s = tl.load(S_block_ptr)
i = off_cblock
q_index = tl.arange(0, CBLOCK) + i * CBLOCK
# Load query values
q = tl.load(Q_block_ptr, mask=block_offset + q_index[:, None] < n, other=0.0).to(
tl.float32
)
# Initialize output accumulator
qkv = tl.zeros([CBLOCK, e], dtype=tl.float32)
# Process all sub-blocks up to and
# including the current one (causal attention)
for j in range(i + 1):
kv_index = tl.arange(0, CBLOCK) + j * CBLOCK
diff = q_index[:, None] - kv_index[None, :]
s_index = s * diff
# Apply causal mask: only attend to positions before the current one
s_index = tl.where(diff >= 0, -s_index, float("-inf"))
decay = tl.exp(s_index)
# Load key and value
k_trans = tl.load(
K_trans_block_ptr,
mask=block_offset + kv_index[None, :] < n,
other=0.0,
).to(tl.float32)
v = tl.load(
V_block_ptr,
mask=block_offset + kv_index[:, None] < n,
other=0.0,
).to(tl.float32)
# Compute attention scores and apply decay
qk = tl.dot(q, k_trans) * decay
# Compute weighted values and accumulate
qkv += tl.dot(qk, v)
# Move to the next sub-block
K_trans_block_ptr += CBLOCK * d
V_block_ptr += CBLOCK * e
# Store the result
tl.store(
O_block_ptr,
qkv.to(O_block_ptr.dtype.element_ty),
mask=block_offset + q_index[:, None] < n,
)
@triton.jit
def _fwd_kv_parallel(
K,
V,
K_decay,
KV,
b: tl.constexpr,
h: tl.constexpr,
n,
d: tl.constexpr,
e: tl.constexpr,
BLOCK: tl.constexpr,
NUM_BLOCK,
D_FBLOCK: tl.constexpr,
E_FBLOCK: tl.constexpr,
NUM_FBLOCK: tl.constexpr,
CBLOCK: tl.constexpr,
NUM_CBLOCK: tl.constexpr,
):
# This kernel computes the key-value outer
# products for each block in parallel
off_bh = tl.program_id(0) # batch-head index
off_block = tl.program_id(1) # block index
off_h = off_bh % h # head index
block_offset = off_block * BLOCK
# Calculate offsets for the current block
k_block_offset = block_offset * d
v_block_offset = block_offset * e
kv_block_offset = off_block * d * e
# Calculate base offsets for the current batch and head
k_offset = off_bh * n * d
v_offset = off_bh * n * e
kv_offset = off_bh * NUM_BLOCK * d * e
# Calculate pointers to the key, value, and key-value tensors
K_trans_block_ptr = (
K
+ k_offset
+ k_block_offset
+ tl.arange(0, CBLOCK)[None, :] * d
+ tl.arange(0, D_FBLOCK)[:, None]
)
V_block_ptr = (
V
+ v_offset
+ v_block_offset
+ tl.arange(0, CBLOCK)[:, None] * e
+ tl.arange(0, E_FBLOCK)[None, :]
)
KV_block_ptr = (
KV
+ kv_offset
+ kv_block_offset
+ tl.arange(0, D_FBLOCK)[:, None] * e
+ tl.arange(0, E_FBLOCK)[None, :]
)
# Load the decay factors for the current head and block
k_decay_ptr = K_decay + off_h * BLOCK + tl.arange(0, CBLOCK)[None, :]
kv_index = tl.arange(0, CBLOCK)
# Initialize the key-value outer product accumulator
kv = tl.zeros([D_FBLOCK, E_FBLOCK], dtype=tl.float32)
# Handle the last block which might be smaller than BLOCK
if off_block == NUM_BLOCK - 1:
split_n = n - (NUM_BLOCK - 1) * BLOCK
else:
split_n = BLOCK
left_shift = tl.cdiv(split_n, CBLOCK) * CBLOCK - split_n
num_blocks = min(tl.cdiv(split_n, CBLOCK), NUM_CBLOCK)
k_decay_ptr += (NUM_CBLOCK - num_blocks) * CBLOCK
# Process all sub-blocks in the current block
for j in range(num_blocks):
left_bound = (1 - j) * left_shift
# Load key and value, handling boundary conditions
k_trans = tl.load(
K_trans_block_ptr - left_shift * d,
mask=kv_index[None, :] >= left_bound,
other=0.0,
)
v = tl.load(
V_block_ptr - left_shift * e,
mask=kv_index[:, None] >= left_bound,
other=0.0,
)
# Load decay factor and compute weighted key-value outer product
k_decay = tl.load(k_decay_ptr)
kv += tl.dot(k_trans * k_decay, v)
# Move to the next sub-block
K_trans_block_ptr += CBLOCK * d
V_block_ptr += CBLOCK * e
k_decay_ptr += CBLOCK
# Store the result
tl.store(KV_block_ptr, kv.to(KV_block_ptr.dtype.element_ty))
@triton.jit
def _fwd_kv_reduce(
S,
KV,
KV_HISTORY,
b: tl.constexpr,
h: tl.constexpr,
n,
d: tl.constexpr,
e: tl.constexpr,
BLOCK: tl.constexpr,
NUM_BLOCK,
D_FBLOCK: tl.constexpr,
E_FBLOCK: tl.constexpr,
):
# This kernel reduces the key-value outer products
# across blocks and updates the KV history
off_bh = tl.program_id(0) # batch-head index
off_h = off_bh % h # head index
kv_offset = off_bh * NUM_BLOCK * d * e
# Calculate pointer to the key-value tensor
KV_block_ptr = (
KV
+ kv_offset
+ tl.arange(0, D_FBLOCK)[:, None] * e
+ tl.arange(0, E_FBLOCK)[None, :]
)
# Load the decay rate for the current head
s_ptrs = S + off_h
s = tl.load(s_ptrs)
# Calculate pointer to the key-value history tensor
kv_history_offset = off_bh * d * e
KV_HISTORY_block_ptr = (
KV_HISTORY
+ kv_history_offset
+ tl.arange(0, D_FBLOCK)[:, None] * e
+ tl.arange(0, E_FBLOCK)[None, :]
)
# Load the previous key-value history
kv_pre = tl.load(KV_HISTORY_block_ptr).to(tl.float32)
# Process all blocks in reverse order to compute the prefix sum
for i in range(NUM_BLOCK):
block_size = min(n - i * BLOCK, BLOCK)
# Compute decay factor for the current block
block_decay = tl.exp(-s.to(tl.float32) * block_size)
# Load the current key-value outer product
kv_cur = tl.load(KV_block_ptr).to(tl.float32)
# Store the previous key-value history to the current block
tl.store(KV_block_ptr, kv_pre.to(KV_block_ptr.dtype.element_ty))
# Update the key-value history with the current block
kv_pre = block_decay * kv_pre + kv_cur
KV_block_ptr += d * e
# Store the updated key-value history
tl.store(KV_HISTORY_block_ptr, kv_pre)
@triton.jit
def _fwd_none_diag_kernel(
Q,
Out,
S,
KV,
b: tl.constexpr,
h: tl.constexpr,
n,
d: tl.constexpr,
e: tl.constexpr,
BLOCK: tl.constexpr,
NUM_BLOCK,
E_FBLOCK: tl.constexpr,
CBLOCK: tl.constexpr,
NUM_CBLOCK: tl.constexpr,
):
# This kernel computes the non-diagonal blocks of the attention matrix
# Each non-diagonal block represents attention
# where queries attend to keys in different blocks
off_bh = tl.program_id(0) # batch-head index
off_h = off_bh % h # head index
off_nc = tl.program_id(1)
off_n = off_nc // NUM_CBLOCK # block index
off_c = off_nc % NUM_CBLOCK # sub-block index
off_e = tl.program_id(2) # output feature block index
n_offset = off_n * BLOCK
c_offset = off_c * CBLOCK
e_offset = off_e * E_FBLOCK
block_offset = n_offset + c_offset
# Calculate offsets for the current batch, head, and block
q_offset = off_bh * n * d + (n_offset + c_offset) * d
o_offset = off_bh * n * e + (n_offset + c_offset) * e + e_offset
kv_offset = off_bh * NUM_BLOCK * d * e + off_n * d * e + e_offset
# Calculate pointers to the query, output, and key-value tensors
Q_block_ptr = (
Q + q_offset + tl.arange(0, CBLOCK)[:, None] * d + tl.arange(0, d)[None, :]
)
O_block_ptr = (
Out
+ o_offset
+ tl.arange(0, CBLOCK)[:, None] * e
+ tl.arange(0, E_FBLOCK)[None, :]
)
KV_block_ptr = (
KV + kv_offset + tl.arange(0, d)[:, None] * e + tl.arange(0, E_FBLOCK)[None, :]
)
# Load the decay rate for the current head
S_block_ptr = S + off_h
s = tl.load(S_block_ptr)
c_array = tl.arange(0, CBLOCK)
# Load the key-value outer product for the current block
kv = tl.load(KV_block_ptr).to(tl.float32)
q_index = block_offset + tl.arange(0, CBLOCK)
# Load query values
q = tl.load(Q_block_ptr, mask=q_index[:, None] < n, other=0.0).to(tl.float32)
# Compute decay factors for the current sub-block
q_decay = tl.exp(-s.to(tl.float32) * (off_c * CBLOCK + c_array[:, None]))
# Compute non-diagonal attention output
qkv_none_diag = tl.dot(q, kv) * q_decay
# Load diagonal attention output (computed by _fwd_diag_kernel)
qkv_diag = tl.load(O_block_ptr, mask=q_index[:, None] < n, other=0.0).to(tl.float32)
# Combine diagonal and non-diagonal attention outputs
qkv = qkv_diag + qkv_none_diag
# Store the result
tl.store(
O_block_ptr, qkv.to(O_block_ptr.dtype.element_ty), mask=q_index[:, None] < n
)
class _attention(torch.autograd.Function):
@staticmethod
def forward(ctx, q, k, v, s, kv_history):
# Forward pass of the lightning attention algorithm
q = q.contiguous()
k = k.contiguous()
v = v.contiguous()
s = s.contiguous()
# Check CUDA compute capability
capability = torch.cuda.get_device_capability()
if capability[0] < 8:
raise RuntimeError(
"Flash attention currently only supported",
"for compute capability >= 80",
)
# Get input dimensions
b, h, n, d = q.shape
e = v.shape[-1]
# Initialize output tensor
o = torch.empty((b, h, n, e), dtype=q.dtype, device=q.device)
# Set block sizes
BLOCK = 256
NUM_BLOCK = triton.cdiv(n, BLOCK)
CBLOCK = 32
NUM_CBLOCK = BLOCK // CBLOCK
assert BLOCK % CBLOCK == 0, "BLOCK must be a multiple of CBLOCK"
# Compute decay factors for keys
array = torch.arange(0, BLOCK, device=q.device) + 1
k_decay = torch.exp(-s * (BLOCK - array.reshape(1, -1)))
# Step 1: Compute diagonal blocks of attention
grid = (b * h * NUM_BLOCK, NUM_CBLOCK)
_fwd_diag_kernel[grid](
q,
k,
v,
o,
s,
b,
h,
n,
d,
e,
BLOCK=BLOCK,
NUM_BLOCK=NUM_BLOCK,
CBLOCK=CBLOCK,
)
# Set feature block sizes
NUM_FBLOCK = 1
D_FBLOCK = d // NUM_FBLOCK
assert d % NUM_FBLOCK == 0
E_FBLOCK = e // NUM_FBLOCK
assert e % NUM_FBLOCK == 0
CBLOCK = 64
NUM_CBLOCK = BLOCK // CBLOCK
assert BLOCK % CBLOCK == 0, "BLOCK must be a multiple of CBLOCK"
# Step 2: Compute key-value outer products for each block in parallel
kv = torch.empty((b, h, NUM_BLOCK, d, e), dtype=torch.float32, device=q.device)
grid = (b * h, NUM_BLOCK)
_fwd_kv_parallel[grid](
k,
v,
k_decay,
kv,
b,
h,
n,
d,
e,
BLOCK=BLOCK,
NUM_BLOCK=NUM_BLOCK,
D_FBLOCK=D_FBLOCK,
E_FBLOCK=E_FBLOCK,
NUM_FBLOCK=NUM_FBLOCK,
CBLOCK=CBLOCK,
NUM_CBLOCK=NUM_CBLOCK,
)
# Step 3: Reduce key-value outer products
# across blocks and update KV history
grid = (b * h, NUM_FBLOCK)
_fwd_kv_reduce[grid](
s,
kv,
kv_history,
b,
h,
n,
d,
e,
BLOCK=BLOCK,
NUM_BLOCK=NUM_BLOCK,
D_FBLOCK=D_FBLOCK,
E_FBLOCK=E_FBLOCK,
)
# Step 4: Compute non-diagonal blocks of attention
grid = (b * h, NUM_BLOCK * NUM_CBLOCK)
_fwd_none_diag_kernel[grid](
q,
o,
s,
kv,
b,
h,
n,
d,
e,
BLOCK=BLOCK,
NUM_BLOCK=NUM_BLOCK,
E_FBLOCK=E_FBLOCK,
CBLOCK=CBLOCK,
NUM_CBLOCK=NUM_CBLOCK,
)
# Save tensors for backward pass
ctx.save_for_backward(q, k, v, s, kv)
ctx.BLOCK = BLOCK
return o, torch.cat([kv, kv_history.unsqueeze(2)], dim=2)
# Apply the lightning attention function
lightning_attention_ = _attention.apply
def lightning_attention(q, k, v, ed, block_size=256, kv_history=None):
"""
Apply lightning attention algorithm
to compute attention efficiently.
Args:
q: Query tensor of shape [batch, heads, seq_len, dim]
k: Key tensor of shape [batch, heads, seq_len, dim]
v: Value tensor of shape [batch, heads, seq_len, dim_v]
ed: Decay rate tensor of shape [heads]
block_size: Size of blocks for block-sparse attention
kv_history: Optional key-value history from previous computations
Returns:
output: Attention output
kv: Updated key-value history
"""
d = q.shape[-1]
e = v.shape[-1]
if ed.dim() == 1:
ed = ed.view(1, -1, 1, 1)
# Split the computation into chunks for better parallelism
m = 128 if d >= 128 else 64
assert d % m == 0, f"Dimension d ({d}) must be divisible by m ({m})"
arr = [m * i for i in range(d // m + 1)]
if arr[-1] != d:
arr.append(d)
n = len(arr)
output = 0
# Initialize or clone key-value history
if kv_history is None:
kv_history = torch.zeros(
(q.shape[0], q.shape[1], d, e), dtype=torch.float32, device=q.device
)
else:
kv_history = kv_history.clone().contiguous()
# Process each chunk and accumulate results
for i in range(n - 1):
s = arr[i]
e = arr[i + 1]
q1 = q[..., s:e]
k1 = k[..., s:e]
o, kv = lightning_attention_(q1, k1, v, ed, kv_history)
output = output + o
return output, kv
@triton.jit
def _linear_attn_decode_kernel(
q_ptr,
k_ptr,
v_ptr,
kv_cache_ptr,
slope_rate,
slot_idx,
output_ptr,
D: tl.constexpr,
qkv_b_stride,
qkv_h_stride,
cache_b_stride,
cache_h_stride,
cache_d0_stride,
cache_d1_stride,
BLOCK_SIZE: tl.constexpr,
):
"""
Kernel for linear attention decoding with KV cache.
This kernel computes attention for a single token using the KV cache.
"""
pid_b = tl.program_id(0) # batch index
pid_h = tl.program_id(1) # head index
pid_d = tl.program_id(2) # dimension block index
# Load slot index for the current batch
slot_id = tl.load(slot_idx + pid_b)
# Skip if slot_id is -1 (padding)
if slot_id == -1:
return
batch_id = pid_b
head_id = pid_h
# Load decay rate for the current head
ratio = tl.load(slope_rate + pid_h)
# Calculate offsets for dimensions
qk_d_offsets = tl.arange(0, D)
v_d_offsets = tl.arange(0, BLOCK_SIZE) + pid_d * BLOCK_SIZE
cache_d_offsets = (
qk_d_offsets[:, None] * cache_d0_stride + v_d_offsets[None, :] * cache_d1_stride
)
# Calculate offsets for the current batch and head
q_offset = batch_id * qkv_b_stride + head_id * qkv_h_stride
k_offset = batch_id * qkv_b_stride + head_id * qkv_h_stride
v_offset = batch_id * qkv_b_stride + head_id * qkv_h_stride
cache_offset = slot_id * cache_b_stride + head_id * cache_h_stride
# Create masks for loading tensors
qk_mask = qk_d_offsets < D
v_mask = v_d_offsets < D
# Load query, key, and value tensors
q = tl.load(q_ptr + q_offset + qk_d_offsets, mask=qk_mask, other=0.0)
k = tl.load(k_ptr + k_offset + qk_d_offsets, mask=qk_mask, other=0.0)
v = tl.load(v_ptr + v_offset + v_d_offsets, mask=v_mask, other=0.0)
# Compute key-value outer product
kv_outer = k[:, None] * v[None, :]
kv_mask = qk_mask[:, None] & v_mask[None, :]
# Apply decay to previous KV cache
ratio = tl.exp(-ratio)
kv_ptr = kv_cache_ptr + cache_offset + cache_d_offsets
kv_cache_old = tl.load(kv_ptr, mask=kv_mask, other=0.0)
kv_outer = kv_outer + ratio * kv_cache_old
# Compute attention output
output = q[:, None].to(tl.float32) * kv_outer
output = tl.sum(output, axis=0)
# Update KV cache and store output
tl.store(kv_ptr, kv_outer, mask=kv_mask)
tl.store(output_ptr + q_offset + v_d_offsets, output, mask=v_mask)
def linear_decode_forward_triton(
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
kv_caches: torch.Tensor,
slope_rate: torch.Tensor,
slot_idx: torch.Tensor,
BLOCK_SIZE: int = 32,
) -> torch.Tensor:
"""
Perform linear attention decoding using Triton kernels.
Args:
q: Query tensor of shape [B, H, 1, D]
k: Key tensor of shape [B, H, 1, D]
v: Value tensor of shape [B, H, 1, D]
kv_caches: Key-value cache tensor
slope_rate: Decay rate tensor
slot_idx: Slot indices for batches
BLOCK_SIZE: Size of blocks for processing
Returns:
output: Attention output tensor
"""
B, H, _, D = q.shape
assert k.shape == (B, H, 1, D)
assert v.shape == (B, H, 1, D)
# Initialize output tensor
output = torch.empty_like(q)
# Set grid dimensions for the kernel
grid = (B, H, D // BLOCK_SIZE)
# Calculate strides for tensors
qkv_b_stride = q.stride(0)
qkv_h_stride = q.stride(1)
cache_b_stride = kv_caches.stride(0)
cache_h_stride = kv_caches.stride(1)
cache_d0_stride = kv_caches.stride(2)
cache_d1_stride = kv_caches.stride(3)
# Launch the kernel
_linear_attn_decode_kernel[grid](
q,
k,
v,
kv_caches,
slope_rate,
slot_idx,
output,
D,
qkv_b_stride,
qkv_h_stride,
cache_b_stride,
cache_h_stride,
cache_d0_stride,
cache_d1_stride,
BLOCK_SIZE=BLOCK_SIZE,
)
# Reshape output and return
output = rearrange(output, "b h n d -> b n (h d)")
return output.squeeze(1).contiguous()
class BailingLinearKernel:
"""
Linear attention kernel implementation for Bailing models.
This class is adapted from MiniMaxText01LinearKernel in vllm:
https://github.com/vllm-project/vllm/blob/a9138e85b14047e06300685b48e3485b995425fb/vllm/model_executor/models/minimax_text_01.py#L289
The implementation maintains the same functionality while being renamed to
match our Bailing model naming convention.
"""
@staticmethod
def jit_linear_forward_prefix(
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
kv_caches: torch.Tensor,
slope_rate: torch.Tensor,
block_size: int,
layer_idx: int = None,
**kwargs,
) -> torch.Tensor:
slope_rate = slope_rate.to(torch.float32)
should_pad_dim = q.dim() == 3
if should_pad_dim:
q = q.unsqueeze(0)
k = k.unsqueeze(0)
v = v.unsqueeze(0)
b, h, n, d = q.shape
e = d
kv_history = kv_caches.reshape(1, h, d, e).contiguous()
output, kv_history = lightning_attention(
q, k, v, slope_rate, block_size=block_size, kv_history=kv_history
)
kv_caches.copy_(kv_history[:, :, -1, :, :].reshape(h, d, e))
assert output.shape[0] == 1, "batch size must be 1"
return output.squeeze(0).transpose(0, 1).reshape([n, h * d]).contiguous()

70
python/sglang/srt/layers/attention/linear/linear_metadata.py Normal file
View File

@@ -0,0 +1,70 @@
from dataclasses import dataclass
import torch
from sglang.srt.layers.attention.mamba.mamba2_metadata import ForwardMetadata
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
@dataclass(kw_only=True)
class BailingLinearMetadata(ForwardMetadata):
num_prefills: int
num_prefill_tokens: int
num_decodes: int
batch_size: int
has_initial_states: torch.Tensor
q_lengths: torch.Tensor
@staticmethod
def prepare_decode(
query_start_loc: torch.Tensor,
mamba_cache_indices: torch.Tensor,
bs: int,
seq_lens: torch.Tensor,
) -> "BailingLinearMetadata":
"""This path is run during CUDA graph capture, i.e. decode only, so `num_prefills` is 0"""
return BailingLinearMetadata(
batch_size=bs,
query_start_loc=query_start_loc,
mamba_cache_indices=mamba_cache_indices,
num_decodes=seq_lens.shape[0],
num_prefills=0,
num_prefill_tokens=0,
has_initial_states=torch.ones_like(seq_lens),
q_lengths=query_start_loc.diff(),
)
@classmethod
def prepare_mixed(
cls,
query_start_loc: torch.Tensor,
mamba_cache_indices: torch.Tensor,
forward_batch: ForwardBatch,
) -> "BailingLinearMetadata":
"""This path cannot run with CUDA graph, as it contains extend requests."""
if forward_batch.extend_num_tokens is None:
return cls.prepare_decode(
query_start_loc=query_start_loc,
mamba_cache_indices=mamba_cache_indices,
bs=forward_batch.batch_size,
seq_lens=forward_batch.seq_lens,
)
num_prefills = len(forward_batch.extend_seq_lens)
num_prefill_tokens = forward_batch.extend_num_tokens
num_decodes = len(forward_batch.seq_lens) - num_prefills
context_lens_tensor = forward_batch.extend_prefix_lens
assert context_lens_tensor is not None
has_initial_states = context_lens_tensor > 0
query_start_loc = query_start_loc[: num_prefills + 1]
return BailingLinearMetadata(
batch_size=forward_batch.batch_size,
query_start_loc=query_start_loc,
mamba_cache_indices=mamba_cache_indices,
num_prefills=num_prefills,
num_prefill_tokens=num_prefill_tokens,
num_decodes=num_decodes,
has_initial_states=has_initial_states,
q_lengths=query_start_loc.diff(),
)

909
python/sglang/srt/layers/attention/linear/seg_la.py Normal file
View File

@@ -0,0 +1,909 @@
# -*- coding: utf-8 -*-
"""
Copyright (c) Ant Financial Service Group and its affiliates.
"""
# Copied from https://code.alipay.com/pia/PainlessInferenceAcceleration/blob/v0.0.6/flood/flood/ops/seg_la.py
from dataclasses import dataclass
from typing import Optional
import torch
import triton
import triton.language as tl
# arg `meta` of `seg_la_fwd` is SegLaMeta
@dataclass
class SegLaMeta:
batch_size: int # batch size, num of requests
max_q_length: int # max(seq_lens)
q_offsets: torch.Tensor # [bs+1], query_start_locations,
s_offsets: torch.Tensor # [bs], slot_ids
q_lengths: torch.Tensor # [bs], query length
s_scales: torch.Tensor # [bs], prefill = 0, decode = 1
s_offsets_stride: int = 0
q_offsets_stride: int = 0
s_scales_stride: int = 0
decay_scales_stride: int = 0
mask: Optional[torch.Tensor] = None # Currently not supported
# fused
@triton.jit
def seg_la_kernel(
Q,
K,
V,
S,
Out,
softmax_scale,
stride_q,
stride_k,
stride_v,
stride_s,
stride_o,
s_offsets,
q_offsets,
q_lengths,
s_scales,
decay_scales,
HEAD_DIM: tl.constexpr,
SPLIT_DIM: tl.constexpr,
BLOCK: tl.constexpr,
EVEN: tl.constexpr,
DECOUPLE: tl.constexpr,
):
bid = tl.program_id(0)
hid = tl.program_id(1)
sid = tl.program_id(2)
# s_scale is 0 (prefill) or 1 (decode)
s_scale = tl.load(s_scales + bid)
q_length = tl.load(q_lengths + bid)
q_offset = tl.load(q_offsets + bid)
s_offset = tl.load(s_offsets + bid)
decay_scale = -tl.load(decay_scales + hid)
offs_b = tl.arange(0, BLOCK)
offs_d = tl.arange(0, HEAD_DIM)
offs_s = tl.arange(0, SPLIT_DIM)
if s_offset == -1:
return
q_ptrs = (
Q
+ q_offset * stride_q
+ hid * HEAD_DIM
+ (offs_b[:, None] * stride_q + offs_d[None, :])
)
k_ptrs = (
K
+ q_offset * stride_k
+ hid * HEAD_DIM
+ (offs_b[:, None] * stride_k + offs_d[None, :])
)
v_ptrs = (
V
+ q_offset * stride_v
+ hid * HEAD_DIM
+ sid * SPLIT_DIM
+ (offs_b[:, None] * stride_v + offs_s[None, :])
)
out_ptrs = (
Out
+ q_offset * stride_o
+ hid * HEAD_DIM
+ sid * SPLIT_DIM
+ (offs_b[:, None] * stride_o + offs_s[None, :])
)
s_ptrs = (
S
+ s_offset * stride_s
+ hid * HEAD_DIM * HEAD_DIM
+ sid * SPLIT_DIM
+ (offs_d[:, None] * HEAD_DIM + offs_s[None, :])
)
state = tl.load(s_ptrs, mask=s_scale > 0).to(tl.float32)
if BLOCK > 1:
for n in range(0, q_length, BLOCK):
n = tl.multiple_of(n, BLOCK)
if EVEN:
q = tl.load(q_ptrs + n * stride_q).to(tl.float32)
k = tl.trans(tl.load(k_ptrs + n * stride_k)).to(tl.float32)
v = tl.load(v_ptrs + n * stride_k).to(tl.float32)
else:
q = tl.load(
q_ptrs + n * stride_q,
mask=(n + offs_b)[:, None] < q_length,
other=0.0,
).to(tl.float32)
k = tl.trans(
tl.load(
k_ptrs + n * stride_k,
mask=(n + offs_b)[:, None] < q_length,
other=0.0,
)
).to(tl.float32)
v = tl.load(
v_ptrs + n * stride_k,
mask=(n + offs_b)[:, None] < q_length,
other=0.0,
).to(tl.float32)
if DECOUPLE:
# only work with small scales
if EVEN:
b = BLOCK
else:
b = min(BLOCK, q_length - n)
b_offs = b - 1 - offs_b
edb = tl.exp(decay_scale * b_offs)
decays = tl.where(b_offs >= 0, edb, 0)
inv_decays = tl.where(b_offs >= 0, 1 / edb, 0)
q = q * inv_decays[:, None]
k = k * decays[None, :]
qk = tl.dot(q, k) * softmax_scale
qk = tl.where(offs_b[None, :] <= offs_b[:, None], qk, 0.0)
o = tl.dot(qk, v)
block_decay = tl.exp(decay_scale * b)
block_decay_plus = block_decay * softmax_scale
o = tl.dot(q, state) * block_decay_plus + o
state = state * block_decay + tl.dot(k, v)
else:
qk = tl.dot(q, k) * softmax_scale
decays = tl.exp(decay_scale * (offs_b[:, None] - offs_b[None, :]))
decays = tl.where(offs_b[None, :] <= offs_b[:, None], decays, 0.0)
qk *= decays
o = tl.dot(qk, v)
decay_arr = tl.exp(decay_scale * (offs_b[:, None] + 1)) * softmax_scale
o = tl.dot(q * decay_arr, state, acc=o)
if EVEN:
b = BLOCK
else:
b = min(BLOCK, q_length - n)
b_offs = b - 1 - offs_b
b_offs = tl.where(b_offs >= 0, b_offs, 10000)
decays = tl.exp(decay_scale * b_offs)
block_decay = tl.exp(decay_scale * b)
state = state * block_decay + tl.dot(k * decays[None, :], v)
if EVEN:
tl.store(out_ptrs + n * stride_o, o.to(Out.dtype.element_ty))
else:
tl.store(
out_ptrs + n * stride_o,
o.to(Out.dtype.element_ty),
mask=(n + offs_b)[:, None] < q_length,
)
tl.store(s_ptrs, state.to(S.dtype.element_ty))
else:
q = tl.trans(tl.load(q_ptrs)).to(tl.float32) * softmax_scale
k = tl.trans(tl.load(k_ptrs)).to(tl.float32)
v = tl.load(v_ptrs).to(tl.float32)
state = state * tl.exp(decay_scale) + k * v
o = tl.sum(q * state, axis=0, keep_dims=True)
tl.store(out_ptrs, o.to(Out.dtype.element_ty))
tl.store(s_ptrs, state.to(S.dtype.element_ty))
# used for prefilling
@triton.jit
def seg_la_p_kernel(
Q,
K,
V,
S,
Out,
softmax_scale,
stride_q,
stride_k,
stride_v,
stride_s,
stride_o,
s_offsets,
q_offsets,
q_lengths,
s_scales,
decay_scales,
HEAD_DIM: tl.constexpr,
K_SPLIT_DIM: tl.constexpr,
V_SPLIT_DIM: tl.constexpr,
BLOCK: tl.constexpr,
EVEN: tl.constexpr,
):
bid = tl.program_id(0)
hid = tl.program_id(1)
kvid = tl.program_id(2)
N = HEAD_DIM // V_SPLIT_DIM
kid = kvid // N
vid = kvid % N
H = tl.num_programs(1)
# s_scale is 0 (first prefill chunk) or 1 (next prefill chunk)
s_scale = tl.load(s_scales + bid)
q_length = tl.load(q_lengths + bid)
q_offset = tl.load(q_offsets + bid)
s_offset = tl.load(s_offsets + bid)
decay_scale = -tl.load(decay_scales + hid)
offs_b = tl.arange(0, BLOCK)
offs_k = tl.arange(0, K_SPLIT_DIM)
offs_v = tl.arange(0, V_SPLIT_DIM)
if s_offset == -1:
return
q_ptrs = (
Q
+ q_offset * stride_q
+ hid * HEAD_DIM
+ kid * K_SPLIT_DIM
+ (offs_b[:, None] * stride_q + offs_k[None, :])
)
k_ptrs = (
K
+ q_offset * stride_k
+ hid * HEAD_DIM
+ kid * K_SPLIT_DIM
+ (offs_b[:, None] * stride_k + offs_k[None, :])
)
v_ptrs = (
V
+ q_offset * stride_v
+ hid * HEAD_DIM
+ vid * V_SPLIT_DIM
+ (offs_b[:, None] * stride_v + offs_v[None, :])
)
# (num_dim_block, length, qo_heads, d)
out_ptrs = (
Out
+ kid * stride_o
+ q_offset * HEAD_DIM * H
+ hid * HEAD_DIM
+ vid * V_SPLIT_DIM
+ (offs_b[:, None] * H * HEAD_DIM + offs_v[None, :])
)
s_ptrs = (
S
+ s_offset * stride_s
+ hid * HEAD_DIM * HEAD_DIM
+ kid * HEAD_DIM * K_SPLIT_DIM
+ vid * V_SPLIT_DIM
+ (offs_k[:, None] * HEAD_DIM + offs_v[None, :])
)
state = tl.load(s_ptrs, mask=s_scale > 0).to(tl.float32)
for n in range(0, q_length, BLOCK):
n = tl.multiple_of(n, BLOCK)
if EVEN:
q = tl.load(q_ptrs + n * stride_q).to(tl.float32)
k = tl.trans(tl.load(k_ptrs + n * stride_k)).to(tl.float32)
v = tl.load(v_ptrs + n * stride_v).to(tl.float32)
b = BLOCK
b_offs = b - 1 - offs_b
decays = tl.exp(decay_scale * b_offs)
inv_decays = 1 / decays
else:
q = tl.load(
q_ptrs + n * stride_q, mask=(n + offs_b)[:, None] < q_length, other=0.0
).to(tl.float32)
k = tl.trans(
tl.load(
k_ptrs + n * stride_k,
mask=(n + offs_b)[:, None] < q_length,
other=0.0,
)
).to(tl.float32)
v = tl.load(
v_ptrs + n * stride_v, mask=(n + offs_b)[:, None] < q_length, other=0.0
).to(tl.float32)
b = min(BLOCK, q_length - n)
b_offs = b - 1 - offs_b
block_decays = tl.exp(decay_scale * b_offs)
decays = tl.where(b_offs >= 0, block_decays, 0)
inv_decays = tl.where(b_offs >= 0, 1 / block_decays, 0)
q = q * inv_decays[:, None]
k = k * decays[None, :]
qk = tl.dot(q, k) * softmax_scale
qk = tl.where(offs_b[None, :] <= offs_b[:, None], qk, 0.0)
o = tl.dot(qk, v)
block_decay = tl.exp(decay_scale * b)
o = tl.dot(q, state) * block_decay * softmax_scale + o
state = state * block_decay + tl.dot(k, v)
if EVEN:
tl.store(out_ptrs + n * H * HEAD_DIM, o.to(Out.dtype.element_ty))
else:
tl.store(
out_ptrs + n * H * HEAD_DIM,
o.to(Out.dtype.element_ty),
mask=(n + offs_b)[:, None] < q_length,
)
tl.store(s_ptrs, state.to(S.dtype.element_ty))
# used for speculative
@triton.jit
def seg_la_s_kernel(
Q,
K,
V,
S,
Out,
Mask,
softmax_scale,
stride_q,
stride_k,
stride_v,
stride_s,
stride_o,
s_offsets,
q_offsets,
q_lengths,
s_scales,
decay_scales,
HEAD_DIM: tl.constexpr,
K_SPLIT_DIM: tl.constexpr,
V_SPLIT_DIM: tl.constexpr,
BLOCK: tl.constexpr,
EVEN: tl.constexpr,
):
bid = tl.program_id(0)
hid = tl.program_id(1)
kvid = tl.program_id(2)
N = HEAD_DIM // V_SPLIT_DIM
kid = kvid // N
vid = kvid % N
H = tl.num_programs(1)
# s_scale is 0 (first prefill chunk) or 1 (next prefill chunk)
s_scale = tl.load(s_scales + bid)
q_length = tl.load(q_lengths + bid)
q_offset = tl.load(q_offsets + bid)
s_offset = tl.load(s_offsets + bid)
decay_scale = -tl.load(decay_scales + hid)
offs_b = tl.arange(0, BLOCK)
offs_k = tl.arange(0, K_SPLIT_DIM)
offs_v = tl.arange(0, V_SPLIT_DIM)
if s_offset == -1:
return
q_ptrs = (
Q
+ q_offset * stride_q
+ hid * HEAD_DIM
+ kid * K_SPLIT_DIM
+ (offs_b[:, None] * stride_q + offs_k[None, :])
)
k_ptrs = (
K
+ q_offset * stride_k
+ hid * HEAD_DIM
+ kid * K_SPLIT_DIM
+ (offs_b[:, None] * stride_k + offs_k[None, :])
)
v_ptrs = (
V
+ q_offset * stride_v
+ hid * HEAD_DIM
+ vid * V_SPLIT_DIM
+ (offs_b[:, None] * stride_v + offs_v[None, :])
)
# (num_dim_block, length, qo_heads, d)
out_ptrs = (
Out
+ kid * stride_o
+ q_offset * HEAD_DIM * H
+ hid * HEAD_DIM
+ vid * V_SPLIT_DIM
+ (offs_b[:, None] * H * HEAD_DIM + offs_v[None, :])
)
s_ptrs = (
S
+ s_offset * stride_s
+ hid * HEAD_DIM * HEAD_DIM
+ kid * HEAD_DIM * K_SPLIT_DIM
+ vid * V_SPLIT_DIM
+ (offs_k[:, None] * HEAD_DIM + offs_v[None, :])
)
state = tl.load(s_ptrs, mask=s_scale > 0).to(tl.float32)
if EVEN:
q = tl.load(q_ptrs).to(tl.float32)
k = tl.trans(tl.load(k_ptrs)).to(tl.float32)
v = tl.load(v_ptrs).to(tl.float32)
mask = tl.load(
Mask
+ bid * BLOCK * BLOCK
+ tl.arange(0, BLOCK)[:, None] * BLOCK
+ tl.arange(0, BLOCK)[None, :]
).to(tl.int32)
positions = tl.sum(mask, 1) - 1
max_pos = tl.max(positions)
b_offs = max_pos - positions
else:
q = tl.load(q_ptrs, mask=offs_b[:, None] < q_length).to(tl.float32)
k = tl.trans(tl.load(k_ptrs, mask=offs_b[:, None] < q_length)).to(tl.float32)
v = tl.load(v_ptrs, mask=offs_b[:, None] < q_length).to(tl.float32)
mask = tl.load(
Mask
+ bid * q_length * q_length
+ tl.arange(0, BLOCK)[:, None] * q_length
+ tl.arange(0, BLOCK)[None, :],
mask=(tl.arange(0, BLOCK)[:, None] < q_length)
& (tl.arange(0, BLOCK)[None, :] < q_length),
).to(tl.int32)
positions = tl.sum(mask, 1) - 1
max_pos = tl.max(positions)
b_offs = max_pos - positions
decays = tl.exp(decay_scale * b_offs)
inv_decays = 1 / decays
q = q * inv_decays[:, None]
k = k * decays[None, :]
qk = tl.dot(q, k) * softmax_scale
qk = qk * mask.to(tl.float32)
o = tl.dot(qk, v)
block_decay = tl.exp(decay_scale * (max_pos + 1))
o = tl.dot(q, state) * block_decay * softmax_scale + o
if EVEN:
tl.store(out_ptrs, o.to(Out.dtype.element_ty))
else:
tl.store(out_ptrs, o.to(Out.dtype.element_ty), mask=offs_b[:, None] < q_length)
# used for decode
@triton.jit
def seg_la_d_kernel(
Q,
K,
V,
S,
Out,
softmax_scale,
stride_q,
stride_k,
stride_v,
stride_s,
stride_o,
s_offsets,
decay_scales,
HEAD_DIM: tl.constexpr,
K_SPLIT_DIM: tl.constexpr,
V_SPLIT_DIM: tl.constexpr,
):
bid = tl.program_id(0)
hid = tl.program_id(1)
kvid = tl.program_id(2)
N = HEAD_DIM // V_SPLIT_DIM
kid = kvid // N
vid = kvid % N
H = tl.num_programs(1)
# s_scale is 0 (first prefill chunk) or 1 (next prefill chunk)
s_offset = tl.load(s_offsets + bid)
if s_offset == -1:
return
decay_scale = -tl.load(decay_scales + hid)
offs_k = tl.arange(0, K_SPLIT_DIM)
offs_v = tl.arange(0, V_SPLIT_DIM)
q_ptrs = Q + bid * stride_q + hid * HEAD_DIM + kid * K_SPLIT_DIM + (offs_k)
k_ptrs = K + bid * stride_k + hid * HEAD_DIM + kid * K_SPLIT_DIM + (offs_k)
v_ptrs = V + bid * stride_v + hid * HEAD_DIM + vid * V_SPLIT_DIM + (offs_v)
# (num_dim_block, length, qo_heads, d)
out_ptrs = (
Out
+ kid * stride_o
+ bid * H * HEAD_DIM
+ hid * HEAD_DIM
+ vid * V_SPLIT_DIM
+ (offs_v)
)
s_ptrs = (
S
+ s_offset * stride_s
+ hid * HEAD_DIM * HEAD_DIM
+ kid * HEAD_DIM * K_SPLIT_DIM
+ vid * V_SPLIT_DIM
+ (offs_k[:, None] * HEAD_DIM + offs_v[None, :])
)
state = tl.load(s_ptrs).to(tl.float32)
k = tl.load(k_ptrs).to(tl.float32)
v = tl.load(v_ptrs).to(tl.float32)
q = tl.load(q_ptrs).to(tl.float32) * softmax_scale
state = state * tl.exp(decay_scale) + k[:, None] * v
o = tl.sum(q[:, None] * state, axis=0)
tl.store(out_ptrs, o.to(Out.dtype.element_ty))
tl.store(s_ptrs, state.to(S.dtype.element_ty))
# used for MTP with only spec-topk=1.
@triton.jit
def seg_la_mtp_kernel(
Q,
K,
V,
S,
CACHES,
Out,
softmax_scale,
stride_q,
stride_k,
stride_v,
stride_s,
stride_c,
stride_o,
s_offsets,
cache_indices,
decay_scales,
step,
HEAD_DIM: tl.constexpr,
K_SPLIT_DIM: tl.constexpr,
V_SPLIT_DIM: tl.constexpr,
):
bid = tl.program_id(0)
hid = tl.program_id(1)
kvid = tl.program_id(2)
N = HEAD_DIM // V_SPLIT_DIM
kid = kvid // N
vid = kvid % N
H = tl.num_programs(1)
s_offset = tl.load(s_offsets + bid)
if s_offset == -1:
return
decay_scale = tl.exp(-tl.load(decay_scales + hid))
offs_k = tl.arange(0, K_SPLIT_DIM)
offs_v = tl.arange(0, V_SPLIT_DIM)
# (length, qo_heads, d)
q_ptrs = Q + bid * step * stride_q + hid * HEAD_DIM + kid * K_SPLIT_DIM + (offs_k)
k_ptrs = K + bid * step * stride_k + hid * HEAD_DIM + kid * K_SPLIT_DIM + (offs_k)
v_ptrs = V + bid * step * stride_v + hid * HEAD_DIM + vid * V_SPLIT_DIM + (offs_v)
# (num_dim_block, length, qo_heads, d)
out_ptrs = (
Out
+ kid * stride_o
+ bid * step * H * HEAD_DIM
+ hid * HEAD_DIM
+ vid * V_SPLIT_DIM
+ (offs_v)
)
# (bs, qo_heads, d, d)
s_ptrs = (
S
+ s_offset * stride_s
+ hid * HEAD_DIM * HEAD_DIM
+ kid * HEAD_DIM * K_SPLIT_DIM
+ vid * V_SPLIT_DIM
+ (offs_k[:, None] * HEAD_DIM + offs_v[None, :])
)
state = tl.load(s_ptrs).to(tl.float32)
# (bs, step, kv_heads, d, d)
cache_indices = tl.load(cache_indices + bid)
c_ptrs = (
CACHES
+ cache_indices * stride_c
+ hid * HEAD_DIM * HEAD_DIM
+ kid * HEAD_DIM * K_SPLIT_DIM
+ vid * V_SPLIT_DIM
+ (offs_k[:, None] * HEAD_DIM + offs_v[None, :])
)
for i in range(step):
q = tl.load(q_ptrs).to(tl.float32) * softmax_scale
k = tl.load(k_ptrs).to(tl.float32)
v = tl.load(v_ptrs).to(tl.float32)
state = state * decay_scale + k[:, None] * v
o = tl.sum(q[:, None] * state, axis=0)
tl.store(out_ptrs, o.to(Out.dtype.element_ty))
tl.store(c_ptrs, state.to(CACHES.dtype.element_ty))
q_ptrs += stride_q
k_ptrs += stride_k
v_ptrs += stride_v
out_ptrs += H * HEAD_DIM
c_ptrs += H * HEAD_DIM * HEAD_DIM
# (k_dim_block, length, qo_heads, d)
@triton.jit
def seg_la_sum_kernel(T, O, DIM: tl.constexpr, NUM_BLOCK: tl.constexpr):
pid = tl.program_id(0)
length = tl.num_programs(0)
x = tl.zeros((DIM,), dtype=tl.float32)
for i in range(NUM_BLOCK):
x += tl.load(T + i * length * DIM + pid * DIM + tl.arange(0, DIM)).to(
tl.float32
)
tl.store(O + pid * DIM + tl.arange(0, DIM), x)
def seg_la_fwd(
q,
k,
v,
s,
decay_scales,
meta,
caches=None,
cache_indices=None,
softmax_scale=None,
decouple=False,
):
length, qo_heads, HEAD_DIM = q.shape
_, kv_heads, _ = k.shape
bs = meta.batch_size
if softmax_scale is None:
softmax_scale = HEAD_DIM ** (-0.5)
# MAX_LENGTH = meta.max_q_length
MAX_LENGTH = triton.cdiv(length, bs)
assert qo_heads == kv_heads, "seg_la does NOT support GQA currently"
if MAX_LENGTH > 1:
# prefill with partitioning q/k/v
# BLOCK should <= 64 with decouple
K_SPLIT_DIM = 32
V_SPLIT_DIM = 32 if bs <= 2 else 64
num_warps = 2 # 2
num_stages = 3 # 3
k_dim_block = HEAD_DIM // K_SPLIT_DIM
v_dim_block = HEAD_DIM // V_SPLIT_DIM
tmp = torch.empty(
(k_dim_block, length, qo_heads, HEAD_DIM), device=q.device, dtype=q.dtype
)
grid = (bs, kv_heads, k_dim_block * v_dim_block)
if caches is not None:
# mtp
EVEN = False
BLOCK = 32
step = length // bs
seg_la_mtp_kernel[grid](
q,
k,
v,
s,
caches,
tmp,
softmax_scale,
q.stride(0),
k.stride(0),
v.stride(0),
s.stride(0),
caches.stride(0),
tmp.stride(0),
meta.s_offsets,
cache_indices,
decay_scales,
step,
HEAD_DIM=HEAD_DIM,
K_SPLIT_DIM=K_SPLIT_DIM,
V_SPLIT_DIM=V_SPLIT_DIM,
num_warps=num_warps,
num_stages=num_stages,
)
elif meta.mask is not None:
# spec
ms = meta.mask.size(-1)
BLOCK = (ms + 15) // 16 * 16
EVEN = BLOCK == ms
seg_la_s_kernel[grid](
q,
k,
v,
s,
tmp,
meta.mask,
softmax_scale,
q.stride(0),
k.stride(0),
v.stride(0),
s.stride(0),
tmp.stride(0),
meta.s_offsets,
meta.q_offsets,
meta.q_lengths,
meta.s_scales,
decay_scales,
HEAD_DIM=HEAD_DIM,
K_SPLIT_DIM=K_SPLIT_DIM,
V_SPLIT_DIM=V_SPLIT_DIM,
BLOCK=BLOCK,
EVEN=EVEN,
num_warps=num_warps,
num_stages=num_stages,
)
else:
# prefill
BLOCK = 32
EVEN = MAX_LENGTH % BLOCK == 0 if bs == 1 else False
seg_la_p_kernel[grid](
q,
k,
v,
s,
tmp,
softmax_scale,
q.stride(0),
k.stride(0),
v.stride(0),
s.stride(0),
tmp.stride(0),
meta.s_offsets,
meta.q_offsets,
meta.q_lengths,
meta.s_scales,
decay_scales,
HEAD_DIM=HEAD_DIM,
K_SPLIT_DIM=K_SPLIT_DIM,
V_SPLIT_DIM=V_SPLIT_DIM,
BLOCK=BLOCK,
EVEN=EVEN,
num_warps=num_warps,
num_stages=num_stages,
)
if k_dim_block > 1:
if length < 2048:
o = tmp.sum(0)
else:
o = torch.empty(
(length, qo_heads, HEAD_DIM), device=q.device, dtype=q.dtype
)
seg_la_sum_kernel[(length,)](
tmp,
o,
DIM=qo_heads * HEAD_DIM,
NUM_BLOCK=k_dim_block,
num_warps=2,
num_stages=3,
)
else:
o = tmp[0]
else:
# decode with partitioning q/k/v
if bs <= 128:
K_SPLIT_DIM = 128 # 128
V_SPLIT_DIM = 32 # 32
num_warps = 2 # 2
num_stages = 2 # 3
else:
K_SPLIT_DIM = 128 # 128
V_SPLIT_DIM = 64 # 32
num_warps = 2 # 2
num_stages = 3 # 3
k_dim_block = HEAD_DIM // K_SPLIT_DIM
v_dim_block = HEAD_DIM // V_SPLIT_DIM
tmp = torch.empty(
(k_dim_block, length, qo_heads, HEAD_DIM), device=q.device, dtype=q.dtype
)
grid = (bs, kv_heads, k_dim_block * v_dim_block)
seg_la_d_kernel[grid](
q,
k,
v,
s,
tmp,
softmax_scale,
q.stride(0),
k.stride(0),
v.stride(0),
s.stride(0),
tmp.stride(0),
meta.s_offsets,
decay_scales,
HEAD_DIM=HEAD_DIM,
K_SPLIT_DIM=K_SPLIT_DIM,
V_SPLIT_DIM=V_SPLIT_DIM,
num_warps=num_warps,
num_stages=num_stages,
)
if k_dim_block > 1:
o = tmp.sum(0)
else:
o = tmp[0]
# if fallback:
# # prefill/decode with partitioning v only
# o = torch.empty(q.shape, device=q.device, dtype=q.dtype)
# if MAX_LENGTH == 1:
# # decode
# BLOCK = 1
# EVEN = False
# SPLIT_DIM = 32
# num_warps = 8
# num_stages = 2
# num_dim_block = HEAD_DIM // SPLIT_DIM
# grid = (batch, kv_heads, num_dim_block)
# else:
# # prefill
# if decouple:
# BLOCK = 64
# SPLIT_DIM = 16
# else:
# BLOCK = HEAD_DIM
# SPLIT_DIM = 32
# # EVEN = all([x % BLOCK == 0 for x in meta.qls])
# EVEN = False
# num_warps = 8
# num_stages = 2
# # prop = torch.cuda.get_device_properties(q.device.index)
# # arch = prop.major * 10 + prop.minor
# # if arch not in (80, 90):
# # num_stages = 1
# num_dim_block = HEAD_DIM // SPLIT_DIM
# grid = (batch, kv_heads, num_dim_block)
# seg_la_kernel[grid](
# q,
# k,
# v,
# s,
# o,
# softmax_scale,
# q.stride(0),
# k.stride(0),
# v.stride(0),
# s.stride(0),
# o.stride(0),
# meta.s_offsets,
# meta.q_offsets,
# meta.q_lengths,
# meta.s_scales,
# decay_scales,
# HEAD_DIM=HEAD_DIM,
# SPLIT_DIM=SPLIT_DIM,
# BLOCK=BLOCK,
# EVEN=EVEN,
# DECOUPLE=decouple,
# num_warps=num_warps,
# num_stages=num_stages
# )
return o

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

@@ -31,6 +31,7 @@ import torch.distributed as dist
from torch import nn
from sglang.srt.configs import (
BailingHybridConfig,
FalconH1Config,
JetNemotronConfig,
JetVLMConfig,
@@ -1549,6 +1550,13 @@ class ModelRunner(ModelRunnerKVCacheMixin):
return config
return None
@property
def hybrid_lightning_config(self):
config = self.model_config.hf_config
if isinstance(config, BailingHybridConfig):
return config
return None
@property
def hybrid_gdn_config(self):
config = self.model_config.hf_config.get_text_config()
@@ -1597,7 +1605,12 @@ class ModelRunner(ModelRunnerKVCacheMixin):
@property
def mambaish_config(self):
return self.mamba2_config or self.hybrid_gdn_config or self.kimi_linear_config
return (
self.mamba2_config
or self.hybrid_gdn_config
or self.kimi_linear_config
or self.hybrid_lightning_config
)
def can_run_piecewise_cuda_graph(self):
if self.is_draft_worker:

8
python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py
View File

@@ -584,7 +584,13 @@ class ModelRunnerKVCacheMixin:
head_dim=self.model_config.head_dim,
# if draft worker, we only need 1 attention layer's kv pool
full_attention_layer_ids=(
[0] if self.is_draft_worker else config.full_attention_layer_ids
[0]
if self.is_draft_worker
else [
i
for i in config.full_attention_layer_ids
if self.start_layer <= i < self.end_layer
]
),
enable_kvcache_transpose=False,
device=self.device,

1571
python/sglang/srt/models/bailing_moe_linear.py Normal file
View File

File diff suppressed because it is too large Load Diff

105
python/sglang/srt/models/bailing_moe_nextn.py
View File

@@ -37,8 +37,13 @@ from sglang.srt.layers.vocab_parallel_embedding import (
)
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
from sglang.srt.models.bailing_moe import BailingMoEBlock, BailingMoEForCausalLM
from sglang.srt.models.bailing_moe_linear import (
BailingMoELinearDecoderLayer,
BailingMoeV2_5ForCausalLM,
)
from sglang.srt.models.utils import WeightsMapper
from sglang.srt.server_args import get_global_server_args
from sglang.srt.utils import add_prefix
from sglang.srt.utils import BumpAllocator, add_prefix
LoraConfig = None
logger = logging.getLogger(__name__)
@@ -52,6 +57,13 @@ class BailingMoEModelNextN(nn.Module):
prefix: str = "",
) -> None:
super().__init__()
self.layer_group_size = 1
self.start_layer = 0
self.end_layer = 1
self.total_num_layers = 1
self.vocab_size = config.vocab_size
config.for_nextn_model = True
if quant_config is not None and quant_config.get_name() == "modelopt_fp4":
logger.warning(
"Overriding DeepseekV3ForCausalLMNextN quant config for modelopt_fp4 Deepseek model."
@@ -63,7 +75,7 @@ class BailingMoEModelNextN(nn.Module):
self.word_embeddings = VocabParallelEmbedding(
config.vocab_size,
config.hidden_size,
use_attn_tp_group=is_dp_attention_enabled(),
enable_tp=not is_dp_attention_enabled(),
prefix=add_prefix("word_embeddings", prefix),
)
@@ -75,16 +87,29 @@ class BailingMoEModelNextN(nn.Module):
config.hidden_size,
bias=False,
quant_config=quant_config,
prefix=add_prefix("eh_proj", prefix),
prefix=add_prefix(f"layers.{config.num_hidden_layers}.eh_proj", prefix),
)
self.decoder = BailingMoEBlock(
config,
0,
quant_config=quant_config,
# is_nextn=True,
prefix=add_prefix("decoder", prefix),
self.is_hybrid = (
hasattr(config, "model_type") and config.model_type == "bailing_hybrid"
)
if self.is_hybrid:
config.attention_type = 1
self.decoder = BailingMoELinearDecoderLayer(
config,
quant_config=quant_config,
layer_id=0,
is_nextn=True,
prefix=add_prefix(f"layers.{config.num_hidden_layers}", prefix),
)
else:
self.decoder = BailingMoEBlock(
config,
0,
quant_config=quant_config,
# is_nextn=True,
prefix=add_prefix("decoder", prefix),
)
self.shared_head = nn.Module()
self.final_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
@@ -107,16 +132,37 @@ class BailingMoEModelNextN(nn.Module):
torch.cat(
(
self.enorm(hidden_states),
self.hnorm(forward_batch.spec_info.hidden_states),
self.hnorm(
forward_batch.spec_info.hidden_states.to(
self.hnorm.weight.dtype
)
),
),
dim=-1,
)
)
residual = None
hidden_states, residual = self.decoder(
positions, hidden_states, forward_batch, residual
)
if self.is_hybrid:
device = input_ids.device
zero_allocator = BumpAllocator(
buffer_size=self.total_num_layers
* 2
* (2 if forward_batch.can_run_tbo else 1),
dtype=torch.float32,
device=device,
)
hidden_states, residual = self.decoder(
hidden_states=hidden_states,
positions=positions,
forward_batch=forward_batch,
residual=residual,
zero_allocator=zero_allocator,
)
else:
hidden_states, residual = self.decoder(
positions, hidden_states, forward_batch, residual
)
if not forward_batch.forward_mode.is_idle():
if residual is not None:
@@ -127,7 +173,18 @@ class BailingMoEModelNextN(nn.Module):
return hidden_states
class BailingMoeForCausalLMNextN(BailingMoEForCausalLM):
class BailingMoeForCausalLMNextN(nn.Module):
packed_modules_mapping = {
"fused_qkv_a_proj_with_mqa": ["q_a_proj", "kv_a_proj_with_mqa"],
"gate_up_proj": ["gate_proj", "up_proj"],
}
# To ensure correct weight loading and mapping.
hf_to_sglang_mapper = WeightsMapper(
orig_to_new_substr={
"attention.dense": "attention.o_proj",
},
)
def __init__(
self,
@@ -154,6 +211,12 @@ class BailingMoeForCausalLMNextN(BailingMoEForCausalLM):
use_attn_tp_group=get_global_server_args().enable_dp_lm_head,
)
self.logits_processor = LogitsProcessor(config)
if hasattr(self.config, "model_type") and config.model_type == "bailing_hybrid":
self.base_load_weights_func = BailingMoeV2_5ForCausalLM.load_weights
self.post_load_weights_func = BailingMoeV2_5ForCausalLM.post_load_weights
else:
self.base_load_weights_func = BailingMoEForCausalLM.load_weights
self.post_load_weights_func = BailingMoEForCausalLM.post_load_weights
@torch.no_grad()
def forward(
@@ -167,8 +230,20 @@ class BailingMoeForCausalLMNextN(BailingMoEForCausalLM):
input_ids, hidden_states, self.lm_head, forward_batch
)
def set_embed_and_head(self, embed, head):
"""Used by the eagle_worker."""
del self.model.word_embeddings.weight
del self.lm_head.weight
self.model.word_embeddings.weight = embed
self.lm_head.weight = head
torch.cuda.empty_cache()
torch.cuda.synchronize()
def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
super().load_weights(weights, is_nextn=True)
self.base_load_weights_func(self, weights, is_nextn=True)
def post_load_weights(self, is_nextn=False, weight_names=None):
self.post_load_weights_func(self, is_nextn=is_nextn, weight_names=weight_names)
EntryClass = [BailingMoeForCausalLMNextN]

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

@@ -1512,7 +1512,7 @@ class ServerArgs:
logger.info(
f"Using {self.attention_backend} as attention backend for {model_arch}."
)
elif model_arch in ["KimiLinearForCausalLM"]:
elif model_arch in ["KimiLinearForCausalLM", "BailingMoeV2_5ForCausalLM"]:
self._handle_mamba_radix_cache(
model_arch=model_arch,
support_mamba_cache=False,
@@ -2339,6 +2339,7 @@ class ServerArgs:
"GlmMoeDsaForCausalLM",
"BailingMoeForCausalLM",
"BailingMoeV2ForCausalLM",
"BailingMoeV2_5ForCausalLM",
"MistralLarge3ForCausalLM",
"PixtralForConditionalGeneration",
]:
@@ -5686,6 +5687,7 @@ def auto_choose_speculative_params(self: ServerArgs):
"GlmMoeDsaForCausalLM",
"BailingMoeForCausalLM",
"BailingMoeV2ForCausalLM",
"BailingMoeV2_5ForCausalLM",
"MistralLarge3ForCausalLM",
"PixtralForConditionalGeneration",
"MiMoV2FlashForCausalLM",

1
python/sglang/srt/speculative/eagle_worker.py
View File

@@ -759,6 +759,7 @@ class EAGLEWorker(TpModelWorker):
if (
self.target_worker.model_runner.hybrid_gdn_config is not None
or self.target_worker.model_runner.mamba2_config is not None
or self.target_worker.model_runner.hybrid_lightning_config is not None
):
self._mamba_verify_update(
batch, res, logits_output, spec_info, seq_lens_pre_verify

2
python/sglang/srt/utils/hf_transformers_utils.py
View File

@@ -45,6 +45,7 @@ from transformers.models.auto.modeling_auto import MODEL_FOR_CAUSAL_LM_MAPPING_N
from sglang.srt.configs import (
AfmoeConfig,
BailingHybridConfig,
ChatGLMConfig,
DbrxConfig,
DeepseekVL2Config,
@@ -77,6 +78,7 @@ from sglang.srt.utils.patch_tokenizer import patch_tokenizer
_CONFIG_REGISTRY: List[Type[PretrainedConfig]] = [
AfmoeConfig,
BailingHybridConfig,
ChatGLMConfig,
DbrxConfig,
ExaoneConfig,