From a34ab87d397f38cf5d53e6e9a344a0fe91994c87 Mon Sep 17 00:00:00 2001 From: "zhiguo.qin" Date: Sun, 23 Aug 2026 15:40:51 +0800 Subject: [PATCH 1/3] compat: add MUSA graph-safe ops --- README.md | 2 +- README_CN.md | 2 +- pyproject.toml | 2 +- src/torchada/__init__.py | 2 +- src/torchada/_patch.py | 45 ++++++ src/torchada/csrc/musa_ops.mu | 279 ++++++++++++++++++++++++++++++++++ tests/test_cuda_patching.py | 56 +++++++ tests/test_log.py | 66 ++++++++ tests/test_multinomial.py | 100 ++++++++++++ tests/test_platform.py | 2 +- 10 files changed, 551 insertions(+), 5 deletions(-) create mode 100644 tests/test_log.py create mode 100644 tests/test_multinomial.py diff --git a/README.md b/README.md index 921ecbf..3199527 100644 --- a/README.md +++ b/README.md @@ -392,7 +392,7 @@ See `src/torchada/_mappings/` for 400+ mapping rules grouped by API domain. ``` # pyproject.toml or requirements.txt -torchada>=0.1.83 +torchada>=0.1.84 ``` ### Step 2: Conditional Import diff --git a/README_CN.md b/README_CN.md index ebec54c..b3040ff 100644 --- a/README_CN.md +++ b/README_CN.md @@ -376,7 +376,7 @@ if torchada.is_gpu_device(device): # 在 CUDA 和 MUSA 上都能工作 ``` # pyproject.toml 或 requirements.txt -torchada>=0.1.83 +torchada>=0.1.84 ``` ### 步骤 2：条件导入 diff --git a/pyproject.toml b/pyproject.toml index b7fe737..836957a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "torchada" -version = "0.1.83" +version = "0.1.84" description = "Adapter package for torch_musa to act exactly like PyTorch CUDA" readme = "README.md" license = {text = "MIT"} diff --git a/src/torchada/__init__.py b/src/torchada/__init__.py index 8eefeb4..001f1ed 100644 --- a/src/torchada/__init__.py +++ b/src/torchada/__init__.py @@ -24,7 +24,7 @@ from torch.utils.cpp_extension import CUDAExtension, BuildExtension, CUDA_HOME """ -__version__ = "0.1.83" +__version__ = "0.1.84" from . import cuda, utils diff --git a/src/torchada/_patch.py b/src/torchada/_patch.py index 1202009..39dafa8 100644 --- a/src/torchada/_patch.py +++ b/src/torchada/_patch.py @@ -65,6 +65,17 @@ def _patch_something(): return func +@patch_function +def _patch_visible_devices_env(): + if "CUDA_VISIBLE_DEVICES" not in os.environ and "MUSA_VISIBLE_DEVICES" in os.environ: + os.environ["CUDA_VISIBLE_DEVICES"] = os.environ["MUSA_VISIBLE_DEVICES"] + elif ( + "MUSA_VISIBLE_DEVICES" not in os.environ + and "CUDA_VISIBLE_DEVICES" in os.environ + ): + os.environ["MUSA_VISIBLE_DEVICES"] = os.environ["CUDA_VISIBLE_DEVICES"] + + def requires_import(*module_names: str) -> Callable[[Callable], Callable]: """ Decorator to guard a patch function with import checks. @@ -108,6 +119,39 @@ def wrapper(*args, **kwargs): return decorator +@patch_function +@requires_import("torch._inductor.template_heuristics.registry") +def _patch_inductor_template_heuristics(): + """Reuse CUDA Inductor template heuristics for CUDA-compatible MUSA templates.""" + if not is_musa_platform(): + return + + import torch._inductor.template_heuristics.registry as registry + + heuristic_registry = getattr(registry, "_TEMPLATE_HEURISTIC_REGISTRY", None) + if not isinstance(heuristic_registry, dict): + return + + changed = False + for key, heuristic_class in list(heuristic_registry.items()): + if len(key) != 3: + continue + template_name, device_type, op_name = key + if device_type != "cuda": + continue + if not isinstance(template_name, str) or not template_name.startswith("triton::"): + continue + musa_key = (template_name, "musa", op_name) + if musa_key not in heuristic_registry: + heuristic_registry[musa_key] = heuristic_class + changed = True + + if changed: + heuristic_cache = getattr(registry, "_HEURISTIC_CACHE", None) + if isinstance(heuristic_cache, dict): + heuristic_cache.clear() + + # Cache for translated device strings - avoids repeated string operations _device_str_cache = {} @@ -2122,6 +2166,7 @@ def apply_patches(): - torch.cuda.nccl -> torch.musa.mccl - torch.amp.autocast(device_type='cuda') -> 'musa' - torch.utils.cpp_extension (CUDAExtension, BuildExtension) -> MUSA versions + - CUDA_VISIBLE_DEVICES -> MUSA_VISIBLE_DEVICES environment fallback - torch._inductor.autotune_process.CUDA_VISIBLE_DEVICES -> MUSA_VISIBLE_DEVICES - torch.accelerator.synchronize() -> torch.musa.synchronize() - torch.accelerator context managers (device_index, stream) for forward compatibility diff --git a/src/torchada/csrc/musa_ops.mu b/src/torchada/csrc/musa_ops.mu index a03f325..e5076bb 100644 --- a/src/torchada/csrc/musa_ops.mu +++ b/src/torchada/csrc/musa_ops.mu @@ -9,6 +9,11 @@ #include "ops.h" #include +#include +#include +#include +#include +#include namespace torchada { @@ -67,6 +72,273 @@ at::Tensor neg_musa_impl(const at::Tensor& self) { return output; } +namespace { + +__device__ unsigned long long multinomial_counter = 0; + +__device__ __forceinline__ unsigned long long splitmix64(unsigned long long x) { + x += 0x9E3779B97F4A7C15ull; + x = (x ^ (x >> 30)) * 0xBF58476D1CE4E5B9ull; + x = (x ^ (x >> 27)) * 0x94D049BB133111EBull; + return x ^ (x >> 31); +} + +__device__ __forceinline__ double uniform01( + unsigned long long seed, + int64_t row, + int64_t sample) { + unsigned long long x = splitmix64( + seed ^ (static_cast(row) * 0xD1B54A32D192ED03ull) ^ + (static_cast(sample) * 0x94D049BB133111EBull)); + constexpr double scale = 1.0 / 9007199254740992.0; + return static_cast(x >> 11) * scale; +} + +template +__device__ __forceinline__ double read_weight( + const scalar_t* input, + int64_t idx) { + double v = static_cast(input[idx]); + return isfinite(v) && v > 0.0 ? v : 0.0; +} + +__device__ __forceinline__ bool already_selected( + const int64_t* output, + int64_t row, + int64_t num_samples, + int64_t current_sample, + int64_t candidate) { + const int64_t base = row * num_samples; + for (int64_t i = 0; i < current_sample; ++i) { + if (output[base + i] == candidate) { + return true; + } + } + return false; +} + +template +__global__ void multinomial_kernel( + const scalar_t* __restrict__ input, + int64_t* __restrict__ output, + int64_t rows, + int64_t cols, + int64_t num_samples, + bool replacement, + unsigned long long seed_base) { + __shared__ double partial[BLOCK]; + __shared__ double prefix[BLOCK]; + __shared__ double total_sum; + __shared__ unsigned long long block_seed; + + int64_t row = static_cast(blockIdx.x); + int tid = threadIdx.x; + if (row >= rows) { + return; + } + + const int64_t row_offset = row * cols; + const int64_t chunk = (cols + BLOCK - 1) / BLOCK; + const int64_t begin = static_cast(tid) * chunk; + const int64_t end = min(begin + chunk, cols); + + if (tid == 0) { + unsigned long long counter = atomicAdd(&multinomial_counter, 1ull); + block_seed = splitmix64( + seed_base ^ counter ^ static_cast(clock64())); + } + __syncthreads(); + unsigned long long seed = block_seed; + + for (int64_t sample = 0; sample < num_samples; ++sample) { + double sum = 0.0; + for (int64_t col = begin; col < end; ++col) { + if (replacement || !already_selected(output, row, num_samples, sample, col)) { + sum += read_weight(input, row_offset + col); + } + } + partial[tid] = sum; + __syncthreads(); + + if (tid == 0) { + double running = 0.0; + for (int i = 0; i < BLOCK; ++i) { + prefix[i] = running; + running += partial[i]; + } + total_sum = running; + } + __syncthreads(); + + double total = total_sum; + int64_t selected = 0; + if (total > 0.0) { + double target = uniform01(seed, row, sample) * total; + double before = prefix[tid]; + double after = before + partial[tid]; + if (target >= before && target < after) { + double running = before; + for (int64_t col = begin; col < end; ++col) { + if (!replacement && already_selected(output, row, num_samples, sample, col)) { + continue; + } + running += read_weight(input, row_offset + col); + if (target < running) { + selected = col; + break; + } + } + } else { + selected = -1; + } + } else { + selected = tid == 0 ? 0 : -1; + } + partial[tid] = static_cast(selected); + __syncthreads(); + + if (tid == 0) { + int64_t chosen = 0; + for (int i = 0; i < BLOCK; ++i) { + int64_t candidate = static_cast(partial[i]); + if (candidate >= 0) { + chosen = candidate; + break; + } + } + output[row * num_samples + sample] = chosen; + } + __syncthreads(); + } +} + +} // namespace + +at::Tensor multinomial_musa_impl( + const at::Tensor& self, + int64_t num_samples, + bool replacement, + c10::optional generator) { + log_op_call("multinomial"); + + TORCH_CHECK(self.dim() == 1 || self.dim() == 2, "prob_dist must be 1 or 2 dim"); + TORCH_CHECK(num_samples >= 0, "cannot sample n_sample < 0 samples"); + + int64_t rows = self.dim() == 1 ? 1 : self.size(0); + int64_t cols = self.dim() == 1 ? self.size(0) : self.size(1); + if (!replacement) { + TORCH_CHECK( + num_samples <= cols, + "cannot sample n_sample > prob_dist.size(-1) samples without replacement"); + } + + auto options = self.options().dtype(at::kLong); + at::Tensor output = self.dim() == 1 + ? at::empty({num_samples}, options) + : at::empty({rows, num_samples}, options); + + if (num_samples == 0 || rows == 0) { + return output; + } + + auto input = self.contiguous(); + constexpr int BLOCK = 256; + musaStream_t stream = at::musa::getCurrentMUSAStream(); + unsigned long long seed_base = generator.has_value() + ? static_cast(generator->current_seed()) + : static_cast(clock()); + + AT_DISPATCH_FLOATING_TYPES_AND2( + at::ScalarType::Half, + at::ScalarType::BFloat16, + input.scalar_type(), + "torchada_multinomial_musa", + [&] { + multinomial_kernel<<>>( + input.data_ptr(), + output.data_ptr(), + rows, + cols, + num_samples, + replacement, + seed_base); + }); + + musaError_t err = musaGetLastError(); + if (err != musaSuccess) { + TORCH_CHECK(false, "MUSA multinomial kernel launch failed: ", musaGetErrorString(err)); + } + + return output; +} + +template +__global__ void log_kernel( + scalar_t* __restrict__ output, + const scalar_t* __restrict__ input, + int64_t numel) { + int64_t idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx < numel) { + double value = static_cast(input[idx]); + output[idx] = static_cast(log(value)); + } +} + +at::Tensor log_musa_impl(const at::Tensor& self) { + log_op_call("log"); + TORCH_CHECK( + at::isFloatingType(self.scalar_type()), + "torchada MUSA log only supports floating point tensors"); + + auto input = self.contiguous(); + auto output = at::empty_like(input); + if (input.numel() == 0) { + return output; + } + + constexpr int threads = 256; + const int64_t numel = input.numel(); + const int blocks = static_cast((numel + threads - 1) / threads); + musaStream_t stream = at::musa::getCurrentMUSAStream(); + + AT_DISPATCH_FLOATING_TYPES_AND2( + at::ScalarType::Half, + at::ScalarType::BFloat16, + input.scalar_type(), + "torchada_log_musa", + [&] { + log_kernel<<>>( + output.data_ptr(), + input.data_ptr(), + numel); + }); + + musaError_t err = musaGetLastError(); + if (err != musaSuccess) { + TORCH_CHECK(false, "MUSA log kernel launch failed: ", musaGetErrorString(err)); + } + + if (!self.is_contiguous()) { + return output.view(self.sizes()); + } + return output; +} + +at::Tensor& log_inplace_musa_impl(at::Tensor& self) { + log_op_call("log_"); + TORCH_CHECK( + at::isFloatingType(self.scalar_type()), + "torchada MUSA log_ only supports floating point tensors"); + + if (self.numel() == 0) { + return self; + } + + auto output = log_musa_impl(self); + self.copy_(output); + return self; +} + } // namespace torchada // ============================================================================ @@ -84,4 +356,11 @@ TORCH_LIBRARY_IMPL(aten, PrivateUse1, m) { // if (torchada::is_override_enabled("neg")) { // m.impl("neg", torchada::neg_musa_impl); // } + if (torchada::is_override_enabled("multinomial")) { + m.impl("multinomial", torchada::multinomial_musa_impl); + } + if (torchada::is_override_enabled("log")) { + m.impl("log", torchada::log_musa_impl); + m.impl("log_", torchada::log_inplace_musa_impl); + } } diff --git a/tests/test_cuda_patching.py b/tests/test_cuda_patching.py index 22a3699..a3df121 100644 --- a/tests/test_cuda_patching.py +++ b/tests/test_cuda_patching.py @@ -4,6 +4,8 @@ These tests verify that torch.cuda.* APIs work transparently on MUSA. """ +import os + import pytest @@ -1082,6 +1084,60 @@ def test_musa_tensor_is_cuda_true(self): raise +class TestVisibleDevicesEnv: + """Test CUDA_VISIBLE_DEVICES and MUSA_VISIBLE_DEVICES fallback.""" + + def test_cuda_visible_devices_env_falls_back_to_musa_visible_devices( + self, monkeypatch + ): + """Test CUDA_VISIBLE_DEVICES is copied to MUSA_VISIBLE_DEVICES.""" + from torchada import _patch + + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "1,3") + monkeypatch.delenv("MUSA_VISIBLE_DEVICES", raising=False) + + _patch._patch_visible_devices_env() + + assert os.environ["MUSA_VISIBLE_DEVICES"] == "1,3" + + def test_musa_visible_devices_env_falls_back_to_cuda_visible_devices( + self, monkeypatch + ): + """Test MUSA_VISIBLE_DEVICES is copied to CUDA_VISIBLE_DEVICES.""" + from torchada import _patch + + monkeypatch.setenv("MUSA_VISIBLE_DEVICES", "0") + monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising=False) + + _patch._patch_visible_devices_env() + + assert os.environ["CUDA_VISIBLE_DEVICES"] == "0" + + def test_existing_visible_devices_envs_are_not_overwritten(self, monkeypatch): + """Test explicit visible device envs have priority.""" + from torchada import _patch + + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "1,3") + monkeypatch.setenv("MUSA_VISIBLE_DEVICES", "0") + + _patch._patch_visible_devices_env() + + assert os.environ["CUDA_VISIBLE_DEVICES"] == "1,3" + assert os.environ["MUSA_VISIBLE_DEVICES"] == "0" + + def test_visible_devices_env_absent_noop(self, monkeypatch): + """Test no visible device envs are added when both are absent.""" + from torchada import _patch + + monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising=False) + monkeypatch.delenv("MUSA_VISIBLE_DEVICES", raising=False) + + _patch._patch_visible_devices_env() + + assert "CUDA_VISIBLE_DEVICES" not in os.environ + assert "MUSA_VISIBLE_DEVICES" not in os.environ + + class TestAutotuneProcess: """Test torch._inductor.autotune_process patching.""" diff --git a/tests/test_log.py b/tests/test_log.py new file mode 100644 index 0000000..50d2dd9 --- /dev/null +++ b/tests/test_log.py @@ -0,0 +1,66 @@ +import pytest +import torch + + +def _require_musa(): + import torchada + + if not torchada.is_musa_platform(): + pytest.skip("MUSA platform required") + + if not hasattr(torch, "musa") or not torch.musa.is_available(): + pytest.skip("MUSA platform required") + + +def test_log_float64_privateuse1_smoke(): + _require_musa() + from torchada import _cpp_ops + + _cpp_ops.load_cpp_ops(force_reload=True) + x = torch.linspace(0.1, 2.0, 128, device="cuda", dtype=torch.float64) + + out = torch.log(x) + torch.cuda.synchronize() + + expected = torch.log(x.cpu()) + torch.testing.assert_close(out.cpu(), expected, rtol=1e-12, atol=1e-12) + + +def test_log_inplace_float64_privateuse1_smoke(): + _require_musa() + from torchada import _cpp_ops + + _cpp_ops.load_cpp_ops(force_reload=True) + x = torch.linspace(0.1, 2.0, 128, device="cuda", dtype=torch.float64) + expected = torch.log(x.cpu()) + + ret = x.log_() + torch.cuda.synchronize() + + assert ret is x + torch.testing.assert_close(x.cpu(), expected, rtol=1e-12, atol=1e-12) + + +def test_log_inplace_float64_graph_capture_replay(): + _require_musa() + from torchada import _cpp_ops + + _cpp_ops.load_cpp_ops(force_reload=True) + x = torch.linspace(0.1, 2.0, 128, device="cuda", dtype=torch.float64) + work = x.clone() + for _ in range(3): + work.copy_(x) + work.log_() + torch.cuda.synchronize() + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + work.copy_(x) + work.log_() + + for _ in range(5): + graph.replay() + torch.cuda.synchronize() + + expected = torch.log(x.cpu()) + torch.testing.assert_close(work.cpu(), expected, rtol=1e-12, atol=1e-12) diff --git a/tests/test_multinomial.py b/tests/test_multinomial.py new file mode 100644 index 0000000..dc30d6a --- /dev/null +++ b/tests/test_multinomial.py @@ -0,0 +1,100 @@ +import pytest +import torch + + +def _require_musa(): + import torchada + + if ( + not torchada.is_musa_platform() + or not hasattr(torch, "musa") + or not torch.musa.is_available() + ): + pytest.skip("MUSA platform required") + + +def test_multinomial_privateuse1_smoke(): + _require_musa() + from torchada import _cpp_ops + + _cpp_ops.load_cpp_ops(force_reload=True) + probs = torch.softmax(torch.randn(4, 64, device="cuda"), dim=-1) + + out = torch.multinomial(probs, 1) + torch.cuda.synchronize() + + assert out.shape == (4, 1) + assert out.dtype == torch.long + assert out.device.type in ("cuda", "musa") + assert int(out.min()) >= 0 + assert int(out.max()) < 64 + + +def test_multinomial_without_replacement_unique(): + _require_musa() + from torchada import _cpp_ops + + _cpp_ops.load_cpp_ops(force_reload=True) + probs = torch.full((128, 16), 1.0 / 16, device="cuda") + + out = torch.multinomial(probs, 8, replacement=False).cpu() + + assert all(len(set(row.tolist())) == 8 for row in out) + + +def test_multinomial_distribution_sanity(): + _require_musa() + from torchada import _cpp_ops + + _cpp_ops.load_cpp_ops(force_reload=True) + rows = 4096 + weights = torch.tensor( + [0.01, 0.02, 0.04, 0.08, 0.16, 0.32, 0.37], + device="cuda", + ).repeat(rows, 1) + + out = torch.multinomial(weights, 1).flatten() + counts = torch.bincount(out.cpu(), minlength=weights.shape[1]) + + assert int(counts.argmax()) == 6 + assert counts[-1] > counts[-2] > counts[-3] + + +def test_multinomial_graph_capture_replay(): + _require_musa() + from torchada import _cpp_ops + + _cpp_ops.load_cpp_ops(force_reload=True) + probs = torch.softmax(torch.randn(2, 128, device="cuda"), dim=-1) + for _ in range(3): + out = torch.multinomial(probs, 1) + torch.cuda.synchronize() + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out = torch.multinomial(probs, 1) + + for _ in range(5): + graph.replay() + torch.cuda.synchronize() + + assert out.shape == (2, 1) + assert int(out.min()) >= 0 + assert int(out.max()) < 128 + + +def test_multinomial_accepts_generator_argument(): + _require_musa() + from torchada import _cpp_ops + + _cpp_ops.load_cpp_ops(force_reload=True) + probs = torch.full((4, 32), 1.0 / 32, device="cuda") + generator = torch.Generator(device="cuda") + generator.manual_seed(1234) + + out = torch.multinomial(probs, 1, generator=generator) + torch.cuda.synchronize() + + assert out.shape == (4, 1) + assert int(out.min()) >= 0 + assert int(out.max()) < 32 diff --git a/tests/test_platform.py b/tests/test_platform.py index 81f72ca..5d6ca61 100644 --- a/tests/test_platform.py +++ b/tests/test_platform.py @@ -59,7 +59,7 @@ def test_get_version(self): version = torchada.get_version() assert version == torchada.__version__ - assert version == "0.1.83" + assert version == "0.1.84" assert isinstance(version, str) def test_project_version_matches_runtime_version(self): From e7aeea19a19a4fb98681bd989d7f5d878556b21f Mon Sep 17 00:00:00 2001 From: "zhiguo.qin" Date: Mon, 24 Aug 2026 19:01:32 +0800 Subject: [PATCH 2/3] test: cover MUSA Inductor heuristic patching --- src/torchada/_patch.py | 9 ++-- src/torchada/csrc/musa_ops.mu | 5 --- tests/test_cuda_patching.py | 82 +++++++++++++++++++++++++++++++---- 3 files changed, 76 insertions(+), 20 deletions(-) diff --git a/src/torchada/_patch.py b/src/torchada/_patch.py index 39dafa8..48ab7e3 100644 --- a/src/torchada/_patch.py +++ b/src/torchada/_patch.py @@ -67,12 +67,9 @@ def _patch_something(): @patch_function def _patch_visible_devices_env(): - if "CUDA_VISIBLE_DEVICES" not in os.environ and "MUSA_VISIBLE_DEVICES" in os.environ: + if "MUSA_VISIBLE_DEVICES" in os.environ: os.environ["CUDA_VISIBLE_DEVICES"] = os.environ["MUSA_VISIBLE_DEVICES"] - elif ( - "MUSA_VISIBLE_DEVICES" not in os.environ - and "CUDA_VISIBLE_DEVICES" in os.environ - ): + elif "CUDA_VISIBLE_DEVICES" in os.environ: os.environ["MUSA_VISIBLE_DEVICES"] = os.environ["CUDA_VISIBLE_DEVICES"] @@ -134,7 +131,7 @@ def _patch_inductor_template_heuristics(): changed = False for key, heuristic_class in list(heuristic_registry.items()): - if len(key) != 3: + if not isinstance(key, tuple) or len(key) != 3: continue template_name, device_type, op_name = key if device_type != "cuda": diff --git a/src/torchada/csrc/musa_ops.mu b/src/torchada/csrc/musa_ops.mu index e5076bb..148bfdc 100644 --- a/src/torchada/csrc/musa_ops.mu +++ b/src/torchada/csrc/musa_ops.mu @@ -348,14 +348,9 @@ at::Tensor& log_inplace_musa_impl(at::Tensor& self) { // time. If set, the override is not registered and torch_musa's default // implementation is used. // -// Uncomment m.impl() lines to activate custom implementations. // ============================================================================ TORCH_LIBRARY_IMPL(aten, PrivateUse1, m) { - // Example: Register neg override only if not disabled - // if (torchada::is_override_enabled("neg")) { - // m.impl("neg", torchada::neg_musa_impl); - // } if (torchada::is_override_enabled("multinomial")) { m.impl("multinomial", torchada::multinomial_musa_impl); } diff --git a/tests/test_cuda_patching.py b/tests/test_cuda_patching.py index a3df121..4ffd8a6 100644 --- a/tests/test_cuda_patching.py +++ b/tests/test_cuda_patching.py @@ -1087,9 +1087,7 @@ def test_musa_tensor_is_cuda_true(self): class TestVisibleDevicesEnv: """Test CUDA_VISIBLE_DEVICES and MUSA_VISIBLE_DEVICES fallback.""" - def test_cuda_visible_devices_env_falls_back_to_musa_visible_devices( - self, monkeypatch - ): + def test_cuda_visible_devices_env_falls_back_to_musa_visible_devices(self, monkeypatch): """Test CUDA_VISIBLE_DEVICES is copied to MUSA_VISIBLE_DEVICES.""" from torchada import _patch @@ -1100,9 +1098,7 @@ def test_cuda_visible_devices_env_falls_back_to_musa_visible_devices( assert os.environ["MUSA_VISIBLE_DEVICES"] == "1,3" - def test_musa_visible_devices_env_falls_back_to_cuda_visible_devices( - self, monkeypatch - ): + def test_musa_visible_devices_env_falls_back_to_cuda_visible_devices(self, monkeypatch): """Test MUSA_VISIBLE_DEVICES is copied to CUDA_VISIBLE_DEVICES.""" from torchada import _patch @@ -1113,8 +1109,8 @@ def test_musa_visible_devices_env_falls_back_to_cuda_visible_devices( assert os.environ["CUDA_VISIBLE_DEVICES"] == "0" - def test_existing_visible_devices_envs_are_not_overwritten(self, monkeypatch): - """Test explicit visible device envs have priority.""" + def test_musa_visible_devices_overrides_cuda_visible_devices(self, monkeypatch): + """Test the MUSA setting wins when both variables are explicit.""" from torchada import _patch monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "1,3") @@ -1122,7 +1118,7 @@ def test_existing_visible_devices_envs_are_not_overwritten(self, monkeypatch): _patch._patch_visible_devices_env() - assert os.environ["CUDA_VISIBLE_DEVICES"] == "1,3" + assert os.environ["CUDA_VISIBLE_DEVICES"] == "0" assert os.environ["MUSA_VISIBLE_DEVICES"] == "0" def test_visible_devices_env_absent_noop(self, monkeypatch): @@ -1138,6 +1134,74 @@ def test_visible_devices_env_absent_noop(self, monkeypatch): assert "MUSA_VISIBLE_DEVICES" not in os.environ +class TestInductorTemplateHeuristics: + """Test CUDA-compatible Triton heuristic registration for MUSA.""" + + def test_copies_only_cuda_triton_heuristics_and_clears_cache(self, monkeypatch): + from torch._inductor.template_heuristics import registry + + from torchada import _patch + + cuda_heuristic = object() + existing_musa_heuristic = object() + heuristic_registry = { + ("triton::bmm", "cuda", None): cuda_heuristic, + ("triton::mm", "cuda", "addmm"): cuda_heuristic, + ("triton::mm", "musa", "addmm"): existing_musa_heuristic, + ("aten::mm", "cuda", None): object(), + ("triton::mm", "cpu", None): object(), + ("malformed", "cuda"): object(), + 1: object(), + } + heuristic_cache = {("cached",): object()} + monkeypatch.setattr(_patch, "is_musa_platform", lambda: True) + monkeypatch.setattr(registry, "_TEMPLATE_HEURISTIC_REGISTRY", heuristic_registry) + monkeypatch.setattr(registry, "_HEURISTIC_CACHE", heuristic_cache) + + _patch._patch_inductor_template_heuristics() + + assert heuristic_registry[("triton::bmm", "musa", None)] is cuda_heuristic + assert heuristic_registry[("triton::mm", "musa", "addmm")] is existing_musa_heuristic + assert ("aten::mm", "musa", None) not in heuristic_registry + assert ("triton::mm", "cpu", None) in heuristic_registry + assert heuristic_cache == {} + + def test_is_idempotent_and_preserves_cache_without_changes(self, monkeypatch): + from torch._inductor.template_heuristics import registry + + from torchada import _patch + + heuristic = object() + heuristic_registry = {("triton::mm", "cuda", None): heuristic} + heuristic_cache = {} + monkeypatch.setattr(_patch, "is_musa_platform", lambda: True) + monkeypatch.setattr(registry, "_TEMPLATE_HEURISTIC_REGISTRY", heuristic_registry) + monkeypatch.setattr(registry, "_HEURISTIC_CACHE", heuristic_cache) + + _patch._patch_inductor_template_heuristics() + heuristic_cache[("after-first-patch",)] = object() + _patch._patch_inductor_template_heuristics() + + assert heuristic_registry[("triton::mm", "musa", None)] is heuristic + assert ("after-first-patch",) in heuristic_cache + + def test_non_musa_platform_is_noop(self, monkeypatch): + from torch._inductor.template_heuristics import registry + + from torchada import _patch + + heuristic_registry = {("triton::mm", "cuda", None): object()} + heuristic_cache = {("cached",): object()} + monkeypatch.setattr(_patch, "is_musa_platform", lambda: False) + monkeypatch.setattr(registry, "_TEMPLATE_HEURISTIC_REGISTRY", heuristic_registry) + monkeypatch.setattr(registry, "_HEURISTIC_CACHE", heuristic_cache) + + _patch._patch_inductor_template_heuristics() + + assert ("triton::mm", "musa", None) not in heuristic_registry + assert ("cached",) in heuristic_cache + + class TestAutotuneProcess: """Test torch._inductor.autotune_process patching.""" From 94212113e13e23b3faaa7b1e31296db75ffec729 Mon Sep 17 00:00:00 2001 From: Xiaodong Ye Date: Tue, 25 Aug 2026 09:23:29 +0800 Subject: [PATCH 3/3] fix: mirror MUSA visible devices to CUDA --- src/torchada/_patch.py | 4 ++-- tests/test_cuda_patching.py | 33 ++++++++++++++++----------------- 2 files changed, 18 insertions(+), 19 deletions(-) diff --git a/src/torchada/_patch.py b/src/torchada/_patch.py index 48ab7e3..00615c8 100644 --- a/src/torchada/_patch.py +++ b/src/torchada/_patch.py @@ -69,8 +69,8 @@ def _patch_something(): def _patch_visible_devices_env(): if "MUSA_VISIBLE_DEVICES" in os.environ: os.environ["CUDA_VISIBLE_DEVICES"] = os.environ["MUSA_VISIBLE_DEVICES"] - elif "CUDA_VISIBLE_DEVICES" in os.environ: - os.environ["MUSA_VISIBLE_DEVICES"] = os.environ["CUDA_VISIBLE_DEVICES"] + else: + os.environ.pop("CUDA_VISIBLE_DEVICES", None) def requires_import(*module_names: str) -> Callable[[Callable], Callable]: diff --git a/tests/test_cuda_patching.py b/tests/test_cuda_patching.py index 4ffd8a6..af5e832 100644 --- a/tests/test_cuda_patching.py +++ b/tests/test_cuda_patching.py @@ -1087,39 +1087,38 @@ def test_musa_tensor_is_cuda_true(self): class TestVisibleDevicesEnv: """Test CUDA_VISIBLE_DEVICES and MUSA_VISIBLE_DEVICES fallback.""" - def test_cuda_visible_devices_env_falls_back_to_musa_visible_devices(self, monkeypatch): - """Test CUDA_VISIBLE_DEVICES is copied to MUSA_VISIBLE_DEVICES.""" + def test_musa_visible_devices_syncs_to_cuda_visible_devices(self, monkeypatch): + """Test CUDA_VISIBLE_DEVICES mirrors MUSA_VISIBLE_DEVICES.""" from torchada import _patch - monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "1,3") - monkeypatch.delenv("MUSA_VISIBLE_DEVICES", raising=False) + monkeypatch.setenv("MUSA_VISIBLE_DEVICES", "1,3") + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "0") _patch._patch_visible_devices_env() - assert os.environ["MUSA_VISIBLE_DEVICES"] == "1,3" + assert os.environ["CUDA_VISIBLE_DEVICES"] == "1,3" - def test_musa_visible_devices_env_falls_back_to_cuda_visible_devices(self, monkeypatch): - """Test MUSA_VISIBLE_DEVICES is copied to CUDA_VISIBLE_DEVICES.""" + def test_cuda_visible_devices_is_cleared_when_musa_is_absent(self, monkeypatch): + """Test CUDA_VISIBLE_DEVICES is removed when MUSA is not configured.""" from torchada import _patch - monkeypatch.setenv("MUSA_VISIBLE_DEVICES", "0") - monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising=False) + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "1,3") + monkeypatch.delenv("MUSA_VISIBLE_DEVICES", raising=False) _patch._patch_visible_devices_env() - assert os.environ["CUDA_VISIBLE_DEVICES"] == "0" + assert "CUDA_VISIBLE_DEVICES" not in os.environ - def test_musa_visible_devices_overrides_cuda_visible_devices(self, monkeypatch): - """Test the MUSA setting wins when both variables are explicit.""" + def test_empty_musa_visible_devices_clears_cuda_value(self, monkeypatch): + """Test an explicitly empty MUSA value is mirrored exactly.""" from torchada import _patch - monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "1,3") - monkeypatch.setenv("MUSA_VISIBLE_DEVICES", "0") + monkeypatch.setenv("MUSA_VISIBLE_DEVICES", "") + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "0") _patch._patch_visible_devices_env() - assert os.environ["CUDA_VISIBLE_DEVICES"] == "0" - assert os.environ["MUSA_VISIBLE_DEVICES"] == "0" + assert os.environ["CUDA_VISIBLE_DEVICES"] == "" def test_visible_devices_env_absent_noop(self, monkeypatch): """Test no visible device envs are added when both are absent.""" @@ -1130,8 +1129,8 @@ def test_visible_devices_env_absent_noop(self, monkeypatch): _patch._patch_visible_devices_env() - assert "CUDA_VISIBLE_DEVICES" not in os.environ assert "MUSA_VISIBLE_DEVICES" not in os.environ + assert "CUDA_VISIBLE_DEVICES" not in os.environ class TestInductorTemplateHeuristics: