Port the PaddlePaddle deltas from paddle/v1.27.0 onto 1.28.0

python/cudnn/__init__.py

@@ -388,6 +388,11 @@ def _load_optional_symbol(name: str) -> Any:
     return value
 
 
+# Import some modules eagerly
+from cudnn import api_base as api_base
+from cudnn.deepseek_sparse_attention import DSA as DSA
+
+
 def __getattr__(name: str) -> Any:
     if name in ("Graph", "wrapper"):
         _wrapper = importlib.import_module(".wrapper", __name__)

python/cudnn/api_base.py

@@ -118,8 +118,12 @@ def __post_init__(self):
         stride = tuple(self.stride)
         stride_order = tuple(self.stride_order)
         device = self.device
-        torch = _torch()
-        if torch is not None and not isinstance(device, (torch.device, Device)):
+        # ``_torch()`` probes ``sys.modules["torch"]``, which on this branch is
+        # either absent or a real PyTorch installed alongside Paddle -- neither
+        # describes a Paddle device.  Probe Paddle directly instead.
+        import paddle as torch
+
+        if not isinstance(device, (torch.device.Device, Device)):
             try:
                 device = torch.device(device)
             except (TypeError, ValueError, RuntimeError) as exc:
@@ -661,6 +665,7 @@ def _is_fp4x2(self, tensor_or_dtype: torch.Tensor | torch.dtype | TensorDesc) ->
         :return: True if tensor/dtype is an FP4x2 packed type
         :rtype: bool
         """
+        return False  # Paddle not support FP4 dtype now
         if tensor_or_dtype is None:
             return False
         if isinstance(tensor_or_dtype, TensorDesc):

python/cudnn/datatypes.py

@@ -28,7 +28,7 @@ def is_torch_available():
     # this condition ensures that datatype mapping is only created once
     if torch_available is None:
         try:
-            import torch
+            import paddle as torch
 
             torch_available = True
             _torch_to_cudnn_data_type_dict = {
@@ -80,7 +80,7 @@ def _is_torch_to_cutlass_available():
     global _torch_to_cutlass_data_type_dict
     if _torch_to_cutlass_data_type_dict is None:
         try:
-            import torch
+            import paddle as torch
             import cutlass
 
             mapping = {
@@ -231,9 +231,12 @@ def _buffer_dtype_to_cudnn(dtype) -> cudnn_data_type:
 
 
 def _torch_to_cutlass_data_type(data_type, interpret_uint8_as_fp4x2: bool = False):
-    # A torch dtype can only be passed in if torch is already imported, so probing
-    # sys.modules avoids importing torch on behalf of other frameworks' dtypes.
-    torch = sys.modules.get("torch")
+    # A framework dtype can only be passed in if the framework is already
+    # imported, so probing sys.modules avoids importing it on behalf of other
+    # frameworks' dtypes.  On this branch the framework is Paddle, and a real
+    # ``torch`` may also be installed alongside it -- probing "torch" would then
+    # hand back PyTorch and every Paddle dtype would miss the mapping below.
+    torch = sys.modules.get("paddle")
     if torch is None or not isinstance(data_type, torch.dtype):
         return None
     if is_cutlass_available() and _is_torch_to_cutlass_available():
@@ -333,7 +336,7 @@ def _library_type(input_type):
 
 def _is_torch_tensor(input_tensor) -> bool:
     if is_torch_available():
-        import torch
+        import paddle as torch
 
         return isinstance(input_tensor, torch.Tensor)
     return False

python/cudnn/deepseek_sparse_attention/__init__.py

@@ -47,10 +47,14 @@ def __getattr__(name):
 
 
 class DSANamespace:
-    def __getattr__(self, name):
-        if name in _SYMBOLS:
-            return _load_symbol(name)
-        raise AttributeError(f"DSA has no attribute {name!r}")
+    # def __getattr__(self, name):
+    #     if name in _SYMBOLS:
+    #         return _load_symbol(name)
+    #     raise AttributeError(f"DSA has no attribute {name!r}")
+    # Make import all symbols eagerly
+    def __init__(self):
+        for symbol in _SYMBOLS:
+            setattr(self, _SYMBOLS[symbol][1], _load_symbol(symbol))
 
 
 DSA = DSANamespace()

python/cudnn/deepseek_sparse_attention/indexer_backward/api.py

@@ -61,9 +61,21 @@ def _validate_indexer_backward_backend(backend: str) -> None:
 def _validate_grad_loss_tensor(grad_loss: torch.Tensor, device: torch.device) -> torch.Tensor:
     if not torch.is_tensor(grad_loss):
         raise TypeError("grad_loss must be a torch.Tensor")
-    if grad_loss.numel() != 1 or grad_loss.dtype != torch.float32 or grad_loss.device != device:
+    # The element count comes from ``shape`` rather than ``numel()`` so the guard
+    # stays host-side. Under a torch-compat proxy (Paddle's
+    # ``enable_compat(scope={"cudnn"})``, which is how PaddleFleet drives these
+    # kernels) ``numel()`` is an *op* returning a 0-D tensor, so ``!= 1`` builds a
+    # device bool and the Python ``or`` forces ``__bool__`` -> a blocking
+    # device-to-host copy on the legacy default stream. That copy is a
+    # full-device barrier there, and it lands between the caller's loss forward
+    # and this backward, serialising the backward against whatever collective is
+    # in flight. ``shape`` is host metadata under both frameworks.
+    grad_loss_numel = 1
+    for dim in grad_loss.shape:
+        grad_loss_numel *= int(dim)
+    if grad_loss_numel != 1 or grad_loss.dtype != torch.float32 or grad_loss.device != device:
         raise ValueError(f"grad_loss must be a single-element float32 tensor on {device}")
-    return grad_loss.detach().view(1)
+    return grad_loss.detach().reshape([1])
 
 
 def _contiguous_input(tensor: torch.Tensor) -> torch.Tensor:
@@ -680,6 +692,20 @@ def check_support(self) -> bool:
         self._value_error_if(self.block_I <= 0, f"block_I must be positive, got {self.block_I}")
         self._value_error_if(self.ratio < 1, f"ratio must be >= 1, got {self.ratio}")
         self._value_error_if(self.heads < 64, f"DenseIndexerBackward requires heads >= 64, got {self.heads}")
+        # The dK epilogue stages a (block_I, head_dim_padded) FP32 tile in SMEM and
+        # ships it with one cp.reduce.async.bulk of
+        # ``actual_rows * head_dim_padded * 4`` bytes, which silently assumes the
+        # global dK row stride equals head_dim_padded.  On top of that the TMEM
+        # load atom (Ld16x256b, Repetition(8)) only tiles the staging buffer
+        # completely at the widths it was tuned for.  Measured on SM100-class
+        # hardware: head_dim 64 and 128 are correct, 96 / 100 / 112 return
+        # silently wrong d_index_k (relative error up to 4e2, and head_dim=100
+        # also corrupts d_index_q with NaN / 1e38), and 192 / 256 abort with
+        # cudaErrorInvalidValue.  Reject instead of returning bad gradients.
+        self._value_error_if(
+            self.head_dim not in (64, 128),
+            f"DenseIndexerBackward supports head_dim 64 or 128, got {self.head_dim}",
+        )
         self._is_supported = True
         return True
 

python/cudnn/deepseek_sparse_attention/indexer_backward/dense_indexer_backward_sm100.py

@@ -83,6 +83,7 @@
     resolve_stream as _resolve_stream,
 )
 from cudnn.deepseek_sparse_attention.utils.seqlen import seqlen_info as _seqlen_info
+from cudnn.deepseek_sparse_attention.utils.tensor_conversion import to_cute_tensor
 
 mul_packed_f32x2 = partial(cute.arch.mul_packed_f32x2, rnd="rn")
 fma_packed_f32x2 = partial(cute.arch.fma_packed_f32x2, rnd="rn")
@@ -430,6 +431,12 @@ def __init__(self, head_dim, heads=64, block_I=128, ratio=1):
             barrier_id=4,
             num_threads=self.WARP_SIZE + 2 * self.WARPGROUP_SIZE,
         )
+        # Reduce warpgroup (warps 8-11) only.  ids 0/3/4 are taken by
+        # sync_threads, compute_sync_barrier and tmem_alloc_barrier.
+        self.dk_reduce_barrier = pipeline.NamedBarrier(
+            barrier_id=5,
+            num_threads=self.WARPGROUP_SIZE,
+        )
 
     @cute.jit
     def __call__(
@@ -1882,9 +1889,14 @@ def _reduce_warpgroup_2q(
             # 3. Signal DK_EMPTY immediately after T2R (single-buffered)
             cute.arch.mbarrier_arrive(mbar + MBAR_2Q_DK_EMPTY)
 
-            # 4. Wait for previous bulk reduce to finish, then signal TMA engine is free
+            # 4. Wait for previous bulk reduce to finish before ANY lane
+            # overwrites the single-buffered staging tile.  cp.async.bulk groups
+            # are per-thread, so only the issuing thread's wait is meaningful —
+            # the barrier is what extends it to the other 127 lanes.
             if bi > 0:
-                cute.arch.cp_async_bulk_wait_group(0, read=True)
+                if wg_tidx == 0:
+                    cute.arch.cp_async_bulk_wait_group(0, read=True)
+                self.dk_reduce_barrier.arrive_and_wait()
 
             # 5. Scatter-write: registers → sdK_reduce
             for pair in cutlass.range(cute.size(tDKrDK) // 2, unroll_full=True):
@@ -1896,7 +1908,16 @@ def _reduce_warpgroup_2q(
                     sdK_reduce[n, d] = tDKrDK[ei] * Float32(sm_scale)
                     sdK_reduce[n, d + 1] = tDKrDK[ei + 1] * Float32(sm_scale)
 
+            # All 128 lanes write disjoint [n, d] slots, but the bulk DMA below
+            # is issued by ONE thread and reads the whole tile.  fence_proxy is
+            # a proxy fence: it only orders the *executing* thread's prior
+            # generic-proxy SMEM writes against the async proxy, and neither
+            # waits for nor publishes any other lane's stores.  A real
+            # warpgroup barrier is required on both sides of it — same pattern
+            # as the dQ epilogue's compute_sync_barrier above.
+            self.dk_reduce_barrier.arrive_and_wait()
             cute.arch.fence_proxy("async.shared", space="cta")
+            self.dk_reduce_barrier.arrive_and_wait()
 
             # 6. Single-thread bulk reduce DMA — only ONE thread in the
             # reduce warpgroup must issue cp.async.bulk; otherwise each
@@ -1991,7 +2012,6 @@ def dense_indexer_backward_sm100(
 
 
 def _build_cute_dsl_kernel(batch, max_seqlen_q, max_seqlen_k, heads, dim, sm_scale, block_I, ratio, is_varlen, has_q_causal_offsets):
-    from cudnn.deepseek_sparse_attention.utils.tensor_conversion import to_cute_tensor
 
     if torch.cuda.get_device_capability()[0] < 10:
         raise RuntimeError("Requires SM100+")

python/cudnn/deepseek_sparse_attention/indexer_backward/indexer_backward_sm100.py

@@ -56,7 +56,6 @@
 from functools import partial
 import torch
 import cuda.bindings.driver as cuda
-
 import cutlass
 import cutlass.cute as cute
 from cutlass import Float32, Int32, const_expr
@@ -79,6 +78,7 @@
     resolve_stream as _resolve_stream,
     torch_stream_context as _torch_stream_context,
 )
+from cudnn.deepseek_sparse_attention.utils.tensor_conversion import to_cute_tensor
 
 mul_packed_f32x2 = partial(cute.arch.mul_packed_f32x2, rnd="rn")
 fma_packed_f32x2 = partial(cute.arch.fma_packed_f32x2, rnd="rn")
@@ -1544,7 +1544,6 @@ class SharedStorage:
 
 
 def _score_grad_inplace_cute(AttnScore, IndexScore, GradLoss, grad_scale, current_stream=None):
-    from cudnn.deepseek_sparse_attention.utils.tensor_conversion import to_cute_tensor
 
     # Kernel reads ``mGradLoss[0]`` so it must be at least 1-D. ``to_cute_tensor``
     # defaults ``leading_dim = ndim - 1`` which collapses to -1 for a 0-D scalar
@@ -1609,7 +1608,6 @@ def _score_grad_inplace(AttnScore, IndexScore, GradLoss, grad_scale, block_I=128
 
 
 def _build_cute_dsl_kernel(heads, dim, topk, sm_scale, block_I, topk_indices_global: bool = True):
-    from cudnn.deepseek_sparse_attention.utils.tensor_conversion import to_cute_tensor
 
     if torch.cuda.get_device_capability()[0] < 10:
         raise RuntimeError("Requires SM100+")
@@ -1682,7 +1680,7 @@ def _run(IndexQ, Weights, IndexK, dIndexQ, dWeights, dIndexK, AttnScore, IndexSc
                 dIndexK_f32 = torch.zeros_like(dIndexK, dtype=torch.float32)
             _run_gemm_only(IndexQ, Weights, IndexK, dIndexQ, dWeights, dIndexK_f32, AttnScore, TopkIndices, current_stream=current_stream)
             with _torch_stream_context(current_stream):
-                dIndexK.copy_(dIndexK_f32)
+                dIndexK.copy_(dIndexK_f32.astype(dIndexK.dtype))
 
     _run.score_grad = partial(_score_grad_inplace, block_I=block_I)
     _run.gemm_only = _run_gemm_only

python/cudnn/deepseek_sparse_attention/indexer_backward/indexer_backward_sm90.py

@@ -1744,7 +1744,7 @@ def _run(
                 current_stream=current_stream,
             )
             with _torch_stream_context(current_stream):
-                dIndexK.copy_(dIndexK_f32)
+                dIndexK.copy_(dIndexK_f32.astype(dIndexK.dtype))
 
     _run.score_grad = partial(_score_grad_inplace, index_is_log=score_input_is_log)
     _run.gemm_only = _run_gemm_only

python/cudnn/deepseek_sparse_attention/indexer_top_k/api.py

@@ -17,12 +17,12 @@
 
 from .local_to_global_dsl import local_to_global as _local_to_global
 from .compactify import compactify as _compactify
+from .indexer_top_k_decode_varlen import cute_dsl_topk_wrapper
 
 _SUPPORTED_DTYPES = (torch.float32, torch.float16, torch.bfloat16)
 
 
 def _get_cute_dsl_topk_wrapper():
-    from .indexer_top_k_decode_varlen import cute_dsl_topk_wrapper
 
     return cute_dsl_topk_wrapper
 

python/cudnn/deepseek_sparse_attention/indexer_top_k/indexer_top_k_decode_varlen.py

@@ -21,7 +21,11 @@
 import cutlass
 import cutlass.cute as cute
 import cutlass.utils as utils
-import torch
+
+# import torch
+import paddle as torch  # why need this? lazy import?
+from cutlass.utils.distributed import atomicAdd
+
 from cudnn.deepseek_sparse_attention.utils.compiler import compile_options
 
 from .block_scan import block_prefix_sum_kernel
@@ -680,17 +684,17 @@ def cute_dsl_topk_wrapper(
     else:
         compiled_kernel = _compile_cache[key]
 
-    output_indices_torch = torch.empty(num_rows, top_k, dtype=torch.int32, device="cuda")
+    output_place = input_values.place
+    output_indices_torch = torch.empty(num_rows, top_k, dtype=torch.int32, device=output_place)
     if return_val:
-        output_values_torch = torch.empty(num_rows, top_k, dtype=torch_dtype, device="cuda")
+        output_values_torch = torch.empty(num_rows, top_k, dtype=torch_dtype, device=output_place)
     else:
         output_values_torch = None
 
     if dtype == cutlass.Float32:
         buffer_numbers = 2
     else:
         buffer_numbers = 1
-
     # Decode-varlen IMA workaround.
     elems_per_row = buffer_numbers * num_cols
     int32_max = (1 << 31) - 1
@@ -703,7 +707,7 @@ def cute_dsl_topk_wrapper(
             buffer_numbers,
             num_cols,
             dtype=torch.int32,
-            device="cuda",
+            device=output_place,
         )
         # TVM FFI uses env stream automatically
         compiled_kernel(
@@ -738,7 +742,7 @@ def cute_dsl_topk_wrapper(
             buffer_numbers,
             num_cols,
             dtype=torch.int32,
-            device="cuda",
+            device=output_place,
         )
         compiled_kernel(
             input_values[row_lo:row_hi],

python/cudnn/deepseek_sparse_attention/score_recompute/_interface_sm100.py

@@ -1081,8 +1081,15 @@ def _dense_indexer_score_recompute(
 
     # Dense indexer score is a raw-logit buffer, so masked/skipped positions
     # must remain outside the softmax/logsumexp domain.
+    #
+    # denom_out needs the same treatment for the same reason: the tile scheduler
+    # only visits q rows below cu_seqlens_q[-1], so in THD every row in
+    # [cu_seqlens_q[-1], total_q) is never written and would otherwise return
+    # whatever the allocator handed back.  0 is the safe value, not -inf:
+    # consumers compute exp(out - denom), and -inf - (-inf) is NaN.
     with _torch_stream_context(current_stream):
         out.fill_(float("-inf"))
+        denom_out.zero_()
     scale_arg = cutlass.Float32(sm_scale)
     max_q_arg = cutlass.Int32(seqlen_q)
     max_k_arg = cutlass.Int32(seqlen_k)
@@ -1402,11 +1409,18 @@ def _dense_attn_score_recompute(
     # region, so skipped masked/padding columns must be pre-filled. MXFP8
     # dense-attention uses mOut as an atomic-add accumulation buffer for valid
     # scores in some paths; those kernels patch invalid/skipped scores to -inf.
+    #
+    # denom_out is zeroed here for every precision, for the same reason as in the
+    # indexer path above: q rows at or beyond cu_seqlens_q[-1] are never visited
+    # by the scheduler, and an L1 norm of 0 is what "no candidate summed" should
+    # read as.  The zero_() at allocation time only covers a subset of the MXFP8
+    # shapes and does not cover a caller-supplied buffer.
     with _torch_stream_context(current_stream):
         if precision == "mxfp8":
             out.zero_()
         else:
             out.fill_(float("-inf"))
+        denom_out.zero_()
     with torch.cuda.nvtx.range("dense_attn_score_recompute"):
         if precision == "mxfp8":
             _dense_attn_score_recompute.compile_cache[compile_key](