[slavaheroes]

Replace the fixed dimension cutoff with GPU-first concatenation and CPU fallback on CUDA out-of-memory errors. Add an optional free-memory margin, a fallback warning, and correctness regression tests.

Refs: #9250

Assisted-by: OpenAI Codex
I reviewed the submitted change

Fixes #9250.

Description

Replace the fixed dimension threshold with GPU-first concatenation. If CUDA concatenation runs out of memory, emit a RuntimeWarning, concatenate the tensors on the CPU.

Add an optional concat_memory_margin argument:

• None (default): attempt GPU concatenation and fall back to the CPU on CUDA OOM.
• A non-negative integer, in bytes: compare the estimated output size plus this margin against free CUDA memory. Use CPU concatenation if there is insufficient free memory.

Add regression tests and update the docstrings and documentation.

Types of changes

• Focused MAISI tests: 54 passed, 1 skipped (requires multiple GPUs).
• Formatting and lint checks passed.
• Focused documentation build passed.
• An actual CUDA OOM under a per-process allocator limit exercised
  the fallback successfully.

For the issue's 512 × 512 × 16 benchmark, median runtime decreased
from 42.40 s to 0.440 s on an RTX 5090. Outputs were bitwise identical.
Peak PyTorch allocated memory increased from approximately
4.02 GiB to 4.25 GiB.

• Non-breaking change (fix or new feature that would not break existing functionality).
• Breaking change (fix or new feature that would cause existing functionality to change).
• New tests added to cover the changes.
• Integration tests passed locally by running ./runtests.sh -f -u --net --coverage.
• Quick tests passed locally by running ./runtests.sh --quick --unittests --disttests.
• In-line docstrings updated.
• Documentation updated, tested make html command in the docs/ folder.

[coderabbitai]

Important

Draft PR not reviewed

Draft PRs are not automatically reviewed by default.

• Trigger a manual review

To automatically review draft PRs, update your CodeRabbit configuration:

reviews:
  auto_review:
    drafts: true

Tip

🎯 Customize your risk assessments

Guide how CodeRabbit assesses merge risk and security architecture risk. Highlight sensitive changes and explain your security boundaries using Merge Risk Instructions and Security Risk Instructions in your review settings. Security architecture risk assessments require the Advanced plan or above.

✅ Automatic approvals, with your rules

Let CodeRabbit approve pull requests once required threads are resolved, pre-merge checks pass, and your risk limits are met. Enable Automatic approvals in review settings, with overrides for teams or file paths.

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

Share

• X
• Mastodon
• Reddit
• LinkedIn