Skip to content

cuda: attn_train and train_attn compile one gated attention backward, attn_back.cuh, SASS unchanged - #1565

Merged
ashhart merged 2 commits into
ashhart:mainfrom
jschmied:prims-1409-attnback
Oct 29, 2026
Merged

ashhart merged 2 commits into
ashhart:mainfrom
jschmied:prims-1409-attnback

Conversation

@jschmied

Copy link
Copy Markdown
Contributor

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 by src/cuda/train_attn.zig) was a copy of the attention part of attn_train.cu (Living Weights, loaded by qwen3_5's cuda_train_ops.zig). It has the same four kernels under tf_gattn_ instead of tf_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 (or tf_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.cu keeps tf_sum_parts and tf_axpy after the include.

3 files, +167 / −324.

Verification

  • SASS: base cuda: sigmoidf_, softplusf_, silu and silu_grad in common/act.cuh, SASS unchanged #1556's head f19336386, head b50786bfb, 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.
    python3 tools/zig/cuda_sass_identity.py --repo . -j 3 --receipt sass.tsv f19336386 HEAD
    
  • Planted control: a reversed warp_sum butterfly changes 30 functions in 2 of 2 images. The fp4pair control cannot be planted, because neither image includes common/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.
  • lean: 425 problems, 425 on the base, 0 new.
  • 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's tf_ck8_reduce with qmmf's reduce_kernel<false>, and DS4.1's cast_bf16_kernel with GLM vision's glm_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.

jschmied and others added 2 commits October 29, 2026 23:40
… 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
ashhart force-pushed the prims-1409-attnback branch from b50786b to 645f0b4 Compare October 29, 2026 23:40
@ashhart
ashhart merged commit e5c30d1 into ashhart:main Oct 29, 2026
9 checks passed
@ashhart

ashhart commented Oct 31, 2026

Copy link
Copy Markdown
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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants