Fix GC-cycle leak: drop frame self-references in CuTe DSL decorators - #3423
Open
thakkarV wants to merge 2 commits into
Open
Fix GC-cycle leak: drop frame self-references in CuTe DSL decorators#3423thakkarV wants to merge 2 commits into
thakkarV wants to merge 2 commits into
Conversation
`BaseDSL.jit`, `BaseDSL.kernel`, the `CuTeDSL.kernel` /
`CuteExperimentalDSL.kernel` overrides, and `DSLBaseError.__init__` bind
`inspect.currentframe()` to a local variable. A frame whose locals
reference the frame itself forms a reference cycle, so the frame can only
be reclaimed by the cyclic garbage collector -- and through `f_back` it
keeps the entire caller stack, including every frame's locals, alive
until a collection runs.
For applications that disable automatic gc (common in CUDA-graph training
loops), every `@cute.jit` / `@cute.kernel` decoration therefore leaks
whatever the calling stack held at decoration time. When kernel-bearing
modules are imported lazily inside a model's forward pass, that includes
live activation tensors: upgrading from 4.4.2 (whose decorators used the
cycle-free one-expression `inspect.currentframe().f_back` spelling) to
4.6.1 cost a measured +3-5 GiB of peak GPU memory per rank in a large
gc-disabled training job, with `CuTeDSL.kernel` frames pinning the
importing stack.
Fix by deleting the local frame reference before returning, following the
existing pattern in `_mlir_helpers/op.py:_get_caller_frame_info()`.
Minimal repro (retains 1 GiB until gc.collect() without this fix):
import gc, torch
import cutlass.cute as cute
gc.disable()
def scope():
big = torch.randn(1 << 28, device="cuda") # 1 GiB, local
@cute.kernel
def knl(): pass
scope()
torch.cuda.synchronize()
print(torch.cuda.memory_allocated()) # 1 GiB before fix, 0 after
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Fixes #3421 and adds a regression test.