Skip to content

[Fix] offload .to(cuda) interception + configurable nd_tiling max_tiles and prefetch lookahead - #61

Open
cennn wants to merge 1 commit into
mainfrom
fix/offload-nd-tiling-for-disagg
Open

[Fix] offload .to(cuda) interception + configurable nd_tiling max_tiles and prefetch lookahead#61
cennn wants to merge 1 commit into
mainfrom
fix/offload-nd-tiling-for-disagg

Conversation

@cennn

@cennn cennn commented Aug 21, 2026

Copy link
Copy Markdown
Collaborator

Summary

Three bug fixes and one config knob exposed during disaggregated multi-model inference on 31 GiB GPUs (RTX 5090):

  1. _patch_cpu_offload_apply now intercepts .to(cuda:X) in addition to .cuda()

    • Callers using model.to(device) (common in multi-model pipelines) bypassed the offload hook entirely, putting all weights on GPU and OOMing on 31 GiB cards.
    • Fix: probe the .to() lambda with a small CPU tensor to detect CUDA target.
  2. nd_tiling max_tiles capped at 2 unconditionally

    • max_tiles=3 overflows the CUDA z-grid limit (65535) on conv-heavy dynamic-shape graphs (e.g. 1080p VAE decode). Grid3D has no z-overflow handling.
    • Previously only guarded on PT>=2.12; now capped on all versions.
  3. max_tiles exposed as MAGI_COMPILE_PASS_CONFIG__ND_TILING_MAX_TILES

    • Allows experimentation without code edits. Default remains 2.
  4. Prefetch lookahead configurable via OffloadConfig.max_prefetch_lookahead

    • Hard-coded 2 layers of lookahead costs GPU memory a small card may not have. Setting to 0 disables prefetch entirely.

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