From d97eb111a368f8f19d0d5aac9b89cdd2243daed3 Mon Sep 17 00:00:00 2001 From: ant-yy Date: Fri, 13 Feb 2026 16:09:15 +0800 Subject: [PATCH] Support LingV2_5 model (#18598) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: zhangkaihong.zkh Co-authored-by: 有禾 Co-authored-by: yudian0504 <138860534+yudian0504@users.noreply.github.com> Co-authored-by: 悠扬 Co-authored-by: xinxingyang Co-authored-by: zmy460290 --- python/sglang/srt/configs/__init__.py | 2 + python/sglang/srt/configs/bailing_hybrid.py | 188 ++ python/sglang/srt/configs/model_config.py | 20 + .../layers/attention/attention_registry.py | 3 + .../layers/attention/fla/layernorm_gated.py | 31 +- .../attention/hybrid_linear_attn_backend.py | 369 ++++ .../layers/attention/linear/lightning_attn.py | 767 ++++++++ .../attention/linear/linear_metadata.py | 70 + .../srt/layers/attention/linear/seg_la.py | 909 ++++++++++ .../sglang/srt/model_executor/model_runner.py | 15 +- .../model_runner_kv_cache_mixin.py | 8 +- .../sglang/srt/models/bailing_moe_linear.py | 1571 +++++++++++++++++ python/sglang/srt/models/bailing_moe_nextn.py | 105 +- python/sglang/srt/server_args.py | 4 +- python/sglang/srt/speculative/eagle_worker.py | 1 + .../sglang/srt/utils/hf_transformers_utils.py | 2 + 16 files changed, 4042 insertions(+), 23 deletions(-) create mode 100644 python/sglang/srt/configs/bailing_hybrid.py create mode 100644 python/sglang/srt/layers/attention/linear/lightning_attn.py create mode 100644 python/sglang/srt/layers/attention/linear/linear_metadata.py create mode 100644 python/sglang/srt/layers/attention/linear/seg_la.py create mode 100644 python/sglang/srt/models/bailing_moe_linear.py diff --git a/python/sglang/srt/configs/__init__.py b/python/sglang/srt/configs/__init__.py index 8e55d7061..965b4e305 100644 --- a/python/sglang/srt/configs/__init__.py +++ b/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", diff --git a/python/sglang/srt/configs/bailing_hybrid.py b/python/sglang/srt/configs/bailing_hybrid.py new file mode 100644 index 000000000..40933d90a --- /dev/null +++ b/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) diff --git a/python/sglang/srt/configs/model_config.py b/python/sglang/srt/configs/model_config.py index 998defbaf..701efb326 100644 --- a/python/sglang/srt/configs/model_config.py +++ b/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 diff --git a/python/sglang/srt/layers/attention/attention_registry.py b/python/sglang/srt/layers/attention/attention_registry.py index 246e2554f..25ffc3f0b 100644 --- a/python/sglang/srt/layers/attention/attention_registry.py +++ b/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." diff --git a/python/sglang/srt/layers/attention/fla/layernorm_gated.py b/python/sglang/srt/layers/attention/fla/layernorm_gated.py index 7bc7b9f47..d75252da4 100644 --- a/python/sglang/srt/layers/attention/fla/layernorm_gated.py +++ b/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) - x *= z * tl.sigmoid(z) + if ACTIVATION == "swish" or ACTIVATION == "silu": + x *= z * tl.sigmoid(z) + elif ACTIVATION == "sigmoid": + x *= tl.sigmoid(z) # Compute mean and variance per row (reduce along axis 1) if not IS_RMS_NORM: @@ -152,7 +156,10 @@ def _layer_norm_fwd_1pass_kernel( if HAS_Z and NORM_BEFORE_GATE: Z_base = Z + rows[:, None] * stride_z_row + col_offsets z = tl.load(Z_base, mask=mask, other=0.0).to(tl.float32) - y *= z * tl.sigmoid(z) + if ACTIVATION == "swish" or ACTIVATION == "silu": + y *= z * tl.sigmoid(z) + elif ACTIVATION == "sigmoid": + y *= tl.sigmoid(z) # Write output tl.store(Y_base, y, mask=mask) @@ -182,6 +189,7 @@ def _layer_norm_fwd( group_size=None, norm_before_gate=True, is_rms_norm=False, + activation: str = "swish", ): M, N = x.shape if group_size is None: @@ -242,6 +250,7 @@ def _layer_norm_fwd( NORM_BEFORE_GATE=norm_before_gate, IS_RMS_NORM=is_rms_norm, num_warps=num_warps, + ACTIVATION=activation, ) return out, mean, rstd @@ -260,6 +269,7 @@ def rms_norm_gated( group_size=None, norm_before_gate=True, is_rms_norm=False, + activation: str = "swish", ): """If z is not None, we do norm(x) * silu(z) if norm_before_gate, else norm(x * silu(z))""" @@ -276,6 +286,8 @@ def rms_norm_gated( weight = weight.contiguous() if bias is not None: bias = bias.contiguous() + if _is_npu: + assert activation == "swish", "NPU only supports swish activation" y, mean, rstd = _layer_norm_fwd( x, weight, @@ -285,6 +297,7 @@ def rms_norm_gated( group_size=group_size, norm_before_gate=norm_before_gate, is_rms_norm=is_rms_norm, + activation=activation, ) return y.reshape(x_shape_og) @@ -302,6 +315,7 @@ class LayerNormFn(torch.autograd.Function): group_size=None, norm_before_gate=True, is_rms_norm=False, + activation: str = "swish", ): return rms_norm_gated( x=x, @@ -312,6 +326,7 @@ class LayerNormFn(torch.autograd.Function): group_size=group_size, norm_before_gate=norm_before_gate, is_rms_norm=is_rms_norm, + activation=activation, ) @@ -324,9 +339,10 @@ def layernorm_fn( group_size=None, norm_before_gate=True, is_rms_norm=False, + activation: str = "swish", ): return LayerNormFn.apply( - x, weight, bias, z, eps, group_size, norm_before_gate, is_rms_norm + x, weight, bias, z, eps, group_size, norm_before_gate, is_rms_norm, activation ) @@ -382,6 +398,7 @@ class RMSNorm(torch.nn.Module): norm_before_gate=True, device=None, dtype=None, + activation: str = "swish", ): """If group_size is not None, we do GroupNorm with each group having group_size elements. group_size=None is equivalent to group_size=hidden_size (i.e. there's only 1 group). @@ -389,6 +406,7 @@ class RMSNorm(torch.nn.Module): factory_kwargs = {"device": device, "dtype": dtype} super().__init__() self.eps = eps + self.activation = activation self.weight = torch.nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) self.register_parameter("bias", None) self.group_size = group_size @@ -402,8 +420,10 @@ class RMSNorm(torch.nn.Module): """If z is not None, we do norm(x) * silu(z) if norm_before_gate, else norm(x * silu(z))""" if _use_cpu: assert ( - self.norm_before_gate and self.group_size is None - ), "CPU rmsnorm_gated currently only supports norm before gate without group size" + self.norm_before_gate + and self.group_size is None + and self.activation == "swish" + ), "CPU rmsnorm_gated currently only supports norm before gate without group size or activation other than swish" return torch.ops.sgl_kernel.fused_rmsnorm_gated_cpu( x, self.weight, z, self.eps ) @@ -417,4 +437,5 @@ class RMSNorm(torch.nn.Module): group_size=self.group_size, norm_before_gate=self.norm_before_gate, is_rms_norm=True, + activation=self.activation, ) diff --git a/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py b/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py index 35b3fd11d..a88278df3 100644 --- a/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py +++ b/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""" diff --git a/python/sglang/srt/layers/attention/linear/lightning_attn.py b/python/sglang/srt/layers/attention/linear/lightning_attn.py new file mode 100644 index 000000000..d415eefc9 --- /dev/null +++ b/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() diff --git a/python/sglang/srt/layers/attention/linear/linear_metadata.py b/python/sglang/srt/layers/attention/linear/linear_metadata.py new file mode 100644 index 000000000..ed2da0409 --- /dev/null +++ b/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(), + ) diff --git a/python/sglang/srt/layers/attention/linear/seg_la.py b/python/sglang/srt/layers/attention/linear/seg_la.py new file mode 100644 index 000000000..4519ff1a3 --- /dev/null +++ b/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 diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index d23d8e9ad..7a93225b4 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/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: diff --git a/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py b/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py index c7dcd5766..9a5a0500b 100644 --- a/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py +++ b/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, diff --git a/python/sglang/srt/models/bailing_moe_linear.py b/python/sglang/srt/models/bailing_moe_linear.py new file mode 100644 index 000000000..e27dd3589 --- /dev/null +++ b/python/sglang/srt/models/bailing_moe_linear.py @@ -0,0 +1,1571 @@ +# coding=utf-8 +# Copyright 2023 Antgroup and The HuggingFace Inc. team. All rights reserved. +import copy +import logging +from typing import Callable, Iterable, Optional, Set, Tuple, Union + +import torch +import torch.nn.functional as F +from torch import nn +from transformers import PretrainedConfig + +from sglang.srt.distributed import ( + get_pp_group, + get_tensor_model_parallel_rank, + get_tensor_model_parallel_world_size, + tensor_model_parallel_all_reduce, +) +from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder +from sglang.srt.layers import deep_gemm_wrapper +from sglang.srt.layers.activation import SiluAndMul +from sglang.srt.layers.attention.fla.layernorm_gated import RMSNorm as RMSNormGated +from sglang.srt.layers.attention.fla.layernorm_gated import layernorm_fn +from sglang.srt.layers.communicator import LayerCommunicator, LayerScatterModes +from sglang.srt.layers.dp_attention import ( + get_attention_tp_rank, + get_attention_tp_size, + is_dp_attention_enabled, +) +from sglang.srt.layers.layernorm import RMSNorm +from sglang.srt.layers.linear import ( + ColumnParallelLinear, + MergedColumnParallelLinear, + QKVParallelLinear, + RowParallelLinear, +) +from sglang.srt.layers.logits_processor import LogitsProcessor +from sglang.srt.layers.moe.ep_moe.layer import DeepEPMoE, get_moe_impl_class +from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE +from sglang.srt.layers.moe.topk import TopK +from sglang.srt.layers.quantization.base_config import QuantizationConfig +from sglang.srt.layers.quantization.fp8_kernel import is_fp8_fnuz +from sglang.srt.layers.quantization.fp8_utils import ( + block_quant_dequant, + block_quant_to_tensor_quant, + channel_quant_to_tensor_quant, + normalize_e4m3fn_to_e4m3fnuz, + requant_weight_ue8m0_inplace, +) +from sglang.srt.layers.quantization.int8_utils import ( + block_dequant as int8_block_dequant, +) +from sglang.srt.layers.radix_attention import RadixAttention +from sglang.srt.layers.rotary_embedding import get_rope_wrapper +from sglang.srt.layers.utils import PPMissingLayer +from sglang.srt.layers.vocab_parallel_embedding import ( + ParallelLMHead, + VocabParallelEmbedding, +) +from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors +from sglang.srt.model_loader.weight_utils import default_weight_loader +from sglang.srt.models.deepseek_v2 import DeepseekV2AttentionMLA, DeepseekV2MLP, _is_hip +from sglang.srt.models.utils import WeightsMapper +from sglang.srt.server_args import get_global_server_args +from sglang.srt.utils import ( + BumpAllocator, + add_prefix, + bind_or_assign, + cpu_has_amx_support, + get_bool_env_var, + get_device_sm, + is_cpu, + is_cuda, + is_flashinfer_available, + is_gfx95_supported, + is_hip, + is_npu, + is_sm100_supported, + make_layers, +) + +_is_hip = is_hip() +_is_cuda = is_cuda() +_is_npu = is_npu() +_is_fp8_fnuz = is_fp8_fnuz() +_use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip +_is_cpu_amx_available = cpu_has_amx_support() +_is_cpu = is_cpu() +_device_sm = get_device_sm() +_is_gfx95_supported = is_gfx95_supported() + +_use_aiter_gfx95 = _use_aiter and _is_gfx95_supported + +if _use_aiter_gfx95: + pass + +if _is_cuda: + from sgl_kernel import awq_dequantize +elif _is_cpu and _is_cpu_amx_available: + pass +elif _is_hip: + from sglang.srt.layers.quantization.awq_triton import ( + awq_dequantize_triton as awq_dequantize, + ) +else: + from vllm._custom_ops import awq_dequantize + +if _is_hip: + pass + +_is_flashinfer_available = is_flashinfer_available() +_is_sm100_supported = is_cuda() and is_sm100_supported() + + +class DsV3MLA(DeepseekV2AttentionMLA): + def __init__(self, **kwargs): + super().__init__(**kwargs) + if kwargs["rope_scaling"]: + self.rotary_emb.forward = self.rotary_emb.forward_cuda + + +LoraConfig = None +logger = logging.getLogger(__name__) +_is_cpu = is_cpu() + + +def is_linear_layer(layer_idx, layer_group_size): + if layer_idx is None: + return False + if layer_group_size > 0: + return (layer_idx + 1) % layer_group_size != 0 + else: + return False + + +def is_pp_missing_parameter( + name: str, + model: torch.nn.Module, +) -> bool: + if isinstance(model, PPMissingLayer): + return True + return False + + +def weight_loader_with_alias(alias: str): + def wrapper(func: Callable): + def inner_func( + param: torch.Tensor, + loaded_weight: torch.Tensor, + *args, + prefix: str = None, + **kwargs, + ): + # pf = "[vLLM][load]" + " " if prefix is None else f"[{prefix}] " + value = func(param, loaded_weight, *args, **kwargs) + return value + + return inner_func + + return wrapper + + +class BailingMLP(nn.Module): + + def __init__( + self, + hidden_size: int, + intermediate_size: int, + reduce_results=True, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.gate_up_proj = MergedColumnParallelLinear( + hidden_size, + [intermediate_size] * 2, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.gate_up_proj", + ) + self.down_proj = RowParallelLinear( + intermediate_size, + hidden_size, + bias=False, + quant_config=quant_config, + reduce_results=reduce_results, + prefix=f"{prefix}.down_proj", + ) + self.act_fn = SiluAndMul() + + def forward( + self, + x, + should_allreduce_fusion: bool = False, + use_reduce_scatter: bool = False, + ): + x, _ = self.gate_up_proj(x) + x = self.act_fn(x) + x, _ = self.down_proj( + x, + skip_all_reduce=use_reduce_scatter or should_allreduce_fusion, + ) + return x + + +class BailingMoEGate(nn.Module): + def __init__( + self, + config, + params_dtype: Optional[torch.dtype] = None, + prefix: str = "", + ): + super().__init__() + if params_dtype is None: + params_dtype = torch.get_default_dtype() + self.params_dtype = params_dtype + self.weight = nn.Parameter( + torch.empty( + (config.num_experts, config.hidden_size), + dtype=self.params_dtype, + ), + ) + if getattr(config, "moe_router_enable_expert_bias", False): + self.expert_bias = nn.Parameter( + torch.empty((config.num_experts,), dtype=torch.float32), + ) + else: + self.expert_bias = None + + def forward(self, hidden_states): + logits = F.linear(hidden_states.to(self.weight.dtype), self.weight, None).to( + hidden_states.dtype + ) + return logits + + +class BailingMoE(nn.Module): + + def __init__( + self, + config: PretrainedConfig, + quant_config: Optional[QuantizationConfig] = None, + layer_id: int = 0, + prefix: str = "moe", + ): + super().__init__() + + self.layer_id = layer_id + + self.tp_size = get_tensor_model_parallel_world_size() + self.tp_rank = get_tensor_model_parallel_rank() + + self.top_k = config.num_experts_per_tok + self.norm_expert_prob = getattr(config, "norm_topk_prob", False) + self.hidden_size = config.hidden_size + self.intermediate_size = config.moe_intermediate_size + self.num_shared_experts = getattr(config, "num_shared_experts", 0) + self.routed_scaling_factor = getattr(config, "routed_scaling_factor", 1.0) + self.score_function = getattr(config, "score_function", None) + + # Gate always runs at half / full precision for now. + router_dtype = getattr(config, "router_dtype", None) + if router_dtype is None: + self.router_dtype = torch.float32 + elif router_dtype == "fp32": + self.router_dtype = torch.float32 + else: + self.router_dtype = torch.bfloat16 + + # check group topk + self.num_expert_group = getattr(config, "n_group", 0) + self.topk_group = getattr(config, "topk_group", 0) + if self.num_expert_group > 0 or self.topk_group > 0: + assert ( + self.num_expert_group > 0 + and 0 < self.topk_group <= self.num_expert_group + ) + self.use_grouped_topk = True + else: + self.num_expert_group = self.topk_group = None + self.use_grouped_topk = False + + self.num_experts = config.num_experts + + self.gate = BailingMoEGate( + config=config, + params_dtype=self.router_dtype, + prefix=add_prefix("gate", prefix), + ) + self.correction_bias = ( + self.gate.expert_bias.data if self.gate.expert_bias is not None else None + ) + + if self.score_function is not None: + assert ( + self.score_function == "softmax" and self.correction_bias is None + ) or ( + self.score_function == "sigmoid" and self.correction_bias is not None + ), "score_function and correction_bias should be in 2 combination (softmax, None) or (sigmoid, not None)" + + self.topk = TopK( + top_k=self.top_k, + use_grouped_topk=self.use_grouped_topk, + renormalize=self.norm_expert_prob, + num_expert_group=self.num_expert_group, + topk_group=self.topk_group, + correction_bias=self.correction_bias, + routed_scaling_factor=self.routed_scaling_factor, + ) + moe_cls = get_moe_impl_class(quant_config) + self.experts = moe_cls( + num_experts=self.num_experts, + top_k=self.top_k, + layer_id=self.layer_id, + hidden_size=self.hidden_size, + intermediate_size=self.intermediate_size, + quant_config=quant_config, + routed_scaling_factor=self.routed_scaling_factor, + prefix=f"{prefix}.experts", + ) + + if self.num_shared_experts > 0: + intermediate_size = self.intermediate_size * self.num_shared_experts + self.shared_experts = BailingMLP( + hidden_size=self.hidden_size, + intermediate_size=intermediate_size, + reduce_results=False, + prefix=f"{prefix}.shared_experts", + quant_config=quant_config, + ) + + def forward( + self, + hidden_states: torch.Tensor, + should_allreduce_fusion: bool = False, + use_reduce_scatter: bool = False, + ) -> torch.Tensor: + num_tokens, hidden_size = hidden_states.shape + hidden_states = hidden_states.view(-1, hidden_size) + if self.num_shared_experts > 0: + shared_output = self.shared_experts(hidden_states) + + router_logits = self.gate(hidden_states) + topk_output = self.topk(hidden_states, router_logits) + final_hidden_states = self.experts(hidden_states, topk_output) + + if self.num_shared_experts > 0: + final_hidden_states = final_hidden_states + shared_output + + if self.tp_size > 1 and not use_reduce_scatter and not should_allreduce_fusion: + final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states) + return final_hidden_states + + +class BailingGroupRMSNormGate(RMSNormGated): + def __init__( + self, + hidden_size, + eps=1e-5, + group_size=None, + norm_before_gate=True, + device=None, + dtype=None, + ): + super().__init__( + hidden_size, + eps=eps, + group_size=group_size, + norm_before_gate=norm_before_gate, + device=device, + dtype=dtype, + activation="sigmoid", + ) + self.weight.weight_loader = self.weight_loader + + @staticmethod + def weight_loader( + param: torch.nn.Parameter, + loaded_weight: torch.Tensor, + ) -> None: + tp_size = get_attention_tp_size() + tp_rank = get_attention_tp_rank() + shard_size = loaded_weight.shape[0] // tp_size + shard = slice(tp_rank * shard_size, (tp_rank + 1) * shard_size) + param.data.copy_(loaded_weight[shard].contiguous()) + return + + +class BailingMoELinearAttention(nn.Module): + def __init__( + self, + config: PretrainedConfig, + quant_config: Optional[QuantizationConfig] = None, + layer_id: int = 0, + prefix: str = "linear_attn", + ): + super().__init__() + + self.layer_id = layer_id + self.hidden_size = config.hidden_size + self.total_num_heads = config.num_attention_heads + self.total_kv_heads = config.num_attention_heads # MHA + + self.head_dim = getattr(config, "head_dim", None) + if self.head_dim is None: + self.head_dim = config.hidden_size // self.total_num_heads + + self.hidden_inner_size = self.head_dim * self.total_num_heads + self.scaling = self.head_dim**-0.5 + self.tp_size = get_attention_tp_size() + self.tp_rank = get_attention_tp_rank() + + assert self.total_num_heads % self.tp_size == 0 + self.tp_heads = self.total_num_heads // self.tp_size + + self.max_position_embeddings = config.max_position_embeddings + self.rope_theta = getattr(config, "rope_theta", 600000) + + self.tp_kv_heads = self.total_kv_heads // self.tp_size + self.q_size_per_rank = self.head_dim * self.tp_heads + self.kv_size_per_rank = self.head_dim * self.tp_kv_heads + + self.use_qk_norm = getattr(config, "use_qk_norm", False) + # minimax / seg_la / fla + # TODO support fla + self.linear_backend = getattr(config, "linear_backend", "seg_la") + logger.info(f"linear_backend in bailing_moe_linear: {self.linear_backend}") + self.linear_scale = True if self.linear_backend == "minimax" else False + self.linear_rope = getattr(config, "linear_rope", True) + if hasattr(config, "use_linear_silu"): + self.linear_silu = config.use_linear_silu + elif hasattr(config, "linear_silu"): + self.linear_silu = config.linear_silu + else: + self.linear_silu = False + + self.query_key_value = QKVParallelLinear( + self.hidden_size, + self.head_dim, + self.total_num_heads, + self.total_kv_heads, + bias=(config.use_bias or config.use_qkv_bias), + quant_config=quant_config, + prefix=f"{prefix}.qkv_proj", + tp_rank=self.tp_rank, + tp_size=self.tp_size, + ) + + if self.use_qk_norm: + self.query_layernorm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) + self.key_layernorm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) + + self.g_proj = ColumnParallelLinear( + self.hidden_size, + self.hidden_inner_size, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.output_gate", + tp_rank=self.tp_rank, + tp_size=self.tp_size, + ) + self.dense = RowParallelLinear( + self.hidden_inner_size, + self.hidden_size, + bias=config.use_bias, + quant_config=quant_config, + prefix=f"{prefix}.out_proj", + tp_rank=self.tp_rank, + tp_size=self.tp_size, + reduce_results=False, + ) + self.attn = RadixAttention( + self.tp_heads, + self.head_dim, + self.scaling, + num_kv_heads=self.tp_kv_heads, + layer_id=layer_id, + quant_config=quant_config, + prefix=f"{prefix}.attn", + ) + + self.group_norm_size = getattr(config, "group_norm_size", 1) + self.rms_norm_eps = float(getattr(config, "rms_norm_eps", 1e-5)) + assert ( + self.tp_size <= self.group_norm_size + ), "tp_size must be less than or equal to group_norm_size that can use local rms norm" + assert ( + self.group_norm_size % self.tp_size == 0 + ), "group_norm_size must be divisible by tp_size" + self.g_norm = BailingGroupRMSNormGate( + hidden_size=self.hidden_inner_size // self.tp_size, + eps=self.rms_norm_eps, + group_size=self.hidden_inner_size // self.group_norm_size, + ) + # use fp32 rotary embedding + if hasattr(config, "rotary_dim"): + rotary_dim = config.rotary_dim + elif hasattr(config, "partial_rotary_factor"): + rotary_dim = int(self.head_dim * config.partial_rotary_factor) + else: + rotary_dim = self.head_dim + + self.rotary_emb = get_rope_wrapper( + self.head_dim, + rotary_dim=rotary_dim, + max_position=self.max_position_embeddings, + base=self.rope_theta, + rope_scaling=config.rope_scaling, + is_neox_style=True, + device=get_global_server_args().device, + dtype=torch.float32, + ) + + @staticmethod + def weight_direct_load(param: torch.Tensor, loaded_weight: torch.Tensor) -> None: + assert param.size() == loaded_weight.size() + param.data.copy_(loaded_weight) + return + + def forward( + self, + hidden_states: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + **kwargs, + ) -> torch.Tensor: + qkv, _ = self.query_key_value(hidden_states) + # logger.warning(f"===={self.layer_id=}, 1-1 {qkv.shape=}") + # use rotary_emb support fp32 + qkv = qkv.to(torch.float32) + if self.linear_silu: + qkv = F.silu(qkv) + + q, k, v = torch.split( + qkv, + [self.q_size_per_rank, self.kv_size_per_rank, self.kv_size_per_rank], + dim=-1, + ) + if self.use_qk_norm: + q = q.reshape(-1, self.tp_heads, self.head_dim) + k = k.reshape(-1, self.tp_kv_heads, self.head_dim) + q = layernorm_fn( + q, + self.query_layernorm.weight.data, + bias=None, + eps=self.rms_norm_eps, + is_rms_norm=True, + ) + k = layernorm_fn( + k, + self.key_layernorm.weight.data, + bias=None, + eps=self.rms_norm_eps, + is_rms_norm=True, + ) + q = q.reshape(-1, self.q_size_per_rank) + k = k.reshape(-1, self.kv_size_per_rank) + + if self.linear_rope: + q, k = self.rotary_emb(positions, q, k) + + q = q.view((qkv.shape[0], self.tp_heads, self.head_dim)) + k = k.view((qkv.shape[0], self.tp_kv_heads, self.head_dim)) + v = v.view((qkv.shape[0], self.tp_kv_heads, self.head_dim)) + # logger.warning(f"===={self.layer_id=}, 1-2 {q.shape=}, {k.shape=}, {v.shape=}") + + if self.linear_scale: + q = q * self.scaling + # q = q.to(torch.float32) + # k = k.to(torch.float32) + # v = v.to(torch.float32) + hidden = self.attn(q, k, v, forward_batch).to(hidden_states.dtype) + gate, _ = self.g_proj(hidden_states) + # logger.warning( + # f"===={self.layer_id=}, 1-3 {gate.shape=}, {hidden.shape=}, {gate.dtype=}, {hidden_states.dtype=}, {hidden.dtype=}" + # ) + if self.group_norm_size > 1: + hidden = self.g_norm(hidden, gate) + else: + hidden = self.g_norm(hidden) + hidden = F.sigmoid(gate) * hidden + # logger.warning(f"===={self.layer_id=}, 1-4 {hidden.shape=}") + hidden = hidden.data.to(hidden_states.dtype) + hidden, _ = self.dense(hidden) + # logger.warning(f"===={self.layer_id=}, 1-5 {hidden.shape=}") + return hidden + + +class BailingMoEAttention(nn.Module): + + def __init__( + self, + config: PretrainedConfig, + quant_config: Optional[QuantizationConfig] = None, + layer_id: int = None, + prefix: str = "mha", + ) -> None: + super().__init__() + self.layer_id = layer_id + + self.hidden_size = config.hidden_size + tp_size = get_attention_tp_size() + self.total_num_heads = config.num_attention_heads + assert self.total_num_heads % tp_size == 0 + self.num_heads = self.total_num_heads // tp_size + self.total_num_kv_heads = config.num_key_value_heads + if self.total_num_kv_heads >= tp_size: + assert self.total_num_kv_heads % tp_size == 0 + else: + assert tp_size % self.total_num_kv_heads == 0 + self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size) + self.head_dim = getattr(config, "head_dim", None) + if self.head_dim is None: + self.head_dim = self.hidden_size // self.total_num_heads + + self.q_size = self.num_heads * self.head_dim + self.kv_size = self.num_kv_heads * self.head_dim + self.scaling = self.head_dim**-0.5 + + self.split_qkv = getattr(config, "using_split_qkv_in_self_attention", False) + assert not self.split_qkv, "split_qkv is not supported for now" + self.use_qk_norm = getattr(config, "use_qk_norm", False) + + self.query_key_value = QKVParallelLinear( + self.hidden_size, + self.head_dim, + self.total_num_heads, + self.total_num_kv_heads, + bias=(config.use_bias or config.use_qkv_bias), + quant_config=quant_config, + prefix=f"{prefix}.qkv_proj", + ) + if self.use_qk_norm: + self.query_layernorm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) + self.key_layernorm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) + + self.dense = RowParallelLinear( + self.total_num_heads * self.head_dim, + self.hidden_size, + bias=config.use_bias, + quant_config=quant_config, + prefix=f"{prefix}.o_proj", + ) + if hasattr(config, "rotary_dim"): + self.rotary_dim = config.rotary_dim + elif hasattr(config, "partial_rotary_factor"): + self.rotary_dim = int(self.head_dim * config.partial_rotary_factor) + else: + self.rotary_dim = self.head_dim + self.max_position_embeddings = config.max_position_embeddings + self.rope_theta = getattr(config, "rope_theta", 600000) + self.rotary_emb = get_rope_wrapper( + self.head_dim, + rotary_dim=self.rotary_dim, + max_position=self.max_position_embeddings, + base=self.rope_theta, + rope_scaling=config.rope_scaling, + device=get_global_server_args().device, + ) + self.attn = RadixAttention( + self.num_heads, + self.head_dim, + self.scaling, + num_kv_heads=self.num_kv_heads, + layer_id=layer_id, + quant_config=quant_config, + prefix=f"{prefix}.attn", + ) + + def _apply_qk_norm( + self, q: torch.Tensor, k: torch.Tensor + ) -> Tuple[torch.Tensor, torch.Tensor]: + q_by_head = q.reshape(-1, self.head_dim) + q_by_head = self.query_layernorm(q_by_head) + q = q_by_head.view(q.shape) + k_by_head = k.reshape(-1, self.head_dim) + k_by_head = self.key_layernorm(k_by_head) + k = k_by_head.view(k.shape) + return q, k + + def forward( + self, + hidden_states: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + **kwargs, + ) -> torch.Tensor: + qkv, _ = self.query_key_value(hidden_states) + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + if self.use_qk_norm: + q, k = self._apply_qk_norm(q, k) + q, k = self.rotary_emb(positions, q, k) + attn_output = self.attn(q, k, v, forward_batch) + output, _ = self.dense(attn_output) + return output + + +class BailingMoELinearDecoderLayer(nn.Module): + + def __init__( + self, + config: PretrainedConfig, + quant_config: Optional[QuantizationConfig] = None, + layer_id: int = 0, + prefix: str = "layer", + is_nextn: bool = False, + ) -> None: + super().__init__() + self.layer_id = layer_id + self.use_mla = getattr(config, "full_attention_type", "mla") == "mla" + alt_stream = None # tptest + # todo nextn + + if config.attention_type == 0: # Linear layer + self.attention = BailingMoELinearAttention( + config, + quant_config=quant_config, + layer_id=self.layer_id, + prefix=prefix + ".attention", + ) + elif config.attention_type == 1: # softmax layer + if self.use_mla: + self.attention = DsV3MLA( + config=config, + hidden_size=config.hidden_size, + num_heads=config.num_attention_heads, + qk_nope_head_dim=config.qk_nope_head_dim, + qk_rope_head_dim=config.qk_rope_head_dim, + v_head_dim=config.v_head_dim, + q_lora_rank=( + config.q_lora_rank if hasattr(config, "q_lora_rank") else None + ), + kv_lora_rank=config.kv_lora_rank, + rope_theta=getattr(config, "rope_theta", 600000), + rope_scaling=config.rope_scaling, + max_position_embeddings=262144, + quant_config=quant_config, + layer_id=layer_id, + reduce_results=False, + prefix=add_prefix("attention", prefix), + alt_stream=alt_stream, + ) + else: + logger.info(f"==={layer_id=} use gqa") + self.attention = BailingMoEAttention( + config, + quant_config=quant_config, + layer_id=self.layer_id, + prefix=prefix + ".attention", + ) + else: + raise ValueError(f"Unsupported attention type: {config.attention_type}") + + self.expert_num = config.num_experts + self.hidden_size = config.hidden_size + is_moe_layer = self._is_layer_sparse(config, self.layer_id) + is_previous_moe_layer = self._is_layer_sparse(config, self.layer_id - 1) + is_next_layer_moe_layer = self._is_layer_sparse(config, self.layer_id + 1) + if self.expert_num == 1: + self.mlp = BailingMLP( + hidden_size=self.hidden_size, + intermediate_size=config.intermediate_size, + quant_config=quant_config, + prefix=add_prefix("mlp", prefix), + ) + else: + if is_nextn or self.layer_id >= config.first_k_dense_replace: + # MoE layer + self.mlp = BailingMoE( + config, + quant_config=quant_config, + layer_id=self.layer_id, + prefix=add_prefix("mlp", prefix), + ) + else: + # dense layer + self.mlp = BailingMLP( + hidden_size=self.hidden_size, + intermediate_size=config.intermediate_size, + quant_config=quant_config, + prefix=add_prefix("mlp", prefix), + ) + rms_norm_eps = float(getattr(config, "rms_norm_eps", 1e-5)) + self.input_layernorm = RMSNorm(self.hidden_size, eps=rms_norm_eps) + self.post_attention_layernorm = RMSNorm(self.hidden_size, eps=rms_norm_eps) + + self.layer_scatter_modes = LayerScatterModes.init_new( + layer_id=layer_id, + num_layers=config.num_hidden_layers, + is_layer_sparse=is_moe_layer, + is_previous_layer_sparse=is_previous_moe_layer, + is_next_layer_sparse=is_next_layer_moe_layer, + ) + + qkv_latent_func = ( + self.attention.prepare_qkv_latent + if config.attention_type == 1 and self.use_mla + else None + ) + self.layer_communicator = LayerCommunicator( + layer_scatter_modes=self.layer_scatter_modes, + input_layernorm=self.input_layernorm, + post_attention_layernorm=self.post_attention_layernorm, + allow_reduce_scatter=False, + qkv_latent_func=qkv_latent_func, + ) + + def _is_layer_sparse( + self, config: PretrainedConfig, layer_id: int, is_nextn: bool = False + ) -> bool: + return is_nextn or ( + config.num_experts is not None and layer_id >= config.first_k_dense_replace + ) + + def forward( + self, + hidden_states: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + residual: Optional[torch.Tensor], + zero_allocator: BumpAllocator, + **kwargs, + ) -> tuple[torch.Tensor, torch.Tensor]: + hidden_states, residual = self.layer_communicator.prepare_attn( + hidden_states, residual, forward_batch + ) + # logger.warning( + # f"===={self.layer_id=}, 1 shape= {hidden_states.shape}, {residual.shape}" + # ) + if not forward_batch.forward_mode.is_idle(): + if self.use_mla: + hidden_states = self.attention( + positions=positions, + hidden_states=hidden_states, + forward_batch=forward_batch, + zero_allocator=zero_allocator, + ) + else: + hidden_states = self.attention( + hidden_states=hidden_states, + positions=positions, + forward_batch=forward_batch, + ) + # logger.warning( + # f"===={self.layer_id=}, 2 shape= {hidden_states.shape}, {residual.shape}" + # ) + hidden_states, residual = self.layer_communicator.prepare_mlp( + hidden_states, residual, forward_batch + ) + # logger.warning( + # f"===={self.layer_id=}, 3 shape= {hidden_states.shape}, {residual.shape}" + # ) + should_allreduce_fusion = ( + self.layer_communicator.should_fuse_mlp_allreduce_with_next_layer( + forward_batch + ) + ) + use_reduce_scatter = self.layer_communicator.should_use_reduce_scatter( + forward_batch + ) + hidden_states = self.mlp( + hidden_states, should_allreduce_fusion, use_reduce_scatter + ) + hidden_states, residual = self.layer_communicator.postprocess_layer( + hidden_states, residual, forward_batch + ) + return hidden_states, residual + + @staticmethod + def shared_moe_coefficient_loader( + param: torch.Tensor, loaded_weight: torch.Tensor + ) -> None: + assert param.size() == loaded_weight.size() + + param.data.copy_(loaded_weight.to(torch.float32)) + return + + +class BailingMoELinearModel(nn.Module): + + def __init__( + self, + config: PretrainedConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.pp_group = get_pp_group() + self.config = config + self.vocab_size = config.vocab_size + self.embed_dim = config.hidden_size + self.num_layers = config.num_hidden_layers + + self.layer_group_size = getattr(config, "layer_group_size", 1) + self.decoder_attention_types = [ + 0 if is_linear_layer(i, self.layer_group_size) else 1 + for i in range(self.num_layers) + ] + logger.info( + f"attention type of layers:{self.decoder_attention_types}, 0 is linear layer and 1 is softmax layer!" + ) + + assert ( + self.num_layers % self.layer_group_size == 0 + ), f"num_layers={self.num_layers} must be divided by layer_group_size={self.layer_group_size}" + + if self.pp_group.is_first_rank: + self.word_embeddings = VocabParallelEmbedding( + self.vocab_size, + self.embed_dim, + enable_tp=not is_dp_attention_enabled(), + org_num_embeddings=self.vocab_size, + ) + else: + self.word_embeddings = PPMissingLayer() + + def layer_fn(idx, prefix): + layer_idx = idx + layer_config = copy.deepcopy(config) + layer_config.attention_type = self.decoder_attention_types[layer_idx] + + decoder_kwargs = {"quant_config": quant_config, "layer_id": layer_idx} + return BailingMoELinearDecoderLayer( + layer_config, **decoder_kwargs, prefix=prefix + ) + + self.layers, self.start_layer, self.end_layer = make_layers( + self.num_layers, + layer_fn, + pp_rank=self.pp_group.rank_in_group, + pp_size=self.pp_group.world_size, + prefix=f"{prefix}.layers", + ) + + linear_layer_nums = sum( + 1 for i in range(self.num_layers) if self.decoder_attention_types[i] == 0 + ) + logger.info(f"linear_layer_nums={linear_layer_nums}") + + norm_kwargs = {} + if hasattr(config, "rms_norm_eps"): + norm_kwargs["eps"] = config.rms_norm_eps + if self.pp_group.is_last_rank: + self.norm = RMSNorm(config.hidden_size, **norm_kwargs) + else: + self.norm = PPMissingLayer() + self.embed_scale = 1.0 + return + + def forward( + self, + input_ids: Optional[torch.Tensor], + positions: torch.Tensor, + forward_batch: Optional[ForwardBatch] = None, + inputs_embeds: Optional[torch.Tensor] = None, + pp_proxy_tensors: Optional[PPProxyTensors] = None, + ) -> Union[torch.Tensor, PPProxyTensors]: + if self.pp_group.is_first_rank: + if inputs_embeds is None: + hidden_states = self.word_embeddings(input_ids) + else: + hidden_states = inputs_embeds + residual = None + else: + assert pp_proxy_tensors is not None + hidden_states = pp_proxy_tensors["hidden_states"] + residual = pp_proxy_tensors["residual"] + + total_num_layers = self.end_layer - self.start_layer + device = inputs_embeds.device if inputs_embeds is not None else input_ids.device + zero_allocator = BumpAllocator( + buffer_size=total_num_layers * 2 * (2 if forward_batch.can_run_tbo else 1), + dtype=torch.float32, + device=device, + ) + + for i in range(self.start_layer, self.end_layer): + with get_global_expert_distribution_recorder().with_current_layer(i): + layer = self.layers[i] + hidden_states, residual = layer( + hidden_states=hidden_states, + positions=positions, + forward_batch=forward_batch, + residual=residual, + zero_allocator=zero_allocator, + ) + if not self.pp_group.is_last_rank: + return PPProxyTensors( + {"hidden_states": hidden_states, "residual": residual} + ) + else: + if not forward_batch.forward_mode.is_idle(): + if residual is None: + hidden_states = self.norm(hidden_states) + else: + hidden_states, _ = self.norm(hidden_states, residual) + return hidden_states + + +class BailingMoELinearForCausalLM(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.out_proj", + "layers.7.attention.out_proj": "layers.7.attention.o_proj", + "layers.15.attention.out_proj": "layers.15.attention.o_proj", + "layers.23.attention.out_proj": "layers.23.attention.o_proj", + "layers.31.attention.out_proj": "layers.31.attention.o_proj", + "layers.39.attention.out_proj": "layers.39.attention.o_proj", + "layers.47.attention.out_proj": "layers.47.attention.o_proj", + "layers.55.attention.out_proj": "layers.55.attention.o_proj", + "layers.63.attention.out_proj": "layers.63.attention.o_proj", + "layers.71.attention.out_proj": "layers.71.attention.o_proj", + "layers.79.attention.out_proj": "layers.79.attention.o_proj", + "attention.query_key_value": "attention.qkv_proj", + "attention.g_proj": "attention.output_gate", + }, + ) + + def __init__( + self, + *, + config, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.pp_group = get_pp_group() + self.config = config + self.quant_config = quant_config + self.model = BailingMoELinearModel( + self.config, quant_config, prefix=add_prefix("model", prefix) + ) + + if self.pp_group.is_last_rank: + self.lm_head = ( + self.word_embeddings + if config.tie_word_embeddings + else ParallelLMHead( + config.vocab_size, + config.hidden_size, + params_dtype=torch.float32, + quant_config=quant_config, + use_attn_tp_group=get_global_server_args().enable_dp_lm_head, + ) + ) + self.logits_processor = LogitsProcessor(config) + else: + self.lm_head = PPMissingLayer() + + @property + def start_layer(self): + return self.model.start_layer + + @property + def end_layer(self): + return self.model.end_layer + + def get_embed_and_head(self): + """Used by the eagle_worker.""" + return self.model.word_embeddings.weight, self.lm_head.weight + + def post_load_weights(self, is_nextn=False, weight_names=None): + + # Perform post-processing after loading weights + if is_nextn: + layer_ids = [self.config.num_hidden_layers] + else: + if weight_names is None: + layer_ids = range(self.model.start_layer, self.model.end_layer) + else: + layer_ids = set() + for name in weight_names: + if "kv_b_proj" in name: + layer_id = int(name.split(".")[2]) + if ( + layer_id < self.model.end_layer + and layer_id >= self.model.start_layer + ): + layer_ids.add(layer_id) + logger.info(f"=====layer_ids {layer_ids}") + + for layer_id in layer_ids: + self_attn = ( + self.model.layers[layer_id].attention + if not is_nextn + else self.model.decoder.attention + ) + if not hasattr(self_attn, "kv_b_proj"): + continue + if hasattr(self_attn.kv_b_proj, "qweight"): + # AWQ compatible + if _is_cuda or _is_hip: + w = awq_dequantize( + self_attn.kv_b_proj.qweight, + self_attn.kv_b_proj.scales, + self_attn.kv_b_proj.qzeros, + ).T + else: + w = awq_dequantize( + self_attn.kv_b_proj.qweight, + self_attn.kv_b_proj.scales, + self_attn.kv_b_proj.qzeros, + 0, + 0, + 0, + ).T + else: + w = self_attn.kv_b_proj.weight + # NOTE(HandH1998): Since `bmm_fp8` only supports per-tensor scale, we have to requantize `self_attn.kv_b_proj`. + # This may affect the accuracy of fp8 model. + # Fix deepseek v3 blockwise bmm by using deep_gemm + use_deep_gemm_bmm = False + + if w.dtype in ( + torch.float8_e4m3fn, + torch.float8_e4m3fnuz, + ): + if ( + hasattr(self.quant_config, "weight_block_size") + and self.quant_config.weight_block_size is not None + ): + weight_block_size = self.quant_config.weight_block_size + assert hasattr(self_attn.kv_b_proj, "weight_scale_inv") + if _is_fp8_fnuz: + weight, weight_scale, _ = normalize_e4m3fn_to_e4m3fnuz( + weight=w, + weight_scale=self_attn.kv_b_proj.weight_scale_inv, + input_scale=None, + ) + else: + weight = w + weight_scale = self_attn.kv_b_proj.weight_scale_inv + + if ( + _is_cuda + and weight_block_size[0] == 128 + and weight_block_size[1] == 128 + ): + if ( + deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM + and not deep_gemm_wrapper.DEEPGEMM_BLACKWELL + and get_bool_env_var("SGL_USE_DEEPGEMM_BMM", "false") + ): + block_scale = weight_scale + use_deep_gemm_bmm = True + else: + w = block_quant_dequant( + weight, + weight_scale, + weight_block_size, + torch.bfloat16, + ) + else: + w, scale = block_quant_to_tensor_quant( + weight, weight_scale, weight_block_size + ) + self_attn.w_scale = scale + else: + if _is_fp8_fnuz: + weight, weight_scale, _ = normalize_e4m3fn_to_e4m3fnuz( + weight=w, + weight_scale=self_attn.kv_b_proj.weight_scale, + input_scale=None, + ) + else: + weight = w + weight_scale = self_attn.kv_b_proj.weight_scale + + w, scale = channel_quant_to_tensor_quant(weight, weight_scale) + self_attn.w_scale = scale + + if w.dtype == torch.int8: + if hasattr(self.quant_config, "weight_block_size"): + # block-wise int8 need it + weight_block_size = self.quant_config.weight_block_size + if weight_block_size is not None: + assert hasattr(self_attn.kv_b_proj, "weight_scale_inv") + weight = w + weight_scale = self_attn.kv_b_proj.weight_scale_inv + w = int8_block_dequant( + weight, weight_scale, weight_block_size + ).to(torch.bfloat16) + else: + # channel-wise int8 need it + w = w.to(torch.bfloat16) * self_attn.kv_b_proj.weight_scale.to( + torch.bfloat16 + ) + + w_kc, w_vc = w.unflatten( + 0, (-1, self_attn.qk_nope_head_dim + self_attn.v_head_dim) + ).split([self_attn.qk_nope_head_dim, self_attn.v_head_dim], dim=1) + if not use_deep_gemm_bmm: + self_attn.w_kc = bind_or_assign( + self_attn.w_kc, w_kc.transpose(1, 2).contiguous().transpose(1, 2) + ) + self_attn.w_vc = bind_or_assign( + self_attn.w_vc, w_vc.contiguous().transpose(1, 2) + ) + if ( + hasattr(self_attn.kv_b_proj, "weight_scale") + and self_attn.w_scale is None + ): + self_attn.w_scale = bind_or_assign( + self_attn.w_scale, self_attn.kv_b_proj.weight_scale + ) + if _is_hip: + self_attn.w_scale *= 2.0 + # TODO: remove this after adding FP8 support in bmm cpu kernel + if _is_cpu and _is_cpu_amx_available and w.dtype == torch.float8_e4m3fn: + self_attn.w_kc = ( + self_attn.w_kc.to(torch.bfloat16) * self_attn.w_scale + ) + self_attn.w_vc = ( + self_attn.w_vc.to(torch.bfloat16) * self_attn.w_scale + ) + else: + num_tiles_k = self_attn.qk_nope_head_dim // weight_block_size[1] + num_tiles_n = self_attn.v_head_dim // weight_block_size[0] + ws_kc, ws_vc = block_scale.unflatten( + 0, (-1, (num_tiles_k + num_tiles_n)) + ).split([num_tiles_k, num_tiles_n], dim=1) + self_attn.w_scale_k = bind_or_assign( + self_attn.w_scale_k, ws_kc.transpose(1, 2).contiguous() + ) + self_attn.w_scale_v = bind_or_assign( + self_attn.w_scale_v, ws_vc.contiguous() + ) + self_attn.w_kc = bind_or_assign( + self_attn.w_kc, w_kc.transpose(1, 2).contiguous() + ) + self_attn.w_vc = bind_or_assign(self_attn.w_vc, w_vc.contiguous()) + self_attn.use_deep_gemm_bmm = True + + if ( + deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM + and deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0 + and hasattr(self.quant_config, "weight_block_size") + and self.quant_config.weight_block_size is not None + ): + self._weight_requant_ue8m0(is_nextn) + + @classmethod + def get_model_config_for_expert_location(cls, config): + num_groups = getattr(config, "n_group", 0) + from sglang.srt.eplb.expert_location import ModelConfigForExpertLocation + + return ModelConfigForExpertLocation( + num_layers=config.num_hidden_layers, + num_logical_experts=config.num_experts, + num_groups=None if num_groups == 0 else num_groups, + ) + + def _weight_requant_ue8m0(self, is_nextn=False): + weight_block_size = self.quant_config.weight_block_size + + moe_layers = list( + range( + self.config.first_k_dense_replace, + self.config.num_hidden_layers, + self.config.moe_layer_freq, + ) + ) + + num_hidden_layers = 1 if is_nextn else self.config.num_hidden_layers + + for layer_id in range(num_hidden_layers): + if is_nextn: + layer = self.model.decoder + else: + layer = self.model.layers[layer_id] + + module_list = [ + layer.self_attn.kv_b_proj, + layer.self_attn.o_proj, + ] + + if self.config.q_lora_rank is not None: + module_list.append(layer.self_attn.fused_qkv_a_proj_with_mqa) + module_list.append(layer.self_attn.q_b_proj) + else: + module_list.append(layer.self_attn.kv_a_proj_with_mqa) + module_list.append(layer.self_attn.q_proj) + + for module in module_list: + requant_weight_ue8m0_inplace( + module.weight, module.weight_scale_inv, weight_block_size + ) + + if layer_id in moe_layers or is_nextn: + shared_experts = getattr(layer.mlp, "shared_experts", None) + if shared_experts is not None: + for module in [ + shared_experts.gate_up_proj, + shared_experts.down_proj, + ]: + requant_weight_ue8m0_inplace( + module.weight, module.weight_scale_inv, weight_block_size + ) + + experts = layer.mlp.experts + if isinstance(experts, DeepEPMoE): + for w in [ + experts.w13_weight_fp8, + experts.w2_weight_fp8, + ]: + requant_weight_ue8m0_inplace(w[0], w[1], weight_block_size) + else: + mlp = layer.mlp + assert isinstance(mlp, DeepseekV2MLP) + for module in [ + mlp.gate_up_proj, + mlp.down_proj, + ]: + requant_weight_ue8m0_inplace( + module.weight, module.weight_scale_inv, weight_block_size + ) + + def get_decoder_attention_types(self): + return self.model.decoder_attention_types + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + inputs_embeds: Optional[torch.Tensor] = None, + pp_proxy_tensors: Optional[PPProxyTensors] = None, + ) -> Union[torch.Tensor, PPProxyTensors]: + hidden_states = self.model( + input_ids=input_ids, + positions=positions, + inputs_embeds=inputs_embeds, + forward_batch=forward_batch, + pp_proxy_tensors=pp_proxy_tensors, + ) + if self.pp_group.is_last_rank: + return self.logits_processor( + input_ids, hidden_states.float(), self.lm_head, forward_batch + ) + else: + return hidden_states + + def load_weights( + self, weights: Iterable[Tuple[str, torch.Tensor]], is_nextn=False + ) -> Set[str]: + def load_linear_attn_weight( + name: str, loaded_weight: torch.Tensor, self + ) -> None: + if is_pp_missing_parameter(name, self): + return + param = params_dict[name] + weight_loader = getattr( + param, "weight_loader", BailingMoELinearAttention.weight_direct_load + ) + weight_loader = weight_loader_with_alias(name)(weight_loader) + weight_loader(param, loaded_weight) + return + + if is_nextn: + if hasattr(self.config, "num_nextn_predict_layers"): + num_nextn_layers = self.config.num_nextn_predict_layers + assert num_nextn_layers == 1, "Only 1 nextn layer is supported" + # compatible with old design + nextn_layer_id = ( + 0 + if self.config.num_hidden_layers == 1 + else self.config.num_hidden_layers + ) + else: + raise ValueError("num nextn_predict_layers is not in the config") + + stacked_params_mapping = [ + # (param_name, shard_name, shard_id) + ("gate_up_proj", "gate_proj", 0), + ("gate_up_proj", "up_proj", 1), + ] + expert_params_mapping = FusedMoE.make_expert_params_mapping( + ckpt_gate_proj_name="gate_proj", + ckpt_down_proj_name="down_proj", + ckpt_up_proj_name="up_proj", + num_experts=self.config.num_experts, + ) + + if is_nextn: + nextn_layer_prefix = f"model.layers.{nextn_layer_id}" + nextn_spec_weight_names = [ + "final_layernorm", + "eh_proj", + "enorm", + "hnorm", + ] + + params_dict = dict(self.named_parameters()) + loaded_params: Set[str] = set() + weight_names = [] + fuse_qkv_a_proj = hasattr(self.config, "q_lora_rank") and ( + self.config.q_lora_rank is not None + ) + cached_a_proj = {} if fuse_qkv_a_proj else None + + for name, loaded_weight in weights: + if name.startswith("model.mtp"): + continue + layer_idx = None + if "model.layers." in name: + layer_idx = int(name.split(".")[2]) + if ( + ("v_head" in name) + or ("inv_freq" in name) + or (self.config.tie_word_embeddings and "lm_head" in name) + ): + continue + + weight_names.append(name) + + if is_nextn: + if not name.startswith(nextn_layer_prefix): + continue + + # Use shared head and embed weights from target model + if "shared_head.head" in name or "embed_tokens" in name: + continue + + is_decoder = True + # For nextn specific weights + for weight_name in nextn_spec_weight_names: + if weight_name in name: + name = name.replace(nextn_layer_prefix, "model") + is_decoder = False + break + # For decoder layer weights + if is_decoder: + name = name.replace(nextn_layer_prefix, "model.decoder") + + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in name: + continue + if "mlp.experts" in name: + continue + + name = name.replace(weight_name, param_name) + if name.endswith(".bias") and name not in params_dict: + continue + if name not in params_dict: + continue + if is_pp_missing_parameter(name, self): + continue + + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) + break + else: + + for mapping in expert_params_mapping: + param_name, weight_name, expert_id, shard_id = mapping + if weight_name not in name: + continue + name = name.replace(weight_name, param_name) + + if name not in params_dict: + continue + if is_pp_missing_parameter(name, self): + continue + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader( + param, + loaded_weight, + name, + shard_id=shard_id, + expert_id=expert_id, + ) + break + else: + + if name.endswith(".bias") and name not in params_dict: + continue + if "slope" in name: + continue + + if fuse_qkv_a_proj and ( + "q_a_proj" in name or "kv_a_proj_with_mqa" in name + ): + cached_a_proj[name] = loaded_weight + q_a_proj_name = ( + name + if "q_a_proj" in name + else name.replace("kv_a_proj_with_mqa", "q_a_proj") + ) + kv_a_proj_name = ( + name + if "kv_a_proj_with_mqa" in name + else name.replace("q_a_proj", "kv_a_proj_with_mqa") + ) + + # When both q_a_proj and kv_a_proj_with_mqa has been cached, load the fused weight to parameter + if ( + q_a_proj_name in cached_a_proj + and kv_a_proj_name in cached_a_proj + ): + q_a_proj_weight = cached_a_proj[q_a_proj_name] + kv_a_proj_weight = cached_a_proj[kv_a_proj_name] + cat_dim = 0 + if self.quant_config is not None and ( + self.quant_config.get_name() == "awq" + or self.quant_config.get_name() == "awq_marlin" + or self.quant_config.get_name() == "moe_wna16" + ): + cat_dim = 1 + fused_weight = torch.cat( + [q_a_proj_weight, kv_a_proj_weight], dim=cat_dim + ) + param_name = ( + name.replace("q_a_proj", "fused_qkv_a_proj_with_mqa") + if "q_a_proj" in name + else name.replace( + "kv_a_proj_with_mqa", + "fused_qkv_a_proj_with_mqa", + ) + ) + if param_name not in params_dict: + continue + param = params_dict[param_name] + weight_loader = getattr( + param, "weight_loader", default_weight_loader + ) + + weight_loader(param, fused_weight) + cached_a_proj.pop(q_a_proj_name) + cached_a_proj.pop(kv_a_proj_name) + else: + + if name not in params_dict: + name = name.replace(".dense.", ".o_proj.") + if name not in params_dict: + continue + if is_pp_missing_parameter(name, self): + continue + if ( + "attention" in name + and "slope" not in name + and is_linear_layer(layer_idx, self.model.layer_group_size) + ): + load_linear_attn_weight(name, loaded_weight, self) + loaded_params.add(name) + continue + + param = params_dict[name] + weight_loader = getattr( + param, "weight_loader", default_weight_loader + ) + weight_loader(param, loaded_weight) + loaded_params.add(name) + self.post_load_weights(is_nextn=is_nextn, weight_names=weight_names) + + return loaded_params + + +class BailingMoeV2_5ForCausalLM(BailingMoELinearForCausalLM): + pass + + +EntryClass = [ + BailingMoeV2_5ForCausalLM, +] diff --git a/python/sglang/srt/models/bailing_moe_nextn.py b/python/sglang/srt/models/bailing_moe_nextn.py index bcfb75149..a365acdff 100644 --- a/python/sglang/srt/models/bailing_moe_nextn.py +++ b/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,16 +87,29 @@ class BailingMoEModelNextN(nn.Module): config.hidden_size, bias=False, quant_config=quant_config, - prefix=add_prefix("eh_proj", prefix), + prefix=add_prefix(f"layers.{config.num_hidden_layers}.eh_proj", prefix), ) - self.decoder = BailingMoEBlock( - config, - 0, - quant_config=quant_config, - # is_nextn=True, - prefix=add_prefix("decoder", prefix), + self.is_hybrid = ( + hasattr(config, "model_type") and config.model_type == "bailing_hybrid" ) + if self.is_hybrid: + config.attention_type = 1 + self.decoder = BailingMoELinearDecoderLayer( + config, + quant_config=quant_config, + layer_id=0, + is_nextn=True, + prefix=add_prefix(f"layers.{config.num_hidden_layers}", prefix), + ) + else: + self.decoder = BailingMoEBlock( + config, + 0, + quant_config=quant_config, + # is_nextn=True, + prefix=add_prefix("decoder", prefix), + ) self.shared_head = nn.Module() self.final_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) @@ -107,16 +132,37 @@ class BailingMoEModelNextN(nn.Module): torch.cat( ( self.enorm(hidden_states), - self.hnorm(forward_batch.spec_info.hidden_states), + self.hnorm( + forward_batch.spec_info.hidden_states.to( + self.hnorm.weight.dtype + ) + ), ), dim=-1, ) ) residual = None - hidden_states, residual = self.decoder( - positions, hidden_states, forward_batch, residual - ) + if self.is_hybrid: + device = input_ids.device + zero_allocator = BumpAllocator( + buffer_size=self.total_num_layers + * 2 + * (2 if forward_batch.can_run_tbo else 1), + dtype=torch.float32, + device=device, + ) + hidden_states, residual = self.decoder( + hidden_states=hidden_states, + positions=positions, + forward_batch=forward_batch, + residual=residual, + zero_allocator=zero_allocator, + ) + else: + hidden_states, residual = self.decoder( + positions, hidden_states, forward_batch, residual + ) if not forward_batch.forward_mode.is_idle(): if residual is not None: @@ -127,7 +173,18 @@ class BailingMoEModelNextN(nn.Module): return hidden_states -class BailingMoeForCausalLMNextN(BailingMoEForCausalLM): +class BailingMoeForCausalLMNextN(nn.Module): + + packed_modules_mapping = { + "fused_qkv_a_proj_with_mqa": ["q_a_proj", "kv_a_proj_with_mqa"], + "gate_up_proj": ["gate_proj", "up_proj"], + } + # To ensure correct weight loading and mapping. + hf_to_sglang_mapper = WeightsMapper( + orig_to_new_substr={ + "attention.dense": "attention.o_proj", + }, + ) def __init__( self, @@ -154,6 +211,12 @@ class BailingMoeForCausalLMNextN(BailingMoEForCausalLM): use_attn_tp_group=get_global_server_args().enable_dp_lm_head, ) self.logits_processor = LogitsProcessor(config) + if hasattr(self.config, "model_type") and config.model_type == "bailing_hybrid": + self.base_load_weights_func = BailingMoeV2_5ForCausalLM.load_weights + self.post_load_weights_func = BailingMoeV2_5ForCausalLM.post_load_weights + else: + self.base_load_weights_func = BailingMoEForCausalLM.load_weights + self.post_load_weights_func = BailingMoEForCausalLM.post_load_weights @torch.no_grad() def forward( @@ -167,8 +230,20 @@ class BailingMoeForCausalLMNextN(BailingMoEForCausalLM): input_ids, hidden_states, self.lm_head, forward_batch ) + def set_embed_and_head(self, embed, head): + """Used by the eagle_worker.""" + del self.model.word_embeddings.weight + del self.lm_head.weight + self.model.word_embeddings.weight = embed + self.lm_head.weight = head + torch.cuda.empty_cache() + torch.cuda.synchronize() + def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): - super().load_weights(weights, is_nextn=True) + self.base_load_weights_func(self, weights, is_nextn=True) + + def post_load_weights(self, is_nextn=False, weight_names=None): + self.post_load_weights_func(self, is_nextn=is_nextn, weight_names=weight_names) EntryClass = [BailingMoeForCausalLMNextN] diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index d15e5238f..cf95d3e1b 100644 --- a/python/sglang/srt/server_args.py +++ b/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", diff --git a/python/sglang/srt/speculative/eagle_worker.py b/python/sglang/srt/speculative/eagle_worker.py index 333c206b2..6d3a76f8f 100644 --- a/python/sglang/srt/speculative/eagle_worker.py +++ b/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 diff --git a/python/sglang/srt/utils/hf_transformers_utils.py b/python/sglang/srt/utils/hf_transformers_utils.py index b4fd6c734..41f70942d 100644 --- a/python/sglang/srt/utils/hf_transformers_utils.py +++ b/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,