Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
mvanhorn wants to merge 2 commits into
base: main
Choose a base branch
from
Open

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
b664ccf
fix: honor sage3 opt-out in native build
mvanhorn Jul 11, 2026
e1bf2e0
fix(omnidreams): make extension cache mode-aware for SAGE3 toggle
mvanhorn Jul 17, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
39 changes: 28 additions & 11 deletions integrations/omnidreams/omnidreams/native/omnidreams_singleview.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -53,11 +53,12 @@
_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"

_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] = []
Expand Down Expand Up @@ -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,
Expand All @@ -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",
Expand All @@ -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",
]
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -525,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:
Expand All @@ -536,6 +552,7 @@ def load_extension(

Comment thread
greptile-apps[bot] marked this conversation as resolved.
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"])
Expand All @@ -554,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),
Expand Down Expand Up @@ -598,7 +615,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"]}\\"',
Expand Down Expand Up @@ -641,7 +658,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=[
Expand All @@ -656,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:
Expand Down
136 changes: 133 additions & 3 deletions integrations/omnidreams/tests/test_omnidreams_singleview_native.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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)

Expand Down Expand Up @@ -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"]
Expand All @@ -282,6 +284,134 @@ 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", {})
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_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,
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,
Expand All @@ -296,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)
Expand Down Expand Up @@ -329,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)
Expand Down