Skip to content

Adapt Online defragmenter for torch compile#986

Open
jwieczorekhabana wants to merge 1 commit intovllm-project:mainfrom
jwieczorekhabana:janwiecz/compile-defragmenter
Open

Adapt Online defragmenter for torch compile#986
jwieczorekhabana wants to merge 1 commit intovllm-project:mainfrom
jwieczorekhabana:janwiecz/compile-defragmenter

Conversation

@jwieczorekhabana
Copy link
Contributor

@jwieczorekhabana jwieczorekhabana commented Feb 17, 2026

[GAUDISW-244575] Reapply Adapt OnlineDefragmenter and CacheSwapUtils for t.compile
Because of the double entrypoint of CacheSwapUtils (forward and swap functions)
torch.compile would process module and forward function while swap's self
would refer to unwrapped module. That results in the function not being run as compiled

Changes made in this patch:

  • Hide CacheSwapUtils entirely in OnlineDefragmenter. Let it be responsible for calling the module correctly
  • Moved warmup_defragmenter to defragmenter itself
  • Removed initialize function of OnlineDefragmenter, fully initialize object in init
  • Adapted unit tests for new implementations

…for t.compile

Because of the double entrypoint of CacheSwapUtils (forward and swap functions)
torch.compile would process module and forward function while swap's self
would refer to unwrapped module. That results in the function not being run as compiled
Changes made in this patch:

 - Hide CacheSwapUtils entirely in OnlineDefragmenter. Let it be responsible for calling the module correctly
 - Moved warmup_defragmenter to defragmenter itself
 - Removed initialize function of OnlineDefragmenter, fully initialize object in init
 - Adapted unit tests for new implementations

Signed-off-by: Jan Wieczorek <jwieczorek@habana.ai>
@jwieczorekhabana jwieczorekhabana force-pushed the janwiecz/compile-defragmenter branch from c17c9fe to 49ba149 Compare February 26, 2026 12:24
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant