[NPU]DeepSeek-V3.2 support npu mlaprolog (#15381)

Co-authored-by: Zhengda Qin <zhengdqin@gmail.com>
Co-authored-by: richhuan <huan_rz@qq.com>
This commit is contained in:
lawtherWu
2026-01-26 20:42:37 +08:00
committed by GitHub
co-authored by Zhengda Qin richhuan
parent 738b1ac988
commit b56366f827
5 changed files with 195 additions and 61 deletions
@@ -1113,6 +1113,7 @@ class Indexer(MultiPlatformOp):
positions: torch.Tensor,
forward_batch: ForwardBatch,
layer_id: int,
dynamic_scale: torch.Tensor = None,
) -> torch.Tensor:
if forward_batch.attn_backend.forward_metadata.seq_lens_cpu_int is None:
actual_seq_lengths_kv = forward_batch.attn_backend.forward_metadata.seq_lens
@@ -1136,6 +1137,9 @@ class Indexer(MultiPlatformOp):
if self.alt_stream is not None:
self.alt_stream.wait_stream(torch.npu.current_stream())
with torch.npu.stream(self.alt_stream):
q_lora = (
(q_lora, dynamic_scale) if dynamic_scale is not None else q_lora
)
q = self.wq_b(q_lora)[
0
] # [bs, 1536] @ [1536, 64 * 128] = [bs, 64 * 128]
@@ -1154,6 +1158,7 @@ class Indexer(MultiPlatformOp):
q.record_stream(self.alt_stream)
q_rope_event = self.alt_stream.record_event()
else:
q_lora = (q_lora, dynamic_scale) if dynamic_scale is not None else q_lora
q = self.wq_b(q_lora)[0] # [bs, 1536] @ [1536, 64 * 128] = [bs, 64 * 128]
q = q.view(bs, self.n_heads, self.head_dim) # [bs, 64, 128]
q_pe, q_nope = torch.split(