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..b5fc8445cfd5 --- /dev/null +++ b/python/sglang/srt/hardware_backend/npu/attention/ascend_gdn_backend.py @@ -0,0 +1,423 @@ +from typing import Optional, Tuple, Union + +import torch +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.fused_gdn_gating import ( + fused_gdn_gating_kernel_without_sigmoid, +) +from sglang.srt.layers.attention.linear.gdn_backend import ( + GDNAttnBackend, + 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 +from sglang.srt.utils import is_cpu + +fused_gdn_gating = fused_gdn_gating_npu +causal_conv1d_fn = causal_conv1d_fn_npu +causal_conv1d_update = causal_conv1d_update_npu + + +class AscendGDNKernelDispatcher(GDNKernelDispatcher): + pass + + +class AscendGDNAttnBackend(GDNAttnBackend): + + def __init__(self, model_runner: ModelRunner): + super().__init__(model_runner) + 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, + 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) + 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: + 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) + 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[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 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 + ) + 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/memory_pool_npu.py b/python/sglang/srt/hardware_backend/npu/memory_pool_npu.py index 82e7ce693e33..92a988186c4f 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 @@ 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 9020e5260af9..e42d652b85df 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,13 @@ 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/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 91194c494396..4643147b27c5 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,19 @@ 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 sgl_kernel_npu.mamba.mamba_state_update_triton import ( + conv_state_rollback, + move_intermediate_cache_dynamic_h_block, + ) + logger = logging.getLogger(__name__) @@ -142,6 +148,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 +416,14 @@ 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 +477,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 +487,17 @@ 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 +515,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 +532,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 +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 + ) + 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 +595,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 +1000,23 @@ 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( + ssm_states, 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 +1049,8 @@ def update_mamba_state_after_mtp_verify( mamba_track_indices, mamba_steps_to_track, ) + + 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..d7131cbba74b 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,13 @@ 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..33f36cbf9cfb 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py @@ -258,6 +258,15 @@ 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..037081431e95 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 +import os from typing import Iterable, Optional, Tuple import torch @@ -31,7 +32,8 @@ 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 +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..9270cacd6796 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 +import os from typing import Iterable, Optional, Tuple import torch @@ -23,6 +24,7 @@ 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 @@ -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