minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
khalilzhk
2026-07-24 21:13:21 +08:00
committed by GitHub
parent 3d91a569ce
commit dfaf75b1a6
1 changed files with 2 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/hardware_backend/npu/graph_runner/npu_cudagraph_backend.py
+2
View File
@@ -53,6 +53,7 @@ class NPUCudaGraphBackend(BaseCudaGraphBackend):
self._outputs: Dict[Any, Any] = {} self._outputs: Dict[Any, Any] = {}
self._pool = None self._pool = None
self._device_module = cuda_graph_runner.device_module 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._tp_group = cuda_graph_runner.model_runner.tp_group
self._capture_stream = None self._capture_stream = None
self._memory_saver_adapter: Optional[Any] = TorchMemorySaverAdapter.create( self._memory_saver_adapter: Optional[Any] = TorchMemorySaverAdapter.create(
@@ -166,6 +167,7 @@ class NPUCudaGraphBackend(BaseCudaGraphBackend):
graph = self._graphs[shape_key] graph = self._graphs[shape_key]
def _update(): def _update():
self._device_module.set_device(self._device_id)
graph.update(cpu_update_input=cpu_update_input) graph.update(cpu_update_input=cpu_update_input)
thread = threading.Thread(target=_update) thread = threading.Thread(target=_update)