Add ctx manager: caching_allocator_disabled to temporarily disable CCA by ColinPeppler · Pull Request #177418 · pytorch/pytorch

ColinPeppler · 2026-03-13T20:10:57Z

Why

An IMA debugging aid to specifically disable CCA on a targeted block of code.
Another option is PYTORCH_NO_CUDA_MEMORY_CACHING=1 but that is set globally.

Usually I'd do this.

torch.cuda.caching_allocator_enable(False)
try:
    ...
finally: # make sure to clean up even on exception
    torch.cuda.caching_allocator_enable(True)

What

Add a utility that

Disables CUDA caching allocator (CCA) when entering the block.
Restores the CCA state when exiting the block (even on exceptions).

with torch.cuda.caching_allocator_disabled():
    ...

Stack from ghstack (oldest at bottom):

cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @ipiszy @kadeng @muchulee8 @amjames @chauhang @aakhundov @coconutruben @jataylo

[ghstack-poisoned]

pytorch-bot · 2026-03-13T20:11:02Z

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/177418

📄 Preview Python docs built from this PR
📄 Preview C++ docs built from this PR
❓ Need help or want to give feedback on the CI? Visit the bot commands wiki

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit 47570c7 with merge base a345892 ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

pytorch-bot · 2026-03-13T20:11:05Z

This PR needs a `release notes:` label

If your changes are user facing and intended to be a part of release notes, please use a label starting with release notes:.

If not, please add the topic: not user facing label.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "topic: not user facing"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.

ghstack-source-id: 166399e Pull Request resolved: #177418

… disable CCA" Useful for IMA debugging. Usually I do this. ``` # Disable CCA ... # Enable CCA ``` Other option is `PYTORCH_NO_CUDA_MEMORY_CACHING=1` but sometimes I like to set it in the script. cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx ipiszy kadeng muchulee8 amjames chauhang aakhundov coconutruben jataylo [ghstack-poisoned]

ghstack-source-id: 5498f42 Pull Request resolved: #177418

ColinPeppler · 2026-03-18T01:41:58Z

Hi @eee4017, can I get your review whenever you get a chance? Thanks!

… disable CCA" ### Why - An IMA debugging aid to specifically disable CCA on a targeted block of code. - Another option is `PYTORCH_NO_CUDA_MEMORY_CACHING=1` but that is set globally. Usually I'd do this. ``` torch.cuda.caching_allocator_enable(False) try: ... finally: # make sure to clean up even on exception torch.cuda.caching_allocator_enable(True) ``` ### What Add a utility that - Disables CUDA caching allocator (CCA) when entering the block. - Restores the CCA state when exiting the block (even on exceptions). ``` with torch.cuda.caching_allocator_disabled(): ... ``` cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx ipiszy kadeng muchulee8 amjames chauhang aakhundov coconutruben jataylo [ghstack-poisoned]

ghstack-source-id: e7ecf46 Pull Request resolved: #177418

eee4017

Since caching_allocator_enable already exists and allocate() bakes the correct deleter (uncached_delete vs local_raw_delete) into each pointer at allocation time, tensors allocated inside the disabled region will always free correctly regardless of whether the allocator is re-enabled. This PR just add a context manager that is a straightforward save/restore wrapper.

ColinPeppler · 2026-03-19T03:05:13Z

@pytorchbot merge

pytorchmergebot · 2026-03-19T03:08:24Z

Merge failed

Reason: Approvers from one of the following sets are needed:

superuser (pytorch/metamates)
Core Reviewers (mruberry, lezcano, Skylion007, ngimel, peterbell10, ...)
Core Maintainers (soumith, gchanan, ezyang, malfet, albanD, ...)

Details for Dev Infra team

Raised by workflow job

Failing merge rule: Core Maintainers

pytorchmergebot · 2026-03-30T23:23:48Z

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here

ghstack-source-id: 82b385f Pull Request resolved: #177418

… disable CCA" ### Why - An IMA debugging aid to specifically disable CCA on a targeted block of code. - Another option is `PYTORCH_NO_CUDA_MEMORY_CACHING=1` but that is set globally. Usually I'd do this. ``` torch.cuda.caching_allocator_enable(False) try: ... finally: # make sure to clean up even on exception torch.cuda.caching_allocator_enable(True) ``` ### What Add a utility that - Disables CUDA caching allocator (CCA) when entering the block. - Restores the CCA state when exiting the block (even on exceptions). ``` with torch.cuda.caching_allocator_disabled(): ... ``` cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx ipiszy kadeng muchulee8 amjames chauhang aakhundov coconutruben jataylo [ghstack-poisoned]

ghstack-source-id: 585da14 Pull Request resolved: #177418

ColinPeppler · 2026-04-03T23:55:15Z

@pytorchbot merge

pytorchmergebot · 2026-04-03T23:58:07Z

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here

pytorchmergebot · 2026-04-04T01:55:35Z

Merge failed

Reason: 1 mandatory check(s) failed. The first few are:

pull / linux-jammy-py3.10-gcc11 / test (distributed, 2, 3, lf.linux.2xlarge.amx)

Dig deeper by viewing the failures on hud

Details for Dev Infra team

Raised by workflow job

Failing merge rule: Core Maintainers

… disable CCA" ### Why - An IMA debugging aid to specifically disable CCA on a targeted block of code. - Another option is `PYTORCH_NO_CUDA_MEMORY_CACHING=1` but that is set globally. Usually I'd do this. ``` torch.cuda.caching_allocator_enable(False) try: ... finally: # make sure to clean up even on exception torch.cuda.caching_allocator_enable(True) ``` ### What Add a utility that - Disables CUDA caching allocator (CCA) when entering the block. - Restores the CCA state when exiting the block (even on exceptions). ``` with torch.cuda.caching_allocator_disabled(): ... ``` cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx ipiszy kadeng muchulee8 amjames chauhang aakhundov coconutruben jataylo [ghstack-poisoned]

ghstack-source-id: 9695999 Pull Request resolved: #177418

… disable CCA" ### Why - An IMA debugging aid to specifically disable CCA on a targeted block of code. - Another option is `PYTORCH_NO_CUDA_MEMORY_CACHING=1` but that is set globally. Usually I'd do this. ``` torch.cuda.caching_allocator_enable(False) try: ... finally: # make sure to clean up even on exception torch.cuda.caching_allocator_enable(True) ``` ### What Add a utility that - Disables CUDA caching allocator (CCA) when entering the block. - Restores the CCA state when exiting the block (even on exceptions). ``` with torch.cuda.caching_allocator_disabled(): ... ``` cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx ipiszy kadeng muchulee8 amjames chauhang aakhundov coconutruben jataylo [ghstack-poisoned]

ghstack-source-id: a4a00fc Pull Request resolved: #177418

… disable CCA" ### Why - An IMA debugging aid to specifically disable CCA on a targeted block of code. - Another option is `PYTORCH_NO_CUDA_MEMORY_CACHING=1` but that is set globally. Usually I'd do this. ``` torch.cuda.caching_allocator_enable(False) try: ... finally: # make sure to clean up even on exception torch.cuda.caching_allocator_enable(True) ``` ### What Add a utility that - Disables CUDA caching allocator (CCA) when entering the block. - Restores the CCA state when exiting the block (even on exceptions). ``` with torch.cuda.caching_allocator_disabled(): ... ``` cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx ipiszy kadeng muchulee8 amjames chauhang aakhundov coconutruben jataylo [ghstack-poisoned]

ghstack-source-id: cba1d6c Pull Request resolved: #177418

ColinPeppler · 2026-04-07T16:57:01Z

@pytorchbot merge

pytorchmergebot · 2026-04-07T16:59:35Z

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here

pytorch#177418) ### Why - An IMA debugging aid to specifically disable CCA on a targeted block of code. - Another option is `PYTORCH_NO_CUDA_MEMORY_CACHING=1` but that is set globally. Usually I'd do this. ``` torch.cuda.caching_allocator_enable(False) try: ... finally: # make sure to clean up even on exception torch.cuda.caching_allocator_enable(True) ``` ### What Add a utility that - Disables CUDA caching allocator (CCA) when entering the block. - Restores the CCA state when exiting the block (even on exceptions). ``` with torch.cuda.caching_allocator_disabled(): ... ``` Pull Request resolved: pytorch#177418 Approved by: https://github.com/eee4017, https://github.com/laithsakka ghstack dependencies: pytorch#177308

…agnostic 8 AOTInductor tests fail on XPU because `caching_allocator_disabled()` (Intorduced by #177418) from `torch.cuda.memory` calls `torch._C._cuda_cudaCachingAllocator_is_enabled()` which doesn't exist in XPU-only builds. Replace the direct import of `torch.cuda.caching_allocator_disabled` with a device-aware wrapper that delegates to the CUDA implementation on CUDA builds and acts as a no-op on other GPU backends (XPU, etc.). ghstack-source-id: 81004c9 Pull-Request: #179659

…agnostic 8 AOTInductor tests fail on XPU because `caching_allocator_disabled()` (Intorduced by #177418) from `torch.cuda.memory` calls `torch._C._cuda_cudaCachingAllocator_is_enabled()` which doesn't exist in XPU-only builds. Replace the direct import of `torch.cuda.caching_allocator_disabled` with a device-aware wrapper that delegates to the CUDA implementation on CUDA builds and acts as a no-op on other GPU backends (XPU, etc.). ghstack-source-id: caf3d25 Pull-Request: #179659

…agnostic 8 AOTInductor tests fail on XPU because `caching_allocator_disabled()` (Intorduced by #177418) from `torch.cuda.memory` calls `torch._C._cuda_cudaCachingAllocator_is_enabled()` which doesn't exist in XPU-only builds. Replace the direct import of `torch.cuda.caching_allocator_disabled` with a device-aware wrapper that delegates to the CUDA implementation on CUDA builds and acts as a no-op on other GPU backends (XPU, etc.). ghstack-source-id: 340ee3f Pull-Request: #179659

…agnostic 8 AOTInductor tests fail on XPU because `caching_allocator_disabled()` (Intorduced by #177418) from `torch.cuda.memory` calls `torch._C._cuda_cudaCachingAllocator_is_enabled()` which doesn't exist in XPU-only builds. Replace the direct import of `torch.cuda.caching_allocator_disabled` with a device-aware wrapper that delegates to the CUDA implementation on CUDA builds and acts as a no-op on other GPU backends (XPU, etc.). ghstack-source-id: f567e0a Pull-Request: #179659

…agnostic 8 AOTInductor tests fail on XPU because `caching_allocator_disabled()` (Intorduced by #177418) from `torch.cuda.memory` calls `torch._C._cuda_cudaCachingAllocator_is_enabled()` which doesn't exist in XPU-only builds. Replace the direct import of `torch.cuda.caching_allocator_disabled` with a device-aware wrapper that delegates to the CUDA implementation on CUDA builds and acts as a no-op on other GPU backends (XPU, etc.). ghstack-source-id: a0c1904 Pull-Request: #179659

…agnostic 8 AOTInductor tests fail on XPU because `caching_allocator_disabled()` (Intorduced by #177418) from `torch.cuda.memory` calls `torch._C._cuda_cudaCachingAllocator_is_enabled()` which doesn't exist in XPU-only builds. Replace the direct import of `torch.cuda.caching_allocator_disabled` with a device-aware wrapper that delegates to the CUDA implementation on CUDA builds and acts as a no-op on other GPU backends (XPU, etc.). ghstack-source-id: 5a4f221 Pull-Request: #179659

Refactor caching allocator disable/enable into context manager

e591e93

[ghstack-poisoned]

ColinPeppler requested review from Aidyn-A, eqy and syed-ahmed as code owners March 13, 2026 20:10

This was referenced Mar 13, 2026

Slicing with backed should produce backed output when possible #175819

Closed

Support negative index slicing with backed symints #177308

Closed

pytorch-bot bot added ciflow/inductor module: inductor labels Mar 13, 2026

ColinPeppler added a commit that referenced this pull request Mar 13, 2026

Refactor caching allocator disable/enable into context manager

a6c8061

ghstack-source-id: 166399e Pull Request resolved: #177418

ColinPeppler changed the title ~~Refactor caching allocator disable/enable into context manager~~ Add ctx manager: caching_allocator_disabled to temporarily disable CCA Mar 13, 2026

ColinPeppler mentioned this pull request Mar 13, 2026

Add ctx manager: caching_allocator_disabled to temporarily disable CCA #177354

Closed

ColinPeppler requested review from eee4017 and ezyang March 13, 2026 20:13

ColinPeppler added the topic: not user facing topic category label Mar 13, 2026

ColinPeppler added a commit that referenced this pull request Mar 14, 2026

Refactor caching allocator disable/enable into context manager

e96f6ed

ghstack-source-id: 5498f42 Pull Request resolved: #177418

ColinPeppler added a commit that referenced this pull request Mar 18, 2026

Refactor caching allocator disable/enable into context manager

0521f30

ghstack-source-id: e7ecf46 Pull Request resolved: #177418

eee4017 approved these changes Mar 18, 2026

View reviewed changes

pytorch-bot bot added the ciflow/trunk Trigger trunk jobs on your pull request label Mar 19, 2026

pytorchmergebot added the merging label Mar 19, 2026

pytorchmergebot removed the merging label Mar 19, 2026

laithsakka approved these changes Mar 30, 2026

View reviewed changes

pytorchmergebot added the merging label Mar 30, 2026

ColinPeppler added a commit that referenced this pull request Apr 2, 2026

Refactor caching allocator disable/enable into context manager

b949519

ghstack-source-id: 82b385f Pull Request resolved: #177418

ColinPeppler added a commit that referenced this pull request Apr 3, 2026

Refactor caching allocator disable/enable into context manager

6e8c490

ghstack-source-id: 585da14 Pull Request resolved: #177418

pytorchmergebot added the merging label Apr 3, 2026

pytorchmergebot removed the merging label Apr 4, 2026

ColinPeppler added a commit that referenced this pull request Apr 6, 2026

Refactor caching allocator disable/enable into context manager

2da1e92

ghstack-source-id: 9695999 Pull Request resolved: #177418

ColinPeppler added a commit that referenced this pull request Apr 6, 2026

Refactor caching allocator disable/enable into context manager

29bc798

ghstack-source-id: a4a00fc Pull Request resolved: #177418

ColinPeppler added a commit that referenced this pull request Apr 6, 2026

Refactor caching allocator disable/enable into context manager

08d8152

ghstack-source-id: cba1d6c Pull Request resolved: #177418

pytorchmergebot added the merging label Apr 7, 2026

pytorchmergebot added the Merged label Apr 7, 2026

pytorchmergebot closed this in 0fb8b82 Apr 7, 2026

pytorchmergebot removed the merging label Apr 7, 2026

etaf mentioned this pull request Apr 8, 2026

[WIP][xpu][fix] test_aot_inductor: Make caching_allocator_disabled device-agnostic #179659

Draft

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add ctx manager: caching_allocator_disabled to temporarily disable CCA#177418

Add ctx manager: caching_allocator_disabled to temporarily disable CCA#177418
ColinPeppler wants to merge 8 commits intogh/ColinPeppler/6/basefrom
gh/ColinPeppler/6/head

ColinPeppler commented Mar 13, 2026 •

edited

Loading

Uh oh!

pytorch-bot bot commented Mar 13, 2026 •

edited

Loading

Uh oh!

pytorch-bot bot commented Mar 13, 2026

Uh oh!

ColinPeppler commented Mar 18, 2026

Uh oh!

eee4017 left a comment

Uh oh!

ColinPeppler commented Mar 19, 2026

Uh oh!

pytorchmergebot commented Mar 19, 2026

Uh oh!

pytorchmergebot commented Mar 30, 2026

Uh oh!

ColinPeppler commented Apr 3, 2026

Uh oh!

pytorchmergebot commented Apr 3, 2026

Uh oh!

pytorchmergebot commented Apr 4, 2026

Uh oh!

ColinPeppler commented Apr 7, 2026

Uh oh!

pytorchmergebot commented Apr 7, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Conversation

ColinPeppler commented Mar 13, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Why

What

Uh oh!

pytorch-bot bot commented Mar 13, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/177418

✅ No Failures

Uh oh!

pytorch-bot bot commented Mar 13, 2026

This PR needs a release notes: label

Uh oh!

ColinPeppler commented Mar 18, 2026

Uh oh!

eee4017 left a comment

Choose a reason for hiding this comment

Uh oh!

ColinPeppler commented Mar 19, 2026

Uh oh!

pytorchmergebot commented Mar 19, 2026

Merge failed

Uh oh!

pytorchmergebot commented Mar 30, 2026

Merge started

Uh oh!

ColinPeppler commented Apr 3, 2026

Uh oh!

pytorchmergebot commented Apr 3, 2026

Merge started

Uh oh!

pytorchmergebot commented Apr 4, 2026

Merge failed

Uh oh!

ColinPeppler commented Apr 7, 2026

Uh oh!

pytorchmergebot commented Apr 7, 2026

Merge started

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

ColinPeppler commented Mar 13, 2026 •

edited

Loading

pytorch-bot bot commented Mar 13, 2026 •

edited

Loading

This PR needs a `release notes:` label