Skip to content

fix(runtime): retry allocation after reclaiming cached pages#1438

Closed
ppodolsky wants to merge 1 commit into
tracel-ai:mainfrom
ppodolsky:fix/allocation-retry-cached-oom
Closed

fix(runtime): retry allocation after reclaiming cached pages#1438
ppodolsky wants to merge 1 commit into
tracel-ai:mainfrom
ppodolsky:fix/allocation-retry-cached-oom

Conversation

@ppodolsky

Copy link
Copy Markdown

Summary

When a new page allocation reports BufferTooBig, reclaim every unused cached page and retry the allocation once before returning the error.

Autotuning can leave one-off pages cached across several size buckets. A later allocation can fail at the device boundary even though those pages have no live slices. The normal reserve path previously surfaced that failure without trying the memory manager's existing explicit cleanup.

The retry also checks the selected pool again after cleanup, since reclamation can make a suitable page immediately reusable without another device allocation.

Validation

  • added allocation_oom_reclaims_cached_pages_and_retries, using a capacity-limited storage implementation to reproduce the cached-page failure deterministically
  • cargo fmt --all -- --check
  • cargo test -p cubecl-runtime allocation_oom_reclaims_cached_pages_and_retries

Compatibility

No public API or successful-allocation behavior changes. Cleanup and retry happen only after the underlying storage returns BufferTooBig.

Downstream validation

This fixes the allocation failure encountered by Hermes during repeated CUDA autotuning/training runs. No CubeK or Burn source change is required for this runtime-only fix; both consume it through their CubeCL revision.

@ppodolsky

Copy link
Copy Markdown
Author

Consolidated into #1440 so CubeCL has one upstream PR for the downstream accelerator revision. The allocation-retry commit and deterministic regression test are now included on #1440's branch.

@ppodolsky ppodolsky closed this Jul 18, 2026
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