perf(pt): save SO2 competition scalars during rotation

deepmd/pt/entrypoints/freeze_pt2.py

@@ -1007,10 +1007,14 @@ def _freeze_sezm_to_pt2(
 
     # Device-dependent Python branches resolve on the CPU tracing inputs, so
     # pin them to the AOTI target. Non-CPU targets bake the block-diagonal SO(2)
-    # contraction, while CUDA targets also bake the fused GIE scatter.
+    # contraction, while CUDA targets also bake the fused GIE scatter and
+    # permit level-2 SO(2) scalar capture despite the CPU tracing inputs.
     force_block_diag = target_device.type != "cpu"
     force_fused_scatter = target_device.type == "cuda"
     for module in model.modules():
+        if isinstance(module, SO2Convolution):
+            module._scalar_capture_retains_rotation = True
+            module._force_scalar_capture_cuda_target = target_device.type == "cuda"
         if isinstance(module, SO2Linear):
             module._force_block_diag_matmul = force_block_diag
         if isinstance(module, GeometricInitialEmbedding):

deepmd/pt_expt/kernels/triton/sezm/so2_value_path.py

@@ -188,6 +188,12 @@ def _rotate_mix_reference(
     dim = (lmax + 1) ** 2
     n_deg = lmax + 1
     reduced = 3 * lmax + 1
+    if _has_no_edges(n_edge):
+        # opt_einsum 3.3 cannot plan zero-cost contractions (0 / 0).
+        # Empty reductions retain the zero derivative to every operand,
+        # including the nonempty node features and shared channel basis.
+        dependency = x[:0].sum() + wigner.sum() + kc.sum() + cb[:0].sum()
+        return x.new_empty((n_focus, n_edge, reduced * focus_dim)) + dependency
     coeff = build_m_major_index(lmax, 1, device=x.device)
     d_to_m = wigner[:, :dim, :dim].index_select(1, coeff)
     x_local = torch.bmm(d_to_m, x.index_select(0, src))  # (E, reduced, C_wide)
@@ -248,6 +254,15 @@ def _rotate_mix_backward_reference(
     dim = (lmax + 1) ** 2
     n_deg = lmax + 1
     reduced = 3 * lmax + 1
+    if _has_no_edges(n_edge):
+        # All three edge cotangents are empty. Preserve their differentiable
+        # dependency without sending zero-size einsums to the path planner.
+        dependency = grad_u.sum() + x[:0].sum() + wigner.sum() + kc.sum() + cb[:0].sum()
+        return (
+            x.new_empty((n_edge, dim, c_wide)) + dependency,
+            wigner.new_empty(wigner.shape) + dependency,
+            kc.new_empty(kc.shape) + dependency,
+        )
     coeff = build_m_major_index(lmax, 1, device=x.device)
     d_to_m = wigner[:, :dim, :dim].index_select(1, coeff)
     x_src = x.index_select(0, src)
@@ -508,6 +523,7 @@ def _rotate_mix_fwd_kernel(
         kc_ptr,  # (E, KSZ * RANK) compact kernel, or (E, L+1, CW) when RANK == 0
         cb_ptr,  # (RANK, CW) channel basis (unread when RANK == 0)
         u_ptr,  # (F, E, ROW) focus-major output
+        scalar_ptr,  # optional independent (E, F, CF), pitch 2*CF
         n_edge,
         x_sn,
         x_sd,
@@ -516,6 +532,7 @@ def _rotate_mix_fwd_kernel(
         CW: tl.constexpr,  # true C_wide; BC = next_power_of_2(CW) lanes with mask
         BC: tl.constexpr,
         RANK: tl.constexpr,
+        SAVE_SCALAR: tl.constexpr = False,
     ):
         """One program per edge, channels vectorized.
 
@@ -576,6 +593,13 @@ def _rotate_mix_fwd_kernel(
                     tl.float32
                 )
                 tl.store(u_ptr + f_off + o * CF, xl[o] * rad, mask=cmask)
+                if SAVE_SCALAR and o == 0:
+                    scalar_off = (
+                        (chan // CF).to(tl.int64) * n_edge * (2 * CF)
+                        + edge * (2 * CF)
+                        + chan % CF
+                    )
+                    tl.store(scalar_ptr + scalar_off, xl[o] * rad, mask=cmask)
             for o in tl.static_range(L):
                 rad = tl.load(rad_base + (o + 1) * CW + chan, mask=cmask, other=0.0).to(
                     tl.float32
@@ -606,6 +630,13 @@ def _rotate_mix_fwd_kernel(
             if RANK == 1:
                 acc = acc * cb[0]
             tl.store(u_ptr + f_off + o * CF, acc, mask=cmask)
+            if SAVE_SCALAR and o == 0:
+                scalar_off = (
+                    (chan // CF).to(tl.int64) * n_edge * (2 * CF)
+                    + edge * (2 * CF)
+                    + chan % CF
+                )
+                tl.store(scalar_ptr + scalar_off, acc, mask=cmask)
         for o in tl.static_range(L):
             accn = tl.zeros((BC,), dtype=tl.float32)
             accq = tl.zeros((BC,), dtype=tl.float32)
@@ -2498,6 +2529,7 @@ def _rotate_mix_impl(
         kc,
         cb,
         u,
+        u,  # unread unless SAVE_SCALAR is enabled
         n_edge,
         x.stride(0),
         x.stride(1),
@@ -2506,12 +2538,82 @@ def _rotate_mix_impl(
         CW=c_wide,
         BC=triton.next_power_of_2(c_wide),
         RANK=int(rank),
+        SAVE_SCALAR=False,
         num_warps=warps,
         num_stages=stages,
     )
     return u
 
 
+def _rotate_mix_scalar_impl(
+    x: Tensor,
+    src: Tensor,
+    src_order: Tensor,
+    src_rowptr: Tensor,
+    wigner: Tensor,
+    kc: Tensor,
+    cb: Tensor,
+    lmax: int,
+    n_focus: int,
+    rank: int,
+) -> tuple[Tensor, Tensor]:
+    """Emit the l=0 competition input without retaining a view of ``u``.
+
+    The extra store uses the existing rotation accumulator; no second gather
+    or rotation is performed.  The padded focus-major scalar allocation keeps
+    the competition reduction layout used by the original compiled path.
+    Its padding is never read.  Both outputs own independent storage.
+    """
+    del src_order, src_rowptr
+    n_edge = src.shape[0]
+    c_wide = int(x.shape[2])
+    focus_dim = c_wide // int(n_focus)
+    row = (3 * int(lmax) + 1) * focus_dim
+    if not _use_triton(x):
+        u = _rotate_mix_reference(x, src, wigner, kc, cb, lmax, n_focus, rank)
+        # CPU export tracing must describe the same independent, padded
+        # scalar storage as the CUDA implementation.
+        scalar = torch.empty_strided(
+            (n_edge, n_focus, focus_dim),
+            (2 * focus_dim, n_edge * 2 * focus_dim, 1),
+            device=x.device,
+            dtype=x.dtype,
+        )
+        scalar.copy_(u[:, :, :focus_dim].permute(1, 0, 2))
+        return u, scalar
+    u = torch.empty(n_focus, n_edge, row, device=x.device, dtype=x.dtype)
+    pitch = 2 * focus_dim
+    scalar = torch.empty_strided(
+        (n_edge, n_focus, focus_dim),
+        (pitch, n_edge * pitch, 1),
+        device=x.device,
+        dtype=x.dtype,
+    )
+    if not _has_no_edges(n_edge):
+        warps, stages = rotate_mix_fwd_config(c_wide, int(lmax))
+        wrap_triton(_rotate_mix_fwd_kernel)[(n_edge,)](
+            x,
+            src,
+            wigner,
+            kc,
+            cb,
+            u,
+            scalar,
+            n_edge,
+            x.stride(0),
+            x.stride(1),
+            L=int(lmax),
+            CF=focus_dim,
+            CW=c_wide,
+            BC=triton.next_power_of_2(c_wide),
+            RANK=int(rank),
+            SAVE_SCALAR=True,
+            num_warps=warps,
+            num_stages=stages,
+        )
+    return u, scalar
+
+
 def _rotate_mix_bwd_impl(
     grad_u: Tensor,
     x: Tensor,
@@ -3608,6 +3710,9 @@ def _mixing_stack_train_bwd_impl(
 _rotate_mix_op = torch.library.triton_op(
     "sezm_triton::so2_rotate_mix", mutates_args=()
 )(_rotate_mix_impl)
+_rotate_mix_scalar_op = torch.library.triton_op(
+    "sezm_triton::so2_rotate_mix_scalar", mutates_args=()
+)(_rotate_mix_scalar_impl)
 _rotate_mix_bwd_op = torch.library.triton_op(
     "sezm_triton::so2_rotate_mix_bwd", mutates_args=()
 )(_rotate_mix_bwd_impl)
@@ -3639,6 +3744,24 @@ def _mixing_stack_train_bwd_impl(
 )(_gated_act_bwd_impl)
 
 
+@_rotate_mix_scalar_op.register_fake
+def _rotate_mix_scalar_fake(
+    x, src, src_order, src_rowptr, wigner, kc, cb, lmax, n_focus, rank
+):
+    """Describe both outputs without specializing symbolic edges in einsum."""
+    focus_dim = x.shape[2] // n_focus
+    n_edge = src.shape[0]
+    return (
+        x.new_empty((n_focus, n_edge, (3 * lmax + 1) * focus_dim)),
+        torch.empty_strided(
+            (n_edge, n_focus, focus_dim),
+            (2 * focus_dim, n_edge * 2 * focus_dim, 1),
+            device=x.device,
+            dtype=x.dtype,
+        ),
+    )
+
+
 @_rotate_mix_op.register_fake
 def _(x, src, src_order, src_rowptr, wigner, kc, cb, lmax, n_focus, rank):
     focus_dim = x.shape[2] // n_focus
@@ -3892,6 +4015,25 @@ def _rotate_mix_backward(ctx, grad_u):
     return grad_x, None, None, None, grad_wigner, grad_kc, grad_cb, None, None, None
 
 
+def _rotate_mix_scalar_setup_context(ctx, inputs, output):
+    _rotate_mix_setup_context(ctx, inputs, output)
+    ctx.set_materialize_grads(False)
+
+
+def _rotate_mix_scalar_backward(ctx, grad_u, grad_scalar):
+    # Exactly the slice/permutation adjoint of the old competition branch.
+    # Keep pad/add visible to Inductor so it can fuse the gradient merge.
+    if grad_scalar is not None:
+        focus_dim = grad_scalar.shape[-1]
+        extra = torch.nn.functional.pad(
+            grad_scalar.permute(1, 0, 2), (0, 3 * ctx.lmax * focus_dim)
+        )
+        grad_u = extra if grad_u is None else grad_u + extra
+    if grad_u is None:
+        return (None,) * 10
+    return _rotate_mix_backward(ctx, grad_u)
+
+
 def _rotate_mix_bwd_setup_context(ctx, inputs, output):
     grad_u, x, src, wigner, kc, cb, lmax, n_focus, rank = inputs
     ctx.save_for_backward(grad_u, x, src, wigner, kc, cb)
@@ -4010,6 +4152,9 @@ def _segment_sum_backward(ctx, grad_out):
 _rotate_mix_op.register_autograd(
     _rotate_mix_backward, setup_context=_rotate_mix_setup_context
 )
+_rotate_mix_scalar_op.register_autograd(
+    _rotate_mix_scalar_backward, setup_context=_rotate_mix_scalar_setup_context
+)
 _rotate_mix_bwd_op.register_autograd(
     _rotate_mix_bwd_backward, setup_context=_rotate_mix_bwd_setup_context
 )
@@ -5176,6 +5321,27 @@ def make_triton_rotate_mix(conv: SO2Convolution) -> _TritonRotateMix | None:
 )
 
 
+def _save_competition_scalar(conv: SO2Convolution, x: Tensor) -> bool:
+    """Capture only where competition backward retains the rotation storage.
+
+    The PT-expt exporter can already release that storage without capture;
+    it pins this internal policy off before tracing. Runtime and PT export
+    retain the independent scalar buffer for the validated level-2 layout.
+    """
+    return (
+        getattr(conv, "_scalar_capture_retains_rotation", True)
+        and not conv.training
+        and conv.triton_infer_level == 2
+        and conv.lmax == 3
+        and conv.n_focus == 2
+        and conv.so2_focus_dim == 32
+        and conv.focus_compete
+        and x.dtype == torch.float32
+        # Export traces CPU inputs before moving the graph to its target.
+        and (x.is_cuda or getattr(conv, "_force_scalar_capture_cuda_target", False))
+    )
+
+
 class _TritonSO2ValuePath:
     """Per-convolution entry running the SO(2) value path through the fused ops.
 
@@ -5331,7 +5497,8 @@ def __call__(
             src_rowptr = torch.cat([counts.new_zeros(1), torch.cumsum(counts, 0)])
         else:
             src_order, src_rowptr = csr
-        u0 = _rotate_mix_op(
+        save_scalar = _save_competition_scalar(conv, x)
+        rotation = (_rotate_mix_scalar_op if save_scalar else _rotate_mix_op)(
             x.contiguous(),
             src,
             src_order,
@@ -5343,13 +5510,21 @@ def __call__(
             conv.n_focus,
             rank,
         )
+        if save_scalar:
+            u0, scalar = rotation
+        else:
+            u0 = rotation
 
         # === Step 3. Cross-focus competition weight from the l = 0 scalars ===
         apply_alpha = bool(conv.focus_compete and conv.n_focus > 1)
         if apply_alpha:
             # The small (E, F, Cf) copy keeps the softmax backward from
             # retaining a view of the whole focus-major activation.
-            gate_src = u0[:, :, : conv.so2_focus_dim].permute(1, 0, 2).contiguous()
+            gate_src = (
+                scalar.contiguous()
+                if save_scalar
+                else u0[:, :, : conv.so2_focus_dim].permute(1, 0, 2).contiguous()
+            )
             alpha = conv._focus_alpha(gate_src).to(u0.dtype).contiguous()
         else:
             alpha = torch.ones(

deepmd/pt_expt/kernels/triton/sezm/sweep_tile_configs.py

@@ -894,6 +894,7 @@ def launch(warps: int, stages: int) -> None:
             kc,
             cb,
             u,
+            u,  # scalar output is disabled during the original rotate sweep
             n_edge,
             x.stride(0),
             x.stride(1),
@@ -902,6 +903,7 @@ def launch(warps: int, stages: int) -> None:
             CW=c_wide,
             BC=triton.next_power_of_2(c_wide),
             RANK=1,
+            SAVE_SCALAR=False,
             num_warps=warps,
             num_stages=stages,
         )

deepmd/pt_expt/utils/serialization.py

@@ -1933,17 +1933,24 @@ def _trace_and_export_impl(
 
     # Device-dependent Python branches resolve on the CPU tracing inputs, so
     # pin them to the AOTI target. Non-CPU targets bake the block-diagonal SO(2)
-    # contraction, while CUDA targets also bake the fused GIE scatter.
+    # contraction, while CUDA targets also bake the fused GIE scatter and
+    # identify the CUDA target despite the CPU tracing inputs. This exporter
+    # already releases the full rotation output: a separate scalar store adds
+    # latency without reducing its peak, so keep the non-capture route.
     from deepmd.pt_expt.descriptor.dpa4_nn.embedding import (
         GeometricInitialEmbedding,
     )
     from deepmd.pt_expt.descriptor.dpa4_nn.so2 import (
+        SO2Convolution,
         SO2Linear,
     )
 
     force_block_diag = target_device.type != "cpu"
     force_fused_scatter = target_device.type == "cuda"
     for module in model.modules():
+        if isinstance(module, SO2Convolution):
+            module._scalar_capture_retains_rotation = False
+            module._force_scalar_capture_cuda_target = target_device.type == "cuda"
         if isinstance(module, SO2Linear):
             module._force_block_diag_matmul = force_block_diag
         if isinstance(module, GeometricInitialEmbedding):