[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:
co-authored by
Zhengda Qin
richhuan
parent
738b1ac988
commit
b56366f827
@@ -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(
|
||||
|
||||
Reference in New Issue
Block a user