Attention optimizations (part 2): channel-aware backward cache

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.

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 .

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 .