diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 4e1088894..f6ca9026e 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -340,6 +340,8 @@ class Envs: SGLANG_NPU_FORWARD_NATIVE_GEMMA_RMS_NORM = EnvBool(False) # Delay all-gather after qlora for better performance for Deepseek v3.2 SGLANG_USE_AG_AFTER_QLORA = EnvBool(False) + # Quantize x to int8 in the dispatch operator + DEEP_NORMAL_MODE_USE_INT8_QUANT = EnvBool(False) SGLANG_NPU_FUSED_MOE_MODE = EnvInt(1) # MTHREADS & MUSA diff --git a/python/sglang/srt/hardware_backend/npu/attention/ascend_gdn_backend.py b/python/sglang/srt/hardware_backend/npu/attention/ascend_gdn_backend.py new file mode 100644 index 000000000..6abe5c2fd --- /dev/null +++ b/python/sglang/srt/hardware_backend/npu/attention/ascend_gdn_backend.py @@ -0,0 +1,425 @@ +from typing import Optional, Tuple, Union + +import torch +from sgl_kernel_npu.fla.fused_gdn_gating import ( + fused_gdn_gating_kernel_without_sigmoid, + fused_gdn_gating_npu, +) +from sgl_kernel_npu.mamba.causal_conv1d import ( + causal_conv1d_fn_npu, + causal_conv1d_update_npu, +) + +from sglang.srt.hardware_backend.npu.attention.ascend_hybrid_linear_attn_backend import ( + AscendMambaAttnBackendBase, +) +from sglang.srt.layers.attention.linear.gdn_backend import GDNKernelDispatcher +from sglang.srt.layers.attention.linear.utils import ( + get_linear_attn_decode_backend, + get_linear_attn_prefill_backend, +) +from sglang.srt.layers.radix_linear_attention import RadixLinearAttention +from sglang.srt.mem_cache.memory_pool import MambaPool +from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode +from sglang.srt.model_executor.model_runner import ModelRunner +from sglang.srt.speculative.eagle_info import EagleDraftInput, EagleVerifyInput + +fused_gdn_gating = fused_gdn_gating_npu +causal_conv1d_fn = causal_conv1d_fn_npu +causal_conv1d_update = causal_conv1d_update_npu + + +class AscendGDNAttnBackend(AscendMambaAttnBackendBase): + + def __init__(self, model_runner: ModelRunner): + super().__init__(model_runner) + self.conv_states_shape = torch.Size( + ( + *model_runner.req_to_token_pool.mamba_pool.mamba_cache.conv[0].shape[ + :-2 + ], + model_runner.req_to_token_pool.mamba_pool.mamba_cache.conv[0].shape[-1], + model_runner.req_to_token_pool.mamba_pool.mamba_cache.conv[0].shape[-2], + ) + ) + decode_backend = get_linear_attn_decode_backend() + prefill_backend = get_linear_attn_prefill_backend() + self.kernel_dispatcher = GDNKernelDispatcher(decode_backend, prefill_backend) + + def prepare_gdn_inputs( + self, + bs: int, + forward_mode: ForwardMode, + spec_info: Optional[Union[EagleDraftInput, EagleVerifyInput]], + ): + cache_indices = self.forward_metadata.mamba_cache_indices + self.num_accepted_tokens = torch.ones( + [bs], dtype=torch.int32, device=cache_indices.device + ) + self.actual_seq_lengths = torch.ones( + [bs], dtype=torch.int32, device=cache_indices.device + ) + if forward_mode.is_target_verify(): + seq_len = spec_info.draft_token_num + self.actual_seq_lengths = self.actual_seq_lengths * seq_len + # indices + self.ssm_state_indices = torch.arange( + cache_indices.shape[0] * seq_len, + dtype=torch.int32, + device=cache_indices.device, + ) + else: + self.ssm_state_indices = cache_indices + + def init_forward_metadata(self, forward_batch: ForwardBatch): + if forward_batch.forward_mode.is_draft_extend(True): + return + super().init_forward_metadata(forward_batch) + self.prepare_gdn_inputs( + forward_batch.batch_size, + forward_batch.forward_mode, + forward_batch.spec_info, + ) + self.graph_mode = False + + 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]], + ): + if forward_mode.is_draft_extend(True): + return + super().init_forward_metadata_capture_cuda_graph( + bs, + num_tokens, + req_pool_indices, + seq_lens, + encoder_lens, + forward_mode, + spec_info, + ) + self.prepare_gdn_inputs(bs, forward_mode, spec_info) + self.graph_mode = True + + 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], + ): + if forward_mode.is_draft_extend(True): + return + super().init_forward_metadata_replay_cuda_graph( + bs, + req_pool_indices, + seq_lens, + seq_lens_sum, + encoder_lens, + forward_mode, + spec_info, + seq_lens_cpu, + ) + self.prepare_gdn_inputs(bs, forward_mode, spec_info) + self.graph_mode = True + + def forward_decode( + self, + layer: RadixLinearAttention, + forward_batch: ForwardBatch, + mixed_qkv: Union[torch.Tensor, Tuple[torch.Tensor, ...]], + a: torch.Tensor, + b: torch.Tensor, + **kwargs, + ): + layer_cache = self.req_to_token_pool.mamba2_layer_cache(layer.layer_id) + conv_states = layer_cache.conv[0] + ssm_states = layer_cache.temporal + query_start_loc = self.forward_metadata.query_start_loc + cache_indices = self.forward_metadata.mamba_cache_indices + + assert isinstance(mixed_qkv, torch.Tensor) + conv_states_tmp = conv_states.transpose(1, 2).clone() + mixed_qkv = causal_conv1d_update( + mixed_qkv, + conv_states_tmp, + layer.conv_weights, + layer.bias, + layer.activation, + conv_state_indices=cache_indices, + ) + conv_states[:] = conv_states_tmp.transpose(1, 2) + + query, key, value = torch.split( + mixed_qkv, + [layer.q_dim, layer.k_dim, layer.v_dim], + dim=-1, + ) + bs = forward_batch.batch_size + query = query.view(1, bs, layer.num_q_heads, layer.head_q_dim) + key = key.view(1, bs, layer.num_k_heads, layer.head_k_dim) + value = value.view(1, bs, layer.num_v_heads, layer.head_v_dim) + + core_attn_out = self.kernel_dispatcher.decode( + q=query, + k=key, + v=value, + a=a, + b=b, + A_log=layer.A_log, + dt_bias=layer.dt_bias, + ssm_states=ssm_states, + cache_indices=cache_indices, + query_start_loc=query_start_loc, + ) + + self._track_mamba_state_decode( + forward_batch, conv_states, ssm_states, cache_indices + ) + return core_attn_out + + def forward_extend( + self, + layer: RadixLinearAttention, + forward_batch: ForwardBatch, + mixed_qkv: Union[torch.Tensor, Tuple[torch.Tensor, ...]], + a: torch.Tensor, + b: torch.Tensor, + **kwargs, + ): + assert isinstance(mixed_qkv, torch.Tensor) + seq_len = mixed_qkv.shape[0] + is_target_verify = forward_batch.forward_mode.is_target_verify() + forward_metadata = self.forward_metadata + + query_start_loc = forward_metadata.query_start_loc + cache_indices = forward_metadata.mamba_cache_indices + retrieve_next_token = forward_metadata.retrieve_next_token + retrieve_next_sibling = forward_metadata.retrieve_next_sibling + retrieve_parent_token = forward_metadata.retrieve_parent_token + + mamba_cache_params = self.req_to_token_pool.mamba2_layer_cache(layer.layer_id) + conv_states = mamba_cache_params.conv[0] + ssm_states = mamba_cache_params.temporal + if is_target_verify: + assert isinstance(mamba_cache_params, MambaPool.SpeculativeState) + intermediate_state_cache = mamba_cache_params.intermediate_ssm + intermediate_conv_window_cache = ( + mamba_cache_params.intermediate_conv_window[0] + ) + has_initial_states = torch.ones( + seq_len // forward_batch.spec_info.draft_token_num, + dtype=torch.bool, + device=forward_batch.input_ids.device, + ) + else: + has_initial_states = forward_batch.extend_prefix_lens > 0 + if is_target_verify: + draft_token_num = forward_batch.spec_info.draft_token_num + num_token_padding = mixed_qkv.shape[0] + batch_size = cache_indices.shape[0] + if ( + not self.graph_mode + and forward_batch.num_token_non_padded_cpu != num_token_padding + ): + mixed_qkv = mixed_qkv[: forward_batch.num_token_non_padded_cpu] + a = a[: forward_batch.num_token_non_padded_cpu] + b = b[: forward_batch.num_token_non_padded_cpu] + seq_len = forward_batch.num_token_non_padded_cpu + + mixed_qkv_reshaped = mixed_qkv.view(batch_size, draft_token_num, -1) + num_accepted_tokens = torch.full( + (batch_size,), + draft_token_num, + dtype=torch.int32, + device=mixed_qkv.device, + ) + mixed_qkv = torch.ops.npu.causal_conv1d_update( + mixed_qkv_reshaped, + layer.conv_weights.transpose(0, 1).contiguous(), + conv_states, + cache_indices, + layer.bias, + num_accepted_tokens, + None, + layer.activation == "silu", + self.pad_slot_id, + ).view(seq_len, -1) + else: + mixed_qkv = mixed_qkv.transpose(0, 1) + if ( + forward_batch.mamba_track_mask is not None + and forward_batch.mamba_track_mask.any() + ): + conv_dst = forward_batch.mamba_track_indices + mixed_qkv_to_track = mixed_qkv[ + :, forward_metadata.track_conv_indices + ].transpose(0, 1) + mask_indices = forward_batch.mamba_track_mask.nonzero(as_tuple=True)[0] + conv_states.transpose(1, 2)[conv_dst[mask_indices]] = mixed_qkv_to_track + kernel_size = layer.conv_weights.shape[-1] + conv_states_for_prefill = conv_states[:, -(kernel_size - 1) :, :] + conv_states_tmp = conv_states_for_prefill.transpose(1, 2).contiguous() + + mixed_qkv = causal_conv1d_fn( + mixed_qkv, + layer.conv_weights, + layer.bias, + activation=layer.activation, + conv_states=conv_states_tmp, + has_initial_state=has_initial_states, + cache_indices=cache_indices, + query_start_loc=query_start_loc, + seq_lens_cpu=forward_batch.extend_seq_lens_cpu, + ).transpose(0, 1)[:seq_len] + conv_states[:, -(kernel_size - 1) :, :] = conv_states_tmp.transpose( + 1, 2 + ).contiguous() + if is_target_verify: + g, beta = fused_gdn_gating_kernel_without_sigmoid( + layer.A_log, a, b, layer.dt_bias + ) + beta = beta.unsqueeze(0) + num_heads, head_k_dim = layer.num_q_heads, layer.head_q_dim + num_value_heads, head_v_dim = layer.num_v_heads, layer.head_v_dim + + mixed_qkv_last_dim = mixed_qkv.shape[-1] + + mixed_qkv = mixed_qkv.view(batch_size, -1, mixed_qkv_last_dim) + beta = beta.view(batch_size, -1, num_value_heads) + g = g.view(batch_size, -1, num_value_heads) + + core_attn_out = self.fused_recurrent_gated_delta_rule_update( + mixed_qkv, + num_heads, + num_value_heads, + head_k_dim, + head_v_dim, + recurrent_state=ssm_states, + beta=beta, + g=g, + cache_indices=cache_indices, + intermediate_state=intermediate_state_cache, + ) + core_attn_out = core_attn_out.view(-1, num_value_heads, head_v_dim) + if (not self.graph_mode) and core_attn_out.shape[0] < num_token_padding: + core_attn_out = torch.cat( + [ + core_attn_out, + core_attn_out.new_zeros( + num_token_padding - core_attn_out.shape[0], + *core_attn_out.shape[1:], + ), + ], + dim=0, + ) + else: + query, key, value = torch.split( + mixed_qkv, + [layer.q_dim, layer.k_dim, layer.v_dim], + dim=-1, + ) + + actual_seq_len = query.shape[0] + query = query.view(1, actual_seq_len, layer.num_q_heads, layer.head_q_dim) + key = key.view(1, actual_seq_len, layer.num_k_heads, layer.head_k_dim) + value = value.view(1, actual_seq_len, layer.num_v_heads, layer.head_v_dim) + + g, beta = fused_gdn_gating(layer.A_log, a, b, layer.dt_bias) + core_attn_out, last_recurrent_state, h = self.kernel_dispatcher.extend( + q=query, + k=key, + v=value, + g=g, + beta=beta, + ssm_states=ssm_states, + cache_indices=cache_indices, + query_start_loc=query_start_loc, + ) + if last_recurrent_state is not None: + last_recurrent_state = last_recurrent_state.to( + ssm_states.dtype, copy=False + ) + ssm_states[cache_indices] = last_recurrent_state + if not forward_batch.spec_algorithm.is_none(): + last_recurrent_state = last_recurrent_state.transpose(-1, -2).to( + ssm_states.dtype, copy=False + ) + else: + last_recurrent_state = last_recurrent_state.to( + ssm_states.dtype, copy=False + ) + ssm_states[cache_indices] = last_recurrent_state + if h is not None: + self._track_mamba_state_extend( + forward_batch, h, ssm_states, forward_metadata + ) + + return core_attn_out + + def fused_recurrent_gated_delta_rule_update( + self, + mix_qkv: torch.Tensor, + num_heads, + num_value_heads, + head_k_dim, + head_v_dim, + recurrent_state: torch.Tensor, + beta: torch.Tensor, + g: torch.Tensor, + cache_indices: torch.Tensor, + intermediate_state: Optional[torch.Tensor] = None, + ): + beta = beta.to(torch.bfloat16) + g = g.to(torch.float32) + batch_size = mix_qkv.shape[0] + seq_len = mix_qkv.shape[1] + scale = 1 / (head_k_dim**0.5) + + if intermediate_state is not None: + intermediate_state = intermediate_state.view( + -1, num_value_heads, head_k_dim, head_v_dim + ) + + if self.graph_mode: + num_accepted_tokens = torch.full( + [batch_size], 1, dtype=torch.int32, device=cache_indices.device + ) + actual_seq_lengths = torch.full( + [batch_size], seq_len, dtype=torch.int32, device=cache_indices.device + ) + ssm_state_indices = self.forward_metadata.mamba_cache_indices_gdn + else: + num_accepted_tokens = self.num_accepted_tokens + actual_seq_lengths = self.actual_seq_lengths + ssm_state_indices = self.ssm_state_indices + + attn_core_out = torch.ops.npu.recurrent_gated_delta_rule( + mix_qkv, + recurrent_state, + beta=beta, + scale=scale, + actual_seq_lengths=actual_seq_lengths, + ssm_state_indices=ssm_state_indices.view(batch_size, seq_len), + nk=num_heads, + nv=num_value_heads, + intermediate_state=intermediate_state, + cache_indices=cache_indices, + num_accepted_tokens=num_accepted_tokens, + g=g, + ) + + if intermediate_state is not None: + intermediate_state = intermediate_state.view( + -1, seq_len, num_value_heads, head_k_dim, head_v_dim + ) + return attn_core_out diff --git a/python/sglang/srt/hardware_backend/npu/attention/ascend_hybrid_linear_attn_backend.py b/python/sglang/srt/hardware_backend/npu/attention/ascend_hybrid_linear_attn_backend.py new file mode 100644 index 000000000..171a612a9 --- /dev/null +++ b/python/sglang/srt/hardware_backend/npu/attention/ascend_hybrid_linear_attn_backend.py @@ -0,0 +1,280 @@ +import logging +from typing import Optional, Union + +import torch +from sgl_kernel_npu.mamba.mamba_state_update_triton import ( + conv_state_rollback, + move_intermediate_cache, +) + +from sglang.srt.layers.attention.base_attn_backend import AttentionBackend +from sglang.srt.layers.attention.hybrid_linear_attn_backend import ( + HybridLinearAttnBackend, + MambaAttnBackendBase, +) +from sglang.srt.layers.attention.mamba.mamba2_metadata import ( + ForwardMetadata, +) +from sglang.srt.model_executor.forward_batch_info import ForwardMode +from sglang.srt.model_executor.model_runner import ModelRunner +from sglang.srt.speculative.eagle_info import EagleDraftInput, EagleVerifyInput +from sglang.srt.speculative.spec_info import SpecInput + +logger = logging.getLogger(__name__) + + +class AscendMambaAttnBackendBase(MambaAttnBackendBase): + def __init__(self, model_runner: ModelRunner): + super().__init__(model_runner) + self.state_indices_list_gdn = [] + + def init_cuda_graph_state(self, max_bs: int, max_num_tokens: int): + assert ( + max_num_tokens % max_bs == 0 + ), f"max_num_tokens={max_num_tokens} must be divisible by max_bs={max_bs}" + draft_token_num = max_num_tokens // max_bs + for i in range(max_bs): + self.state_indices_list.append( + torch.full( + (i + 1,), self.pad_slot_id, dtype=torch.int32, device=self.device + ) + ) + self.state_indices_list_gdn.append( + torch.full( + ((i + 1) * draft_token_num,), + self.pad_slot_id, + dtype=torch.int32, + device=self.device, + ) + ) + self.query_start_loc_list.append( + torch.zeros((i + 2,), dtype=torch.int32, device=self.device) + ) + self.retrieve_next_token_list.append( + torch.zeros( + (i + 1, draft_token_num), dtype=torch.int32, device=self.device + ) + ) + self.retrieve_next_sibling_list.append( + torch.zeros( + (i + 1, draft_token_num), dtype=torch.int32, device=self.device + ) + ) + self.retrieve_parent_token_list.append( + torch.zeros( + (i + 1, draft_token_num), dtype=torch.int32, device=self.device + ) + ) + self.cached_cuda_graph_decode_query_start_loc = torch.arange( + 0, max_bs + 1, dtype=torch.int32, device=self.device + ) + self.cached_cuda_graph_verify_query_start_loc = torch.arange( + 0, + max_bs * draft_token_num + 1, + step=draft_token_num, + dtype=torch.int32, + device=self.device, + ) + + def _capture_metadata( + self, + bs: int, + req_pool_indices: torch.Tensor, + forward_mode: ForwardMode, + spec_info: Optional[Union[EagleDraftInput, EagleVerifyInput]], + ): + mamba_indices = self.req_to_token_pool.get_mamba_indices(req_pool_indices) + self.state_indices_list[bs - 1][: len(mamba_indices)].copy_(mamba_indices) + if forward_mode.is_decode_or_idle(): + self.query_start_loc_list[bs - 1].copy_( + self.cached_cuda_graph_decode_query_start_loc[: bs + 1] + ) + elif forward_mode.is_target_verify(): + self.query_start_loc_list[bs - 1].copy_( + self.cached_cuda_graph_verify_query_start_loc[: bs + 1] + ) + ssm_state_indices = torch.arange( + mamba_indices.shape[0] * spec_info.draft_token_num, + dtype=torch.int32, + device=mamba_indices.device, + ) + self.state_indices_list_gdn[bs - 1][ + : len(mamba_indices) * spec_info.draft_token_num + ].copy_(ssm_state_indices) + else: + raise ValueError(f"Invalid forward mode: {forward_mode=}") + + # If topk > 1, we need to use retrieve_next_token and retrieve_next_sibling to handle the eagle tree custom attention mask + if forward_mode.is_target_verify() and spec_info.topk > 1: + # They are None during cuda graph capture so skip the copy_... + # self.retrieve_next_token_list[bs - 1].copy_(spec_info.retrive_next_token) + # self.retrieve_next_sibling_list[bs - 1].copy_(spec_info.retrive_next_sibling) + return ForwardMetadata( + query_start_loc=self.query_start_loc_list[bs - 1], + mamba_cache_indices=self.state_indices_list[bs - 1], + retrieve_next_token=self.retrieve_next_token_list[bs - 1], + retrieve_next_sibling=self.retrieve_next_sibling_list[bs - 1], + retrieve_parent_token=self.retrieve_parent_token_list[bs - 1], + ) + else: + return ForwardMetadata( + query_start_loc=self.query_start_loc_list[bs - 1], + mamba_cache_indices=self.state_indices_list[bs - 1], + mamba_cache_indices_gdn=self.state_indices_list_gdn[bs - 1], + ) + + def _replay_metadata( + self, + bs: int, + req_pool_indices: torch.Tensor, + forward_mode: ForwardMode, + spec_info: Optional[SpecInput], + seq_lens_cpu: Optional[torch.Tensor], + ): + num_padding = torch.count_nonzero( + seq_lens_cpu == self.get_cuda_graph_seq_len_fill_value() + ) + # Make sure forward metadata is correctly handled for padding reqs + req_pool_indices[bs - num_padding :] = 0 + mamba_indices = self.req_to_token_pool.get_mamba_indices(req_pool_indices) + mamba_indices[bs - num_padding :] = 0 + self.state_indices_list[bs - 1][: len(mamba_indices)].copy_(mamba_indices) + if forward_mode.is_decode_or_idle(): + if num_padding == 0: + self.query_start_loc_list[bs - 1].copy_( + self.cached_cuda_graph_decode_query_start_loc[: bs + 1] + ) + else: + self.query_start_loc_list[bs - 1][: bs - num_padding].copy_( + self.cached_cuda_graph_decode_query_start_loc[: bs - num_padding] + ) + self.query_start_loc_list[bs - 1][bs - num_padding :].fill_( + bs - num_padding + ) + elif forward_mode.is_target_verify(): + ssm_state_indices = torch.arange( + len(mamba_indices[: bs - num_padding]) * spec_info.draft_token_num, + dtype=torch.int32, + device=mamba_indices.device, + ) + self.state_indices_list_gdn[bs - 1][ + : len(mamba_indices[: bs - num_padding]) * spec_info.draft_token_num + ].copy_(ssm_state_indices) + self.state_indices_list_gdn[bs - 1][ + len(mamba_indices[: bs - num_padding]) * spec_info.draft_token_num : + ] = 0 + if num_padding == 0: + self.query_start_loc_list[bs - 1].copy_( + self.cached_cuda_graph_verify_query_start_loc[: bs + 1] + ) + else: + self.query_start_loc_list[bs - 1][: bs - num_padding].copy_( + self.cached_cuda_graph_verify_query_start_loc[: bs - num_padding] + ) + self.query_start_loc_list[bs - 1][bs - num_padding :].fill_( + (bs - num_padding) * spec_info.draft_token_num + ) + else: + raise ValueError(f"Invalid forward mode: {forward_mode=}") + + # If topk > 1, we need to use retrieve_next_token and retrieve_next_sibling to handle the eagle tree custom attention mask + if forward_mode.is_target_verify() and spec_info.topk > 1: + bs_without_pad = spec_info.retrive_next_token.shape[0] + self.retrieve_next_token_list[bs - 1][:bs_without_pad].copy_( + spec_info.retrive_next_token + ) + self.retrieve_next_sibling_list[bs - 1][:bs_without_pad].copy_( + spec_info.retrive_next_sibling + ) + return ForwardMetadata( + query_start_loc=self.query_start_loc_list[bs - 1], + mamba_cache_indices=self.state_indices_list[bs - 1], + retrieve_next_token=self.retrieve_next_token_list[bs - 1], + retrieve_next_sibling=self.retrieve_next_sibling_list[bs - 1], + retrieve_parent_token=self.retrieve_parent_token_list[bs - 1], + ) + else: + return ForwardMetadata( + query_start_loc=self.query_start_loc_list[bs - 1], + mamba_cache_indices=self.state_indices_list[bs - 1], + mamba_cache_indices_gdn=self.state_indices_list_gdn[bs - 1], + ) + + def get_cuda_graph_seq_len_fill_value(self): + return 0 # Mamba attn does not use seq lens to index kv cache + + +class AscendMamba2AttnBackend(AscendMambaAttnBackendBase): + pass + + +class AscendHybridLinearAttnBackend(HybridLinearAttnBackend): + def __init__( + self, + full_attn_backend: AttentionBackend, + linear_attn_backend: AscendMambaAttnBackendBase, + full_attn_layers: list[int], + ): + super().__init__(full_attn_backend, linear_attn_backend, full_attn_layers) + + def update_mamba_state_after_mtp_verify( + self, + accepted_steps: torch.Tensor, + mamba_track_indices: Optional[torch.Tensor], + mamba_steps_to_track: Optional[torch.Tensor], + model, + ): + """ + Update mamba states after MTP verify using fully fused Triton kernel. + + This replaces the original advanced indexing operations with a single fused + gather-scatter kernel that also handles masking internally, avoiding: + - index_elementwise_kernel from tensor[bool_mask] + - index_select kernel launches + - nonzero kernel launches + """ + request_number = accepted_steps.shape[0] + + state_indices_tensor = ( + self.linear_attn_backend.forward_metadata.mamba_cache_indices[ + :request_number + ] + ) + + mamba_caches = ( + self.linear_attn_backend.req_to_token_pool.get_speculative_mamba2_params_all_layers() + ) + + conv_states = mamba_caches.conv[0] + ssm_states = mamba_caches.temporal + intermediate_state_cache = mamba_caches.intermediate_ssm + dst_indices_tensor = state_indices_tensor.to(torch.int64) # [N] + src_indices_tensor = torch.arange( + dst_indices_tensor.shape[0], + device=dst_indices_tensor.device, + dtype=torch.int64, + ) + last_steps = accepted_steps.to(torch.int64) # [N] + + move_intermediate_cache( + ssm_states, + intermediate_state_cache, + dst_indices_tensor, + src_indices_tensor, + last_steps, + ) + + draft_token_num = intermediate_state_cache.shape[2] + if dst_indices_tensor.numel() > 0: + conv_state_rollback( + conv_states, + dst_indices_tensor, + last_steps, + draft_token_num, + ) + return + + def update_verify_buffers_to_fill_after_draft( + self, spec_info: SpecInput, cuda_graph_bs: Optional[int] + ): + pass diff --git a/python/sglang/srt/hardware_backend/npu/memory_pool_npu.py b/python/sglang/srt/hardware_backend/npu/memory_pool_npu.py index ea81f4e58..e4f319fa5 100644 --- a/python/sglang/srt/hardware_backend/npu/memory_pool_npu.py +++ b/python/sglang/srt/hardware_backend/npu/memory_pool_npu.py @@ -15,6 +15,30 @@ if TYPE_CHECKING: from sglang.srt.layers.radix_attention import RadixAttention +def _init_npu_conv_state( + conv_state_in, conv_state_shape, speculative_num_draft_tokens: Optional[int] = None +): + extra_conv_len = 0 + if speculative_num_draft_tokens is not None: + extra_conv_len = speculative_num_draft_tokens - 1 + + # conv_state shape (layers, pool_size, conv_wind + draft_step, dim) for conv1d ascendc ops require dim as last dim + conv_state = [ + torch.zeros( + size=( + conv_state_in.shape[0], + conv_state_in.shape[1], + conv_shape[1] + extra_conv_len, + conv_shape[0], + ), + dtype=conv_state_in.dtype, + device=conv_state_in.device, + ) + for conv_shape in conv_state_shape + ] + return conv_state + + class NPUMHATokenToKVPool(MHATokenToKVPool): def __init__( diff --git a/python/sglang/srt/layers/attention/attention_registry.py b/python/sglang/srt/layers/attention/attention_registry.py index b9b860409..e870d8f77 100644 --- a/python/sglang/srt/layers/attention/attention_registry.py +++ b/python/sglang/srt/layers/attention/attention_registry.py @@ -200,11 +200,6 @@ def attn_backend_wrapper(runner: "ModelRunner", full_attn_backend: "AttentionBac if cfg := runner.mambaish_config: from sglang.srt.layers.attention.fla.utils import check_environments - from sglang.srt.layers.attention.hybrid_linear_attn_backend import ( - HybridLinearAttnBackend, - Mamba2AttnBackend, - ) - from sglang.srt.layers.attention.linear.gdn_backend import GDNAttnBackend from sglang.srt.layers.attention.linear.kda_backend import KDAAttnBackend from sglang.srt.layers.attention.linear.lightning_backend import ( LightningAttentionBackend, @@ -214,6 +209,23 @@ def attn_backend_wrapper(runner: "ModelRunner", full_attn_backend: "AttentionBac ) from sglang.srt.utils import is_blackwell, is_npu + if not is_npu(): + from sglang.srt.layers.attention.hybrid_linear_attn_backend import ( + HybridLinearAttnBackend, + Mamba2AttnBackend, + ) + from sglang.srt.layers.attention.linear.gdn_backend import GDNAttnBackend + else: + from sglang.srt.hardware_backend.npu.attention.ascend_gdn_backend import ( + AscendGDNAttnBackend as GDNAttnBackend, + ) + from sglang.srt.hardware_backend.npu.attention.ascend_hybrid_linear_attn_backend import ( + AscendHybridLinearAttnBackend as HybridLinearAttnBackend, + ) + from sglang.srt.hardware_backend.npu.attention.ascend_hybrid_linear_attn_backend import ( + AscendMamba2AttnBackend as Mamba2AttnBackend, + ) + check_environments() initialize_linear_attn_config(runner.server_args) if runner.hybrid_gdn_config is not None: diff --git a/python/sglang/srt/layers/attention/mamba/mamba2_metadata.py b/python/sglang/srt/layers/attention/mamba/mamba2_metadata.py index 3c1548e9b..35d8abaa8 100644 --- a/python/sglang/srt/layers/attention/mamba/mamba2_metadata.py +++ b/python/sglang/srt/layers/attention/mamba/mamba2_metadata.py @@ -27,6 +27,7 @@ from sglang.srt.model_executor.forward_batch_info import ForwardBatch class ForwardMetadata: query_start_loc: torch.Tensor mamba_cache_indices: torch.Tensor + mamba_cache_indices_gdn: Optional[torch.Tensor] = None # For topk > 1 eagle retrieve_next_token: Optional[torch.Tensor] = None retrieve_next_sibling: Optional[torch.Tensor] = None diff --git a/python/sglang/srt/layers/layernorm.py b/python/sglang/srt/layers/layernorm.py index 8c4f92969..a92469a98 100644 --- a/python/sglang/srt/layers/layernorm.py +++ b/python/sglang/srt/layers/layernorm.py @@ -110,6 +110,7 @@ logger = logging.getLogger(__name__) if _is_npu: import torch_npu + from sgl_kernel_npu.norm.add_rmsnorm_bias import add_gemma_rms_norm def _forward_with_allreduce_fusion( @@ -672,11 +673,13 @@ class GemmaRMSNorm(MultiPlatformOp): if residual is not None: if post_residual_addition is not None: residual = residual + post_residual_addition - x = x + residual - residual = x + norm_out, residual = add_gemma_rms_norm( + x, self.weight, residual, self.variance_epsilon + ) + return norm_out, residual x, _ = torch_npu.npu_gemma_rms_norm(x, self.weight, self.variance_epsilon) - return x if residual is None else (x, residual) + return x def forward_xpu( self, diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index df4ca52a5..e7dbb03dd 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py @@ -260,6 +260,15 @@ class MambaPool: for conv_shape in conv_state_shape ] + if _is_npu: + from sglang.srt.hardware_backend.npu.memory_pool_npu import ( + _init_npu_conv_state, + ) + + conv_state = _init_npu_conv_state( + conv_state[0], conv_state_shape, speculative_num_draft_tokens + ) + if _is_cpu and _cpu_has_amx_support: from sglang.srt.layers.amx_utils import _init_amx_conv_state diff --git a/python/sglang/srt/models/qwen3_5_mtp.py b/python/sglang/srt/models/qwen3_5_mtp.py index 3fa89fcda..c60737c56 100644 --- a/python/sglang/srt/models/qwen3_5_mtp.py +++ b/python/sglang/srt/models/qwen3_5_mtp.py @@ -15,6 +15,7 @@ """Inference-only Qwen3_5 MTP model.""" import logging +from contextlib import ExitStack from typing import Iterable, Optional, Tuple import torch @@ -22,6 +23,7 @@ from torch import nn from transformers import PretrainedConfig from sglang.srt.distributed import get_pp_group, get_tensor_model_parallel_world_size +from sglang.srt.environ import envs from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder from sglang.srt.eplb.expert_location import ModelConfigForExpertLocation from sglang.srt.layers.layernorm import GemmaRMSNorm @@ -31,7 +33,8 @@ from sglang.srt.layers.vocab_parallel_embedding import ParallelLMHead from sglang.srt.model_executor.forward_batch_info import ForwardBatch from sglang.srt.model_loader.weight_utils import default_weight_loader from sglang.srt.models.qwen3_5 import Qwen3_5ForCausalLM -from sglang.srt.utils import add_prefix +from sglang.srt.server_args import get_global_server_args +from sglang.srt.utils import add_prefix, is_npu logger = logging.getLogger(__name__) @@ -53,6 +56,11 @@ class Qwen3_5ForCausalLMMTP(nn.Module): # The MTP model is unquantized in the nvfp4 checkpoint. if quant_config and quant_config.get_name() == "modelopt_fp4": quant_config = None + if ( + is_npu() + and get_global_server_args().speculative_draft_model_quantization is None + ): + quant_config = None self.config = config self.tp_size = get_tensor_model_parallel_world_size() @@ -118,6 +126,18 @@ class Qwen3_5ForCausalLMMTP(nn.Module): input_embeds: Optional[torch.Tensor] = None, **kwargs, ): + exit_stack = ExitStack() + if ( + is_npu() + and self.quant_config is None + and get_global_server_args().quantization is not None + ): + # ascend mtp unquant + exit_stack.enter_context(envs.SGLANG_DEEPEP_BF16_DISPATCH.override(True)) + exit_stack.enter_context( + envs.DEEP_NORMAL_MODE_USE_INT8_QUANT.override(False) + ) + assert input_embeds is None input_embeds = forward_batch.mm_input_embeds if ( @@ -150,6 +170,8 @@ class Qwen3_5ForCausalLMMTP(nn.Module): hidden_states, ) + exit_stack.close() + return self.logits_processor( input_ids, hidden_states, self.lm_head, forward_batch ) diff --git a/python/sglang/srt/models/qwen3_next_mtp.py b/python/sglang/srt/models/qwen3_next_mtp.py index b2bdbbbe8..5f0dcb5e0 100644 --- a/python/sglang/srt/models/qwen3_next_mtp.py +++ b/python/sglang/srt/models/qwen3_next_mtp.py @@ -15,6 +15,7 @@ """Inference-only Qwen3Next MTP Speculative Decoding.""" import logging +from contextlib import ExitStack from typing import Iterable, Optional, Tuple import torch @@ -22,6 +23,7 @@ from torch import nn from transformers import PretrainedConfig from sglang.srt.distributed import get_pp_group, get_tensor_model_parallel_world_size +from sglang.srt.environ import envs from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder from sglang.srt.layers.layernorm import GemmaRMSNorm from sglang.srt.layers.logits_processor import LogitsProcessor @@ -30,7 +32,7 @@ from sglang.srt.layers.vocab_parallel_embedding import ParallelLMHead from sglang.srt.model_executor.forward_batch_info import ForwardBatch from sglang.srt.models.qwen3_next import Qwen3NextForCausalLM, Qwen3NextModel from sglang.srt.server_args import get_global_server_args -from sglang.srt.utils import add_prefix +from sglang.srt.utils import add_prefix, is_npu logger = logging.getLogger(__name__) @@ -46,6 +48,11 @@ class Qwen3NextForCausalLMMTP(Qwen3NextForCausalLM): nn.Module.__init__(self) self.config = config self.tp_size = get_tensor_model_parallel_world_size() + if ( + is_npu() + and get_global_server_args().speculative_draft_model_quantization is None + ): + quant_config = None self.quant_config = quant_config # if not set, model load will be broken in Qwen3NextForCausalLM load_weights() self.pp_group = get_pp_group() @@ -86,6 +93,18 @@ class Qwen3NextForCausalLMMTP(Qwen3NextForCausalLM): input_embeds: Optional[torch.Tensor] = None, **kwargs, ): + exit_stack = ExitStack() + if ( + is_npu() + and self.quant_config is None + and get_global_server_args().quantization is not None + ): + # ascend mtp unquant + exit_stack.enter_context(envs.SGLANG_DEEPEP_BF16_DISPATCH.override(True)) + exit_stack.enter_context( + envs.DEEP_NORMAL_MODE_USE_INT8_QUANT.override(False) + ) + if input_embeds is None: input_embeds = self.model.embed_tokens(input_ids) @@ -104,6 +123,8 @@ class Qwen3NextForCausalLMMTP(Qwen3NextForCausalLM): hidden_states, ) + exit_stack.close() + return self.logits_processor( input_ids, hidden_states, self.lm_head, forward_batch )