diff --git a/README.md b/README.md index e8b73ca..b7b7f3c 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.81 +torchada>=0.1.82 ``` ### Step 2: Conditional Import diff --git a/README_CN.md b/README_CN.md index f5e0e47..db0d447 100644 --- a/README_CN.md +++ b/README_CN.md @@ -375,7 +375,7 @@ if torchada.is_gpu_device(device): # 在 CUDA 和 MUSA 上都能工作 ``` # pyproject.toml 或 requirements.txt -torchada>=0.1.81 +torchada>=0.1.82 ``` ### 步骤 2：条件导入 diff --git a/benchmarks/benchmark_history.json b/benchmarks/benchmark_history.json index e357bfd..b404a83 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.81", + "version": "0.1.82", "date": "2026-01-29", "platform": "MUSA", "pytorch_version": "2.7.1", diff --git a/pyproject.toml b/pyproject.toml index 6f64cad..e887f98 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "torchada" -version = "0.1.81" +version = "0.1.82" 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 e2a2b76..3b871a6 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.81" +__version__ = "0.1.82" from . import cuda, utils diff --git a/src/torchada/_patch.py b/src/torchada/_patch.py index b008553..22b93de 100644 --- a/src/torchada/_patch.py +++ b/src/torchada/_patch.py @@ -1587,6 +1587,55 @@ def _accepts_only_qv(func: Callable) -> bool: continue setattr(flash_attn_interface, _name, _drop_only_qv(_func)) + # Ring Attention imports the private forward API for its softmax LSE. + # Adapt the public output+LSE contract only when the provider omits it. + if not hasattr(flash_attn_interface, "_flash_attn_forward"): + _public_flash_attn = getattr(flash_attn_interface, "flash_attn_func", None) + try: + _public_parameters = inspect.signature(_public_flash_attn).parameters + except (TypeError, ValueError): + _public_parameters = {} + + if callable(_public_flash_attn) and "return_softmax_lse" in _public_parameters: + + def _flash_attn_forward( + q, + k, + v, + *, + softmax_scale=None, + causal=False, + window_size_left=-1, + window_size_right=-1, + softcap=0.0, + ): + """Adapt public FA3 output+LSE to its low-level Ring contract.""" + result = _public_flash_attn( + q, + k, + v, + softmax_scale=softmax_scale, + causal=causal, + window_size=(window_size_left, window_size_right), + softcap=softcap, + return_softmax_lse=True, + ) + if not isinstance(result, (tuple, list)) or len(result) < 2: + raise RuntimeError( + "flash_attn_func(return_softmax_lse=True) must return " + "(output, softmax_lse)" + ) + output, softmax_lse = result[:2] + # Inference-only Ring consumers use output and LSE; the public + # API does not expose the private dropout auxiliaries. + return output, softmax_lse, None, None + + _flash_attn_forward.__name__ = "_flash_attn_forward" + _flash_attn_forward.__qualname__ = "_flash_attn_forward" + _flash_attn_forward.__module__ = flash_attn_interface.__name__ + _flash_attn_forward._torchada_compat_shim = True + flash_attn_interface._flash_attn_forward = _flash_attn_forward + class _CDLLWrapper: """ diff --git a/tests/test_cuda_patching.py b/tests/test_cuda_patching.py index 67f904d..2341b12 100644 --- a/tests/test_cuda_patching.py +++ b/tests/test_cuda_patching.py @@ -2755,6 +2755,128 @@ def flash_attn_kwargs_func(q, **kwargs): if mod is not None: sys.modules[name] = mod + @staticmethod + def _install_fa3_test_provider(monkeypatch, flash_attn_func, native_forward=None): + import sys + from types import ModuleType + + flash_attn = ModuleType("flash_attn_interface") + flash_attn.flash_attn_func = flash_attn_func + if native_forward is not None: + flash_attn._flash_attn_forward = native_forward + sgl_kernel = ModuleType("sgl_kernel") + sgl_kernel.__path__ = [] + sgl_kernel.flash_attn = flash_attn + monkeypatch.setitem(sys.modules, "flash_attn_interface", flash_attn) + monkeypatch.setitem(sys.modules, "sgl_kernel", sgl_kernel) + monkeypatch.setitem(sys.modules, "sgl_kernel.flash_attn", flash_attn) + return flash_attn + + def test_missing_fa3_private_forward_is_adapted(self, monkeypatch): + """Adapt the public output+LSE API to Ring's private FA3 contract.""" + from torchada._patch import _patch_flash_attn + + calls = [] + + def flash_attn_func( + q, + k, + v, + *, + softmax_scale=None, + causal=False, + window_size=(-1, -1), + softcap=0.0, + return_softmax_lse=False, + ): + calls.append( + { + "q": q, + "k": k, + "v": v, + "softmax_scale": softmax_scale, + "causal": causal, + "window_size": window_size, + "softcap": softcap, + "return_softmax_lse": return_softmax_lse, + } + ) + return "output", "softmax_lse" + + flash_attn = self._install_fa3_test_provider(monkeypatch, flash_attn_func) + + _patch_flash_attn() + from flash_attn_interface import _flash_attn_forward + + assert _flash_attn_forward is flash_attn._flash_attn_forward + result = flash_attn._flash_attn_forward( + "q", + "k", + "v", + softmax_scale=0.125, + causal=True, + window_size_left=32, + window_size_right=16, + softcap=4.0, + ) + + assert result == ("output", "softmax_lse", None, None) + assert calls == [ + { + "q": "q", + "k": "k", + "v": "v", + "softmax_scale": 0.125, + "causal": True, + "window_size": (32, 16), + "softcap": 4.0, + "return_softmax_lse": True, + } + ] + assert flash_attn._flash_attn_forward.__name__ == "_flash_attn_forward" + assert flash_attn._flash_attn_forward._torchada_compat_shim is True + first = flash_attn._flash_attn_forward + _patch_flash_attn() + assert flash_attn._flash_attn_forward is first + + def native_forward(q, k, v, **kwargs): + return q, k, v, kwargs + + native_provider = self._install_fa3_test_provider( + monkeypatch, flash_attn_func, native_forward + ) + _patch_flash_attn() + assert native_provider._flash_attn_forward is native_forward + + def test_fa3_private_forward_fails_closed_without_lse(self, monkeypatch): + """Require both a public LSE parameter and a valid LSE result.""" + from torchada._patch import _patch_flash_attn + + def no_lse_support(q, k, v): + return q + + flash_attn = self._install_fa3_test_provider(monkeypatch, no_lse_support) + _patch_flash_attn() + assert not hasattr(flash_attn, "_flash_attn_forward") + + def invalid_lse_result( + q, + k, + v, + *, + softmax_scale=None, + causal=False, + window_size=(-1, -1), + softcap=0.0, + return_softmax_lse=False, + ): + return q + + flash_attn = self._install_fa3_test_provider(monkeypatch, invalid_lse_result) + _patch_flash_attn() + with pytest.raises(RuntimeError, match="must return"): + flash_attn._flash_attn_forward("q", "k", "v") + class TestAcceleratorModuleWrapper: """Test the _AcceleratorModuleWrapper priority / fallback logic in isolation. diff --git a/tests/test_platform.py b/tests/test_platform.py index 56b0be3..5847fb2 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.81" + assert version == "0.1.82" assert isinstance(version, str) def test_project_version_matches_runtime_version(self):