diff --git a/ecenet/edge_frame_kernel.py b/ecenet/edge_frame_kernel.py index 6083ce3..f935bc6 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 @@ -346,6 +348,177 @@ 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 _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, @@ -391,6 +564,83 @@ 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)) + +# 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): + """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() + + +_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). + 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) + dD[:, :, cc] = ddc[:, :, pc] + dD[:, :, sc] = dds[:, :, ps] + 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 @@ -492,7 +742,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 +758,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) @@ -533,18 +784,28 @@ 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(), 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 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 @@ -558,7 +819,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) @@ -575,6 +837,16 @@ 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 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,)]( + 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, 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(), @@ -601,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 @@ -626,17 +899,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 @@ -651,7 +935,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. @@ -838,7 +1123,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. @@ -866,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, @@ -885,28 +1173,48 @@ 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) + 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, 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] @@ -919,13 +1227,25 @@ 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: dm_cos = torch.empty(E, n_ch, n_ang, dtype=dh.dtype, device=dh.device) dm_sin = torch.empty_like(dm_cos) + 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,)]( + 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, S, n_ang), + 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/ecenet/equivariant.py b/ecenet/equivariant.py index e2f793a..93f4614 100644 --- a/ecenet/equivariant.py +++ b/ecenet/equivariant.py @@ -26,6 +26,27 @@ 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. 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): @@ -34,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 @@ -41,14 +63,44 @@ 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 _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.""" + 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): + 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() + # 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/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) 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, diff --git a/tests/test_edge_frame_kernel.py b/tests/test_edge_frame_kernel.py index b9a8b5a..c06da6e 100644 --- a/tests/test_edge_frame_kernel.py +++ b/tests/test_edge_frame_kernel.py @@ -441,6 +441,24 @@ def test_triton_paths(): print("test_triton_paths[pack_unrotate]: OK") +def test_triton_packed_d(): + """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 + try: + 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 + + def test_model_integration(): """Full ECENet (no MP): energy and autograd forces identical flag on/off.""" import ecenet as _ecenet @@ -524,6 +542,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.") diff --git a/tests/test_equivariant_linear.py b/tests/test_equivariant_linear.py new file mode 100644 index 0000000..e59727d --- /dev/null +++ b/tests/test_equivariant_linear.py @@ -0,0 +1,158 @@ +"""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) + 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 + + +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})") + + +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.") 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 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')