[NPU] enable mla prepare fused kernel only when being mla attn (#22024)

2026-04-08 00:49:16 +08:00
parent be42fbbbd7
commit 6131fb5882
1 changed files with 2 additions and 2 deletions
@@ -860,7 +860,7 @@ class AscendAttnBackend(AttentionBackend):
        sinks: Optional[torch.Tensor] = None,
        slopes: Optional[torch.Tensor] = None,
    ):
-        if is_mla_preprocess_enabled():
+        if is_mla_preprocess_enabled() and self.use_mla:
            # MLAPO and MLAPROLOG do save kv_cache
            save_kv_cache = False
        if self.is_dllm_model:
@@ -1773,7 +1773,7 @@ class AscendAttnBackend(AttentionBackend):
        sinks: Optional[torch.Tensor] = None,
        slopes: Optional[torch.Tensor] = None,
    ):
-        if is_mla_preprocess_enabled():
+        if is_mla_preprocess_enabled() and self.use_mla:
            # MLAPO does saving kv_cache
            save_kv_cache = False
        if topk_indices is not None: