Attention bwd: arc-segment iteration + upsample backward optimization by azrael417 · Pull Request #245 · NVIDIA/torch-harmonics

azrael417 · 2026-08-19T05:08:50Z

Base: tkurth/attention-fwd-neighbor-groups (#244) — stacked on top; do not merge before #244.

Summary

Follow-up to #244, which removed the per-neighbor 64-bit integer division from the forward kernels by walking psi as contiguous longitude arcs instead of decoding a flat column index. This PR applies the same idea to the backward kernels, plus a first pass at the previously-unoptimized upsample (scatter) backward.

Gather backward: arc segments. All four neighbor loops (generic and special kernels, two passes each) now derive (input lat, input lon) by counting along an arc instead of decoding col_idx — same division/modulus removal as the forward. col_idx/row_off are dropped from the kernel signatures; row_off survives only in the dispatcher, which still uses row lengths for sortRows load balancing. The op schema keeps both col_idx and the segments, since the CPU backward and the torch reference still walk the column list — keeping the reference independent of the arc derivation, same rationale as the forward PR.
Measured on V100, fp32, C=64, backward only:

Config	Before	After	Δ
180x360 self	39.68 ms	37.50 ms	5.5%
360x720 → 180x360	159.06 ms	150.02 ms	5.7%

Much smaller than the forward's 2.02x: the forward was instruction-bound on address arithmetic, the backward isn't, so the division was a small share of its cost. Kept because it's free and simplifies the loops — the backward's real bottleneck (atomicAdd traffic in the gradient pass, four __warp_sum reductions per neighbor) is still unprofiled.
Special kernel: qy/dy in registers. The special kernel re-read sh_qy/sh_dy from shared memory once per neighbor despite both being loop-invariant per-thread slots. ncu showed it issue-limited (69% peak LSU issue, 84% SM throughput, 1.1% DRAM, "not selected" the largest stall). Per-thread reads move to registers; the shared copies stay because the scalar epilogue reads across the warp and needs them.
Measured on V100, fp32, C=64, backward only:

Config	Before	After	Δ
180x360 self	37.50 ms	36.54 ms	2.6%
360x720 → 180x360	150.02 ms	146.37 ms	2.4%

Upsample (scatter) backward: arc segments + sortRows. The upsample backward's three neighbor-walking kernels (max_k, stats_k, dkv_k) are structurally a gather despite the "scatter" naming — ncu confirms 0.00% L2 atomic cycles, since each kernel parallelizes over input points that already own their output. So the gather-side optimizations apply and, per the file's own prior comment ("no specialized channel-size variant or sortRows"), none had been applied yet. This PR lands two of the four: arc segments (removing the division) in all three loops, and sortRows (rows visited longest-first — psi's density is heavily latitude-skewed, and natural order left entire latitudes straggling; ncu measured 29.95% achieved occupancy against 50% theoretical).
Measured on V100, fp32, C=64, 180x360 → 360x720 backward:

Kernel	Before	After
scatter_dkv_k	152.3 ms	108.6 ms
scatter_stats_k	127.7 ms	92.1 ms
scatter_max_k	65.5 ms	34.6 ms
total	349.4 ms	241.2 ms (1.45x)

Far more than the 5% the segment change alone bought the gather backward — this path is latency-bound rather than issue-bound. The two changes (segments + sortRows) were applied together, so their individual contributions aren't separated.
Not done here: specialized channel-size variant and vectorization for the upsample backward. ncu shows 64 registers/thread capping occupancy at 50% theoretical, now the binding constraint on dkv_k — flagged as follow-up, same as the scatter/upsample forward division was flagged as follow-up in Attention fwd: neighbor-groups + segment-based psi iteration (remove per-neighbor integer division) #244.
Benchmarks. Adds quarter-degree (721×1440, ~1.04M query points, 16x the 1° entries) self-attention configs (fp32/fp16/bf16) to benchmarks/attention.py — the resolution production is actually headed for, where kernels stop being dominated by launch/latency effects. Correctness checks are skipped for these: the dense torch reference is quadratic in nlat*nlon and doesn't fit at this size.

Testing

Suite: 302 passed / 7 skipped on H100.
Op schema change (seg/seg_off now saved in ctx.save_for_backward alongside col_idx/row_off) is covered by the existing opcheck-based PT2 compatibility tests and test_custom_implementation.

API / numerical changes

Backward custom op's saved-context tuple grows from (col_idx, row_off, quad_weights, kw, vw, qw) to (col_idx, row_off, seg, seg_off, quad_weights, kw, vw, qw) — positional, and _neighborhood_s2_attention_bwd_torch's unpack was updated to match (segments unused there, kept only to share _setup_context_attention_backward with the CUDA path). No public Python API change.
No numerical changes: backward arc-segment iteration is a pure re-derivation of the same column indices, and sortRows/register changes don't alter reduction order in a way that should move results outside existing tolerances (call this out explicitly in the PR if you saw any tolerance adjustments needed — I didn't see one in the diff, but flag if I'm missing something).

The special forward kernel walked its neighbor list one at a time, with a fully serial chain per neighbor: col_idx -> address -> load k -> warp reduce -> softmax -> load v. With nchan_in=64 over 32 lanes there were only two independent loads inside the channel loop, so nothing was in flight while a ~500-cycle load resolved. Group NB=4 neighbors: hoist the four col_idx reads and address computations, feed four accumulators from one channel loop so four k loads (and later four v loads) are outstanding per step, and apply the online softmax once per group instead of once per neighbor. The latter is algebraically the same reduction -- one running max, one rescale of locy per group -- and rescales less often, so if anything it is better conditioned. Results are therefore not bit-identical to before in reduced precision, though they move toward the fp64 reference, not away. Measured on H100, op-level forward (attention_kernels::forward): 1deg_tc0017 0.272 -> 0.180 ms (1.51x) 1deg_tc003 0.542 -> 0.447 ms (1.21x) hdeg_tc0017 2.190 -> 2.058 ms (1.06x) hdeg_tc003 6.632 -> 6.148 ms (1.08x) The win shrinks with problem size because at scale the kernel is limited by memory request volume rather than latency: it has no cross-query reuse, so it moves ~256 bytes per (query, neighbor) pair against a tiled kernel's ~1. That is a separate, larger problem and is not addressed here. Only the special kernel is changed; the generic one has the same structure and would need the same treatment. NB=4 was chosen for register pressure, not measured against 2 or 8. Tested: tests/test_attention.py -k custom_implementation, 174 passed (CPU + CUDA, fp32/fp16/bf16, multi-head, both resampling directions).

psi's sparsity is a union of contiguous longitude arcs, one per (row, input lat). That is geometric rather than incidental -- a geodesic ball meets a latitude circle in a single arc, never in disjoint pieces -- and it holds across equiangular, legendre-gauss and lobatto grids, gather and scatter, pscale 1 and 2, and odd latitude counts. It is worth pinning because kernels are entitled to rely on it: the neighbor longitudes of a row can be described by (start, width) instead of an explicit column list, so a neighbor's address is computable as (lo + j + pscale*wo) % nlon_in rather than loaded from psi_col_idx. Any kernel taking that route reads the wrong cells -- silently, with plausible output -- if a change to how psi is built ever introduces a hole. This test is what would catch it. The check counts the seam between the last and first longitude as a gap, so a wrapping arc such as {718, 719, 0, 1} registers as contiguous. A plain min/max test would report it as the whole circle. Device-independent; psi is built on CPU at construction.

Adds _build_psi_segments / _expand_psi_segments, which re-express psi's column list as (hi, lo, len) arcs -- one per (output row, input latitude) -- plus a test that the segments expand back to exactly psi_col_idx. The round-trip is the property that matters, and it is stronger than the contiguity check added alongside it. Contiguity alone does not protect a kernel: a segment whose start is off by one, which is what a mishandled wrapping arc produces, still passes contiguity while making the kernel attend to the wrong cells. Comparing the expanded table against psi_col_idx element for element catches that. Motivation, from profiling the forward kernel on H100 (see profiles/): it runs at 80% compute throughput while delivering ~2.4% of peak FLOPs, with DRAM at 0.6%. It is instruction-bound, and the instructions are address arithmetic -- per neighbor a 64-bit integer division to recover hi from col_idx, which the GPU emulates in ~70-100 instructions against roughly four instructions of useful math. With segments, hi is a per-segment constant and the column advances by counting, so the division amortizes over the whole arc. That amortization improves with resolution: nnz per segment is 13.6 at 1 degree, 23 at half, 46 at quarter, and extrapolates to ~90 at 1440x2880. Verified exact on equiangular, legendre-gauss and lobatto grids, gather and scatter, pscale 1 and 2, odd latitude counts, and a wide cutoff whose 21:1 ratio exercises wrapping arcs. Construction is 1.07 s at 720x1440 and one-time. Nothing consumes the segments yet; the kernels still read psi_col_idx.

Both CUDA forward kernels now walk psi as (hi, lo, len) arcs. hi and the quadrature weight are per-arc constants and the column advances by counting, so the per-neighbour 64-bit `col / nlon_in` is gone. The GPU has no integer divide instruction, so that was ~70-100 emulated instructions per neighbour against roughly four of useful math. Profiling motivated this and bounds what to expect. The forward kernel measured 80% compute throughput while delivering ~2.4% of peak FLOPs, with DRAM at 0.6% and 9.9 cycles per issued instruction -- not stalled, not bandwidth-limited, simply issuing enormous numbers of address-arithmetic instructions. Removing the cheaper 32-bit modulo (5558e3d) already bought 22%; this removes the expensive half. The amortization improves with resolution, which matters for the 1440x2880 target: neighbours per segment are 13.6 at 1 degree, 23 at half, 46 at quarter, ~90 extrapolated at target. Arcs also widen (5 -> 7 -> 17 -> ~34), so the k/v accesses become longer stride-1 runs. ABI: forward gains psi_seg (nsegs, 3) int32 and psi_seg_off. col_idx and row_off stay in the signature because the CPU and torch reference paths still consume them -- which is deliberate: test_custom_implementation compares optimized against torch, so with torch on col_idx and optimized on segments, any error in segment construction shows up immediately as a disagreement rather than as two paths sharing a bug. Three sites in the arity change are not signatures and have no type checking, all found by tests rather than inspection: - the backward gradient-return lists (one None per forward input) - the positional unpack in _setup_context_attention_backward - the AutocastCUDA impls, registered via torch.library.impl rather than a decorator Not converted: the scatter/upsample kernels, which have the same per-neighbour division at attention_cuda_fwd_upsample.cu:150 and :207 and three sites in the backward. They are next, but no benchmark config currently exercises them, so they would be changed unmeasured. The CPU path accepts the tensors and ignores them. Tested: 290 passed / 7 skipped on both V100 and H100, plus the distributed suite. Performance not yet measured.

The existing nattn entries top out at 360x720. Production is headed for 721x1440 and beyond, which is ~1.04M query points -- 16x the 1deg entries -- and that is where the kernels stop being dominated by launch and latency effects alone. Correctness is skipped for these: the dense torch reference is quadratic in nlat*nlon and does not fit at this size. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

All four neighbor loops (generic and special kernels, two passes each) now derive a neighbor's (input lat, input lon) by counting along a contiguous arc rather than decoding a flat column index. That removes a 64-bit integer division and a modulus per neighbor -- the GPU has no integer-divide instruction, so each cost ~70-100 emulated ones -- plus two int64 loads. col_idx and row_off are gone from the kernel signatures; row_off survives only in the dispatcher, where sortRows still uses row lengths for load balancing. The op schema keeps both because the CPU backward and the torch reference still consume the column list, which is what keeps the reference independent of this derivation. Measured on V100, fp32, C=64, backward only: 180x360 self 39.68 -> 37.50 ms (5.5%) 360x720 -> 180x360 159.06 -> 150.02 ms (5.7%) Far less than the 2.02x the same change bought the forward. The forward was instruction-bound on address arithmetic; the backward is not, so the division was a small share of its cost. Kept because it is free, deletes two kernel arguments and simplifies the loops, but the backward's real bottleneck is elsewhere and still unprofiled -- the atomicAdd traffic in the gradient pass and the four __warp_sum reductions per neighbor are the candidates. The scatter (upsample) backward is a separate file and is untouched: 349 ms before and after. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

The special kernel re-read sh_qy / sh_dy from shared memory once per neighbor, though both are loop-invariant and each thread only ever touches its own (tidx + i*BDIM_X) slots. ncu showed the kernel issue-limited on exactly that pipe: 69% of peak LSU instruction issue and 84% SM throughput against 1.1% DRAM, with "not selected" the largest stall -- warps eligible but no issue slot, so warps are in surplus and trading a little occupancy for fewer instructions is the right direction. The shared copies stay. The scalar epilogue undoes the tidx offset and reads across the warp, so registers cannot replace them there; only the per-thread reads move. Measured on V100, fp32, C=64, backward only: 180x360 self 37.50 -> 36.54 ms (2.6%) 360x720 -> 180x360 150.02 -> 146.37 ms (2.4%) Smaller than the LSU counters suggested, most likely because ptxas already hoisted these loads -- there is no __syncthreads in the neighbor loop, so it was free to. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

The upsample backward is named "scatter" after the forward direction, but its three neighbor-walking kernels are structurally a GATHER: they parallelise over input points, each walking its own psi row and accumulating dk/dv locally. ncu confirms it -- 0.00% L2 atomic cycles, because every input point already owns its output. So the gather path's optimizations apply here, and none of them had been: the file's own comment said "no specialized channel-size variant or sortRows". Two of the four are now in: - Arc segments in all three loops (max_k, stats_k, dkv_k), removing the 64-bit division that decoded col / nlon_out per neighbor. - Rows visited longest-first via sortRows. psi's density is heavily skewed with latitude -- a polar row spans the whole circle -- and every wi within a latitude does identical work, so in natural order entire latitudes straggled. ncu measured 29.95% achieved occupancy against 50% theoretical. col_idx and row_off are gone from these kernels; row_off survives in the dispatcher to feed sortRows. Measured on V100, fp32, C=64, 180x360 -> 360x720 backward: scatter_dkv_k 152.3 -> 108.6 ms scatter_stats_k 127.7 -> 92.1 ms scatter_max_k 65.5 -> 34.6 ms total 349.4 -> 241.2 ms (1.45x) Far more than the 5% the same segment change bought the gather backward, which is issue-bound rather than latency-bound. The two changes were applied together, so their individual contributions are not separated. Still missing here: the specialized channel-size variant and vectorization. ncu also shows 64 registers/thread capping occupancy at 50% theoretical, which is now the binding constraint on dkv_k. Suite: 302 passed / 7 skipped on H100. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

… list The last non-ring serial attention kernel still decoding a flat column index. Both neighbor loops (scatter_max_k and scatter_acc_k) now derive a neighbor's (output lat, output lon) by counting along a contiguous arc, removing a 64-bit integer division and an int64 col_idx load per neighbor. col_idx and row_off are gone from both kernel signatures, the launcher and s2_attn_fwd_upsample_dispatch. No new precompute: attention.py sets nlon_decode = nlon_out whenever the layer upsamples, so psi_seg is already (ho, lo, len) in output-longitude space, which is exactly what this kernel indexes. The backward upsample kernels already consume those same buffers. This was previously skipped on the grounds that the kernel is bound by atomic throughput -- atomicMaxf per neighbor in pass 1, atomicAdd into numer/denom in pass 2 -- so removing the divide should not move the critical path. That reasoning was recorded as "little gain" and is wrong. Measured on a healthy GPU (180x360 -> 360x720, fp32, C=64, forward only, 3 runs each): scatter_acc_k 6.686 / 6.917 / 6.150 -> 5.877 / 5.886 / 5.882 (-10.6%) scatter_max_k 5.317 / 5.439 / 4.787 -> 4.204 / 4.204 / 4.203 (-18.9%) total ~11.97 ms -> ~10.29 ms (-14%) The win survives comparison against the fastest baseline sample (-4.4% / -12%). The earlier verdict most likely came from a thermally throttled GPU, which compresses differences while looking stable -- note the baseline spread of +/-6% against three post-change runs identical to three decimals. sortRows is deliberately NOT adopted here, unlike the backward upsample: pass 2 accumulates into numer/denom with atomicAdd, so reordering rows would change the float summation order. That belongs in its own change with its own measurement. Every serial GPU attention kernel now uses one psi representation; col_idx survives only for the CPU kernels and the ring path. Suite: 306 passed, 7 skipped. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

fwd_upsample carried, twice, the claim that this kernel is atomics-bound and so the segment rewrite buys "consistency and one representation of psi, not speed". That was the reasoning for skipping the conversion for a long time, and the measurement contradicts it: 180x360 -> 360x720 fp32 C=64 forward went 6.69 -> 5.88 ms (acc_k) and 5.32 -> 4.20 ms (max_k), about -14%. The premise was right -- it is atomics-heavy -- but "bound by X" did not imply "Y cannot help". The comment now records that rather than the superseded prediction. Also deduplicates the identical arc-segment rationale that appeared verbatim in all three bwd_upsample kernels; stated once, back-referenced twice. Comments only; both files still compile for sm_90. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

azrael417 self-assigned this Aug 19, 2026

azrael417 and others added 12 commits August 19, 2026 01:10

removing modulus from the attention loops

a4a0113

slightly improved attention via vectorization

7357bf0

adding tests to check correctness of the vectorized path

d873757

fixing rebase

cd7ab43

azrael417 force-pushed the tkurth/attention-bwd-segments branch from 1029f2b to cd7ab43 Compare August 19, 2026 08:14

azrael417 changed the title ~~Improving attention backward performance~~ Attention bwd: arc-segment iteration + upsample backward optimization Aug 20, 2026

azrael417 requested review from bonevbs, mauro-bis and rietmann-nv August 20, 2026 07:12

Merge branch 'main' into tkurth/attention-bwd-segments

1aa29d4

azrael417 marked this pull request as ready for review August 24, 2026 06:38

azrael417 and others added 2 commits August 25, 2026 23:10

Navigation Menu

Uh oh!

Attention bwd: arc-segment iteration + upsample backward optimization - #245

Attention bwd: arc-segment iteration + upsample backward optimization#245
azrael417 wants to merge 15 commits into
mainfrom
tkurth/attention-bwd-segments

azrael417 commented Aug 19, 2026 •

edited

Loading

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Uh oh!

Conversation

azrael417 commented Aug 19, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

azrael417 commented Aug 19, 2026 •

edited

Loading