[Feature] Support DeepSeek MTP on NPU (#11897)

Co-authored-by: liupeng374 <liupeng374@huawei.com>
This commit is contained in:
Even Zhou
2025-10-30 21:56:56 +08:00
committed by GitHub
co-authored by liupeng374
parent cafebef154
commit ce6b17c0f9
16 changed files with 850 additions and 117 deletions
+11 -2
View File
@@ -38,12 +38,13 @@ from sglang.srt.models.deepseek_v2 import (
enable_nextn_moe_bf16_cast_to_fp8,
)
from sglang.srt.server_args import get_global_server_args
from sglang.srt.utils import BumpAllocator, add_prefix, is_cuda
from sglang.srt.utils import BumpAllocator, add_prefix, is_cuda, is_npu
logger = logging.getLogger(__name__)
_is_cuda = is_cuda()
_is_npu = is_npu()
class DeepseekModelNextN(nn.Module):
@@ -85,13 +86,21 @@ class DeepseekModelNextN(nn.Module):
self.eh_proj = nn.Linear(2 * config.hidden_size, config.hidden_size, bias=False)
self.alt_stream = torch.cuda.Stream() if _is_cuda else None
layer_name = "decoder"
if _is_npu and (
get_global_server_args().speculative_draft_model_path
== get_global_server_args().model_path
):
layer_name = "layers." + str(config.num_hidden_layers)
self.decoder = DeepseekV2DecoderLayer(
config,
0,
quant_config=quant_config,
moe_quant_config=moe_quant_config,
is_nextn=True,
prefix=add_prefix("decoder", prefix),
prefix=add_prefix(layer_name, prefix),
alt_stream=self.alt_stream,
)