diff --git a/README.md b/README.md index 921ecbf..746cfa0 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.81 ``` ### Step 2: Conditional Import diff --git a/README_CN.md b/README_CN.md index ebec54c..f23d739 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.81 ``` ### 步骤 2：条件导入 diff --git a/benchmarks/benchmark_history.json b/benchmarks/benchmark_history.json index 15880af..e357bfd 100644 --- a/benchmarks/benchmark_history.json +++ b/benchmarks/benchmark_history.json @@ -3,7 +3,7 @@ "description": "Historical benchmark results for torchada performance tracking", "results": [ { - "version": "0.1.83", + "version": "0.1.81", "date": "2026-01-29", "platform": "MUSA", "pytorch_version": "2.7.1", diff --git a/pyproject.toml b/pyproject.toml index b7fe737..6f64cad 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.81" 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..e2a2b76 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.81" from . import cuda, utils diff --git a/src/torchada/_patch.py b/src/torchada/_patch.py index 1202009..9ab63c5 100644 --- a/src/torchada/_patch.py +++ b/src/torchada/_patch.py @@ -418,7 +418,9 @@ def _configure_cuda_graph_debug_dump_dir() -> Optional[str]: try: os.makedirs(dump_dir, exist_ok=True) except Exception as exc: # noqa: BLE001 - warnings.warn(f"TORCHADA_CUDA_GRAPH_DEBUG_DUMP_PATH={path_setting!r} is unusable: {exc!r}") + warnings.warn( + f"TORCHADA_CUDA_GRAPH_DEBUG_DUMP_PATH={path_setting!r} is unusable: {exc!r}" + ) _cuda_graph_debug_dump_dir = None return None @@ -678,7 +680,9 @@ def _discover_factory_functions() -> List[str]: if not names: names.update(_FALLBACK_FACTORY_FUNCTIONS) # ``*_like`` variants accept device= but are not device-injected by torch. - names |= {n + "_like" for n in tuple(names) if callable(getattr(torch, n + "_like", None))} + names |= { + n + "_like" for n in tuple(names) if callable(getattr(torch, n + "_like", None)) + } for extra in _EXTRA_FACTORY_FUNCTIONS: if callable(getattr(torch, extra, None)): names.add(extra) @@ -1212,7 +1216,9 @@ class ProfileWrapper: def __init__(self, *args, activities=None, **kwargs): translated_activities = _translate_activities(activities) - self._profiler = original_profile(*args, activities=translated_activities, **kwargs) + self._profiler = original_profile( + *args, activities=translated_activities, **kwargs + ) def __enter__(self): return self._profiler.__enter__() @@ -1303,7 +1309,9 @@ def patched_impl(self, *args, **kwargs): bound.apply_defaults() if bound.arguments.get("dispatch_key") in cuda_dispatch_key_map: - bound.arguments["dispatch_key"] = cuda_dispatch_key_map[bound.arguments["dispatch_key"]] + bound.arguments["dispatch_key"] = cuda_dispatch_key_map[ + bound.arguments["dispatch_key"] + ] return original_impl(*bound.args, **bound.kwargs) @@ -1566,7 +1574,9 @@ def _accepts_only_qv(func: Callable) -> bool: # module-level ``__getattr__`` (PEP 562), which would keep them out of # ``dir()``. flash_attn_varlen_func and flash_attn_with_kvcache are the ones # sglang's FA3 path forwards ``only_qv`` into. - candidate_names = {n for n in dir(flash_attn_interface) if n.startswith("flash_attn")} + candidate_names = { + n for n in dir(flash_attn_interface) if n.startswith("flash_attn") + } candidate_names.update( ( "flash_attn_func", @@ -1874,7 +1884,9 @@ def __getattr__(self, name): elif name in self._REMAP_ATTRS: value = getattr(self._musa_module, self._REMAP_ATTRS[name]) else: - raise AttributeError(f"module 'torch.accelerator' has no attribute '{name}'") + raise AttributeError( + f"module 'torch.accelerator' has no attribute '{name}'" + ) object.__setattr__(self, name, value) return value @@ -2002,7 +2014,9 @@ def _patch_torch_accelerator(): wrapper = _AcceleratorModuleWrapper(_original_torch_accelerator, torch.musa) - wrapper._set_override("synchronize", _make_patched_accelerator_synchronize(torch.musa)) + wrapper._set_override( + "synchronize", _make_patched_accelerator_synchronize(torch.musa) + ) device_index_cm, stream_cm = _make_accelerator_context_managers(wrapper) if not hasattr(_original_torch_accelerator, "device_index"): wrapper._set_override("device_index", device_index_cm) @@ -2030,7 +2044,9 @@ def gdc_wait(): raise NotImplementedError("tl.extra.cuda.gdc_wait is not supported on MUSA") def gdc_launch_dependents(): - raise NotImplementedError("tl.extra.cuda.gdc_launch_dependents is not supported on MUSA") + raise NotImplementedError( + "tl.extra.cuda.gdc_launch_dependents is not supported on MUSA" + ) if not hasattr(tl.extra.cuda, "gdc_wait"): tl.extra.cuda.gdc_wait = gdc_wait 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..ebcca5d 100644 --- a/tests/test_cuda_patching.py +++ b/tests/test_cuda_patching.py @@ -835,15 +835,15 @@ def test_memory_pool_functions_available_on_musa(self): pytest.skip("C++ ops extension not loaded") musa_memory = torch.musa.memory - assert hasattr( - musa_memory, "_cuda_beginAllocateCurrentThreadToPool" - ), "_cuda_beginAllocateCurrentThreadToPool not found in torch.musa.memory" - assert hasattr( - musa_memory, "_cuda_endAllocateToPool" - ), "_cuda_endAllocateToPool not found in torch.musa.memory" - assert hasattr( - musa_memory, "_cuda_releasePool" - ), "_cuda_releasePool not found in torch.musa.memory" + assert hasattr(musa_memory, "_cuda_beginAllocateCurrentThreadToPool"), ( + "_cuda_beginAllocateCurrentThreadToPool not found in torch.musa.memory" + ) + assert hasattr(musa_memory, "_cuda_endAllocateToPool"), ( + "_cuda_endAllocateToPool not found in torch.musa.memory" + ) + assert hasattr(musa_memory, "_cuda_releasePool"), ( + "_cuda_releasePool not found in torch.musa.memory" + ) def test_memory_pool_functions_importable_from_cuda_memory(self): """Test that CUDA memory pool functions can be imported from torch.cuda.memory. @@ -955,14 +955,18 @@ def test_generator_isinstance_check(self): gen_cuda = torch.Generator(device="cuda") # isinstance should work for both CPU and CUDA/MUSA generators - assert isinstance(gen_cpu, torch.Generator), "CPU generator isinstance check failed" - assert isinstance(gen_cuda, torch.Generator), "CUDA/MUSA generator isinstance check failed" + assert isinstance(gen_cpu, torch.Generator), ( + "CPU generator isinstance check failed" + ) + assert isinstance(gen_cuda, torch.Generator), ( + "CUDA/MUSA generator isinstance check failed" + ) # Test with a list (like sglang's generator_or_list_generators does) gens = [gen_cpu, gen_cuda] - assert all( - isinstance(g, torch.Generator) for g in gens - ), "List of generators isinstance check failed" + assert all(isinstance(g, torch.Generator) for g in gens), ( + "List of generators isinstance check failed" + ) def test_generator_isinstance_negative(self): """Test isinstance returns False for non-Generator objects.""" @@ -1097,7 +1101,9 @@ def test_cuda_visible_devices_patched(self): # CUDA_VISIBLE_DEVICES constant was added in PyTorch 2.2.0 # It does NOT exist in PyTorch 2.1.x and earlier versions if not hasattr(autotune_process, "CUDA_VISIBLE_DEVICES"): - pytest.skip("CUDA_VISIBLE_DEVICES not available (requires PyTorch >= 2.2.0)") + pytest.skip( + "CUDA_VISIBLE_DEVICES not available (requires PyTorch >= 2.2.0)" + ) if torchada.is_musa_platform(): # On MUSA platform, CUDA_VISIBLE_DEVICES should be patched to MUSA_VISIBLE_DEVICES @@ -1117,7 +1123,9 @@ def test_cuda_visible_devices_is_string(self): # CUDA_VISIBLE_DEVICES constant was added in PyTorch 2.2.0 # It does NOT exist in PyTorch 2.1.x and earlier versions if not hasattr(autotune_process, "CUDA_VISIBLE_DEVICES"): - pytest.skip("CUDA_VISIBLE_DEVICES not available (requires PyTorch >= 2.2.0)") + pytest.skip( + "CUDA_VISIBLE_DEVICES not available (requires PyTorch >= 2.2.0)" + ) assert isinstance(autotune_process.CUDA_VISIBLE_DEVICES, str) @@ -1132,7 +1140,9 @@ def test_cuda_visible_devices_env_var_format(self): # CUDA_VISIBLE_DEVICES constant was added in PyTorch 2.2.0 # It does NOT exist in PyTorch 2.1.x and earlier versions if not hasattr(autotune_process, "CUDA_VISIBLE_DEVICES"): - pytest.skip("CUDA_VISIBLE_DEVICES not available (requires PyTorch >= 2.2.0)") + pytest.skip( + "CUDA_VISIBLE_DEVICES not available (requires PyTorch >= 2.2.0)" + ) env_var = autotune_process.CUDA_VISIBLE_DEVICES # Env var should be uppercase and use underscores @@ -1340,7 +1350,9 @@ def test_torch_backends_cuda_matmul_fp32_precision(self): try: _ = torch.backends.cuda.matmul.fp32_precision except (AttributeError, AssertionError): - pytest.skip("fp32_precision not available (torchada MUSA-specific attribute)") + pytest.skip( + "fp32_precision not available (torchada MUSA-specific attribute)" + ) if torch.__version__ >= torch.torch_version.TorchVersion("2.9.0"): # PyTorch 2.9+: Only use the new API. Do NOT call torch.get_float32_matmul_precision() @@ -1354,9 +1366,9 @@ def test_torch_backends_cuda_matmul_fp32_precision(self): # Save original state original = torch.backends.cuda.matmul.fp32_precision - assert ( - original in valid_precisions - ), f"fp32_precision value '{original}' not in expected set {valid_precisions}" + assert original in valid_precisions, ( + f"fp32_precision value '{original}' not in expected set {valid_precisions}" + ) # Test setting values for test_value in test_values: @@ -1696,14 +1708,20 @@ def test_library_impl_with_keyset_and_with_allow_override(self): lib_name = f"test_lib_{uuid.uuid4().hex[:8]}" test_lib = torch.library.Library(lib_name, "DEF") - def identity_allow_override_with_keyset_1(keyset, x: torch.Tensor) -> torch.Tensor: + def identity_allow_override_with_keyset_1( + keyset, x: torch.Tensor + ) -> torch.Tensor: return x - def doubled_allow_override_with_keyset_2(keyset, x: torch.Tensor) -> torch.Tensor: + def doubled_allow_override_with_keyset_2( + keyset, x: torch.Tensor + ) -> torch.Tensor: return 2 * x test_lib.define("test_op(Tensor x) -> Tensor") - test_lib.impl("test_op", identity_allow_override_with_keyset_1, "CPU", with_keyset=True) + test_lib.impl( + "test_op", identity_allow_override_with_keyset_1, "CPU", with_keyset=True + ) test_lib.impl( "test_op", doubled_allow_override_with_keyset_2, @@ -1792,7 +1810,9 @@ def identity_with_keyset(keyset, x: torch.Tensor) -> torch.Tensor: test_lib.define("identity_op(Tensor x) -> Tensor") # Use CUDA with with_keyset=True - test_lib.impl("identity_op", identity_with_keyset, dispatch_key="CUDA", with_keyset=True) + test_lib.impl( + "identity_op", identity_with_keyset, dispatch_key="CUDA", with_keyset=True + ) x = torch.randn(3).musa() op = getattr(torch.ops, lib_name) @@ -2148,7 +2168,9 @@ def test_factory_functions_wrapped_and_cacheable(self): safelist = None if safelist is not None: for name in names: - assert f"torch.{name}" in safelist, f"torch.{name} not registered cacheable" + assert f"torch.{name}" in safelist, ( + f"torch.{name} not registered cacheable" + ) def test_factory_set_discovered_at_runtime(self): """Wrapping is driven by torch's device-constructor registry at runtime, @@ -2163,10 +2185,19 @@ def test_factory_set_discovered_at_runtime(self): registry_only = [ n - for n in ("sparse_coo_tensor", "scalar_tensor", "vander", "logspace", "tril_indices") - if n not in _FALLBACK_FACTORY_FUNCTIONS and callable(getattr(torch, n, None)) + for n in ( + "sparse_coo_tensor", + "scalar_tensor", + "vander", + "logspace", + "tril_indices", + ) + if n not in _FALLBACK_FACTORY_FUNCTIONS + and callable(getattr(torch, n, None)) ] - assert registry_only, "no known registry-only factory present on this torch build" + assert registry_only, ( + "no known registry-only factory present on this torch build" + ) assert any(hasattr(getattr(torch, n), "__wrapped__") for n in registry_only), ( f"none of {registry_only} wrapped — discovery fell back to the static list " "instead of reading torch._device_constructors()" @@ -2210,7 +2241,10 @@ def test_compiled_factory_graph_is_cacheable(self): pytest.skip("Only applicable on MUSA platform") if os.environ.get("TORCHDYNAMO_DISABLE") == "1": pytest.skip("Dynamo disabled in environment") - if getattr(getattr(torch, "compiler", None), "save_cache_artifacts", None) is None: + if ( + getattr(getattr(torch, "compiler", None), "save_cache_artifacts", None) + is None + ): pytest.skip("torch.compiler.save_cache_artifacts not available") # Other tests in this module leak torch.library state that breaks any @@ -2267,7 +2301,9 @@ def test_zeros_with_cuda_device(self): except RuntimeError as e: # Skip if MUDNN kernel execution fails (expected in test containers) if "MUDNN" in str(e) or "invalid device function" in str(e): - pytest.skip("MUDNN kernel execution failed (expected in test containers)") + pytest.skip( + "MUDNN kernel execution failed (expected in test containers)" + ) raise @pytest.mark.gpu @@ -2286,7 +2322,9 @@ def test_ones_with_cuda_device(self): except RuntimeError as e: # Skip if MUDNN kernel execution fails (expected in test containers) if "MUDNN" in str(e) or "invalid device function" in str(e): - pytest.skip("MUDNN kernel execution failed (expected in test containers)") + pytest.skip( + "MUDNN kernel execution failed (expected in test containers)" + ) raise @pytest.mark.gpu @@ -2305,7 +2343,9 @@ def test_randn_with_cuda_device(self): except RuntimeError as e: # Skip if MUDNN kernel execution fails (expected in test containers) if "MUDNN" in str(e) or "invalid device function" in str(e): - pytest.skip("MUDNN kernel execution failed (expected in test containers)") + pytest.skip( + "MUDNN kernel execution failed (expected in test containers)" + ) raise @pytest.mark.gpu @@ -2404,7 +2444,9 @@ def test_musa_device_should_pass(self): v = torch.randn(2, 2, device="musa") except RuntimeError as e: if "MUDNN" in str(e) or "invalid device function" in str(e): - pytest.skip("MUDNN kernel execution failed (expected in test containers)") + pytest.skip( + "MUDNN kernel execution failed (expected in test containers)" + ) raise validator = self.get_validator() @@ -2652,7 +2694,9 @@ def test_multiple_flash_attn_functions_accessible(self): ] for func_name in expected_funcs: assert hasattr(sgl_flash_attn, func_name), f"Missing {func_name}" - assert callable(getattr(sgl_flash_attn, func_name)), f"{func_name} not callable" + assert callable(getattr(sgl_flash_attn, func_name)), ( + f"{func_name} not callable" + ) def test_only_qv_argument_is_dropped(self): """Newer sglang FA3 callers forward an ``only_qv`` keyword down into the @@ -2719,7 +2763,8 @@ def flash_attn_kwargs_func(q, **kwargs): # varlen kernel had no only_qv -> wrapped, argument dropped, the # rest of the call forwarded unchanged. assert ( - fake_fai.flash_attn_varlen_func(1, 2, 3, causal=True, only_qv=True) == "varlen-ok" + fake_fai.flash_attn_varlen_func(1, 2, 3, causal=True, only_qv=True) + == "varlen-ok" ) assert calls["varlen"] == {"q": 1, "k": 2, "v": 3, "causal": True} @@ -3375,7 +3420,9 @@ def test_dispatch_shim_header_present(self): from torchada.utils.cpp_extension import stable_compat_include_dir - p = os.path.join(stable_compat_include_dir(), "torch", "headeronly", "core", "Dispatch.h") + p = os.path.join( + stable_compat_include_dir(), "torch", "headeronly", "core", "Dispatch.h" + ) assert os.path.isfile(p), f"Dispatch.h shim missing: {p}" text = open(p, encoding="utf-8").read() # The THO_DISPATCH_* macros vLLM/SGLang stable kernels use. @@ -3389,7 +3436,9 @@ def test_device_shim_header_present(self): from torchada.utils.cpp_extension import stable_compat_include_dir - p = os.path.join(stable_compat_include_dir(), "torch", "csrc", "stable", "device.h") + p = os.path.join( + stable_compat_include_dir(), "torch", "csrc", "stable", "device.h" + ) assert os.path.isfile(p), f"device.h shim missing: {p}" text = open(p, encoding="utf-8").read() assert "#include_next " in text @@ -3408,7 +3457,9 @@ def test_device_shim_members_are_initialized(self): from torchada.utils.cpp_extension import stable_compat_include_dir text = open( - os.path.join(stable_compat_include_dir(), "torch", "csrc", "stable", "device.h"), + os.path.join( + stable_compat_include_dir(), "torch", "csrc", "stable", "device.h" + ), encoding="utf-8", ).read() assert "int32_t type_;" not in text, "type_ left uninitialized" @@ -3422,7 +3473,9 @@ def test_macros_shim_header_present(self): from torchada.utils.cpp_extension import stable_compat_include_dir - p = os.path.join(stable_compat_include_dir(), "torch", "csrc", "stable", "macros.h") + p = os.path.join( + stable_compat_include_dir(), "torch", "csrc", "stable", "macros.h" + ) assert os.path.isfile(p), f"macros.h shim missing: {p}" assert "library.h" in open(p, encoding="utf-8").read() @@ -3530,7 +3583,10 @@ def test_cuda_bf16_shim_aliases_nv_types(self): def test_include_paths_appends_stable_compat_on_musa(self): """include_paths() auto-appends the stable_compat dir for device builds on MUSA.""" import torchada - from torchada.utils.cpp_extension import include_paths, stable_compat_include_dir + from torchada.utils.cpp_extension import ( + include_paths, + stable_compat_include_dir, + ) if not torchada.is_musa_platform(): pytest.skip("Only applicable on MUSA platform") @@ -3543,7 +3599,10 @@ def test_include_paths_appends_stable_compat_on_musa(self): def test_include_paths_omits_stable_compat_for_cpu_only(self): """include_paths(device_type='cpu') does not add the device stable_compat dir.""" import torchada - from torchada.utils.cpp_extension import include_paths, stable_compat_include_dir + from torchada.utils.cpp_extension import ( + include_paths, + stable_compat_include_dir, + ) if not torchada.is_musa_platform(): pytest.skip("Only applicable on MUSA platform") @@ -3697,3 +3756,37 @@ def test_ensure_stable_headers_patched_noop_off_musa(self): # Should return immediately (is_musa_platform() guard) without error. ce._ensure_stable_headers_patched() + + +class TestVisibleDevicesEnv: + """Test visible device env var translation.""" + + def test_musa_visible_devices_backfills_cuda(self, monkeypatch): + import importlib + import os + + monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising=False) + monkeypatch.setenv("MUSA_VISIBLE_DEVICES", "4,5") + + import torchada._patch as torchada_patch + + importlib.reload(torchada_patch) + torchada_patch._patch_visible_devices_env() + + assert os.environ["MUSA_VISIBLE_DEVICES"] == "4,5" + assert os.environ["CUDA_VISIBLE_DEVICES"] == "4,5" + + def test_cuda_visible_devices_backfills_musa(self, monkeypatch): + import importlib + import os + + monkeypatch.delenv("MUSA_VISIBLE_DEVICES", raising=False) + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "2,7") + + import torchada._patch as torchada_patch + + importlib.reload(torchada_patch) + torchada_patch._patch_visible_devices_env() + + assert os.environ["CUDA_VISIBLE_DEVICES"] == "2,7" + assert os.environ["MUSA_VISIBLE_DEVICES"] == "2,7" 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..4084c24 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.81" assert isinstance(version, str) def test_project_version_matches_runtime_version(self): @@ -68,7 +68,9 @@ def test_project_version_matches_runtime_version(self): pyproject = Path(__file__).resolve().parents[1] / "pyproject.toml" match = re.search( - r'^version = "([^"]+)"$', pyproject.read_text(encoding="utf-8"), re.MULTILINE + r'^version = "([^"]+)"$', + pyproject.read_text(encoding="utf-8"), + re.MULTILINE, ) assert match is not None