[NPU][Bugfix] fix MTP accuracy regression on Qwen3 hybrid models (#27202)

This commit is contained in:
AndyLi429
2026-06-09 15:47:23 +08:00
committed by GitHub
parent 2218622f50
commit cd6efcb947
2 changed files with 38 additions and 7 deletions
@@ -157,16 +157,11 @@ class AscendMambaAttnBackendBase(MambaAttnBackendBase):
)
elif forward_mode.is_target_verify():
ssm_state_indices = torch.arange(
len(mamba_indices[: bs - num_padding]) * spec_info.draft_token_num,
bs * spec_info.draft_token_num,
dtype=torch.int32,
device=mamba_indices.device,
)
self.state_indices_list_gdn[bs - 1][
: len(mamba_indices[: bs - num_padding]) * spec_info.draft_token_num
].copy_(ssm_state_indices)
self.state_indices_list_gdn[bs - 1][
len(mamba_indices[: bs - num_padding]) * spec_info.draft_token_num :
] = 0
self.state_indices_list_gdn[bs - 1].copy_(ssm_state_indices)
if num_padding == 0:
self.query_start_loc_list[bs - 1].copy_(
self.cached_cuda_graph_verify_query_start_loc[: bs + 1]
@@ -269,6 +264,28 @@ class AscendHybridLinearAttnBackend(HybridLinearAttnBackend):
)
draft_token_num = intermediate_state_cache.shape[2]
if mamba_track_indices is not None:
assert mamba_steps_to_track is not None
mamba_track_indices = mamba_track_indices.to(torch.int64)
mamba_steps_to_track = mamba_steps_to_track.to(torch.int64)
move_intermediate_cache(
ssm_states,
intermediate_state_cache,
mamba_track_indices,
src_indices_tensor,
mamba_steps_to_track,
)
track_mask = mamba_steps_to_track >= 0
# Track conv state from the verify-time window before rolling back
# the working slot; NPU does not keep per-step conv intermediates.
track_indices = mamba_track_indices[track_mask]
if track_indices.numel() > 0:
conv_states[:, track_indices] = conv_states[
:, dst_indices_tensor[track_mask]
]
if dst_indices_tensor.numel() > 0:
conv_state_rollback(
conv_states,
@@ -276,6 +293,15 @@ class AscendHybridLinearAttnBackend(HybridLinearAttnBackend):
last_steps,
draft_token_num,
)
if mamba_track_indices is not None and mamba_track_indices.numel() > 0:
conv_state_rollback(
conv_states,
mamba_track_indices,
mamba_steps_to_track,
draft_token_num,
)
return
def update_verify_buffers_to_fill_after_draft(
@@ -109,6 +109,11 @@ class NPUMHATokenToKVPool(MHATokenToKVPool):
self.k_buffer.append(k_buffer_layer)
self.v_buffer.append(v_buffer_layer)
def _init_kv_copy_and_warmup(self):
# implementation relies on self.data_strides / self.data_ptrs, which the
# NPU paged buffer layout never builds.
self._kv_copy_config = None
# for disagg
def get_contiguous_buf_infos(self):
# layer_num x [seq_len, head_num, head_dim]