[Feature] Support DeepSeek MTP on NPU (#11897)
Co-authored-by: liupeng374 <liupeng374@huawei.com>
This commit is contained in:
@@ -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,
|
||||
)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user