From 0e5c871f4b6557b26eb3f7ad2e00b888cd22cbc6 Mon Sep 17 00:00:00 2001 From: Thorsten Kurth Date: Mon, 13 Jul 2026 16:24:24 +0200 Subject: [PATCH 01/14] distributed attention upsample implementation --- setup.py | 2 + tests/test_attention.py | 151 +++++ tests/test_distributed_attention.py | 27 +- .../optimized/attention_interface.cpp | 33 + .../optimized/attention_optimized.py | 69 ++ .../optimized/kernels_cuda/attention_cuda.cuh | 24 + .../kernels_cuda/attention_cuda_bwd.cu | 3 + .../kernels_cuda/attention_cuda_bwd_ring.cu | 1 + .../attention_cuda_bwd_ring_upsample.cu | 393 +++++++++++ .../attention_cuda_bwd_upsample.cu | 3 + .../kernels_cuda/attention_cuda_fwd.cu | 1 + .../kernels_cuda/attention_cuda_fwd_ring.cu | 1 + .../attention_cuda_fwd_ring_upsample.cu | 355 ++++++++++ .../attention_cuda_fwd_upsample.cu | 3 + .../kernels_cuda/attention_cuda_utils.cu | 6 + .../optimized/kernels_cuda/disco_cuda_bwd.cu | 3 + .../optimized/kernels_cuda/disco_cuda_fwd.cu | 3 + .../disco_cuda_fwd_dense_kpacked_sm100.cu | 4 + .../disco_cuda_fwd_dense_kpacked_sm90.cu | 4 + .../distributed/distributed_attention.py | 616 ++++++++++++++++-- 20 files changed, 1644 insertions(+), 58 deletions(-) create mode 100644 torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_bwd_ring_upsample.cu create mode 100644 torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_fwd_ring_upsample.cu diff --git a/setup.py b/setup.py index 7da8a625..4a1bcbf2 100644 --- a/setup.py +++ b/setup.py @@ -195,6 +195,8 @@ def get_ext_modules(): "torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_bwd_upsample.cu", "torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_fwd_ring.cu", "torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_bwd_ring.cu", + "torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_fwd_ring_upsample.cu", + "torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_bwd_ring_upsample.cu", ] ) ext_modules.append(CUDAExtension("torch_harmonics.attention._C", attention_sources, extra_compile_args=get_compile_args("attention"))) diff --git a/tests/test_attention.py b/tests/test_attention.py index f9ddb0a2..dba10289 100644 --- a/tests/test_attention.py +++ b/tests/test_attention.py @@ -964,6 +964,157 @@ def test_ring_kernels_pt2_compatibility(self, batch_size, channels, heads, in_sh opcheck(torch.ops.attention_kernels.backward_ring_step_pass2, bwd2_inputs) + @parameterized.expand( + [ + # Format: [batch_size, channels, heads, in_shape, out_shape, grid_in, grid_out] + # upsampling (scatter) cases, pscale_out=2 + [4, 4, 1, (6, 12), (12, 24), "equiangular", "equiangular"], + [4, 8, 4, (6, 12), (12, 24), "equiangular", "equiangular"], + ], + skip_on_empty=True, + ) + def test_ring_upsample_kernels_pt2_compatibility(self, batch_size, channels, heads, in_shape, out_shape, grid_in, grid_out, verbose=False): + """Tests whether the upsample (scatter) ring-step CUDA kernels (used by + DistributedNeighborhoodAttentionS2 in the upsample direction) are PyTorch 2 compatible. + + Only the local CUDA kernels are exercised — the ring exchange itself (NCCL P2P) is not + tested here. With az_size = polar_size = 1 a single ring step covers the full longitude + and the serial scatter psi coincides with the local psi built by + _build_local_psi_upsample (lat_lo_out = lon_lo_out = 0, no halo padding), so we can call + the kernels in single-rank mode (lon_lo_kx=0, lat_halo_start=0). + + opcheck only verifies the op contract (schema, fake tensors, AOT dispatch); the input + tensor values do not need to be numerically meaningful, so kw/vw/qw and the gradient / + state buffers are allocated directly with the right shapes.""" + + if self.device.type != "cuda": + raise unittest.SkipTest("ring kernels are only registered for CUDA") + if not cuda_kernels_is_available(): + raise unittest.SkipTest("skipping test because CUDA kernels are not available") + + set_seed(333) + + nlat_in, nlon_in = in_shape + nlat_out, nlon_out = out_shape + pscale_out = nlon_out // nlon_in + + # Build the module just to get a consistent (quad_weights, psi_col_idx, psi_roff_idx) + # for the chosen grid; we do not exercise its forward. For upsample shapes the module + # builds the scatter psi (rows keyed by hi, cols encoding ho * nlon_out + wo). + att = NeighborhoodAttentionS2( + in_channels=channels, + num_heads=heads, + in_shape=in_shape, + out_shape=out_shape, + grid_in=grid_in, + grid_out=grid_out, + bias=False, + optimized_kernel=True, + ).to(self.device) + self.assertTrue(att.upsample) + + # Channel counts after the (head-folded) projections in DistributedNeighborhoodAttentionS2. + Bnh = batch_size * heads + C_k = channels // heads + C_v = channels // heads + + # Synthetic projected k/v/q with the correct kernel-side shapes. + kw = torch.randn(Bnh, C_k, nlat_in, nlon_in, device=self.device, dtype=torch.float32) + vw = torch.randn(Bnh, C_v, nlat_in, nlon_in, device=self.device, dtype=torch.float32) + qw = torch.randn(Bnh, C_k, nlat_out, nlon_out, device=self.device, dtype=torch.float32) + + # ---- forward ring step ---- + # State buffers in channels-last layout, as expected by the CUDA kernels. + y_acc = torch.zeros(Bnh, nlat_out, nlon_out, C_v, device=self.device, dtype=torch.float32) + alpha_sum = torch.zeros(Bnh, nlat_out, nlon_out, device=self.device, dtype=torch.float32) + qdotk_max = torch.full((Bnh, nlat_out, nlon_out), float("-inf"), device=self.device, dtype=torch.float32) + + fwd_inputs = ( + kw, + vw, + qw, + y_acc, + alpha_sum, + qdotk_max, + att.quad_weights, + att.psi_col_idx, + att.psi_roff_idx, + nlon_in, + nlon_out, + pscale_out, + 0, + 0, + nlat_out, + nlon_out, + ) + + opcheck(torch.ops.attention_kernels.forward_ring_step_upsample, fwd_inputs) + + # ---- backward pass 1: scatter integral / alpha_k / alpha_kvw stats ---- + # Uses the forward-final qdotk_max; synthetic but well-formed values (no -inf) so the + # kernel doesn't produce NaNs (opcheck doesn't check numerics, but NaNs can interact + # badly with AOT dispatch comparisons). + dy = torch.randn(Bnh, C_v, nlat_out, nlon_out, device=self.device, dtype=torch.float32) + + fwd_qdotk_max = torch.zeros(Bnh, nlat_out, nlon_out, device=self.device, dtype=torch.float32) + integral_buf = torch.zeros(Bnh, nlat_out, nlon_out, device=self.device, dtype=torch.float32) + alpha_k_buf = torch.zeros(Bnh, nlat_out, nlon_out, C_k, device=self.device, dtype=torch.float32) + alpha_kvw_buf = torch.zeros(Bnh, nlat_out, nlon_out, C_k, device=self.device, dtype=torch.float32) + + bwd1_inputs = ( + kw, + vw, + qw, + dy, + fwd_qdotk_max, + integral_buf, + alpha_k_buf, + alpha_kvw_buf, + att.quad_weights, + att.psi_col_idx, + att.psi_roff_idx, + nlon_in, + nlon_out, + pscale_out, + 0, + 0, + nlat_out, + nlon_out, + ) + + opcheck(torch.ops.attention_kernels.backward_ring_step_upsample_pass1, bwd1_inputs) + + # ---- backward pass 2: accumulate chunk-local dkx/dvx using finalized stats ---- + fwd_alpha_sum = torch.ones(Bnh, nlat_out, nlon_out, device=self.device, dtype=torch.float32) + integral_norm = torch.zeros(Bnh, nlat_out, nlon_out, device=self.device, dtype=torch.float32) + + dkw = torch.zeros(Bnh, nlat_in, nlon_in, C_k, device=self.device, dtype=torch.float32) + dvw = torch.zeros(Bnh, nlat_in, nlon_in, C_v, device=self.device, dtype=torch.float32) + + bwd2_inputs = ( + kw, + vw, + qw, + dy, + fwd_alpha_sum, + fwd_qdotk_max, + integral_norm, + dkw, + dvw, + att.quad_weights, + att.psi_col_idx, + att.psi_roff_idx, + nlon_in, + nlon_out, + pscale_out, + 0, + 0, + nlat_out, + nlon_out, + ) + + opcheck(torch.ops.attention_kernels.backward_ring_step_upsample_pass2, bwd2_inputs) + @parameterized.expand( [ # self attention diff --git a/tests/test_distributed_attention.py b/tests/test_distributed_attention.py index 980a9981..ee386891 100644 --- a/tests/test_distributed_attention.py +++ b/tests/test_distributed_attention.py @@ -234,7 +234,28 @@ def _gather_helper_bwd(self, tensor, attn_dist, use_out_shapes=False): # # BDIM_X=1024 (per-head 8193..16384). This also stresses the dynamic-shmem opt-in # # (ensure_dyn_shmem) on the TMA path: ~5*nchans*4 B exceeds the default 48 KiB limit. # [8, 16, 8, 16, 2, 8704, 1, None, None, "equiangular", "equiangular", False, torch.float32, 1e-5, 1e-4], - # upsampling is not supported by the kernel yet (serial layer asserts nlon_in % nlon_out == 0) + # upsampling tests (scatter ring kernels), pscale_out=2 (lat+lon) + [32, 64, 64, 128, 2, 16, 1, None, None, "equiangular", "equiangular", False, torch.float32, 1e-5, 1e-4], + [33, 64, 65, 128, 2, 16, 1, None, None, "equiangular", "equiangular", False, torch.float32, 1e-5, 1e-4], + # pscale_out=3 upsampling (exercises pscale_out*wi kernel arithmetic) + [32, 32, 64, 96, 2, 16, 1, None, None, "equiangular", "equiangular", False, torch.float32, 1e-5, 1e-4], + # pscale_out=4 upsampling + [32, 32, 64, 128, 2, 16, 1, None, None, "equiangular", "equiangular", False, torch.float32, 1e-5, 1e-4], + # lon-only upsampling (same nlat; isolates the azimuth ring scatter, no lat halo) + [64, 64, 64, 128, 2, 16, 1, None, None, "equiangular", "equiangular", False, torch.float32, 1e-5, 1e-4], + # odd nlat_in -> even nlat_out, pscale_out=2 + [33, 64, 64, 128, 2, 16, 1, None, None, "equiangular", "equiangular", False, torch.float32, 1e-5, 1e-4], + # legendre-gauss / mixed grid upsampling + [32, 64, 64, 128, 2, 16, 1, None, None, "legendre-gauss", "legendre-gauss", False, torch.float32, 1e-5, 1e-4], + [32, 64, 64, 128, 2, 16, 1, None, None, "legendre-gauss", "equiangular", False, torch.float32, 1e-5, 1e-4], + # heads=4 with asymmetric channels (k=32, out=16; in=16), upsampling + [32, 64, 64, 128, 2, 16, 4, 32, 16, "equiangular", "equiangular", False, torch.float32, 1e-5, 1e-4], + # upsampling with QK norm + [32, 64, 64, 128, 2, 16, 1, None, None, "equiangular", "equiangular", True, torch.float32, 1e-5, 1e-4], + [32, 64, 64, 128, 2, 16, 2, None, None, "equiangular", "equiangular", True, torch.float32, 1e-5, 1e-4], + # upsampling AMP coverage (see AMP tolerance note below) + [32, 64, 64, 128, 2, 16, 1, None, None, "equiangular", "equiangular", False, torch.float16, 5e-2, 1e-2], + [32, 64, 64, 128, 2, 16, 1, None, None, "equiangular", "equiangular", False, torch.bfloat16, 3e-1, 5e-2], # AMP coverage — one row per dtype on a downsample config. Tolerances # are looser than the fp32 rows because bf16/fp16 cuBLAS rounding at # the einsum output dominates the abs error, and the distributed path @@ -380,6 +401,10 @@ def test_distributed_neighborhood_attention( [64, 128, 32, 64, 2, 16, 1, None, None, "equiangular", "equiangular", "k"], [64, 128, 32, 64, 2, 16, 1, None, None, "equiangular", "equiangular", "v"], [64, 128, 32, 64, 2, 16, 1, None, None, "equiangular", "equiangular", "q"], + # upsample config (pscale_out=2), each branch frozen in turn + [32, 64, 64, 128, 2, 16, 1, None, None, "equiangular", "equiangular", "k"], + [32, 64, 64, 128, 2, 16, 1, None, None, "equiangular", "equiangular", "v"], + [32, 64, 64, 128, 2, 16, 1, None, None, "equiangular", "equiangular", "q"], ], skip_on_empty=True, ) diff --git a/torch_harmonics/attention/optimized/attention_interface.cpp b/torch_harmonics/attention/optimized/attention_interface.cpp index 9192ae4f..e8d54ab1 100644 --- a/torch_harmonics/attention/optimized/attention_interface.cpp +++ b/torch_harmonics/attention/optimized/attention_interface.cpp @@ -123,6 +123,39 @@ namespace attention_kernels "col_idx, Tensor row_off, Tensor row_idx, int nlon_in, int pscale, int lon_lo_kx, int lat_halo_start, " "int nlat_out, int nlon_out, int n_long_rows, int max_row_len, int mid_row_len) -> ()", {at::Tag::pt2_compliant_tag}); + + // ---- Ring-step variants for the UPSAMPLE (input-keyed scatter) direction ---- + // Used by DistributedNeighborhoodAttentionS2 when nlon_out % nlon_in == 0. + // K/V live on the coarse input grid and rotate along the azimuth ring; Q and + // the softmax state buffers live on the fine output grid and stay local. + // psi convention (see _build_local_psi_upsample in distributed_attention.py): + // row_off : indexed by hi_local in [0, nlat_halo] (halo-padded local input + // rows; hi_global = lat_halo_start + hi_local) + // col_idx : ho_local * nlon_out_global + wo_shifted, with + // wo_shifted = (wo_canonical - lon_lo_out) mod nlon_out_global. + // The kernel maps wo_shifted -> (wo_shifted + pscale_out * (lon_lo_kx + wi_local)) + // mod nlon_out_global and treats the cell as local iff the result < nlon_out + // (the LOCAL output width). pscale_out is the GLOBAL nlon_out / nlon_in. + // A single forward step runs the 3-phase max/rescale/accumulate scheme so the + // online softmax stays consistent across ring steps despite the scatter form. + m.def("forward_ring_step_upsample(Tensor kx, Tensor vx, Tensor qy, Tensor(a!) y_acc, Tensor(b!) " + "alpha_sum_buf, Tensor(c!) qdotk_max_buf, Tensor quad_weights, Tensor col_idx, Tensor row_off, int " + "nlon_in, int nlon_out_global, int pscale_out, int lon_lo_kx, int lat_halo_start, int nlat_out, int " + "nlon_out) -> ()", + {at::Tag::pt2_compliant_tag}); + // Backward reuses the forward-final alpha_sum / qdotk_max (no max recompute): + // pass1 scatters the per-output stats (integral, alpha_k, alpha_kvw) needed for + // dqy; pass2 accumulates chunk-local dkx/dvx (allreduced in Python). + m.def("backward_ring_step_upsample_pass1(Tensor kx, Tensor vx, Tensor qy, Tensor dy, Tensor qdotk_max_buf, " + "Tensor(a!) integral_buf, Tensor(b!) alpha_k_buf, Tensor(c!) alpha_kvw_buf, Tensor quad_weights, Tensor " + "col_idx, Tensor row_off, int nlon_in, int nlon_out_global, int pscale_out, int lon_lo_kx, int " + "lat_halo_start, int nlat_out, int nlon_out) -> ()", + {at::Tag::pt2_compliant_tag}); + m.def("backward_ring_step_upsample_pass2(Tensor kx, Tensor vx, Tensor qy, Tensor dy, Tensor alpha_sum_buf, " + "Tensor qdotk_max_buf, Tensor integral_norm_buf, Tensor(a!) dkx, Tensor(b!) dvx, Tensor quad_weights, " + "Tensor col_idx, Tensor row_off, int nlon_in, int nlon_out_global, int pscale_out, int lon_lo_kx, int " + "lat_halo_start, int nlat_out, int nlon_out) -> ()", + {at::Tag::pt2_compliant_tag}); } } // namespace attention_kernels diff --git a/torch_harmonics/attention/optimized/attention_optimized.py b/torch_harmonics/attention/optimized/attention_optimized.py index ec8b0936..b23d6f0c 100644 --- a/torch_harmonics/attention/optimized/attention_optimized.py +++ b/torch_harmonics/attention/optimized/attention_optimized.py @@ -145,6 +145,75 @@ def _( ) -> None: pass + # fake implementations for the upsample (scatter) ring step ops + @torch.library.register_fake("attention_kernels::forward_ring_step_upsample") + def _( + kx: torch.Tensor, + vx: torch.Tensor, + qy: torch.Tensor, + y_acc: torch.Tensor, + alpha_sum_buf: torch.Tensor, + qdotk_max_buf: torch.Tensor, + quad_weights: torch.Tensor, + col_idx: torch.Tensor, + row_off: torch.Tensor, + nlon_in: int, + nlon_out_global: int, + pscale_out: int, + lon_lo_kx: int, + lat_halo_start: int, + nlat_out: int, + nlon_out: int, + ) -> None: + pass + + @torch.library.register_fake("attention_kernels::backward_ring_step_upsample_pass1") + def _( + kx: torch.Tensor, + vx: torch.Tensor, + qy: torch.Tensor, + dy: torch.Tensor, + qdotk_max_buf: torch.Tensor, + integral_buf: torch.Tensor, + alpha_k_buf: torch.Tensor, + alpha_kvw_buf: torch.Tensor, + quad_weights: torch.Tensor, + col_idx: torch.Tensor, + row_off: torch.Tensor, + nlon_in: int, + nlon_out_global: int, + pscale_out: int, + lon_lo_kx: int, + lat_halo_start: int, + nlat_out: int, + nlon_out: int, + ) -> None: + pass + + @torch.library.register_fake("attention_kernels::backward_ring_step_upsample_pass2") + def _( + kx: torch.Tensor, + vx: torch.Tensor, + qy: torch.Tensor, + dy: torch.Tensor, + alpha_sum_buf: torch.Tensor, + qdotk_max_buf: torch.Tensor, + integral_norm_buf: torch.Tensor, + dkx: torch.Tensor, + dvx: torch.Tensor, + quad_weights: torch.Tensor, + col_idx: torch.Tensor, + row_off: torch.Tensor, + nlon_in: int, + nlon_out_global: int, + pscale_out: int, + lon_lo_kx: int, + lat_halo_start: int, + nlat_out: int, + nlon_out: int, + ) -> None: + pass + # forward @torch.library.custom_op("attention_kernels::_neighborhood_s2_attention_optimized", mutates_args=()) def _neighborhood_s2_attention_optimized( diff --git a/torch_harmonics/attention/optimized/kernels_cuda/attention_cuda.cuh b/torch_harmonics/attention/optimized/kernels_cuda/attention_cuda.cuh index 30a77f21..7452f47b 100644 --- a/torch_harmonics/attention/optimized/kernels_cuda/attention_cuda.cuh +++ b/torch_harmonics/attention/optimized/kernels_cuda/attention_cuda.cuh @@ -75,4 +75,28 @@ namespace attention_kernels int64_t lat_halo_start, int64_t nlat_out, int64_t nlon_out, int64_t n_long_rows, int64_t max_row_len, int64_t mid_row_len); + // ring-step variants for the upsample (input-keyed scatter) direction + void s2_attention_fwd_ring_step_upsample_cuda(at::Tensor kx, at::Tensor vx, at::Tensor qy, at::Tensor y_acc, + at::Tensor alpha_sum_buf, at::Tensor qdotk_max_buf, + at::Tensor quad_weights, at::Tensor psi_col_idx, + at::Tensor psi_row_off, int64_t nlon_in, int64_t nlon_out_global, + int64_t pscale_out, int64_t lon_lo_kx, int64_t lat_halo_start, + int64_t nlat_out, int64_t nlon_out); + + void s2_attention_bwd_ring_step_upsample_pass1_cuda(at::Tensor kx, at::Tensor vx, at::Tensor qy, at::Tensor dy, + at::Tensor qdotk_max_buf, at::Tensor integral_buf, + at::Tensor alpha_k_buf, at::Tensor alpha_kvw_buf, + at::Tensor quad_weights, at::Tensor psi_col_idx, + at::Tensor psi_row_off, int64_t nlon_in, + int64_t nlon_out_global, int64_t pscale_out, int64_t lon_lo_kx, + int64_t lat_halo_start, int64_t nlat_out, int64_t nlon_out); + + void s2_attention_bwd_ring_step_upsample_pass2_cuda(at::Tensor kx, at::Tensor vx, at::Tensor qy, at::Tensor dy, + at::Tensor alpha_sum_buf, at::Tensor qdotk_max_buf, + at::Tensor integral_norm_buf, at::Tensor dkx, at::Tensor dvx, + at::Tensor quad_weights, at::Tensor psi_col_idx, + at::Tensor psi_row_off, int64_t nlon_in, + int64_t nlon_out_global, int64_t pscale_out, int64_t lon_lo_kx, + int64_t lat_halo_start, int64_t nlat_out, int64_t nlon_out); + } // namespace attention_kernels diff --git a/torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_bwd.cu b/torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_bwd.cu index 0314b703..ad8e1b9a 100644 --- a/torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_bwd.cu +++ b/torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_bwd.cu @@ -38,6 +38,7 @@ #include #include #include +#include #include #include @@ -952,6 +953,8 @@ namespace attention_kernels } }); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + // convert precision back to starting dtype (no-op for fp32; narrows for fp16/bf16) dkx = dkx.to(kx_type); dvx = dvx.to(vx_type); diff --git a/torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_bwd_ring.cu b/torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_bwd_ring.cu index d715603a..7dd76d0c 100644 --- a/torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_bwd_ring.cu +++ b/torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_bwd_ring.cu @@ -35,6 +35,7 @@ #include #include #include +#include #include diff --git a/torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_bwd_ring_upsample.cu b/torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_bwd_ring_upsample.cu new file mode 100644 index 00000000..3bd2801a --- /dev/null +++ b/torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_bwd_ring_upsample.cu @@ -0,0 +1,393 @@ +// coding=utf-8 +// +// SPDX-FileCopyrightText: Copyright (c) 2026 The torch-harmonics Authors. All rights reserved. +// SPDX-License-Identifier: BSD-3-Clause +// +// Redistribution and use in source and binary forms, with or without +// modification, are permitted provided that the following conditions are met: +// +// 1. Redistributions of source code must retain the above copyright notice, this +// list of conditions and the following disclaimer. +// +// 2. Redistributions in binary form must reproduce the above copyright notice, +// this list of conditions and the following disclaimer in the documentation +// and/or other materials provided with the distribution. +// +// 3. Neither the name of the copyright holder nor the names of its +// contributors may be used to endorse or promote products derived from +// this software without specific prior written permission. +// +// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +// DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +// FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +// DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +// SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +// CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +// OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +// ===================================================================================== +// Upsample (scatter-style) attention backward, RING-STEP variants — CUDA +// ===================================================================================== +// +// Ring counterpart of attention_cuda_bwd_upsample.cu, used by +// DistributedNeighborhoodAttentionS2 in the upsample direction. See +// attention_cuda_fwd_ring_upsample.cu for the psi / index conventions (local +// halo-keyed rows, local-output-keyed cols, wo pre-shifted by -lon_lo_out). +// +// Unlike the serial upsample backward, no max-recomputation pass is needed: +// the forward saves the FINAL (globally reduced) alpha_sum and qdotk_max +// buffers, which are passed back in here. That leaves two ring sweeps: +// +// pass 1 (stats): per coarse chunk cell, scatter with the fixed forward max +// alpha -> (implicitly known: forward alpha_sum) +// alpha*g -> atomicAdd integral_buf[b,ho,wo] (g = dy.v) +// alpha*k -> atomicAdd alpha_k_buf[b,ho,wo,:] +// alpha*g*k -> atomicAdd alpha_kvw_buf[b,ho,wo,:] +// dqy is finalized in Python from these buffers: +// dqy = (alpha_sum*alpha_kvw - integral*alpha_k) / alpha_sum^2 +// +// pass 2 (dkdv): per coarse chunk cell, accumulate locally over its row +// dvx += (alpha/S)*dy and dkx += q*(g - integral_norm)*(alpha/S) +// with integral_norm = integral/S precomputed in Python. +// dkx/dvx are chunk-shaped, direct write (no atomics); Python +// accumulates chunks into a full-width buffer and allreduces +// over the azimuth group. +// +// Generic-only (scalar loads): correctness path. +// ===================================================================================== + +#include "attention_cuda.cuh" +#include +#include +#include + +#include +#include + +#include "cudamacro.h" +#include "attention_cuda_utils.cuh" + +#define THREADS (64) + +namespace attention_kernels +{ + + // map (wo_shifted, wi_global) -> local output longitude; returns -1 if the + // target cell is not owned by this rank. + __device__ __forceinline__ int bwd_ring_up_local_wo(int wo_shifted, int wi_global, int pscale_out, + int nlon_out_global, int nlon_out_local) + { + int w = wo_shifted + pscale_out * wi_global; // < 2*nlon_out_global + if (w >= nlon_out_global) { w -= nlon_out_global; } + return (w < nlon_out_local) ? w : -1; + } + + // pass 1: per coarse chunk cell, scatter the softmax stats (integral, alpha_k, + // alpha_kvw) to the local fine output cells, using the saved forward max. + template + __global__ __launch_bounds__(THREADS_PER_BLOCK) void s2_attn_bwd_ring_upsample_stats_k( + int nchan_in, int nchan_out, int nlat_halo, int nlon_kx, int nlon_out_global, int pscale_out, int lon_lo_kx, + int lat_halo_start, int nlat_out, int nlon_out, const STORAGE_T *__restrict__ kx, + const STORAGE_T *__restrict__ vx, const STORAGE_T *__restrict__ qy, const STORAGE_T *__restrict__ dy, + const int64_t *__restrict__ row_off, const int64_t *__restrict__ col_idx, + const float *__restrict__ quad_weights, const float *__restrict__ qdotk_max_buf, + float *__restrict__ integral_buf, float *__restrict__ alpha_k_buf, float *__restrict__ alpha_kvw_buf) + { + extern __shared__ float shext[]; + float *sh_k = shext + threadIdx.y * (nchan_in + nchan_out); + float *sh_v = sh_k + nchan_in; + + const int batch = blockIdx.y; + const int wid = blockIdx.x * blockDim.y + threadIdx.y; + if (wid >= nlat_halo * nlon_kx) { return; } + + const int tidx = threadIdx.x; + const int hi = wid / nlon_kx; // LOCAL halo row + const int wi_local = wid - hi * nlon_kx; + const int wi_global = lon_lo_kx + wi_local; + + const int64_t rbeg = row_off[hi]; + const int rlen = static_cast(row_off[hi + 1] - rbeg); + if (rlen == 0) { return; } // empty row (e.g. pole padding) + const int64_t *col_hi = col_idx + rbeg; + + kx += int64_t(batch) * nlat_halo * nlon_kx * nchan_in + (int64_t(hi) * nlon_kx + wi_local) * nchan_in; + vx += int64_t(batch) * nlat_halo * nlon_kx * nchan_out + (int64_t(hi) * nlon_kx + wi_local) * nchan_out; + qy += int64_t(batch) * nlat_out * nlon_out * nchan_in; + dy += int64_t(batch) * nlat_out * nlon_out * nchan_out; + qdotk_max_buf += int64_t(batch) * nlat_out * nlon_out; + integral_buf += int64_t(batch) * nlat_out * nlon_out; + alpha_k_buf += int64_t(batch) * nlat_out * nlon_out * nchan_in; + alpha_kvw_buf += int64_t(batch) * nlat_out * nlon_out * nchan_in; + + for (int chan = tidx; chan < nchan_in; chan += WARP_SIZE) { sh_k[chan] = vload(kx, chan); } + for (int chan = tidx; chan < nchan_out; chan += WARP_SIZE) { sh_v[chan] = vload(vx, chan); } + + // quad_weights is indexed by the GLOBAL input latitude + const float qw = quad_weights[lat_halo_start + hi]; + + for (int off = 0; off < rlen; off++) { + const int64_t col = col_hi[off]; + const int ho = static_cast(col / nlon_out_global); // LOCAL output row + const int wo = bwd_ring_up_local_wo(static_cast(col - int64_t(ho) * nlon_out_global), wi_global, + pscale_out, nlon_out_global, nlon_out); + if (wo < 0) { continue; } // target cell not owned by this rank + + const int64_t cell = int64_t(ho) * nlon_out + wo; + const STORAGE_T *_qy = qy + cell * nchan_in; + const STORAGE_T *_dy = dy + cell * nchan_out; + + float qd = 0.f, gd = 0.f; + for (int chan = tidx; chan < nchan_in; chan += WARP_SIZE) { qd += sh_k[chan] * vload(_qy, chan); } + for (int chan = tidx; chan < nchan_out; chan += WARP_SIZE) { gd += sh_v[chan] * vload(_dy, chan); } + qd = __warp_sum(qd); + gd = __warp_sum(gd); + + const float alpha = expf(qd - qdotk_max_buf[cell]) * qw; + const float ag = alpha * gd; + if (tidx == 0) { atomicAdd(&integral_buf[cell], ag); } + float *_alpha_k = alpha_k_buf + cell * nchan_in; + float *_alpha_kvw = alpha_kvw_buf + cell * nchan_in; + for (int chan = tidx; chan < nchan_in; chan += WARP_SIZE) { + atomicAdd(&_alpha_k[chan], alpha * sh_k[chan]); + atomicAdd(&_alpha_kvw[chan], ag * sh_k[chan]); + } + } + } + + // pass 2: per coarse chunk cell, accumulate dkx/dvx locally over its row and + // write into the chunk-shaped gradient buffers (no atomics). + template + __global__ __launch_bounds__(THREADS_PER_BLOCK) void s2_attn_bwd_ring_upsample_dkv_k( + int nchan_in, int nchan_out, int nlat_halo, int nlon_kx, int nlon_out_global, int pscale_out, int lon_lo_kx, + int lat_halo_start, int nlat_out, int nlon_out, const STORAGE_T *__restrict__ kx, + const STORAGE_T *__restrict__ vx, const STORAGE_T *__restrict__ qy, const STORAGE_T *__restrict__ dy, + const int64_t *__restrict__ row_off, const int64_t *__restrict__ col_idx, const float *__restrict__ quad_weights, + const float *__restrict__ alpha_sum_buf, const float *__restrict__ qdotk_max_buf, + const float *__restrict__ integral_norm_buf, float *__restrict__ dkx, float *__restrict__ dvx) + { + extern __shared__ float shext[]; + float *sh_k = shext + threadIdx.y * (2 * nchan_in + 2 * nchan_out); + float *sh_v = sh_k + nchan_in; + float *sh_dk = sh_v + nchan_out; + float *sh_dv = sh_dk + nchan_in; + + const int batch = blockIdx.y; + const int wid = blockIdx.x * blockDim.y + threadIdx.y; + if (wid >= nlat_halo * nlon_kx) { return; } + + const int tidx = threadIdx.x; + const int hi = wid / nlon_kx; // LOCAL halo row + const int wi_local = wid - hi * nlon_kx; + const int wi_global = lon_lo_kx + wi_local; + + const int64_t rbeg = row_off[hi]; + const int rlen = static_cast(row_off[hi + 1] - rbeg); + if (rlen == 0) { return; } // empty row: dkx/dvx stay zero (buffers pre-zeroed) + const int64_t *col_hi = col_idx + rbeg; + + kx += int64_t(batch) * nlat_halo * nlon_kx * nchan_in + (int64_t(hi) * nlon_kx + wi_local) * nchan_in; + vx += int64_t(batch) * nlat_halo * nlon_kx * nchan_out + (int64_t(hi) * nlon_kx + wi_local) * nchan_out; + qy += int64_t(batch) * nlat_out * nlon_out * nchan_in; + dy += int64_t(batch) * nlat_out * nlon_out * nchan_out; + alpha_sum_buf += int64_t(batch) * nlat_out * nlon_out; + qdotk_max_buf += int64_t(batch) * nlat_out * nlon_out; + integral_norm_buf += int64_t(batch) * nlat_out * nlon_out; + dkx += int64_t(batch) * nlat_halo * nlon_kx * nchan_in + (int64_t(hi) * nlon_kx + wi_local) * nchan_in; + dvx += int64_t(batch) * nlat_halo * nlon_kx * nchan_out + (int64_t(hi) * nlon_kx + wi_local) * nchan_out; + + for (int chan = tidx; chan < nchan_in; chan += WARP_SIZE) { + sh_k[chan] = vload(kx, chan); + sh_dk[chan] = 0.f; + } + for (int chan = tidx; chan < nchan_out; chan += WARP_SIZE) { + sh_v[chan] = vload(vx, chan); + sh_dv[chan] = 0.f; + } + + // quad_weights is indexed by the GLOBAL input latitude + const float qw = quad_weights[lat_halo_start + hi]; + + for (int off = 0; off < rlen; off++) { + const int64_t col = col_hi[off]; + const int ho = static_cast(col / nlon_out_global); // LOCAL output row + const int wo = bwd_ring_up_local_wo(static_cast(col - int64_t(ho) * nlon_out_global), wi_global, + pscale_out, nlon_out_global, nlon_out); + if (wo < 0) { continue; } // target cell not owned by this rank + + const int64_t cell = int64_t(ho) * nlon_out + wo; + const STORAGE_T *_qy = qy + cell * nchan_in; + const STORAGE_T *_dy = dy + cell * nchan_out; + + float qd = 0.f, gd = 0.f; + for (int chan = tidx; chan < nchan_in; chan += WARP_SIZE) { qd += sh_k[chan] * vload(_qy, chan); } + for (int chan = tidx; chan < nchan_out; chan += WARP_SIZE) { gd += sh_v[chan] * vload(_dy, chan); } + qd = __warp_sum(qd); + gd = __warp_sum(gd); + + const float alpha_norm = expf(qd - qdotk_max_buf[cell]) * qw / alpha_sum_buf[cell]; + const float scale_dk = (gd - integral_norm_buf[cell]) * alpha_norm; + + for (int chan = tidx; chan < nchan_out; chan += WARP_SIZE) { sh_dv[chan] += alpha_norm * vload(_dy, chan); } + for (int chan = tidx; chan < nchan_in; chan += WARP_SIZE) { sh_dk[chan] += scale_dk * vload(_qy, chan); } + } + + for (int chan = tidx; chan < nchan_in; chan += WARP_SIZE) { dkx[chan] = sh_dk[chan]; } + for (int chan = tidx; chan < nchan_out; chan += WARP_SIZE) { dvx[chan] = sh_dv[chan]; } + } + + void s2_attention_bwd_ring_step_upsample_pass1_cuda(at::Tensor kx, at::Tensor vx, at::Tensor qy, at::Tensor dy, + at::Tensor qdotk_max_buf, at::Tensor integral_buf, + at::Tensor alpha_k_buf, at::Tensor alpha_kvw_buf, + at::Tensor quad_weights, at::Tensor psi_col_idx, + at::Tensor psi_row_off, int64_t nlon_in, + int64_t nlon_out_global, int64_t pscale_out, int64_t lon_lo_kx, + int64_t lat_halo_start, int64_t nlat_out, int64_t nlon_out) + { + CHECK_CUDA_INPUT_TENSOR(kx); + CHECK_CUDA_INPUT_TENSOR(vx); + CHECK_CUDA_INPUT_TENSOR(qy); + CHECK_CUDA_INPUT_TENSOR(dy); + CHECK_CUDA_TENSOR(qdotk_max_buf); + CHECK_CUDA_TENSOR(integral_buf); + CHECK_CUDA_TENSOR(alpha_k_buf); + CHECK_CUDA_TENSOR(alpha_kvw_buf); + CHECK_CUDA_TENSOR(quad_weights); + CHECK_CUDA_TENSOR(psi_col_idx); + CHECK_CUDA_TENSOR(psi_row_off); + + const int batch_size = kx.size(0); + const int nlat_halo = kx.size(2); + const int nlon_kx = kx.size(3); + const size_t nchans_in = qy.size(1); + const size_t nchans_out = vx.size(1); + + auto stream = at::cuda::getCurrentCUDAStream().stream(); + + int64_t *_row_off = reinterpret_cast(psi_row_off.data_ptr()); + int64_t *_col_idx = reinterpret_cast(psi_col_idx.data_ptr()); + float *_quad_weights = reinterpret_cast(quad_weights.data_ptr()); + float *_qdotk_max = reinterpret_cast(qdotk_max_buf.data_ptr()); + float *_integral = reinterpret_cast(integral_buf.data_ptr()); + float *_alpha_k = reinterpret_cast(alpha_k_buf.data_ptr()); + float *_alpha_kvw = reinterpret_cast(alpha_kvw_buf.data_ptr()); + + AT_DISPATCH_FLOATING_TYPES_AND2( + at::kHalf, at::kBFloat16, qy.scalar_type(), "s2_attention_bwd_ring_step_upsample_pass1_cuda", [&] { + using storage_t = scalar_t; + + torch::Tensor kxP = kx; + torch::Tensor vxP = vx; + torch::Tensor qyP = qy; + torch::Tensor dyP = dy; + + if (kxP.strides()[1] != 1) { kxP = permute_4D_to0231(kxP); } + if (vxP.strides()[1] != 1) { vxP = permute_4D_to0231(vxP); } + if (qyP.strides()[1] != 1) { qyP = permute_4D_to0231(qyP); } + if (dyP.strides()[1] != 1) { dyP = permute_4D_to0231(dyP); } + + storage_t *_kxp = reinterpret_cast(kxP.data_ptr()); + storage_t *_vxp = reinterpret_cast(vxP.data_ptr()); + storage_t *_qyp = reinterpret_cast(qyP.data_ptr()); + storage_t *_dyp = reinterpret_cast(dyP.data_ptr()); + + dim3 block(WARP_SIZE, THREADS / WARP_SIZE); + dim3 grid_in(DIV_UP(nlat_halo * nlon_kx, block.y), batch_size); + + const size_t sh_stats = sizeof(float) * (nchans_in + nchans_out) * block.y; + + s2_attn_bwd_ring_upsample_stats_k<<>>( + static_cast(nchans_in), static_cast(nchans_out), nlat_halo, nlon_kx, + static_cast(nlon_out_global), static_cast(pscale_out), static_cast(lon_lo_kx), + static_cast(lat_halo_start), static_cast(nlat_out), static_cast(nlon_out), _kxp, + _vxp, _qyp, _dyp, _row_off, _col_idx, _quad_weights, _qdotk_max, _integral, _alpha_k, _alpha_kvw); + CHECK_ERROR("s2_attn_bwd_ring_upsample_stats_k"); + }); + + C10_CUDA_KERNEL_LAUNCH_CHECK(); + } + + void s2_attention_bwd_ring_step_upsample_pass2_cuda(at::Tensor kx, at::Tensor vx, at::Tensor qy, at::Tensor dy, + at::Tensor alpha_sum_buf, at::Tensor qdotk_max_buf, + at::Tensor integral_norm_buf, at::Tensor dkx, at::Tensor dvx, + at::Tensor quad_weights, at::Tensor psi_col_idx, + at::Tensor psi_row_off, int64_t nlon_in, + int64_t nlon_out_global, int64_t pscale_out, int64_t lon_lo_kx, + int64_t lat_halo_start, int64_t nlat_out, int64_t nlon_out) + { + CHECK_CUDA_INPUT_TENSOR(kx); + CHECK_CUDA_INPUT_TENSOR(vx); + CHECK_CUDA_INPUT_TENSOR(qy); + CHECK_CUDA_INPUT_TENSOR(dy); + CHECK_CUDA_TENSOR(alpha_sum_buf); + CHECK_CUDA_TENSOR(qdotk_max_buf); + CHECK_CUDA_TENSOR(integral_norm_buf); + CHECK_CUDA_TENSOR(dkx); + CHECK_CUDA_TENSOR(dvx); + CHECK_CUDA_TENSOR(quad_weights); + CHECK_CUDA_TENSOR(psi_col_idx); + CHECK_CUDA_TENSOR(psi_row_off); + + const int batch_size = kx.size(0); + const int nlat_halo = kx.size(2); + const int nlon_kx = kx.size(3); + const size_t nchans_in = qy.size(1); + const size_t nchans_out = vx.size(1); + + auto stream = at::cuda::getCurrentCUDAStream().stream(); + + int64_t *_row_off = reinterpret_cast(psi_row_off.data_ptr()); + int64_t *_col_idx = reinterpret_cast(psi_col_idx.data_ptr()); + float *_quad_weights = reinterpret_cast(quad_weights.data_ptr()); + float *_alpha_sum = reinterpret_cast(alpha_sum_buf.data_ptr()); + float *_qdotk_max = reinterpret_cast(qdotk_max_buf.data_ptr()); + float *_integral_norm = reinterpret_cast(integral_norm_buf.data_ptr()); + // gradient chunk buffers are always fp32, channels-last [B, nlat_halo, nlon_kx, C] + float *_dkx = reinterpret_cast(dkx.data_ptr()); + float *_dvx = reinterpret_cast(dvx.data_ptr()); + + AT_DISPATCH_FLOATING_TYPES_AND2( + at::kHalf, at::kBFloat16, qy.scalar_type(), "s2_attention_bwd_ring_step_upsample_pass2_cuda", [&] { + using storage_t = scalar_t; + + torch::Tensor kxP = kx; + torch::Tensor vxP = vx; + torch::Tensor qyP = qy; + torch::Tensor dyP = dy; + + if (kxP.strides()[1] != 1) { kxP = permute_4D_to0231(kxP); } + if (vxP.strides()[1] != 1) { vxP = permute_4D_to0231(vxP); } + if (qyP.strides()[1] != 1) { qyP = permute_4D_to0231(qyP); } + if (dyP.strides()[1] != 1) { dyP = permute_4D_to0231(dyP); } + + storage_t *_kxp = reinterpret_cast(kxP.data_ptr()); + storage_t *_vxp = reinterpret_cast(vxP.data_ptr()); + storage_t *_qyp = reinterpret_cast(qyP.data_ptr()); + storage_t *_dyp = reinterpret_cast(dyP.data_ptr()); + + dim3 block(WARP_SIZE, THREADS / WARP_SIZE); + dim3 grid_in(DIV_UP(nlat_halo * nlon_kx, block.y), batch_size); + + const size_t sh_dkv = sizeof(float) * (2 * nchans_in + 2 * nchans_out) * block.y; + + s2_attn_bwd_ring_upsample_dkv_k<<>>( + static_cast(nchans_in), static_cast(nchans_out), nlat_halo, nlon_kx, + static_cast(nlon_out_global), static_cast(pscale_out), static_cast(lon_lo_kx), + static_cast(lat_halo_start), static_cast(nlat_out), static_cast(nlon_out), _kxp, _vxp, + _qyp, _dyp, _row_off, _col_idx, _quad_weights, _alpha_sum, _qdotk_max, _integral_norm, _dkx, _dvx); + CHECK_ERROR("s2_attn_bwd_ring_upsample_dkv_k"); + }); + + C10_CUDA_KERNEL_LAUNCH_CHECK(); + } + + TORCH_LIBRARY_IMPL(attention_kernels, CUDA, m) + { + m.impl("backward_ring_step_upsample_pass1", &s2_attention_bwd_ring_step_upsample_pass1_cuda); + m.impl("backward_ring_step_upsample_pass2", &s2_attention_bwd_ring_step_upsample_pass2_cuda); + } + +} // namespace attention_kernels diff --git a/torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_bwd_upsample.cu b/torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_bwd_upsample.cu index 57dc4269..8745cc55 100644 --- a/torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_bwd_upsample.cu +++ b/torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_bwd_upsample.cu @@ -58,6 +58,7 @@ #include "attention_cuda.cuh" #include #include +#include #include #include @@ -401,6 +402,8 @@ namespace attention_kernels static_cast(nlat_out), static_cast(nlon_out), _kxp, _vxp, _qyp, _dyp, _row_off, _col_idx, _quad_weights, _dkxp, _dvxp, _dqyp, stream); }); + + C10_CUDA_KERNEL_LAUNCH_CHECK(); } } // namespace attention_kernels diff --git a/torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_fwd.cu b/torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_fwd.cu index c102e675..0fee7df0 100644 --- a/torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_fwd.cu +++ b/torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_fwd.cu @@ -35,6 +35,7 @@ #include #include #include +#include #include diff --git a/torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_fwd_ring.cu b/torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_fwd_ring.cu index ddf7c059..84862e49 100644 --- a/torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_fwd_ring.cu +++ b/torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_fwd_ring.cu @@ -35,6 +35,7 @@ #include #include #include +#include #include diff --git a/torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_fwd_ring_upsample.cu b/torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_fwd_ring_upsample.cu new file mode 100644 index 00000000..08d83727 --- /dev/null +++ b/torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_fwd_ring_upsample.cu @@ -0,0 +1,355 @@ +// coding=utf-8 +// +// SPDX-FileCopyrightText: Copyright (c) 2026 The torch-harmonics Authors. All rights reserved. +// SPDX-License-Identifier: BSD-3-Clause +// +// Redistribution and use in source and binary forms, with or without +// modification, are permitted provided that the following conditions are met: +// +// 1. Redistributions of source code must retain the above copyright notice, this +// list of conditions and the following disclaimer. +// +// 2. Redistributions in binary form must reproduce the above copyright notice, +// this list of conditions and the following disclaimer in the documentation +// and/or other materials provided with the distribution. +// +// 3. Neither the name of the copyright holder nor the names of its +// contributors may be used to endorse or promote products derived from +// this software without specific prior written permission. +// +// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +// DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +// FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +// DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +// SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +// CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +// OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +// ===================================================================================== +// Upsample (scatter-style) attention forward, RING-STEP variant — CUDA +// ===================================================================================== +// +// Used by DistributedNeighborhoodAttentionS2 in the upsample direction +// (nlon_out % nlon_in == 0). K/V live on the coarse INPUT grid and are sharded +// along longitude across an azimuth process group; Q and the output live on the +// fine OUTPUT grid and stay local. Each call processes one rotating K/V chunk. +// +// psi convention (built by _build_local_psi_upsample in distributed_attention.py): +// row_off : indexed by hi_local in [0, nlat_halo], the halo-padded LOCAL +// input latitude rows; hi_global = lat_halo_start + hi_local. +// Pole-padding rows (hi_global outside the global grid) are empty. +// col_idx : ho_local * nlon_out_global + wo_shifted, where ho_local indexes +// the LOCAL output latitude rows and +// wo_shifted = (wo_canonical - lon_lo_out) mod nlon_out_global. +// The kernel evaluates +// w = (wo_shifted + pscale_out * wi_global) mod nlon_out_global +// which equals (wo_global - lon_lo_out) mod nlon_out_global; the +// target output cell is local iff w < nlon_out (the LOCAL width). +// +// Because the softmax-normalization cell (fine output) is NOT the row key +// (coarse input), the per-ring-step online softmax cannot be done warp-locally +// like in the gather ring kernel. Instead each step runs the same 3-phase +// scheme the gather ring uses for its long rows: +// phase 1 (max): scatter q.k -> atomicMax into qdotk_max_curr (a copy of +// the running qdotk_max_buf) +// phase 2 (rescale): per local output cell, y_acc / alpha_sum are rescaled by +// exp(old_max - new_max) and qdotk_max_buf := qdotk_max_curr +// phase 3 (acc): scatter alpha = exp(q.k - new_max)*w_quad -> atomicAdd +// into alpha_sum_buf, and alpha * v -> atomicAdd into y_acc +// Finalization (y = y_acc / alpha_sum) happens once in Python after the last step. +// +// Generic-only (scalar loads, no long/short row split): correctness path +// mirroring the serial upsample scatter kernels in attention_cuda_fwd_upsample.cu. +// ===================================================================================== + +#include "attention_cuda.cuh" +#include +#include +#include + +#include +#include + +#include "cudamacro.h" +#include "attention_cuda_utils.cuh" + +#define THREADS (64) + +namespace attention_kernels +{ + + // float atomicMax (no native float overload); CAS loop, correct for any sign. + __device__ __forceinline__ float ring_up_atomicMaxf(float *addr, float val) + { + int *ai = reinterpret_cast(addr); + int old = *ai; + while (val > __int_as_float(old)) { + const int assumed = old; + old = atomicCAS(ai, assumed, __float_as_int(val)); + if (old == assumed) { break; } + } + return __int_as_float(old); + } + + // map (wo_shifted, wi_global) -> local output longitude; returns -1 if the + // target cell is not owned by this rank. + __device__ __forceinline__ int ring_up_local_wo(int wo_shifted, int wi_global, int pscale_out, int nlon_out_global, + int nlon_out_local) + { + int w = wo_shifted + pscale_out * wi_global; // < 2*nlon_out_global + if (w >= nlon_out_global) { w -= nlon_out_global; } + return (w < nlon_out_local) ? w : -1; + } + + // phase 1: per coarse chunk cell, scatter q.k into the per-output-cell running max. + template + __global__ __launch_bounds__(THREADS_PER_BLOCK) void s2_attn_fwd_ring_upsample_max_k( + int nchan_in, int nlat_halo, int nlon_kx, int nlon_out_global, int pscale_out, int lon_lo_kx, int nlat_out, + int nlon_out, const STORAGE_T *__restrict__ kx, const STORAGE_T *__restrict__ qy, + const int64_t *__restrict__ row_off, const int64_t *__restrict__ col_idx, float *__restrict__ qdotk_max_curr) + { + extern __shared__ float shext[]; + float *sh_k = shext + threadIdx.y * nchan_in; + + const int batch = blockIdx.y; + const int wid = blockIdx.x * blockDim.y + threadIdx.y; + if (wid >= nlat_halo * nlon_kx) { return; } + + const int tidx = threadIdx.x; + const int hi = wid / nlon_kx; // LOCAL halo row + const int wi_local = wid - hi * nlon_kx; + const int wi_global = lon_lo_kx + wi_local; + + const int64_t rbeg = row_off[hi]; + const int rlen = static_cast(row_off[hi + 1] - rbeg); + if (rlen == 0) { return; } // empty row (e.g. pole padding) + const int64_t *col_hi = col_idx + rbeg; + + kx += int64_t(batch) * nlat_halo * nlon_kx * nchan_in + (int64_t(hi) * nlon_kx + wi_local) * nchan_in; + qy += int64_t(batch) * nlat_out * nlon_out * nchan_in; + qdotk_max_curr += int64_t(batch) * nlat_out * nlon_out; + + for (int chan = tidx; chan < nchan_in; chan += WARP_SIZE) { sh_k[chan] = vload(kx, chan); } + + for (int off = 0; off < rlen; off++) { + const int64_t col = col_hi[off]; + const int ho = static_cast(col / nlon_out_global); // LOCAL output row + const int wo = ring_up_local_wo(static_cast(col - int64_t(ho) * nlon_out_global), wi_global, + pscale_out, nlon_out_global, nlon_out); + if (wo < 0) { continue; } // target cell not owned by this rank + + const STORAGE_T *_qy = qy + (int64_t(ho) * nlon_out + wo) * nchan_in; + + float qd = 0.f; + for (int chan = tidx; chan < nchan_in; chan += WARP_SIZE) { qd += sh_k[chan] * vload(_qy, chan); } + qd = __warp_sum(qd); + if (tidx == 0) { ring_up_atomicMaxf(&qdotk_max_curr[int64_t(ho) * nlon_out + wo], qd); } + } + } + + // phase 2: per LOCAL output cell, rescale the running state by exp(old_max - new_max) + // and commit the new max into the persistent buffer. + template + __global__ __launch_bounds__(THREADS_PER_BLOCK) void s2_attn_fwd_ring_upsample_rescale_k( + int nchan_out, int nlat_out, int nlon_out, const float *__restrict__ qdotk_max_curr, + float *__restrict__ qdotk_max_buf, float *__restrict__ alpha_sum_buf, float *__restrict__ y_acc) + { + const int batch = blockIdx.y; + const int wid = blockIdx.x * blockDim.y + threadIdx.y; + if (wid >= nlat_out * nlon_out) { return; } + const int tidx = threadIdx.x; + + const int64_t cell = int64_t(batch) * nlat_out * nlon_out + wid; + const float qdotk_max_old = qdotk_max_buf[cell]; + const float qdotk_max_new = qdotk_max_curr[cell]; + + // covers the untouched case, including -inf == -inf (exp would yield NaN) + if (qdotk_max_old == qdotk_max_new) { return; } + + const float corr = expf(qdotk_max_old - qdotk_max_new); + + float *_y_acc = y_acc + cell * nchan_out; + for (int chan = tidx; chan < nchan_out; chan += WARP_SIZE) { _y_acc[chan] *= corr; } + + if (tidx == 0) { + alpha_sum_buf[cell] *= corr; + qdotk_max_buf[cell] = qdotk_max_new; + } + } + + // phase 3: per coarse chunk cell, scatter exp(q.k - max)*w into alpha_sum and + // exp(...)*w*v into y_acc. + template + __global__ __launch_bounds__(THREADS_PER_BLOCK) void s2_attn_fwd_ring_upsample_acc_k( + int nchan_in, int nchan_out, int nlat_halo, int nlon_kx, int nlon_out_global, int pscale_out, int lon_lo_kx, + int lat_halo_start, int nlat_out, int nlon_out, const STORAGE_T *__restrict__ kx, + const STORAGE_T *__restrict__ vx, const STORAGE_T *__restrict__ qy, const int64_t *__restrict__ row_off, + const int64_t *__restrict__ col_idx, const float *__restrict__ quad_weights, + const float *__restrict__ qdotk_max_buf, float *__restrict__ alpha_sum_buf, float *__restrict__ y_acc) + { + extern __shared__ float shext[]; + float *sh_k = shext + threadIdx.y * (nchan_in + nchan_out); + float *sh_v = sh_k + nchan_in; + + const int batch = blockIdx.y; + const int wid = blockIdx.x * blockDim.y + threadIdx.y; + if (wid >= nlat_halo * nlon_kx) { return; } + + const int tidx = threadIdx.x; + const int hi = wid / nlon_kx; // LOCAL halo row + const int wi_local = wid - hi * nlon_kx; + const int wi_global = lon_lo_kx + wi_local; + + const int64_t rbeg = row_off[hi]; + const int rlen = static_cast(row_off[hi + 1] - rbeg); + if (rlen == 0) { return; } // empty row (e.g. pole padding) + const int64_t *col_hi = col_idx + rbeg; + + kx += int64_t(batch) * nlat_halo * nlon_kx * nchan_in + (int64_t(hi) * nlon_kx + wi_local) * nchan_in; + vx += int64_t(batch) * nlat_halo * nlon_kx * nchan_out + (int64_t(hi) * nlon_kx + wi_local) * nchan_out; + qy += int64_t(batch) * nlat_out * nlon_out * nchan_in; + qdotk_max_buf += int64_t(batch) * nlat_out * nlon_out; + alpha_sum_buf += int64_t(batch) * nlat_out * nlon_out; + y_acc += int64_t(batch) * nlat_out * nlon_out * nchan_out; + + for (int chan = tidx; chan < nchan_in; chan += WARP_SIZE) { sh_k[chan] = vload(kx, chan); } + for (int chan = tidx; chan < nchan_out; chan += WARP_SIZE) { sh_v[chan] = vload(vx, chan); } + + // quad_weights is indexed by the GLOBAL input latitude + const float qw = quad_weights[lat_halo_start + hi]; + + for (int off = 0; off < rlen; off++) { + const int64_t col = col_hi[off]; + const int ho = static_cast(col / nlon_out_global); // LOCAL output row + const int wo = ring_up_local_wo(static_cast(col - int64_t(ho) * nlon_out_global), wi_global, + pscale_out, nlon_out_global, nlon_out); + if (wo < 0) { continue; } // target cell not owned by this rank + + const int64_t cell = int64_t(ho) * nlon_out + wo; + const STORAGE_T *_qy = qy + cell * nchan_in; + + float qd = 0.f; + for (int chan = tidx; chan < nchan_in; chan += WARP_SIZE) { qd += sh_k[chan] * vload(_qy, chan); } + qd = __warp_sum(qd); + + const float alpha = expf(qd - qdotk_max_buf[cell]) * qw; + if (tidx == 0) { atomicAdd(&alpha_sum_buf[cell], alpha); } + float *_y_acc = y_acc + cell * nchan_out; + for (int chan = tidx; chan < nchan_out; chan += WARP_SIZE) { atomicAdd(&_y_acc[chan], alpha * sh_v[chan]); } + } + } + + // host launcher for one ring step (allocates the temporary max buffer, runs + // the three phases). STORAGE_T deduces from the activation pointers. + template + static void launch_attn_fwd_ring_upsample_step(int batch_size, int nchans_in, int nchans_out, int nlat_halo, + int nlon_kx, int nlon_out_global, int pscale_out, int lon_lo_kx, + int lat_halo_start, int nlat_out, int nlon_out, STORAGE_T *_kxp, + STORAGE_T *_vxp, STORAGE_T *_qyp, int64_t *_row_off, + int64_t *_col_idx, float *_quad_weights, float *_y_acc, + float *_alpha_sum, float *_qdotk_max, cudaStream_t stream) + { + // temporary buffer holding the updated max; starts as a copy of the + // running max so untouched cells stay unchanged (rescale is a no-op there) + torch::Tensor qdotk_max_curr + = torch::from_blob(_qdotk_max, {int64_t(batch_size) * nlat_out * nlon_out}, + torch::TensorOptions().dtype(torch::kFloat32).device(torch::kCUDA)) + .clone(); + float *_qdotk_max_curr = reinterpret_cast(qdotk_max_curr.data_ptr()); + + dim3 block(WARP_SIZE, THREADS / WARP_SIZE); + dim3 grid_in(DIV_UP(nlat_halo * nlon_kx, block.y), batch_size); + dim3 grid_out(DIV_UP(nlat_out * nlon_out, block.y), batch_size); + + const size_t sh_max = sizeof(float) * nchans_in * block.y; + const size_t sh_acc = sizeof(float) * (nchans_in + nchans_out) * block.y; + + s2_attn_fwd_ring_upsample_max_k + <<>>(nchans_in, nlat_halo, nlon_kx, nlon_out_global, pscale_out, lon_lo_kx, + nlat_out, nlon_out, _kxp, _qyp, _row_off, _col_idx, _qdotk_max_curr); + CHECK_ERROR("s2_attn_fwd_ring_upsample_max_k"); + + // also commits qdotk_max_curr into the persistent _qdotk_max buffer + s2_attn_fwd_ring_upsample_rescale_k<<>>( + nchans_out, nlat_out, nlon_out, _qdotk_max_curr, _qdotk_max, _alpha_sum, _y_acc); + CHECK_ERROR("s2_attn_fwd_ring_upsample_rescale_k"); + + s2_attn_fwd_ring_upsample_acc_k<<>>( + nchans_in, nchans_out, nlat_halo, nlon_kx, nlon_out_global, pscale_out, lon_lo_kx, lat_halo_start, nlat_out, + nlon_out, _kxp, _vxp, _qyp, _row_off, _col_idx, _quad_weights, _qdotk_max, _alpha_sum, _y_acc); + CHECK_ERROR("s2_attn_fwd_ring_upsample_acc_k"); + } + + void s2_attention_fwd_ring_step_upsample_cuda(at::Tensor kx, at::Tensor vx, at::Tensor qy, at::Tensor y_acc, + at::Tensor alpha_sum_buf, at::Tensor qdotk_max_buf, + at::Tensor quad_weights, at::Tensor psi_col_idx, + at::Tensor psi_row_off, int64_t nlon_in, int64_t nlon_out_global, + int64_t pscale_out, int64_t lon_lo_kx, int64_t lat_halo_start, + int64_t nlat_out, int64_t nlon_out) + { + CHECK_CUDA_INPUT_TENSOR(kx); + CHECK_CUDA_INPUT_TENSOR(vx); + CHECK_CUDA_INPUT_TENSOR(qy); + CHECK_CUDA_TENSOR(y_acc); + CHECK_CUDA_TENSOR(alpha_sum_buf); + CHECK_CUDA_TENSOR(qdotk_max_buf); + CHECK_CUDA_TENSOR(quad_weights); + CHECK_CUDA_TENSOR(psi_col_idx); + CHECK_CUDA_TENSOR(psi_row_off); + + const int batch_size = kx.size(0); + const int nlat_halo = kx.size(2); // kx is [B,C,H,W], H is dim 2 + const int nlon_kx = kx.size(3); // W is dim 3 + const size_t nchans_in = qy.size(1); + const size_t nchans_out = vx.size(1); + + auto stream = at::cuda::getCurrentCUDAStream().stream(); + + int64_t *_row_off = reinterpret_cast(psi_row_off.data_ptr()); + int64_t *_col_idx = reinterpret_cast(psi_col_idx.data_ptr()); + float *_quad_weights = reinterpret_cast(quad_weights.data_ptr()); + float *_y_acc = reinterpret_cast(y_acc.data_ptr()); + float *_alpha_sum = reinterpret_cast(alpha_sum_buf.data_ptr()); + float *_qdotk_max = reinterpret_cast(qdotk_max_buf.data_ptr()); + + // ATen dispatch over the input dtype. Tier B: native storage — kernels + // widen to fp32 at load; the state buffers stay fp32. + AT_DISPATCH_FLOATING_TYPES_AND2( + at::kHalf, at::kBFloat16, qy.scalar_type(), "s2_attention_fwd_ring_step_upsample_cuda", [&] { + using storage_t = scalar_t; + + torch::Tensor kxP = kx; + torch::Tensor vxP = vx; + torch::Tensor qyP = qy; + + bool kx_is_channels_last = kxP.strides()[1] == 1; + bool vx_is_channels_last = vxP.strides()[1] == 1; + bool qy_is_channels_last = qyP.strides()[1] == 1; + + if (!kx_is_channels_last) { kxP = permute_4D_to0231(kxP); } + if (!vx_is_channels_last) { vxP = permute_4D_to0231(vxP); } + if (!qy_is_channels_last) { qyP = permute_4D_to0231(qyP); } + + storage_t *_kxp = reinterpret_cast(kxP.data_ptr()); + storage_t *_vxp = reinterpret_cast(vxP.data_ptr()); + storage_t *_qyp = reinterpret_cast(qyP.data_ptr()); + + launch_attn_fwd_ring_upsample_step( + batch_size, static_cast(nchans_in), static_cast(nchans_out), nlat_halo, nlon_kx, + static_cast(nlon_out_global), static_cast(pscale_out), static_cast(lon_lo_kx), + static_cast(lat_halo_start), static_cast(nlat_out), static_cast(nlon_out), _kxp, + _vxp, _qyp, _row_off, _col_idx, _quad_weights, _y_acc, _alpha_sum, _qdotk_max, stream); + }); + + C10_CUDA_KERNEL_LAUNCH_CHECK(); + } + + TORCH_LIBRARY_IMPL(attention_kernels, CUDA, m) + { + m.impl("forward_ring_step_upsample", &s2_attention_fwd_ring_step_upsample_cuda); + } + +} // namespace attention_kernels diff --git a/torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_fwd_upsample.cu b/torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_fwd_upsample.cu index bdbbdcf6..32025dc6 100644 --- a/torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_fwd_upsample.cu +++ b/torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_fwd_upsample.cu @@ -54,6 +54,7 @@ #include "attention_cuda.cuh" #include #include +#include #include #include @@ -288,6 +289,8 @@ namespace attention_kernels static_cast(nlat_out), static_cast(nlon_out), _kxp, _vxp, _qyp, _row_off, _col_idx, _quad_weights, _yp, stream); }); + + C10_CUDA_KERNEL_LAUNCH_CHECK(); } } // namespace attention_kernels diff --git a/torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_utils.cu b/torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_utils.cu index 35073677..930eb324 100644 --- a/torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_utils.cu +++ b/torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_utils.cu @@ -33,6 +33,7 @@ #include #include #include +#include #include @@ -103,6 +104,8 @@ namespace attention_kernels CHECK_CUDA(cub::DeviceRadixSort::SortPairsDescending(_temp_storage_d, temp_storage_bytes, _rlen_d, _rlen_sort_d, _rids_d, _rids_sort_d, nlat_out, 0, sizeof(*_rlen_d) * 8, stream)); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return rids_sort_d; } // END - CSR rows sorting kernels and functions @@ -137,6 +140,7 @@ namespace attention_kernels ([&] { launch_permute_to0231(src, dst); })); CHECK_ERROR("permute_to0231_k_tile_sm100"); } + C10_CUDA_KERNEL_LAUNCH_CHECK(); return dst; } @@ -161,6 +165,7 @@ namespace attention_kernels ([&] { launch_permute_to0312(src, dst); })); CHECK_ERROR("permute_to0312_k_tile_sm100"); } + C10_CUDA_KERNEL_LAUNCH_CHECK(); return dst; } @@ -244,6 +249,7 @@ namespace attention_kernels auto stream = at::cuda::getCurrentCUDAStream().stream(); get_rlen_boundary_k<<<1, 1024, 0, stream>>>(thres, split_len, nrows, row_idx, row_off, num_lr, max_rl0, max_rl1); + C10_CUDA_KERNEL_LAUNCH_CHECK(); at::Tensor tmp_h = tmp_d.cpu(); int64_t *tmp_ptr_h = tmp_h.data_ptr(); diff --git a/torch_harmonics/disco/optimized/kernels_cuda/disco_cuda_bwd.cu b/torch_harmonics/disco/optimized/kernels_cuda/disco_cuda_bwd.cu index c7ca08ff..eb1b5217 100644 --- a/torch_harmonics/disco/optimized/kernels_cuda/disco_cuda_bwd.cu +++ b/torch_harmonics/disco/optimized/kernels_cuda/disco_cuda_bwd.cu @@ -32,6 +32,7 @@ #include "disco_cuda.cuh" #include +#include #include namespace disco_kernels @@ -303,6 +304,8 @@ namespace disco_kernels exit(EXIT_FAILURE); } + C10_CUDA_KERNEL_LAUNCH_CHECK(); + // convert type if requested out = out.to(inp.dtype()); diff --git a/torch_harmonics/disco/optimized/kernels_cuda/disco_cuda_fwd.cu b/torch_harmonics/disco/optimized/kernels_cuda/disco_cuda_fwd.cu index 03442fe2..92d2c381 100644 --- a/torch_harmonics/disco/optimized/kernels_cuda/disco_cuda_fwd.cu +++ b/torch_harmonics/disco/optimized/kernels_cuda/disco_cuda_fwd.cu @@ -32,6 +32,7 @@ #include "disco_cuda.cuh" #include +#include #include namespace disco_kernels @@ -304,6 +305,8 @@ namespace disco_kernels exit(EXIT_FAILURE); } + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return out; } diff --git a/torch_harmonics/disco/optimized/kernels_cuda/disco_cuda_fwd_dense_kpacked_sm100.cu b/torch_harmonics/disco/optimized/kernels_cuda/disco_cuda_fwd_dense_kpacked_sm100.cu index d2e91d4d..e5bab389 100644 --- a/torch_harmonics/disco/optimized/kernels_cuda/disco_cuda_fwd_dense_kpacked_sm100.cu +++ b/torch_harmonics/disco/optimized/kernels_cuda/disco_cuda_fwd_dense_kpacked_sm100.cu @@ -65,6 +65,7 @@ #include "disco_cuda_ptx.cuh" #include +#include #include #include #include @@ -338,6 +339,9 @@ namespace disco_kernels else launch(&disco_fwd_dense_kpacked_tcgen05_blk_k, __half {}); } + + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return out; } diff --git a/torch_harmonics/disco/optimized/kernels_cuda/disco_cuda_fwd_dense_kpacked_sm90.cu b/torch_harmonics/disco/optimized/kernels_cuda/disco_cuda_fwd_dense_kpacked_sm90.cu index ac50f38d..7c199f50 100644 --- a/torch_harmonics/disco/optimized/kernels_cuda/disco_cuda_fwd_dense_kpacked_sm90.cu +++ b/torch_harmonics/disco/optimized/kernels_cuda/disco_cuda_fwd_dense_kpacked_sm90.cu @@ -98,6 +98,7 @@ #include "disco_cuda_ptx.cuh" #include +#include #include #include #include @@ -421,6 +422,9 @@ namespace disco_kernels else launch(&disco_fwd_dense_kpacked_wgmma_blk_k, __half {}); } + + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return out; } diff --git a/torch_harmonics/distributed/distributed_attention.py b/torch_harmonics/distributed/distributed_attention.py index e721be4d..f067ee95 100644 --- a/torch_harmonics/distributed/distributed_attention.py +++ b/torch_harmonics/distributed/distributed_attention.py @@ -438,6 +438,350 @@ def backward(ctx, dy, _dalpha_sum, _dqdotk_max): return (dkw, dvw, dqy, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None) +class _RingNeighborhoodAttentionUpsampleFn(torch.autograd.Function): + """Forward ring attention + backward ring for the UPSAMPLE (scatter) direction. + + K/V live on the coarse input grid (sharded, halo-padded in lat, rotating + around the azimuth ring); Q and the output live on the fine output grid and + stay local. + + kw, vw : [B*nh, C_k/C_v, H_halo, W_in_local] channels-first, lat-halo-padded + qw : [B*nh, C_k, H_out_local, W_out_local] channels-first + + State buffers use channels-last layout as required by the CUDA kernels: + y_acc : [B, H_out, W_out, C_v] + alpha_k/kvw : [B, H_out, W_out, C_k] + alpha_sum/qdotk_max/integral : [B, H_out, W_out] + + The local psi is built by _build_local_psi_upsample: rows are keyed by the + halo-padded LOCAL input latitude, cols encode (ho_local, wo_shifted) on the + fine output grid with wo pre-shifted by -lon_lo_out. + """ + + @staticmethod + @torch.amp.custom_fwd(device_type="cuda") + def forward( + kw, + vw, + qw, + psi_col_idx, + psi_roff_idx, + quad_weights, + nlon_in: int, + nlon_out_global: int, + pscale_out: int, + lon_chunk_starts: list, + nlon_kx_list: list, + lat_halo_start: int, + nlat_out_local: int, + nlon_out_local: int, + r_lat: int, + az_group, + az_rank: int, + az_size: int, + ): + B, _, _, _ = kw.shape + _, C_v, _, _ = vw.shape + device = kw.device + + # Capture input dtype so we can cast the user-visible output (y_out) + # back at the end of forward. The internal accumulators stay fp32 — + # softmax stability requires that, and the saved alpha_sum/qdotk_max + # feed fp32 math in backward. + inp_dtype = kw.dtype + + # Allocate state buffers in formats expected by the CUDA kernels: + # y_acc: channels-last [B, H, W, C_v]; scalars: [B, H, W] + y_acc = torch.zeros(B, nlat_out_local, nlon_out_local, C_v, device=device, dtype=torch.float32) + alpha_sum = torch.zeros(B, nlat_out_local, nlon_out_local, device=device, dtype=torch.float32) + qdotk_max = torch.full((B, nlat_out_local, nlon_out_local), float("-inf"), device=device, dtype=torch.float32) + + kw_chunk = kw.contiguous() + vw_chunk = vw.contiguous() + + for step in range(az_size): + src_rank = (az_rank + step) % az_size + lon_lo_kx = lon_chunk_starts[src_rank] + + # Pre-allocate receive buffers for the NEXT chunk (correct shape) + if step < az_size - 1: + next_src = (az_rank + step + 1) % az_size + recv_kw, recv_vw, reqs = _ring_kv(kw_chunk, vw_chunk, az_group, nlon_kx_list[next_src], nlon_kx_list[next_src]) + + attention_kernels.forward_ring_step_upsample.default( + kw_chunk, + vw_chunk, + qw, + y_acc, + alpha_sum, + qdotk_max, + quad_weights, + psi_col_idx, + psi_roff_idx, + nlon_in, + nlon_out_global, + pscale_out, + lon_lo_kx, + lat_halo_start, + nlat_out_local, + nlon_out_local, + ) + + if step < az_size - 1: + for req in reqs: + req.wait() + kw_chunk = recv_kw.clone() + vw_chunk = recv_vw.clone() + + # Finalize: y = y_acc / alpha_sum (both channels-last layout). Cast + # back to the input dtype to keep the op faithful to its input dtype. + y_out = y_acc / alpha_sum.unsqueeze(-1) # [B, H, W, C_v] + y_out = y_out.permute(0, 3, 1, 2).to(dtype=inp_dtype).contiguous() # [B, C_v, H, W] + + # alpha_sum and qdotk_max are returned so setup_context can save them; + # they are marked non-differentiable there, so backward still only + # receives one gradient argument (dy for y_out). + return y_out, alpha_sum, qdotk_max + + @staticmethod + @_custom_setup_context(device_type="cuda") + def setup_context(ctx, inputs, output): + ( + kw, + vw, + qw, + psi_col_idx, + psi_roff_idx, + quad_weights, + nlon_in, + nlon_out_global, + pscale_out, + lon_chunk_starts, + nlon_kx_list, + lat_halo_start, + nlat_out_local, + nlon_out_local, + r_lat, + az_group, + az_rank, + az_size, + ) = inputs + y_out, alpha_sum, qdotk_max = output + # alpha_sum and qdotk_max are internal accumulators, not true outputs; + # marking them non-differentiable keeps backward's signature as (ctx, dy). + ctx.mark_non_differentiable(alpha_sum, qdotk_max) + ctx.save_for_backward(kw, vw, qw, psi_col_idx, psi_roff_idx, quad_weights, alpha_sum, qdotk_max) + ctx.nlon_in = nlon_in + ctx.nlon_out_global = nlon_out_global + ctx.pscale_out = pscale_out + ctx.lon_chunk_starts = lon_chunk_starts + ctx.nlon_kx_list = nlon_kx_list + ctx.lat_halo_start = lat_halo_start + ctx.nlat_out_local = nlat_out_local + ctx.nlon_out_local = nlon_out_local + ctx.az_group = az_group + ctx.az_rank = az_rank + ctx.az_size = az_size + + @staticmethod + @torch.amp.custom_bwd(device_type="cuda") + def backward(ctx, dy, _dalpha_sum, _dqdotk_max): + # _dalpha_sum and _dqdotk_max are always None (non-differentiable outputs) + (kw, vw, qw, psi_col_idx, psi_roff_idx, quad_weights, fwd_alpha_sum, fwd_qdotk_max) = ctx.saved_tensors + + nlon_in = ctx.nlon_in + nlon_out_global = ctx.nlon_out_global + pscale_out = ctx.pscale_out + lon_chunk_starts = ctx.lon_chunk_starts + nlon_kx_list = ctx.nlon_kx_list + lat_halo_start = ctx.lat_halo_start + nlat_out_local = ctx.nlat_out_local + nlon_out_local = ctx.nlon_out_local + az_group = ctx.az_group + az_rank = ctx.az_rank + az_size = ctx.az_size + + # Autograd contract: skip per-branch work (kernel calls, allreduces) for any + # of {kw, vw, qw} that doesn't need a gradient, and return None in those slots. + kw_needs_grad = ctx.needs_input_grad[0] + vw_needs_grad = ctx.needs_input_grad[1] + qw_needs_grad = ctx.needs_input_grad[2] + + # Defensive: if somehow none of (kw, vw, qw) need grad, there's nothing to compute. + if not (kw_needs_grad or vw_needs_grad or qw_needs_grad): + return (None,) * 18 + + B, C_k, H_halo, _ = kw.shape + _, C_v, _, _ = vw.shape + device = kw.device + + # Capture input dtypes so the returned grads can be cast back. The + # backward kernels consume kw/vw/qw/dy in their native dtype (widen at + # load, fp32 compute/accumulation), keeping the backward ring exchange + # at 16-bit under AMP. + kw_dtype = kw.dtype + vw_dtype = vw.dtype + qw_dtype = qw.dtype + dy_cf = dy.contiguous() # channels-first [B, C_v, H, W], native dtype + + # ---------------------------------------------------------------- + # Backward pass 1: re-accumulate {integral, alpha_k, alpha_kvw} via + # ring, using the SAVED forward alpha_sum/qdotk_max (no max recompute + # is needed in the upsample direction — the forward-final softmax + # stats are authoritative). Required whenever any of (kw, vw, qw) + # needs grad: integral feeds pass-2's integral_norm and dqy reads + # alpha_k/alpha_kvw. The kernel writes all buffers in one call, so + # pass-1 cannot be pruned per-branch. + # ---------------------------------------------------------------- + integral_buf = torch.zeros(B, nlat_out_local, nlon_out_local, device=device, dtype=torch.float32) + alpha_k_buf = torch.zeros(B, nlat_out_local, nlon_out_local, C_k, device=device, dtype=torch.float32) + alpha_kvw_buf = torch.zeros_like(alpha_k_buf) + + kw_chunk = kw.contiguous() + vw_chunk = vw.contiguous() + + for step in range(az_size): + src_rank = (az_rank + step) % az_size + lon_lo_kx = lon_chunk_starts[src_rank] + + if step < az_size - 1: + next_src = (az_rank + step + 1) % az_size + recv_kw, recv_vw, reqs = _ring_kv(kw_chunk, vw_chunk, az_group, nlon_kx_list[next_src], nlon_kx_list[next_src]) + + attention_kernels.backward_ring_step_upsample_pass1.default( + kw_chunk, + vw_chunk, + qw, + dy_cf, + fwd_qdotk_max, + integral_buf, + alpha_k_buf, + alpha_kvw_buf, + quad_weights, + psi_col_idx, + psi_roff_idx, + nlon_in, + nlon_out_global, + pscale_out, + lon_lo_kx, + lat_halo_start, + nlat_out_local, + nlon_out_local, + ) + + if step < az_size - 1: + for req in reqs: + req.wait() + kw_chunk = recv_kw.clone() + vw_chunk = recv_vw.clone() + + # ---------------------------------------------------------------- + # Finalize pass-1 outputs. + # ---------------------------------------------------------------- + alpha_sum_inv = 1.0 / fwd_alpha_sum # [B, H, W] + + # integral_norm only feeds pass-2; skip if neither kw nor vw needs grad. + if kw_needs_grad or vw_needs_grad: + integral_norm = integral_buf * alpha_sum_inv # [B, H, W] + + # dqy[b,h,w,c] = inv_sq*(alpha_sum*alpha_kvw - integral*alpha_k) + if qw_needs_grad: + alpha_sum_inv_sq = alpha_sum_inv**2 + dqy_cl = alpha_sum_inv_sq.unsqueeze(-1) * (fwd_alpha_sum.unsqueeze(-1) * alpha_kvw_buf - integral_buf.unsqueeze(-1) * alpha_k_buf) # [B, H, W, C_k] + dqy = dqy_cl.permute(0, 3, 1, 2).to(dtype=qw_dtype).contiguous() # [B, C_k, H, W] + else: + dqy = None + + # ---------------------------------------------------------------- + # Backward pass 2: accumulate dkw/dvw contributions per chunk. + # Each GPU computes its LOCAL outputs' contribution to every lon chunk + # it visits; then allreduce across azimuth ranks, extract local chunk. + # TODO: replace allreduce with ring reduce-scatter for efficiency. + # ---------------------------------------------------------------- + if kw_needs_grad or vw_needs_grad: + kw_chunk = kw.contiguous() + vw_chunk = vw.contiguous() + nlon_in_total = sum(nlon_kx_list) + dkw_full_cl = torch.zeros(B, H_halo, nlon_in_total, C_k, device=device, dtype=torch.float32) if kw_needs_grad else None + dvw_full_cl = torch.zeros(B, H_halo, nlon_in_total, C_v, device=device, dtype=torch.float32) if vw_needs_grad else None + + for step in range(az_size): + src_rank = (az_rank + step) % az_size + lon_lo_kx = lon_chunk_starts[src_rank] + nlon_kx = nlon_kx_list[src_rank] + + # Channels-last gradient buffers for this chunk (both required by the + # fused kernel signature; we discard the one we don't need). + dkw_chunk_cl = torch.zeros(B, H_halo, nlon_kx, C_k, device=device, dtype=torch.float32) + dvw_chunk_cl = torch.zeros(B, H_halo, nlon_kx, C_v, device=device, dtype=torch.float32) + + attention_kernels.backward_ring_step_upsample_pass2.default( + kw_chunk, + vw_chunk, + qw, + dy_cf, + fwd_alpha_sum, + fwd_qdotk_max, + integral_norm, + dkw_chunk_cl, + dvw_chunk_cl, + quad_weights, + psi_col_idx, + psi_roff_idx, + nlon_in, + nlon_out_global, + pscale_out, + lon_lo_kx, + lat_halo_start, + nlat_out_local, + nlon_out_local, + ) + + if kw_needs_grad: + dkw_full_cl[:, :, lon_lo_kx : lon_lo_kx + nlon_kx, :].add_(dkw_chunk_cl) + if vw_needs_grad: + dvw_full_cl[:, :, lon_lo_kx : lon_lo_kx + nlon_kx, :].add_(dvw_chunk_cl) + + if step < az_size - 1: + next_src = (az_rank + step + 1) % az_size + recv_kw, recv_vw, reqs = _ring_kv(kw_chunk, vw_chunk, az_group, nlon_kx_list[next_src], nlon_kx_list[next_src]) + for req in reqs: + req.wait() + kw_chunk = recv_kw.clone() + vw_chunk = recv_vw.clone() + + # Per-branch allreduce — only the branches we'll return. + if az_size > 1 and az_group is not None: + if kw_needs_grad: + dist.all_reduce(dkw_full_cl, group=az_group) + if vw_needs_grad: + dist.all_reduce(dvw_full_cl, group=az_group) + + my_lo = lon_chunk_starts[az_rank] + my_nlon = nlon_kx_list[az_rank] + # Extract local chunk and convert channels-last → channels-first. + # No halo stripping: dkw/dvw must match kw/vw shape (= key_halo/value_halo). + if kw_needs_grad: + dkw_cl = dkw_full_cl[:, :, my_lo : my_lo + my_nlon, :].contiguous() + dkw = dkw_cl.permute(0, 3, 1, 2).to(dtype=kw_dtype).contiguous() # [B, C_k, H_halo, W_local] + else: + dkw = None + if vw_needs_grad: + dvw_cl = dvw_full_cl[:, :, my_lo : my_lo + my_nlon, :].contiguous() + dvw = dvw_cl.permute(0, 3, 1, 2).to(dtype=vw_dtype).contiguous() # [B, C_v, H_halo, W_local] + else: + dvw = None + else: + dkw = None + dvw = None + + # Return grads for (kw, vw, qw, psi_col, psi_roff, quad_weights, + # nlon_in, nlon_out_global, pscale_out, lon_chunk_starts, + # nlon_kx_list, lat_halo_start, nlat_out_local, nlon_out_local, + # r_lat, az_group, az_rank, az_size) + return (dkw, dvw, dqy, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None) + + # --------------------------------------------------------------------------- # Distributed Neighborhood Attention on the 2-sphere # --------------------------------------------------------------------------- @@ -454,6 +798,12 @@ class DistributedNeighborhoodAttentionS2(NeighborhoodAttentionS2): of key/value chunks over the azimuth group so that every output point can attend to its full spherical neighborhood. + All three directions of the serial layer are supported: self-attention + (in_shape == out_shape), downsampling cross-attention (gather kernels, + nlon_in % nlon_out == 0) and upsampling cross-attention (scatter kernels, + nlon_out % nlon_in == 0). In all cases K/V (which live on the input grid) + rotate around the azimuth ring while Q and the softmax state stay local. + Inherits learnable parameters from :class:`NeighborhoodAttentionS2`. """ @@ -511,44 +861,71 @@ def __init__( self.nlat_out_local = self.lat_out_shapes[self.comm_rank_polar] self.nlon_out_local = self.lon_out_shapes[self.comm_rank_azimuth] - # Downsampling invariant: every azimuth rank must carry the same lon pscale. - # The global `nlon_in % nlon_out == 0` check is inherited from the serial + # Uniform-pscale invariant: every azimuth rank must carry the same lon pscale. + # The global divisibility check is inherited from the serial # NeighborhoodAttentionS2.__init__, but that is not sufficient in distributed: # if compute_split_shapes hands different ranks different local pscales # (e.g. nlon_in=12, nlon_out=4, comm_size_azimuth=3 -> [4,4,4] vs [2,1,1]), # the p-shift mapping in the ring exchange is ill-defined. - pscale_lon = self.nlon_in // self.nlon_out - for r, (lon_in_r, lon_out_r) in enumerate(zip(self.lon_in_shapes, self.lon_out_shapes)): - if lon_in_r != pscale_lon * lon_out_r: - raise ValueError( - f"DistributedNeighborhoodAttentionS2: inconsistent azimuth split at rank {r}: " - f"nlon_in_local={lon_in_r}, nlon_out_local={lon_out_r}. " - f"Every azimuth rank must satisfy nlon_in_local == (nlon_in // nlon_out) * nlon_out_local " - f"= {pscale_lon} * nlon_out_local. " - f"Choose (nlon_in, nlon_out, comm_size_azimuth) so that compute_split_shapes " - f"produces uniform local pscale." - ) - - # global lon offsets + if self.upsample: + pscale_lon = self.nlon_out // self.nlon_in + for r, (lon_in_r, lon_out_r) in enumerate(zip(self.lon_in_shapes, self.lon_out_shapes)): + if lon_out_r != pscale_lon * lon_in_r: + raise ValueError( + f"DistributedNeighborhoodAttentionS2: inconsistent azimuth split at rank {r}: " + f"nlon_in_local={lon_in_r}, nlon_out_local={lon_out_r}. " + f"Every azimuth rank must satisfy nlon_out_local == (nlon_out // nlon_in) * nlon_in_local " + f"= {pscale_lon} * nlon_in_local. " + f"Choose (nlon_in, nlon_out, comm_size_azimuth) so that compute_split_shapes " + f"produces uniform local pscale." + ) + else: + pscale_lon = self.nlon_in // self.nlon_out + for r, (lon_in_r, lon_out_r) in enumerate(zip(self.lon_in_shapes, self.lon_out_shapes)): + if lon_in_r != pscale_lon * lon_out_r: + raise ValueError( + f"DistributedNeighborhoodAttentionS2: inconsistent azimuth split at rank {r}: " + f"nlon_in_local={lon_in_r}, nlon_out_local={lon_out_r}. " + f"Every azimuth rank must satisfy nlon_in_local == (nlon_in // nlon_out) * nlon_out_local " + f"= {pscale_lon} * nlon_out_local. " + f"Choose (nlon_in, nlon_out, comm_size_azimuth) so that compute_split_shapes " + f"produces uniform local pscale." + ) + + # global lon/lat offsets self.lon_in_starts = list(accumulate([0] + self.lon_in_shapes[:-1])) self.lon_out_starts = list(accumulate([0] + self.lon_out_shapes[:-1])) + self.lat_in_starts = list(accumulate([0] + self.lat_in_shapes[:-1])) self.lat_out_starts = list(accumulate([0] + self.lat_out_shapes[:-1])) self.lon_lo_out = self.lon_out_starts[self.comm_rank_azimuth] self.lat_lo_out = self.lat_out_starts[self.comm_rank_polar] - # ---- build local psi ---- - # The global psi built by the base class covers all output lat rows. - # We filter to only the rows owned by this rank and shift the wi - # component of col_idx by lon_lo_out so that the kernel can use - # local wo directly without knowing the global lon offset. - self._build_local_psi() # also precomputes self.psi_{n_long_rows,max_row_len,mid_row_len} - - # ---- lat halo size ---- - # Compute r_lat from the global psi: maximum |hi_global - ho_global| - # over all (ho, hi) pairs in the neighbourhood. - # Use the lat_out_lo of our polar rank to compute ho_global. - self.r_lat = self._compute_r_lat() + if self.upsample: + # ---- lat halo size ---- + # For the scatter direction psi rows are keyed by hi, so the halo + # radius must be known BEFORE the local psi (whose rows span the + # halo-padded input range) can be built. + self.r_lat = self._compute_r_lat_upsample() + + # ---- build local psi ---- + # Rows are re-keyed to the halo-padded local input lat range, cols + # are filtered to the local output lat rows and the wo component is + # pre-shifted by -lon_lo_out (see _build_local_psi_upsample). + self._build_local_psi_upsample() + else: + # ---- build local psi ---- + # The global psi built by the base class covers all output lat rows. + # We filter to only the rows owned by this rank and shift the wi + # component of col_idx by lon_lo_out so that the kernel can use + # local wo directly without knowing the global lon offset. + self._build_local_psi() # also precomputes self.psi_{n_long_rows,max_row_len,mid_row_len} + + # ---- lat halo size ---- + # Compute r_lat from the global psi: maximum |hi_global - ho_global| + # over all (ho, hi) pairs in the neighbourhood. + # Use the lat_out_lo of our polar rank to compute ho_global. + self.r_lat = self._compute_r_lat() # ----------------------------------------------------------------------- @@ -623,12 +1000,11 @@ def _compute_r_lat(self) -> int: if col_idx.numel() == 0: return 0 - lat_in_starts = list(accumulate([0] + self.lat_in_shapes[:-1])) roff = self.psi_roff_idx r = 0 for rank in range(self.comm_size_polar): - lat_in_lo = lat_in_starts[rank] + lat_in_lo = self.lat_in_starts[rank] lat_in_hi = lat_in_lo + self.lat_in_shapes[rank] lat_out_lo = self.lat_out_starts[rank] lat_out_hi = lat_out_lo + self.lat_out_shapes[rank] @@ -645,6 +1021,108 @@ def _compute_r_lat(self) -> int: return r + # ----------------------------------------------------------------------- + # upsample (scatter) direction helpers + # ----------------------------------------------------------------------- + + def _compute_r_lat_upsample(self) -> int: + """Max lat halo radius needed across all polar ranks, upsample direction. + + In the scatter psi (rows keyed by input lat hi, cols encoding output + cells), the entries relevant to a polar rank are those whose OUTPUT row + ho falls into its local output shard; the halo is then determined by how + far the corresponding INPUT rows hi reach outside its local input shard. + Computed locally from the global psi (built identically on every rank + by the base class), so no communication is required. + """ + + if polar_group_size() == 1: + return 0 + + col_idx = self.psi_col_idx # global, rows = nlat_in, cols = ho * nlon_out + wo + if col_idx.numel() == 0: + return 0 + + roff = self.psi_roff_idx + # input-lat row index of every nonzero entry + nnz_per_row = roff[1:] - roff[:-1] + hi_of_nz = torch.repeat_interleave(torch.arange(self.nlat_in, dtype=torch.int64, device=col_idx.device), nnz_per_row) + ho = (col_idx // self.nlon_out).long() + + r = 0 + for rank in range(self.comm_size_polar): + lat_in_lo = self.lat_in_starts[rank] + lat_in_hi = lat_in_lo + self.lat_in_shapes[rank] + lat_out_lo = self.lat_out_starts[rank] + lat_out_hi = lat_out_lo + self.lat_out_shapes[rank] + + mask = (ho >= lat_out_lo) & (ho < lat_out_hi) + if not bool(mask.any()): + continue + + hi = hi_of_nz[mask] + r_top = max(0, lat_in_lo - int(hi.min().item())) + r_bot = max(0, int(hi.max().item()) - (lat_in_hi - 1)) + r = max(r, r_top, r_bot) + + return r + + def _build_local_psi_upsample(self): + """Build the local scatter psi for the upsample ring kernels. + + The global psi built by the base class has rows keyed by the input lat + hi in [0, nlat_in) and cols encoding ho * nlon_out + wo_canonical on the + fine output grid (canonical at wi = 0). The local psi + + * re-keys the rows to the halo-padded LOCAL input lat range + [lat_halo_start, lat_halo_start + nlat_halo); pole-padding rows + (hi outside the global grid) are empty, + * keeps only entries whose output row ho falls into the local output + shard and re-keys ho to ho_local = ho - lat_lo_out, + * pre-shifts the wo component by -lon_lo_out (mod nlon_out) so the + kernel's mapping w = (wo_shifted + pscale_out * wi_global) mod + nlon_out directly yields the LOCAL output longitude, with the + locality test w < nlon_out_local. + """ + + nlon_out = self.nlon_out + lat_lo_out = self.lat_lo_out + lat_hi_out = lat_lo_out + self.nlat_out_local + lon_lo_out = self.lon_lo_out + + col_idx_global = self.psi_col_idx # [nnz] int64 + roff_global = self.psi_roff_idx # [nlat_in+1] int64 + + nlat_halo = self.nlat_in_local + 2 * self.r_lat + lat_halo_start = self.lat_in_starts[self.comm_rank_polar] - self.r_lat + + # input-lat row index of every nonzero entry + nnz_per_row = roff_global[1:] - roff_global[:-1] + hi_of_nz = torch.repeat_interleave(torch.arange(self.nlat_in, dtype=torch.int64, device=col_idx_global.device), nnz_per_row) + + ho = col_idx_global // nlon_out + wo = col_idx_global - ho * nlon_out + + # keep entries whose input row lies in the halo-padded local range and + # whose output row is owned by this polar rank + hi_local = hi_of_nz - lat_halo_start + mask = (ho >= lat_lo_out) & (ho < lat_hi_out) & (hi_local >= 0) & (hi_local < nlat_halo) + + hi_sel = hi_local[mask] + ho_sel = ho[mask] - lat_lo_out + wo_sel = (wo[mask] - lon_lo_out) % nlon_out + col_idx_local = ho_sel * nlon_out + wo_sel + + # rebuild the CSR row offsets over the halo-padded local rows; masked + # selection preserves the global row-major order, so col_idx_local is + # already CSR-consistent with roff_local + counts = torch.bincount(hi_sel, minlength=nlat_halo) + roff_local = torch.zeros(nlat_halo + 1, dtype=roff_global.dtype, device=roff_global.device) + roff_local[1:] = torch.cumsum(counts, dim=0) + + self.register_buffer("psi_col_idx_local", col_idx_local.contiguous(), persistent=False) + self.register_buffer("psi_roff_idx_local", roff_local.contiguous(), persistent=False) + # ----------------------------------------------------------------------- def forward( @@ -710,41 +1188,65 @@ def forward( value_halo = value_proj # global lat index of first halo row - lat_in_starts = list(accumulate([0] + self.lat_in_shapes[:-1])) - lat_halo_start = lat_in_starts[self.comm_rank_polar] - self.r_lat + lat_halo_start = self.lat_in_starts[self.comm_rank_polar] - self.r_lat # ---- 3. ring attention ---- - # Global pscale — the kernel must not infer this from local shapes, - # because kernel `nlon_out` is nlon_out_local which differs when az_size > 1. - pscale = self.nlon_in // self.nlon_out # Under autocast, cast k/v/q to the autocast dtype before .apply() — # mirrors PyTorch's autocast-eligible-op dataflow. Upstream Linear # projections under autocast already produce bf16, so this is usually # a no-op; covers the case where upstream is fp32-producing. key_halo, value_halo, query_proj = _cast_to_autocast_dtype(key_halo, value_halo, query_proj) - out, _, _ = _RingNeighborhoodAttentionFn.apply( - key_halo, - value_halo, - query_proj, - self.psi_col_idx_local, - self.psi_roff_idx_local, - self.psi_row_idx_local, - self.quad_weights, - self.nlon_in, - pscale, - self.lon_in_starts, # lon chunk starts for kv (same as lon_in) - self.lon_in_shapes, # lon chunk sizes for kv - lat_halo_start, - self.nlat_out_local, - self.nlon_out_local, - self.r_lat, - azimuth_group(), - self.comm_rank_azimuth, - self.comm_size_azimuth, - self.psi_n_long_rows, - self.psi_max_row_len, - self.psi_mid_row_len, - ) # [Bnh, C_v, H_out_local, W_out_local] + if self.upsample: + # Global pscale_out — the kernel must not infer this from local shapes, + # because kernel `nlon_out` is nlon_out_local which differs when az_size > 1. + pscale_out = self.nlon_out // self.nlon_in + out, _, _ = _RingNeighborhoodAttentionUpsampleFn.apply( + key_halo, + value_halo, + query_proj, + self.psi_col_idx_local, + self.psi_roff_idx_local, + self.quad_weights, + self.nlon_in, + self.nlon_out, + pscale_out, + self.lon_in_starts, # lon chunk starts for kv (same as lon_in) + self.lon_in_shapes, # lon chunk sizes for kv + lat_halo_start, + self.nlat_out_local, + self.nlon_out_local, + self.r_lat, + azimuth_group(), + self.comm_rank_azimuth, + self.comm_size_azimuth, + ) # [Bnh, C_v, H_out_local, W_out_local] + else: + # Global pscale — the kernel must not infer this from local shapes, + # because kernel `nlon_out` is nlon_out_local which differs when az_size > 1. + pscale = self.nlon_in // self.nlon_out + out, _, _ = _RingNeighborhoodAttentionFn.apply( + key_halo, + value_halo, + query_proj, + self.psi_col_idx_local, + self.psi_roff_idx_local, + self.psi_row_idx_local, + self.quad_weights, + self.nlon_in, + pscale, + self.lon_in_starts, # lon chunk starts for kv (same as lon_in) + self.lon_in_shapes, # lon chunk sizes for kv + lat_halo_start, + self.nlat_out_local, + self.nlon_out_local, + self.r_lat, + azimuth_group(), + self.comm_rank_azimuth, + self.comm_size_azimuth, + self.psi_n_long_rows, + self.psi_max_row_len, + self.psi_mid_row_len, + ) # [Bnh, C_v, H_out_local, W_out_local] # unfold num_heads B_nh, C_v, H_out, W_out = out.shape From b1e08fcc6601c4060f020df9ff5653d3d9ae17bf Mon Sep 17 00:00:00 2001 From: Thorsten Kurth Date: Mon, 13 Jul 2026 16:29:05 +0200 Subject: [PATCH 02/14] amending changelog --- Changelog.md | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/Changelog.md b/Changelog.md index 546a79f7..11afab9f 100644 --- a/Changelog.md +++ b/Changelog.md @@ -2,6 +2,11 @@ ## Versioning +### v0.9.2c + +* Added upsampling support to `DistributedNeighborhoodAttentionS2` (`nlon_out % nlon_in == 0`): new upsample (scatter) ring-step CUDA kernels for forward and backward, matching the serial upsample attention. K/V rotate around the azimuth ring while queries and the softmax state stay local; all three directions (self-attention, downsampling, upsampling) are now supported by the distributed layer. +* Consistent CUDA kernel launch error checking: all attention and DISCO CUDA host wrappers now call `C10_CUDA_KERNEL_LAUNCH_CHECK()` and explicitly include `` instead of relying on transitive includes. + ### v0.9.2b * Added `benchmarks/` suite covering SHT, DISCO convolution (self and downsampling), and spherical attention (global, neighborhood self, neighborhood cross-resolution) across resolutions, dtypes, and channel counts. Entry point: `python benchmarks/run.py`. Supports baseline CSV comparison for regression tracking in performance PRs. From 73e178595a75b9e49c10e95e15e238152c71cac5 Mon Sep 17 00:00:00 2001 From: Thorsten Kurth Date: Mon, 13 Jul 2026 16:30:38 +0200 Subject: [PATCH 03/14] bumping version number to 0.9.2c --- pyproject.toml | 2 +- torch_harmonics/__init__.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index aa009ed5..4167b27d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -43,7 +43,7 @@ dependencies = [ [tool.setuptools_scm] -fallback_version = "0.9.2a" +fallback_version = "0.9.2c" [tool.setuptools.packages.find] include = ["torch_harmonics*"] diff --git a/torch_harmonics/__init__.py b/torch_harmonics/__init__.py index 7891a0b3..4199860e 100644 --- a/torch_harmonics/__init__.py +++ b/torch_harmonics/__init__.py @@ -29,7 +29,7 @@ # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. # -__version__ = "0.9.2a" +__version__ = "0.9.2c" from . import examples, quadrature, random_fields from .attention import AttentionS2, NeighborhoodAttentionS2 From 143b23a90fa517fd7dd6a109e1fdc422ffed49d9 Mon Sep 17 00:00:00 2001 From: Thorsten Kurth Date: Tue, 14 Jul 2026 02:57:56 -0700 Subject: [PATCH 04/14] fixed wgma for attention but slow --- Dockerfile | 2 +- Dockerfile.examples | 2 +- setup.py | 1 + .../optimized/attention_interface.cpp | 6 + .../attention_cuda_fwd_wgmma_sm90.cu | 530 ++++++++++++++++++ .../kernels_cuda/attention_cuda_ptx.cuh | 177 ++++++ 6 files changed, 716 insertions(+), 2 deletions(-) create mode 100644 torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_fwd_wgmma_sm90.cu create mode 100644 torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_ptx.cuh diff --git a/Dockerfile b/Dockerfile index 2bc08e6f..938a5992 100644 --- a/Dockerfile +++ b/Dockerfile @@ -38,6 +38,6 @@ RUN pip install parameterized # The custom CUDA extension does not suppport architerctures < 7.0 ENV FORCE_CUDA_EXTENSION=1 ENV TORCH_HARMONICS_ENABLE_OPENMP=1 -ENV TORCH_CUDA_ARCH_LIST="8.0 8.6 8.7 9.0 10.0+PTX" +ENV TORCH_CUDA_ARCH_LIST="8.0 8.6 8.7 9.0a 10.0a+PTX" COPY . /workspace/torch_harmonics RUN cd /workspace/torch_harmonics && NVCC_THREADS=4 MAX_JOBS=8 pip install --no-build-isolation . diff --git a/Dockerfile.examples b/Dockerfile.examples index 9c6cc2b6..fb0df26e 100644 --- a/Dockerfile.examples +++ b/Dockerfile.examples @@ -57,5 +57,5 @@ COPY . /workspace/torch_harmonics # The custom CUDA extension does not suppport architerctures < 7.0 ENV FORCE_CUDA_EXTENSION=1 ENV TORCH_HARMONICS_ENABLE_OPENMP=1 -ENV TORCH_CUDA_ARCH_LIST="8.0 8.6 8.7 9.0 10.0+PTX" +ENV TORCH_CUDA_ARCH_LIST="8.0 8.6 8.7 9.0a 10.0a+PTX" RUN cd /workspace/torch_harmonics && NVCC_THREADS=4 MAX_JOBS=8 pip install --no-deps --no-build-isolation . diff --git a/setup.py b/setup.py index 4a1bcbf2..b94e2805 100644 --- a/setup.py +++ b/setup.py @@ -197,6 +197,7 @@ def get_ext_modules(): "torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_bwd_ring.cu", "torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_fwd_ring_upsample.cu", "torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_bwd_ring_upsample.cu", + "torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_fwd_wgmma_sm90.cu", ] ) ext_modules.append(CUDAExtension("torch_harmonics.attention._C", attention_sources, extra_compile_args=get_compile_args("attention"))) diff --git a/torch_harmonics/attention/optimized/attention_interface.cpp b/torch_harmonics/attention/optimized/attention_interface.cpp index e8d54ab1..3a1e06ed 100644 --- a/torch_harmonics/attention/optimized/attention_interface.cpp +++ b/torch_harmonics/attention/optimized/attention_interface.cpp @@ -83,6 +83,12 @@ namespace attention_kernels m.def("forward(Tensor kx, Tensor vx, Tensor qy, Tensor quad_weights, Tensor col_idx, Tensor row_off, int " "nlon_in, int nlat_out, int nlon_out) -> Tensor", {at::Tag::pt2_compliant_tag}); + // Standalone Hopper WGMMA (sm_90a) forward, for isolated benchmarking of the + // tensor-core neighborhood-attention path. Not routed by the module; call it + // directly via torch.ops.attention_kernels.forward_wgmma. Drop-in signature + // match for `forward` (row_idx is derived internally via sortRows). + m.def("forward_wgmma(Tensor kx, Tensor vx, Tensor qy, Tensor quad_weights, Tensor col_idx, Tensor row_off, " + "int nlon_in, int nlat_out, int nlon_out) -> Tensor"); m.def("backward(Tensor kx, Tensor vx, Tensor qy, Tensor dy, Tensor quad_weights, Tensor col_idx, Tensor " "row_off, int nlon_in, int nlat_out, int nlon_out) -> (Tensor, Tensor, Tensor)", {at::Tag::pt2_compliant_tag}); diff --git a/torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_fwd_wgmma_sm90.cu b/torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_fwd_wgmma_sm90.cu new file mode 100644 index 00000000..f21fde62 --- /dev/null +++ b/torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_fwd_wgmma_sm90.cu @@ -0,0 +1,530 @@ +// coding=utf-8 +// +// SPDX-FileCopyrightText: Copyright (c) 2026 The torch-harmonics Authors. All rights reserved. +// SPDX-License-Identifier: BSD-3-Clause +// +// Redistribution and use in source and binary forms, with or without +// modification, are permitted provided that the following conditions are met: +// +// 1. Redistributions of source code must retain the above copyright notice, this +// list of conditions and the following disclaimer. +// +// 2. Redistributions in binary form must reproduce the above copyright notice, +// this list of conditions and the following disclaimer in the documentation +// and/or other materials provided with the distribution. +// +// 3. Neither the name of the copyright holder nor the names of its +// contributors may be used to endorse or promote products derived from +// this software without specific prior written permission. +// +// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +// DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +// FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +// DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +// SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +// CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +// OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +// ===================================================================================== +// Gather / self attention forward — Hopper (sm_90a) WGMMA, fp16, EXPERIMENTAL v2 +// ===================================================================================== +// +// v2: N-TILED ONLINE SOFTMAX (FlashAttention-shaped). Handles arbitrary neighbor +// counts — crucially the polar rows, whose neighborhood spans the full longitude +// dimension (rlen ~ W, and a few x W across rings) because longitudes converge at +// the poles. The earlier single-pass v1 (rlen <= NHALO) could not. +// +// Built on the VALIDATED WGMMA primitives from the DISCO tensor-core branch +// (attention_cuda_ptx.cuh): descriptor, m64n16k16 fp16 mma (predicate scale_D), +// Major::MN A/B core layouts, and the accumulator -> (m,n) epilogue. +// +// Structure (one warpgroup per (ho-row tile of TM queries, batch)): +// O[TM x Cph] (fp32, shared), running max m[TM], running denom l[TM]. +// for n0 in 0..rlen step NT: +// QK^T : S_tile[TM x NT] = Q . Khalo_tile^T (WGMMA, contract Cph in k=16) +// online: m_new = max(m, rowmax(S_tile)); corr = exp(m - m_new); +// l = l*corr + sum exp(S_tile - m_new)*w; O *= corr; P = exp(S - m_new)*w +// PV : O[:, d] += P_tile . Vhalo_tile[:, d] (WGMMA, channel-tiled DTILE) +// finalize: O /= l, narrow to fp16, store. +// +// PROVEN here: WGMMA mechanics + the online-softmax / N-tiling / channel-tiling +// control flow. STILL "VALIDATE ON HOPPER": the per-query longitude-shift HALO + +// MASK (v2 stages neighbors at the tile-base longitude, exact only for rows whose +// neighbor longitudes are wo-independent; the general union-band halo + per-(m,n) +// mask is the remaining keying item). O resident in shared caps Cph ~ 800 on +// Hopper; larger C wants the 2-pass variant (pass A: stats; pass B: PV). +// +// BUILD: TORCH_CUDA_ARCH_LIST="9.0a+PTX". Gated by TORCH_HARMONICS_ATTN_WGMMA=1. +// ===================================================================================== + +#include "attention_cuda.cuh" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "cudamacro.h" +#include "attention_cuda_utils.cuh" +#include "attention_cuda_ptx.cuh" + +namespace attention_kernels +{ + + static constexpr int WG_THREADS = 128; + static constexpr int TM = 64; // queries per row-tile (WGMMA M) + static constexpr int NT = 16; // neighbor tile width (WGMMA N for QK^T) + static constexpr int DTILE = 16; // channel tile for PV (WGMMA N of the PV GEMM) + static constexpr int KTILE = 16; // WGMMA K + static constexpr int CPH_MAX = 1024; + + // Major::MN core-matrix element indexers (see attention_cuda_ptx.cuh): + // A [M,K] M-fast: (m/8)*128 + (m%8) + 8*k ; B [K,N] N-fast: (n/8)*128 + k*8 + (n%8) + __host__ __device__ __forceinline__ int aIdx(int m, int k) { return (m / 8) * 128 + (m % 8) + 8 * k; } + __host__ __device__ __forceinline__ int bIdx(int k, int n) { return (n / 8) * 128 + k * 8 + (n % 8); } + +#if defined(__CUDA_ARCH_FEAT_SM90_ALL) + + // descriptors are dtype-agnostic (fp16 and bf16 are both 2-byte; the byte strides + // 128/256 are identical), so a void* address is all that's needed. + __device__ __forceinline__ uint64_t descA(const void *p) { return make_wgmma_desc(p, 8 * 16, 16 * 16); } + __device__ __forceinline__ uint64_t descB(const void *p) { return make_wgmma_desc(p, 8 * 16, 16 * 16); } + + // float -> element conversion and the m64n16k16 mma, selected by the storage type. + template __device__ __forceinline__ T f2e(float x); + template <> __device__ __forceinline__ __half f2e<__half>(float x) { return __float2half(x); } + template <> __device__ __forceinline__ __nv_bfloat16 f2e<__nv_bfloat16>(float x) { return __float2bfloat16(x); } + + template __device__ __forceinline__ void wgmma_n16(float (&d)[8], uint64_t da, uint64_t db) + { + if constexpr (std::is_same::value) { + wgmma_m64n16k16_acc_fp16(d, da, db); + } else { + wgmma_m64n16k16_acc_bf16(d, da, db); + } + } + + // m64n16k16 accumulator -> (m,n) writeback (NHALO/NT=16 -> 8 cells/thread). + template + __device__ __forceinline__ void epilogue_n16(const float (&acc)[8], int warp_id, int lane, OP op) + { + const int m01 = warp_id * 16 + (lane >> 2); + const int m23 = m01 + 8; + const int n_a = (lane & 3) * 2; + const int n_b = n_a + 1; + const int ng = 0; // single n-group for N=16 + op(m01, n_a + 8 * ng, acc[ng * 4 + 0]); + op(m01, n_b + 8 * ng, acc[ng * 4 + 1]); + op(m23, n_a + 8 * ng, acc[ng * 4 + 2]); + op(m23, n_b + 8 * ng, acc[ng * 4 + 3]); + const int ng1 = 1; + op(m01, n_a + 8 * ng1, acc[ng1 * 4 + 0]); + op(m01, n_b + 8 * ng1, acc[ng1 * 4 + 1]); + op(m23, n_a + 8 * ng1, acc[ng1 * 4 + 2]); + op(m23, n_b + 8 * ng1, acc[ng1 * 4 + 3]); + } + + // NOTE: no __restrict__ on the parameters — nvcc/cudafe device-stub generation for + // a TEMPLATED __global__ mishandles __restrict__ qualifiers ("template-id ... does + // not match any template declaration"). The non-templated kernels keep __restrict__. + template + __global__ __launch_bounds__(WG_THREADS) void s2_attn_fwd_wgmma_sm90_k(int nchan, int nlat_in, int nlon_in, + int nlat_out, int nlon_out, + int tiles_per_row, const T *kx, const T *vx, + const T *qy, const int32_t *row_idx, + const int64_t *row_off, const int64_t *col_idx, + const float *quad_weights, T *y) + { + const int tid = threadIdx.x; + const int warp_id = tid / 32; + const int lane = tid - warp_id * 32; + + const int batch = blockIdx.y; + const int row = blockIdx.x / tiles_per_row; + const int tile = blockIdx.x - row * tiles_per_row; + const int wo_base = tile * TM; + if (row >= nlat_out || wo_base >= nlon_out) { return; } + + const int ho = row_idx[row]; + const int pscale = nlon_in / nlon_out; + const int tm = min(TM, nlon_out - wo_base); + const int64_t rbeg = row_off[ho]; + const int rlen = static_cast(row_off[ho + 1] - rbeg); + + // ---- shared ---- + extern __shared__ __align__(128) char smem_raw[]; + float *shO = reinterpret_cast(smem_raw); // [TM x nchan] fp32 running output + float *shM = shO + TM * nchan; // [TM] running max + float *shL = shM + TM; // [TM] running denom + float *shStile = shL + TM; // [TM x NT] fp32 scores + T *shQ = reinterpret_cast(shStile + TM * NT); // [TM x KTILE] + T *shK = shQ + TM * KTILE; // [KTILE x NT] + T *shP = shK + KTILE * NT; // [TM x NT] probs (aIdx) + T *shV = shP + TM * NT; // [NT x DTILE] + + for (int i = tid; i < TM * nchan; i += WG_THREADS) { shO[i] = 0.f; } + for (int i = tid; i < TM; i += WG_THREADS) { + shM[i] = -FLT_MAX; + shL[i] = 0.f; + } + __syncthreads(); + + // ---- build the canonical membership mask + latitude band (the HALO). The + // neighborhood spans latitudes [hi_min, hi_max]; per query the longitudes shift + // by pscale*wo but the canonical (wo=0) membership is fixed. We stage the + // full-longitude band halo (absolute input cells, no shift) and mask each query + // by de-shifting a halo column back to canonical and testing membership. ---- + char *member = reinterpret_cast(shV + NT * DTILE); // [nhi <= nlat_in][nlon_in] (after the T tiles) + + __shared__ int sh_himin; + __shared__ int sh_himax; + if (tid == 0) { + sh_himin = nlat_in; + sh_himax = -1; + } + __syncthreads(); + for (int off = tid; off < rlen; off += WG_THREADS) { + const int hi = static_cast(col_idx[rbeg + off] / nlon_in); + // ::atomic* — the global int built-ins; the namespace has a custom + // atomicMax(float*,float) that would otherwise shadow these. + ::atomicMin(&sh_himin, hi); + ::atomicMax(&sh_himax, hi); + } + __syncthreads(); + const int hi_min = sh_himin; + const int nhi = sh_himax - sh_himin + 1; + const int Nhalo = nhi * nlon_in; + + for (int i = tid; i < Nhalo; i += WG_THREADS) { member[i] = 0; } + __syncthreads(); + for (int off = tid; off < rlen; off += WG_THREADS) { + const int64_t col = col_idx[rbeg + off]; + const int hi = static_cast(col / nlon_in); + const int wi = static_cast(col - int64_t(hi) * nlon_in); + member[(hi - hi_min) * nlon_in + wi] = 1; + } + __syncthreads(); + + // halo staging: absolute input cell (hi_min+hi_idx, wi); no per-query shift. + auto stage_halo = [&](const T *src, int nchan_src, int hc, int c) -> T { + const int hi_idx = hc / nlon_in; + const int wi = hc - hi_idx * nlon_in; + const int64_t g = int64_t(batch) * nlat_in * nlon_in * nchan_src + + (int64_t(hi_min + hi_idx) * nlon_in + wi) * nchan_src + c; + return src[g]; + }; + + // ============================ halo N-tile loop (online softmax) =============== + for (int n0 = 0; n0 < Nhalo; n0 += NT) { + const int nt = min(NT, Nhalo - n0); + + // ---- QK^T: S[TM x NT] = Q . Khalo^T (contract nchan, k=16) ---- + float sacc[8]; +#pragma unroll + for (int i = 0; i < 8; i++) { sacc[i] = 0.f; } + for (int k0 = 0; k0 < nchan; k0 += KTILE) { + for (int idx = tid; idx < TM * KTILE; idx += WG_THREADS) { + const int m = idx / KTILE, k = idx - m * KTILE; + T v = f2e(0.f); + if (m < tm) { + const int wo = wo_base + m; + const int64_t g = int64_t(batch) * nlat_out * nlon_out * nchan + + (int64_t(ho) * nlon_out + wo) * nchan + (k0 + k); + v = qy[g]; + } + shQ[aIdx(m, k)] = v; + } + for (int idx = tid; idx < KTILE * NT; idx += WG_THREADS) { + const int k = idx / NT, n = idx - k * NT; + shK[bIdx(k, n)] = (n < nt) ? stage_halo(kx, nchan, n0 + n, k0 + k) : f2e(0.f); + } + // promote the generic-proxy shQ/shK stores to the WGMMA async proxy + // (required; wgmma.fence does not cover shared operands). + fence_proxy_async_shared_cta(); + __syncthreads(); + // synchronous per k-tile: the async mma must finish reading shQ/shK + // (commit+wait) before the next iteration overwrites them. + wgmma_fence(); + wgmma_n16(sacc, descA(shQ), descB(shK)); + wgmma_commit_group(); + wgmma_wait_group<0>(); + __syncthreads(); + } + + for (int i = tid; i < TM * NT; i += WG_THREADS) { shStile[i] = -FLT_MAX; } + __syncthreads(); + epilogue_n16(sacc, warp_id, lane, [&](int m, int n, float vv) { + if (m < tm && n < nt) { shStile[m * NT + n] = vv; } + }); + __syncthreads(); + + // ---- masked online softmax (per query: de-shift halo col -> canonical, + // test membership). The score Q.K at the absolute cell is already correct; + // the mask just selects which halo cells are THIS query's neighbors. ---- + for (int m = tid; m < tm; m += WG_THREADS) { + const int wo = wo_base + m; + const int shift = (pscale * wo) % nlon_in; + float tmax = -FLT_MAX; + for (int n = 0; n < nt; n++) { + const int hc = n0 + n; + const int hi_idx = hc / nlon_in; + int cw = (hc - hi_idx * nlon_in) - shift; + if (cw < 0) cw += nlon_in; + if (member[hi_idx * nlon_in + cw]) { tmax = fmaxf(tmax, shStile[m * NT + n]); } + } + const float m_new = fmaxf(shM[m], tmax); + const float corr = expf(shM[m] - m_new); + float lsum = 0.f; + for (int n = 0; n < nt; n++) { + const int hc = n0 + n; + const int hi_idx = hc / nlon_in; + int cw = (hc - hi_idx * nlon_in) - shift; + if (cw < 0) cw += nlon_in; + float p = 0.f; + if (member[hi_idx * nlon_in + cw]) { + p = expf(shStile[m * NT + n] - m_new) * quad_weights[hi_min + hi_idx]; + } + shP[aIdx(m, n)] = f2e(p); + lsum += p; + } + for (int n = nt; n < NT; n++) { shP[aIdx(m, n)] = f2e(0.f); } + shL[m] = shL[m] * corr + lsum; + for (int c = 0; c < nchan; c++) { shO[m * nchan + c] *= corr; } + shM[m] = m_new; + } + __syncthreads(); + + // ---- PV: O[:, d] += P . Vhalo[:, d] (contract NT halo cells, channel-tiled) ---- + for (int d0 = 0; d0 < nchan; d0 += DTILE) { + for (int idx = tid; idx < NT * DTILE; idx += WG_THREADS) { + const int k = idx / DTILE, n = idx - k * DTILE; + shV[bIdx(k, n)] = (k < nt && (d0 + n) < nchan) ? stage_halo(vx, nchan, n0 + k, d0 + n) : f2e(0.f); + } + // promote the generic-proxy shP (softmax) + shV stores to the WGMMA + // async proxy before the PV mma reads them. + fence_proxy_async_shared_cta(); + __syncthreads(); + float oacc[8]; +#pragma unroll + for (int i = 0; i < 8; i++) { oacc[i] = 0.f; } + wgmma_fence(); + wgmma_n16(oacc, descA(shP), descB(shV)); + wgmma_commit_group(); + wgmma_wait_group<0>(); + epilogue_n16(oacc, warp_id, lane, [&](int m, int n, float vv) { + const int c = d0 + n; + if (m < tm && c < nchan) { shO[m * nchan + c] += vv; } + }); + __syncthreads(); + } + } + + // ---- finalize: O /= l, store ---- + for (int m = tid; m < tm; m += WG_THREADS) { + const float inv = (shL[m] > 0.f) ? 1.f / shL[m] : 0.f; + for (int c = 0; c < nchan; c++) { + const int wo = wo_base + m; + const int64_t g + = int64_t(batch) * nlat_out * nlon_out * nchan + (int64_t(ho) * nlon_out + wo) * nchan + c; + y[g] = f2e(shO[m * nchan + c] * inv); + } + } + } + + // --------------------------------------------------------------------------------- + // DEBUG micro-kernel: D[64x16] = A[64x16] @ B[16x16], one m64n16k16, using the EXACT + // descriptor / aIdx-bIdx staging / epilogue the real kernel uses. A,B,D are plain + // row-major fp16 in global memory. Isolates the WGMMA mechanics from all attention + // logic so they can be unit-tested against a host A@B. One warpgroup, one block. + // --------------------------------------------------------------------------------- + __global__ __launch_bounds__(WG_THREADS) void s2_wgmma_gemm_debug_k(const __half *__restrict__ A, + const __half *__restrict__ B, + __half *__restrict__ D) + { + const int tid = threadIdx.x; + const int warp_id = tid / 32; + const int lane = tid - warp_id * 32; + + // dynamic shared (matches DISCO + the real kernel — GMMA descriptors must + // address dynamic smem the same way; static __shared__ here read as zero). + extern __shared__ __align__(128) char smem_raw[]; + __half *shA = reinterpret_cast<__half *>(smem_raw); // [64 x 16] in aIdx layout + __half *shB = shA + TM * KTILE; // [16 x 16] in bIdx layout + + for (int idx = tid; idx < TM * KTILE; idx += WG_THREADS) { + const int m = idx / KTILE, k = idx - m * KTILE; + shA[aIdx(m, k)] = A[m * KTILE + k]; // A row-major [64,16] + } + for (int idx = tid; idx < KTILE * NT; idx += WG_THREADS) { + const int k = idx / NT, n = idx - k * NT; + shB[bIdx(k, n)] = B[k * NT + n]; // B row-major [16,16] + } + fence_proxy_async_shared_cta(); + __syncthreads(); + + float acc[8]; +#pragma unroll + for (int i = 0; i < 8; i++) { acc[i] = 0.f; } + wgmma_fence(); + wgmma_m64n16k16_acc_fp16(acc, descA(shA), descB(shB)); + wgmma_commit_group(); + wgmma_wait_group<0>(); + + epilogue_n16(acc, warp_id, lane, [&](int m, int n, float v) { D[m * NT + n] = __float2half(v); }); + } + +#else + template + __global__ void s2_attn_fwd_wgmma_sm90_k(int, int, int, int, int, int, const T *, const T *, const T *, + const int32_t *, const int64_t *, const int64_t *, const float *, T *) + { + // matches the real kernel's (unqualified) parameter signature. + } + __global__ void s2_wgmma_gemm_debug_k(const __half *, const __half *, __half *) { } +#endif + + // ----------------------------------------------------------------------------- + // host dispatch — true if launched. Gated: env, sm_90+, fp16, self/gather, + // Cph%16==0, Cph<=CPH_MAX. v2 N-tiles the neighbors so there is NO rlen cap. + // ----------------------------------------------------------------------------- + bool s2_attn_fwd_wgmma_dispatch(int64_t batch_size, int64_t nchans_in, int64_t nchans_out, int64_t nlat_in, + int64_t nlon_in, int64_t nlat_out, int64_t nlon_out, at::Tensor kxP, at::Tensor vxP, + at::Tensor qyP, at::Tensor row_idx, at::Tensor row_off, at::Tensor col_idx, + at::Tensor quad_weights, at::Tensor yP) + { + // diagnostics gated behind TORCH_HARMONICS_WGMMA_VERBOSE so bench runs are quiet. + const char *vv = std::getenv("TORCH_HARMONICS_WGMMA_VERBOSE"); + const bool verbose = (vv != nullptr && std::atoi(vv) != 0); + const auto dt = qyP.scalar_type(); + const bool is_half = (dt == at::kHalf); + const bool is_bf16 = (dt == at::kBFloat16); + if (verbose) { + std::fprintf(stderr, "[torch-harmonics] wgmma_dispatch reached: half=%d bf16=%d nin=%ld nout=%ld major=%d\n", + static_cast(is_half), static_cast(is_bf16), static_cast(nchans_in), + static_cast(nchans_out), at::cuda::getCurrentDeviceProperties()->major); + } + if (!is_half && !is_bf16) { return false; } + if (nchans_in != nchans_out) { return false; } + if ((nchans_in % KTILE) != 0 || nchans_in > CPH_MAX) { return false; } + if (at::cuda::getCurrentDeviceProperties()->major < 9) { return false; } + + if (verbose) { + const int64_t max_rlen = row_off.diff().max().item(); + std::fprintf( + stderr, + "[torch-harmonics] attention WGMMA sm90 kernel LAUNCHED (dtype=%s head_dim=%ld max_row_len=%ld)\n", + is_half ? "fp16" : "bf16", static_cast(nchans_in), static_cast(max_rlen)); + } + + const int nchan = static_cast(nchans_in); + const int tiles_per_row = static_cast(DIV_UP(nlon_out, TM)); + auto stream = at::cuda::getCurrentCUDAStream().stream(); + + const int32_t *_row_idx = reinterpret_cast(row_idx.data_ptr()); + const int64_t *_row_off = reinterpret_cast(row_off.data_ptr()); + const int64_t *_col_idx = reinterpret_cast(col_idx.data_ptr()); + const float *_qw = reinterpret_cast(quad_weights.data_ptr()); + + dim3 block(WG_THREADS); + dim3 grid(nlat_out * tiles_per_row, batch_size); + + // fp16 and bf16 are both 2-byte; the K/V/Q/P tiles are 2 bytes/elem. + membership + // mask (halo): nlat_in*nlon_in bytes (worst-case latitude band). + const size_t sh_bytes = sizeof(float) * (TM * nchan + 2 * TM + TM * NT) + + size_t(2) * (TM * KTILE + KTILE * NT + TM * NT + NT * DTILE) + + static_cast(nlat_in) * static_cast(nlon_in); + // Hopper smem ceiling ~227 KB; fall back if the (nlat_in*nlon_in) membership + // mask makes the tile too large (large grids — the full-lon halo is the v1 + // limit; an arc-band halo would shrink this). + if (sh_bytes > 220000) { return false; } + + // launch the kernel instantiated for the storage type (fp16 or bf16). Written + // out explicitly (not via a generic lambda) — nvcc's device-stub generation + // mishandles launching a templated __global__ from inside an `auto` lambda. + const int nli = static_cast(nlat_in), nlonI = static_cast(nlon_in); + const int nlo = static_cast(nlat_out), nlonO = static_cast(nlon_out); + if (is_half) { + cudaFuncSetAttribute(reinterpret_cast(s2_attn_fwd_wgmma_sm90_k<__half>), + cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast(sh_bytes)); + s2_attn_fwd_wgmma_sm90_k<__half><<>>( + nchan, nli, nlonI, nlo, nlonO, tiles_per_row, reinterpret_cast(kxP.data_ptr()), + reinterpret_cast(vxP.data_ptr()), reinterpret_cast(qyP.data_ptr()), + _row_idx, _row_off, _col_idx, _qw, reinterpret_cast<__half *>(yP.data_ptr())); + } else { + cudaFuncSetAttribute(reinterpret_cast(s2_attn_fwd_wgmma_sm90_k<__nv_bfloat16>), + cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast(sh_bytes)); + s2_attn_fwd_wgmma_sm90_k<__nv_bfloat16><<>>( + nchan, nli, nlonI, nlo, nlonO, tiles_per_row, reinterpret_cast(kxP.data_ptr()), + reinterpret_cast(vxP.data_ptr()), + reinterpret_cast(qyP.data_ptr()), _row_idx, _row_off, _col_idx, _qw, + reinterpret_cast<__nv_bfloat16 *>(yP.data_ptr())); + } + CHECK_ERROR("s2_attn_fwd_wgmma_sm90_k"); + return true; + } + + // ----------------------------------------------------------------------------- + // Standalone forward op for isolated benchmarking of the WGMMA (sm_90a) path. + // Mirrors s2_attention_fwd_cuda's channels-last handling, allocates the output, + // and launches the WGMMA kernel via the dispatch above. Registered as + // attention_kernels::forward_wgmma. Not routed by the module. + // ----------------------------------------------------------------------------- + torch::Tensor s2_attn_fwd_wgmma_op(at::Tensor kx, at::Tensor vx, at::Tensor qy, at::Tensor quad_weights, + at::Tensor col_idx, at::Tensor row_off, int64_t nlon_in, int64_t nlat_out, + int64_t nlon_out) + { + CHECK_CUDA_INPUT_TENSOR(kx); + CHECK_CUDA_INPUT_TENSOR(vx); + CHECK_CUDA_INPUT_TENSOR(qy); + CHECK_CUDA_TENSOR(quad_weights); + CHECK_CUDA_TENSOR(col_idx); + CHECK_CUDA_TENSOR(row_off); + + const int batch_size = kx.size(0); + const int64_t nlat_in = kx.size(2); + const size_t nchans_in = qy.size(1); + const size_t nchans_out = vx.size(1); + auto qy_type = qy.dtype(); + auto stream = at::cuda::getCurrentCUDAStream().stream(); + + // compacted-row -> ho map, sorted by row length for load balance — derived + // internally exactly like the production forward (s2_attention_fwd_cuda). + at::Tensor row_idx = sortRows(static_cast(nlat_out), row_off, stream); + + // The WGMMA kernel indexes activations as channels-last (NHWC) contiguous. + torch::Tensor kxP = kx, vxP = vx, qyP = qy; + const bool qy_is_channels_last = (qyP.strides()[1] == 1); + if (kxP.strides()[1] != 1) { kxP = permute_4D_to0231(kxP); } + if (vxP.strides()[1] != 1) { vxP = permute_4D_to0231(vxP); } + if (!qy_is_channels_last) { qyP = permute_4D_to0231(qyP); } + + const int64_t out_dims[] = {batch_size, nlat_out, nlon_out, static_cast(nchans_out)}; + torch::Tensor yP = torch::empty(out_dims, kxP.options()); + + const bool ok = s2_attn_fwd_wgmma_dispatch(batch_size, nchans_in, nchans_out, nlat_in, nlon_in, nlat_out, + nlon_out, kxP, vxP, qyP, row_idx, row_off, col_idx, quad_weights, yP); + TORCH_CHECK(ok, + "forward_wgmma: WGMMA kernel not launched (needs sm_90a build + fp16/bf16 + " + "nchans_in==nchans_out + nchans%16==0 + nchans<=CPH_MAX + arc-band fits smem)"); + + // PyTorch C10 launch check — surfaces any async launch error from the kernel + // (matches s2_attention_fwd_cuda). CHECK_ERROR inside the dispatch already checks + // cudaGetLastError post-launch; this is the canonical torch-side guard. + C10_CUDA_KERNEL_LAUNCH_CHECK(); + + torch::Tensor y = yP; + if (!qy_is_channels_last) { y = permute_4D_to0312(y); } + return y.to(qy_type); + } + + TORCH_LIBRARY_IMPL(attention_kernels, CUDA, m) { m.impl("forward_wgmma", &s2_attn_fwd_wgmma_op); } + +} // namespace attention_kernels diff --git a/torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_ptx.cuh b/torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_ptx.cuh new file mode 100644 index 00000000..8717fe86 --- /dev/null +++ b/torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_ptx.cuh @@ -0,0 +1,177 @@ +// coding=utf-8 +// +// SPDX-FileCopyrightText: Copyright (c) 2026 The torch-harmonics Authors. All rights reserved. +// SPDX-License-Identifier: BSD-3-Clause +// +// Redistribution and use in source and binary forms, with or without +// modification, are permitted provided that the following conditions are met: +// +// 1. Redistributions of source code must retain the above copyright notice, this +// list of conditions and the following disclaimer. +// +// 2. Redistributions in binary form must reproduce the above copyright notice, +// this list of conditions and the following disclaimer in the documentation +// and/or other materials provided with the distribution. +// +// 3. Neither the name of the copyright holder nor the names of its +// contributors may be used to endorse or promote products derived from +// this software without specific prior written permission. +// +// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +// DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +// FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +// DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +// SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +// CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +// OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +// ===================================================================================== +// PTX inline-asm wrappers for the attention WGMMA path (Hopper sm_90a). +// +// These are extracted VERBATIM (semantics-preserving) from the validated DISCO +// tensor-core implementation (branch tkurth/disco-tc, disco_cuda_ptx.cuh), which +// was debugged on Hopper. Re-derive nothing here — the layout/encoding decisions +// below are the ones that produced correct DISCO results. +// +// Hardware gating: +// cp.async : SM_80+, used only from the Hopper kernel. +// wgmma.* : SM_90a. ptxas REJECTS WGMMA against plain .target sm_90 — it needs +// .target sm_90a. NVCC defines __CUDA_ARCH_FEAT_SM90_ALL only for +// sm_90a, so we gate on that. Build TORCH_CUDA_ARCH_LIST="9.0a+PTX". +// +// Layout conventions (Major::MN, no swizzle): +// A [M=64, K=16] M-fast in the 8x8 core: byte(m,k) = (m/8)*256 + (m%8 + 8*k)*2 +// B [K=16, N] N-fast in the 8x8 core: byte(k,n) = (n/8)*256 + k*16 + (n%8)*2 +// descriptor LBO/SBO are byte strides between 8x8 core matrices in 16-byte units; +// make_wgmma_desc lands `leading` in bits 16-29 and `stride` in bits 32-45. +// A: leading(K-outer)=8 (128B), stride(M-outer)=16 (256B) +// B: leading(K-outer)=8 (128B), stride(N-outer)= 0 (N=8) or 16 (256B, N>=16) +// +// Accumulator fragment (m64nNk16.f32, PTX ISA 9.7.16.5.4): N/2 fp32 cells/thread, +// in n-groups of 4 cells per 8 N-cols. warp w in [0,4), lane l in [0,32): +// cell 4*ng+0: m = w*16 + l/4, n = (l%4)*2 + 8*ng +// cell 4*ng+1: m = w*16 + l/4, n = (l%4)*2 + 1 + 8*ng +// cell 4*ng+2: m = w*16 + l/4 + 8, n = (l%4)*2 + 8*ng +// cell 4*ng+3: m = w*16 + l/4 + 8, n = (l%4)*2 + 1 + 8*ng +// ===================================================================================== + +#pragma once + +#include + +namespace attention_kernels +{ + +#if __CUDA_ARCH__ >= 800 + __device__ __forceinline__ void cp_async_16B(void *smem_dst, const void *gmem_src) + { + unsigned smem_addr = __cvta_generic_to_shared(smem_dst); + asm volatile("cp.async.cg.shared.global [%0], [%1], 16;\n" ::"r"(smem_addr), "l"(gmem_src)); + } + __device__ __forceinline__ void cp_async_commit() { asm volatile("cp.async.commit_group;\n" ::); } + template __device__ __forceinline__ void cp_async_wait_group() + { + asm volatile("cp.async.wait_group %0;\n" ::"n"(N)); + } + __device__ __forceinline__ void cp_async_wait_all() { asm volatile("cp.async.wait_all;\n" ::); } +#endif // __CUDA_ARCH__ >= 800 + +#if defined(__CUDA_ARCH_FEAT_SM90_ALL) + + // 64-bit shared-memory matrix descriptor (PTX ISA 9.7.16.5.1). No swizzle. + __device__ __forceinline__ uint64_t make_wgmma_desc(const void *smem_ptr, uint32_t leading_byte_offset, + uint32_t stride_byte_offset, uint32_t swizzle = 0) + { + unsigned smem_addr = __cvta_generic_to_shared(smem_ptr); + uint64_t desc = 0; + desc |= ((uint64_t)((smem_addr >> 4) & 0x3fffu)) << 0; // start address + desc |= ((uint64_t)((leading_byte_offset >> 4) & 0x3fffu)) << 16; // LBO + desc |= ((uint64_t)((stride_byte_offset >> 4) & 0x3fffu)) << 32; // SBO + desc |= ((uint64_t)(swizzle & 0x3u)) << 62; // swizzle mode + return desc; + } + + __device__ __forceinline__ void wgmma_fence() { asm volatile("wgmma.fence.sync.aligned;\n" ::); } + + // Make shared-memory writes issued through the GENERIC proxy (normal smem stores when + // staging shQ/shK/shP/shV) visible to WGMMA's ASYNC proxy. Producer threads MUST + // execute this after staging and before the __syncthreads that releases the + // WGMMA-consuming warpgroup. wgmma.fence only orders the accumulator registers, NOT + // the shared operands — omitting this is a nondeterministic race (garbage results). + // Verbatim from the debugged DISCO path (disco_cuda_ptx.cuh). + __device__ __forceinline__ void fence_proxy_async_shared_cta() + { + asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); + } + + __device__ __forceinline__ void wgmma_commit_group() { asm volatile("wgmma.commit_group.sync.aligned;\n" ::); } + template __device__ __forceinline__ void wgmma_wait_group() + { + asm volatile("wgmma.wait_group.sync.aligned %0;\n" ::"n"(N)); + } + + // fp16 wrappers (.f32.f16.f16). scale_D is a PREDICATE (not a literal immediate); + // scale_A=scale_B=1, tnspA=tnspB=1 (Major::MN). Accumulating (D += A*B). + __device__ __forceinline__ void wgmma_m64n8k16_acc_fp16(float (&d)[4], uint64_t desc_a, uint64_t desc_b) + { + int32_t scale_D = 1; + asm volatile("{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %6, 0;\n" + "wgmma.mma_async.sync.aligned.m64n8k16.f32.f16.f16 " + "{%0, %1, %2, %3}, %4, %5, p, %7, %8, %9, %10;\n" + "}\n" + : "+f"(d[0]), "+f"(d[1]), "+f"(d[2]), "+f"(d[3]) + : "l"(desc_a), "l"(desc_b), "r"(scale_D), "n"(1), "n"(1), "n"(1), "n"(1)); + } + + __device__ __forceinline__ void wgmma_m64n16k16_acc_fp16(float (&d)[8], uint64_t desc_a, uint64_t desc_b) + { + int32_t scale_D = 1; + asm volatile("{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %10, 0;\n" + "wgmma.mma_async.sync.aligned.m64n16k16.f32.f16.f16 " + "{%0, %1, %2, %3, %4, %5, %6, %7}, %8, %9, p, %11, %12, %13, %14;\n" + "}\n" + : "+f"(d[0]), "+f"(d[1]), "+f"(d[2]), "+f"(d[3]), "+f"(d[4]), "+f"(d[5]), "+f"(d[6]), "+f"(d[7]) + : "l"(desc_a), "l"(desc_b), "r"(scale_D), "n"(1), "n"(1), "n"(1), "n"(1)); + } + + // bf16 wrapper (.f32.bf16.bf16) — identical encoding to fp16 except the input dtype. + // Native on sm_90a (Hopper bf16 tensor cores). scale_D predicate, tnspA=tnspB=1. + __device__ __forceinline__ void wgmma_m64n16k16_acc_bf16(float (&d)[8], uint64_t desc_a, uint64_t desc_b) + { + int32_t scale_D = 1; + asm volatile("{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %10, 0;\n" + "wgmma.mma_async.sync.aligned.m64n16k16.f32.bf16.bf16 " + "{%0, %1, %2, %3, %4, %5, %6, %7}, %8, %9, p, %11, %12, %13, %14;\n" + "}\n" + : "+f"(d[0]), "+f"(d[1]), "+f"(d[2]), "+f"(d[3]), "+f"(d[4]), "+f"(d[5]), "+f"(d[6]), "+f"(d[7]) + : "l"(desc_a), "l"(desc_b), "r"(scale_D), "n"(1), "n"(1), "n"(1), "n"(1)); + } + + __device__ __forceinline__ void wgmma_m64n32k16_acc_fp16(float (&d)[16], uint64_t desc_a, uint64_t desc_b) + { + int32_t scale_D = 1; + asm volatile("{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %18, 0;\n" + "wgmma.mma_async.sync.aligned.m64n32k16.f32.f16.f16 " + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15}, " + "%16, %17, p, %19, %20, %21, %22;\n" + "}\n" + : "+f"(d[0]), "+f"(d[1]), "+f"(d[2]), "+f"(d[3]), "+f"(d[4]), "+f"(d[5]), "+f"(d[6]), "+f"(d[7]), + "+f"(d[8]), "+f"(d[9]), "+f"(d[10]), "+f"(d[11]), "+f"(d[12]), "+f"(d[13]), "+f"(d[14]), + "+f"(d[15]) + : "l"(desc_a), "l"(desc_b), "r"(scale_D), "n"(1), "n"(1), "n"(1), "n"(1)); + } + +#endif // __CUDA_ARCH_FEAT_SM90_ALL + +} // namespace attention_kernels From eba650e20e6e870f278c236bfcd433cc8e16b7f8 Mon Sep 17 00:00:00 2001 From: Thorsten Kurth Date: Tue, 14 Jul 2026 06:47:04 -0700 Subject: [PATCH 05/14] adding temp kernels --- setup.py | 3 + .../optimized/attention_interface.cpp | 11 + .../kernels_cuda/attention_cuda_bwd_fast.cu | 982 ++++++++++++++++++ .../kernels_cuda/attention_cuda_fwd_coop.cu | 321 ++++++ .../kernels_cuda/attention_cuda_fwd_fast.cu | 622 +++++++++++ .../kernels_cuda/attention_cuda_fwd_half.cu | 287 +++++ .../kernels_cuda/attention_cuda_utils.cuh | 18 + 7 files changed, 2244 insertions(+) create mode 100644 torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_bwd_fast.cu create mode 100644 torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_fwd_coop.cu create mode 100644 torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_fwd_fast.cu create mode 100644 torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_fwd_half.cu diff --git a/setup.py b/setup.py index b94e2805..395b911e 100644 --- a/setup.py +++ b/setup.py @@ -198,6 +198,9 @@ def get_ext_modules(): "torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_fwd_ring_upsample.cu", "torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_bwd_ring_upsample.cu", "torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_fwd_wgmma_sm90.cu", + "torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_fwd_fast.cu", + "torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_bwd_fast.cu", + "torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_fwd_coop.cu", ] ) ext_modules.append(CUDAExtension("torch_harmonics.attention._C", attention_sources, extra_compile_args=get_compile_args("attention"))) diff --git a/torch_harmonics/attention/optimized/attention_interface.cpp b/torch_harmonics/attention/optimized/attention_interface.cpp index 3a1e06ed..a3ceebe7 100644 --- a/torch_harmonics/attention/optimized/attention_interface.cpp +++ b/torch_harmonics/attention/optimized/attention_interface.cpp @@ -92,6 +92,17 @@ namespace attention_kernels m.def("backward(Tensor kx, Tensor vx, Tensor qy, Tensor dy, Tensor quad_weights, Tensor col_idx, Tensor " "row_off, int nlon_in, int nlat_out, int nlon_out) -> (Tensor, Tensor, Tensor)", {at::Tag::pt2_compliant_tag}); + // Optimized (FMA + __expf) gather forward/backward copies for A/B benchmarking + // vs the baseline forward/backward. Same signatures; call directly via + // torch.ops.attention_kernels.forward_fast / backward_fast. + m.def("forward_fast(Tensor kx, Tensor vx, Tensor qy, Tensor quad_weights, Tensor col_idx, Tensor row_off, int " + "nlon_in, int nlat_out, int nlon_out) -> Tensor"); + m.def("backward_fast(Tensor kx, Tensor vx, Tensor qy, Tensor dy, Tensor quad_weights, Tensor col_idx, Tensor " + "row_off, int nlon_in, int nlat_out, int nlon_out) -> (Tensor, Tensor, Tensor)"); + // Cross-wo cooperative-tiling forward (shared K/V arc reuse). Gather/self only, + // nchan_in==nchan_out. Same signature as forward; call via forward_coop. + m.def("forward_coop(Tensor kx, Tensor vx, Tensor qy, Tensor quad_weights, Tensor col_idx, Tensor row_off, int " + "nlon_in, int nlat_out, int nlon_out) -> Tensor"); // ---- Ring-step variants for DistributedNeighborhoodAttentionS2 ---- // K/V are sharded along longitude across an azimuth process group; each diff --git a/torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_bwd_fast.cu b/torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_bwd_fast.cu new file mode 100644 index 00000000..4df9569a --- /dev/null +++ b/torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_bwd_fast.cu @@ -0,0 +1,982 @@ +// coding=utf-8 +// +// SPDX-FileCopyrightText: Copyright (c) 2025 The torch-harmonics Authors. All rights reserved. +// SPDX-License-Identifier: BSD-3-Clause +// +// Redistribution and use in source and binary forms, with or without +// modification, are permitted provided that the following conditions are met: +// +// 1. Redistributions of source code must retain the above copyright notice, this +// list of conditions and the following disclaimer. +// +// 2. Redistributions in binary form must reproduce the above copyright notice, +// this list of conditions and the following disclaimer in the documentation +// and/or other materials provided with the distribution. +// +// 3. Neither the name of the copyright holder nor the names of its +// contributors may be used to endorse or promote products derived from +// this software without specific prior written permission. +// +// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +// DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +// FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +// DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +// SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +// CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +// OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +#include "attention_cuda.cuh" +#include +#include +#include "c10/core/MemoryFormat.h" + +#include +#include +#include +#include +#include +#include + +#include +#include +#include + +#include "cudamacro.h" +#include "attention_cuda_utils.cuh" + +#include +#include +#include + +#define THREADS (64) + +#define MAX_LOCAL_ARR_LEN (16) + +namespace attention_kernels +{ + + // scatter-direction dispatcher, defined in attention_cuda_bwd_upsample.cu; + // called by s2_attention_bwd_dkvq_cuda when nlon_out % nlon_in == 0. + void s2_attn_bwd_upsample_dispatch(int batch_size, size_t nchans_in, size_t nchans_out, int64_t nlon_in, + int64_t nlat_in, int64_t nlat_out, int64_t nlon_out, torch::Tensor kxP, + torch::Tensor vxP, torch::Tensor qyP, torch::Tensor dyP, + torch::Tensor psi_row_off, torch::Tensor psi_col_idx, torch::Tensor quad_weights, + torch::Tensor dkxP, torch::Tensor dvxP, torch::Tensor dqyP); + + // ===================================================================================== + // Optimized (FMA + __expf) copy of the GATHER backward, exposed as the ``backward_fast`` + // op for A/B benchmarking against the baseline ``backward``. Nested ``fast`` namespace + // keeps the template kernels' symbols distinct from the baseline (no ODR clash); the + // upsample path delegates to attention_kernels::s2_attn_bwd_upsample_dispatch via + // enclosing-namespace lookup. Kept until the fast path is validated as the default. + // ===================================================================================== + namespace fast + { + + // BEGIN backward kernels and functions + + // called with (blockDim.x=32 and blockDim.y>1, BDIM=blockDim.x*blockDim.y) + // + // STORAGE_T is the global-memory element type of the INPUTS (kx/vx/qy/dy): + // float4 for the fp32 vectorized path, float / c10::Half / c10::BFloat16 for + // the scalar path. COMPUTE_T (float4 or float) is the arithmetic type and the + // type of the gradient OUTPUTS dkx/dvx/dqy — these stay fp32 because dkx/dvx + // are atomically scatter-accumulated across overlapping neighborhoods (fp16 + // atomics would lose precision / aren't well supported). The wrapper allocates + // dkx/dvx/dqy in fp32 and casts to the input dtype at the end. vload widens + // STORAGE_T inputs to COMPUTE_T at the load site. + template + __global__ __launch_bounds__(BDIM_X) void s2_attn_bwd_generic_vec_k( + int nchans_in, // no. of elements along channel dim + int nchans_out, // no. of elements along channel dim + int nlat_in, int nlon_in, int nlat_out, int nlon_out, + const STORAGE_T *__restrict__ kx, // [batch][nlat_in][nlon_in][nchan_in] + const STORAGE_T *__restrict__ vx, // [batch][nlat_in][nlon_in][nchan_out] + const STORAGE_T *__restrict__ qy, // [batch][nlat_out][nlon_out][nchan_in] + const STORAGE_T *__restrict__ dy, // [batch][nlat_out][nlon_out][nchan_out] + const int32_t *__restrict__ row_idx, const int64_t *__restrict__ row_off, + const int64_t *__restrict__ col_idx, const float *__restrict__ quad_weights, + typename vec_traits::compute_t *__restrict__ dkx, // [batch][nlat_in][nlon_in][nchan_in] + typename vec_traits::compute_t *__restrict__ dvx, // [batch][nlat_in][nlon_in][nchan_out] + typename vec_traits::compute_t *__restrict__ dqy) + { // [batch][nlat_out][nlon_out][nchan_in] + using COMPUTE_T = typename vec_traits::compute_t; + + extern __shared__ __align__(sizeof(float4)) float shext[]; + + // for dqy + COMPUTE_T *sh_alpha_k__ = reinterpret_cast(shext) + threadIdx.y * (nchans_in * 4 + nchans_out); + COMPUTE_T *sh_alpha_vw_ = sh_alpha_k__ + nchans_in; + COMPUTE_T *sh_alpha_kvw = sh_alpha_vw_ + nchans_in; + + COMPUTE_T *sh_dy = sh_alpha_kvw + nchans_in; + COMPUTE_T *sh_qy = sh_dy + nchans_out; + // sh_alpha_k__[nchan_in], sh_alpha_vw_[nchan_in], sh_alpha_kvw[nchan_in] + // sh_dy[nchan_out], sh_qy[nchan_in] + + const int batch = blockIdx.y; + + const uint64_t wid = uint64_t(blockIdx.x) * blockDim.y + threadIdx.y; + if (wid >= uint64_t(nlat_out) * nlon_out) { return; } + + const int tidx = threadIdx.x; + + // use permuted rows + const int h = wid / nlon_out; + const int wo = wid - (h * nlon_out); + const int ho = row_idx[h]; + + // one output lon step corresponds to pscale input lon steps (requires nlon_in % nlon_out == 0) + const int pscale = nlon_in / nlon_out; + + // offset input tensors + kx += int64_t(batch) * nlat_in * nlon_in * nchans_in; + qy += int64_t(batch) * nlat_out * nlon_out * nchans_in + int64_t(ho) * nlon_out * nchans_in + + int64_t(wo) * nchans_in; + + vx += int64_t(batch) * nlat_in * nlon_in * nchans_out; + dy += int64_t(batch) * nlat_out * nlon_out * nchans_out + int64_t(ho) * nlon_out * nchans_out + + int64_t(wo) * nchans_out; + + // offset output tensors + dkx += int64_t(batch) * nlat_in * nlon_in * nchans_in; + dvx += int64_t(batch) * nlat_in * nlon_in * nchans_out; + dqy += int64_t(batch) * nlat_out * nlon_out * nchans_in + int64_t(ho) * nlon_out * nchans_in + + int64_t(wo) * nchans_in; + + // zero/init shared memory + for (int chan = tidx; chan < nchans_in; chan += WARP_SIZE) { + sh_alpha_k__[chan] = __vset(0.0f); + sh_alpha_vw_[chan] = __vset(0.0f); + sh_alpha_kvw[chan] = __vset(0.0f); + + sh_qy[chan] = vload(qy, chan); + } + for (int chan = tidx; chan < nchans_out; chan += WARP_SIZE) { sh_dy[chan] = vload(dy, chan); } + +#if __CUDA_ARCH__ < 900 + // for architectures < 9.0, sh_dy and sh_qy will be read + // as individual floats at the end of the kernel, which + // breaks the assumption that each COMPUTE_T location is + // written to and read by the same thread throughout the + // kernel, in the case COMPUTE_T==float4 + if constexpr (std::is_same::value) { __syncwarp(); } +#endif + + // for dkx, dvx, dqy + float alpha_sum = 0.0f; + float qdotk_max = -FLT_MAX; + + // for dkx + float integral = 0.0f; + + const int64_t rbeg = row_off[ho]; + const int64_t rend = row_off[ho + 1]; + + col_idx += rbeg; + + const int rlen = rend - rbeg; + + // accumulate alpha_sum, integral, and shared stats, + // along with a progressively computed qdotk_max. + for (int off = 0; off < rlen; off++) { + + const int64_t col = col_idx[off]; + + const int hi = col / nlon_in; + const int wi = col - (hi * nlon_in); + const int wi_wo = wi + pscale * wo; + const int wip = wi_wo - (wi_wo / nlon_in) * nlon_in; + + const STORAGE_T *_kx = kx + int64_t(hi) * nlon_in * nchans_in + int64_t(wip) * nchans_in; + const STORAGE_T *_vx = vx + int64_t(hi) * nlon_in * nchans_out + int64_t(wip) * nchans_out; + + COMPUTE_T qdotk_v = __vset(0.0f); + COMPUTE_T gdotv_v = __vset(0.0f); + + for (int chan = tidx; chan < nchans_in; chan += WARP_SIZE) { + qdotk_v = __vadd(qdotk_v, __vmul(sh_qy[chan], vload(_kx, chan))); + } + for (int chan = tidx; chan < nchans_out; chan += WARP_SIZE) { + gdotv_v = __vadd(gdotv_v, __vmul(sh_dy[chan], vload(_vx, chan))); + } + + const float qdotk = __warp_sum(__vred(qdotk_v)); + const float gdotv = __warp_sum(__vred(gdotv_v)); + + const float qdotk_max_tmp = max(qdotk_max, qdotk); + const float alpha_inz = __expf(qdotk - qdotk_max_tmp) * quad_weights[hi]; + const float max_correction = __expf(qdotk_max - qdotk_max_tmp); + alpha_sum = alpha_sum * max_correction + alpha_inz; + + integral = integral * max_correction + alpha_inz * gdotv; + + const float ainz_gdotv = alpha_inz * gdotv; + + for (int chan = tidx; chan < nchans_in; chan += WARP_SIZE) { + + const COMPUTE_T kxval = vload(_kx, chan); + + sh_alpha_k__[chan] = __vfma_scale(alpha_inz, kxval, __vscale(max_correction, sh_alpha_k__[chan])); + sh_alpha_vw_[chan] = __vfma_scale(max_correction, sh_alpha_vw_[chan], __vset(ainz_gdotv)); + sh_alpha_kvw[chan] = __vfma_scale(ainz_gdotv, kxval, __vscale(max_correction, sh_alpha_kvw[chan])); + } + qdotk_max = qdotk_max_tmp; + } + + const float alpha_sum_inv = 1.0f / alpha_sum; + + integral *= alpha_sum_inv; + + // Write dqy (fp32 output) + for (int chan = tidx; chan < nchans_in; chan += WARP_SIZE) { + + dqy[chan] = __vscale( + alpha_sum_inv * alpha_sum_inv, + __vsub(__vscale(alpha_sum, sh_alpha_kvw[chan]), __vmul(sh_alpha_vw_[chan], sh_alpha_k__[chan]))); + } + + // accumulate gradients for k and v + for (int off = 0; off < rlen; off++) { + + const int64_t col = col_idx[off]; + const int hi = col / nlon_in; + const int wi = col - (hi * nlon_in); + const int wi_wo = wi + pscale * wo; + const int wip = wi_wo - (wi_wo / nlon_in) * nlon_in; + + const STORAGE_T *_kx = kx + int64_t(hi) * nlon_in * nchans_in + int64_t(wip) * nchans_in; + const STORAGE_T *_vx = vx + int64_t(hi) * nlon_in * nchans_out + int64_t(wip) * nchans_out; + + COMPUTE_T qdotk_v = __vset(0.0f); + COMPUTE_T gdotv_v = __vset(0.0f); + + for (int chan = tidx; chan < nchans_in; chan += WARP_SIZE) { + qdotk_v = __vadd(qdotk_v, __vmul(sh_qy[chan], vload(_kx, chan))); + } + for (int chan = tidx; chan < nchans_out; chan += WARP_SIZE) { + gdotv_v = __vadd(gdotv_v, __vmul(sh_dy[chan], vload(_vx, chan))); + } + + const float qdotk = __warp_sum(__vred(qdotk_v)); + const float gdotv = __warp_sum(__vred(gdotv_v)); + + const float alpha_inz = __expf(qdotk - qdotk_max) * quad_weights[hi]; + + // _dkx / _dvx are COMPUTE_T (fp32) gradient buffers, accumulated atomically. + COMPUTE_T *_dkx = dkx + int64_t(hi) * nlon_in * nchans_in + int64_t(wip) * nchans_in; + COMPUTE_T *_dvx = dvx + int64_t(hi) * nlon_in * nchans_out + int64_t(wip) * nchans_out; + + const float alpha_mul = alpha_inz * alpha_sum_inv; + + const float scale_fact_qy = (gdotv - integral) * alpha_mul; + const float scale_fact_dy = alpha_mul; + + // float4, 128-bit atomics are only supported by devices of compute + // capability 9.x+, so on older devices we resort to 32-bit atomics + +#if __CUDA_ARCH__ < 900 + // to use 32-bit operations on consecutve addresses + float *sh_qy_scl = reinterpret_cast(sh_qy); + float *sh_dy_scl = reinterpret_cast(sh_dy); + + float *_dkx_scl = reinterpret_cast(_dkx); + float *_dvx_scl = reinterpret_cast(_dvx); + + constexpr int VEC_SIZE = sizeof(COMPUTE_T) / sizeof(float); + + // 32-bit, consecutive atomics to glmem; + // strided atomics results in a severe slowdown + for (int chan = tidx; chan < nchans_in * VEC_SIZE; chan += WARP_SIZE) { + atomicAdd(_dkx_scl + chan, scale_fact_qy * sh_qy_scl[chan]); + } + for (int chan = tidx; chan < nchans_out * VEC_SIZE; chan += WARP_SIZE) { + atomicAdd(_dvx_scl + chan, scale_fact_dy * sh_dy_scl[chan]); + } +#else + // 128-bit, consecutive atomics to glmem + for (int chan = tidx; chan < nchans_in; chan += WARP_SIZE) { + atomicAdd(_dkx + chan, __vscale(scale_fact_qy, sh_qy[chan])); + } + for (int chan = tidx; chan < nchans_out; chan += WARP_SIZE) { + atomicAdd(_dvx + chan, __vscale(scale_fact_dy, sh_dy[chan])); + } +#endif + } + + return; + } + + // called with either (BDIM_X=32 and BDIM_Y>1) || (2^K=BDIM_X > 32 and BDIM_Y=1) + template = nchan_in + typename STORAGE_T> + __global__ __launch_bounds__(BDIM_X *BDIM_Y) void s2_attn_bwd_special_vec_k( + int nchan_in, // no. of elements along channel dim + int nchan_out, // no. of elements along channel dim + int nlat_in, int nlon_in, int nlat_out, int nlon_out, + const STORAGE_T *__restrict__ kx, // [batch][nlat_in][nlon_in][nchan_in] + const STORAGE_T *__restrict__ vx, // [batch][nlat_in][nlon_in][nchan_out] + const STORAGE_T *__restrict__ qy, // [batch][nlat_out][nlon_out][nchan_in] + const STORAGE_T *__restrict__ dy, // [batch][nlat_out][nlon_out][nchan_out] + const int32_t *__restrict__ row_idx, const int64_t *__restrict__ row_off, + const int64_t *__restrict__ col_idx, const float *__restrict__ quad_weights, + typename vec_traits::compute_t *__restrict__ dkx, // [batch][nlat_in][nlon_in][nchan_in] + typename vec_traits::compute_t *__restrict__ dvx, // [batch][nlat_in][nlon_in][nchan_out] + typename vec_traits::compute_t *__restrict__ dqy) + { // [batch][nlat_out][nlon_out][nchan_in] + using COMPUTE_T = typename vec_traits::compute_t; + + static_assert(0 == (BDIM_X & (BDIM_X - 1))); + static_assert(0 == (BDIM_Y & (BDIM_Y - 1))); + static_assert((BDIM_X == 32 && BDIM_Y > 1) || (BDIM_X > 32 && BDIM_Y == 1)); + + constexpr int NLOC_M1 = NLOC - 1; + + const int tidx = threadIdx.x; + const int batch = blockIdx.y; + const uint64_t ctaid = uint64_t(blockIdx.x) * blockDim.y + threadIdx.y; + + if (ctaid >= uint64_t(nlat_out) * nlon_out) { return; } + + extern __shared__ __align__(sizeof(float4)) float shext[]; + + // sh_dy[nchan_out], sh_qy[nchan_in] + COMPUTE_T *sh_dy = reinterpret_cast(shext) + threadIdx.y * (nchan_in + nchan_out); // + tidx; + COMPUTE_T *sh_qy = sh_dy + nchan_out + tidx; + + if constexpr (CHOUT_AS_IN) { sh_dy += tidx; } + + // for dqy + COMPUTE_T loc_k__[NLOC]; + COMPUTE_T loc_vw_[NLOC]; + COMPUTE_T loc_kvw[NLOC]; + + // use permuted rows + const int h = ctaid / nlon_out; + const int wo = ctaid - (h * nlon_out); + const int ho = row_idx[h]; + + // one output lon step corresponds to pscale input lon steps (requires nlon_in % nlon_out == 0) + const int pscale = nlon_in / nlon_out; + + // offset input tensors + kx += int64_t(batch) * nlat_in * nlon_in * nchan_in + tidx; + qy += int64_t(batch) * nlat_out * nlon_out * nchan_in + int64_t(ho) * nlon_out * nchan_in + + int64_t(wo) * nchan_in + tidx; + + vx += int64_t(batch) * nlat_in * nlon_in * nchan_out; // + tidx; + dy += int64_t(batch) * nlat_out * nlon_out * nchan_out + int64_t(ho) * nlon_out * nchan_out + + int64_t(wo) * nchan_out; // + tidx; + if constexpr (CHOUT_AS_IN) { + vx += tidx; + dy += tidx; + } + + // offset output tensors + dkx += int64_t(batch) * nlat_in * nlon_in * nchan_in + tidx; + dvx += int64_t(batch) * nlat_in * nlon_in * nchan_out; // + tidx; + if constexpr (CHOUT_AS_IN) { dvx += tidx; } + dqy += int64_t(batch) * nlat_out * nlon_out * nchan_in + int64_t(ho) * nlon_out * nchan_in + + int64_t(wo) * nchan_in + tidx; + +#pragma unroll + for (int i = 0; i < NLOC; i++) { + loc_k__[i] = __vset(0.0f); + loc_vw_[i] = __vset(0.0f); + loc_kvw[i] = __vset(0.0f); + } + +#pragma unroll + for (int i = 0; i < NLOC_M1; i++) { sh_qy[i * BDIM_X] = vload(qy, i * BDIM_X); } + if (NLOC_M1 * BDIM_X + tidx < nchan_in) { sh_qy[NLOC_M1 * BDIM_X] = vload(qy, NLOC_M1 * BDIM_X); } + + if constexpr (CHOUT_AS_IN) { +#pragma unroll + for (int i = 0; i < NLOC_M1; i++) { sh_dy[i * BDIM_X] = vload(dy, i * BDIM_X); } + if (NLOC_M1 * BDIM_X + tidx < nchan_out) { sh_dy[NLOC_M1 * BDIM_X] = vload(dy, NLOC_M1 * BDIM_X); } + } else { + for (int chan = tidx; chan < nchan_out; chan += BDIM_X) { sh_dy[chan] = vload(dy, chan); } + } + +#if __CUDA_ARCH__ < 900 + // for architectures < 9.0, sh_dy and sh_qy will be read + // as individual floats at the end of the kernel, which + // breaks the assumption that each COMPUTE_T location is + // written to and read by the same thread throughout the + // kernel, in the case COMPUTE_T==float4 + if constexpr (std::is_same::value) { + if constexpr (BDIM_X == 32) { + __syncwarp(); + } else { + __syncthreads(); + } + } +#endif + + // for dkx, dvx, dqy + float alpha_sum = 0.0f; + float qdotk_max = -FLT_MAX; + + // for dkx + float integral = 0.0f; + + const int64_t rbeg = row_off[ho]; + const int64_t rend = row_off[ho + 1]; + + col_idx += rbeg; + + const int rlen = rend - rbeg; + + // accumulate alpha_sum, integral, and shared stats, + // along with a progressively computed qdotk_max. + for (int off = 0; off < rlen; off++) { + + const int64_t col = col_idx[off]; + + const int hi = col / nlon_in; + const int wi = col - (hi * nlon_in); + const int wi_wo = wi + pscale * wo; + const int wip = wi_wo - (wi_wo / nlon_in) * nlon_in; + + const STORAGE_T *_kx = kx + int64_t(hi) * nlon_in * nchan_in + int64_t(wip) * nchan_in; + const STORAGE_T *_vx = vx + int64_t(hi) * nlon_in * nchan_out + int64_t(wip) * nchan_out; + + COMPUTE_T qdotk_v = __vset(0.0f); + COMPUTE_T gdotv_v = __vset(0.0f); + +#pragma unroll + for (int i = 0; i < NLOC_M1; i++) { + qdotk_v = __vadd(qdotk_v, __vmul(sh_qy[i * BDIM_X], vload(_kx, i * BDIM_X))); + } + if (NLOC_M1 * BDIM_X + tidx < nchan_in) { + qdotk_v = __vadd(qdotk_v, __vmul(sh_qy[NLOC_M1 * BDIM_X], vload(_kx, NLOC_M1 * BDIM_X))); + } + if constexpr (CHOUT_AS_IN) { +#pragma unroll + for (int i = 0; i < NLOC_M1; i++) { + gdotv_v = __vadd(gdotv_v, __vmul(sh_dy[i * BDIM_X], vload(_vx, i * BDIM_X))); + } + if (NLOC_M1 * BDIM_X + tidx < nchan_out) { + gdotv_v = __vadd(gdotv_v, __vmul(sh_dy[NLOC_M1 * BDIM_X], vload(_vx, NLOC_M1 * BDIM_X))); + } + } else { + for (int chan = tidx; chan < nchan_out; chan += BDIM_X) { + gdotv_v = __vadd(gdotv_v, __vmul(sh_dy[chan], vload(_vx, chan))); + } + } + + float qdotk = __vred(qdotk_v); + float gdotv = __vred(gdotv_v); + + if constexpr (BDIM_X == 32) { + qdotk = __warp_sum(qdotk); + gdotv = __warp_sum(gdotv); + } else { + qdotk = __block_sum(qdotk); + gdotv = __block_sum(gdotv); + } + + const float qdotk_max_tmp = max(qdotk_max, qdotk); + const float alpha_inz = __expf(qdotk - qdotk_max_tmp) * quad_weights[hi]; + const float max_correction = __expf(qdotk_max - qdotk_max_tmp); + + alpha_sum = alpha_sum * max_correction + alpha_inz; + integral = integral * max_correction + alpha_inz * gdotv; + + const float ainz_gdotv = alpha_inz * gdotv; + +#pragma unroll + for (int i = 0; i < NLOC_M1; i++) { + const COMPUTE_T kxval = vload(_kx, i * BDIM_X); + loc_k__[i] = __vfma_scale(alpha_inz, kxval, __vscale(max_correction, loc_k__[i])); + loc_vw_[i] = __vfma_scale(max_correction, loc_vw_[i], __vset(ainz_gdotv)); + loc_kvw[i] = __vfma_scale(ainz_gdotv, kxval, __vscale(max_correction, loc_kvw[i])); + } + if (NLOC_M1 * BDIM_X + tidx < nchan_in) { + const COMPUTE_T kxval = vload(_kx, NLOC_M1 * BDIM_X); + loc_k__[NLOC_M1] = __vfma_scale(alpha_inz, kxval, __vscale(max_correction, loc_k__[NLOC_M1])); + loc_vw_[NLOC_M1] = __vfma_scale(max_correction, loc_vw_[NLOC_M1], __vset(ainz_gdotv)); + loc_kvw[NLOC_M1] = __vfma_scale(ainz_gdotv, kxval, __vscale(max_correction, loc_kvw[NLOC_M1])); + } + + qdotk_max = qdotk_max_tmp; + } + + const float alpha_sum_inv = 1.0f / alpha_sum; + + integral *= alpha_sum_inv; + + // Write dqy + const float alpha_sum_inv_sq = alpha_sum_inv * alpha_sum_inv; + +#pragma unroll + for (int i = 0; i < NLOC_M1; i++) { + dqy[i * BDIM_X] = __vscale(alpha_sum_inv_sq, + __vsub(__vscale(alpha_sum, loc_kvw[i]), __vmul(loc_vw_[i], loc_k__[i]))); + } + if (NLOC_M1 * BDIM_X + tidx < nchan_in) { + dqy[NLOC_M1 * BDIM_X] + = __vscale(alpha_sum_inv_sq, + __vsub(__vscale(alpha_sum, loc_kvw[NLOC_M1]), __vmul(loc_vw_[NLOC_M1], loc_k__[NLOC_M1]))); + } + + // accumulate gradients for k and v + for (int off = 0; off < rlen; off++) { + + const int64_t col = col_idx[off]; + + const int hi = col / nlon_in; + const int wi = col - (hi * nlon_in); + const int wi_wo = wi + pscale * wo; + const int wip = wi_wo - (wi_wo / nlon_in) * nlon_in; + + const STORAGE_T *_kx = kx + int64_t(hi) * nlon_in * nchan_in + int64_t(wip) * nchan_in; + const STORAGE_T *_vx = vx + int64_t(hi) * nlon_in * nchan_out + int64_t(wip) * nchan_out; + + COMPUTE_T qdotk_v = __vset(0.0f); + COMPUTE_T gdotv_v = __vset(0.0f); + +#pragma unroll + for (int i = 0; i < NLOC_M1; i++) { + qdotk_v = __vadd(qdotk_v, __vmul(sh_qy[i * BDIM_X], vload(_kx, i * BDIM_X))); + } + if (NLOC_M1 * BDIM_X + tidx < nchan_in) { + qdotk_v = __vadd(qdotk_v, __vmul(sh_qy[NLOC_M1 * BDIM_X], vload(_kx, NLOC_M1 * BDIM_X))); + } + if constexpr (CHOUT_AS_IN) { +#pragma unroll + for (int i = 0; i < NLOC_M1; i++) { + gdotv_v = __vadd(gdotv_v, __vmul(sh_dy[i * BDIM_X], vload(_vx, i * BDIM_X))); + } + if (NLOC_M1 * BDIM_X + tidx < nchan_out) { + gdotv_v = __vadd(gdotv_v, __vmul(sh_dy[NLOC_M1 * BDIM_X], vload(_vx, NLOC_M1 * BDIM_X))); + } + } else { + for (int chan = tidx; chan < nchan_out; chan += BDIM_X) { + gdotv_v = __vadd(gdotv_v, __vmul(sh_dy[chan], vload(_vx, chan))); + } + } + + float qdotk = __vred(qdotk_v); + float gdotv = __vred(gdotv_v); + + if constexpr (BDIM_X == 32) { + qdotk = __warp_sum(qdotk); + gdotv = __warp_sum(gdotv); + } else { + qdotk = __block_sum(qdotk); + gdotv = __block_sum(gdotv); + } + + const float alpha_inz = __expf(qdotk - qdotk_max) * quad_weights[hi]; + + COMPUTE_T *_dkx = dkx + int64_t(hi) * nlon_in * nchan_in + int64_t(wip) * nchan_in; + COMPUTE_T *_dvx = dvx + int64_t(hi) * nlon_in * nchan_out + int64_t(wip) * nchan_out; + + const float alpha_mul = alpha_inz * alpha_sum_inv; + + const float scale_fact_qy = (gdotv - integral) * alpha_mul; + const float scale_fact_dy = alpha_mul; + + // float4, 128-bit atomics are only supported by devices of compute + // capability 9.x+, so on older devices we resort to 32-bit atomics + +#if __CUDA_ARCH__ < 900 + constexpr int VEC_SIZE = sizeof(COMPUTE_T) / sizeof(float); + + // making the loop count known at compile time doesn't seem + // to make any difference here so let's keep this (much) + // simpler version + float *sh_qy_scl = reinterpret_cast(sh_qy); + float *sh_dy_scl = reinterpret_cast(sh_dy); + + float *_dkx_scl = reinterpret_cast(_dkx); + float *_dvx_scl = reinterpret_cast(_dvx); + + sh_qy_scl -= tidx * VEC_SIZE; + _dkx_scl -= tidx * VEC_SIZE; + if constexpr (CHOUT_AS_IN) { + sh_dy_scl -= tidx * VEC_SIZE; + _dvx_scl -= tidx * VEC_SIZE; + } + + // 32-bit, consecutive atomics to glmem + // strided atomics results in a severe slowdown + for (int chan = tidx; chan < nchan_in * VEC_SIZE; chan += BDIM_X) { + atomicAdd(_dkx_scl + chan, scale_fact_qy * sh_qy_scl[chan]); + } + for (int chan = tidx; chan < nchan_out * VEC_SIZE; chan += BDIM_X) { + atomicAdd(_dvx_scl + chan, scale_fact_dy * sh_dy_scl[chan]); + } +#else +#pragma unroll + for (int i = 0; i < NLOC_M1; i++) { + atomicAdd(_dkx + i * BDIM_X, __vscale(scale_fact_qy, sh_qy[i * BDIM_X])); + } + if (NLOC_M1 * BDIM_X + tidx < nchan_in) { + atomicAdd(_dkx + NLOC_M1 * BDIM_X, __vscale(scale_fact_qy, sh_qy[NLOC_M1 * BDIM_X])); + } + if constexpr (CHOUT_AS_IN) { +#pragma unroll + for (int i = 0; i < NLOC_M1; i++) { + atomicAdd(_dvx + i * BDIM_X, __vscale(scale_fact_dy, sh_dy[i * BDIM_X])); + } + if (NLOC_M1 * BDIM_X + tidx < nchan_out) { + atomicAdd(_dvx + NLOC_M1 * BDIM_X, __vscale(scale_fact_dy, sh_dy[NLOC_M1 * BDIM_X])); + } + } else { + for (int chan = tidx; chan < nchan_out; chan += BDIM_X) { + atomicAdd(_dvx + chan, __vscale(scale_fact_dy, sh_dy[chan])); + } + } +#endif + } + + return; + } + + template + void launch_gen_attn_bwd(int batch_size, int nchans_in, int nchans_out, int nlat_in, int nlon_in, int nlat_out, + int nlon_out, STORAGE_T *_kxp, STORAGE_T *_vxp, STORAGE_T *_qyp, STORAGE_T *_dyp, + int32_t *_row_idx, int64_t *_row_off, int64_t *_col_idx, float *_quad_weights, + typename vec_traits::compute_t *_dkxp, + typename vec_traits::compute_t *_dvxp, + typename vec_traits::compute_t *_dqyp, cudaStream_t stream) + { + + dim3 block(WARP_SIZE, THREADS / WARP_SIZE); + dim3 grid(DIV_UP(nlat_out * nlon_out, block.y), batch_size); + + // shared memory holds compute-type (COMPUTE_T) data, not STORAGE_T. 5 arrays per warp. + size_t shsize = sizeof(typename vec_traits::compute_t) * (nchans_in * 4 + nchans_out) * block.y; + + s2_attn_bwd_generic_vec_k<<>>( + nchans_in, nchans_out, nlat_in, nlon_in, nlat_out, nlon_out, _kxp, _vxp, _qyp, _dyp, _row_idx, _row_off, + _col_idx, _quad_weights, _dkxp, _dvxp, _dqyp); + CHECK_ERROR("s2_attn_bwd_generic_vec_k"); + + return; + } + + template + void launch_spc_attn_bwd(int nloc, // "BDIM_X*nloc" >= nchans_out + int batch_size, int nchans_in, int nchans_out, int nlat_in, int nlon_in, int nlat_out, + int nlon_out, STORAGE_T *_kxp, STORAGE_T *_vxp, STORAGE_T *_qyp, STORAGE_T *_dyp, + int32_t *_row_idx, int64_t *_row_off, int64_t *_col_idx, float *_quad_weights, + typename vec_traits::compute_t *_dkxp, + typename vec_traits::compute_t *_dvxp, + typename vec_traits::compute_t *_dqyp, cudaStream_t stream) + { + + if (CUR_LOC_SIZE == nloc) { + + dim3 block(BDIM_X, BDIM_Y); + dim3 grid(DIV_UP(nlat_out * nlon_out, block.y), batch_size); + + // shared memory holds compute-type (COMPUTE_T) data, not STORAGE_T. + // 2 arrays per cta, block.y > 1 iif block.x==32 + size_t shsize = sizeof(typename vec_traits::compute_t) * (nchans_in + nchans_out) * block.y; + + // nloc determines the size of local arrays used to store + // temporary buffers loc_k__[], loc_vw_[] and loc_kvw[], + // of size nchans_in each; + // if nchans_out is >= BDIM_X*(nloc-1) and <= BDIM_X*nloc + // then we can use the same compile-time known loops used + // for input channels, with the exception of testing + // whether to execute the last iteration based on "nchans_out" + // instead of "nchans_in"; in this way as long as the + // difference between the number of input and output channels + // is <= BDIM_X we can use the faster path + if (nchans_out >= BDIM_X * (CUR_LOC_SIZE - 1) && nchans_out <= BDIM_X * CUR_LOC_SIZE) { + s2_attn_bwd_special_vec_k<<>>( + nchans_in, nchans_out, nlat_in, nlon_in, nlat_out, nlon_out, _kxp, _vxp, _qyp, _dyp, _row_idx, + _row_off, _col_idx, _quad_weights, _dkxp, _dvxp, _dqyp); + } else { + s2_attn_bwd_special_vec_k<<>>( + nchans_in, nchans_out, nlat_in, nlon_in, nlat_out, nlon_out, _kxp, _vxp, _qyp, _dyp, _row_idx, + _row_off, _col_idx, _quad_weights, _dkxp, _dvxp, _dqyp); + } + CHECK_ERROR("s2_attn_bwd_special_vec_k"); + + return; + } + if constexpr (CUR_LOC_SIZE < MAX_LOC_SIZE) { + launch_spc_attn_bwd( + nloc, batch_size, nchans_in, nchans_out, nlat_in, nlon_in, nlat_out, nlon_out, _kxp, _vxp, _qyp, + _dyp, _row_idx, _row_off, _col_idx, _quad_weights, _dkxp, _dvxp, _dqyp, stream); + } + return; + } + + // Picks the block size (BDIM_X) instance and launches the backward gather + // kernel for a given input storage vector type SV. Inputs are SV*; gradient + // outputs are COMPUTE_T* (fp32) — see s2_attn_bwd_*_vec_k. Backward uses the + // special kernel only up to BDIM_X=512 (1024 spills); larger falls to generic. + template + static void bwd_dispatch_bdimx(int bdimx, int nloc, int64_t batch_size, int64_t nchans_in, int64_t nchans_out, + int nlat_in, int64_t nlon_in, int64_t nlat_out, int64_t nlon_out, SV *_kxp, + SV *_vxp, SV *_qyp, SV *_dyp, int32_t *_row_idx, int64_t *_row_off, + int64_t *_col_idx, float *_quad_weights, typename vec_traits::compute_t *_dkxp, + typename vec_traits::compute_t *_dvxp, + typename vec_traits::compute_t *_dqyp, cudaStream_t stream) + { + switch (bdimx) { + case 32: + launch_spc_attn_bwd<32, 2, 1, MAX_LOC>(nloc, batch_size, nchans_in, nchans_out, nlat_in, nlon_in, + nlat_out, nlon_out, _kxp, _vxp, _qyp, _dyp, _row_idx, _row_off, + _col_idx, _quad_weights, _dkxp, _dvxp, _dqyp, stream); + break; + case 64: + launch_spc_attn_bwd<64, 1, MIN_LOC, MAX_LOC>( + nloc, batch_size, nchans_in, nchans_out, nlat_in, nlon_in, nlat_out, nlon_out, _kxp, _vxp, _qyp, + _dyp, _row_idx, _row_off, _col_idx, _quad_weights, _dkxp, _dvxp, _dqyp, stream); + break; + case 128: + launch_spc_attn_bwd<128, 1, MIN_LOC, MAX_LOC>( + nloc, batch_size, nchans_in, nchans_out, nlat_in, nlon_in, nlat_out, nlon_out, _kxp, _vxp, _qyp, + _dyp, _row_idx, _row_off, _col_idx, _quad_weights, _dkxp, _dvxp, _dqyp, stream); + break; + case 256: + launch_spc_attn_bwd<256, 1, MIN_LOC, MAX_LOC>( + nloc, batch_size, nchans_in, nchans_out, nlat_in, nlon_in, nlat_out, nlon_out, _kxp, _vxp, _qyp, + _dyp, _row_idx, _row_off, _col_idx, _quad_weights, _dkxp, _dvxp, _dqyp, stream); + break; + case 512: + launch_spc_attn_bwd<512, 1, MIN_LOC, MAX_LOC>( + nloc, batch_size, nchans_in, nchans_out, nlat_in, nlon_in, nlat_out, nlon_out, _kxp, _vxp, _qyp, + _dyp, _row_idx, _row_off, _col_idx, _quad_weights, _dkxp, _dvxp, _dqyp, stream); + break; + default: + launch_gen_attn_bwd(batch_size, nchans_in, nchans_out, nlat_in, nlon_in, nlat_out, nlon_out, _kxp, _vxp, + _qyp, _dyp, _row_idx, _row_off, _col_idx, _quad_weights, _dkxp, _dvxp, _dqyp, stream); + break; + } + } + + // Templated on the input storage element type (float / c10::Half / c10::BFloat16). + // Gradient outputs dkx/dvx/dqy are always fp32 (COMPUTE_T) — dkx/dvx are atomic + // scatter-accumulated, so they cannot be reduced precision; the wrapper casts + // them to the input dtype at the end. fp32 keeps the float4 vectorized path; + // fp16/bf16 (and unaligned fp32) take the scalar STORAGE_T path. + template + static void s2_attn_bwd_dispatch(int64_t batch_size, int64_t nchans_in, int64_t nchans_out, int64_t nlon_in, + int64_t nlat_out, int64_t nlon_out, at::Tensor kxP, at::Tensor vxP, + at::Tensor qyP, at::Tensor dyP, at::Tensor row_off, at::Tensor col_idx, + at::Tensor quad_weights, at::Tensor dkxP, at::Tensor dvxP, at::Tensor dqyP) + { + + static_assert(0 == (MAX_LOCAL_ARR_LEN & (MAX_LOCAL_ARR_LEN - 1))); + + // get stream + auto stream = at::cuda::getCurrentCUDAStream().stream(); + + // sort row indices (ho-s) in descending order + // based on (row_off[ho+1]-row_off[ho]) + at::Tensor row_idx = sortRows(nlat_out, row_off, stream); + + const int nlat_in = kxP.size(1); + + // smallest power of two "bdimx" (>=32) s.t. bdimx*MAX_LOCAL_ARR_LEN >= nchans_in + int bdimx; + bdimx = DIV_UP(nchans_in, MAX_LOCAL_ARR_LEN); + bdimx = max(bdimx, WARP_SIZE); + bdimx = next_pow2(bdimx); + + scalar_t *_kxp = reinterpret_cast(kxP.data_ptr()); + scalar_t *_vxp = reinterpret_cast(vxP.data_ptr()); + scalar_t *_qyp = reinterpret_cast(qyP.data_ptr()); + scalar_t *_dyp = reinterpret_cast(dyP.data_ptr()); + + // gradient outputs are fp32 + float *_dkxp = reinterpret_cast(dkxP.data_ptr()); + float *_dvxp = reinterpret_cast(dvxP.data_ptr()); + float *_dqyp = reinterpret_cast(dqyP.data_ptr()); + + int32_t *_row_idx = reinterpret_cast(row_idx.data_ptr()); + int64_t *_row_off = reinterpret_cast(row_off.data_ptr()); + int64_t *_col_idx = reinterpret_cast(col_idx.data_ptr()); + float *_quad_weights = reinterpret_cast(quad_weights.data_ptr()); + + constexpr int MIN_LOC_ARR_LEN = MAX_LOCAL_ARR_LEN / 2 + 1; + + if constexpr (std::is_same::value) { + // fp32: float4 vectorized when 16B-aligned + 4-divisible, else scalar. + constexpr int VEC_SIZE = sizeof(float4) / sizeof(float); // 4 + const bool use_vec = is_aligned<16>(_kxp) && is_aligned<16>(_vxp) && is_aligned<16>(_qyp) + && is_aligned<16>(_dyp) && is_aligned<16>(_dkxp) && is_aligned<16>(_dvxp) && is_aligned<16>(_dqyp) + && (nchans_in % VEC_SIZE) == 0 && (nchans_out % VEC_SIZE) == 0; + + if (use_vec) { + constexpr int MAX_VEC = MAX_LOCAL_ARR_LEN / VEC_SIZE; + constexpr int MIN_VEC = MAX_VEC / 2 + 1; + const int64_t nci = nchans_in / VEC_SIZE; + const int64_t nco = nchans_out / VEC_SIZE; + bwd_dispatch_bdimx( + bdimx, DIV_UP(nci, bdimx), batch_size, nci, nco, nlat_in, nlon_in, nlat_out, nlon_out, + reinterpret_cast(_kxp), reinterpret_cast(_vxp), + reinterpret_cast(_qyp), reinterpret_cast(_dyp), _row_idx, _row_off, + _col_idx, _quad_weights, reinterpret_cast(_dkxp), reinterpret_cast(_dvxp), + reinterpret_cast(_dqyp), stream); + } else { + bwd_dispatch_bdimx( + bdimx, DIV_UP(nchans_in, bdimx), batch_size, nchans_in, nchans_out, nlat_in, nlon_in, nlat_out, + nlon_out, _kxp, _vxp, _qyp, _dyp, _row_idx, _row_off, _col_idx, _quad_weights, _dkxp, _dvxp, + _dqyp, stream); + } + } else { + // fp16/bf16: scalar STORAGE_T inputs, fp32 outputs; fp32 compute/accumulation. + bwd_dispatch_bdimx( + bdimx, DIV_UP(nchans_in, bdimx), batch_size, nchans_in, nchans_out, nlat_in, nlon_in, nlat_out, + nlon_out, _kxp, _vxp, _qyp, _dyp, _row_idx, _row_off, _col_idx, _quad_weights, _dkxp, _dvxp, _dqyp, + stream); + } + + return; + } + + // END backward kernels and functions + + std::tuple + s2_attention_bwd_dkvq_cuda(at::Tensor kx, at::Tensor vx, at::Tensor qy, at::Tensor dy, at::Tensor quad_weights, + at::Tensor psi_col_idx, at::Tensor psi_row_off, int64_t nlon_in, int64_t nlat_out, + int64_t nlon_out) + { + + CHECK_CUDA_INPUT_TENSOR(kx); + CHECK_CUDA_INPUT_TENSOR(vx); + CHECK_CUDA_INPUT_TENSOR(qy); + CHECK_CUDA_INPUT_TENSOR(dy); + CHECK_CUDA_TENSOR(quad_weights); + CHECK_CUDA_TENSOR(psi_col_idx); + CHECK_CUDA_TENSOR(psi_row_off); + + // direction selection: gather (self / downsample) iff nlon_in is an integer + // multiple of nlon_out; scatter (upsample) iff nlon_out is an integer multiple + // of nlon_in. Self-attention satisfies both and routes through the gather path. + const bool downsample = (nlon_in % nlon_out == 0); + const bool upsample = (nlon_out % nlon_in == 0); + TORCH_CHECK(downsample || upsample, "either nlon_in (", nlon_in, + ") must be an integer multiple of nlon_out (", nlon_out, "), or vice versa"); + + // const size_t uo_num_channels = kx.size(1); + size_t nchans_in = qy.size(1); // or kx.size(1) + size_t nchans_out = vx.size(1); + + const int batch_size = kx.size(0); + const int64_t nlat_in = kx.size(2); + + // extract dtype + auto kx_type = kx.dtype(); // nchans_in + auto qy_type = qy.dtype(); + auto vx_type = vx.dtype(); // ncahn_out + auto dy_type = dy.dtype(); + + torch::Tensor dkx, dvx, dqy; + + // ATen dispatch over the input dtype. + // + // Gather (downsample / self) path: native storage. kx/vx/qy/dy keep their + // dtype; the kernel widens to fp32 at load (Tier B). Gradient buffers + // dkx/dvx/dqy are allocated fp32 because dkx/dvx are atomically + // scatter-accumulated (reduced-precision atomics would lose precision); + // they are cast back to the input dtype at the end. + // + // Scatter (upsample) path: native storage as well — inputs stay in their + // dtype (widened at load inside the dispatch), gradients stay fp32. + AT_DISPATCH_FLOATING_TYPES_AND2(at::kHalf, at::kBFloat16, qy.scalar_type(), "s2_attention_bwd_dkvq_cuda", [&] { + using storage_t = scalar_t; + + const auto f32_like + = [](const torch::Tensor &t) { return torch::zeros_like(t, t.options().dtype(torch::kFloat32)); }; + + if (downsample) { + // native-storage gather path + torch::Tensor kxP = kx; + torch::Tensor vxP = vx; + torch::Tensor qyP = qy; + torch::Tensor dyP = dy; + + // safer than is_contiguous(ChannelsLast), which fails for num_channels == 1 + bool kx_is_channels_last = kxP.strides()[1] == 1; + bool vx_is_channels_last = vxP.strides()[1] == 1; + bool qy_is_channels_last = qyP.strides()[1] == 1; + bool dy_is_channels_last = dyP.strides()[1] == 1; + + if (!kx_is_channels_last) { kxP = permute_4D_to0231(kxP); } + if (!vx_is_channels_last) { vxP = permute_4D_to0231(vxP); } + if (!qy_is_channels_last) { qyP = permute_4D_to0231(qyP); } + if (!dy_is_channels_last) { dyP = permute_4D_to0231(dyP); } + + // fp32 gradient buffers (same shape/layout as the native inputs) + torch::Tensor dkxP = f32_like(kxP); + torch::Tensor dvxP = f32_like(vxP); + torch::Tensor dqyP = f32_like(qyP); + + s2_attn_bwd_dispatch(batch_size, nchans_in, nchans_out, nlon_in, nlat_out, nlon_out, kxP, + vxP, qyP, dyP, psi_row_off, psi_col_idx, quad_weights, dkxP, dvxP, + dqyP); + + dkx = dkxP; + dvx = dvxP; + dqy = dqyP; + + if (!kx_is_channels_last) { dkx = permute_4D_to0312(dkx); } + if (!vx_is_channels_last) { dvx = permute_4D_to0312(dvx); } + if (!qy_is_channels_last) { dqy = permute_4D_to0312(dqy); } + } else { + // native-storage scatter (upsample) path + torch::Tensor kxP = kx; + torch::Tensor vxP = vx; + torch::Tensor qyP = qy; + torch::Tensor dyP = dy; + + bool kx_is_channels_last = kxP.strides()[1] == 1; + bool vx_is_channels_last = vxP.strides()[1] == 1; + bool qy_is_channels_last = qyP.strides()[1] == 1; + bool dy_is_channels_last = dyP.strides()[1] == 1; + + if (!kx_is_channels_last) { kxP = permute_4D_to0231(kxP); } + if (!vx_is_channels_last) { vxP = permute_4D_to0231(vxP); } + if (!qy_is_channels_last) { qyP = permute_4D_to0231(qyP); } + if (!dy_is_channels_last) { dyP = permute_4D_to0231(dyP); } + + torch::Tensor dkxP = f32_like(kxP); + torch::Tensor dvxP = f32_like(vxP); + torch::Tensor dqyP = f32_like(qyP); + + s2_attn_bwd_upsample_dispatch(batch_size, nchans_in, nchans_out, nlon_in, nlat_in, nlat_out, + nlon_out, kxP, vxP, qyP, dyP, psi_row_off, psi_col_idx, quad_weights, + dkxP, dvxP, dqyP); + + dkx = dkxP; + dvx = dvxP; + dqy = dqyP; + + if (!kx_is_channels_last) { dkx = permute_4D_to0312(dkx); } + if (!vx_is_channels_last) { dvx = permute_4D_to0312(dvx); } + if (!qy_is_channels_last) { dqy = permute_4D_to0312(dqy); } + } + }); + + C10_CUDA_KERNEL_LAUNCH_CHECK(); + + // convert precision back to starting dtype (no-op for fp32; narrows for fp16/bf16) + dkx = dkx.to(kx_type); + dvx = dvx.to(vx_type); + dqy = dqy.to(qy_type); + + return std::make_tuple(dkx, dvx, dqy); + } + + TORCH_LIBRARY_IMPL(attention_kernels, CUDA, m) { m.impl("backward_fast", &s2_attention_bwd_dkvq_cuda); } + + } // namespace fast + +} // namespace attention_kernels diff --git a/torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_fwd_coop.cu b/torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_fwd_coop.cu new file mode 100644 index 00000000..5f28c1cf --- /dev/null +++ b/torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_fwd_coop.cu @@ -0,0 +1,321 @@ +// coding=utf-8 +// +// SPDX-FileCopyrightText: Copyright (c) 2026 The torch-harmonics Authors. All rights reserved. +// SPDX-License-Identifier: BSD-3-Clause +// +// Redistribution and use in source and binary forms, with or without +// modification, are permitted provided that the following conditions are met: +// +// 1. Redistributions of source code must retain the above copyright notice, this +// list of conditions and the following disclaimer. +// +// 2. Redistributions in binary form must reproduce the above copyright notice, +// this list of conditions and the following disclaimer in the documentation +// and/or other materials provided with the distribution. +// +// 3. Neither the name of the copyright holder nor the names of its +// contributors may be used to endorse or promote products derived from +// this software without specific prior written permission. +// +// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +// DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +// FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +// DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +// SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +// CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +// OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +// ===================================================================================== +// Gather / self attention forward — cross-wo COOPERATIVE TILING (sm_80+, no WGMMA/TMA). +// +// Motivation (H100 ncu, 180x360 C=256 tc=0.03): the baseline s2_attn_fwd_special_vec_k +// is latency-bound on the col_idx-driven K/V gather (long_scoreboard #1, L1 hit ~52%). +// Adjacent output wo attend nearly the SAME input K/V window (shifted by pscale*wo), so +// each (ho,wo) warp re-gathers redundantly. This kernel tiles WO_TILE consecutive wo of +// one output row into a block; the block stages the UNION K/V longitude arc once into +// shared and every wo reuses it -> ~WO_TILE fewer global loads on the polar rows that +// dominate the work. +// +// Occupancy-conscious (Mauro's TMA attempt tanked occupancy -> much worse): accumulators +// (locy) + query (locq) stay in REGISTERS; only the K/V arc is shared, streamed in +// CHUNK-wide longitude tiles so the shared footprint is small. +// +// Geometry: col_idx is sorted by input lat hi; each (ho,hi) neighbor set is a contiguous +// signed longitude arc [-r, +r] centered on the output longitude (a geodesic disc, no +// holes). For output wo the arc shifts by pscale*wo. Per hi we reduce r over the row's +// neighbors, stage the tile-union arc, and each wo tests membership by de-shifting a +// staged cell back to canonical signed wi and checking |swi| <= r. +// +// Uses accurate expf (NOT __expf) so an A/B vs baseline isolates the TILING effect. +// Assumes nchan_in == nchan_out (self / square-channel gather). fp16/bf16/fp32 storage, +// fp32 compute. Registered as the ``forward_coop`` op. +// ===================================================================================== + +#include "attention_cuda.cuh" +#include +#include +#include +#include + +#include +#include +#include + +#include "cudamacro.h" +#include "attention_cuda_utils.cuh" + +namespace attention_kernels +{ + + // compacted-row -> ho map, defined in attention_cuda_utils.cu. + // (declared in attention_cuda_utils.cuh: sortRows) + + namespace coop + { + + static constexpr int WARP = 32; + static constexpr int WO_TILE = 4; // output wo positions (warps) per block + static constexpr int CHUNK = 16; // input longitudes staged per shared tile + + __device__ __forceinline__ int cmod(int a, int n) { return ((a % n) + n) % n; } + // map a longitude in [0,n) to signed [-n/2, n/2) so a pole-centred window is contiguous + __device__ __forceinline__ int csigned(int wi, int n) { return (wi >= (n >> 1)) ? (wi - n) : wi; } + + // block-wide max over ints (blockDim.x = WO_TILE*WARP, WO_TILE <= 32). + __device__ __forceinline__ int block_max_int(int v, int *sh) + { + const int lane = threadIdx.x & (WARP - 1); + const int wid = threadIdx.x >> 5; +#pragma unroll + for (int o = 16; o > 0; o >>= 1) { v = max(v, __shfl_down_sync(0xffffffffu, v, o)); } + if (lane == 0) { sh[wid] = v; } + __syncthreads(); + v = (threadIdx.x < (blockDim.x >> 5)) ? sh[lane] : INT_MIN; + if (wid == 0) { +#pragma unroll + for (int o = 16; o > 0; o >>= 1) { v = max(v, __shfl_down_sync(0xffffffffu, v, o)); } + if (lane == 0) { sh[0] = v; } + } + __syncthreads(); + return sh[0]; + } + + // NLOC = ceil(nchan / WARP): per-lane register slots for query / output accumulator. + template + __global__ __launch_bounds__(WO_TILE *WARP) void s2_attn_fwd_coop_k( + int nchan, int nlat_in, int nlon_in, int nlat_out, int nlon_out, int tiles_per_row, + const STORAGE_T *__restrict__ kx, const STORAGE_T *__restrict__ vx, const STORAGE_T *__restrict__ qy, + const int32_t *__restrict__ row_idx, const int64_t *__restrict__ row_off, + const int64_t *__restrict__ col_idx, const float *__restrict__ quad_weights, STORAGE_T *__restrict__ y) + { + using COMPUTE_T = typename vec_traits::compute_t; // float for the scalar path + + const int lane = threadIdx.x & (WARP - 1); + const int warp = threadIdx.x >> 5; + + const int batch = blockIdx.y; + const int brow = blockIdx.x / tiles_per_row; + const int wtile = blockIdx.x - brow * tiles_per_row; + if (brow >= nlat_out) { return; } + + const int ho = row_idx[brow]; + const int pscale = nlon_in / nlon_out; + const int wo = wtile * WO_TILE + warp; + const bool valid = (wo < nlon_out); + + extern __shared__ __align__(16) char smem[]; + STORAGE_T *shK = reinterpret_cast(smem); // [CHUNK * nchan] + STORAGE_T *shV = shK + size_t(CHUNK) * nchan; // [CHUNK * nchan] + int *shred = reinterpret_cast(shV + size_t(CHUNK) * nchan); // [WO_TILE] reduction scratch + + // query for this wo, resident in registers (widened once). + COMPUTE_T locq[NLOC]; + COMPUTE_T locy[NLOC]; +#pragma unroll + for (int i = 0; i < NLOC; i++) { locy[i] = 0.f; } + if (valid) { + const STORAGE_T *_qy = qy + (int64_t(batch) * nlat_out * nlon_out + int64_t(ho) * nlon_out + wo) * nchan; +#pragma unroll + for (int i = 0; i < NLOC; i++) { + const int ch = lane + i * WARP; + locq[i] = (ch < nchan) ? vload(_qy, ch) : COMPUTE_T(0.f); + } + } + + float qmax = -FLT_MAX; + float asum = 0.f; + + const int64_t rbeg = row_off[ho]; + const int64_t rend = row_off[ho + 1]; + const int rlen = int(rend - rbeg); + if (rlen <= 0) { return; } + + const int hi_min = int(col_idx[rbeg] / nlon_in); + const int hi_max = int(col_idx[rend - 1] / nlon_in); + + for (int hi = hi_min; hi <= hi_max; hi++) { + + // radius r for this hi = max |signed wi| over the row's neighbors at this hi. + int loc = -1; + for (int off = threadIdx.x; off < rlen; off += blockDim.x) { + const int64_t c = col_idx[rbeg + off]; + if (int(c / nlon_in) == hi) { loc = max(loc, abs(csigned(int(c % nlon_in), nlon_in))); } + } + const int r = block_max_int(loc, shred); + if (r < 0) { continue; } // no neighbors at this hi (shouldn't happen for a contiguous band) + + const int arc_lo = -r + pscale * (wtile * WO_TILE); // signed abs longitude of arc start + const int arc_w = min(nlon_in, 2 * r + pscale * (WO_TILE - 1) + 1); + const float qw = quad_weights[hi]; + const STORAGE_T *kx_hi = kx + (int64_t(batch) * nlat_in + hi) * nlon_in * nchan; + const STORAGE_T *vx_hi = vx + (int64_t(batch) * nlat_in + hi) * nlon_in * nchan; + + for (int c0 = 0; c0 < arc_w; c0 += CHUNK) { + const int cw = min(CHUNK, arc_w - c0); + + // stage cw input longitudes (K and V) into shared, coalesced along channel. + for (int idx = threadIdx.x; idx < cw * nchan; idx += blockDim.x) { + const int j = idx / nchan; + const int ch = idx - j * nchan; + const int lon = cmod(arc_lo + c0 + j, nlon_in); + const int64_t g = int64_t(lon) * nchan + ch; + shK[j * nchan + ch] = kx_hi[g]; + shV[j * nchan + ch] = vx_hi[g]; + } + __syncthreads(); + + if (valid) { + for (int j = 0; j < cw; j++) { + const int lon = cmod(arc_lo + c0 + j, nlon_in); + const int swi = csigned(cmod(lon - pscale * wo, nlon_in), nlon_in); + if (abs(swi) > r) { continue; } // not a neighbor of this wo (disc membership) + + const STORAGE_T *shk = shK + j * nchan; + const STORAGE_T *shv = shV + j * nchan; + + COMPUTE_T acc = 0.f; +#pragma unroll + for (int i = 0; i < NLOC; i++) { + const int ch = lane + i * WARP; + if (ch < nchan) { acc = __vfma(locq[i], vload(shk, ch), acc); } + } + float qdotk = __warp_sum(acc); + + const float qmax2 = max(qmax, qdotk); + const float alpha = expf(qdotk - qmax2) * qw; + const float corr = expf(qmax - qmax2); + asum = alpha + asum * corr; +#pragma unroll + for (int i = 0; i < NLOC; i++) { + const int ch = lane + i * WARP; + if (ch < nchan) { + locy[i] = __vfma_scale(alpha, vload(shv, ch), __vscale(corr, locy[i])); + } + } + qmax = qmax2; + } + } + __syncthreads(); + } + } + + if (valid) { + const float inv = (asum > 0.f) ? 1.f / asum : 0.f; + STORAGE_T *_y = y + (int64_t(batch) * nlat_out * nlon_out + int64_t(ho) * nlon_out + wo) * nchan; +#pragma unroll + for (int i = 0; i < NLOC; i++) { + const int ch = lane + i * WARP; + if (ch < nchan) { vstore(_y, ch, __vscale(inv, locy[i])); } + } + } + } + + // ---- launcher: pick NLOC = ceil(nchan/WARP) via a compile-time recursion ---- + template + void launch_coop(int nloc, int nchan, int nlat_in, int nlon_in, int nlat_out, int nlon_out, int tiles_per_row, + int batch_size, const STORAGE_T *kx, const STORAGE_T *vx, const STORAGE_T *qy, + const int32_t *row_idx, const int64_t *row_off, const int64_t *col_idx, const float *qw, + STORAGE_T *y, cudaStream_t stream) + { + if (NLOC == nloc) { + dim3 block(WO_TILE * WARP); + dim3 grid(nlat_out * tiles_per_row, batch_size); + size_t shbytes = size_t(2) * CHUNK * nchan * sizeof(STORAGE_T) + WO_TILE * sizeof(int); + s2_attn_fwd_coop_k + <<>>(nchan, nlat_in, nlon_in, nlat_out, nlon_out, tiles_per_row, kx, + vx, qy, row_idx, row_off, col_idx, qw, y); + CHECK_ERROR("s2_attn_fwd_coop_k"); + return; + } + if constexpr (NLOC < 32) { + launch_coop(nloc, nchan, nlat_in, nlon_in, nlat_out, nlon_out, tiles_per_row, + batch_size, kx, vx, qy, row_idx, row_off, col_idx, qw, y, stream); + } + } + + torch::Tensor s2_attention_fwd_coop_cuda(at::Tensor kx, at::Tensor vx, at::Tensor qy, at::Tensor quad_weights, + at::Tensor col_idx, at::Tensor row_off, int64_t nlon_in, + int64_t nlat_out, int64_t nlon_out) + { + CHECK_CUDA_INPUT_TENSOR(kx); + CHECK_CUDA_INPUT_TENSOR(vx); + CHECK_CUDA_INPUT_TENSOR(qy); + CHECK_CUDA_TENSOR(quad_weights); + CHECK_CUDA_TENSOR(col_idx); + CHECK_CUDA_TENSOR(row_off); + + const int64_t nchans_in = qy.size(1); + const int64_t nchans_out = vx.size(1); + TORCH_CHECK(nchans_in == nchans_out, "forward_coop: requires nchan_in == nchan_out (self/square-channel)"); + TORCH_CHECK(nlon_in % nlon_out == 0, "forward_coop: gather path only (nlon_in % nlon_out == 0)"); + const int nchan = int(nchans_in); + TORCH_CHECK(nchan <= 32 * 32, "forward_coop: nchan too large for register accumulator"); + + const int batch_size = kx.size(0); + const int64_t nlat_in = kx.size(2); + auto qy_type = qy.dtype(); + auto stream = at::cuda::getCurrentCUDAStream().stream(); + + at::Tensor row_idx = sortRows(int(nlat_out), row_off, stream); + + // channels-last (NHWC) as the kernels index it. + torch::Tensor kxP = kx, vxP = vx, qyP = qy; + const bool qy_cl = (qyP.strides()[1] == 1); + if (kxP.strides()[1] != 1) { kxP = permute_4D_to0231(kxP); } + if (vxP.strides()[1] != 1) { vxP = permute_4D_to0231(vxP); } + if (!qy_cl) { qyP = permute_4D_to0231(qyP); } + + const int64_t out_dims[] = {batch_size, nlat_out, nlon_out, nchans_out}; + torch::Tensor yP = torch::empty(out_dims, kxP.options()); + + const int tiles_per_row = int((nlon_out + WO_TILE - 1) / WO_TILE); + const int nloc = (nchan + WARP - 1) / WARP; + + const int32_t *_row_idx = reinterpret_cast(row_idx.data_ptr()); + const int64_t *_row_off = reinterpret_cast(row_off.data_ptr()); + const int64_t *_col_idx = reinterpret_cast(col_idx.data_ptr()); + const float *_qw = reinterpret_cast(quad_weights.data_ptr()); + + AT_DISPATCH_FLOATING_TYPES_AND2(at::kHalf, at::kBFloat16, qyP.scalar_type(), "s2_attention_fwd_coop_cuda", [&] { + launch_coop<1, scalar_t>(nloc, nchan, int(nlat_in), int(nlon_in), int(nlat_out), int(nlon_out), + tiles_per_row, batch_size, reinterpret_cast(kxP.data_ptr()), + reinterpret_cast(vxP.data_ptr()), + reinterpret_cast(qyP.data_ptr()), _row_idx, _row_off, + _col_idx, _qw, reinterpret_cast(yP.data_ptr()), stream); + }); + + C10_CUDA_KERNEL_LAUNCH_CHECK(); + + torch::Tensor yout = yP; + if (!qy_cl) { yout = permute_4D_to0312(yout); } + return yout.to(qy_type); + } + + TORCH_LIBRARY_IMPL(attention_kernels, CUDA, m) { m.impl("forward_coop", &s2_attention_fwd_coop_cuda); } + + } // namespace coop + +} // namespace attention_kernels diff --git a/torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_fwd_fast.cu b/torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_fwd_fast.cu new file mode 100644 index 00000000..918bae19 --- /dev/null +++ b/torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_fwd_fast.cu @@ -0,0 +1,622 @@ +// coding=utf-8 +// +// SPDX-FileCopyrightText: Copyright (c) 2025 The torch-harmonics Authors. All rights reserved. +// SPDX-License-Identifier: BSD-3-Clause +// +// Redistribution and use in source and binary forms, with or without +// modification, are permitted provided that the following conditions are met: +// +// 1. Redistributions of source code must retain the above copyright notice, this +// list of conditions and the following disclaimer. +// +// 2. Redistributions in binary form must reproduce the above copyright notice, +// this list of conditions and the following disclaimer in the documentation +// and/or other materials provided with the distribution. +// +// 3. Neither the name of the copyright holder nor the names of its +// contributors may be used to endorse or promote products derived from +// this software without specific prior written permission. +// +// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +// DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +// FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +// DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +// SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +// CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +// OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +#include "attention_cuda.cuh" +#include +#include +#include +#include +#include +#include +#include + +#include + +#include +#include + +#include "cudamacro.h" +#include "attention_cuda_utils.cuh" + +#define THREADS (64) + +#define MAX_LOCAL_ARR_LEN (16) + +// BEGIN - forward kernels and functions + +namespace attention_kernels +{ + + // scatter-direction launcher, defined in attention_cuda_fwd_upsample.cu; + // called by s2_attention_fwd_cuda when nlon_out % nlon_in == 0. + void s2_attn_fwd_upsample_dispatch(int batch_size, size_t nchans_in, size_t nchans_out, int64_t nlon_in, + int64_t nlat_in, int64_t nlat_out, int64_t nlon_out, torch::Tensor kxP, + torch::Tensor vxP, torch::Tensor qyP, torch::Tensor psi_row_off, + torch::Tensor psi_col_idx, torch::Tensor quad_weights, torch::Tensor yP); + + // ===================================================================================== + // Optimized (FMA + __expf) copy of the GATHER forward, exposed as the ``forward_fast`` + // op for A/B benchmarking against the baseline ``forward``. Everything lives in the + // nested ``fast`` namespace so the template kernels get distinct symbols (no ODR clash + // with the baseline copy). The upsample path delegates to the shared + // attention_kernels::s2_attn_fwd_upsample_dispatch (resolved by enclosing-namespace + // lookup). Kept until the fast path is validated as the default. + // ===================================================================================== + namespace fast + { + + // called with (blockDim.x=32 and blockDim.y>1, BDIM_X=blockDim.x*blockDim.y) + // + // STORAGE_T is the global-memory element type (float4 for the fp32 vectorized + // path; float / c10::Half / c10::BFloat16 for the scalar path). COMPUTE_T is + // the arithmetic type (float4 for the vectorized path, float otherwise) — all + // dot products, softmax and accumulation happen in COMPUTE_T. vload/vstore + // widen/narrow at the memory boundary. + + template + __global__ __launch_bounds__(BDIM_X) void s2_attn_fwd_generic_vec_k( + int nchan_in, // no. of STORAGE_T elements along channel dim + int nchan_out, // no. of STORAGE_T elements along channel dim + int nlat_in, int nlon_in, int nlat_out, int nlon_out, const STORAGE_T *__restrict__ kx, + const STORAGE_T *__restrict__ vx, const STORAGE_T *__restrict__ qy, const int32_t *__restrict__ row_idx, + const int64_t *__restrict__ row_off, const int64_t *__restrict__ col_idx, + const float *__restrict__ quad_weights, STORAGE_T *__restrict__ y) + { + using COMPUTE_T = typename vec_traits::compute_t; + + extern __shared__ __align__(sizeof(float4)) float shext[]; + COMPUTE_T *shy = reinterpret_cast(shext) + threadIdx.y * nchan_out; + + const int batch = blockIdx.y; + const int wid = blockIdx.x * blockDim.y + threadIdx.y; + + if (wid >= nlat_out * nlon_out) { return; } + + const int tidx = threadIdx.x; + + const int h = wid / nlon_out; + const int wo = wid - (h * nlon_out); + const int ho = row_idx[h]; + + // one output lon step corresponds to pscale input lon steps (requires nlon_in % nlon_out == 0) + const int pscale = nlon_in / nlon_out; + + for (int chan = tidx; chan < nchan_out; chan += WARP_SIZE) { shy[chan] = __vset(0.f); } + + kx += int64_t(batch) * nlat_in * nlon_in * nchan_in; + qy += int64_t(batch) * nlat_out * nlon_out * nchan_in + int64_t(ho) * nchan_in * nlon_out + + int64_t(wo) * nchan_in; + + vx += int64_t(batch) * nlat_in * nlon_in * nchan_out; + y += int64_t(batch) * nlat_out * nlon_out * nchan_out + int64_t(ho) * nchan_out * nlon_out + + int64_t(wo) * nchan_out; + + float alpha_sum = 0.0f; + float qdotk_max = -FLT_MAX; + + const int64_t rbeg = row_off[ho]; + const int64_t rend = row_off[ho + 1]; + + col_idx += rbeg; + + const int rlen = rend - rbeg; + + for (int off = 0; off < rlen; off++) { + + const int64_t col = col_idx[off]; + + const int hi = col / nlon_in; + const int wi = col - (hi * nlon_in); + const int wi_wo = wi + pscale * wo; + const int wip = wi_wo - (wi_wo / nlon_in) * nlon_in; + + const STORAGE_T *_kx = kx + int64_t(hi) * nlon_in * nchan_in + int64_t(wip) * nchan_in; + const STORAGE_T *_vx = vx + int64_t(hi) * nlon_in * nchan_out + int64_t(wip) * nchan_out; + + COMPUTE_T qdotkv = __vset(0.f); + + for (int chan = tidx; chan < nchan_in; chan += WARP_SIZE) { + qdotkv = __vfma(vload(qy, chan), vload(_kx, chan), qdotkv); + } + + float qdotk = __warp_sum(__vred(qdotkv)); + + float qdotk_max_tmp; + float alpha; + float exp_save; + + qdotk_max_tmp = max(qdotk_max, qdotk); + alpha = __expf(qdotk - qdotk_max_tmp) * quad_weights[hi]; + exp_save = __expf(qdotk_max - qdotk_max_tmp); + + alpha_sum = alpha + alpha_sum * exp_save; + + for (int chan = tidx; chan < nchan_out; chan += WARP_SIZE) { + shy[chan] = __vfma_scale(alpha, vload(_vx, chan), __vscale(exp_save, shy[chan])); + } + qdotk_max = qdotk_max_tmp; + } + + alpha_sum = 1.0f / alpha_sum; + for (int chan = tidx; chan < nchan_out; chan += WARP_SIZE) { + vstore(y, chan, __vscale(alpha_sum, shy[chan])); + } + + return; + } + + // called with either (BDIM_X=32 and BDIM_Y>1) || (2^K=BDIM_X > 32 and BDIM_Y=1) + template = nchan_out + typename STORAGE_T> + __global__ __launch_bounds__(BDIM_X *BDIM_Y) void s2_attn_fwd_special_vec_k( + int nchan_in, // no. of STORAGE_T elements along channel dim + int nchan_out, // no. of STORAGE_T elements along channel dim + int nlat_in, int nlon_in, int nlat_out, int nlon_out, const STORAGE_T *__restrict__ kx, + const STORAGE_T *__restrict__ vx, const STORAGE_T *__restrict__ qy, const int32_t *__restrict__ row_idx, + const int64_t *__restrict__ row_off, const int64_t *__restrict__ col_idx, + const float *__restrict__ quad_weights, STORAGE_T *__restrict__ y) + { + using COMPUTE_T = typename vec_traits::compute_t; + + static_assert(0 == (BDIM_X & (BDIM_X - 1))); + static_assert(0 == (BDIM_Y & (BDIM_Y - 1))); + static_assert((BDIM_X == 32 && BDIM_Y > 1) || (BDIM_X > 32 && BDIM_Y == 1)); + + constexpr int NLOC_M1 = NLOC - 1; + + const int tidx = threadIdx.x; + const int batch = blockIdx.y; + const int ctaid = blockIdx.x * blockDim.y + threadIdx.y; + + if (ctaid >= nlat_out * nlon_out) { return; } + + COMPUTE_T locy[NLOC]; + + // shq holds q already widened to COMPUTE_T (converted once on load, reused + // across the neighbor loop). + extern __shared__ __align__(sizeof(float4)) float shext[]; + COMPUTE_T *shq = reinterpret_cast(shext) + threadIdx.y * nchan_in; + + if constexpr (CHIN_AS_OUT) { shq += tidx; } + + const int h = ctaid / nlon_out; + const int wo = ctaid - (h * nlon_out); + const int ho = row_idx[h]; + + // one output lon step corresponds to pscale input lon steps (requires nlon_in % nlon_out == 0) + const int pscale = nlon_in / nlon_out; + + kx += int64_t(batch) * nlat_in * nlon_in * nchan_in; + qy += int64_t(batch) * nlat_out * nlon_out * nchan_in + int64_t(ho) * nlon_out * nchan_in + + int64_t(wo) * nchan_in; + if constexpr (CHIN_AS_OUT) { + kx += tidx; + qy += tidx; + } + + vx += int64_t(batch) * nlat_in * nlon_in * nchan_out + tidx; + y += int64_t(batch) * nlat_out * nlon_out * nchan_out + int64_t(ho) * nlon_out * nchan_out + + int64_t(wo) * nchan_out + tidx; + +#pragma unroll + for (int i = 0; i < NLOC; i++) { locy[i] = __vset(0.f); } + + if constexpr (CHIN_AS_OUT) { +#pragma unroll + for (int i = 0; i < NLOC_M1; i++) { shq[i * BDIM_X] = vload(qy, i * BDIM_X); } + if (NLOC_M1 * BDIM_X + tidx < nchan_in) { shq[NLOC_M1 * BDIM_X] = vload(qy, NLOC_M1 * BDIM_X); } + } else { + for (int chan = tidx; chan < nchan_in; chan += BDIM_X) { shq[chan] = vload(qy, chan); } + } + + float alpha_sum = 0.0f; + float qdotk_max = -FLT_MAX; + + const int64_t rbeg = row_off[ho]; + const int64_t rend = row_off[ho + 1]; + + col_idx += rbeg; + + const int rlen = rend - rbeg; + + for (int off = 0; off < rlen; off++) { + + const int64_t col = col_idx[off]; + + const int hi = col / nlon_in; + const int wi = col - (hi * nlon_in); + const int wi_wo = wi + pscale * wo; + const int wip = wi_wo - (wi_wo / nlon_in) * nlon_in; + + const STORAGE_T *_kx = kx + int64_t(hi) * nlon_in * nchan_in + int64_t(wip) * nchan_in; + const STORAGE_T *_vx = vx + int64_t(hi) * nlon_in * nchan_out + int64_t(wip) * nchan_out; + + COMPUTE_T qdotkv = __vset(0.f); + + if constexpr (CHIN_AS_OUT) { +#pragma unroll + for (int i = 0; i < NLOC_M1; i++) { + qdotkv = __vfma(shq[i * BDIM_X], vload(_kx, i * BDIM_X), qdotkv); + } + if (NLOC_M1 * BDIM_X + tidx < nchan_in) { + qdotkv = __vfma(shq[NLOC_M1 * BDIM_X], vload(_kx, NLOC_M1 * BDIM_X), qdotkv); + } + } else { + for (int chan = tidx; chan < nchan_in; chan += BDIM_X) { + qdotkv = __vfma(shq[chan], vload(_kx, chan), qdotkv); + } + } + + float qdotk = __vred(qdotkv); + if constexpr (BDIM_X == 32) { + qdotk = __warp_sum(qdotk); + } else { + qdotk = __block_sum(qdotk); + } + + float qdotk_max_tmp; + float alpha; + float exp_save; + + qdotk_max_tmp = max(qdotk_max, qdotk); + alpha = __expf(qdotk - qdotk_max_tmp) * quad_weights[hi]; + exp_save = __expf(qdotk_max - qdotk_max_tmp); + + alpha_sum = alpha + alpha_sum * exp_save; + +#pragma unroll + for (int i = 0; i < NLOC_M1; i++) { + locy[i] = __vfma_scale(alpha, vload(_vx, i * BDIM_X), __vscale(exp_save, locy[i])); + } + if (NLOC_M1 * BDIM_X + tidx < nchan_out) { + locy[NLOC_M1] = __vfma_scale(alpha, vload(_vx, NLOC_M1 * BDIM_X), __vscale(exp_save, locy[NLOC_M1])); + } + + qdotk_max = qdotk_max_tmp; + } + + alpha_sum = 1.0f / alpha_sum; + +#pragma unroll + for (int i = 0; i < NLOC_M1; i++) { vstore(y, i * BDIM_X, __vscale(alpha_sum, locy[i])); } + if (NLOC_M1 * BDIM_X + tidx < nchan_out) { + vstore(y, NLOC_M1 * BDIM_X, __vscale(alpha_sum, locy[NLOC_M1])); + } + + return; + } + + template + void launch_gen_attn_fwd(int batch_size, int nchans_in, int nchans_out, int nlat_in, int nlon_in, int nlat_out, + int nlon_out, STORAGE_T *__restrict__ _kxp, STORAGE_T *__restrict__ _vxp, + STORAGE_T *__restrict__ _qyp, int32_t *_row_idx, int64_t *_row_off, int64_t *_col_idx, + float *_quad_weights, STORAGE_T *__restrict__ _yp, cudaStream_t stream) + { + + dim3 block(WARP_SIZE, THREADS / WARP_SIZE); + dim3 grid(DIV_UP(nlat_out * nlon_out, block.y), batch_size); + + // shared memory holds compute-type (COMPUTE_T) data, not STORAGE_T. + size_t shsize = sizeof(typename vec_traits::compute_t) * nchans_out * block.y; + + s2_attn_fwd_generic_vec_k + <<>>(nchans_in, nchans_out, nlat_in, nlon_in, nlat_out, nlon_out, _kxp, + _vxp, _qyp, _row_idx, _row_off, _col_idx, _quad_weights, _yp); + CHECK_ERROR("s2_attn_fwd_generic_vec_k"); + + return; + } + + template + void launch_spc_attn_fwd(int nloc, // "BDIM_X*nloc" >= nchans_out + int batch_size, int nchans_in, int nchans_out, int nlat_in, int nlon_in, int nlat_out, + int nlon_out, STORAGE_T *__restrict__ _kxp, STORAGE_T *__restrict__ _vxp, + STORAGE_T *__restrict__ _qyp, int32_t *_row_idx, int64_t *_row_off, int64_t *_col_idx, + float *_quad_weights, STORAGE_T *__restrict__ _yp, cudaStream_t stream) + { + + if (CUR_LOC_SIZE == nloc) { + + dim3 block(BDIM_X, BDIM_Y); + dim3 grid(DIV_UP(nlat_out * nlon_out, block.y), batch_size); + + // shared memory holds compute-type (COMPUTE_T) data, not STORAGE_T. + // block.y > 1 iif block.x==32 + size_t shsize = sizeof(typename vec_traits::compute_t) * nchans_in * block.y; + + // nloc determines the size of local arrays used to store + // y vectors, of length nchans_out; + // if nchans_in is >= BDIM_X*(nloc-1) and <= BDIM_X*nloc + // then we can use the same compile-time known loops used + // for output channels, with the execpetion of testing + // whether to execute the last iteration based on "nchans_in" + // rather than on "nchans_out"; in this way as long as the + // difference between the number of input and output channels + // is <= BDIM_X we can use the faster path + if (nchans_in >= BDIM_X * (CUR_LOC_SIZE - 1) && nchans_in <= BDIM_X * CUR_LOC_SIZE) { + + s2_attn_fwd_special_vec_k<<>>( + nchans_in, nchans_out, nlat_in, nlon_in, nlat_out, nlon_out, _kxp, _vxp, _qyp, _row_idx, + _row_off, _col_idx, _quad_weights, _yp); + } else { + + s2_attn_fwd_special_vec_k<<>>( + nchans_in, nchans_out, nlat_in, nlon_in, nlat_out, nlon_out, _kxp, _vxp, _qyp, _row_idx, + _row_off, _col_idx, _quad_weights, _yp); + } + CHECK_ERROR("s2_attn_fwd_special_vec_k"); + + return; + } + if constexpr (CUR_LOC_SIZE < MAX_LOC_SIZE) { + launch_spc_attn_fwd( + nloc, batch_size, nchans_in, nchans_out, nlat_in, nlon_in, nlat_out, nlon_out, _kxp, _vxp, _qyp, + _row_idx, _row_off, _col_idx, _quad_weights, _yp, stream); + } + return; + } + + // Picks the block size (BDIM_X) instance and launches the gather kernel for a + // given storage vector type SV. MAX_LOC / MIN_LOC bound the per-thread local + // array length (in COMPUTE_T units). nci / nco are channel counts in SV units. + template + static void fwd_dispatch_bdimx(int bdimx, int nloc, int64_t batch_size, int64_t nci, int64_t nco, int nlat_in, + int64_t nlon_in, int64_t nlat_out, int64_t nlon_out, SV *_kxp, SV *_vxp, + SV *_qyp, int32_t *_row_idx, int64_t *_row_off, int64_t *_col_idx, + float *_quad_weights, SV *_yp, cudaStream_t stream) + { + // use 2D blocks only if 32 threads are enough + switch (bdimx) { + case 32: + launch_spc_attn_fwd<32, 2, 1, MAX_LOC>(nloc, batch_size, nci, nco, nlat_in, nlon_in, nlat_out, nlon_out, + _kxp, _vxp, _qyp, _row_idx, _row_off, _col_idx, _quad_weights, + _yp, stream); + break; + case 64: + launch_spc_attn_fwd<64, 1, MIN_LOC, MAX_LOC>(nloc, batch_size, nci, nco, nlat_in, nlon_in, nlat_out, + nlon_out, _kxp, _vxp, _qyp, _row_idx, _row_off, _col_idx, + _quad_weights, _yp, stream); + break; + case 128: + launch_spc_attn_fwd<128, 1, MIN_LOC, MAX_LOC>(nloc, batch_size, nci, nco, nlat_in, nlon_in, nlat_out, + nlon_out, _kxp, _vxp, _qyp, _row_idx, _row_off, _col_idx, + _quad_weights, _yp, stream); + break; + case 256: + launch_spc_attn_fwd<256, 1, MIN_LOC, MAX_LOC>(nloc, batch_size, nci, nco, nlat_in, nlon_in, nlat_out, + nlon_out, _kxp, _vxp, _qyp, _row_idx, _row_off, _col_idx, + _quad_weights, _yp, stream); + break; + case 512: + launch_spc_attn_fwd<512, 1, MIN_LOC, MAX_LOC>(nloc, batch_size, nci, nco, nlat_in, nlon_in, nlat_out, + nlon_out, _kxp, _vxp, _qyp, _row_idx, _row_off, _col_idx, + _quad_weights, _yp, stream); + break; + case 1024: + launch_spc_attn_fwd<1024, 1, MIN_LOC, MAX_LOC>(nloc, batch_size, nci, nco, nlat_in, nlon_in, nlat_out, + nlon_out, _kxp, _vxp, _qyp, _row_idx, _row_off, _col_idx, + _quad_weights, _yp, stream); + break; + default: + launch_gen_attn_fwd(batch_size, nci, nco, nlat_in, nlon_in, nlat_out, nlon_out, _kxp, _vxp, _qyp, + _row_idx, _row_off, _col_idx, _quad_weights, _yp, stream); + break; + } + } + + // Templated on the storage element type (float / c10::Half / c10::BFloat16). + // Path selection (compute / accumulation are fp32 in every case): + // - fp32, 16B-aligned, nchans % 4 == 0 -> float4 vectorized (LDG.128) + // - fp16/bf16, 16B-aligned, nchans % 8 == 0 -> half8/bf168 vectorized (LDG.128) + // - otherwise -> scalar STORAGE_T path + template + static void s2_attn_fwd_dispatch(int64_t batch_size, int64_t nchans_in, int64_t nchans_out, int64_t nlon_in, + int64_t nlat_out, int64_t nlon_out, at::Tensor kxP, at::Tensor vxP, + at::Tensor qyP, at::Tensor row_off, at::Tensor col_idx, + at::Tensor quad_weights, at::Tensor yP) + { + + static_assert(0 == (MAX_LOCAL_ARR_LEN & (MAX_LOCAL_ARR_LEN - 1))); + + // get stream + auto stream = at::cuda::getCurrentCUDAStream().stream(); + + // sort row indices (ho-s) in descending order + // based on (row_off[ho+1]-row_off[ho]) + at::Tensor row_idx = sortRows(nlat_out, row_off, stream); + + const int nlat_in = kxP.size(1); + + // smallest power of two "bdimx" (>=32) s.t. bdimx*MAX_LOCAL_ARR_LEN >= nchans_out + int bdimx; + bdimx = DIV_UP(nchans_out, MAX_LOCAL_ARR_LEN); + bdimx = max(bdimx, WARP_SIZE); + bdimx = next_pow2(bdimx); + + scalar_t *_kxp = reinterpret_cast(kxP.data_ptr()); + scalar_t *_vxp = reinterpret_cast(vxP.data_ptr()); + scalar_t *_qyp = reinterpret_cast(qyP.data_ptr()); + scalar_t *_yp = reinterpret_cast(yP.data_ptr()); + + int32_t *_row_idx = reinterpret_cast(row_idx.data_ptr()); + int64_t *_row_off = reinterpret_cast(row_off.data_ptr()); + int64_t *_col_idx = reinterpret_cast(col_idx.data_ptr()); + float *_quad_weights = reinterpret_cast(quad_weights.data_ptr()); + + constexpr int MIN_LOC_ARR_LEN = MAX_LOCAL_ARR_LEN / 2 + 1; + + if constexpr (std::is_same::value) { + // fp32: float4 vectorized when 16B-aligned + 4-divisible, else scalar. + constexpr int VEC_SIZE = sizeof(float4) / sizeof(float); // 4 + const bool use_vec = is_aligned<16>(_kxp) && is_aligned<16>(_vxp) && is_aligned<16>(_qyp) + && is_aligned<16>(_yp) && (nchans_in % VEC_SIZE) == 0 && (nchans_out % VEC_SIZE) == 0; + + if (use_vec) { + constexpr int MAX_VEC = MAX_LOCAL_ARR_LEN / VEC_SIZE; + constexpr int MIN_VEC = MAX_VEC / 2 + 1; + const int64_t nci = nchans_in / VEC_SIZE; + const int64_t nco = nchans_out / VEC_SIZE; + fwd_dispatch_bdimx( + bdimx, DIV_UP(nco, bdimx), batch_size, nci, nco, nlat_in, nlon_in, nlat_out, nlon_out, + reinterpret_cast(_kxp), reinterpret_cast(_vxp), + reinterpret_cast(_qyp), _row_idx, _row_off, _col_idx, _quad_weights, + reinterpret_cast(_yp), stream); + } else { + fwd_dispatch_bdimx( + bdimx, DIV_UP(nchans_out, bdimx), batch_size, nchans_in, nchans_out, nlat_in, nlon_in, nlat_out, + nlon_out, _kxp, _vxp, _qyp, _row_idx, _row_off, _col_idx, _quad_weights, _yp, stream); + } + } else { + // fp16/bf16: scalar STORAGE_T path (widen at load, narrow at store; fp32 + // compute/accumulation). A vectorized 8-wide path was tried and reverted: + // it raised register pressure and lowered occupancy, and ncu shows this + // kernel is latency/occupancy-bound (DRAM ~25%), not bandwidth-bound, so + // vectorizing reduced precision only hurt. See the AMP refactor notes. + fwd_dispatch_bdimx( + bdimx, DIV_UP(nchans_out, bdimx), batch_size, nchans_in, nchans_out, nlat_in, nlon_in, nlat_out, + nlon_out, _kxp, _vxp, _qyp, _row_idx, _row_off, _col_idx, _quad_weights, _yp, stream); + } + + return; + } + + // END - forward kernels and functions + + torch::Tensor s2_attention_fwd_cuda(at::Tensor kx, at::Tensor vx, at::Tensor qy, at::Tensor quad_weights, + at::Tensor psi_col_idx, at::Tensor psi_row_off, int64_t nlon_in, + int64_t nlat_out, int64_t nlon_out) + { + CHECK_CUDA_INPUT_TENSOR(kx); + CHECK_CUDA_INPUT_TENSOR(vx); + CHECK_CUDA_INPUT_TENSOR(qy); + CHECK_CUDA_TENSOR(quad_weights); + CHECK_CUDA_TENSOR(psi_col_idx); + CHECK_CUDA_TENSOR(psi_row_off); + + // direction selection: gather (self / downsample) iff nlon_in is an integer + // multiple of nlon_out; scatter (upsample) iff nlon_out is an integer multiple + // of nlon_in. Self-attention satisfies both and routes through the gather path. + const bool downsample = (nlon_in % nlon_out == 0); + const bool upsample = (nlon_out % nlon_in == 0); + TORCH_CHECK(downsample || upsample, "either nlon_in (", nlon_in, + ") must be an integer multiple of nlon_out (", nlon_out, "), or vice versa"); + + size_t nchans_in = qy.size(1); // or kx.size(1) + size_t nchans_out = vx.size(1); + + const int batch_size = kx.size(0); + const int64_t nlat_in = kx.size(2); + + // extract dtype + auto qy_type = qy.dtype(); + + const int64_t out_dims[] = {batch_size, nlat_out, nlon_out, nchans_out}; + torch::Tensor y; + + // ATen dispatch over the input dtype. + // + // Gather (downsample / self) path: native storage. kx/vx/qy keep their + // dtype; the kernel widens to fp32 at load and narrows back at store + // (Tier B), so there is no whole-tensor fp32 copy and the read bandwidth + // for fp16/bf16 is halved. Compute/accumulation are fp32 in-kernel. + // + // Scatter (upsample) path: the scatter kernels are not yet on the native + // storage path, so we still upcast to fp32 here and downcast the result + // at the end (the trailing y.to(qy_type)). fp32 inputs are unaffected + // either way (the .to(kFloat32) is a no-op). + AT_DISPATCH_FLOATING_TYPES_AND2(at::kHalf, at::kBFloat16, qy.scalar_type(), "s2_attention_fwd_cuda", [&] { + using storage_t = scalar_t; + + if (downsample) { + // native-storage gather path + torch::Tensor kxP = kx; + torch::Tensor vxP = vx; + torch::Tensor qyP = qy; + + // safer than is_contiguous(ChannelsLast), which fails for num_channels == 1 + bool kx_is_channels_last = kxP.strides()[1] == 1; + bool vx_is_channels_last = vxP.strides()[1] == 1; + bool qy_is_channels_last = qyP.strides()[1] == 1; + + if (!kx_is_channels_last) { kxP = permute_4D_to0231(kxP); } + if (!vx_is_channels_last) { vxP = permute_4D_to0231(vxP); } + if (!qy_is_channels_last) { qyP = permute_4D_to0231(qyP); } + + torch::Tensor yP = torch::empty(out_dims, kxP.options()); // native dtype + + s2_attn_fwd_dispatch(batch_size, nchans_in, nchans_out, nlon_in, nlat_out, nlon_out, kxP, + vxP, qyP, psi_row_off, psi_col_idx, quad_weights, yP); + + y = yP; + if (!qy_is_channels_last) { y = permute_4D_to0312(y); } + } else { + // upsample (scatter) path: native storage. s2_attn_fwd_upsample_dispatch + // does its own AT_DISPATCH and widens fp16/bf16 at load (fp32 compute), + // narrowing the output at store — same as the gather path. + torch::Tensor kxP = kx; + torch::Tensor vxP = vx; + torch::Tensor qyP = qy; + + bool kx_is_channels_last = kxP.strides()[1] == 1; + bool vx_is_channels_last = vxP.strides()[1] == 1; + bool qy_is_channels_last = qyP.strides()[1] == 1; + + if (!kx_is_channels_last) { kxP = permute_4D_to0231(kxP); } + if (!vx_is_channels_last) { vxP = permute_4D_to0231(vxP); } + if (!qy_is_channels_last) { qyP = permute_4D_to0231(qyP); } + + torch::Tensor yP = torch::empty(out_dims, kxP.options()); // native dtype + + s2_attn_fwd_upsample_dispatch(batch_size, nchans_in, nchans_out, nlon_in, nlat_in, nlat_out, + nlon_out, kxP, vxP, qyP, psi_row_off, psi_col_idx, quad_weights, yP); + + y = yP; + if (!qy_is_channels_last) { y = permute_4D_to0312(y); } + } + }); + + // convert precision back to starting dtype. No-op now that both the gather + // and upsample paths produce native-dtype output; kept as a safety net. + y = y.to(qy_type); + + C10_CUDA_KERNEL_LAUNCH_CHECK(); + + return y; + } + + TORCH_LIBRARY_IMPL(attention_kernels, CUDA, m) { m.impl("forward_fast", &s2_attention_fwd_cuda); } + + } // namespace fast + +} // namespace attention_kernels diff --git a/torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_fwd_half.cu b/torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_fwd_half.cu new file mode 100644 index 00000000..c20aef20 --- /dev/null +++ b/torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_fwd_half.cu @@ -0,0 +1,287 @@ +// coding=utf-8 +// +// SPDX-FileCopyrightText: Copyright (c) 2026 The torch-harmonics Authors. All rights reserved. +// SPDX-License-Identifier: BSD-3-Clause +// +// Redistribution and use in source and binary forms, with or without +// modification, are permitted provided that the following conditions are met: +// +// 1. Redistributions of source code must retain the above copyright notice, this +// list of conditions and the following disclaimer. +// +// 2. Redistributions in binary form must reproduce the above copyright notice, +// this list of conditions and the following disclaimer in the documentation +// and/or other materials provided with the distribution. +// +// 3. Neither the name of the copyright holder nor the names of its +// contributors may be used to endorse or promote products derived from +// this software without specific prior written permission. +// +// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +// DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +// FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +// DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +// SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +// CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +// OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +// ===================================================================================== +// Gather (downsample / self) attention forward — fp16/bf16 PACKED specialization +// ===================================================================================== +// +// A bf16/fp16-specialized variant of the generic gather forward (warp-per-output +// cell). Motivation: stored as a scalar, each c10::Half lands in its own 32-bit +// register (top 16 bits wasted), so a wide fp16 load doubles register pressure +// versus fp32 and tips this (register-limited) kernel into spills — which is why +// the earlier scalar-fp16 "float8" experiment regressed. Here the activations are +// loaded 8-wide (LDG.128) as 4x __half2 / __nv_bfloat162 and the q.k dot product +// runs on packed __hfma2, so two halves share one register. +// +// Precision is unchanged from the scalar path: the q.k accumulator is packed +// half2 only for the handful of per-lane terms (nchan/8/32), then reduced in fp32; +// the softmax and the per-channel value accumulator `shy` stay fp32. Only the +// loads + the dot product are packed; alpha*v accumulates in fp32 as before. +// +// bf16 packed FMA is native on sm_80+; on older arches the traits fall back to a +// float-based fma2 so the TU still compiles, and the host dispatch only routes +// bf16 here when the device is sm_80+ (fp16 packs on Volta+). +// ===================================================================================== + +#include "attention_cuda.cuh" + +#include +#include +#include +#include + +#include "cudamacro.h" +#include "attention_cuda_utils.cuh" + +#define THREADS (64) + +namespace attention_kernels +{ + + // packed-half traits: maps a c10 element type to its 2-wide vector type, an + // 8-wide (16-byte, LDG.128) load struct, and the pack/unpack/fma2 primitives. + template struct pack8; + + template <> struct pack8 { + using v2_t = __half2; + struct v8_t { + __half2 a, b, c, d; + }; + static __device__ __forceinline__ v2_t zero() { return __floats2half2_rn(0.f, 0.f); } + static __device__ __forceinline__ v2_t fma2(v2_t x, v2_t y, v2_t acc) { return __hfma2(x, y, acc); } + static __device__ __forceinline__ float2 to_f2(v2_t v) { return __half22float2(v); } + static __device__ __forceinline__ v2_t pack(float lo, float hi) { return __floats2half2_rn(lo, hi); } + }; + + template <> struct pack8 { + using v2_t = __nv_bfloat162; + struct v8_t { + __nv_bfloat162 a, b, c, d; + }; + static __device__ __forceinline__ v2_t zero() { return __floats2bfloat162_rn(0.f, 0.f); } + static __device__ __forceinline__ v2_t fma2(v2_t x, v2_t y, v2_t acc) + { +#if __CUDA_ARCH__ >= 800 + return __hfma2(x, y, acc); +#else + // pre-sm_80: bf16 packed FMA is not native. This path exists only so the + // TU compiles; the host dispatch never routes bf16 here below sm_80. + const float2 xf = __bfloat1622float2(x); + const float2 yf = __bfloat1622float2(y); + const float2 af = __bfloat1622float2(acc); + return __floats2bfloat162_rn(xf.x * yf.x + af.x, xf.y * yf.y + af.y); +#endif + } + static __device__ __forceinline__ float2 to_f2(v2_t v) { return __bfloat1622float2(v); } + static __device__ __forceinline__ v2_t pack(float lo, float hi) { return __floats2bfloat162_rn(lo, hi); } + }; + + // warp-per-output-cell gather forward, packed fp16/bf16. nchan_in/nchan_out are + // in ELEM units and must be multiples of 8 (checked by the host dispatch). + template + __global__ __launch_bounds__(BDIM_X) void s2_attn_fwd_generic_half2_k( + int nchan_in, int nchan_out, int nlat_in, int nlon_in, int nlat_out, int nlon_out, const ELEM *__restrict__ kx, + const ELEM *__restrict__ vx, const ELEM *__restrict__ qy, const int32_t *__restrict__ row_idx, + const int64_t *__restrict__ row_off, const int64_t *__restrict__ col_idx, + const float *__restrict__ quad_weights, ELEM *__restrict__ y) + { + using PK = pack8; + using v2_t = typename PK::v2_t; + using v8_t = typename PK::v8_t; + + const int nchan8_in = nchan_in >> 3; + const int nchan8_out = nchan_out >> 3; + + // fp32 per-channel value accumulator (precision unchanged from scalar path). + // Each lane owns the 8 contiguous channels of one v8_t; storing them at + // stride SHY_STRIDE (=9) per c8 block makes the per-j cross-lane access + // bank-conflict free (gcd(9,32)==1), since lane t writes shy[c8*9 + j] and + // c8 = t + 32*m -> bank (9*t + j) mod 32 is a bijection over the warp. + constexpr int SHY_STRIDE = 9; + extern __shared__ __align__(sizeof(float4)) float shext[]; + float *shy = shext + threadIdx.y * (SHY_STRIDE * nchan8_out); + + const int batch = blockIdx.y; + const int wid = blockIdx.x * blockDim.y + threadIdx.y; + if (wid >= nlat_out * nlon_out) { return; } + + const int tidx = threadIdx.x; + const int h = wid / nlon_out; + const int wo = wid - (h * nlon_out); + const int ho = row_idx[h]; + const int pscale = nlon_in / nlon_out; + + for (int i = tidx; i < SHY_STRIDE * nchan8_out; i += WARP_SIZE) { shy[i] = 0.f; } + + kx += int64_t(batch) * nlat_in * nlon_in * nchan_in; + qy += int64_t(batch) * nlat_out * nlon_out * nchan_in + int64_t(ho) * nchan_in * nlon_out + + int64_t(wo) * nchan_in; + vx += int64_t(batch) * nlat_in * nlon_in * nchan_out; + y += int64_t(batch) * nlat_out * nlon_out * nchan_out + int64_t(ho) * nchan_out * nlon_out + + int64_t(wo) * nchan_out; + + const v8_t *qy8 = reinterpret_cast(qy); + + float alpha_sum = 0.0f; + float qdotk_max = -FLT_MAX; + + const int64_t rbeg = row_off[ho]; + const int64_t rend = row_off[ho + 1]; + col_idx += rbeg; + const int rlen = static_cast(rend - rbeg); + + for (int off = 0; off < rlen; off++) { + const int64_t col = col_idx[off]; + const int hi = col / nlon_in; + const int wi = col - (hi * nlon_in); + const int wi_wo = wi + pscale * wo; + const int wip = wi_wo - (wi_wo / nlon_in) * nlon_in; + + const v8_t *kx8 + = reinterpret_cast(kx + int64_t(hi) * nlon_in * nchan_in + int64_t(wip) * nchan_in); + const v8_t *vx8 + = reinterpret_cast(vx + int64_t(hi) * nlon_in * nchan_out + int64_t(wip) * nchan_out); + + // q.k: packed half2 FMA per lane, reduced to fp32 across the warp. + v2_t acc = PK::zero(); + for (int c8 = tidx; c8 < nchan8_in; c8 += WARP_SIZE) { + const v8_t q = qy8[c8]; + const v8_t k = kx8[c8]; + acc = PK::fma2(q.a, k.a, acc); + acc = PK::fma2(q.b, k.b, acc); + acc = PK::fma2(q.c, k.c, acc); + acc = PK::fma2(q.d, k.d, acc); + } + const float2 accf = PK::to_f2(acc); + const float qdotk = __warp_sum(accf.x + accf.y); + + const float qdotk_max_tmp = max(qdotk_max, qdotk); + const float alpha = expf(qdotk - qdotk_max_tmp) * quad_weights[hi]; + const float exp_save = expf(qdotk_max - qdotk_max_tmp); + alpha_sum = alpha + alpha_sum * exp_save; + + // value accumulate: load v packed, convert to fp32, accumulate in fp32. + for (int c8 = tidx; c8 < nchan8_out; c8 += WARP_SIZE) { + const v8_t v = vx8[c8]; + const float2 a = PK::to_f2(v.a); + const float2 b = PK::to_f2(v.b); + const float2 c = PK::to_f2(v.c); + const float2 d = PK::to_f2(v.d); + float *sh = shy + c8 * SHY_STRIDE; + sh[0] = exp_save * sh[0] + alpha * a.x; + sh[1] = exp_save * sh[1] + alpha * a.y; + sh[2] = exp_save * sh[2] + alpha * b.x; + sh[3] = exp_save * sh[3] + alpha * b.y; + sh[4] = exp_save * sh[4] + alpha * c.x; + sh[5] = exp_save * sh[5] + alpha * c.y; + sh[6] = exp_save * sh[6] + alpha * d.x; + sh[7] = exp_save * sh[7] + alpha * d.y; + } + qdotk_max = qdotk_max_tmp; + } + + const float inv = 1.0f / alpha_sum; + v8_t *y8 = reinterpret_cast(y); + for (int c8 = tidx; c8 < nchan8_out; c8 += WARP_SIZE) { + const float *sh = shy + c8 * SHY_STRIDE; + v8_t out; + out.a = PK::pack(sh[0] * inv, sh[1] * inv); + out.b = PK::pack(sh[2] * inv, sh[3] * inv); + out.c = PK::pack(sh[4] * inv, sh[5] * inv); + out.d = PK::pack(sh[6] * inv, sh[7] * inv); + y8[c8] = out; + } + } + + // ----------------------------------------------------------------------------- + // host dispatch — called from s2_attn_fwd_dispatch (attention_cuda_fwd.cu) for + // the fp16/bf16 gather/self path. Returns true if it launched the packed kernel, + // false when not eligible (caller then takes the scalar fallback). Eligibility: + // dtype in {Half, BFloat16}, nchans % 8 == 0, 16-byte aligned, and (for bf16) + // device sm_80+. + // ----------------------------------------------------------------------------- + bool s2_attn_fwd_half2_dispatch(int64_t batch_size, int64_t nchans_in, int64_t nchans_out, int64_t nlat_in, + int64_t nlon_in, int64_t nlat_out, int64_t nlon_out, at::Tensor kxP, at::Tensor vxP, + at::Tensor qyP, at::Tensor row_idx, at::Tensor row_off, at::Tensor col_idx, + at::Tensor quad_weights, at::Tensor yP) + { + if ((nchans_in % 8) != 0 || (nchans_out % 8) != 0) { return false; } + + const auto dt = qyP.scalar_type(); + const bool is_half = (dt == at::kHalf); + const bool is_bf16 = (dt == at::kBFloat16); + if (!is_half && !is_bf16) { return false; } + + // bf16 packed FMA is native only on sm_80+; below that, defer to the scalar path. + if (is_bf16 && at::cuda::getCurrentDeviceProperties()->major < 8) { return false; } + + void *kxp = kxP.data_ptr(); + void *vxp = vxP.data_ptr(); + void *qyp = qyP.data_ptr(); + void *yp = yP.data_ptr(); + if (!is_aligned<16>(kxp) || !is_aligned<16>(vxp) || !is_aligned<16>(qyp) || !is_aligned<16>(yp)) { + return false; + } + + auto stream = at::cuda::getCurrentCUDAStream().stream(); + + int32_t *_row_idx = reinterpret_cast(row_idx.data_ptr()); + int64_t *_row_off = reinterpret_cast(row_off.data_ptr()); + int64_t *_col_idx = reinterpret_cast(col_idx.data_ptr()); + float *_quad_weights = reinterpret_cast(quad_weights.data_ptr()); + + dim3 block(WARP_SIZE, THREADS / WARP_SIZE); + dim3 grid(DIV_UP(nlat_out * nlon_out, block.y), batch_size); + // padded fp32 accumulator: 9 floats per 8-channel block (bank-conflict free) + const size_t shsize = sizeof(float) * 9 * (nchans_out / 8) * block.y; + + const int nci = static_cast(nchans_in); + const int nco = static_cast(nchans_out); + const int nli = static_cast(nlat_in); + const int nlonI = static_cast(nlon_in); + const int nlo = static_cast(nlat_out); + const int nlonO = static_cast(nlon_out); + + if (is_half) { + s2_attn_fwd_generic_half2_k<<>>( + nci, nco, nli, nlonI, nlo, nlonO, reinterpret_cast(kxp), + reinterpret_cast(vxp), reinterpret_cast(qyp), _row_idx, _row_off, + _col_idx, _quad_weights, reinterpret_cast(yp)); + } else { + s2_attn_fwd_generic_half2_k<<>>( + nci, nco, nli, nlonI, nlo, nlonO, reinterpret_cast(kxp), + reinterpret_cast(vxp), reinterpret_cast(qyp), _row_idx, + _row_off, _col_idx, _quad_weights, reinterpret_cast(yp)); + } + CHECK_ERROR("s2_attn_fwd_generic_half2_k"); + return true; + } + +} // namespace attention_kernels diff --git a/torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_utils.cuh b/torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_utils.cuh index 105587f8..8050b297 100644 --- a/torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_utils.cuh +++ b/torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_utils.cuh @@ -159,6 +159,24 @@ namespace attention_kernels return make_float4(s * v.x, s * v.y, s * v.z, s * v.w); } + // Fused multiply-add helpers (single-rounding fmaf). Explicit fusion for the hot + // dot-product and online-softmax accumulate loops (ncu flagged non-fused fp32 pairs). + // __vfma(a,b,c) = a*b + c (elementwise) + // __vfma_scale(s,v,c)= s*v + c (scalar s broadcast over vector v) + __device__ float __forceinline__ __vfma(float a, float b, float c) { return fmaf(a, b, c); } + + __device__ float4 __forceinline__ __vfma(float4 a, float4 b, float4 c) + { + return make_float4(fmaf(a.x, b.x, c.x), fmaf(a.y, b.y, c.y), fmaf(a.z, b.z, c.z), fmaf(a.w, b.w, c.w)); + } + + __device__ float __forceinline__ __vfma_scale(float s, float v, float c) { return fmaf(s, v, c); } + + __device__ float4 __forceinline__ __vfma_scale(float s, float4 v, float4 c) + { + return make_float4(fmaf(s, v.x, c.x), fmaf(s, v.y, c.y), fmaf(s, v.z, c.z), fmaf(s, v.w, c.w)); + } + __device__ float4 __forceinline__ __vdiv(float s, float4 v) { return make_float4(s / v.x, s / v.y, s / v.z, s / v.w); From 805e10898108b0408c0035220c50b87e2d8ec927 Mon Sep 17 00:00:00 2001 From: Thorsten Kurth Date: Tue, 14 Jul 2026 07:17:32 -0700 Subject: [PATCH 06/14] adding coop bwd kernel --- setup.py | 1 + .../optimized/attention_interface.cpp | 3 + .../kernels_cuda/attention_cuda_bwd_coop.cu | 1022 +++++++++++++++++ 3 files changed, 1026 insertions(+) create mode 100644 torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_bwd_coop.cu diff --git a/setup.py b/setup.py index 395b911e..637c8d08 100644 --- a/setup.py +++ b/setup.py @@ -201,6 +201,7 @@ def get_ext_modules(): "torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_fwd_fast.cu", "torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_bwd_fast.cu", "torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_fwd_coop.cu", + "torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_bwd_coop.cu", ] ) ext_modules.append(CUDAExtension("torch_harmonics.attention._C", attention_sources, extra_compile_args=get_compile_args("attention"))) diff --git a/torch_harmonics/attention/optimized/attention_interface.cpp b/torch_harmonics/attention/optimized/attention_interface.cpp index a3ceebe7..d9236d50 100644 --- a/torch_harmonics/attention/optimized/attention_interface.cpp +++ b/torch_harmonics/attention/optimized/attention_interface.cpp @@ -103,6 +103,9 @@ namespace attention_kernels // nchan_in==nchan_out. Same signature as forward; call via forward_coop. m.def("forward_coop(Tensor kx, Tensor vx, Tensor qy, Tensor quad_weights, Tensor col_idx, Tensor row_off, int " "nlon_in, int nlat_out, int nlon_out) -> Tensor"); + // Backward with pass-2 recompute elimination (caches qdotk/gdotv in shared). + m.def("backward_coop(Tensor kx, Tensor vx, Tensor qy, Tensor dy, Tensor quad_weights, Tensor col_idx, Tensor " + "row_off, int nlon_in, int nlat_out, int nlon_out) -> (Tensor, Tensor, Tensor)"); // ---- Ring-step variants for DistributedNeighborhoodAttentionS2 ---- // K/V are sharded along longitude across an azimuth process group; each diff --git a/torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_bwd_coop.cu b/torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_bwd_coop.cu new file mode 100644 index 00000000..776bfbe5 --- /dev/null +++ b/torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_bwd_coop.cu @@ -0,0 +1,1022 @@ +// coding=utf-8 +// +// SPDX-FileCopyrightText: Copyright (c) 2025 The torch-harmonics Authors. All rights reserved. +// SPDX-License-Identifier: BSD-3-Clause +// +// Redistribution and use in source and binary forms, with or without +// modification, are permitted provided that the following conditions are met: +// +// 1. Redistributions of source code must retain the above copyright notice, this +// list of conditions and the following disclaimer. +// +// 2. Redistributions in binary form must reproduce the above copyright notice, +// this list of conditions and the following disclaimer in the documentation +// and/or other materials provided with the distribution. +// +// 3. Neither the name of the copyright holder nor the names of its +// contributors may be used to endorse or promote products derived from +// this software without specific prior written permission. +// +// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +// DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +// FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +// DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +// SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +// CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +// OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +#include "attention_cuda.cuh" +#include +#include +#include "c10/core/MemoryFormat.h" + +#include +#include +#include +#include +#include +#include + +#include +#include +#include + +#include "cudamacro.h" +#include "attention_cuda_utils.cuh" + +#include +#include +#include + +#define THREADS (64) + +#define MAX_LOCAL_ARR_LEN (16) + +namespace attention_kernels +{ + + // scatter-direction dispatcher, defined in attention_cuda_bwd_upsample.cu; + // called by s2_attention_bwd_dkvq_cuda when nlon_out % nlon_in == 0. + void s2_attn_bwd_upsample_dispatch(int batch_size, size_t nchans_in, size_t nchans_out, int64_t nlon_in, + int64_t nlat_in, int64_t nlat_out, int64_t nlon_out, torch::Tensor kxP, + torch::Tensor vxP, torch::Tensor qyP, torch::Tensor dyP, + torch::Tensor psi_row_off, torch::Tensor psi_col_idx, torch::Tensor quad_weights, + torch::Tensor dkxP, torch::Tensor dvxP, torch::Tensor dqyP); + + // ===================================================================================== + // Backward with PASS-2 RECOMPUTE ELIMINATION, exposed as ``backward_coop``. + // The baseline backward gathers K/V twice: pass 1 computes qdotk/gdotv + dqy stats; + // pass 2 RE-gathers K/V to recompute qdotk/gdotv purely to reproduce two scalars per + // neighbor, then scatters sh_qy/sh_dy into dk/dv. Here pass 1 caches qdotk/gdotv in + // shared and pass 2 reads them back -> pass 2 does NO K/V gather (attacks the + // long_scoreboard #1 stall, which pass 2 is half of). Rows with rlen > CACHE_CAP fall + // back to the recompute path. Nested ``coop`` namespace -> distinct symbols (no ODR + // clash with the baseline). Kept until validated. + // ===================================================================================== + namespace coop + { + + // per-warp shared cache capacity for qdotk/gdotv (floats each). Covers the polar + // rows of the target grids (rlen<=712). Longer rows fall back to recompute. + static constexpr int CACHE_CAP = 768; + + // BEGIN backward kernels and functions + + // called with (blockDim.x=32 and blockDim.y>1, BDIM=blockDim.x*blockDim.y) + // + // STORAGE_T is the global-memory element type of the INPUTS (kx/vx/qy/dy): + // float4 for the fp32 vectorized path, float / c10::Half / c10::BFloat16 for + // the scalar path. COMPUTE_T (float4 or float) is the arithmetic type and the + // type of the gradient OUTPUTS dkx/dvx/dqy — these stay fp32 because dkx/dvx + // are atomically scatter-accumulated across overlapping neighborhoods (fp16 + // atomics would lose precision / aren't well supported). The wrapper allocates + // dkx/dvx/dqy in fp32 and casts to the input dtype at the end. vload widens + // STORAGE_T inputs to COMPUTE_T at the load site. + template + __global__ __launch_bounds__(BDIM_X) void s2_attn_bwd_generic_vec_k( + int nchans_in, // no. of elements along channel dim + int nchans_out, // no. of elements along channel dim + int nlat_in, int nlon_in, int nlat_out, int nlon_out, + const STORAGE_T *__restrict__ kx, // [batch][nlat_in][nlon_in][nchan_in] + const STORAGE_T *__restrict__ vx, // [batch][nlat_in][nlon_in][nchan_out] + const STORAGE_T *__restrict__ qy, // [batch][nlat_out][nlon_out][nchan_in] + const STORAGE_T *__restrict__ dy, // [batch][nlat_out][nlon_out][nchan_out] + const int32_t *__restrict__ row_idx, const int64_t *__restrict__ row_off, + const int64_t *__restrict__ col_idx, const float *__restrict__ quad_weights, + typename vec_traits::compute_t *__restrict__ dkx, // [batch][nlat_in][nlon_in][nchan_in] + typename vec_traits::compute_t *__restrict__ dvx, // [batch][nlat_in][nlon_in][nchan_out] + typename vec_traits::compute_t *__restrict__ dqy) + { // [batch][nlat_out][nlon_out][nchan_in] + using COMPUTE_T = typename vec_traits::compute_t; + + extern __shared__ __align__(sizeof(float4)) float shext[]; + + // for dqy + COMPUTE_T *sh_alpha_k__ = reinterpret_cast(shext) + threadIdx.y * (nchans_in * 4 + nchans_out); + COMPUTE_T *sh_alpha_vw_ = sh_alpha_k__ + nchans_in; + COMPUTE_T *sh_alpha_kvw = sh_alpha_vw_ + nchans_in; + + COMPUTE_T *sh_dy = sh_alpha_kvw + nchans_in; + COMPUTE_T *sh_qy = sh_dy + nchans_out; + // sh_alpha_k__[nchan_in], sh_alpha_vw_[nchan_in], sh_alpha_kvw[nchan_in] + // sh_dy[nchan_out], sh_qy[nchan_in] + + const int batch = blockIdx.y; + + const uint64_t wid = uint64_t(blockIdx.x) * blockDim.y + threadIdx.y; + if (wid >= uint64_t(nlat_out) * nlon_out) { return; } + + const int tidx = threadIdx.x; + + // use permuted rows + const int h = wid / nlon_out; + const int wo = wid - (h * nlon_out); + const int ho = row_idx[h]; + + // one output lon step corresponds to pscale input lon steps (requires nlon_in % nlon_out == 0) + const int pscale = nlon_in / nlon_out; + + // offset input tensors + kx += int64_t(batch) * nlat_in * nlon_in * nchans_in; + qy += int64_t(batch) * nlat_out * nlon_out * nchans_in + int64_t(ho) * nlon_out * nchans_in + + int64_t(wo) * nchans_in; + + vx += int64_t(batch) * nlat_in * nlon_in * nchans_out; + dy += int64_t(batch) * nlat_out * nlon_out * nchans_out + int64_t(ho) * nlon_out * nchans_out + + int64_t(wo) * nchans_out; + + // offset output tensors + dkx += int64_t(batch) * nlat_in * nlon_in * nchans_in; + dvx += int64_t(batch) * nlat_in * nlon_in * nchans_out; + dqy += int64_t(batch) * nlat_out * nlon_out * nchans_in + int64_t(ho) * nlon_out * nchans_in + + int64_t(wo) * nchans_in; + + // zero/init shared memory + for (int chan = tidx; chan < nchans_in; chan += WARP_SIZE) { + sh_alpha_k__[chan] = __vset(0.0f); + sh_alpha_vw_[chan] = __vset(0.0f); + sh_alpha_kvw[chan] = __vset(0.0f); + + sh_qy[chan] = vload(qy, chan); + } + for (int chan = tidx; chan < nchans_out; chan += WARP_SIZE) { sh_dy[chan] = vload(dy, chan); } + +#if __CUDA_ARCH__ < 900 + // for architectures < 9.0, sh_dy and sh_qy will be read + // as individual floats at the end of the kernel, which + // breaks the assumption that each COMPUTE_T location is + // written to and read by the same thread throughout the + // kernel, in the case COMPUTE_T==float4 + if constexpr (std::is_same::value) { __syncwarp(); } +#endif + + // for dkx, dvx, dqy + float alpha_sum = 0.0f; + float qdotk_max = -FLT_MAX; + + // for dkx + float integral = 0.0f; + + const int64_t rbeg = row_off[ho]; + const int64_t rend = row_off[ho + 1]; + + col_idx += rbeg; + + const int rlen = rend - rbeg; + + // accumulate alpha_sum, integral, and shared stats, + // along with a progressively computed qdotk_max. + for (int off = 0; off < rlen; off++) { + + const int64_t col = col_idx[off]; + + const int hi = col / nlon_in; + const int wi = col - (hi * nlon_in); + const int wi_wo = wi + pscale * wo; + const int wip = wi_wo - (wi_wo / nlon_in) * nlon_in; + + const STORAGE_T *_kx = kx + int64_t(hi) * nlon_in * nchans_in + int64_t(wip) * nchans_in; + const STORAGE_T *_vx = vx + int64_t(hi) * nlon_in * nchans_out + int64_t(wip) * nchans_out; + + COMPUTE_T qdotk_v = __vset(0.0f); + COMPUTE_T gdotv_v = __vset(0.0f); + + for (int chan = tidx; chan < nchans_in; chan += WARP_SIZE) { + qdotk_v = __vadd(qdotk_v, __vmul(sh_qy[chan], vload(_kx, chan))); + } + for (int chan = tidx; chan < nchans_out; chan += WARP_SIZE) { + gdotv_v = __vadd(gdotv_v, __vmul(sh_dy[chan], vload(_vx, chan))); + } + + const float qdotk = __warp_sum(__vred(qdotk_v)); + const float gdotv = __warp_sum(__vred(gdotv_v)); + + const float qdotk_max_tmp = max(qdotk_max, qdotk); + const float alpha_inz = expf(qdotk - qdotk_max_tmp) * quad_weights[hi]; + const float max_correction = expf(qdotk_max - qdotk_max_tmp); + alpha_sum = alpha_sum * max_correction + alpha_inz; + + integral = integral * max_correction + alpha_inz * gdotv; + + const float ainz_gdotv = alpha_inz * gdotv; + + for (int chan = tidx; chan < nchans_in; chan += WARP_SIZE) { + + const COMPUTE_T kxval = vload(_kx, chan); + + sh_alpha_k__[chan] = __vadd(__vscale(max_correction, sh_alpha_k__[chan]), __vscale(alpha_inz, kxval)); + sh_alpha_vw_[chan] + = __vadd(__vscale(max_correction, sh_alpha_vw_[chan]), __vset(ainz_gdotv)); + sh_alpha_kvw[chan] + = __vadd(__vscale(max_correction, sh_alpha_kvw[chan]), __vscale(ainz_gdotv, kxval)); + } + qdotk_max = qdotk_max_tmp; + } + + const float alpha_sum_inv = 1.0f / alpha_sum; + + integral *= alpha_sum_inv; + + // Write dqy (fp32 output) + for (int chan = tidx; chan < nchans_in; chan += WARP_SIZE) { + + dqy[chan] = __vscale( + alpha_sum_inv * alpha_sum_inv, + __vsub(__vscale(alpha_sum, sh_alpha_kvw[chan]), __vmul(sh_alpha_vw_[chan], sh_alpha_k__[chan]))); + } + + // accumulate gradients for k and v + for (int off = 0; off < rlen; off++) { + + const int64_t col = col_idx[off]; + const int hi = col / nlon_in; + const int wi = col - (hi * nlon_in); + const int wi_wo = wi + pscale * wo; + const int wip = wi_wo - (wi_wo / nlon_in) * nlon_in; + + const STORAGE_T *_kx = kx + int64_t(hi) * nlon_in * nchans_in + int64_t(wip) * nchans_in; + const STORAGE_T *_vx = vx + int64_t(hi) * nlon_in * nchans_out + int64_t(wip) * nchans_out; + + COMPUTE_T qdotk_v = __vset(0.0f); + COMPUTE_T gdotv_v = __vset(0.0f); + + for (int chan = tidx; chan < nchans_in; chan += WARP_SIZE) { + qdotk_v = __vadd(qdotk_v, __vmul(sh_qy[chan], vload(_kx, chan))); + } + for (int chan = tidx; chan < nchans_out; chan += WARP_SIZE) { + gdotv_v = __vadd(gdotv_v, __vmul(sh_dy[chan], vload(_vx, chan))); + } + + const float qdotk = __warp_sum(__vred(qdotk_v)); + const float gdotv = __warp_sum(__vred(gdotv_v)); + + const float alpha_inz = expf(qdotk - qdotk_max) * quad_weights[hi]; + + // _dkx / _dvx are COMPUTE_T (fp32) gradient buffers, accumulated atomically. + COMPUTE_T *_dkx = dkx + int64_t(hi) * nlon_in * nchans_in + int64_t(wip) * nchans_in; + COMPUTE_T *_dvx = dvx + int64_t(hi) * nlon_in * nchans_out + int64_t(wip) * nchans_out; + + const float alpha_mul = alpha_inz * alpha_sum_inv; + + const float scale_fact_qy = (gdotv - integral) * alpha_mul; + const float scale_fact_dy = alpha_mul; + + // float4, 128-bit atomics are only supported by devices of compute + // capability 9.x+, so on older devices we resort to 32-bit atomics + +#if __CUDA_ARCH__ < 900 + // to use 32-bit operations on consecutve addresses + float *sh_qy_scl = reinterpret_cast(sh_qy); + float *sh_dy_scl = reinterpret_cast(sh_dy); + + float *_dkx_scl = reinterpret_cast(_dkx); + float *_dvx_scl = reinterpret_cast(_dvx); + + constexpr int VEC_SIZE = sizeof(COMPUTE_T) / sizeof(float); + + // 32-bit, consecutive atomics to glmem; + // strided atomics results in a severe slowdown + for (int chan = tidx; chan < nchans_in * VEC_SIZE; chan += WARP_SIZE) { + atomicAdd(_dkx_scl + chan, scale_fact_qy * sh_qy_scl[chan]); + } + for (int chan = tidx; chan < nchans_out * VEC_SIZE; chan += WARP_SIZE) { + atomicAdd(_dvx_scl + chan, scale_fact_dy * sh_dy_scl[chan]); + } +#else + // 128-bit, consecutive atomics to glmem + for (int chan = tidx; chan < nchans_in; chan += WARP_SIZE) { + atomicAdd(_dkx + chan, __vscale(scale_fact_qy, sh_qy[chan])); + } + for (int chan = tidx; chan < nchans_out; chan += WARP_SIZE) { + atomicAdd(_dvx + chan, __vscale(scale_fact_dy, sh_dy[chan])); + } +#endif + } + + return; + } + + // called with either (BDIM_X=32 and BDIM_Y>1) || (2^K=BDIM_X > 32 and BDIM_Y=1) + template = nchan_in + typename STORAGE_T> + __global__ __launch_bounds__(BDIM_X *BDIM_Y) void s2_attn_bwd_special_vec_k( + int nchan_in, // no. of elements along channel dim + int nchan_out, // no. of elements along channel dim + int nlat_in, int nlon_in, int nlat_out, int nlon_out, + const STORAGE_T *__restrict__ kx, // [batch][nlat_in][nlon_in][nchan_in] + const STORAGE_T *__restrict__ vx, // [batch][nlat_in][nlon_in][nchan_out] + const STORAGE_T *__restrict__ qy, // [batch][nlat_out][nlon_out][nchan_in] + const STORAGE_T *__restrict__ dy, // [batch][nlat_out][nlon_out][nchan_out] + const int32_t *__restrict__ row_idx, const int64_t *__restrict__ row_off, + const int64_t *__restrict__ col_idx, const float *__restrict__ quad_weights, + typename vec_traits::compute_t *__restrict__ dkx, // [batch][nlat_in][nlon_in][nchan_in] + typename vec_traits::compute_t *__restrict__ dvx, // [batch][nlat_in][nlon_in][nchan_out] + typename vec_traits::compute_t *__restrict__ dqy) + { // [batch][nlat_out][nlon_out][nchan_in] + using COMPUTE_T = typename vec_traits::compute_t; + + static_assert(0 == (BDIM_X & (BDIM_X - 1))); + static_assert(0 == (BDIM_Y & (BDIM_Y - 1))); + static_assert((BDIM_X == 32 && BDIM_Y > 1) || (BDIM_X > 32 && BDIM_Y == 1)); + + constexpr int NLOC_M1 = NLOC - 1; + + const int tidx = threadIdx.x; + const int batch = blockIdx.y; + const uint64_t ctaid = uint64_t(blockIdx.x) * blockDim.y + threadIdx.y; + + if (ctaid >= uint64_t(nlat_out) * nlon_out) { return; } + + extern __shared__ __align__(sizeof(float4)) float shext[]; + + // sh_dy[nchan_out], sh_qy[nchan_in] + COMPUTE_T *sh_dy = reinterpret_cast(shext) + threadIdx.y * (nchan_in + nchan_out); // + tidx; + COMPUTE_T *sh_qy = sh_dy + nchan_out + tidx; + + if constexpr (CHOUT_AS_IN) { sh_dy += tidx; } + + // per-warp qdotk/gdotv cache (float), placed after all warps' sh_dy/sh_qy region. + constexpr int VECF = sizeof(COMPUTE_T) / sizeof(float); + float *qdc = shext + BDIM_Y * (nchan_in + nchan_out) * VECF + threadIdx.y * (2 * CACHE_CAP); + float *gdc = qdc + CACHE_CAP; + + // for dqy + COMPUTE_T loc_k__[NLOC]; + COMPUTE_T loc_vw_[NLOC]; + COMPUTE_T loc_kvw[NLOC]; + + // use permuted rows + const int h = ctaid / nlon_out; + const int wo = ctaid - (h * nlon_out); + const int ho = row_idx[h]; + + // one output lon step corresponds to pscale input lon steps (requires nlon_in % nlon_out == 0) + const int pscale = nlon_in / nlon_out; + + // offset input tensors + kx += int64_t(batch) * nlat_in * nlon_in * nchan_in + tidx; + qy += int64_t(batch) * nlat_out * nlon_out * nchan_in + int64_t(ho) * nlon_out * nchan_in + + int64_t(wo) * nchan_in + tidx; + + vx += int64_t(batch) * nlat_in * nlon_in * nchan_out; // + tidx; + dy += int64_t(batch) * nlat_out * nlon_out * nchan_out + int64_t(ho) * nlon_out * nchan_out + + int64_t(wo) * nchan_out; // + tidx; + if constexpr (CHOUT_AS_IN) { + vx += tidx; + dy += tidx; + } + + // offset output tensors + dkx += int64_t(batch) * nlat_in * nlon_in * nchan_in + tidx; + dvx += int64_t(batch) * nlat_in * nlon_in * nchan_out; // + tidx; + if constexpr (CHOUT_AS_IN) { dvx += tidx; } + dqy += int64_t(batch) * nlat_out * nlon_out * nchan_in + int64_t(ho) * nlon_out * nchan_in + + int64_t(wo) * nchan_in + tidx; + +#pragma unroll + for (int i = 0; i < NLOC; i++) { + loc_k__[i] = __vset(0.0f); + loc_vw_[i] = __vset(0.0f); + loc_kvw[i] = __vset(0.0f); + } + +#pragma unroll + for (int i = 0; i < NLOC_M1; i++) { sh_qy[i * BDIM_X] = vload(qy, i * BDIM_X); } + if (NLOC_M1 * BDIM_X + tidx < nchan_in) { sh_qy[NLOC_M1 * BDIM_X] = vload(qy, NLOC_M1 * BDIM_X); } + + if constexpr (CHOUT_AS_IN) { +#pragma unroll + for (int i = 0; i < NLOC_M1; i++) { sh_dy[i * BDIM_X] = vload(dy, i * BDIM_X); } + if (NLOC_M1 * BDIM_X + tidx < nchan_out) { sh_dy[NLOC_M1 * BDIM_X] = vload(dy, NLOC_M1 * BDIM_X); } + } else { + for (int chan = tidx; chan < nchan_out; chan += BDIM_X) { sh_dy[chan] = vload(dy, chan); } + } + +#if __CUDA_ARCH__ < 900 + // for architectures < 9.0, sh_dy and sh_qy will be read + // as individual floats at the end of the kernel, which + // breaks the assumption that each COMPUTE_T location is + // written to and read by the same thread throughout the + // kernel, in the case COMPUTE_T==float4 + if constexpr (std::is_same::value) { + if constexpr (BDIM_X == 32) { + __syncwarp(); + } else { + __syncthreads(); + } + } +#endif + + // for dkx, dvx, dqy + float alpha_sum = 0.0f; + float qdotk_max = -FLT_MAX; + + // for dkx + float integral = 0.0f; + + const int64_t rbeg = row_off[ho]; + const int64_t rend = row_off[ho + 1]; + + col_idx += rbeg; + + const int rlen = rend - rbeg; + // cache pass-1 qdotk/gdotv iff the row fits; long rows fall back to recompute. + const bool use_cache = (rlen <= CACHE_CAP); + + // accumulate alpha_sum, integral, and shared stats, + // along with a progressively computed qdotk_max. + for (int off = 0; off < rlen; off++) { + + const int64_t col = col_idx[off]; + + const int hi = col / nlon_in; + const int wi = col - (hi * nlon_in); + const int wi_wo = wi + pscale * wo; + const int wip = wi_wo - (wi_wo / nlon_in) * nlon_in; + + const STORAGE_T *_kx = kx + int64_t(hi) * nlon_in * nchan_in + int64_t(wip) * nchan_in; + const STORAGE_T *_vx = vx + int64_t(hi) * nlon_in * nchan_out + int64_t(wip) * nchan_out; + + COMPUTE_T qdotk_v = __vset(0.0f); + COMPUTE_T gdotv_v = __vset(0.0f); + +#pragma unroll + for (int i = 0; i < NLOC_M1; i++) { + qdotk_v = __vadd(qdotk_v, __vmul(sh_qy[i * BDIM_X], vload(_kx, i * BDIM_X))); + } + if (NLOC_M1 * BDIM_X + tidx < nchan_in) { + qdotk_v = __vadd(qdotk_v, __vmul(sh_qy[NLOC_M1 * BDIM_X], vload(_kx, NLOC_M1 * BDIM_X))); + } + if constexpr (CHOUT_AS_IN) { +#pragma unroll + for (int i = 0; i < NLOC_M1; i++) { + gdotv_v = __vadd(gdotv_v, __vmul(sh_dy[i * BDIM_X], vload(_vx, i * BDIM_X))); + } + if (NLOC_M1 * BDIM_X + tidx < nchan_out) { + gdotv_v = __vadd(gdotv_v, __vmul(sh_dy[NLOC_M1 * BDIM_X], vload(_vx, NLOC_M1 * BDIM_X))); + } + } else { + for (int chan = tidx; chan < nchan_out; chan += BDIM_X) { + gdotv_v = __vadd(gdotv_v, __vmul(sh_dy[chan], vload(_vx, chan))); + } + } + + float qdotk = __vred(qdotk_v); + float gdotv = __vred(gdotv_v); + + if constexpr (BDIM_X == 32) { + qdotk = __warp_sum(qdotk); + gdotv = __warp_sum(gdotv); + } else { + qdotk = __block_sum(qdotk); + gdotv = __block_sum(gdotv); + } + + // cache the two per-neighbor scalars so pass 2 need not re-gather K/V. + if (use_cache) { + qdc[off] = qdotk; + gdc[off] = gdotv; + } + + const float qdotk_max_tmp = max(qdotk_max, qdotk); + const float alpha_inz = expf(qdotk - qdotk_max_tmp) * quad_weights[hi]; + const float max_correction = expf(qdotk_max - qdotk_max_tmp); + + alpha_sum = alpha_sum * max_correction + alpha_inz; + integral = integral * max_correction + alpha_inz * gdotv; + + const float ainz_gdotv = alpha_inz * gdotv; + +#pragma unroll + for (int i = 0; i < NLOC_M1; i++) { + const COMPUTE_T kxval = vload(_kx, i * BDIM_X); + loc_k__[i] = __vadd(__vscale(max_correction, loc_k__[i]), __vscale(alpha_inz, kxval)); + loc_vw_[i] = __vadd(__vscale(max_correction, loc_vw_[i]), __vset(ainz_gdotv)); + loc_kvw[i] = __vadd(__vscale(max_correction, loc_kvw[i]), __vscale(ainz_gdotv, kxval)); + } + if (NLOC_M1 * BDIM_X + tidx < nchan_in) { + const COMPUTE_T kxval = vload(_kx, NLOC_M1 * BDIM_X); + loc_k__[NLOC_M1] = __vadd(__vscale(max_correction, loc_k__[NLOC_M1]), __vscale(alpha_inz, kxval)); + loc_vw_[NLOC_M1] = __vadd(__vscale(max_correction, loc_vw_[NLOC_M1]), __vset(ainz_gdotv)); + loc_kvw[NLOC_M1] = __vadd(__vscale(max_correction, loc_kvw[NLOC_M1]), __vscale(ainz_gdotv, kxval)); + } + + qdotk_max = qdotk_max_tmp; + } + + const float alpha_sum_inv = 1.0f / alpha_sum; + + integral *= alpha_sum_inv; + + // Write dqy + const float alpha_sum_inv_sq = alpha_sum_inv * alpha_sum_inv; + +#pragma unroll + for (int i = 0; i < NLOC_M1; i++) { + dqy[i * BDIM_X] = __vscale(alpha_sum_inv_sq, + __vsub(__vscale(alpha_sum, loc_kvw[i]), __vmul(loc_vw_[i], loc_k__[i]))); + } + if (NLOC_M1 * BDIM_X + tidx < nchan_in) { + dqy[NLOC_M1 * BDIM_X] + = __vscale(alpha_sum_inv_sq, + __vsub(__vscale(alpha_sum, loc_kvw[NLOC_M1]), __vmul(loc_vw_[NLOC_M1], loc_k__[NLOC_M1]))); + } + + // make the pass-1 qdotk/gdotv cache writes visible before pass 2 reads them. + if (use_cache) { + if constexpr (BDIM_X == 32) { + __syncwarp(); + } else { + __syncthreads(); + } + } + + // accumulate gradients for k and v + for (int off = 0; off < rlen; off++) { + + const int64_t col = col_idx[off]; + + const int hi = col / nlon_in; + const int wi = col - (hi * nlon_in); + const int wi_wo = wi + pscale * wo; + const int wip = wi_wo - (wi_wo / nlon_in) * nlon_in; + + float qdotk, gdotv; + if (use_cache) { + // pass 2: reuse pass-1 scalars from shared — NO K/V gather. + qdotk = qdc[off]; + gdotv = gdc[off]; + } else { + const STORAGE_T *_kx = kx + int64_t(hi) * nlon_in * nchan_in + int64_t(wip) * nchan_in; + const STORAGE_T *_vx = vx + int64_t(hi) * nlon_in * nchan_out + int64_t(wip) * nchan_out; + + COMPUTE_T qdotk_v = __vset(0.0f); + COMPUTE_T gdotv_v = __vset(0.0f); + +#pragma unroll + for (int i = 0; i < NLOC_M1; i++) { + qdotk_v = __vadd(qdotk_v, __vmul(sh_qy[i * BDIM_X], vload(_kx, i * BDIM_X))); + } + if (NLOC_M1 * BDIM_X + tidx < nchan_in) { + qdotk_v = __vadd(qdotk_v, __vmul(sh_qy[NLOC_M1 * BDIM_X], vload(_kx, NLOC_M1 * BDIM_X))); + } + if constexpr (CHOUT_AS_IN) { +#pragma unroll + for (int i = 0; i < NLOC_M1; i++) { + gdotv_v = __vadd(gdotv_v, __vmul(sh_dy[i * BDIM_X], vload(_vx, i * BDIM_X))); + } + if (NLOC_M1 * BDIM_X + tidx < nchan_out) { + gdotv_v = __vadd(gdotv_v, __vmul(sh_dy[NLOC_M1 * BDIM_X], vload(_vx, NLOC_M1 * BDIM_X))); + } + } else { + for (int chan = tidx; chan < nchan_out; chan += BDIM_X) { + gdotv_v = __vadd(gdotv_v, __vmul(sh_dy[chan], vload(_vx, chan))); + } + } + + qdotk = __vred(qdotk_v); + gdotv = __vred(gdotv_v); + + if constexpr (BDIM_X == 32) { + qdotk = __warp_sum(qdotk); + gdotv = __warp_sum(gdotv); + } else { + qdotk = __block_sum(qdotk); + gdotv = __block_sum(gdotv); + } + } + + const float alpha_inz = expf(qdotk - qdotk_max) * quad_weights[hi]; + + COMPUTE_T *_dkx = dkx + int64_t(hi) * nlon_in * nchan_in + int64_t(wip) * nchan_in; + COMPUTE_T *_dvx = dvx + int64_t(hi) * nlon_in * nchan_out + int64_t(wip) * nchan_out; + + const float alpha_mul = alpha_inz * alpha_sum_inv; + + const float scale_fact_qy = (gdotv - integral) * alpha_mul; + const float scale_fact_dy = alpha_mul; + + // float4, 128-bit atomics are only supported by devices of compute + // capability 9.x+, so on older devices we resort to 32-bit atomics + +#if __CUDA_ARCH__ < 900 + constexpr int VEC_SIZE = sizeof(COMPUTE_T) / sizeof(float); + + // making the loop count known at compile time doesn't seem + // to make any difference here so let's keep this (much) + // simpler version + float *sh_qy_scl = reinterpret_cast(sh_qy); + float *sh_dy_scl = reinterpret_cast(sh_dy); + + float *_dkx_scl = reinterpret_cast(_dkx); + float *_dvx_scl = reinterpret_cast(_dvx); + + sh_qy_scl -= tidx * VEC_SIZE; + _dkx_scl -= tidx * VEC_SIZE; + if constexpr (CHOUT_AS_IN) { + sh_dy_scl -= tidx * VEC_SIZE; + _dvx_scl -= tidx * VEC_SIZE; + } + + // 32-bit, consecutive atomics to glmem + // strided atomics results in a severe slowdown + for (int chan = tidx; chan < nchan_in * VEC_SIZE; chan += BDIM_X) { + atomicAdd(_dkx_scl + chan, scale_fact_qy * sh_qy_scl[chan]); + } + for (int chan = tidx; chan < nchan_out * VEC_SIZE; chan += BDIM_X) { + atomicAdd(_dvx_scl + chan, scale_fact_dy * sh_dy_scl[chan]); + } +#else +#pragma unroll + for (int i = 0; i < NLOC_M1; i++) { + atomicAdd(_dkx + i * BDIM_X, __vscale(scale_fact_qy, sh_qy[i * BDIM_X])); + } + if (NLOC_M1 * BDIM_X + tidx < nchan_in) { + atomicAdd(_dkx + NLOC_M1 * BDIM_X, __vscale(scale_fact_qy, sh_qy[NLOC_M1 * BDIM_X])); + } + if constexpr (CHOUT_AS_IN) { +#pragma unroll + for (int i = 0; i < NLOC_M1; i++) { + atomicAdd(_dvx + i * BDIM_X, __vscale(scale_fact_dy, sh_dy[i * BDIM_X])); + } + if (NLOC_M1 * BDIM_X + tidx < nchan_out) { + atomicAdd(_dvx + NLOC_M1 * BDIM_X, __vscale(scale_fact_dy, sh_dy[NLOC_M1 * BDIM_X])); + } + } else { + for (int chan = tidx; chan < nchan_out; chan += BDIM_X) { + atomicAdd(_dvx + chan, __vscale(scale_fact_dy, sh_dy[chan])); + } + } +#endif + } + + return; + } + + template + void launch_gen_attn_bwd(int batch_size, int nchans_in, int nchans_out, int nlat_in, int nlon_in, int nlat_out, + int nlon_out, STORAGE_T *_kxp, STORAGE_T *_vxp, STORAGE_T *_qyp, STORAGE_T *_dyp, + int32_t *_row_idx, int64_t *_row_off, int64_t *_col_idx, float *_quad_weights, + typename vec_traits::compute_t *_dkxp, + typename vec_traits::compute_t *_dvxp, + typename vec_traits::compute_t *_dqyp, cudaStream_t stream) + { + + dim3 block(WARP_SIZE, THREADS / WARP_SIZE); + dim3 grid(DIV_UP(nlat_out * nlon_out, block.y), batch_size); + + // shared memory holds compute-type (COMPUTE_T) data, not STORAGE_T. 5 arrays per warp. + size_t shsize = sizeof(typename vec_traits::compute_t) * (nchans_in * 4 + nchans_out) * block.y; + + s2_attn_bwd_generic_vec_k<<>>( + nchans_in, nchans_out, nlat_in, nlon_in, nlat_out, nlon_out, _kxp, _vxp, _qyp, _dyp, _row_idx, _row_off, + _col_idx, _quad_weights, _dkxp, _dvxp, _dqyp); + CHECK_ERROR("s2_attn_bwd_generic_vec_k"); + + return; + } + + template + void launch_spc_attn_bwd(int nloc, // "BDIM_X*nloc" >= nchans_out + int batch_size, int nchans_in, int nchans_out, int nlat_in, int nlon_in, int nlat_out, + int nlon_out, STORAGE_T *_kxp, STORAGE_T *_vxp, STORAGE_T *_qyp, STORAGE_T *_dyp, + int32_t *_row_idx, int64_t *_row_off, int64_t *_col_idx, float *_quad_weights, + typename vec_traits::compute_t *_dkxp, + typename vec_traits::compute_t *_dvxp, + typename vec_traits::compute_t *_dqyp, cudaStream_t stream) + { + + if (CUR_LOC_SIZE == nloc) { + + dim3 block(BDIM_X, BDIM_Y); + dim3 grid(DIV_UP(nlat_out * nlon_out, block.y), batch_size); + + // shared memory holds compute-type (COMPUTE_T) data, not STORAGE_T. + // 2 arrays per cta, block.y > 1 iif block.x==32 + size_t shsize = sizeof(typename vec_traits::compute_t) * (nchans_in + nchans_out) * block.y; + // + per-warp qdotk/gdotv cache (float) for the pass-2 recompute elimination. + shsize += size_t(block.y) * (2 * CACHE_CAP) * sizeof(float); + + // nloc determines the size of local arrays used to store + // temporary buffers loc_k__[], loc_vw_[] and loc_kvw[], + // of size nchans_in each; + // if nchans_out is >= BDIM_X*(nloc-1) and <= BDIM_X*nloc + // then we can use the same compile-time known loops used + // for input channels, with the exception of testing + // whether to execute the last iteration based on "nchans_out" + // instead of "nchans_in"; in this way as long as the + // difference between the number of input and output channels + // is <= BDIM_X we can use the faster path + if (nchans_out >= BDIM_X * (CUR_LOC_SIZE - 1) && nchans_out <= BDIM_X * CUR_LOC_SIZE) { + s2_attn_bwd_special_vec_k<<>>( + nchans_in, nchans_out, nlat_in, nlon_in, nlat_out, nlon_out, _kxp, _vxp, _qyp, _dyp, _row_idx, + _row_off, _col_idx, _quad_weights, _dkxp, _dvxp, _dqyp); + } else { + s2_attn_bwd_special_vec_k<<>>( + nchans_in, nchans_out, nlat_in, nlon_in, nlat_out, nlon_out, _kxp, _vxp, _qyp, _dyp, _row_idx, + _row_off, _col_idx, _quad_weights, _dkxp, _dvxp, _dqyp); + } + CHECK_ERROR("s2_attn_bwd_special_vec_k"); + + return; + } + if constexpr (CUR_LOC_SIZE < MAX_LOC_SIZE) { + launch_spc_attn_bwd( + nloc, batch_size, nchans_in, nchans_out, nlat_in, nlon_in, nlat_out, nlon_out, _kxp, _vxp, _qyp, + _dyp, _row_idx, _row_off, _col_idx, _quad_weights, _dkxp, _dvxp, _dqyp, stream); + } + return; + } + + // Picks the block size (BDIM_X) instance and launches the backward gather + // kernel for a given input storage vector type SV. Inputs are SV*; gradient + // outputs are COMPUTE_T* (fp32) — see s2_attn_bwd_*_vec_k. Backward uses the + // special kernel only up to BDIM_X=512 (1024 spills); larger falls to generic. + template + static void bwd_dispatch_bdimx(int bdimx, int nloc, int64_t batch_size, int64_t nchans_in, int64_t nchans_out, + int nlat_in, int64_t nlon_in, int64_t nlat_out, int64_t nlon_out, SV *_kxp, + SV *_vxp, SV *_qyp, SV *_dyp, int32_t *_row_idx, int64_t *_row_off, + int64_t *_col_idx, float *_quad_weights, typename vec_traits::compute_t *_dkxp, + typename vec_traits::compute_t *_dvxp, + typename vec_traits::compute_t *_dqyp, cudaStream_t stream) + { + switch (bdimx) { + case 32: + launch_spc_attn_bwd<32, 2, 1, MAX_LOC>(nloc, batch_size, nchans_in, nchans_out, nlat_in, nlon_in, + nlat_out, nlon_out, _kxp, _vxp, _qyp, _dyp, _row_idx, _row_off, + _col_idx, _quad_weights, _dkxp, _dvxp, _dqyp, stream); + break; + case 64: + launch_spc_attn_bwd<64, 1, MIN_LOC, MAX_LOC>( + nloc, batch_size, nchans_in, nchans_out, nlat_in, nlon_in, nlat_out, nlon_out, _kxp, _vxp, _qyp, + _dyp, _row_idx, _row_off, _col_idx, _quad_weights, _dkxp, _dvxp, _dqyp, stream); + break; + case 128: + launch_spc_attn_bwd<128, 1, MIN_LOC, MAX_LOC>( + nloc, batch_size, nchans_in, nchans_out, nlat_in, nlon_in, nlat_out, nlon_out, _kxp, _vxp, _qyp, + _dyp, _row_idx, _row_off, _col_idx, _quad_weights, _dkxp, _dvxp, _dqyp, stream); + break; + case 256: + launch_spc_attn_bwd<256, 1, MIN_LOC, MAX_LOC>( + nloc, batch_size, nchans_in, nchans_out, nlat_in, nlon_in, nlat_out, nlon_out, _kxp, _vxp, _qyp, + _dyp, _row_idx, _row_off, _col_idx, _quad_weights, _dkxp, _dvxp, _dqyp, stream); + break; + case 512: + launch_spc_attn_bwd<512, 1, MIN_LOC, MAX_LOC>( + nloc, batch_size, nchans_in, nchans_out, nlat_in, nlon_in, nlat_out, nlon_out, _kxp, _vxp, _qyp, + _dyp, _row_idx, _row_off, _col_idx, _quad_weights, _dkxp, _dvxp, _dqyp, stream); + break; + default: + launch_gen_attn_bwd(batch_size, nchans_in, nchans_out, nlat_in, nlon_in, nlat_out, nlon_out, _kxp, _vxp, + _qyp, _dyp, _row_idx, _row_off, _col_idx, _quad_weights, _dkxp, _dvxp, _dqyp, stream); + break; + } + } + + // Templated on the input storage element type (float / c10::Half / c10::BFloat16). + // Gradient outputs dkx/dvx/dqy are always fp32 (COMPUTE_T) — dkx/dvx are atomic + // scatter-accumulated, so they cannot be reduced precision; the wrapper casts + // them to the input dtype at the end. fp32 keeps the float4 vectorized path; + // fp16/bf16 (and unaligned fp32) take the scalar STORAGE_T path. + template + static void s2_attn_bwd_dispatch(int64_t batch_size, int64_t nchans_in, int64_t nchans_out, int64_t nlon_in, + int64_t nlat_out, int64_t nlon_out, at::Tensor kxP, at::Tensor vxP, + at::Tensor qyP, at::Tensor dyP, at::Tensor row_off, at::Tensor col_idx, + at::Tensor quad_weights, at::Tensor dkxP, at::Tensor dvxP, at::Tensor dqyP) + { + + static_assert(0 == (MAX_LOCAL_ARR_LEN & (MAX_LOCAL_ARR_LEN - 1))); + + // get stream + auto stream = at::cuda::getCurrentCUDAStream().stream(); + + // sort row indices (ho-s) in descending order + // based on (row_off[ho+1]-row_off[ho]) + at::Tensor row_idx = sortRows(nlat_out, row_off, stream); + + const int nlat_in = kxP.size(1); + + // smallest power of two "bdimx" (>=32) s.t. bdimx*MAX_LOCAL_ARR_LEN >= nchans_in + int bdimx; + bdimx = DIV_UP(nchans_in, MAX_LOCAL_ARR_LEN); + bdimx = max(bdimx, WARP_SIZE); + bdimx = next_pow2(bdimx); + + scalar_t *_kxp = reinterpret_cast(kxP.data_ptr()); + scalar_t *_vxp = reinterpret_cast(vxP.data_ptr()); + scalar_t *_qyp = reinterpret_cast(qyP.data_ptr()); + scalar_t *_dyp = reinterpret_cast(dyP.data_ptr()); + + // gradient outputs are fp32 + float *_dkxp = reinterpret_cast(dkxP.data_ptr()); + float *_dvxp = reinterpret_cast(dvxP.data_ptr()); + float *_dqyp = reinterpret_cast(dqyP.data_ptr()); + + int32_t *_row_idx = reinterpret_cast(row_idx.data_ptr()); + int64_t *_row_off = reinterpret_cast(row_off.data_ptr()); + int64_t *_col_idx = reinterpret_cast(col_idx.data_ptr()); + float *_quad_weights = reinterpret_cast(quad_weights.data_ptr()); + + constexpr int MIN_LOC_ARR_LEN = MAX_LOCAL_ARR_LEN / 2 + 1; + + if constexpr (std::is_same::value) { + // fp32: float4 vectorized when 16B-aligned + 4-divisible, else scalar. + constexpr int VEC_SIZE = sizeof(float4) / sizeof(float); // 4 + const bool use_vec = is_aligned<16>(_kxp) && is_aligned<16>(_vxp) && is_aligned<16>(_qyp) + && is_aligned<16>(_dyp) && is_aligned<16>(_dkxp) && is_aligned<16>(_dvxp) && is_aligned<16>(_dqyp) + && (nchans_in % VEC_SIZE) == 0 && (nchans_out % VEC_SIZE) == 0; + + if (use_vec) { + constexpr int MAX_VEC = MAX_LOCAL_ARR_LEN / VEC_SIZE; + constexpr int MIN_VEC = MAX_VEC / 2 + 1; + const int64_t nci = nchans_in / VEC_SIZE; + const int64_t nco = nchans_out / VEC_SIZE; + bwd_dispatch_bdimx( + bdimx, DIV_UP(nci, bdimx), batch_size, nci, nco, nlat_in, nlon_in, nlat_out, nlon_out, + reinterpret_cast(_kxp), reinterpret_cast(_vxp), + reinterpret_cast(_qyp), reinterpret_cast(_dyp), _row_idx, _row_off, + _col_idx, _quad_weights, reinterpret_cast(_dkxp), reinterpret_cast(_dvxp), + reinterpret_cast(_dqyp), stream); + } else { + bwd_dispatch_bdimx( + bdimx, DIV_UP(nchans_in, bdimx), batch_size, nchans_in, nchans_out, nlat_in, nlon_in, nlat_out, + nlon_out, _kxp, _vxp, _qyp, _dyp, _row_idx, _row_off, _col_idx, _quad_weights, _dkxp, _dvxp, + _dqyp, stream); + } + } else { + // fp16/bf16: scalar STORAGE_T inputs, fp32 outputs; fp32 compute/accumulation. + bwd_dispatch_bdimx( + bdimx, DIV_UP(nchans_in, bdimx), batch_size, nchans_in, nchans_out, nlat_in, nlon_in, nlat_out, + nlon_out, _kxp, _vxp, _qyp, _dyp, _row_idx, _row_off, _col_idx, _quad_weights, _dkxp, _dvxp, _dqyp, + stream); + } + + return; + } + + // END backward kernels and functions + + std::tuple + s2_attention_bwd_dkvq_cuda(at::Tensor kx, at::Tensor vx, at::Tensor qy, at::Tensor dy, at::Tensor quad_weights, + at::Tensor psi_col_idx, at::Tensor psi_row_off, int64_t nlon_in, int64_t nlat_out, + int64_t nlon_out) + { + + CHECK_CUDA_INPUT_TENSOR(kx); + CHECK_CUDA_INPUT_TENSOR(vx); + CHECK_CUDA_INPUT_TENSOR(qy); + CHECK_CUDA_INPUT_TENSOR(dy); + CHECK_CUDA_TENSOR(quad_weights); + CHECK_CUDA_TENSOR(psi_col_idx); + CHECK_CUDA_TENSOR(psi_row_off); + + // direction selection: gather (self / downsample) iff nlon_in is an integer + // multiple of nlon_out; scatter (upsample) iff nlon_out is an integer multiple + // of nlon_in. Self-attention satisfies both and routes through the gather path. + const bool downsample = (nlon_in % nlon_out == 0); + const bool upsample = (nlon_out % nlon_in == 0); + TORCH_CHECK(downsample || upsample, "either nlon_in (", nlon_in, + ") must be an integer multiple of nlon_out (", nlon_out, "), or vice versa"); + + // const size_t uo_num_channels = kx.size(1); + size_t nchans_in = qy.size(1); // or kx.size(1) + size_t nchans_out = vx.size(1); + + const int batch_size = kx.size(0); + const int64_t nlat_in = kx.size(2); + + // extract dtype + auto kx_type = kx.dtype(); // nchans_in + auto qy_type = qy.dtype(); + auto vx_type = vx.dtype(); // ncahn_out + auto dy_type = dy.dtype(); + + torch::Tensor dkx, dvx, dqy; + + // ATen dispatch over the input dtype. + // + // Gather (downsample / self) path: native storage. kx/vx/qy/dy keep their + // dtype; the kernel widens to fp32 at load (Tier B). Gradient buffers + // dkx/dvx/dqy are allocated fp32 because dkx/dvx are atomically + // scatter-accumulated (reduced-precision atomics would lose precision); + // they are cast back to the input dtype at the end. + // + // Scatter (upsample) path: native storage as well — inputs stay in their + // dtype (widened at load inside the dispatch), gradients stay fp32. + AT_DISPATCH_FLOATING_TYPES_AND2(at::kHalf, at::kBFloat16, qy.scalar_type(), "s2_attention_bwd_dkvq_cuda", [&] { + using storage_t = scalar_t; + + const auto f32_like + = [](const torch::Tensor &t) { return torch::zeros_like(t, t.options().dtype(torch::kFloat32)); }; + + if (downsample) { + // native-storage gather path + torch::Tensor kxP = kx; + torch::Tensor vxP = vx; + torch::Tensor qyP = qy; + torch::Tensor dyP = dy; + + // safer than is_contiguous(ChannelsLast), which fails for num_channels == 1 + bool kx_is_channels_last = kxP.strides()[1] == 1; + bool vx_is_channels_last = vxP.strides()[1] == 1; + bool qy_is_channels_last = qyP.strides()[1] == 1; + bool dy_is_channels_last = dyP.strides()[1] == 1; + + if (!kx_is_channels_last) { kxP = permute_4D_to0231(kxP); } + if (!vx_is_channels_last) { vxP = permute_4D_to0231(vxP); } + if (!qy_is_channels_last) { qyP = permute_4D_to0231(qyP); } + if (!dy_is_channels_last) { dyP = permute_4D_to0231(dyP); } + + // fp32 gradient buffers (same shape/layout as the native inputs) + torch::Tensor dkxP = f32_like(kxP); + torch::Tensor dvxP = f32_like(vxP); + torch::Tensor dqyP = f32_like(qyP); + + s2_attn_bwd_dispatch(batch_size, nchans_in, nchans_out, nlon_in, nlat_out, nlon_out, kxP, + vxP, qyP, dyP, psi_row_off, psi_col_idx, quad_weights, dkxP, dvxP, + dqyP); + + dkx = dkxP; + dvx = dvxP; + dqy = dqyP; + + if (!kx_is_channels_last) { dkx = permute_4D_to0312(dkx); } + if (!vx_is_channels_last) { dvx = permute_4D_to0312(dvx); } + if (!qy_is_channels_last) { dqy = permute_4D_to0312(dqy); } + } else { + // native-storage scatter (upsample) path + torch::Tensor kxP = kx; + torch::Tensor vxP = vx; + torch::Tensor qyP = qy; + torch::Tensor dyP = dy; + + bool kx_is_channels_last = kxP.strides()[1] == 1; + bool vx_is_channels_last = vxP.strides()[1] == 1; + bool qy_is_channels_last = qyP.strides()[1] == 1; + bool dy_is_channels_last = dyP.strides()[1] == 1; + + if (!kx_is_channels_last) { kxP = permute_4D_to0231(kxP); } + if (!vx_is_channels_last) { vxP = permute_4D_to0231(vxP); } + if (!qy_is_channels_last) { qyP = permute_4D_to0231(qyP); } + if (!dy_is_channels_last) { dyP = permute_4D_to0231(dyP); } + + torch::Tensor dkxP = f32_like(kxP); + torch::Tensor dvxP = f32_like(vxP); + torch::Tensor dqyP = f32_like(qyP); + + s2_attn_bwd_upsample_dispatch(batch_size, nchans_in, nchans_out, nlon_in, nlat_in, nlat_out, + nlon_out, kxP, vxP, qyP, dyP, psi_row_off, psi_col_idx, quad_weights, + dkxP, dvxP, dqyP); + + dkx = dkxP; + dvx = dvxP; + dqy = dqyP; + + if (!kx_is_channels_last) { dkx = permute_4D_to0312(dkx); } + if (!vx_is_channels_last) { dvx = permute_4D_to0312(dvx); } + if (!qy_is_channels_last) { dqy = permute_4D_to0312(dqy); } + } + }); + + C10_CUDA_KERNEL_LAUNCH_CHECK(); + + // convert precision back to starting dtype (no-op for fp32; narrows for fp16/bf16) + dkx = dkx.to(kx_type); + dvx = dvx.to(vx_type); + dqy = dqy.to(qy_type); + + return std::make_tuple(dkx, dvx, dqy); + } + + TORCH_LIBRARY_IMPL(attention_kernels, CUDA, m) { m.impl("backward_coop", &s2_attention_bwd_dkvq_cuda); } + + } // namespace coop + +} // namespace attention_kernels From d44393971dc7dfd5a3b8fb52c88cb240c8d503f2 Mon Sep 17 00:00:00 2001 From: Thorsten Kurth Date: Tue, 14 Jul 2026 07:59:04 -0700 Subject: [PATCH 07/14] adding split kernel --- .../kernels_cuda/attention_cuda_bwd_coop.cu | 115 +++++++++++------- 1 file changed, 70 insertions(+), 45 deletions(-) diff --git a/torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_bwd_coop.cu b/torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_bwd_coop.cu index 776bfbe5..6a7fffb2 100644 --- a/torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_bwd_coop.cu +++ b/torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_bwd_coop.cu @@ -336,7 +336,10 @@ namespace attention_kernels const int64_t *__restrict__ col_idx, const float *__restrict__ quad_weights, typename vec_traits::compute_t *__restrict__ dkx, // [batch][nlat_in][nlon_in][nchan_in] typename vec_traits::compute_t *__restrict__ dvx, // [batch][nlat_in][nlon_in][nchan_out] - typename vec_traits::compute_t *__restrict__ dqy) + typename vec_traits::compute_t *__restrict__ dqy, + // split params: process rows [row_ofs, row_ofs+nrows); cache_on gates the pass-2 + // recompute elimination (must be 0 when launched with the small no-cache shsize). + int64_t row_ofs, int64_t nrows, int cache_on) { // [batch][nlat_out][nlon_out][nchan_in] using COMPUTE_T = typename vec_traits::compute_t; @@ -350,7 +353,7 @@ namespace attention_kernels const int batch = blockIdx.y; const uint64_t ctaid = uint64_t(blockIdx.x) * blockDim.y + threadIdx.y; - if (ctaid >= uint64_t(nlat_out) * nlon_out) { return; } + if (ctaid >= uint64_t(nrows) * nlon_out) { return; } extern __shared__ __align__(sizeof(float4)) float shext[]; @@ -370,9 +373,10 @@ namespace attention_kernels COMPUTE_T loc_vw_[NLOC]; COMPUTE_T loc_kvw[NLOC]; - // use permuted rows - const int h = ctaid / nlon_out; - const int wo = ctaid - (h * nlon_out); + // use permuted rows (offset into the sorted row list for the split launch) + const int h_local = ctaid / nlon_out; + const int wo = ctaid - (h_local * nlon_out); + const int h = h_local + int(row_ofs); const int ho = row_idx[h]; // one output lon step corresponds to pscale input lon steps (requires nlon_in % nlon_out == 0) @@ -445,8 +449,8 @@ namespace attention_kernels col_idx += rbeg; const int rlen = rend - rbeg; - // cache pass-1 qdotk/gdotv iff the row fits; long rows fall back to recompute. - const bool use_cache = (rlen <= CACHE_CAP); + // cache pass-1 qdotk/gdotv iff enabled for this launch AND the row fits. + const bool use_cache = (cache_on != 0) && (rlen <= CACHE_CAP); // accumulate alpha_sum, integral, and shared stats, // along with a progressively computed qdotk_max. @@ -710,38 +714,30 @@ namespace attention_kernels int32_t *_row_idx, int64_t *_row_off, int64_t *_col_idx, float *_quad_weights, typename vec_traits::compute_t *_dkxp, typename vec_traits::compute_t *_dvxp, - typename vec_traits::compute_t *_dqyp, cudaStream_t stream) + typename vec_traits::compute_t *_dqyp, cudaStream_t stream, + int64_t row_ofs = 0, int64_t nrows = -1, int cache_on = 1) { if (CUR_LOC_SIZE == nloc) { + // process rows [row_ofs, row_ofs + grid_rows); nrows<0 means the whole grid. + const int64_t grid_rows = (nrows < 0) ? int64_t(nlat_out) : nrows; dim3 block(BDIM_X, BDIM_Y); - dim3 grid(DIV_UP(nlat_out * nlon_out, block.y), batch_size); + dim3 grid(DIV_UP(grid_rows * nlon_out, block.y), batch_size); // shared memory holds compute-type (COMPUTE_T) data, not STORAGE_T. - // 2 arrays per cta, block.y > 1 iif block.x==32 size_t shsize = sizeof(typename vec_traits::compute_t) * (nchans_in + nchans_out) * block.y; - // + per-warp qdotk/gdotv cache (float) for the pass-2 recompute elimination. - shsize += size_t(block.y) * (2 * CACHE_CAP) * sizeof(float); - - // nloc determines the size of local arrays used to store - // temporary buffers loc_k__[], loc_vw_[] and loc_kvw[], - // of size nchans_in each; - // if nchans_out is >= BDIM_X*(nloc-1) and <= BDIM_X*nloc - // then we can use the same compile-time known loops used - // for input channels, with the exception of testing - // whether to execute the last iteration based on "nchans_out" - // instead of "nchans_in"; in this way as long as the - // difference between the number of input and output channels - // is <= BDIM_X we can use the faster path + // + per-warp qdotk/gdotv cache (float) ONLY when caching is enabled for this launch. + if (cache_on) { shsize += size_t(block.y) * (2 * CACHE_CAP) * sizeof(float); } + if (nchans_out >= BDIM_X * (CUR_LOC_SIZE - 1) && nchans_out <= BDIM_X * CUR_LOC_SIZE) { s2_attn_bwd_special_vec_k<<>>( nchans_in, nchans_out, nlat_in, nlon_in, nlat_out, nlon_out, _kxp, _vxp, _qyp, _dyp, _row_idx, - _row_off, _col_idx, _quad_weights, _dkxp, _dvxp, _dqyp); + _row_off, _col_idx, _quad_weights, _dkxp, _dvxp, _dqyp, row_ofs, grid_rows, cache_on); } else { s2_attn_bwd_special_vec_k<<>>( nchans_in, nchans_out, nlat_in, nlon_in, nlat_out, nlon_out, _kxp, _vxp, _qyp, _dyp, _row_idx, - _row_off, _col_idx, _quad_weights, _dkxp, _dvxp, _dqyp); + _row_off, _col_idx, _quad_weights, _dkxp, _dvxp, _dqyp, row_ofs, grid_rows, cache_on); } CHECK_ERROR("s2_attn_bwd_special_vec_k"); @@ -749,8 +745,8 @@ namespace attention_kernels } if constexpr (CUR_LOC_SIZE < MAX_LOC_SIZE) { launch_spc_attn_bwd( - nloc, batch_size, nchans_in, nchans_out, nlat_in, nlon_in, nlat_out, nlon_out, _kxp, _vxp, _qyp, - _dyp, _row_idx, _row_off, _col_idx, _quad_weights, _dkxp, _dvxp, _dqyp, stream); + nloc, batch_size, nchans_in, nchans_out, nlat_in, nlon_in, nlat_out, nlon_out, _kxp, _vxp, _qyp, _dyp, + _row_idx, _row_off, _col_idx, _quad_weights, _dkxp, _dvxp, _dqyp, stream, row_ofs, nrows, cache_on); } return; } @@ -765,35 +761,42 @@ namespace attention_kernels SV *_vxp, SV *_qyp, SV *_dyp, int32_t *_row_idx, int64_t *_row_off, int64_t *_col_idx, float *_quad_weights, typename vec_traits::compute_t *_dkxp, typename vec_traits::compute_t *_dvxp, - typename vec_traits::compute_t *_dqyp, cudaStream_t stream) + typename vec_traits::compute_t *_dqyp, cudaStream_t stream, + int64_t row_ofs = 0, int64_t nrows = -1, int cache_on = 1) { switch (bdimx) { case 32: launch_spc_attn_bwd<32, 2, 1, MAX_LOC>(nloc, batch_size, nchans_in, nchans_out, nlat_in, nlon_in, nlat_out, nlon_out, _kxp, _vxp, _qyp, _dyp, _row_idx, _row_off, - _col_idx, _quad_weights, _dkxp, _dvxp, _dqyp, stream); + _col_idx, _quad_weights, _dkxp, _dvxp, _dqyp, stream, row_ofs, + nrows, cache_on); break; case 64: - launch_spc_attn_bwd<64, 1, MIN_LOC, MAX_LOC>( - nloc, batch_size, nchans_in, nchans_out, nlat_in, nlon_in, nlat_out, nlon_out, _kxp, _vxp, _qyp, - _dyp, _row_idx, _row_off, _col_idx, _quad_weights, _dkxp, _dvxp, _dqyp, stream); + launch_spc_attn_bwd<64, 1, MIN_LOC, MAX_LOC>(nloc, batch_size, nchans_in, nchans_out, nlat_in, nlon_in, + nlat_out, nlon_out, _kxp, _vxp, _qyp, _dyp, _row_idx, + _row_off, _col_idx, _quad_weights, _dkxp, _dvxp, _dqyp, + stream, row_ofs, nrows, cache_on); break; case 128: - launch_spc_attn_bwd<128, 1, MIN_LOC, MAX_LOC>( - nloc, batch_size, nchans_in, nchans_out, nlat_in, nlon_in, nlat_out, nlon_out, _kxp, _vxp, _qyp, - _dyp, _row_idx, _row_off, _col_idx, _quad_weights, _dkxp, _dvxp, _dqyp, stream); + launch_spc_attn_bwd<128, 1, MIN_LOC, MAX_LOC>(nloc, batch_size, nchans_in, nchans_out, nlat_in, nlon_in, + nlat_out, nlon_out, _kxp, _vxp, _qyp, _dyp, _row_idx, + _row_off, _col_idx, _quad_weights, _dkxp, _dvxp, _dqyp, + stream, row_ofs, nrows, cache_on); break; case 256: - launch_spc_attn_bwd<256, 1, MIN_LOC, MAX_LOC>( - nloc, batch_size, nchans_in, nchans_out, nlat_in, nlon_in, nlat_out, nlon_out, _kxp, _vxp, _qyp, - _dyp, _row_idx, _row_off, _col_idx, _quad_weights, _dkxp, _dvxp, _dqyp, stream); + launch_spc_attn_bwd<256, 1, MIN_LOC, MAX_LOC>(nloc, batch_size, nchans_in, nchans_out, nlat_in, nlon_in, + nlat_out, nlon_out, _kxp, _vxp, _qyp, _dyp, _row_idx, + _row_off, _col_idx, _quad_weights, _dkxp, _dvxp, _dqyp, + stream, row_ofs, nrows, cache_on); break; case 512: - launch_spc_attn_bwd<512, 1, MIN_LOC, MAX_LOC>( - nloc, batch_size, nchans_in, nchans_out, nlat_in, nlon_in, nlat_out, nlon_out, _kxp, _vxp, _qyp, - _dyp, _row_idx, _row_off, _col_idx, _quad_weights, _dkxp, _dvxp, _dqyp, stream); + launch_spc_attn_bwd<512, 1, MIN_LOC, MAX_LOC>(nloc, batch_size, nchans_in, nchans_out, nlat_in, nlon_in, + nlat_out, nlon_out, _kxp, _vxp, _qyp, _dyp, _row_idx, + _row_off, _col_idx, _quad_weights, _dkxp, _dvxp, _dqyp, + stream, row_ofs, nrows, cache_on); break; default: + // generic path: no caching, whole grid (split params ignored). launch_gen_attn_bwd(batch_size, nchans_in, nchans_out, nlat_in, nlon_in, nlat_out, nlon_out, _kxp, _vxp, _qyp, _dyp, _row_idx, _row_off, _col_idx, _quad_weights, _dkxp, _dvxp, _dqyp, stream); break; @@ -872,10 +875,32 @@ namespace attention_kernels } } else { // fp16/bf16: scalar STORAGE_T inputs, fp32 outputs; fp32 compute/accumulation. - bwd_dispatch_bdimx( - bdimx, DIV_UP(nchans_in, bdimx), batch_size, nchans_in, nchans_out, nlat_in, nlon_in, nlat_out, - nlon_out, _kxp, _vxp, _qyp, _dyp, _row_idx, _row_off, _col_idx, _quad_weights, _dkxp, _dvxp, _dqyp, - stream); + // ROW SPLIT: long/polar rows get the pass-2 recompute-elimination cache; cheap + // short rows run the pristine recompute path at full occupancy (no cache tax). + const int nloc_s = DIV_UP(nchans_in, bdimx); + const bool special = (bdimx <= 512); + int64_t n_long = 0, m0 = 0, m1 = 0; + if (special) { + // reuse the ring's relative 0.1*max_rlen threshold but drop the 1024 floor + // (split_len=1) so our polar rows (rlen<1024) are classified "long". + split_csr_rows(SPLIT_ROW_LENGTH_THRES, 1, nlat_out, _row_idx, _row_off, &n_long, &m0, &m1); + } + if (special && n_long > 0 && n_long < nlat_out) { + bwd_dispatch_bdimx( + bdimx, nloc_s, batch_size, nchans_in, nchans_out, nlat_in, nlon_in, nlat_out, nlon_out, _kxp, + _vxp, _qyp, _dyp, _row_idx, _row_off, _col_idx, _quad_weights, _dkxp, _dvxp, _dqyp, stream, + /*row_ofs=*/0, /*nrows=*/n_long, /*cache_on=*/1); + bwd_dispatch_bdimx( + bdimx, nloc_s, batch_size, nchans_in, nchans_out, nlat_in, nlon_in, nlat_out, nlon_out, _kxp, + _vxp, _qyp, _dyp, _row_idx, _row_off, _col_idx, _quad_weights, _dkxp, _dvxp, _dqyp, stream, + /*row_ofs=*/n_long, /*nrows=*/nlat_out - n_long, /*cache_on=*/0); + } else { + // no split: whole grid; cache on for the special path, off for generic (bdimx>512). + bwd_dispatch_bdimx( + bdimx, nloc_s, batch_size, nchans_in, nchans_out, nlat_in, nlon_in, nlat_out, nlon_out, _kxp, + _vxp, _qyp, _dyp, _row_idx, _row_off, _col_idx, _quad_weights, _dkxp, _dvxp, _dqyp, stream, 0, + -1, special ? 1 : 0); + } } return; From 372ffa31115bec6bc14398b5a378f72c1484f5b5 Mon Sep 17 00:00:00 2001 From: Thorsten Kurth Date: Tue, 14 Jul 2026 22:09:45 -0700 Subject: [PATCH 08/14] removing old kernels --- setup.py | 5 - .../optimized/attention_interface.cpp | 20 - .../kernels_cuda/attention_cuda_bwd.cu | 107 +- .../kernels_cuda/attention_cuda_bwd_coop.cu | 1047 ----------------- .../kernels_cuda/attention_cuda_bwd_fast.cu | 982 ---------------- .../kernels_cuda/attention_cuda_fwd_coop.cu | 321 ----- .../kernels_cuda/attention_cuda_fwd_fast.cu | 622 ---------- .../attention_cuda_fwd_wgmma_sm90.cu | 530 --------- .../kernels_cuda/attention_cuda_ptx.cuh | 177 --- .../kernels_cuda/attention_cuda_utils.cuh | 18 - 10 files changed, 75 insertions(+), 3754 deletions(-) delete mode 100644 torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_bwd_coop.cu delete mode 100644 torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_bwd_fast.cu delete mode 100644 torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_fwd_coop.cu delete mode 100644 torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_fwd_fast.cu delete mode 100644 torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_fwd_wgmma_sm90.cu delete mode 100644 torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_ptx.cuh diff --git a/setup.py b/setup.py index 637c8d08..4a1bcbf2 100644 --- a/setup.py +++ b/setup.py @@ -197,11 +197,6 @@ def get_ext_modules(): "torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_bwd_ring.cu", "torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_fwd_ring_upsample.cu", "torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_bwd_ring_upsample.cu", - "torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_fwd_wgmma_sm90.cu", - "torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_fwd_fast.cu", - "torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_bwd_fast.cu", - "torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_fwd_coop.cu", - "torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_bwd_coop.cu", ] ) ext_modules.append(CUDAExtension("torch_harmonics.attention._C", attention_sources, extra_compile_args=get_compile_args("attention"))) diff --git a/torch_harmonics/attention/optimized/attention_interface.cpp b/torch_harmonics/attention/optimized/attention_interface.cpp index d9236d50..e8d54ab1 100644 --- a/torch_harmonics/attention/optimized/attention_interface.cpp +++ b/torch_harmonics/attention/optimized/attention_interface.cpp @@ -83,29 +83,9 @@ namespace attention_kernels m.def("forward(Tensor kx, Tensor vx, Tensor qy, Tensor quad_weights, Tensor col_idx, Tensor row_off, int " "nlon_in, int nlat_out, int nlon_out) -> Tensor", {at::Tag::pt2_compliant_tag}); - // Standalone Hopper WGMMA (sm_90a) forward, for isolated benchmarking of the - // tensor-core neighborhood-attention path. Not routed by the module; call it - // directly via torch.ops.attention_kernels.forward_wgmma. Drop-in signature - // match for `forward` (row_idx is derived internally via sortRows). - m.def("forward_wgmma(Tensor kx, Tensor vx, Tensor qy, Tensor quad_weights, Tensor col_idx, Tensor row_off, " - "int nlon_in, int nlat_out, int nlon_out) -> Tensor"); m.def("backward(Tensor kx, Tensor vx, Tensor qy, Tensor dy, Tensor quad_weights, Tensor col_idx, Tensor " "row_off, int nlon_in, int nlat_out, int nlon_out) -> (Tensor, Tensor, Tensor)", {at::Tag::pt2_compliant_tag}); - // Optimized (FMA + __expf) gather forward/backward copies for A/B benchmarking - // vs the baseline forward/backward. Same signatures; call directly via - // torch.ops.attention_kernels.forward_fast / backward_fast. - m.def("forward_fast(Tensor kx, Tensor vx, Tensor qy, Tensor quad_weights, Tensor col_idx, Tensor row_off, int " - "nlon_in, int nlat_out, int nlon_out) -> Tensor"); - m.def("backward_fast(Tensor kx, Tensor vx, Tensor qy, Tensor dy, Tensor quad_weights, Tensor col_idx, Tensor " - "row_off, int nlon_in, int nlat_out, int nlon_out) -> (Tensor, Tensor, Tensor)"); - // Cross-wo cooperative-tiling forward (shared K/V arc reuse). Gather/self only, - // nchan_in==nchan_out. Same signature as forward; call via forward_coop. - m.def("forward_coop(Tensor kx, Tensor vx, Tensor qy, Tensor quad_weights, Tensor col_idx, Tensor row_off, int " - "nlon_in, int nlat_out, int nlon_out) -> Tensor"); - // Backward with pass-2 recompute elimination (caches qdotk/gdotv in shared). - m.def("backward_coop(Tensor kx, Tensor vx, Tensor qy, Tensor dy, Tensor quad_weights, Tensor col_idx, Tensor " - "row_off, int nlon_in, int nlat_out, int nlon_out) -> (Tensor, Tensor, Tensor)"); // ---- Ring-step variants for DistributedNeighborhoodAttentionS2 ---- // K/V are sharded along longitude across an azimuth process group; each diff --git a/torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_bwd.cu b/torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_bwd.cu index ad8e1b9a..7a086b00 100644 --- a/torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_bwd.cu +++ b/torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_bwd.cu @@ -66,6 +66,13 @@ namespace attention_kernels torch::Tensor psi_row_off, torch::Tensor psi_col_idx, torch::Tensor quad_weights, torch::Tensor dkxP, torch::Tensor dvxP, torch::Tensor dqyP); + // Per-warp shared cache capacity (floats each) for the special-kernel pass-2 recompute + // elimination: pass 1 caches qdotk/gdotv, pass 2 reads them back instead of re-gathering + // K/V (halves the backward gather, the #1 long_scoreboard stall). Rows with rlen > cap + // fall back to recompute. Arch-agnostic (plain shared mem + syncwarp). 768 covers the + // polar rows of ERA5-class grids; ~12 KB/block at BDIM_Y=2, well under the 48 KB default. + static constexpr int BWD_CACHE_CAP = 768; + // BEGIN backward kernels and functions // called with (blockDim.x=32 and blockDim.y>1, BDIM=blockDim.x*blockDim.y) @@ -317,8 +324,9 @@ namespace attention_kernels const float *__restrict__ quad_weights, typename vec_traits::compute_t *__restrict__ dkx, // [batch][nlat_in][nlon_in][nchan_in] typename vec_traits::compute_t *__restrict__ dvx, // [batch][nlat_in][nlon_in][nchan_out] - typename vec_traits::compute_t *__restrict__ dqy) - { // [batch][nlat_out][nlon_out][nchan_in] + typename vec_traits::compute_t *__restrict__ dqy, + int cache_on) // 1 iff the launcher allocated the qdotk/gdotv cache (fits in shared) + { // [batch][nlat_out][nlon_out][nchan_in] using COMPUTE_T = typename vec_traits::compute_t; static_assert(0 == (BDIM_X & (BDIM_X - 1))); @@ -341,6 +349,11 @@ namespace attention_kernels if constexpr (CHOUT_AS_IN) { sh_dy += tidx; } + // per-warp qdotk/gdotv cache (float), after all warps' sh_dy/sh_qy region. + constexpr int VECF = sizeof(COMPUTE_T) / sizeof(float); + float *qdc = shext + BDIM_Y * (nchan_in + nchan_out) * VECF + threadIdx.y * (2 * BWD_CACHE_CAP); + float *gdc = qdc + BWD_CACHE_CAP; + // for dqy COMPUTE_T loc_k__[NLOC]; COMPUTE_T loc_vw_[NLOC]; @@ -421,6 +434,8 @@ namespace attention_kernels col_idx += rbeg; const int rlen = rend - rbeg; + // cache pass-1 qdotk/gdotv iff the launcher allocated the cache AND the row fits. + const bool use_cache = (cache_on != 0) && (rlen <= BWD_CACHE_CAP); // accumulate alpha_sum, integral, and shared stats, // along with a progressively computed qdotk_max. @@ -471,6 +486,12 @@ namespace attention_kernels gdotv = __block_sum(gdotv); } + // cache the two per-neighbor scalars so pass 2 need not re-gather K/V. + if (use_cache) { + qdc[off] = qdotk; + gdc[off] = gdotv; + } + const float qdotk_max_tmp = max(qdotk_max, qdotk); const float alpha_inz = expf(qdotk - qdotk_max_tmp) * quad_weights[hi]; const float max_correction = expf(qdotk_max - qdotk_max_tmp); @@ -515,6 +536,15 @@ namespace attention_kernels __vsub(__vscale(alpha_sum, loc_kvw[NLOC_M1]), __vmul(loc_vw_[NLOC_M1], loc_k__[NLOC_M1]))); } + // make the pass-1 qdotk/gdotv cache writes visible before pass 2 reads them. + if (use_cache) { + if constexpr (BDIM_X == 32) { + __syncwarp(); + } else { + __syncthreads(); + } + } + // accumulate gradients for k and v for (int off = 0; off < rlen; off++) { @@ -525,42 +555,49 @@ namespace attention_kernels const int wi_wo = wi + pscale * wo; const int wip = wi_wo - (wi_wo / nlon_in) * nlon_in; - const STORAGE_T *_kx = kx + int64_t(hi) * nlon_in * nchan_in + int64_t(wip) * nchan_in; - const STORAGE_T *_vx = vx + int64_t(hi) * nlon_in * nchan_out + int64_t(wip) * nchan_out; + float qdotk, gdotv; + if (use_cache) { + // pass 2: reuse pass-1 scalars from shared — NO K/V gather. + qdotk = qdc[off]; + gdotv = gdc[off]; + } else { + const STORAGE_T *_kx = kx + int64_t(hi) * nlon_in * nchan_in + int64_t(wip) * nchan_in; + const STORAGE_T *_vx = vx + int64_t(hi) * nlon_in * nchan_out + int64_t(wip) * nchan_out; - COMPUTE_T qdotk_v = __vset(0.0f); - COMPUTE_T gdotv_v = __vset(0.0f); + COMPUTE_T qdotk_v = __vset(0.0f); + COMPUTE_T gdotv_v = __vset(0.0f); -#pragma unroll - for (int i = 0; i < NLOC_M1; i++) { - qdotk_v = __vadd(qdotk_v, __vmul(sh_qy[i * BDIM_X], vload(_kx, i * BDIM_X))); - } - if (NLOC_M1 * BDIM_X + tidx < nchan_in) { - qdotk_v = __vadd(qdotk_v, __vmul(sh_qy[NLOC_M1 * BDIM_X], vload(_kx, NLOC_M1 * BDIM_X))); - } - if constexpr (CHOUT_AS_IN) { #pragma unroll for (int i = 0; i < NLOC_M1; i++) { - gdotv_v = __vadd(gdotv_v, __vmul(sh_dy[i * BDIM_X], vload(_vx, i * BDIM_X))); + qdotk_v = __vadd(qdotk_v, __vmul(sh_qy[i * BDIM_X], vload(_kx, i * BDIM_X))); } - if (NLOC_M1 * BDIM_X + tidx < nchan_out) { - gdotv_v = __vadd(gdotv_v, __vmul(sh_dy[NLOC_M1 * BDIM_X], vload(_vx, NLOC_M1 * BDIM_X))); + if (NLOC_M1 * BDIM_X + tidx < nchan_in) { + qdotk_v = __vadd(qdotk_v, __vmul(sh_qy[NLOC_M1 * BDIM_X], vload(_kx, NLOC_M1 * BDIM_X))); } - } else { - for (int chan = tidx; chan < nchan_out; chan += BDIM_X) { - gdotv_v = __vadd(gdotv_v, __vmul(sh_dy[chan], vload(_vx, chan))); + if constexpr (CHOUT_AS_IN) { +#pragma unroll + for (int i = 0; i < NLOC_M1; i++) { + gdotv_v = __vadd(gdotv_v, __vmul(sh_dy[i * BDIM_X], vload(_vx, i * BDIM_X))); + } + if (NLOC_M1 * BDIM_X + tidx < nchan_out) { + gdotv_v = __vadd(gdotv_v, __vmul(sh_dy[NLOC_M1 * BDIM_X], vload(_vx, NLOC_M1 * BDIM_X))); + } + } else { + for (int chan = tidx; chan < nchan_out; chan += BDIM_X) { + gdotv_v = __vadd(gdotv_v, __vmul(sh_dy[chan], vload(_vx, chan))); + } } - } - float qdotk = __vred(qdotk_v); - float gdotv = __vred(gdotv_v); + qdotk = __vred(qdotk_v); + gdotv = __vred(gdotv_v); - if constexpr (BDIM_X == 32) { - qdotk = __warp_sum(qdotk); - gdotv = __warp_sum(gdotv); - } else { - qdotk = __block_sum(qdotk); - gdotv = __block_sum(gdotv); + if constexpr (BDIM_X == 32) { + qdotk = __warp_sum(qdotk); + gdotv = __warp_sum(gdotv); + } else { + qdotk = __block_sum(qdotk); + gdotv = __block_sum(gdotv); + } } const float alpha_inz = expf(qdotk - qdotk_max) * quad_weights[hi]; @@ -672,7 +709,13 @@ namespace attention_kernels // shared memory holds compute-type (COMPUTE_T) data, not STORAGE_T. // 2 arrays per cta, block.y > 1 iif block.x==32 - size_t shsize = sizeof(typename vec_traits::compute_t) * (nchans_in + nchans_out) * block.y; + const size_t base_sh = sizeof(typename vec_traits::compute_t) * (nchans_in + nchans_out) * block.y; + // pass-2 recompute-elimination cache (qdotk/gdotv, float). Enable only when the + // total fits the 48 KB default opt-out shared window; otherwise fall back to the + // recompute path (cache_on = 0) so huge-nchan launches never overflow. + const size_t cache_sh = size_t(block.y) * (2 * BWD_CACHE_CAP) * sizeof(float); + const int cache_on = (base_sh + cache_sh <= 48u * 1024u) ? 1 : 0; + const size_t shsize = base_sh + (cache_on ? cache_sh : 0); // nloc determines the size of local arrays used to store // temporary buffers loc_k__[], loc_vw_[] and loc_kvw[], @@ -687,11 +730,11 @@ namespace attention_kernels if (nchans_out >= BDIM_X * (CUR_LOC_SIZE - 1) && nchans_out <= BDIM_X * CUR_LOC_SIZE) { s2_attn_bwd_special_vec_k<<>>( nchans_in, nchans_out, nlat_in, nlon_in, nlat_out, nlon_out, _kxp, _vxp, _qyp, _dyp, _row_idx, - _row_off, _col_idx, _quad_weights, _dkxp, _dvxp, _dqyp); + _row_off, _col_idx, _quad_weights, _dkxp, _dvxp, _dqyp, cache_on); } else { s2_attn_bwd_special_vec_k<<>>( nchans_in, nchans_out, nlat_in, nlon_in, nlat_out, nlon_out, _kxp, _vxp, _qyp, _dyp, _row_idx, - _row_off, _col_idx, _quad_weights, _dkxp, _dvxp, _dqyp); + _row_off, _col_idx, _quad_weights, _dkxp, _dvxp, _dqyp, cache_on); } CHECK_ERROR("s2_attn_bwd_special_vec_k"); diff --git a/torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_bwd_coop.cu b/torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_bwd_coop.cu deleted file mode 100644 index 6a7fffb2..00000000 --- a/torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_bwd_coop.cu +++ /dev/null @@ -1,1047 +0,0 @@ -// coding=utf-8 -// -// SPDX-FileCopyrightText: Copyright (c) 2025 The torch-harmonics Authors. All rights reserved. -// SPDX-License-Identifier: BSD-3-Clause -// -// Redistribution and use in source and binary forms, with or without -// modification, are permitted provided that the following conditions are met: -// -// 1. Redistributions of source code must retain the above copyright notice, this -// list of conditions and the following disclaimer. -// -// 2. Redistributions in binary form must reproduce the above copyright notice, -// this list of conditions and the following disclaimer in the documentation -// and/or other materials provided with the distribution. -// -// 3. Neither the name of the copyright holder nor the names of its -// contributors may be used to endorse or promote products derived from -// this software without specific prior written permission. -// -// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" -// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE -// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE -// DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE -// FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL -// DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR -// SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER -// CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, -// OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE -// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. - -#include "attention_cuda.cuh" -#include -#include -#include "c10/core/MemoryFormat.h" - -#include -#include -#include -#include -#include -#include - -#include -#include -#include - -#include "cudamacro.h" -#include "attention_cuda_utils.cuh" - -#include -#include -#include - -#define THREADS (64) - -#define MAX_LOCAL_ARR_LEN (16) - -namespace attention_kernels -{ - - // scatter-direction dispatcher, defined in attention_cuda_bwd_upsample.cu; - // called by s2_attention_bwd_dkvq_cuda when nlon_out % nlon_in == 0. - void s2_attn_bwd_upsample_dispatch(int batch_size, size_t nchans_in, size_t nchans_out, int64_t nlon_in, - int64_t nlat_in, int64_t nlat_out, int64_t nlon_out, torch::Tensor kxP, - torch::Tensor vxP, torch::Tensor qyP, torch::Tensor dyP, - torch::Tensor psi_row_off, torch::Tensor psi_col_idx, torch::Tensor quad_weights, - torch::Tensor dkxP, torch::Tensor dvxP, torch::Tensor dqyP); - - // ===================================================================================== - // Backward with PASS-2 RECOMPUTE ELIMINATION, exposed as ``backward_coop``. - // The baseline backward gathers K/V twice: pass 1 computes qdotk/gdotv + dqy stats; - // pass 2 RE-gathers K/V to recompute qdotk/gdotv purely to reproduce two scalars per - // neighbor, then scatters sh_qy/sh_dy into dk/dv. Here pass 1 caches qdotk/gdotv in - // shared and pass 2 reads them back -> pass 2 does NO K/V gather (attacks the - // long_scoreboard #1 stall, which pass 2 is half of). Rows with rlen > CACHE_CAP fall - // back to the recompute path. Nested ``coop`` namespace -> distinct symbols (no ODR - // clash with the baseline). Kept until validated. - // ===================================================================================== - namespace coop - { - - // per-warp shared cache capacity for qdotk/gdotv (floats each). Covers the polar - // rows of the target grids (rlen<=712). Longer rows fall back to recompute. - static constexpr int CACHE_CAP = 768; - - // BEGIN backward kernels and functions - - // called with (blockDim.x=32 and blockDim.y>1, BDIM=blockDim.x*blockDim.y) - // - // STORAGE_T is the global-memory element type of the INPUTS (kx/vx/qy/dy): - // float4 for the fp32 vectorized path, float / c10::Half / c10::BFloat16 for - // the scalar path. COMPUTE_T (float4 or float) is the arithmetic type and the - // type of the gradient OUTPUTS dkx/dvx/dqy — these stay fp32 because dkx/dvx - // are atomically scatter-accumulated across overlapping neighborhoods (fp16 - // atomics would lose precision / aren't well supported). The wrapper allocates - // dkx/dvx/dqy in fp32 and casts to the input dtype at the end. vload widens - // STORAGE_T inputs to COMPUTE_T at the load site. - template - __global__ __launch_bounds__(BDIM_X) void s2_attn_bwd_generic_vec_k( - int nchans_in, // no. of elements along channel dim - int nchans_out, // no. of elements along channel dim - int nlat_in, int nlon_in, int nlat_out, int nlon_out, - const STORAGE_T *__restrict__ kx, // [batch][nlat_in][nlon_in][nchan_in] - const STORAGE_T *__restrict__ vx, // [batch][nlat_in][nlon_in][nchan_out] - const STORAGE_T *__restrict__ qy, // [batch][nlat_out][nlon_out][nchan_in] - const STORAGE_T *__restrict__ dy, // [batch][nlat_out][nlon_out][nchan_out] - const int32_t *__restrict__ row_idx, const int64_t *__restrict__ row_off, - const int64_t *__restrict__ col_idx, const float *__restrict__ quad_weights, - typename vec_traits::compute_t *__restrict__ dkx, // [batch][nlat_in][nlon_in][nchan_in] - typename vec_traits::compute_t *__restrict__ dvx, // [batch][nlat_in][nlon_in][nchan_out] - typename vec_traits::compute_t *__restrict__ dqy) - { // [batch][nlat_out][nlon_out][nchan_in] - using COMPUTE_T = typename vec_traits::compute_t; - - extern __shared__ __align__(sizeof(float4)) float shext[]; - - // for dqy - COMPUTE_T *sh_alpha_k__ = reinterpret_cast(shext) + threadIdx.y * (nchans_in * 4 + nchans_out); - COMPUTE_T *sh_alpha_vw_ = sh_alpha_k__ + nchans_in; - COMPUTE_T *sh_alpha_kvw = sh_alpha_vw_ + nchans_in; - - COMPUTE_T *sh_dy = sh_alpha_kvw + nchans_in; - COMPUTE_T *sh_qy = sh_dy + nchans_out; - // sh_alpha_k__[nchan_in], sh_alpha_vw_[nchan_in], sh_alpha_kvw[nchan_in] - // sh_dy[nchan_out], sh_qy[nchan_in] - - const int batch = blockIdx.y; - - const uint64_t wid = uint64_t(blockIdx.x) * blockDim.y + threadIdx.y; - if (wid >= uint64_t(nlat_out) * nlon_out) { return; } - - const int tidx = threadIdx.x; - - // use permuted rows - const int h = wid / nlon_out; - const int wo = wid - (h * nlon_out); - const int ho = row_idx[h]; - - // one output lon step corresponds to pscale input lon steps (requires nlon_in % nlon_out == 0) - const int pscale = nlon_in / nlon_out; - - // offset input tensors - kx += int64_t(batch) * nlat_in * nlon_in * nchans_in; - qy += int64_t(batch) * nlat_out * nlon_out * nchans_in + int64_t(ho) * nlon_out * nchans_in - + int64_t(wo) * nchans_in; - - vx += int64_t(batch) * nlat_in * nlon_in * nchans_out; - dy += int64_t(batch) * nlat_out * nlon_out * nchans_out + int64_t(ho) * nlon_out * nchans_out - + int64_t(wo) * nchans_out; - - // offset output tensors - dkx += int64_t(batch) * nlat_in * nlon_in * nchans_in; - dvx += int64_t(batch) * nlat_in * nlon_in * nchans_out; - dqy += int64_t(batch) * nlat_out * nlon_out * nchans_in + int64_t(ho) * nlon_out * nchans_in - + int64_t(wo) * nchans_in; - - // zero/init shared memory - for (int chan = tidx; chan < nchans_in; chan += WARP_SIZE) { - sh_alpha_k__[chan] = __vset(0.0f); - sh_alpha_vw_[chan] = __vset(0.0f); - sh_alpha_kvw[chan] = __vset(0.0f); - - sh_qy[chan] = vload(qy, chan); - } - for (int chan = tidx; chan < nchans_out; chan += WARP_SIZE) { sh_dy[chan] = vload(dy, chan); } - -#if __CUDA_ARCH__ < 900 - // for architectures < 9.0, sh_dy and sh_qy will be read - // as individual floats at the end of the kernel, which - // breaks the assumption that each COMPUTE_T location is - // written to and read by the same thread throughout the - // kernel, in the case COMPUTE_T==float4 - if constexpr (std::is_same::value) { __syncwarp(); } -#endif - - // for dkx, dvx, dqy - float alpha_sum = 0.0f; - float qdotk_max = -FLT_MAX; - - // for dkx - float integral = 0.0f; - - const int64_t rbeg = row_off[ho]; - const int64_t rend = row_off[ho + 1]; - - col_idx += rbeg; - - const int rlen = rend - rbeg; - - // accumulate alpha_sum, integral, and shared stats, - // along with a progressively computed qdotk_max. - for (int off = 0; off < rlen; off++) { - - const int64_t col = col_idx[off]; - - const int hi = col / nlon_in; - const int wi = col - (hi * nlon_in); - const int wi_wo = wi + pscale * wo; - const int wip = wi_wo - (wi_wo / nlon_in) * nlon_in; - - const STORAGE_T *_kx = kx + int64_t(hi) * nlon_in * nchans_in + int64_t(wip) * nchans_in; - const STORAGE_T *_vx = vx + int64_t(hi) * nlon_in * nchans_out + int64_t(wip) * nchans_out; - - COMPUTE_T qdotk_v = __vset(0.0f); - COMPUTE_T gdotv_v = __vset(0.0f); - - for (int chan = tidx; chan < nchans_in; chan += WARP_SIZE) { - qdotk_v = __vadd(qdotk_v, __vmul(sh_qy[chan], vload(_kx, chan))); - } - for (int chan = tidx; chan < nchans_out; chan += WARP_SIZE) { - gdotv_v = __vadd(gdotv_v, __vmul(sh_dy[chan], vload(_vx, chan))); - } - - const float qdotk = __warp_sum(__vred(qdotk_v)); - const float gdotv = __warp_sum(__vred(gdotv_v)); - - const float qdotk_max_tmp = max(qdotk_max, qdotk); - const float alpha_inz = expf(qdotk - qdotk_max_tmp) * quad_weights[hi]; - const float max_correction = expf(qdotk_max - qdotk_max_tmp); - alpha_sum = alpha_sum * max_correction + alpha_inz; - - integral = integral * max_correction + alpha_inz * gdotv; - - const float ainz_gdotv = alpha_inz * gdotv; - - for (int chan = tidx; chan < nchans_in; chan += WARP_SIZE) { - - const COMPUTE_T kxval = vload(_kx, chan); - - sh_alpha_k__[chan] = __vadd(__vscale(max_correction, sh_alpha_k__[chan]), __vscale(alpha_inz, kxval)); - sh_alpha_vw_[chan] - = __vadd(__vscale(max_correction, sh_alpha_vw_[chan]), __vset(ainz_gdotv)); - sh_alpha_kvw[chan] - = __vadd(__vscale(max_correction, sh_alpha_kvw[chan]), __vscale(ainz_gdotv, kxval)); - } - qdotk_max = qdotk_max_tmp; - } - - const float alpha_sum_inv = 1.0f / alpha_sum; - - integral *= alpha_sum_inv; - - // Write dqy (fp32 output) - for (int chan = tidx; chan < nchans_in; chan += WARP_SIZE) { - - dqy[chan] = __vscale( - alpha_sum_inv * alpha_sum_inv, - __vsub(__vscale(alpha_sum, sh_alpha_kvw[chan]), __vmul(sh_alpha_vw_[chan], sh_alpha_k__[chan]))); - } - - // accumulate gradients for k and v - for (int off = 0; off < rlen; off++) { - - const int64_t col = col_idx[off]; - const int hi = col / nlon_in; - const int wi = col - (hi * nlon_in); - const int wi_wo = wi + pscale * wo; - const int wip = wi_wo - (wi_wo / nlon_in) * nlon_in; - - const STORAGE_T *_kx = kx + int64_t(hi) * nlon_in * nchans_in + int64_t(wip) * nchans_in; - const STORAGE_T *_vx = vx + int64_t(hi) * nlon_in * nchans_out + int64_t(wip) * nchans_out; - - COMPUTE_T qdotk_v = __vset(0.0f); - COMPUTE_T gdotv_v = __vset(0.0f); - - for (int chan = tidx; chan < nchans_in; chan += WARP_SIZE) { - qdotk_v = __vadd(qdotk_v, __vmul(sh_qy[chan], vload(_kx, chan))); - } - for (int chan = tidx; chan < nchans_out; chan += WARP_SIZE) { - gdotv_v = __vadd(gdotv_v, __vmul(sh_dy[chan], vload(_vx, chan))); - } - - const float qdotk = __warp_sum(__vred(qdotk_v)); - const float gdotv = __warp_sum(__vred(gdotv_v)); - - const float alpha_inz = expf(qdotk - qdotk_max) * quad_weights[hi]; - - // _dkx / _dvx are COMPUTE_T (fp32) gradient buffers, accumulated atomically. - COMPUTE_T *_dkx = dkx + int64_t(hi) * nlon_in * nchans_in + int64_t(wip) * nchans_in; - COMPUTE_T *_dvx = dvx + int64_t(hi) * nlon_in * nchans_out + int64_t(wip) * nchans_out; - - const float alpha_mul = alpha_inz * alpha_sum_inv; - - const float scale_fact_qy = (gdotv - integral) * alpha_mul; - const float scale_fact_dy = alpha_mul; - - // float4, 128-bit atomics are only supported by devices of compute - // capability 9.x+, so on older devices we resort to 32-bit atomics - -#if __CUDA_ARCH__ < 900 - // to use 32-bit operations on consecutve addresses - float *sh_qy_scl = reinterpret_cast(sh_qy); - float *sh_dy_scl = reinterpret_cast(sh_dy); - - float *_dkx_scl = reinterpret_cast(_dkx); - float *_dvx_scl = reinterpret_cast(_dvx); - - constexpr int VEC_SIZE = sizeof(COMPUTE_T) / sizeof(float); - - // 32-bit, consecutive atomics to glmem; - // strided atomics results in a severe slowdown - for (int chan = tidx; chan < nchans_in * VEC_SIZE; chan += WARP_SIZE) { - atomicAdd(_dkx_scl + chan, scale_fact_qy * sh_qy_scl[chan]); - } - for (int chan = tidx; chan < nchans_out * VEC_SIZE; chan += WARP_SIZE) { - atomicAdd(_dvx_scl + chan, scale_fact_dy * sh_dy_scl[chan]); - } -#else - // 128-bit, consecutive atomics to glmem - for (int chan = tidx; chan < nchans_in; chan += WARP_SIZE) { - atomicAdd(_dkx + chan, __vscale(scale_fact_qy, sh_qy[chan])); - } - for (int chan = tidx; chan < nchans_out; chan += WARP_SIZE) { - atomicAdd(_dvx + chan, __vscale(scale_fact_dy, sh_dy[chan])); - } -#endif - } - - return; - } - - // called with either (BDIM_X=32 and BDIM_Y>1) || (2^K=BDIM_X > 32 and BDIM_Y=1) - template = nchan_in - typename STORAGE_T> - __global__ __launch_bounds__(BDIM_X *BDIM_Y) void s2_attn_bwd_special_vec_k( - int nchan_in, // no. of elements along channel dim - int nchan_out, // no. of elements along channel dim - int nlat_in, int nlon_in, int nlat_out, int nlon_out, - const STORAGE_T *__restrict__ kx, // [batch][nlat_in][nlon_in][nchan_in] - const STORAGE_T *__restrict__ vx, // [batch][nlat_in][nlon_in][nchan_out] - const STORAGE_T *__restrict__ qy, // [batch][nlat_out][nlon_out][nchan_in] - const STORAGE_T *__restrict__ dy, // [batch][nlat_out][nlon_out][nchan_out] - const int32_t *__restrict__ row_idx, const int64_t *__restrict__ row_off, - const int64_t *__restrict__ col_idx, const float *__restrict__ quad_weights, - typename vec_traits::compute_t *__restrict__ dkx, // [batch][nlat_in][nlon_in][nchan_in] - typename vec_traits::compute_t *__restrict__ dvx, // [batch][nlat_in][nlon_in][nchan_out] - typename vec_traits::compute_t *__restrict__ dqy, - // split params: process rows [row_ofs, row_ofs+nrows); cache_on gates the pass-2 - // recompute elimination (must be 0 when launched with the small no-cache shsize). - int64_t row_ofs, int64_t nrows, int cache_on) - { // [batch][nlat_out][nlon_out][nchan_in] - using COMPUTE_T = typename vec_traits::compute_t; - - static_assert(0 == (BDIM_X & (BDIM_X - 1))); - static_assert(0 == (BDIM_Y & (BDIM_Y - 1))); - static_assert((BDIM_X == 32 && BDIM_Y > 1) || (BDIM_X > 32 && BDIM_Y == 1)); - - constexpr int NLOC_M1 = NLOC - 1; - - const int tidx = threadIdx.x; - const int batch = blockIdx.y; - const uint64_t ctaid = uint64_t(blockIdx.x) * blockDim.y + threadIdx.y; - - if (ctaid >= uint64_t(nrows) * nlon_out) { return; } - - extern __shared__ __align__(sizeof(float4)) float shext[]; - - // sh_dy[nchan_out], sh_qy[nchan_in] - COMPUTE_T *sh_dy = reinterpret_cast(shext) + threadIdx.y * (nchan_in + nchan_out); // + tidx; - COMPUTE_T *sh_qy = sh_dy + nchan_out + tidx; - - if constexpr (CHOUT_AS_IN) { sh_dy += tidx; } - - // per-warp qdotk/gdotv cache (float), placed after all warps' sh_dy/sh_qy region. - constexpr int VECF = sizeof(COMPUTE_T) / sizeof(float); - float *qdc = shext + BDIM_Y * (nchan_in + nchan_out) * VECF + threadIdx.y * (2 * CACHE_CAP); - float *gdc = qdc + CACHE_CAP; - - // for dqy - COMPUTE_T loc_k__[NLOC]; - COMPUTE_T loc_vw_[NLOC]; - COMPUTE_T loc_kvw[NLOC]; - - // use permuted rows (offset into the sorted row list for the split launch) - const int h_local = ctaid / nlon_out; - const int wo = ctaid - (h_local * nlon_out); - const int h = h_local + int(row_ofs); - const int ho = row_idx[h]; - - // one output lon step corresponds to pscale input lon steps (requires nlon_in % nlon_out == 0) - const int pscale = nlon_in / nlon_out; - - // offset input tensors - kx += int64_t(batch) * nlat_in * nlon_in * nchan_in + tidx; - qy += int64_t(batch) * nlat_out * nlon_out * nchan_in + int64_t(ho) * nlon_out * nchan_in - + int64_t(wo) * nchan_in + tidx; - - vx += int64_t(batch) * nlat_in * nlon_in * nchan_out; // + tidx; - dy += int64_t(batch) * nlat_out * nlon_out * nchan_out + int64_t(ho) * nlon_out * nchan_out - + int64_t(wo) * nchan_out; // + tidx; - if constexpr (CHOUT_AS_IN) { - vx += tidx; - dy += tidx; - } - - // offset output tensors - dkx += int64_t(batch) * nlat_in * nlon_in * nchan_in + tidx; - dvx += int64_t(batch) * nlat_in * nlon_in * nchan_out; // + tidx; - if constexpr (CHOUT_AS_IN) { dvx += tidx; } - dqy += int64_t(batch) * nlat_out * nlon_out * nchan_in + int64_t(ho) * nlon_out * nchan_in - + int64_t(wo) * nchan_in + tidx; - -#pragma unroll - for (int i = 0; i < NLOC; i++) { - loc_k__[i] = __vset(0.0f); - loc_vw_[i] = __vset(0.0f); - loc_kvw[i] = __vset(0.0f); - } - -#pragma unroll - for (int i = 0; i < NLOC_M1; i++) { sh_qy[i * BDIM_X] = vload(qy, i * BDIM_X); } - if (NLOC_M1 * BDIM_X + tidx < nchan_in) { sh_qy[NLOC_M1 * BDIM_X] = vload(qy, NLOC_M1 * BDIM_X); } - - if constexpr (CHOUT_AS_IN) { -#pragma unroll - for (int i = 0; i < NLOC_M1; i++) { sh_dy[i * BDIM_X] = vload(dy, i * BDIM_X); } - if (NLOC_M1 * BDIM_X + tidx < nchan_out) { sh_dy[NLOC_M1 * BDIM_X] = vload(dy, NLOC_M1 * BDIM_X); } - } else { - for (int chan = tidx; chan < nchan_out; chan += BDIM_X) { sh_dy[chan] = vload(dy, chan); } - } - -#if __CUDA_ARCH__ < 900 - // for architectures < 9.0, sh_dy and sh_qy will be read - // as individual floats at the end of the kernel, which - // breaks the assumption that each COMPUTE_T location is - // written to and read by the same thread throughout the - // kernel, in the case COMPUTE_T==float4 - if constexpr (std::is_same::value) { - if constexpr (BDIM_X == 32) { - __syncwarp(); - } else { - __syncthreads(); - } - } -#endif - - // for dkx, dvx, dqy - float alpha_sum = 0.0f; - float qdotk_max = -FLT_MAX; - - // for dkx - float integral = 0.0f; - - const int64_t rbeg = row_off[ho]; - const int64_t rend = row_off[ho + 1]; - - col_idx += rbeg; - - const int rlen = rend - rbeg; - // cache pass-1 qdotk/gdotv iff enabled for this launch AND the row fits. - const bool use_cache = (cache_on != 0) && (rlen <= CACHE_CAP); - - // accumulate alpha_sum, integral, and shared stats, - // along with a progressively computed qdotk_max. - for (int off = 0; off < rlen; off++) { - - const int64_t col = col_idx[off]; - - const int hi = col / nlon_in; - const int wi = col - (hi * nlon_in); - const int wi_wo = wi + pscale * wo; - const int wip = wi_wo - (wi_wo / nlon_in) * nlon_in; - - const STORAGE_T *_kx = kx + int64_t(hi) * nlon_in * nchan_in + int64_t(wip) * nchan_in; - const STORAGE_T *_vx = vx + int64_t(hi) * nlon_in * nchan_out + int64_t(wip) * nchan_out; - - COMPUTE_T qdotk_v = __vset(0.0f); - COMPUTE_T gdotv_v = __vset(0.0f); - -#pragma unroll - for (int i = 0; i < NLOC_M1; i++) { - qdotk_v = __vadd(qdotk_v, __vmul(sh_qy[i * BDIM_X], vload(_kx, i * BDIM_X))); - } - if (NLOC_M1 * BDIM_X + tidx < nchan_in) { - qdotk_v = __vadd(qdotk_v, __vmul(sh_qy[NLOC_M1 * BDIM_X], vload(_kx, NLOC_M1 * BDIM_X))); - } - if constexpr (CHOUT_AS_IN) { -#pragma unroll - for (int i = 0; i < NLOC_M1; i++) { - gdotv_v = __vadd(gdotv_v, __vmul(sh_dy[i * BDIM_X], vload(_vx, i * BDIM_X))); - } - if (NLOC_M1 * BDIM_X + tidx < nchan_out) { - gdotv_v = __vadd(gdotv_v, __vmul(sh_dy[NLOC_M1 * BDIM_X], vload(_vx, NLOC_M1 * BDIM_X))); - } - } else { - for (int chan = tidx; chan < nchan_out; chan += BDIM_X) { - gdotv_v = __vadd(gdotv_v, __vmul(sh_dy[chan], vload(_vx, chan))); - } - } - - float qdotk = __vred(qdotk_v); - float gdotv = __vred(gdotv_v); - - if constexpr (BDIM_X == 32) { - qdotk = __warp_sum(qdotk); - gdotv = __warp_sum(gdotv); - } else { - qdotk = __block_sum(qdotk); - gdotv = __block_sum(gdotv); - } - - // cache the two per-neighbor scalars so pass 2 need not re-gather K/V. - if (use_cache) { - qdc[off] = qdotk; - gdc[off] = gdotv; - } - - const float qdotk_max_tmp = max(qdotk_max, qdotk); - const float alpha_inz = expf(qdotk - qdotk_max_tmp) * quad_weights[hi]; - const float max_correction = expf(qdotk_max - qdotk_max_tmp); - - alpha_sum = alpha_sum * max_correction + alpha_inz; - integral = integral * max_correction + alpha_inz * gdotv; - - const float ainz_gdotv = alpha_inz * gdotv; - -#pragma unroll - for (int i = 0; i < NLOC_M1; i++) { - const COMPUTE_T kxval = vload(_kx, i * BDIM_X); - loc_k__[i] = __vadd(__vscale(max_correction, loc_k__[i]), __vscale(alpha_inz, kxval)); - loc_vw_[i] = __vadd(__vscale(max_correction, loc_vw_[i]), __vset(ainz_gdotv)); - loc_kvw[i] = __vadd(__vscale(max_correction, loc_kvw[i]), __vscale(ainz_gdotv, kxval)); - } - if (NLOC_M1 * BDIM_X + tidx < nchan_in) { - const COMPUTE_T kxval = vload(_kx, NLOC_M1 * BDIM_X); - loc_k__[NLOC_M1] = __vadd(__vscale(max_correction, loc_k__[NLOC_M1]), __vscale(alpha_inz, kxval)); - loc_vw_[NLOC_M1] = __vadd(__vscale(max_correction, loc_vw_[NLOC_M1]), __vset(ainz_gdotv)); - loc_kvw[NLOC_M1] = __vadd(__vscale(max_correction, loc_kvw[NLOC_M1]), __vscale(ainz_gdotv, kxval)); - } - - qdotk_max = qdotk_max_tmp; - } - - const float alpha_sum_inv = 1.0f / alpha_sum; - - integral *= alpha_sum_inv; - - // Write dqy - const float alpha_sum_inv_sq = alpha_sum_inv * alpha_sum_inv; - -#pragma unroll - for (int i = 0; i < NLOC_M1; i++) { - dqy[i * BDIM_X] = __vscale(alpha_sum_inv_sq, - __vsub(__vscale(alpha_sum, loc_kvw[i]), __vmul(loc_vw_[i], loc_k__[i]))); - } - if (NLOC_M1 * BDIM_X + tidx < nchan_in) { - dqy[NLOC_M1 * BDIM_X] - = __vscale(alpha_sum_inv_sq, - __vsub(__vscale(alpha_sum, loc_kvw[NLOC_M1]), __vmul(loc_vw_[NLOC_M1], loc_k__[NLOC_M1]))); - } - - // make the pass-1 qdotk/gdotv cache writes visible before pass 2 reads them. - if (use_cache) { - if constexpr (BDIM_X == 32) { - __syncwarp(); - } else { - __syncthreads(); - } - } - - // accumulate gradients for k and v - for (int off = 0; off < rlen; off++) { - - const int64_t col = col_idx[off]; - - const int hi = col / nlon_in; - const int wi = col - (hi * nlon_in); - const int wi_wo = wi + pscale * wo; - const int wip = wi_wo - (wi_wo / nlon_in) * nlon_in; - - float qdotk, gdotv; - if (use_cache) { - // pass 2: reuse pass-1 scalars from shared — NO K/V gather. - qdotk = qdc[off]; - gdotv = gdc[off]; - } else { - const STORAGE_T *_kx = kx + int64_t(hi) * nlon_in * nchan_in + int64_t(wip) * nchan_in; - const STORAGE_T *_vx = vx + int64_t(hi) * nlon_in * nchan_out + int64_t(wip) * nchan_out; - - COMPUTE_T qdotk_v = __vset(0.0f); - COMPUTE_T gdotv_v = __vset(0.0f); - -#pragma unroll - for (int i = 0; i < NLOC_M1; i++) { - qdotk_v = __vadd(qdotk_v, __vmul(sh_qy[i * BDIM_X], vload(_kx, i * BDIM_X))); - } - if (NLOC_M1 * BDIM_X + tidx < nchan_in) { - qdotk_v = __vadd(qdotk_v, __vmul(sh_qy[NLOC_M1 * BDIM_X], vload(_kx, NLOC_M1 * BDIM_X))); - } - if constexpr (CHOUT_AS_IN) { -#pragma unroll - for (int i = 0; i < NLOC_M1; i++) { - gdotv_v = __vadd(gdotv_v, __vmul(sh_dy[i * BDIM_X], vload(_vx, i * BDIM_X))); - } - if (NLOC_M1 * BDIM_X + tidx < nchan_out) { - gdotv_v = __vadd(gdotv_v, __vmul(sh_dy[NLOC_M1 * BDIM_X], vload(_vx, NLOC_M1 * BDIM_X))); - } - } else { - for (int chan = tidx; chan < nchan_out; chan += BDIM_X) { - gdotv_v = __vadd(gdotv_v, __vmul(sh_dy[chan], vload(_vx, chan))); - } - } - - qdotk = __vred(qdotk_v); - gdotv = __vred(gdotv_v); - - if constexpr (BDIM_X == 32) { - qdotk = __warp_sum(qdotk); - gdotv = __warp_sum(gdotv); - } else { - qdotk = __block_sum(qdotk); - gdotv = __block_sum(gdotv); - } - } - - const float alpha_inz = expf(qdotk - qdotk_max) * quad_weights[hi]; - - COMPUTE_T *_dkx = dkx + int64_t(hi) * nlon_in * nchan_in + int64_t(wip) * nchan_in; - COMPUTE_T *_dvx = dvx + int64_t(hi) * nlon_in * nchan_out + int64_t(wip) * nchan_out; - - const float alpha_mul = alpha_inz * alpha_sum_inv; - - const float scale_fact_qy = (gdotv - integral) * alpha_mul; - const float scale_fact_dy = alpha_mul; - - // float4, 128-bit atomics are only supported by devices of compute - // capability 9.x+, so on older devices we resort to 32-bit atomics - -#if __CUDA_ARCH__ < 900 - constexpr int VEC_SIZE = sizeof(COMPUTE_T) / sizeof(float); - - // making the loop count known at compile time doesn't seem - // to make any difference here so let's keep this (much) - // simpler version - float *sh_qy_scl = reinterpret_cast(sh_qy); - float *sh_dy_scl = reinterpret_cast(sh_dy); - - float *_dkx_scl = reinterpret_cast(_dkx); - float *_dvx_scl = reinterpret_cast(_dvx); - - sh_qy_scl -= tidx * VEC_SIZE; - _dkx_scl -= tidx * VEC_SIZE; - if constexpr (CHOUT_AS_IN) { - sh_dy_scl -= tidx * VEC_SIZE; - _dvx_scl -= tidx * VEC_SIZE; - } - - // 32-bit, consecutive atomics to glmem - // strided atomics results in a severe slowdown - for (int chan = tidx; chan < nchan_in * VEC_SIZE; chan += BDIM_X) { - atomicAdd(_dkx_scl + chan, scale_fact_qy * sh_qy_scl[chan]); - } - for (int chan = tidx; chan < nchan_out * VEC_SIZE; chan += BDIM_X) { - atomicAdd(_dvx_scl + chan, scale_fact_dy * sh_dy_scl[chan]); - } -#else -#pragma unroll - for (int i = 0; i < NLOC_M1; i++) { - atomicAdd(_dkx + i * BDIM_X, __vscale(scale_fact_qy, sh_qy[i * BDIM_X])); - } - if (NLOC_M1 * BDIM_X + tidx < nchan_in) { - atomicAdd(_dkx + NLOC_M1 * BDIM_X, __vscale(scale_fact_qy, sh_qy[NLOC_M1 * BDIM_X])); - } - if constexpr (CHOUT_AS_IN) { -#pragma unroll - for (int i = 0; i < NLOC_M1; i++) { - atomicAdd(_dvx + i * BDIM_X, __vscale(scale_fact_dy, sh_dy[i * BDIM_X])); - } - if (NLOC_M1 * BDIM_X + tidx < nchan_out) { - atomicAdd(_dvx + NLOC_M1 * BDIM_X, __vscale(scale_fact_dy, sh_dy[NLOC_M1 * BDIM_X])); - } - } else { - for (int chan = tidx; chan < nchan_out; chan += BDIM_X) { - atomicAdd(_dvx + chan, __vscale(scale_fact_dy, sh_dy[chan])); - } - } -#endif - } - - return; - } - - template - void launch_gen_attn_bwd(int batch_size, int nchans_in, int nchans_out, int nlat_in, int nlon_in, int nlat_out, - int nlon_out, STORAGE_T *_kxp, STORAGE_T *_vxp, STORAGE_T *_qyp, STORAGE_T *_dyp, - int32_t *_row_idx, int64_t *_row_off, int64_t *_col_idx, float *_quad_weights, - typename vec_traits::compute_t *_dkxp, - typename vec_traits::compute_t *_dvxp, - typename vec_traits::compute_t *_dqyp, cudaStream_t stream) - { - - dim3 block(WARP_SIZE, THREADS / WARP_SIZE); - dim3 grid(DIV_UP(nlat_out * nlon_out, block.y), batch_size); - - // shared memory holds compute-type (COMPUTE_T) data, not STORAGE_T. 5 arrays per warp. - size_t shsize = sizeof(typename vec_traits::compute_t) * (nchans_in * 4 + nchans_out) * block.y; - - s2_attn_bwd_generic_vec_k<<>>( - nchans_in, nchans_out, nlat_in, nlon_in, nlat_out, nlon_out, _kxp, _vxp, _qyp, _dyp, _row_idx, _row_off, - _col_idx, _quad_weights, _dkxp, _dvxp, _dqyp); - CHECK_ERROR("s2_attn_bwd_generic_vec_k"); - - return; - } - - template - void launch_spc_attn_bwd(int nloc, // "BDIM_X*nloc" >= nchans_out - int batch_size, int nchans_in, int nchans_out, int nlat_in, int nlon_in, int nlat_out, - int nlon_out, STORAGE_T *_kxp, STORAGE_T *_vxp, STORAGE_T *_qyp, STORAGE_T *_dyp, - int32_t *_row_idx, int64_t *_row_off, int64_t *_col_idx, float *_quad_weights, - typename vec_traits::compute_t *_dkxp, - typename vec_traits::compute_t *_dvxp, - typename vec_traits::compute_t *_dqyp, cudaStream_t stream, - int64_t row_ofs = 0, int64_t nrows = -1, int cache_on = 1) - { - - if (CUR_LOC_SIZE == nloc) { - - // process rows [row_ofs, row_ofs + grid_rows); nrows<0 means the whole grid. - const int64_t grid_rows = (nrows < 0) ? int64_t(nlat_out) : nrows; - dim3 block(BDIM_X, BDIM_Y); - dim3 grid(DIV_UP(grid_rows * nlon_out, block.y), batch_size); - - // shared memory holds compute-type (COMPUTE_T) data, not STORAGE_T. - size_t shsize = sizeof(typename vec_traits::compute_t) * (nchans_in + nchans_out) * block.y; - // + per-warp qdotk/gdotv cache (float) ONLY when caching is enabled for this launch. - if (cache_on) { shsize += size_t(block.y) * (2 * CACHE_CAP) * sizeof(float); } - - if (nchans_out >= BDIM_X * (CUR_LOC_SIZE - 1) && nchans_out <= BDIM_X * CUR_LOC_SIZE) { - s2_attn_bwd_special_vec_k<<>>( - nchans_in, nchans_out, nlat_in, nlon_in, nlat_out, nlon_out, _kxp, _vxp, _qyp, _dyp, _row_idx, - _row_off, _col_idx, _quad_weights, _dkxp, _dvxp, _dqyp, row_ofs, grid_rows, cache_on); - } else { - s2_attn_bwd_special_vec_k<<>>( - nchans_in, nchans_out, nlat_in, nlon_in, nlat_out, nlon_out, _kxp, _vxp, _qyp, _dyp, _row_idx, - _row_off, _col_idx, _quad_weights, _dkxp, _dvxp, _dqyp, row_ofs, grid_rows, cache_on); - } - CHECK_ERROR("s2_attn_bwd_special_vec_k"); - - return; - } - if constexpr (CUR_LOC_SIZE < MAX_LOC_SIZE) { - launch_spc_attn_bwd( - nloc, batch_size, nchans_in, nchans_out, nlat_in, nlon_in, nlat_out, nlon_out, _kxp, _vxp, _qyp, _dyp, - _row_idx, _row_off, _col_idx, _quad_weights, _dkxp, _dvxp, _dqyp, stream, row_ofs, nrows, cache_on); - } - return; - } - - // Picks the block size (BDIM_X) instance and launches the backward gather - // kernel for a given input storage vector type SV. Inputs are SV*; gradient - // outputs are COMPUTE_T* (fp32) — see s2_attn_bwd_*_vec_k. Backward uses the - // special kernel only up to BDIM_X=512 (1024 spills); larger falls to generic. - template - static void bwd_dispatch_bdimx(int bdimx, int nloc, int64_t batch_size, int64_t nchans_in, int64_t nchans_out, - int nlat_in, int64_t nlon_in, int64_t nlat_out, int64_t nlon_out, SV *_kxp, - SV *_vxp, SV *_qyp, SV *_dyp, int32_t *_row_idx, int64_t *_row_off, - int64_t *_col_idx, float *_quad_weights, typename vec_traits::compute_t *_dkxp, - typename vec_traits::compute_t *_dvxp, - typename vec_traits::compute_t *_dqyp, cudaStream_t stream, - int64_t row_ofs = 0, int64_t nrows = -1, int cache_on = 1) - { - switch (bdimx) { - case 32: - launch_spc_attn_bwd<32, 2, 1, MAX_LOC>(nloc, batch_size, nchans_in, nchans_out, nlat_in, nlon_in, - nlat_out, nlon_out, _kxp, _vxp, _qyp, _dyp, _row_idx, _row_off, - _col_idx, _quad_weights, _dkxp, _dvxp, _dqyp, stream, row_ofs, - nrows, cache_on); - break; - case 64: - launch_spc_attn_bwd<64, 1, MIN_LOC, MAX_LOC>(nloc, batch_size, nchans_in, nchans_out, nlat_in, nlon_in, - nlat_out, nlon_out, _kxp, _vxp, _qyp, _dyp, _row_idx, - _row_off, _col_idx, _quad_weights, _dkxp, _dvxp, _dqyp, - stream, row_ofs, nrows, cache_on); - break; - case 128: - launch_spc_attn_bwd<128, 1, MIN_LOC, MAX_LOC>(nloc, batch_size, nchans_in, nchans_out, nlat_in, nlon_in, - nlat_out, nlon_out, _kxp, _vxp, _qyp, _dyp, _row_idx, - _row_off, _col_idx, _quad_weights, _dkxp, _dvxp, _dqyp, - stream, row_ofs, nrows, cache_on); - break; - case 256: - launch_spc_attn_bwd<256, 1, MIN_LOC, MAX_LOC>(nloc, batch_size, nchans_in, nchans_out, nlat_in, nlon_in, - nlat_out, nlon_out, _kxp, _vxp, _qyp, _dyp, _row_idx, - _row_off, _col_idx, _quad_weights, _dkxp, _dvxp, _dqyp, - stream, row_ofs, nrows, cache_on); - break; - case 512: - launch_spc_attn_bwd<512, 1, MIN_LOC, MAX_LOC>(nloc, batch_size, nchans_in, nchans_out, nlat_in, nlon_in, - nlat_out, nlon_out, _kxp, _vxp, _qyp, _dyp, _row_idx, - _row_off, _col_idx, _quad_weights, _dkxp, _dvxp, _dqyp, - stream, row_ofs, nrows, cache_on); - break; - default: - // generic path: no caching, whole grid (split params ignored). - launch_gen_attn_bwd(batch_size, nchans_in, nchans_out, nlat_in, nlon_in, nlat_out, nlon_out, _kxp, _vxp, - _qyp, _dyp, _row_idx, _row_off, _col_idx, _quad_weights, _dkxp, _dvxp, _dqyp, stream); - break; - } - } - - // Templated on the input storage element type (float / c10::Half / c10::BFloat16). - // Gradient outputs dkx/dvx/dqy are always fp32 (COMPUTE_T) — dkx/dvx are atomic - // scatter-accumulated, so they cannot be reduced precision; the wrapper casts - // them to the input dtype at the end. fp32 keeps the float4 vectorized path; - // fp16/bf16 (and unaligned fp32) take the scalar STORAGE_T path. - template - static void s2_attn_bwd_dispatch(int64_t batch_size, int64_t nchans_in, int64_t nchans_out, int64_t nlon_in, - int64_t nlat_out, int64_t nlon_out, at::Tensor kxP, at::Tensor vxP, - at::Tensor qyP, at::Tensor dyP, at::Tensor row_off, at::Tensor col_idx, - at::Tensor quad_weights, at::Tensor dkxP, at::Tensor dvxP, at::Tensor dqyP) - { - - static_assert(0 == (MAX_LOCAL_ARR_LEN & (MAX_LOCAL_ARR_LEN - 1))); - - // get stream - auto stream = at::cuda::getCurrentCUDAStream().stream(); - - // sort row indices (ho-s) in descending order - // based on (row_off[ho+1]-row_off[ho]) - at::Tensor row_idx = sortRows(nlat_out, row_off, stream); - - const int nlat_in = kxP.size(1); - - // smallest power of two "bdimx" (>=32) s.t. bdimx*MAX_LOCAL_ARR_LEN >= nchans_in - int bdimx; - bdimx = DIV_UP(nchans_in, MAX_LOCAL_ARR_LEN); - bdimx = max(bdimx, WARP_SIZE); - bdimx = next_pow2(bdimx); - - scalar_t *_kxp = reinterpret_cast(kxP.data_ptr()); - scalar_t *_vxp = reinterpret_cast(vxP.data_ptr()); - scalar_t *_qyp = reinterpret_cast(qyP.data_ptr()); - scalar_t *_dyp = reinterpret_cast(dyP.data_ptr()); - - // gradient outputs are fp32 - float *_dkxp = reinterpret_cast(dkxP.data_ptr()); - float *_dvxp = reinterpret_cast(dvxP.data_ptr()); - float *_dqyp = reinterpret_cast(dqyP.data_ptr()); - - int32_t *_row_idx = reinterpret_cast(row_idx.data_ptr()); - int64_t *_row_off = reinterpret_cast(row_off.data_ptr()); - int64_t *_col_idx = reinterpret_cast(col_idx.data_ptr()); - float *_quad_weights = reinterpret_cast(quad_weights.data_ptr()); - - constexpr int MIN_LOC_ARR_LEN = MAX_LOCAL_ARR_LEN / 2 + 1; - - if constexpr (std::is_same::value) { - // fp32: float4 vectorized when 16B-aligned + 4-divisible, else scalar. - constexpr int VEC_SIZE = sizeof(float4) / sizeof(float); // 4 - const bool use_vec = is_aligned<16>(_kxp) && is_aligned<16>(_vxp) && is_aligned<16>(_qyp) - && is_aligned<16>(_dyp) && is_aligned<16>(_dkxp) && is_aligned<16>(_dvxp) && is_aligned<16>(_dqyp) - && (nchans_in % VEC_SIZE) == 0 && (nchans_out % VEC_SIZE) == 0; - - if (use_vec) { - constexpr int MAX_VEC = MAX_LOCAL_ARR_LEN / VEC_SIZE; - constexpr int MIN_VEC = MAX_VEC / 2 + 1; - const int64_t nci = nchans_in / VEC_SIZE; - const int64_t nco = nchans_out / VEC_SIZE; - bwd_dispatch_bdimx( - bdimx, DIV_UP(nci, bdimx), batch_size, nci, nco, nlat_in, nlon_in, nlat_out, nlon_out, - reinterpret_cast(_kxp), reinterpret_cast(_vxp), - reinterpret_cast(_qyp), reinterpret_cast(_dyp), _row_idx, _row_off, - _col_idx, _quad_weights, reinterpret_cast(_dkxp), reinterpret_cast(_dvxp), - reinterpret_cast(_dqyp), stream); - } else { - bwd_dispatch_bdimx( - bdimx, DIV_UP(nchans_in, bdimx), batch_size, nchans_in, nchans_out, nlat_in, nlon_in, nlat_out, - nlon_out, _kxp, _vxp, _qyp, _dyp, _row_idx, _row_off, _col_idx, _quad_weights, _dkxp, _dvxp, - _dqyp, stream); - } - } else { - // fp16/bf16: scalar STORAGE_T inputs, fp32 outputs; fp32 compute/accumulation. - // ROW SPLIT: long/polar rows get the pass-2 recompute-elimination cache; cheap - // short rows run the pristine recompute path at full occupancy (no cache tax). - const int nloc_s = DIV_UP(nchans_in, bdimx); - const bool special = (bdimx <= 512); - int64_t n_long = 0, m0 = 0, m1 = 0; - if (special) { - // reuse the ring's relative 0.1*max_rlen threshold but drop the 1024 floor - // (split_len=1) so our polar rows (rlen<1024) are classified "long". - split_csr_rows(SPLIT_ROW_LENGTH_THRES, 1, nlat_out, _row_idx, _row_off, &n_long, &m0, &m1); - } - if (special && n_long > 0 && n_long < nlat_out) { - bwd_dispatch_bdimx( - bdimx, nloc_s, batch_size, nchans_in, nchans_out, nlat_in, nlon_in, nlat_out, nlon_out, _kxp, - _vxp, _qyp, _dyp, _row_idx, _row_off, _col_idx, _quad_weights, _dkxp, _dvxp, _dqyp, stream, - /*row_ofs=*/0, /*nrows=*/n_long, /*cache_on=*/1); - bwd_dispatch_bdimx( - bdimx, nloc_s, batch_size, nchans_in, nchans_out, nlat_in, nlon_in, nlat_out, nlon_out, _kxp, - _vxp, _qyp, _dyp, _row_idx, _row_off, _col_idx, _quad_weights, _dkxp, _dvxp, _dqyp, stream, - /*row_ofs=*/n_long, /*nrows=*/nlat_out - n_long, /*cache_on=*/0); - } else { - // no split: whole grid; cache on for the special path, off for generic (bdimx>512). - bwd_dispatch_bdimx( - bdimx, nloc_s, batch_size, nchans_in, nchans_out, nlat_in, nlon_in, nlat_out, nlon_out, _kxp, - _vxp, _qyp, _dyp, _row_idx, _row_off, _col_idx, _quad_weights, _dkxp, _dvxp, _dqyp, stream, 0, - -1, special ? 1 : 0); - } - } - - return; - } - - // END backward kernels and functions - - std::tuple - s2_attention_bwd_dkvq_cuda(at::Tensor kx, at::Tensor vx, at::Tensor qy, at::Tensor dy, at::Tensor quad_weights, - at::Tensor psi_col_idx, at::Tensor psi_row_off, int64_t nlon_in, int64_t nlat_out, - int64_t nlon_out) - { - - CHECK_CUDA_INPUT_TENSOR(kx); - CHECK_CUDA_INPUT_TENSOR(vx); - CHECK_CUDA_INPUT_TENSOR(qy); - CHECK_CUDA_INPUT_TENSOR(dy); - CHECK_CUDA_TENSOR(quad_weights); - CHECK_CUDA_TENSOR(psi_col_idx); - CHECK_CUDA_TENSOR(psi_row_off); - - // direction selection: gather (self / downsample) iff nlon_in is an integer - // multiple of nlon_out; scatter (upsample) iff nlon_out is an integer multiple - // of nlon_in. Self-attention satisfies both and routes through the gather path. - const bool downsample = (nlon_in % nlon_out == 0); - const bool upsample = (nlon_out % nlon_in == 0); - TORCH_CHECK(downsample || upsample, "either nlon_in (", nlon_in, - ") must be an integer multiple of nlon_out (", nlon_out, "), or vice versa"); - - // const size_t uo_num_channels = kx.size(1); - size_t nchans_in = qy.size(1); // or kx.size(1) - size_t nchans_out = vx.size(1); - - const int batch_size = kx.size(0); - const int64_t nlat_in = kx.size(2); - - // extract dtype - auto kx_type = kx.dtype(); // nchans_in - auto qy_type = qy.dtype(); - auto vx_type = vx.dtype(); // ncahn_out - auto dy_type = dy.dtype(); - - torch::Tensor dkx, dvx, dqy; - - // ATen dispatch over the input dtype. - // - // Gather (downsample / self) path: native storage. kx/vx/qy/dy keep their - // dtype; the kernel widens to fp32 at load (Tier B). Gradient buffers - // dkx/dvx/dqy are allocated fp32 because dkx/dvx are atomically - // scatter-accumulated (reduced-precision atomics would lose precision); - // they are cast back to the input dtype at the end. - // - // Scatter (upsample) path: native storage as well — inputs stay in their - // dtype (widened at load inside the dispatch), gradients stay fp32. - AT_DISPATCH_FLOATING_TYPES_AND2(at::kHalf, at::kBFloat16, qy.scalar_type(), "s2_attention_bwd_dkvq_cuda", [&] { - using storage_t = scalar_t; - - const auto f32_like - = [](const torch::Tensor &t) { return torch::zeros_like(t, t.options().dtype(torch::kFloat32)); }; - - if (downsample) { - // native-storage gather path - torch::Tensor kxP = kx; - torch::Tensor vxP = vx; - torch::Tensor qyP = qy; - torch::Tensor dyP = dy; - - // safer than is_contiguous(ChannelsLast), which fails for num_channels == 1 - bool kx_is_channels_last = kxP.strides()[1] == 1; - bool vx_is_channels_last = vxP.strides()[1] == 1; - bool qy_is_channels_last = qyP.strides()[1] == 1; - bool dy_is_channels_last = dyP.strides()[1] == 1; - - if (!kx_is_channels_last) { kxP = permute_4D_to0231(kxP); } - if (!vx_is_channels_last) { vxP = permute_4D_to0231(vxP); } - if (!qy_is_channels_last) { qyP = permute_4D_to0231(qyP); } - if (!dy_is_channels_last) { dyP = permute_4D_to0231(dyP); } - - // fp32 gradient buffers (same shape/layout as the native inputs) - torch::Tensor dkxP = f32_like(kxP); - torch::Tensor dvxP = f32_like(vxP); - torch::Tensor dqyP = f32_like(qyP); - - s2_attn_bwd_dispatch(batch_size, nchans_in, nchans_out, nlon_in, nlat_out, nlon_out, kxP, - vxP, qyP, dyP, psi_row_off, psi_col_idx, quad_weights, dkxP, dvxP, - dqyP); - - dkx = dkxP; - dvx = dvxP; - dqy = dqyP; - - if (!kx_is_channels_last) { dkx = permute_4D_to0312(dkx); } - if (!vx_is_channels_last) { dvx = permute_4D_to0312(dvx); } - if (!qy_is_channels_last) { dqy = permute_4D_to0312(dqy); } - } else { - // native-storage scatter (upsample) path - torch::Tensor kxP = kx; - torch::Tensor vxP = vx; - torch::Tensor qyP = qy; - torch::Tensor dyP = dy; - - bool kx_is_channels_last = kxP.strides()[1] == 1; - bool vx_is_channels_last = vxP.strides()[1] == 1; - bool qy_is_channels_last = qyP.strides()[1] == 1; - bool dy_is_channels_last = dyP.strides()[1] == 1; - - if (!kx_is_channels_last) { kxP = permute_4D_to0231(kxP); } - if (!vx_is_channels_last) { vxP = permute_4D_to0231(vxP); } - if (!qy_is_channels_last) { qyP = permute_4D_to0231(qyP); } - if (!dy_is_channels_last) { dyP = permute_4D_to0231(dyP); } - - torch::Tensor dkxP = f32_like(kxP); - torch::Tensor dvxP = f32_like(vxP); - torch::Tensor dqyP = f32_like(qyP); - - s2_attn_bwd_upsample_dispatch(batch_size, nchans_in, nchans_out, nlon_in, nlat_in, nlat_out, - nlon_out, kxP, vxP, qyP, dyP, psi_row_off, psi_col_idx, quad_weights, - dkxP, dvxP, dqyP); - - dkx = dkxP; - dvx = dvxP; - dqy = dqyP; - - if (!kx_is_channels_last) { dkx = permute_4D_to0312(dkx); } - if (!vx_is_channels_last) { dvx = permute_4D_to0312(dvx); } - if (!qy_is_channels_last) { dqy = permute_4D_to0312(dqy); } - } - }); - - C10_CUDA_KERNEL_LAUNCH_CHECK(); - - // convert precision back to starting dtype (no-op for fp32; narrows for fp16/bf16) - dkx = dkx.to(kx_type); - dvx = dvx.to(vx_type); - dqy = dqy.to(qy_type); - - return std::make_tuple(dkx, dvx, dqy); - } - - TORCH_LIBRARY_IMPL(attention_kernels, CUDA, m) { m.impl("backward_coop", &s2_attention_bwd_dkvq_cuda); } - - } // namespace coop - -} // namespace attention_kernels diff --git a/torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_bwd_fast.cu b/torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_bwd_fast.cu deleted file mode 100644 index 4df9569a..00000000 --- a/torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_bwd_fast.cu +++ /dev/null @@ -1,982 +0,0 @@ -// coding=utf-8 -// -// SPDX-FileCopyrightText: Copyright (c) 2025 The torch-harmonics Authors. All rights reserved. -// SPDX-License-Identifier: BSD-3-Clause -// -// Redistribution and use in source and binary forms, with or without -// modification, are permitted provided that the following conditions are met: -// -// 1. Redistributions of source code must retain the above copyright notice, this -// list of conditions and the following disclaimer. -// -// 2. Redistributions in binary form must reproduce the above copyright notice, -// this list of conditions and the following disclaimer in the documentation -// and/or other materials provided with the distribution. -// -// 3. Neither the name of the copyright holder nor the names of its -// contributors may be used to endorse or promote products derived from -// this software without specific prior written permission. -// -// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" -// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE -// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE -// DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE -// FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL -// DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR -// SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER -// CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, -// OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE -// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. - -#include "attention_cuda.cuh" -#include -#include -#include "c10/core/MemoryFormat.h" - -#include -#include -#include -#include -#include -#include - -#include -#include -#include - -#include "cudamacro.h" -#include "attention_cuda_utils.cuh" - -#include -#include -#include - -#define THREADS (64) - -#define MAX_LOCAL_ARR_LEN (16) - -namespace attention_kernels -{ - - // scatter-direction dispatcher, defined in attention_cuda_bwd_upsample.cu; - // called by s2_attention_bwd_dkvq_cuda when nlon_out % nlon_in == 0. - void s2_attn_bwd_upsample_dispatch(int batch_size, size_t nchans_in, size_t nchans_out, int64_t nlon_in, - int64_t nlat_in, int64_t nlat_out, int64_t nlon_out, torch::Tensor kxP, - torch::Tensor vxP, torch::Tensor qyP, torch::Tensor dyP, - torch::Tensor psi_row_off, torch::Tensor psi_col_idx, torch::Tensor quad_weights, - torch::Tensor dkxP, torch::Tensor dvxP, torch::Tensor dqyP); - - // ===================================================================================== - // Optimized (FMA + __expf) copy of the GATHER backward, exposed as the ``backward_fast`` - // op for A/B benchmarking against the baseline ``backward``. Nested ``fast`` namespace - // keeps the template kernels' symbols distinct from the baseline (no ODR clash); the - // upsample path delegates to attention_kernels::s2_attn_bwd_upsample_dispatch via - // enclosing-namespace lookup. Kept until the fast path is validated as the default. - // ===================================================================================== - namespace fast - { - - // BEGIN backward kernels and functions - - // called with (blockDim.x=32 and blockDim.y>1, BDIM=blockDim.x*blockDim.y) - // - // STORAGE_T is the global-memory element type of the INPUTS (kx/vx/qy/dy): - // float4 for the fp32 vectorized path, float / c10::Half / c10::BFloat16 for - // the scalar path. COMPUTE_T (float4 or float) is the arithmetic type and the - // type of the gradient OUTPUTS dkx/dvx/dqy — these stay fp32 because dkx/dvx - // are atomically scatter-accumulated across overlapping neighborhoods (fp16 - // atomics would lose precision / aren't well supported). The wrapper allocates - // dkx/dvx/dqy in fp32 and casts to the input dtype at the end. vload widens - // STORAGE_T inputs to COMPUTE_T at the load site. - template - __global__ __launch_bounds__(BDIM_X) void s2_attn_bwd_generic_vec_k( - int nchans_in, // no. of elements along channel dim - int nchans_out, // no. of elements along channel dim - int nlat_in, int nlon_in, int nlat_out, int nlon_out, - const STORAGE_T *__restrict__ kx, // [batch][nlat_in][nlon_in][nchan_in] - const STORAGE_T *__restrict__ vx, // [batch][nlat_in][nlon_in][nchan_out] - const STORAGE_T *__restrict__ qy, // [batch][nlat_out][nlon_out][nchan_in] - const STORAGE_T *__restrict__ dy, // [batch][nlat_out][nlon_out][nchan_out] - const int32_t *__restrict__ row_idx, const int64_t *__restrict__ row_off, - const int64_t *__restrict__ col_idx, const float *__restrict__ quad_weights, - typename vec_traits::compute_t *__restrict__ dkx, // [batch][nlat_in][nlon_in][nchan_in] - typename vec_traits::compute_t *__restrict__ dvx, // [batch][nlat_in][nlon_in][nchan_out] - typename vec_traits::compute_t *__restrict__ dqy) - { // [batch][nlat_out][nlon_out][nchan_in] - using COMPUTE_T = typename vec_traits::compute_t; - - extern __shared__ __align__(sizeof(float4)) float shext[]; - - // for dqy - COMPUTE_T *sh_alpha_k__ = reinterpret_cast(shext) + threadIdx.y * (nchans_in * 4 + nchans_out); - COMPUTE_T *sh_alpha_vw_ = sh_alpha_k__ + nchans_in; - COMPUTE_T *sh_alpha_kvw = sh_alpha_vw_ + nchans_in; - - COMPUTE_T *sh_dy = sh_alpha_kvw + nchans_in; - COMPUTE_T *sh_qy = sh_dy + nchans_out; - // sh_alpha_k__[nchan_in], sh_alpha_vw_[nchan_in], sh_alpha_kvw[nchan_in] - // sh_dy[nchan_out], sh_qy[nchan_in] - - const int batch = blockIdx.y; - - const uint64_t wid = uint64_t(blockIdx.x) * blockDim.y + threadIdx.y; - if (wid >= uint64_t(nlat_out) * nlon_out) { return; } - - const int tidx = threadIdx.x; - - // use permuted rows - const int h = wid / nlon_out; - const int wo = wid - (h * nlon_out); - const int ho = row_idx[h]; - - // one output lon step corresponds to pscale input lon steps (requires nlon_in % nlon_out == 0) - const int pscale = nlon_in / nlon_out; - - // offset input tensors - kx += int64_t(batch) * nlat_in * nlon_in * nchans_in; - qy += int64_t(batch) * nlat_out * nlon_out * nchans_in + int64_t(ho) * nlon_out * nchans_in - + int64_t(wo) * nchans_in; - - vx += int64_t(batch) * nlat_in * nlon_in * nchans_out; - dy += int64_t(batch) * nlat_out * nlon_out * nchans_out + int64_t(ho) * nlon_out * nchans_out - + int64_t(wo) * nchans_out; - - // offset output tensors - dkx += int64_t(batch) * nlat_in * nlon_in * nchans_in; - dvx += int64_t(batch) * nlat_in * nlon_in * nchans_out; - dqy += int64_t(batch) * nlat_out * nlon_out * nchans_in + int64_t(ho) * nlon_out * nchans_in - + int64_t(wo) * nchans_in; - - // zero/init shared memory - for (int chan = tidx; chan < nchans_in; chan += WARP_SIZE) { - sh_alpha_k__[chan] = __vset(0.0f); - sh_alpha_vw_[chan] = __vset(0.0f); - sh_alpha_kvw[chan] = __vset(0.0f); - - sh_qy[chan] = vload(qy, chan); - } - for (int chan = tidx; chan < nchans_out; chan += WARP_SIZE) { sh_dy[chan] = vload(dy, chan); } - -#if __CUDA_ARCH__ < 900 - // for architectures < 9.0, sh_dy and sh_qy will be read - // as individual floats at the end of the kernel, which - // breaks the assumption that each COMPUTE_T location is - // written to and read by the same thread throughout the - // kernel, in the case COMPUTE_T==float4 - if constexpr (std::is_same::value) { __syncwarp(); } -#endif - - // for dkx, dvx, dqy - float alpha_sum = 0.0f; - float qdotk_max = -FLT_MAX; - - // for dkx - float integral = 0.0f; - - const int64_t rbeg = row_off[ho]; - const int64_t rend = row_off[ho + 1]; - - col_idx += rbeg; - - const int rlen = rend - rbeg; - - // accumulate alpha_sum, integral, and shared stats, - // along with a progressively computed qdotk_max. - for (int off = 0; off < rlen; off++) { - - const int64_t col = col_idx[off]; - - const int hi = col / nlon_in; - const int wi = col - (hi * nlon_in); - const int wi_wo = wi + pscale * wo; - const int wip = wi_wo - (wi_wo / nlon_in) * nlon_in; - - const STORAGE_T *_kx = kx + int64_t(hi) * nlon_in * nchans_in + int64_t(wip) * nchans_in; - const STORAGE_T *_vx = vx + int64_t(hi) * nlon_in * nchans_out + int64_t(wip) * nchans_out; - - COMPUTE_T qdotk_v = __vset(0.0f); - COMPUTE_T gdotv_v = __vset(0.0f); - - for (int chan = tidx; chan < nchans_in; chan += WARP_SIZE) { - qdotk_v = __vadd(qdotk_v, __vmul(sh_qy[chan], vload(_kx, chan))); - } - for (int chan = tidx; chan < nchans_out; chan += WARP_SIZE) { - gdotv_v = __vadd(gdotv_v, __vmul(sh_dy[chan], vload(_vx, chan))); - } - - const float qdotk = __warp_sum(__vred(qdotk_v)); - const float gdotv = __warp_sum(__vred(gdotv_v)); - - const float qdotk_max_tmp = max(qdotk_max, qdotk); - const float alpha_inz = __expf(qdotk - qdotk_max_tmp) * quad_weights[hi]; - const float max_correction = __expf(qdotk_max - qdotk_max_tmp); - alpha_sum = alpha_sum * max_correction + alpha_inz; - - integral = integral * max_correction + alpha_inz * gdotv; - - const float ainz_gdotv = alpha_inz * gdotv; - - for (int chan = tidx; chan < nchans_in; chan += WARP_SIZE) { - - const COMPUTE_T kxval = vload(_kx, chan); - - sh_alpha_k__[chan] = __vfma_scale(alpha_inz, kxval, __vscale(max_correction, sh_alpha_k__[chan])); - sh_alpha_vw_[chan] = __vfma_scale(max_correction, sh_alpha_vw_[chan], __vset(ainz_gdotv)); - sh_alpha_kvw[chan] = __vfma_scale(ainz_gdotv, kxval, __vscale(max_correction, sh_alpha_kvw[chan])); - } - qdotk_max = qdotk_max_tmp; - } - - const float alpha_sum_inv = 1.0f / alpha_sum; - - integral *= alpha_sum_inv; - - // Write dqy (fp32 output) - for (int chan = tidx; chan < nchans_in; chan += WARP_SIZE) { - - dqy[chan] = __vscale( - alpha_sum_inv * alpha_sum_inv, - __vsub(__vscale(alpha_sum, sh_alpha_kvw[chan]), __vmul(sh_alpha_vw_[chan], sh_alpha_k__[chan]))); - } - - // accumulate gradients for k and v - for (int off = 0; off < rlen; off++) { - - const int64_t col = col_idx[off]; - const int hi = col / nlon_in; - const int wi = col - (hi * nlon_in); - const int wi_wo = wi + pscale * wo; - const int wip = wi_wo - (wi_wo / nlon_in) * nlon_in; - - const STORAGE_T *_kx = kx + int64_t(hi) * nlon_in * nchans_in + int64_t(wip) * nchans_in; - const STORAGE_T *_vx = vx + int64_t(hi) * nlon_in * nchans_out + int64_t(wip) * nchans_out; - - COMPUTE_T qdotk_v = __vset(0.0f); - COMPUTE_T gdotv_v = __vset(0.0f); - - for (int chan = tidx; chan < nchans_in; chan += WARP_SIZE) { - qdotk_v = __vadd(qdotk_v, __vmul(sh_qy[chan], vload(_kx, chan))); - } - for (int chan = tidx; chan < nchans_out; chan += WARP_SIZE) { - gdotv_v = __vadd(gdotv_v, __vmul(sh_dy[chan], vload(_vx, chan))); - } - - const float qdotk = __warp_sum(__vred(qdotk_v)); - const float gdotv = __warp_sum(__vred(gdotv_v)); - - const float alpha_inz = __expf(qdotk - qdotk_max) * quad_weights[hi]; - - // _dkx / _dvx are COMPUTE_T (fp32) gradient buffers, accumulated atomically. - COMPUTE_T *_dkx = dkx + int64_t(hi) * nlon_in * nchans_in + int64_t(wip) * nchans_in; - COMPUTE_T *_dvx = dvx + int64_t(hi) * nlon_in * nchans_out + int64_t(wip) * nchans_out; - - const float alpha_mul = alpha_inz * alpha_sum_inv; - - const float scale_fact_qy = (gdotv - integral) * alpha_mul; - const float scale_fact_dy = alpha_mul; - - // float4, 128-bit atomics are only supported by devices of compute - // capability 9.x+, so on older devices we resort to 32-bit atomics - -#if __CUDA_ARCH__ < 900 - // to use 32-bit operations on consecutve addresses - float *sh_qy_scl = reinterpret_cast(sh_qy); - float *sh_dy_scl = reinterpret_cast(sh_dy); - - float *_dkx_scl = reinterpret_cast(_dkx); - float *_dvx_scl = reinterpret_cast(_dvx); - - constexpr int VEC_SIZE = sizeof(COMPUTE_T) / sizeof(float); - - // 32-bit, consecutive atomics to glmem; - // strided atomics results in a severe slowdown - for (int chan = tidx; chan < nchans_in * VEC_SIZE; chan += WARP_SIZE) { - atomicAdd(_dkx_scl + chan, scale_fact_qy * sh_qy_scl[chan]); - } - for (int chan = tidx; chan < nchans_out * VEC_SIZE; chan += WARP_SIZE) { - atomicAdd(_dvx_scl + chan, scale_fact_dy * sh_dy_scl[chan]); - } -#else - // 128-bit, consecutive atomics to glmem - for (int chan = tidx; chan < nchans_in; chan += WARP_SIZE) { - atomicAdd(_dkx + chan, __vscale(scale_fact_qy, sh_qy[chan])); - } - for (int chan = tidx; chan < nchans_out; chan += WARP_SIZE) { - atomicAdd(_dvx + chan, __vscale(scale_fact_dy, sh_dy[chan])); - } -#endif - } - - return; - } - - // called with either (BDIM_X=32 and BDIM_Y>1) || (2^K=BDIM_X > 32 and BDIM_Y=1) - template = nchan_in - typename STORAGE_T> - __global__ __launch_bounds__(BDIM_X *BDIM_Y) void s2_attn_bwd_special_vec_k( - int nchan_in, // no. of elements along channel dim - int nchan_out, // no. of elements along channel dim - int nlat_in, int nlon_in, int nlat_out, int nlon_out, - const STORAGE_T *__restrict__ kx, // [batch][nlat_in][nlon_in][nchan_in] - const STORAGE_T *__restrict__ vx, // [batch][nlat_in][nlon_in][nchan_out] - const STORAGE_T *__restrict__ qy, // [batch][nlat_out][nlon_out][nchan_in] - const STORAGE_T *__restrict__ dy, // [batch][nlat_out][nlon_out][nchan_out] - const int32_t *__restrict__ row_idx, const int64_t *__restrict__ row_off, - const int64_t *__restrict__ col_idx, const float *__restrict__ quad_weights, - typename vec_traits::compute_t *__restrict__ dkx, // [batch][nlat_in][nlon_in][nchan_in] - typename vec_traits::compute_t *__restrict__ dvx, // [batch][nlat_in][nlon_in][nchan_out] - typename vec_traits::compute_t *__restrict__ dqy) - { // [batch][nlat_out][nlon_out][nchan_in] - using COMPUTE_T = typename vec_traits::compute_t; - - static_assert(0 == (BDIM_X & (BDIM_X - 1))); - static_assert(0 == (BDIM_Y & (BDIM_Y - 1))); - static_assert((BDIM_X == 32 && BDIM_Y > 1) || (BDIM_X > 32 && BDIM_Y == 1)); - - constexpr int NLOC_M1 = NLOC - 1; - - const int tidx = threadIdx.x; - const int batch = blockIdx.y; - const uint64_t ctaid = uint64_t(blockIdx.x) * blockDim.y + threadIdx.y; - - if (ctaid >= uint64_t(nlat_out) * nlon_out) { return; } - - extern __shared__ __align__(sizeof(float4)) float shext[]; - - // sh_dy[nchan_out], sh_qy[nchan_in] - COMPUTE_T *sh_dy = reinterpret_cast(shext) + threadIdx.y * (nchan_in + nchan_out); // + tidx; - COMPUTE_T *sh_qy = sh_dy + nchan_out + tidx; - - if constexpr (CHOUT_AS_IN) { sh_dy += tidx; } - - // for dqy - COMPUTE_T loc_k__[NLOC]; - COMPUTE_T loc_vw_[NLOC]; - COMPUTE_T loc_kvw[NLOC]; - - // use permuted rows - const int h = ctaid / nlon_out; - const int wo = ctaid - (h * nlon_out); - const int ho = row_idx[h]; - - // one output lon step corresponds to pscale input lon steps (requires nlon_in % nlon_out == 0) - const int pscale = nlon_in / nlon_out; - - // offset input tensors - kx += int64_t(batch) * nlat_in * nlon_in * nchan_in + tidx; - qy += int64_t(batch) * nlat_out * nlon_out * nchan_in + int64_t(ho) * nlon_out * nchan_in - + int64_t(wo) * nchan_in + tidx; - - vx += int64_t(batch) * nlat_in * nlon_in * nchan_out; // + tidx; - dy += int64_t(batch) * nlat_out * nlon_out * nchan_out + int64_t(ho) * nlon_out * nchan_out - + int64_t(wo) * nchan_out; // + tidx; - if constexpr (CHOUT_AS_IN) { - vx += tidx; - dy += tidx; - } - - // offset output tensors - dkx += int64_t(batch) * nlat_in * nlon_in * nchan_in + tidx; - dvx += int64_t(batch) * nlat_in * nlon_in * nchan_out; // + tidx; - if constexpr (CHOUT_AS_IN) { dvx += tidx; } - dqy += int64_t(batch) * nlat_out * nlon_out * nchan_in + int64_t(ho) * nlon_out * nchan_in - + int64_t(wo) * nchan_in + tidx; - -#pragma unroll - for (int i = 0; i < NLOC; i++) { - loc_k__[i] = __vset(0.0f); - loc_vw_[i] = __vset(0.0f); - loc_kvw[i] = __vset(0.0f); - } - -#pragma unroll - for (int i = 0; i < NLOC_M1; i++) { sh_qy[i * BDIM_X] = vload(qy, i * BDIM_X); } - if (NLOC_M1 * BDIM_X + tidx < nchan_in) { sh_qy[NLOC_M1 * BDIM_X] = vload(qy, NLOC_M1 * BDIM_X); } - - if constexpr (CHOUT_AS_IN) { -#pragma unroll - for (int i = 0; i < NLOC_M1; i++) { sh_dy[i * BDIM_X] = vload(dy, i * BDIM_X); } - if (NLOC_M1 * BDIM_X + tidx < nchan_out) { sh_dy[NLOC_M1 * BDIM_X] = vload(dy, NLOC_M1 * BDIM_X); } - } else { - for (int chan = tidx; chan < nchan_out; chan += BDIM_X) { sh_dy[chan] = vload(dy, chan); } - } - -#if __CUDA_ARCH__ < 900 - // for architectures < 9.0, sh_dy and sh_qy will be read - // as individual floats at the end of the kernel, which - // breaks the assumption that each COMPUTE_T location is - // written to and read by the same thread throughout the - // kernel, in the case COMPUTE_T==float4 - if constexpr (std::is_same::value) { - if constexpr (BDIM_X == 32) { - __syncwarp(); - } else { - __syncthreads(); - } - } -#endif - - // for dkx, dvx, dqy - float alpha_sum = 0.0f; - float qdotk_max = -FLT_MAX; - - // for dkx - float integral = 0.0f; - - const int64_t rbeg = row_off[ho]; - const int64_t rend = row_off[ho + 1]; - - col_idx += rbeg; - - const int rlen = rend - rbeg; - - // accumulate alpha_sum, integral, and shared stats, - // along with a progressively computed qdotk_max. - for (int off = 0; off < rlen; off++) { - - const int64_t col = col_idx[off]; - - const int hi = col / nlon_in; - const int wi = col - (hi * nlon_in); - const int wi_wo = wi + pscale * wo; - const int wip = wi_wo - (wi_wo / nlon_in) * nlon_in; - - const STORAGE_T *_kx = kx + int64_t(hi) * nlon_in * nchan_in + int64_t(wip) * nchan_in; - const STORAGE_T *_vx = vx + int64_t(hi) * nlon_in * nchan_out + int64_t(wip) * nchan_out; - - COMPUTE_T qdotk_v = __vset(0.0f); - COMPUTE_T gdotv_v = __vset(0.0f); - -#pragma unroll - for (int i = 0; i < NLOC_M1; i++) { - qdotk_v = __vadd(qdotk_v, __vmul(sh_qy[i * BDIM_X], vload(_kx, i * BDIM_X))); - } - if (NLOC_M1 * BDIM_X + tidx < nchan_in) { - qdotk_v = __vadd(qdotk_v, __vmul(sh_qy[NLOC_M1 * BDIM_X], vload(_kx, NLOC_M1 * BDIM_X))); - } - if constexpr (CHOUT_AS_IN) { -#pragma unroll - for (int i = 0; i < NLOC_M1; i++) { - gdotv_v = __vadd(gdotv_v, __vmul(sh_dy[i * BDIM_X], vload(_vx, i * BDIM_X))); - } - if (NLOC_M1 * BDIM_X + tidx < nchan_out) { - gdotv_v = __vadd(gdotv_v, __vmul(sh_dy[NLOC_M1 * BDIM_X], vload(_vx, NLOC_M1 * BDIM_X))); - } - } else { - for (int chan = tidx; chan < nchan_out; chan += BDIM_X) { - gdotv_v = __vadd(gdotv_v, __vmul(sh_dy[chan], vload(_vx, chan))); - } - } - - float qdotk = __vred(qdotk_v); - float gdotv = __vred(gdotv_v); - - if constexpr (BDIM_X == 32) { - qdotk = __warp_sum(qdotk); - gdotv = __warp_sum(gdotv); - } else { - qdotk = __block_sum(qdotk); - gdotv = __block_sum(gdotv); - } - - const float qdotk_max_tmp = max(qdotk_max, qdotk); - const float alpha_inz = __expf(qdotk - qdotk_max_tmp) * quad_weights[hi]; - const float max_correction = __expf(qdotk_max - qdotk_max_tmp); - - alpha_sum = alpha_sum * max_correction + alpha_inz; - integral = integral * max_correction + alpha_inz * gdotv; - - const float ainz_gdotv = alpha_inz * gdotv; - -#pragma unroll - for (int i = 0; i < NLOC_M1; i++) { - const COMPUTE_T kxval = vload(_kx, i * BDIM_X); - loc_k__[i] = __vfma_scale(alpha_inz, kxval, __vscale(max_correction, loc_k__[i])); - loc_vw_[i] = __vfma_scale(max_correction, loc_vw_[i], __vset(ainz_gdotv)); - loc_kvw[i] = __vfma_scale(ainz_gdotv, kxval, __vscale(max_correction, loc_kvw[i])); - } - if (NLOC_M1 * BDIM_X + tidx < nchan_in) { - const COMPUTE_T kxval = vload(_kx, NLOC_M1 * BDIM_X); - loc_k__[NLOC_M1] = __vfma_scale(alpha_inz, kxval, __vscale(max_correction, loc_k__[NLOC_M1])); - loc_vw_[NLOC_M1] = __vfma_scale(max_correction, loc_vw_[NLOC_M1], __vset(ainz_gdotv)); - loc_kvw[NLOC_M1] = __vfma_scale(ainz_gdotv, kxval, __vscale(max_correction, loc_kvw[NLOC_M1])); - } - - qdotk_max = qdotk_max_tmp; - } - - const float alpha_sum_inv = 1.0f / alpha_sum; - - integral *= alpha_sum_inv; - - // Write dqy - const float alpha_sum_inv_sq = alpha_sum_inv * alpha_sum_inv; - -#pragma unroll - for (int i = 0; i < NLOC_M1; i++) { - dqy[i * BDIM_X] = __vscale(alpha_sum_inv_sq, - __vsub(__vscale(alpha_sum, loc_kvw[i]), __vmul(loc_vw_[i], loc_k__[i]))); - } - if (NLOC_M1 * BDIM_X + tidx < nchan_in) { - dqy[NLOC_M1 * BDIM_X] - = __vscale(alpha_sum_inv_sq, - __vsub(__vscale(alpha_sum, loc_kvw[NLOC_M1]), __vmul(loc_vw_[NLOC_M1], loc_k__[NLOC_M1]))); - } - - // accumulate gradients for k and v - for (int off = 0; off < rlen; off++) { - - const int64_t col = col_idx[off]; - - const int hi = col / nlon_in; - const int wi = col - (hi * nlon_in); - const int wi_wo = wi + pscale * wo; - const int wip = wi_wo - (wi_wo / nlon_in) * nlon_in; - - const STORAGE_T *_kx = kx + int64_t(hi) * nlon_in * nchan_in + int64_t(wip) * nchan_in; - const STORAGE_T *_vx = vx + int64_t(hi) * nlon_in * nchan_out + int64_t(wip) * nchan_out; - - COMPUTE_T qdotk_v = __vset(0.0f); - COMPUTE_T gdotv_v = __vset(0.0f); - -#pragma unroll - for (int i = 0; i < NLOC_M1; i++) { - qdotk_v = __vadd(qdotk_v, __vmul(sh_qy[i * BDIM_X], vload(_kx, i * BDIM_X))); - } - if (NLOC_M1 * BDIM_X + tidx < nchan_in) { - qdotk_v = __vadd(qdotk_v, __vmul(sh_qy[NLOC_M1 * BDIM_X], vload(_kx, NLOC_M1 * BDIM_X))); - } - if constexpr (CHOUT_AS_IN) { -#pragma unroll - for (int i = 0; i < NLOC_M1; i++) { - gdotv_v = __vadd(gdotv_v, __vmul(sh_dy[i * BDIM_X], vload(_vx, i * BDIM_X))); - } - if (NLOC_M1 * BDIM_X + tidx < nchan_out) { - gdotv_v = __vadd(gdotv_v, __vmul(sh_dy[NLOC_M1 * BDIM_X], vload(_vx, NLOC_M1 * BDIM_X))); - } - } else { - for (int chan = tidx; chan < nchan_out; chan += BDIM_X) { - gdotv_v = __vadd(gdotv_v, __vmul(sh_dy[chan], vload(_vx, chan))); - } - } - - float qdotk = __vred(qdotk_v); - float gdotv = __vred(gdotv_v); - - if constexpr (BDIM_X == 32) { - qdotk = __warp_sum(qdotk); - gdotv = __warp_sum(gdotv); - } else { - qdotk = __block_sum(qdotk); - gdotv = __block_sum(gdotv); - } - - const float alpha_inz = __expf(qdotk - qdotk_max) * quad_weights[hi]; - - COMPUTE_T *_dkx = dkx + int64_t(hi) * nlon_in * nchan_in + int64_t(wip) * nchan_in; - COMPUTE_T *_dvx = dvx + int64_t(hi) * nlon_in * nchan_out + int64_t(wip) * nchan_out; - - const float alpha_mul = alpha_inz * alpha_sum_inv; - - const float scale_fact_qy = (gdotv - integral) * alpha_mul; - const float scale_fact_dy = alpha_mul; - - // float4, 128-bit atomics are only supported by devices of compute - // capability 9.x+, so on older devices we resort to 32-bit atomics - -#if __CUDA_ARCH__ < 900 - constexpr int VEC_SIZE = sizeof(COMPUTE_T) / sizeof(float); - - // making the loop count known at compile time doesn't seem - // to make any difference here so let's keep this (much) - // simpler version - float *sh_qy_scl = reinterpret_cast(sh_qy); - float *sh_dy_scl = reinterpret_cast(sh_dy); - - float *_dkx_scl = reinterpret_cast(_dkx); - float *_dvx_scl = reinterpret_cast(_dvx); - - sh_qy_scl -= tidx * VEC_SIZE; - _dkx_scl -= tidx * VEC_SIZE; - if constexpr (CHOUT_AS_IN) { - sh_dy_scl -= tidx * VEC_SIZE; - _dvx_scl -= tidx * VEC_SIZE; - } - - // 32-bit, consecutive atomics to glmem - // strided atomics results in a severe slowdown - for (int chan = tidx; chan < nchan_in * VEC_SIZE; chan += BDIM_X) { - atomicAdd(_dkx_scl + chan, scale_fact_qy * sh_qy_scl[chan]); - } - for (int chan = tidx; chan < nchan_out * VEC_SIZE; chan += BDIM_X) { - atomicAdd(_dvx_scl + chan, scale_fact_dy * sh_dy_scl[chan]); - } -#else -#pragma unroll - for (int i = 0; i < NLOC_M1; i++) { - atomicAdd(_dkx + i * BDIM_X, __vscale(scale_fact_qy, sh_qy[i * BDIM_X])); - } - if (NLOC_M1 * BDIM_X + tidx < nchan_in) { - atomicAdd(_dkx + NLOC_M1 * BDIM_X, __vscale(scale_fact_qy, sh_qy[NLOC_M1 * BDIM_X])); - } - if constexpr (CHOUT_AS_IN) { -#pragma unroll - for (int i = 0; i < NLOC_M1; i++) { - atomicAdd(_dvx + i * BDIM_X, __vscale(scale_fact_dy, sh_dy[i * BDIM_X])); - } - if (NLOC_M1 * BDIM_X + tidx < nchan_out) { - atomicAdd(_dvx + NLOC_M1 * BDIM_X, __vscale(scale_fact_dy, sh_dy[NLOC_M1 * BDIM_X])); - } - } else { - for (int chan = tidx; chan < nchan_out; chan += BDIM_X) { - atomicAdd(_dvx + chan, __vscale(scale_fact_dy, sh_dy[chan])); - } - } -#endif - } - - return; - } - - template - void launch_gen_attn_bwd(int batch_size, int nchans_in, int nchans_out, int nlat_in, int nlon_in, int nlat_out, - int nlon_out, STORAGE_T *_kxp, STORAGE_T *_vxp, STORAGE_T *_qyp, STORAGE_T *_dyp, - int32_t *_row_idx, int64_t *_row_off, int64_t *_col_idx, float *_quad_weights, - typename vec_traits::compute_t *_dkxp, - typename vec_traits::compute_t *_dvxp, - typename vec_traits::compute_t *_dqyp, cudaStream_t stream) - { - - dim3 block(WARP_SIZE, THREADS / WARP_SIZE); - dim3 grid(DIV_UP(nlat_out * nlon_out, block.y), batch_size); - - // shared memory holds compute-type (COMPUTE_T) data, not STORAGE_T. 5 arrays per warp. - size_t shsize = sizeof(typename vec_traits::compute_t) * (nchans_in * 4 + nchans_out) * block.y; - - s2_attn_bwd_generic_vec_k<<>>( - nchans_in, nchans_out, nlat_in, nlon_in, nlat_out, nlon_out, _kxp, _vxp, _qyp, _dyp, _row_idx, _row_off, - _col_idx, _quad_weights, _dkxp, _dvxp, _dqyp); - CHECK_ERROR("s2_attn_bwd_generic_vec_k"); - - return; - } - - template - void launch_spc_attn_bwd(int nloc, // "BDIM_X*nloc" >= nchans_out - int batch_size, int nchans_in, int nchans_out, int nlat_in, int nlon_in, int nlat_out, - int nlon_out, STORAGE_T *_kxp, STORAGE_T *_vxp, STORAGE_T *_qyp, STORAGE_T *_dyp, - int32_t *_row_idx, int64_t *_row_off, int64_t *_col_idx, float *_quad_weights, - typename vec_traits::compute_t *_dkxp, - typename vec_traits::compute_t *_dvxp, - typename vec_traits::compute_t *_dqyp, cudaStream_t stream) - { - - if (CUR_LOC_SIZE == nloc) { - - dim3 block(BDIM_X, BDIM_Y); - dim3 grid(DIV_UP(nlat_out * nlon_out, block.y), batch_size); - - // shared memory holds compute-type (COMPUTE_T) data, not STORAGE_T. - // 2 arrays per cta, block.y > 1 iif block.x==32 - size_t shsize = sizeof(typename vec_traits::compute_t) * (nchans_in + nchans_out) * block.y; - - // nloc determines the size of local arrays used to store - // temporary buffers loc_k__[], loc_vw_[] and loc_kvw[], - // of size nchans_in each; - // if nchans_out is >= BDIM_X*(nloc-1) and <= BDIM_X*nloc - // then we can use the same compile-time known loops used - // for input channels, with the exception of testing - // whether to execute the last iteration based on "nchans_out" - // instead of "nchans_in"; in this way as long as the - // difference between the number of input and output channels - // is <= BDIM_X we can use the faster path - if (nchans_out >= BDIM_X * (CUR_LOC_SIZE - 1) && nchans_out <= BDIM_X * CUR_LOC_SIZE) { - s2_attn_bwd_special_vec_k<<>>( - nchans_in, nchans_out, nlat_in, nlon_in, nlat_out, nlon_out, _kxp, _vxp, _qyp, _dyp, _row_idx, - _row_off, _col_idx, _quad_weights, _dkxp, _dvxp, _dqyp); - } else { - s2_attn_bwd_special_vec_k<<>>( - nchans_in, nchans_out, nlat_in, nlon_in, nlat_out, nlon_out, _kxp, _vxp, _qyp, _dyp, _row_idx, - _row_off, _col_idx, _quad_weights, _dkxp, _dvxp, _dqyp); - } - CHECK_ERROR("s2_attn_bwd_special_vec_k"); - - return; - } - if constexpr (CUR_LOC_SIZE < MAX_LOC_SIZE) { - launch_spc_attn_bwd( - nloc, batch_size, nchans_in, nchans_out, nlat_in, nlon_in, nlat_out, nlon_out, _kxp, _vxp, _qyp, - _dyp, _row_idx, _row_off, _col_idx, _quad_weights, _dkxp, _dvxp, _dqyp, stream); - } - return; - } - - // Picks the block size (BDIM_X) instance and launches the backward gather - // kernel for a given input storage vector type SV. Inputs are SV*; gradient - // outputs are COMPUTE_T* (fp32) — see s2_attn_bwd_*_vec_k. Backward uses the - // special kernel only up to BDIM_X=512 (1024 spills); larger falls to generic. - template - static void bwd_dispatch_bdimx(int bdimx, int nloc, int64_t batch_size, int64_t nchans_in, int64_t nchans_out, - int nlat_in, int64_t nlon_in, int64_t nlat_out, int64_t nlon_out, SV *_kxp, - SV *_vxp, SV *_qyp, SV *_dyp, int32_t *_row_idx, int64_t *_row_off, - int64_t *_col_idx, float *_quad_weights, typename vec_traits::compute_t *_dkxp, - typename vec_traits::compute_t *_dvxp, - typename vec_traits::compute_t *_dqyp, cudaStream_t stream) - { - switch (bdimx) { - case 32: - launch_spc_attn_bwd<32, 2, 1, MAX_LOC>(nloc, batch_size, nchans_in, nchans_out, nlat_in, nlon_in, - nlat_out, nlon_out, _kxp, _vxp, _qyp, _dyp, _row_idx, _row_off, - _col_idx, _quad_weights, _dkxp, _dvxp, _dqyp, stream); - break; - case 64: - launch_spc_attn_bwd<64, 1, MIN_LOC, MAX_LOC>( - nloc, batch_size, nchans_in, nchans_out, nlat_in, nlon_in, nlat_out, nlon_out, _kxp, _vxp, _qyp, - _dyp, _row_idx, _row_off, _col_idx, _quad_weights, _dkxp, _dvxp, _dqyp, stream); - break; - case 128: - launch_spc_attn_bwd<128, 1, MIN_LOC, MAX_LOC>( - nloc, batch_size, nchans_in, nchans_out, nlat_in, nlon_in, nlat_out, nlon_out, _kxp, _vxp, _qyp, - _dyp, _row_idx, _row_off, _col_idx, _quad_weights, _dkxp, _dvxp, _dqyp, stream); - break; - case 256: - launch_spc_attn_bwd<256, 1, MIN_LOC, MAX_LOC>( - nloc, batch_size, nchans_in, nchans_out, nlat_in, nlon_in, nlat_out, nlon_out, _kxp, _vxp, _qyp, - _dyp, _row_idx, _row_off, _col_idx, _quad_weights, _dkxp, _dvxp, _dqyp, stream); - break; - case 512: - launch_spc_attn_bwd<512, 1, MIN_LOC, MAX_LOC>( - nloc, batch_size, nchans_in, nchans_out, nlat_in, nlon_in, nlat_out, nlon_out, _kxp, _vxp, _qyp, - _dyp, _row_idx, _row_off, _col_idx, _quad_weights, _dkxp, _dvxp, _dqyp, stream); - break; - default: - launch_gen_attn_bwd(batch_size, nchans_in, nchans_out, nlat_in, nlon_in, nlat_out, nlon_out, _kxp, _vxp, - _qyp, _dyp, _row_idx, _row_off, _col_idx, _quad_weights, _dkxp, _dvxp, _dqyp, stream); - break; - } - } - - // Templated on the input storage element type (float / c10::Half / c10::BFloat16). - // Gradient outputs dkx/dvx/dqy are always fp32 (COMPUTE_T) — dkx/dvx are atomic - // scatter-accumulated, so they cannot be reduced precision; the wrapper casts - // them to the input dtype at the end. fp32 keeps the float4 vectorized path; - // fp16/bf16 (and unaligned fp32) take the scalar STORAGE_T path. - template - static void s2_attn_bwd_dispatch(int64_t batch_size, int64_t nchans_in, int64_t nchans_out, int64_t nlon_in, - int64_t nlat_out, int64_t nlon_out, at::Tensor kxP, at::Tensor vxP, - at::Tensor qyP, at::Tensor dyP, at::Tensor row_off, at::Tensor col_idx, - at::Tensor quad_weights, at::Tensor dkxP, at::Tensor dvxP, at::Tensor dqyP) - { - - static_assert(0 == (MAX_LOCAL_ARR_LEN & (MAX_LOCAL_ARR_LEN - 1))); - - // get stream - auto stream = at::cuda::getCurrentCUDAStream().stream(); - - // sort row indices (ho-s) in descending order - // based on (row_off[ho+1]-row_off[ho]) - at::Tensor row_idx = sortRows(nlat_out, row_off, stream); - - const int nlat_in = kxP.size(1); - - // smallest power of two "bdimx" (>=32) s.t. bdimx*MAX_LOCAL_ARR_LEN >= nchans_in - int bdimx; - bdimx = DIV_UP(nchans_in, MAX_LOCAL_ARR_LEN); - bdimx = max(bdimx, WARP_SIZE); - bdimx = next_pow2(bdimx); - - scalar_t *_kxp = reinterpret_cast(kxP.data_ptr()); - scalar_t *_vxp = reinterpret_cast(vxP.data_ptr()); - scalar_t *_qyp = reinterpret_cast(qyP.data_ptr()); - scalar_t *_dyp = reinterpret_cast(dyP.data_ptr()); - - // gradient outputs are fp32 - float *_dkxp = reinterpret_cast(dkxP.data_ptr()); - float *_dvxp = reinterpret_cast(dvxP.data_ptr()); - float *_dqyp = reinterpret_cast(dqyP.data_ptr()); - - int32_t *_row_idx = reinterpret_cast(row_idx.data_ptr()); - int64_t *_row_off = reinterpret_cast(row_off.data_ptr()); - int64_t *_col_idx = reinterpret_cast(col_idx.data_ptr()); - float *_quad_weights = reinterpret_cast(quad_weights.data_ptr()); - - constexpr int MIN_LOC_ARR_LEN = MAX_LOCAL_ARR_LEN / 2 + 1; - - if constexpr (std::is_same::value) { - // fp32: float4 vectorized when 16B-aligned + 4-divisible, else scalar. - constexpr int VEC_SIZE = sizeof(float4) / sizeof(float); // 4 - const bool use_vec = is_aligned<16>(_kxp) && is_aligned<16>(_vxp) && is_aligned<16>(_qyp) - && is_aligned<16>(_dyp) && is_aligned<16>(_dkxp) && is_aligned<16>(_dvxp) && is_aligned<16>(_dqyp) - && (nchans_in % VEC_SIZE) == 0 && (nchans_out % VEC_SIZE) == 0; - - if (use_vec) { - constexpr int MAX_VEC = MAX_LOCAL_ARR_LEN / VEC_SIZE; - constexpr int MIN_VEC = MAX_VEC / 2 + 1; - const int64_t nci = nchans_in / VEC_SIZE; - const int64_t nco = nchans_out / VEC_SIZE; - bwd_dispatch_bdimx( - bdimx, DIV_UP(nci, bdimx), batch_size, nci, nco, nlat_in, nlon_in, nlat_out, nlon_out, - reinterpret_cast(_kxp), reinterpret_cast(_vxp), - reinterpret_cast(_qyp), reinterpret_cast(_dyp), _row_idx, _row_off, - _col_idx, _quad_weights, reinterpret_cast(_dkxp), reinterpret_cast(_dvxp), - reinterpret_cast(_dqyp), stream); - } else { - bwd_dispatch_bdimx( - bdimx, DIV_UP(nchans_in, bdimx), batch_size, nchans_in, nchans_out, nlat_in, nlon_in, nlat_out, - nlon_out, _kxp, _vxp, _qyp, _dyp, _row_idx, _row_off, _col_idx, _quad_weights, _dkxp, _dvxp, - _dqyp, stream); - } - } else { - // fp16/bf16: scalar STORAGE_T inputs, fp32 outputs; fp32 compute/accumulation. - bwd_dispatch_bdimx( - bdimx, DIV_UP(nchans_in, bdimx), batch_size, nchans_in, nchans_out, nlat_in, nlon_in, nlat_out, - nlon_out, _kxp, _vxp, _qyp, _dyp, _row_idx, _row_off, _col_idx, _quad_weights, _dkxp, _dvxp, _dqyp, - stream); - } - - return; - } - - // END backward kernels and functions - - std::tuple - s2_attention_bwd_dkvq_cuda(at::Tensor kx, at::Tensor vx, at::Tensor qy, at::Tensor dy, at::Tensor quad_weights, - at::Tensor psi_col_idx, at::Tensor psi_row_off, int64_t nlon_in, int64_t nlat_out, - int64_t nlon_out) - { - - CHECK_CUDA_INPUT_TENSOR(kx); - CHECK_CUDA_INPUT_TENSOR(vx); - CHECK_CUDA_INPUT_TENSOR(qy); - CHECK_CUDA_INPUT_TENSOR(dy); - CHECK_CUDA_TENSOR(quad_weights); - CHECK_CUDA_TENSOR(psi_col_idx); - CHECK_CUDA_TENSOR(psi_row_off); - - // direction selection: gather (self / downsample) iff nlon_in is an integer - // multiple of nlon_out; scatter (upsample) iff nlon_out is an integer multiple - // of nlon_in. Self-attention satisfies both and routes through the gather path. - const bool downsample = (nlon_in % nlon_out == 0); - const bool upsample = (nlon_out % nlon_in == 0); - TORCH_CHECK(downsample || upsample, "either nlon_in (", nlon_in, - ") must be an integer multiple of nlon_out (", nlon_out, "), or vice versa"); - - // const size_t uo_num_channels = kx.size(1); - size_t nchans_in = qy.size(1); // or kx.size(1) - size_t nchans_out = vx.size(1); - - const int batch_size = kx.size(0); - const int64_t nlat_in = kx.size(2); - - // extract dtype - auto kx_type = kx.dtype(); // nchans_in - auto qy_type = qy.dtype(); - auto vx_type = vx.dtype(); // ncahn_out - auto dy_type = dy.dtype(); - - torch::Tensor dkx, dvx, dqy; - - // ATen dispatch over the input dtype. - // - // Gather (downsample / self) path: native storage. kx/vx/qy/dy keep their - // dtype; the kernel widens to fp32 at load (Tier B). Gradient buffers - // dkx/dvx/dqy are allocated fp32 because dkx/dvx are atomically - // scatter-accumulated (reduced-precision atomics would lose precision); - // they are cast back to the input dtype at the end. - // - // Scatter (upsample) path: native storage as well — inputs stay in their - // dtype (widened at load inside the dispatch), gradients stay fp32. - AT_DISPATCH_FLOATING_TYPES_AND2(at::kHalf, at::kBFloat16, qy.scalar_type(), "s2_attention_bwd_dkvq_cuda", [&] { - using storage_t = scalar_t; - - const auto f32_like - = [](const torch::Tensor &t) { return torch::zeros_like(t, t.options().dtype(torch::kFloat32)); }; - - if (downsample) { - // native-storage gather path - torch::Tensor kxP = kx; - torch::Tensor vxP = vx; - torch::Tensor qyP = qy; - torch::Tensor dyP = dy; - - // safer than is_contiguous(ChannelsLast), which fails for num_channels == 1 - bool kx_is_channels_last = kxP.strides()[1] == 1; - bool vx_is_channels_last = vxP.strides()[1] == 1; - bool qy_is_channels_last = qyP.strides()[1] == 1; - bool dy_is_channels_last = dyP.strides()[1] == 1; - - if (!kx_is_channels_last) { kxP = permute_4D_to0231(kxP); } - if (!vx_is_channels_last) { vxP = permute_4D_to0231(vxP); } - if (!qy_is_channels_last) { qyP = permute_4D_to0231(qyP); } - if (!dy_is_channels_last) { dyP = permute_4D_to0231(dyP); } - - // fp32 gradient buffers (same shape/layout as the native inputs) - torch::Tensor dkxP = f32_like(kxP); - torch::Tensor dvxP = f32_like(vxP); - torch::Tensor dqyP = f32_like(qyP); - - s2_attn_bwd_dispatch(batch_size, nchans_in, nchans_out, nlon_in, nlat_out, nlon_out, kxP, - vxP, qyP, dyP, psi_row_off, psi_col_idx, quad_weights, dkxP, dvxP, - dqyP); - - dkx = dkxP; - dvx = dvxP; - dqy = dqyP; - - if (!kx_is_channels_last) { dkx = permute_4D_to0312(dkx); } - if (!vx_is_channels_last) { dvx = permute_4D_to0312(dvx); } - if (!qy_is_channels_last) { dqy = permute_4D_to0312(dqy); } - } else { - // native-storage scatter (upsample) path - torch::Tensor kxP = kx; - torch::Tensor vxP = vx; - torch::Tensor qyP = qy; - torch::Tensor dyP = dy; - - bool kx_is_channels_last = kxP.strides()[1] == 1; - bool vx_is_channels_last = vxP.strides()[1] == 1; - bool qy_is_channels_last = qyP.strides()[1] == 1; - bool dy_is_channels_last = dyP.strides()[1] == 1; - - if (!kx_is_channels_last) { kxP = permute_4D_to0231(kxP); } - if (!vx_is_channels_last) { vxP = permute_4D_to0231(vxP); } - if (!qy_is_channels_last) { qyP = permute_4D_to0231(qyP); } - if (!dy_is_channels_last) { dyP = permute_4D_to0231(dyP); } - - torch::Tensor dkxP = f32_like(kxP); - torch::Tensor dvxP = f32_like(vxP); - torch::Tensor dqyP = f32_like(qyP); - - s2_attn_bwd_upsample_dispatch(batch_size, nchans_in, nchans_out, nlon_in, nlat_in, nlat_out, - nlon_out, kxP, vxP, qyP, dyP, psi_row_off, psi_col_idx, quad_weights, - dkxP, dvxP, dqyP); - - dkx = dkxP; - dvx = dvxP; - dqy = dqyP; - - if (!kx_is_channels_last) { dkx = permute_4D_to0312(dkx); } - if (!vx_is_channels_last) { dvx = permute_4D_to0312(dvx); } - if (!qy_is_channels_last) { dqy = permute_4D_to0312(dqy); } - } - }); - - C10_CUDA_KERNEL_LAUNCH_CHECK(); - - // convert precision back to starting dtype (no-op for fp32; narrows for fp16/bf16) - dkx = dkx.to(kx_type); - dvx = dvx.to(vx_type); - dqy = dqy.to(qy_type); - - return std::make_tuple(dkx, dvx, dqy); - } - - TORCH_LIBRARY_IMPL(attention_kernels, CUDA, m) { m.impl("backward_fast", &s2_attention_bwd_dkvq_cuda); } - - } // namespace fast - -} // namespace attention_kernels diff --git a/torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_fwd_coop.cu b/torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_fwd_coop.cu deleted file mode 100644 index 5f28c1cf..00000000 --- a/torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_fwd_coop.cu +++ /dev/null @@ -1,321 +0,0 @@ -// coding=utf-8 -// -// SPDX-FileCopyrightText: Copyright (c) 2026 The torch-harmonics Authors. All rights reserved. -// SPDX-License-Identifier: BSD-3-Clause -// -// Redistribution and use in source and binary forms, with or without -// modification, are permitted provided that the following conditions are met: -// -// 1. Redistributions of source code must retain the above copyright notice, this -// list of conditions and the following disclaimer. -// -// 2. Redistributions in binary form must reproduce the above copyright notice, -// this list of conditions and the following disclaimer in the documentation -// and/or other materials provided with the distribution. -// -// 3. Neither the name of the copyright holder nor the names of its -// contributors may be used to endorse or promote products derived from -// this software without specific prior written permission. -// -// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" -// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE -// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE -// DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE -// FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL -// DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR -// SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER -// CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, -// OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE -// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. - -// ===================================================================================== -// Gather / self attention forward — cross-wo COOPERATIVE TILING (sm_80+, no WGMMA/TMA). -// -// Motivation (H100 ncu, 180x360 C=256 tc=0.03): the baseline s2_attn_fwd_special_vec_k -// is latency-bound on the col_idx-driven K/V gather (long_scoreboard #1, L1 hit ~52%). -// Adjacent output wo attend nearly the SAME input K/V window (shifted by pscale*wo), so -// each (ho,wo) warp re-gathers redundantly. This kernel tiles WO_TILE consecutive wo of -// one output row into a block; the block stages the UNION K/V longitude arc once into -// shared and every wo reuses it -> ~WO_TILE fewer global loads on the polar rows that -// dominate the work. -// -// Occupancy-conscious (Mauro's TMA attempt tanked occupancy -> much worse): accumulators -// (locy) + query (locq) stay in REGISTERS; only the K/V arc is shared, streamed in -// CHUNK-wide longitude tiles so the shared footprint is small. -// -// Geometry: col_idx is sorted by input lat hi; each (ho,hi) neighbor set is a contiguous -// signed longitude arc [-r, +r] centered on the output longitude (a geodesic disc, no -// holes). For output wo the arc shifts by pscale*wo. Per hi we reduce r over the row's -// neighbors, stage the tile-union arc, and each wo tests membership by de-shifting a -// staged cell back to canonical signed wi and checking |swi| <= r. -// -// Uses accurate expf (NOT __expf) so an A/B vs baseline isolates the TILING effect. -// Assumes nchan_in == nchan_out (self / square-channel gather). fp16/bf16/fp32 storage, -// fp32 compute. Registered as the ``forward_coop`` op. -// ===================================================================================== - -#include "attention_cuda.cuh" -#include -#include -#include -#include - -#include -#include -#include - -#include "cudamacro.h" -#include "attention_cuda_utils.cuh" - -namespace attention_kernels -{ - - // compacted-row -> ho map, defined in attention_cuda_utils.cu. - // (declared in attention_cuda_utils.cuh: sortRows) - - namespace coop - { - - static constexpr int WARP = 32; - static constexpr int WO_TILE = 4; // output wo positions (warps) per block - static constexpr int CHUNK = 16; // input longitudes staged per shared tile - - __device__ __forceinline__ int cmod(int a, int n) { return ((a % n) + n) % n; } - // map a longitude in [0,n) to signed [-n/2, n/2) so a pole-centred window is contiguous - __device__ __forceinline__ int csigned(int wi, int n) { return (wi >= (n >> 1)) ? (wi - n) : wi; } - - // block-wide max over ints (blockDim.x = WO_TILE*WARP, WO_TILE <= 32). - __device__ __forceinline__ int block_max_int(int v, int *sh) - { - const int lane = threadIdx.x & (WARP - 1); - const int wid = threadIdx.x >> 5; -#pragma unroll - for (int o = 16; o > 0; o >>= 1) { v = max(v, __shfl_down_sync(0xffffffffu, v, o)); } - if (lane == 0) { sh[wid] = v; } - __syncthreads(); - v = (threadIdx.x < (blockDim.x >> 5)) ? sh[lane] : INT_MIN; - if (wid == 0) { -#pragma unroll - for (int o = 16; o > 0; o >>= 1) { v = max(v, __shfl_down_sync(0xffffffffu, v, o)); } - if (lane == 0) { sh[0] = v; } - } - __syncthreads(); - return sh[0]; - } - - // NLOC = ceil(nchan / WARP): per-lane register slots for query / output accumulator. - template - __global__ __launch_bounds__(WO_TILE *WARP) void s2_attn_fwd_coop_k( - int nchan, int nlat_in, int nlon_in, int nlat_out, int nlon_out, int tiles_per_row, - const STORAGE_T *__restrict__ kx, const STORAGE_T *__restrict__ vx, const STORAGE_T *__restrict__ qy, - const int32_t *__restrict__ row_idx, const int64_t *__restrict__ row_off, - const int64_t *__restrict__ col_idx, const float *__restrict__ quad_weights, STORAGE_T *__restrict__ y) - { - using COMPUTE_T = typename vec_traits::compute_t; // float for the scalar path - - const int lane = threadIdx.x & (WARP - 1); - const int warp = threadIdx.x >> 5; - - const int batch = blockIdx.y; - const int brow = blockIdx.x / tiles_per_row; - const int wtile = blockIdx.x - brow * tiles_per_row; - if (brow >= nlat_out) { return; } - - const int ho = row_idx[brow]; - const int pscale = nlon_in / nlon_out; - const int wo = wtile * WO_TILE + warp; - const bool valid = (wo < nlon_out); - - extern __shared__ __align__(16) char smem[]; - STORAGE_T *shK = reinterpret_cast(smem); // [CHUNK * nchan] - STORAGE_T *shV = shK + size_t(CHUNK) * nchan; // [CHUNK * nchan] - int *shred = reinterpret_cast(shV + size_t(CHUNK) * nchan); // [WO_TILE] reduction scratch - - // query for this wo, resident in registers (widened once). - COMPUTE_T locq[NLOC]; - COMPUTE_T locy[NLOC]; -#pragma unroll - for (int i = 0; i < NLOC; i++) { locy[i] = 0.f; } - if (valid) { - const STORAGE_T *_qy = qy + (int64_t(batch) * nlat_out * nlon_out + int64_t(ho) * nlon_out + wo) * nchan; -#pragma unroll - for (int i = 0; i < NLOC; i++) { - const int ch = lane + i * WARP; - locq[i] = (ch < nchan) ? vload(_qy, ch) : COMPUTE_T(0.f); - } - } - - float qmax = -FLT_MAX; - float asum = 0.f; - - const int64_t rbeg = row_off[ho]; - const int64_t rend = row_off[ho + 1]; - const int rlen = int(rend - rbeg); - if (rlen <= 0) { return; } - - const int hi_min = int(col_idx[rbeg] / nlon_in); - const int hi_max = int(col_idx[rend - 1] / nlon_in); - - for (int hi = hi_min; hi <= hi_max; hi++) { - - // radius r for this hi = max |signed wi| over the row's neighbors at this hi. - int loc = -1; - for (int off = threadIdx.x; off < rlen; off += blockDim.x) { - const int64_t c = col_idx[rbeg + off]; - if (int(c / nlon_in) == hi) { loc = max(loc, abs(csigned(int(c % nlon_in), nlon_in))); } - } - const int r = block_max_int(loc, shred); - if (r < 0) { continue; } // no neighbors at this hi (shouldn't happen for a contiguous band) - - const int arc_lo = -r + pscale * (wtile * WO_TILE); // signed abs longitude of arc start - const int arc_w = min(nlon_in, 2 * r + pscale * (WO_TILE - 1) + 1); - const float qw = quad_weights[hi]; - const STORAGE_T *kx_hi = kx + (int64_t(batch) * nlat_in + hi) * nlon_in * nchan; - const STORAGE_T *vx_hi = vx + (int64_t(batch) * nlat_in + hi) * nlon_in * nchan; - - for (int c0 = 0; c0 < arc_w; c0 += CHUNK) { - const int cw = min(CHUNK, arc_w - c0); - - // stage cw input longitudes (K and V) into shared, coalesced along channel. - for (int idx = threadIdx.x; idx < cw * nchan; idx += blockDim.x) { - const int j = idx / nchan; - const int ch = idx - j * nchan; - const int lon = cmod(arc_lo + c0 + j, nlon_in); - const int64_t g = int64_t(lon) * nchan + ch; - shK[j * nchan + ch] = kx_hi[g]; - shV[j * nchan + ch] = vx_hi[g]; - } - __syncthreads(); - - if (valid) { - for (int j = 0; j < cw; j++) { - const int lon = cmod(arc_lo + c0 + j, nlon_in); - const int swi = csigned(cmod(lon - pscale * wo, nlon_in), nlon_in); - if (abs(swi) > r) { continue; } // not a neighbor of this wo (disc membership) - - const STORAGE_T *shk = shK + j * nchan; - const STORAGE_T *shv = shV + j * nchan; - - COMPUTE_T acc = 0.f; -#pragma unroll - for (int i = 0; i < NLOC; i++) { - const int ch = lane + i * WARP; - if (ch < nchan) { acc = __vfma(locq[i], vload(shk, ch), acc); } - } - float qdotk = __warp_sum(acc); - - const float qmax2 = max(qmax, qdotk); - const float alpha = expf(qdotk - qmax2) * qw; - const float corr = expf(qmax - qmax2); - asum = alpha + asum * corr; -#pragma unroll - for (int i = 0; i < NLOC; i++) { - const int ch = lane + i * WARP; - if (ch < nchan) { - locy[i] = __vfma_scale(alpha, vload(shv, ch), __vscale(corr, locy[i])); - } - } - qmax = qmax2; - } - } - __syncthreads(); - } - } - - if (valid) { - const float inv = (asum > 0.f) ? 1.f / asum : 0.f; - STORAGE_T *_y = y + (int64_t(batch) * nlat_out * nlon_out + int64_t(ho) * nlon_out + wo) * nchan; -#pragma unroll - for (int i = 0; i < NLOC; i++) { - const int ch = lane + i * WARP; - if (ch < nchan) { vstore(_y, ch, __vscale(inv, locy[i])); } - } - } - } - - // ---- launcher: pick NLOC = ceil(nchan/WARP) via a compile-time recursion ---- - template - void launch_coop(int nloc, int nchan, int nlat_in, int nlon_in, int nlat_out, int nlon_out, int tiles_per_row, - int batch_size, const STORAGE_T *kx, const STORAGE_T *vx, const STORAGE_T *qy, - const int32_t *row_idx, const int64_t *row_off, const int64_t *col_idx, const float *qw, - STORAGE_T *y, cudaStream_t stream) - { - if (NLOC == nloc) { - dim3 block(WO_TILE * WARP); - dim3 grid(nlat_out * tiles_per_row, batch_size); - size_t shbytes = size_t(2) * CHUNK * nchan * sizeof(STORAGE_T) + WO_TILE * sizeof(int); - s2_attn_fwd_coop_k - <<>>(nchan, nlat_in, nlon_in, nlat_out, nlon_out, tiles_per_row, kx, - vx, qy, row_idx, row_off, col_idx, qw, y); - CHECK_ERROR("s2_attn_fwd_coop_k"); - return; - } - if constexpr (NLOC < 32) { - launch_coop(nloc, nchan, nlat_in, nlon_in, nlat_out, nlon_out, tiles_per_row, - batch_size, kx, vx, qy, row_idx, row_off, col_idx, qw, y, stream); - } - } - - torch::Tensor s2_attention_fwd_coop_cuda(at::Tensor kx, at::Tensor vx, at::Tensor qy, at::Tensor quad_weights, - at::Tensor col_idx, at::Tensor row_off, int64_t nlon_in, - int64_t nlat_out, int64_t nlon_out) - { - CHECK_CUDA_INPUT_TENSOR(kx); - CHECK_CUDA_INPUT_TENSOR(vx); - CHECK_CUDA_INPUT_TENSOR(qy); - CHECK_CUDA_TENSOR(quad_weights); - CHECK_CUDA_TENSOR(col_idx); - CHECK_CUDA_TENSOR(row_off); - - const int64_t nchans_in = qy.size(1); - const int64_t nchans_out = vx.size(1); - TORCH_CHECK(nchans_in == nchans_out, "forward_coop: requires nchan_in == nchan_out (self/square-channel)"); - TORCH_CHECK(nlon_in % nlon_out == 0, "forward_coop: gather path only (nlon_in % nlon_out == 0)"); - const int nchan = int(nchans_in); - TORCH_CHECK(nchan <= 32 * 32, "forward_coop: nchan too large for register accumulator"); - - const int batch_size = kx.size(0); - const int64_t nlat_in = kx.size(2); - auto qy_type = qy.dtype(); - auto stream = at::cuda::getCurrentCUDAStream().stream(); - - at::Tensor row_idx = sortRows(int(nlat_out), row_off, stream); - - // channels-last (NHWC) as the kernels index it. - torch::Tensor kxP = kx, vxP = vx, qyP = qy; - const bool qy_cl = (qyP.strides()[1] == 1); - if (kxP.strides()[1] != 1) { kxP = permute_4D_to0231(kxP); } - if (vxP.strides()[1] != 1) { vxP = permute_4D_to0231(vxP); } - if (!qy_cl) { qyP = permute_4D_to0231(qyP); } - - const int64_t out_dims[] = {batch_size, nlat_out, nlon_out, nchans_out}; - torch::Tensor yP = torch::empty(out_dims, kxP.options()); - - const int tiles_per_row = int((nlon_out + WO_TILE - 1) / WO_TILE); - const int nloc = (nchan + WARP - 1) / WARP; - - const int32_t *_row_idx = reinterpret_cast(row_idx.data_ptr()); - const int64_t *_row_off = reinterpret_cast(row_off.data_ptr()); - const int64_t *_col_idx = reinterpret_cast(col_idx.data_ptr()); - const float *_qw = reinterpret_cast(quad_weights.data_ptr()); - - AT_DISPATCH_FLOATING_TYPES_AND2(at::kHalf, at::kBFloat16, qyP.scalar_type(), "s2_attention_fwd_coop_cuda", [&] { - launch_coop<1, scalar_t>(nloc, nchan, int(nlat_in), int(nlon_in), int(nlat_out), int(nlon_out), - tiles_per_row, batch_size, reinterpret_cast(kxP.data_ptr()), - reinterpret_cast(vxP.data_ptr()), - reinterpret_cast(qyP.data_ptr()), _row_idx, _row_off, - _col_idx, _qw, reinterpret_cast(yP.data_ptr()), stream); - }); - - C10_CUDA_KERNEL_LAUNCH_CHECK(); - - torch::Tensor yout = yP; - if (!qy_cl) { yout = permute_4D_to0312(yout); } - return yout.to(qy_type); - } - - TORCH_LIBRARY_IMPL(attention_kernels, CUDA, m) { m.impl("forward_coop", &s2_attention_fwd_coop_cuda); } - - } // namespace coop - -} // namespace attention_kernels diff --git a/torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_fwd_fast.cu b/torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_fwd_fast.cu deleted file mode 100644 index 918bae19..00000000 --- a/torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_fwd_fast.cu +++ /dev/null @@ -1,622 +0,0 @@ -// coding=utf-8 -// -// SPDX-FileCopyrightText: Copyright (c) 2025 The torch-harmonics Authors. All rights reserved. -// SPDX-License-Identifier: BSD-3-Clause -// -// Redistribution and use in source and binary forms, with or without -// modification, are permitted provided that the following conditions are met: -// -// 1. Redistributions of source code must retain the above copyright notice, this -// list of conditions and the following disclaimer. -// -// 2. Redistributions in binary form must reproduce the above copyright notice, -// this list of conditions and the following disclaimer in the documentation -// and/or other materials provided with the distribution. -// -// 3. Neither the name of the copyright holder nor the names of its -// contributors may be used to endorse or promote products derived from -// this software without specific prior written permission. -// -// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" -// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE -// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE -// DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE -// FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL -// DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR -// SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER -// CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, -// OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE -// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. - -#include "attention_cuda.cuh" -#include -#include -#include -#include -#include -#include -#include - -#include - -#include -#include - -#include "cudamacro.h" -#include "attention_cuda_utils.cuh" - -#define THREADS (64) - -#define MAX_LOCAL_ARR_LEN (16) - -// BEGIN - forward kernels and functions - -namespace attention_kernels -{ - - // scatter-direction launcher, defined in attention_cuda_fwd_upsample.cu; - // called by s2_attention_fwd_cuda when nlon_out % nlon_in == 0. - void s2_attn_fwd_upsample_dispatch(int batch_size, size_t nchans_in, size_t nchans_out, int64_t nlon_in, - int64_t nlat_in, int64_t nlat_out, int64_t nlon_out, torch::Tensor kxP, - torch::Tensor vxP, torch::Tensor qyP, torch::Tensor psi_row_off, - torch::Tensor psi_col_idx, torch::Tensor quad_weights, torch::Tensor yP); - - // ===================================================================================== - // Optimized (FMA + __expf) copy of the GATHER forward, exposed as the ``forward_fast`` - // op for A/B benchmarking against the baseline ``forward``. Everything lives in the - // nested ``fast`` namespace so the template kernels get distinct symbols (no ODR clash - // with the baseline copy). The upsample path delegates to the shared - // attention_kernels::s2_attn_fwd_upsample_dispatch (resolved by enclosing-namespace - // lookup). Kept until the fast path is validated as the default. - // ===================================================================================== - namespace fast - { - - // called with (blockDim.x=32 and blockDim.y>1, BDIM_X=blockDim.x*blockDim.y) - // - // STORAGE_T is the global-memory element type (float4 for the fp32 vectorized - // path; float / c10::Half / c10::BFloat16 for the scalar path). COMPUTE_T is - // the arithmetic type (float4 for the vectorized path, float otherwise) — all - // dot products, softmax and accumulation happen in COMPUTE_T. vload/vstore - // widen/narrow at the memory boundary. - - template - __global__ __launch_bounds__(BDIM_X) void s2_attn_fwd_generic_vec_k( - int nchan_in, // no. of STORAGE_T elements along channel dim - int nchan_out, // no. of STORAGE_T elements along channel dim - int nlat_in, int nlon_in, int nlat_out, int nlon_out, const STORAGE_T *__restrict__ kx, - const STORAGE_T *__restrict__ vx, const STORAGE_T *__restrict__ qy, const int32_t *__restrict__ row_idx, - const int64_t *__restrict__ row_off, const int64_t *__restrict__ col_idx, - const float *__restrict__ quad_weights, STORAGE_T *__restrict__ y) - { - using COMPUTE_T = typename vec_traits::compute_t; - - extern __shared__ __align__(sizeof(float4)) float shext[]; - COMPUTE_T *shy = reinterpret_cast(shext) + threadIdx.y * nchan_out; - - const int batch = blockIdx.y; - const int wid = blockIdx.x * blockDim.y + threadIdx.y; - - if (wid >= nlat_out * nlon_out) { return; } - - const int tidx = threadIdx.x; - - const int h = wid / nlon_out; - const int wo = wid - (h * nlon_out); - const int ho = row_idx[h]; - - // one output lon step corresponds to pscale input lon steps (requires nlon_in % nlon_out == 0) - const int pscale = nlon_in / nlon_out; - - for (int chan = tidx; chan < nchan_out; chan += WARP_SIZE) { shy[chan] = __vset(0.f); } - - kx += int64_t(batch) * nlat_in * nlon_in * nchan_in; - qy += int64_t(batch) * nlat_out * nlon_out * nchan_in + int64_t(ho) * nchan_in * nlon_out - + int64_t(wo) * nchan_in; - - vx += int64_t(batch) * nlat_in * nlon_in * nchan_out; - y += int64_t(batch) * nlat_out * nlon_out * nchan_out + int64_t(ho) * nchan_out * nlon_out - + int64_t(wo) * nchan_out; - - float alpha_sum = 0.0f; - float qdotk_max = -FLT_MAX; - - const int64_t rbeg = row_off[ho]; - const int64_t rend = row_off[ho + 1]; - - col_idx += rbeg; - - const int rlen = rend - rbeg; - - for (int off = 0; off < rlen; off++) { - - const int64_t col = col_idx[off]; - - const int hi = col / nlon_in; - const int wi = col - (hi * nlon_in); - const int wi_wo = wi + pscale * wo; - const int wip = wi_wo - (wi_wo / nlon_in) * nlon_in; - - const STORAGE_T *_kx = kx + int64_t(hi) * nlon_in * nchan_in + int64_t(wip) * nchan_in; - const STORAGE_T *_vx = vx + int64_t(hi) * nlon_in * nchan_out + int64_t(wip) * nchan_out; - - COMPUTE_T qdotkv = __vset(0.f); - - for (int chan = tidx; chan < nchan_in; chan += WARP_SIZE) { - qdotkv = __vfma(vload(qy, chan), vload(_kx, chan), qdotkv); - } - - float qdotk = __warp_sum(__vred(qdotkv)); - - float qdotk_max_tmp; - float alpha; - float exp_save; - - qdotk_max_tmp = max(qdotk_max, qdotk); - alpha = __expf(qdotk - qdotk_max_tmp) * quad_weights[hi]; - exp_save = __expf(qdotk_max - qdotk_max_tmp); - - alpha_sum = alpha + alpha_sum * exp_save; - - for (int chan = tidx; chan < nchan_out; chan += WARP_SIZE) { - shy[chan] = __vfma_scale(alpha, vload(_vx, chan), __vscale(exp_save, shy[chan])); - } - qdotk_max = qdotk_max_tmp; - } - - alpha_sum = 1.0f / alpha_sum; - for (int chan = tidx; chan < nchan_out; chan += WARP_SIZE) { - vstore(y, chan, __vscale(alpha_sum, shy[chan])); - } - - return; - } - - // called with either (BDIM_X=32 and BDIM_Y>1) || (2^K=BDIM_X > 32 and BDIM_Y=1) - template = nchan_out - typename STORAGE_T> - __global__ __launch_bounds__(BDIM_X *BDIM_Y) void s2_attn_fwd_special_vec_k( - int nchan_in, // no. of STORAGE_T elements along channel dim - int nchan_out, // no. of STORAGE_T elements along channel dim - int nlat_in, int nlon_in, int nlat_out, int nlon_out, const STORAGE_T *__restrict__ kx, - const STORAGE_T *__restrict__ vx, const STORAGE_T *__restrict__ qy, const int32_t *__restrict__ row_idx, - const int64_t *__restrict__ row_off, const int64_t *__restrict__ col_idx, - const float *__restrict__ quad_weights, STORAGE_T *__restrict__ y) - { - using COMPUTE_T = typename vec_traits::compute_t; - - static_assert(0 == (BDIM_X & (BDIM_X - 1))); - static_assert(0 == (BDIM_Y & (BDIM_Y - 1))); - static_assert((BDIM_X == 32 && BDIM_Y > 1) || (BDIM_X > 32 && BDIM_Y == 1)); - - constexpr int NLOC_M1 = NLOC - 1; - - const int tidx = threadIdx.x; - const int batch = blockIdx.y; - const int ctaid = blockIdx.x * blockDim.y + threadIdx.y; - - if (ctaid >= nlat_out * nlon_out) { return; } - - COMPUTE_T locy[NLOC]; - - // shq holds q already widened to COMPUTE_T (converted once on load, reused - // across the neighbor loop). - extern __shared__ __align__(sizeof(float4)) float shext[]; - COMPUTE_T *shq = reinterpret_cast(shext) + threadIdx.y * nchan_in; - - if constexpr (CHIN_AS_OUT) { shq += tidx; } - - const int h = ctaid / nlon_out; - const int wo = ctaid - (h * nlon_out); - const int ho = row_idx[h]; - - // one output lon step corresponds to pscale input lon steps (requires nlon_in % nlon_out == 0) - const int pscale = nlon_in / nlon_out; - - kx += int64_t(batch) * nlat_in * nlon_in * nchan_in; - qy += int64_t(batch) * nlat_out * nlon_out * nchan_in + int64_t(ho) * nlon_out * nchan_in - + int64_t(wo) * nchan_in; - if constexpr (CHIN_AS_OUT) { - kx += tidx; - qy += tidx; - } - - vx += int64_t(batch) * nlat_in * nlon_in * nchan_out + tidx; - y += int64_t(batch) * nlat_out * nlon_out * nchan_out + int64_t(ho) * nlon_out * nchan_out - + int64_t(wo) * nchan_out + tidx; - -#pragma unroll - for (int i = 0; i < NLOC; i++) { locy[i] = __vset(0.f); } - - if constexpr (CHIN_AS_OUT) { -#pragma unroll - for (int i = 0; i < NLOC_M1; i++) { shq[i * BDIM_X] = vload(qy, i * BDIM_X); } - if (NLOC_M1 * BDIM_X + tidx < nchan_in) { shq[NLOC_M1 * BDIM_X] = vload(qy, NLOC_M1 * BDIM_X); } - } else { - for (int chan = tidx; chan < nchan_in; chan += BDIM_X) { shq[chan] = vload(qy, chan); } - } - - float alpha_sum = 0.0f; - float qdotk_max = -FLT_MAX; - - const int64_t rbeg = row_off[ho]; - const int64_t rend = row_off[ho + 1]; - - col_idx += rbeg; - - const int rlen = rend - rbeg; - - for (int off = 0; off < rlen; off++) { - - const int64_t col = col_idx[off]; - - const int hi = col / nlon_in; - const int wi = col - (hi * nlon_in); - const int wi_wo = wi + pscale * wo; - const int wip = wi_wo - (wi_wo / nlon_in) * nlon_in; - - const STORAGE_T *_kx = kx + int64_t(hi) * nlon_in * nchan_in + int64_t(wip) * nchan_in; - const STORAGE_T *_vx = vx + int64_t(hi) * nlon_in * nchan_out + int64_t(wip) * nchan_out; - - COMPUTE_T qdotkv = __vset(0.f); - - if constexpr (CHIN_AS_OUT) { -#pragma unroll - for (int i = 0; i < NLOC_M1; i++) { - qdotkv = __vfma(shq[i * BDIM_X], vload(_kx, i * BDIM_X), qdotkv); - } - if (NLOC_M1 * BDIM_X + tidx < nchan_in) { - qdotkv = __vfma(shq[NLOC_M1 * BDIM_X], vload(_kx, NLOC_M1 * BDIM_X), qdotkv); - } - } else { - for (int chan = tidx; chan < nchan_in; chan += BDIM_X) { - qdotkv = __vfma(shq[chan], vload(_kx, chan), qdotkv); - } - } - - float qdotk = __vred(qdotkv); - if constexpr (BDIM_X == 32) { - qdotk = __warp_sum(qdotk); - } else { - qdotk = __block_sum(qdotk); - } - - float qdotk_max_tmp; - float alpha; - float exp_save; - - qdotk_max_tmp = max(qdotk_max, qdotk); - alpha = __expf(qdotk - qdotk_max_tmp) * quad_weights[hi]; - exp_save = __expf(qdotk_max - qdotk_max_tmp); - - alpha_sum = alpha + alpha_sum * exp_save; - -#pragma unroll - for (int i = 0; i < NLOC_M1; i++) { - locy[i] = __vfma_scale(alpha, vload(_vx, i * BDIM_X), __vscale(exp_save, locy[i])); - } - if (NLOC_M1 * BDIM_X + tidx < nchan_out) { - locy[NLOC_M1] = __vfma_scale(alpha, vload(_vx, NLOC_M1 * BDIM_X), __vscale(exp_save, locy[NLOC_M1])); - } - - qdotk_max = qdotk_max_tmp; - } - - alpha_sum = 1.0f / alpha_sum; - -#pragma unroll - for (int i = 0; i < NLOC_M1; i++) { vstore(y, i * BDIM_X, __vscale(alpha_sum, locy[i])); } - if (NLOC_M1 * BDIM_X + tidx < nchan_out) { - vstore(y, NLOC_M1 * BDIM_X, __vscale(alpha_sum, locy[NLOC_M1])); - } - - return; - } - - template - void launch_gen_attn_fwd(int batch_size, int nchans_in, int nchans_out, int nlat_in, int nlon_in, int nlat_out, - int nlon_out, STORAGE_T *__restrict__ _kxp, STORAGE_T *__restrict__ _vxp, - STORAGE_T *__restrict__ _qyp, int32_t *_row_idx, int64_t *_row_off, int64_t *_col_idx, - float *_quad_weights, STORAGE_T *__restrict__ _yp, cudaStream_t stream) - { - - dim3 block(WARP_SIZE, THREADS / WARP_SIZE); - dim3 grid(DIV_UP(nlat_out * nlon_out, block.y), batch_size); - - // shared memory holds compute-type (COMPUTE_T) data, not STORAGE_T. - size_t shsize = sizeof(typename vec_traits::compute_t) * nchans_out * block.y; - - s2_attn_fwd_generic_vec_k - <<>>(nchans_in, nchans_out, nlat_in, nlon_in, nlat_out, nlon_out, _kxp, - _vxp, _qyp, _row_idx, _row_off, _col_idx, _quad_weights, _yp); - CHECK_ERROR("s2_attn_fwd_generic_vec_k"); - - return; - } - - template - void launch_spc_attn_fwd(int nloc, // "BDIM_X*nloc" >= nchans_out - int batch_size, int nchans_in, int nchans_out, int nlat_in, int nlon_in, int nlat_out, - int nlon_out, STORAGE_T *__restrict__ _kxp, STORAGE_T *__restrict__ _vxp, - STORAGE_T *__restrict__ _qyp, int32_t *_row_idx, int64_t *_row_off, int64_t *_col_idx, - float *_quad_weights, STORAGE_T *__restrict__ _yp, cudaStream_t stream) - { - - if (CUR_LOC_SIZE == nloc) { - - dim3 block(BDIM_X, BDIM_Y); - dim3 grid(DIV_UP(nlat_out * nlon_out, block.y), batch_size); - - // shared memory holds compute-type (COMPUTE_T) data, not STORAGE_T. - // block.y > 1 iif block.x==32 - size_t shsize = sizeof(typename vec_traits::compute_t) * nchans_in * block.y; - - // nloc determines the size of local arrays used to store - // y vectors, of length nchans_out; - // if nchans_in is >= BDIM_X*(nloc-1) and <= BDIM_X*nloc - // then we can use the same compile-time known loops used - // for output channels, with the execpetion of testing - // whether to execute the last iteration based on "nchans_in" - // rather than on "nchans_out"; in this way as long as the - // difference between the number of input and output channels - // is <= BDIM_X we can use the faster path - if (nchans_in >= BDIM_X * (CUR_LOC_SIZE - 1) && nchans_in <= BDIM_X * CUR_LOC_SIZE) { - - s2_attn_fwd_special_vec_k<<>>( - nchans_in, nchans_out, nlat_in, nlon_in, nlat_out, nlon_out, _kxp, _vxp, _qyp, _row_idx, - _row_off, _col_idx, _quad_weights, _yp); - } else { - - s2_attn_fwd_special_vec_k<<>>( - nchans_in, nchans_out, nlat_in, nlon_in, nlat_out, nlon_out, _kxp, _vxp, _qyp, _row_idx, - _row_off, _col_idx, _quad_weights, _yp); - } - CHECK_ERROR("s2_attn_fwd_special_vec_k"); - - return; - } - if constexpr (CUR_LOC_SIZE < MAX_LOC_SIZE) { - launch_spc_attn_fwd( - nloc, batch_size, nchans_in, nchans_out, nlat_in, nlon_in, nlat_out, nlon_out, _kxp, _vxp, _qyp, - _row_idx, _row_off, _col_idx, _quad_weights, _yp, stream); - } - return; - } - - // Picks the block size (BDIM_X) instance and launches the gather kernel for a - // given storage vector type SV. MAX_LOC / MIN_LOC bound the per-thread local - // array length (in COMPUTE_T units). nci / nco are channel counts in SV units. - template - static void fwd_dispatch_bdimx(int bdimx, int nloc, int64_t batch_size, int64_t nci, int64_t nco, int nlat_in, - int64_t nlon_in, int64_t nlat_out, int64_t nlon_out, SV *_kxp, SV *_vxp, - SV *_qyp, int32_t *_row_idx, int64_t *_row_off, int64_t *_col_idx, - float *_quad_weights, SV *_yp, cudaStream_t stream) - { - // use 2D blocks only if 32 threads are enough - switch (bdimx) { - case 32: - launch_spc_attn_fwd<32, 2, 1, MAX_LOC>(nloc, batch_size, nci, nco, nlat_in, nlon_in, nlat_out, nlon_out, - _kxp, _vxp, _qyp, _row_idx, _row_off, _col_idx, _quad_weights, - _yp, stream); - break; - case 64: - launch_spc_attn_fwd<64, 1, MIN_LOC, MAX_LOC>(nloc, batch_size, nci, nco, nlat_in, nlon_in, nlat_out, - nlon_out, _kxp, _vxp, _qyp, _row_idx, _row_off, _col_idx, - _quad_weights, _yp, stream); - break; - case 128: - launch_spc_attn_fwd<128, 1, MIN_LOC, MAX_LOC>(nloc, batch_size, nci, nco, nlat_in, nlon_in, nlat_out, - nlon_out, _kxp, _vxp, _qyp, _row_idx, _row_off, _col_idx, - _quad_weights, _yp, stream); - break; - case 256: - launch_spc_attn_fwd<256, 1, MIN_LOC, MAX_LOC>(nloc, batch_size, nci, nco, nlat_in, nlon_in, nlat_out, - nlon_out, _kxp, _vxp, _qyp, _row_idx, _row_off, _col_idx, - _quad_weights, _yp, stream); - break; - case 512: - launch_spc_attn_fwd<512, 1, MIN_LOC, MAX_LOC>(nloc, batch_size, nci, nco, nlat_in, nlon_in, nlat_out, - nlon_out, _kxp, _vxp, _qyp, _row_idx, _row_off, _col_idx, - _quad_weights, _yp, stream); - break; - case 1024: - launch_spc_attn_fwd<1024, 1, MIN_LOC, MAX_LOC>(nloc, batch_size, nci, nco, nlat_in, nlon_in, nlat_out, - nlon_out, _kxp, _vxp, _qyp, _row_idx, _row_off, _col_idx, - _quad_weights, _yp, stream); - break; - default: - launch_gen_attn_fwd(batch_size, nci, nco, nlat_in, nlon_in, nlat_out, nlon_out, _kxp, _vxp, _qyp, - _row_idx, _row_off, _col_idx, _quad_weights, _yp, stream); - break; - } - } - - // Templated on the storage element type (float / c10::Half / c10::BFloat16). - // Path selection (compute / accumulation are fp32 in every case): - // - fp32, 16B-aligned, nchans % 4 == 0 -> float4 vectorized (LDG.128) - // - fp16/bf16, 16B-aligned, nchans % 8 == 0 -> half8/bf168 vectorized (LDG.128) - // - otherwise -> scalar STORAGE_T path - template - static void s2_attn_fwd_dispatch(int64_t batch_size, int64_t nchans_in, int64_t nchans_out, int64_t nlon_in, - int64_t nlat_out, int64_t nlon_out, at::Tensor kxP, at::Tensor vxP, - at::Tensor qyP, at::Tensor row_off, at::Tensor col_idx, - at::Tensor quad_weights, at::Tensor yP) - { - - static_assert(0 == (MAX_LOCAL_ARR_LEN & (MAX_LOCAL_ARR_LEN - 1))); - - // get stream - auto stream = at::cuda::getCurrentCUDAStream().stream(); - - // sort row indices (ho-s) in descending order - // based on (row_off[ho+1]-row_off[ho]) - at::Tensor row_idx = sortRows(nlat_out, row_off, stream); - - const int nlat_in = kxP.size(1); - - // smallest power of two "bdimx" (>=32) s.t. bdimx*MAX_LOCAL_ARR_LEN >= nchans_out - int bdimx; - bdimx = DIV_UP(nchans_out, MAX_LOCAL_ARR_LEN); - bdimx = max(bdimx, WARP_SIZE); - bdimx = next_pow2(bdimx); - - scalar_t *_kxp = reinterpret_cast(kxP.data_ptr()); - scalar_t *_vxp = reinterpret_cast(vxP.data_ptr()); - scalar_t *_qyp = reinterpret_cast(qyP.data_ptr()); - scalar_t *_yp = reinterpret_cast(yP.data_ptr()); - - int32_t *_row_idx = reinterpret_cast(row_idx.data_ptr()); - int64_t *_row_off = reinterpret_cast(row_off.data_ptr()); - int64_t *_col_idx = reinterpret_cast(col_idx.data_ptr()); - float *_quad_weights = reinterpret_cast(quad_weights.data_ptr()); - - constexpr int MIN_LOC_ARR_LEN = MAX_LOCAL_ARR_LEN / 2 + 1; - - if constexpr (std::is_same::value) { - // fp32: float4 vectorized when 16B-aligned + 4-divisible, else scalar. - constexpr int VEC_SIZE = sizeof(float4) / sizeof(float); // 4 - const bool use_vec = is_aligned<16>(_kxp) && is_aligned<16>(_vxp) && is_aligned<16>(_qyp) - && is_aligned<16>(_yp) && (nchans_in % VEC_SIZE) == 0 && (nchans_out % VEC_SIZE) == 0; - - if (use_vec) { - constexpr int MAX_VEC = MAX_LOCAL_ARR_LEN / VEC_SIZE; - constexpr int MIN_VEC = MAX_VEC / 2 + 1; - const int64_t nci = nchans_in / VEC_SIZE; - const int64_t nco = nchans_out / VEC_SIZE; - fwd_dispatch_bdimx( - bdimx, DIV_UP(nco, bdimx), batch_size, nci, nco, nlat_in, nlon_in, nlat_out, nlon_out, - reinterpret_cast(_kxp), reinterpret_cast(_vxp), - reinterpret_cast(_qyp), _row_idx, _row_off, _col_idx, _quad_weights, - reinterpret_cast(_yp), stream); - } else { - fwd_dispatch_bdimx( - bdimx, DIV_UP(nchans_out, bdimx), batch_size, nchans_in, nchans_out, nlat_in, nlon_in, nlat_out, - nlon_out, _kxp, _vxp, _qyp, _row_idx, _row_off, _col_idx, _quad_weights, _yp, stream); - } - } else { - // fp16/bf16: scalar STORAGE_T path (widen at load, narrow at store; fp32 - // compute/accumulation). A vectorized 8-wide path was tried and reverted: - // it raised register pressure and lowered occupancy, and ncu shows this - // kernel is latency/occupancy-bound (DRAM ~25%), not bandwidth-bound, so - // vectorizing reduced precision only hurt. See the AMP refactor notes. - fwd_dispatch_bdimx( - bdimx, DIV_UP(nchans_out, bdimx), batch_size, nchans_in, nchans_out, nlat_in, nlon_in, nlat_out, - nlon_out, _kxp, _vxp, _qyp, _row_idx, _row_off, _col_idx, _quad_weights, _yp, stream); - } - - return; - } - - // END - forward kernels and functions - - torch::Tensor s2_attention_fwd_cuda(at::Tensor kx, at::Tensor vx, at::Tensor qy, at::Tensor quad_weights, - at::Tensor psi_col_idx, at::Tensor psi_row_off, int64_t nlon_in, - int64_t nlat_out, int64_t nlon_out) - { - CHECK_CUDA_INPUT_TENSOR(kx); - CHECK_CUDA_INPUT_TENSOR(vx); - CHECK_CUDA_INPUT_TENSOR(qy); - CHECK_CUDA_TENSOR(quad_weights); - CHECK_CUDA_TENSOR(psi_col_idx); - CHECK_CUDA_TENSOR(psi_row_off); - - // direction selection: gather (self / downsample) iff nlon_in is an integer - // multiple of nlon_out; scatter (upsample) iff nlon_out is an integer multiple - // of nlon_in. Self-attention satisfies both and routes through the gather path. - const bool downsample = (nlon_in % nlon_out == 0); - const bool upsample = (nlon_out % nlon_in == 0); - TORCH_CHECK(downsample || upsample, "either nlon_in (", nlon_in, - ") must be an integer multiple of nlon_out (", nlon_out, "), or vice versa"); - - size_t nchans_in = qy.size(1); // or kx.size(1) - size_t nchans_out = vx.size(1); - - const int batch_size = kx.size(0); - const int64_t nlat_in = kx.size(2); - - // extract dtype - auto qy_type = qy.dtype(); - - const int64_t out_dims[] = {batch_size, nlat_out, nlon_out, nchans_out}; - torch::Tensor y; - - // ATen dispatch over the input dtype. - // - // Gather (downsample / self) path: native storage. kx/vx/qy keep their - // dtype; the kernel widens to fp32 at load and narrows back at store - // (Tier B), so there is no whole-tensor fp32 copy and the read bandwidth - // for fp16/bf16 is halved. Compute/accumulation are fp32 in-kernel. - // - // Scatter (upsample) path: the scatter kernels are not yet on the native - // storage path, so we still upcast to fp32 here and downcast the result - // at the end (the trailing y.to(qy_type)). fp32 inputs are unaffected - // either way (the .to(kFloat32) is a no-op). - AT_DISPATCH_FLOATING_TYPES_AND2(at::kHalf, at::kBFloat16, qy.scalar_type(), "s2_attention_fwd_cuda", [&] { - using storage_t = scalar_t; - - if (downsample) { - // native-storage gather path - torch::Tensor kxP = kx; - torch::Tensor vxP = vx; - torch::Tensor qyP = qy; - - // safer than is_contiguous(ChannelsLast), which fails for num_channels == 1 - bool kx_is_channels_last = kxP.strides()[1] == 1; - bool vx_is_channels_last = vxP.strides()[1] == 1; - bool qy_is_channels_last = qyP.strides()[1] == 1; - - if (!kx_is_channels_last) { kxP = permute_4D_to0231(kxP); } - if (!vx_is_channels_last) { vxP = permute_4D_to0231(vxP); } - if (!qy_is_channels_last) { qyP = permute_4D_to0231(qyP); } - - torch::Tensor yP = torch::empty(out_dims, kxP.options()); // native dtype - - s2_attn_fwd_dispatch(batch_size, nchans_in, nchans_out, nlon_in, nlat_out, nlon_out, kxP, - vxP, qyP, psi_row_off, psi_col_idx, quad_weights, yP); - - y = yP; - if (!qy_is_channels_last) { y = permute_4D_to0312(y); } - } else { - // upsample (scatter) path: native storage. s2_attn_fwd_upsample_dispatch - // does its own AT_DISPATCH and widens fp16/bf16 at load (fp32 compute), - // narrowing the output at store — same as the gather path. - torch::Tensor kxP = kx; - torch::Tensor vxP = vx; - torch::Tensor qyP = qy; - - bool kx_is_channels_last = kxP.strides()[1] == 1; - bool vx_is_channels_last = vxP.strides()[1] == 1; - bool qy_is_channels_last = qyP.strides()[1] == 1; - - if (!kx_is_channels_last) { kxP = permute_4D_to0231(kxP); } - if (!vx_is_channels_last) { vxP = permute_4D_to0231(vxP); } - if (!qy_is_channels_last) { qyP = permute_4D_to0231(qyP); } - - torch::Tensor yP = torch::empty(out_dims, kxP.options()); // native dtype - - s2_attn_fwd_upsample_dispatch(batch_size, nchans_in, nchans_out, nlon_in, nlat_in, nlat_out, - nlon_out, kxP, vxP, qyP, psi_row_off, psi_col_idx, quad_weights, yP); - - y = yP; - if (!qy_is_channels_last) { y = permute_4D_to0312(y); } - } - }); - - // convert precision back to starting dtype. No-op now that both the gather - // and upsample paths produce native-dtype output; kept as a safety net. - y = y.to(qy_type); - - C10_CUDA_KERNEL_LAUNCH_CHECK(); - - return y; - } - - TORCH_LIBRARY_IMPL(attention_kernels, CUDA, m) { m.impl("forward_fast", &s2_attention_fwd_cuda); } - - } // namespace fast - -} // namespace attention_kernels diff --git a/torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_fwd_wgmma_sm90.cu b/torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_fwd_wgmma_sm90.cu deleted file mode 100644 index f21fde62..00000000 --- a/torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_fwd_wgmma_sm90.cu +++ /dev/null @@ -1,530 +0,0 @@ -// coding=utf-8 -// -// SPDX-FileCopyrightText: Copyright (c) 2026 The torch-harmonics Authors. All rights reserved. -// SPDX-License-Identifier: BSD-3-Clause -// -// Redistribution and use in source and binary forms, with or without -// modification, are permitted provided that the following conditions are met: -// -// 1. Redistributions of source code must retain the above copyright notice, this -// list of conditions and the following disclaimer. -// -// 2. Redistributions in binary form must reproduce the above copyright notice, -// this list of conditions and the following disclaimer in the documentation -// and/or other materials provided with the distribution. -// -// 3. Neither the name of the copyright holder nor the names of its -// contributors may be used to endorse or promote products derived from -// this software without specific prior written permission. -// -// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" -// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE -// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE -// DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE -// FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL -// DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR -// SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER -// CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, -// OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE -// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. - -// ===================================================================================== -// Gather / self attention forward — Hopper (sm_90a) WGMMA, fp16, EXPERIMENTAL v2 -// ===================================================================================== -// -// v2: N-TILED ONLINE SOFTMAX (FlashAttention-shaped). Handles arbitrary neighbor -// counts — crucially the polar rows, whose neighborhood spans the full longitude -// dimension (rlen ~ W, and a few x W across rings) because longitudes converge at -// the poles. The earlier single-pass v1 (rlen <= NHALO) could not. -// -// Built on the VALIDATED WGMMA primitives from the DISCO tensor-core branch -// (attention_cuda_ptx.cuh): descriptor, m64n16k16 fp16 mma (predicate scale_D), -// Major::MN A/B core layouts, and the accumulator -> (m,n) epilogue. -// -// Structure (one warpgroup per (ho-row tile of TM queries, batch)): -// O[TM x Cph] (fp32, shared), running max m[TM], running denom l[TM]. -// for n0 in 0..rlen step NT: -// QK^T : S_tile[TM x NT] = Q . Khalo_tile^T (WGMMA, contract Cph in k=16) -// online: m_new = max(m, rowmax(S_tile)); corr = exp(m - m_new); -// l = l*corr + sum exp(S_tile - m_new)*w; O *= corr; P = exp(S - m_new)*w -// PV : O[:, d] += P_tile . Vhalo_tile[:, d] (WGMMA, channel-tiled DTILE) -// finalize: O /= l, narrow to fp16, store. -// -// PROVEN here: WGMMA mechanics + the online-softmax / N-tiling / channel-tiling -// control flow. STILL "VALIDATE ON HOPPER": the per-query longitude-shift HALO + -// MASK (v2 stages neighbors at the tile-base longitude, exact only for rows whose -// neighbor longitudes are wo-independent; the general union-band halo + per-(m,n) -// mask is the remaining keying item). O resident in shared caps Cph ~ 800 on -// Hopper; larger C wants the 2-pass variant (pass A: stats; pass B: PV). -// -// BUILD: TORCH_CUDA_ARCH_LIST="9.0a+PTX". Gated by TORCH_HARMONICS_ATTN_WGMMA=1. -// ===================================================================================== - -#include "attention_cuda.cuh" - -#include -#include -#include -#include -#include -#include -#include -#include -#include - -#include "cudamacro.h" -#include "attention_cuda_utils.cuh" -#include "attention_cuda_ptx.cuh" - -namespace attention_kernels -{ - - static constexpr int WG_THREADS = 128; - static constexpr int TM = 64; // queries per row-tile (WGMMA M) - static constexpr int NT = 16; // neighbor tile width (WGMMA N for QK^T) - static constexpr int DTILE = 16; // channel tile for PV (WGMMA N of the PV GEMM) - static constexpr int KTILE = 16; // WGMMA K - static constexpr int CPH_MAX = 1024; - - // Major::MN core-matrix element indexers (see attention_cuda_ptx.cuh): - // A [M,K] M-fast: (m/8)*128 + (m%8) + 8*k ; B [K,N] N-fast: (n/8)*128 + k*8 + (n%8) - __host__ __device__ __forceinline__ int aIdx(int m, int k) { return (m / 8) * 128 + (m % 8) + 8 * k; } - __host__ __device__ __forceinline__ int bIdx(int k, int n) { return (n / 8) * 128 + k * 8 + (n % 8); } - -#if defined(__CUDA_ARCH_FEAT_SM90_ALL) - - // descriptors are dtype-agnostic (fp16 and bf16 are both 2-byte; the byte strides - // 128/256 are identical), so a void* address is all that's needed. - __device__ __forceinline__ uint64_t descA(const void *p) { return make_wgmma_desc(p, 8 * 16, 16 * 16); } - __device__ __forceinline__ uint64_t descB(const void *p) { return make_wgmma_desc(p, 8 * 16, 16 * 16); } - - // float -> element conversion and the m64n16k16 mma, selected by the storage type. - template __device__ __forceinline__ T f2e(float x); - template <> __device__ __forceinline__ __half f2e<__half>(float x) { return __float2half(x); } - template <> __device__ __forceinline__ __nv_bfloat16 f2e<__nv_bfloat16>(float x) { return __float2bfloat16(x); } - - template __device__ __forceinline__ void wgmma_n16(float (&d)[8], uint64_t da, uint64_t db) - { - if constexpr (std::is_same::value) { - wgmma_m64n16k16_acc_fp16(d, da, db); - } else { - wgmma_m64n16k16_acc_bf16(d, da, db); - } - } - - // m64n16k16 accumulator -> (m,n) writeback (NHALO/NT=16 -> 8 cells/thread). - template - __device__ __forceinline__ void epilogue_n16(const float (&acc)[8], int warp_id, int lane, OP op) - { - const int m01 = warp_id * 16 + (lane >> 2); - const int m23 = m01 + 8; - const int n_a = (lane & 3) * 2; - const int n_b = n_a + 1; - const int ng = 0; // single n-group for N=16 - op(m01, n_a + 8 * ng, acc[ng * 4 + 0]); - op(m01, n_b + 8 * ng, acc[ng * 4 + 1]); - op(m23, n_a + 8 * ng, acc[ng * 4 + 2]); - op(m23, n_b + 8 * ng, acc[ng * 4 + 3]); - const int ng1 = 1; - op(m01, n_a + 8 * ng1, acc[ng1 * 4 + 0]); - op(m01, n_b + 8 * ng1, acc[ng1 * 4 + 1]); - op(m23, n_a + 8 * ng1, acc[ng1 * 4 + 2]); - op(m23, n_b + 8 * ng1, acc[ng1 * 4 + 3]); - } - - // NOTE: no __restrict__ on the parameters — nvcc/cudafe device-stub generation for - // a TEMPLATED __global__ mishandles __restrict__ qualifiers ("template-id ... does - // not match any template declaration"). The non-templated kernels keep __restrict__. - template - __global__ __launch_bounds__(WG_THREADS) void s2_attn_fwd_wgmma_sm90_k(int nchan, int nlat_in, int nlon_in, - int nlat_out, int nlon_out, - int tiles_per_row, const T *kx, const T *vx, - const T *qy, const int32_t *row_idx, - const int64_t *row_off, const int64_t *col_idx, - const float *quad_weights, T *y) - { - const int tid = threadIdx.x; - const int warp_id = tid / 32; - const int lane = tid - warp_id * 32; - - const int batch = blockIdx.y; - const int row = blockIdx.x / tiles_per_row; - const int tile = blockIdx.x - row * tiles_per_row; - const int wo_base = tile * TM; - if (row >= nlat_out || wo_base >= nlon_out) { return; } - - const int ho = row_idx[row]; - const int pscale = nlon_in / nlon_out; - const int tm = min(TM, nlon_out - wo_base); - const int64_t rbeg = row_off[ho]; - const int rlen = static_cast(row_off[ho + 1] - rbeg); - - // ---- shared ---- - extern __shared__ __align__(128) char smem_raw[]; - float *shO = reinterpret_cast(smem_raw); // [TM x nchan] fp32 running output - float *shM = shO + TM * nchan; // [TM] running max - float *shL = shM + TM; // [TM] running denom - float *shStile = shL + TM; // [TM x NT] fp32 scores - T *shQ = reinterpret_cast(shStile + TM * NT); // [TM x KTILE] - T *shK = shQ + TM * KTILE; // [KTILE x NT] - T *shP = shK + KTILE * NT; // [TM x NT] probs (aIdx) - T *shV = shP + TM * NT; // [NT x DTILE] - - for (int i = tid; i < TM * nchan; i += WG_THREADS) { shO[i] = 0.f; } - for (int i = tid; i < TM; i += WG_THREADS) { - shM[i] = -FLT_MAX; - shL[i] = 0.f; - } - __syncthreads(); - - // ---- build the canonical membership mask + latitude band (the HALO). The - // neighborhood spans latitudes [hi_min, hi_max]; per query the longitudes shift - // by pscale*wo but the canonical (wo=0) membership is fixed. We stage the - // full-longitude band halo (absolute input cells, no shift) and mask each query - // by de-shifting a halo column back to canonical and testing membership. ---- - char *member = reinterpret_cast(shV + NT * DTILE); // [nhi <= nlat_in][nlon_in] (after the T tiles) - - __shared__ int sh_himin; - __shared__ int sh_himax; - if (tid == 0) { - sh_himin = nlat_in; - sh_himax = -1; - } - __syncthreads(); - for (int off = tid; off < rlen; off += WG_THREADS) { - const int hi = static_cast(col_idx[rbeg + off] / nlon_in); - // ::atomic* — the global int built-ins; the namespace has a custom - // atomicMax(float*,float) that would otherwise shadow these. - ::atomicMin(&sh_himin, hi); - ::atomicMax(&sh_himax, hi); - } - __syncthreads(); - const int hi_min = sh_himin; - const int nhi = sh_himax - sh_himin + 1; - const int Nhalo = nhi * nlon_in; - - for (int i = tid; i < Nhalo; i += WG_THREADS) { member[i] = 0; } - __syncthreads(); - for (int off = tid; off < rlen; off += WG_THREADS) { - const int64_t col = col_idx[rbeg + off]; - const int hi = static_cast(col / nlon_in); - const int wi = static_cast(col - int64_t(hi) * nlon_in); - member[(hi - hi_min) * nlon_in + wi] = 1; - } - __syncthreads(); - - // halo staging: absolute input cell (hi_min+hi_idx, wi); no per-query shift. - auto stage_halo = [&](const T *src, int nchan_src, int hc, int c) -> T { - const int hi_idx = hc / nlon_in; - const int wi = hc - hi_idx * nlon_in; - const int64_t g = int64_t(batch) * nlat_in * nlon_in * nchan_src - + (int64_t(hi_min + hi_idx) * nlon_in + wi) * nchan_src + c; - return src[g]; - }; - - // ============================ halo N-tile loop (online softmax) =============== - for (int n0 = 0; n0 < Nhalo; n0 += NT) { - const int nt = min(NT, Nhalo - n0); - - // ---- QK^T: S[TM x NT] = Q . Khalo^T (contract nchan, k=16) ---- - float sacc[8]; -#pragma unroll - for (int i = 0; i < 8; i++) { sacc[i] = 0.f; } - for (int k0 = 0; k0 < nchan; k0 += KTILE) { - for (int idx = tid; idx < TM * KTILE; idx += WG_THREADS) { - const int m = idx / KTILE, k = idx - m * KTILE; - T v = f2e(0.f); - if (m < tm) { - const int wo = wo_base + m; - const int64_t g = int64_t(batch) * nlat_out * nlon_out * nchan - + (int64_t(ho) * nlon_out + wo) * nchan + (k0 + k); - v = qy[g]; - } - shQ[aIdx(m, k)] = v; - } - for (int idx = tid; idx < KTILE * NT; idx += WG_THREADS) { - const int k = idx / NT, n = idx - k * NT; - shK[bIdx(k, n)] = (n < nt) ? stage_halo(kx, nchan, n0 + n, k0 + k) : f2e(0.f); - } - // promote the generic-proxy shQ/shK stores to the WGMMA async proxy - // (required; wgmma.fence does not cover shared operands). - fence_proxy_async_shared_cta(); - __syncthreads(); - // synchronous per k-tile: the async mma must finish reading shQ/shK - // (commit+wait) before the next iteration overwrites them. - wgmma_fence(); - wgmma_n16(sacc, descA(shQ), descB(shK)); - wgmma_commit_group(); - wgmma_wait_group<0>(); - __syncthreads(); - } - - for (int i = tid; i < TM * NT; i += WG_THREADS) { shStile[i] = -FLT_MAX; } - __syncthreads(); - epilogue_n16(sacc, warp_id, lane, [&](int m, int n, float vv) { - if (m < tm && n < nt) { shStile[m * NT + n] = vv; } - }); - __syncthreads(); - - // ---- masked online softmax (per query: de-shift halo col -> canonical, - // test membership). The score Q.K at the absolute cell is already correct; - // the mask just selects which halo cells are THIS query's neighbors. ---- - for (int m = tid; m < tm; m += WG_THREADS) { - const int wo = wo_base + m; - const int shift = (pscale * wo) % nlon_in; - float tmax = -FLT_MAX; - for (int n = 0; n < nt; n++) { - const int hc = n0 + n; - const int hi_idx = hc / nlon_in; - int cw = (hc - hi_idx * nlon_in) - shift; - if (cw < 0) cw += nlon_in; - if (member[hi_idx * nlon_in + cw]) { tmax = fmaxf(tmax, shStile[m * NT + n]); } - } - const float m_new = fmaxf(shM[m], tmax); - const float corr = expf(shM[m] - m_new); - float lsum = 0.f; - for (int n = 0; n < nt; n++) { - const int hc = n0 + n; - const int hi_idx = hc / nlon_in; - int cw = (hc - hi_idx * nlon_in) - shift; - if (cw < 0) cw += nlon_in; - float p = 0.f; - if (member[hi_idx * nlon_in + cw]) { - p = expf(shStile[m * NT + n] - m_new) * quad_weights[hi_min + hi_idx]; - } - shP[aIdx(m, n)] = f2e(p); - lsum += p; - } - for (int n = nt; n < NT; n++) { shP[aIdx(m, n)] = f2e(0.f); } - shL[m] = shL[m] * corr + lsum; - for (int c = 0; c < nchan; c++) { shO[m * nchan + c] *= corr; } - shM[m] = m_new; - } - __syncthreads(); - - // ---- PV: O[:, d] += P . Vhalo[:, d] (contract NT halo cells, channel-tiled) ---- - for (int d0 = 0; d0 < nchan; d0 += DTILE) { - for (int idx = tid; idx < NT * DTILE; idx += WG_THREADS) { - const int k = idx / DTILE, n = idx - k * DTILE; - shV[bIdx(k, n)] = (k < nt && (d0 + n) < nchan) ? stage_halo(vx, nchan, n0 + k, d0 + n) : f2e(0.f); - } - // promote the generic-proxy shP (softmax) + shV stores to the WGMMA - // async proxy before the PV mma reads them. - fence_proxy_async_shared_cta(); - __syncthreads(); - float oacc[8]; -#pragma unroll - for (int i = 0; i < 8; i++) { oacc[i] = 0.f; } - wgmma_fence(); - wgmma_n16(oacc, descA(shP), descB(shV)); - wgmma_commit_group(); - wgmma_wait_group<0>(); - epilogue_n16(oacc, warp_id, lane, [&](int m, int n, float vv) { - const int c = d0 + n; - if (m < tm && c < nchan) { shO[m * nchan + c] += vv; } - }); - __syncthreads(); - } - } - - // ---- finalize: O /= l, store ---- - for (int m = tid; m < tm; m += WG_THREADS) { - const float inv = (shL[m] > 0.f) ? 1.f / shL[m] : 0.f; - for (int c = 0; c < nchan; c++) { - const int wo = wo_base + m; - const int64_t g - = int64_t(batch) * nlat_out * nlon_out * nchan + (int64_t(ho) * nlon_out + wo) * nchan + c; - y[g] = f2e(shO[m * nchan + c] * inv); - } - } - } - - // --------------------------------------------------------------------------------- - // DEBUG micro-kernel: D[64x16] = A[64x16] @ B[16x16], one m64n16k16, using the EXACT - // descriptor / aIdx-bIdx staging / epilogue the real kernel uses. A,B,D are plain - // row-major fp16 in global memory. Isolates the WGMMA mechanics from all attention - // logic so they can be unit-tested against a host A@B. One warpgroup, one block. - // --------------------------------------------------------------------------------- - __global__ __launch_bounds__(WG_THREADS) void s2_wgmma_gemm_debug_k(const __half *__restrict__ A, - const __half *__restrict__ B, - __half *__restrict__ D) - { - const int tid = threadIdx.x; - const int warp_id = tid / 32; - const int lane = tid - warp_id * 32; - - // dynamic shared (matches DISCO + the real kernel — GMMA descriptors must - // address dynamic smem the same way; static __shared__ here read as zero). - extern __shared__ __align__(128) char smem_raw[]; - __half *shA = reinterpret_cast<__half *>(smem_raw); // [64 x 16] in aIdx layout - __half *shB = shA + TM * KTILE; // [16 x 16] in bIdx layout - - for (int idx = tid; idx < TM * KTILE; idx += WG_THREADS) { - const int m = idx / KTILE, k = idx - m * KTILE; - shA[aIdx(m, k)] = A[m * KTILE + k]; // A row-major [64,16] - } - for (int idx = tid; idx < KTILE * NT; idx += WG_THREADS) { - const int k = idx / NT, n = idx - k * NT; - shB[bIdx(k, n)] = B[k * NT + n]; // B row-major [16,16] - } - fence_proxy_async_shared_cta(); - __syncthreads(); - - float acc[8]; -#pragma unroll - for (int i = 0; i < 8; i++) { acc[i] = 0.f; } - wgmma_fence(); - wgmma_m64n16k16_acc_fp16(acc, descA(shA), descB(shB)); - wgmma_commit_group(); - wgmma_wait_group<0>(); - - epilogue_n16(acc, warp_id, lane, [&](int m, int n, float v) { D[m * NT + n] = __float2half(v); }); - } - -#else - template - __global__ void s2_attn_fwd_wgmma_sm90_k(int, int, int, int, int, int, const T *, const T *, const T *, - const int32_t *, const int64_t *, const int64_t *, const float *, T *) - { - // matches the real kernel's (unqualified) parameter signature. - } - __global__ void s2_wgmma_gemm_debug_k(const __half *, const __half *, __half *) { } -#endif - - // ----------------------------------------------------------------------------- - // host dispatch — true if launched. Gated: env, sm_90+, fp16, self/gather, - // Cph%16==0, Cph<=CPH_MAX. v2 N-tiles the neighbors so there is NO rlen cap. - // ----------------------------------------------------------------------------- - bool s2_attn_fwd_wgmma_dispatch(int64_t batch_size, int64_t nchans_in, int64_t nchans_out, int64_t nlat_in, - int64_t nlon_in, int64_t nlat_out, int64_t nlon_out, at::Tensor kxP, at::Tensor vxP, - at::Tensor qyP, at::Tensor row_idx, at::Tensor row_off, at::Tensor col_idx, - at::Tensor quad_weights, at::Tensor yP) - { - // diagnostics gated behind TORCH_HARMONICS_WGMMA_VERBOSE so bench runs are quiet. - const char *vv = std::getenv("TORCH_HARMONICS_WGMMA_VERBOSE"); - const bool verbose = (vv != nullptr && std::atoi(vv) != 0); - const auto dt = qyP.scalar_type(); - const bool is_half = (dt == at::kHalf); - const bool is_bf16 = (dt == at::kBFloat16); - if (verbose) { - std::fprintf(stderr, "[torch-harmonics] wgmma_dispatch reached: half=%d bf16=%d nin=%ld nout=%ld major=%d\n", - static_cast(is_half), static_cast(is_bf16), static_cast(nchans_in), - static_cast(nchans_out), at::cuda::getCurrentDeviceProperties()->major); - } - if (!is_half && !is_bf16) { return false; } - if (nchans_in != nchans_out) { return false; } - if ((nchans_in % KTILE) != 0 || nchans_in > CPH_MAX) { return false; } - if (at::cuda::getCurrentDeviceProperties()->major < 9) { return false; } - - if (verbose) { - const int64_t max_rlen = row_off.diff().max().item(); - std::fprintf( - stderr, - "[torch-harmonics] attention WGMMA sm90 kernel LAUNCHED (dtype=%s head_dim=%ld max_row_len=%ld)\n", - is_half ? "fp16" : "bf16", static_cast(nchans_in), static_cast(max_rlen)); - } - - const int nchan = static_cast(nchans_in); - const int tiles_per_row = static_cast(DIV_UP(nlon_out, TM)); - auto stream = at::cuda::getCurrentCUDAStream().stream(); - - const int32_t *_row_idx = reinterpret_cast(row_idx.data_ptr()); - const int64_t *_row_off = reinterpret_cast(row_off.data_ptr()); - const int64_t *_col_idx = reinterpret_cast(col_idx.data_ptr()); - const float *_qw = reinterpret_cast(quad_weights.data_ptr()); - - dim3 block(WG_THREADS); - dim3 grid(nlat_out * tiles_per_row, batch_size); - - // fp16 and bf16 are both 2-byte; the K/V/Q/P tiles are 2 bytes/elem. + membership - // mask (halo): nlat_in*nlon_in bytes (worst-case latitude band). - const size_t sh_bytes = sizeof(float) * (TM * nchan + 2 * TM + TM * NT) - + size_t(2) * (TM * KTILE + KTILE * NT + TM * NT + NT * DTILE) - + static_cast(nlat_in) * static_cast(nlon_in); - // Hopper smem ceiling ~227 KB; fall back if the (nlat_in*nlon_in) membership - // mask makes the tile too large (large grids — the full-lon halo is the v1 - // limit; an arc-band halo would shrink this). - if (sh_bytes > 220000) { return false; } - - // launch the kernel instantiated for the storage type (fp16 or bf16). Written - // out explicitly (not via a generic lambda) — nvcc's device-stub generation - // mishandles launching a templated __global__ from inside an `auto` lambda. - const int nli = static_cast(nlat_in), nlonI = static_cast(nlon_in); - const int nlo = static_cast(nlat_out), nlonO = static_cast(nlon_out); - if (is_half) { - cudaFuncSetAttribute(reinterpret_cast(s2_attn_fwd_wgmma_sm90_k<__half>), - cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast(sh_bytes)); - s2_attn_fwd_wgmma_sm90_k<__half><<>>( - nchan, nli, nlonI, nlo, nlonO, tiles_per_row, reinterpret_cast(kxP.data_ptr()), - reinterpret_cast(vxP.data_ptr()), reinterpret_cast(qyP.data_ptr()), - _row_idx, _row_off, _col_idx, _qw, reinterpret_cast<__half *>(yP.data_ptr())); - } else { - cudaFuncSetAttribute(reinterpret_cast(s2_attn_fwd_wgmma_sm90_k<__nv_bfloat16>), - cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast(sh_bytes)); - s2_attn_fwd_wgmma_sm90_k<__nv_bfloat16><<>>( - nchan, nli, nlonI, nlo, nlonO, tiles_per_row, reinterpret_cast(kxP.data_ptr()), - reinterpret_cast(vxP.data_ptr()), - reinterpret_cast(qyP.data_ptr()), _row_idx, _row_off, _col_idx, _qw, - reinterpret_cast<__nv_bfloat16 *>(yP.data_ptr())); - } - CHECK_ERROR("s2_attn_fwd_wgmma_sm90_k"); - return true; - } - - // ----------------------------------------------------------------------------- - // Standalone forward op for isolated benchmarking of the WGMMA (sm_90a) path. - // Mirrors s2_attention_fwd_cuda's channels-last handling, allocates the output, - // and launches the WGMMA kernel via the dispatch above. Registered as - // attention_kernels::forward_wgmma. Not routed by the module. - // ----------------------------------------------------------------------------- - torch::Tensor s2_attn_fwd_wgmma_op(at::Tensor kx, at::Tensor vx, at::Tensor qy, at::Tensor quad_weights, - at::Tensor col_idx, at::Tensor row_off, int64_t nlon_in, int64_t nlat_out, - int64_t nlon_out) - { - CHECK_CUDA_INPUT_TENSOR(kx); - CHECK_CUDA_INPUT_TENSOR(vx); - CHECK_CUDA_INPUT_TENSOR(qy); - CHECK_CUDA_TENSOR(quad_weights); - CHECK_CUDA_TENSOR(col_idx); - CHECK_CUDA_TENSOR(row_off); - - const int batch_size = kx.size(0); - const int64_t nlat_in = kx.size(2); - const size_t nchans_in = qy.size(1); - const size_t nchans_out = vx.size(1); - auto qy_type = qy.dtype(); - auto stream = at::cuda::getCurrentCUDAStream().stream(); - - // compacted-row -> ho map, sorted by row length for load balance — derived - // internally exactly like the production forward (s2_attention_fwd_cuda). - at::Tensor row_idx = sortRows(static_cast(nlat_out), row_off, stream); - - // The WGMMA kernel indexes activations as channels-last (NHWC) contiguous. - torch::Tensor kxP = kx, vxP = vx, qyP = qy; - const bool qy_is_channels_last = (qyP.strides()[1] == 1); - if (kxP.strides()[1] != 1) { kxP = permute_4D_to0231(kxP); } - if (vxP.strides()[1] != 1) { vxP = permute_4D_to0231(vxP); } - if (!qy_is_channels_last) { qyP = permute_4D_to0231(qyP); } - - const int64_t out_dims[] = {batch_size, nlat_out, nlon_out, static_cast(nchans_out)}; - torch::Tensor yP = torch::empty(out_dims, kxP.options()); - - const bool ok = s2_attn_fwd_wgmma_dispatch(batch_size, nchans_in, nchans_out, nlat_in, nlon_in, nlat_out, - nlon_out, kxP, vxP, qyP, row_idx, row_off, col_idx, quad_weights, yP); - TORCH_CHECK(ok, - "forward_wgmma: WGMMA kernel not launched (needs sm_90a build + fp16/bf16 + " - "nchans_in==nchans_out + nchans%16==0 + nchans<=CPH_MAX + arc-band fits smem)"); - - // PyTorch C10 launch check — surfaces any async launch error from the kernel - // (matches s2_attention_fwd_cuda). CHECK_ERROR inside the dispatch already checks - // cudaGetLastError post-launch; this is the canonical torch-side guard. - C10_CUDA_KERNEL_LAUNCH_CHECK(); - - torch::Tensor y = yP; - if (!qy_is_channels_last) { y = permute_4D_to0312(y); } - return y.to(qy_type); - } - - TORCH_LIBRARY_IMPL(attention_kernels, CUDA, m) { m.impl("forward_wgmma", &s2_attn_fwd_wgmma_op); } - -} // namespace attention_kernels diff --git a/torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_ptx.cuh b/torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_ptx.cuh deleted file mode 100644 index 8717fe86..00000000 --- a/torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_ptx.cuh +++ /dev/null @@ -1,177 +0,0 @@ -// coding=utf-8 -// -// SPDX-FileCopyrightText: Copyright (c) 2026 The torch-harmonics Authors. All rights reserved. -// SPDX-License-Identifier: BSD-3-Clause -// -// Redistribution and use in source and binary forms, with or without -// modification, are permitted provided that the following conditions are met: -// -// 1. Redistributions of source code must retain the above copyright notice, this -// list of conditions and the following disclaimer. -// -// 2. Redistributions in binary form must reproduce the above copyright notice, -// this list of conditions and the following disclaimer in the documentation -// and/or other materials provided with the distribution. -// -// 3. Neither the name of the copyright holder nor the names of its -// contributors may be used to endorse or promote products derived from -// this software without specific prior written permission. -// -// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" -// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE -// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE -// DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE -// FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL -// DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR -// SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER -// CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, -// OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE -// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. - -// ===================================================================================== -// PTX inline-asm wrappers for the attention WGMMA path (Hopper sm_90a). -// -// These are extracted VERBATIM (semantics-preserving) from the validated DISCO -// tensor-core implementation (branch tkurth/disco-tc, disco_cuda_ptx.cuh), which -// was debugged on Hopper. Re-derive nothing here — the layout/encoding decisions -// below are the ones that produced correct DISCO results. -// -// Hardware gating: -// cp.async : SM_80+, used only from the Hopper kernel. -// wgmma.* : SM_90a. ptxas REJECTS WGMMA against plain .target sm_90 — it needs -// .target sm_90a. NVCC defines __CUDA_ARCH_FEAT_SM90_ALL only for -// sm_90a, so we gate on that. Build TORCH_CUDA_ARCH_LIST="9.0a+PTX". -// -// Layout conventions (Major::MN, no swizzle): -// A [M=64, K=16] M-fast in the 8x8 core: byte(m,k) = (m/8)*256 + (m%8 + 8*k)*2 -// B [K=16, N] N-fast in the 8x8 core: byte(k,n) = (n/8)*256 + k*16 + (n%8)*2 -// descriptor LBO/SBO are byte strides between 8x8 core matrices in 16-byte units; -// make_wgmma_desc lands `leading` in bits 16-29 and `stride` in bits 32-45. -// A: leading(K-outer)=8 (128B), stride(M-outer)=16 (256B) -// B: leading(K-outer)=8 (128B), stride(N-outer)= 0 (N=8) or 16 (256B, N>=16) -// -// Accumulator fragment (m64nNk16.f32, PTX ISA 9.7.16.5.4): N/2 fp32 cells/thread, -// in n-groups of 4 cells per 8 N-cols. warp w in [0,4), lane l in [0,32): -// cell 4*ng+0: m = w*16 + l/4, n = (l%4)*2 + 8*ng -// cell 4*ng+1: m = w*16 + l/4, n = (l%4)*2 + 1 + 8*ng -// cell 4*ng+2: m = w*16 + l/4 + 8, n = (l%4)*2 + 8*ng -// cell 4*ng+3: m = w*16 + l/4 + 8, n = (l%4)*2 + 1 + 8*ng -// ===================================================================================== - -#pragma once - -#include - -namespace attention_kernels -{ - -#if __CUDA_ARCH__ >= 800 - __device__ __forceinline__ void cp_async_16B(void *smem_dst, const void *gmem_src) - { - unsigned smem_addr = __cvta_generic_to_shared(smem_dst); - asm volatile("cp.async.cg.shared.global [%0], [%1], 16;\n" ::"r"(smem_addr), "l"(gmem_src)); - } - __device__ __forceinline__ void cp_async_commit() { asm volatile("cp.async.commit_group;\n" ::); } - template __device__ __forceinline__ void cp_async_wait_group() - { - asm volatile("cp.async.wait_group %0;\n" ::"n"(N)); - } - __device__ __forceinline__ void cp_async_wait_all() { asm volatile("cp.async.wait_all;\n" ::); } -#endif // __CUDA_ARCH__ >= 800 - -#if defined(__CUDA_ARCH_FEAT_SM90_ALL) - - // 64-bit shared-memory matrix descriptor (PTX ISA 9.7.16.5.1). No swizzle. - __device__ __forceinline__ uint64_t make_wgmma_desc(const void *smem_ptr, uint32_t leading_byte_offset, - uint32_t stride_byte_offset, uint32_t swizzle = 0) - { - unsigned smem_addr = __cvta_generic_to_shared(smem_ptr); - uint64_t desc = 0; - desc |= ((uint64_t)((smem_addr >> 4) & 0x3fffu)) << 0; // start address - desc |= ((uint64_t)((leading_byte_offset >> 4) & 0x3fffu)) << 16; // LBO - desc |= ((uint64_t)((stride_byte_offset >> 4) & 0x3fffu)) << 32; // SBO - desc |= ((uint64_t)(swizzle & 0x3u)) << 62; // swizzle mode - return desc; - } - - __device__ __forceinline__ void wgmma_fence() { asm volatile("wgmma.fence.sync.aligned;\n" ::); } - - // Make shared-memory writes issued through the GENERIC proxy (normal smem stores when - // staging shQ/shK/shP/shV) visible to WGMMA's ASYNC proxy. Producer threads MUST - // execute this after staging and before the __syncthreads that releases the - // WGMMA-consuming warpgroup. wgmma.fence only orders the accumulator registers, NOT - // the shared operands — omitting this is a nondeterministic race (garbage results). - // Verbatim from the debugged DISCO path (disco_cuda_ptx.cuh). - __device__ __forceinline__ void fence_proxy_async_shared_cta() - { - asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); - } - - __device__ __forceinline__ void wgmma_commit_group() { asm volatile("wgmma.commit_group.sync.aligned;\n" ::); } - template __device__ __forceinline__ void wgmma_wait_group() - { - asm volatile("wgmma.wait_group.sync.aligned %0;\n" ::"n"(N)); - } - - // fp16 wrappers (.f32.f16.f16). scale_D is a PREDICATE (not a literal immediate); - // scale_A=scale_B=1, tnspA=tnspB=1 (Major::MN). Accumulating (D += A*B). - __device__ __forceinline__ void wgmma_m64n8k16_acc_fp16(float (&d)[4], uint64_t desc_a, uint64_t desc_b) - { - int32_t scale_D = 1; - asm volatile("{\n" - ".reg .pred p;\n" - "setp.ne.b32 p, %6, 0;\n" - "wgmma.mma_async.sync.aligned.m64n8k16.f32.f16.f16 " - "{%0, %1, %2, %3}, %4, %5, p, %7, %8, %9, %10;\n" - "}\n" - : "+f"(d[0]), "+f"(d[1]), "+f"(d[2]), "+f"(d[3]) - : "l"(desc_a), "l"(desc_b), "r"(scale_D), "n"(1), "n"(1), "n"(1), "n"(1)); - } - - __device__ __forceinline__ void wgmma_m64n16k16_acc_fp16(float (&d)[8], uint64_t desc_a, uint64_t desc_b) - { - int32_t scale_D = 1; - asm volatile("{\n" - ".reg .pred p;\n" - "setp.ne.b32 p, %10, 0;\n" - "wgmma.mma_async.sync.aligned.m64n16k16.f32.f16.f16 " - "{%0, %1, %2, %3, %4, %5, %6, %7}, %8, %9, p, %11, %12, %13, %14;\n" - "}\n" - : "+f"(d[0]), "+f"(d[1]), "+f"(d[2]), "+f"(d[3]), "+f"(d[4]), "+f"(d[5]), "+f"(d[6]), "+f"(d[7]) - : "l"(desc_a), "l"(desc_b), "r"(scale_D), "n"(1), "n"(1), "n"(1), "n"(1)); - } - - // bf16 wrapper (.f32.bf16.bf16) — identical encoding to fp16 except the input dtype. - // Native on sm_90a (Hopper bf16 tensor cores). scale_D predicate, tnspA=tnspB=1. - __device__ __forceinline__ void wgmma_m64n16k16_acc_bf16(float (&d)[8], uint64_t desc_a, uint64_t desc_b) - { - int32_t scale_D = 1; - asm volatile("{\n" - ".reg .pred p;\n" - "setp.ne.b32 p, %10, 0;\n" - "wgmma.mma_async.sync.aligned.m64n16k16.f32.bf16.bf16 " - "{%0, %1, %2, %3, %4, %5, %6, %7}, %8, %9, p, %11, %12, %13, %14;\n" - "}\n" - : "+f"(d[0]), "+f"(d[1]), "+f"(d[2]), "+f"(d[3]), "+f"(d[4]), "+f"(d[5]), "+f"(d[6]), "+f"(d[7]) - : "l"(desc_a), "l"(desc_b), "r"(scale_D), "n"(1), "n"(1), "n"(1), "n"(1)); - } - - __device__ __forceinline__ void wgmma_m64n32k16_acc_fp16(float (&d)[16], uint64_t desc_a, uint64_t desc_b) - { - int32_t scale_D = 1; - asm volatile("{\n" - ".reg .pred p;\n" - "setp.ne.b32 p, %18, 0;\n" - "wgmma.mma_async.sync.aligned.m64n32k16.f32.f16.f16 " - "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15}, " - "%16, %17, p, %19, %20, %21, %22;\n" - "}\n" - : "+f"(d[0]), "+f"(d[1]), "+f"(d[2]), "+f"(d[3]), "+f"(d[4]), "+f"(d[5]), "+f"(d[6]), "+f"(d[7]), - "+f"(d[8]), "+f"(d[9]), "+f"(d[10]), "+f"(d[11]), "+f"(d[12]), "+f"(d[13]), "+f"(d[14]), - "+f"(d[15]) - : "l"(desc_a), "l"(desc_b), "r"(scale_D), "n"(1), "n"(1), "n"(1), "n"(1)); - } - -#endif // __CUDA_ARCH_FEAT_SM90_ALL - -} // namespace attention_kernels diff --git a/torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_utils.cuh b/torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_utils.cuh index 8050b297..105587f8 100644 --- a/torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_utils.cuh +++ b/torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_utils.cuh @@ -159,24 +159,6 @@ namespace attention_kernels return make_float4(s * v.x, s * v.y, s * v.z, s * v.w); } - // Fused multiply-add helpers (single-rounding fmaf). Explicit fusion for the hot - // dot-product and online-softmax accumulate loops (ncu flagged non-fused fp32 pairs). - // __vfma(a,b,c) = a*b + c (elementwise) - // __vfma_scale(s,v,c)= s*v + c (scalar s broadcast over vector v) - __device__ float __forceinline__ __vfma(float a, float b, float c) { return fmaf(a, b, c); } - - __device__ float4 __forceinline__ __vfma(float4 a, float4 b, float4 c) - { - return make_float4(fmaf(a.x, b.x, c.x), fmaf(a.y, b.y, c.y), fmaf(a.z, b.z, c.z), fmaf(a.w, b.w, c.w)); - } - - __device__ float __forceinline__ __vfma_scale(float s, float v, float c) { return fmaf(s, v, c); } - - __device__ float4 __forceinline__ __vfma_scale(float s, float4 v, float4 c) - { - return make_float4(fmaf(s, v.x, c.x), fmaf(s, v.y, c.y), fmaf(s, v.z, c.z), fmaf(s, v.w, c.w)); - } - __device__ float4 __forceinline__ __vdiv(float s, float4 v) { return make_float4(s / v.x, s / v.y, s / v.z, s / v.w); From 86f117112b3eed771e064403ebba86e02faf1583 Mon Sep 17 00:00:00 2001 From: Thorsten Kurth Date: Tue, 14 Jul 2026 22:11:08 -0700 Subject: [PATCH 09/14] removing half.cu --- .../kernels_cuda/attention_cuda_fwd_half.cu | 287 ------------------ 1 file changed, 287 deletions(-) delete mode 100644 torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_fwd_half.cu diff --git a/torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_fwd_half.cu b/torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_fwd_half.cu deleted file mode 100644 index c20aef20..00000000 --- a/torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_fwd_half.cu +++ /dev/null @@ -1,287 +0,0 @@ -// coding=utf-8 -// -// SPDX-FileCopyrightText: Copyright (c) 2026 The torch-harmonics Authors. All rights reserved. -// SPDX-License-Identifier: BSD-3-Clause -// -// Redistribution and use in source and binary forms, with or without -// modification, are permitted provided that the following conditions are met: -// -// 1. Redistributions of source code must retain the above copyright notice, this -// list of conditions and the following disclaimer. -// -// 2. Redistributions in binary form must reproduce the above copyright notice, -// this list of conditions and the following disclaimer in the documentation -// and/or other materials provided with the distribution. -// -// 3. Neither the name of the copyright holder nor the names of its -// contributors may be used to endorse or promote products derived from -// this software without specific prior written permission. -// -// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" -// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE -// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE -// DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE -// FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL -// DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR -// SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER -// CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, -// OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE -// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. - -// ===================================================================================== -// Gather (downsample / self) attention forward — fp16/bf16 PACKED specialization -// ===================================================================================== -// -// A bf16/fp16-specialized variant of the generic gather forward (warp-per-output -// cell). Motivation: stored as a scalar, each c10::Half lands in its own 32-bit -// register (top 16 bits wasted), so a wide fp16 load doubles register pressure -// versus fp32 and tips this (register-limited) kernel into spills — which is why -// the earlier scalar-fp16 "float8" experiment regressed. Here the activations are -// loaded 8-wide (LDG.128) as 4x __half2 / __nv_bfloat162 and the q.k dot product -// runs on packed __hfma2, so two halves share one register. -// -// Precision is unchanged from the scalar path: the q.k accumulator is packed -// half2 only for the handful of per-lane terms (nchan/8/32), then reduced in fp32; -// the softmax and the per-channel value accumulator `shy` stay fp32. Only the -// loads + the dot product are packed; alpha*v accumulates in fp32 as before. -// -// bf16 packed FMA is native on sm_80+; on older arches the traits fall back to a -// float-based fma2 so the TU still compiles, and the host dispatch only routes -// bf16 here when the device is sm_80+ (fp16 packs on Volta+). -// ===================================================================================== - -#include "attention_cuda.cuh" - -#include -#include -#include -#include - -#include "cudamacro.h" -#include "attention_cuda_utils.cuh" - -#define THREADS (64) - -namespace attention_kernels -{ - - // packed-half traits: maps a c10 element type to its 2-wide vector type, an - // 8-wide (16-byte, LDG.128) load struct, and the pack/unpack/fma2 primitives. - template struct pack8; - - template <> struct pack8 { - using v2_t = __half2; - struct v8_t { - __half2 a, b, c, d; - }; - static __device__ __forceinline__ v2_t zero() { return __floats2half2_rn(0.f, 0.f); } - static __device__ __forceinline__ v2_t fma2(v2_t x, v2_t y, v2_t acc) { return __hfma2(x, y, acc); } - static __device__ __forceinline__ float2 to_f2(v2_t v) { return __half22float2(v); } - static __device__ __forceinline__ v2_t pack(float lo, float hi) { return __floats2half2_rn(lo, hi); } - }; - - template <> struct pack8 { - using v2_t = __nv_bfloat162; - struct v8_t { - __nv_bfloat162 a, b, c, d; - }; - static __device__ __forceinline__ v2_t zero() { return __floats2bfloat162_rn(0.f, 0.f); } - static __device__ __forceinline__ v2_t fma2(v2_t x, v2_t y, v2_t acc) - { -#if __CUDA_ARCH__ >= 800 - return __hfma2(x, y, acc); -#else - // pre-sm_80: bf16 packed FMA is not native. This path exists only so the - // TU compiles; the host dispatch never routes bf16 here below sm_80. - const float2 xf = __bfloat1622float2(x); - const float2 yf = __bfloat1622float2(y); - const float2 af = __bfloat1622float2(acc); - return __floats2bfloat162_rn(xf.x * yf.x + af.x, xf.y * yf.y + af.y); -#endif - } - static __device__ __forceinline__ float2 to_f2(v2_t v) { return __bfloat1622float2(v); } - static __device__ __forceinline__ v2_t pack(float lo, float hi) { return __floats2bfloat162_rn(lo, hi); } - }; - - // warp-per-output-cell gather forward, packed fp16/bf16. nchan_in/nchan_out are - // in ELEM units and must be multiples of 8 (checked by the host dispatch). - template - __global__ __launch_bounds__(BDIM_X) void s2_attn_fwd_generic_half2_k( - int nchan_in, int nchan_out, int nlat_in, int nlon_in, int nlat_out, int nlon_out, const ELEM *__restrict__ kx, - const ELEM *__restrict__ vx, const ELEM *__restrict__ qy, const int32_t *__restrict__ row_idx, - const int64_t *__restrict__ row_off, const int64_t *__restrict__ col_idx, - const float *__restrict__ quad_weights, ELEM *__restrict__ y) - { - using PK = pack8; - using v2_t = typename PK::v2_t; - using v8_t = typename PK::v8_t; - - const int nchan8_in = nchan_in >> 3; - const int nchan8_out = nchan_out >> 3; - - // fp32 per-channel value accumulator (precision unchanged from scalar path). - // Each lane owns the 8 contiguous channels of one v8_t; storing them at - // stride SHY_STRIDE (=9) per c8 block makes the per-j cross-lane access - // bank-conflict free (gcd(9,32)==1), since lane t writes shy[c8*9 + j] and - // c8 = t + 32*m -> bank (9*t + j) mod 32 is a bijection over the warp. - constexpr int SHY_STRIDE = 9; - extern __shared__ __align__(sizeof(float4)) float shext[]; - float *shy = shext + threadIdx.y * (SHY_STRIDE * nchan8_out); - - const int batch = blockIdx.y; - const int wid = blockIdx.x * blockDim.y + threadIdx.y; - if (wid >= nlat_out * nlon_out) { return; } - - const int tidx = threadIdx.x; - const int h = wid / nlon_out; - const int wo = wid - (h * nlon_out); - const int ho = row_idx[h]; - const int pscale = nlon_in / nlon_out; - - for (int i = tidx; i < SHY_STRIDE * nchan8_out; i += WARP_SIZE) { shy[i] = 0.f; } - - kx += int64_t(batch) * nlat_in * nlon_in * nchan_in; - qy += int64_t(batch) * nlat_out * nlon_out * nchan_in + int64_t(ho) * nchan_in * nlon_out - + int64_t(wo) * nchan_in; - vx += int64_t(batch) * nlat_in * nlon_in * nchan_out; - y += int64_t(batch) * nlat_out * nlon_out * nchan_out + int64_t(ho) * nchan_out * nlon_out - + int64_t(wo) * nchan_out; - - const v8_t *qy8 = reinterpret_cast(qy); - - float alpha_sum = 0.0f; - float qdotk_max = -FLT_MAX; - - const int64_t rbeg = row_off[ho]; - const int64_t rend = row_off[ho + 1]; - col_idx += rbeg; - const int rlen = static_cast(rend - rbeg); - - for (int off = 0; off < rlen; off++) { - const int64_t col = col_idx[off]; - const int hi = col / nlon_in; - const int wi = col - (hi * nlon_in); - const int wi_wo = wi + pscale * wo; - const int wip = wi_wo - (wi_wo / nlon_in) * nlon_in; - - const v8_t *kx8 - = reinterpret_cast(kx + int64_t(hi) * nlon_in * nchan_in + int64_t(wip) * nchan_in); - const v8_t *vx8 - = reinterpret_cast(vx + int64_t(hi) * nlon_in * nchan_out + int64_t(wip) * nchan_out); - - // q.k: packed half2 FMA per lane, reduced to fp32 across the warp. - v2_t acc = PK::zero(); - for (int c8 = tidx; c8 < nchan8_in; c8 += WARP_SIZE) { - const v8_t q = qy8[c8]; - const v8_t k = kx8[c8]; - acc = PK::fma2(q.a, k.a, acc); - acc = PK::fma2(q.b, k.b, acc); - acc = PK::fma2(q.c, k.c, acc); - acc = PK::fma2(q.d, k.d, acc); - } - const float2 accf = PK::to_f2(acc); - const float qdotk = __warp_sum(accf.x + accf.y); - - const float qdotk_max_tmp = max(qdotk_max, qdotk); - const float alpha = expf(qdotk - qdotk_max_tmp) * quad_weights[hi]; - const float exp_save = expf(qdotk_max - qdotk_max_tmp); - alpha_sum = alpha + alpha_sum * exp_save; - - // value accumulate: load v packed, convert to fp32, accumulate in fp32. - for (int c8 = tidx; c8 < nchan8_out; c8 += WARP_SIZE) { - const v8_t v = vx8[c8]; - const float2 a = PK::to_f2(v.a); - const float2 b = PK::to_f2(v.b); - const float2 c = PK::to_f2(v.c); - const float2 d = PK::to_f2(v.d); - float *sh = shy + c8 * SHY_STRIDE; - sh[0] = exp_save * sh[0] + alpha * a.x; - sh[1] = exp_save * sh[1] + alpha * a.y; - sh[2] = exp_save * sh[2] + alpha * b.x; - sh[3] = exp_save * sh[3] + alpha * b.y; - sh[4] = exp_save * sh[4] + alpha * c.x; - sh[5] = exp_save * sh[5] + alpha * c.y; - sh[6] = exp_save * sh[6] + alpha * d.x; - sh[7] = exp_save * sh[7] + alpha * d.y; - } - qdotk_max = qdotk_max_tmp; - } - - const float inv = 1.0f / alpha_sum; - v8_t *y8 = reinterpret_cast(y); - for (int c8 = tidx; c8 < nchan8_out; c8 += WARP_SIZE) { - const float *sh = shy + c8 * SHY_STRIDE; - v8_t out; - out.a = PK::pack(sh[0] * inv, sh[1] * inv); - out.b = PK::pack(sh[2] * inv, sh[3] * inv); - out.c = PK::pack(sh[4] * inv, sh[5] * inv); - out.d = PK::pack(sh[6] * inv, sh[7] * inv); - y8[c8] = out; - } - } - - // ----------------------------------------------------------------------------- - // host dispatch — called from s2_attn_fwd_dispatch (attention_cuda_fwd.cu) for - // the fp16/bf16 gather/self path. Returns true if it launched the packed kernel, - // false when not eligible (caller then takes the scalar fallback). Eligibility: - // dtype in {Half, BFloat16}, nchans % 8 == 0, 16-byte aligned, and (for bf16) - // device sm_80+. - // ----------------------------------------------------------------------------- - bool s2_attn_fwd_half2_dispatch(int64_t batch_size, int64_t nchans_in, int64_t nchans_out, int64_t nlat_in, - int64_t nlon_in, int64_t nlat_out, int64_t nlon_out, at::Tensor kxP, at::Tensor vxP, - at::Tensor qyP, at::Tensor row_idx, at::Tensor row_off, at::Tensor col_idx, - at::Tensor quad_weights, at::Tensor yP) - { - if ((nchans_in % 8) != 0 || (nchans_out % 8) != 0) { return false; } - - const auto dt = qyP.scalar_type(); - const bool is_half = (dt == at::kHalf); - const bool is_bf16 = (dt == at::kBFloat16); - if (!is_half && !is_bf16) { return false; } - - // bf16 packed FMA is native only on sm_80+; below that, defer to the scalar path. - if (is_bf16 && at::cuda::getCurrentDeviceProperties()->major < 8) { return false; } - - void *kxp = kxP.data_ptr(); - void *vxp = vxP.data_ptr(); - void *qyp = qyP.data_ptr(); - void *yp = yP.data_ptr(); - if (!is_aligned<16>(kxp) || !is_aligned<16>(vxp) || !is_aligned<16>(qyp) || !is_aligned<16>(yp)) { - return false; - } - - auto stream = at::cuda::getCurrentCUDAStream().stream(); - - int32_t *_row_idx = reinterpret_cast(row_idx.data_ptr()); - int64_t *_row_off = reinterpret_cast(row_off.data_ptr()); - int64_t *_col_idx = reinterpret_cast(col_idx.data_ptr()); - float *_quad_weights = reinterpret_cast(quad_weights.data_ptr()); - - dim3 block(WARP_SIZE, THREADS / WARP_SIZE); - dim3 grid(DIV_UP(nlat_out * nlon_out, block.y), batch_size); - // padded fp32 accumulator: 9 floats per 8-channel block (bank-conflict free) - const size_t shsize = sizeof(float) * 9 * (nchans_out / 8) * block.y; - - const int nci = static_cast(nchans_in); - const int nco = static_cast(nchans_out); - const int nli = static_cast(nlat_in); - const int nlonI = static_cast(nlon_in); - const int nlo = static_cast(nlat_out); - const int nlonO = static_cast(nlon_out); - - if (is_half) { - s2_attn_fwd_generic_half2_k<<>>( - nci, nco, nli, nlonI, nlo, nlonO, reinterpret_cast(kxp), - reinterpret_cast(vxp), reinterpret_cast(qyp), _row_idx, _row_off, - _col_idx, _quad_weights, reinterpret_cast(yp)); - } else { - s2_attn_fwd_generic_half2_k<<>>( - nci, nco, nli, nlonI, nlo, nlonO, reinterpret_cast(kxp), - reinterpret_cast(vxp), reinterpret_cast(qyp), _row_idx, - _row_off, _col_idx, _quad_weights, reinterpret_cast(yp)); - } - CHECK_ERROR("s2_attn_fwd_generic_half2_k"); - return true; - } - -} // namespace attention_kernels From 15eaac21149b5b3004bc6ba4d82de1e557333f13 Mon Sep 17 00:00:00 2001 From: Thorsten Kurth Date: Tue, 14 Jul 2026 22:37:53 -0700 Subject: [PATCH 10/14] make backward channel caching aware --- .../kernels_cuda/attention_cuda_bwd.cu | 21 +++++++++++++++---- 1 file changed, 17 insertions(+), 4 deletions(-) diff --git a/torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_bwd.cu b/torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_bwd.cu index 7a086b00..9eccbc3e 100644 --- a/torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_bwd.cu +++ b/torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_bwd.cu @@ -73,6 +73,12 @@ namespace attention_kernels // polar rows of ERA5-class grids; ~12 KB/block at BDIM_Y=2, well under the 48 KB default. static constexpr int BWD_CACHE_CAP = 768; + // Minimum ACTUAL channel count for the pass-2 cache to pay off. Below this, the fixed + // ~12 KB cache shared costs more occupancy than the (channel-scaled) gather it saves. + // Measured H100 break-even ~150 (C=64 regresses ~10%, C=256 wins ~12%); 192 keeps a + // safety margin. Tunable — revisit with more (nchan, dtype) data points. + static constexpr int BWD_CACHE_MIN_NCHAN = 192; + // BEGIN backward kernels and functions // called with (blockDim.x=32 and blockDim.y>1, BDIM=blockDim.x*blockDim.y) @@ -710,11 +716,18 @@ namespace attention_kernels // shared memory holds compute-type (COMPUTE_T) data, not STORAGE_T. // 2 arrays per cta, block.y > 1 iif block.x==32 const size_t base_sh = sizeof(typename vec_traits::compute_t) * (nchans_in + nchans_out) * block.y; - // pass-2 recompute-elimination cache (qdotk/gdotv, float). Enable only when the - // total fits the 48 KB default opt-out shared window; otherwise fall back to the - // recompute path (cache_on = 0) so huge-nchan launches never overflow. const size_t cache_sh = size_t(block.y) * (2 * BWD_CACHE_CAP) * sizeof(float); - const int cache_on = (base_sh + cache_sh <= 48u * 1024u) ? 1 : 0; + + // Enable the pass-2 recompute-elimination cache only when it pays. The benefit + // (gather loads removed) scales with the channel count, but the cache's shared + // footprint (~12 KB) is FIXED -> at small nchan the occupancy tax outweighs the + // saving (measured: H100 C=64 fp16/bf16 ~0.90x regression; C=256 ~1.12x win). + // Gate on the ACTUAL channel count (nchans_in is in STORAGE_T units: x4 for the + // fp32 float4 path). Also require the total to fit the 48 KB default window so + // huge-nchan launches never overflow shared. + constexpr int sv_elems = std::is_same::value ? 4 : 1; + const int64_t actual_nchan = int64_t(nchans_in) * sv_elems; + const int cache_on = (actual_nchan >= BWD_CACHE_MIN_NCHAN && base_sh + cache_sh <= 48u * 1024u) ? 1 : 0; const size_t shsize = base_sh + (cache_on ? cache_sh : 0); // nloc determines the size of local arrays used to store From 1b0342af30b8e1c3a54e877fe99fed96cde75075 Mon Sep 17 00:00:00 2001 From: Thorsten Kurth Date: Tue, 14 Jul 2026 23:04:14 -0700 Subject: [PATCH 11/14] templating cached kernel --- .../kernels_cuda/attention_cuda_bwd.cu | 57 +++++++++++++------ 1 file changed, 41 insertions(+), 16 deletions(-) diff --git a/torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_bwd.cu b/torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_bwd.cu index 9eccbc3e..13096e05 100644 --- a/torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_bwd.cu +++ b/torch_harmonics/attention/optimized/kernels_cuda/attention_cuda_bwd.cu @@ -317,7 +317,10 @@ namespace attention_kernels template = nchan_in - typename STORAGE_T> + typename STORAGE_T, + bool USE_CACHE> // compile-time: true iff the launcher allocated the qdotk/gdotv cache. + // When false, ALL caching code is dead-stripped -> byte-identical to + // the pre-cache baseline (no residual register/occupancy cost). __global__ __launch_bounds__(BDIM_X *BDIM_Y) void s2_attn_bwd_special_vec_k( int nchan_in, // no. of elements along channel dim int nchan_out, // no. of elements along channel dim @@ -330,9 +333,8 @@ namespace attention_kernels const float *__restrict__ quad_weights, typename vec_traits::compute_t *__restrict__ dkx, // [batch][nlat_in][nlon_in][nchan_in] typename vec_traits::compute_t *__restrict__ dvx, // [batch][nlat_in][nlon_in][nchan_out] - typename vec_traits::compute_t *__restrict__ dqy, - int cache_on) // 1 iff the launcher allocated the qdotk/gdotv cache (fits in shared) - { // [batch][nlat_out][nlon_out][nchan_in] + typename vec_traits::compute_t *__restrict__ dqy) + { // [batch][nlat_out][nlon_out][nchan_in] using COMPUTE_T = typename vec_traits::compute_t; static_assert(0 == (BDIM_X & (BDIM_X - 1))); @@ -356,9 +358,14 @@ namespace attention_kernels if constexpr (CHOUT_AS_IN) { sh_dy += tidx; } // per-warp qdotk/gdotv cache (float), after all warps' sh_dy/sh_qy region. - constexpr int VECF = sizeof(COMPUTE_T) / sizeof(float); - float *qdc = shext + BDIM_Y * (nchan_in + nchan_out) * VECF + threadIdx.y * (2 * BWD_CACHE_CAP); - float *gdc = qdc + BWD_CACHE_CAP; + // Guarded by USE_CACHE so the pointers (and all downstream cache code) vanish + // entirely when caching is off — no residual register cost on the baseline path. + float *qdc = nullptr, *gdc = nullptr; + if constexpr (USE_CACHE) { + constexpr int VECF = sizeof(COMPUTE_T) / sizeof(float); + qdc = shext + BDIM_Y * (nchan_in + nchan_out) * VECF + threadIdx.y * (2 * BWD_CACHE_CAP); + gdc = qdc + BWD_CACHE_CAP; + } // for dqy COMPUTE_T loc_k__[NLOC]; @@ -440,8 +447,10 @@ namespace attention_kernels col_idx += rbeg; const int rlen = rend - rbeg; - // cache pass-1 qdotk/gdotv iff the launcher allocated the cache AND the row fits. - const bool use_cache = (cache_on != 0) && (rlen <= BWD_CACHE_CAP); + // cache pass-1 qdotk/gdotv iff enabled at compile time AND the row fits. When + // USE_CACHE is false this folds to a compile-time false -> all `if (use_cache)` + // blocks below (store, sync, pass-2 read) are dead-stripped. + const bool use_cache = USE_CACHE && (rlen <= BWD_CACHE_CAP); // accumulate alpha_sum, integral, and shared stats, // along with a progressively computed qdotk_max. @@ -740,14 +749,30 @@ namespace attention_kernels // instead of "nchans_in"; in this way as long as the // difference between the number of input and output channels // is <= BDIM_X we can use the faster path - if (nchans_out >= BDIM_X * (CUR_LOC_SIZE - 1) && nchans_out <= BDIM_X * CUR_LOC_SIZE) { - s2_attn_bwd_special_vec_k<<>>( - nchans_in, nchans_out, nlat_in, nlon_in, nlat_out, nlon_out, _kxp, _vxp, _qyp, _dyp, _row_idx, - _row_off, _col_idx, _quad_weights, _dkxp, _dvxp, _dqyp, cache_on); + // 4-way instantiation: CHOUT_AS_IN {1,0} x USE_CACHE {true,false}. USE_CACHE is a + // COMPILE-TIME param so the false path dead-strips all caching code (byte-identical + // to the pre-cache baseline -> no residual cost on the common small-nchan configs). + const bool chout = (nchans_out >= BDIM_X * (CUR_LOC_SIZE - 1) && nchans_out <= BDIM_X * CUR_LOC_SIZE); + if (chout && cache_on) { + s2_attn_bwd_special_vec_k + <<>>(nchans_in, nchans_out, nlat_in, nlon_in, nlat_out, nlon_out, _kxp, + _vxp, _qyp, _dyp, _row_idx, _row_off, _col_idx, _quad_weights, + _dkxp, _dvxp, _dqyp); + } else if (chout) { + s2_attn_bwd_special_vec_k + <<>>(nchans_in, nchans_out, nlat_in, nlon_in, nlat_out, nlon_out, _kxp, + _vxp, _qyp, _dyp, _row_idx, _row_off, _col_idx, _quad_weights, + _dkxp, _dvxp, _dqyp); + } else if (cache_on) { + s2_attn_bwd_special_vec_k + <<>>(nchans_in, nchans_out, nlat_in, nlon_in, nlat_out, nlon_out, _kxp, + _vxp, _qyp, _dyp, _row_idx, _row_off, _col_idx, _quad_weights, + _dkxp, _dvxp, _dqyp); } else { - s2_attn_bwd_special_vec_k<<>>( - nchans_in, nchans_out, nlat_in, nlon_in, nlat_out, nlon_out, _kxp, _vxp, _qyp, _dyp, _row_idx, - _row_off, _col_idx, _quad_weights, _dkxp, _dvxp, _dqyp, cache_on); + s2_attn_bwd_special_vec_k + <<>>(nchans_in, nchans_out, nlat_in, nlon_in, nlat_out, nlon_out, _kxp, + _vxp, _qyp, _dyp, _row_idx, _row_off, _col_idx, _quad_weights, + _dkxp, _dvxp, _dqyp); } CHECK_ERROR("s2_attn_bwd_special_vec_k"); From 3bac9ee73f3ca1c0950da116e53ce692fd7bde98 Mon Sep 17 00:00:00 2001 From: Thorsten Kurth Date: Tue, 14 Jul 2026 23:39:56 -0700 Subject: [PATCH 12/14] adding bigger embed dim attention examples --- benchmarks/attention.py | 770 ++++++++++++++++++++++++++++++++++++++++ 1 file changed, 770 insertions(+) diff --git a/benchmarks/attention.py b/benchmarks/attention.py index b208bf58..a8bbf4ff 100644 --- a/benchmarks/attention.py +++ b/benchmarks/attention.py @@ -578,6 +578,776 @@ def _nattn_reference(state): skip_correctness=True, tags=["attention", "neighborhood", "cross"], ), + # ---- C=128 head-dim variants (mirror the C=64 CUDA nattn configs) ---- + dict( + name="nattn_s2_opt_1deg_b1_c128_h1_tc0017_float32_cuda", + device="cuda", + dtype=torch.float32, + batch=1, + channels=128, + num_heads=1, + nlat_in=180, + nlon_in=360, + nlat_out=180, + nlon_out=360, + theta_cutoff=0.017, + optimized=True, + skip_correctness=False, + tags=["attention", "neighborhood", "self"], + ), + dict( + name="nattn_s2_opt_1deg_b1_c128_h1_tc0017_float16_cuda", + device="cuda", + dtype=torch.float16, + batch=1, + channels=128, + num_heads=1, + nlat_in=180, + nlon_in=360, + nlat_out=180, + nlon_out=360, + theta_cutoff=0.017, + optimized=True, + skip_correctness=False, + tags=["attention", "neighborhood", "self"], + ), + dict( + name="nattn_s2_opt_1deg_b1_c128_h1_tc0017_bfloat16_cuda", + device="cuda", + dtype=torch.bfloat16, + batch=1, + channels=128, + num_heads=1, + nlat_in=180, + nlon_in=360, + nlat_out=180, + nlon_out=360, + theta_cutoff=0.017, + optimized=True, + skip_correctness=False, + tags=["attention", "neighborhood", "self"], + ), + dict( + name="nattn_s2_opt_1deg_b1_c128_h1_tc003_float32_cuda", + device="cuda", + dtype=torch.float32, + batch=1, + channels=128, + num_heads=1, + nlat_in=180, + nlon_in=360, + nlat_out=180, + nlon_out=360, + theta_cutoff=0.03, + optimized=True, + skip_correctness=False, + tags=["attention", "neighborhood", "self"], + ), + dict( + name="nattn_s2_opt_1deg_b1_c128_h1_tc003_float16_cuda", + device="cuda", + dtype=torch.float16, + batch=1, + channels=128, + num_heads=1, + nlat_in=180, + nlon_in=360, + nlat_out=180, + nlon_out=360, + theta_cutoff=0.03, + optimized=True, + skip_correctness=False, + tags=["attention", "neighborhood", "self"], + ), + dict( + name="nattn_s2_opt_1deg_b1_c128_h1_tc003_bfloat16_cuda", + device="cuda", + dtype=torch.bfloat16, + batch=1, + channels=128, + num_heads=1, + nlat_in=180, + nlon_in=360, + nlat_out=180, + nlon_out=360, + theta_cutoff=0.03, + optimized=True, + skip_correctness=False, + tags=["attention", "neighborhood", "self"], + ), + dict( + name="nattn_s2_opt_hdeg_b1_c128_h1_tc0017_float32_cuda", + device="cuda", + dtype=torch.float32, + batch=1, + channels=128, + num_heads=1, + nlat_in=360, + nlon_in=720, + nlat_out=360, + nlon_out=720, + theta_cutoff=0.017, + optimized=True, + skip_correctness=True, + tags=["attention", "neighborhood", "self"], + ), + dict( + name="nattn_s2_opt_hdeg_b1_c128_h1_tc0017_float16_cuda", + device="cuda", + dtype=torch.float16, + batch=1, + channels=128, + num_heads=1, + nlat_in=360, + nlon_in=720, + nlat_out=360, + nlon_out=720, + theta_cutoff=0.017, + optimized=True, + skip_correctness=True, + tags=["attention", "neighborhood", "self"], + ), + dict( + name="nattn_s2_opt_hdeg_b1_c128_h1_tc0017_bfloat16_cuda", + device="cuda", + dtype=torch.bfloat16, + batch=1, + channels=128, + num_heads=1, + nlat_in=360, + nlon_in=720, + nlat_out=360, + nlon_out=720, + theta_cutoff=0.017, + optimized=True, + skip_correctness=True, + tags=["attention", "neighborhood", "self"], + ), + dict( + name="nattn_s2_opt_hdeg_b1_c128_h1_tc003_float32_cuda", + device="cuda", + dtype=torch.float32, + batch=1, + channels=128, + num_heads=1, + nlat_in=360, + nlon_in=720, + nlat_out=360, + nlon_out=720, + theta_cutoff=0.03, + optimized=True, + skip_correctness=True, + tags=["attention", "neighborhood", "self"], + ), + dict( + name="nattn_s2_opt_hdeg_b1_c128_h1_tc003_float16_cuda", + device="cuda", + dtype=torch.float16, + batch=1, + channels=128, + num_heads=1, + nlat_in=360, + nlon_in=720, + nlat_out=360, + nlon_out=720, + theta_cutoff=0.03, + optimized=True, + skip_correctness=True, + tags=["attention", "neighborhood", "self"], + ), + dict( + name="nattn_s2_opt_hdeg_b1_c128_h1_tc003_bfloat16_cuda", + device="cuda", + dtype=torch.bfloat16, + batch=1, + channels=128, + num_heads=1, + nlat_in=360, + nlon_in=720, + nlat_out=360, + nlon_out=720, + theta_cutoff=0.03, + optimized=True, + skip_correctness=True, + tags=["attention", "neighborhood", "self"], + ), + dict( + name="nattn_s2_opt_h1deg_b1_c128_h1_tc0017_float32_cuda", + device="cuda", + dtype=torch.float32, + batch=1, + channels=128, + num_heads=1, + nlat_in=360, + nlon_in=720, + nlat_out=180, + nlon_out=360, + theta_cutoff=0.017, + optimized=True, + skip_correctness=True, + tags=["attention", "neighborhood", "cross"], + ), + dict( + name="nattn_s2_opt_h1deg_b1_c128_h1_tc0017_float16_cuda", + device="cuda", + dtype=torch.float16, + batch=1, + channels=128, + num_heads=1, + nlat_in=360, + nlon_in=720, + nlat_out=180, + nlon_out=360, + theta_cutoff=0.017, + optimized=True, + skip_correctness=True, + tags=["attention", "neighborhood", "cross"], + ), + dict( + name="nattn_s2_opt_h1deg_b1_c128_h1_tc0017_bfloat16_cuda", + device="cuda", + dtype=torch.bfloat16, + batch=1, + channels=128, + num_heads=1, + nlat_in=360, + nlon_in=720, + nlat_out=180, + nlon_out=360, + theta_cutoff=0.017, + optimized=True, + skip_correctness=True, + tags=["attention", "neighborhood", "cross"], + ), + dict( + name="nattn_s2_opt_h1deg_b1_c128_h1_tc003_float32_cuda", + device="cuda", + dtype=torch.float32, + batch=1, + channels=128, + num_heads=1, + nlat_in=360, + nlon_in=720, + nlat_out=180, + nlon_out=360, + theta_cutoff=0.03, + optimized=True, + skip_correctness=True, + tags=["attention", "neighborhood", "cross"], + ), + dict( + name="nattn_s2_opt_h1deg_b1_c128_h1_tc003_float16_cuda", + device="cuda", + dtype=torch.float16, + batch=1, + channels=128, + num_heads=1, + nlat_in=360, + nlon_in=720, + nlat_out=180, + nlon_out=360, + theta_cutoff=0.03, + optimized=True, + skip_correctness=True, + tags=["attention", "neighborhood", "cross"], + ), + dict( + name="nattn_s2_opt_h1deg_b1_c128_h1_tc003_bfloat16_cuda", + device="cuda", + dtype=torch.bfloat16, + batch=1, + channels=128, + num_heads=1, + nlat_in=360, + nlon_in=720, + nlat_out=180, + nlon_out=360, + theta_cutoff=0.03, + optimized=True, + skip_correctness=True, + tags=["attention", "neighborhood", "cross"], + ), + dict( + name="nattn_s2_opt_1hdeg_b1_c128_h1_tc0017_float32_cuda", + device="cuda", + dtype=torch.float32, + batch=1, + channels=128, + num_heads=1, + nlat_in=180, + nlon_in=360, + nlat_out=360, + nlon_out=720, + theta_cutoff=0.017, + optimized=True, + skip_correctness=True, + tags=["attention", "neighborhood", "cross"], + ), + dict( + name="nattn_s2_opt_1hdeg_b1_c128_h1_tc0017_float16_cuda", + device="cuda", + dtype=torch.float16, + batch=1, + channels=128, + num_heads=1, + nlat_in=180, + nlon_in=360, + nlat_out=360, + nlon_out=720, + theta_cutoff=0.017, + optimized=True, + skip_correctness=True, + tags=["attention", "neighborhood", "cross"], + ), + dict( + name="nattn_s2_opt_1hdeg_b1_c128_h1_tc0017_bfloat16_cuda", + device="cuda", + dtype=torch.bfloat16, + batch=1, + channels=128, + num_heads=1, + nlat_in=180, + nlon_in=360, + nlat_out=360, + nlon_out=720, + theta_cutoff=0.017, + optimized=True, + skip_correctness=True, + tags=["attention", "neighborhood", "cross"], + ), + dict( + name="nattn_s2_opt_1hdeg_b1_c128_h1_tc003_float32_cuda", + device="cuda", + dtype=torch.float32, + batch=1, + channels=128, + num_heads=1, + nlat_in=180, + nlon_in=360, + nlat_out=360, + nlon_out=720, + theta_cutoff=0.03, + optimized=True, + skip_correctness=True, + tags=["attention", "neighborhood", "cross"], + ), + dict( + name="nattn_s2_opt_1hdeg_b1_c128_h1_tc003_float16_cuda", + device="cuda", + dtype=torch.float16, + batch=1, + channels=128, + num_heads=1, + nlat_in=180, + nlon_in=360, + nlat_out=360, + nlon_out=720, + theta_cutoff=0.03, + optimized=True, + skip_correctness=True, + tags=["attention", "neighborhood", "cross"], + ), + dict( + name="nattn_s2_opt_1hdeg_b1_c128_h1_tc003_bfloat16_cuda", + device="cuda", + dtype=torch.bfloat16, + batch=1, + channels=128, + num_heads=1, + nlat_in=180, + nlon_in=360, + nlat_out=360, + nlon_out=720, + theta_cutoff=0.03, + optimized=True, + skip_correctness=True, + tags=["attention", "neighborhood", "cross"], + ), + # ---- C=256 head-dim variants (mirror the C=64 CUDA nattn configs) ---- + dict( + name="nattn_s2_opt_1deg_b1_c256_h1_tc0017_float32_cuda", + device="cuda", + dtype=torch.float32, + batch=1, + channels=256, + num_heads=1, + nlat_in=180, + nlon_in=360, + nlat_out=180, + nlon_out=360, + theta_cutoff=0.017, + optimized=True, + skip_correctness=False, + tags=["attention", "neighborhood", "self"], + ), + dict( + name="nattn_s2_opt_1deg_b1_c256_h1_tc0017_float16_cuda", + device="cuda", + dtype=torch.float16, + batch=1, + channels=256, + num_heads=1, + nlat_in=180, + nlon_in=360, + nlat_out=180, + nlon_out=360, + theta_cutoff=0.017, + optimized=True, + skip_correctness=False, + tags=["attention", "neighborhood", "self"], + ), + dict( + name="nattn_s2_opt_1deg_b1_c256_h1_tc0017_bfloat16_cuda", + device="cuda", + dtype=torch.bfloat16, + batch=1, + channels=256, + num_heads=1, + nlat_in=180, + nlon_in=360, + nlat_out=180, + nlon_out=360, + theta_cutoff=0.017, + optimized=True, + skip_correctness=False, + tags=["attention", "neighborhood", "self"], + ), + dict( + name="nattn_s2_opt_1deg_b1_c256_h1_tc003_float32_cuda", + device="cuda", + dtype=torch.float32, + batch=1, + channels=256, + num_heads=1, + nlat_in=180, + nlon_in=360, + nlat_out=180, + nlon_out=360, + theta_cutoff=0.03, + optimized=True, + skip_correctness=False, + tags=["attention", "neighborhood", "self"], + ), + dict( + name="nattn_s2_opt_1deg_b1_c256_h1_tc003_float16_cuda", + device="cuda", + dtype=torch.float16, + batch=1, + channels=256, + num_heads=1, + nlat_in=180, + nlon_in=360, + nlat_out=180, + nlon_out=360, + theta_cutoff=0.03, + optimized=True, + skip_correctness=False, + tags=["attention", "neighborhood", "self"], + ), + dict( + name="nattn_s2_opt_1deg_b1_c256_h1_tc003_bfloat16_cuda", + device="cuda", + dtype=torch.bfloat16, + batch=1, + channels=256, + num_heads=1, + nlat_in=180, + nlon_in=360, + nlat_out=180, + nlon_out=360, + theta_cutoff=0.03, + optimized=True, + skip_correctness=False, + tags=["attention", "neighborhood", "self"], + ), + dict( + name="nattn_s2_opt_hdeg_b1_c256_h1_tc0017_float32_cuda", + device="cuda", + dtype=torch.float32, + batch=1, + channels=256, + num_heads=1, + nlat_in=360, + nlon_in=720, + nlat_out=360, + nlon_out=720, + theta_cutoff=0.017, + optimized=True, + skip_correctness=True, + tags=["attention", "neighborhood", "self"], + ), + dict( + name="nattn_s2_opt_hdeg_b1_c256_h1_tc0017_float16_cuda", + device="cuda", + dtype=torch.float16, + batch=1, + channels=256, + num_heads=1, + nlat_in=360, + nlon_in=720, + nlat_out=360, + nlon_out=720, + theta_cutoff=0.017, + optimized=True, + skip_correctness=True, + tags=["attention", "neighborhood", "self"], + ), + dict( + name="nattn_s2_opt_hdeg_b1_c256_h1_tc0017_bfloat16_cuda", + device="cuda", + dtype=torch.bfloat16, + batch=1, + channels=256, + num_heads=1, + nlat_in=360, + nlon_in=720, + nlat_out=360, + nlon_out=720, + theta_cutoff=0.017, + optimized=True, + skip_correctness=True, + tags=["attention", "neighborhood", "self"], + ), + dict( + name="nattn_s2_opt_hdeg_b1_c256_h1_tc003_float32_cuda", + device="cuda", + dtype=torch.float32, + batch=1, + channels=256, + num_heads=1, + nlat_in=360, + nlon_in=720, + nlat_out=360, + nlon_out=720, + theta_cutoff=0.03, + optimized=True, + skip_correctness=True, + tags=["attention", "neighborhood", "self"], + ), + dict( + name="nattn_s2_opt_hdeg_b1_c256_h1_tc003_float16_cuda", + device="cuda", + dtype=torch.float16, + batch=1, + channels=256, + num_heads=1, + nlat_in=360, + nlon_in=720, + nlat_out=360, + nlon_out=720, + theta_cutoff=0.03, + optimized=True, + skip_correctness=True, + tags=["attention", "neighborhood", "self"], + ), + dict( + name="nattn_s2_opt_hdeg_b1_c256_h1_tc003_bfloat16_cuda", + device="cuda", + dtype=torch.bfloat16, + batch=1, + channels=256, + num_heads=1, + nlat_in=360, + nlon_in=720, + nlat_out=360, + nlon_out=720, + theta_cutoff=0.03, + optimized=True, + skip_correctness=True, + tags=["attention", "neighborhood", "self"], + ), + dict( + name="nattn_s2_opt_h1deg_b1_c256_h1_tc0017_float32_cuda", + device="cuda", + dtype=torch.float32, + batch=1, + channels=256, + num_heads=1, + nlat_in=360, + nlon_in=720, + nlat_out=180, + nlon_out=360, + theta_cutoff=0.017, + optimized=True, + skip_correctness=True, + tags=["attention", "neighborhood", "cross"], + ), + dict( + name="nattn_s2_opt_h1deg_b1_c256_h1_tc0017_float16_cuda", + device="cuda", + dtype=torch.float16, + batch=1, + channels=256, + num_heads=1, + nlat_in=360, + nlon_in=720, + nlat_out=180, + nlon_out=360, + theta_cutoff=0.017, + optimized=True, + skip_correctness=True, + tags=["attention", "neighborhood", "cross"], + ), + dict( + name="nattn_s2_opt_h1deg_b1_c256_h1_tc0017_bfloat16_cuda", + device="cuda", + dtype=torch.bfloat16, + batch=1, + channels=256, + num_heads=1, + nlat_in=360, + nlon_in=720, + nlat_out=180, + nlon_out=360, + theta_cutoff=0.017, + optimized=True, + skip_correctness=True, + tags=["attention", "neighborhood", "cross"], + ), + dict( + name="nattn_s2_opt_h1deg_b1_c256_h1_tc003_float32_cuda", + device="cuda", + dtype=torch.float32, + batch=1, + channels=256, + num_heads=1, + nlat_in=360, + nlon_in=720, + nlat_out=180, + nlon_out=360, + theta_cutoff=0.03, + optimized=True, + skip_correctness=True, + tags=["attention", "neighborhood", "cross"], + ), + dict( + name="nattn_s2_opt_h1deg_b1_c256_h1_tc003_float16_cuda", + device="cuda", + dtype=torch.float16, + batch=1, + channels=256, + num_heads=1, + nlat_in=360, + nlon_in=720, + nlat_out=180, + nlon_out=360, + theta_cutoff=0.03, + optimized=True, + skip_correctness=True, + tags=["attention", "neighborhood", "cross"], + ), + dict( + name="nattn_s2_opt_h1deg_b1_c256_h1_tc003_bfloat16_cuda", + device="cuda", + dtype=torch.bfloat16, + batch=1, + channels=256, + num_heads=1, + nlat_in=360, + nlon_in=720, + nlat_out=180, + nlon_out=360, + theta_cutoff=0.03, + optimized=True, + skip_correctness=True, + tags=["attention", "neighborhood", "cross"], + ), + dict( + name="nattn_s2_opt_1hdeg_b1_c256_h1_tc0017_float32_cuda", + device="cuda", + dtype=torch.float32, + batch=1, + channels=256, + num_heads=1, + nlat_in=180, + nlon_in=360, + nlat_out=360, + nlon_out=720, + theta_cutoff=0.017, + optimized=True, + skip_correctness=True, + tags=["attention", "neighborhood", "cross"], + ), + dict( + name="nattn_s2_opt_1hdeg_b1_c256_h1_tc0017_float16_cuda", + device="cuda", + dtype=torch.float16, + batch=1, + channels=256, + num_heads=1, + nlat_in=180, + nlon_in=360, + nlat_out=360, + nlon_out=720, + theta_cutoff=0.017, + optimized=True, + skip_correctness=True, + tags=["attention", "neighborhood", "cross"], + ), + dict( + name="nattn_s2_opt_1hdeg_b1_c256_h1_tc0017_bfloat16_cuda", + device="cuda", + dtype=torch.bfloat16, + batch=1, + channels=256, + num_heads=1, + nlat_in=180, + nlon_in=360, + nlat_out=360, + nlon_out=720, + theta_cutoff=0.017, + optimized=True, + skip_correctness=True, + tags=["attention", "neighborhood", "cross"], + ), + dict( + name="nattn_s2_opt_1hdeg_b1_c256_h1_tc003_float32_cuda", + device="cuda", + dtype=torch.float32, + batch=1, + channels=256, + num_heads=1, + nlat_in=180, + nlon_in=360, + nlat_out=360, + nlon_out=720, + theta_cutoff=0.03, + optimized=True, + skip_correctness=True, + tags=["attention", "neighborhood", "cross"], + ), + dict( + name="nattn_s2_opt_1hdeg_b1_c256_h1_tc003_float16_cuda", + device="cuda", + dtype=torch.float16, + batch=1, + channels=256, + num_heads=1, + nlat_in=180, + nlon_in=360, + nlat_out=360, + nlon_out=720, + theta_cutoff=0.03, + optimized=True, + skip_correctness=True, + tags=["attention", "neighborhood", "cross"], + ), + dict( + name="nattn_s2_opt_1hdeg_b1_c256_h1_tc003_bfloat16_cuda", + device="cuda", + dtype=torch.bfloat16, + batch=1, + channels=256, + num_heads=1, + nlat_in=180, + nlon_in=360, + nlat_out=360, + nlon_out=720, + theta_cutoff=0.03, + optimized=True, + skip_correctness=True, + tags=["attention", "neighborhood", "cross"], + ), ] for cfg in _ATTN_CONFIGS: From c0af9b422d3a46f25095f829ef120beab8f5abfe Mon Sep 17 00:00:00 2001 From: Thorsten Kurth Date: Wed, 15 Jul 2026 00:04:30 -0700 Subject: [PATCH 13/14] adding many-channel tests --- tests/test_attention.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/tests/test_attention.py b/tests/test_attention.py index dba10289..210bb043 100644 --- a/tests/test_attention.py +++ b/tests/test_attention.py @@ -173,6 +173,18 @@ def setUp(self): # gather with QK norm enabled [4, 4, 4, 1, (6, 12), (6, 12), "equiangular", "equiangular", True, torch.float16, 2e-2, 1e-2], [4, 4, 4, 1, (6, 12), (6, 12), "equiangular", "equiangular", True, torch.bfloat16, 5e-2, 5e-2], + # ---- large head-dim: exercises the backward pass-2 cache (USE_CACHE=true, gated + # at per-head nchan >= 192). C=64 configs above all skip it, so these are the ONLY + # coverage of the cached code path. Loose fp16/bf16 tol: the cache is a memory + # optimization (bit-identical math to recompute), so any cache bug corrupts grossly. + [2, 256, 256, 1, (6, 12), (6, 12), "equiangular", "equiangular", False, torch.float32, 1e-5, 1e-3], # cache, fp32 float4 path (VECF=4) + [2, 256, 256, 1, (6, 12), (6, 12), "equiangular", "equiangular", False, torch.float16, 3e-2, 2e-2], # cache, fp16 scalar path (VECF=1) + [2, 256, 256, 1, (6, 12), (6, 12), "equiangular", "equiangular", False, torch.bfloat16, 6e-2, 5e-2], # cache, bf16 scalar path + [2, 192, 192, 1, (6, 12), (6, 12), "equiangular", "equiangular", False, torch.float16, 3e-2, 2e-2], # cache threshold edge (nchan == 192) + [2, 256, 256, 1, (12, 24), (6, 12), "equiangular", "equiangular", False, torch.float16, 3e-2, 2e-2], # cache + downsample (pscale=2) + [2, 256, 192, 1, (6, 12), (6, 12), "equiangular", "equiangular", False, torch.float16, 3e-2, 2e-2], # cache + asym channels (CHOUT_AS_IN=0) + [2, 512, 512, 2, (6, 12), (6, 12), "equiangular", "equiangular", False, torch.float16, 3e-2, 2e-2], # cache + multi-head (per-head nchan=256) + [2, 256, 256, 1, (6, 12), (6, 12), "equiangular", "equiangular", True, torch.float16, 3e-2, 2e-2], # cache + qk-norm ], skip_on_empty=True, ) From 10829d04d351c8f5fc5b16f803f30db4c8c3122d Mon Sep 17 00:00:00 2001 From: Thorsten Kurth Date: Wed, 15 Jul 2026 00:44:05 -0700 Subject: [PATCH 14/14] adding new attention benchmarks --- benchmarks/reference_results.csv | 96 ++++++++++++++++++++++++++++++++ 1 file changed, 96 insertions(+) diff --git a/benchmarks/reference_results.csv b/benchmarks/reference_results.csv index e56db71f..b47061dd 100644 --- a/benchmarks/reference_results.csv +++ b/benchmarks/reference_results.csv @@ -30,6 +30,54 @@ nattn_s2_opt_1hdeg_b1_c64_h1_tc0017_bfloat16_cuda,NVIDIA H100 80GB HBM3,cuda:0,t nattn_s2_opt_1hdeg_b1_c64_h1_tc003_float32_cuda,NVIDIA H100 80GB HBM3,cuda:0,torch.float32,10.986873626708984,19.279705810546876, nattn_s2_opt_1hdeg_b1_c64_h1_tc003_float16_cuda,NVIDIA H100 80GB HBM3,cuda:0,torch.float16,10.898742485046387,19.370265579223634, nattn_s2_opt_1hdeg_b1_c64_h1_tc003_bfloat16_cuda,NVIDIA H100 80GB HBM3,cuda:0,torch.bfloat16,10.873712062835693,19.312300872802734, +nattn_s2_opt_1deg_b1_c128_h1_tc0017_float32_cuda,NVIDIA H100 80GB HBM3,cuda:0,torch.float32,0.8038655936717987,1.5477919936180116, +nattn_s2_opt_1deg_b1_c128_h1_tc0017_float16_cuda,NVIDIA H100 80GB HBM3,cuda:0,torch.float16,0.7235487997531891,1.3159232020378113, +nattn_s2_opt_1deg_b1_c128_h1_tc0017_bfloat16_cuda,NVIDIA H100 80GB HBM3,cuda:0,torch.bfloat16,0.7283968031406403,1.3183295965194701, +nattn_s2_opt_1deg_b1_c128_h1_tc003_float32_cuda,NVIDIA H100 80GB HBM3,cuda:0,torch.float32,1.0556959867477418,2.3022400379180907, +nattn_s2_opt_1deg_b1_c128_h1_tc003_float16_cuda,NVIDIA H100 80GB HBM3,cuda:0,torch.float16,0.9775615990161896,2.1314303874969482, +nattn_s2_opt_1deg_b1_c128_h1_tc003_bfloat16_cuda,NVIDIA H100 80GB HBM3,cuda:0,torch.bfloat16,0.988239997625351,2.1374304056167603, +nattn_s2_opt_hdeg_b1_c128_h1_tc0017_float32_cuda,NVIDIA H100 80GB HBM3,cuda:0,torch.float32,4.015561580657959,9.460294342041015, +nattn_s2_opt_hdeg_b1_c128_h1_tc0017_float16_cuda,NVIDIA H100 80GB HBM3,cuda:0,torch.float16,3.8374847888946535,9.538758468627929, +nattn_s2_opt_hdeg_b1_c128_h1_tc0017_bfloat16_cuda,NVIDIA H100 80GB HBM3,cuda:0,torch.bfloat16,3.8582112073898314,9.542521572113037, +nattn_s2_opt_hdeg_b1_c128_h1_tc003_float32_cuda,NVIDIA H100 80GB HBM3,cuda:0,torch.float32,8.636419200897217,23.62946548461914, +nattn_s2_opt_hdeg_b1_c128_h1_tc003_float16_cuda,NVIDIA H100 80GB HBM3,cuda:0,torch.float16,8.662201595306396,22.204838371276857, +nattn_s2_opt_hdeg_b1_c128_h1_tc003_bfloat16_cuda,NVIDIA H100 80GB HBM3,cuda:0,torch.bfloat16,8.699862384796143,22.21078701019287, +nattn_s2_opt_h1deg_b1_c128_h1_tc0017_float32_cuda,NVIDIA H100 80GB HBM3,cuda:0,torch.float32,2.305340790748596,4.6428159236907955, +nattn_s2_opt_h1deg_b1_c128_h1_tc0017_float16_cuda,NVIDIA H100 80GB HBM3,cuda:0,torch.float16,2.1442016124725343,4.43514232635498, +nattn_s2_opt_h1deg_b1_c128_h1_tc0017_bfloat16_cuda,NVIDIA H100 80GB HBM3,cuda:0,torch.bfloat16,2.1407423973083497,4.435984039306641, +nattn_s2_opt_h1deg_b1_c128_h1_tc003_float32_cuda,NVIDIA H100 80GB HBM3,cuda:0,torch.float32,3.47881920337677,7.917999982833862, +nattn_s2_opt_h1deg_b1_c128_h1_tc003_float16_cuda,NVIDIA H100 80GB HBM3,cuda:0,torch.float16,3.1986655712127687,7.394899225234985, +nattn_s2_opt_h1deg_b1_c128_h1_tc003_bfloat16_cuda,NVIDIA H100 80GB HBM3,cuda:0,torch.bfloat16,3.2089983940124513,7.381545639038086, +nattn_s2_opt_1hdeg_b1_c128_h1_tc0017_float32_cuda,NVIDIA H100 80GB HBM3,cuda:0,torch.float32,5.811430406570435,10.52003526687622, +nattn_s2_opt_1hdeg_b1_c128_h1_tc0017_float16_cuda,NVIDIA H100 80GB HBM3,cuda:0,torch.float16,5.645350408554077,10.68530559539795, +nattn_s2_opt_1hdeg_b1_c128_h1_tc0017_bfloat16_cuda,NVIDIA H100 80GB HBM3,cuda:0,torch.bfloat16,5.651510429382324,10.686198520660401, +nattn_s2_opt_1hdeg_b1_c128_h1_tc003_float32_cuda,NVIDIA H100 80GB HBM3,cuda:0,torch.float32,11.85411205291748,21.88383045196533, +nattn_s2_opt_1hdeg_b1_c128_h1_tc003_float16_cuda,NVIDIA H100 80GB HBM3,cuda:0,torch.float16,11.78884153366089,22.22824935913086, +nattn_s2_opt_1hdeg_b1_c128_h1_tc003_bfloat16_cuda,NVIDIA H100 80GB HBM3,cuda:0,torch.bfloat16,11.767302513122559,22.188777351379393, +nattn_s2_opt_1deg_b1_c256_h1_tc0017_float32_cuda,NVIDIA H100 80GB HBM3,cuda:0,torch.float32,1.3151296019554137,2.448527979850769, +nattn_s2_opt_1deg_b1_c256_h1_tc0017_float16_cuda,NVIDIA H100 80GB HBM3,cuda:0,torch.float16,1.1952351927757263,2.2989312171936036, +nattn_s2_opt_1deg_b1_c256_h1_tc0017_bfloat16_cuda,NVIDIA H100 80GB HBM3,cuda:0,torch.bfloat16,1.201964795589447,2.300534391403198, +nattn_s2_opt_1deg_b1_c256_h1_tc003_float32_cuda,NVIDIA H100 80GB HBM3,cuda:0,torch.float32,1.6019168019294738,3.9608543872833253, +nattn_s2_opt_1deg_b1_c256_h1_tc003_float16_cuda,NVIDIA H100 80GB HBM3,cuda:0,torch.float16,1.5105119943618774,3.6577247858047484, +nattn_s2_opt_1deg_b1_c256_h1_tc003_bfloat16_cuda,NVIDIA H100 80GB HBM3,cuda:0,torch.bfloat16,1.5166207909584046,3.6637247800827026, +nattn_s2_opt_hdeg_b1_c256_h1_tc0017_float32_cuda,NVIDIA H100 80GB HBM3,cuda:0,torch.float32,6.04332160949707,17.621888160705566, +nattn_s2_opt_hdeg_b1_c256_h1_tc0017_float16_cuda,NVIDIA H100 80GB HBM3,cuda:0,torch.float16,5.865270376205444,17.606397247314455, +nattn_s2_opt_hdeg_b1_c256_h1_tc0017_bfloat16_cuda,NVIDIA H100 80GB HBM3,cuda:0,torch.bfloat16,5.865472030639649,17.605395126342774, +nattn_s2_opt_hdeg_b1_c256_h1_tc003_float32_cuda,NVIDIA H100 80GB HBM3,cuda:0,torch.float32,11.208694362640381,42.684652709960936, +nattn_s2_opt_hdeg_b1_c256_h1_tc003_float16_cuda,NVIDIA H100 80GB HBM3,cuda:0,torch.float16,11.872768020629882,41.85445442199707, +nattn_s2_opt_hdeg_b1_c256_h1_tc003_bfloat16_cuda,NVIDIA H100 80GB HBM3,cuda:0,torch.bfloat16,11.852147006988526,41.879529571533205, +nattn_s2_opt_h1deg_b1_c256_h1_tc0017_float32_cuda,NVIDIA H100 80GB HBM3,cuda:0,torch.float32,3.494000029563904,7.7221856117248535, +nattn_s2_opt_h1deg_b1_c256_h1_tc0017_float16_cuda,NVIDIA H100 80GB HBM3,cuda:0,torch.float16,3.5219136476516724,8.236246395111085, +nattn_s2_opt_h1deg_b1_c256_h1_tc0017_bfloat16_cuda,NVIDIA H100 80GB HBM3,cuda:0,torch.bfloat16,3.5142367839813233,8.209958171844482, +nattn_s2_opt_h1deg_b1_c256_h1_tc003_float32_cuda,NVIDIA H100 80GB HBM3,cuda:0,torch.float32,4.861337614059448,13.771324825286865, +nattn_s2_opt_h1deg_b1_c256_h1_tc003_float16_cuda,NVIDIA H100 80GB HBM3,cuda:0,torch.float16,4.984585618972778,13.794476795196534, +nattn_s2_opt_h1deg_b1_c256_h1_tc003_bfloat16_cuda,NVIDIA H100 80GB HBM3,cuda:0,torch.bfloat16,4.950995206832886,13.829500865936279, +nattn_s2_opt_1hdeg_b1_c256_h1_tc0017_float32_cuda,NVIDIA H100 80GB HBM3,cuda:0,torch.float32,8.32130241394043,16.14897289276123, +nattn_s2_opt_1hdeg_b1_c256_h1_tc0017_float16_cuda,NVIDIA H100 80GB HBM3,cuda:0,torch.float16,7.883110380172729,15.865196704864502, +nattn_s2_opt_1hdeg_b1_c256_h1_tc0017_bfloat16_cuda,NVIDIA H100 80GB HBM3,cuda:0,torch.bfloat16,7.878681612014771,15.881715106964112, +nattn_s2_opt_1hdeg_b1_c256_h1_tc003_float32_cuda,NVIDIA H100 80GB HBM3,cuda:0,torch.float32,15.970498943328858,32.8452091217041, +nattn_s2_opt_1hdeg_b1_c256_h1_tc003_float16_cuda,NVIDIA H100 80GB HBM3,cuda:0,torch.float16,15.494704055786134,32.05101089477539, +nattn_s2_opt_1hdeg_b1_c256_h1_tc003_bfloat16_cuda,NVIDIA H100 80GB HBM3,cuda:0,torch.bfloat16,15.523945617675782,32.103951644897464, nattn_s2_opt_1deg_b1_c64_h1_tc0017_float32_cpu,Neoverse-V2,cpu,torch.float32,12.204264802858233,50.22553245071322, nattn_s2_opt_1deg_b1_c64_h1_tc0017_float32_cuda,NVIDIA GB200,cuda:0,torch.float32,0.642623993754387,0.8532544016838074, nattn_s2_opt_1deg_b1_c64_h1_tc0017_float16_cuda,NVIDIA GB200,cuda:0,torch.float16,0.712169599533081,0.8436768054962158, @@ -55,6 +103,54 @@ nattn_s2_opt_1hdeg_b1_c64_h1_tc0017_bfloat16_cuda,NVIDIA GB200,cuda:0,torch.bflo nattn_s2_opt_1hdeg_b1_c64_h1_tc003_float32_cuda,NVIDIA GB200,cuda:0,torch.float32,9.419171237945557,17.70709114074707, nattn_s2_opt_1hdeg_b1_c64_h1_tc003_float16_cuda,NVIDIA GB200,cuda:0,torch.float16,9.33152961730957,17.731641578674317, nattn_s2_opt_1hdeg_b1_c64_h1_tc003_bfloat16_cuda,NVIDIA GB200,cuda:0,torch.bfloat16,9.338521575927734,17.681894302368164, +nattn_s2_opt_1deg_b1_c128_h1_tc0017_float32_cuda,NVIDIA GB200,cuda:0,torch.float32,0.6464288115501404,0.8982656002044678, +nattn_s2_opt_1deg_b1_c128_h1_tc0017_float16_cuda,NVIDIA GB200,cuda:0,torch.float16,0.8072128057479858,0.8403103947639465, +nattn_s2_opt_1deg_b1_c128_h1_tc0017_bfloat16_cuda,NVIDIA GB200,cuda:0,torch.bfloat16,0.8065024018287659,0.8422367990016937, +nattn_s2_opt_1deg_b1_c128_h1_tc003_float32_cuda,NVIDIA GB200,cuda:0,torch.float32,0.8657280027866363,1.490118396282196, +nattn_s2_opt_1deg_b1_c128_h1_tc003_float16_cuda,NVIDIA GB200,cuda:0,torch.float16,1.160150396823883,1.458784008026123, +nattn_s2_opt_1deg_b1_c128_h1_tc003_bfloat16_cuda,NVIDIA GB200,cuda:0,torch.bfloat16,1.162390410900116,1.4687040209770204, +nattn_s2_opt_hdeg_b1_c128_h1_tc0017_float32_cuda,NVIDIA GB200,cuda:0,torch.float32,2.678755211830139,6.113699245452881, +nattn_s2_opt_hdeg_b1_c128_h1_tc0017_float16_cuda,NVIDIA GB200,cuda:0,torch.float16,2.567913627624512,6.780736017227173, +nattn_s2_opt_hdeg_b1_c128_h1_tc0017_bfloat16_cuda,NVIDIA GB200,cuda:0,torch.bfloat16,2.560611200332642,6.8271135807037355, +nattn_s2_opt_hdeg_b1_c128_h1_tc003_float32_cuda,NVIDIA GB200,cuda:0,torch.float32,6.623935985565185,16.6205472946167, +nattn_s2_opt_hdeg_b1_c128_h1_tc003_float16_cuda,NVIDIA GB200,cuda:0,torch.float16,6.6969247341156,18.05959987640381, +nattn_s2_opt_hdeg_b1_c128_h1_tc003_bfloat16_cuda,NVIDIA GB200,cuda:0,torch.bfloat16,6.660614347457885,18.19893741607666, +nattn_s2_opt_h1deg_b1_c128_h1_tc0017_float32_cuda,NVIDIA GB200,cuda:0,torch.float32,1.7558143854141235,2.926527976989746, +nattn_s2_opt_h1deg_b1_c128_h1_tc0017_float16_cuda,NVIDIA GB200,cuda:0,torch.float16,2.160470414161682,3.103455996513367, +nattn_s2_opt_h1deg_b1_c128_h1_tc0017_bfloat16_cuda,NVIDIA GB200,cuda:0,torch.bfloat16,2.178147220611572,3.120028781890869, +nattn_s2_opt_h1deg_b1_c128_h1_tc003_float32_cuda,NVIDIA GB200,cuda:0,torch.float32,2.7188128232955933,5.337728071212768, +nattn_s2_opt_h1deg_b1_c128_h1_tc003_float16_cuda,NVIDIA GB200,cuda:0,torch.float16,3.490959978103638,5.482310390472412, +nattn_s2_opt_h1deg_b1_c128_h1_tc003_bfloat16_cuda,NVIDIA GB200,cuda:0,torch.bfloat16,3.4898272275924684,5.509193563461304, +nattn_s2_opt_1hdeg_b1_c128_h1_tc0017_float32_cuda,NVIDIA GB200,cuda:0,torch.float32,4.631680011749268,8.67250566482544, +nattn_s2_opt_1hdeg_b1_c128_h1_tc0017_float16_cuda,NVIDIA GB200,cuda:0,torch.float16,4.4586175918579105,8.38065595626831, +nattn_s2_opt_1hdeg_b1_c128_h1_tc0017_bfloat16_cuda,NVIDIA GB200,cuda:0,torch.bfloat16,4.444348764419556,8.37009925842285, +nattn_s2_opt_1hdeg_b1_c128_h1_tc003_float32_cuda,NVIDIA GB200,cuda:0,torch.float32,9.22985610961914,18.115558433532716, +nattn_s2_opt_1hdeg_b1_c128_h1_tc003_float16_cuda,NVIDIA GB200,cuda:0,torch.float16,9.224611186981202,17.674252891540526, +nattn_s2_opt_1hdeg_b1_c128_h1_tc003_bfloat16_cuda,NVIDIA GB200,cuda:0,torch.bfloat16,9.224793624877929,17.694447898864745, +nattn_s2_opt_1deg_b1_c256_h1_tc0017_float32_cuda,NVIDIA GB200,cuda:0,torch.float32,0.7627680003643036,1.2071712017059326, +nattn_s2_opt_1deg_b1_c256_h1_tc0017_float16_cuda,NVIDIA GB200,cuda:0,torch.float16,0.782147204875946,1.1415551900863647, +nattn_s2_opt_1deg_b1_c256_h1_tc0017_bfloat16_cuda,NVIDIA GB200,cuda:0,torch.bfloat16,0.7767360031604766,1.1477919936180114, +nattn_s2_opt_1deg_b1_c256_h1_tc003_float32_cuda,NVIDIA GB200,cuda:0,torch.float32,0.9833951950073242,2.1040703535079954, +nattn_s2_opt_1deg_b1_c256_h1_tc003_float16_cuda,NVIDIA GB200,cuda:0,torch.float16,1.066864001750946,1.916915202140808, +nattn_s2_opt_1deg_b1_c256_h1_tc003_bfloat16_cuda,NVIDIA GB200,cuda:0,torch.bfloat16,1.0341727793216706,1.912873589992523, +nattn_s2_opt_hdeg_b1_c256_h1_tc0017_float32_cuda,NVIDIA GB200,cuda:0,torch.float32,3.557459235191345,9.973465633392333, +nattn_s2_opt_hdeg_b1_c256_h1_tc0017_float16_cuda,NVIDIA GB200,cuda:0,torch.float16,3.4665535926818847,9.693532752990723, +nattn_s2_opt_hdeg_b1_c256_h1_tc0017_bfloat16_cuda,NVIDIA GB200,cuda:0,torch.bfloat16,3.4779200077056887,9.70391025543213, +nattn_s2_opt_hdeg_b1_c256_h1_tc003_float32_cuda,NVIDIA GB200,cuda:0,torch.float32,7.845561552047729,27.506601333618164, +nattn_s2_opt_hdeg_b1_c256_h1_tc003_float16_cuda,NVIDIA GB200,cuda:0,torch.float16,8.293366527557373,25.728153800964357, +nattn_s2_opt_hdeg_b1_c256_h1_tc003_bfloat16_cuda,NVIDIA GB200,cuda:0,torch.bfloat16,8.328464126586914,25.73126087188721, +nattn_s2_opt_h1deg_b1_c256_h1_tc0017_float32_cuda,NVIDIA GB200,cuda:0,torch.float32,2.2152352333068848,4.246713590621948, +nattn_s2_opt_h1deg_b1_c256_h1_tc0017_float16_cuda,NVIDIA GB200,cuda:0,torch.float16,2.0254112124443053,4.589411211013794, +nattn_s2_opt_h1deg_b1_c256_h1_tc0017_bfloat16_cuda,NVIDIA GB200,cuda:0,torch.bfloat16,2.025459182262421,4.591772747039795, +nattn_s2_opt_h1deg_b1_c256_h1_tc003_float32_cuda,NVIDIA GB200,cuda:0,torch.float32,3.304819202423096,8.204819297790527, +nattn_s2_opt_h1deg_b1_c256_h1_tc003_float16_cuda,NVIDIA GB200,cuda:0,torch.float16,3.133948826789856,8.17946262359619, +nattn_s2_opt_h1deg_b1_c256_h1_tc003_bfloat16_cuda,NVIDIA GB200,cuda:0,torch.bfloat16,3.1385760068893434,8.190579223632813, +nattn_s2_opt_1hdeg_b1_c256_h1_tc0017_float32_cuda,NVIDIA GB200,cuda:0,torch.float32,5.6309216022491455,12.69488000869751, +nattn_s2_opt_1hdeg_b1_c256_h1_tc0017_float16_cuda,NVIDIA GB200,cuda:0,torch.float16,5.436556768417359,10.752636909484863, +nattn_s2_opt_1hdeg_b1_c256_h1_tc0017_bfloat16_cuda,NVIDIA GB200,cuda:0,torch.bfloat16,5.43755202293396,10.742163181304932, +nattn_s2_opt_1hdeg_b1_c256_h1_tc003_float32_cuda,NVIDIA GB200,cuda:0,torch.float32,10.92185935974121,26.535648155212403, +nattn_s2_opt_1hdeg_b1_c256_h1_tc003_float16_cuda,NVIDIA GB200,cuda:0,torch.float16,10.896281623840332,22.283491325378417, +nattn_s2_opt_1hdeg_b1_c256_h1_tc003_bfloat16_cuda,NVIDIA GB200,cuda:0,torch.bfloat16,10.84849271774292,22.339695930480957, disco_s2_opt_1deg_b1_c64_tc0017_float32_cuda,NVIDIA H100 80GB HBM3,cuda:0,torch.float32,0.3701696008443832,0.4438847988843918, disco_s2_opt_1deg_b1_c64_tc0017_float32_cuda_fused,NVIDIA H100 80GB HBM3,cuda:0,torch.float32,0.3692095994949341,0.67985919713974, disco_s2_opt_1deg_b1_c64_tc0017_float16_cuda,NVIDIA H100 80GB HBM3,cuda:0,torch.float16,0.36078080236911775,0.4198303997516632,