From f8029945c32cc53ae5b6f6479f8a5ab4f731cb69 Mon Sep 17 00:00:00 2001 From: Li Wang Date: Tue, 17 Jun 2025 21:25:40 +0800 Subject: [PATCH] [Bugfix] Remove cuda related lines and add additional pip mirror (#1252) ### What this PR does / why we need it? - For npu environment, we should use `PYTORCH_NPU_ALLOC_CONF ` rather than `PYTORCH_CUDA_ALLOC_CONF` - Add `PIP_EXTRA_INDEX_URL` to make nightly_benchmarks happy --------- Signed-off-by: wangli --- .github/workflows/nightly_benchmarks.yaml | 2 ++ vllm_ascend/device_allocator/camem.py | 2 +- 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/.github/workflows/nightly_benchmarks.yaml b/.github/workflows/nightly_benchmarks.yaml index 6e03c1e..2b9c062 100644 --- a/.github/workflows/nightly_benchmarks.yaml +++ b/.github/workflows/nightly_benchmarks.yaml @@ -167,6 +167,8 @@ jobs: cp -r benchmarks/* /github/home/benchmarks/ - name: Run benchmark iteration + env: + PIP_EXTRA_INDEX_URL: https://mirrors.huaweicloud.com/ascend/repos/pypi if: github.event_name != 'pull_request' run: | while IFS= read -r line || [[ -n "$line" ]]; do diff --git a/vllm_ascend/device_allocator/camem.py b/vllm_ascend/device_allocator/camem.py index 7592a04..7156602 100644 --- a/vllm_ascend/device_allocator/camem.py +++ b/vllm_ascend/device_allocator/camem.py @@ -144,7 +144,7 @@ class CaMemAllocator: return CaMemAllocator.instance def __init__(self): - conf = os.environ.get("PYTORCH_CUDA_ALLOC_CONF", "") + conf = os.environ.get("PYTORCH_NPU_ALLOC_CONF", "") assert "expandable_segments:True" not in conf, \ ("Expandable segments are not compatible with memory pool. " "Please track https://github.com/pytorch/pytorch/issues/147851 "