[NPU] bugfix for Base class add mamba_track_indices parameter (#29999)
This commit is contained in:
@@ -132,6 +132,8 @@ class AscendMambaAttnBackendBase(MambaAttnBackendBase):
|
||||
seq_lens_cpu: Optional[torch.Tensor],
|
||||
num_padding: Optional[int] = None,
|
||||
in_capture: bool = False,
|
||||
*args,
|
||||
**kwargs,
|
||||
):
|
||||
# out_graph passes seq_lens_cpu=None at capture; mirror the base guard.
|
||||
if seq_lens_cpu is None:
|
||||
|
||||
@@ -284,7 +284,7 @@ def npu_fused_moe_without_routing_weights_bf16(
|
||||
# gmm1: gate_up_proj
|
||||
hidden_states = torch.ops.npu.npu_grouped_matmul(
|
||||
x=[hidden_states],
|
||||
weight=[layer.w13_weight],
|
||||
weight=[layer.w13_weight.transpose(1, 2)],
|
||||
split_item=2,
|
||||
group_list_type=group_list_type,
|
||||
group_type=0,
|
||||
@@ -297,7 +297,7 @@ def npu_fused_moe_without_routing_weights_bf16(
|
||||
# gmm2: down_proj
|
||||
hidden_states = torch.ops.npu.npu_grouped_matmul(
|
||||
x=[hidden_states],
|
||||
weight=[layer.w2_weight],
|
||||
weight=[layer.w2_weight.transpose(1, 2)],
|
||||
split_item=2,
|
||||
group_list_type=group_list_type,
|
||||
group_type=0,
|
||||
|
||||
Reference in New Issue
Block a user