Support LingV2_5 model (#18598)

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>

python/sglang/srt/configs/__init__.py

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

python/sglang/srt/configs/bailing_hybrid.py

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

python/sglang/srt/configs/model_config.py

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

python/sglang/srt/layers/attention/attention_registry.py

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

python/sglang/srt/layers/attention/fla/layernorm_gated.py

@@ -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)
+        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)
+        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
+                self.norm_before_gate
-            ), "CPU rmsnorm_gated currently only supports norm before gate without group size"
+                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,
             )

python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py

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

python/sglang/srt/layers/attention/linear/lightning_attn.py

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

python/sglang/srt/layers/attention/linear/linear_metadata.py

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

python/sglang/srt/layers/attention/linear/seg_la.py

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

python/sglang/srt/model_executor/model_runner.py

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

python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py

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

python/sglang/srt/models/bailing_moe_nextn.py

@@ -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,9 +87,22 @@ 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.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,
@@ -107,13 +132,34 @@ 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
+        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
             )
@@ -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]

python/sglang/srt/server_args.py

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

python/sglang/srt/speculative/eagle_worker.py

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

python/sglang/srt/utils/hf_transformers_utils.py

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