Repository navigation
learning: fused group norm + scale/shift + SiLU, and 3D neighbourhood attention, as CUDA passes - #167
Merged
Merged
Conversation
NeighborhoodAttention3d lets each voxel of a 3D feature map attend to the window^3 voxels around it (the window shifted inwards at an edge), with a learned bias per relative position. Its cost is linear in the volume, so it fits a denoiser at 64^3 where full attention would need over 100 GB. On a card the operation and its gradient run as three passes compiled into the library (src/csrc/ops/attention.cu): gather_dot, gather_sum and a scatter_sum written as a gather, so no atomics and repeatable sums. Float32, float16 and bfloat16, 8 to 64 channels per head, windows up to 9, sums in float32. A build without CUDA exports the entry points and reports them absent; anywhere else a plain-torch path serves. scripts/check_attention.py compares the passes and their gradients with the plain-torch path and times them; --standalone builds attention.cu alone with nvcc, so the kernels can be checked without building BART. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01MFB6QPF5rgL8sKizNYzB8D
…hout importing bartorch Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01MFB6QPF5rgL8sKizNYzB8D
…al; documented; exact reference The plain-torch path computes in float64 for float64 inputs, so the tests hold it to the voxel-by-voxel attention to round-off. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01MFB6QPF5rgL8sKizNYzB8D
…hift and SiLU in one CUDA pass Two short passes take a group's float32 statistics (Welford, a group spread over several blocks and merged), and one more reads the activations once more and writes them once, in their own element type, with the norm's affine, the per-item scale and shift and SiLU folded into one multiply-add per row. Eager torch reads and writes them four or five times, with casts around a float32 norm, which in MRI2Qmap's teacher costs more than its convolutions. Inference only: under autograd, without a card or without CUDA in the library, plain torch computes the same. scripts/check_groupnorm.py --standalone builds groupnorm.cu alone with nvcc to check and time it. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01MFB6QPoLKTK5gZEQeb1W8D
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
… the host _neighborhood._on gives the card's stream with that card current, and no stream off a card; the group norm's call moves into _launch. A stand-in library doing the passes' arithmetic in torch on the host memory it is handed then checks everything on the Python side of the ABI without a card: pointers, shapes, workspace, the error report, and the attention's autograd function with its gradients against the plain-torch path. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01MFB6QPoLKTK5gZEQeb1W8D
The standalone checks run nvcc with a fixed argument list on a file of this checkout; nothing reaches it from outside. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01MFB6QPoLKTK5gZEQeb1W8D
Collaborator
Author
|
CodeFactor reports 8 minor issues on this PR. I can't see which ones, because its issue list only renders in a signed-in browser. My guess was bandit's subprocess warnings in the two Coverage: the CUDA calls are now exercised on the host through a stand-in library in 371f622 ( Generated by Claude Code |
cuDNN's fastest 3-D convolutions run channels last; a norm that only reads NCDHW forces a transpose on either side. The pass now takes channels-last input as it lies (statistics per channel strip, groups merged after) and writes channels-last output. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01MFB6QPF5rgL8sKizNYzB8D
The statistics pass split each item into as many runs as a single group's size called for, so a deep level (512 channels at 8^3, three items) ran on three blocks; the A40 showed it 3-12x slower than the NCDHW pass below 64^3 and 2.5-4x at 64^3. Runs are now sized by the item's whole extent, at least 4096 elements each and at most 256 per item. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01MFB6QPF5rgL8sKizNYzB8D
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What changed
learning.group_norm_silu(x, groups, weight, bias, scale, shift, eps, silu): a diffusion UNet's block in one pass over memory.src/csrc/ops/groupnorm.cu: Welford fp32 statistics spread over several blocks, then one apply pass in the input's own dtype (fp32, fp16 or bf16).bartorch_gn_built,bartorch_gn_workspaceandbartorch_gn_silu, with stubs ingroupnorm.c.learning.NeighborhoodAttention3d/neighborhood_attention3d: each voxel attends to a shifted window^3 around it, through three CUDA passes (attention.cu) with their gradients, and a plain-torch fallback.scripts/check_groupnorm.pyandscripts/check_attention.py:--standalonebuilds the one.cuwith nvcc and checks and times it on a card, without building BART.Why
MRI2Qmap's teacher denoiser spent about three quarters of its time outside its convolutions: fp32 group norms, casts and unfused attention. Prebuilt kernels in the wheel are allowed on the targets; compiling at run time is not. Together with the mri2qmap-bartorch changes, the teacher's recon goes from 181 s to 140 s on an A40, with identical maps.
Validation
./scripts/run_tests.sh(CPU build, clang): 2550 passed, 177 skipped. The CUDA tests skip without a card.check_groupnorm.py --standalone --arch 86on an A40 (nvcc 12.8): max relative error against float64 1.7e-7 (fp32), 3.0e-4 (fp16), 3.5e-3 (bf16). Fused vs eager fp32 norm: (128, 64³) 0.58 vs 5.21 ms, (256, 32³) 0.38 vs 1.39 ms.check_attention.py --standaloneon an RTX 4060: max rel error fp32 1.4e-6, fp16 1.5e-3, bf16 1.2e-2. 64³, 2 heads, window 5: 23 ms forward against 0.91 s for plain torch..cufiles were compiled standalone with nvcc on the cards above.Documentation impact
New entries in
docs/api/learning.mdforgroup_norm_silu,NeighborhoodAttention3dandneighborhood_attention3d. A release is needed for mri2qmap-bartorch to use the fused norm.Generated by Claude Code