From 5b5445dbd3059c5133c840ddfafb90e66af7fecf Mon Sep 17 00:00:00 2001 From: iridiumine <42236072+iridiumine@users.noreply.github.com> Date: Wed, 18 Mar 2026 16:09:23 +0800 Subject: [PATCH 1/7] feat: support Ascend NPU MTP adaptation for GDN attention backend --- .../npu/attention/ascend_gdn_backend.py | 352 ++++++++++++++++++ .../hardware_backend/npu/memory_pool_npu.py | 15 + .../layers/attention/attention_registry.py | 6 +- .../attention/hybrid_linear_attn_backend.py | 57 ++- .../layers/attention/mamba/mamba2_metadata.py | 1 + python/sglang/srt/layers/layernorm.py | 7 +- python/sglang/srt/mem_cache/memory_pool.py | 4 + python/sglang/srt/models/qwen3_5_mtp.py | 15 +- python/sglang/srt/models/qwen3_next_mtp.py | 12 + 9 files changed, 459 insertions(+), 10 deletions(-) create mode 100644 python/sglang/srt/hardware_backend/npu/attention/ascend_gdn_backend.py 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 000000000000..57b380a7e167 --- /dev/null +++ b/python/sglang/srt/hardware_backend/npu/attention/ascend_gdn_backend.py @@ -0,0 +1,352 @@ +from typing import Tuple, Union, Optional +import torch +import torch.nn.functional as F +import torch_npu + +from sglang.srt.layers.attention.fla.fused_gdn_gating import fused_gdn_gating +from sglang.srt.layers.attention.fla.l2norm import l2norm_fwd +from sglang.srt.layers.attention.hybrid_linear_attn_backend import MambaAttnBackendBase +from sglang.srt.layers.attention.linear.gdn_backend import ( + GDNKernelDispatcher, + GDNAttnBackend, +) +from sglang.srt.layers.attention.linear.kernels.gdn_triton import TritonGDNKernel +from sglang.srt.layers.attention.mamba.causal_conv1d_triton import ( + causal_conv1d_fn, + causal_conv1d_update, +) +from sglang.srt.layers.radix_linear_attention import RadixLinearAttention +from sglang.srt.mem_cache.memory_pool import MambaPool +from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode +from sglang.srt.model_executor.model_runner import ModelRunner +from sglang.srt.speculative.eagle_info import EagleDraftInput, EagleVerifyInput +from sglang.srt.utils import is_cpu, is_cuda, is_npu +from sglang.srt.utils.common import rank0_log + +def npu_causal_conv1d_update( + hidden_state: torch.Tensor, + weight: torch.Tensor, + conv_state: torch.Tensor, + bias: Optional[torch.Tensor] = None, + silu_activation: bool = True, +): + bsz, hidden_size, seq_len = hidden_state.shape + kernel_size = weight.shape[-1] + target_state_len = (kernel_size - 1) + (seq_len - 1) + full_context = torch.cat([conv_state, hidden_state], dim=-1).to(weight.dtype) + computation_input = full_context[:, :, -(kernel_size - 1 + seq_len):] + windows = computation_input.unfold(-1, kernel_size, 1) + out = (windows * weight[None, :, None, :]).sum(dim=-1) + + if bias is not None: + out = out + bias[None, :, None] + if silu_activation: + out = F.silu(out) + out = out.to(hidden_state.dtype) + + if target_state_len > 0: + new_conv_state = full_context[:, :, -target_state_len:] + else: + new_conv_state = torch.empty(bsz, hidden_size, 0, device=hidden_state.device, dtype=hidden_state.dtype) + return out, new_conv_state + +class AscendGDNKernelDispatcher(GDNKernelDispatcher): + pass + +class AscendGDNAttnBackend(GDNAttnBackend): + + def __init__(self, model_runner: ModelRunner): + super().__init__(model_runner) + # self.graph_mode = False + + 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 + start_indices = cache_indices * seq_len + offset = torch.arange(seq_len, device=start_indices.device) + ranges = start_indices.unsqueeze(1) + offset + self.ssm_state_indices = ranges.flatten().to(torch.int32) + self.forward_metadata.ssm_state_indices = self.ssm_state_indices + else: + self.ssm_state_indices = cache_indices + self.forward_metadata.ssm_state_indices = self.ssm_state_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 + intermediate_state_cache = None + + if not forward_batch.spec_algorithm.is_none(): + intermediate_state_cache = mamba_cache_params.intermediate_ssm + if is_target_verify: + assert isinstance(mamba_cache_params, MambaPool.SpeculativeState) + intermediate_conv_window_cache = mamba_cache_params.intermediate_conv_window[0] + has_initial_states = torch.ones( + seq_len // forward_batch.spec_info.draft_token_num, + dtype=torch.bool, device=forward_batch.input_ids.device + ) + intermediate_state_indices = torch.arange( + cache_indices.shape[0], dtype=torch.int32, device=cache_indices.device + ) + else: + has_initial_states = forward_batch.extend_prefix_lens > 0 + + if is_target_verify: + batch_size = seq_len // forward_batch.spec_info.draft_token_num + draft_token_num = forward_batch.spec_info.draft_token_num + 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) + conv_states_to_use = conv_states[cache_indices] + mixed_qkv_processed, new_conv_state = npu_causal_conv1d_update( + mixed_qkv_reshaped.transpose(1, 2).contiguous(), + layer.conv_weights, + conv_states_to_use.transpose(1, 2).contiguous(), + layer.bias, True + ) + mixed_qkv = mixed_qkv_processed.transpose(1, 2).contiguous().view(seq_len, -1) + conv_states[cache_indices] = new_conv_state.transpose(1, 2).contiguous() + else: + mixed_qkv = mixed_qkv.transpose(0, 1) + if forward_batch.mamba_track_mask is not None and forward_batch.mamba_track_mask.any(): + conv_dst = forward_batch.mamba_track_indices + mixed_qkv_to_track = mixed_qkv[:, forward_metadata.track_conv_indices].transpose(0, 1) + mask_indices = forward_batch.mamba_track_mask.nonzero(as_tuple=True)[0] + conv_states[conv_dst[mask_indices]] = mixed_qkv_to_track + + 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() + + query, key, value = torch.split(mixed_qkv, [layer.q_dim, layer.k_dim, layer.v_dim], dim=-1) + actual_seq_len = query.shape[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 + + query = query.view(1, actual_seq_len, num_heads, head_k_dim) + key = key.view(1, actual_seq_len, num_heads, head_k_dim) + value = value.view(1, actual_seq_len, num_value_heads, head_v_dim) + g, beta = fused_gdn_gating(layer.A_log, a, b, layer.dt_bias) + + if is_target_verify: + query = query.view(-1, num_heads, head_k_dim) + key = key.view(-1, num_heads, head_k_dim) + value = value.view(-1, num_value_heads, head_v_dim) + query = l2norm_fwd(query.contiguous(), eps=1e-6, output_dtype=torch.bfloat16) + key = l2norm_fwd(key.contiguous(), eps=1e-6, output_dtype=torch.bfloat16) + + core_attn_out = self.fused_recurrent_gated_delta_rule_update_npu( + query, key, value, 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: + 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 not forward_batch.spec_algorithm.is_none() and last_recurrent_state is not None: + last_recurrent_state = last_recurrent_state.transpose(-1, -2).to(ssm_states.dtype, copy=False) + intermediate_state_cache[cache_indices, 0] = last_recurrent_state + elif 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 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_npu( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + recurrent_state: torch.Tensor, + beta: torch.Tensor, + g: torch.Tensor, + cache_indices: torch.Tensor, + intermediate_state: Optional[torch.Tensor] = None, + ): + _, num_heads, head_k_dim = query.shape + _, num_value_heads, head_v_dim = value.shape + beta = beta.view(-1, num_value_heads).to(torch.bfloat16) + g = g.view(-1, num_value_heads).to(torch.float32) + batch_size = cache_indices.shape[0] + seq_len = query.shape[0] // batch_size + scale = 1 / (head_k_dim ** 0.5) + + if intermediate_state is not None: + ssm_state = intermediate_state.view(-1, num_value_heads, head_k_dim, head_v_dim) + else: + ssm_state = recurrent_state[cache_indices] + + if self.graph_mode: + num_accepted_tokens = torch.ones([batch_size], dtype=torch.int32, device=cache_indices.device) + actual_seq_lengths = torch.ones([batch_size], dtype=torch.int32, device=cache_indices.device) * seq_len + 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.forward_metadata.mamba_cache_indices + + attn_core_out = torch_npu.npu_recurrent_gated_delta_rule( + query, key, value, ssm_state, beta=beta, scale=scale, + actual_seq_lengths=actual_seq_lengths, ssm_state_indices=ssm_state_indices, + num_accepted_tokens=num_accepted_tokens, g=g, + ) + + if intermediate_state is not None: + intermediate_state = ssm_state.view(-1, seq_len, num_value_heads, head_k_dim, head_v_dim) + return attn_core_out \ No newline at end of file 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 82e7ce693e33..db36a2c3f77c 100644 --- a/python/sglang/srt/hardware_backend/npu/memory_pool_npu.py +++ b/python/sglang/srt/hardware_backend/npu/memory_pool_npu.py @@ -14,6 +14,21 @@ 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): diff --git a/python/sglang/srt/layers/attention/attention_registry.py b/python/sglang/srt/layers/attention/attention_registry.py index 9020e5260af9..bf13ea2e149a 100644 --- a/python/sglang/srt/layers/attention/attention_registry.py +++ b/python/sglang/srt/layers/attention/attention_registry.py @@ -192,7 +192,6 @@ def attn_backend_wrapper(runner: "ModelRunner", full_attn_backend: "AttentionBac 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, @@ -202,6 +201,11 @@ def attn_backend_wrapper(runner: "ModelRunner", full_attn_backend: "AttentionBac ) from sglang.srt.utils import is_blackwell, is_npu + if is_npu(): + from sglang.srt.hardware_backend.npu.attention.ascend_gdn_backend import AscendGDNAttnBackend as GDNAttnBackend + else: + from sglang.srt.layers.attention.linear.gdn_backend import GDNAttnBackend + 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/hybrid_linear_attn_backend.py b/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py index 91194c494396..6ac003a0bce2 100644 --- a/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py +++ b/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py @@ -22,13 +22,16 @@ from sglang.srt.server_args import get_global_server_args from sglang.srt.speculative.eagle_info import EagleDraftInput, EagleVerifyInput from sglang.srt.speculative.spec_info import SpecInput -from sglang.srt.utils import is_cpu +from sglang.srt.utils import is_cpu, is_npu if not is_cpu(): from sglang.srt.layers.attention.fla.chunk_delta_h import ( CHUNK_SIZE as FLA_CHUNK_SIZE, ) +if is_npu(): + from sglang.srt.layers.attention.mamba.mamba_state_scatter_triton import conv_state_rollback, move_intermediate_cache_dynamic_h_block_v1 + logger = logging.getLogger(__name__) @@ -142,6 +145,7 @@ def __init__(self, model_runner: ModelRunner): self.req_to_token_pool: HybridReqToTokenPool = model_runner.req_to_token_pool self.forward_metadata: ForwardMetadata = None self.state_indices_list = [] + self.state_indices_list_gdn = [] self.query_start_loc_list = [] self.retrieve_next_token_list = [] self.retrieve_next_sibling_list = [] @@ -409,6 +413,11 @@ def init_cuda_graph_state(self, max_bs: int, max_num_tokens: int): (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) ) @@ -462,6 +471,8 @@ def _capture_metadata( 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] @@ -470,10 +481,14 @@ def _capture_metadata( self.query_start_loc_list[bs - 1].copy_( self.cached_cuda_graph_verify_query_start_loc[: bs + 1] ) + start_indices = mamba_indices * spec_info.draft_token_num + offset = torch.arange(spec_info.draft_token_num, device=start_indices.device) + ranges = start_indices.unsqueeze(1) + offset + ssm_state_indices = ranges.flatten().to(torch.int32) + 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=}") - 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 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: @@ -491,6 +506,7 @@ def _capture_metadata( 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( @@ -507,7 +523,7 @@ def _replay_metadata( # 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 :] = -1 + 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: @@ -522,6 +538,13 @@ def _replay_metadata( bs - num_padding ) elif forward_mode.is_target_verify(): + start_indices = mamba_indices[:bs - num_padding] * spec_info.draft_token_num + offset = torch.arange(spec_info.draft_token_num, device=start_indices.device) + ranges = start_indices.unsqueeze(1) + offset + ssm_state_indices = ranges.flatten().to(torch.int32) + 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] @@ -556,10 +579,11 @@ def _replay_metadata( 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 1 # Mamba attn does not use seq lens to index kv cache + return 0 # Mamba attn does not use seq lens to index kv cache def get_cpu_graph_seq_len_fill_value(self): return 1 @@ -960,6 +984,21 @@ def update_mamba_state_after_mtp_verify( ssm_states = mamba_caches.temporal intermediate_state_cache = mamba_caches.intermediate_ssm intermediate_conv_window_cache = mamba_caches.intermediate_conv_window[0] + if is_npu(): + valid_state_indices = state_indices_tensor.to(torch.int64) # [N] + last_steps = accepted_steps.to(torch.int64) # [N] + + move_intermediate_cache_dynamic_h_block_v1(intermediate_state_cache, valid_state_indices, last_steps) + + draft_token_num = intermediate_state_cache.shape[2] + if valid_state_indices.numel() > 0: + conv_state_rollback( + conv_states, + valid_state_indices, + last_steps, + draft_token_num, + ) + return # Use fully fused kernel that handles masking internally # This avoids separate nonzero() and index_select() calls @@ -992,3 +1031,11 @@ def update_mamba_state_after_mtp_verify( mamba_track_indices, mamba_steps_to_track, ) + + def get_verify_buffers_to_fill_after_draft(self): + return [None, None] + + def update_verify_buffers_to_fill_after_draft( + self, spec_info: SpecInput, cuda_graph_bs: Optional[int] + ): + pass diff --git a/python/sglang/srt/layers/attention/mamba/mamba2_metadata.py b/python/sglang/srt/layers/attention/mamba/mamba2_metadata.py index 5eeb2b65e307..5d34e1e111d3 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 @@ 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 614be8f0ec61..50c21ba41f97 100644 --- a/python/sglang/srt/layers/layernorm.py +++ b/python/sglang/srt/layers/layernorm.py @@ -84,6 +84,7 @@ if _is_npu: import torch_npu + from sgl_kernel_npu.norm.add_rmsnorm_bias import add_gemma_rms_norm class RMSNorm(MultiPlatformOp): @@ -512,11 +513,11 @@ def forward_npu( 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 40fe210df2d8..481249dc374d 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py @@ -258,6 +258,10 @@ def __init__( 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 3fa89fcda0a9..1266b8096aac 100644 --- a/python/sglang/srt/models/qwen3_5_mtp.py +++ b/python/sglang/srt/models/qwen3_5_mtp.py @@ -17,10 +17,12 @@ import logging from typing import Iterable, Optional, Tuple +import os import torch from torch import nn from transformers import PretrainedConfig +from sglang.srt.server_args import get_global_server_args from sglang.srt.distributed import get_pp_group, get_tensor_model_parallel_world_size from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder from sglang.srt.eplb.expert_location import ModelConfigForExpertLocation @@ -31,7 +33,7 @@ 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.utils import add_prefix, is_npu logger = logging.getLogger(__name__) @@ -53,6 +55,9 @@ def __init__( # The MTP model is unquantized in the nvfp4 checkpoint. if quant_config and quant_config.get_name() == "modelopt_fp4": quant_config = None + if get_global_server_args().speculative_draft_model_quantization is None: + quant_config = None + self.quant_config = quant_config self.config = config self.tp_size = get_tensor_model_parallel_world_size() @@ -118,6 +123,10 @@ def forward( input_embeds: Optional[torch.Tensor] = None, **kwargs, ): + if is_npu() and self.quant_config is None: + # ascend mtp unquant + os.environ["SGLANG_DEEPEP_BF16_DISPATCH"] = "1" + os.environ["DEEP_NORMAL_MODE_USE_INT8_QUANT"] = "0" assert input_embeds is None input_embeds = forward_batch.mm_input_embeds if ( @@ -149,6 +158,10 @@ def forward( forward_batch, hidden_states, ) + if is_npu() and self.quant_config is None: + # ascend mtp unquant + os.environ["SGLANG_DEEPEP_BF16_DISPATCH"] = "0" + os.environ["DEEP_NORMAL_MODE_USE_INT8_QUANT"] = "1" 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 b2bdbbbe8705..9bbdb9d8afc6 100644 --- a/python/sglang/srt/models/qwen3_next_mtp.py +++ b/python/sglang/srt/models/qwen3_next_mtp.py @@ -17,10 +17,12 @@ import logging from typing import Iterable, Optional, Tuple +import os import torch from torch import nn from transformers import PretrainedConfig +from sglang.srt.hardware_backend.npu.graph_runner.npu_graph_runner import is_npu from sglang.srt.distributed import get_pp_group, get_tensor_model_parallel_world_size from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder from sglang.srt.layers.layernorm import GemmaRMSNorm @@ -46,6 +48,8 @@ def __init__( nn.Module.__init__(self) self.config = config self.tp_size = get_tensor_model_parallel_world_size() + if 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 +90,10 @@ def forward( input_embeds: Optional[torch.Tensor] = None, **kwargs, ): + if is_npu() and self.quant_config is None: + # ascend mtp unquant + os.environ["SGLANG_DEEPEP_BF16_DISPATCH"] = "1" + os.environ["DEEP_NORMAL_MODE_USE_INT8_QUANT"] = "0" if input_embeds is None: input_embeds = self.model.embed_tokens(input_ids) @@ -103,6 +111,10 @@ def forward( forward_batch, hidden_states, ) + if is_npu() and self.quant_config is None: + # ascend mtp unquant + os.environ["SGLANG_DEEPEP_BF16_DISPATCH"] = "0" + os.environ["DEEP_NORMAL_MODE_USE_INT8_QUANT"] = "1" return self.logits_processor( input_ids, hidden_states, self.lm_head, forward_batch From e11d7d283418296236b6474f706f7cd4ef7ebcfd Mon Sep 17 00:00:00 2001 From: iridiumine <42236072+iridiumine@users.noreply.github.com> Date: Thu, 19 Mar 2026 15:37:05 +0800 Subject: [PATCH 2/7] fix: optimize and fix bugs for Ascend NPU MTP GDN backend --- .../npu/attention/ascend_gdn_backend.py | 266 ++++++++++++------ 1 file changed, 181 insertions(+), 85 deletions(-) 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 index 57b380a7e167..9cce5c5c4a82 100644 --- a/python/sglang/srt/hardware_backend/npu/attention/ascend_gdn_backend.py +++ b/python/sglang/srt/hardware_backend/npu/attention/ascend_gdn_backend.py @@ -3,52 +3,59 @@ import torch.nn.functional as F import torch_npu -from sglang.srt.layers.attention.fla.fused_gdn_gating import fused_gdn_gating from sglang.srt.layers.attention.fla.l2norm import l2norm_fwd -from sglang.srt.layers.attention.hybrid_linear_attn_backend import MambaAttnBackendBase from sglang.srt.layers.attention.linear.gdn_backend import ( GDNKernelDispatcher, GDNAttnBackend, ) -from sglang.srt.layers.attention.linear.kernels.gdn_triton import TritonGDNKernel -from sglang.srt.layers.attention.mamba.causal_conv1d_triton import ( - causal_conv1d_fn, - causal_conv1d_update, +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 -from sglang.srt.utils import is_cpu, is_cuda, is_npu -from sglang.srt.utils.common import rank0_log +from sglang.srt.utils import is_cpu + +from sgl_kernel_npu.fla.fused_gdn_gating import fused_gdn_gating_npu +from sgl_kernel_npu.mamba.causal_conv1d import ( + causal_conv1d_fn_npu, + causal_conv1d_update_npu, +) + +fused_gdn_gating = fused_gdn_gating_npu +causal_conv1d_fn = causal_conv1d_fn_npu +causal_conv1d_update = causal_conv1d_update_npu def npu_causal_conv1d_update( - hidden_state: torch.Tensor, - weight: torch.Tensor, - conv_state: torch.Tensor, - bias: Optional[torch.Tensor] = None, - silu_activation: bool = True, -): - bsz, hidden_size, seq_len = hidden_state.shape - kernel_size = weight.shape[-1] - target_state_len = (kernel_size - 1) + (seq_len - 1) - full_context = torch.cat([conv_state, hidden_state], dim=-1).to(weight.dtype) - computation_input = full_context[:, :, -(kernel_size - 1 + seq_len):] - windows = computation_input.unfold(-1, kernel_size, 1) - out = (windows * weight[None, :, None, :]).sum(dim=-1) - - if bias is not None: - out = out + bias[None, :, None] - if silu_activation: - out = F.silu(out) - out = out.to(hidden_state.dtype) - - if target_state_len > 0: - new_conv_state = full_context[:, :, -target_state_len:] - else: - new_conv_state = torch.empty(bsz, hidden_size, 0, device=hidden_state.device, dtype=hidden_state.dtype) - return out, new_conv_state + hidden_state: torch.Tensor, + weight: torch.Tensor, + conv_state: torch.Tensor, + bias: Optional[torch.Tensor] = None, + silu_activation: bool = True, + ): + bsz, hidden_size, seq_len = hidden_state.shape + kernel_size = weight.shape[-1] + target_state_len = (kernel_size - 1) + (seq_len - 1) + full_context = torch.cat([conv_state, hidden_state], dim=-1).to(weight.dtype) + computation_input = full_context[:, :, -(kernel_size - 1 + seq_len):] + windows = computation_input.unfold(-1, kernel_size, 1) + + out = (windows * weight[None, :, None, :]).sum(dim=-1) + if bias is not None: + out = out + bias[None, :, None] + if silu_activation: + out = F.silu(out) + out = out.to(hidden_state.dtype) + + if target_state_len > 0: + new_conv_state = full_context[:, :, -target_state_len:] + else: + new_conv_state = torch.empty(bsz, hidden_size, 0, device=hidden_state.device, dtype=hidden_state.dtype) + + return out, new_conv_state class AscendGDNKernelDispatcher(GDNKernelDispatcher): pass @@ -57,7 +64,13 @@ class AscendGDNAttnBackend(GDNAttnBackend): def __init__(self, model_runner: ModelRunner): super().__init__(model_runner) - # self.graph_mode = False + self.conv_states_shape = ( + model_runner.req_to_token_pool.mamba_pool.mamba_cache.conv[0].shape + ) + + decode_backend = get_linear_attn_decode_backend() + prefill_backend = get_linear_attn_prefill_backend() + self.kernel_dispatcher = AscendGDNKernelDispatcher(decode_backend, prefill_backend) def prepare_gdn_inputs( self, @@ -108,7 +121,13 @@ def init_forward_metadata_capture_cuda_graph( 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 + 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 @@ -127,8 +146,14 @@ def init_forward_metadata_replay_cuda_graph( 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 + 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 @@ -161,7 +186,9 @@ def forward_decode( 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 + 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) @@ -169,9 +196,15 @@ def forward_decode( 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, + 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, ) @@ -203,112 +236,158 @@ def forward_extend( 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 - intermediate_state_cache = None - if not forward_batch.spec_algorithm.is_none(): intermediate_state_cache = mamba_cache_params.intermediate_ssm if is_target_verify: assert isinstance(mamba_cache_params, MambaPool.SpeculativeState) - intermediate_conv_window_cache = mamba_cache_params.intermediate_conv_window[0] + 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 + dtype=torch.bool, + device=forward_batch.input_ids.device, ) intermediate_state_indices = torch.arange( cache_indices.shape[0], dtype=torch.int32, device=cache_indices.device ) else: has_initial_states = forward_batch.extend_prefix_lens > 0 - if is_target_verify: batch_size = seq_len // forward_batch.spec_info.draft_token_num draft_token_num = forward_batch.spec_info.draft_token_num 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: + 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) + mixed_qkv_reshaped = mixed_qkv.view( + batch_size, draft_token_num, -1 + ) conv_states_to_use = conv_states[cache_indices] mixed_qkv_processed, new_conv_state = npu_causal_conv1d_update( mixed_qkv_reshaped.transpose(1, 2).contiguous(), layer.conv_weights, conv_states_to_use.transpose(1, 2).contiguous(), - layer.bias, True + layer.bias, + True ) mixed_qkv = mixed_qkv_processed.transpose(1, 2).contiguous().view(seq_len, -1) conv_states[cache_indices] = new_conv_state.transpose(1, 2).contiguous() else: mixed_qkv = mixed_qkv.transpose(0, 1) - if forward_batch.mamba_track_mask is not None and forward_batch.mamba_track_mask.any(): + 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) + mixed_qkv_to_track = mixed_qkv[ + :, forward_metadata.track_conv_indices + ].transpose(0, 1) mask_indices = forward_batch.mamba_track_mask.nonzero(as_tuple=True)[0] conv_states[conv_dst[mask_indices]] = mixed_qkv_to_track - 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, + 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() - query, key, value = torch.split(mixed_qkv, [layer.q_dim, layer.k_dim, layer.v_dim], dim=-1) + query, key, value = torch.split( + mixed_qkv, + [layer.q_dim, layer.k_dim, layer.v_dim], + dim=-1, + ) + actual_seq_len = query.shape[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 + 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) - query = query.view(1, actual_seq_len, num_heads, head_k_dim) - key = key.view(1, actual_seq_len, num_heads, head_k_dim) - value = value.view(1, actual_seq_len, num_value_heads, head_v_dim) g, beta = fused_gdn_gating(layer.A_log, a, b, layer.dt_bias) if is_target_verify: query = query.view(-1, num_heads, head_k_dim) key = key.view(-1, num_heads, head_k_dim) value = value.view(-1, num_value_heads, head_v_dim) - query = l2norm_fwd(query.contiguous(), eps=1e-6, output_dtype=torch.bfloat16) + query = l2norm_fwd( + query.contiguous(), eps=1e-6, output_dtype=torch.bfloat16 + ) key = l2norm_fwd(key.contiguous(), eps=1e-6, output_dtype=torch.bfloat16) - core_attn_out = self.fused_recurrent_gated_delta_rule_update_npu( - query, key, value, recurrent_state=ssm_states, - beta=beta, g=g, cache_indices=cache_indices, + core_attn_out = self.fused_recurrent_gated_delta_rule_update( + query, + key, + value, + 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) + 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: + 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, + 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 not forward_batch.spec_algorithm.is_none() and last_recurrent_state is not None: + if is_cpu() and 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) intermediate_state_cache[cache_indices, 0] = last_recurrent_state - elif last_recurrent_state is not None: - last_recurrent_state = last_recurrent_state.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) + self._track_mamba_state_extend( + forward_batch, h, ssm_states, forward_metadata + ) + return core_attn_out - def fused_recurrent_gated_delta_rule_update_npu( + def fused_recurrent_gated_delta_rule_update( self, query: torch.Tensor, key: torch.Tensor, @@ -319,7 +398,7 @@ def fused_recurrent_gated_delta_rule_update_npu( cache_indices: torch.Tensor, intermediate_state: Optional[torch.Tensor] = None, ): - _, num_heads, head_k_dim = query.shape + _, num_heads, head_k_dim = query.shape # T, N, D _, num_value_heads, head_v_dim = value.shape beta = beta.view(-1, num_value_heads).to(torch.bfloat16) g = g.view(-1, num_value_heads).to(torch.float32) @@ -328,13 +407,21 @@ def fused_recurrent_gated_delta_rule_update_npu( scale = 1 / (head_k_dim ** 0.5) if intermediate_state is not None: - ssm_state = intermediate_state.view(-1, num_value_heads, head_k_dim, head_v_dim) + # MTP intermediate_state + # intermediate_state[cache_indices, 0] = recurrent_state[cache_indices] # update indexput slow + ssm_state = intermediate_state.view( + -1, num_value_heads, head_k_dim, head_v_dim + ) else: ssm_state = recurrent_state[cache_indices] if self.graph_mode: - num_accepted_tokens = torch.ones([batch_size], dtype=torch.int32, device=cache_indices.device) - actual_seq_lengths = torch.ones([batch_size], dtype=torch.int32, device=cache_indices.device) * seq_len + num_accepted_tokens = torch.ones( + [batch_size], dtype=torch.int32, device=cache_indices.device + ) + actual_seq_lengths = torch.ones( + [batch_size], dtype=torch.int32, device=cache_indices.device + ) * seq_len ssm_state_indices = self.forward_metadata.mamba_cache_indices_gdn else: num_accepted_tokens = self.num_accepted_tokens @@ -342,11 +429,20 @@ def fused_recurrent_gated_delta_rule_update_npu( ssm_state_indices = self.forward_metadata.mamba_cache_indices attn_core_out = torch_npu.npu_recurrent_gated_delta_rule( - query, key, value, ssm_state, beta=beta, scale=scale, - actual_seq_lengths=actual_seq_lengths, ssm_state_indices=ssm_state_indices, - num_accepted_tokens=num_accepted_tokens, g=g, + query, + key, + value, + ssm_state, # for shape: (BlockNum, Nv, Dv, Dk) + beta=beta, + scale=scale, + actual_seq_lengths=actual_seq_lengths, + ssm_state_indices=ssm_state_indices, + num_accepted_tokens=num_accepted_tokens, + g=g, ) if intermediate_state is not None: - intermediate_state = ssm_state.view(-1, seq_len, num_value_heads, head_k_dim, head_v_dim) + intermediate_state = ssm_state.view( + -1, seq_len, num_value_heads, head_k_dim, head_v_dim + ) return attn_core_out \ No newline at end of file From eb9232109b3161dbb81181c47a5ef6e2228dd566 Mon Sep 17 00:00:00 2001 From: iridiumine <42236072+iridiumine@users.noreply.github.com> Date: Thu, 19 Mar 2026 17:28:04 +0800 Subject: [PATCH 3/7] style: fix lint issues --- .../npu/attention/ascend_gdn_backend.py | 115 ++++++++++-------- .../hardware_backend/npu/memory_pool_npu.py | 13 +- .../layers/attention/attention_registry.py | 4 +- .../attention/hybrid_linear_attn_backend.py | 38 ++++-- python/sglang/srt/layers/layernorm.py | 4 +- python/sglang/srt/mem_cache/memory_pool.py | 9 +- python/sglang/srt/models/qwen3_5_mtp.py | 4 +- python/sglang/srt/models/qwen3_next_mtp.py | 4 +- 8 files changed, 119 insertions(+), 72 deletions(-) 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 index 9cce5c5c4a82..562e0b17fce2 100644 --- a/python/sglang/srt/hardware_backend/npu/attention/ascend_gdn_backend.py +++ b/python/sglang/srt/hardware_backend/npu/attention/ascend_gdn_backend.py @@ -1,12 +1,18 @@ -from typing import Tuple, Union, Optional +from typing import Optional, Tuple, Union + import torch import torch.nn.functional as F import torch_npu +from sgl_kernel_npu.fla.fused_gdn_gating import fused_gdn_gating_npu +from sgl_kernel_npu.mamba.causal_conv1d import ( + causal_conv1d_fn_npu, + causal_conv1d_update_npu, +) from sglang.srt.layers.attention.fla.l2norm import l2norm_fwd from sglang.srt.layers.attention.linear.gdn_backend import ( - GDNKernelDispatcher, GDNAttnBackend, + GDNKernelDispatcher, ) from sglang.srt.layers.attention.linear.utils import ( get_linear_attn_decode_backend, @@ -19,47 +25,46 @@ from sglang.srt.speculative.eagle_info import EagleDraftInput, EagleVerifyInput from sglang.srt.utils import is_cpu -from sgl_kernel_npu.fla.fused_gdn_gating import fused_gdn_gating_npu -from sgl_kernel_npu.mamba.causal_conv1d import ( - causal_conv1d_fn_npu, - causal_conv1d_update_npu, -) - fused_gdn_gating = fused_gdn_gating_npu causal_conv1d_fn = causal_conv1d_fn_npu causal_conv1d_update = causal_conv1d_update_npu + def npu_causal_conv1d_update( - hidden_state: torch.Tensor, - weight: torch.Tensor, - conv_state: torch.Tensor, - bias: Optional[torch.Tensor] = None, - silu_activation: bool = True, - ): - bsz, hidden_size, seq_len = hidden_state.shape - kernel_size = weight.shape[-1] - target_state_len = (kernel_size - 1) + (seq_len - 1) - full_context = torch.cat([conv_state, hidden_state], dim=-1).to(weight.dtype) - computation_input = full_context[:, :, -(kernel_size - 1 + seq_len):] - windows = computation_input.unfold(-1, kernel_size, 1) - - out = (windows * weight[None, :, None, :]).sum(dim=-1) - if bias is not None: - out = out + bias[None, :, None] - if silu_activation: - out = F.silu(out) - out = out.to(hidden_state.dtype) - - if target_state_len > 0: - new_conv_state = full_context[:, :, -target_state_len:] - else: - new_conv_state = torch.empty(bsz, hidden_size, 0, device=hidden_state.device, dtype=hidden_state.dtype) + hidden_state: torch.Tensor, + weight: torch.Tensor, + conv_state: torch.Tensor, + bias: Optional[torch.Tensor] = None, + silu_activation: bool = True, +): + bsz, hidden_size, seq_len = hidden_state.shape + kernel_size = weight.shape[-1] + target_state_len = (kernel_size - 1) + (seq_len - 1) + full_context = torch.cat([conv_state, hidden_state], dim=-1).to(weight.dtype) + computation_input = full_context[:, :, -(kernel_size - 1 + seq_len) :] + windows = computation_input.unfold(-1, kernel_size, 1) + + out = (windows * weight[None, :, None, :]).sum(dim=-1) + if bias is not None: + out = out + bias[None, :, None] + if silu_activation: + out = F.silu(out) + out = out.to(hidden_state.dtype) + + if target_state_len > 0: + new_conv_state = full_context[:, :, -target_state_len:] + else: + new_conv_state = torch.empty( + bsz, hidden_size, 0, device=hidden_state.device, dtype=hidden_state.dtype + ) + + return out, new_conv_state - return out, new_conv_state class AscendGDNKernelDispatcher(GDNKernelDispatcher): pass + class AscendGDNAttnBackend(GDNAttnBackend): def __init__(self, model_runner: ModelRunner): @@ -67,10 +72,12 @@ def __init__(self, model_runner: ModelRunner): self.conv_states_shape = ( model_runner.req_to_token_pool.mamba_pool.mamba_cache.conv[0].shape ) - + decode_backend = get_linear_attn_decode_backend() prefill_backend = get_linear_attn_prefill_backend() - self.kernel_dispatcher = AscendGDNKernelDispatcher(decode_backend, prefill_backend) + self.kernel_dispatcher = AscendGDNKernelDispatcher( + decode_backend, prefill_backend + ) def prepare_gdn_inputs( self, @@ -104,7 +111,7 @@ def init_forward_metadata(self, forward_batch: ForwardBatch): self.prepare_gdn_inputs( forward_batch.batch_size, forward_batch.forward_mode, - forward_batch.spec_info + forward_batch.spec_info, ) self.graph_mode = False @@ -127,7 +134,7 @@ def init_forward_metadata_capture_cuda_graph( seq_lens, encoder_lens, forward_mode, - spec_info + spec_info, ) self.prepare_gdn_inputs(bs, forward_mode, spec_info) self.graph_mode = True @@ -153,7 +160,7 @@ def init_forward_metadata_replay_cuda_graph( encoder_lens, forward_mode, spec_info, - seq_lens_cpu + seq_lens_cpu, ) self.prepare_gdn_inputs(bs, forward_mode, spec_info) self.graph_mode = True @@ -267,18 +274,18 @@ def forward_extend( 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 - ) + mixed_qkv_reshaped = mixed_qkv.view(batch_size, draft_token_num, -1) conv_states_to_use = conv_states[cache_indices] mixed_qkv_processed, new_conv_state = npu_causal_conv1d_update( mixed_qkv_reshaped.transpose(1, 2).contiguous(), layer.conv_weights, conv_states_to_use.transpose(1, 2).contiguous(), layer.bias, - True + True, + ) + mixed_qkv = ( + mixed_qkv_processed.transpose(1, 2).contiguous().view(seq_len, -1) ) - mixed_qkv = mixed_qkv_processed.transpose(1, 2).contiguous().view(seq_len, -1) conv_states[cache_indices] = new_conv_state.transpose(1, 2).contiguous() else: mixed_qkv = mixed_qkv.transpose(0, 1) @@ -293,7 +300,7 @@ def forward_extend( mask_indices = forward_batch.mamba_track_mask.nonzero(as_tuple=True)[0] conv_states[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_for_prefill = conv_states[:, -(kernel_size - 1) :, :] conv_states_tmp = conv_states_for_prefill.transpose(1, 2).contiguous() mixed_qkv = causal_conv1d_fn( @@ -307,7 +314,9 @@ def forward_extend( 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() + conv_states[:, -(kernel_size - 1) :, :] = conv_states_tmp.transpose( + 1, 2 + ).contiguous() query, key, value = torch.split( mixed_qkv, @@ -356,7 +365,6 @@ def forward_extend( dim=0, ) else: - 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, @@ -373,7 +381,9 @@ def forward_extend( ) 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) + last_recurrent_state = last_recurrent_state.transpose(-1, -2).to( + ssm_states.dtype, copy=False + ) intermediate_state_cache[cache_indices, 0] = last_recurrent_state else: last_recurrent_state = last_recurrent_state.to( @@ -404,7 +414,7 @@ def fused_recurrent_gated_delta_rule_update( g = g.view(-1, num_value_heads).to(torch.float32) batch_size = cache_indices.shape[0] seq_len = query.shape[0] // batch_size - scale = 1 / (head_k_dim ** 0.5) + scale = 1 / (head_k_dim**0.5) if intermediate_state is not None: # MTP intermediate_state @@ -419,9 +429,10 @@ def fused_recurrent_gated_delta_rule_update( num_accepted_tokens = torch.ones( [batch_size], dtype=torch.int32, device=cache_indices.device ) - actual_seq_lengths = torch.ones( - [batch_size], dtype=torch.int32, device=cache_indices.device - ) * seq_len + actual_seq_lengths = ( + torch.ones([batch_size], dtype=torch.int32, device=cache_indices.device) + * seq_len + ) ssm_state_indices = self.forward_metadata.mamba_cache_indices_gdn else: num_accepted_tokens = self.num_accepted_tokens @@ -445,4 +456,4 @@ def fused_recurrent_gated_delta_rule_update( intermediate_state = ssm_state.view( -1, seq_len, num_value_heads, head_k_dim, head_v_dim ) - return attn_core_out \ No newline at end of file + return attn_core_out 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 db36a2c3f77c..92a988186c4f 100644 --- a/python/sglang/srt/hardware_backend/npu/memory_pool_npu.py +++ b/python/sglang/srt/hardware_backend/npu/memory_pool_npu.py @@ -14,7 +14,10 @@ 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): + +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 @@ -22,7 +25,12 @@ def _init_npu_conv_state(conv_state_in, conv_state_shape, speculative_num_draft_ # 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]), + 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, ) @@ -30,6 +38,7 @@ def _init_npu_conv_state(conv_state_in, conv_state_shape, speculative_num_draft_ ] 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 bf13ea2e149a..e42d652b85df 100644 --- a/python/sglang/srt/layers/attention/attention_registry.py +++ b/python/sglang/srt/layers/attention/attention_registry.py @@ -202,7 +202,9 @@ def attn_backend_wrapper(runner: "ModelRunner", full_attn_backend: "AttentionBac from sglang.srt.utils import is_blackwell, is_npu if is_npu(): - from sglang.srt.hardware_backend.npu.attention.ascend_gdn_backend import AscendGDNAttnBackend as GDNAttnBackend + from sglang.srt.hardware_backend.npu.attention.ascend_gdn_backend import ( + AscendGDNAttnBackend as GDNAttnBackend, + ) else: from sglang.srt.layers.attention.linear.gdn_backend import GDNAttnBackend 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 6ac003a0bce2..ab38f58814f7 100644 --- a/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py +++ b/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py @@ -30,7 +30,10 @@ ) if is_npu(): - from sglang.srt.layers.attention.mamba.mamba_state_scatter_triton import conv_state_rollback, move_intermediate_cache_dynamic_h_block_v1 + from sglang.srt.layers.attention.mamba.mamba_state_scatter_triton import ( + conv_state_rollback, + move_intermediate_cache_dynamic_h_block_v1, + ) logger = logging.getLogger(__name__) @@ -415,7 +418,10 @@ def init_cuda_graph_state(self, max_bs: int, max_num_tokens: int): ) self.state_indices_list_gdn.append( torch.full( - ((i + 1) * draft_token_num,), self.pad_slot_id, dtype=torch.int32, device=self.device + ((i + 1) * draft_token_num,), + self.pad_slot_id, + dtype=torch.int32, + device=self.device, ) ) self.query_start_loc_list.append( @@ -482,11 +488,14 @@ def _capture_metadata( self.cached_cuda_graph_verify_query_start_loc[: bs + 1] ) start_indices = mamba_indices * spec_info.draft_token_num - offset = torch.arange(spec_info.draft_token_num, device=start_indices.device) + offset = torch.arange( + spec_info.draft_token_num, device=start_indices.device + ) ranges = start_indices.unsqueeze(1) + offset ssm_state_indices = ranges.flatten().to(torch.int32) - self.state_indices_list_gdn[bs - 1][:len(mamba_indices) * spec_info.draft_token_num].copy_( - ssm_state_indices) + 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=}") @@ -538,13 +547,20 @@ def _replay_metadata( bs - num_padding ) elif forward_mode.is_target_verify(): - start_indices = mamba_indices[:bs - num_padding] * spec_info.draft_token_num - offset = torch.arange(spec_info.draft_token_num, device=start_indices.device) + start_indices = ( + mamba_indices[: bs - num_padding] * spec_info.draft_token_num + ) + offset = torch.arange( + spec_info.draft_token_num, device=start_indices.device + ) ranges = start_indices.unsqueeze(1) + offset ssm_state_indices = ranges.flatten().to(torch.int32) 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 + : 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] @@ -988,7 +1004,9 @@ def update_mamba_state_after_mtp_verify( valid_state_indices = state_indices_tensor.to(torch.int64) # [N] last_steps = accepted_steps.to(torch.int64) # [N] - move_intermediate_cache_dynamic_h_block_v1(intermediate_state_cache, valid_state_indices, last_steps) + move_intermediate_cache_dynamic_h_block_v1( + intermediate_state_cache, valid_state_indices, last_steps + ) draft_token_num = intermediate_state_cache.shape[2] if valid_state_indices.numel() > 0: diff --git a/python/sglang/srt/layers/layernorm.py b/python/sglang/srt/layers/layernorm.py index 50c21ba41f97..d7131cbba74b 100644 --- a/python/sglang/srt/layers/layernorm.py +++ b/python/sglang/srt/layers/layernorm.py @@ -513,7 +513,9 @@ def forward_npu( if residual is not None: if post_residual_addition is not None: residual = residual + post_residual_addition - norm_out, residual = add_gemma_rms_norm(x, self.weight, residual, self.variance_epsilon) + 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) diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index 481249dc374d..33f36cbf9cfb 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py @@ -259,8 +259,13 @@ def __init__( ] 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) + 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 1266b8096aac..037081431e95 100644 --- a/python/sglang/srt/models/qwen3_5_mtp.py +++ b/python/sglang/srt/models/qwen3_5_mtp.py @@ -15,14 +15,13 @@ """Inference-only Qwen3_5 MTP model.""" import logging +import os from typing import Iterable, Optional, Tuple -import os import torch from torch import nn from transformers import PretrainedConfig -from sglang.srt.server_args import get_global_server_args from sglang.srt.distributed import get_pp_group, get_tensor_model_parallel_world_size from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder from sglang.srt.eplb.expert_location import ModelConfigForExpertLocation @@ -33,6 +32,7 @@ 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.server_args import get_global_server_args from sglang.srt.utils import add_prefix, is_npu logger = logging.getLogger(__name__) diff --git a/python/sglang/srt/models/qwen3_next_mtp.py b/python/sglang/srt/models/qwen3_next_mtp.py index 9bbdb9d8afc6..9270cacd6796 100644 --- a/python/sglang/srt/models/qwen3_next_mtp.py +++ b/python/sglang/srt/models/qwen3_next_mtp.py @@ -15,16 +15,16 @@ """Inference-only Qwen3Next MTP Speculative Decoding.""" import logging +import os from typing import Iterable, Optional, Tuple -import os import torch from torch import nn from transformers import PretrainedConfig -from sglang.srt.hardware_backend.npu.graph_runner.npu_graph_runner import is_npu from sglang.srt.distributed import get_pp_group, get_tensor_model_parallel_world_size from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder +from sglang.srt.hardware_backend.npu.graph_runner.npu_graph_runner import is_npu from sglang.srt.layers.layernorm import GemmaRMSNorm from sglang.srt.layers.logits_processor import LogitsProcessor from sglang.srt.layers.quantization.base_config import QuantizationConfig From 53b68c62756e600066f8793ed3cf0412db22878e Mon Sep 17 00:00:00 2001 From: iridiumine <42236072+iridiumine@users.noreply.github.com> Date: Tue, 24 Mar 2026 16:03:09 +0800 Subject: [PATCH 4/7] feat: add Ascend NPU support for GDN attention backend --- .../npu/attention/ascend_gdn_backend.py | 141 ++++++++++++------ .../layers/attention/fla/fused_gdn_gating.py | 66 ++++++++ .../attention/hybrid_linear_attn_backend.py | 9 +- 3 files changed, 167 insertions(+), 49 deletions(-) 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 index 562e0b17fce2..3fc962828563 100644 --- a/python/sglang/srt/hardware_backend/npu/attention/ascend_gdn_backend.py +++ b/python/sglang/srt/hardware_backend/npu/attention/ascend_gdn_backend.py @@ -9,7 +9,9 @@ causal_conv1d_update_npu, ) -from sglang.srt.layers.attention.fla.l2norm import l2norm_fwd +from sglang.srt.layers.attention.fla.fused_gdn_gating import ( + fused_gdn_gating_kernel_without_sigmoid, +) from sglang.srt.layers.attention.linear.gdn_backend import ( GDNAttnBackend, GDNKernelDispatcher, @@ -69,10 +71,6 @@ class AscendGDNAttnBackend(GDNAttnBackend): def __init__(self, model_runner: ModelRunner): super().__init__(model_runner) - self.conv_states_shape = ( - model_runner.req_to_token_pool.mamba_pool.mamba_cache.conv[0].shape - ) - decode_backend = get_linear_attn_decode_backend() prefill_backend = get_linear_attn_prefill_backend() self.kernel_dispatcher = AscendGDNKernelDispatcher( @@ -99,10 +97,8 @@ def prepare_gdn_inputs( offset = torch.arange(seq_len, device=start_indices.device) ranges = start_indices.unsqueeze(1) + offset self.ssm_state_indices = ranges.flatten().to(torch.int32) - self.forward_metadata.ssm_state_indices = self.ssm_state_indices else: self.ssm_state_indices = cache_indices - self.forward_metadata.ssm_state_indices = self.ssm_state_indices def init_forward_metadata(self, forward_batch: ForwardBatch): if forward_batch.forward_mode.is_draft_extend(True): @@ -243,10 +239,9 @@ def forward_extend( 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 not forward_batch.spec_algorithm.is_none(): - intermediate_state_cache = mamba_cache_params.intermediate_ssm 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] ) @@ -255,9 +250,6 @@ def forward_extend( dtype=torch.bool, device=forward_batch.input_ids.device, ) - intermediate_state_indices = torch.arange( - cache_indices.shape[0], dtype=torch.int32, device=cache_indices.device - ) else: has_initial_states = forward_batch.extend_prefix_lens > 0 if is_target_verify: @@ -318,34 +310,26 @@ def forward_extend( 1, 2 ).contiguous() - query, key, value = torch.split( - mixed_qkv, - [layer.q_dim, layer.k_dim, layer.v_dim], - dim=-1, - ) - - actual_seq_len = query.shape[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 - query = query.view(1, actual_seq_len, layer.num_q_heads, layer.head_q_dim) - key = key.view(1, actual_seq_len, layer.num_k_heads, layer.head_k_dim) - value = value.view(1, actual_seq_len, layer.num_v_heads, layer.head_v_dim) - - g, beta = fused_gdn_gating(layer.A_log, a, b, layer.dt_bias) - if is_target_verify: - query = query.view(-1, num_heads, head_k_dim) - key = key.view(-1, num_heads, head_k_dim) - value = value.view(-1, num_value_heads, head_v_dim) - query = l2norm_fwd( - query.contiguous(), eps=1e-6, output_dtype=torch.bfloat16 + g, beta = fused_gdn_gating_kernel_without_sigmoid( + layer.A_log, a, b, layer.dt_bias ) - key = l2norm_fwd(key.contiguous(), eps=1e-6, output_dtype=torch.bfloat16) + 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 - core_attn_out = self.fused_recurrent_gated_delta_rule_update( - query, - key, - value, + 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_npu_v2( + mixed_qkv, + num_heads, + num_value_heads, + head_k_dim, + head_v_dim, recurrent_state=ssm_states, beta=beta, g=g, @@ -365,6 +349,18 @@ def forward_extend( 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, @@ -384,12 +380,11 @@ def forward_extend( last_recurrent_state = last_recurrent_state.transpose(-1, -2).to( ssm_states.dtype, copy=False ) - intermediate_state_cache[cache_indices, 0] = last_recurrent_state else: last_recurrent_state = last_recurrent_state.to( ssm_states.dtype, copy=False ) - ssm_states[cache_indices] = last_recurrent_state + ssm_states[cache_indices] = last_recurrent_state if h is not None: self._track_mamba_state_extend( forward_batch, h, ssm_states, forward_metadata @@ -397,7 +392,7 @@ def forward_extend( return core_attn_out - def fused_recurrent_gated_delta_rule_update( + def fused_recurrent_gated_delta_rule_update_npu( self, query: torch.Tensor, key: torch.Tensor, @@ -418,12 +413,14 @@ def fused_recurrent_gated_delta_rule_update( if intermediate_state is not None: # MTP intermediate_state - # intermediate_state[cache_indices, 0] = recurrent_state[cache_indices] # update indexput slow + intermediate_state[cache_indices, 0] = recurrent_state[ + cache_indices + ] # update indexput slow ssm_state = intermediate_state.view( -1, num_value_heads, head_k_dim, head_v_dim ) else: - ssm_state = recurrent_state[cache_indices] + ssm_state = recurrent_state if self.graph_mode: num_accepted_tokens = torch.ones( @@ -437,7 +434,7 @@ def fused_recurrent_gated_delta_rule_update( else: num_accepted_tokens = self.num_accepted_tokens actual_seq_lengths = self.actual_seq_lengths - ssm_state_indices = self.forward_metadata.mamba_cache_indices + ssm_state_indices = self.ssm_state_indices attn_core_out = torch_npu.npu_recurrent_gated_delta_rule( query, @@ -457,3 +454,61 @@ def fused_recurrent_gated_delta_rule_update( -1, seq_len, num_value_heads, head_k_dim, head_v_dim ) return attn_core_out + + def fused_recurrent_gated_delta_rule_update_npu_v2( + 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/layers/attention/fla/fused_gdn_gating.py b/python/sglang/srt/layers/attention/fla/fused_gdn_gating.py index 6e92208ec130..a82c18ad9abb 100644 --- a/python/sglang/srt/layers/attention/fla/fused_gdn_gating.py +++ b/python/sglang/srt/layers/attention/fla/fused_gdn_gating.py @@ -67,3 +67,69 @@ def fused_gdn_gating( num_warps=1, ) return g, beta_output + + +@triton.jit +def fused_gdn_gating_kernel_without_sigmoid_kernel( + g, + A_log, + a, + dt_bias, + batch, + seq_len, + NUM_HEADS: tl.constexpr, + beta: tl.constexpr, + threshold: tl.constexpr, + BLK_BATCHES: tl.constexpr, + BLK_HEADS: tl.constexpr, +): + i_b, i_s, i_d = tl.program_id(0), tl.program_id(1), tl.program_id(2) + batch_off = i_b * BLK_BATCHES + tl.arange(0, BLK_BATCHES) + head_off = i_d * BLK_HEADS + tl.arange(0, BLK_HEADS) + head_mask = head_off < NUM_HEADS + a_off = ( + batch_off[:, None] * seq_len * NUM_HEADS + i_s * NUM_HEADS + head_off[None, :] + ) + a_mask = (batch_off[:, None] < batch) & head_mask[None, :] + blk_A_log = tl.load(A_log + head_off, mask=head_mask) + blk_bias = tl.load(dt_bias + head_off, mask=head_mask) + blk_a = tl.load(a + a_off, mask=a_mask) + x = blk_a.to(tl.float32) + blk_bias.to(tl.float32) + softplus_x = tl.where( + beta * x <= threshold, (1 / beta) * tl.log(1 + tl.exp(beta * x)), x + ) + blk_g = -tl.exp(blk_A_log.to(tl.float32)) * softplus_x + tl.store(g + a_off, blk_g.to(g.dtype.element_ty), mask=a_mask) + + +def fused_gdn_gating_kernel_without_sigmoid( + A_log: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + dt_bias: torch.Tensor, + beta: float = 1.0, + threshold: float = 20.0, +) -> Tuple[torch.Tensor, torch.Tensor]: + batch, num_heads = a.shape + seq_len = 1 + g = torch.empty_like(a, dtype=torch.float32) + num_cores = 48 # num_vectorcore of NPU + NUM_BLK_BATCHES = triton.cdiv(num_cores, triton.cdiv(num_heads, 8)) + BLK_BATCHES = triton.cdiv(batch, NUM_BLK_BATCHES) + grid = (NUM_BLK_BATCHES, seq_len, triton.cdiv(num_heads, 8)) + fused_gdn_gating_kernel_without_sigmoid_kernel[grid]( + g, + A_log, + a, + dt_bias, + batch, + seq_len, + num_heads, + beta, + threshold, + BLK_BATCHES=BLK_BATCHES, + BLK_HEADS=8, + num_warps=1, + ) + g = g.unsqueeze(0) + return g, b 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 ab38f58814f7..b6e2c8414e80 100644 --- a/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py +++ b/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py @@ -32,7 +32,7 @@ if is_npu(): from sglang.srt.layers.attention.mamba.mamba_state_scatter_triton import ( conv_state_rollback, - move_intermediate_cache_dynamic_h_block_v1, + move_intermediate_cache_dynamic_h_block, ) logger = logging.getLogger(__name__) @@ -1004,8 +1004,8 @@ def update_mamba_state_after_mtp_verify( valid_state_indices = state_indices_tensor.to(torch.int64) # [N] last_steps = accepted_steps.to(torch.int64) # [N] - move_intermediate_cache_dynamic_h_block_v1( - intermediate_state_cache, valid_state_indices, last_steps + move_intermediate_cache_dynamic_h_block( + ssm_states, intermediate_state_cache, valid_state_indices, last_steps ) draft_token_num = intermediate_state_cache.shape[2] @@ -1050,9 +1050,6 @@ def update_mamba_state_after_mtp_verify( mamba_steps_to_track, ) - def get_verify_buffers_to_fill_after_draft(self): - return [None, None] - def update_verify_buffers_to_fill_after_draft( self, spec_info: SpecInput, cuda_graph_bs: Optional[int] ): From acf82841212fa3385bdebb627701d615946b44e0 Mon Sep 17 00:00:00 2001 From: iridiumine <42236072+iridiumine@users.noreply.github.com> Date: Tue, 24 Mar 2026 17:37:59 +0800 Subject: [PATCH 5/7] fix: delete unuse func and change v2 to it --- .../npu/attention/ascend_gdn_backend.py | 68 +------------------ 1 file changed, 2 insertions(+), 66 deletions(-) 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 index 3fc962828563..8b5228a7a66c 100644 --- a/python/sglang/srt/hardware_backend/npu/attention/ascend_gdn_backend.py +++ b/python/sglang/srt/hardware_backend/npu/attention/ascend_gdn_backend.py @@ -2,7 +2,6 @@ import torch import torch.nn.functional as F -import torch_npu from sgl_kernel_npu.fla.fused_gdn_gating import fused_gdn_gating_npu from sgl_kernel_npu.mamba.causal_conv1d import ( causal_conv1d_fn_npu, @@ -324,7 +323,7 @@ def forward_extend( 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_npu_v2( + core_attn_out = self.fused_recurrent_gated_delta_rule_update( mixed_qkv, num_heads, num_value_heads, @@ -392,70 +391,7 @@ def forward_extend( return core_attn_out - def fused_recurrent_gated_delta_rule_update_npu( - self, - query: torch.Tensor, - key: torch.Tensor, - value: torch.Tensor, - recurrent_state: torch.Tensor, - beta: torch.Tensor, - g: torch.Tensor, - cache_indices: torch.Tensor, - intermediate_state: Optional[torch.Tensor] = None, - ): - _, num_heads, head_k_dim = query.shape # T, N, D - _, num_value_heads, head_v_dim = value.shape - beta = beta.view(-1, num_value_heads).to(torch.bfloat16) - g = g.view(-1, num_value_heads).to(torch.float32) - batch_size = cache_indices.shape[0] - seq_len = query.shape[0] // batch_size - scale = 1 / (head_k_dim**0.5) - - if intermediate_state is not None: - # MTP intermediate_state - intermediate_state[cache_indices, 0] = recurrent_state[ - cache_indices - ] # update indexput slow - ssm_state = intermediate_state.view( - -1, num_value_heads, head_k_dim, head_v_dim - ) - else: - ssm_state = recurrent_state - - if self.graph_mode: - num_accepted_tokens = torch.ones( - [batch_size], dtype=torch.int32, device=cache_indices.device - ) - actual_seq_lengths = ( - torch.ones([batch_size], dtype=torch.int32, device=cache_indices.device) - * seq_len - ) - 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_npu.npu_recurrent_gated_delta_rule( - query, - key, - value, - ssm_state, # for shape: (BlockNum, Nv, Dv, Dk) - beta=beta, - scale=scale, - actual_seq_lengths=actual_seq_lengths, - ssm_state_indices=ssm_state_indices, - num_accepted_tokens=num_accepted_tokens, - g=g, - ) - - if intermediate_state is not None: - intermediate_state = ssm_state.view( - -1, seq_len, num_value_heads, head_k_dim, head_v_dim - ) - return attn_core_out - - def fused_recurrent_gated_delta_rule_update_npu_v2( + def fused_recurrent_gated_delta_rule_update( self, mix_qkv: torch.Tensor, num_heads, From da8bd670fbe4d85b666ef4058a92cc1ed396a1d0 Mon Sep 17 00:00:00 2001 From: iridiumine <42236072+iridiumine@users.noreply.github.com> Date: Thu, 26 Mar 2026 16:26:53 +0800 Subject: [PATCH 6/7] fix: use causal_conv1d_update --- .../npu/attention/ascend_gdn_backend.py | 59 +++++-------------- 1 file changed, 16 insertions(+), 43 deletions(-) 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 index 8b5228a7a66c..b5fc8445cfd5 100644 --- a/python/sglang/srt/hardware_backend/npu/attention/ascend_gdn_backend.py +++ b/python/sglang/srt/hardware_backend/npu/attention/ascend_gdn_backend.py @@ -1,7 +1,6 @@ from typing import Optional, Tuple, Union import torch -import torch.nn.functional as F from sgl_kernel_npu.fla.fused_gdn_gating import fused_gdn_gating_npu from sgl_kernel_npu.mamba.causal_conv1d import ( causal_conv1d_fn_npu, @@ -31,37 +30,6 @@ causal_conv1d_update = causal_conv1d_update_npu -def npu_causal_conv1d_update( - hidden_state: torch.Tensor, - weight: torch.Tensor, - conv_state: torch.Tensor, - bias: Optional[torch.Tensor] = None, - silu_activation: bool = True, -): - bsz, hidden_size, seq_len = hidden_state.shape - kernel_size = weight.shape[-1] - target_state_len = (kernel_size - 1) + (seq_len - 1) - full_context = torch.cat([conv_state, hidden_state], dim=-1).to(weight.dtype) - computation_input = full_context[:, :, -(kernel_size - 1 + seq_len) :] - windows = computation_input.unfold(-1, kernel_size, 1) - - out = (windows * weight[None, :, None, :]).sum(dim=-1) - if bias is not None: - out = out + bias[None, :, None] - if silu_activation: - out = F.silu(out) - out = out.to(hidden_state.dtype) - - if target_state_len > 0: - new_conv_state = full_context[:, :, -target_state_len:] - else: - new_conv_state = torch.empty( - bsz, hidden_size, 0, device=hidden_state.device, dtype=hidden_state.dtype - ) - - return out, new_conv_state - - class AscendGDNKernelDispatcher(GDNKernelDispatcher): pass @@ -266,18 +234,23 @@ def forward_extend( seq_len = forward_batch.num_token_non_padded_cpu mixed_qkv_reshaped = mixed_qkv.view(batch_size, draft_token_num, -1) - conv_states_to_use = conv_states[cache_indices] - mixed_qkv_processed, new_conv_state = npu_causal_conv1d_update( - mixed_qkv_reshaped.transpose(1, 2).contiguous(), - layer.conv_weights, - conv_states_to_use.transpose(1, 2).contiguous(), - layer.bias, - True, - ) - mixed_qkv = ( - mixed_qkv_processed.transpose(1, 2).contiguous().view(seq_len, -1) + num_accepted_tokens = torch.full( + (batch_size,), + draft_token_num, + dtype=torch.int32, + device=mixed_qkv.device, ) - conv_states[cache_indices] = new_conv_state.transpose(1, 2).contiguous() + 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 ( From 067c0a184c6dc0df9fc0fe79f11dd6e89d8bbd97 Mon Sep 17 00:00:00 2001 From: iridiumine <42236072+iridiumine@users.noreply.github.com> Date: Mon, 30 Mar 2026 19:51:14 +0800 Subject: [PATCH 7/7] fix: use sgl_kernel_npu.mamba.mamba_state_update_triton --- .../sglang/srt/layers/attention/hybrid_linear_attn_backend.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 b6e2c8414e80..4643147b27c5 100644 --- a/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py +++ b/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py @@ -30,7 +30,7 @@ ) if is_npu(): - from sglang.srt.layers.attention.mamba.mamba_state_scatter_triton import ( + from sgl_kernel_npu.mamba.mamba_state_update_triton import ( conv_state_rollback, move_intermediate_cache_dynamic_h_block, )