Distributed neighborhood-attention upsampling + consistent CUDA launch checking

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 `<c10/cuda/CUDAException.h>` 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.

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*"]

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")))

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

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,
     )

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

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

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(