From b664ccf26b92e9003e0389ae29eaa89dca662409 Mon Sep 17 00:00:00 2001 From: Matt Van Horn <455140+mvanhorn@users.noreply.github.com> Date: Sat, 11 Jul 2026 01:37:41 -0700 Subject: [PATCH 1/2] fix: honor sage3 opt-out in native build --- .../native/omnidreams_singleview.py | 28 ++++-- .../test_omnidreams_singleview_native.py | 93 +++++++++++++++++++ 2 files changed, 115 insertions(+), 6 deletions(-) diff --git a/integrations/omnidreams/omnidreams/native/omnidreams_singleview.py b/integrations/omnidreams/omnidreams/native/omnidreams_singleview.py index db394b749..dc1863fdd 100644 --- a/integrations/omnidreams/omnidreams/native/omnidreams_singleview.py +++ b/integrations/omnidreams/omnidreams/native/omnidreams_singleview.py @@ -53,6 +53,7 @@ _PYTORCH_MAX_JOBS_ENV = "MAX_JOBS" _DEFAULT_MAX_JOBS_CAP = 8 _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" @@ -365,7 +366,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 _extension_sources() -> list[Path]: + disable_sage3 = _sage3_disabled() return [ _EXTENSION_SOURCE, _NATIVE_PRIMITIVES_SOURCE, @@ -378,8 +384,14 @@ def _extension_sources() -> list[Path]: _VAE_STREAMING_DIR / "lightvae_fp8_warp_mma_stages.cu", _VAE_STREAMING_DIR / "lightvae_fp8_attention.cu", _DIT_STREAMING_PYEXT_DIR / "streaming_dit_bridge.cu", - _DIT_STREAMING_PYEXT_DIR / "sage3_blackwell_api_shim.cu", - _DIT_STREAMING_PYEXT_DIR / "sage3_fp4_quant_shim.cu", + *( + [] + if disable_sage3 + else [ + _DIT_STREAMING_PYEXT_DIR / "sage3_blackwell_api_shim.cu", + _DIT_STREAMING_PYEXT_DIR / "sage3_fp4_quant_shim.cu", + ] + ), _DIT_STREAMING_KERNEL_DIR / "attention.cu", _DIT_STREAMING_KERNEL_DIR / "block_quant.cu", _DIT_STREAMING_KERNEL_DIR / "cosmos_adaln_lora.cu", @@ -391,7 +403,8 @@ def _extension_sources() -> list[Path]: _DIT_STREAMING_KERNEL_DIR / "cosmos_gemm_bf16.cu", _DIT_STREAMING_KERNEL_DIR / "cosmos_modulate.cu", _DIT_STREAMING_KERNEL_DIR / "ops.cu", - _DIT_STREAMING_KERNEL_DIR / "sage3_attention.cu", + _DIT_STREAMING_KERNEL_DIR + / ("sage3_attention_stub.cu" if disable_sage3 else "sage3_attention.cu"), _DIT_STREAMING_KERNEL_DIR / "sparge_attention_sm89_inst.cu", _DIT_STREAMING_KERNEL_DIR / "transformer_block.cu", ] @@ -421,10 +434,12 @@ def _source_fingerprint() -> str: def _extension_name(thirdparty_info: dict[str, Any]) -> str: + has_sage3 = int(not _sage3_disabled()) digest = hashlib.sha256() digest.update(_source_fingerprint().encode("ascii")) digest.update(json.dumps(thirdparty_info, sort_keys=True).encode("utf-8")) - return f"omnidreams_singleview_native_{digest.hexdigest()[:12]}" + digest.update(f"sage3={has_sage3}".encode("ascii")) + return f"omnidreams_singleview_native_sage3_{has_sage3}_{digest.hexdigest()[:12]}" def _validate_max_jobs(value: int | str) -> str: @@ -536,6 +551,7 @@ def load_extension( thirdparty_info = validate_thirdparty() extension_name = _extension_name(thirdparty_info) + has_sage3 = int(not _sage3_disabled()) cutlass_dir = Path(thirdparty_info["cutlass"]["path"]) cutlass_include = cutlass_dir / "include" sage_attention_dir = Path(thirdparty_info["SageAttention"]["path"]) @@ -598,7 +614,7 @@ def load_extension( "-std=c++20", "-DOMNIDREAMS_SINGLEVIEW_WITH_CUDA", "-DOMNIDREAMS_SINGLEVIEW_USE_CUTLASS", - "-DOMNIDREAMS_SINGLEVIEW_HAS_SAGE3=1", + f"-DOMNIDREAMS_SINGLEVIEW_HAS_SAGE3={has_sage3}", "-DOMNIDREAMS_SINGLEVIEW_HAS_SPARGE=1", "-DOMNIDREAMS_SINGLEVIEW_CUTLASS_SHA=" f'\\"{thirdparty_info["cutlass"]["commit"]}\\"', @@ -641,7 +657,7 @@ def load_extension( "-DCUTLASS_ENABLE_TENSOR_CORE_MMA=1", "-DOMNIDREAMS_SINGLEVIEW_WITH_CUDA", "-DOMNIDREAMS_SINGLEVIEW_USE_CUTLASS", - "-DOMNIDREAMS_SINGLEVIEW_HAS_SAGE3=1", + f"-DOMNIDREAMS_SINGLEVIEW_HAS_SAGE3={has_sage3}", "-DOMNIDREAMS_SINGLEVIEW_HAS_SPARGE=1", ], extra_ldflags=[ diff --git a/integrations/omnidreams/tests/test_omnidreams_singleview_native.py b/integrations/omnidreams/tests/test_omnidreams_singleview_native.py index a891bb4be..007a50a22 100644 --- a/integrations/omnidreams/tests/test_omnidreams_singleview_native.py +++ b/integrations/omnidreams/tests/test_omnidreams_singleview_native.py @@ -156,6 +156,7 @@ def fake_load_torch_extension(**kwargs: object) -> ModuleType: monkeypatch.delenv("MAX_JOBS", raising=False) monkeypatch.delenv("TORCH_CUDA_ARCH_LIST", raising=False) monkeypatch.delenv("OMNIDREAMS_SINGLEVIEW_CUDA_ARCH_LIST", raising=False) + monkeypatch.delenv("OMNIDREAMS_SINGLEVIEW_DISABLE_SAGE3", raising=False) extension = native.load_extension(build_root=build_root) @@ -264,6 +265,7 @@ 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_SPARGE_ATTN_SHA=\\"sparge-test-sha\\"' in captured["extra_cflags"] @@ -282,6 +284,97 @@ def fake_load_torch_extension(**kwargs: object) -> ModuleType: assert "TORCH_CUDA_ARCH_LIST" not in os.environ +@pytest.mark.ci_cpu +@pytest.mark.parametrize( + ("value", "expected"), + [ + (None, False), + ("", False), + ("0", False), + ("false", False), + ("yes", False), + ("1", True), + ("true", True), + ("TRUE", True), + ], +) +def test_sage3_build_opt_out_parses_affirmative_values( + value: str | None, + expected: bool, + monkeypatch: pytest.MonkeyPatch, +) -> None: + if value is None: + monkeypatch.delenv("OMNIDREAMS_SINGLEVIEW_DISABLE_SAGE3", raising=False) + else: + monkeypatch.setenv("OMNIDREAMS_SINGLEVIEW_DISABLE_SAGE3", value) + + assert native._sage3_disabled() is expected + sources = {source.name for source in native._extension_sources()} + if expected: + assert "sage3_attention_stub.cu" in sources + assert "sage3_blackwell_api_shim.cu" not in sources + assert "sage3_fp4_quant_shim.cu" not in sources + assert "sage3_attention.cu" not in sources + else: + assert "sage3_attention_stub.cu" not in sources + assert "sage3_blackwell_api_shim.cu" in sources + assert "sage3_fp4_quant_shim.cu" in sources + assert "sage3_attention.cu" in sources + + +@pytest.mark.ci_cpu +def test_load_extension_uses_sage3_stub_when_disabled( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + import torch.utils.cpp_extension as cpp_extension + + captured: dict[str, Any] = {} + + def fake_load_torch_extension(**kwargs: object) -> ModuleType: + captured.update(kwargs) + return _fake_extension_module() + + monkeypatch.setattr(native, "_extension", None) + monkeypatch.setattr(native, "_extension_load_error", None) + monkeypatch.setattr( + native, "validate_thirdparty", lambda: _fake_thirdparty_info(tmp_path) + ) + monkeypatch.setattr(cpp_extension, "load", fake_load_torch_extension) + monkeypatch.setattr(native, "_python_package_dir", lambda package: None) + monkeypatch.setenv("OMNIDREAMS_SINGLEVIEW_DISABLE_SAGE3", "1") + + extension = native.load_extension(build_root=tmp_path / "native-build") + + assert extension is not None + sources = {Path(str(source)).name for source in captured["sources"]} + assert "sage3_blackwell_api_shim.cu" not in sources + assert "sage3_fp4_quant_shim.cu" not in sources + assert "sage3_attention.cu" not in sources + assert "sage3_attention_stub.cu" in sources + assert "lightvae_fp8_attention.cu" in sources + assert "-DOMNIDREAMS_SINGLEVIEW_HAS_SAGE3=0" in captured["extra_cflags"] + assert "-DOMNIDREAMS_SINGLEVIEW_HAS_SAGE3=0" in captured["extra_cuda_cflags"] + + +@pytest.mark.ci_cpu +def test_extension_name_isolated_by_sage3_build_opt_out( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + thirdparty_info = _fake_thirdparty_info(tmp_path) + monkeypatch.setattr(native, "_source_fingerprint", lambda: "fixed-fingerprint") + monkeypatch.delenv("OMNIDREAMS_SINGLEVIEW_DISABLE_SAGE3", raising=False) + full_name = native._extension_name(thirdparty_info) + + monkeypatch.setenv("OMNIDREAMS_SINGLEVIEW_DISABLE_SAGE3", "true") + stubbed_name = native._extension_name(thirdparty_info) + + assert stubbed_name != full_name + assert "_sage3_1_" in full_name + assert "_sage3_0_" in stubbed_name + + @pytest.mark.ci_cpu def test_load_extension_respects_existing_max_jobs( tmp_path: Path, From e1bf2e01554c3bd41e7fa28ecb99df2698d5d9db Mon Sep 17 00:00:00 2001 From: Matt Van Horn <455140+mvanhorn@users.noreply.github.com> Date: Fri, 17 Jul 2026 11:16:26 -0700 Subject: [PATCH 2/2] fix(omnidreams): make extension cache mode-aware for SAGE3 toggle Key the _extension cache on the resolved OMNIDREAMS_SINGLEVIEW_DISABLE_SAGE3 mode so toggling it at runtime loads the correct binary instead of a stale cached one. --- .../native/omnidreams_singleview.py | 13 +++--- .../test_omnidreams_singleview_native.py | 45 +++++++++++++++++-- 2 files changed, 48 insertions(+), 10 deletions(-) diff --git a/integrations/omnidreams/omnidreams/native/omnidreams_singleview.py b/integrations/omnidreams/omnidreams/native/omnidreams_singleview.py index dc1863fdd..3ccaea426 100644 --- a/integrations/omnidreams/omnidreams/native/omnidreams_singleview.py +++ b/integrations/omnidreams/omnidreams/native/omnidreams_singleview.py @@ -58,7 +58,7 @@ _DEFAULT_CUDA_ARCH_LIST = "12.0a" _native_build_module: ModuleType | None = None -_extension: ModuleType | None = None +_extension: dict[bool, ModuleType] = {} _extension_load_error: Exception | None = None _state_lock = threading.RLock() _dll_directory_handles: list[object] = [] @@ -540,8 +540,9 @@ def load_extension( global _extension, _extension_load_error with _state_lock: - if _extension is not None: - return _extension + sage3_disabled = _sage3_disabled() + if (extension := _extension.get(sage3_disabled)) is not None: + return extension _extension_load_error = None try: @@ -551,7 +552,7 @@ def load_extension( thirdparty_info = validate_thirdparty() extension_name = _extension_name(thirdparty_info) - has_sage3 = int(not _sage3_disabled()) + has_sage3 = int(not sage3_disabled) cutlass_dir = Path(thirdparty_info["cutlass"]["path"]) cutlass_include = cutlass_dir / "include" sage_attention_dir = Path(thirdparty_info["SageAttention"]["path"]) @@ -570,7 +571,7 @@ def load_extension( _add_windows_cuda_dll_directories(cudnn_package_dir) with _scoped_torch_max_jobs(max_jobs), _scoped_cuda_arch_list(): - _extension = load_torch_extension( + _extension[sage3_disabled] = load_torch_extension( name=extension_name, sources=[str(source) for source in _extension_sources()], build_directory=str(extension_build_dir), @@ -672,7 +673,7 @@ def load_extension( except Exception as exc: # pragma: no cover - environment-specific build path _extension_load_error = exc return None - return _extension + return _extension[sage3_disabled] 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 007a50a22..a073d95d1 100644 --- a/integrations/omnidreams/tests/test_omnidreams_singleview_native.py +++ b/integrations/omnidreams/tests/test_omnidreams_singleview_native.py @@ -147,7 +147,7 @@ def fake_load_torch_extension(**kwargs: object) -> ModuleType: captured["cuda_arch_list_env"] = os.environ.get("TORCH_CUDA_ARCH_LIST") return _fake_extension_module() - monkeypatch.setattr(native, "_extension", None) + 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) @@ -335,7 +335,7 @@ def fake_load_torch_extension(**kwargs: object) -> ModuleType: captured.update(kwargs) return _fake_extension_module() - monkeypatch.setattr(native, "_extension", None) + monkeypatch.setattr(native, "_extension", {}) monkeypatch.setattr(native, "_extension_load_error", None) monkeypatch.setattr( native, "validate_thirdparty", lambda: _fake_thirdparty_info(tmp_path) @@ -357,6 +357,43 @@ def fake_load_torch_extension(**kwargs: object) -> ModuleType: assert "-DOMNIDREAMS_SINGLEVIEW_HAS_SAGE3=0" in captured["extra_cuda_cflags"] +@pytest.mark.ci_cpu +def test_load_extension_caches_separate_sage3_modes( + 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 + + monkeypatch.setattr(native, "_extension", {}) + monkeypatch.setattr(native, "_extension_load_error", None) + monkeypatch.setattr( + native, "validate_thirdparty", lambda: _fake_thirdparty_info(tmp_path) + ) + monkeypatch.setattr(cpp_extension, "load", fake_load_torch_extension) + monkeypatch.setattr(native, "_python_package_dir", lambda package: None) + monkeypatch.delenv("OMNIDREAMS_SINGLEVIEW_DISABLE_SAGE3", raising=False) + + sage3_extension = native.load_extension(build_root=tmp_path / "native-build") + monkeypatch.setenv("OMNIDREAMS_SINGLEVIEW_DISABLE_SAGE3", "1") + stub_extension = native.load_extension(build_root=tmp_path / "native-build") + + assert sage3_extension is extensions[0] + assert stub_extension is extensions[1] + assert native.load_extension(build_root=tmp_path / "native-build") is stub_extension + monkeypatch.delenv("OMNIDREAMS_SINGLEVIEW_DISABLE_SAGE3", raising=False) + assert ( + native.load_extension(build_root=tmp_path / "native-build") is sage3_extension + ) + assert len(extensions) == 2 + + @pytest.mark.ci_cpu def test_extension_name_isolated_by_sage3_build_opt_out( tmp_path: Path, @@ -389,7 +426,7 @@ def fake_load_torch_extension(**kwargs: object) -> ModuleType: captured["cuda_arch_list_env"] = os.environ.get("TORCH_CUDA_ARCH_LIST") return _fake_extension_module() - monkeypatch.setattr(native, "_extension", None) + monkeypatch.setattr(native, "_extension", {}) monkeypatch.setattr(native, "_extension_load_error", None) monkeypatch.setattr( native, "validate_thirdparty", lambda: _fake_thirdparty_info(tmp_path) @@ -422,7 +459,7 @@ def fake_load_torch_extension(**_: object) -> ModuleType: return _fake_extension_module() thirdparty_info = _fake_thirdparty_info(tmp_path) - monkeypatch.setattr(native, "_extension", None) + monkeypatch.setattr(native, "_extension", {}) monkeypatch.setattr(native, "_extension_load_error", None) monkeypatch.setattr(native, "validate_thirdparty", lambda: thirdparty_info) monkeypatch.setattr("torch.utils.cpp_extension.load", fake_load_torch_extension)