diff --git a/src/paddlefleet/cudnn_ops/attn/csa_sparse_attn_bwd_cudnn.py b/src/paddlefleet/cudnn_ops/attn/csa_sparse_attn_bwd_cudnn.py index 80c570c4da..28373a3307 100644 --- a/src/paddlefleet/cudnn_ops/attn/csa_sparse_attn_bwd_cudnn.py +++ b/src/paddlefleet/cudnn_ops/attn/csa_sparse_attn_bwd_cudnn.py @@ -21,6 +21,130 @@ def _require_cudnn_frontend(): raise ImportError(CUDNN_FRONTEND_HINT) +def _patch_paddle_memory_format(): + """Patch paddle with torch.memory_format sentinel attributes for H-Card. + + The cudnn package's api_base.TensorDesc.is_contiguous() uses + torch.contiguous_format / preserve_format / channels_last / channels_last_3d + for memory layout checks. When Paddle's torch-proxy intercepts `import torch`, + these attributes are missing. We inject them as unique sentinel objects so that + identity/equality checks in cudnn work correctly. + """ + import paddle + + class _MemoryFormat: + """Minimal stand-in for torch.memory_format enum values.""" + + def __init__(self, name): + self._name = name + + def __repr__(self): + return f"paddle.{self._name}" + + def __hash__(self): + return hash(self._name) + + def __eq__(self, other): + return self is other + + for name in ( + "contiguous_format", + "preserve_format", + "channels_last", + "channels_last_3d", + ): + if not hasattr(paddle, name): + setattr(paddle, name, _MemoryFormat(name)) + + # Also need paddle.memory_format type for type annotations + if not hasattr(paddle, "memory_format"): + paddle.memory_format = type(_MemoryFormat("contiguous_format")) + + +_patch_paddle_memory_format() + + +def _patch_cutlass_nvgpu(): + """Patch cutlass.cute.nvgpu to expose OperandMajorMode at top level for H-Card. + + The cudnn SM90 backward kernel (dsa_bwd_sm90.py) references + cute.nvgpu.OperandMajorMode, but in cutlass 4.4.x this symbol lives + under cute.nvgpu.warpgroup. Promote it so the kernel can find it. + """ + try: + from cutlass.cute import nvgpu + + if not hasattr(nvgpu, "OperandMajorMode"): + from cutlass.cute.nvgpu.warpgroup import OperandMajorMode + + nvgpu.OperandMajorMode = OperandMajorMode + except (ImportError, AttributeError): + pass + + +_patch_cutlass_nvgpu() + + +def _patch_paddle_nvtx_range(): + """Patch paddle nvtx classes with a ``range`` context-manager for H-Card. + + The cudnn SM90 backward kernel uses ``torch.cuda.nvtx.range(...)`` as a + context manager. When Paddle's torch-proxy is active, ``torch.cuda.nvtx`` + resolves to ``paddle.cuda.nvtx`` (NOT paddle.device.nvtx). Both classes + only expose ``range_push`` / ``range_pop``. We inject a compatible ``range`` + into both to be safe. + """ + from contextlib import contextmanager + + import paddle + + def _inject_range(nvtx_cls): + if nvtx_cls is None or hasattr(nvtx_cls, "range"): + return + + @contextmanager + def _nvtx_range(msg: str, *args, **kwargs): + nvtx_cls.range_push(msg.format(*args, **kwargs)) + try: + yield + finally: + nvtx_cls.range_pop() + + nvtx_cls.range = staticmethod(_nvtx_range) + + _inject_range(getattr(paddle.device, "nvtx", None)) + _inject_range(getattr(paddle.cuda, "nvtx", None)) + + +_patch_paddle_nvtx_range() + + +def _patch_paddle_stream_cuda_stream(): + """Give paddle.device.Stream a torch-like ``cuda_stream`` property for H-Card. + + cudnn's resolve_stream() defaults to torch.cuda.current_stream().cuda_stream. + Under Paddle's torch-proxy that returns a paddle Stream, which exposes the + raw stream as stream_base.cuda_stream instead. Bridging the attribute makes + resolve_stream() pick up Paddle's current stream on both SM90 and SM100. + """ + import paddle + from paddle.base import core + + S = paddle.device.Stream + if not hasattr(S, "cuda_stream"): + + def _cuda_stream(self): + assert isinstance(self.stream_base, core.CUDAStream), ( + "cuda_stream is only available for CUDA streams" + ) + return self.stream_base.cuda_stream + + S.cuda_stream = property(_cuda_stream) + + +_patch_paddle_stream_cuda_stream() + + def csa_sparse_attn_bwd_cudnn( q, # (total_sq, H, D) bf16 kv, # (total_skv, D) bf16 diff --git a/tests/single_card_tests/ai_edited_test/fusions/test_csa_sparse_attn_backends.py b/tests/single_card_tests/ai_edited_test/fusions/test_csa_sparse_attn_backends.py index d03d70b8a7..0e7a5ad3e7 100644 --- a/tests/single_card_tests/ai_edited_test/fusions/test_csa_sparse_attn_backends.py +++ b/tests/single_card_tests/ai_edited_test/fusions/test_csa_sparse_attn_backends.py @@ -14,7 +14,11 @@ """Verify TileLang and cuDNN CSA sparse attention backends agree.""" +import importlib.util +import os +import sys import unittest +from unittest.mock import MagicMock, patch import paddle @@ -189,3 +193,284 @@ def test_tilelang_and_cudnn_backends_match(self): if __name__ == "__main__": unittest.main() + + +# --------------------------------------------------------------------------- +# Unit tests for import-time patches in csa_sparse_attn_bwd_cudnn.py +# (No GPU kernel dependency — always runs regardless of cudnn availability) +# --------------------------------------------------------------------------- + +# Load module under test from source tree. +# CI does editable install so coverage tracks the source file automatically. +# We use spec_from_file_location + register in sys.modules so coverage +# (source=paddlefleet) recognizes the executed file path. +_bwd_mod = None +_BWD_CUDNN_PATH = os.path.join( + os.path.dirname(__file__), + "../../../../src/paddlefleet/cudnn_ops/attn/csa_sparse_attn_bwd_cudnn.py", +) +_BWD_CUDNN_PATH = os.path.abspath(_BWD_CUDNN_PATH) +_BWD_MODULE_NAME = "paddlefleet.cudnn_ops.attn.csa_sparse_attn_bwd_cudnn" + + +def _load_bwd_module(): + """Load the bwd_cudnn module, mocking paddlefleet_ops if needed.""" + spec = importlib.util.spec_from_file_location( + _BWD_MODULE_NAME, _BWD_CUDNN_PATH + ) + mod = importlib.util.module_from_spec(spec) + sys.modules[_BWD_MODULE_NAME] = mod + spec.loader.exec_module(mod) + return mod + + +try: + _bwd_mod = _load_bwd_module() +except ImportError: + try: + _mock_pflops = MagicMock() + _mock_pflops.CUDNN_FRONTEND_HINT = "cudnn frontend not available" + _mock_pflops.is_cudnn_frontend_available = MagicMock(return_value=False) + with patch.dict(sys.modules, {"paddlefleet_ops": _mock_pflops}): + _bwd_mod = _load_bwd_module() + except (RuntimeError, AttributeError, ModuleNotFoundError, ImportError): + _bwd_mod = None +except (RuntimeError, AttributeError, ModuleNotFoundError): + _bwd_mod = None + +if _bwd_mod is not None: + _patch_cutlass_nvgpu = _bwd_mod._patch_cutlass_nvgpu + _patch_paddle_nvtx_range = _bwd_mod._patch_paddle_nvtx_range + _patch_paddle_stream_cuda_stream = _bwd_mod._patch_paddle_stream_cuda_stream + _HAS_BWD_CUDNN_MODULE = True +else: + _HAS_BWD_CUDNN_MODULE = False + _patch_cutlass_nvgpu = None + _patch_paddle_nvtx_range = None + _patch_paddle_stream_cuda_stream = None + + +@unittest.skipUnless( + _HAS_BWD_CUDNN_MODULE, "csa_sparse_attn_bwd_cudnn module unavailable" +) +class TestMemoryFormatPatch(unittest.TestCase): + """Test _MemoryFormat sentinel objects (lines 42/45/48).""" + + def test_repr(self): + fmt = paddle.contiguous_format + self.assertEqual(repr(fmt), "paddle.contiguous_format") + + def test_hash(self): + fmt = paddle.contiguous_format + self.assertEqual(hash(fmt), hash("contiguous_format")) + + def test_eq_identity(self): + fmt = paddle.contiguous_format + self.assertEqual(fmt, fmt) + + def test_eq_different_object(self): + self.assertNotEqual(paddle.contiguous_format, paddle.preserve_format) + + def test_all_formats_exist(self): + for name in ( + "contiguous_format", + "preserve_format", + "channels_last", + "channels_last_3d", + ): + self.assertTrue(hasattr(paddle, name), f"paddle.{name} missing") + + +@unittest.skipUnless( + _HAS_BWD_CUDNN_MODULE, "csa_sparse_attn_bwd_cudnn module unavailable" +) +class TestPatchCutlassNvgpuException(unittest.TestCase): + """Test _patch_cutlass_nvgpu exception path (lines 81-82).""" + + def test_import_error_is_silenced(self): + with patch.dict( + sys.modules, + { + "cutlass": None, + "cutlass.cute": None, + "cutlass.cute.nvgpu": None, + }, + ): + _patch_cutlass_nvgpu() + + def test_attribute_error_is_silenced(self): + mock_nvgpu = MagicMock(spec=[]) + mock_cute = MagicMock() + mock_cute.nvgpu = mock_nvgpu + mock_cutlass = MagicMock() + mock_cutlass.cute = mock_cute + + with patch.dict( + sys.modules, + { + "cutlass": mock_cutlass, + "cutlass.cute": mock_cute, + "cutlass.cute.nvgpu": mock_nvgpu, + "cutlass.cute.nvgpu.warpgroup": None, + }, + ): + _patch_cutlass_nvgpu() + + +@unittest.skipUnless( + _HAS_BWD_CUDNN_MODULE, "csa_sparse_attn_bwd_cudnn module unavailable" +) +class TestPatchPaddleNvtxRange(unittest.TestCase): + """Test _patch_paddle_nvtx_range logic (lines 103/107-109/111).""" + + def test_skip_when_range_exists(self): + class FakeNvtx: + range = staticmethod(lambda msg: None) + range_push = staticmethod(lambda msg: None) + range_pop = staticmethod(lambda: None) + + original_range = FakeNvtx.range + with ( + patch.object(paddle.device, "nvtx", FakeNvtx, create=True), + patch.object(paddle.cuda, "nvtx", None, create=True), + ): + _patch_paddle_nvtx_range() + self.assertIs(FakeNvtx.range, original_range) + + def test_inject_range_formats_and_pushpop(self): + class FakeNvtx: + pushed = [] + popped = 0 + + @staticmethod + def range_push(msg): + FakeNvtx.pushed.append(msg) + + @staticmethod + def range_pop(): + FakeNvtx.popped += 1 + + with ( + patch.object(paddle.device, "nvtx", FakeNvtx, create=True), + patch.object(paddle.cuda, "nvtx", None, create=True), + ): + _patch_paddle_nvtx_range() + + with FakeNvtx.range("layer_{}", 3): + pass + + self.assertEqual(FakeNvtx.pushed, ["layer_3"]) + self.assertEqual(FakeNvtx.popped, 1) + + def test_range_pop_on_exception(self): + class FakeNvtx: + popped = 0 + + @staticmethod + def range_push(msg): + pass + + @staticmethod + def range_pop(): + FakeNvtx.popped += 1 + + with ( + patch.object(paddle.device, "nvtx", FakeNvtx, create=True), + patch.object(paddle.cuda, "nvtx", None, create=True), + ): + _patch_paddle_nvtx_range() + + with self.assertRaises(RuntimeError), FakeNvtx.range("test"): + raise RuntimeError("boom") + + self.assertEqual(FakeNvtx.popped, 1) + + +@unittest.skipUnless( + _HAS_BWD_CUDNN_MODULE, "csa_sparse_attn_bwd_cudnn module unavailable" +) +class TestPatchPaddleStreamCudaStream(unittest.TestCase): + """Test _patch_paddle_stream_cuda_stream assert guard (lines 136-140).""" + + def test_cuda_stream_property_returns_int(self): + _patch_paddle_stream_cuda_stream() + s = paddle.device.Stream() + val = s.cuda_stream + self.assertIsInstance(val, int) + + def test_non_cuda_stream_raises_assertion(self): + from paddle.base import core + + _patch_paddle_stream_cuda_stream() + s = paddle.device.Stream() + + mock_base = MagicMock(spec=[]) + self.assertNotIsInstance(mock_base, core.CUDAStream) + + with patch.object( + type(s), + "stream_base", + new=property(lambda self: mock_base), + create=True, + ): + with self.assertRaises(AssertionError) as ctx: + _ = s.cuda_stream + self.assertIn("only available for CUDA streams", str(ctx.exception)) + + +@unittest.skipUnless( + _HAS_BWD_CUDNN_MODULE, "csa_sparse_attn_bwd_cudnn module unavailable" +) +class TestCsaSparseAttnBwdCudnnFunction(unittest.TestCase): + """Test csa_sparse_attn_bwd_cudnn function body (lines 148-176).""" + + def test_raises_when_cudnn_unavailable(self): + """Cover _require_cudnn_frontend raising ImportError (line 159-160).""" + csa_fn = _bwd_mod.csa_sparse_attn_bwd_cudnn + with ( + patch.object( + _bwd_mod, "is_cudnn_frontend_available", return_value=False + ), + self.assertRaises(ImportError), + ): + csa_fn(None, None, None, None, None, None, None) + + def test_calls_wrapper_and_returns_tuple(self): + """Cover the import + call + return path (lines 160-176).""" + fake_result = { + "dq": "mock_dq", + "dkv": "mock_dkv", + "d_sink": "mock_d_sink", + } + csa_fn = _bwd_mod.csa_sparse_attn_bwd_cudnn + + with ( + patch.object( + _bwd_mod, "is_cudnn_frontend_available", return_value=True + ), + patch.dict( + sys.modules, + { + "paddlefleet_ops.cudnn.deepseek_sparse_attention.sparse_attention_backward.api": MagicMock( + sparse_attention_backward_wrapper=MagicMock( + return_value=fake_result + ) + ), + }, + ), + ): + dq, dkv, d_sink = csa_fn( + "q", + "kv", + "out", + "dout", + "lse", + "attn_sink", + "topk", + softmax_scale=0.1, + topk_length=64, + ) + + self.assertEqual(dq, "mock_dq") + self.assertEqual(dkv, "mock_dkv") + self.assertEqual(d_sink, "mock_d_sink")