Revert "perf(dpmodel): contract the DPA4 grid-branch router with matmul"

This reverts commit 7518a41.

Measurement did not support it.  An op-level profile attributed the 45.6 ms
reduce to ExpandBackward0, not to this multiply, and re-benchmarking after
the change moved DPA4 eager training by nothing (1.514 -> 1.552 s/step, i.e.
run-to-run noise) while the offending kernel stayed byte-identical at
410.7 ms.  The GridBranch product is well under the size that would matter.

Since matmul is autocast-listed where mul/sum are not, keeping it would have
silently moved this contraction into bf16 under the autocast region for no
measured gain.  The actual site is the broadcast weight in so3.py, fixed
separately.

Author: Han Wang

deepmd/dpmodel/descriptor/dpa4_nn/grid_net.py

@@ -401,22 +401,8 @@ def call(
         router = self.router(scalar_pair)
         router = xp.exp(router - xp.max(router, axis=-1, keepdims=True))
         router = router / xp.sum(router, axis=-1, keepdims=True)
-        # einsum "ngfhc,nfh->ngfc", expressed as a batched matmul.
-        #
-        # NOT as ``xp.sum(value * router[:, None, :, :, None], axis=3)``: that
-        # broadcast-then-reduce materialises the whole (N, G, F, H, C) product
-        # -- ~0.8 GB at this example's grid resolution -- and reads it straight
-        # back, which measured as the single most expensive kernel of a DPA4
-        # training step. ``matmul`` contracts H in place, so only the
-        # (N, G, F, C) result is written. ``matmul`` broadcasts its leading
-        # batch axes, so the router's (N, 1, F, 1, H) view lines up with
-        # value's (N, G, F, H, C) without permuting (a permute here would
-        # reintroduce the very copy this avoids).
-        router_row = xp.reshape(
-            router, (n_batch, 1, n_focus, 1, self.n_branches)
-        )  # (N, 1, F, 1, H)
-        out = xp.matmul(router_row, value)  # (N, G, F, 1, C)
-        out = xp.reshape(out, (n_batch, n_grid, n_focus, self.channels))
+        # einsum "ngfhc,nfh->ngfc" as a broadcast sum over the branch axis
+        out = xp.sum(value * router[:, None, :, :, None], axis=3)  # (N, G, F, C)
 
         # === Step 3. Project back to coefficients and mix output channels ===
         return _project_frames(from_grid(out), self.out_proj, self.n_frames)