fix(npu): align colocate with pr285 container validation

Split TMS env for NPU torch vs GPU preload, set PYTORCH_NPU_ALLOC_CONF=expandable_segments:False on colocate, guard empty_cache under TMS, and add vllm/vllm-ascend patches for flash_attn and NPU IPC engine API.

Signed-off-by: kaiyuan <kyxiezju@163.com>

Author: CalvinXKY

docker/npu_patch/vllm-ascend.patch

@@ -37,3 +37,18 @@ index 63ee94c..4e81943 100644
          self.available_kv_cache_memory_bytes = (
              self.requested_memory - profile_result.non_kv_cache_memory - npugraph_memory_estimate_applied
          )
+diff --git a/vllm_ascend/distributed/weight_transfer/npu_ipc_engine.py b/vllm_ascend/distributed/weight_transfer/npu_ipc_engine.py
+index 0000000..0000000 100644
+--- a/vllm_ascend/distributed/weight_transfer/npu_ipc_engine.py
++++ b/vllm_ascend/distributed/weight_transfer/npu_ipc_engine.py
+@@ -116,8 +116,8 @@ class NPUIPCWeightTransferEngine(WeightTransferEngine[NPUIPCWeightTransferInitIn
+     init_info_cls = NPUIPCWeightTransferInitInfo
+     update_info_cls = NPUIPCWeightTransferUpdateInfo
+ 
+-    def __init__(self, config: WeightTransferConfig, parallel_config: ParallelConfig) -> None:
+-        super().__init__(config, parallel_config)
++    def __init__(self, config: WeightTransferConfig, parallel_config: ParallelConfig, model: torch.nn.Module = None) -> None:
++        super().__init__(config, parallel_config, model)
+ 
+     def parse_update_info(self, update_dict: dict[str, Any]) -> NPUIPCWeightTransferUpdateInfo:
+         """Parse update dict, deserializing pickled IPC handles if present.

docker/npu_patch/vllm.patch

@@ -11,3 +11,16 @@ index cb61bcabd..5c076d521 100644
 
          # Cache the new tokens. Preempted requests should be skipped.
          if status_before_update == RequestStatus.RUNNING:
+diff --git a/vllm/model_executor/layers/rotary_embedding/common.py b/vllm/model_executor/layers/rotary_embedding/common.py
+index 0000000..0000000 100644
+--- a/vllm/model_executor/layers/rotary_embedding/common.py
++++ b/vllm/model_executor/layers/rotary_embedding/common.py
+@@ -135,7 +135,9 @@ class ApplyRotaryEmb(CustomOp):
+         self.enable_fp32_compute = enable_fp32_compute
+ 
+         self.apply_rotary_emb_flash_attn = None
+-        if not current_platform.is_cpu() and find_spec("flash_attn") is not None:
++        if not current_platform.is_cpu() and find_spec("flash_attn") is not None and not hasattr(current_platform, "is_npu"):
+             from flash_attn.ops.triton.rotary import apply_rotary
+ 
+             self.apply_rotary_emb_flash_attn = apply_rotary

vime/backends/megatron_utils/actor.py

@@ -12,6 +12,17 @@
 
 if is_npu():
     from mindspeed.megatron_adaptor import repatch
+
+    _orig_npu_empty_cache = torch.npu.empty_cache
+
+    def _safe_empty_cache():
+        try:
+            _orig_npu_empty_cache()
+        except RuntimeError:
+            pass
+
+    torch.npu.empty_cache = _safe_empty_cache
+    torch.cuda.empty_cache = _safe_empty_cache
 from megatron.core import mpu
 from torch_memory_saver import torch_memory_saver
 from transformers import AutoConfig, AutoTokenizer

vime/backends/vllm_utils/vllm_engine.py

@@ -366,7 +366,7 @@ def build_vllm_subprocess_env(server_args: dict[str, Any]) -> dict[str, str]:
         if cann_python_path is not None:
             prepend_pythonpath(env, cann_python_path)
         if getattr(args, "colocate", False):
-            env.pop("PYTORCH_NPU_ALLOC_CONF", None)
+            env["PYTORCH_NPU_ALLOC_CONF"] = "expandable_segments:False"
     else:
         env["CUDA_VISIBLE_DEVICES"] = server_args["visible_devices"]
     env.setdefault("VLLM_SERVER_DEV_MODE", "1")

vime/ray/actor_group.py

@@ -63,34 +63,34 @@ def _allocate_gpus_for_actor(self, pg, num_gpus_per_actor):
         if self.args.offload_train and self.args.train_backend == "megatron":
             import torch_memory_saver
 
-            for path in [
-                "torch_memory_saver_hook_mode_preload_cu12.abi3.so",
-                "torch_memory_saver_hook_mode_preload.abi3.so",
-            ]:
-                dynlib_path = os.path.join(
-                    os.path.dirname(os.path.dirname(torch_memory_saver.__file__)),
-                    path,
-                )
-                if os.path.exists(dynlib_path):
-                    break
-            else:
-                raise FileNotFoundError(
-                    "Cannot find torch_memory_saver dynamic library. Please make sure torch_memory_saver is properly installed."
-                )
-
-            env_vars["LD_PRELOAD"] = dynlib_path
-            env_vars["TMS_INIT_ENABLE"] = "1"
-            env_vars["TMS_INIT_ENABLE_CPU_BACKUP"] = "1"
-
             if is_npu():
                 env_vars["TMS_HOOK_MODE"] = "torch"
                 env_vars["TMS_REGION_TAG"] = "training"
                 env_vars["TMS_ENABLE_CPU_BACKUP"] = "1"
                 if self.args.colocate:
-                    env_vars.pop("PYTORCH_NPU_ALLOC_CONF", None)
+                    env_vars["PYTORCH_NPU_ALLOC_CONF"] = "expandable_segments:False"
                 cann_python_path = get_cann_python_site_packages()
                 if cann_python_path is not None:
                     prepend_pythonpath(env_vars, cann_python_path)
+            else:
+                for path in [
+                    "torch_memory_saver_hook_mode_preload_cu12.abi3.so",
+                    "torch_memory_saver_hook_mode_preload.abi3.so",
+                ]:
+                    dynlib_path = os.path.join(
+                        os.path.dirname(os.path.dirname(torch_memory_saver.__file__)),
+                        path,
+                    )
+                    if os.path.exists(dynlib_path):
+                        break
+                else:
+                    raise FileNotFoundError(
+                        "Cannot find torch_memory_saver dynamic library. Please make sure torch_memory_saver is properly installed."
+                    )
+
+                env_vars["LD_PRELOAD"] = dynlib_path
+                env_vars["TMS_INIT_ENABLE"] = "1"
+                env_vars["TMS_INIT_ENABLE_CPU_BACKUP"] = "1"
 
         # We cannot do routing replay for critic.
         if self.args.use_routing_replay and self.role == "actor":

vime/ray/rollout.py

@@ -131,7 +131,7 @@ def start_engines(self, port_cursors: dict[int, int] | None = None) -> tuple[lis
                 if cann_python_path is not None:
                     prepend_pythonpath(env_vars, cann_python_path)
                 if self.args.colocate:
-                    env_vars["PYTORCH_NPU_ALLOC_CONF"] = ""
+                    env_vars["PYTORCH_NPU_ALLOC_CONF"] = "expandable_segments:False"
                 env_vars["VLLM_USE_AOT_COMPILE"] = "0"
             rollout_engine = RolloutRayActor.options(
                 num_cpus=num_cpus,

vime/utils/memory_utils.py

@@ -15,7 +15,8 @@ def clear_memory(clear_host_memory: bool = False):
     else:
         torch.cuda.synchronize()
     gc.collect()
-    torch.cuda.empty_cache()
+    if not is_npu():
+        torch.cuda.empty_cache()
     if is_npu():
         try:
             torch.npu.empty_cache()