Skip to content

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

Merged
jiahy0825 merged 1 commit into
mainfrom
fix/offload-nd-tiling-for-disagg
Aug 24, 2026
Merged

[Fix] offload .to(cuda) interception + configurable nd_tiling max_tiles and prefetch lookahead#61
jiahy0825 merged 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.

@cennn
cennn force-pushed the fix/offload-nd-tiling-for-disagg branch 2 times, most recently from 2f096c0 to 44a0e94 Compare August 24, 2026 08:06
@cennn
cennn force-pushed the fix/offload-nd-tiling-for-disagg branch from 44a0e94 to f6e2c6a Compare August 24, 2026 08:08

@jiahy0825 jiahy0825 left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

@jiahy0825
jiahy0825 merged commit ee85c8c into main Aug 24, 2026
8 of 10 checks passed
@jiahy0825
jiahy0825 deleted the fix/offload-nd-tiling-for-disagg branch August 24, 2026 12:35
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.

2 participants