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. 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/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: 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, 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/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..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, ) @@ -964,6 +976,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/__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 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..13096e05 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 @@ -65,6 +66,19 @@ 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; + + // 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) @@ -303,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 @@ -340,6 +357,16 @@ 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. + // 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]; COMPUTE_T loc_vw_[NLOC]; @@ -420,6 +447,10 @@ namespace attention_kernels col_idx += rbeg; const int rlen = rend - rbeg; + // 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. @@ -470,6 +501,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); @@ -514,6 +551,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++) { @@ -524,42 +570,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]; @@ -671,7 +724,20 @@ 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; + const size_t cache_sh = size_t(block.y) * (2 * BWD_CACHE_CAP) * sizeof(float); + + // 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 // temporary buffers loc_k__[], loc_vw_[] and loc_kvw[], @@ -683,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); + // 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); + 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"); @@ -952,6 +1034,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