[NPU] bugfix for Base class add mamba_track_indices parameter (#29999)

This commit is contained in:
McZyWu
2026-07-03 15:10:41 +08:00
committed by GitHub
parent fe60764f54
commit 8416544ab0
2 changed files with 4 additions and 2 deletions
@@ -132,6 +132,8 @@ class AscendMambaAttnBackendBase(MambaAttnBackendBase):
seq_lens_cpu: Optional[torch.Tensor], seq_lens_cpu: Optional[torch.Tensor],
num_padding: Optional[int] = None, num_padding: Optional[int] = None,
in_capture: bool = False, in_capture: bool = False,
*args,
**kwargs,
): ):
# out_graph passes seq_lens_cpu=None at capture; mirror the base guard. # out_graph passes seq_lens_cpu=None at capture; mirror the base guard.
if seq_lens_cpu is None: if seq_lens_cpu is None:
@@ -284,7 +284,7 @@ def npu_fused_moe_without_routing_weights_bf16(
# gmm1: gate_up_proj # gmm1: gate_up_proj
hidden_states = torch.ops.npu.npu_grouped_matmul( hidden_states = torch.ops.npu.npu_grouped_matmul(
x=[hidden_states], x=[hidden_states],
weight=[layer.w13_weight], weight=[layer.w13_weight.transpose(1, 2)],
split_item=2, split_item=2,
group_list_type=group_list_type, group_list_type=group_list_type,
group_type=0, group_type=0,
@@ -297,7 +297,7 @@ def npu_fused_moe_without_routing_weights_bf16(
# gmm2: down_proj # gmm2: down_proj
hidden_states = torch.ops.npu.npu_grouped_matmul( hidden_states = torch.ops.npu.npu_grouped_matmul(
x=[hidden_states], x=[hidden_states],
weight=[layer.w2_weight], weight=[layer.w2_weight.transpose(1, 2)],
split_item=2, split_item=2,
group_list_type=group_list_type, group_list_type=group_list_type,
group_type=0, group_type=0,