[NPU] bugfix for extra device memory on Ascend (#30112)
This commit is contained in:
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user