From 16c642f2805a891d783ab7f29ca4340cb4424795 Mon Sep 17 00:00:00 2001 From: alacour Date: Thu, 3 Sep 2026 18:27:51 -0700 Subject: [PATCH 01/12] =?UTF-8?q?EquivariantLinear:=20one=20row-major=20bl?= =?UTF-8?q?ock-diagonal=20GEMM=20=E2=80=94=20no=20layout=20churn?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 Claude-Session: https://claude.ai/code/session_019WyADHSVvKbQnYo1PAjGPK --- ecenet/equivariant.py | 42 ++++++++-- tests/test_equivariant_linear.py | 129 +++++++++++++++++++++++++++++++ 2 files changed, 164 insertions(+), 7 deletions(-) create mode 100644 tests/test_equivariant_linear.py diff --git a/ecenet/equivariant.py b/ecenet/equivariant.py index e2f793a..6e3a0f2 100644 --- a/ecenet/equivariant.py +++ b/ecenet/equivariant.py @@ -26,6 +26,20 @@ class EquivariantLinear(nn.Module): Same weights for cos/sin parts. Bias only on m=0 (invariant). Angular channels: m = 0, 1, ..., m_max (index 0 is m=0). + + The map y[e,o,m] = Σ_i x[e,i,m]·W[m,o,i] never mixes angular modes, so it + is evaluated as ONE dense row-major GEMM over the flattened (feature, + angular) axis, against a (in·n_ang, out·n_ang) weight that is + block-diagonal per m (assembled from the parameter each call — ~MB-scale, + autograd flows through the assembly; the zero blocks cost n_ang× the + strictly needed flops, which TF32/tensor cores absorb). The earlier + einsum formulation lowered to a batched bmm over the angular axis, whose + batch-contiguity requirement copied the full per-edge activation on every + call and emitted an m-major output that downstream consumers (the fused + nonlinearity, residual adds, the next einsum) paid to re-normalize — + profiling showed that layout churn at ~29% of the whole force step. Here + input and output are both plain row-major views: no permutes, no copies, + and the layout every consumer wants. """ def __init__(self, in_features, out_features, n_angular, m_max): @@ -41,14 +55,28 @@ def __init__(self, in_features, out_features, n_angular, m_max): self.bias = nn.Parameter(torch.zeros(out_features)) - def forward(self, A_cos, A_sin): - A_cos_out = torch.einsum('...id,doi->...od', A_cos, self.weights) - A_sin_out = torch.einsum('...id,doi->...od', A_sin, self.weights) - - # Bias only on m=0 (index 0) - A_cos_out[..., 0] = A_cos_out[..., 0] + self.bias + def _dense_weight(self): + """(in·n_ang, out·n_ang) block-diagonal-per-m weight, row-major flat + index c = feature·n_ang + m on both sides.""" + na, Fo, Fi = self.weights.shape + Wd = self.weights.new_zeros(Fi, na, Fo, na) + m = torch.arange(na, device=self.weights.device) + Wd[:, m, :, m] = self.weights.permute(0, 2, 1) # (na, Fi, Fo) slot + return Wd.reshape(Fi * na, Fo * na) - return A_cos_out, A_sin_out + def forward(self, A_cos, A_sin): + na, Fo = self.n_angular, self.out_features + lead = A_cos.shape[:-2] + Wd = self._dense_weight() + # Bias enters the GEMM epilogue (addmm), on the m=0 slots of the cos + # part only — sin has no invariant channel. + bias_full = self.bias.new_zeros(Fo, na) + bias_full[:, 0] = self.bias + x_cos = A_cos.reshape(-1, self.in_features * na) + x_sin = A_sin.reshape(-1, self.in_features * na) + y_cos = torch.addmm(bias_full.reshape(-1), x_cos, Wd) + y_sin = x_sin @ Wd + return y_cos.view(*lead, Fo, na), y_sin.view(*lead, Fo, na) class RealSpaceNonlinearity(nn.Module): diff --git a/tests/test_equivariant_linear.py b/tests/test_equivariant_linear.py new file mode 100644 index 0000000..c4598b7 --- /dev/null +++ b/tests/test_equivariant_linear.py @@ -0,0 +1,129 @@ +"""Tests for the block-diagonal-GEMM EquivariantLinear (``ecenet/equivariant.py``). + +The layer computes y[e,o,m] = Σ_i x[e,i,m]·W[m,o,i] (bias on the m=0 cos slot +only) as one dense row-major GEMM against a per-m block-diagonal weight. These +tests pin it to the original einsum formulation: + + 1. forward equivalence vs the einsum reference (several shapes, fp64); + 2. gradient equivalence — input, weight, and bias grads vs the reference; + 3. non-contiguous (m-major, einsum-layout) inputs give the same result; + 4. batched leading dims (B, n_e, F, n_ang) round-trip unchanged. + +Pure PyTorch on CPU (fp64). Run: python tests/test_equivariant_linear.py +""" + +import os +import sys # noqa: E402 — repo root on path for `import ecenet` when run as a script + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + + +import torch + +from ecenet.equivariant import EquivariantLinear + +torch.manual_seed(0) +DTYPE = torch.float64 + + +def reference(A_cos, A_sin, weights, bias): + """The original einsum formulation — the spec.""" + oc = torch.einsum('...id,doi->...od', A_cos, weights) + os_ = torch.einsum('...id,doi->...od', A_sin, weights) + oc = oc.clone() + oc[..., 0] = oc[..., 0] + bias + return oc, os_ + + +def make_layer(Fi=12, Fo=9, m_max=2, seed=0): + torch.manual_seed(seed) + lin = EquivariantLinear(Fi, Fo, m_max + 1, m_max).to(DTYPE) + with torch.no_grad(): # non-trivial bias + lin.bias.add_(torch.randn_like(lin.bias)) + return lin + + +def make_inputs(n_e=40, Fi=12, m_max=2, seed=1): + g = torch.Generator().manual_seed(seed) + mk = lambda: torch.randn(n_e, Fi, m_max + 1, generator=g, dtype=DTYPE) # noqa: E731 + return mk(), mk() + + +def test_forward_equivalence(): + worst = 0.0 + for Fi, Fo, m_max in ((12, 9, 2), (7, 7, 3), (1, 5, 0), (16, 3, 1)): + lin = make_layer(Fi, Fo, m_max) + A_cos, A_sin = make_inputs(Fi=Fi, m_max=m_max) + oc, os_ = lin(A_cos, A_sin) + rc, rs = reference(A_cos, A_sin, lin.weights, lin.bias) + e = max((oc - rc).abs().max(), (os_ - rs).abs().max()).item() + worst = max(worst, e) + assert e < 1e-12, f"forward Fi={Fi},Fo={Fo},m={m_max}: {e:.2e}" + print(f" forward equivalence vs einsum reference (worst {worst:.1e})") + + +def test_gradient_equivalence(): + lin = make_layer() + A_cos, A_sin = make_inputs(seed=2) + goc = torch.randn_like(A_cos[:, :9]) # (n_e, Fo, n_ang) + gos = torch.randn_like(A_cos[:, :9]) + + def grads(fn): + a = A_cos.detach().clone().requires_grad_(True) + b = A_sin.detach().clone().requires_grad_(True) + lin.zero_grad() + oc, os_ = fn(a, b) + (oc * goc + os_ * gos).sum().backward() + return a.grad, b.grad, lin.weights.grad.clone(), lin.bias.grad.clone() + + new = grads(lambda a, b: lin(a, b)) + ref = grads(lambda a, b: reference(a, b, lin.weights, lin.bias)) + worst = 0.0 + for name, n, r in zip(("dA_cos", "dA_sin", "dW", "db"), new, ref): + e = (n - r).abs().max().item() + worst = max(worst, e) + assert e < 1e-12, f"{name} mismatch: {e:.2e}" + print(f" gradient equivalence: input/weight/bias grads match (worst {worst:.1e})") + + +def test_noncontiguous_input(): + """m-major (einsum-layout) inputs — strides (F, 1, n_e·F) — must give the + same result as contiguous ones (reshape copies as needed internally).""" + lin = make_layer() + A_cos, A_sin = make_inputs(seed=3) + + def as_m_major(t): + n_e, F, na = t.shape + out = torch.empty(na, n_e, F, dtype=t.dtype).permute(1, 2, 0) + out.copy_(t) + assert not out.is_contiguous() + return out + + oc, os_ = lin(A_cos, A_sin) + mc, ms = lin(as_m_major(A_cos), as_m_major(A_sin)) + e = max((oc - mc).abs().max(), (os_ - ms).abs().max()).item() + assert e < 1e-15, f"m-major input mismatch: {e:.2e}" + print(f" non-contiguous (m-major) input matches contiguous ({e:.1e})") + + +def test_batched_leading_dims(): + lin = make_layer() + A_cos, A_sin = make_inputs(n_e=24, seed=4) + B_cos = A_cos.reshape(4, 6, 12, 3) + B_sin = A_sin.reshape(4, 6, 12, 3) + oc, os_ = lin(A_cos, A_sin) + bc, bs = lin(B_cos, B_sin) + assert bc.shape == (4, 6, 9, 3) and bs.shape == (4, 6, 9, 3) + e = max((bc.reshape(24, 9, 3) - oc).abs().max(), + (bs.reshape(24, 9, 3) - os_).abs().max()).item() + assert e < 1e-15, f"batched mismatch: {e:.2e}" + print(f" batched leading dims match flat ({e:.1e})") + + +if __name__ == "__main__": + print("EquivariantLinear block-diagonal-GEMM tests") + test_forward_equivalence() + test_gradient_equivalence() + test_noncontiguous_input() + test_batched_leading_dims() + print("All tests passed.") From 8fe30b63fef7415ec9990f1d48fd5decb396dd9c Mon Sep 17 00:00:00 2001 From: alacour Date: Thu, 3 Sep 2026 18:37:37 -0700 Subject: [PATCH 02/12] EquivariantLinear: auto-dispatch dense GEMM only where tensor cores absorb it MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 Claude-Session: https://claude.ai/code/session_019WyADHSVvKbQnYo1PAjGPK --- ecenet/equivariant.py | 30 +++++++++++++++++++++++++++--- tests/test_equivariant_linear.py | 29 +++++++++++++++++++++++++++++ 2 files changed, 56 insertions(+), 3 deletions(-) diff --git a/ecenet/equivariant.py b/ecenet/equivariant.py index 6e3a0f2..93f4614 100644 --- a/ecenet/equivariant.py +++ b/ecenet/equivariant.py @@ -37,9 +37,16 @@ class EquivariantLinear(nn.Module): batch-contiguity requirement copied the full per-edge activation on every call and emitted an m-major output that downstream consumers (the fused nonlinearity, residual adds, the next einsum) paid to re-normalize — - profiling showed that layout churn at ~29% of the whole force step. Here - input and output are both plain row-major views: no permutes, no copies, - and the layout every consumer wants. + profiling showed that layout churn at ~29% of the whole force step. The + dense form's input and output are both plain row-major views: no + permutes, no copies, and the layout every consumer wants. + + The trade only pays where the padded GEMM is (near-)free: on CUDA tensor + cores. On fp32 CUDA cores (TF32 off) and in float64 the 3x padding costs + exactly 3x — measured slower than the einsum path despite the copies — + so dispatch is automatic: dense for CUDA fp16/bf16, and for CUDA fp32 + with TF32 enabled; einsum otherwise. ``dense_gemm`` (True/False) + overrides the automatic choice. """ def __init__(self, in_features, out_features, n_angular, m_max): @@ -48,6 +55,7 @@ def __init__(self, in_features, out_features, n_angular, m_max): self.out_features = out_features self.n_angular = n_angular self.m_max = m_max + self.dense_gemm = None # None = auto (see _use_dense) # (n_angular, out_features, in_features) std = (2.0 / (in_features + out_features)) ** 0.5 @@ -55,6 +63,16 @@ def __init__(self, in_features, out_features, n_angular, m_max): self.bias = nn.Parameter(torch.zeros(out_features)) + def _use_dense(self, x): + if self.dense_gemm is not None: + return self.dense_gemm + if not x.is_cuda: + return False + if x.dtype in (torch.float16, torch.bfloat16): + return True + return (x.dtype == torch.float32 + and torch.backends.cuda.matmul.allow_tf32) + def _dense_weight(self): """(in·n_ang, out·n_ang) block-diagonal-per-m weight, row-major flat index c = feature·n_ang + m on both sides.""" @@ -65,6 +83,12 @@ def _dense_weight(self): return Wd.reshape(Fi * na, Fo * na) def forward(self, A_cos, A_sin): + if not self._use_dense(A_cos): + A_cos_out = torch.einsum('...id,doi->...od', A_cos, self.weights) + A_sin_out = torch.einsum('...id,doi->...od', A_sin, self.weights) + A_cos_out[..., 0] = A_cos_out[..., 0] + self.bias # bias on m=0 only + return A_cos_out, A_sin_out + na, Fo = self.n_angular, self.out_features lead = A_cos.shape[:-2] Wd = self._dense_weight() diff --git a/tests/test_equivariant_linear.py b/tests/test_equivariant_linear.py index c4598b7..e59727d 100644 --- a/tests/test_equivariant_linear.py +++ b/tests/test_equivariant_linear.py @@ -38,6 +38,7 @@ def reference(A_cos, A_sin, weights, bias): def make_layer(Fi=12, Fo=9, m_max=2, seed=0): torch.manual_seed(seed) lin = EquivariantLinear(Fi, Fo, m_max + 1, m_max).to(DTYPE) + lin.dense_gemm = True # force the dense path (auto picks einsum on CPU) with torch.no_grad(): # non-trivial bias lin.bias.add_(torch.randn_like(lin.bias)) return lin @@ -120,10 +121,38 @@ def test_batched_leading_dims(): print(f" batched leading dims match flat ({e:.1e})") +def test_dispatch(): + """Auto-dispatch: einsum on CPU / strict fp32 / fp64; dense only where the + padded GEMM rides tensor cores (CUDA fp16/bf16, or fp32 with TF32 on). + dense_gemm=True/False overrides both ways.""" + lin = EquivariantLinear(4, 4, 3, 2) + x64 = torch.randn(2, 4, 3, dtype=torch.float64) + assert not lin._use_dense(x64), "CPU/fp64 must take the einsum path" + lin.dense_gemm = True + assert lin._use_dense(x64), "dense_gemm=True must override" + lin.dense_gemm = False + assert not lin._use_dense(x64), "dense_gemm=False must override" + lin.dense_gemm = None + if torch.cuda.is_available(): + x32 = torch.randn(2, 4, 3, dtype=torch.float32, device='cuda') + old = torch.backends.cuda.matmul.allow_tf32 + try: + torch.backends.cuda.matmul.allow_tf32 = False + assert not lin._use_dense(x32), "CUDA fp32 without TF32 → einsum" + torch.backends.cuda.matmul.allow_tf32 = True + assert lin._use_dense(x32), "CUDA fp32 with TF32 → dense" + finally: + torch.backends.cuda.matmul.allow_tf32 = old + print(" dispatch: einsum on CPU/strict-fp32, dense on TF32; overrides work") + else: + print(" dispatch: einsum on CPU, overrides work (CUDA cases skipped)") + + if __name__ == "__main__": print("EquivariantLinear block-diagonal-GEMM tests") test_forward_equivalence() test_gradient_equivalence() test_noncontiguous_input() test_batched_leading_dims() + test_dispatch() print("All tests passed.") From 1e609a4bc191911c6cae02b624dc1c2939c09a20 Mon Sep 17 00:00:00 2001 From: alacour Date: Thu, 3 Sep 2026 18:45:26 -0700 Subject: [PATCH 03/12] benchmark_calculator: --tf32 flag, recorded in the CSV dtype column MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 Claude-Session: https://claude.ai/code/session_019WyADHSVvKbQnYo1PAjGPK --- tools/benchmark_calculator.py | 20 +++++++++++++++++++- 1 file changed, 19 insertions(+), 1 deletion(-) diff --git a/tools/benchmark_calculator.py b/tools/benchmark_calculator.py index 93e0dc4..9ea0413 100644 --- a/tools/benchmark_calculator.py +++ b/tools/benchmark_calculator.py @@ -169,6 +169,13 @@ def main(): help='replicate the box into an AxBxC supercell (size scaling)') p.add_argument('--float32', action='store_true', help='float32 where the factory supports it (state it in the table)') + p.add_argument('--tf32', action='store_true', + help='enable TF32 matmuls process-wide (torch backends flag; ' + 'float32 only). ecenet: also required for the dense ' + 'EquivariantLinear GEMM path to auto-dispatch. Recorded ' + 'in the CSV dtype column. Note --calc dpa ignores torch ' + "flags — DeePMD's precision is set via DP_TF32_INFER " + 'and friends in its own environment.') p.add_argument('--fuse', action='store_true', help="enable the model's fast paths (ecenet: edge-frame + " 'activation fusion; mace: enable_cueq; no-op for ' @@ -186,6 +193,16 @@ def main(): p.add_argument('--csv', default=None, help='append one result row (collate runs)') args = p.parse_args() + if args.tf32: + if torch is None: + sys.exit('--tf32 needs torch installed') + torch.backends.cuda.matmul.allow_tf32 = True + torch.backends.cudnn.allow_tf32 = True + torch.set_float32_matmul_precision('high') + if not args.float32: + print('[tf32] set, but without --float32 most factories run float64 ' + '— TF32 will not apply there') + atoms = read(args.box, index=args.frame_idx) if args.repeat: atoms = atoms.repeat(tuple(args.repeat)) @@ -248,7 +265,8 @@ def main(): 'checkpoint': args.checkpoint or '', 'box': os.path.basename(args.box), 'repeat': 'x'.join(map(str, args.repeat)) if args.repeat else '', 'n_atoms': n_atoms, - 'dtype': 'float32' if args.float32 else 'default/float64', + 'dtype': (('float32' if args.float32 else 'default/float64') + + ('+tf32' if args.tf32 else '')), 'fuse': args.fuse, 'device': gpu, 'host': os.uname().nodename, # login vs compute node — shared-GPU # login timings are contended garbage From 66484e690fd01972d32e3419f20354dcbf460e69 Mon Sep 17 00:00:00 2001 From: alacour Date: Thu, 3 Sep 2026 19:04:08 -0700 Subject: [PATCH 04/12] eval_spice tools: --tf32 flag so accuracy is measurable at deployment 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 Claude-Session: https://claude.ai/code/session_019WyADHSVvKbQnYo1PAjGPK --- tools/eval_spice_bec.py | 11 +++++++++++ tools/eval_spice_dipoles.py | 11 +++++++++++ 2 files changed, 22 insertions(+) diff --git a/tools/eval_spice_bec.py b/tools/eval_spice_bec.py index df74c92..c389e6c 100644 --- a/tools/eval_spice_bec.py +++ b/tools/eval_spice_bec.py @@ -167,12 +167,23 @@ def main(): ap.add_argument('--max_frames', type=int, default=None, help='cap frames per file (smoke runs)') ap.add_argument('--device', default='cpu') + ap.add_argument('--tf32', action='store_true', + help='enable TF32 matmuls (CUDA + float32 checkpoints only) ' + '— evaluate at the same precision the model is ' + 'benchmarked/deployed at') ap.add_argument('--save', default=None, help='write per-atom sign-aligned results to this file: ' '.csv (one row per atom: subset, frame, symbol, 9 ' 'pred + 9 ref components) or .npz (arrays)') args = ap.parse_args() + if args.tf32: + torch.backends.cuda.matmul.allow_tf32 = True + torch.backends.cudnn.allow_tf32 = True + torch.set_float32_matmul_precision('high') + if not args.device.startswith('cuda'): + print('[tf32] requested but device is not CUDA — no effect') + data_dir = os.path.expanduser(args.data_dir) if args.files: xyz_files = [os.path.join(data_dir, f if f.endswith('.xyz') diff --git a/tools/eval_spice_dipoles.py b/tools/eval_spice_dipoles.py index cd54c86..acad63f 100644 --- a/tools/eval_spice_dipoles.py +++ b/tools/eval_spice_dipoles.py @@ -171,12 +171,23 @@ def main(): "(reference frames are neutral; removal is upstream " "BEC's remove_mean=True convention)") ap.add_argument('--device', default='cpu') + ap.add_argument('--tf32', action='store_true', + help='enable TF32 matmuls (CUDA + float32 checkpoints only) ' + '— evaluate at the same precision the model is ' + 'benchmarked/deployed at') ap.add_argument('--save', default=None, help='write per-frame sign-aligned results to this file: ' '.csv (one row per frame, ready for a correlation ' 'plot) or .npz (arrays)') args = ap.parse_args() + if args.tf32: + torch.backends.cuda.matmul.allow_tf32 = True + torch.backends.cudnn.allow_tf32 = True + torch.set_float32_matmul_precision('high') + if not args.device.startswith('cuda'): + print('[tf32] requested but device is not CUDA — no effect') + data_dir = os.path.expanduser(args.data_dir) if args.files: xyz_files = [os.path.join(data_dir, f if f.endswith('.xyz') From 1aa6de5ab89606b43ca63383dd146df9507b0a32 Mon Sep 17 00:00:00 2001 From: alacour Date: Thu, 3 Sep 2026 19:08:16 -0700 Subject: [PATCH 05/12] =?UTF-8?q?eval=5Fspice:=20--tf32=20flag=20=E2=80=94?= =?UTF-8?q?=20measure=20test=20MAEs=20at=20deployment=20precision?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 Claude-Session: https://claude.ai/code/session_019WyADHSVvKbQnYo1PAjGPK --- scripts/eval_spice.py | 14 +++++++++++++- 1 file changed, 13 insertions(+), 1 deletion(-) diff --git a/scripts/eval_spice.py b/scripts/eval_spice.py index 8897c16..a135ca7 100644 --- a/scripts/eval_spice.py +++ b/scripts/eval_spice.py @@ -174,6 +174,10 @@ def main(): # Eval options parser.add_argument('--batch_size', type=int, default=8) parser.add_argument('--float32', action='store_true') + parser.add_argument('--tf32', action='store_true', + help='Enable TF32 matmuls (CUDA + --float32 only) — ' + 'measure the test MAEs at the same precision the ' + 'model is benchmarked/deployed at.') parser.add_argument('--device', default=None) parser.add_argument('--ignore_les', action='store_true', help='Evaluate the short-range part of an LES checkpoint ' @@ -183,7 +187,15 @@ def main(): dtype = torch.float32 if args.float32 else torch.float64 device = torch.device(args.device if args.device else ('cuda' if torch.cuda.is_available() else 'cpu')) - print(f"Device: {device}, dtype: {dtype}") + if args.tf32: + torch.backends.cuda.matmul.allow_tf32 = True + torch.backends.cudnn.allow_tf32 = True + torch.set_float32_matmul_precision('high') + if dtype == torch.float64: + print('[tf32] requested but dtype=float64 → no effect ' + '(TF32 is float32-only); add --float32 to use it') + print(f"Device: {device}, dtype: {dtype}" + f"{' +tf32' if args.tf32 and dtype == torch.float32 else ''}") # ── Load checkpoint ──────────────────────────────────────────────────── ckpt = torch.load(args.checkpoint, map_location=device, weights_only=False) From b1102589e1c5f5d7e71b310382000588d8374d7f Mon Sep 17 00:00:00 2001 From: alacour Date: Thu, 3 Sep 2026 19:20:40 -0700 Subject: [PATCH 06/12] =?UTF-8?q?edge=5Fframe=5Fkernel:=20ECENET=5FEF=5FWA?= =?UTF-8?q?RPS=20=E2=80=94=20sweepable=20warps=20on=20every=20launch?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 Claude-Session: https://claude.ai/code/session_019WyADHSVvKbQnYo1PAjGPK --- ecenet/edge_frame_kernel.py | 29 ++++++++++++++++++++++------- 1 file changed, 22 insertions(+), 7 deletions(-) diff --git a/ecenet/edge_frame_kernel.py b/ecenet/edge_frame_kernel.py index 6083ce3..979d7e1 100644 --- a/ecenet/edge_frame_kernel.py +++ b/ecenet/edge_frame_kernel.py @@ -40,6 +40,8 @@ repeat; ``cos_valid``/``sin_valid`` zero the |m| > l (and m=0 sin) slots. """ +import os + import torch from torch.profiler import record_function @@ -391,6 +393,13 @@ def _e2n_fwd_kernel(gc_ptr, gs_ptr, d_ptr, perm_ptr, aptr_ptr, _EF_TABLES: dict = {} +# Warps per program for every edge-frame kernel launch, overridable per GPU +# without a code change (read once at import): ECENET_EF_WARPS=2 python ... +# The kernels' tiles are small ((BLOCK_R, 16)-ish), so fewer warps than the +# Triton default of 4 may win; sweep alongside profile_step. +_EF_WARPS = int(os.environ.get('ECENET_EF_WARPS', 4)) + + def _next_pow2(x: int) -> int: """Smallest power of two ≥ x, floored at 16. tl.arange REQUIRES a power of two (a multiple of 16 like 48 or 96 compiles to "arange's range must @@ -492,7 +501,7 @@ def _e2n_forward_triton(g_cos, g_sin, edge_dst, D_block, n_atoms, n_base): g_cos.contiguous(), g_sin.contiguous(), D_block.contiguous(), perm, aptr, srcoff, okc, oks, Delta, n_base, S, P, - SP=_next_pow2(S), RB=_next_pow2(n_base)) + SP=_next_pow2(S), RB=_next_pow2(n_base), num_warps=_EF_WARPS) return Delta # Per-edge: pack+unrotate is exactly the dx backward kernel's math @@ -508,7 +517,8 @@ def _e2n_forward_triton(g_cos, g_sin, edge_dst, D_block, n_atoms, n_base): g_cos.contiguous(), g_sin.contiguous(), D_block.contiguous(), cos_col, cos_ok, sin_col, sin_ok, h_global, n_base, n_base, S, P, - SP=_next_pow2(S), P16=_next_pow2(P), BLOCK_R=block_r) + SP=_next_pow2(S), P16=_next_pow2(P), BLOCK_R=block_r, + num_warps=_EF_WARPS) Delta = torch.zeros(n_atoms, n_base, S, dtype=g_cos.dtype, device=g_cos.device) Delta.index_add_(0, edge_dst, h_global) @@ -539,7 +549,8 @@ def _ef_forward_triton(A_emb, edge_i, edge_j, D_block, n_ch, n_ang): cos_col, cos_ok, sin_col, sin_ok, A_cos, A_sin, C, R, S, P, - SP=_next_pow2(S), P16=_next_pow2(P), BLOCK_R=block_r) + SP=_next_pow2(S), P16=_next_pow2(P), BLOCK_R=block_r, + num_warps=_EF_WARPS) return A_cos, A_sin @@ -558,7 +569,8 @@ def _ef_backward_triton(dA_cos, dA_sin, A_emb, edge_i, edge_j, D_block, A_emb_c = A_emb.contiguous() args = (cos_col, cos_ok, sin_col, sin_ok) block_r = _ef_block_r(R) - kw = dict(SP=_next_pow2(S), P16=_next_pow2(P), BLOCK_R=block_r) + kw = dict(SP=_next_pow2(S), P16=_next_pow2(P), BLOCK_R=block_r, + num_warps=_EF_WARPS) def _scatter(dA_both): # scatter back to atoms (torch: well-optimized, visible to compile) @@ -838,7 +850,8 @@ def backward(ctx, dDelta): g_cos.contiguous(), g_sin.contiguous(), cos_col, cos_ok, sin_col, sin_ok, dD, n_base, n_base, n_sph, P, - SP=_next_pow2(n_sph), P16=_next_pow2(P), BLOCK_R=block_r) + SP=_next_pow2(n_sph), P16=_next_pow2(P), BLOCK_R=block_r, + num_warps=_EF_WARPS) return dg_cos, dg_sin, None, dD, None, None, None, None, None # Eager (double-differentiable) path. @@ -893,7 +906,8 @@ def forward(ctx, m_cos, m_sin, D_block, D_block.contiguous(), cos_col, cos_ok, sin_col, sin_ok, h_global, n_base, n_base, S, P, - SP=_next_pow2(S), P16=_next_pow2(P), BLOCK_R=block_r) + SP=_next_pow2(S), P16=_next_pow2(P), BLOCK_R=block_r, + num_warps=_EF_WARPS) else: h = _pack_grads(m_cos, m_sin, cos_flat_idx, sin_flat_idx, cos_valid, sin_valid, @@ -919,7 +933,8 @@ def backward(ctx, dh): P = (n_ch // n_base) * n_ang block_r = _ef_block_r(n_base) tabs = _ef_tables(S, n_ang, dh.device) - kw = dict(SP=_next_pow2(S), P16=_next_pow2(P), BLOCK_R=block_r) + kw = dict(SP=_next_pow2(S), P16=_next_pow2(P), BLOCK_R=block_r, + num_warps=_EF_WARPS) # Merged path: dh (the dominant read) loaded once → dm AND dD. if need_dm and need_dd and block_r >= n_base: From acbd1367a38b2c79b2f9b5431f0d3b8eeb224174 Mon Sep 17 00:00:00 2001 From: alacour Date: Thu, 3 Sep 2026 19:33:51 -0700 Subject: [PATCH 07/12] edge_frame_kernel: packed-D merged-backward variants (ECENET_EF_PACKD) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 Claude-Session: https://claude.ai/code/session_019WyADHSVvKbQnYo1PAjGPK --- ecenet/edge_frame_kernel.py | 151 ++++++++++++++++++++++++++++++++ tests/test_edge_frame_kernel.py | 20 +++++ 2 files changed, 171 insertions(+) diff --git a/ecenet/edge_frame_kernel.py b/ecenet/edge_frame_kernel.py index 979d7e1..5967d54 100644 --- a/ecenet/edge_frame_kernel.py +++ b/ecenet/edge_frame_kernel.py @@ -348,6 +348,109 @@ def _pu_bwd_merged_kernel(dh_ptr, mc_ptr, ms_ptr, d_ptr, tl.store(dd_base + sin_col[None, :], acc_s, mask=k_ok[:, None] & (sin_ok[None, :] > 0)) + # ── Packed-D variants (ECENET_EF_PACKD=1) ──────────────────────────────── + # ncu on the merged backward kernels (A100, 44k edges): L1/TEX throughput + # 75-86%, DRAM 9-18%, compute ~30% — L1-bound, with MIO short-scoreboard + # stalls. The scalar per-element gathers of D through the column tables + # (and the matching dD scatter-stores) are the amplification: hundreds of + # non-vectorizable L1 transactions per program while DRAM idles. These + # variants trade that for DRAM headroom: the wrapper pre-packs D into + # dense (E, S, P) / (E, P, S) tensors once per call (plain torch indexing) + # and unpacks dD afterwards, so the kernel does only vectorized coalesced + # tile IO. Math and masks identical to the table-gather kernels. + + @triton.jit + def _ef_bwd_merged_packed_kernel(dc_ptr, ds_ptr, a_ptr, ei_ptr, ej_ptr, + dct_ptr, dst_ptr, cosk_ptr, sink_ptr, + dab_ptr, ddc_ptr, dds_ptr, + C, R, S, P, + SP: tl.constexpr, P16: tl.constexpr, + BLOCK_R: tl.constexpr): + e = tl.program_id(0) + offs_r = tl.arange(0, BLOCK_R) + offs_k = tl.arange(0, SP) + offs_p = tl.arange(0, P16) + row_ok = offs_r < R + p_ok = offs_p < P + k_ok = offs_k < S + + cos_ok = tl.load(cosk_ptr + offs_p, mask=p_ok, other=0) + sin_ok = tl.load(sink_ptr + offs_p, mask=p_ok, other=0) + + g_off = e * (R * P) + offs_r[:, None] * P + offs_p[None, :] + dC = tl.load(dc_ptr + g_off, + mask=row_ok[:, None] & (cos_ok[None, :] > 0), other=0.0) + dS = tl.load(ds_ptr + g_off, + mask=row_ok[:, None] & (sin_ok[None, :] > 0), other=0.0) + + # dA_both = dC @ DcT + dS @ DsT — DxT pre-packed (E, P, S): plain tiles + t_off = e * (P * S) + offs_p[:, None] * S + offs_k[None, :] + t_mask = p_ok[:, None] & k_ok[None, :] + DcT = tl.load(dct_ptr + t_off, mask=t_mask, other=0.0) + DsT = tl.load(dst_ptr + t_off, mask=t_mask, other=0.0) + dA = tl.dot(dC, DcT, input_precision="ieee") \ + + tl.dot(dS, DsT, input_precision="ieee") + dab_off = e * (R * S) + offs_r[:, None] * S + offs_k[None, :] + tl.store(dab_ptr + dab_off, dA, mask=row_ok[:, None] & k_ok[None, :]) + + # dD = Aᵀ @ grads, stored packed (E, S, P); wrapper unpacks to columns + ei = tl.load(ei_ptr + e) + ej = tl.load(ej_ptr + e) + atom = tl.where(offs_r < C, ei, ej) + ch = tl.where(offs_r < C, offs_r, offs_r - C) + a_ptrs = a_ptr + (atom * C + ch)[:, None] * S + offs_k[None, :] + A = tl.load(a_ptrs, mask=row_ok[:, None] & k_ok[None, :], other=0.0) + acc_c = tl.dot(tl.trans(A), dC, input_precision="ieee") # (SP, P16) + acc_s = tl.dot(tl.trans(A), dS, input_precision="ieee") + o_off = e * (S * P) + offs_k[:, None] * P + offs_p[None, :] + o_mask = k_ok[:, None] & p_ok[None, :] + tl.store(ddc_ptr + o_off, acc_c, mask=o_mask) + tl.store(dds_ptr + o_off, acc_s, mask=o_mask) + + @triton.jit + def _pu_bwd_merged_packed_kernel(dh_ptr, mc_ptr, ms_ptr, + dcp_ptr, dsp_ptr, cosk_ptr, sink_ptr, + dmc_ptr, dms_ptr, ddc_ptr, dds_ptr, + R, S, P, + SP: tl.constexpr, P16: tl.constexpr, + BLOCK_R: tl.constexpr): + e = tl.program_id(0) + offs_r = tl.arange(0, BLOCK_R) + offs_k = tl.arange(0, SP) + offs_p = tl.arange(0, P16) + row_ok = offs_r < R + p_ok = offs_p < P + k_ok = offs_k < S + + cos_ok = tl.load(cosk_ptr + offs_p, mask=p_ok, other=0) + sin_ok = tl.load(sink_ptr + offs_p, mask=p_ok, other=0) + + dh_off = e * (R * S) + offs_r[:, None] * S + offs_k[None, :] + dh = tl.load(dh_ptr + dh_off, + mask=row_ok[:, None] & k_ok[None, :], other=0.0) + + # dm = dh @ D-cols — Dx pre-packed (E, S, P): plain tiles + d_off = e * (S * P) + offs_k[:, None] * P + offs_p[None, :] + d_mask = k_ok[:, None] & p_ok[None, :] + Dc = tl.load(dcp_ptr + d_off, mask=d_mask, other=0.0) + Ds = tl.load(dsp_ptr + d_off, mask=d_mask, other=0.0) + dmc = tl.dot(dh, Dc, input_precision="ieee") # (BLOCK_R, P16) + dms = tl.dot(dh, Ds, input_precision="ieee") + out_off = e * (R * P) + offs_r[:, None] * P + offs_p[None, :] + st_mask = row_ok[:, None] & p_ok[None, :] + tl.store(dmc_ptr + out_off, dmc, mask=st_mask) + tl.store(dms_ptr + out_off, dms, mask=st_mask) + + # dD = dhᵀ @ h, stored packed (E, S, P); wrapper unpacks to columns + mC = tl.load(mc_ptr + out_off, + mask=row_ok[:, None] & (cos_ok[None, :] > 0), other=0.0) + mS = tl.load(ms_ptr + out_off, + mask=row_ok[:, None] & (sin_ok[None, :] > 0), other=0.0) + acc_c = tl.dot(tl.trans(dh), mC, input_precision="ieee") # (SP, P16) + acc_s = tl.dot(tl.trans(dh), mS, input_precision="ieee") + tl.store(ddc_ptr + d_off, acc_c, mask=d_mask) + tl.store(dds_ptr + d_off, acc_s, mask=d_mask) + @triton.jit def _e2n_fwd_kernel(gc_ptr, gs_ptr, d_ptr, perm_ptr, aptr_ptr, srcoff_ptr, okc_ptr, oks_ptr, @@ -399,6 +502,31 @@ def _e2n_fwd_kernel(gc_ptr, gs_ptr, d_ptr, perm_ptr, aptr_ptr, # Triton default of 4 may win; sweep alongside profile_step. _EF_WARPS = int(os.environ.get('ECENET_EF_WARPS', 4)) +# Packed-D variants of the merged backward kernels (see the kernel comment +# block): trades the scalar table-gathered D loads / dD scatter-stores +# (L1-bound per ncu) for dense pre-packed tensors and vectorized tile IO. +# Off by default until benchmarked; module-level so tests can toggle it. +_EF_PACKD = os.environ.get('ECENET_EF_PACKD', '0') == '1' + + +def _pack_D(D_block, cos_col, cos_ok, sin_col, sin_ok): + """Dense per-edge packed D: Dx[e, k, p] = D[e, k, col(p)]·ok(p), (E, S, P).""" + Dc = D_block[:, :, cos_col.long()] * cos_ok.to(D_block.dtype) + Ds = D_block[:, :, sin_col.long()] * sin_ok.to(D_block.dtype) + return Dc.contiguous(), Ds.contiguous() + + +def _unpack_dD(ddc, dds, cos_col, cos_ok, sin_col, sin_ok, S): + """Inverse of _pack_D for the gradient: scatter packed (E, S, P) columns + back into (E, S, S). cos and sin column sets are disjoint; slots no p + touches stay 0 (matches the zeros-init of the table-scatter path).""" + dD = ddc.new_zeros(ddc.shape[0], S, S) + vc = cos_ok.bool() + vs = sin_ok.bool() + dD[:, :, cos_col[vc].long()] = ddc[:, :, vc] + dD[:, :, sin_col[vs].long()] = dds[:, :, vs] + return dD + def _next_pow2(x: int) -> int: """Smallest power of two ≥ x, floored at 16. tl.arange REQUIRES a power @@ -587,6 +715,18 @@ def _scatter(dA_both): # to cover the edge (block_r ≥ R; true for n_base/2C ≤ 128). if need_dx and need_dd and block_r >= R: dA_both = torch.empty(E, R, S, dtype=dA_cos.dtype, device=dA_cos.device) + if _EF_PACKD: + Dc, Ds = _pack_D(D_block, *args) + DcT = Dc.transpose(1, 2).contiguous() + DsT = Ds.transpose(1, 2).contiguous() + ddc = torch.empty(E, S, P, dtype=dA_cos.dtype, device=dA_cos.device) + dds = torch.empty_like(ddc) + _ef_bwd_merged_packed_kernel[(E,)]( + dA_cos, dA_sin, A_emb_c, edge_i.contiguous(), + edge_j.contiguous(), DcT, DsT, args[1], args[3], + dA_both, ddc, dds, + C, R, S, P, **kw) + return _scatter(dA_both), _unpack_dD(ddc, dds, *args, S) dD = torch.zeros_like(D_block) # columns no p touches stay 0 _ef_bwd_merged_kernel[(E,)]( dA_cos, dA_sin, A_emb_c, edge_i.contiguous(), edge_j.contiguous(), @@ -941,6 +1081,17 @@ def backward(ctx, dh): dm_cos = torch.empty(E, n_ch, n_ang, dtype=dh.dtype, device=dh.device) dm_sin = torch.empty_like(dm_cos) + if _EF_PACKD: + Dc, Ds = _pack_D(D_block, *tabs) + ddc = torch.empty(E, S, P, dtype=dh.dtype, device=dh.device) + dds = torch.empty_like(ddc) + _pu_bwd_merged_packed_kernel[(E,)]( + dh_c, m_cos.contiguous(), m_sin.contiguous(), + Dc, Ds, tabs[1], tabs[3], + dm_cos, dm_sin, ddc, dds, + n_base, S, P, **kw) + return (dm_cos, dm_sin, _unpack_dD(ddc, dds, *tabs, S), + None, None, None, None) dD = torch.zeros_like(D_block) _pu_bwd_merged_kernel[(E,)]( dh_c, m_cos.contiguous(), m_sin.contiguous(), diff --git a/tests/test_edge_frame_kernel.py b/tests/test_edge_frame_kernel.py index b9a8b5a..d632c8f 100644 --- a/tests/test_edge_frame_kernel.py +++ b/tests/test_edge_frame_kernel.py @@ -441,6 +441,25 @@ def test_triton_paths(): print("test_triton_paths[pack_unrotate]: OK") +def test_triton_packed_d(): + """CUDA-only: the ECENET_EF_PACKD merged-backward variants (dense packed-D + tile IO instead of per-element table gathers) pass the same fp64 + comparisons as the default kernels — test_triton_paths rerun with the + module flag on covers both EdgeFrameFused and PackUnrotateFused merged + backwards.""" + if not torch.cuda.is_available(): + print("test_triton_packed_d: SKIP (no CUDA)") + return + import ecenet.edge_frame_kernel as efk + old = efk._EF_PACKD + efk._EF_PACKD = True + try: + test_triton_paths() + finally: + efk._EF_PACKD = old + print("test_triton_packed_d: OK (packed-D merged backward)") + + def test_model_integration(): """Full ECENet (no MP): energy and autograd forces identical flag on/off.""" import ecenet as _ecenet @@ -524,6 +543,7 @@ def run(): test_pack_unrotate_matches_mp_ops() test_pack_unrotate_gradchecks() test_triton_paths() + test_triton_packed_d() test_model_integration() test_mp_integration() print("\nAll tests passed.") From 31c2ff3b71fd301c4856c492bc8607eecc030d60 Mon Sep 17 00:00:00 2001 From: alacour Date: Thu, 3 Sep 2026 19:54:48 -0700 Subject: [PATCH 08/12] edge_frame_kernel: amortize D packing once per step via the shared D_block MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 Claude-Session: https://claude.ai/code/session_019WyADHSVvKbQnYo1PAjGPK --- ecenet/edge_frame_kernel.py | 66 ++++++++++++++++++++++++++++++------- 1 file changed, 54 insertions(+), 12 deletions(-) diff --git a/ecenet/edge_frame_kernel.py b/ecenet/edge_frame_kernel.py index 5967d54..91aee5d 100644 --- a/ecenet/edge_frame_kernel.py +++ b/ecenet/edge_frame_kernel.py @@ -528,6 +528,30 @@ def _unpack_dD(ddc, dds, cos_col, cos_ok, sin_col, sin_ok, S): return dD +def _get_packed_D(D_block, n_ang): + """Once-per-step packed D — (Dc, Ds, DcT, DsT), each (E, S, P)/(E, P, S). + + Cached on the D_block PYTHON OBJECT: the model builds D_block once per + step and hands the same object to every fused edge-frame op (4x + EdgeFrameFused + 3x PackUnrotateFused at n_mp=4), so the packing cost is + paid once, not per backward — per-call packing was measured to eat the + packed kernels' entire ~5 ms/step win. A new step's fresh D_block starts + with no attribute, so there is no cross-step staleness. The attribute + does NOT survive save_for_backward's re-wrapping, so Functions must carry + the packed tensors through ctx themselves.""" + packed = getattr(D_block, '_ecenet_packed', None) + if packed is None or packed[0].shape[-1] != _ef_tables( + D_block.shape[-1], n_ang, D_block.device)[0].shape[0]: + with torch.no_grad(): + tabs = _ef_tables(D_block.shape[-1], n_ang, D_block.device) + Dc, Ds = _pack_D(D_block, *tabs) + packed = (Dc, Ds, + Dc.transpose(1, 2).contiguous(), + Ds.transpose(1, 2).contiguous()) + D_block._ecenet_packed = packed + return packed + + def _next_pow2(x: int) -> int: """Smallest power of two ≥ x, floored at 16. tl.arange REQUIRES a power of two (a multiple of 16 like 48 or 96 compiles to "arange's range must @@ -683,7 +707,7 @@ def _ef_forward_triton(A_emb, edge_i, edge_j, D_block, n_ch, n_ang): def _ef_backward_triton(dA_cos, dA_sin, A_emb, edge_i, edge_j, D_block, - n_ang, single, need_dx, need_dd): + n_ang, single, need_dx, need_dd, packedT=None): E = edge_i.shape[0] C = A_emb.shape[1] R = C if single else 2 * C @@ -715,10 +739,8 @@ def _scatter(dA_both): # to cover the edge (block_r ≥ R; true for n_base/2C ≤ 128). if need_dx and need_dd and block_r >= R: dA_both = torch.empty(E, R, S, dtype=dA_cos.dtype, device=dA_cos.device) - if _EF_PACKD: - Dc, Ds = _pack_D(D_block, *args) - DcT = Dc.transpose(1, 2).contiguous() - DsT = Ds.transpose(1, 2).contiguous() + if packedT is not None: + DcT, DsT = packedT ddc = torch.empty(E, S, P, dtype=dA_cos.dtype, device=dA_cos.device) dds = torch.empty_like(ddc) _ef_bwd_merged_packed_kernel[(E,)]( @@ -778,17 +800,28 @@ def forward(ctx, A_emb, edge_i, edge_j, D_block, .view(-1, n_ch, n_ang)) * cos_valid A_sin = (A_flat.index_select(1, sin_flat_idx) .view(-1, n_ch, n_ang)) * sin_valid + # Packed-D path: fetch the per-step packed D (amortized on the shared + # D_block object) in the FORWARD and carry it through ctx — attributes + # do not survive save_for_backward's re-wrapping. + extra = () + ctx.packed = _EF_PACKD and _ef_triton_ok(A_emb, edge_i.shape[0]) + if ctx.packed: + _, _, DcT, DsT = _get_packed_D(D_block, n_ang) + extra = (DcT, DsT) ctx.save_for_backward(A_emb, edge_i, edge_i if single else edge_j, D_block, - cos_flat_idx, sin_flat_idx, cos_valid, sin_valid) + cos_flat_idx, sin_flat_idx, cos_valid, sin_valid, + *extra) ctx.n_ang = n_ang ctx.single = single return A_cos, A_sin @staticmethod def backward(ctx, dA_cos, dA_sin): + saved = ctx.saved_tensors (A_emb, edge_i, edge_j, D_block, - cos_flat_idx, sin_flat_idx, cos_valid, sin_valid) = ctx.saved_tensors + cos_flat_idx, sin_flat_idx, cos_valid, sin_valid) = saved[:8] + packedT = saved[8:10] if ctx.packed else None E = dA_cos.shape[0] C = A_emb.shape[1] single = ctx.single @@ -803,7 +836,8 @@ def backward(ctx, dA_cos, dA_sin): dA_cos, dA_sin, A_emb, edge_i, edge_j, D_block, ctx.n_ang, single=single, need_dx=ctx.needs_input_grad[0], - need_dd=ctx.needs_input_grad[3]) + need_dd=ctx.needs_input_grad[3], + packedT=packedT) return dA_emb, None, None, dD, None, None, None, None # Adjoint of select: scatter masked grads into the rotated layout. @@ -1053,14 +1087,22 @@ def forward(ctx, m_cos, m_sin, D_block, cos_valid, sin_valid, n_base * S).view(E, n_base, S) h_global = torch.bmm(h, D_block.transpose(-1, -2)) + extra = () + ctx.packed = _EF_PACKD and _ef_triton_ok(m_cos, E) + if ctx.packed: + Dc, Ds, _, _ = _get_packed_D(D_block, n_ang) + extra = (Dc, Ds) ctx.save_for_backward(m_cos, m_sin, D_block, - cos_flat_idx, sin_flat_idx, cos_valid, sin_valid) + cos_flat_idx, sin_flat_idx, cos_valid, sin_valid, + *extra) return h_global @staticmethod def backward(ctx, dh): + saved = ctx.saved_tensors (m_cos, m_sin, D_block, - cos_flat_idx, sin_flat_idx, cos_valid, sin_valid) = ctx.saved_tensors + cos_flat_idx, sin_flat_idx, cos_valid, sin_valid) = saved[:7] + packed = saved[7:9] if ctx.packed else None E = m_cos.shape[0] n_ch, n_ang = cos_valid.shape S = D_block.shape[-1] @@ -1081,8 +1123,8 @@ def backward(ctx, dh): dm_cos = torch.empty(E, n_ch, n_ang, dtype=dh.dtype, device=dh.device) dm_sin = torch.empty_like(dm_cos) - if _EF_PACKD: - Dc, Ds = _pack_D(D_block, *tabs) + if packed is not None: + Dc, Ds = packed ddc = torch.empty(E, S, P, dtype=dh.dtype, device=dh.device) dds = torch.empty_like(ddc) _pu_bwd_merged_packed_kernel[(E,)]( From d06578017c7ed8c9b6e454bbb0eea0abec4082bd Mon Sep 17 00:00:00 2001 From: alacour Date: Thu, 3 Sep 2026 20:02:33 -0700 Subject: [PATCH 09/12] =?UTF-8?q?edge=5Fframe=5Fkernel:=20sync-free=20dD?= =?UTF-8?q?=20unpack=20=E2=80=94=20cached=20integer=20indices?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 Claude-Session: https://claude.ai/code/session_019WyADHSVvKbQnYo1PAjGPK --- ecenet/edge_frame_kernel.py | 34 ++++++++++++++++++++++++++-------- 1 file changed, 26 insertions(+), 8 deletions(-) diff --git a/ecenet/edge_frame_kernel.py b/ecenet/edge_frame_kernel.py index 91aee5d..24791a7 100644 --- a/ecenet/edge_frame_kernel.py +++ b/ecenet/edge_frame_kernel.py @@ -516,15 +516,33 @@ def _pack_D(D_block, cos_col, cos_ok, sin_col, sin_ok): return Dc.contiguous(), Ds.contiguous() -def _unpack_dD(ddc, dds, cos_col, cos_ok, sin_col, sin_ok, S): +_UNPACK_IDX: dict = {} + + +def _unpack_indices(S, n_ang, device): + """Static integer index pairs for _unpack_dD, cached per (S, n_ang, + device). MUST be integer tensors: boolean-mask indexing at call time + forces a host-device sync per backward (measured ~10 ms of CPU stall per + EdgeFrameFusedBackward — it erased the packed kernels' entire win).""" + key = (S, n_ang, str(device)) + if key not in _UNPACK_IDX: + cos_col, cos_ok, sin_col, sin_ok = _ef_tables(S, n_ang, device) + vc = cos_ok.bool() + vs = sin_ok.bool() + _UNPACK_IDX[key] = (vc.nonzero().flatten(), cos_col[vc].long(), + vs.nonzero().flatten(), sin_col[vs].long()) + return _UNPACK_IDX[key] + + +def _unpack_dD(ddc, dds, S, n_ang): """Inverse of _pack_D for the gradient: scatter packed (E, S, P) columns back into (E, S, S). cos and sin column sets are disjoint; slots no p - touches stay 0 (matches the zeros-init of the table-scatter path).""" + touches stay 0 (matches the zeros-init of the table-scatter path). + Integer-index ops only — fully async, no host sync.""" + pc, cc, ps, sc = _unpack_indices(S, n_ang, ddc.device) dD = ddc.new_zeros(ddc.shape[0], S, S) - vc = cos_ok.bool() - vs = sin_ok.bool() - dD[:, :, cos_col[vc].long()] = ddc[:, :, vc] - dD[:, :, sin_col[vs].long()] = dds[:, :, vs] + dD[:, :, cc] = ddc[:, :, pc] + dD[:, :, sc] = dds[:, :, ps] return dD @@ -748,7 +766,7 @@ def _scatter(dA_both): edge_j.contiguous(), DcT, DsT, args[1], args[3], dA_both, ddc, dds, C, R, S, P, **kw) - return _scatter(dA_both), _unpack_dD(ddc, dds, *args, S) + return _scatter(dA_both), _unpack_dD(ddc, dds, S, n_ang) dD = torch.zeros_like(D_block) # columns no p touches stay 0 _ef_bwd_merged_kernel[(E,)]( dA_cos, dA_sin, A_emb_c, edge_i.contiguous(), edge_j.contiguous(), @@ -1132,7 +1150,7 @@ def backward(ctx, dh): Dc, Ds, tabs[1], tabs[3], dm_cos, dm_sin, ddc, dds, n_base, S, P, **kw) - return (dm_cos, dm_sin, _unpack_dD(ddc, dds, *tabs, S), + return (dm_cos, dm_sin, _unpack_dD(ddc, dds, S, n_ang), None, None, None, None) dD = torch.zeros_like(D_block) _pu_bwd_merged_kernel[(E,)]( From 8c61f76d8b321e95450adac1d56ad5cb1473225e Mon Sep 17 00:00:00 2001 From: alacour Date: Thu, 3 Sep 2026 20:06:50 -0700 Subject: [PATCH 10/12] edge_frame_kernel: packed-D forward kernels; flip ECENET_EF_PACKD default 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 Claude-Session: https://claude.ai/code/session_019WyADHSVvKbQnYo1PAjGPK --- ecenet/edge_frame_kernel.py | 121 ++++++++++++++++++++++++++++---- tests/test_edge_frame_kernel.py | 15 ++-- 2 files changed, 113 insertions(+), 23 deletions(-) diff --git a/ecenet/edge_frame_kernel.py b/ecenet/edge_frame_kernel.py index 24791a7..0b80432 100644 --- a/ecenet/edge_frame_kernel.py +++ b/ecenet/edge_frame_kernel.py @@ -451,6 +451,74 @@ def _pu_bwd_merged_packed_kernel(dh_ptr, mc_ptr, ms_ptr, tl.store(ddc_ptr + d_off, acc_c, mask=d_mask) tl.store(dds_ptr + d_off, acc_s, mask=d_mask) + @triton.jit + def _ef_fwd_packed_kernel(a_ptr, ei_ptr, ej_ptr, dcp_ptr, dsp_ptr, + outc_ptr, outs_ptr, + C, R, S, P, + SP: tl.constexpr, P16: tl.constexpr, + BLOCK_R: tl.constexpr): + # _ef_fwd_kernel with the column-gathered D loads replaced by tiles of + # the pre-packed (E, S, P) Dc/Ds (ok-zeros baked in by _pack_D). + e = tl.program_id(0) + offs_r = tl.program_id(1) * BLOCK_R + tl.arange(0, BLOCK_R) + offs_k = tl.arange(0, SP) + offs_p = tl.arange(0, P16) + row_ok = offs_r < R + + ei = tl.load(ei_ptr + e) + ej = tl.load(ej_ptr + e) + atom = tl.where(offs_r < C, ei, ej) + ch = tl.where(offs_r < C, offs_r, offs_r - C) + a_ptrs = a_ptr + (atom * C + ch)[:, None] * S + offs_k[None, :] + A = tl.load(a_ptrs, mask=row_ok[:, None] & (offs_k[None, :] < S), other=0.0) + + d_off = e * (S * P) + offs_k[:, None] * P + offs_p[None, :] + d_mask = (offs_k[:, None] < S) & (offs_p[None, :] < P) + Dc = tl.load(dcp_ptr + d_off, mask=d_mask, other=0.0) + Ds = tl.load(dsp_ptr + d_off, mask=d_mask, other=0.0) + + OC = tl.dot(A, Dc, input_precision="ieee") # (BLOCK_R, P16) + OS = tl.dot(A, Ds, input_precision="ieee") + out_off = e * (R * P) + offs_r[:, None] * P + offs_p[None, :] + st_mask = row_ok[:, None] & (offs_p[None, :] < P) + tl.store(outc_ptr + out_off, OC, mask=st_mask) + tl.store(outs_ptr + out_off, OS, mask=st_mask) + + @triton.jit + def _ef_bwd_dx_packed_kernel(dc_ptr, ds_ptr, dct_ptr, dst_ptr, + cosk_ptr, sink_ptr, dab_ptr, + C, R, S, P, + SP: tl.constexpr, P16: tl.constexpr, + BLOCK_R: tl.constexpr): + # _ef_bwd_dx_kernel on the pre-packed (E, P, S) transposed D. The + # ok-masks on the grad loads (the eager path's ·valid) stay. + e = tl.program_id(0) + offs_r = tl.program_id(1) * BLOCK_R + tl.arange(0, BLOCK_R) + offs_k = tl.arange(0, SP) + offs_p = tl.arange(0, P16) + row_ok = offs_r < R + p_ok = offs_p < P + k_ok = offs_k < S + + cos_ok = tl.load(cosk_ptr + offs_p, mask=p_ok, other=0) + sin_ok = tl.load(sink_ptr + offs_p, mask=p_ok, other=0) + + g_off = e * (R * P) + offs_r[:, None] * P + offs_p[None, :] + dC = tl.load(dc_ptr + g_off, + mask=row_ok[:, None] & (cos_ok[None, :] > 0), other=0.0) + dS = tl.load(ds_ptr + g_off, + mask=row_ok[:, None] & (sin_ok[None, :] > 0), other=0.0) + + t_off = e * (P * S) + offs_p[:, None] * S + offs_k[None, :] + t_mask = p_ok[:, None] & k_ok[None, :] + DcT = tl.load(dct_ptr + t_off, mask=t_mask, other=0.0) + DsT = tl.load(dst_ptr + t_off, mask=t_mask, other=0.0) + + dA = tl.dot(dC, DcT, input_precision="ieee") \ + + tl.dot(dS, DsT, input_precision="ieee") # (BLOCK_R, SP) + dab_off = e * (R * S) + offs_r[:, None] * S + offs_k[None, :] + tl.store(dab_ptr + dab_off, dA, mask=row_ok[:, None] & k_ok[None, :]) + @triton.jit def _e2n_fwd_kernel(gc_ptr, gs_ptr, d_ptr, perm_ptr, aptr_ptr, srcoff_ptr, okc_ptr, oks_ptr, @@ -502,11 +570,14 @@ def _e2n_fwd_kernel(gc_ptr, gs_ptr, d_ptr, perm_ptr, aptr_ptr, # Triton default of 4 may win; sweep alongside profile_step. _EF_WARPS = int(os.environ.get('ECENET_EF_WARPS', 4)) -# Packed-D variants of the merged backward kernels (see the kernel comment -# block): trades the scalar table-gathered D loads / dD scatter-stores -# (L1-bound per ncu) for dense pre-packed tensors and vectorized tile IO. -# Off by default until benchmarked; module-level so tests can toggle it. -_EF_PACKD = os.environ.get('ECENET_EF_PACKD', '0') == '1' +# Packed-D kernel variants (see the kernel comment block): trade the scalar +# table-gathered D loads / dD scatter-stores (L1-bound per ncu, 75-86% L1 +# with DRAM at 9-18%) for dense pre-packed tensors and vectorized tile IO, +# packed once per step on the shared D_block object. Measured on the A100 +# (512-atom box): merged backwards 2.33→1.35 / 1.98→1.65 ms per call, joint +# force step 111.4→107.8 ms. Default ON; ECENET_EF_PACKD=0 restores the +# table-gather kernels. Module-level so tests can toggle it. +_EF_PACKD = os.environ.get('ECENET_EF_PACKD', '1') == '1' def _pack_D(D_block, cos_col, cos_ok, sin_col, sin_ok): @@ -713,6 +784,15 @@ def _ef_forward_triton(A_emb, edge_i, edge_j, D_block, n_ch, n_ang): A_sin = torch.empty_like(A_cos) block_r = _ef_block_r(R) grid = (E, triton.cdiv(R, block_r)) + if _EF_PACKD: + Dc, Ds, _, _ = _get_packed_D(D_block, n_ang) + _ef_fwd_packed_kernel[grid]( + A_emb.contiguous(), edge_i.contiguous(), ej.contiguous(), + Dc, Ds, A_cos, A_sin, + C, R, S, P, + SP=_next_pow2(S), P16=_next_pow2(P), BLOCK_R=block_r, + num_warps=_EF_WARPS) + return A_cos, A_sin _ef_fwd_kernel[grid]( A_emb.contiguous(), edge_i.contiguous(), ej.contiguous(), D_block.contiguous(), @@ -1090,16 +1170,27 @@ def forward(ctx, m_cos, m_sin, D_block, device=m_cos.device) block_r = _ef_block_r(n_base) grid = (E, triton.cdiv(n_base, block_r)) - _ef_bwd_dx_kernel[grid]( - m_cos.contiguous(), m_sin.contiguous(), - # forward wants h @ Dᵀ where the dx kernel computes - # Σ_p g[·,p]·D[k, col(p)] — i.e. it contracts against D's - # COLUMNS, which is exactly the transpose we need. - D_block.contiguous(), - cos_col, cos_ok, sin_col, sin_ok, h_global, - n_base, n_base, S, P, - SP=_next_pow2(S), P16=_next_pow2(P), BLOCK_R=block_r, - num_warps=_EF_WARPS) + if _EF_PACKD: + # forward wants h @ Dᵀ; the packed dx kernel contracts + # against DcT/DsT (E, P, S), exactly that transpose. + _, _, DcT, DsT = _get_packed_D(D_block, n_ang) + _ef_bwd_dx_packed_kernel[grid]( + m_cos.contiguous(), m_sin.contiguous(), + DcT, DsT, cos_ok, sin_ok, h_global, + n_base, n_base, S, P, + SP=_next_pow2(S), P16=_next_pow2(P), BLOCK_R=block_r, + num_warps=_EF_WARPS) + else: + _ef_bwd_dx_kernel[grid]( + m_cos.contiguous(), m_sin.contiguous(), + # forward wants h @ Dᵀ where the dx kernel computes + # Σ_p g[·,p]·D[k, col(p)] — i.e. it contracts against + # D's COLUMNS, which is exactly the transpose we need. + D_block.contiguous(), + cos_col, cos_ok, sin_col, sin_ok, h_global, + n_base, n_base, S, P, + SP=_next_pow2(S), P16=_next_pow2(P), BLOCK_R=block_r, + num_warps=_EF_WARPS) else: h = _pack_grads(m_cos, m_sin, cos_flat_idx, sin_flat_idx, cos_valid, sin_valid, diff --git a/tests/test_edge_frame_kernel.py b/tests/test_edge_frame_kernel.py index d632c8f..c06da6e 100644 --- a/tests/test_edge_frame_kernel.py +++ b/tests/test_edge_frame_kernel.py @@ -442,22 +442,21 @@ def test_triton_paths(): def test_triton_packed_d(): - """CUDA-only: the ECENET_EF_PACKD merged-backward variants (dense packed-D - tile IO instead of per-element table gathers) pass the same fp64 - comparisons as the default kernels — test_triton_paths rerun with the - module flag on covers both EdgeFrameFused and PackUnrotateFused merged - backwards.""" + """CUDA-only: run test_triton_paths under BOTH _EF_PACKD settings, so the + packed-D kernels (the default) and the table-gather kernels (the + ECENET_EF_PACKD=0 fallback) both stay pinned to the fp64 reference.""" if not torch.cuda.is_available(): print("test_triton_packed_d: SKIP (no CUDA)") return import ecenet.edge_frame_kernel as efk old = efk._EF_PACKD - efk._EF_PACKD = True try: - test_triton_paths() + for packed in (True, False): + efk._EF_PACKD = packed + test_triton_paths() + print(f"test_triton_packed_d[packed={packed}]: OK") finally: efk._EF_PACKD = old - print("test_triton_packed_d: OK (packed-D merged backward)") def test_model_integration(): From 852d6cf5039e2ff9ec0d149fbb94a19927255e5d Mon Sep 17 00:00:00 2001 From: alacour Date: Thu, 3 Sep 2026 20:15:49 -0700 Subject: [PATCH 11/12] =?UTF-8?q?train=5Fecenet:=20tf32=20switch=20?= =?UTF-8?q?=E2=80=94=20parity=20with=20the=20spice/xyz=20trainers?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 Claude-Session: https://claude.ai/code/session_019WyADHSVvKbQnYo1PAjGPK --- scripts/train_ecenet.py | 20 ++++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/scripts/train_ecenet.py b/scripts/train_ecenet.py index 9156cbc..a28a136 100644 --- a/scripts/train_ecenet.py +++ b/scripts/train_ecenet.py @@ -214,6 +214,7 @@ def train_ecenet( eval_batch_size=32, seed=42, dtype=torch.float64, + tf32=False, # route float32 matmuls to TF32 tensor cores (Ampere+) device=None, checkpoint_path=None, reset_optimizer=False, @@ -248,6 +249,25 @@ def loss_fn(pred, tgt): elif isinstance(device, str): device = torch.device(device) + # TF32: on Ampere+ the step is dominated by fp32 matmuls (forward and the + # force double-backward). Routing those to TF32 tensor cores is the cheapest + # large speedup, but TF32 keeps only ~10 mantissa bits — A/B the val + # force/energy MAE before trusting it. No effect under float64 (TF32 is a + # float32-only mode), so warn rather than silently do nothing. (Same + # switch as train_ecenet_spice / train_ecenet_xyz.) + if tf32: + if dtype == torch.float64: + if verbose: + print_flush(" [tf32] requested but dtype=float64 → no effect " + "(TF32 is float32-only); use dtype=torch.float32") + else: + torch.backends.cuda.matmul.allow_tf32 = True + torch.backends.cudnn.allow_tf32 = True + torch.set_float32_matmul_precision('high') + if verbose: + print_flush(" [tf32] enabled: float32 matmuls → TF32 tensor cores " + "(A/B the val MAE against a tf32=False run)") + # ── Data ────────────────────────────────────────────────────────────── (pos_train, frc_train, eng_train, pos_val, frc_val, eng_val, From 8bff38897f0b40643a067cf671388f4ee7b29ed3 Mon Sep 17 00:00:00 2001 From: alacour Date: Thu, 3 Sep 2026 20:24:31 -0700 Subject: [PATCH 12/12] edge_frame_kernel: refresh docstrings for the packed-D default Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_019WyADHSVvKbQnYo1PAjGPK --- ecenet/edge_frame_kernel.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/ecenet/edge_frame_kernel.py b/ecenet/edge_frame_kernel.py index 0b80432..f935bc6 100644 --- a/ecenet/edge_frame_kernel.py +++ b/ecenet/edge_frame_kernel.py @@ -873,7 +873,8 @@ def _scatter(dA_both): class EdgeFrameFused(torch.autograd.Function): """Fused gather → rotate → select with analytic, double-differentiable - backward. Saves (A_emb, D_block, indices) — NOT the (E, R, n_sph) + backward. Saves (A_emb, D_block, indices; plus the per-step packed D + under _EF_PACKD) — NOT the (E, R, n_sph) intermediates; the gathered rows are re-gathered in the backward. edge_j=None → single-source mode (MP steps 5-6): rows are A_emb[edge_i]'s @@ -1151,9 +1152,11 @@ class PackUnrotateFused(torch.autograd.Function): No gather, no scatter, no gate — the attention/message weighting and the node accumulation stay eager in the node frame exactly as the unfused path. The win is dropping the packed intermediate h from HBM (forward) and - from the bmm's saved set (backward). All three kernels are the existing - ones: forward = _ef_bwd_dx (its "grads" input is any per-edge (E, R, P) - tensor), backward dm = _ef_fwd with identity indices, dD = _ef_bwd_dd.""" + from the bmm's saved set (backward). The kernels are shared with the edge + frame: forward = _ef_bwd_dx (its "grads" input is any per-edge (E, R, P) + tensor) — or its packed-D variant under _EF_PACKD (default) — backward = + the merged _pu kernel (packed or table-gather), with dm = _ef_fwd with + identity indices / dD = _ef_bwd_dd as the non-merged fallbacks.""" @staticmethod def forward(ctx, m_cos, m_sin, D_block,