fix: add cuDNN-only build option that skips the SageAttention-3 shim

[mvanhorn]

Summary

Adds a build-time option to produce a cuDNN-only OmniDreams single-view native extension that skips the SageAttention-3 Blackwell shim. The heavy Sage3 path needs >62 GB in a single cudafe++ invocation even when the manifest selects the cuDNN backend and Sage3 is never used at runtime; the opt-out builds against the existing stubs instead.

Why this matters

Reported in #334. A collaborator confirmed on 2026-07-03 that current main still compiles the Sage3 (and LightVAE FP8) sources unconditionally, and that the right fix is a build-time cuDNN-only / no-Sage3 mode that uses the existing Sage3 stubs and makes sage3_is_built() return false. A ready-made stub translation unit already exists at integrations/omnidreams/omnidreams_singleview/src/dit_streaming/kernels/sage3_attention_stub.cu (it defines sage3_is_built() -> false, sage3_is_runtime_supported() -> false, and TORCH_CHECK(false, ...) bodies for every run_sage3_* entry point), but nothing wired it into the build.

Changes

• Adds a build flag that swaps sage3_blackwell_api_shim.cu for the existing sage3_attention_stub.cu and drops the CUTLASS sm_120 kernel sources, so a cuDNN-only extension compiles without the >62 GB Sage3 compile. Default behavior (Sage3 built) is unchanged; the flag is opt-in.

Testing

Built the extension with the opt-out flag set and confirmed the Sage3 shim and CUTLASS sm_120 sources are excluded, sage3_is_built() returns false, and the cuDNN attention path runs. Default build still compiles Sage3.

Fixes #334

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[greptile-apps]

Greptile Summary

This PR adds a cuDNN-only build mode for the OmniDreams single-view native extension.

• Adds an environment flag to disable Sage3 during native extension builds.
• Swaps Sage3 shim and kernel sources for the existing Sage3 stub source when disabled.
• Gives Sage3-enabled and Sage3-disabled builds separate extension names and cache entries.
• Adds tests for flag parsing, source selection, compile flags, extension naming, and cache separation.

Confidence Score: 5/5

This looks safe to merge.

• No blocking issues found in the changed code.

Important Files Changed

Filename	Overview
integrations/omnidreams/omnidreams/native/omnidreams_singleview.py	Adds Sage3 opt-out handling, conditional source selection, mode-specific extension naming, and separate cached modules per build mode.
integrations/omnidreams/tests/test_omnidreams_singleview_native.py	Adds coverage for the new Sage3 opt-out build path and updates existing cache setup for the new per-mode cache shape.

_{Reviews (2): Last reviewed commit: "fix(omnidreams): make extension cache mo..." | Re-trigger Greptile}

[jarcherNV]

Thanks, this looks like the right direction overall. The source selection, compile flags, extension name isolation, and Sage3 availability checks all make sense to me.

I think we need one cache fix before merging, though. load_extension() still returns the first loaded module from the module-level _extension cache before checking the current OMNIDREAMS_SINGLEVIEW_DISABLE_SAGE3 mode. That means a long-running process that loads the Sage3 build first, then sets OMNIDREAMS_SINGLEVIEW_DISABLE_SAGE3=1, will keep using the Sage3-enabled binary. The reverse case has the same problem.

Could you update the cache to be keyed by the extension name/build mode, or store the loaded extension’s cache key and only return it when it matches the current mode? It would also be good to add a regression test that loads once in one mode, toggles the env var, calls load_extension() again, and verifies the second mode is loaded separately.

I think once that's done along with a rebase, we can trigger CI and attempt to merge.

[mvanhorn]

Good catch — fixed in e1bf2e0. The _extension cache is now keyed on the resolved OMNIDREAMS_SINGLEVIEW_DISABLE_SAGE3 mode, so toggling it at runtime loads the correct binary instead of returning the first-loaded module. Added a test covering both modes.