minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[NPU] Support MTP for Qwen3.5 (#20918)

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