[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
parent cafebef154
commit ce6b17c0f9
16 changed files with 850 additions and 117 deletions

View File

@@ -75,9 +75,13 @@ class NPUGraphRunner(CudaGraphRunner):
# Replay
if not is_deepseek_nsa(self.model_runner.model_config.hf_config):
seq_lens = forward_batch.seq_lens.cpu().tolist() + [0] * (
self.bs - self.raw_bs
)
if forward_batch.forward_mode.is_target_verify():
seq_lens_cpu = forward_batch.seq_lens.cpu() + self.num_tokens_per_bs
seq_lens = seq_lens_cpu.tolist() + [0] * (self.bs - self.raw_bs)
else:
seq_lens = forward_batch.seq_lens.cpu().tolist() + [0] * (
self.bs - self.raw_bs
)
thread = threading.Thread(target=self._update_inputs, args=(seq_lens,))
thread.start()
self.graphs[self.bs].replay()