[NPU] bugfix for extra device memory on Ascend (#30112)

This commit is contained in:
khalilzhk
2026-07-24 21:13:21 +08:00
committed by GitHub
parent 3d91a569ce
commit dfaf75b1a6
@@ -53,6 +53,7 @@ class NPUCudaGraphBackend(BaseCudaGraphBackend):
self._outputs: Dict[Any, Any] = {}
self._pool = None
self._device_module = cuda_graph_runner.device_module
self._device_id = self._device_module.current_device()
self._tp_group = cuda_graph_runner.model_runner.tp_group
self._capture_stream = None
self._memory_saver_adapter: Optional[Any] = TorchMemorySaverAdapter.create(
@@ -166,6 +167,7 @@ class NPUCudaGraphBackend(BaseCudaGraphBackend):
graph = self._graphs[shape_key]
def _update():
self._device_module.set_device(self._device_id)
graph.update(cpu_update_input=cpu_update_input)
thread = threading.Thread(target=_update)