[NPU] fix dcp break (#29140)

2026-06-24 17:33:02 +08:00
parent 73d976e375
commit 10e0bcd622
1 changed files with 1 additions and 0 deletions
@@ -177,6 +177,7 @@ class NPUMHATokenToKVPool(MHATokenToKVPool):
        k_scale: Optional[float] = None,
        v_scale: Optional[float] = None,
        layer_id_override: Optional[int] = None,
+        dcp_kv_mask: Optional[torch.Tensor] = None,
    ):
        loc, _ = unwrap_write_loc(loc_info)
        if layer_id_override is not None: