fix(ascend): patch accelerator cache cleanup

Author: hooray03

tests/unit_tests/dispatch/test_ascend_patch.py

@@ -0,0 +1,21 @@
+# Copyright (c) 2026 BAAI. All rights reserved.
+
+from types import SimpleNamespace
+
+import torch
+
+from vllm_fl.dispatch.backends.vendor.ascend.patch import (
+    patch_accelerator_empty_cache,
+)
+
+
+def test_patch_accelerator_empty_cache(monkeypatch):
+    replacement = lambda: None
+    monkeypatch.setattr(
+        torch, "npu", SimpleNamespace(empty_cache=replacement), raising=False
+    )
+    monkeypatch.setattr(torch.accelerator, "empty_cache", lambda: None)
+
+    patch_accelerator_empty_cache()
+
+    assert torch.accelerator.empty_cache is replacement

vllm_fl/dispatch/backends/vendor/ascend/__init__.py

@@ -5,8 +5,9 @@
 """
 
 from .ascend import AscendBackend
-from .patch import patch_mamba_config
+from .patch import patch_accelerator_empty_cache, patch_mamba_config
 
+patch_accelerator_empty_cache()
 patch_mamba_config()
 
 __all__ = ["AscendBackend"]

vllm_fl/dispatch/backends/vendor/ascend/patch.py

@@ -9,6 +9,23 @@
 logger = logging.getLogger(__name__)
 _patches_applied = False
 
+
+def patch_accelerator_empty_cache():
+    """Redirect generic accelerator cache cleanup to torch-npu.
+
+    ``torch.accelerator.empty_cache`` uses an incompatible allocator on NPU.
+    Patch it during Ascend backend import so cleanup also works in the engine
+    parent process, not only in worker processes.
+    """
+    npu = getattr(torch, "npu", None)
+    accelerator = getattr(torch, "accelerator", None)
+    if npu is None or accelerator is None or not hasattr(npu, "empty_cache"):
+        return
+
+    accelerator.empty_cache = npu.empty_cache
+    logger.info("Patched torch.accelerator.empty_cache for Ascend NPU")
+
+
 def apply_ascend_patches():
     """Apply all Ascend-specific patches."""
     global _patches_applied