# coding=utf-8
# SPDX-FileCopyrightText: Copyright (c) 2022 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.
from itertools import accumulate
from typing import Optional, Tuple, Union
import torch
import torch.distributed as dist
import torch.nn as nn
from attention_helpers import optimized_kernels_is_available
from torch_harmonics.attention import attention_kernels
from torch_harmonics.attention._attention_utils import _check_extent, _check_ndim
from torch_harmonics.attention._layout import to_nchw, to_nhwc
from torch_harmonics.attention.attention import NeighborhoodAttentionS2
from torch_harmonics.distributed._amp_utils import _cast_to_autocast_dtype, _custom_fwd, _custom_setup_context
from torch_harmonics.quadrature import effective_theta_cutoff, precompute_latitudes
from .primitives import compute_polar_halo_radius, compute_split_shapes, get_group_neighbors, polar_halo_exchange
from .utils import azimuth_group, azimuth_group_rank, azimuth_group_size, polar_group_rank, polar_group_size
# ---------------------------------------------------------------------------
# autograd.Function for the ring-step attention kernel calls
# ---------------------------------------------------------------------------
@torch.compiler.disable()
def _ring_kv(kw_chunk, vw_chunk, az_group, next_nlon_kw, next_nlon_kv):
"""
Async send current chunks, receive next chunks with known shapes.
Chunks are physical NHWC ``[B, H, W, C]``. The ring carries them in the layout
the kernels consume, so a received chunk is usable directly and no step pays a
conversion; only the local chunk is converted, once, before the loop. The sent
tensors have to be contiguous for NCCL, which is why the ring exchanges the NHWC
tensors themselves rather than channels-first views of them.
"""
send_to, recv_from = get_group_neighbors(az_group)
B, H, _, C_k = kw_chunk.shape
B, H, _, C_v = vw_chunk.shape
recv_kw = torch.empty(B, H, next_nlon_kw, C_k, device=kw_chunk.device, dtype=kw_chunk.dtype)
recv_vw = torch.empty(B, H, next_nlon_kv, C_v, device=vw_chunk.device, dtype=vw_chunk.dtype)
ops = [
dist.P2POp(dist.isend, kw_chunk, send_to, az_group),
dist.P2POp(dist.irecv, recv_kw, recv_from, az_group),
dist.P2POp(dist.isend, vw_chunk, send_to, az_group),
dist.P2POp(dist.irecv, recv_vw, recv_from, az_group),
]
reqs = dist.batch_isend_irecv(ops)
return recv_kw, recv_vw, reqs
@torch.compiler.disable()
def _ring_grad(dkw_acc, dvw_acc, az_group, next_nlon):
"""Rotate the dkw/dvw accumulators one hop along the same ring as :func:`_ring_kv`.
Each accumulator travels *with* the lon chunk it belongs to. At step ``s`` rank ``r``
holds chunk ``(r + s) % P`` and adds its contribution to the accumulator it currently
carries; the next hop goes to ``send_to == r - 1``, which holds that same chunk at step
``s + 1``. So after the ``P - 1`` in-loop hops plus one final hop, every accumulator is
back at the rank owning its chunk, fully reduced.
This is a ring reduce-scatter, and it replaces accumulating into a buffer spanning the
*whole global longitude axis* and all-reducing that. The buffer was the one allocation
here that did not shrink as azimuth ranks were added -- ``B x H_halo x nlon_global x C``
in fp32, twice -- so per-rank cost went up with the group size while the slice actually
kept went down. It also moves less: ``P`` chunk-sized hops against the all_reduce's
``2 (P - 1)`` chunk-equivalents.
A ``None`` accumulator (that branch needs no gradient) is skipped, so the rotation costs
nothing for branches that were pruned by the autograd contract.
Shapes follow ``_ring_kv`` exactly: the accumulator's ``W`` is always the current
chunk's ``nlon``, so ``next_nlon`` is the same value passed there.
"""
send_to, recv_from = get_group_neighbors(az_group)
ops = []
recv_dkw = recv_dvw = None
if dkw_acc is not None:
B, H, _, C_k = dkw_acc.shape
recv_dkw = torch.empty(B, H, next_nlon, C_k, device=dkw_acc.device, dtype=dkw_acc.dtype)
ops += [
dist.P2POp(dist.isend, dkw_acc, send_to, az_group),
dist.P2POp(dist.irecv, recv_dkw, recv_from, az_group),
]
if dvw_acc is not None:
B, H, _, C_v = dvw_acc.shape
recv_dvw = torch.empty(B, H, next_nlon, C_v, device=dvw_acc.device, dtype=dvw_acc.dtype)
ops += [
dist.P2POp(dist.isend, dvw_acc, send_to, az_group),
dist.P2POp(dist.irecv, recv_dvw, recv_from, az_group),
]
reqs = dist.batch_isend_irecv(ops) if ops else []
return recv_dkw, recv_dvw, reqs
class _RingNeighborhoodAttentionFn(torch.autograd.Function):
"""Forward ring attention + backward ring for one attention head group.
kw, vw : [B*nh, C_k/C_v, H_halo, W_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]
"""
@staticmethod
@_custom_fwd(device_type="cuda")
def forward(
kw,
vw,
qw,
psi_col_idx,
psi_roff_idx,
psi_row_idx,
quad_weights,
nlon_in: int,
pscale: 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,
psi_n_long_rows: int,
psi_max_row_len: int,
psi_mid_row_len: 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 (y_acc,
# alpha_sum, qdotk_max) 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)
# Convert to the kernels' NHWC ABI once, here, and keep the ring in that
# layout for every step: received chunks arrive ready to use. Previously the
# launcher converted all three tensors on every one of the az_size steps --
# including qw, which is loop-invariant and so was converted identically
# each time.
kw_chunk = to_nhwc(kw.contiguous())
vw_chunk = to_nhwc(vw.contiguous())
qw_nhwc = to_nhwc(qw.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.default(
kw_chunk,
vw_chunk,
qw_nhwc,
y_acc,
alpha_sum,
qdotk_max,
quad_weights,
psi_col_idx,
psi_roff_idx,
psi_row_idx,
nlon_in,
pscale,
lon_lo_kx,
lat_halo_start,
nlat_out_local,
nlon_out_local,
psi_n_long_rows,
psi_max_row_len,
psi_mid_row_len,
)
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 = to_nchw(y_out).to(dtype=inp_dtype) # [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,
psi_row_idx,
quad_weights,
nlon_in,
pscale,
lon_chunk_starts,
nlon_kx_list,
lat_halo_start,
nlat_out_local,
nlon_out_local,
r_lat,
az_group,
az_rank,
az_size,
psi_n_long_rows,
psi_max_row_len,
psi_mid_row_len,
) = 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, psi_row_idx, quad_weights, alpha_sum, qdotk_max)
ctx.nlon_in = nlon_in
ctx.pscale = pscale
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
ctx.psi_n_long_rows = psi_n_long_rows
ctx.psi_max_row_len = psi_max_row_len
ctx.psi_mid_row_len = psi_mid_row_len
@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, psi_row_idx, quad_weights, fwd_alpha_sum, fwd_qdotk_max = ctx.saved_tensors
nlon_in = ctx.nlon_in
pscale = ctx.pscale
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
psi_n_long_rows = ctx.psi_n_long_rows
psi_max_row_len = ctx.psi_max_row_len
psi_mid_row_len = ctx.psi_mid_row_len
# 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.
# This is what lets torch.compile / AOTAutograd prune dead subgraphs (including
# the NCCL allreduces) from the compiled backward.
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 (e.g., user wired
# requires_grad onto one of the index buffers), there's nothing to compute.
if not (kw_needs_grad or vw_needs_grad or qw_needs_grad):
return (None,) * 21
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 ring kernels consume kw/vw/qw/dy in their native dtype (Tier B:
# widen fp16/bf16 at load, fp32 compute/accumulation). Keeping them native
# — instead of upcasting to fp32 here — also keeps the backward ring
# exchange at 16-bit under AMP (halved K/V comm volume), matching the
# forward ring. The fp32 accumulators (integral_buf, alpha_k/kvw_buf,
# dkw/dvw_acc) are unaffected; the returned grads are cast back to the
# captured input dtypes at the end.
kw_dtype = kw.dtype
vw_dtype = vw.dtype
qw_dtype = qw.dtype
# dy and qw are loop-invariant: converted once here, not once per ring step.
dy_nhwc = to_nhwc(dy.contiguous()) # NHWC [B, H, W, C_v], native dtype
qw_nhwc = to_nhwc(qw.contiguous())
# ----------------------------------------------------------------
# Backward pass 1: re-accumulate {alpha_sum, qdotk_max, integral,
# alpha_k, alpha_kvw} via ring.
# 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 three buffers in one call, so pass-1 cannot be pruned
# per-branch.
# ----------------------------------------------------------------
bwd_alpha_sum = torch.zeros(B, nlat_out_local, nlon_out_local, device=device, dtype=torch.float32)
bwd_qdotk_max = torch.full((B, nlat_out_local, nlon_out_local), float("-inf"), device=device, dtype=torch.float32)
integral_buf = torch.zeros_like(bwd_alpha_sum)
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_nhwc = to_nhwc(kw.contiguous())
vw_nhwc = to_nhwc(vw.contiguous())
kw_chunk, vw_chunk = kw_nhwc, vw_nhwc
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_pass1.default(
kw_chunk,
vw_chunk,
qw_nhwc,
dy_nhwc,
bwd_alpha_sum,
bwd_qdotk_max,
integral_buf,
alpha_k_buf,
alpha_kvw_buf,
quad_weights,
psi_col_idx,
psi_roff_idx,
psi_row_idx,
nlon_in,
pscale,
lon_lo_kx,
lat_halo_start,
nlat_out_local,
nlon_out_local,
psi_n_long_rows,
psi_max_row_len,
psi_mid_row_len,
)
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.
# Use the SAVED forward alpha_sum/qdotk_max (same values, but authoritative).
# ----------------------------------------------------------------
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 = to_nchw(dqy_cl).to(dtype=qw_dtype) # [B, C_k, H, W]
else:
dqy = None
# ----------------------------------------------------------------
# Backward pass 2: scatter dkw/dvw contributions.
# Each GPU computes its contribution to every lon chunk it visits, and a ring
# reduce-scatter carries each chunk's accumulator back to the rank owning it
# (see _ring_grad). The accumulator is chunk-sized, so this is the whole local
# gradient at the end -- no allreduce and no slice.
# Skip entirely if neither kw nor vw needs grad; a branch that does not
# need one still gets a scratch output, since the fused kernel writes
# both in a single call.
# ----------------------------------------------------------------
if kw_needs_grad or vw_needs_grad:
# pass 1 rotated kw_chunk/vw_chunk; reset to the local chunk, which is
# already converted -- no second conversion needed.
kw_chunk, vw_chunk = kw_nhwc, vw_nhwc
# the accumulator starts on this rank's own chunk, which is the one it holds
# at step 0, and is re-sized by each hop to match the chunk it then carries
my_nlon = nlon_kx_list[az_rank]
dkw_acc = torch.zeros(B, H_halo, my_nlon, C_k, device=device, dtype=torch.float32) if kw_needs_grad else None
dvw_acc = torch.zeros(B, H_halo, my_nlon, 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]
# This kernel accumulates into its gradient outputs with atomicAdd and never
# clears them, so the accumulator goes in directly -- no per-step temp and
# no add_ afterwards. Its width is the current chunk's by construction.
# Both arguments are required by the fused signature, so a branch that needs
# no gradient still gets a scratch buffer, which is then discarded.
# NOTE: the upsample direction's kernel assigns rather than accumulates, so
# it must keep the per-step temp; do not mirror this there.
dkw_out = dkw_acc if kw_needs_grad else torch.zeros(B, H_halo, nlon_kx, C_k, device=device, dtype=torch.float32)
dvw_out = dvw_acc if vw_needs_grad else torch.zeros(B, H_halo, nlon_kx, C_v, device=device, dtype=torch.float32)
attention_kernels.backward_ring_step_pass2.default(
kw_chunk,
vw_chunk,
qw_nhwc,
dy_nhwc,
fwd_alpha_sum,
fwd_qdotk_max,
integral_norm,
dkw_out,
dvw_out,
quad_weights,
psi_col_idx,
psi_roff_idx,
psi_row_idx,
nlon_in,
pscale,
lon_lo_kx,
lat_halo_start,
nlat_out_local,
nlon_out_local,
psi_n_long_rows,
psi_max_row_len,
psi_mid_row_len,
)
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])
# the accumulators follow the chunks they belong to, so they rotate in
# the same direction and to the same per-chunk width
recv_dkw, recv_dvw, grad_reqs = _ring_grad(dkw_acc, dvw_acc, az_group, nlon_kx_list[next_src])
for req in reqs:
req.wait()
for req in grad_reqs:
req.wait()
kw_chunk = recv_kw.clone()
vw_chunk = recv_vw.clone()
# Cloned for the same reason as kw/vw, and with more at stake: the next
# step's kernel atomicAdds *into* the accumulator, so handing it the
# irecv destination directly would mutate a buffer the collective owns.
# The clone is chunk-sized, so it does not give back the saving above.
dkw_acc = recv_dkw.clone() if recv_dkw is not None else None
dvw_acc = recv_dvw.clone() if recv_dvw is not None else None
# One final hop returns each accumulator to the rank owning its chunk: after the
# loop a rank carries the accumulator for chunk az_rank - 1, which is send_to.
if az_size > 1:
recv_dkw, recv_dvw, grad_reqs = _ring_grad(dkw_acc, dvw_acc, az_group, my_nlon)
for req in grad_reqs:
req.wait()
# No clone here, deliberately: nothing writes into these again. The only
# remaining use is to_nchw, which always materializes a new tensor, so the
# returned gradient never aliases the irecv destination.
dkw_acc, dvw_acc = recv_dkw, recv_dvw
# The accumulator IS the local chunk now, so only the layout conversion is left.
# No halo stripping: dkw/dvw must match kw/vw shape (= key_halo/value_halo).
# The autograd through torch.cat in _exchange_lat_halo extracts the
# middle H_in rows as the gradient for key_proj/value_proj.
dkw = to_nchw(dkw_acc).to(dtype=kw_dtype) if kw_needs_grad else None # [B, C_k, H_halo, W_local]
dvw = to_nchw(dvw_acc).to(dtype=vw_dtype) if vw_needs_grad else None # [B, C_v, H_halo, W_local]
else:
dkw = None
dvw = None
# Return grads for (kw, vw, qw, psi_col, psi_roff, psi_row, quad_weights,
# nlon_in, pscale, lon_chunk_starts, nlon_kx_list, lat_halo_start,
# nlat_out_local, nlon_out_local, r_lat,
# az_group, az_rank, az_size,
# psi_n_long_rows, psi_max_row_len, psi_mid_row_len)
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
@_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)
# Converted once here, not once per ring step; qw is loop-invariant and the
# rotating chunks stay NHWC across the exchange.
kw_chunk = to_nhwc(kw.contiguous())
vw_chunk = to_nhwc(vw.contiguous())
qw_nhwc = to_nhwc(qw.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_nhwc,
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 = to_nchw(y_out).to(dtype=inp_dtype) # [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 and qw are loop-invariant: converted once here, not once per ring step.
dy_nhwc = to_nhwc(dy.contiguous()) # NHWC [B, H, W, C_v], native dtype
qw_nhwc = to_nhwc(qw.contiguous())
# ----------------------------------------------------------------
# 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_nhwc = to_nhwc(kw.contiguous())
vw_nhwc = to_nhwc(vw.contiguous())
kw_chunk, vw_chunk = kw_nhwc, vw_nhwc
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_nhwc,
dy_nhwc,
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 = to_nchw(dqy_cl).to(dtype=qw_dtype) # [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,
# and a ring reduce-scatter carries each chunk's accumulator back to its owner
# (see _ring_grad) -- chunk-sized throughout, so no allreduce and no slice.
# ----------------------------------------------------------------
if kw_needs_grad or vw_needs_grad:
# pass 1 rotated kw_chunk/vw_chunk; reset to the local chunk, which is
# already converted -- no second conversion needed.
kw_chunk, vw_chunk = kw_nhwc, vw_nhwc
# starts on this rank's own chunk (the one held at step 0), re-sized by each hop
my_nlon = nlon_kx_list[az_rank]
dkw_acc = torch.zeros(B, H_halo, my_nlon, C_k, device=device, dtype=torch.float32) if kw_needs_grad else None
dvw_acc = torch.zeros(B, H_halo, my_nlon, 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]
# Unlike the gather direction, this kernel *assigns* its gradient outputs
# (`dkx[chan] = sh_dk[chan]` in attention_cuda_bwd_ring_upsample.cu) rather
# than accumulating with atomicAdd: in the scatter direction each input cell
# is written by exactly one row, so it can store instead of reduce. The
# accumulator therefore cannot be passed in directly -- each step would
# overwrite the previous ones -- so it keeps a zeroed per-step temp and adds.
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_nhwc,
dy_nhwc,
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_acc.add_(dkw_chunk_cl)
if vw_needs_grad:
dvw_acc.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])
# accumulators follow the chunks they belong to: same direction, same width
recv_dkw, recv_dvw, grad_reqs = _ring_grad(dkw_acc, dvw_acc, az_group, nlon_kx_list[next_src])
for req in reqs:
req.wait()
for req in grad_reqs:
req.wait()
kw_chunk = recv_kw.clone()
vw_chunk = recv_vw.clone()
# Cloned for the same reason as kw/vw, and with more at stake: the next
# step's add_ writes *into* the accumulator, so keeping the irecv
# destination would mutate a buffer the collective owns.
# The clone is chunk-sized, so it does not give back the saving above.
dkw_acc = recv_dkw.clone() if recv_dkw is not None else None
dvw_acc = recv_dvw.clone() if recv_dvw is not None else None
# Final hop home: after the loop a rank carries the accumulator for chunk
# az_rank - 1, which is exactly send_to.
if az_size > 1:
recv_dkw, recv_dvw, grad_reqs = _ring_grad(dkw_acc, dvw_acc, az_group, my_nlon)
for req in grad_reqs:
req.wait()
# No clone here, deliberately: nothing writes into these again. The only
# remaining use is to_nchw, which always materializes a new tensor, so the
# returned gradient never aliases the irecv destination.
dkw_acc, dvw_acc = recv_dkw, recv_dvw
# The accumulator IS the local chunk now, so only the layout conversion is left.
# No halo stripping: dkw/dvw must match kw/vw shape (= key_halo/value_halo).
dkw = to_nchw(dkw_acc).to(dtype=kw_dtype) if kw_needs_grad else None # [B, C_k, H_halo, W_local]
dvw = to_nchw(dvw_acc).to(dtype=vw_dtype) if vw_needs_grad else None # [B, C_v, H_halo, W_local]
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
# ---------------------------------------------------------------------------
[docs]
class DistributedNeighborhoodAttentionS2(NeighborhoodAttentionS2):
"""
Distributed neighborhood attention on the 2-sphere using a ring exchange
strategy for the longitude dimension and halo exchange for the latitude
dimension.
Data is assumed to be split along both the latitude (polar group) and
longitude (azimuth group) dimensions. The forward pass uses ring exchange
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:`torch_harmonics.NeighborhoodAttentionS2`.
.. seealso::
:class:`torch_harmonics.NeighborhoodAttentionS2`
Serial counterpart with full parameter documentation.
"""
def __init__(
self,
in_channels: int,
in_shape: Tuple[int, int],
out_shape: Tuple[int, int],
grid_in: Optional[str] = "equiangular",
grid_out: Optional[str] = "equiangular",
num_heads: Optional[int] = 1,
scale: Optional[Union[torch.Tensor, float]] = None,
use_qknorm: Optional[bool] = False,
bias: Optional[bool] = True,
theta_cutoff: Optional[float] = None,
k_channels: Optional[int] = None,
out_channels: Optional[int] = None,
optimized_kernel: Optional[bool] = True,
):
if not optimized_kernels_is_available():
raise RuntimeError("Optimized kernels are required to run DistributedNeighborhoodAttentionS2.")
# initialise base class (builds global psi, creates parameters)
super().__init__(
in_channels,
in_shape,
out_shape,
grid_in=grid_in,
grid_out=grid_out,
num_heads=num_heads,
scale=scale,
use_qknorm=use_qknorm,
bias=bias,
theta_cutoff=theta_cutoff,
k_channels=k_channels,
out_channels=out_channels,
optimized_kernel=True,
)
# ---- distributed info ----
self.comm_size_polar = polar_group_size()
self.comm_rank_polar = polar_group_rank()
self.comm_size_azimuth = azimuth_group_size()
self.comm_rank_azimuth = azimuth_group_rank()
# split shapes
self.lat_in_shapes = compute_split_shapes(self.nlat_in, self.comm_size_polar)
self.lon_in_shapes = compute_split_shapes(self.nlon_in, self.comm_size_azimuth)
self.lat_out_shapes = compute_split_shapes(self.nlat_out, self.comm_size_polar)
self.lon_out_shapes = compute_split_shapes(self.nlon_out, self.comm_size_azimuth)
# local sizes for this rank
self.nlat_in_local = self.lat_in_shapes[self.comm_rank_polar]
self.nlon_in_local = self.lon_in_shapes[self.comm_rank_azimuth]
self.nlat_out_local = self.lat_out_shapes[self.comm_rank_polar]
self.nlon_out_local = self.lon_out_shapes[self.comm_rank_azimuth]
# 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.
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]
# ---- lat halo size ----
# Derived from the grid geometry rather than measured off the psi: an output latitude
# can only reach input latitudes within theta_cutoff of it, and the halo is how far
# that band runs past a rank's own input range. The criterion is symmetric in the two
# grids, so the same call covers the gather and scatter directions -- the upsample path
# builds its psi with the shapes swapped, but the latitudes it needs are the same ones.
# It also raises if the halo outgrows a local chunk, which the immediate-neighbour
# exchange could not serve.
# the grid types are constructor arguments the base class does not retain, so they are
# read from the local parameters rather than off self
lats_in, _ = precompute_latitudes(self.nlat_in, grid=grid_in)
lats_out, _ = precompute_latitudes(self.nlat_out, grid=grid_out)
self.r_lat = compute_polar_halo_radius(
lats_in,
lats_out,
effective_theta_cutoff(self.theta_cutoff),
self.lat_in_shapes,
self.lat_out_shapes,
)
if self.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). This needs r_lat,
# which is why the halo size is settled above rather than after the build.
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}
# the local psi supersedes the global one the base class built; drop it
self._drop_global_psi()
# -----------------------------------------------------------------------
def _drop_global_psi(self):
"""Free the global sparsity buffers registered by the serial base class.
``_build_local_psi`` / ``_build_local_psi_upsample`` are their only readers and both
run during ``__init__``; the ring kernels take the ``*_local`` buffers exclusively.
Left in place they would sit on the device for the lifetime of the module alongside
the local copies that replace them -- and being keyed to the *global* grid, they do
not shrink as ranks are added. ``psi_row_idx``, ``psi_seg`` and ``psi_seg_off`` are
never read here at all: the local row order is rebuilt in ``_build_local_psi`` and
the arc segments are only consumed by the serial kernels.
All five are ``persistent=False``, so no checkpoint content changes -- only
``named_buffers()`` differs from the serial module. Deleting rather than setting
them to ``None`` means a stale read raises ``AttributeError`` naming the buffer,
instead of a ``None`` propagating into index arithmetic and failing a frame later.
``del`` is the supported route: ``nn.Module.__delattr__`` removes the entry from
``_buffers`` and discards the name from ``_non_persistent_buffers_set``. The base
class registers all five unconditionally, so none of these can be missing.
"""
del self.psi_row_idx
del self.psi_col_idx
del self.psi_roff_idx
del self.psi_seg
del self.psi_seg_off
# -----------------------------------------------------------------------
def _build_local_psi(self):
"""Filter global psi to local output lat rows and shift col_idx wi."""
lat_lo = self.lat_lo_out
lat_hi = lat_lo + self.nlat_out_local
# global psi from the base class (built over all nlat_out rows)
col_idx_global = self.psi_col_idx # [nnz] int64
roff_global = self.psi_roff_idx # [nlat_out+1] int64
# psi_row_idx stores the sorted permutation: value is the row index.
# psi_roff_idx[ho] .. psi_roff_idx[ho+1] gives entries for row ho.
# (The row_idx buffer is the *sort order*, not the row indices directly.)
# For the distributed case we rebuild roff for the local rows only.
# Build local roff: select rows lat_lo..lat_hi-1
roff_local = roff_global[lat_lo : lat_hi + 1] - roff_global[lat_lo] # offset by first entry
# Select the corresponding col_idx entries
start = roff_global[lat_lo].item()
end = roff_global[lat_hi].item()
col_idx_local = col_idx_global[start:end].clone()
# Shift wi by pscale * lon_lo_out so the kernel can reconstruct wip from wo_local:
# col stores hi_global * nlon_in + wi_canonical. For global wo_global = lon_lo_out + wo_local,
# the target input column is (wi_canonical + pscale * wo_global) % nlon_in. The kernel evaluates
# (wi_shifted + pscale * wo_local) % nlon_in, so pre-shifting by pscale * lon_lo_out absorbs
# the rank-offset piece. pscale = 1 when nlon_in == nlon_out (same-shape case).
nlon_in = self.nlon_in
lon_lo = self.lon_lo_out
pscale = self.nlon_in // self.nlon_out
hi_global = col_idx_local // nlon_in
wi_canon = col_idx_local - hi_global * nlon_in
wi_shifted = (wi_canon + pscale * lon_lo) % nlon_in
col_idx_shifted = hi_global * nlon_in + wi_shifted
# Build sorted row_idx for local output rows (0-indexed within local range)
# Reuse the serial sort order: just re-sort by nnz per local row
nnz_per_row = (roff_local[1:] - roff_local[:-1]).cpu()
row_idx_local = torch.argsort(nnz_per_row, descending=True).to(torch.int32)
self.register_buffer("psi_col_idx_local", col_idx_shifted, persistent=False)
self.register_buffer("psi_roff_idx_local", roff_local, persistent=False)
self.register_buffer("psi_row_idx_local", row_idx_local, persistent=False)
# Precompute the CSR long/short row split once, here in the constructor,
# on the still-on-CPU local psi buffers (split_csr_rows has a CPU path).
# The split depends only on the psi sparsity geometry, which is fixed
# after init, so it is identical on every ring step / iteration. Computing
# it once keeps it off the per-step hot path (it otherwise cost a 24-byte
# D2H sync per ring step) and off any compiled forward. Stored as plain
# Python ints and threaded into the ring-step ops.
n_long_rows, max_row_len, mid_row_len = attention_kernels.split_csr_rows.default(row_idx_local, roff_local, self.nlat_out_local)
self.psi_n_long_rows = int(n_long_rows)
self.psi_max_row_len = int(max_row_len)
self.psi_mid_row_len = int(mid_row_len)
# -----------------------------------------------------------------------
# upsample (scatter) direction helpers
# -----------------------------------------------------------------------
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(
self,
query: torch.Tensor,
key: Optional[torch.Tensor] = None,
value: Optional[torch.Tensor] = None,
) -> torch.Tensor:
if key is None:
key = query
if value is None:
value = query
_check_ndim(query, 4, "query")
_check_ndim(key, 4, "key")
_check_ndim(value, 4, "value")
_check_extent(query, -2, self.nlat_out_local, "query latitudes")
_check_extent(query, -1, self.nlon_out_local, "query longitudes")
_check_extent(key, -2, self.nlat_in_local, "key latitudes")
_check_extent(key, -1, self.nlon_in_local, "key longitudes")
_check_extent(value, -2, self.nlat_in_local, "value latitudes")
_check_extent(value, -1, self.nlon_in_local, "value longitudes")
# ---- 1. project to k/v/q ----
key_proj = nn.functional.conv2d(key, self.k_weights, bias=self.k_bias)
value_proj = nn.functional.conv2d(value, self.v_weights, bias=self.v_bias)
query_proj = nn.functional.conv2d(query, self.q_weights, bias=self.q_bias)
# QK normalization (must come before scale)
if self.q_norm_weights is not None:
B, C, H, W = query_proj.shape
query_proj = query_proj.reshape(B, self.num_heads, -1, H, W).permute(0, 1, 3, 4, 2)
query_proj = nn.functional.rms_norm(query_proj, normalized_shape=self.q_norm_weights.shape, weight=1 + self.q_norm_weights)
query_proj = query_proj.permute(0, 1, 4, 2, 3).reshape(B, C, H, W).contiguous()
if self.k_norm_weights is not None:
B, C, H, W = key_proj.shape
key_proj = key_proj.reshape(B, self.num_heads, -1, H, W).permute(0, 1, 3, 4, 2)
key_proj = nn.functional.rms_norm(key_proj, normalized_shape=self.k_norm_weights.shape, weight=1 + self.k_norm_weights)
key_proj = key_proj.permute(0, 1, 4, 2, 3).reshape(B, C, H, W).contiguous()
# scale after normalization
query_proj = query_proj * self.scale
# fold num_heads into batch
B, _, H, W = key_proj.shape
key_proj = key_proj.reshape(B * self.num_heads, -1, H, W)
B, _, H, W = value_proj.shape
value_proj = value_proj.reshape(B * self.num_heads, -1, H, W)
B, _, H, W = query_proj.shape
query_proj = query_proj.reshape(B * self.num_heads, -1, H, W)
# ---- 2. lat halo exchange ----
# key_proj/value_proj: [Bnh, C, H_in_local, W_in_local]
# Use differentiable halo exchange when there is an actual polar split;
# otherwise fall through to the identity (no-op).
if self.r_lat > 0 and self.comm_size_polar > 1:
key_halo = polar_halo_exchange(key_proj, self.r_lat)
value_halo = polar_halo_exchange(value_proj, self.r_lat)
else:
key_halo = key_proj
value_halo = value_proj
# global lat index of first halo row
lat_halo_start = self.lat_in_starts[self.comm_rank_polar] - self.r_lat
# ---- 3. ring attention ----
# 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)
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
B_orig = B_nh // self.num_heads
out = out.reshape(B_orig, self.num_heads * C_v, H_out, W_out)
# ---- 4. output projection ----
out = nn.functional.conv2d(out, self.proj_weights, bias=self.proj_bias)
return out
