【NPU】fix decode MTP + eagle shape error (#32791)
This commit is contained in:
@@ -497,6 +497,31 @@ class AscendAttnBackend(AttentionBackend):
|
|||||||
):
|
):
|
||||||
self.forward_metadata.seq_lens_cpu_int += self.speculative_step_id + 1
|
self.forward_metadata.seq_lens_cpu_int += self.speculative_step_id + 1
|
||||||
|
|
||||||
|
# Set actual_seq_lengths_q from the pre-pad batch size so that the DSA
|
||||||
|
# indexer reads a value consistent with actual_seq_lengths_kv /
|
||||||
|
# block_tables (which are also built from the pre-pad batch here).
|
||||||
|
# Without this, eager decode under DP attention pads q to the global
|
||||||
|
# max while kv stays local, causing a shape mismatch in
|
||||||
|
# npu_lightning_indexer.
|
||||||
|
if (
|
||||||
|
forward_batch.forward_mode.is_target_verify()
|
||||||
|
or forward_batch.forward_mode.is_draft_extend_v2()
|
||||||
|
):
|
||||||
|
self.forward_metadata.actual_seq_lengths_q = torch.arange(
|
||||||
|
self.speculative_num_draft_tokens,
|
||||||
|
self.speculative_num_draft_tokens
|
||||||
|
+ forward_batch.seq_lens.shape[0] * self.speculative_num_draft_tokens,
|
||||||
|
self.speculative_num_draft_tokens,
|
||||||
|
dtype=torch.int32,
|
||||||
|
device=self.device,
|
||||||
|
)
|
||||||
|
elif forward_batch.forward_mode.is_decode_or_idle():
|
||||||
|
self.forward_metadata.actual_seq_lengths_q = torch.tensor(
|
||||||
|
[1 + i for i in range(forward_batch.seq_lens.shape[0])],
|
||||||
|
dtype=torch.int32,
|
||||||
|
device=self.device,
|
||||||
|
)
|
||||||
|
|
||||||
if (
|
if (
|
||||||
self.use_mla
|
self.use_mla
|
||||||
and forward_batch.forward_mode.is_extend()
|
and forward_batch.forward_mode.is_extend()
|
||||||
@@ -721,6 +746,27 @@ class AscendAttnBackend(AttentionBackend):
|
|||||||
|
|
||||||
self.graph_mode = True
|
self.graph_mode = True
|
||||||
|
|
||||||
|
def _pad_topk_indices(
|
||||||
|
self, topk_indices: torch.Tensor, num_tokens: int
|
||||||
|
) -> torch.Tensor:
|
||||||
|
current_tokens = topk_indices.shape[0]
|
||||||
|
if current_tokens == num_tokens:
|
||||||
|
return topk_indices
|
||||||
|
|
||||||
|
assert current_tokens <= num_tokens, (
|
||||||
|
f"topk_indices rows ({current_tokens}) > num_tokens ({num_tokens}); "
|
||||||
|
"this indicates a mismatch between indexer output and q layout."
|
||||||
|
)
|
||||||
|
|
||||||
|
pad_size = num_tokens - current_tokens
|
||||||
|
padding = torch.full(
|
||||||
|
(pad_size, topk_indices.shape[1]),
|
||||||
|
-1,
|
||||||
|
dtype=topk_indices.dtype,
|
||||||
|
device=topk_indices.device,
|
||||||
|
)
|
||||||
|
return torch.cat([topk_indices, padding], dim=0)
|
||||||
|
|
||||||
def get_cuda_graph_seq_len_fill_value(self):
|
def get_cuda_graph_seq_len_fill_value(self):
|
||||||
return 0
|
return 0
|
||||||
|
|
||||||
@@ -1072,6 +1118,8 @@ class AscendAttnBackend(AttentionBackend):
|
|||||||
actual_seq_lengths_kv,
|
actual_seq_lengths_kv,
|
||||||
)
|
)
|
||||||
else:
|
else:
|
||||||
|
if topk_indices is not None:
|
||||||
|
topk_indices = self._pad_topk_indices(topk_indices, q_nope.shape[0])
|
||||||
topk_indices = _expand_dsa_sparse_indices(topk_indices)
|
topk_indices = _expand_dsa_sparse_indices(topk_indices)
|
||||||
attn_out, _, _ = torch_npu.npu_sparse_flash_attention(
|
attn_out, _, _ = torch_npu.npu_sparse_flash_attention(
|
||||||
query=q_nope,
|
query=q_nope,
|
||||||
|
|||||||
Reference in New Issue
Block a user