perf(dpmodel): stop broadcasting DPA4 so3 linear weights across nodes

Both so3 channel mixers spelled their einsum as a batched matmul with the
NODE/EDGE axis as the matmul BATCH and the weight carrying a dummy leading
axis:

    matmul(x[:, :, :, None, :], weight_expanded[None, ...])

matmul broadcasts batch axes, so this expands the weight to
(N, D, F, Cin, Cout).  For examples/water/dpa4 that turns a 165K-element
parameter into 191M elements -- about 0.8 GB -- on every call, and autograd
must then reduce the whole expanded gradient back to the parameter shape.

An op-level CUDA profile of a DPA4 training step attributed 45.6 ms per
call to that ExpandBackward0 reduce over a [1152, 9, 1, 32, 576] operand,
three calls per step, making it the most expensive kernel in the run; the
ChannelLinear twin cost a further ~7-9 ms per call over [102510, 1, 32, 64].
The pt backend spells the same contraction as torch.einsum and never
expands the weight.

Batch over the small (D, F) / (F,) axes instead, which keeps N as matmul
ROWS.  The weight is then used in place and its gradient is an ordinary
matmul.  The transposes this adds touch only the (N, D, F, C) operands,
which are orders of magnitude smaller than the expanded weight.

Author: Han Wang

deepmd/dpmodel/descriptor/dpa4_nn/so3.py

@@ -131,10 +131,19 @@ def call(self, x: Any) -> Any:
             xp, self.weight[...], device=array_api_compat.device(x)
         )
         weight = xp.reshape(weight, (self.in_channels, self.n_focus, self.out_channels))
-        # einsum "bfi,ifo->bfo" as a broadcast batched matmul:
-        # (B, F, 1, Cin) @ (1, F, Cin, Cout) -> (B, F, 1, Cout)
+        # einsum "bfi,ifo->bfo" as a matmul batched over the FOCUS axis.
+        #
+        # NOT as ``matmul(x[:, :, None, :], weight[None, ...])``: that makes B a
+        # batch axis, so matmul broadcasts the weight to (B, F, Cin, Cout) --
+        # inflating a few-hundred-KB parameter into hundreds of millions of
+        # elements per call, whose gradient autograd must then reduce back down
+        # (an ``ExpandBackward0`` reduce that measured as the single most
+        # expensive kernel of a DPA4 training step). Batching over F instead
+        # keeps B as matmul ROWS, so the weight is used in place and its
+        # gradient is an ordinary matmul.
         weight = xp.permute_dims(weight, (1, 0, 2))  # (F, Cin, Cout)
-        out = xp.matmul(x[:, :, None, :], weight[None, ...])[..., 0, :]
+        out = xp.matmul(xp.permute_dims(x, (1, 0, 2)), weight)  # (F, B, Cout)
+        out = xp.permute_dims(out, (1, 0, 2))  # (B, F, Cout)
         if self.use_bias:
             bias = xp_asarray_nodetach(
                 xp, self.bias[...], device=array_api_compat.device(x)
@@ -439,12 +448,21 @@ def call(self, x: Any) -> Any:
         weight_expanded = xp.take(weight, expand_index, axis=0)  # (D, Cin, F, Cout)
 
         # === Step 2. Per-focus, per-degree channel mixing ===
-        # einsum "ndfi,difo->ndfo" as a broadcast batched matmul:
-        # (N, D, F, 1, Cin) @ (1, D, F, Cin, Cout) -> (N, D, F, 1, Cout)
+        # einsum "ndfi,difo->ndfo" as a matmul batched over the (D, F) axes.
+        #
+        # NOT as ``matmul(x[:, :, :, None, :], weight_expanded[None, ...])``:
+        # that makes N a batch axis, so matmul broadcasts the weight to
+        # (N, D, F, Cin, Cout) -- for the water DPA4 example, a 165K-element
+        # parameter expanded to 191M elements (~0.8 GB) on every call, whose
+        # gradient autograd then reduces back down. That ``ExpandBackward0``
+        # reduce measured at 45.6 ms per call, three calls per training step:
+        # the most expensive kernel in the run. Batching over the small (D, F)
+        # axes keeps N as matmul ROWS, so the weight is never expanded.
         weight_expanded = xp.permute_dims(
             weight_expanded, (0, 2, 1, 3)
         )  # (D, F, Cin, Cout)
-        out = xp.matmul(x[:, :, :, None, :], weight_expanded[None, ...])[..., 0, :]
+        out = xp.matmul(xp.permute_dims(x, (1, 2, 0, 3)), weight_expanded)
+        out = xp.permute_dims(out, (2, 0, 1, 3))  # (N, D, F, Cout)
 
         # === Step 3. Add l=0 bias ===
         if self.mlp_bias: