Fix OmniDreams native CUDA architecture selection on GB300

[fangjunzhou-nv]

Summary

Previously, the native extension defaulted to 12.0a whenever the user did not explicitly set a CUDA architecture. GB300 devices report compute capability 10.3 and require an architecture-specific 10.3a build. As a result, the extension could be compiled for an incompatible target and fail to load or execute on GB300.

This PR fixed the OmniDreams single-view native extension to compile for the current GPU architecture instead of always defaulting to 12.0a.

Background

The native extension includes SageAttention-3 by default. Its FP4 path uses architecture-specific PTX instructions:

flashdreams/integrations/omnidreams/omnidreams_singleview/src/dit_streaming/kernels/sage3_attention.cu

Lines 85 to 107 in ac214dd

	inline __device__ uint32_t fp32_vec_to_e2m1_local(float2* array) {
	#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)
	uint32_t val;
	asm volatile(
	"{\n"
	".reg .b8 byte0;\n"
	".reg .b8 byte1;\n"
	".reg .b8 byte2;\n"
	".reg .b8 byte3;\n"
	"cvt.rn.satfinite.e2m1x2.f32 byte0, %2, %1;\n"
	"cvt.rn.satfinite.e2m1x2.f32 byte1, %4, %3;\n"
	"cvt.rn.satfinite.e2m1x2.f32 byte2, %6, %5;\n"
	"cvt.rn.satfinite.e2m1x2.f32 byte3, %8, %7;\n"
	"mov.b32 %0, {byte0, byte1, byte2, byte3};\n"
	"}"
	: "=r"(val)
	: "f"(array[0].x), "f"(array[0].y), "f"(array[1].x), "f"(array[1].y),
	"f"(array[2].x), "f"(array[2].y), "f"(array[3].x), "f"(array[3].y));
	return val;
	#else
	return 0;
	#endif
	}

SageAttention’s own build configuration compiles compute capability 12.0 devices with:

-gencode=arch=compute_120a,code=sm_120a

The a suffix enables the architecture-specific features required by these FP4 kernels.

[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

The PR changes OmniDreams native-extension architecture selection so validated SM120a GPUs explicitly target 12.0a, while GB300 and other devices use PyTorch’s default selection with SageAttention-3 stubs.

• Adds conservative SM120a device detection and architecture-aware extension cache keys.
• Disables architecture-specific SageAttention-3 sources unless the effective target is exactly 12.0a.
• Documents the selection and override behavior and expands native-loader tests.

Confidence Score: 5/5

The PR appears safe to merge.

No blocking failure remains.

Important Files Changed

Filename	Overview
integrations/omnidreams/omnidreams/native/omnidreams_singleview.py	Introduces conservative device detection, architecture-sensitive SageAttention-3 source selection, scoped build configuration, and architecture-isolated extension caching.
integrations/omnidreams/tests/test_omnidreams_singleview_native.py	Updates default-build expectations and adds coverage for device detection, override precedence, SageAttention-3 gating, and cache isolation.
integrations/omnidreams/README.md	Documents automatic architecture selection, SageAttention-3 stub behavior, overrides, and cache separation.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
  A[Load native extension] --> B{Architecture override set?}
  B -->|TORCH_CUDA_ARCH_LIST| C[Use PyTorch override]
  B -->|OmniDreams override| D[Use OmniDreams override]
  B -->|No override| E{Validated capability 12.0 device?}
  E -->|Yes| F[Target 12.0a]
  E -->|No| G[Use PyTorch default]
  C --> H{Effective target exactly 12.0a?}
  D --> H
  F --> H
  G --> H
  H -->|Yes| I[Build SageAttention-3 sources]
  H -->|No| J[Build SageAttention-3 stubs]
  I --> K[Cache by Sage3 mode and architecture identity]
  J --> K

Loading

_{Reviews (3): Last reviewed commit: "Merge branch 'main' into dev/fangjun/fix..." | Re-trigger Greptile}

[ArielG-NV]

This fix looks like a hack. If PyTorch does not recognize/list a CUDA arch something is going very wrong (something is already a hack?).

I prefer if we do not merge this and instead figure out the root cause of the problem and address it in the issue description for future maintainers of the code before we consider one off hacks.

[fangjunzhou-nv]

This fix looks like a hack. If PyTorch does not recognize/list a CUDA arch something is going very wrong (something is already a hack?).

I prefer if we do not merge this and instead figure out the root cause of the problem and address it in the issue description for future maintainers of the code before we consider one off hacks.

The omnidreams singleview was using PyTorch CUDA extension that compiled specifically targeting CUDA architecture 12.0a (which is the arch of RTX PRO 6000). What I'm adding in this PR is using torch.cuda.get_device_capability() to detect the corresponding arch when compiling PyTorch extension.

[ArielG-NV]

Question

[jarcherNV]

Regarding Ariel's comments, I dug into this a bit.

From what I can tell, major in {10, 12} → always "a" is actually correct for every chip that exists today. So this isn't a blind guess, it matches NVIDIA's published tables. That said, it's not actually guaranteed for a future chip NVIDIA ships.

Perhaps we swap _ARCH_SPECIFIC_CUDA_MAJORS for an explicit table of the exact (major, minor) pairs we've verified need a (i.e. {(10,0), (10,3), (12,0), (12,1)} or whatever's actually been tested), and raise a clear error for anything outside that table? That keeps today's fix working exactly as-is for GB300, which I think addresses the "this feels like an unbounded hack" concern without requiring a full root-cause investigation before this can land.

This has the downside of needing to be updated when we want to support newer hardware, but after checking, we seem to already have that requirement, just much more silently, and this would instead give a more clear, exact, error message.

Separately, non-blocking nit: _resolved_cuda_arch_list() (around line 499) looks like dead code after this refactor — its only caller, _scoped_cuda_arch_list, is now always invoked with an explicit cuda_arch_list argument from load_extension, so the None-fallback branch that calls it never executes, and no test references it either. Worth deleting as part of this PR so the next reader isn't left wondering which of the two "resolve the arch" functions is actually live.

Maybe there is a better way of handling this? What do you guys think @ArielG-NV and @fangjunzhou-nv ?

[ArielG-NV]

I think its about time to overhaul and unhackify our cuda compilation path (flashdreams has enough hacks as is)

[fangjunzhou-nv]

I spent some time actually dig into the root cause and it turns out the sage attention implementation in the CUDA kernel used some intrinsics that are ONLY available on sm120a machine (see the PR desc for detail).

At this point I've just spent too much time on this. My understanding is we want to ditch the custom CUDA kernel for omnidreams interactive drive at some point. We either use the new flashdreams.accelerated Triton kernel library I'm currently building or even if we want to use CUDA we want to do it properly.

I was originally submitting this PR just as a in place minor fix to get omnidreams CUDA kernel working on GB300 to get some perf numbers.

I do think if we also want to gather the perf numbers for omnidreams CUDA kernel in our incoming ci benchmark we need this fix to at least collect the numbers.

[fangjunzhou-nv]

Since I've found the root cause is sage3 requiring sm120a to compile, I've updated the omnidreams native to disable sage3 when targeting other platforms. This helps the omnidreams CUDA extension to compile on all platforms while keeping the SAGE3 impl compiled on sm120a platforms.

[ArielG-NV]

It looks like (given the PR desc provided) that the correct solution is to remove all hacky CUDA-Arch detection code and instead modify the Sage3 kernel (that we have inside the flashdreams code-base).

We should modify it such that there is an impl for CUDA arch sm120+ and sm100<= (macro should allow detecting at compile time the target arch being compiled for).

This way we may even get perf benefits from SageAttention-3 while removing hacks in the code, although I do not know if this is the case (I just prefer we don't start injecting hacks into code when the correct solution is cleaner).

[ArielG-NV]

(Refer to comment above)