diff --git a/integrations/omnidreams/README.md b/integrations/omnidreams/README.md index be0561bce..83ace7717 100644 --- a/integrations/omnidreams/README.md +++ b/integrations/omnidreams/README.md @@ -178,6 +178,15 @@ explicitly to opt into Sparge/SageAttention-3 experiments. Use `native_dit_sparge_hybrid_period > 1` with `"sparge"` to enable the FP8 Sparge/SageAttention-3 hybrid schedule when the extension and GPU support it. +The native extension explicitly targets `12.0a` on validated compute capability +12.0 GPUs that support the architecture-specific SageAttention-3 FP4 path. On +other GPUs, including GB300, it leaves architecture selection to PyTorch and +builds SageAttention-3 stubs so its SM120a-only FP4 instructions are excluded. +Set `OMNIDREAMS_SINGLEVIEW_CUDA_ARCH_LIST` to override this behavior, or set +`TORCH_CUDA_ARCH_LIST` to use PyTorch's standard override (which takes +precedence). Explicit `12.0a` and PyTorch-default builds use separate extension +caches so an incompatible kernel image is not reused between them. + ## Run (shared demo API) From the repository root on a CUDA machine: diff --git a/integrations/omnidreams/omnidreams/native/omnidreams_singleview.py b/integrations/omnidreams/omnidreams/native/omnidreams_singleview.py index 0e9607e84..c0d082370 100644 --- a/integrations/omnidreams/omnidreams/native/omnidreams_singleview.py +++ b/integrations/omnidreams/omnidreams/native/omnidreams_singleview.py @@ -55,10 +55,17 @@ _NATIVE_CUDA_ARCH_LIST_ENV = "OMNIDREAMS_SINGLEVIEW_CUDA_ARCH_LIST" _DISABLE_SAGE3_ENV = "OMNIDREAMS_SINGLEVIEW_DISABLE_SAGE3" _PYTORCH_CUDA_ARCH_LIST_ENV = "TORCH_CUDA_ARCH_LIST" -_DEFAULT_CUDA_ARCH_LIST = "12.0a" +_PYTORCH_DEFAULT_CUDA_ARCH_LIST = "pytorch-default" +# CUDA reports capability 12.0 without the "a" suffix, so mirror the +# conservative device allowlist used by sage3_is_runtime_supported(). +_SM120A_DEVICE_NAME_MARKERS = ( + "GeForce RTX 5090", + "RTX PRO 6000", + "RTX 6000", +) _native_build_module: ModuleType | None = None -_extension: dict[bool, ModuleType] = {} +_extension: dict[tuple[bool, str], ModuleType] = {} _extension_load_error: Exception | None = None _state_lock = threading.RLock() _dll_directory_handles: list[object] = [] @@ -366,8 +373,12 @@ def _file_sha256(path: Path) -> str: return digest.hexdigest() -def _sage3_disabled() -> bool: - return os.environ.get(_DISABLE_SAGE3_ENV, "").strip().lower() in {"1", "true"} +def _sage3_disabled(cuda_arch_list: str | None = None) -> bool: + if os.environ.get(_DISABLE_SAGE3_ENV, "").strip().lower() in {"1", "true"}: + return True + if cuda_arch_list is None: + cuda_arch_list = _effective_cuda_arch_list() + return cuda_arch_list != "12.0a" def _extension_sources() -> list[Path]: @@ -433,12 +444,20 @@ def _source_fingerprint() -> str: return digest.hexdigest() -def _extension_name(thirdparty_info: dict[str, Any]) -> str: - has_sage3 = int(not _sage3_disabled()) +def _extension_name( + thirdparty_info: dict[str, Any], + *, + cuda_arch_list: str | None = None, +) -> str: + cuda_arch_list = _cuda_arch_identity( + _effective_cuda_arch_list() if cuda_arch_list is None else cuda_arch_list + ) + has_sage3 = int(not _sage3_disabled(cuda_arch_list)) digest = hashlib.sha256() digest.update(_source_fingerprint().encode("ascii")) digest.update(json.dumps(thirdparty_info, sort_keys=True).encode("utf-8")) digest.update(f"sage3={has_sage3}".encode("ascii")) + digest.update(f"cuda_arch_list={cuda_arch_list}".encode("ascii")) return f"omnidreams_singleview_native_sage3_{has_sage3}_{digest.hexdigest()[:12]}" @@ -463,19 +482,34 @@ def _resolved_max_jobs(max_jobs: int | str | None) -> str | None: return str(min(os.cpu_count() or 1, _DEFAULT_MAX_JOBS_CAP)) -def _resolved_cuda_arch_list() -> str | None: - if os.environ.get(_PYTORCH_CUDA_ARCH_LIST_ENV): +def _detected_cuda_arch_list() -> str | None: + try: + import torch + + if not torch.cuda.is_available(): + return None + if torch.cuda.get_device_capability() != (12, 0): + return None + device_name = torch.cuda.get_device_name() + if not any(marker in device_name for marker in _SM120A_DEVICE_NAME_MARKERS): + return None + return "12.0a" + except Exception: return None - return os.environ.get(_NATIVE_CUDA_ARCH_LIST_ENV, _DEFAULT_CUDA_ARCH_LIST) -def _effective_cuda_arch_list() -> str: - return os.environ.get( - _PYTORCH_CUDA_ARCH_LIST_ENV, - os.environ.get(_NATIVE_CUDA_ARCH_LIST_ENV, _DEFAULT_CUDA_ARCH_LIST), +def _effective_cuda_arch_list() -> str | None: + return ( + os.environ.get(_PYTORCH_CUDA_ARCH_LIST_ENV) + or os.environ.get(_NATIVE_CUDA_ARCH_LIST_ENV) + or _detected_cuda_arch_list() ) +def _cuda_arch_identity(cuda_arch_list: str | None) -> str: + return cuda_arch_list or _PYTORCH_DEFAULT_CUDA_ARCH_LIST + + def _python_package_dir(package: str) -> Path | None: spec = importlib.util.find_spec(package) if spec is None or spec.submodule_search_locations is None: @@ -505,14 +539,13 @@ def _scoped_torch_max_jobs(max_jobs: int | str | None) -> Iterator[None]: @contextlib.contextmanager -def _scoped_cuda_arch_list() -> Iterator[None]: - resolved = _resolved_cuda_arch_list() - if resolved is None: +def _scoped_cuda_arch_list(cuda_arch_list: str | None) -> Iterator[None]: + if cuda_arch_list is None: yield return previous = os.environ.get(_PYTORCH_CUDA_ARCH_LIST_ENV) - os.environ[_PYTORCH_CUDA_ARCH_LIST_ENV] = resolved + os.environ[_PYTORCH_CUDA_ARCH_LIST_ENV] = cuda_arch_list try: yield finally: @@ -540,8 +573,11 @@ def load_extension( global _extension, _extension_load_error with _state_lock: - sage3_disabled = _sage3_disabled() - if (extension := _extension.get(sage3_disabled)) is not None: + cuda_arch_list = _effective_cuda_arch_list() + cuda_arch_identity = _cuda_arch_identity(cuda_arch_list) + sage3_disabled = _sage3_disabled(cuda_arch_identity) + extension_key = (sage3_disabled, cuda_arch_identity) + if (extension := _extension.get(extension_key)) is not None: return extension _extension_load_error = None @@ -551,7 +587,10 @@ def load_extension( from torch.utils.cpp_extension import load as load_torch_extension thirdparty_info = validate_thirdparty() - extension_name = _extension_name(thirdparty_info) + extension_name = _extension_name( + thirdparty_info, + cuda_arch_list=cuda_arch_identity, + ) has_sage3 = int(not sage3_disabled) cutlass_dir = Path(thirdparty_info["cutlass"]["path"]) cutlass_include = cutlass_dir / "include" @@ -570,8 +609,11 @@ def load_extension( extension_build_dir.mkdir(parents=True, exist_ok=True) _add_windows_cuda_dll_directories(cudnn_package_dir) - with _scoped_torch_max_jobs(max_jobs), _scoped_cuda_arch_list(): - _extension[sage3_disabled] = load_torch_extension( + with ( + _scoped_torch_max_jobs(max_jobs), + _scoped_cuda_arch_list(cuda_arch_list), + ): + _extension[extension_key] = load_torch_extension( name=extension_name, sources=[str(source) for source in _extension_sources()], build_directory=str(extension_build_dir), @@ -636,7 +678,7 @@ def load_extension( "-DOMNIDREAMS_SINGLEVIEW_SPARGE_ATTN_SHA=" f'\\"{thirdparty_info["SpargeAttn"]["commit"]}\\"', "-DOMNIDREAMS_SINGLEVIEW_CUDA_ARCH_LIST=" - f'\\"{_effective_cuda_arch_list()}\\"', + f'\\"{cuda_arch_identity}\\"', ], extra_cuda_cflags=[ # Assume MSVC for Windows @@ -677,7 +719,7 @@ def load_extension( except Exception as exc: # pragma: no cover - environment-specific build path _extension_load_error = exc return None - return _extension[sage3_disabled] + return _extension[extension_key] def extension_load_error() -> Exception | None: diff --git a/integrations/omnidreams/tests/test_omnidreams_singleview_native.py b/integrations/omnidreams/tests/test_omnidreams_singleview_native.py index 66193fc70..a9ed6ab43 100644 --- a/integrations/omnidreams/tests/test_omnidreams_singleview_native.py +++ b/integrations/omnidreams/tests/test_omnidreams_singleview_native.py @@ -153,6 +153,7 @@ def fake_load_torch_extension(**kwargs: object) -> ModuleType: monkeypatch.setattr(cpp_extension, "load", fake_load_torch_extension) monkeypatch.setattr(native.os, "cpu_count", lambda: 48) monkeypatch.setattr(native, "_python_package_dir", lambda package: None) + monkeypatch.setattr(native, "_detected_cuda_arch_list", lambda: None) monkeypatch.delenv("MAX_JOBS", raising=False) monkeypatch.delenv("TORCH_CUDA_ARCH_LIST", raising=False) monkeypatch.delenv("OMNIDREAMS_SINGLEVIEW_CUDA_ARCH_LIST", raising=False) @@ -211,8 +212,6 @@ def fake_load_torch_extension(**kwargs: object) -> ModuleType: "lightvae_fp8_warp_mma_stages.cu", "lightvae_fp8_attention.cu", "streaming_dit_bridge.cu", - "sage3_blackwell_api_shim.cu", - "sage3_fp4_quant_shim.cu", "attention.cu", "block_quant.cu", "cosmos_adaln_lora.cu", @@ -224,7 +223,7 @@ def fake_load_torch_extension(**kwargs: object) -> ModuleType: "cosmos_gemm_bf16.cu", "cosmos_modulate.cu", "ops.cu", - "sage3_attention.cu", + "sage3_attention_stub.cu", "sparge_attention_sm89_inst.cu", "transformer_block.cu", ] @@ -267,27 +266,81 @@ def fake_load_torch_extension(**kwargs: object) -> ModuleType: '-DOMNIDREAMS_SINGLEVIEW_SAGE_ATTENTION_SHA=\\"sage-test-sha\\"' in captured["extra_cflags"] ) - assert "-DOMNIDREAMS_SINGLEVIEW_HAS_SAGE3=1" in captured["extra_cflags"] + assert "-DOMNIDREAMS_SINGLEVIEW_HAS_SAGE3=0" in captured["extra_cflags"] assert ( '-DOMNIDREAMS_SINGLEVIEW_SPARGE_ATTN_SHA=\\"sparge-test-sha\\"' in captured["extra_cflags"] ) assert "-DOMNIDREAMS_SINGLEVIEW_HAS_SPARGE=1" in captured["extra_cflags"] assert ( - '-DOMNIDREAMS_SINGLEVIEW_CUDA_ARCH_LIST=\\"12.0a\\"' in captured["extra_cflags"] + '-DOMNIDREAMS_SINGLEVIEW_CUDA_ARCH_LIST=\\"pytorch-default\\"' + in captured["extra_cflags"] ) assert "-DOMNIDREAMS_SINGLEVIEW_WITH_CUDA" in captured["extra_cuda_cflags"] if os.name == "nt": assert "-Xcompiler=/Zc:preprocessor" in captured["extra_cuda_cflags"] - assert "-DOMNIDREAMS_SINGLEVIEW_HAS_SAGE3=1" in captured["extra_cuda_cflags"] + assert "-DOMNIDREAMS_SINGLEVIEW_HAS_SAGE3=0" in captured["extra_cuda_cflags"] assert "-DOMNIDREAMS_SINGLEVIEW_HAS_SPARGE=1" in captured["extra_cuda_cflags"] assert captured["with_cuda"] is True assert captured["max_jobs_env"] == "8" - assert captured["cuda_arch_list_env"] == "12.0a" + assert captured["cuda_arch_list_env"] is None assert "MAX_JOBS" not in os.environ assert "TORCH_CUDA_ARCH_LIST" not in os.environ +@pytest.mark.ci_cpu +@pytest.mark.parametrize( + ("capability", "device_name", "expected"), + [ + ((12, 0), "NVIDIA GeForce RTX 5090", "12.0a"), + ((12, 0), "NVIDIA RTX PRO 6000 Blackwell", "12.0a"), + ((12, 0), "Unvalidated Compute Capability 12.0 GPU", None), + ((10, 3), "NVIDIA GB300", None), + ((8, 9), "NVIDIA RTX 6000 Ada Generation", None), + ], +) +def test_detected_cuda_arch_list_only_selects_validated_sm120a_devices( + capability: tuple[int, int], + device_name: str, + expected: str | None, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda: capability) + monkeypatch.setattr(torch.cuda, "get_device_name", lambda: device_name) + + assert native._detected_cuda_arch_list() == expected + + +@pytest.mark.ci_cpu +def test_effective_cuda_arch_list_precedence( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(native, "_detected_cuda_arch_list", lambda: "12.0a") + monkeypatch.delenv("TORCH_CUDA_ARCH_LIST", raising=False) + monkeypatch.delenv("OMNIDREAMS_SINGLEVIEW_CUDA_ARCH_LIST", raising=False) + + assert native._effective_cuda_arch_list() == "12.0a" + + monkeypatch.setenv("OMNIDREAMS_SINGLEVIEW_CUDA_ARCH_LIST", "10.3a") + assert native._effective_cuda_arch_list() == "10.3a" + + monkeypatch.setenv("TORCH_CUDA_ARCH_LIST", "8.9") + assert native._effective_cuda_arch_list() == "8.9" + + +@pytest.mark.ci_cpu +def test_effective_cuda_arch_list_uses_pytorch_default_without_sm120a( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(native, "_detected_cuda_arch_list", lambda: None) + monkeypatch.delenv("TORCH_CUDA_ARCH_LIST", raising=False) + monkeypatch.delenv("OMNIDREAMS_SINGLEVIEW_CUDA_ARCH_LIST", raising=False) + + assert native._effective_cuda_arch_list() is None + assert native._cuda_arch_identity(None) == "pytorch-default" + + @pytest.mark.ci_cpu @pytest.mark.parametrize( ("value", "expected"), @@ -307,6 +360,7 @@ def test_sage3_build_opt_out_parses_affirmative_values( expected: bool, monkeypatch: pytest.MonkeyPatch, ) -> None: + monkeypatch.setenv("OMNIDREAMS_SINGLEVIEW_CUDA_ARCH_LIST", "12.0a") if value is None: monkeypatch.delenv("OMNIDREAMS_SINGLEVIEW_DISABLE_SAGE3", raising=False) else: @@ -326,6 +380,23 @@ def test_sage3_build_opt_out_parses_affirmative_values( assert "sage3_attention.cu" in sources +@pytest.mark.ci_cpu +@pytest.mark.parametrize( + ("cuda_arch_list", "expected"), + [ + ("12.0a", False), + ("pytorch-default", True), + ("10.3a", True), + ("12.0", True), + ], +) +def test_sage3_build_requires_exact_sm120a_target( + cuda_arch_list: str, + expected: bool, +) -> None: + assert native._sage3_disabled(cuda_arch_list) is expected + + @pytest.mark.ci_cpu def test_load_extension_uses_sage3_stub_when_disabled( tmp_path: Path, @@ -381,6 +452,7 @@ def fake_load_torch_extension(**_: object) -> ModuleType: monkeypatch.setattr(native, "validate_thirdparty", lambda: thirdparty_info) monkeypatch.setattr(cpp_extension, "load", fake_load_torch_extension) monkeypatch.setattr(native, "_python_package_dir", lambda package: None) + monkeypatch.setenv("OMNIDREAMS_SINGLEVIEW_CUDA_ARCH_LIST", "12.0a") monkeypatch.delenv("OMNIDREAMS_SINGLEVIEW_DISABLE_SAGE3", raising=False) sage3_extension = native.load_extension(build_root=tmp_path / "native-build") @@ -397,6 +469,45 @@ def fake_load_torch_extension(**_: object) -> ModuleType: assert len(extensions) == 2 +@pytest.mark.ci_cpu +def test_load_extension_caches_separate_cuda_architectures( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + import torch.utils.cpp_extension as cpp_extension + + extensions: list[ModuleType] = [] + + def fake_load_torch_extension(**_: object) -> ModuleType: + extension = _fake_extension_module() + extensions.append(extension) + return extension + + thirdparty_info = _fake_thirdparty_info(tmp_path) + monkeypatch.setattr(native, "_extension", {}) + monkeypatch.setattr(native, "_extension_load_error", None) + monkeypatch.setattr(native, "validate_thirdparty", lambda: thirdparty_info) + monkeypatch.setattr(cpp_extension, "load", fake_load_torch_extension) + monkeypatch.setattr(native, "_python_package_dir", lambda package: None) + monkeypatch.delenv("TORCH_CUDA_ARCH_LIST", raising=False) + + monkeypatch.setenv("OMNIDREAMS_SINGLEVIEW_CUDA_ARCH_LIST", "10.3a") + sm103_extension = native.load_extension(build_root=tmp_path / "native-build") + monkeypatch.setenv("OMNIDREAMS_SINGLEVIEW_CUDA_ARCH_LIST", "12.0a") + sm120_extension = native.load_extension(build_root=tmp_path / "native-build") + + assert sm103_extension is extensions[0] + assert sm120_extension is extensions[1] + assert ( + native.load_extension(build_root=tmp_path / "native-build") is sm120_extension + ) + monkeypatch.setenv("OMNIDREAMS_SINGLEVIEW_CUDA_ARCH_LIST", "10.3a") + assert ( + native.load_extension(build_root=tmp_path / "native-build") is sm103_extension + ) + assert len(extensions) == 2 + + @pytest.mark.ci_cpu def test_extension_name_isolated_by_sage3_build_opt_out( tmp_path: Path, @@ -404,6 +515,7 @@ def test_extension_name_isolated_by_sage3_build_opt_out( ) -> None: thirdparty_info = _fake_thirdparty_info(tmp_path) monkeypatch.setattr(native, "_source_fingerprint", lambda: "fixed-fingerprint") + monkeypatch.setenv("OMNIDREAMS_SINGLEVIEW_CUDA_ARCH_LIST", "12.0a") monkeypatch.delenv("OMNIDREAMS_SINGLEVIEW_DISABLE_SAGE3", raising=False) full_name = native._extension_name(thirdparty_info) @@ -415,6 +527,26 @@ def test_extension_name_isolated_by_sage3_build_opt_out( assert "_sage3_0_" in stubbed_name +@pytest.mark.ci_cpu +def test_extension_name_isolated_by_cuda_architecture( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + thirdparty_info = _fake_thirdparty_info(tmp_path) + monkeypatch.setattr(native, "_source_fingerprint", lambda: "fixed-fingerprint") + + sm103_name = native._extension_name( + thirdparty_info, + cuda_arch_list="10.3a", + ) + sm120_name = native._extension_name( + thirdparty_info, + cuda_arch_list="12.0a", + ) + + assert sm103_name != sm120_name + + @pytest.mark.ci_cpu def test_load_extension_respects_existing_max_jobs( tmp_path: Path, @@ -937,10 +1069,7 @@ def test_cuda_native_extension_builds(tmp_path: Path) -> None: assert extension.is_available() build_info = extension.build_info() assert build_info["with_cuda"] is True - expected_arch = os.environ.get( - "TORCH_CUDA_ARCH_LIST", - os.environ.get("OMNIDREAMS_SINGLEVIEW_CUDA_ARCH_LIST", "12.0a"), - ) + expected_arch = native._cuda_arch_identity(native._effective_cuda_arch_list()) assert build_info["cuda_arch_list"] == expected_arch assert hasattr(extension, "native_tensor_descriptor") assert hasattr(extension, "native_tensor_ref_descriptor")