Exclude unused kernels from compilation, when cuDNN backend is selected?

[erwincoumans]

Some kernels are very heavy (seem to require >64GB of memory), and they may not be used if some manifest selects cuDNN backend.

Could we excluding the Sage Blackwell source from the build, if cuDNN is selected?

Excuses if above assumption are wrong, not very familiar with how things work, it could be more complicated:

Unfortunately it may not be a quick flag. The entire 558-byte shim is:

#undef PYBIND11_MODULE
#define PYBIND11_MODULE(name, variable) static void native_sage3_blackwell_pybind_stub(...)
#include "sageattention3_blackwell/sageattn3/blackwell/api.cu"   // ← the 62 GB compile

That one #include pulls in all of SageAttention‑3's Blackwell api.cu (the full CUTLASS sm_120 kernel library) and instantiates it. 
That's your 62 GB cudafe++.

[erwincoumans]

According to Claude: What a proper "skip Sage3" option actually takes (a real change, touching the model's hot path):

• Gate the shim's #include "...api.cu" behind a new macro, and provide stub definitions of the symbols it exports when off.
• Add #ifdef OMNIDREAMS_SINGLEVIEW_HAS_SAGE3 guards + fallbacks around every Sage3 reference in streaming_dit_bridge.cu/.h, streaming_dit_bindings.cpp, and sage3_attention.cu (so sage3_is_built() → false, dispatch never calls the missing kernels).
• Plumb a Python build flag (env/manifest, default = today's behavior) that, when set, drops the 3 Sage3 sources from _extension_sources() and omits -DHAS_SAGE3.
• Verify inference still works on the cuDNN path (you don't use Sage3 at runtime, so functionally fine — but it must link and not regress).

This could make the 62 GB compile disappear permanently — but it's an invasive, multi-file edit to vendor/research code on the model's core forward path, with real risk of a subtle break, and it needs an actual inference test to validate. That's a meaningfully bigger lift than the one-time compile.

[erwincoumans]

It seems this kernel: lightvae_fp8_attention.cu causing out-of-memory error during compilation.

[jarcherNV]

Thanks for the report. I checked current main, and this is still a valid native-build footprint issue.

The single-view native extension currently still builds the heavy Sage3 and LightVAE FP8 sources unconditionally, even when the intended runtime path is cuDNN. We do cap native build parallelism through MAX_JOBS, which helps with aggregate build memory, but it does not remove a single large compile unit.

The right fix is likely a build-time option for a cuDNN-only/no-Sage3 native extension, using the existing Sage3 stubs and making sage3_is_built() return false in that mode. We also need to confirm whether the current OOM is from the Sage3 shim or lightvae_fp8_attention.cu, since those may need separate gates.