EquivariantLinear: one row-major block-diagonal GEMM — no layout churn - #14
Merged
Merged
Conversation
torch.profiler on the joint force step (512-atom box, A100) showed aten::copy_ at ~45 ms/step — 29% of everything, 312 copies — traced to this layer's einsum: it lowers to a batched bmm over the angular axis, whose batch-contiguity requirement copies the full per-edge activation on every call and emits an m-major output that every consumer (the fused nonlinearity's wrapper, residual adds, the next einsum) pays again to re-normalize. The map y[e,o,m] = sum_i x[e,i,m] W[m,o,i] never mixes angular modes, so it is now evaluated as one dense row-major GEMM over the flattened (feature, angular) axis against a per-m block-diagonal (in*n_ang, out*n_ang) weight assembled from the parameter each call (MB-scale; autograd flows through the assembly; bias enters the addmm epilogue on the m=0 cos slots). The zero blocks cost n_ang-times the strictly needed flops, which TF32/tensor cores absorb; in exchange there are no permutes, no copies, and the row-major output is exactly the layout the fused RealSpace kernel's fast path wants, so its fallback copies disappear too. Parameters and state_dict are unchanged — existing checkpoints load as before. tests/test_equivariant_linear.py pins the new forward to the einsum reference: forward, input/weight/bias gradients, non-contiguous m-major inputs, batched leading dims — all exact at fp64. Existing suites (realspace kernel, bottleneck, attention MP, ecenet, edge-frame kernel, calculator, LES) pass unchanged. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_019WyADHSVvKbQnYo1PAjGPK
…bsorb it The dense block-diagonal GEMM pays 3x flops for its zero blocks. With TF32 that padding is free and removing the einsum layout churn nets 159.6 -> 111.4 ms on the joint step; on fp32 CUDA cores with TF32 off it costs exactly 3x (linear_down 2.1 -> 5.6 ms measured) and regressed the strict-fp32 step to 244.8 ms, and float64 — the calculator default — would pay the same. Dispatch is now automatic: dense for CUDA fp16/bf16 and for CUDA fp32 with TF32 enabled, the original einsum path otherwise; dense_gemm=True/False overrides. Tests force the dense path (auto picks einsum on CPU) and cover the dispatch logic on both CPU and CUDA. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_019WyADHSVvKbQnYo1PAjGPK
Collaborator
Author
|
Follow-up commit 8fe30b6: the dense GEMM's 3× zero-block padding is only free on tensor cores — with TF32 off it cost exactly 3× on the fp32 CUDA-core path (joint step 218.6 → 244.8 ms regression), and float64 (the calculator default) would pay the same. Dispatch is now automatic: dense for CUDA fp16/bf16 and for CUDA fp32 with TF32 enabled; the original einsum path otherwise ( Verification on GPU: |
Without it a float32 ecenet benchmark silently takes the einsum path (the dense EquivariantLinear GEMM auto-dispatches only when allow_tf32 is set in-process) and the CSV could not distinguish TF32 runs. The flag sets the torch backend switches process-wide and stamps '+tf32' into the dtype column. DPA runs are unaffected by torch flags (DeePMD precision is set via DP_TF32_INFER in its own environment) — noted in the help text. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_019WyADHSVvKbQnYo1PAjGPK
… precision The dipole/BEC evals could only run at the checkpoint's native matmul precision; with the TF32 benchmarking path there was no way to measure what TF32 costs those observables. The flag mirrors profile_step/ benchmark_calculator: sets the torch backend switches, warns when the device is not CUDA. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_019WyADHSVvKbQnYo1PAjGPK
Completes the TF32 story: the energy/force test-set evaluation (the headline accuracy axis) can now run at the same precision as the benchmarks. Same backend switches as benchmark_calculator/profile_step; warns when dtype is float64, and the printed dtype line records +tf32. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_019WyADHSVvKbQnYo1PAjGPK
The edge-frame family is the largest remaining kernel cost (~25 ms/step at 512 atoms: _ef_bwd_merged 2.33 ms/call, _pu_bwd_merged 1.98, fwd ~1.0), all running ~3.5x above their traffic bound. Unlike the RealSpace case the loads are already coalesced tiles, so the gap is some mix of padded ieee tl.dot arithmetic (9 valid of 16 in both dot dims), 36-byte row misalignment, and per-edge program overhead — not separable without measurement. First step: make warps-per-program env-tunable (default 4, the previous implicit value) on all ten launch sites so the cheap axis can be swept before any redesign. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_019WyADHSVvKbQnYo1PAjGPK
ncu on the merged backward kernels (A100, 512-atom box, 44k edges) settles the question the flat warps sweep left open: L1/TEX throughput 75-86% with DRAM at 9-18% and compute ~30% — the kernels are L1-bound, stalling on MIO short-scoreboard, while DRAM idles. The amplification is the per-element gathered D loads through the (l,m) column tables and the matching dD scatter-stores: hundreds of non-vectorizable L1 transactions per program. The packed variants spend the idle DRAM instead: the wrapper pre-packs D into dense (E, S, P)/(E, P, S) tensors once per call (plain torch indexing, ~14 MB at 44k edges) and unpacks the packed dD afterwards, so the kernels do only vectorized coalesced tile IO. Math and masks identical. Gated off by default behind ECENET_EF_PACKD=1 (module flag, monkeypatchable) pending an A/B on the A100; ncu's stall-fix estimate is ~33-37% on these kernels. test_triton_packed_d reruns the full test_triton_paths comparison set (both merged backwards, vs fp64 eager truth) with the flag on. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_019WyADHSVvKbQnYo1PAjGPK
…block The packed kernels won at the kernel level (ncu: _ef_bwd_merged 2.94 -> 1.68 ms, _pu_bwd_merged 2.49 -> 1.95 ms) but the joint step stayed flat: per-call packing/transposing in the backward wrappers (~8 small torch ops x 7 backward calls/step) ate the entire ~5 ms win. The model builds ONE D_block per step and hands the same Python object to every fused edge-frame op, so _get_packed_D now caches the packed tensors as an attribute on that object — packing runs once per step. Function forwards fetch it and carry the tensors through save_for_backward (attributes do not survive the re-wrapping); the backward wrappers receive them instead of packing. A fresh step's fresh D_block starts clean, so no cross-step staleness. Still gated behind ECENET_EF_PACKD=1 pending the A/B. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_019WyADHSVvKbQnYo1PAjGPK
torch.profiler with the packed path on showed the GPU win arriving (packed kernels 1.35/1.65 ms per call, CUDA total down) while wall time stayed flat, and named the thief: EdgeFrameFusedBackward at 10.6 ms CPU per call, with aten::index at 44% of CPU total. _unpack_dD's boolean-mask indexing (cos_col[vc], dD[:, :, col[vc]]) forces a host-device sync on every backward call — seven pipeline stalls per step. The masks are static, so the unpack now uses integer index tensors precomputed once per (S, n_ang, device); every op in the unpack is async. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_019WyADHSVvKbQnYo1PAjGPK
…ault on With the sync-free unpack the packed backward finally reached the wall clock: joint force step 111.4 -> 107.8 ms on the A100 512-atom box (EdgeFrameFusedBackward CPU 10.6 ms -> 0.29 ms per call, aten::index off the CPU hot list, packed kernels at 1.35/1.65 ms in-app). This extends the same treatment to the forward side, which carries the identical table-gather pattern: _ef_fwd_packed_kernel (EdgeFrameFused forward, ~1.04 ms x 4/step) and _ef_bwd_dx_packed_kernel (PackUnrotate's forward contraction, ~1.17 ms x 3/step) read the pre-packed Dc/Ds / DcT/DsT tiles instead of gathering through the column tables; amortization is free since the packed tensors already live on the shared per-step D_block. Default flipped ON (ECENET_EF_PACKD=0 restores the table-gather kernels); test_triton_packed_d now runs the full fp64 comparison suite under both settings. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_019WyADHSVvKbQnYo1PAjGPK
Same kwarg and enable block as train_ecenet_spice / train_ecenet_xyz: routes float32 matmuls to TF32 tensor cores, warns under float64, and reminds to A/B the val MAE. Completes the TF32 story across training, evaluation, and benchmarking. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_019WyADHSVvKbQnYo1PAjGPK
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_019WyADHSVvKbQnYo1PAjGPK
Edge-frame kernels: packed-D tile IO (measurement-driven), sweepable warps
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
1 participant
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.
Summary
torch.profiler (via #13's
--torch_profile) showedaten::copy_at ~45 ms/step — 29% of the joint force step — traced toEquivariantLinear's einsum: it lowers to a batched bmm over the angular axis, which copies the full per-edge activation for batch contiguity on every call and emits an m-major output that every consumer pays again to re-normalize.Since the map
y[e,o,m] = Σᵢ x[e,i,m]·W[m,o,i]never mixes angular modes, it is now one dense row-major GEMM over the flattened (feature, angular) axis against a per-m block-diagonal weight assembled from the unchanged parameter each call (bias in the addmm epilogue on the m=0 cos slots). No permutes, no copies, and the row-major output is exactly the layout the fused RealSpace kernel's fast path wants. The zero blocks cost n_ang× the strictly needed flops, absorbed by TF32 tensor cores. Parameters/state_dict unchanged — existing checkpoints load as before.Measured (A100, 512-atom diamond box, 44k edges, float32+TF32, all fusion flags):
aten::copy_out of the profiler top-25 (was Sync from dev #1);aten::bmm890 → 40 calls/stepTest plan
tests/test_equivariant_linear.py: forward, input/weight/bias gradient, non-contiguous m-major input, and batched-leading-dim equivalence vs the einsum reference — exact at fp64.🤖 Generated with Claude Code
https://claude.ai/code/session_019WyADHSVvKbQnYo1PAjGPK