From 6caf67876a269214ba80b9c66f3d512e86c64908 Mon Sep 17 00:00:00 2001 From: alacour Date: Thu, 3 Sep 2026 16:41:50 -0700 Subject: [PATCH 1/5] realspace kernel: coalesced tile IO + tunable launch; profile_step: time fused paths MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The RealSpace Triton kernels loaded/stored each angular column separately at stride n_ang — N_ANG transactions per operand walking DRAM at one-third sector utilization, which capped this bandwidth-bound op well below the A100's bandwidth (measured ~15% at 44k edges x 256 features). Both kernels now do one dense (BLOCK, N_ANG) coalesced tile load/store per operand and pull columns out of the register tile with a multiply-mask reduction (free next to the loads it replaces; no 3D intermediates, so no spill pressure). Launch config is overridable per GPU via ECENET_RS_BLOCK / ECENET_RS_WARPS for a quick sweep. Numerics unchanged: tests/test_realspace_kernel.py passes (CPU paths; run the Triton test on a GPU box). profile_step.py previously timed the *unfused reference* ops for the edge-frame rotate and the MP pack/unrotate/rotate-back/unpack even when the fused flags were on, so the sub-component lines neither summed to the real forward nor showed the fused kernels' cost. The reference lines are now labeled "(unfused ref)" when a fused path replaces them, and the fused ops (edge_frame_fused, pack_unrotate_fused, edge_frame_fused_single) get their own timing lines; nonlin lines are tagged "(fused)" when that path is engaged. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_019WyADHSVvKbQnYo1PAjGPK --- ecenet/realspace_kernel.py | 87 ++++++++++++++++++++++++++++---------- tools/profile_step.py | 51 +++++++++++++++++----- 2 files changed, 106 insertions(+), 32 deletions(-) diff --git a/ecenet/realspace_kernel.py b/ecenet/realspace_kernel.py index 09a0599..aa5f676 100644 --- a/ecenet/realspace_kernel.py +++ b/ecenet/realspace_kernel.py @@ -32,6 +32,8 @@ rms-norm / channel-mix variants) from silently taking the wrong path. """ +import os + import torch try: @@ -41,7 +43,11 @@ except ImportError: # CPU-only / no-triton env → PyTorch path _HAS_TRITON = False -_RS_BLOCK = 128 # rows (edge·feature) per program; tunable +# Launch configuration. Rows (edge·feature) per program and warps per program; +# overridable per GPU without a code change (read once at import): +# ECENET_RS_BLOCK=256 ECENET_RS_WARPS=8 python ... +_RS_BLOCK = int(os.environ.get('ECENET_RS_BLOCK', 128)) +_RS_WARPS = int(os.environ.get('ECENET_RS_WARPS', 4)) def is_fusible(nl): @@ -80,6 +86,16 @@ def realspace_reference(A_cos, A_sin, cos_synth, sin_synth, # store. The fat grid tensor never touches HBM, and the backward recomputes it (so # it's never saved) — that's the memory win. # +# Memory access: the op is bandwidth-bound (a handful of flops per loaded float), +# so all global traffic goes through ONE (BLOCK, N_ANG) tile load/store per +# operand — a dense row-major region, fully coalesced. The earlier per-m column +# loads walked DRAM at stride n_ang (a third of each sector wasted) and issued +# N_ANG separate transactions per operand; that access pattern, not arithmetic, +# capped the kernel well below bandwidth. Columns are then pulled out of the +# register tile with a multiply-mask reduction — arithmetically free next to the +# loads it replaces, and it keeps the live set to 2D tiles (no (BLOCK, N_ANG, +# N_GRID) intermediate to spill). +# # n_ang (3-4) and n_grid (9-13) are far below tl.dot's 16-min, so synthesis and # analysis are written as explicit outer-product accumulations / axis reductions # over the (compile-time) angular and grid dims — small 2D tiles only, no tl.dot, @@ -96,30 +112,42 @@ def _rs_fwd_kernel(acos_ptr, asin_ptr, cs_ptr, ss_ptr, ca_ptr, sa_ptr, BLOCK: tl.constexpr, N_ANG: tl.constexpr, N_GRID: tl.constexpr): offs = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) # row = edge·feature mask = offs < R + a_idx = tl.arange(0, N_ANG) offs_g = tl.arange(0, N_GRID) g_mask = offs_g < n_grid + # One coalesced (BLOCK, N_ANG) tile load per operand. + tile = offs[:, None] * s_row + a_idx[None, :] * s_col + t_mask = mask[:, None] & (a_idx < n_ang)[None, :] + ac_all = tl.load(acos_ptr + tile, mask=t_mask, other=0.0) + as_all = tl.load(asin_ptr + tile, mask=t_mask, other=0.0) + # Synthesis: f(grid) = Σ_m A_cos[:,m]·cos_synth[m,:] + A_sin[:,m]·sin_synth[m,:] + # (column m extracted from the register tile by multiply-mask reduction) f = tl.zeros((BLOCK, N_GRID), dtype=tl.float32) for m in tl.static_range(N_ANG): - col = mask & (m < n_ang) - ac = tl.load(acos_ptr + offs * s_row + m * s_col, mask=col, other=0.0) - as_ = tl.load(asin_ptr + offs * s_row + m * s_col, mask=col, other=0.0) + sel = (a_idx == m).to(tl.float32) + ac = tl.sum(ac_all * sel[None, :], axis=1) + as_ = tl.sum(as_all * sel[None, :], axis=1) cs = tl.load(cs_ptr + m * n_grid + offs_g, mask=g_mask & (m < n_ang), other=0.0) ss = tl.load(ss_ptr + m * n_grid + offs_g, mask=g_mask & (m < n_ang), other=0.0) f += ac[:, None] * cs[None, :] + as_[:, None] * ss[None, :] h = f * tl.sigmoid(f) # silu - # Analysis: out[:,a] = Σ_k h[:,k]·cos_analysis[k,a] (store each column directly) + # Analysis: out[:,a] = Σ_k h[:,k]·cos_analysis[k,a] — accumulated into + # (BLOCK, N_ANG) register tiles, one coalesced store per operand. + oc_all = tl.zeros((BLOCK, N_ANG), dtype=tl.float32) + os_all = tl.zeros((BLOCK, N_ANG), dtype=tl.float32) for a in tl.static_range(N_ANG): col = g_mask & (a < n_ang) ca = tl.load(ca_ptr + offs_g * n_ang + a, mask=col, other=0.0) sa = tl.load(sa_ptr + offs_g * n_ang + a, mask=col, other=0.0) - oc = tl.sum(h * ca[None, :], axis=1) - os = tl.sum(h * sa[None, :], axis=1) - tl.store(oc_ptr + offs * s_row + a * s_col, oc, mask=mask & (a < n_ang)) - tl.store(os_ptr + offs * s_row + a * s_col, os, mask=mask & (a < n_ang)) + sel = (a_idx == a).to(tl.float32) + oc_all += tl.sum(h * ca[None, :], axis=1)[:, None] * sel[None, :] + os_all += tl.sum(h * sa[None, :], axis=1)[:, None] * sel[None, :] + tl.store(oc_ptr + tile, oc_all, mask=t_mask) + tl.store(os_ptr + tile, os_all, mask=t_mask) @triton.jit def _rs_bwd_kernel(acos_ptr, asin_ptr, cs_ptr, ss_ptr, ca_ptr, sa_ptr, @@ -128,15 +156,24 @@ def _rs_bwd_kernel(acos_ptr, asin_ptr, cs_ptr, ss_ptr, ca_ptr, sa_ptr, BLOCK: tl.constexpr, N_ANG: tl.constexpr, N_GRID: tl.constexpr): offs = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) mask = offs < R + a_idx = tl.arange(0, N_ANG) offs_g = tl.arange(0, N_GRID) g_mask = offs_g < n_grid + # One coalesced (BLOCK, N_ANG) tile load per operand (see forward). + tile = offs[:, None] * s_row + a_idx[None, :] * s_col + t_mask = mask[:, None] & (a_idx < n_ang)[None, :] + ac_all = tl.load(acos_ptr + tile, mask=t_mask, other=0.0) + as_all = tl.load(asin_ptr + tile, mask=t_mask, other=0.0) + goc_all = tl.load(goc_ptr + tile, mask=t_mask, other=0.0) + gos_all = tl.load(gos_ptr + tile, mask=t_mask, other=0.0) + # Recompute f (never stored in the forward). f = tl.zeros((BLOCK, N_GRID), dtype=tl.float32) for m in tl.static_range(N_ANG): - col = mask & (m < n_ang) - ac = tl.load(acos_ptr + offs * s_row + m * s_col, mask=col, other=0.0) - as_ = tl.load(asin_ptr + offs * s_row + m * s_col, mask=col, other=0.0) + sel = (a_idx == m).to(tl.float32) + ac = tl.sum(ac_all * sel[None, :], axis=1) + as_ = tl.sum(as_all * sel[None, :], axis=1) cs = tl.load(cs_ptr + m * n_grid + offs_g, mask=g_mask & (m < n_ang), other=0.0) ss = tl.load(ss_ptr + m * n_grid + offs_g, mask=g_mask & (m < n_ang), other=0.0) f += ac[:, None] * cs[None, :] + as_[:, None] * ss[None, :] @@ -146,24 +183,28 @@ def _rs_bwd_kernel(acos_ptr, asin_ptr, cs_ptr, ss_ptr, ca_ptr, sa_ptr, # dh[:,k] = Σ_a g_out_cos[:,a]·cos_analysis[k,a] + g_out_sin[:,a]·sin_analysis[k,a] dh = tl.zeros((BLOCK, N_GRID), dtype=tl.float32) for a in tl.static_range(N_ANG): - col = mask & (a < n_ang) - goc = tl.load(goc_ptr + offs * s_row + a * s_col, mask=col, other=0.0) - gos = tl.load(gos_ptr + offs * s_row + a * s_col, mask=col, other=0.0) + sel = (a_idx == a).to(tl.float32) + goc = tl.sum(goc_all * sel[None, :], axis=1) + gos = tl.sum(gos_all * sel[None, :], axis=1) ca = tl.load(ca_ptr + offs_g * n_ang + a, mask=g_mask & (a < n_ang), other=0.0) sa = tl.load(sa_ptr + offs_g * n_ang + a, mask=g_mask & (a < n_ang), other=0.0) dh += goc[:, None] * ca[None, :] + gos[:, None] * sa[None, :] df = dh * silu_prime # through silu - # dA_cos[:,m] = Σ_k df[:,k]·cos_synth[m,k] + # dA_cos[:,m] = Σ_k df[:,k]·cos_synth[m,k] — accumulated into (BLOCK, + # N_ANG) register tiles, one coalesced store per operand. + dac_all = tl.zeros((BLOCK, N_ANG), dtype=tl.float32) + das_all = tl.zeros((BLOCK, N_ANG), dtype=tl.float32) for m in tl.static_range(N_ANG): col = g_mask & (m < n_ang) cs = tl.load(cs_ptr + m * n_grid + offs_g, mask=col, other=0.0) ss = tl.load(ss_ptr + m * n_grid + offs_g, mask=col, other=0.0) - dac = tl.sum(df * cs[None, :], axis=1) - das = tl.sum(df * ss[None, :], axis=1) - tl.store(dac_ptr + offs * s_row + m * s_col, dac, mask=mask & (m < n_ang)) - tl.store(das_ptr + offs * s_row + m * s_col, das, mask=mask & (m < n_ang)) + sel = (a_idx == m).to(tl.float32) + dac_all += tl.sum(df * cs[None, :], axis=1)[:, None] * sel[None, :] + das_all += tl.sum(df * ss[None, :], axis=1)[:, None] * sel[None, :] + tl.store(dac_ptr + tile, dac_all, mask=t_mask) + tl.store(das_ptr + tile, das_all, mask=t_mask) def _can_use_triton(A_cos, activation): @@ -189,7 +230,8 @@ def _realspace_forward_triton(A_cos, A_sin, cos_synth, sin_synth, grid = (triton.cdiv(R, _RS_BLOCK),) _rs_fwd_kernel[grid](acos, asin, cs, ss, ca, sa, oc, os, R, n_ang, n_grid, acos.stride(0), acos.stride(1), - BLOCK=_RS_BLOCK, N_ANG=N_ANG, N_GRID=N_GRID) + BLOCK=_RS_BLOCK, N_ANG=N_ANG, N_GRID=N_GRID, + num_warps=_RS_WARPS) out = lambda t: t.reshape(n_e, F, n_ang).to(A_cos.dtype) # noqa: E731 return out(oc), out(os) @@ -209,7 +251,8 @@ def _realspace_backward_triton(g_out_cos, g_out_sin, A_cos, A_sin, grid = (triton.cdiv(R, _RS_BLOCK),) _rs_bwd_kernel[grid](acos, asin, cs, ss, ca, sa, goc, gos, dac, das, R, n_ang, n_grid, acos.stride(0), acos.stride(1), - BLOCK=_RS_BLOCK, N_ANG=N_ANG, N_GRID=N_GRID) + BLOCK=_RS_BLOCK, N_ANG=N_ANG, N_GRID=N_GRID, + num_warps=_RS_WARPS) out = lambda t: t.reshape(n_e, F, n_ang).to(A_cos.dtype) # noqa: E731 return out(dac), out(das) diff --git a/tools/profile_step.py b/tools/profile_step.py index 7392e54..5596d87 100644 --- a/tools/profile_step.py +++ b/tools/profile_step.py @@ -165,23 +165,34 @@ def fresh_pos(): A_emb = model._embed(A, types) A_both = torch.cat([A_emb[edge_i], A_emb[edge_j]], dim=1) -# Wigner rotation breakdown +# Wigner rotation breakdown. When --edge_frame_fused is on, the real forward +# replaces the gather+cat, the bmm rotate, and sph_to_angular with one fused op +# — those lines are then labeled "(unfused ref)" and the fused op is timed too, +# so the sub-components reflect the path that actually runs. +_ef = ' (unfused ref)' if args.edge_frame_fused else '' _, D_list = time_fn(" Wigner D (recursive_wigner_D)", lambda: recursive_wigner_D(r_hat, model.l_max)) _, D_block_main = time_fn(" Wigner D (build_D_block_from_list)", lambda: build_D_block_from_list(D_list, len(r_hat), model.l_max, r_hat.device, r_hat.dtype)) -time_fn(" Wigner bmm (A_both @ D)", +time_fn(f" Wigner bmm (A_both @ D){_ef}", lambda: torch.bmm(A_both, D_block_main)) -time_fn("Wigner rotate total (main, w/ cached D)", +time_fn(f"Wigner rotate total (main, w/ cached D){_ef}", lambda: wigner_rotate(A_both, D_block_main)) A_rot = wigner_rotate(A_both, D_block_main) # sph_to_angular -time_fn("sph_to_angular (repeat_interleave + gather)", +time_fn(f"sph_to_angular (repeat_interleave + gather){_ef}", lambda: model.sph_to_angular(A_rot)) +if args.edge_frame_fused: + from ecenet.edge_frame_kernel import edge_frame_fused, edge_frame_fused_single, \ + pack_unrotate_fused + time_fn("edge-frame fused (gather+rotate+reshape)", + lambda: edge_frame_fused(A_emb, edge_i, edge_j, D_block_main, + model.sph_to_angular)) + A_cos, A_sin = model.sph_to_angular(A_rot) type_i = types[edge_i] @@ -200,12 +211,16 @@ def fresh_pos(): # Low-rank (bottleneck_dim) layers have linear_down/linear_up instead # of a single full-width linear; the nonlinearity runs at the # bottleneck width, between the two. + # The nonlin lines call the real module, so they already time the fused + # path when --fuse_nonlin is on — the suffix just says which path ran. + nl_tag = (' (fused)' if layer.use_nonlinearity + and getattr(layer.nonlin, 'fused', False) else '') if layer.bottleneck_dim is not None: time_fn(f" EquivariantLinear down [{gi},{li}]", lambda l=layer: l.linear_down(A_cos, A_sin)) lin_out = layer.linear_down(A_cos, A_sin) if layer.use_nonlinearity: - time_fn(f" RealSpaceNonlinearity [{gi},{li}]", + time_fn(f" RealSpaceNonlinearity [{gi},{li}]{nl_tag}", lambda l=layer, c=lin_out: l.nonlin(*c)) nl_out = layer.nonlin(*lin_out) if layer.use_nonlinearity else lin_out time_fn(f" EquivariantLinear up [{gi},{li}]", @@ -215,7 +230,7 @@ def fresh_pos(): lambda l=layer: l.linear(A_cos, A_sin)) lin_out = layer.linear(A_cos, A_sin) if layer.use_nonlinearity: - time_fn(f" RealSpaceNonlinearity [{gi},{li}]", + time_fn(f" RealSpaceNonlinearity [{gi},{li}]{nl_tag}", lambda l=layer, c=lin_out: l.nonlin(*c)) A_cos, A_sin = layer(A_cos, A_sin) @@ -236,12 +251,20 @@ def _trunk(): m_sin = u_sin[:, :mp.n_ch] + A_sin s = u_cos[:, mp.n_ch:mp.n_ch + mp.n_scores, 0] - _, h_packed = time_fn(" MP pack cos/sin → n_sph", + # With e2n fusion the real forward runs pack+unrotate as one fused op; + # the two reference lines are kept (labeled) for comparison and the + # fused op is timed as its own line. + _e2n = ' (unfused ref)' if getattr(mp, 'edge_frame_fused_e2n', False) else '' + _, h_packed = time_fn(f" MP pack cos/sin → n_sph{_e2n}", lambda: mp._pack(m_cos, m_sin)) D_block_main_T = D_block_main.transpose(-1, -2) - time_fn(" MP bmm unrotate (h @ D^T)", + time_fn(f" MP bmm unrotate (h @ D^T){_e2n}", lambda: torch.bmm(h_packed, D_block_main_T)) + if getattr(mp, 'edge_frame_fused_e2n', False): + time_fn(" MP pack+unrotate (fused e2n)", + lambda: pack_unrotate_fused(m_cos, m_sin, D_block_main, + mp.l_max, mp.m_max)) h_global = torch.bmm(h_packed, D_block_main_T) def _weights(): @@ -275,10 +298,18 @@ def _weights(): Delta = torch.zeros(n_atoms, H, hb, mp.n_sph, device=device, dtype=dtype ).scatter_add(0, idx, contrib).reshape(n_atoms, mp.n_base, mp.n_sph) - time_fn(" MP bmm rotate back (Delta @ D)", + # Same for the return leg: gather+rotate-back+unpack is one fused op in + # the real forward when edge_frame_fused is on. + _eff = ' (unfused ref)' if getattr(mp, 'edge_frame_fused', False) else '' + time_fn(f" MP bmm rotate back (Delta @ D){_eff}", lambda: torch.bmm(Delta[edge_i], D_block_main)) v_rot = torch.bmm(Delta[edge_i], D_block_main) - _, d = time_fn(" MP unpack n_sph → cos/sin", lambda: mp._unpack(v_rot, n_e)) + _, d = time_fn(f" MP unpack n_sph → cos/sin{_eff}", + lambda: mp._unpack(v_rot, n_e)) + if getattr(mp, 'edge_frame_fused', False): + time_fn(" MP gather+rotate+unpack (fused)", + lambda: edge_frame_fused_single(Delta, edge_i, D_block_main, + mp.l_max, mp.m_max)) time_fn(" MP receiver block", lambda: mp.receiver(*d)) A_cos, A_sin = mp(A_cos, A_sin, r_hat, dist_ij, edge_i, edge_j, From 085098fad0e3fcb217ca742b27b1133f54488634 Mon Sep 17 00:00:00 2001 From: alacour Date: Thu, 3 Sep 2026 17:51:21 -0700 Subject: [PATCH 2/5] =?UTF-8?q?realspace=20kernel:=20read=20einsum's=20m-m?= =?UTF-8?q?ajor=20layout=20in=20place=20=E2=80=94=20no=20contiguous=20copi?= =?UTF-8?q?es?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The launch-config sweep on the A100 showed a ~1.4 ms floor that BLOCK/WARPS could not move: the cost was outside the kernel. EquivariantLinear's einsum emits (n_e, F, n_ang) with strides (F, 1, n_e*F), so the wrapper's .contiguous() was paying a full read+write copy of both inputs every forward call (and of inputs plus both incoming grads in the backward) before the kernel launched. That layout is itself perfectly coalesced (rows are unit-stride), so the copies bought nothing. Each operand now carries its own (row, col) strides into the kernel, and the _rows() helper flattens (e, f) to rows without copying whenever stride(0) == F*stride(1) — covering both dense row-major and the einsum m-major layout, with a contiguous fallback for anything else. GPU test extended with an m-major case (forward and backward) against the contiguous result. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_019WyADHSVvKbQnYo1PAjGPK --- ecenet/realspace_kernel.py | 84 ++++++++++++++++++++++++---------- tests/test_realspace_kernel.py | 26 +++++++++++ 2 files changed, 85 insertions(+), 25 deletions(-) diff --git a/ecenet/realspace_kernel.py b/ecenet/realspace_kernel.py index aa5f676..384b82f 100644 --- a/ecenet/realspace_kernel.py +++ b/ecenet/realspace_kernel.py @@ -108,7 +108,8 @@ def realspace_reference(A_cos, A_sin, cos_synth, sin_synth, @triton.jit def _rs_fwd_kernel(acos_ptr, asin_ptr, cs_ptr, ss_ptr, ca_ptr, sa_ptr, oc_ptr, os_ptr, - R, n_ang, n_grid, s_row, s_col, + R, n_ang, n_grid, + sc_row, sc_col, ss_row, ss_col, so_row, so_col, BLOCK: tl.constexpr, N_ANG: tl.constexpr, N_GRID: tl.constexpr): offs = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) # row = edge·feature mask = offs < R @@ -116,11 +117,15 @@ def _rs_fwd_kernel(acos_ptr, asin_ptr, cs_ptr, ss_ptr, ca_ptr, sa_ptr, offs_g = tl.arange(0, N_GRID) g_mask = offs_g < n_grid - # One coalesced (BLOCK, N_ANG) tile load per operand. - tile = offs[:, None] * s_row + a_idx[None, :] * s_col + # One coalesced (BLOCK, N_ANG) tile load per operand. Each operand + # carries its own (row, col) strides so both supported layouts — dense + # row-major and einsum's m-major (rows unit-stride, see _rows) — read + # straight from where the producer left them, no .contiguous() copies. t_mask = mask[:, None] & (a_idx < n_ang)[None, :] - ac_all = tl.load(acos_ptr + tile, mask=t_mask, other=0.0) - as_all = tl.load(asin_ptr + tile, mask=t_mask, other=0.0) + ac_all = tl.load(acos_ptr + offs[:, None] * sc_row + a_idx[None, :] * sc_col, + mask=t_mask, other=0.0) + as_all = tl.load(asin_ptr + offs[:, None] * ss_row + a_idx[None, :] * ss_col, + mask=t_mask, other=0.0) # Synthesis: f(grid) = Σ_m A_cos[:,m]·cos_synth[m,:] + A_sin[:,m]·sin_synth[m,:] # (column m extracted from the register tile by multiply-mask reduction) @@ -146,13 +151,16 @@ def _rs_fwd_kernel(acos_ptr, asin_ptr, cs_ptr, ss_ptr, ca_ptr, sa_ptr, sel = (a_idx == a).to(tl.float32) oc_all += tl.sum(h * ca[None, :], axis=1)[:, None] * sel[None, :] os_all += tl.sum(h * sa[None, :], axis=1)[:, None] * sel[None, :] - tl.store(oc_ptr + tile, oc_all, mask=t_mask) - tl.store(os_ptr + tile, os_all, mask=t_mask) + tile_o = offs[:, None] * so_row + a_idx[None, :] * so_col + tl.store(oc_ptr + tile_o, oc_all, mask=t_mask) + tl.store(os_ptr + tile_o, os_all, mask=t_mask) @triton.jit def _rs_bwd_kernel(acos_ptr, asin_ptr, cs_ptr, ss_ptr, ca_ptr, sa_ptr, goc_ptr, gos_ptr, dac_ptr, das_ptr, - R, n_ang, n_grid, s_row, s_col, + R, n_ang, n_grid, + sc_row, sc_col, ss_row, ss_col, + gc_row, gc_col, gs_row, gs_col, so_row, so_col, BLOCK: tl.constexpr, N_ANG: tl.constexpr, N_GRID: tl.constexpr): offs = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) mask = offs < R @@ -160,13 +168,18 @@ def _rs_bwd_kernel(acos_ptr, asin_ptr, cs_ptr, ss_ptr, ca_ptr, sa_ptr, offs_g = tl.arange(0, N_GRID) g_mask = offs_g < n_grid - # One coalesced (BLOCK, N_ANG) tile load per operand (see forward). - tile = offs[:, None] * s_row + a_idx[None, :] * s_col + # One coalesced (BLOCK, N_ANG) tile load per operand, each with its own + # strides (see forward — inputs and incoming grads have independent + # producers, so their layouts are independent). t_mask = mask[:, None] & (a_idx < n_ang)[None, :] - ac_all = tl.load(acos_ptr + tile, mask=t_mask, other=0.0) - as_all = tl.load(asin_ptr + tile, mask=t_mask, other=0.0) - goc_all = tl.load(goc_ptr + tile, mask=t_mask, other=0.0) - gos_all = tl.load(gos_ptr + tile, mask=t_mask, other=0.0) + ac_all = tl.load(acos_ptr + offs[:, None] * sc_row + a_idx[None, :] * sc_col, + mask=t_mask, other=0.0) + as_all = tl.load(asin_ptr + offs[:, None] * ss_row + a_idx[None, :] * ss_col, + mask=t_mask, other=0.0) + goc_all = tl.load(goc_ptr + offs[:, None] * gc_row + a_idx[None, :] * gc_col, + mask=t_mask, other=0.0) + gos_all = tl.load(gos_ptr + offs[:, None] * gs_row + a_idx[None, :] * gs_col, + mask=t_mask, other=0.0) # Recompute f (never stored in the forward). f = tl.zeros((BLOCK, N_GRID), dtype=tl.float32) @@ -203,8 +216,9 @@ def _rs_bwd_kernel(acos_ptr, asin_ptr, cs_ptr, ss_ptr, ca_ptr, sa_ptr, sel = (a_idx == m).to(tl.float32) dac_all += tl.sum(df * cs[None, :], axis=1)[:, None] * sel[None, :] das_all += tl.sum(df * ss[None, :], axis=1)[:, None] * sel[None, :] - tl.store(dac_ptr + tile, dac_all, mask=t_mask) - tl.store(das_ptr + tile, das_all, mask=t_mask) + tile_o = offs[:, None] * so_row + a_idx[None, :] * so_col + tl.store(dac_ptr + tile_o, dac_all, mask=t_mask) + tl.store(das_ptr + tile_o, das_all, mask=t_mask) def _can_use_triton(A_cos, activation): @@ -216,20 +230,37 @@ def _fp32(*ts): return [t.to(torch.float32).contiguous() for t in ts] +def _rows(t): + """View a (n_e, F, n_ang) operand as R = n_e·F rows of n_ang coefficients + WITHOUT copying, returning (tensor, s_row, s_col) for the kernel's stride + arithmetic. Rows flatten whenever stride(0) == F·stride(1), which covers + both layouts the model produces: dense row-major, and the m-major output + of EquivariantLinear's einsum (strides (F, 1, R)) — the layout the old + `.contiguous()` copied at full read+write cost on every call. Both are + fully coalesced under the kernel's per-operand strides (m-major rows are + unit-stride). Anything else falls back to one contiguous copy.""" + t = t.to(torch.float32) + n_e, F, n_ang = t.shape + if t.stride(0) != F * t.stride(1): + t = t.contiguous() + return t, t.stride(1), t.stride(2) + + def _realspace_forward_triton(A_cos, A_sin, cos_synth, sin_synth, cos_analysis, sin_analysis): n_e, F, n_ang = A_cos.shape n_grid = cos_synth.shape[1] R = n_e * F - acos, asin, cs, ss, ca, sa = _fp32( - A_cos.reshape(R, n_ang), A_sin.reshape(R, n_ang), - cos_synth, sin_synth, cos_analysis, sin_analysis) + acos, sc_r, sc_c = _rows(A_cos) + asin, ss_r, ss_c = _rows(A_sin) + cs, ss, ca, sa = _fp32(cos_synth, sin_synth, cos_analysis, sin_analysis) oc = torch.empty(R, n_ang, device=A_cos.device, dtype=torch.float32) os = torch.empty_like(oc) N_ANG, N_GRID = triton.next_power_of_2(n_ang), triton.next_power_of_2(n_grid) grid = (triton.cdiv(R, _RS_BLOCK),) _rs_fwd_kernel[grid](acos, asin, cs, ss, ca, sa, oc, os, - R, n_ang, n_grid, acos.stride(0), acos.stride(1), + R, n_ang, n_grid, + sc_r, sc_c, ss_r, ss_c, oc.stride(0), oc.stride(1), BLOCK=_RS_BLOCK, N_ANG=N_ANG, N_GRID=N_GRID, num_warps=_RS_WARPS) out = lambda t: t.reshape(n_e, F, n_ang).to(A_cos.dtype) # noqa: E731 @@ -241,16 +272,19 @@ def _realspace_backward_triton(g_out_cos, g_out_sin, A_cos, A_sin, n_e, F, n_ang = A_cos.shape n_grid = cos_synth.shape[1] R = n_e * F - acos, asin, cs, ss, ca, sa, goc, gos = _fp32( - A_cos.reshape(R, n_ang), A_sin.reshape(R, n_ang), - cos_synth, sin_synth, cos_analysis, sin_analysis, - g_out_cos.reshape(R, n_ang), g_out_sin.reshape(R, n_ang)) + acos, sc_r, sc_c = _rows(A_cos) + asin, ss_r, ss_c = _rows(A_sin) + goc, gc_r, gc_c = _rows(g_out_cos) + gos, gs_r, gs_c = _rows(g_out_sin) + cs, ss, ca, sa = _fp32(cos_synth, sin_synth, cos_analysis, sin_analysis) dac = torch.empty(R, n_ang, device=A_cos.device, dtype=torch.float32) das = torch.empty_like(dac) N_ANG, N_GRID = triton.next_power_of_2(n_ang), triton.next_power_of_2(n_grid) grid = (triton.cdiv(R, _RS_BLOCK),) _rs_bwd_kernel[grid](acos, asin, cs, ss, ca, sa, goc, gos, dac, das, - R, n_ang, n_grid, acos.stride(0), acos.stride(1), + R, n_ang, n_grid, + sc_r, sc_c, ss_r, ss_c, + gc_r, gc_c, gs_r, gs_c, dac.stride(0), dac.stride(1), BLOCK=_RS_BLOCK, N_ANG=N_ANG, N_GRID=N_GRID, num_warps=_RS_WARPS) out = lambda t: t.reshape(n_e, F, n_ang).to(A_cos.dtype) # noqa: E731 diff --git a/tests/test_realspace_kernel.py b/tests/test_realspace_kernel.py index c37df24..2c10fa5 100644 --- a/tests/test_realspace_kernel.py +++ b/tests/test_realspace_kernel.py @@ -197,6 +197,32 @@ def run(fn): print(f" triton vs fp64 ref (F={F}, m_max={m_max}): " f"fwd {f_err:.1e}, dA {b_err:.1e}") + # m-major layout (strides (F, 1, n_e·F) — what EquivariantLinear's + # einsum emits): the kernel reads it via strides, no contiguous copy. + # Must match the contiguous-input result bit-for-bit-ish. + def as_m_major(t): + out = torch.empty(t.shape[2], t.shape[0], t.shape[1], + device=t.device, dtype=t.dtype).permute(1, 2, 0) + out.copy_(t) + return out + + def run_mm(fn): + a = as_m_major(A_cos).requires_grad_(True) + b = as_m_major(A_sin).requires_grad_(True) + oc, os = fn(a, b) + (oc * goc + os * gos).sum().backward() + return oc, os, a.grad, b.grad + + mm = run_mm(lambda a, b: RealSpaceFused.apply( + a, b, nl.cos_synth, nl.sin_synth, nl.cos_analysis, nl.sin_analysis, + nl.activation)) + f_mm = max((ker[0] - mm[0]).abs().max(), (ker[1] - mm[1]).abs().max()).item() + b_mm = max((ker[2] - mm[2]).abs().max(), (ker[3] - mm[3]).abs().max()).item() + assert f_mm < 1e-6, f"m-major fwd F={F},m={m_max}: {f_mm:.2e}" + assert b_mm < 1e-6, f"m-major grad F={F},m={m_max}: {b_mm:.2e}" + print(f" triton m-major layout (F={F}, m_max={m_max}): " + f"fwd {f_mm:.1e}, dA {b_mm:.1e}") + def _structure(n=7, n_types=4, seed=0): g = torch.Generator().manual_seed(seed) From 83028f19424c3c8fcde5d4e05069353b8a2215e4 Mon Sep 17 00:00:00 2001 From: alacour Date: Thu, 3 Sep 2026 18:00:11 -0700 Subject: [PATCH 3/5] =?UTF-8?q?realspace=20kernel:=20revert=20in-place=20m?= =?UTF-8?q?-major=20reads=20=E2=80=94=20measured=204-6x=20slower=20on=20A1?= =?UTF-8?q?00?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The previous commit read EquivariantLinear's m-major einsum output in place (strides (F, 1, R)) to skip the wrapper's contiguous copies. Profiling on the A100 showed the opposite of the intent: the fused nonlinearity went from 1.38 ms to ~6.0 ms per instance (512-atom box, 44k edges), with the same regression inside the MP trunk and receiver lines and a proportional hit in the backward — the (BLOCK, N_ANG) tile load degrades to strided scalar accesses across three ~45 MB-apart streams instead of one vectorized coalesced read. _rows() now skips the copy only for genuine row-major layouts (unit-stride coefficient axis); the m-major layout goes back to the contiguous copy the kernel is fast on. Per-operand stride plumbing stays for a future transposed-tile experiment. Restores the prior 159.6 ms joint-step timing. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_019WyADHSVvKbQnYo1PAjGPK --- ecenet/realspace_kernel.py | 21 ++++++++++++--------- 1 file changed, 12 insertions(+), 9 deletions(-) diff --git a/ecenet/realspace_kernel.py b/ecenet/realspace_kernel.py index 384b82f..a5332fb 100644 --- a/ecenet/realspace_kernel.py +++ b/ecenet/realspace_kernel.py @@ -231,17 +231,20 @@ def _fp32(*ts): def _rows(t): - """View a (n_e, F, n_ang) operand as R = n_e·F rows of n_ang coefficients - WITHOUT copying, returning (tensor, s_row, s_col) for the kernel's stride - arithmetic. Rows flatten whenever stride(0) == F·stride(1), which covers - both layouts the model produces: dense row-major, and the m-major output - of EquivariantLinear's einsum (strides (F, 1, R)) — the layout the old - `.contiguous()` copied at full read+write cost on every call. Both are - fully coalesced under the kernel's per-operand strides (m-major rows are - unit-stride). Anything else falls back to one contiguous copy.""" + """View a (n_e, F, n_ang) operand as R = n_e·F rows of n_ang coefficients, + returning (tensor, s_row, s_col) for the kernel's stride arithmetic. + + The copy is skipped only for row-major layouts (rows collapse AND the + coefficient axis is unit-stride) — the case the kernel loads as one + vectorized coalesced tile. Reading EquivariantLinear's m-major einsum + output (strides (F, 1, R)) in place was tried and MEASURED ~4-6x slower + on an A100 than copying it contiguous first, despite its unit-stride + rows — the tile load degrades to strided scalar accesses across three + ~45 MB-apart streams. So the m-major layout takes the contiguous copy; + revisit only with a transposed-tile load variant benchmarked in hand.""" t = t.to(torch.float32) n_e, F, n_ang = t.shape - if t.stride(0) != F * t.stride(1): + if t.stride(2) != 1 or t.stride(0) != F * t.stride(1): t = t.contiguous() return t, t.stride(1), t.stride(2) From 768d2c3d418680e7baa9c7c7a0ea0ad5f4043aab Mon Sep 17 00:00:00 2001 From: alacour Date: Thu, 3 Sep 2026 18:03:15 -0700 Subject: [PATCH 4/5] =?UTF-8?q?profile=5Fstep:=20--torch=5Fprofile=20?= =?UTF-8?q?=E2=80=94=20per-op=20forward/backward/memory=20breakdown?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The per-section timers only see the forward, op by op; the backward (the larger half of a force call — currently ~61% of the joint step) and the kernel-level cost of composite blocks (receiver, trunk) are invisible to them. --torch_profile reruns the same Totals closures under torch.profiler and prints three aggregate tables: forward-only by self CUDA time, the joint E+F step by self CUDA time (rows that grow between the two are the backward; record_function tags mark the Triton regions), and the joint step by self CUDA memory for the saved-tensor footprint behind large-box OOMs. --profile_rows controls table length. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_019WyADHSVvKbQnYo1PAjGPK --- tools/profile_step.py | 59 +++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 59 insertions(+) diff --git a/tools/profile_step.py b/tools/profile_step.py index 5596d87..401d544 100644 --- a/tools/profile_step.py +++ b/tools/profile_step.py @@ -38,6 +38,15 @@ parser.add_argument('--edge_frame_e2n', action='store_true', help="With --edge_frame_fused: ALSO fuse the MP layers' " 'pack+unrotate (their step 3).') +parser.add_argument('--torch_profile', action='store_true', + help='After the timing sections, run the forward and the ' + 'joint E+F step under torch.profiler and print the top ' + 'ops by self CUDA time (plus a top-memory table). This ' + 'is the view the per-section timers cannot give: the ' + 'backward decomposed by op, and record_function-tagged ' + 'kernel regions.') +parser.add_argument('--profile_rows', type=int, default=25, + help='--torch_profile: rows per table (default 25)') args = parser.parse_args() dtype = torch.float32 if args.float32 else torch.float64 @@ -383,3 +392,53 @@ def fwd_forces_les(): time_fn("forward_pbc + E_lr + forces (joint graph)", fwd_forces_les) print() + +# ── 5. torch.profiler breakdown (--torch_profile) ───────────────────────────── +# The per-section timers above can only see the forward, op by op; the backward +# (the larger half of a force call) and the true kernel-level cost of composite +# blocks are invisible to them. This profiles the same closures used in Totals +# and prints per-op aggregates: +# * forward-only table — kernel costs of the real forward path; +# * joint-step table — forward + backward together: rows that grow between +# the two tables (aten::mm vs MmBackward0 etc.) are the backward; the +# record_function tags (edge_frame_fused, pack_unrotate_fused, ...) mark +# our Triton regions in both; +# * joint-step memory table — top self-CUDA allocators, for the saved-tensor +# footprint behind large-box OOMs. +if args.torch_profile: + from torch.profiler import ProfilerActivity, profile + + acts = [ProfilerActivity.CPU] + if device.type == 'cuda': + acts.append(ProfilerActivity.CUDA) + sort_key = ('self_cuda_time_total' if device.type == 'cuda' + else 'self_cpu_time_total') + + def fwd_only(): + return model.forward_pbc(fresh_pos(), types, edge_i, edge_j, + shift_e, nb_src, nb_dst, shift_n) + + joint_fn = fwd_forces_les if les_module is not None else fwd_forces + joint_label = ('joint E_sr + E_lr + forces' if les_module is not None + else 'joint E + forces') + + def run_profile(fn, label, n=args.n_time): + fn() # warm (allocator, autotune caches) + sync() + with profile(activities=acts, profile_memory=True) as prof: + for _ in range(n): + fn() + sync() + ka = prof.key_averages() + sep(f"torch.profiler: {label} (x{n}, top {args.profile_rows} by self " + f"{'CUDA' if device.type == 'cuda' else 'CPU'} time)") + print(ka.table(sort_by=sort_key, row_limit=args.profile_rows)) + return ka + + run_profile(fwd_only, "forward only") + ka_joint = run_profile(joint_fn, joint_label) + + if device.type == 'cuda': + sep(f"torch.profiler: {joint_label} (top 10 by self CUDA memory)") + print(ka_joint.table(sort_by='self_cuda_memory_usage', row_limit=10)) +print() From 0bc8bfb0986d585b573d403171f4be57d47c0ac3 Mon Sep 17 00:00:00 2001 From: alacour Date: Thu, 3 Sep 2026 18:21:14 -0700 Subject: [PATCH 5/5] =?UTF-8?q?profile=5Fstep:=20ruff=20I001=20=E2=80=94?= =?UTF-8?q?=20reformat=20conditional=20import?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_019WyADHSVvKbQnYo1PAjGPK --- tools/profile_step.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/tools/profile_step.py b/tools/profile_step.py index 401d544..2e19e40 100644 --- a/tools/profile_step.py +++ b/tools/profile_step.py @@ -196,8 +196,7 @@ def fresh_pos(): lambda: model.sph_to_angular(A_rot)) if args.edge_frame_fused: - from ecenet.edge_frame_kernel import edge_frame_fused, edge_frame_fused_single, \ - pack_unrotate_fused + from ecenet.edge_frame_kernel import edge_frame_fused, edge_frame_fused_single, pack_unrotate_fused time_fn("edge-frame fused (gather+rotate+reshape)", lambda: edge_frame_fused(A_emb, edge_i, edge_j, D_block_main, model.sph_to_angular))