Repository navigation
cuda: attn_train and train_attn compile one gated attention backward, attn_back.cuh, SASS unchanged - #1565
Merged
Conversation
… attn_back.cuh, SASS unchanged ashhart#1409. train_attn.cu (4-bit training, src/cuda/train_attn.zig) was a copy of attn_train.cu's attention part (Living Weights, qwen3_5's cuda_train_ops.zig): the same four kernels under tf_gattn_ instead of tf_attn_ names, and the same helpers, differing only in comments (a SASS census of every image at sm_121 found the four functions identical). The text now lives once in attn_back.cuh; each unit names its kernels with #define TF_ATTN_BACK(name) before the include, so both images keep their symbols and the loaders, contract records and tests stay as they are. attn_train.cu keeps tf_sum_parts and tf_axpy after the include.
ashhart
force-pushed
the
prims-1409-attnback
branch
from
October 29, 2026 23:40
b50786b to
645f0b4
Compare
Owner
|
Thank you @jschmied. This is in TensorFold 1.2.0, released today: https://github.com/ashhart/TensorFold/releases/tag/v1.2.0 It landed with your commits under your own name, so you're listed as a TensorFold contributor, and it's credited to you in the release notes and the changelog. |
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
Part of #1409. Stacked on #1556, with one commit on top, because #1556 edits the same two files (
sigm→sigmoidf_).train_attn.cu(4-bit training, loaded bysrc/cuda/train_attn.zig) was a copy of the attention part ofattn_train.cu(Living Weights, loaded by qwen3_5'scuda_train_ops.zig). It has the same four kernels undertf_gattn_instead oftf_attn_names (prep,q,gate_back,prep_back) and the same helpers; only the comments differ. A SASS census of every image at sm_121 found the four functions identical across the two images. That census is how this copy showed up: K0's text comparison missed it because of the names.The text now lives once, in
attn_back.cuh. Each unit names its kernels with#define TF_ATTN_BACK(name) tf_attn_##name(ortf_gattn_##name) before the include. Both images keep their symbols, so the loaders, contract records (gradient_pending.zig), tests and receipts stay as they are.attn_train.cukeepstf_sum_partsandtf_axpyafter the include.3 files, +167 / −324.
Verification
f19336386, headb50786bfb, nvcc 13.0, sm_80/86/89/120/121. Both images (attn_train,train_attn), all 50 functions: 0 changed, 0 added, 0 removed, 0 unstable.warp_sumbutterfly changes 30 functions in 2 of 2 images. Thefp4paircontrol cannot be planted, because neither image includescommon/qmm_frag.cuh, so the tool exits 1 on that line alone.zig build test-gradient-serving: the same set of differing pins as main. Neither file is pinned.zig build test -Dnvcc=/usr/local/cuda/bin/nvcc: passes (RC 0). It ran before a comment-only fix (two header comment lines folded into one), and the SASS was re-run after it.Not included
The census also found two pairs of tiny kernels that compile to the same code:
ck_lane'stf_ck8_reducewithqmmf'sreduce_kernel<false>, and DS4.1'scast_bf16_kernelwith GLM vision'sglm_vis_to_bf16. Each is 4 to 6 lines, with different parameter types in different families, so a shared body would add more text than it removes. I left them alone.