From 07f75b5f9ea4fc7c515e26878b18285c3a900461 Mon Sep 17 00:00:00 2001
From: OutisLi <LTC201806070316@gmail.com>
Date: Fri, 11 Sep 2026 12:47:18 +0800
Subject: [PATCH 1/8] feat(dpa4): fixed radial basis, single envelope and
 descriptor-declared AdamW routing

- `basis_type` accepts `bessel/fix` and `gaussian/fix`, which keep the Bessel
  frequencies or Gaussian centres at their initial values instead of training
  them: the dpmodel parser resolves the family and the flag, the PT
  `RadialBasis` and the pt-expt parameter promotion of DPA4 and DPA4C freeze
  the parameter, and it keeps its name and shape so checkpoints load under
  either form.
- `env_exp` accepts a single integer: one C^3 envelope on the message-passing
  edge weights and a bare radial basis. The fused CUDA edge-radial kernel
  accepts an empty basis-envelope series.
- `HybridMuonOptimizer` takes `adam_patterns`; the model bases of both
  backends compose the patterns their descriptor declares through
  `adam_route_patterns()` (DPA4: the first radial-embedding layer and the
  env-seed radial projection; DPA4C: the first radial-embedding layer), spin
  models delegate to their backbone, and both trainers pass the patterns to
  the optimizer. Nothing is written in the input.
- Presets `v20260911`: DPA4 with `env_exp` 5 and `gaussian/fix` on the
  `v20260901` normalization settings; DPA4C with `gaussian/fix`.

The HybridMuon routing patterns are read from the unwrapped model wrapper,
so that DDP and FSDP training build the optimizer.
---
 deepmd/dpmodel/descriptor/dpa4.py             | 34 +++++---
 deepmd/dpmodel/descriptor/dpa4_nn/radial.py   | 82 +++++++++++++------
 deepmd/dpmodel/descriptor/dpa4c.py            |  6 +-
 deepmd/dpmodel/model/dp_model.py              | 20 +++++
 deepmd/pt/model/descriptor/sezm.py            | 47 ++++++++---
 deepmd/pt/model/descriptor/sezm_nn/radial.py  | 55 ++++++++-----
 deepmd/pt/model/model/model.py                | 20 +++++
 deepmd/pt/model/model/spin_model.py           |  4 +
 deepmd/pt/optimizer/hybrid_muon.py            | 54 +++++++++++-
 deepmd/pt/train/training.py                   |  6 ++
 deepmd/pt_expt/descriptor/dpa4.py             | 21 +++++
 deepmd/pt_expt/descriptor/dpa4c.py            | 21 +++++
 .../pt_expt/kernels/cuda/dpa4/edge_radial.py  | 19 +++--
 deepmd/pt_expt/train/training.py              |  5 ++
 deepmd/utils/argcheck.py                      | 16 ++--
 deepmd/utils/model_preset.py                  | 16 ++++
 doc/model/dpa4.md                             | 32 +++++---
 doc/model/dpa4c.md                            | 10 ++-
 source/op/pt/dpa4/edge_radial.cu              | 19 ++++-
 source/tests/common/test_model_preset.py      | 22 +++++
 .../tests/consistent/descriptor/test_dpa4.py  |  2 +
 source/tests/pt/model/test_descriptor_sezm.py | 43 ++++++++++
 .../pt/model/test_dpa4_dpmodel_parity.py      | 54 +++++++++---
 .../pt/model/test_dpa4_ptexpt_grad_parity.py  | 33 ++++++++
 source/tests/pt/model/test_sezm_model.py      | 28 +++++++
 source/tests/pt/model/test_sezm_spin_model.py | 15 ++++
 source/tests/pt/test_hybrid_muon.py           | 38 +++++++++
 .../descriptor/test_dpa4_accelerated.py       | 22 ++++-
 .../pt_expt/descriptor/test_dpa4c_cpu.py      | 39 +++++++++
 .../tests/pt_expt/model/test_dpa4_export.py   | 15 +++-
 .../tests/pt_expt/model/test_dpa4_interop.py  | 34 ++++++++
 .../pt_expt/model/test_dpa4_native_spin.py    | 13 +++
 32 files changed, 725 insertions(+), 120 deletions(-)

diff --git a/deepmd/dpmodel/descriptor/dpa4.py b/deepmd/dpmodel/descriptor/dpa4.py
index f1106cd395..6501bbf644 100644
--- a/deepmd/dpmodel/descriptor/dpa4.py
+++ b/deepmd/dpmodel/descriptor/dpa4.py
@@ -307,14 +307,18 @@ class DescrptDPA4(NativeOP, BaseDescriptor):
     rcut
         Cutoff radius in Å.
     env_exp
-        C^3 cutoff envelope exponents `[rbf_env_exp, edge_env_exp]`.
-        - `rbf_env_exp`: Controls radial basis function envelope decay.
-        - `edge_env_exp`: Controls message passing edge weight envelope decay.
+        C^3 cutoff envelope exponents. A list `[rbf_env_exp, edge_env_exp]`
+        specifies the radial-basis and message-passing envelopes separately.
+        A zero radial-basis exponent disables that envelope.
+        An integer specifies only the message-passing envelope exponent and
+        disables the radial-basis envelope.
         Larger values give weaker suppression (values stay near 1.0 longer).
     channels
         Total channels per (l,m) coefficient.
     basis_type
-        Radial basis type. Supported values are ``"bessel"`` and ``"gaussian"``.
+        Radial basis type. Supported values are ``"bessel"``, ``"gaussian"``,
+        ``"bessel/fix"`` and ``"gaussian/fix"``; the ``/fix`` forms keep the
+        frequencies or centres fixed during training.
     n_radial
         Number of radial basis functions.
     radial_mlp
@@ -607,7 +611,7 @@ def __init__(
         ntypes: int,
         sel: list[int] | int,
         rcut: float = 6.0,
-        env_exp: list[int] | None = None,
+        env_exp: int | list[int] | None = None,
         channels: int = 64,
         basis_type: str = "bessel",
         n_radial: int = 16,
@@ -675,11 +679,17 @@ def __init__(
         self.rcut = float(rcut)
         if env_exp is None:
             env_exp = [7, 5]
-        if len(env_exp) != 2:
-            raise ValueError(
-                "`env_exp` must be a list of two integers: [rbf_env_exp, edge_env_exp]"
-            )
-        self.env_exp = [int(x) for x in env_exp]
+        if isinstance(env_exp, int):
+            self.env_exp = env_exp
+            edge_env_exp = env_exp
+        else:
+            if len(env_exp) != 2:
+                raise ValueError(
+                    "`env_exp` must be an integer or a list of two integers: "
+                    "[rbf_env_exp, edge_env_exp]"
+                )
+            self.env_exp = [int(x) for x in env_exp]
+            edge_env_exp = self.env_exp[1]
         self.eps = float(eps)
         # Floor for the envelope-squared degree normalization (GIE / env_seed).
         # version < 1.1 keeps the tiny ``eps`` floor (legacy path, untouched);
@@ -1068,7 +1078,7 @@ def __init__(
             basis_type=self.basis_type,
             n_radial=self.n_radial,
             precision=self.compute_precision,  # force fp32+
-            exponent=self.env_exp[0],
+            exponent=0 if isinstance(self.env_exp, int) else self.env_exp[0],
         )
 
         # === Shared radial embedding: RBF -> per-l radial features ===
@@ -1088,7 +1098,7 @@ def __init__(
         )
 
         # === C^3 cutoff envelope for edge weight ===
-        self.edge_envelope = C3CutoffEnvelope(rcut=self.rcut, exponent=self.env_exp[1])
+        self.edge_envelope = C3CutoffEnvelope(rcut=self.rcut, exponent=edge_env_exp)
 
         # === Edge-aligned Wigner-D calculator ===
         # Cartesian blocks (degree 1 or 2) skip the SO(2) rotations, so the full
diff --git a/deepmd/dpmodel/descriptor/dpa4_nn/radial.py b/deepmd/dpmodel/descriptor/dpa4_nn/radial.py
index aada21eda5..375e284123 100644
--- a/deepmd/dpmodel/descriptor/dpa4_nn/radial.py
+++ b/deepmd/dpmodel/descriptor/dpa4_nn/radial.py
@@ -81,8 +81,8 @@ class RadialMLP(NativeOP):
     compile and non-compile paths.
 
     The hidden RMSNorm normalizes each edge's radial features by their own RMS.
-    The input ``edge_rbf`` carries the C^3 cutoff envelope and therefore
-    vanishes at ``rcut``; the RMSNorm divides that envelope out, and its ``eps``
+    When the input ``edge_rbf`` includes the C^3 cutoff envelope, it vanishes
+    at ``rcut``. The RMSNorm divides that envelope out, and its ``eps``
     floor is crossed as the edge approaches ``rcut``. On a sparse neighborhood
     (e.g. a dimer) this floor-crossing produces a sharp kink in the potential
     energy surface just inside the cutoff. Setting ``radial_norm=False`` drops
@@ -424,9 +424,39 @@ def call(self, r: Any) -> Any:
         return t4 * (35.0 + t * (-84.0 + t * (70.0 - 20.0 * t)))
 
 
+def parse_basis_type(basis_type: str) -> tuple[str, bool]:
+    """
+    Split a radial basis type into its family and its ``/fix`` flag.
+
+    Parameters
+    ----------
+    basis_type : str
+        One of ``"bessel"``, ``"gaussian"``, ``"bessel/fix"`` or
+        ``"gaussian/fix"`` (case-insensitive).
+
+    Returns
+    -------
+    tuple[str, bool]
+        The basis family (``"bessel"`` or ``"gaussian"``) and whether the
+        basis parameters are held fixed during training.
+
+    Raises
+    ------
+    ValueError
+        If the basis type is not one of the four supported values.
+    """
+    family, _, suffix = str(basis_type).lower().partition("/")
+    if family not in ("bessel", "gaussian") or suffix not in ("", "fix"):
+        raise ValueError(
+            "`basis_type` must be 'bessel', 'gaussian', 'bessel/fix' or "
+            f"'gaussian/fix', got '{basis_type}'"
+        )
+    return family, suffix == "fix"
+
+
 class RadialBasis(NativeOP):
     """
-    Radial basis with C^3 cutoff envelope.
+    Radial basis with an optional C^3 cutoff envelope.
 
     The trainable radial parameters are stored in ``adam_freqs`` so HybridMuon
     routes them to Adam without weight decay.
@@ -450,8 +480,8 @@ class RadialBasis(NativeOP):
 
         w_n = n * π / rcut, for n = 1..n_radial (in 1/Å)
 
-    The C^3 cutoff envelope is multiplied directly into the output to ensure
-    strict smoothness at ``rcut``.
+    A positive ``exponent`` multiplies the C^3 cutoff envelope directly into
+    the output. Zero selects the raw basis without constructing an envelope.
 
     Parameters
     ----------
@@ -460,16 +490,15 @@ class RadialBasis(NativeOP):
     n_radial : int
         Number of basis functions.
     basis_type : str, optional
-        Radial basis type. Supported values are ``"bessel"`` and ``"gaussian"``.
+        Radial basis type. Supported values are ``"bessel"``, ``"gaussian"``,
+        ``"bessel/fix"`` and ``"gaussian/fix"``; the ``/fix`` forms are
+        evaluated like their family and differ only in training, where the
+        PT backend keeps their frequencies or centres fixed.
     precision : str
         Floating-point precision for the radial basis frequencies and outputs.
     exponent : int, optional
-        Exponent for the C^3 cutoff envelope polynomial. Default is 7.
-    apply_envelope : bool, optional
-        Whether :meth:`call` multiplies the raw basis by the C³ envelope.
-        The default ``True`` preserves the DPA4 radial contract. Consumers that
-        apply one shared envelope after combining radial and type features may
-        request the raw basis with ``False``.
+        Exponent for the C^3 cutoff envelope polynomial. Zero disables the
+        envelope. Default is 7.
     """
 
     def __init__(
@@ -479,7 +508,6 @@ def __init__(
         n_radial: int = 10,
         precision: str = DEFAULT_PRECISION,
         exponent: int = 7,
-        apply_envelope: bool = True,
     ) -> None:
         self.rcut = float(rcut)
         if self.rcut <= 0.0:
@@ -488,17 +516,17 @@ def __init__(
         if self.n_radial <= 0:
             raise ValueError("`n_radial` must be positive")
         self.basis_type = str(basis_type).lower()
-        if self.basis_type not in ("bessel", "gaussian"):
-            raise ValueError("`basis_type` must be either 'bessel' or 'gaussian'")
+        # The ``/fix`` suffix governs training only: the PT backend freezes the
+        # basis parameters, the array-API basis evaluates either form alike.
+        self.basis_family, _ = parse_basis_type(self.basis_type)
         self.precision = precision
         self.exponent = int(exponent)
-        self.apply_envelope = bool(apply_envelope)
         prec = PRECISION_DICT[self.precision.lower()]
         self.pi_tensor = math.pi
 
         # Frequencies: n*π/rcut, n=1..n_radial
         # Shape: (1, n_radial), stored as a trainable array.
-        if self.basis_type == "bessel":
+        if self.basis_family == "bessel":
             freqs = np.arange(1, self.n_radial + 1, dtype=prec) * (math.pi / self.rcut)
         else:
             freqs = np.linspace(0.0, self.rcut, self.n_radial, dtype=prec)
@@ -506,10 +534,14 @@ def __init__(
         gaussian_width = self.rcut / max(self.n_radial - 1, 1)
         self.gaussian_coeff = -0.5 / (gaussian_width * gaussian_width)
 
-        self.envelope = C3CutoffEnvelope(
-            rcut=self.rcut,
-            exponent=self.exponent,
-            precision=self.precision,
+        self.envelope = (
+            C3CutoffEnvelope(
+                rcut=self.rcut,
+                exponent=self.exponent,
+                precision=self.precision,
+            )
+            if self.exponent != 0
+            else None
         )
 
     def call(self, r: Any) -> Any:
@@ -525,7 +557,7 @@ def call(self, r: Any) -> Any:
         -------
         Array
             Radial basis with shape ``(N, n_radial)``. When
-            ``apply_envelope=True``, the output includes the C³ envelope and
+            ``exponent > 0``, the output includes the C³ envelope and
             vanishes smoothly at ``rcut``; otherwise it is the raw basis.
         """
         xp = array_api_compat.array_namespace(r)
@@ -534,7 +566,7 @@ def call(self, r: Any) -> Any:
         )
         # === Step 1. Radial basis ===
         # Shape: (N, 1) * (1, n_radial) -> (N, n_radial)
-        if self.basis_type == "bessel":
+        if self.basis_family == "bessel":
             # phi_n(r) = w_n * sinc(w_n * r / π)
             x = r * freqs  # (N, n_rbf)
             # torch.sinc(z) = sin(π z) / (π z) with sinc(0) = 1. The array API
@@ -551,7 +583,7 @@ def call(self, r: Any) -> Any:
             raw = xp.exp(dr * dr * self.gaussian_coeff)  # (N, n_rbf)
 
         # === Step 2. Apply the optional C³ envelope ===
-        if self.apply_envelope:
+        if self.envelope is not None:
             return raw * self.envelope(r)
         return raw
 
@@ -565,7 +597,6 @@ def serialize(self) -> dict[str, Any]:
                 "basis_type": self.basis_type,
                 "n_radial": self.n_radial,
                 "exponent": self.exponent,
-                "apply_envelope": self.apply_envelope,
                 "precision": np.dtype(PRECISION_DICT[self.precision]).name,
             },
             "@variables": {"adam_freqs": to_numpy_array(self.adam_freqs)},
@@ -588,7 +619,6 @@ def deserialize(cls, data: dict[str, Any]) -> RadialBasis:
             n_radial=int(config["n_radial"]),
             basis_type=str(config.get("basis_type", "bessel")),
             exponent=int(config.get("exponent", 7)),
-            apply_envelope=bool(config.get("apply_envelope", True)),
             precision=precision,
         )
         if variables is not None:
diff --git a/deepmd/dpmodel/descriptor/dpa4c.py b/deepmd/dpmodel/descriptor/dpa4c.py
index e5a985b146..e38d8374da 100644
--- a/deepmd/dpmodel/descriptor/dpa4c.py
+++ b/deepmd/dpmodel/descriptor/dpa4c.py
@@ -150,7 +150,8 @@ class DescrptDPA4C(NativeOP, BaseDescriptor):
     lmax
         Maximum angular degree. Supported values are 2, 3, and 4.
     basis_type
-        DPA4 radial basis type: ``"bessel"`` or ``"gaussian"``.
+        DPA4 radial basis type: ``"bessel"``, ``"gaussian"`` or their ``/fix``
+        forms, which keep the basis parameters fixed during training.
     n_radial
         Number of DPA4 radial basis functions forming the fixed analytic
         radial input.
@@ -322,8 +323,7 @@ def __init__(
             basis_type=self.basis_type,
             n_radial=self.n_radial,
             precision=self.precision,
-            exponent=self._ENVELOPE_EXPONENT,
-            apply_envelope=False,
+            exponent=0,
         )
         self.radial_embedding = SwiGLUMLP(
             [self.n_radial, radial_hidden, self.channels],
diff --git a/deepmd/dpmodel/model/dp_model.py b/deepmd/dpmodel/model/dp_model.py
index 5a1a253596..c1c390eda3 100644
--- a/deepmd/dpmodel/model/dp_model.py
+++ b/deepmd/dpmodel/model/dp_model.py
@@ -20,6 +20,26 @@ class DPModelCommon:
     neighbor selection updates and fitting network access.
     """
 
+    def adam_route_patterns(self) -> list[str]:
+        """
+        Name patterns of the parameters that take the AdamW path under HybridMuon.
+
+        Returns
+        -------
+        list[str]
+            Substrings of parameter names, composed from the tensors the
+            descriptor declares through its own ``adam_route_patterns``: rows
+            of these matrices that correspond to rarely visited inputs receive
+            almost no gradient, and Adam moves each row with its own gradient
+            history, whereas Muon's orthogonalized update moves every row of a
+            matrix at the same rate.
+        """
+        descriptor = getattr(getattr(self, "atomic_model", None), "descriptor", None)
+        declared = getattr(descriptor, "adam_route_patterns", None)
+        if declared is None:
+            return []
+        return [f"descriptor.{p}" for p in declared()]
+
     @classmethod
     def update_sel(
         cls,
diff --git a/deepmd/pt/model/descriptor/sezm.py b/deepmd/pt/model/descriptor/sezm.py
index 3988cd29eb..b770a9edc9 100644
--- a/deepmd/pt/model/descriptor/sezm.py
+++ b/deepmd/pt/model/descriptor/sezm.py
@@ -148,14 +148,18 @@ class DescrptSeZM(BaseDescriptor, nn.Module):
     rcut
         Cutoff radius in Å.
     env_exp
-        C^3 cutoff envelope exponents `[rbf_env_exp, edge_env_exp]`.
-        - `rbf_env_exp`: Controls radial basis function envelope decay.
-        - `edge_env_exp`: Controls message passing edge weight envelope decay.
+        C^3 cutoff envelope exponents. A list `[rbf_env_exp, edge_env_exp]`
+        specifies the radial-basis and message-passing envelopes separately.
+        A zero radial-basis exponent disables that envelope.
+        An integer specifies only the message-passing envelope exponent and
+        disables the radial-basis envelope.
         Larger values give weaker suppression (values stay near 1.0 longer).
     channels
         Total channels per (l,m) coefficient.
     basis_type
-        Radial basis type. Supported values are ``"bessel"`` and ``"gaussian"``.
+        Radial basis type. Supported values are ``"bessel"``, ``"gaussian"``,
+        ``"bessel/fix"`` and ``"gaussian/fix"``; the ``/fix`` forms keep the
+        frequencies or centres fixed during training.
     n_radial
         Number of radial basis functions.
     radial_mlp
@@ -449,7 +453,7 @@ def __init__(
         ntypes: int,
         sel: list[int] | int,
         rcut: float = 6.0,
-        env_exp: list[int] | None = None,
+        env_exp: int | list[int] | None = None,
         channels: int = 64,
         basis_type: str = "bessel",
         n_radial: int = 16,
@@ -523,11 +527,17 @@ def __init__(
         self.rcut = float(rcut)
         if env_exp is None:
             env_exp = [7, 5]
-        if len(env_exp) != 2:
-            raise ValueError(
-                "`env_exp` must be a list of two integers: [rbf_env_exp, edge_env_exp]"
-            )
-        self.env_exp = [int(x) for x in env_exp]
+        if isinstance(env_exp, int):
+            self.env_exp = env_exp
+            edge_env_exp = env_exp
+        else:
+            if len(env_exp) != 2:
+                raise ValueError(
+                    "`env_exp` must be an integer or a list of two integers: "
+                    "[rbf_env_exp, edge_env_exp]"
+                )
+            self.env_exp = [int(x) for x in env_exp]
+            edge_env_exp = self.env_exp[1]
         self.eps = float(eps)
         # Floor for the envelope-squared degree normalization (GIE / env_seed).
         # version < 1.1 keeps the tiny ``eps`` floor (legacy path, untouched);
@@ -929,7 +939,7 @@ def __init__(
             basis_type=self.basis_type,
             n_radial=self.n_radial,
             dtype=self.compute_dtype,  # force fp32+
-            exponent=self.env_exp[0],
+            exponent=0 if isinstance(self.env_exp, int) else self.env_exp[0],
             trainable=self.trainable,
         )
 
@@ -950,7 +960,7 @@ def __init__(
         )
 
         # === C^3 cutoff envelope for edge weight ===
-        self.edge_envelope = C3CutoffEnvelope(rcut=self.rcut, exponent=self.env_exp[1])
+        self.edge_envelope = C3CutoffEnvelope(rcut=self.rcut, exponent=edge_env_exp)
 
         # === Edge-aligned Wigner-D calculator ===
         # Cartesian blocks (degree 1 or 2) skip the SO(2) rotations, so the full
@@ -2344,6 +2354,19 @@ def _compute_mode_ctx(self, device: torch.device) -> Generator[None, None, None]
             yield
 
     # === DeePMD descriptor interface ===
+    def adam_route_patterns(self) -> list[str]:
+        """
+        Name patterns, relative to the descriptor, of the tensors that take the
+        AdamW path under HybridMuon: the first layer of the radial embedding and
+        the radial projection of the environment seed, which read the radial
+        basis and whose rows for rarely visited separations receive almost no
+        gradient.
+        """
+        return [
+            "radial_embedding.net.0.",
+            "env_seed_embedding.rbf_proj_layer1.",
+        ]
+
     def get_rcut(self) -> float:
         return self.rcut
 
diff --git a/deepmd/pt/model/descriptor/sezm_nn/radial.py b/deepmd/pt/model/descriptor/sezm_nn/radial.py
index 7448ff9fe1..43776e2210 100644
--- a/deepmd/pt/model/descriptor/sezm_nn/radial.py
+++ b/deepmd/pt/model/descriptor/sezm_nn/radial.py
@@ -21,6 +21,9 @@
     rearrange,
 )
 
+from deepmd.dpmodel.descriptor.dpa4_nn.radial import (
+    parse_basis_type,
+)
 from deepmd.dpmodel.utils.seed import (
     child_seed,
 )
@@ -83,8 +86,8 @@ class RadialMLP(nn.Module):
     compile and non-compile paths.
 
     The hidden RMSNorm normalizes each edge's radial features by their own RMS.
-    The input ``edge_rbf`` carries the C^3 cutoff envelope and therefore
-    vanishes at ``rcut``; the RMSNorm divides that envelope out, and its ``eps``
+    When the input ``edge_rbf`` includes the C^3 cutoff envelope, it vanishes
+    at ``rcut``. The RMSNorm divides that envelope out, and its ``eps``
     floor is crossed as the edge approaches ``rcut``. On a sparse neighborhood
     (e.g. a dimer) this floor-crossing produces a sharp kink in the potential
     energy surface just inside the cutoff. Setting ``radial_norm=False`` drops
@@ -422,7 +425,7 @@ def forward(self, r: torch.Tensor) -> torch.Tensor:
 
 class RadialBasis(nn.Module):
     """
-    Radial basis with C^3 cutoff envelope.
+    Radial basis with an optional C^3 cutoff envelope.
 
     The trainable radial parameters are stored in ``adam_freqs`` so HybridMuon
     routes them to Adam without weight decay.
@@ -446,8 +449,8 @@ class RadialBasis(nn.Module):
 
         w_n = n * π / rcut, for n = 1..n_radial (in 1/Å)
 
-    The C^3 cutoff envelope is multiplied directly into the output to ensure
-    strict smoothness at ``rcut``.
+    A positive ``exponent`` multiplies the C^3 cutoff envelope directly into
+    the output. Zero selects the raw basis without constructing an envelope.
 
     Parameters
     ----------
@@ -456,11 +459,16 @@ class RadialBasis(nn.Module):
     n_radial : int
         Number of basis functions.
     basis_type : str, optional
-        Radial basis type. Supported values are ``"bessel"`` and ``"gaussian"``.
+        Radial basis type. Supported values are ``"bessel"``, ``"gaussian"``,
+        ``"bessel/fix"`` and ``"gaussian/fix"``. The ``/fix`` forms keep the
+        frequencies or centres at their initial values: an isolated pair or a
+        compressed contact drives no data gradient into the basis, and a
+        trainable basis drifts there under the optimizer's momentum.
     dtype : torch.dtype
         Floating-point dtype for the radial basis frequencies and outputs.
     exponent : int, optional
-        Exponent for the C^3 cutoff envelope polynomial. Default is 7.
+        Exponent for the C^3 cutoff envelope polynomial. Zero disables the
+        envelope. Default is 7.
     trainable : bool, optional
         Whether the basis frequencies are trainable. Default is True.
     """
@@ -482,8 +490,7 @@ def __init__(
         if self.n_radial <= 0:
             raise ValueError("`n_radial` must be positive")
         self.basis_type = str(basis_type).lower()
-        if self.basis_type not in ("bessel", "gaussian"):
-            raise ValueError("`basis_type` must be either 'bessel' or 'gaussian'")
+        self.basis_family, fixed = parse_basis_type(self.basis_type)
         self.dtype = dtype
         self.device = env.DEVICE
         self.precision = RESERVED_PRECISION_DICT[self.dtype]
@@ -496,7 +503,7 @@ def __init__(
 
         # Frequencies: n*π/rcut, n=1..n_radial
         # Shape: (1, n_radial), stored as trainable nn.Parameter.
-        if self.basis_type == "bessel":
+        if self.basis_family == "bessel":
             freqs = torch.arange(
                 1,
                 self.n_radial + 1,
@@ -511,7 +518,7 @@ def __init__(
                 device=self.device,
                 dtype=self.dtype,
             )
-        self.trainable = bool(trainable)
+        self.trainable = bool(trainable) and not fixed
         self.adam_freqs = nn.Parameter(
             rearrange(freqs, "n_radial -> 1 n_radial"),
             requires_grad=self.trainable,
@@ -527,10 +534,14 @@ def __init__(
             persistent=False,
         )
 
-        self.envelope = C3CutoffEnvelope(
-            rcut=self.rcut,
-            exponent=self.exponent,
-            dtype=self.dtype,
+        self.envelope = (
+            C3CutoffEnvelope(
+                rcut=self.rcut,
+                exponent=self.exponent,
+                dtype=self.dtype,
+            )
+            if self.exponent != 0
+            else None
         )
 
     def forward(self, r: torch.Tensor) -> torch.Tensor:
@@ -545,12 +556,13 @@ def forward(self, r: torch.Tensor) -> torch.Tensor:
         Returns
         -------
         torch.Tensor
-            Radial basis multiplied by C^3 cutoff envelope with shape (N, n_rbf).
-            The output is smoothly truncated to zero at r = rcut.
+            Radial basis with shape (N, n_radial). When
+            ``exponent > 0``, the output includes the C^3 envelope and
+            vanishes smoothly at ``rcut``; otherwise it is the raw basis.
         """
         # === Step 1. Radial basis ===
         # Shape: (N, 1) * (1, n_radial) -> (N, n_radial)
-        if self.basis_type == "bessel":
+        if self.basis_family == "bessel":
             # phi_n(r) = w_n * sinc(w_n * r / π)
             x = r * self.adam_freqs  # (N, n_rbf)
             raw = self.adam_freqs * torch.sinc(x / self.pi_tensor)  # (N, n_rbf)
@@ -558,9 +570,10 @@ def forward(self, r: torch.Tensor) -> torch.Tensor:
             dr = r - self.adam_freqs  # (N, n_rbf)
             raw = torch.exp(dr * dr * self.gaussian_coeff)  # (N, n_rbf)
 
-        # === Step 2. Apply C^3 envelope for smooth cutoff ===
-        envelope = self.envelope(r)  # (N, 1)
-        return raw * envelope
+        # === Step 2. Apply the optional C^3 envelope ===
+        if self.envelope is not None:
+            return raw * self.envelope(r)
+        return raw
 
     def serialize(self) -> dict[str, Any]:
         """Serialize RadialBasis including trainable frequencies."""
diff --git a/deepmd/pt/model/model/model.py b/deepmd/pt/model/model/model.py
index 02ecedbce7..a4cd4bb689 100644
--- a/deepmd/pt/model/model/model.py
+++ b/deepmd/pt/model/model/model.py
@@ -128,3 +128,23 @@ def get_min_nbor_dist(self) -> float | None:
     def get_ntypes(self) -> int:
         """Returns the number of element types."""
         return len(self.get_type_map())
+
+    def adam_route_patterns(self) -> list[str]:
+        """
+        Name patterns of the parameters that take the AdamW path under HybridMuon.
+
+        Returns
+        -------
+        list[str]
+            Substrings of parameter names, composed from the tensors the
+            descriptor declares through its own ``adam_route_patterns``: rows
+            of these matrices that correspond to rarely visited inputs receive
+            almost no gradient, and Adam moves each row with its own gradient
+            history, whereas Muon's orthogonalized update moves every row of a
+            matrix at the same rate.
+        """
+        descriptor = getattr(getattr(self, "atomic_model", None), "descriptor", None)
+        declared = getattr(descriptor, "adam_route_patterns", None)
+        if declared is None:
+            return []
+        return [f"descriptor.{p}" for p in declared()]
diff --git a/deepmd/pt/model/model/spin_model.py b/deepmd/pt/model/model/spin_model.py
index 029f6c1dd0..d87983abaf 100644
--- a/deepmd/pt/model/model/spin_model.py
+++ b/deepmd/pt/model/model/spin_model.py
@@ -375,6 +375,10 @@ def expand_aparam(aparam: torch.Tensor, nloc: int) -> torch.Tensor:
         return aparam
 
     @torch.jit.export
+    def adam_route_patterns(self) -> list[str]:
+        """Route the backbone's declared tensors; the wrapper adds no parameters of its own."""
+        return self.backbone_model.adam_route_patterns()
+
     def get_type_map(self) -> list[str]:
         """Get the type map."""
         tmap = self.backbone_model.get_type_map()
diff --git a/deepmd/pt/optimizer/hybrid_muon.py b/deepmd/pt/optimizer/hybrid_muon.py
index ce3a538cec..653e42b71f 100644
--- a/deepmd/pt/optimizer/hybrid_muon.py
+++ b/deepmd/pt/optimizer/hybrid_muon.py
@@ -135,6 +135,7 @@
     from collections.abc import (
         Callable,
         Iterable,
+        Sequence,
     )
 
 # ============================================================================
@@ -706,6 +707,29 @@ def _compute_muon_nesterov_updates(
     ]
 
 
+def adam_route_patterns(models: Iterable[Any]) -> list[str]:
+    """
+    Collect the AdamW name patterns declared by the task models.
+
+    Parameters
+    ----------
+    models : Iterable[Any]
+        The task models; each may implement ``adam_route_patterns``.
+
+    Returns
+    -------
+    list[str]
+        Sorted union of the declared patterns.
+    """
+    return sorted(
+        {
+            pattern
+            for model in models
+            for pattern in getattr(model, "adam_route_patterns", list)()
+        }
+    )
+
+
 def get_adam_route(
     param_name: str | None,
 ) -> str:
@@ -829,6 +853,8 @@ class HybridMuonOptimizer(Optimizer):
       update.
     - Parameters with final effective name segment starting with ``adamw_``
       (case-insensitive): Adam with decoupled weight decay (AdamW-style).
+    - Matrix parameters whose full name contains one of ``adam_patterns``:
+      Adam with decoupled weight decay instead of Muon.
     - 1D parameters: standard Adam update.
     - Parameters are routed by effective shape (singleton dimensions removed).
     - ``muon_mode="2d"``:
@@ -881,6 +907,13 @@ class HybridMuonOptimizer(Optimizer):
         with AdamW-style decoupled decay. Not applied to 1D Adam parameters.
     adam_betas : tuple[float, float]
         Adam beta coefficients with default (0.9, 0.95).
+    adam_patterns : Sequence[str]
+        Case-insensitive substrings of full parameter names. Matrix parameters
+        whose name contains one of them take the AdamW path (decoupled weight
+        decay) instead of Muon; vector parameters are on Adam regardless.
+        The trainer fills this from the model's ``adam_route_patterns`` — the
+        tensors that only a small part of the data constrains, such as the
+        layers reading a descriptor's radial basis or the fitting network.
     lr_adjust : float
         Learning rate adjustment mode for Muon scaling and Adam learning rate.
         - If lr_adjust <= 0: use match-RMS scaling for Muon,
@@ -949,6 +982,7 @@ def __init__(
         flash_muon: bool = True,
         magma_muon: bool = True,
         use_foreach: bool | None = None,
+        adam_patterns: Sequence[str] = (),
     ) -> None:
         # === Step 1. Validate routing mode ===
         muon_mode = str(muon_mode).lower()
@@ -972,6 +1006,9 @@ def __init__(
         super().__init__(params, defaults)
 
         # === Step 3. Build parameter id -> name mapping ===
+        # Name patterns, like the names themselves, are routing configuration
+        # of this run and stay out of ``defaults`` and the optimizer state.
+        self._adam_patterns = tuple(str(p).lower() for p in adam_patterns)
         self._param_name_map: dict[int, str] = {}
         if named_parameters is not None:
             self.set_param_names(named_parameters)
@@ -1020,6 +1057,13 @@ def __init__(
         self._per_parameter_adam_clock = False
         self._bias_corrections_migrated = False
 
+    def _matches_adam_pattern(self, param_name: str | None) -> bool:
+        """Return whether the full parameter name contains one of ``adam_patterns``."""
+        if param_name is None or not self._adam_patterns:
+            return False
+        lowered = param_name.lower()
+        return any(pattern in lowered for pattern in self._adam_patterns)
+
     def set_param_names(
         self, named_parameters: Iterable[tuple[str, torch.Tensor]]
     ) -> None:
@@ -1607,9 +1651,15 @@ def _build_param_routing(self) -> None:
                     adam_no_decay.append({"param": p, "name": param_name})
                     continue
 
-                # === Step 3. Non-matrix effective shape in current mode → AdamW-style ===
+                # === Step 3. Pattern-routed or non-matrix effective shape → AdamW-style ===
+                # A parameter whose full name contains one of ``adam_patterns``
+                # takes the Adam path with decoupled decay instead of Muon: the
+                # orthogonalized Muon step moves every row of a matrix at the
+                # same rate, including rows that only a small part of the data
+                # constrains, whereas Adam moves a row in proportion to its own
+                # gradient history.
                 matrix_shape = get_matrix_view_shape(effective_shape, muon_mode)
-                if matrix_shape is None:
+                if matrix_shape is None or self._matches_adam_pattern(param_name):
                     adam_decay.append({"param": p, "name": param_name})
                     continue
 
diff --git a/deepmd/pt/train/training.py b/deepmd/pt/train/training.py
index dd515b1a64..bb5248d89d 100644
--- a/deepmd/pt/train/training.py
+++ b/deepmd/pt/train/training.py
@@ -71,6 +71,9 @@
     KFOptimizerWrapper,
     LKFOptimizer,
 )
+from deepmd.pt.optimizer.hybrid_muon import (
+    adam_route_patterns,
+)
 from deepmd.pt.train.wrapper import (
     ModelWrapper,
 )
@@ -1106,6 +1109,9 @@ def update_finetune_bias(
                     "enable_gram": bool(self.opt_param.get("enable_gram")),
                     "flash_muon": bool(self.opt_param.get("flash_muon")),
                     "magma_muon": bool(self.opt_param.get("magma_muon")),
+                    "adam_patterns": adam_route_patterns(
+                        self._get_inner_module().model.values()
+                    ),
                     # FSDP2 shards parameters as DTensor; several torch._foreach_*
                     # ops lack DTensor sharding propagation on older PyTorch, so
                     # fall back to the per-tensor path under zero_stage >= 2.
diff --git a/deepmd/pt_expt/descriptor/dpa4.py b/deepmd/pt_expt/descriptor/dpa4.py
index 2205c6b849..0de6466d90 100644
--- a/deepmd/pt_expt/descriptor/dpa4.py
+++ b/deepmd/pt_expt/descriptor/dpa4.py
@@ -16,6 +16,9 @@
     C3CutoffEnvelope as C3CutoffEnvelopeDP,
 )
 from deepmd.dpmodel.descriptor.dpa4_nn.radial import InnerClamp as InnerClampDP
+from deepmd.dpmodel.descriptor.dpa4_nn.radial import (
+    parse_basis_type,
+)
 from deepmd.pt_expt.common import (
     register_dpmodel_mapping,
     torch_module,
@@ -188,6 +191,11 @@ def _promote_trainable_tree(module: torch.nn.Module) -> torch.nn.Module:
         if getattr(sub, "trainable", True) is False:
             for p in sub.parameters(recurse=True):
                 p.requires_grad_(False)
+    # A ``/fix`` radial basis keeps its frequencies or centres, as in the pt
+    # backend.
+    for sub in module.modules():
+        if type(sub).__name__ == "RadialBasis" and parse_basis_type(sub.basis_type)[1]:
+            sub.adam_freqs.requires_grad_(False)
     return module
 
 
@@ -199,6 +207,19 @@ def _promote_trainable_tree(module: torch.nn.Module) -> torch.nn.Module:
 class DescrptDPA4(DescrptDPA4DP):
     _update_sel_cls = UpdateSel
 
+    def adam_route_patterns(self) -> list[str]:
+        """
+        Name patterns, relative to the descriptor, of the tensors that take the
+        AdamW path under HybridMuon: the first layer of the radial embedding and
+        the radial projection of the environment seed, which read the radial
+        basis and whose rows for rarely visited separations receive almost no
+        gradient.
+        """
+        return [
+            "radial_embedding.net.0.",
+            "env_seed_embedding.rbf_proj_layer1.",
+        ]
+
     def __init__(self, *args: Any, **kwargs: Any) -> None:
         super().__init__(*args, **kwargs)
         # The fused convolution paths consume only the three structural rows of
diff --git a/deepmd/pt_expt/descriptor/dpa4c.py b/deepmd/pt_expt/descriptor/dpa4c.py
index bc5fd1e324..5eb3ef3020 100644
--- a/deepmd/pt_expt/descriptor/dpa4c.py
+++ b/deepmd/pt_expt/descriptor/dpa4c.py
@@ -15,6 +15,9 @@
 
 import torch
 
+from deepmd.dpmodel.descriptor.dpa4_nn.radial import (
+    parse_basis_type,
+)
 from deepmd.dpmodel.descriptor.dpa4c import DescrptDPA4C as DescrptDPA4CDP
 from deepmd.pt_expt.common import (
     torch_module,
@@ -79,6 +82,15 @@ def _promote_trainable_tree(module: torch.nn.Module) -> torch.nn.Module:
         if not getattr(submodule, "trainable", True):
             for parameter in submodule.parameters(recurse=True):
                 parameter.requires_grad_(False)
+    # A ``/fix`` radial basis keeps its frequencies or centres, as in the pt
+    # backend's DPA4.  The parameter keeps its name, so checkpoints of either
+    # form load under the other.
+    for submodule in module.modules():
+        if (
+            type(submodule).__name__ == "RadialBasis"
+            and parse_basis_type(submodule.basis_type)[1]
+        ):
+            submodule.adam_freqs.requires_grad_(False)
     return module
 
 
@@ -103,6 +115,15 @@ class DescrptDPA4C(DescrptDPA4CDP):
 
     _update_sel_cls = UpdateSel
 
+    def adam_route_patterns(self) -> list[str]:
+        """
+        Name patterns, relative to the descriptor, of the tensors that take the
+        AdamW path under HybridMuon: the first layer of the radial embedding,
+        which reads the radial basis and whose rows for rarely visited
+        separations receive almost no gradient.
+        """
+        return ["radial_embedding.layers.0."]
+
     #: Artifacts whose element type is not the ``float32`` the kernel consumes.
     _COMPRESSION_BUFFER_DTYPES: ClassVar[dict[str, torch.dtype]] = {
         "info": torch.float64,
diff --git a/deepmd/pt_expt/kernels/cuda/dpa4/edge_radial.py b/deepmd/pt_expt/kernels/cuda/dpa4/edge_radial.py
index b77338e4b4..1c2f1bdffa 100644
--- a/deepmd/pt_expt/kernels/cuda/dpa4/edge_radial.py
+++ b/deepmd/pt_expt/kernels/cuda/dpa4/edge_radial.py
@@ -8,6 +8,8 @@
     env[e] = keep[e] * E_p1(r)
     rbf[e, n] = keep[e] * phi_n(r) * E_p2(r)
 
+An empty basis-envelope series selects the raw basis, without ``E_p2``.
+
 Written as tensor operations this chain is cheap enough that the compiler
 inlines it into every consumer of ``env`` and ``rbf`` and re-evaluates it there,
 so a 96 MB pass is paid several times over. Behind an operator boundary it runs
@@ -62,8 +64,10 @@ def series_coefficients(exponent: int) -> tuple[float, ...]:
 
 
 def supported(exponent_env: int, exponent_rbf: int) -> bool:
-    """Whether both envelope orders fit the staged series limit."""
-    return 2 <= exponent_env <= _MAX_SERIES and 2 <= exponent_rbf <= _MAX_SERIES
+    """Whether the envelope orders fit the staged limit, including a bare basis."""
+    return 2 <= exponent_env <= _MAX_SERIES and (
+        exponent_rbf == 0 or 2 <= exponent_rbf <= _MAX_SERIES
+    )
 
 
 def _forward_fake(
@@ -166,7 +170,8 @@ def edge_radial(
     env_series : torch.Tensor
         Horner coefficients of the edge envelope with shape (p1,).
     rbf_series : torch.Tensor
-        Horner coefficients of the basis envelope with shape (p2,).
+        Horner coefficients of the basis envelope with shape (p2,). An empty
+        tensor disables the basis envelope.
     rcut : float
         Cutoff radius in Å.
     gaussian_coeff : float
@@ -198,9 +203,9 @@ def __init__(self, envelope: Any, basis: Any) -> None:
         self._envelope = envelope
         self._basis = basis
         self._rcut = float(envelope.rcut)
-        self._basis_type = BESSEL if basis.basis_type == "bessel" else GAUSSIAN
+        self._basis_type = BESSEL if basis.basis_family == "bessel" else GAUSSIAN
         self._env = series_coefficients(envelope.p)
-        self._rbf = series_coefficients(basis.envelope.p)
+        self._rbf = series_coefficients(basis.exponent)
         self._series: tuple[torch.Tensor, torch.Tensor] | None = None
 
     def series(self, device: torch.device) -> tuple[torch.Tensor, torch.Tensor]:
@@ -260,8 +265,8 @@ def make_cuda_edge_radial(envelope: Any, basis: Any) -> EdgeRadialCuda | None:
         return None
     if basis.adam_freqs.dtype is not torch.float32:
         return None
-    if not supported(int(envelope.p), int(basis.envelope.p)):
+    if not supported(int(envelope.p), int(basis.exponent)):
         return None
-    if basis.basis_type not in ("bessel", "gaussian"):
+    if basis.basis_family not in ("bessel", "gaussian"):
         return None
     return EdgeRadialCuda(envelope, basis)
diff --git a/deepmd/pt_expt/train/training.py b/deepmd/pt_expt/train/training.py
index aecfc56b6a..062971c6bd 100644
--- a/deepmd/pt_expt/train/training.py
+++ b/deepmd/pt_expt/train/training.py
@@ -76,6 +76,9 @@
 from deepmd.pt.optimizer import (
     HybridMuonOptimizer,
 )
+from deepmd.pt.optimizer.hybrid_muon import (
+    adam_route_patterns,
+)
 from deepmd.pt.utils.compile_compat import (
     apply_global_compile_patches,
     build_inductor_compile_options,
@@ -2255,6 +2258,7 @@ def update_finetune_bias(
                 weight_decay=weight_decay,
             )
         else:
+            adam_patterns = adam_route_patterns(self.models.values())
             self.optimizer = self._create_optimizer(
                 HybridMuonOptimizer,
                 lr=initial_lr,
@@ -2267,6 +2271,7 @@ def update_finetune_bias(
                 enable_gram=bool(optimizer_params["enable_gram"]),
                 flash_muon=bool(optimizer_params["flash_muon"]),
                 magma_muon=bool(optimizer_params["magma_muon"]),
+                adam_patterns=adam_patterns,
                 # Sharded parameters are DTensors, and several torch._foreach_*
                 # ops lack sharding propagation, so the per-tensor path applies.
                 use_foreach=False if self.sharding.shards_parameters else None,
diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py
index 65ab8a85a2..cd781c4f84 100644
--- a/deepmd/utils/argcheck.py
+++ b/deepmd/utils/argcheck.py
@@ -608,13 +608,19 @@ def descrpt_se_zm_args() -> list[Argument]:
     - `str`: "auto" or "auto:factor" sizes `sel` from the training data via neighbor statistics (`factor` larger than 1, rounded up to a multiple of 4; "auto" equals "auto:1.1"). This requires the neighbor-statistics pass and is therefore unavailable under `--skip-neighbor-stat`.'
     doc_rcut = "The cut-off radius."
     doc_env_exp = (
-        "C^3 cutoff envelope exponents `[rbf_env_exp, edge_env_exp]`. "
-        "`rbf_env_exp` controls radial basis function envelope decay; "
-        "`edge_env_exp` controls message passing edge weight envelope decay. "
+        "C^3 cutoff envelope exponents. A list `[rbf_env_exp, edge_env_exp]` "
+        "specifies the radial-basis and message-passing envelopes separately. "
+        "A zero radial-basis exponent disables that envelope. "
+        "An integer specifies only the message-passing envelope exponent and "
+        "disables the radial-basis envelope. "
         "Larger values give weaker suppression."
     )
     doc_channels = "Total channels per (l,m) coefficient."
-    doc_basis_type = "Radial basis type. Supported values are `bessel` and `gaussian`."
+    doc_basis_type = (
+        "Radial basis type. Supported values are `bessel`, `gaussian`, `bessel/fix` and `gaussian/fix`. "
+        "The `/fix` forms keep the Bessel frequencies or Gaussian centres at their initial values instead of training them, "
+        "so that separations no training frame constrains cannot move them."
+    )
     doc_n_radial = "Number of radial basis functions."
     doc_radial_mlp = "Hidden layer sizes for radial networks. An output layer of size (l_schedule[0]+extra_node_l+1)*channels will be automatically appended. Use 0 as a placeholder to be replaced by channels."
     doc_edge_norm = "Channel RMSNorm on the cutoff-vanishing feature branches. A bool switches every site together: `false` removes the RMSNorm from the radial-network hidden layers, the environment-seed FiLM scale/shift logits and the cross-focus competition scalars, and uses unit-floor residual scaling for post-SO(2) messages. A list of three bools `[radial, film, focus]` switches the sites individually; the post-SO(2) treatment follows the first (radial) entry. Recommended: `[false, true, false]` — the radial-site norms amplify noise where the radial features vanish at the cutoff and produce a spurious long-range force step, while the FiLM and focus norms are safe to keep."
@@ -919,7 +925,7 @@ def descrpt_se_zm_args() -> list[Argument]:
         Argument("rcut", float, optional=True, default=6.0, doc=doc_rcut),
         Argument(
             "env_exp",
-            list[int],
+            [int, list[int]],
             optional=True,
             default=[7, 5],
             doc=doc_env_exp,
diff --git a/deepmd/utils/model_preset.py b/deepmd/utils/model_preset.py
index d00d791b2f..7fae9a2714 100644
--- a/deepmd/utils/model_preset.py
+++ b/deepmd/utils/model_preset.py
@@ -210,6 +210,17 @@
         },
         "grades": ("nano", "mini", "neo", "air", "plus", "pro", "max", "ultra"),
     },
+    # One C^3 envelope on the messages with no envelope on the radial basis,
+    # and fixed Gaussian centres in place of trainable Bessel frequencies.
+    "v20260911": {
+        "descriptor": {
+            "edge_norm": [False, True, True],
+            "sandwich_norm": [True, False, True, False],
+            "env_exp": 5,
+            "basis_type": "gaussian/fix",
+        },
+        "grades": ("nano", "mini", "neo", "air", "plus", "pro", "max", "ultra"),
+    },
 }
 
 # === DPA4C ===
@@ -249,6 +260,11 @@
 }
 _DPA4C_VERSIONS: dict[str, dict[str, Any]] = {
     "v20260901": {"grades": ("nano", "mini", "neo", "air", "plus")},
+    # Fixed Gaussian centres in place of trainable Bessel frequencies.
+    "v20260911": {
+        "descriptor": {"basis_type": "gaussian/fix"},
+        "grades": ("nano", "mini", "neo", "air", "plus"),
+    },
 }
 
 
diff --git a/doc/model/dpa4.md b/doc/model/dpa4.md
index c3fb7fe54d..e67e06831b 100644
--- a/doc/model/dpa4.md
+++ b/doc/model/dpa4.md
@@ -138,10 +138,16 @@ release a preset reproduces: a later release with different settings gets a new
 version, and existing presets are never changed. The available DPA4 presets
 are, in ascending cost:
 
-- `v20260901`, the current release grades: `dpa4-nano-v20260901`,
-  `dpa4-mini-v20260901`, `dpa4-neo-v20260901`, `dpa4-air-v20260901`,
-  `dpa4-plus-v20260901`, `dpa4-pro-v20260901`, `dpa4-max-v20260901` and
-  `dpa4-ultra-v20260901`.
+- `v20260911`, the current release grades: `dpa4-nano-v20260911`,
+  `dpa4-mini-v20260911`, `dpa4-neo-v20260911`, `dpa4-air-v20260911`,
+  `dpa4-plus-v20260911`, `dpa4-pro-v20260911`, `dpa4-max-v20260911` and
+  `dpa4-ultra-v20260911`. They expand the radial basis on fixed Gaussian
+  centres (`basis_type` `gaussian/fix`) and apply a single cutoff envelope to
+  the message-passing edge weights (`env_exp` 5).
+- `v20260901`, the previous release grades with trainable Bessel functions and
+  two envelopes: `dpa4-nano-v20260901`, `dpa4-mini-v20260901`,
+  `dpa4-neo-v20260901`, `dpa4-air-v20260901`, `dpa4-plus-v20260901`,
+  `dpa4-pro-v20260901`, `dpa4-max-v20260901` and `dpa4-ultra-v20260901`.
 - `v20260820`, the earlier baseline grades: `dpa4-nano-v20260820`,
   `dpa4-mini-v20260820`, `dpa4-neo-v20260820`, `dpa4-air-v20260820`,
   `dpa4-plus-v20260820` and `dpa4-pro-v20260820`.
@@ -912,12 +918,18 @@ only the `l = 0` scalar channels are read out and passed to the fitting network:
 
 ### Radial basis and smooth cutoff
 
-Every edge uses a radial basis (`basis_type`, with `n_radial` functions)
-multiplied by a smooth envelope whose value and first three derivatives vanish
-at `rcut`. This smoothness matters for MD because nonsmooth descriptor cutoffs
-would be inherited by the force derivatives. The two `env_exp` exponents control
-the radial-basis envelope and the message-passing edge weights respectively;
-larger values keep an envelope closer to one for more of the cutoff range.
+Every edge uses a radial basis (`basis_type`, with `n_radial` functions):
+Bessel functions (`bessel`) or Gaussians (`gaussian`), whose frequencies or
+centres are trained by default. The `bessel/fix` and `gaussian/fix` forms keep
+them at their initial values, so that separations no training frame constrains
+cannot move them. The message-passing edge weights carry a smooth envelope
+whose value and first three derivatives vanish at `rcut`. This smoothness
+matters for MD because nonsmooth descriptor cutoffs would be inherited by the
+force derivatives. `env_exp` written as one integer sets the exponent of that
+envelope and leaves the radial basis bare; written as a list
+`[rbf_env_exp, edge_env_exp]` it applies a second envelope to the radial basis
+itself. Larger values keep an envelope closer to one for more of the cutoff
+range.
 
 ### Attention and focus streams
 
diff --git a/doc/model/dpa4c.md b/doc/model/dpa4c.md
index 446b2386ec..bd3ab45d5b 100644
--- a/doc/model/dpa4c.md
+++ b/doc/model/dpa4c.md
@@ -121,7 +121,9 @@ carry the accuracy–cost trade-off:
 - **Radial basis** —
   {ref}`basis_type <model[standard]/descriptor[dpa4c]/basis_type>` and
   {ref}`n_radial <model[standard]/descriptor[dpa4c]/n_radial>` select the
-  analytic basis that feeds the radial network.
+  analytic basis that feeds the radial network; the `bessel/fix` and
+  `gaussian/fix` forms keep the frequencies or centres at their initial
+  values instead of training them.
 
 > [!IMPORTANT]
 > The compressed CUDA path is compiled for `channels` in `{8, 16, 32, 64, 128}`,
@@ -135,8 +137,10 @@ carry the accuracy–cost trade-off:
 The released grades, Nano, Mini, Neo, Air and Plus in ascending cost, pair each
 descriptor width with a fitting width sized against it. They are good starting
 points; `Neo` is the general-purpose default. Each grade is available as a named
-model preset, `dpa4c-nano-v20260901`, `dpa4c-mini-v20260901`,
-`dpa4c-neo-v20260901`, `dpa4c-air-v20260901` and `dpa4c-plus-v20260901`:
+model preset: `dpa4c-nano-v20260911`, `dpa4c-mini-v20260911`,
+`dpa4c-neo-v20260911`, `dpa4c-air-v20260911` and `dpa4c-plus-v20260911` expand
+the radial basis on fixed Gaussian centres (`basis_type` `gaussian/fix`), and
+the `v20260901` presets of the same grades keep the trainable Bessel basis:
 setting `model.preset` fills in `type_map` (all 118 elements), `descriptor` and
 `fitting_net` from the release configuration, and entries written next to the
 preset take precedence, as a whole for `type_map` and key by key inside
diff --git a/source/op/pt/dpa4/edge_radial.cu b/source/op/pt/dpa4/edge_radial.cu
index e75ad9e666..0300f61061 100644
--- a/source/op/pt/dpa4/edge_radial.cu
+++ b/source/op/pt/dpa4/edge_radial.cu
@@ -8,6 +8,8 @@
 //   env[e]     = keep[e] * E_p1(r)
 //   rbf[e, n]  = keep[e] * phi_n(r) * E_p2(r)
 //
+// An empty basis-envelope series selects the raw basis without E_p2.
+//
 // with the C3 cutoff envelope written in its cancellation-free factorization
 //
 //   u = clamp((rcut - r) / rcut, 0, 1),  x = 1 - u,  E_p(r) = u^4 * S_p(x)
@@ -45,8 +47,9 @@ constexpr int kMaxSeries = 16;
 /// Basis families with an implementation.
 enum BasisType : int { kBessel = 0, kGaussian = 1 };
 
-/// The C3 envelope and its derivative with respect to the distance.
+/// The optional C3 envelope and its derivative with respect to the distance.
 ///
+/// An empty series denotes the identity factor with zero derivative.
 /// ``u`` saturates outside the cutoff, where both the value and the derivative
 /// are identically zero, which is what makes the potential energy surface C3
 /// continuous at ``rcut``.
@@ -56,6 +59,11 @@ __device__ __forceinline__ void envelope_pair(float r,
                                               int order,
                                               float& value,
                                               float& derivative) {
+  if (order == 0) {
+    value = 1.f;
+    derivative = 0.f;
+    return;
+  }
   const float u = fminf(fmaxf((1.f - r * inv_rcut), 0.f), 1.f);
   const float x = 1.f - u;
   float s = series[order - 1];
@@ -211,8 +219,13 @@ void check_inputs(const torch::Tensor& edge_len,
   TORCH_CHECK(
       env_series.numel() <= kMaxSeries && rbf_series.numel() <= kMaxSeries,
       "dpa4_edge_radial: envelope order beyond the staged limit");
-  TORCH_CHECK(env_series.numel() >= 2 && rbf_series.numel() >= 2,
-              "dpa4_edge_radial: the envelope series needs at least two terms");
+  TORCH_CHECK(
+      env_series.numel() >= 2,
+      "dpa4_edge_radial: the edge envelope series needs at least two terms");
+  TORCH_CHECK(
+      rbf_series.numel() == 0 || rbf_series.numel() >= 2,
+      "dpa4_edge_radial: the basis envelope series must be empty or have "
+      "at least two terms");
   TORCH_CHECK(freqs.numel() > 0,
               "dpa4_edge_radial: the basis must be non-empty");
 }
diff --git a/source/tests/common/test_model_preset.py b/source/tests/common/test_model_preset.py
index 7bd3d6e310..e96a9820c7 100644
--- a/source/tests/common/test_model_preset.py
+++ b/source/tests/common/test_model_preset.py
@@ -168,6 +168,28 @@ def test_dpa4_versions_differ_only_in_normalization_options() -> None:
         assert f"dpa4-{grade}-v20260901" in MODEL_PRESETS
 
 
+def test_v20260911_fixes_the_basis_and_uses_one_envelope() -> None:
+    for grade in ("nano", "mini", "neo", "air", "plus", "pro", "max", "ultra"):
+        old = get_model_preset(f"dpa4-{grade}-v20260901")
+        new = get_model_preset(f"dpa4-{grade}-v20260911")
+        assert old["type_map"] == new["type_map"]
+        assert old["fitting_net"] == new["fitting_net"]
+        assert {
+            key: new["descriptor"][key]
+            for key in set(old["descriptor"]) | set(new["descriptor"])
+            if old["descriptor"].get(key) != new["descriptor"].get(key)
+        } == {"env_exp": 5, "basis_type": "gaussian/fix"}
+    for grade in ("nano", "mini", "neo", "air", "plus"):
+        old = get_model_preset(f"dpa4c-{grade}-v20260901")
+        new = get_model_preset(f"dpa4c-{grade}-v20260911")
+        assert old["fitting_net"] == new["fitting_net"]
+        assert {
+            key: new["descriptor"][key]
+            for key in set(old["descriptor"]) | set(new["descriptor"])
+            if old["descriptor"].get(key) != new["descriptor"].get(key)
+        } == {"basis_type": "gaussian/fix"}
+
+
 def test_presets_carry_no_runtime_options() -> None:
     for name, preset in MODEL_PRESETS.items():
         for region in ("descriptor", "fitting_net"):
diff --git a/source/tests/consistent/descriptor/test_dpa4.py b/source/tests/consistent/descriptor/test_dpa4.py
index 38666106ec..929bcfec24 100644
--- a/source/tests/consistent/descriptor/test_dpa4.py
+++ b/source/tests/consistent/descriptor/test_dpa4.py
@@ -99,6 +99,8 @@ def dpa4_case(**overrides: Any) -> tuple:
     dpa4_case(s2_activation=[False, False]),
     # gaussian radial basis
     dpa4_case(basis_type="gaussian"),
+    # the /fix suffix is accepted by every backend and evaluates like its family
+    dpa4_case(basis_type="gaussian/fix"),
     # float32 baseline
     dpa4_case(precision="float32"),
     # float32 mixed high-risk path
diff --git a/source/tests/pt/model/test_descriptor_sezm.py b/source/tests/pt/model/test_descriptor_sezm.py
index 4d069cd75f..0a98f7e664 100644
--- a/source/tests/pt/model/test_descriptor_sezm.py
+++ b/source/tests/pt/model/test_descriptor_sezm.py
@@ -646,6 +646,49 @@ def test_forward_with_descriptor_variants(self) -> None:
             with self.subTest(mode=name):
                 self._assert_forward_backward_smoke(**model_kwargs)
 
+    def test_fixed_basis_types_freeze_the_basis(self) -> None:
+        """``/fix`` basis types keep the basis parameters out of training only."""
+        for family in ("bessel", "gaussian"):
+            with self.subTest(family=family):
+                torch.manual_seed(7)
+                free = self._assert_forward_backward_smoke(
+                    **_descriptor_kwargs(channels=4, basis_type=family)
+                )
+                torch.manual_seed(7)
+                fixed = self._assert_forward_backward_smoke(
+                    **_descriptor_kwargs(channels=4, basis_type=f"{family}/fix")
+                )
+                self.assertEqual(fixed.radial_basis.basis_type, f"{family}/fix")
+                self.assertEqual(fixed.radial_basis.basis_family, family)
+                self.assertFalse(fixed.radial_basis.adam_freqs.requires_grad)
+                self.assertTrue(free.radial_basis.adam_freqs.requires_grad)
+                torch.testing.assert_close(
+                    fixed.radial_basis.adam_freqs, free.radial_basis.adam_freqs
+                )
+                # Every other parameter is unaffected by the suffix.
+                frozen = [
+                    name for name, p in fixed.named_parameters() if not p.requires_grad
+                ]
+                self.assertEqual(frozen, ["radial_basis.adam_freqs"])
+                # The same weights give the same descriptor, and the suffix
+                # survives a serialization round trip.
+                fixed.load_state_dict(free.state_dict())
+                coord, atype, nlist = _tiny_two_atom_system(
+                    self.device, dtype=torch.float32
+                )
+                extended_coord = coord.reshape(1, -1)
+                torch.testing.assert_close(
+                    fixed(extended_coord, atype, nlist, mapping=None, comm_dict=None)[
+                        0
+                    ],
+                    free(extended_coord, atype, nlist, mapping=None, comm_dict=None)[0],
+                )
+                self.assertEqual(
+                    fixed.serialize()["config"]["basis_type"], f"{family}/fix"
+                )
+        with self.assertRaises(ValueError):
+            DescrptSeZM(**_descriptor_kwargs(channels=4, basis_type="gaussian-frozen"))
+
     def test_forward_with_attention_variants(self) -> None:
         """Test forward/backward smoke paths for attention-based variants."""
         cases = {
diff --git a/source/tests/pt/model/test_dpa4_dpmodel_parity.py b/source/tests/pt/model/test_dpa4_dpmodel_parity.py
index 74f703fbc6..60c80ccbde 100644
--- a/source/tests/pt/model/test_dpa4_dpmodel_parity.py
+++ b/source/tests/pt/model/test_dpa4_dpmodel_parity.py
@@ -397,8 +397,8 @@ def test_envelope(self, exponent) -> None:
         np.testing.assert_array_equal(np.asarray(res)[r[:, 0] >= self.rcut], 0.0)
 
     @pytest.mark.parametrize("basis_type", ["bessel", "gaussian"])  # both bases
-    @pytest.mark.parametrize("exponent", [5, 7])  # envelope exponent
-    def test_radial_basis(self, basis_type, exponent) -> None:
+    @pytest.mark.parametrize("exponent", [0, 5, 7])
+    def test_radial_basis(self, basis_type: str, exponent: int) -> None:
         from deepmd.dpmodel.descriptor.dpa4_nn.radial import (
             RadialBasis as DPRadialBasis,
         )
@@ -423,11 +423,16 @@ def test_radial_basis(self, basis_type, exponent) -> None:
         # pt state_dict key contract: only the trainable frequencies
         assert list(serialized["@variables"]) == ["adam_freqs"]
         dp_mod = DPRadialBasis.deserialize(serialized)
+        assert dp_mod.exponent == exponent
+        assert (dp_mod.envelope is None) is (exponent == 0)
         r = self._r_grid()
         assert_parity(dp_mod.call(r), pt_mod(to_pt(r)))
 
     @pytest.mark.parametrize("trainable", [True, False])
-    def test_radial_basis_roundtrip_preserves_trainable(self, trainable: bool) -> None:
+    @pytest.mark.parametrize("exponent", [0, 7])
+    def test_radial_basis_roundtrip_preserves_trainable(
+        self, trainable: bool, exponent: int
+    ) -> None:
         from deepmd.pt.model.descriptor.sezm_nn.radial import (
             RadialBasis as PTRadialBasis,
         )
@@ -437,11 +442,14 @@ def test_radial_basis_roundtrip_preserves_trainable(self, trainable: bool) -> No
             n_radial=8,
             dtype=torch.float64,
             trainable=trainable,
+            exponent=exponent,
         )
         restored = PTRadialBasis.deserialize(radial_basis.serialize())
 
         assert restored.trainable is trainable
         assert restored.adam_freqs.requires_grad is trainable
+        assert restored.exponent == exponent
+        assert (restored.envelope is None) is (exponent == 0)
 
     def test_radial_basis_deserializes_version_one_without_trainable(self) -> None:
         from deepmd.pt.model.descriptor.sezm_nn.radial import (
@@ -460,10 +468,12 @@ def test_radial_basis_deserializes_version_one_without_trainable(self) -> None:
 
         assert restored.trainable is True
         assert restored.adam_freqs.requires_grad is True
+        assert restored.exponent == 7
+        assert restored.envelope is not None
 
     @pytest.mark.parametrize("basis_type", ["bessel", "gaussian"])  # both bases
-    @pytest.mark.parametrize("apply_envelope", [True, False])  # both envelope modes
-    def test_radial_basis_roundtrip(self, basis_type, apply_envelope) -> None:
+    @pytest.mark.parametrize("exponent", [0, 7])
+    def test_radial_basis_roundtrip(self, basis_type: str, exponent: int) -> None:
         from deepmd.dpmodel.descriptor.dpa4_nn.radial import (
             RadialBasis as DPRadialBasis,
         )
@@ -473,11 +483,11 @@ def test_radial_basis_roundtrip(self, basis_type, apply_envelope) -> None:
             basis_type=basis_type,
             n_radial=12,
             precision="float64",
-            exponent=7,
-            apply_envelope=apply_envelope,
+            exponent=exponent,
         )
         dp_mod2 = DPRadialBasis.deserialize(dp_mod.serialize())
-        assert dp_mod2.apply_envelope is apply_envelope
+        assert dp_mod2.exponent == exponent
+        assert (dp_mod2.envelope is None) is (exponent == 0)
         r = self._r_grid()
         np.testing.assert_array_equal(
             np.asarray(dp_mod.call(r)), np.asarray(dp_mod2.call(r))
@@ -489,17 +499,17 @@ def test_radial_basis_envelope_modes(self, basis_type) -> None:
             RadialBasis as DPRadialBasis,
         )
 
-        def make(apply_envelope: bool) -> DPRadialBasis:
+        def make(exponent: int) -> DPRadialBasis:
             return DPRadialBasis(
                 rcut=self.rcut,
                 basis_type=basis_type,
                 n_radial=12,
                 precision="float64",
-                exponent=7,
-                apply_envelope=apply_envelope,
+                exponent=exponent,
             )
 
-        enveloped, raw = make(True), make(False)
+        enveloped, raw = make(7), make(0)
+        assert raw.envelope is None
         r = self._r_grid()
         enveloped_value = np.asarray(enveloped.call(r))
         raw_value = np.asarray(raw.call(r))
@@ -3930,6 +3940,26 @@ def test_descriptor(self, use_env_seed, n_blocks) -> None:
         )
         self._assert_descr_parity(pt_mod, dp_mod)
 
+    @pytest.mark.parametrize("basis_type", ["bessel", "gaussian"])
+    @pytest.mark.parametrize("env_exp", [5, [0, 5], [7, 5]])
+    def test_descriptor_env_exp(
+        self, basis_type: str, env_exp: int | list[int]
+    ) -> None:
+        pt_mod, dp_mod, _ = self._build_descr_pair(
+            basis_type=basis_type, env_exp=env_exp
+        )
+        for module in (pt_mod, dp_mod):
+            assert module.env_exp == env_exp
+            exponent = env_exp[0] if isinstance(env_exp, list) else 0
+            assert module.radial_basis.exponent == exponent
+            assert (module.radial_basis.envelope is None) is (exponent == 0)
+            assert module.edge_envelope.p == 5
+        self._assert_descr_parity(pt_mod, dp_mod)
+        pt_restored = type(pt_mod).deserialize(pt_mod.serialize())
+        dp_restored = type(dp_mod).deserialize(dp_mod.serialize())
+        assert pt_restored.env_exp == dp_restored.env_exp == env_exp
+        self._assert_descr_parity(pt_restored, dp_restored)
+
     @pytest.mark.parametrize(
         "edge_norm", [False, True, [False, True, False]]
     )  # cutoff-vanishing normalization modes: all off, all on, per-site
diff --git a/source/tests/pt/model/test_dpa4_ptexpt_grad_parity.py b/source/tests/pt/model/test_dpa4_ptexpt_grad_parity.py
index e836bcf8b7..b209169faf 100644
--- a/source/tests/pt/model/test_dpa4_ptexpt_grad_parity.py
+++ b/source/tests/pt/model/test_dpa4_ptexpt_grad_parity.py
@@ -165,6 +165,11 @@ def _inputs(self, seed=2151):
             # ReducedEquivariantRMSNorm is reachable only through so2_norm and
             # is the one module that sizes its forward off a stored index array
             pytest.param({"so2_norm": True}, id="so2_norm"),
+            pytest.param({"env_exp": 5}, id="single_envelope_bessel"),
+            pytest.param(
+                {"env_exp": 5, "basis_type": "gaussian"},
+                id="single_envelope_gaussian",
+            ),
         ],
     )
     def test_descriptor_grad_parity(self, overrides) -> None:
@@ -200,6 +205,34 @@ def test_descriptor_grad_parity(self, overrides) -> None:
         # math, where fp64 accumulation-order drift reaches ~3e-11 rel
         _assert_grad_trees_match(pt_mod, expt_mod, rtol=1e-10, atol=1e-12)
 
+    @pytest.mark.parametrize("basis_type", ["bessel", "gaussian"])
+    def test_single_envelope_coordinate_gradient_loss(self, basis_type: str) -> None:
+        """Differentiate a coordinate-gradient loss through both raw-basis backends."""
+        pt_mod, expt_mod = self._build_pair(env_exp=5, basis_type=basis_type)
+        inp = self._inputs()
+        outputs, derivatives = [], []
+        for module in (pt_mod, expt_mod):
+            coord = to_pt(inp["coord"]).reshape(self.nf, -1).requires_grad_(True)
+            output = module(
+                coord,
+                to_pt(inp["atype_ext"]),
+                to_pt(inp["nlist"]),
+                mapping=to_pt(inp["mapping"]),
+            )[0]
+            # Mean reductions keep the loss scale independent of the number
+            # of descriptor and coordinate components.
+            energy = output.square().mean()
+            derivative = torch.autograd.grad(energy, coord, create_graph=True)[0]
+            assert derivative.abs().max().item() > 1e-10
+            (energy + derivative.square().mean()).backward()
+            outputs.append(output.detach().cpu().numpy())
+            derivatives.append(derivative.detach().cpu().numpy())
+        np.testing.assert_allclose(outputs[0], outputs[1], rtol=1e-10, atol=1e-12)
+        np.testing.assert_allclose(
+            derivatives[0], derivatives[1], rtol=1e-10, atol=1e-12
+        )
+        _assert_grad_trees_match(pt_mod, expt_mod, rtol=1e-10, atol=1e-12)
+
     def test_descriptor_grad_parity_native_spin(self) -> None:
         # Native per-atom spin (``use_spin``) adds trainable Parameters that
         # pt_expt must promote from dpmodel numpy->buffer:
diff --git a/source/tests/pt/model/test_sezm_model.py b/source/tests/pt/model/test_sezm_model.py
index 6d979223c1..73f9baccdf 100644
--- a/source/tests/pt/model/test_sezm_model.py
+++ b/source/tests/pt/model/test_sezm_model.py
@@ -3467,3 +3467,31 @@ def test_forward_and_backward_match_eager(self) -> None:
                 rtol=force_grad_rtol,
                 msg=f"force-grad-sq mismatch at {name}",
             )
+
+
+class TestSeZMModelAdamRouting(unittest.TestCase):
+    """The model composes the AdamW routing patterns its descriptor declares."""
+
+    def test_adam_route_patterns_match_parameters(self) -> None:
+        """Every declared AdamW pattern names existing matrices; the base declares none."""
+        model = get_model(_build_lora_sezm_model_params())
+        patterns = model.adam_route_patterns()
+        self.assertEqual(len(patterns), 2)
+        names = [name for name, _ in model.named_parameters()]
+        for pattern in patterns:
+            matched = [n for n in names if pattern in n]
+            self.assertTrue(matched, pattern)
+            self.assertTrue(
+                any(dict(model.named_parameters())[n].dim() >= 2 for n in matched),
+                pattern,
+            )
+        from types import (
+            SimpleNamespace,
+        )
+
+        from deepmd.pt.model.model.model import (
+            BaseModel,
+        )
+
+        plain = SimpleNamespace(atomic_model=SimpleNamespace(descriptor=object()))
+        self.assertEqual(BaseModel.adam_route_patterns(plain), [])
diff --git a/source/tests/pt/model/test_sezm_spin_model.py b/source/tests/pt/model/test_sezm_spin_model.py
index c98c80f30c..a86baa7464 100644
--- a/source/tests/pt/model/test_sezm_spin_model.py
+++ b/source/tests/pt/model/test_sezm_spin_model.py
@@ -203,6 +203,21 @@ def _build_model_params(
             "use_compile": use_compile,
         }
 
+    def test_adam_route_patterns_match_parameters(self) -> None:
+        """The native-spin SeZM model routes the same tensors as the plain model."""
+        model = get_model(self._build_model_params()).to(self.device)
+        patterns = sorted(model.adam_route_patterns())
+        self.assertEqual(
+            patterns,
+            [
+                "descriptor.env_seed_embedding.rbf_proj_layer1.",
+                "descriptor.radial_embedding.net.0.",
+            ],
+        )
+        names = [name for name, _ in model.named_parameters()]
+        for pattern in patterns:
+            self.assertTrue(any(pattern in name for name in names), pattern)
+
     def test_factory_shapes_and_masks(self) -> None:
         """Factory should build SeZMSpinModel with public real-type metadata."""
         model = get_model(self._build_model_params()).to(self.device)
diff --git a/source/tests/pt/test_hybrid_muon.py b/source/tests/pt/test_hybrid_muon.py
index c143337a43..7cd4c790a9 100644
--- a/source/tests/pt/test_hybrid_muon.py
+++ b/source/tests/pt/test_hybrid_muon.py
@@ -311,6 +311,44 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         self.assertIn("momentum_buffer", optimizer.state[model.bias_proj.weight])
         self.assertNotIn("exp_avg", optimizer.state[model.bias_proj.weight])
 
+    def test_adam_patterns_route_matrices_to_adamw(self) -> None:
+        """Matrices whose name contains an ``adam_patterns`` entry leave Muon."""
+        torch.manual_seed(42)
+
+        class ToyModel(torch.nn.Module):
+            def __init__(self, device: torch.device) -> None:
+                super().__init__()
+                self.descriptor = torch.nn.Linear(4, 6, bias=False, device=device)
+                self.fitting_net = torch.nn.Linear(6, 3, device=device)
+
+            def forward(self, x: torch.Tensor) -> torch.Tensor:
+                return self.fitting_net(self.descriptor(x)).sum()
+
+        model = ToyModel(self.device)
+        optimizer = HybridMuonOptimizer(
+            model.parameters(),
+            lr=0.02,
+            weight_decay=0.01,
+            named_parameters=tuple(model.named_parameters()),
+            adam_patterns=["Fitting_Net"],
+        )
+        model(torch.randn(4, 4, device=self.device)).backward()
+        optimizer.step()
+
+        # The matching matrix takes the AdamW path, its vector plain Adam.
+        self.assertIn("exp_avg", optimizer.state[model.fitting_net.weight])
+        self.assertNotIn("momentum_buffer", optimizer.state[model.fitting_net.weight])
+        routes = optimizer._routing[0]
+        self.assertTrue(
+            any(e["param"] is model.fitting_net.weight for e in routes["adam_decay"])
+        )
+        self.assertTrue(
+            any(e["param"] is model.fitting_net.bias for e in routes["adam_no_decay"])
+        )
+        # A matrix outside the patterns stays on Muon.
+        self.assertIn("momentum_buffer", optimizer.state[model.descriptor.weight])
+        self.assertNotIn("exp_avg", optimizer.state[model.descriptor.weight])
+
     def test_2d_mode_routes_3d_weight_to_adam(self) -> None:
         """Test muon_mode='2d' routes 3D matrix weights to Adam."""
         torch.manual_seed(42)
diff --git a/source/tests/pt_expt/descriptor/test_dpa4_accelerated.py b/source/tests/pt_expt/descriptor/test_dpa4_accelerated.py
index 820a29834e..e4f7b6c37e 100644
--- a/source/tests/pt_expt/descriptor/test_dpa4_accelerated.py
+++ b/source/tests/pt_expt/descriptor/test_dpa4_accelerated.py
@@ -48,6 +48,9 @@
 from ...common.test_mixins import (
     TestCaseSingleFrameWithNlist,
 )
+from ...dpa4_fixtures import (
+    jitter_zero_arrays,
+)
 
 
 def _make_descriptor(
@@ -57,11 +60,13 @@ def _make_descriptor(
     precision: str = "float32",
     *,
     source_gated: bool = False,
+    env_exp: int | list[int] | None = None,
 ) -> DescrptDPA4:
     return DescrptDPA4(
         ntypes=ntypes,
         sel=sel,
         rcut=rcut,
+        env_exp=env_exp,
         channels=32,
         n_radial=8,
         lmax=2,
@@ -85,10 +90,12 @@ def _make_descriptor(
     ("precision", "expected_bound"),
     [("float32", True), ("float64", False)],
 )
+@pytest.mark.parametrize("env_exp", [5, [7, 5]])
 def test_fp32_only_cuda_bindings(
     monkeypatch,
     precision: str,
     expected_bound: bool,
+    env_exp: int | list[int],
 ) -> None:
     """Bind the handwritten CUDA path only for its supported precision."""
     for name in (
@@ -103,7 +110,9 @@ def test_fp32_only_cuda_bindings(
     monkeypatch.setattr(grid_pair, "op_available", lambda: True)
     monkeypatch.setattr(zonal_scatter, "op_available", lambda: True)
 
-    descriptor = _make_descriptor(2, [20], 4.0, precision=precision).eval()
+    descriptor = _make_descriptor(
+        2, [20], 4.0, precision=precision, env_exp=env_exp
+    ).eval()
     initial_embeddings = [
         module
         for module in descriptor.modules()
@@ -231,7 +240,10 @@ def test_source_gated_flash_retains_dense_rotations(self, monkeypatch) -> None:
         np.testing.assert_allclose(gradient, dense_gradient, rtol=2e-4, atol=2e-5)
 
     @pytest.mark.parametrize("backend", ["triton", "cuda", "cutile"])
-    def test_forward_and_coordinate_gradient(self, monkeypatch, backend) -> None:
+    @pytest.mark.parametrize("env_exp", [None, 5])
+    def test_forward_and_coordinate_gradient(
+        self, monkeypatch: pytest.MonkeyPatch, backend: str, env_exp: int | None
+    ) -> None:
         if backend == "triton" and not FORCE_ASSEMBLY_TRITON_AVAILABLE:
             pytest.skip("Triton is unavailable")
         if backend == "cutile" and not CUTILE_AVAILABLE:
@@ -244,7 +256,10 @@ def test_forward_and_coordinate_gradient(self, monkeypatch, backend) -> None:
             "DP_CUTE_INFER",
         ):
             monkeypatch.setenv(name, "0")
-        data = _make_descriptor(self.nt, self.sel_mix, self.rcut).serialize()
+        data = _make_descriptor(
+            self.nt, self.sel_mix, self.rcut, env_exp=env_exp
+        ).serialize()
+        data = jitter_zero_arrays(data, np.random.default_rng(73))
         reference = DescrptDPA4.deserialize(data).to(self.device).eval()
 
         levels = {
@@ -284,6 +299,7 @@ def test_forward_and_coordinate_gradient(self, monkeypatch, backend) -> None:
         coord_ref, atype, nlist = self._inputs()
         output_ref = reference(coord_ref, atype, nlist)[0]
         grad_ref = torch.autograd.grad(output_ref.sum(), coord_ref)[0]
+        assert grad_ref.abs().max().item() > 1e-10
 
         coord, atype, nlist = self._inputs()
         output = accelerated(coord, atype, nlist)[0]
diff --git a/source/tests/pt_expt/descriptor/test_dpa4c_cpu.py b/source/tests/pt_expt/descriptor/test_dpa4c_cpu.py
index a04c651f87..d7c7c0c23b 100644
--- a/source/tests/pt_expt/descriptor/test_dpa4c_cpu.py
+++ b/source/tests/pt_expt/descriptor/test_dpa4c_cpu.py
@@ -69,6 +69,45 @@ def _build_descriptor(
     ).eval()
 
 
+def test_fixed_basis_types_freeze_the_basis() -> None:
+    """``/fix`` basis types keep the DPA4C basis parameters out of training."""
+    for family in ("bessel", "gaussian"):
+        fixed = DescrptDPA4C(
+            rcut=3.0,
+            ntypes=2,
+            channels=8,
+            lmax=2,
+            n_radial=8,
+            precision="float32",
+            seed=17,
+            basis_type=f"{family}/fix",
+        )
+        assert fixed.radial_basis.basis_family == family
+        assert not fixed.radial_basis.adam_freqs.requires_grad
+        frozen = [name for name, p in fixed.named_parameters() if not p.requires_grad]
+        assert frozen == ["radial_basis.adam_freqs"]
+        assert fixed.serialize()["basis_type"] == f"{family}/fix"
+
+
+def test_adam_route_patterns_name_the_first_radial_layer() -> None:
+    """The descriptor routes the layer that reads the radial basis to AdamW."""
+    descriptor = DescrptDPA4C(
+        rcut=3.0,
+        ntypes=2,
+        channels=8,
+        lmax=2,
+        n_radial=8,
+        radial_modes=2,
+        precision="float32",
+        seed=17,
+    )
+    patterns = descriptor.adam_route_patterns()
+    assert patterns == ["radial_embedding.layers.0."]
+    names = [name for name, _ in descriptor.named_parameters()]
+    routed = [name for name in names if any(pattern in name for pattern in patterns)]
+    assert routed == ["radial_embedding.layers.0.w"]
+
+
 def _build_graph(descriptor: DescrptDPA4C, canonical: bool, node_count: int = 24):
     generator = torch.Generator().manual_seed(23)
     coordinate = 5.0 * torch.rand(
diff --git a/source/tests/pt_expt/model/test_dpa4_export.py b/source/tests/pt_expt/model/test_dpa4_export.py
index 7fdc206ce2..a164c5080f 100644
--- a/source/tests/pt_expt/model/test_dpa4_export.py
+++ b/source/tests/pt_expt/model/test_dpa4_export.py
@@ -168,7 +168,10 @@ def test_dpa4_fp32_cpu_export_runs_without_cuda_only_ops(monkeypatch) -> None:
 
 
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required")
-def test_dpa4_fp32_cuda_export_runs_with_fast_ops(monkeypatch) -> None:
+@pytest.mark.parametrize("env_exp", [None, 5])
+def test_dpa4_fp32_cuda_export_runs_with_fast_ops(
+    monkeypatch: pytest.MonkeyPatch, env_exp: int | None
+) -> None:
     """CPU tracing preserves the CUDA fast operators for a CUDA target."""
     try:
         import deepmd.pt.cxx_op  # noqa: F401
@@ -204,6 +207,7 @@ def test_dpa4_fp32_cuda_export_runs_with_fast_ops(monkeypatch) -> None:
     config = copy.deepcopy(_DPA4_CONFIG)
     config["descriptor"]["precision"] = "float32"
     config["descriptor"]["channels"] = 32
+    config["descriptor"]["env_exp"] = env_exp
     config["fitting_net"]["precision"] = "float32"
     model = get_model(config).to("cpu").eval()
     data = {"model": model.serialize()}
@@ -294,10 +298,15 @@ def test_dpa4_triton_force_assembly_survives_cpu_trace(monkeypatch) -> None:
         pytest.param(torch.device("cuda"), id="cuda"),
     ],
 )
+@pytest.mark.parametrize(
+    "env_exp",
+    [pytest.param(None, id="double_envelope"), pytest.param(5, id="single_envelope")],
+)
 def test_dpa4_fp32_aoti_package_runs_on_target(
     monkeypatch,
     tmp_path,
     target_device,
+    env_exp: int | None,
 ) -> None:
     """CPU tracing produces runnable CPU and CUDA packages for their target."""
     if target_device.type == "cuda" and not torch.cuda.is_available():
@@ -346,6 +355,7 @@ def test_dpa4_fp32_aoti_package_runs_on_target(
     config = copy.deepcopy(_DPA4_CONFIG)
     config["descriptor"]["precision"] = "float32"
     config["descriptor"]["channels"] = 32
+    config["descriptor"]["env_exp"] = env_exp
     config["fitting_net"]["precision"] = "float32"
     model = get_model(config).to("cpu").eval()
     data = {"model": model.serialize()}
@@ -392,6 +402,9 @@ def test_dpa4_fp32_aoti_package_runs_on_target(
     assert output
     assert all(bool(torch.isfinite(value).all()) for value in output.values())
     assert torch.max(torch.abs(output["force"])).item() > 1e-6
+    expected = exported.module()(*sample)
+    for key, value in output.items():
+        torch.testing.assert_close(value, expected[key], rtol=1e-4, atol=1e-5)
 
 
 @pytest.mark.skipif(
diff --git a/source/tests/pt_expt/model/test_dpa4_interop.py b/source/tests/pt_expt/model/test_dpa4_interop.py
index 9f4829d1e1..49761ebe1d 100644
--- a/source/tests/pt_expt/model/test_dpa4_interop.py
+++ b/source/tests/pt_expt/model/test_dpa4_interop.py
@@ -109,6 +109,40 @@ def _forward_smoke(model: EnergyModel) -> dict:
 
 
 class TestDPA4Interop:
+    @pytest.mark.parametrize("basis_type", ["bessel", "gaussian"])
+    def test_single_envelope_normalization_and_roundtrip(self, basis_type: str) -> None:
+        """Preserve the integer envelope configuration and its energy/force function."""
+        config = copy.deepcopy(_DPA4_RAW_CONFIG)
+        config["descriptor"].update(env_exp=5, basis_type=basis_type)
+        model_params = _normalize_model(config)
+        assert model_params["descriptor"]["env_exp"] == 5
+        pt_model = pt_get_model(model_params).to(env.DEVICE).eval()
+        generator = torch.Generator(device=env.DEVICE).manual_seed(29)
+        with torch.no_grad():
+            for parameter in pt_model.parameters():
+                parameter.add_(
+                    0.01
+                    * torch.randn(
+                        parameter.shape,
+                        dtype=parameter.dtype,
+                        device=parameter.device,
+                        generator=generator,
+                    )
+                )
+        expt_model = BaseModel.deserialize(pt_model.serialize()).to(env.DEVICE).eval()
+        for model in (pt_model, expt_model):
+            descriptor = model.atomic_model.descriptor
+            assert descriptor.env_exp == 5
+            assert descriptor.radial_basis.exponent == 0
+            assert descriptor.radial_basis.envelope is None
+        expected = _forward_smoke(pt_model)
+        actual = _forward_smoke(expt_model)
+        assert expected["force"].abs().max().item() > 1e-10
+        for key in ("energy", "force", "virial"):
+            torch.testing.assert_close(
+                actual[key], expected[key], rtol=1e-9, atol=1e-10
+            )
+
     def test_serialize_layout(self, pt_dpa4_model) -> None:
         """The pt serialize layout matches the interop override's expectations."""
         ser = pt_dpa4_model.serialize()
diff --git a/source/tests/pt_expt/model/test_dpa4_native_spin.py b/source/tests/pt_expt/model/test_dpa4_native_spin.py
index 8e0b1447de..795abc18f5 100644
--- a/source/tests/pt_expt/model/test_dpa4_native_spin.py
+++ b/source/tests/pt_expt/model/test_dpa4_native_spin.py
@@ -300,6 +300,19 @@ def _jittered_wrapper(seed: int = 11) -> NativeSpinEnergyModel:
 class TestNativeSpinEnergyModelPtExpt:
     """Public ``forward()`` contract of the pt_expt ``NativeSpinEnergyModel``."""
 
+    def test_adam_route_patterns_match_parameters(self) -> None:
+        """The DPA4 backbone and its native-spin model declare the same routed tensors."""
+        expected = [
+            "descriptor.env_seed_embedding.rbf_proj_layer1.",
+            "descriptor.radial_embedding.net.0.",
+        ]
+        for model in (_build_jittered_backbone(), _jittered_wrapper()):
+            patterns = sorted(model.adam_route_patterns())
+            assert patterns == expected
+            names = [name for name, _ in model.named_parameters()]
+            for pattern in patterns:
+                assert any(pattern in name for name in names), pattern
+
     def setup_method(self) -> None:
         self.device = _env.DEVICE
         self.model = _jittered_wrapper(seed=11)

From 906d1cd2301e7155f0b5e825d5a4c10633844eb3 Mon Sep 17 00:00:00 2001
From: OutisLi <LTC201806070316@gmail.com>
Date: Fri, 11 Sep 2026 17:52:29 +0800
Subject: [PATCH 2/8] fix(pt): edge-free frames, fused-kernel caches and the
 read-out under export

The padded forward skipped the radial embedding, the environment seed,
the geometric initial embedding and the interaction blocks whenever a
frame held no valid edge, so an isolated atom was a different function
of its features in an edge-free frame than in a frame with other edges,
and the descriptor jumped when the last edge of a frame left the cutoff.
The dpmodel and sparse-edge paths never had the shortcut. The padded
path now takes the same route for any edge count, the special empty
cache is gone, and the Triton radial mixer reshapes with the explicit
rank so that an empty edge set is well defined.

The fused radial function and the Wigner table builders cache constant
tensors on first use. Under the freeze the first make_fx trace built
them as fake tensors bound to that trace, the with-comm trace reused
them, and torch.export rejected the mixed fake modes, so `dp --pt freeze`
of a DPA4 model failed on a CUDA target. Constants built under a tracing
mode are now returned without being cached.

The scalar SO(3) read-out product is written as a weighted product-sum
instead of a three-operand einsum: the contraction-path search of the
latter guarded on the symbolic atom count and broke the with-comm export
whenever the fused kernels are off (CPU targets included).
---
 deepmd/pt/model/descriptor/sezm.py            | 54 +++++++--------
 .../pt/model/descriptor/sezm_nn/edge_cache.py | 67 -------------------
 .../pt/model/descriptor/sezm_nn/grid_net.py   |  9 ++-
 .../pt_expt/kernels/cuda/dpa4/edge_radial.py  | 10 ++-
 deepmd/pt_expt/kernels/cuda/dpa4/so2_conv.py  |  8 ++-
 .../pt_expt/kernels/cuda/dpa4/wigner_dense.py |  8 ++-
 .../pt_expt/kernels/triton/sezm/radial_mix.py |  6 +-
 source/tests/pt/model/test_descriptor_sezm.py | 55 +++++++++++++--
 .../pt/model/test_dpa4_dpmodel_parity.py      | 10 +++
 source/tests/pt/model/test_sezm_model.py      |  1 -
 10 files changed, 113 insertions(+), 115 deletions(-)

diff --git a/deepmd/pt/model/descriptor/sezm.py b/deepmd/pt/model/descriptor/sezm.py
index b770a9edc9..377b837fe8 100644
--- a/deepmd/pt/model/descriptor/sezm.py
+++ b/deepmd/pt/model/descriptor/sezm.py
@@ -1346,7 +1346,6 @@ def forward(
                 ),
                 edge_envelope=self.edge_envelope,
                 radial_basis=self.radial_basis,
-                n_radial=self.radial_basis.n_radial,
                 # Random local-Z roll is a training-only augmentation;
                 # the model is roll-equivariant, so inference fixes gamma.
                 random_gamma=self.random_gamma and self.training,
@@ -1360,21 +1359,19 @@ def forward(
 
         # === Step 5. Compute radial features once (fp32+) ===
         # Shape: (E, (node_init_lmax+1)*C) -> (E, node_init_lmax+1, C)
-        radial_feat = None
         with nvtx_range("radial_embedding"):
-            if edge_cache.src.numel() > 0:
-                radial_feat = rearrange(
-                    self.radial_embedding(edge_cache.edge_rbf),
-                    "E (L C) -> E L C",
-                    L=self.node_init_lmax + 1,
-                    C=self.channels,
-                )  # (E, node_init_lmax+1, C)
-                if self.version >= 1.1:
-                    radial_feat = radial_feat * edge_cache.edge_env.reshape(-1, 1, 1)
+            radial_feat = rearrange(
+                self.radial_embedding(edge_cache.edge_rbf),
+                "E (L C) -> E L C",
+                L=self.node_init_lmax + 1,
+                C=self.channels,
+            )  # (E, node_init_lmax+1, C)
+            if self.version >= 1.1:
+                radial_feat = radial_feat * edge_cache.edge_env.reshape(-1, 1, 1)
 
         # === Step 6. Env FiLM conditioning (optional, fp32+) ===
         with nvtx_range("env_film"):
-            if self.use_env_seed and edge_cache.src.numel() > 0:
+            if self.use_env_seed:
                 atype_flat = atype_loc.reshape(-1)  # (N,)
                 spin_flat = (
                     spin.reshape(n_nodes, 3)
@@ -1403,7 +1400,7 @@ def forward(
 
         # === Step 8. Geometric Initial Embedding (+ neighbor spin l=1) ===
         with nvtx_range("gie"):
-            if self.use_gie and radial_feat is not None:
+            if self.use_gie:
                 # GIE only needs l>=1, slice radial_feat[:, 1:, :]
                 zonal_coupling = self._build_gie_zonal_coupling(edge_cache)
                 spin_l1_message = (
@@ -1431,26 +1428,25 @@ def forward(
 
         # === Step 10. Fuse edge type features into radial features (fp32+) ===
         with nvtx_range("radial_fuse"):
-            if radial_feat is not None:
-                radial_feat = radial_feat + rearrange(
-                    edge_cache.edge_type_feat, "E C -> E 1 C"
-                )
-                radial_feat = radial_feat.to(dtype=self.dtype)
-                rad_feat_per_block = [
-                    radial_feat[:, :rad_len, :] for rad_len in self.rad_sizes_per_block
-                ]  # list of (E, lmax+1, C)
-            else:
-                rad_feat_per_block = []
+            radial_feat = radial_feat + rearrange(
+                edge_cache.edge_type_feat, "E C -> E 1 C"
+            )
+            radial_feat = radial_feat.to(dtype=self.dtype)
+            rad_feat_per_block = [
+                radial_feat[:, :rad_len, :] for rad_len in self.rad_sizes_per_block
+            ]  # list of (E, lmax+1, C)
 
         # === Step 11. Convert to self.dtype and run blocks ===
-        # The block stage is skipped entirely when there are no interaction
-        # blocks (zero-block descriptor) or no valid edges, sparing the working
-        # edge-cache dtype cast that only the blocks consume.
+        # The block stage is skipped entirely for the zero-block descriptor,
+        # sparing the working edge-cache dtype cast that only the blocks consume.
+        # A frame without valid edges takes the same path as any other, so an
+        # isolated atom is one function of its features whether or not the
+        # frame holds other edges.
         with nvtx_range("blocks"):
             x = x.to(dtype=self.dtype)  # (N, D, 1, C)
             if force_embedding is not None:
                 x = x + force_embedding.to(dtype=self.dtype)
-            if self.blocks and edge_cache.src.numel() > 0:
+            if self.blocks:
                 edge_cache = edge_cache_to_dtype(edge_cache, self.dtype)
                 with self._compute_mode_ctx(extended_coord.device):
                     x = self._forward_blocks(x, edge_cache, rad_feat_per_block)
@@ -1858,8 +1854,8 @@ def _apply_readout(self, x: torch.Tensor, n_rows: int) -> torch.Tensor:
         ----------
         x
             Node features with shape ``(n_rows, D, 1, channels)``. With the
-            blocks skipped (zero-block or empty-edge path) ``D`` is the initial
-            degree; otherwise the pyramid has shrunk it, so the read-out slice to
+            blocks skipped (zero-block descriptor) ``D`` is the initial degree;
+            otherwise the pyramid has shrunk it, so the read-out slice to
             ``node_readout_dim`` is a no-op there.
         n_rows
             Number of node rows fed to the read-out.
diff --git a/deepmd/pt/model/descriptor/sezm_nn/edge_cache.py b/deepmd/pt/model/descriptor/sezm_nn/edge_cache.py
index 97bf9dab63..1953c2f620 100644
--- a/deepmd/pt/model/descriptor/sezm_nn/edge_cache.py
+++ b/deepmd/pt/model/descriptor/sezm_nn/edge_cache.py
@@ -296,7 +296,6 @@ def build_edge_cache(
     deg_norm_floor: float,
     edge_envelope: Callable[[torch.Tensor], torch.Tensor],
     radial_basis: Callable[[torch.Tensor], torch.Tensor],
-    n_radial: int,
     random_gamma: bool,
     wigner_calc: WignerCalculatorFn,
     build_wigner: bool = True,
@@ -353,8 +352,6 @@ def build_edge_cache(
         C^3 edge envelope module.
     radial_basis
         Radial basis module.
-    n_radial
-        Number of radial basis channels used for empty-cache allocation.
     random_gamma
         Whether to apply a random roll around the local +Z axis before
         constructing Wigner-D blocks.
@@ -384,15 +381,6 @@ def build_edge_cache(
             nall=nall,
         )
 
-    if src.numel() == 0:
-        return _get_empty_edge_cache(
-            n_nodes=n_nodes,
-            n_radial=n_radial,
-            n_channel=type_ebed.shape[1],
-            device=extended_coord.device,
-            dtype=extended_coord.dtype,
-        )
-
     # === Step 3-5. Edge geometry/RBF chain ===
     #   gather -> edge_vec -> edge_len -> edge_env -> edge_rbf
     coord_flat = coord.reshape(nf * nall, 3)
@@ -750,61 +738,6 @@ def _finalize_edge_cache(
     )
 
 
-def _get_empty_edge_cache(
-    *,
-    n_nodes: int,
-    n_radial: int,
-    n_channel: int,
-    device: torch.device,
-    dtype: torch.dtype,
-) -> EdgeFeatureCache:
-    """
-    Allocate an empty edge cache for one SeZM forward pass.
-
-    Parameters
-    ----------
-    n_nodes
-        Number of local nodes in the flattened frame-major layout.
-    n_radial
-        Number of radial basis channels.
-    n_channel
-        Edge type feature width.
-    device
-        Target device for the cache tensors.
-    dtype
-        Target floating-point dtype for the cache tensors.
-
-    Returns
-    -------
-    EdgeFeatureCache
-        Empty cache with valid tensor shapes and neutral degree normalization.
-    """
-    empty_long = torch.empty(0, dtype=torch.long, device=device)
-    empty_vec = torch.empty(0, 3, dtype=dtype, device=device)
-    empty_quat = torch.empty(0, 4, dtype=dtype, device=device)
-    empty_rbf = torch.empty(0, n_radial, dtype=dtype, device=device)
-    empty_type_feat = torch.empty(0, n_channel, dtype=dtype, device=device)
-    deg = torch.zeros(n_nodes, dtype=dtype, device=device)
-    inv_sqrt_deg = torch.ones(n_nodes, 1, 1, dtype=dtype, device=device)
-    return EdgeFeatureCache(
-        src=empty_long,
-        dst=empty_long,
-        edge_type_feat=empty_type_feat,
-        edge_vec=empty_vec,
-        edge_rbf=empty_rbf,
-        edge_env=torch.empty(0, 1, dtype=dtype, device=device),
-        deg=deg,
-        inv_sqrt_deg=inv_sqrt_deg,
-        D_full=None,
-        Dt_full=None,
-        D_to_m_cache={},
-        Dt_from_m_cache={},
-        csr_cache={},
-        edge_src_gate=None,
-        edge_quat=empty_quat,
-    )
-
-
 def _build_standard_edge_index(
     *,
     nlist: torch.Tensor,
diff --git a/deepmd/pt/model/descriptor/sezm_nn/grid_net.py b/deepmd/pt/model/descriptor/sezm_nn/grid_net.py
index 0e14ec85b3..c04673d298 100644
--- a/deepmd/pt/model/descriptor/sezm_nn/grid_net.py
+++ b/deepmd/pt/model/descriptor/sezm_nn/grid_net.py
@@ -1138,11 +1138,10 @@ def _scalar_so3_product(
         n_batch, coeff_dim, n_focus, _ = left.shape
         left_view = left.reshape(n_batch, coeff_dim, n_focus, self.n_frames, -1)
         right_view = right.reshape_as(left_view)
-        scalar = torch.einsum(
-            "ndfkc,dk,ndfkc->nfc",
-            left_view,
-            weight,
-            right_view,
+        # A weighted diagonal product-sum over (d, k); written out so that no
+        # contraction-path search runs on the symbolic node count under export.
+        scalar = (left_view * weight[None, :, None, :, None] * right_view).sum(
+            dim=(1, 3)
         )
         return scalar[:, None, :, :]
 
diff --git a/deepmd/pt_expt/kernels/cuda/dpa4/edge_radial.py b/deepmd/pt_expt/kernels/cuda/dpa4/edge_radial.py
index 1c2f1bdffa..ab44f426d0 100644
--- a/deepmd/pt_expt/kernels/cuda/dpa4/edge_radial.py
+++ b/deepmd/pt_expt/kernels/cuda/dpa4/edge_radial.py
@@ -32,6 +32,9 @@
 )
 
 import torch
+from torch._subclasses.fake_tensor import (
+    FakeTensor,
+)
 
 __all__ = [
     "BESSEL",
@@ -211,10 +214,15 @@ def __init__(self, envelope: Any, basis: Any) -> None:
     def series(self, device: torch.device) -> tuple[torch.Tensor, torch.Tensor]:
         """The two Horner series on the compute device."""
         if self._series is None or self._series[0].device != device:
-            self._series = (
+            series = (
                 torch.tensor(self._env, dtype=torch.float32, device=device),
                 torch.tensor(self._rbf, dtype=torch.float32, device=device),
             )
+            # Series built under a tracing mode are fake tensors bound to that
+            # trace; only real series are kept for later calls.
+            if isinstance(series[0], FakeTensor):
+                return series
+            self._series = series
         return self._series
 
     def __call__(
diff --git a/deepmd/pt_expt/kernels/cuda/dpa4/so2_conv.py b/deepmd/pt_expt/kernels/cuda/dpa4/so2_conv.py
index 7a86943666..de407189ed 100644
--- a/deepmd/pt_expt/kernels/cuda/dpa4/so2_conv.py
+++ b/deepmd/pt_expt/kernels/cuda/dpa4/so2_conv.py
@@ -61,6 +61,9 @@
 )
 
 import torch
+from torch._subclasses.fake_tensor import (
+    FakeTensor,
+)
 
 __all__ = [
     "SO2ConvCuda",
@@ -183,7 +186,10 @@ def wigner_run_tables(lmax: int) -> tuple[torch.Tensor, ...]:
         exps.to(torch.int8).contiguous(),
         dexps.to(torch.int8).contiguous(),
     )
-    _RUN_TABLE_CACHE[lmax] = tables
+    # Tables built under a tracing mode are fake tensors bound to that trace;
+    # only real tables are shared across calls.
+    if not any(isinstance(table, FakeTensor) for table in tables):
+        _RUN_TABLE_CACHE[lmax] = tables
     return tables
 
 
diff --git a/deepmd/pt_expt/kernels/cuda/dpa4/wigner_dense.py b/deepmd/pt_expt/kernels/cuda/dpa4/wigner_dense.py
index f5a1f55d90..3980ed898d 100644
--- a/deepmd/pt_expt/kernels/cuda/dpa4/wigner_dense.py
+++ b/deepmd/pt_expt/kernels/cuda/dpa4/wigner_dense.py
@@ -31,6 +31,9 @@
 )
 
 import torch
+from torch._subclasses.fake_tensor import (
+    FakeTensor,
+)
 
 from .so2_conv import (
     _monomial_exponents,
@@ -133,7 +136,10 @@ def wigner_dense_tables(lmax: int) -> tuple[torch.Tensor, ...]:
         torch.tensor(coeffs, dtype=torch.float32, device="cpu"),
         torch.tensor(monos, dtype=torch.int32, device="cpu"),
     )
-    _DENSE_TABLE_CACHE[lmax] = tables
+    # Tables built under a tracing mode are fake tensors bound to that trace;
+    # only real tables are shared across calls.
+    if not any(isinstance(table, FakeTensor) for table in tables):
+        _DENSE_TABLE_CACHE[lmax] = tables
     return tables
 
 
diff --git a/deepmd/pt_expt/kernels/triton/sezm/radial_mix.py b/deepmd/pt_expt/kernels/triton/sezm/radial_mix.py
index 8fdf8bd00e..2a95cb6727 100644
--- a/deepmd/pt_expt/kernels/triton/sezm/radial_mix.py
+++ b/deepmd/pt_expt/kernels/triton/sezm/radial_mix.py
@@ -123,7 +123,7 @@ def radial_mix_reference(
     for coeff0, comp0, num_l in _block_layout(int(lmax)):
         # K[e, o, i, r] = compact[e, comp0 + i * num_l + o, r]
         block = compact[:, comp0 : comp0 + num_l * num_l, :].reshape(
-            n_edge, num_l, num_l, -1
+            n_edge, num_l, num_l, compact.shape[-1]
         )
         block = block.permute(0, 2, 1, 3)  # (E, o, i, R)
         x_block = x_local[:, coeff0 : coeff0 + num_l, :]  # (E, i, C)
@@ -176,7 +176,7 @@ def _radial_mix_backward_reference(
         #   with K[e, o, i, r] = compact[e, comp0 + i * num_l + o, r].
         k_block = (
             compact[:, comp0 : comp0 + num_l * num_l, :]
-            .reshape(n_edge, num_l, num_l, -1)
+            .reshape(n_edge, num_l, num_l, compact.shape[-1])
             .permute(0, 2, 1, 3)
         )  # (E, o, i, R)
         x_block = x_local[:, coeff0 : coeff0 + num_l, :]  # (E, i, C)
@@ -826,7 +826,7 @@ def channel_basis_grad(
         # degree instead would leave the rank axis innermost and force a
         # transposing copy of the whole edge tensor before the reduction.
         kernel = compact[:, comp0 : comp0 + num_l * num_l, :].reshape(
-            n_edge, num_l, num_l, -1
+            n_edge, num_l, num_l, compact.shape[-1]
         )  # (E, i, o, R)
         x_block = x_local[:, coeff0 : coeff0 + num_l, :]  # (E, i, C)
         g_block = grad_out[:, coeff0 : coeff0 + num_l, :]  # (E, o, C)
diff --git a/source/tests/pt/model/test_descriptor_sezm.py b/source/tests/pt/model/test_descriptor_sezm.py
index 0a98f7e664..a077c623f0 100644
--- a/source/tests/pt/model/test_descriptor_sezm.py
+++ b/source/tests/pt/model/test_descriptor_sezm.py
@@ -124,6 +124,22 @@ def _descriptor_kwargs(**overrides) -> dict:
     return kwargs
 
 
+def _perturb_parameters(model: torch.nn.Module, seed: int, scale: float = 0.1) -> None:
+    """Move every parameter off its initial value; fresh output projections are zero."""
+    generator = torch.Generator(device=env.DEVICE).manual_seed(seed)
+    with torch.no_grad():
+        for parameter in model.parameters():
+            parameter.add_(
+                scale
+                * torch.randn(
+                    parameter.shape,
+                    dtype=parameter.dtype,
+                    device=parameter.device,
+                    generator=generator,
+                )
+            )
+
+
 def _attention_descriptor_kwargs(
     *,
     precision: str = "float32",
@@ -332,17 +348,15 @@ def test_cartesian_rotation_invariance(self) -> None:
                 torch.testing.assert_close(desc, desc_rot, atol=1e-10, rtol=1e-10)
 
     def test_so3_readout_empty_edge_shrinking_schedule(self) -> None:
-        """so3_readout glu/mlp must handle the empty-edge path.
+        """so3_readout glu/mlp must handle a frame without edges.
 
-        With a shrinking ``l_schedule`` and no edges (every atom isolated),
-        ``_forward_blocks`` is skipped so ``x`` keeps the *initial* node degree
-        ``node_ebed_dims[0]``; the readout must truncate it to the final degree
-        ``node_ebed_dims[-1]`` (what ``output_ffn`` is built for) before the FFN.
-        Regression for the readout shape mismatch on isolated atoms.
+        With a shrinking ``l_schedule`` and no edges (every atom isolated), the
+        blocks run on an empty edge set and the readout receives the final node
+        degree ``node_ebed_dims[-1]`` (what ``output_ffn`` is built for).
         """
         coord, atype, _ = _tiny_two_atom_system(self.device, dtype=torch.float32)
         extended_coord = coord.reshape(1, -1).detach().requires_grad_(True)
-        # all neighbors masked out -> edge_cache.src.numel() == 0 -> blocks skipped
+        # all neighbors masked out -> an empty edge set
         nlist = torch.full((1, 2, 2), -1, dtype=torch.int64, device=self.device)
         for readout in ("glu", "mlp"):
             with self.subTest(so3_readout=readout):
@@ -355,6 +369,33 @@ def test_so3_readout_empty_edge_shrinking_schedule(self) -> None:
                 self.assertEqual(desc.shape, (1, 2, 4))
                 self.assertTrue(torch.all(torch.isfinite(desc)))
 
+    def test_edge_free_frame_continues_the_cutoff_limit(self) -> None:
+        """A frame without edges is the limit of a frame whose last edge leaves the cutoff."""
+        model = DescrptSeZM(**_descriptor_kwargs(precision="float64", seed=5))
+        model = model.to(self.device).eval()
+        _perturb_parameters(model, seed=5)
+        atype = torch.tensor([[0, 1]], dtype=torch.int32, device=self.device)
+        empty_nlist = torch.full((1, 2, 2), -1, dtype=torch.int64, device=self.device)
+        pair_nlist = torch.tensor(
+            [[[1, -1], [0, -1]]], dtype=torch.int64, device=self.device
+        )
+
+        def descriptor(distance: float, nlist: torch.Tensor) -> torch.Tensor:
+            coord = torch.tensor(
+                [[0.0, 0.0, 0.0], [distance, 0.0, 0.0]],
+                dtype=torch.float64,
+                device=self.device,
+            ).reshape(1, -1)
+            return model(coord, atype, nlist, mapping=None, comm_dict=None)[0]
+
+        isolated = descriptor(10.0, empty_nlist)
+        torch.testing.assert_close(
+            descriptor(model.rcut - 1e-6, pair_nlist), isolated, rtol=0.0, atol=1e-12
+        )
+        self.assertFalse(
+            torch.allclose(descriptor(model.rcut - 0.5, pair_nlist), isolated)
+        )
+
     def test_so3_readout_scalar_path_matches_full_output(self) -> None:
         """The scalar-specialized final FFN matches slicing its full output."""
         dtype = torch.float64
diff --git a/source/tests/pt/model/test_dpa4_dpmodel_parity.py b/source/tests/pt/model/test_dpa4_dpmodel_parity.py
index 60c80ccbde..80c8119c1f 100644
--- a/source/tests/pt/model/test_dpa4_dpmodel_parity.py
+++ b/source/tests/pt/model/test_dpa4_dpmodel_parity.py
@@ -4016,6 +4016,16 @@ def test_descriptor_exclude_types(self, exclude_types) -> None:
         pt_mod, dp_mod, _ = self._build_descr_pair(exclude_types=exclude_types)
         self._assert_descr_parity(pt_mod, dp_mod)
 
+    def test_descriptor_without_edges(self) -> None:
+        pt_mod, dp_mod, _ = self._build_descr_pair()
+        inp = self._inputs()
+        nlist = np.full_like(inp["nlist"], -1)
+        out_dp = dp_mod.call(
+            inp["coord"].reshape(self.nf, -1), inp["atype_ext"], nlist, mapping=None
+        )
+        out_pt = pt_mod(to_pt(inp["coord"]), to_pt(inp["atype_ext"]), to_pt(nlist))
+        assert_parity(out_dp[0], out_pt[0], rtol=1e-10, atol=1e-12)
+
     def test_descriptor_no_mapping(self) -> None:
         # pt forward accepts mapping=None when neighbor indices are local;
         # mapping is NOT required by either backend
diff --git a/source/tests/pt/model/test_sezm_model.py b/source/tests/pt/model/test_sezm_model.py
index 73f9baccdf..5808353967 100644
--- a/source/tests/pt/model/test_sezm_model.py
+++ b/source/tests/pt/model/test_sezm_model.py
@@ -690,7 +690,6 @@ def test_fixed_edge_geometry_matches_standard_cache(self) -> None:
             deg_norm_floor=descriptor.deg_norm_floor,
             edge_envelope=descriptor.edge_envelope,
             radial_basis=descriptor.radial_basis,
-            n_radial=descriptor.radial_basis.n_radial,
             random_gamma=False,
             wigner_calc=descriptor.wigner_calc,
         )

From e07229159e0cfc92d7fdaf97b55fb2be49cf797f Mon Sep 17 00:00:00 2001
From: OutisLi <LTC201806070316@gmail.com>
Date: Fri, 11 Sep 2026 18:58:49 +0800
Subject: [PATCH 3/8] fix(pt-expt): match CUDA fake output layouts

---
 deepmd/pt_expt/kernels/cuda/dpa4/grid_pair.py |  5 ++-
 .../kernels/cuda/dpa4/zonal_scatter.py        |  3 +-
 .../pt/model/test_descriptor_sezm_cuda.py     | 39 +++++++++++++++++++
 3 files changed, 44 insertions(+), 3 deletions(-)

diff --git a/deepmd/pt_expt/kernels/cuda/dpa4/grid_pair.py b/deepmd/pt_expt/kernels/cuda/dpa4/grid_pair.py
index 31dfa013b6..5e83cc550e 100644
--- a/deepmd/pt_expt/kernels/cuda/dpa4/grid_pair.py
+++ b/deepmd/pt_expt/kernels/cuda/dpa4/grid_pair.py
@@ -54,7 +54,8 @@ def _forward_fake(
     from_grid: torch.Tensor,
 ) -> torch.Tensor:
     del right, to_grid, from_grid
-    return torch.empty_like(left)
+    # The CUDA implementation allocates outputs from contiguous operands.
+    return left.new_empty(left.shape)
 
 
 def _backward_fake(
@@ -65,7 +66,7 @@ def _backward_fake(
     from_grid: torch.Tensor,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     del grad_out, to_grid, from_grid
-    return torch.empty_like(left), torch.empty_like(right)
+    return left.new_empty(left.shape), right.new_empty(right.shape)
 
 
 def _setup_context(ctx: Any, inputs: tuple[Any, ...], output: torch.Tensor) -> None:
diff --git a/deepmd/pt_expt/kernels/cuda/dpa4/zonal_scatter.py b/deepmd/pt_expt/kernels/cuda/dpa4/zonal_scatter.py
index fe2419bc3e..7011f215ee 100644
--- a/deepmd/pt_expt/kernels/cuda/dpa4/zonal_scatter.py
+++ b/deepmd/pt_expt/kernels/cuda/dpa4/zonal_scatter.py
@@ -65,7 +65,8 @@ def _backward_fake(
     node_scale: torch.Tensor,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     del grad_out, dst, node_scale
-    return torch.empty_like(zonal), torch.empty_like(radial)
+    # The CUDA implementation allocates gradients from contiguous operands.
+    return zonal.new_empty(zonal.shape), radial.new_empty(radial.shape)
 
 
 def _setup_context(ctx: Any, inputs: tuple[Any, ...], output: torch.Tensor) -> None:
diff --git a/source/tests/pt/model/test_descriptor_sezm_cuda.py b/source/tests/pt/model/test_descriptor_sezm_cuda.py
index 8f8cc468f0..6791f99640 100644
--- a/source/tests/pt/model/test_descriptor_sezm_cuda.py
+++ b/source/tests/pt/model/test_descriptor_sezm_cuda.py
@@ -556,6 +556,25 @@ def test_backward_matches_autograd_on_every_slot_count(self) -> None:
                     scale = ref.abs().max()
                     self.assertLess(((got - ref).abs().max() / scale).item(), 5e-6)
 
+    def test_fake_layout_matches_noncontiguous_inputs(self) -> None:
+        """Forward and backward fake layouts match the CUDA allocations for views."""
+        left, right, to_grid, from_grid_t = self._case(2, 12, 32, 24)
+        left = left.transpose(1, 2).contiguous().transpose(1, 2)
+        right = right.transpose(1, 2).contiguous().transpose(1, 2)
+        self.assertFalse(left.is_contiguous())
+        self.assertFalse(right.is_contiguous())
+        inputs = (left, right, to_grid, from_grid_t)
+        output = grid_pair(*inputs)
+        for op, args in (
+            (torch.ops.deepmd.dpa4_grid_pair.default, inputs),
+            (
+                torch.ops.deepmd.dpa4_grid_pair_backward.default,
+                (torch.ones_like(output), *inputs),
+            ),
+        ):
+            with self.subTest(operator=str(op)):
+                torch.library.opcheck(op, args, test_utils=("test_faketensor",))
+
 
 @unittest.skipUnless(CUDA_ZONAL, "requires the CUDA dpa4_zonal_scatter operator")
 class TestSeZMZonalScatterCuda(unittest.TestCase):
@@ -644,6 +663,26 @@ def test_backward_reaches_the_degree_normalization(self) -> None:
         ).item()
         self.assertLess(rel, 5e-6)
 
+    def test_fake_layout_matches_noncontiguous_inputs(self) -> None:
+        """Forward and backward fake layouts match the CUDA allocations for views."""
+        n_node = 2
+        zonal, radial, dst, order, row_ptr, scale, _ = self._case(2, n_node, 2, 32)
+        zonal = zonal.T.contiguous().T
+        radial = radial.transpose(1, 2).contiguous().transpose(1, 2)
+        self.assertFalse(zonal.is_contiguous())
+        self.assertFalse(radial.is_contiguous())
+        inputs = (zonal, radial, dst, order, row_ptr, scale, n_node)
+        output = zonal_scatter(*inputs)
+        for op, args in (
+            (torch.ops.deepmd.dpa4_zonal_scatter.default, inputs),
+            (
+                torch.ops.deepmd.dpa4_zonal_scatter_backward.default,
+                (torch.ones_like(output), zonal, radial, dst, scale),
+            ),
+        ):
+            with self.subTest(operator=str(op)):
+                torch.library.opcheck(op, args, test_utils=("test_faketensor",))
+
     def test_channel_widths_beyond_one_block(self) -> None:
         # A lane owns one channel of a 32-wide block, so wider features sweep
         # the edge list more than once.

From 5580ebd3f291aeeca68b66a4afb3dc8577f4447c Mon Sep 17 00:00:00 2001
From: OutisLi <LTC201806070316@gmail.com>
Date: Fri, 11 Sep 2026 21:14:51 +0800
Subject: [PATCH 4/8] fix(schema): support model preset overrides

---
 deepmd/utils/argcheck.py                |   6 +-
 deepmd/utils/json_schema.py             | 221 ++++++++++++++++++++++++
 doc/train/train-input.rst               |  20 +++
 pyproject.toml                          |   1 +
 source/tests/common/test_json_schema.py | 193 +++++++++++++++++++++
 5 files changed, 440 insertions(+), 1 deletion(-)
 create mode 100644 deepmd/utils/json_schema.py
 create mode 100644 source/tests/common/test_json_schema.py

diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py
index cd781c4f84..0de0d1f7ab 100644
--- a/deepmd/utils/argcheck.py
+++ b/deepmd/utils/argcheck.py
@@ -6460,13 +6460,17 @@ def gen_json_schema(multi_task: bool = False) -> str:
     str
         JSON schema.
     """
+    from deepmd.utils.json_schema import (
+        with_model_presets,
+    )
+
     arg = Argument(
         "DeePMD-kit",
         dict,
         gen_args(multi_task=multi_task),
         doc=f"DeePMD-kit {__version__}",
     )
-    return json.dumps(generate_json_schema(arg))
+    return json.dumps(with_model_presets(generate_json_schema(arg), multi_task))
 
 
 def _check_dpa3_chg_spin_migration(data: dict[str, Any]) -> None:
diff --git a/deepmd/utils/json_schema.py b/deepmd/utils/json_schema.py
new file mode 100644
index 0000000000..027d31a2d1
--- /dev/null
+++ b/deepmd/utils/json_schema.py
@@ -0,0 +1,221 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""JSON Schema for model inputs before preset expansion."""
+
+import json
+import re
+from typing import (
+    Any,
+)
+
+from deepmd.utils.model_preset import (
+    MODEL_PRESETS,
+)
+
+_REGIONS = ("type_map", "descriptor", "fitting_net")
+
+
+def _preset_pattern(names: list[str]) -> str:
+    """Match complete preset names case-insensitively with ECMA-262 syntax."""
+    alternatives = [
+        "".join(
+            f"[{char.lower()}{char.upper()}]"
+            if char.isalpha()
+            else f"[{re.escape(char)}]"
+            for char in name
+        )
+        for name in names
+    ]
+    # Unlike '$', the final assertion also excludes a trailing newline.
+    return "^(?:" + "|".join(alternatives) + r")(?![\s\S])"
+
+
+def _with_defaults(
+    schema: dict[str, Any],
+    defaults: dict[str, Any],
+    ref: str,
+    *,
+    merge_regions: bool = True,
+    multi_task: bool = False,
+) -> dict[str, Any]:
+    """Validate overrides while accounting for fields supplied by a preset.
+
+    Model regions inherit missing fields; descriptor and fitting overrides
+    merge one level only. Other nested objects retain their full contracts.
+    Variant conditions use the inherited tag only when the input omits it.
+    Unchanged field definitions are referenced to keep the schema compact.
+    """
+    result = {
+        key: value
+        for key, value in schema.items()
+        if key not in ("properties", "required", "allOf")
+    }
+    properties = {}
+    for name, field in schema.get("properties", {}).items():
+        pointer = ref + "/properties/" + name.replace("~", "~0").replace("/", "~1")
+        if merge_regions and name in ("descriptor", "fitting_net"):
+            properties[name] = _with_defaults(
+                field, defaults.get(name, {}), pointer, merge_regions=False
+            )
+        elif name == "type" and name in defaults:
+            properties[name] = {**field, "default": defaults[name]}
+        else:
+            properties[name] = {"$ref": pointer}
+        if multi_task and merge_regions and name in _REGIONS:
+            properties[name] = {"anyOf": [properties[name], {"type": "string"}]}
+    result["properties"] = properties
+    result["required"] = [
+        name for name in schema.get("required", []) if name not in defaults
+    ]
+    conditions = []
+    for index, clause in enumerate(schema.get("allOf", [])):
+        pointer = f"{ref}/allOf/{index}"
+        if "if" in clause:
+            condition = clause["if"]
+            cases = condition["oneOf"]
+            flag = next(iter(cases[0]["properties"]))
+            if flag in defaults:
+                tags = [case["properties"][flag]["const"] for case in cases]
+                condition = {
+                    "properties": {flag: {"enum": tags}},
+                    "required": [] if defaults[flag] in tags else [flag],
+                }
+            conditions.append(
+                {
+                    "if": condition,
+                    "then": _with_defaults(
+                        clause["then"],
+                        defaults,
+                        pointer + "/then",
+                        merge_regions=merge_regions,
+                        multi_task=multi_task,
+                    ),
+                }
+            )
+        elif not any(
+            name in defaults
+            for choice in clause["oneOf"]
+            for name in choice["required"]
+        ):
+            conditions.append({"$ref": pointer})
+    if conditions:
+        result["allOf"] = conditions
+    return result
+
+
+def _when_preset(
+    choices: list[tuple[list[str], dict[str, Any]]],
+    fallback: dict[str, Any],
+) -> dict[str, Any]:
+    """Select the matching preset schema, or the caller's inherited default."""
+    return {
+        "if": {"required": ["preset"]},
+        "then": {
+            "allOf": [
+                {
+                    "if": {
+                        "properties": {
+                            "preset": {
+                                "type": "string",
+                                "pattern": _preset_pattern(names),
+                            }
+                        }
+                    },
+                    "then": selected,
+                }
+                for names, selected in choices
+            ]
+        },
+        "else": fallback,
+    }
+
+
+def _model_selector(
+    choices: list[tuple[list[str], dict[str, Any]]],
+    fallback: dict[str, Any],
+) -> dict[str, Any]:
+    return {
+        "type": "object",
+        "properties": {"preset": {"$ref": "#/$defs/model_preset_name"}},
+        **_when_preset(choices, fallback),
+    }
+
+
+def with_model_presets(schema: dict[str, Any], multi_task: bool) -> dict[str, Any]:
+    """Add raw preset input forms to a generated training schema.
+
+    Parameters
+    ----------
+    schema : dict
+        Training schema generated from the argument definitions, updated in place.
+    multi_task : bool
+        Whether model entries live in a multi-task ``model_dict``.
+
+    Returns
+    -------
+    dict
+        Schema accepting presets and partial overrides alongside ordinary inputs.
+    """
+    model = schema["properties"]["model"]
+    if multi_task:
+        branches = model["properties"]["model_dict"]
+        base = branches.pop("items")
+    else:
+        base = model
+
+    names = sorted(MODEL_PRESETS)
+    definitions = {
+        "model": base,
+        "model_preset_name": {
+            "description": "Named model architecture. Explicit model fields override its defaults.",
+            "type": "string",
+            "anyOf": [{"enum": names}, {"pattern": _preset_pattern(names)}],
+        },
+    }
+    choices = []
+    shared_schemas = {}
+    for name in names:
+        selected = _with_defaults(
+            base, MODEL_PRESETS[name], "#/$defs/model", multi_task=multi_task
+        )
+        key = json.dumps(selected, sort_keys=True)
+        if key not in shared_schemas:
+            definition = f"model_preset_{len(choices)}"
+            definitions[definition] = selected
+            shared_schemas[key] = []
+            choices.append((shared_schemas[key], {"$ref": f"#/$defs/{definition}"}))
+        shared_schemas[key].append(name)
+    schema["$defs"] = definitions
+
+    if not multi_task:
+        schema["properties"]["model"] = _model_selector(
+            choices, {"$ref": "#/$defs/model"}
+        )
+        return schema
+
+    # Object-valued repeat arguments use additionalProperties, not array items.
+    # Non-preset branches retain the existing multi-task schema's permissiveness;
+    # their shared references and cascaded fields are resolved at training time.
+    branches["additionalProperties"] = _model_selector(choices, {})
+    model["properties"]["preset"] = {"$ref": "#/$defs/model_preset_name"}
+    inherited = [
+        (
+            group,
+            {
+                "allOf": [
+                    selected,
+                    {
+                        "properties": {
+                            "model_dict": {
+                                "additionalProperties": _model_selector(
+                                    choices, selected
+                                )
+                            }
+                        }
+                    },
+                ]
+            },
+        )
+        for group, selected in choices
+    ]
+    model.update(_when_preset(inherited, {}))
+    return schema
diff --git a/doc/train/train-input.rst b/doc/train/train-input.rst
index 5f92ab73e6..9141603615 100644
--- a/doc/train/train-input.rst
+++ b/doc/train/train-input.rst
@@ -25,6 +25,26 @@ To do so, in a VS Code workspace, one can generate a JSON schema file for the in
 
    dp doc-train-input --out-type json_schema > deepmd.json
 
+The schema recognizes named model presets. When ``model.preset`` is present,
+the preset supplies the model, descriptor, and fitting types, so only the
+overrides need to be written:
+
+.. code-block:: json
+
+   {
+      "model": {
+         "preset": "dpa4-ultra-v20260911",
+         "descriptor": {"use_amp": true, "seed": 42},
+         "fitting_net": {"seed": 42},
+         "use_compile": true
+      }
+   }
+
+The editor validates the override fields and offers preset names and
+model-specific completions. Regenerate the schema after updating DeePMD-kit
+to include additional presets or input options. Multi-task schemas also
+recognize the top-level default preset and presets inside ``model_dict``.
+
 Then one can `map the schema <https://code.visualstudio.com/docs/languages/json#_mapping-to-a-schema-in-the-workspace>`_
 by updating the workspace settings in the `.vscode/settings.json` file as follows:
 
diff --git a/pyproject.toml b/pyproject.toml
index b0cb3b67be..5da7f9b1c6 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -88,6 +88,7 @@ test = [
     "pytest-sugar",
     "pytest-split",
     "pytest-timeout",
+    "jsonschema>=4.18",
     "dpgui",
     # DPA-ADAPT tests import sklearn via dpa_adapt.cv at module load time.
     "scikit-learn",
diff --git a/source/tests/common/test_json_schema.py b/source/tests/common/test_json_schema.py
new file mode 100644
index 0000000000..843af94da8
--- /dev/null
+++ b/source/tests/common/test_json_schema.py
@@ -0,0 +1,193 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""JSON Schema validation of input files before preset expansion."""
+
+import copy
+import json
+from typing import (
+    Any,
+)
+
+import pytest
+from jsonschema import (
+    Draft202012Validator,
+)
+
+from deepmd.utils.argcheck import (
+    gen_json_schema,
+    normalize,
+)
+from deepmd.utils.compat import (
+    update_deepmd_input,
+)
+from deepmd.utils.model_preset import (
+    MODEL_PRESETS,
+)
+
+
+@pytest.fixture(scope="module")
+def validators() -> dict[bool, Draft202012Validator]:
+    result = {}
+    for multi_task in (False, True):
+        schema = json.loads(gen_json_schema(multi_task=multi_task))
+        Draft202012Validator.check_schema(schema)
+        result[multi_task] = Draft202012Validator(schema)
+    return result
+
+
+def _input(model: dict[str, Any], multi_task: bool = False) -> dict[str, Any]:
+    training = {"training_data": {"systems": ["dummy"]}, "numb_steps": 10}
+    if multi_task:
+        training = {
+            "data_dict": {"a": {"training_data": {"systems": ["dummy"]}}},
+            "numb_steps": 10,
+        }
+    return {"model": model, "training": training}
+
+
+def _assert_valid(validator: Draft202012Validator, data: dict[str, Any]) -> None:
+    errors = list(validator.iter_errors(data))
+    assert not errors, [(list(error.path), error.message) for error in errors]
+
+
+def test_preset_user_example(validators: dict[bool, Draft202012Validator]) -> None:
+    data = _input(
+        {
+            "preset": "dpa4-ultra-v20260911",
+            "descriptor": {"use_amp": True, "seed": 42},
+            "fitting_net": {"seed": 42},
+            "use_compile": True,
+            "enable_tf32": True,
+            "_comment": "that's all",
+        }
+    )
+    data["learning_rate"] = {
+        "type": "cosine",
+        "start_lr": 1.2e-4,
+        "stop_lr": 1e-6,
+        "warmup_ratio": 0.003,
+        "warmup_start_factor": 0.2,
+    }
+    normalized = normalize(update_deepmd_input(copy.deepcopy(data), warning=False))
+    assert normalized["model"]["descriptor"]["type"] == "dpa4"
+    _assert_valid(validators[False], data)
+
+
+@pytest.mark.parametrize("name", sorted(MODEL_PRESETS))
+def test_every_preset_allows_partial_overrides(
+    validators: dict[bool, Draft202012Validator], name: str
+) -> None:
+    for model in (
+        {"preset": name},
+        {
+            "preset": name.upper(),
+            "descriptor": {"use_amp": True, "seed": 42},
+            "fitting_net": {"seed": 42},
+        },
+    ):
+        _assert_valid(validators[False], _input(model))
+
+
+@pytest.mark.parametrize(
+    "overrides",
+    [
+        {"preset": "dpa4-unknown-v20260911"},
+        {"preset": "dpa4-nano-v20260911\n"},
+        {"preset": 42},
+        {"descriptor": {"seed": "forty-two"}},
+        {"descriptor": {"use_amp": "true"}},
+        {"descriptor": {"type": "unknown"}},
+        {"fitting_net": {"seed": "forty-two"}},
+        {"fitting_net": {"neuron": "wide"}},
+        {"fitting_net": {"type": "unknown"}},
+        {
+            "fitting_net": {
+                "type": "property",
+                "property_name": "charge",
+                "task_dim": "four",
+            }
+        },
+        {"use_compile": "true"},
+        {"type": "unknown"},
+        {"type": "frozen"},
+        {"spin": {"scheme": "native"}},
+    ],
+)
+def test_presets_preserve_field_validation(
+    validators: dict[bool, Draft202012Validator], overrides: dict[str, Any]
+) -> None:
+    model = {"preset": "dpa4-nano-v20260911", **overrides}
+    assert not validators[False].is_valid(_input(model))
+
+
+def test_explicit_types_and_aliases(
+    validators: dict[bool, Draft202012Validator],
+) -> None:
+    _assert_valid(
+        validators[False],
+        _input(
+            {
+                "preset": "dpa4-nano-v20260911",
+                "type": "SeZM",
+                "descriptor": {"type": "SeZM", "so2_layers": 3},
+                "fitting_net": {
+                    "type": "property",
+                    "property_name": "charge",
+                    "task_dim": 4,
+                },
+            }
+        ),
+    )
+
+
+def test_plain_models_keep_required_fields(
+    validators: dict[bool, Draft202012Validator],
+) -> None:
+    model = {
+        "type_map": ["O"],
+        "descriptor": {"type": "se_e2_a", "sel": [10]},
+        "fitting_net": {"neuron": [4]},
+    }
+    _assert_valid(validators[False], _input(model))
+    del model["descriptor"]["type"]
+    assert not validators[False].is_valid(_input(model))
+
+
+def test_multi_task_preset_inheritance(
+    validators: dict[bool, Draft202012Validator],
+) -> None:
+    model = {
+        "preset": "DPA4-Nano-v20260911",
+        "descriptor": {"use_amp": True},
+        "fitting_net": {"seed": 42},
+        "model_dict": {
+            "a": {},
+            "b": {"preset": "dpa4c-mini-v20260911", "descriptor": {"seed": 42}},
+        },
+    }
+    _assert_valid(validators[True], _input(model, multi_task=True))
+    model["model_dict"]["b"]["descriptor"]["seed"] = "forty-two"
+    assert not validators[True].is_valid(_input(model, multi_task=True))
+
+
+def test_multi_task_branch_preset_and_shared_references(
+    validators: dict[bool, Draft202012Validator],
+) -> None:
+    model = {
+        "shared_dict": {
+            "type_map": ["O", "H"],
+            "descriptor": {"type": "dpa4", "use_amp": True},
+        },
+        "model_dict": {
+            "a": {
+                "preset": "dpa4-nano-v20260911",
+                "type_map": "type_map",
+                "descriptor": "descriptor:1",
+                "fitting_net": {"seed": 42},
+            }
+        },
+    }
+    _assert_valid(validators[True], _input(model, multi_task=True))
+    model["preset"] = "dpa4-mini-v20260911"
+    _assert_valid(validators[True], _input(model, multi_task=True))
+    model["model_dict"]["a"]["preset"] = "unknown"
+    assert not validators[True].is_valid(_input(model, multi_task=True))

From 1a17b1583f0678d4cc53bc41bbc5e1af3b61235a Mon Sep 17 00:00:00 2001
From: OutisLi <LTC201806070316@gmail.com>
Date: Fri, 11 Sep 2026 23:34:41 +0800
Subject: [PATCH 5/8] fix(dpa4): address review and CI regressions

---
 deepmd/dpmodel/descriptor/dpa4_nn/radial.py   |   8 +-
 deepmd/jax/descriptor/dpa4.py                 |   3 +-
 deepmd/pt_expt/descriptor/dpa4.py             |  13 +-
 deepmd/pt_expt/descriptor/dpa4c.py            |  17 +-
 .../pt_expt/kernels/cuda/dpa4/edge_radial.py  |  10 +-
 deepmd/utils/argcheck.py                      |   6 +-
 deepmd/utils/json_schema.py                   |  22 +-
 deepmd/utils/model_preset.py                  | 260 +----------------
 deepmd/utils/model_preset_data.py             | 262 ++++++++++++++++++
 source/op/pt/dpa4/edge_radial.cu              |  19 +-
 source/tests/jax/test_dpa4.py                 |  16 +-
 .../pt/model/test_descriptor_sezm_cuda.py     |  47 ++++
 .../pt/model/test_dpa4_dpmodel_parity.py      |   1 -
 source/tests/pt_expt/descriptor/test_dpa4.py  |   5 +-
 source/tests/tf2/test_dpa4.py                 |  18 +-
 15 files changed, 392 insertions(+), 315 deletions(-)
 create mode 100644 deepmd/utils/model_preset_data.py

diff --git a/deepmd/dpmodel/descriptor/dpa4_nn/radial.py b/deepmd/dpmodel/descriptor/dpa4_nn/radial.py
index 375e284123..6e20c41ed5 100644
--- a/deepmd/dpmodel/descriptor/dpa4_nn/radial.py
+++ b/deepmd/dpmodel/descriptor/dpa4_nn/radial.py
@@ -493,7 +493,7 @@ class RadialBasis(NativeOP):
         Radial basis type. Supported values are ``"bessel"``, ``"gaussian"``,
         ``"bessel/fix"`` and ``"gaussian/fix"``; the ``/fix`` forms are
         evaluated like their family and differ only in training, where the
-        PT backend keeps their frequencies or centres fixed.
+        backends keep their frequencies or centres fixed.
     precision : str
         Floating-point precision for the radial basis frequencies and outputs.
     exponent : int, optional
@@ -516,9 +516,9 @@ def __init__(
         if self.n_radial <= 0:
             raise ValueError("`n_radial` must be positive")
         self.basis_type = str(basis_type).lower()
-        # The ``/fix`` suffix governs training only: the PT backend freezes the
-        # basis parameters, the array-API basis evaluates either form alike.
-        self.basis_family, _ = parse_basis_type(self.basis_type)
+        # Parameter promotion in every backend consumes this trainability flag.
+        self.basis_family, fixed = parse_basis_type(self.basis_type)
+        self.trainable = not fixed
         self.precision = precision
         self.exponent = int(exponent)
         prec = PRECISION_DICT[self.precision.lower()]
diff --git a/deepmd/jax/descriptor/dpa4.py b/deepmd/jax/descriptor/dpa4.py
index 3db2141498..704b038770 100644
--- a/deepmd/jax/descriptor/dpa4.py
+++ b/deepmd/jax/descriptor/dpa4.py
@@ -245,8 +245,7 @@ def _promote_parameter_lists(
 def _promote_trainable_tree(module: Any) -> Any:
     root_trainable = bool(getattr(module, "trainable", True))
     for submodule in _iter_object_tree(module):
-        # A frozen descriptor freezes every descendant, including helper
-        # modules such as RadialBasis that do not carry a local flag.
+        # A frozen descriptor freezes every descendant regardless of its local flag.
         trainable = root_trainable and bool(getattr(submodule, "trainable", True))
         names = _TRAINABLE_ATTRS.get(type(submodule).__name__)
         if names is not None:
diff --git a/deepmd/pt_expt/descriptor/dpa4.py b/deepmd/pt_expt/descriptor/dpa4.py
index 0de6466d90..db843ffb01 100644
--- a/deepmd/pt_expt/descriptor/dpa4.py
+++ b/deepmd/pt_expt/descriptor/dpa4.py
@@ -16,9 +16,6 @@
     C3CutoffEnvelope as C3CutoffEnvelopeDP,
 )
 from deepmd.dpmodel.descriptor.dpa4_nn.radial import InnerClamp as InnerClampDP
-from deepmd.dpmodel.descriptor.dpa4_nn.radial import (
-    parse_basis_type,
-)
 from deepmd.pt_expt.common import (
     register_dpmodel_mapping,
     torch_module,
@@ -162,14 +159,13 @@ def forward(self, *args: Any, **kwargs: Any) -> Any:
 
 def _promote_trainable(module: torch.nn.Module, names: tuple[str, ...]) -> None:
     """Re-register the given float buffers of *module* as Parameters."""
-    if not getattr(module, "trainable", True):
-        return
+    trainable = bool(getattr(module, "trainable", True))
     for name in names:
         buf = module._buffers.get(name)
         if buf is None or not buf.is_floating_point():
             continue
         del module._buffers[name]
-        setattr(module, name, torch.nn.Parameter(buf, requires_grad=True))
+        setattr(module, name, torch.nn.Parameter(buf, requires_grad=trainable))
 
 
 def _promote_trainable_tree(module: torch.nn.Module) -> torch.nn.Module:
@@ -191,11 +187,6 @@ def _promote_trainable_tree(module: torch.nn.Module) -> torch.nn.Module:
         if getattr(sub, "trainable", True) is False:
             for p in sub.parameters(recurse=True):
                 p.requires_grad_(False)
-    # A ``/fix`` radial basis keeps its frequencies or centres, as in the pt
-    # backend.
-    for sub in module.modules():
-        if type(sub).__name__ == "RadialBasis" and parse_basis_type(sub.basis_type)[1]:
-            sub.adam_freqs.requires_grad_(False)
     return module
 
 
diff --git a/deepmd/pt_expt/descriptor/dpa4c.py b/deepmd/pt_expt/descriptor/dpa4c.py
index 5eb3ef3020..4fa1b17d66 100644
--- a/deepmd/pt_expt/descriptor/dpa4c.py
+++ b/deepmd/pt_expt/descriptor/dpa4c.py
@@ -15,9 +15,6 @@
 
 import torch
 
-from deepmd.dpmodel.descriptor.dpa4_nn.radial import (
-    parse_basis_type,
-)
 from deepmd.dpmodel.descriptor.dpa4c import DescrptDPA4C as DescrptDPA4CDP
 from deepmd.pt_expt.common import (
     torch_module,
@@ -69,28 +66,18 @@ def _promote_trainable_tree(module: torch.nn.Module) -> torch.nn.Module:
         The same module, for use as an expression.
     """
     for submodule in module.modules():
-        if not getattr(submodule, "trainable", True):
-            continue
+        trainable = bool(getattr(submodule, "trainable", True))
         for name in _TRAINABLE_ATTRS.get(type(submodule).__name__, ()):
             value = submodule._buffers.get(name)
             if value is None or not value.is_floating_point():
                 continue
             del submodule._buffers[name]
-            setattr(submodule, name, torch.nn.Parameter(value, requires_grad=True))
+            setattr(submodule, name, torch.nn.Parameter(value, requires_grad=trainable))
 
     for submodule in module.modules():
         if not getattr(submodule, "trainable", True):
             for parameter in submodule.parameters(recurse=True):
                 parameter.requires_grad_(False)
-    # A ``/fix`` radial basis keeps its frequencies or centres, as in the pt
-    # backend's DPA4.  The parameter keeps its name, so checkpoints of either
-    # form load under the other.
-    for submodule in module.modules():
-        if (
-            type(submodule).__name__ == "RadialBasis"
-            and parse_basis_type(submodule.basis_type)[1]
-        ):
-            submodule.adam_freqs.requires_grad_(False)
     return module
 
 
diff --git a/deepmd/pt_expt/kernels/cuda/dpa4/edge_radial.py b/deepmd/pt_expt/kernels/cuda/dpa4/edge_radial.py
index ab44f426d0..6cb21b4cb9 100644
--- a/deepmd/pt_expt/kernels/cuda/dpa4/edge_radial.py
+++ b/deepmd/pt_expt/kernels/cuda/dpa4/edge_radial.py
@@ -68,9 +68,7 @@ def series_coefficients(exponent: int) -> tuple[float, ...]:
 
 def supported(exponent_env: int, exponent_rbf: int) -> bool:
     """Whether the envelope orders fit the staged limit, including a bare basis."""
-    return 2 <= exponent_env <= _MAX_SERIES and (
-        exponent_rbf == 0 or 2 <= exponent_rbf <= _MAX_SERIES
-    )
+    return 1 <= exponent_env <= _MAX_SERIES and 0 <= exponent_rbf <= _MAX_SERIES
 
 
 def _forward_fake(
@@ -203,9 +201,11 @@ class EdgeRadialCuda:
     """
 
     def __init__(self, envelope: Any, basis: Any) -> None:
+        ensure_registered()
         self._envelope = envelope
         self._basis = basis
         self._rcut = float(envelope.rcut)
+        self._gaussian_coeff = float(basis.gaussian_coeff)
         self._basis_type = BESSEL if basis.basis_family == "bessel" else GAUSSIAN
         self._env = series_coefficients(envelope.p)
         self._rbf = series_coefficients(basis.exponent)
@@ -245,14 +245,14 @@ def __call__(
             ``C3CutoffEnvelope`` and ``RadialBasis`` scaled by ``keep``.
         """
         env_series, rbf_series = self.series(edge_len.device)
-        return edge_radial(
+        return torch.ops.deepmd.dpa4_edge_radial(
             edge_len,
             keep,
             self._basis.adam_freqs,
             env_series,
             rbf_series,
             self._rcut,
-            float(self._basis.gaussian_coeff),
+            self._gaussian_coeff,
             self._basis_type,
         )
 
diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py
index 0de0d1f7ab..4d2a99fbba 100644
--- a/deepmd/utils/argcheck.py
+++ b/deepmd/utils/argcheck.py
@@ -6470,7 +6470,11 @@ def gen_json_schema(multi_task: bool = False) -> str:
         gen_args(multi_task=multi_task),
         doc=f"DeePMD-kit {__version__}",
     )
-    return json.dumps(with_model_presets(generate_json_schema(arg), multi_task))
+    return json.dumps(
+        with_model_presets(
+            generate_json_schema(arg), generate_json_schema(model_args()), multi_task
+        )
+    )
 
 
 def _check_dpa3_chg_spin_migration(data: dict[str, Any]) -> None:
diff --git a/deepmd/utils/json_schema.py b/deepmd/utils/json_schema.py
index 027d31a2d1..1a080ec777 100644
--- a/deepmd/utils/json_schema.py
+++ b/deepmd/utils/json_schema.py
@@ -7,7 +7,7 @@
     Any,
 )
 
-from deepmd.utils.model_preset import (
+from deepmd.utils.model_preset_data import (
     MODEL_PRESETS,
 )
 
@@ -140,13 +140,17 @@ def _model_selector(
     }
 
 
-def with_model_presets(schema: dict[str, Any], multi_task: bool) -> dict[str, Any]:
+def with_model_presets(
+    schema: dict[str, Any], model_schema: dict[str, Any], multi_task: bool
+) -> dict[str, Any]:
     """Add raw preset input forms to a generated training schema.
 
     Parameters
     ----------
     schema : dict
         Training schema generated from the argument definitions, updated in place.
+    model_schema : dict
+        Schema of one non-repeating model argument.
     multi_task : bool
         Whether model entries live in a multi-task ``model_dict``.
 
@@ -156,11 +160,17 @@ def with_model_presets(schema: dict[str, Any], multi_task: bool) -> dict[str, An
         Schema accepting presets and partial overrides alongside ordinary inputs.
     """
     model = schema["properties"]["model"]
+    base = {
+        key: value
+        for key, value in model_schema.items()
+        if key not in ("$schema", "$id", "title")
+    }
     if multi_task:
-        branches = model["properties"]["model_dict"]
-        base = branches.pop("items")
-    else:
-        base = model
+        branches = {
+            "type": "object",
+            "description": model["properties"]["model_dict"].get("description", ""),
+        }
+        model["properties"]["model_dict"] = branches
 
     names = sorted(MODEL_PRESETS)
     definitions = {
diff --git a/deepmd/utils/model_preset.py b/deepmd/utils/model_preset.py
index 7fae9a2714..66d59ffe29 100644
--- a/deepmd/utils/model_preset.py
+++ b/deepmd/utils/model_preset.py
@@ -39,6 +39,10 @@
     descrpt_args_plugin,
     fitting_args_plugin,
 )
+from deepmd.utils.model_preset_data import (
+    MODEL_PRESETS,
+    PERIODIC_TABLE,
+)
 
 log = logging.getLogger(__name__)
 
@@ -49,22 +53,6 @@
     "get_model_preset",
 ]
 
-# fmt: off
-PERIODIC_TABLE: tuple[str, ...] = (
-    "H", "He", "Li", "Be", "B", "C", "N", "O", "F", "Ne", "Na", "Mg", "Al",
-    "Si", "P", "S", "Cl", "Ar", "K", "Ca", "Sc", "Ti", "V", "Cr", "Mn", "Fe",
-    "Co", "Ni", "Cu", "Zn", "Ga", "Ge", "As", "Se", "Br", "Kr", "Rb", "Sr",
-    "Y", "Zr", "Nb", "Mo", "Tc", "Ru", "Rh", "Pd", "Ag", "Cd", "In", "Sn",
-    "Sb", "Te", "I", "Xe", "Cs", "Ba", "La", "Ce", "Pr", "Nd", "Pm", "Sm",
-    "Eu", "Gd", "Tb", "Dy", "Ho", "Er", "Tm", "Yb", "Lu", "Hf", "Ta", "W",
-    "Re", "Os", "Ir", "Pt", "Au", "Hg", "Tl", "Pb", "Bi", "Po", "At", "Rn",
-    "Fr", "Ra", "Ac", "Th", "Pa", "U", "Np", "Pu", "Am", "Cm", "Bk", "Cf",
-    "Es", "Fm", "Md", "No", "Lr", "Rf", "Db", "Sg", "Bh", "Hs", "Mt", "Ds",
-    "Rg", "Cn", "Nh", "Fl", "Mc", "Lv", "Ts", "Og",
-)
-"""The 118 element symbols in atomic-number order, the ``type_map`` of every preset."""
-# fmt: on
-
 # Regions a preset may define. An expanded configuration lists the regions its
 # preset defines first, in this order, followed by the remaining explicit
 # entries in their original order.
@@ -77,246 +65,6 @@
     "fitting_net": fitting_args_plugin,
 }
 
-# === DPA4 (SeZM) ===
-# Descriptor and fitting options shared by every DPA4 grade and version.
-_DPA4_DESCRIPTOR: dict[str, Any] = {
-    "type": "dpa4",
-    "rcut": 6.0,
-    "n_radial": 16,
-    "use_env_seed": True,
-    "mmax": 1,
-    "radial_so2_mode": "degree_channel",
-    "focus_dim": 0,
-    "n_atten_head": 1,
-    "message_node_so3": True,
-    "ffn_neurons": 0,
-    "ffn_so3_grid": True,
-    "grid_mlp": False,
-    "grid_branch": [0, 0, 1],
-    "ffn_blocks": 1,
-    "so3_readout": "mlp",
-    "precision": "float32",
-}
-_DPA4_FITTING: dict[str, Any] = {
-    "type": "dpa4_ener",
-    "neuron": [0],
-    "precision": "float32",
-}
-# Scaling knobs of each grade.
-_DPA4_GRADES: dict[str, dict[str, dict[str, Any]]] = {
-    "nano": {
-        "descriptor": {
-            "channels": 32,
-            "lmax": 1,
-            "n_blocks": 2,
-            "mixing_layers": 3,
-            "radial_so2_mode": "none",
-            "n_focus": 1,
-        },
-    },
-    "mini": {
-        "descriptor": {
-            "channels": 32,
-            "lmax": 2,
-            "n_blocks": 2,
-            "mixing_layers": 3,
-            "radial_so2_rank": 1,
-            "n_focus": 1,
-        },
-    },
-    "neo": {
-        "descriptor": {
-            "channels": 32,
-            "lmax": 3,
-            "n_blocks": 2,
-            "mixing_layers": 3,
-            "radial_so2_rank": 1,
-            "n_focus": 2,
-        },
-    },
-    "air": {
-        "descriptor": {
-            "channels": 64,
-            "lmax": 3,
-            "n_blocks": 3,
-            "mixing_layers": 4,
-            "radial_so2_rank": 1,
-            "n_focus": 1,
-        },
-    },
-    "plus": {
-        "descriptor": {
-            "channels": 64,
-            "lmax": 4,
-            "n_blocks": 4,
-            "mixing_layers": 4,
-            "radial_so2_rank": 2,
-            "n_focus": 1,
-        },
-    },
-    "pro": {
-        "descriptor": {
-            "channels": 64,
-            "lmax": 5,
-            "n_blocks": 6,
-            "mixing_layers": 4,
-            "radial_so2_rank": 2,
-            "n_focus": 2,
-            "so3_readout": "none",
-        },
-    },
-    "max": {
-        "descriptor": {
-            "channels": 96,
-            "lmax": 6,
-            "n_blocks": 8,
-            "mixing_layers": 4,
-            "radial_so2_rank": 4,
-            "n_focus": 2,
-            "so3_readout": "none",
-        },
-    },
-    "ultra": {
-        "descriptor": {
-            "channels": 128,
-            "lmax": 6,
-            "n_blocks": 10,
-            "mixing_layers": 4,
-            "radial_so2_rank": 4,
-            "n_focus": 3,
-            "message_node_so3": False,
-            "ffn_so3_grid": False,
-            "grid_branch": 0,
-        },
-    },
-}
-# Options that changed with each version, and the grades the version ships.
-_DPA4_VERSIONS: dict[str, dict[str, Any]] = {
-    # Channel RMSNorm on every cutoff-vanishing branch; post-norm after the
-    # SO(2) branch and pre-norm before the FFN branch.
-    "v20260820": {
-        "descriptor": {
-            "edge_norm": True,
-            "sandwich_norm": [False, True, True, False],
-        },
-        "grades": ("nano", "mini", "neo", "air", "plus", "pro"),
-    },
-    # Radial-site RMSNorm removed; pre-norm before both the SO(2) and the FFN
-    # branch.
-    "v20260901": {
-        "descriptor": {
-            "edge_norm": [False, True, True],
-            "sandwich_norm": [True, False, True, False],
-        },
-        "grades": ("nano", "mini", "neo", "air", "plus", "pro", "max", "ultra"),
-    },
-    # One C^3 envelope on the messages with no envelope on the radial basis,
-    # and fixed Gaussian centres in place of trainable Bessel frequencies.
-    "v20260911": {
-        "descriptor": {
-            "edge_norm": [False, True, True],
-            "sandwich_norm": [True, False, True, False],
-            "env_exp": 5,
-            "basis_type": "gaussian/fix",
-        },
-        "grades": ("nano", "mini", "neo", "air", "plus", "pro", "max", "ultra"),
-    },
-}
-
-# === DPA4C ===
-# DPA4C is a descriptor of the standard model, so no model ``type`` is set.
-_DPA4C_DESCRIPTOR: dict[str, Any] = {
-    "type": "dpa4c",
-    "rcut": 6.0,
-    "precision": "float32",
-}
-_DPA4C_FITTING: dict[str, Any] = {
-    "type": "ener",
-    "resnet_dt": False,
-    "activation_function": "silu",
-    "precision": "float32",
-}
-_DPA4C_GRADES: dict[str, dict[str, dict[str, Any]]] = {
-    "nano": {
-        "descriptor": {"channels": 8, "lmax": 2, "radial_modes": 0},
-        "fitting_net": {"neuron": [96, 96, 96]},
-    },
-    "mini": {
-        "descriptor": {"channels": 32, "lmax": 2, "radial_modes": 0},
-        "fitting_net": {"neuron": [192, 192, 192]},
-    },
-    "neo": {
-        "descriptor": {"channels": 64, "lmax": 2, "radial_modes": 0},
-        "fitting_net": {"neuron": [256, 256, 256]},
-    },
-    "air": {
-        "descriptor": {"channels": 64, "lmax": 3, "radial_modes": 4},
-        "fitting_net": {"neuron": [256, 256, 256]},
-    },
-    "plus": {
-        "descriptor": {"channels": 128, "lmax": 3, "radial_modes": 4},
-        "fitting_net": {"neuron": [384, 384, 384]},
-    },
-}
-_DPA4C_VERSIONS: dict[str, dict[str, Any]] = {
-    "v20260901": {"grades": ("nano", "mini", "neo", "air", "plus")},
-    # Fixed Gaussian centres in place of trainable Bessel frequencies.
-    "v20260911": {
-        "descriptor": {"basis_type": "gaussian/fix"},
-        "grades": ("nano", "mini", "neo", "air", "plus"),
-    },
-}
-
-
-def _build_family(
-    family: str,
-    model_options: dict[str, Any],
-    descriptor: dict[str, Any],
-    fitting_net: dict[str, Any],
-    grades: dict[str, dict[str, dict[str, Any]]],
-    versions: dict[str, dict[str, Any]],
-) -> dict[str, dict[str, Any]]:
-    """
-    Compose the presets ``<family>-<grade>-<version>`` of one family.
-
-    ``model_options`` holds the model-level entries shared by every preset of
-    the family (the model ``type`` for DPA4, nothing for DPA4C). ``descriptor``
-    and ``fitting_net`` hold the options shared by every grade and version.
-    Each grade in ``grades`` adds its own ``descriptor`` and ``fitting_net``
-    options, and each version in ``versions`` adds the ``descriptor`` options
-    it changed and names the grades it ships.
-    """
-    presets: dict[str, dict[str, Any]] = {}
-    for version, spec in versions.items():
-        for grade in spec["grades"]:
-            presets[f"{family}-{grade}-{version}"] = {
-                **model_options,
-                "type_map": list(PERIODIC_TABLE),
-                "descriptor": {
-                    **descriptor,
-                    **spec.get("descriptor", {}),
-                    **grades[grade].get("descriptor", {}),
-                },
-                "fitting_net": {**fitting_net, **grades[grade].get("fitting_net", {})},
-            }
-    return presets
-
-
-MODEL_PRESETS: dict[str, dict[str, Any]] = {
-    **_build_family(
-        "dpa4",
-        {"type": "dpa4"},
-        _DPA4_DESCRIPTOR,
-        _DPA4_FITTING,
-        _DPA4_GRADES,
-        _DPA4_VERSIONS,
-    ),
-    **_build_family(
-        "dpa4c", {}, _DPA4C_DESCRIPTOR, _DPA4C_FITTING, _DPA4C_GRADES, _DPA4C_VERSIONS
-    ),
-}
-"""All presets keyed by name; each value holds the regions the preset defines."""
-
 
 def get_model_preset(name: str) -> dict[str, Any]:
     """
diff --git a/deepmd/utils/model_preset_data.py b/deepmd/utils/model_preset_data.py
new file mode 100644
index 0000000000..5363a75324
--- /dev/null
+++ b/deepmd/utils/model_preset_data.py
@@ -0,0 +1,262 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Declarative architecture tables for named DPA4-family model presets."""
+
+from typing import (
+    Any,
+)
+
+# fmt: off
+PERIODIC_TABLE: tuple[str, ...] = (
+    "H", "He", "Li", "Be", "B", "C", "N", "O", "F", "Ne", "Na", "Mg", "Al",
+    "Si", "P", "S", "Cl", "Ar", "K", "Ca", "Sc", "Ti", "V", "Cr", "Mn", "Fe",
+    "Co", "Ni", "Cu", "Zn", "Ga", "Ge", "As", "Se", "Br", "Kr", "Rb", "Sr",
+    "Y", "Zr", "Nb", "Mo", "Tc", "Ru", "Rh", "Pd", "Ag", "Cd", "In", "Sn",
+    "Sb", "Te", "I", "Xe", "Cs", "Ba", "La", "Ce", "Pr", "Nd", "Pm", "Sm",
+    "Eu", "Gd", "Tb", "Dy", "Ho", "Er", "Tm", "Yb", "Lu", "Hf", "Ta", "W",
+    "Re", "Os", "Ir", "Pt", "Au", "Hg", "Tl", "Pb", "Bi", "Po", "At", "Rn",
+    "Fr", "Ra", "Ac", "Th", "Pa", "U", "Np", "Pu", "Am", "Cm", "Bk", "Cf",
+    "Es", "Fm", "Md", "No", "Lr", "Rf", "Db", "Sg", "Bh", "Hs", "Mt", "Ds",
+    "Rg", "Cn", "Nh", "Fl", "Mc", "Lv", "Ts", "Og",
+)
+"""The 118 element symbols in atomic-number order, the ``type_map`` of every preset."""
+# fmt: on
+
+# === DPA4 (SeZM) ===
+# Descriptor and fitting options shared by every DPA4 grade and version.
+_DPA4_DESCRIPTOR: dict[str, Any] = {
+    "type": "dpa4",
+    "rcut": 6.0,
+    "n_radial": 16,
+    "use_env_seed": True,
+    "mmax": 1,
+    "radial_so2_mode": "degree_channel",
+    "focus_dim": 0,
+    "n_atten_head": 1,
+    "message_node_so3": True,
+    "ffn_neurons": 0,
+    "ffn_so3_grid": True,
+    "grid_mlp": False,
+    "grid_branch": [0, 0, 1],
+    "ffn_blocks": 1,
+    "so3_readout": "mlp",
+    "precision": "float32",
+}
+_DPA4_FITTING: dict[str, Any] = {
+    "type": "dpa4_ener",
+    "neuron": [0],
+    "precision": "float32",
+}
+# Scaling knobs of each grade.
+_DPA4_GRADES: dict[str, dict[str, dict[str, Any]]] = {
+    "nano": {
+        "descriptor": {
+            "channels": 32,
+            "lmax": 1,
+            "n_blocks": 2,
+            "mixing_layers": 3,
+            "radial_so2_mode": "none",
+            "n_focus": 1,
+        },
+    },
+    "mini": {
+        "descriptor": {
+            "channels": 32,
+            "lmax": 2,
+            "n_blocks": 2,
+            "mixing_layers": 3,
+            "radial_so2_rank": 1,
+            "n_focus": 1,
+        },
+    },
+    "neo": {
+        "descriptor": {
+            "channels": 32,
+            "lmax": 3,
+            "n_blocks": 2,
+            "mixing_layers": 3,
+            "radial_so2_rank": 1,
+            "n_focus": 2,
+        },
+    },
+    "air": {
+        "descriptor": {
+            "channels": 64,
+            "lmax": 3,
+            "n_blocks": 3,
+            "mixing_layers": 4,
+            "radial_so2_rank": 1,
+            "n_focus": 1,
+        },
+    },
+    "plus": {
+        "descriptor": {
+            "channels": 64,
+            "lmax": 4,
+            "n_blocks": 4,
+            "mixing_layers": 4,
+            "radial_so2_rank": 2,
+            "n_focus": 1,
+        },
+    },
+    "pro": {
+        "descriptor": {
+            "channels": 64,
+            "lmax": 5,
+            "n_blocks": 6,
+            "mixing_layers": 4,
+            "radial_so2_rank": 2,
+            "n_focus": 2,
+            "so3_readout": "none",
+        },
+    },
+    "max": {
+        "descriptor": {
+            "channels": 96,
+            "lmax": 6,
+            "n_blocks": 8,
+            "mixing_layers": 4,
+            "radial_so2_rank": 4,
+            "n_focus": 2,
+            "so3_readout": "none",
+        },
+    },
+    "ultra": {
+        "descriptor": {
+            "channels": 128,
+            "lmax": 6,
+            "n_blocks": 10,
+            "mixing_layers": 4,
+            "radial_so2_rank": 4,
+            "n_focus": 3,
+            "message_node_so3": False,
+            "ffn_so3_grid": False,
+            "grid_branch": 0,
+        },
+    },
+}
+# Options that changed with each version, and the grades the version ships.
+_DPA4_VERSIONS: dict[str, dict[str, Any]] = {
+    # Channel RMSNorm on every cutoff-vanishing branch; post-norm after the
+    # SO(2) branch and pre-norm before the FFN branch.
+    "v20260820": {
+        "descriptor": {
+            "edge_norm": True,
+            "sandwich_norm": [False, True, True, False],
+        },
+        "grades": ("nano", "mini", "neo", "air", "plus", "pro"),
+    },
+    # Radial-site RMSNorm removed; pre-norm before both the SO(2) and the FFN
+    # branch.
+    "v20260901": {
+        "descriptor": {
+            "edge_norm": [False, True, True],
+            "sandwich_norm": [True, False, True, False],
+        },
+        "grades": ("nano", "mini", "neo", "air", "plus", "pro", "max", "ultra"),
+    },
+    # One C^3 envelope on the messages with no envelope on the radial basis,
+    # and fixed Gaussian centres in place of trainable Bessel frequencies.
+    "v20260911": {
+        "descriptor": {
+            "edge_norm": [False, True, True],
+            "sandwich_norm": [True, False, True, False],
+            "env_exp": 5,
+            "basis_type": "gaussian/fix",
+        },
+        "grades": ("nano", "mini", "neo", "air", "plus", "pro", "max", "ultra"),
+    },
+}
+
+# === DPA4C ===
+# DPA4C is a descriptor of the standard model, so no model ``type`` is set.
+_DPA4C_DESCRIPTOR: dict[str, Any] = {
+    "type": "dpa4c",
+    "rcut": 6.0,
+    "precision": "float32",
+}
+_DPA4C_FITTING: dict[str, Any] = {
+    "type": "ener",
+    "resnet_dt": False,
+    "activation_function": "silu",
+    "precision": "float32",
+}
+_DPA4C_GRADES: dict[str, dict[str, dict[str, Any]]] = {
+    "nano": {
+        "descriptor": {"channels": 8, "lmax": 2, "radial_modes": 0},
+        "fitting_net": {"neuron": [96, 96, 96]},
+    },
+    "mini": {
+        "descriptor": {"channels": 32, "lmax": 2, "radial_modes": 0},
+        "fitting_net": {"neuron": [192, 192, 192]},
+    },
+    "neo": {
+        "descriptor": {"channels": 64, "lmax": 2, "radial_modes": 0},
+        "fitting_net": {"neuron": [256, 256, 256]},
+    },
+    "air": {
+        "descriptor": {"channels": 64, "lmax": 3, "radial_modes": 4},
+        "fitting_net": {"neuron": [256, 256, 256]},
+    },
+    "plus": {
+        "descriptor": {"channels": 128, "lmax": 3, "radial_modes": 4},
+        "fitting_net": {"neuron": [384, 384, 384]},
+    },
+}
+_DPA4C_VERSIONS: dict[str, dict[str, Any]] = {
+    "v20260901": {"grades": ("nano", "mini", "neo", "air", "plus")},
+    # Fixed Gaussian centres in place of trainable Bessel frequencies.
+    "v20260911": {
+        "descriptor": {"basis_type": "gaussian/fix"},
+        "grades": ("nano", "mini", "neo", "air", "plus"),
+    },
+}
+
+
+def _build_family(
+    family: str,
+    model_options: dict[str, Any],
+    descriptor: dict[str, Any],
+    fitting_net: dict[str, Any],
+    grades: dict[str, dict[str, dict[str, Any]]],
+    versions: dict[str, dict[str, Any]],
+) -> dict[str, dict[str, Any]]:
+    """
+    Compose the presets ``<family>-<grade>-<version>`` of one family.
+
+    ``model_options`` holds the model-level entries shared by every preset of
+    the family (the model ``type`` for DPA4, nothing for DPA4C). ``descriptor``
+    and ``fitting_net`` hold the options shared by every grade and version.
+    Each grade in ``grades`` adds its own ``descriptor`` and ``fitting_net``
+    options, and each version in ``versions`` adds the ``descriptor`` options
+    it changed and names the grades it ships.
+    """
+    presets: dict[str, dict[str, Any]] = {}
+    for version, spec in versions.items():
+        for grade in spec["grades"]:
+            presets[f"{family}-{grade}-{version}"] = {
+                **model_options,
+                "type_map": list(PERIODIC_TABLE),
+                "descriptor": {
+                    **descriptor,
+                    **spec.get("descriptor", {}),
+                    **grades[grade].get("descriptor", {}),
+                },
+                "fitting_net": {**fitting_net, **grades[grade].get("fitting_net", {})},
+            }
+    return presets
+
+
+MODEL_PRESETS: dict[str, dict[str, Any]] = {
+    **_build_family(
+        "dpa4",
+        {"type": "dpa4"},
+        _DPA4_DESCRIPTOR,
+        _DPA4_FITTING,
+        _DPA4_GRADES,
+        _DPA4_VERSIONS,
+    ),
+    **_build_family(
+        "dpa4c", {}, _DPA4C_DESCRIPTOR, _DPA4C_FITTING, _DPA4C_GRADES, _DPA4C_VERSIONS
+    ),
+}
+"""All presets keyed by name; each value holds the regions the preset defines."""
diff --git a/source/op/pt/dpa4/edge_radial.cu b/source/op/pt/dpa4/edge_radial.cu
index 0300f61061..1bd73a7d55 100644
--- a/source/op/pt/dpa4/edge_radial.cu
+++ b/source/op/pt/dpa4/edge_radial.cu
@@ -126,9 +126,11 @@ __global__ __launch_bounds__(kThreads) void edge_radial_fwd_kernel(
     const float scale = mask * e2;
     float* row = rbf + e * static_cast<long>(n_radial);
     if (basis == kBessel) {
-      const float inv_r = 1.f / r;
+      // The sinc limit at zero distance is phi(0) = f.
+      const float inv_r = r == 0.f ? 0.f : 1.f / r;
       for (int n = 0; n < n_radial; ++n) {
-        row[n] = scale * sinf(r * s_freq[n]) * inv_r;
+        const float phi = r == 0.f ? s_freq[n] : sinf(r * s_freq[n]) * inv_r;
+        row[n] = scale * phi;
       }
     } else {
       for (int n = 0; n < n_radial; ++n) {
@@ -183,7 +185,7 @@ __global__ __launch_bounds__(kThreads) void edge_radial_bwd_kernel(
     float d2 = 0.f;
     envelope_pair(r, inv_rcut, s_rbf, rbf_order, e2, d2);
     const float* row = grad_rbf + e * static_cast<long>(n_radial);
-    const float inv_r = 1.f / r;
+    const float inv_r = r == 0.f ? 0.f : 1.f / r;
     for (int n = 0; n < n_radial; ++n) {
       float phi = 0.f;
       float dphi = 0.f;
@@ -191,7 +193,8 @@ __global__ __launch_bounds__(kThreads) void edge_radial_bwd_kernel(
         float sine = 0.f;
         float cosine = 0.f;
         sincosf(r * s_freq[n], &sine, &cosine);
-        phi = sine * inv_r;
+        // With phi(0) = f and inv_r = 0, the derivative limit is zero.
+        phi = r == 0.f ? s_freq[n] : sine * inv_r;
         // d/dr [sin(r f) / r] = (f cos(r f) - sin(r f) / r) / r. The two terms
         // cancel to leading order at large ``r f``, so the difference is formed
         // with a fused multiply-add to keep the rounding to one step.
@@ -220,12 +223,8 @@ void check_inputs(const torch::Tensor& edge_len,
       env_series.numel() <= kMaxSeries && rbf_series.numel() <= kMaxSeries,
       "dpa4_edge_radial: envelope order beyond the staged limit");
   TORCH_CHECK(
-      env_series.numel() >= 2,
-      "dpa4_edge_radial: the edge envelope series needs at least two terms");
-  TORCH_CHECK(
-      rbf_series.numel() == 0 || rbf_series.numel() >= 2,
-      "dpa4_edge_radial: the basis envelope series must be empty or have "
-      "at least two terms");
+      env_series.numel() >= 1,
+      "dpa4_edge_radial: the edge envelope series needs at least one term");
   TORCH_CHECK(freqs.numel() > 0,
               "dpa4_edge_radial: the basis must be non-empty");
 }
diff --git a/source/tests/jax/test_dpa4.py b/source/tests/jax/test_dpa4.py
index 1733f6d83d..18a50dac21 100644
--- a/source/tests/jax/test_dpa4.py
+++ b/source/tests/jax/test_dpa4.py
@@ -1,6 +1,8 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 """Focused tests for JAX DPA4 descriptor trainable-state conversion."""
 
+import pytest
+
 from deepmd.jax.descriptor.dpa4 import (
     DescrptDPA4,
     _iter_object_tree,
@@ -13,7 +15,7 @@
 )
 
 
-def _make_trainable_descriptor() -> DescrptDPA4:
+def _make_trainable_descriptor(basis_type: str = "bessel") -> DescrptDPA4:
     """Build a small descriptor that enables the optional trainable leaves."""
     return DescrptDPA4(
         ntypes=2,
@@ -21,6 +23,7 @@ def _make_trainable_descriptor() -> DescrptDPA4:
         rcut=4.0,
         channels=4,
         n_radial=4,
+        basis_type=basis_type,
         lmax=1,
         mmax=1,
         n_blocks=1,
@@ -61,6 +64,17 @@ def test_optional_dpa4_weights_are_jax_parameters() -> None:
     )
 
 
+@pytest.mark.parametrize("family", ["bessel", "gaussian"])
+def test_fixed_basis_is_not_an_optimizer_parameter(family: str) -> None:
+    """Fixed basis arrays remain serializable without entering the parameter tree."""
+    descriptor = _make_trainable_descriptor(f"{family}/fix")
+    restored = DescrptDPA4.deserialize(descriptor.serialize())
+    for model in (descriptor, restored):
+        assert not model.radial_basis.trainable
+        assert not isinstance(model.radial_basis.adam_freqs, nnx.Param)
+        assert len(nnx.to_flat_state(nnx.state(model, nnx.Param))) > 0
+
+
 def test_frozen_descriptor_has_no_optimizer_visible_parameters() -> None:
     """The root freeze flag must demote every descendant parameter."""
     descriptor = DescrptDPA4(
diff --git a/source/tests/pt/model/test_descriptor_sezm_cuda.py b/source/tests/pt/model/test_descriptor_sezm_cuda.py
index 6791f99640..7500d15f4b 100644
--- a/source/tests/pt/model/test_descriptor_sezm_cuda.py
+++ b/source/tests/pt/model/test_descriptor_sezm_cuda.py
@@ -875,6 +875,53 @@ def test_declines_a_mismatched_cutoff(self) -> None:
         )
         self.assertIsNone(make_cuda_edge_radial(envelope, basis))
 
+    def test_warmed_series_cache_survives_compile(self) -> None:
+        """Dynamo lifts cached eager constants into the same-device compiled graph."""
+        _, _, fused = self._modules("gaussian")
+        r, keep = self._distances(32, self.RCUT)
+        expected = fused(r, keep)
+        cached = fused.series(r.device)
+        compiled = torch.compile(fused, fullgraph=True)
+        actual = compiled(r, keep)
+        for got, want in zip(actual, expected, strict=True):
+            torch.testing.assert_close(got, want)
+        self.assertIs(fused.series(r.device), cached)
+
+    def test_one_term_envelopes_match_the_reference(self) -> None:
+        """Every positive envelope order is eligible, including one-term series."""
+        for basis_type in ("bessel", "gaussian"):
+            for basis_exponent in (0, 1, 7):
+                with self.subTest(basis=basis_type, exponent=basis_exponent):
+                    envelope = C3CutoffEnvelope(
+                        rcut=self.RCUT, exponent=1, dtype=torch.float32
+                    ).cuda()
+                    basis = RadialBasis(
+                        rcut=self.RCUT,
+                        basis_type=basis_type,
+                        n_radial=8,
+                        exponent=basis_exponent,
+                        dtype=torch.float32,
+                    ).cuda()
+                    fused = make_cuda_edge_radial(envelope, basis)
+                    self.assertIsNotNone(fused)
+                    r, keep = self._distances(128, self.RCUT)
+                    r[0] = 0.0
+                    keep[0] = 1.0
+                    r.requires_grad_(True)
+                    expected = (envelope(r) * keep, basis(r) * keep)
+                    actual = fused(r, keep)
+                    for got, want in zip(actual, expected, strict=True):
+                        torch.testing.assert_close(got, want, rtol=5e-5, atol=5e-6)
+                    expected_grad = torch.autograd.grad(
+                        sum(value.sum() for value in expected), r
+                    )[0]
+                    actual_grad = torch.autograd.grad(
+                        sum(value.sum() for value in actual), r
+                    )[0]
+                    torch.testing.assert_close(
+                        actual_grad, expected_grad, rtol=5e-5, atol=5e-6
+                    )
+
 
 @unittest.skipUnless(_IMPORT_OK, "requires the pt_expt CUDA bindings")
 class TestSeZMConvCudaGate(unittest.TestCase):
diff --git a/source/tests/pt/model/test_dpa4_dpmodel_parity.py b/source/tests/pt/model/test_dpa4_dpmodel_parity.py
index 80c8119c1f..edc7872910 100644
--- a/source/tests/pt/model/test_dpa4_dpmodel_parity.py
+++ b/source/tests/pt/model/test_dpa4_dpmodel_parity.py
@@ -3234,7 +3234,6 @@ def _build_real_edge_caches(
         deg_norm_floor=deg_norm_floor,
         edge_envelope=pt_env,
         radial_basis=pt_rb,
-        n_radial=n_radial,
         random_gamma=random_gamma,
         wigner_calc=pt_wig,
     )
diff --git a/source/tests/pt_expt/descriptor/test_dpa4.py b/source/tests/pt_expt/descriptor/test_dpa4.py
index 20a649408b..1504708050 100644
--- a/source/tests/pt_expt/descriptor/test_dpa4.py
+++ b/source/tests/pt_expt/descriptor/test_dpa4.py
@@ -346,10 +346,11 @@ def test_trainable_parameters(self) -> None:
             assert name not in param_names
         # wigner tables must never be trainable
         assert not any("wigner" in n.lower() for n in param_names)
-        # all promoted parameters are float and trainable
+        # Declared weights remain parameters when their owning module freezes them.
         for name, p in param_names.items():
             assert p.is_floating_point(), name
-            assert p.requires_grad, name
+            owner = dd0.get_submodule(name.rpartition(".")[0])
+            assert p.requires_grad == bool(getattr(owner, "trainable", True)), name
 
     @pytest.mark.parametrize(
         "via_deserialize", [False, True]
diff --git a/source/tests/tf2/test_dpa4.py b/source/tests/tf2/test_dpa4.py
index 8dd1670449..6c5e2bf686 100644
--- a/source/tests/tf2/test_dpa4.py
+++ b/source/tests/tf2/test_dpa4.py
@@ -33,7 +33,7 @@
 )
 
 
-def _make_trainable_descriptor() -> DescrptDPA4:
+def _make_trainable_descriptor(basis_type: str = "bessel") -> DescrptDPA4:
     """Build a small descriptor that enables the optional trainable leaves."""
     return DescrptDPA4(
         ntypes=2,
@@ -41,6 +41,7 @@ def _make_trainable_descriptor() -> DescrptDPA4:
         rcut=4.0,
         channels=4,
         n_radial=4,
+        basis_type=basis_type,
         lmax=1,
         mmax=1,
         n_blocks=1,
@@ -93,6 +94,21 @@ def test_optional_dpa4_weights_are_tf2_trainable_variables() -> None:
     _assert_optional_weights_are_tracked(_make_trainable_descriptor())
 
 
+@pytest.mark.parametrize("family", ["bessel", "gaussian"])
+def test_fixed_basis_is_not_an_optimizer_parameter(family: str) -> None:
+    """Fixed basis variables remain trackable without entering optimization."""
+    descriptor = _make_trainable_descriptor(f"{family}/fix")
+    restored = DescrptDPA4.deserialize(descriptor.serialize())
+    for model in (descriptor, restored):
+        variable = object.__getattribute__(
+            model.radial_basis, "_tf2_adam_freqs_variable"
+        )
+        assert not model.radial_basis.trainable
+        assert not variable.trainable
+        assert any(candidate is variable for candidate in model.variables)
+        assert all(candidate is not variable for candidate in model.trainable_variables)
+
+
 def test_dpa4_deserialize_refreshes_trackable_state() -> None:
     """Serialization must preserve values and nested TensorFlow trackables."""
     descriptor = _make_trainable_descriptor()

From 7afe4f6f1fe5ad1120a9a79c587deda801304035 Mon Sep 17 00:00:00 2001
From: OutisLi <LTC201806070316@gmail.com>
Date: Sat, 12 Sep 2026 10:48:36 +0800
Subject: [PATCH 6/8] test(pt-expt): isolate CUDA binding checks from
 registration

---
 source/tests/pt_expt/descriptor/test_dpa4_accelerated.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/source/tests/pt_expt/descriptor/test_dpa4_accelerated.py b/source/tests/pt_expt/descriptor/test_dpa4_accelerated.py
index e4f7b6c37e..ec4e3b5630 100644
--- a/source/tests/pt_expt/descriptor/test_dpa4_accelerated.py
+++ b/source/tests/pt_expt/descriptor/test_dpa4_accelerated.py
@@ -92,7 +92,7 @@ def _make_descriptor(
 )
 @pytest.mark.parametrize("env_exp", [5, [7, 5]])
 def test_fp32_only_cuda_bindings(
-    monkeypatch,
+    monkeypatch: pytest.MonkeyPatch,
     precision: str,
     expected_bound: bool,
     env_exp: int | list[int],
@@ -107,6 +107,8 @@ def test_fp32_only_cuda_bindings(
         monkeypatch.setenv(name, "0")
     monkeypatch.setenv("DP_CUDA_INFER", "1")
     monkeypatch.setattr(edge_radial, "op_available", lambda: True)
+    # Binding eligibility is independent of native operator registration.
+    monkeypatch.setattr(edge_radial, "ensure_registered", lambda: None)
     monkeypatch.setattr(grid_pair, "op_available", lambda: True)
     monkeypatch.setattr(zonal_scatter, "op_available", lambda: True)
 

From b77ae9508ced4400479f4eb9db31db2b5ea4fb81 Mon Sep 17 00:00:00 2001
From: OutisLi <LTC201806070316@gmail.com>
Date: Sat, 12 Sep 2026 11:38:47 +0800
Subject: [PATCH 7/8] refactor(presets): align supported versions and examples

---
 deepmd/utils/model_preset.py                  |   9 +-
 deepmd/utils/model_preset_data.py             |  19 +--
 doc/model/dpa4.md                             |  46 ++----
 doc/model/dpa4c.md                            |  49 +++---
 doc/model/overall.md                          |   2 +-
 examples/water/dpa4/README.md                 |   4 +-
 .../water/dpa4/input_multitask_preset.json    |   2 +-
 examples/water/dpa4/input_preset.json         |   2 +-
 examples/water/dpa4c/README.md                |  10 +-
 examples/water/dpa4c/input.json               |  22 +--
 source/tests/common/test_model_preset.py      | 140 ++++++++++--------
 11 files changed, 135 insertions(+), 170 deletions(-)

diff --git a/deepmd/utils/model_preset.py b/deepmd/utils/model_preset.py
index 66d59ffe29..1502fc2d48 100644
--- a/deepmd/utils/model_preset.py
+++ b/deepmd/utils/model_preset.py
@@ -11,11 +11,10 @@
 architecture (``use_amp``, ``seed``, ``sel``, charge and spin conditioning,
 ...) are supplied alongside the preset.
 
-The tables below are organised per family. Shared descriptor and fitting
-options are written once, every grade lists only its scaling knobs, and every
-version lists the options it changed together with the grades it ships. A new
-version therefore adds one entry to the family's version table, and a new
-grade adds one entry to its grade table. Existing versions are never edited.
+Preset tables in :mod:`deepmd.utils.model_preset_data` are organised per
+family. Shared descriptor and fitting options are written once, every grade
+lists its scaling knobs, and every version defines its descriptor options
+and available grades.
 
 In the multi-task layout a preset next to ``model_dict`` is the base of every
 branch and of the ``shared_dict`` entries the branches reference as
diff --git a/deepmd/utils/model_preset_data.py b/deepmd/utils/model_preset_data.py
index 5363a75324..271929c112 100644
--- a/deepmd/utils/model_preset_data.py
+++ b/deepmd/utils/model_preset_data.py
@@ -134,7 +134,7 @@
         },
     },
 }
-# Options that changed with each version, and the grades the version ships.
+# Descriptor options and available grades for each version.
 _DPA4_VERSIONS: dict[str, dict[str, Any]] = {
     # Channel RMSNorm on every cutoff-vanishing branch; post-norm after the
     # SO(2) branch and pre-norm before the FFN branch.
@@ -145,17 +145,10 @@
         },
         "grades": ("nano", "mini", "neo", "air", "plus", "pro"),
     },
-    # Radial-site RMSNorm removed; pre-norm before both the SO(2) and the FFN
-    # branch.
-    "v20260901": {
-        "descriptor": {
-            "edge_norm": [False, True, True],
-            "sandwich_norm": [True, False, True, False],
-        },
-        "grades": ("nano", "mini", "neo", "air", "plus", "pro", "max", "ultra"),
-    },
+    # Channel RMSNorm on FiLM and focus features, with no radial-site RMSNorm.
+    # Pre-norm before SO(2) and FFN branches, without post-norm.
     # One C^3 envelope on the messages with no envelope on the radial basis,
-    # and fixed Gaussian centres in place of trainable Bessel frequencies.
+    # and fixed Gaussian centres.
     "v20260911": {
         "descriptor": {
             "edge_norm": [False, True, True],
@@ -227,8 +220,8 @@ def _build_family(
     the family (the model ``type`` for DPA4, nothing for DPA4C). ``descriptor``
     and ``fitting_net`` hold the options shared by every grade and version.
     Each grade in ``grades`` adds its own ``descriptor`` and ``fitting_net``
-    options, and each version in ``versions`` adds the ``descriptor`` options
-    it changed and names the grades it ships.
+    options, and each version in ``versions`` defines its ``descriptor``
+    options and available grades.
     """
     presets: dict[str, dict[str, Any]] = {}
     for version, spec in versions.items():
diff --git a/doc/model/dpa4.md b/doc/model/dpa4.md
index e67e06831b..265bc133f4 100644
--- a/doc/model/dpa4.md
+++ b/doc/model/dpa4.md
@@ -85,16 +85,22 @@ unnecessary and not recommended (see [Hardware selection](#hardware-selection)).
 
 ### Presets
 
-The released DPA4 grades are available as named model presets. Setting
-`model.preset` fills in the four architecture-defining regions of the model
-section, `type`, `type_map` (all 118 elements), `descriptor` and
-`fitting_net`, from the release configuration, so an input only names the grade
-and adds what is specific to the run:
+DPA4 preset names use `dpa4-<size>-<version>`. Available sizes are listed in
+ascending computational cost:
+
+| Version     | Available sizes                                             |
+| ----------- | ----------------------------------------------------------- |
+| `v20260911` | `nano`, `mini`, `neo`, `air`, `plus`, `pro`, `max`, `ultra` |
+| `v20260820` | `nano`, `mini`, `neo`, `air`, `plus`, `pro`                 |
+
+Setting `model.preset` supplies `type`, `type_map` (all 118 elements),
+`descriptor` and `fitting_net`. An input names the preset and adds only
+run-specific settings:
 
 ```json
 {
   "model": {
-    "preset": "dpa4-nano-v20260901",
+    "preset": "dpa4-nano-v20260911",
     "type_map": [
       "O",
       "H"
@@ -118,12 +124,9 @@ take precedence over it:
 
 - `type` and `type_map` are replaced as a whole. The two-element `type_map`
   above replaces the 118-element periodic table of the preset.
-- Inside `descriptor` and `fitting_net` the merge is key by key: a key that the
-  preset defines takes the explicit value (`rcut` above, written here with the
-  value the preset has anyway), and a key it does not define is added. Options
-  that are not part of an architecture are meant to be added this way:
-  `use_amp`, `seed`, `sel`, `trainable`, and the charge and spin conditioning
-  pair `add_chg_spin_ebd` / `default_chg_spin` for molecular datasets.
+- Inside `descriptor` and `fitting_net`, explicit keys replace preset values
+  or add settings such as `use_amp`, `seed`, `sel`, `trainable`, and the
+  charge and spin conditioning pair `add_chg_spin_ebd` / `default_chg_spin`.
 
 Every explicit entry that changes a preset value is reported in the log. In
 multi-task training a `preset` next to `model_dict` is the base of every branch
@@ -133,25 +136,6 @@ a `preset` inside a branch applies to that branch alone, and shared-dictionary
 references written in a branch keep precedence over the preset. See
 `examples/water/dpa4/input_multitask_preset.json`.
 
-Preset names are `<family>-<grade>-<version>`. The version tag identifies the
-release a preset reproduces: a later release with different settings gets a new
-version, and existing presets are never changed. The available DPA4 presets
-are, in ascending cost:
-
-- `v20260911`, the current release grades: `dpa4-nano-v20260911`,
-  `dpa4-mini-v20260911`, `dpa4-neo-v20260911`, `dpa4-air-v20260911`,
-  `dpa4-plus-v20260911`, `dpa4-pro-v20260911`, `dpa4-max-v20260911` and
-  `dpa4-ultra-v20260911`. They expand the radial basis on fixed Gaussian
-  centres (`basis_type` `gaussian/fix`) and apply a single cutoff envelope to
-  the message-passing edge weights (`env_exp` 5).
-- `v20260901`, the previous release grades with trainable Bessel functions and
-  two envelopes: `dpa4-nano-v20260901`, `dpa4-mini-v20260901`,
-  `dpa4-neo-v20260901`, `dpa4-air-v20260901`, `dpa4-plus-v20260901`,
-  `dpa4-pro-v20260901`, `dpa4-max-v20260901` and `dpa4-ultra-v20260901`.
-- `v20260820`, the earlier baseline grades: `dpa4-nano-v20260820`,
-  `dpa4-mini-v20260820`, `dpa4-neo-v20260820`, `dpa4-air-v20260820`,
-  `dpa4-plus-v20260820` and `dpa4-pro-v20260820`.
-
 `examples/water/dpa4/input_preset.json` is a water example that names a preset
 instead of spelling out the architecture.
 
diff --git a/doc/model/dpa4c.md b/doc/model/dpa4c.md
index bd3ab45d5b..7d5fbc0ec5 100644
--- a/doc/model/dpa4c.md
+++ b/doc/model/dpa4c.md
@@ -125,6 +125,13 @@ carry the accuracy–cost trade-off:
   `gaussian/fix` forms keep the frequencies or centres at their initial
   values instead of training them.
 
+The fitting network is sized against the descriptor because the invariant
+output grows with `channels`. Unlike `radial_modes`, fitting width is not a free
+trade against memory: it adds per-atom activations and the derivatives saved for
+the force backward pass, so widening or deepening it costs throughput and
+capacity together. Widen it only when validation error is limited by fitting
+capacity rather than by the descriptor.
+
 > [!IMPORTANT]
 > The compressed CUDA path is compiled for `channels` in `{8, 16, 32, 64, 128}`,
 > `lmax` in `{2, 3, 4}`, and `radial_modes` in `{0, 2, 4, 8}` only. A model
@@ -132,25 +139,23 @@ carry the accuracy–cost trade-off:
 > `dp --pt-expt compress` rejects it. Choose these values with deployment in
 > mind.
 
-### Recommended configurations and presets
-
-The released grades, Nano, Mini, Neo, Air and Plus in ascending cost, pair each
-descriptor width with a fitting width sized against it. They are good starting
-points; `Neo` is the general-purpose default. Each grade is available as a named
-model preset: `dpa4c-nano-v20260911`, `dpa4c-mini-v20260911`,
-`dpa4c-neo-v20260911`, `dpa4c-air-v20260911` and `dpa4c-plus-v20260911` expand
-the radial basis on fixed Gaussian centres (`basis_type` `gaussian/fix`), and
-the `v20260901` presets of the same grades keep the trainable Bessel basis:
-setting `model.preset` fills in `type_map` (all 118 elements), `descriptor` and
-`fitting_net` from the release configuration, and entries written next to the
-preset take precedence, as a whole for `type_map` and key by key inside
-`descriptor` and `fitting_net`. Run-specific options such as `use_amp` and
-`seed` are added alongside:
+### Presets
+
+DPA4C preset names use `dpa4c-<size>-<version>`. Available sizes are listed in
+ascending computational cost; `neo` is the general-purpose starting point:
+
+| Version     | Available sizes                      |
+| ----------- | ------------------------------------ |
+| `v20260911` | `nano`, `mini`, `neo`, `air`, `plus` |
+| `v20260901` | `nano`, `mini`, `neo`, `air`, `plus` |
+
+Setting `model.preset` supplies `type_map` (all 118 elements), `descriptor`
+and `fitting_net`. Run-specific settings are written alongside the preset:
 
 ```json
 {
   "model": {
-    "preset": "dpa4c-neo-v20260901",
+    "preset": "dpa4c-neo-v20260911",
     "type_map": [
       "O",
       "H"
@@ -165,17 +170,9 @@ preset take precedence, as a whole for `type_map` and key by key inside
 }
 ```
 
-The version tag identifies the release a preset reproduces: a later release
-with different settings gets a new version, and existing presets are never
-changed. The expansion and merge rules are described on the
-[DPA4 page](dpa4.md#presets).
-
-The fitting network is sized against the descriptor because the invariant
-output grows with `channels`. Unlike `radial_modes`, fitting width is not a free
-trade against memory: it adds per-atom activations and the derivatives saved for
-the force backward pass, so widening or deepening it costs throughput and
-capacity together. Widen it only when validation error is limited by fitting
-capacity rather than by the descriptor.
+Explicit settings override the preset: `type_map` is replaced as a whole,
+while `descriptor` and `fitting_net` are merged key by key. See the
+[DPA4 preset rules](dpa4.md#presets) for expansion details and multi-task use.
 
 ## Training
 
diff --git a/doc/model/overall.md b/doc/model/overall.md
index 7b6d26cb22..2ff909ac55 100644
--- a/doc/model/overall.md
+++ b/doc/model/overall.md
@@ -47,7 +47,7 @@ The two subsections, {ref}`descriptor <model[standard]/descriptor>` and {ref}`fi
 
 The {ref}`type_map <model/type_map>` is optional, which provides the element names (but not necessarily same as the actual name of the element) of the corresponding atom types. A water model, as in this example, has two kinds of atoms. The atom types are internally recorded as integers, e.g., `0` for oxygen and `1` for hydrogen here. A mapping from the atom type to their names is provided by {ref}`type_map <model/type_map>`.
 
-Some model families ship named presets of their released architectures. Setting `preset` in the {ref}`model <model>` section fills in `type`, `type_map`, `descriptor` and `fitting_net` from the named architecture, and entries written next to it take precedence. See [DPA4](dpa4.md#presets) and [DPA4C](dpa4c.md) for the available presets and the merge rules.
+Some model families ship named presets of their released architectures. Setting `preset` in the {ref}`model <model>` section fills in `type`, `type_map`, `descriptor` and `fitting_net` from the named architecture, and entries written next to it take precedence. See [DPA4](dpa4.md#presets) and [DPA4C](dpa4c.md#presets) for the available presets and the merge rules.
 
 DeePMD-kit implements the following descriptors:
 
diff --git a/examples/water/dpa4/README.md b/examples/water/dpa4/README.md
index 8abdf05ffc..a0708cf36c 100644
--- a/examples/water/dpa4/README.md
+++ b/examples/water/dpa4/README.md
@@ -9,14 +9,14 @@ Input files:
 - `input.json`: baseline conservative energy training, using a compact
   DPA4-Mini-style parameter set.
 - `input_preset.json`: energy training with the model architecture taken from
-  the named preset `dpa4-nano-v20260901` instead of being written out.
+  the named preset `dpa4-nano-v20260911` instead of being written out.
 - `input-zbl.json`: energy training with ZBL zone bridging.
 - `input_dens.json`: direct-force denoising training.
 - `input_multitask.json`: multitask training with a shared descriptor and
   case-conditioned shared fitting network.
 - `input_multitask_preset.json`: the same multitask training with the shared
   descriptor and fitting network taken from the named preset
-  `dpa4-nano-v20260901`.
+  `dpa4-nano-v20260911`.
 - `lora_ft.json`: LoRA fine-tuning.
 - `lmp/`: compact checkpoint and LAMMPS smoke-test files.
 
diff --git a/examples/water/dpa4/input_multitask_preset.json b/examples/water/dpa4/input_multitask_preset.json
index 9bb2e5167f..4c66f581cb 100644
--- a/examples/water/dpa4/input_multitask_preset.json
+++ b/examples/water/dpa4/input_multitask_preset.json
@@ -1,7 +1,7 @@
 {
   "_comment": "DPA4-Nano multitask example with a shared descriptor and case-conditioned shared fitting network, both taken from the named preset; only run-specific keys are written out.",
   "model": {
-    "preset": "dpa4-nano-v20260901",
+    "preset": "dpa4-nano-v20260911",
     "use_compile": false,
     "enable_tf32": true,
     "shared_dict": {
diff --git a/examples/water/dpa4/input_preset.json b/examples/water/dpa4/input_preset.json
index f7fcac91a0..326af82d9e 100644
--- a/examples/water/dpa4/input_preset.json
+++ b/examples/water/dpa4/input_preset.json
@@ -1,7 +1,7 @@
 {
   "_comment": "DPA4-Nano energy-training example for the water dataset. The model architecture comes from the named preset; the explicit type_map and rcut override it and use_amp and seed supplement it.",
   "model": {
-    "preset": "dpa4-nano-v20260901",
+    "preset": "dpa4-nano-v20260911",
     "type_map": [
       "O",
       "H"
diff --git a/examples/water/dpa4c/README.md b/examples/water/dpa4c/README.md
index 1ca3aaf9a1..3537f4b048 100644
--- a/examples/water/dpa4c/README.md
+++ b/examples/water/dpa4c/README.md
@@ -8,12 +8,10 @@ backend:
 dp --pt-expt train input.json
 ```
 
-The model section spells out the `Neo` grade. The same architecture is
-available as the named preset `dpa4c-neo-v20260901`: writing
-`"preset": "dpa4c-neo-v20260901"` in the model section fills in the
-118-element `type_map`, `descriptor` and `fitting_net`, and the run-specific
-entries written next to it take precedence: the water `type_map`, `use_amp`
-and `seed`.
+The model section selects the `Nano` grade with the named preset
+`dpa4c-nano-v20260911`. The preset supplies the 118-element `type_map`,
+`descriptor` and `fitting_net`; the explicit water `type_map` replaces the
+element list, while `use_amp` and `seed` supply run-specific settings.
 
 DPA4C is built for extreme-speed molecular dynamics, so its arguments are best
 read as a budget split between two quantities: inference throughput and the
diff --git a/examples/water/dpa4c/input.json b/examples/water/dpa4c/input.json
index 6e72b2a4fc..15b8a1de87 100644
--- a/examples/water/dpa4c/input.json
+++ b/examples/water/dpa4c/input.json
@@ -1,34 +1,16 @@
 {
-  "_comment": "DPA4C energy-training example for the water dataset.",
+  "_comment": "DPA4C-Nano energy-training example for the water dataset.",
   "model": {
+    "preset": "dpa4c-nano-v20260911",
     "type_map": [
       "O",
       "H"
     ],
     "descriptor": {
-      "type": "dpa4c",
-      "_comment": "The Neo grade, the architecture of the preset dpa4c-neo-v20260901: channels and lmax fix every derived width, and radial_modes 0 leaves each ordered atom-type pair with a rescaled copy of one shared radial function.",
-      "rcut": 6.0,
-      "channels": 64,
-      "lmax": 2,
-      "basis_type": "bessel",
-      "n_radial": 16,
-      "radial_modes": 0,
-      "_comment_precision": "float32 is required by the compressed CUDA path.",
-      "precision": "float32",
       "use_amp": false,
       "seed": 42
     },
     "fitting_net": {
-      "_comment": "The fitting width paired with the Neo descriptor width.",
-      "neuron": [
-        256,
-        256,
-        256
-      ],
-      "resnet_dt": false,
-      "activation_function": "silu",
-      "precision": "float32",
       "seed": 42
     }
   },
diff --git a/source/tests/common/test_model_preset.py b/source/tests/common/test_model_preset.py
index e96a9820c7..05b92e388f 100644
--- a/source/tests/common/test_model_preset.py
+++ b/source/tests/common/test_model_preset.py
@@ -52,7 +52,7 @@ def test_every_preset_expands_to_a_valid_model(name: str) -> None:
 
 def test_explicit_entries_override_and_supplement(caplog) -> None:
     model = {
-        "preset": "dpa4-nano-v20260901",
+        "preset": "dpa4-nano-v20260911",
         "type_map": ["O", "H"],
         "descriptor": {"rcut": 5.0, "use_amp": True, "seed": 1},
         "fitting_net": {"seed": 1},
@@ -62,7 +62,7 @@ def test_explicit_entries_override_and_supplement(caplog) -> None:
     with caplog.at_level(logging.INFO, logger="deepmd.utils.model_preset"):
         expanded = expand_model_preset(model)
     assert model == original
-    preset = get_model_preset("dpa4-nano-v20260901")
+    preset = get_model_preset("dpa4-nano-v20260911")
 
     assert "preset" not in expanded
     assert expanded["type"] == "dpa4"
@@ -92,26 +92,26 @@ def test_expansion_is_idempotent_and_a_noop_without_preset() -> None:
     assert expand_model_preset(plain) is plain
     multi = {"shared_dict": {}, "model_dict": {"a": copy.deepcopy(plain)}}
     assert expand_model_preset(multi) is multi
-    expanded = expand_model_preset({"preset": "dpa4c-neo-v20260901"})
+    expanded = expand_model_preset({"preset": "dpa4c-neo-v20260911"})
     assert expand_model_preset(expanded) is expanded
 
 
 def test_preset_name_is_case_insensitive() -> None:
-    assert expand_model_preset({"preset": "DPA4-Neo-v20260901"}) == expand_model_preset(
-        {"preset": "dpa4-neo-v20260901"}
+    assert expand_model_preset({"preset": "DPA4-Neo-v20260911"}) == expand_model_preset(
+        {"preset": "dpa4-neo-v20260911"}
     )
 
 
 def test_unknown_or_malformed_preset_raises() -> None:
     with pytest.raises(ValueError, match="Unknown model preset"):
-        expand_model_preset({"preset": "dpa4-huge-v20260901"})
+        expand_model_preset({"preset": "dpa4-huge-v20260911"})
     with pytest.raises(ValueError, match="must be a string"):
         expand_model_preset({"preset": 3})
 
 
 def test_multi_task_branches_expand_with_shared_references() -> None:
     model = {
-        "preset": "dpa4-mini-v20260901",
+        "preset": "dpa4-mini-v20260911",
         "shared_dict": {
             "type_map": ["O", "H"],
             "descriptor": {"type": "dpa4", "rcut": 6.0},
@@ -123,7 +123,7 @@ def test_multi_task_branches_expand_with_shared_references() -> None:
                 "fitting_net": {"seed": 2},
             },
             "water_2": {
-                "preset": "dpa4-neo-v20260901",
+                "preset": "dpa4-neo-v20260911",
                 "type_map": "type_map",
             },
         },
@@ -131,7 +131,7 @@ def test_multi_task_branches_expand_with_shared_references() -> None:
     expanded = expand_model_preset(model)
     assert "preset" not in expanded
     # The shared descriptor is the mini descriptor with the explicit keys on top.
-    mini = get_model_preset("dpa4-mini-v20260901")
+    mini = get_model_preset("dpa4-mini-v20260911")
     assert expanded["shared_dict"]["descriptor"] == {**mini["descriptor"], "rcut": 6.0}
     assert expanded["shared_dict"]["type_map"] == ["O", "H"]
 
@@ -150,44 +150,56 @@ def test_multi_task_branches_expand_with_shared_references() -> None:
     assert water_2["descriptor"]["n_focus"] == 2
 
 
-def test_dpa4_versions_differ_only_in_normalization_options() -> None:
-    for grade in ("nano", "mini", "neo", "air", "plus", "pro"):
-        old = get_model_preset(f"dpa4-{grade}-v20260820")
-        new = get_model_preset(f"dpa4-{grade}-v20260901")
-        assert old["type"] == new["type"]
-        assert old["type_map"] == new["type_map"]
-        assert old["fitting_net"] == new["fitting_net"]
-        changed = {
-            key
-            for key in set(old["descriptor"]) | set(new["descriptor"])
-            if old["descriptor"].get(key) != new["descriptor"].get(key)
-        }
-        assert changed == {"edge_norm", "sandwich_norm"}
-    for grade in ("max", "ultra"):
-        assert f"dpa4-{grade}-v20260820" not in MODEL_PRESETS
-        assert f"dpa4-{grade}-v20260901" in MODEL_PRESETS
-
-
-def test_v20260911_fixes_the_basis_and_uses_one_envelope() -> None:
-    for grade in ("nano", "mini", "neo", "air", "plus", "pro", "max", "ultra"):
-        old = get_model_preset(f"dpa4-{grade}-v20260901")
-        new = get_model_preset(f"dpa4-{grade}-v20260911")
-        assert old["type_map"] == new["type_map"]
-        assert old["fitting_net"] == new["fitting_net"]
-        assert {
-            key: new["descriptor"][key]
-            for key in set(old["descriptor"]) | set(new["descriptor"])
-            if old["descriptor"].get(key) != new["descriptor"].get(key)
-        } == {"env_exp": 5, "basis_type": "gaussian/fix"}
-    for grade in ("nano", "mini", "neo", "air", "plus"):
-        old = get_model_preset(f"dpa4c-{grade}-v20260901")
-        new = get_model_preset(f"dpa4c-{grade}-v20260911")
-        assert old["fitting_net"] == new["fitting_net"]
-        assert {
-            key: new["descriptor"][key]
-            for key in set(old["descriptor"]) | set(new["descriptor"])
-            if old["descriptor"].get(key) != new["descriptor"].get(key)
-        } == {"basis_type": "gaussian/fix"}
+def test_preset_catalog() -> None:
+    catalog = {
+        ("dpa4", "v20260820"): ("nano", "mini", "neo", "air", "plus", "pro"),
+        ("dpa4", "v20260911"): (
+            "nano",
+            "mini",
+            "neo",
+            "air",
+            "plus",
+            "pro",
+            "max",
+            "ultra",
+        ),
+        ("dpa4c", "v20260901"): ("nano", "mini", "neo", "air", "plus"),
+        ("dpa4c", "v20260911"): ("nano", "mini", "neo", "air", "plus"),
+    }
+    assert set(MODEL_PRESETS) == {
+        f"{family}-{grade}-{version}"
+        for (family, version), grades in catalog.items()
+        for grade in grades
+    }
+
+
+@pytest.mark.parametrize(
+    ("family", "version", "options"),
+    [
+        (
+            "dpa4",
+            "v20260820",
+            {"edge_norm": True, "sandwich_norm": [False, True, True, False]},
+        ),
+        (
+            "dpa4",
+            "v20260911",
+            {
+                "edge_norm": [False, True, True],
+                "sandwich_norm": [True, False, True, False],
+                "env_exp": 5,
+                "basis_type": "gaussian/fix",
+            },
+        ),
+        ("dpa4c", "v20260901", {"basis_type": "bessel"}),
+        ("dpa4c", "v20260911", {"basis_type": "gaussian/fix"}),
+    ],
+)
+def test_version_descriptor_options(family: str, version: str, options: dict) -> None:
+    for name in MODEL_PRESETS:
+        if name.startswith(f"{family}-") and name.endswith(f"-{version}"):
+            descriptor = _normalize_model(get_model_preset(name))["descriptor"]
+            assert {key: descriptor[key] for key in options} == options
 
 
 def test_presets_carry_no_runtime_options() -> None:
@@ -209,7 +221,7 @@ def test_periodic_table_matches_econf_type_map() -> None:
 
 def test_leftover_preset_fails_argument_check() -> None:
     config = {
-        "model": {"preset": "dpa4-nano-v20260901"},
+        "model": {"preset": "dpa4-nano-v20260911"},
         "training": copy.deepcopy(TRAINING),
     }
     with pytest.raises(ArgumentKeyError, match="preset"):
@@ -219,7 +231,7 @@ def test_leftover_preset_fails_argument_check() -> None:
 def test_multi_task_preset_passes_shared_param_preprocessing() -> None:
     """The expansion runs before multi-task preprocessing and argument check."""
     model = {
-        "preset": "dpa4-nano-v20260901",
+        "preset": "dpa4-nano-v20260911",
         "shared_dict": {
             "type_map": ["O", "H"],
             "descriptor": {"type": "dpa4", "rcut": 6.0},
@@ -241,7 +253,7 @@ def test_multi_task_preset_passes_shared_param_preprocessing() -> None:
         expand_model_preset(model), lambda item_key, params: dict
     )
     assert set(shared_links) == {"descriptor"}
-    nano = get_model_preset("dpa4-nano-v20260901")
+    nano = get_model_preset("dpa4-nano-v20260911")
     for branch in processed["model_dict"].values():
         assert branch["type"] == "dpa4"
         assert branch["type_map"] == ["O", "H"]
@@ -265,14 +277,14 @@ def test_multi_task_preset_passes_shared_param_preprocessing() -> None:
 
 def test_multi_task_top_level_regions_are_branch_defaults() -> None:
     model = {
-        "preset": "dpa4-nano-v20260901",
+        "preset": "dpa4-nano-v20260911",
         "descriptor": {"rcut": 7.0},
         "fitting_net": {"seed": 3},
         "shared_dict": {"type_map": ["O", "H"]},
         "model_dict": {
             "water_1": {"type_map": "type_map"},
             "water_2": {
-                "preset": "dpa4-mini-v20260901",
+                "preset": "dpa4-mini-v20260911",
                 "type_map": "type_map",
                 "descriptor": {"rcut": 5.0},
             },
@@ -282,7 +294,7 @@ def test_multi_task_top_level_regions_are_branch_defaults() -> None:
     water_1 = expanded["model_dict"]["water_1"]
     assert water_1["descriptor"]["rcut"] == 7.0
     assert water_1["descriptor"]["lmax"] == 1
-    nano = get_model_preset("dpa4-nano-v20260901")
+    nano = get_model_preset("dpa4-nano-v20260911")
     assert water_1["fitting_net"] == {**nano["fitting_net"], "seed": 3}
     # A branch entry replaces the top-level default as a whole.
     water_2 = expanded["model_dict"]["water_2"]
@@ -302,7 +314,7 @@ def test_multi_task_top_level_regions_are_branch_defaults() -> None:
 
 def test_shared_dict_entries_take_the_top_level_preset_as_base() -> None:
     model = {
-        "preset": "dpa4-nano-v20260901",
+        "preset": "dpa4-nano-v20260911",
         "shared_dict": {
             "type_map": ["O", "H"],
             "descriptor": {"use_amp": True, "seed": 42},
@@ -316,14 +328,14 @@ def test_shared_dict_entries_take_the_top_level_preset_as_base() -> None:
                 "fitting_net": "shared_fit",
             },
             "water_2": {
-                "preset": "dpa4-mini-v20260901",
+                "preset": "dpa4-mini-v20260911",
                 "type_map": "type_map",
                 "descriptor": "descriptor",
             },
         },
     }
     expanded = expand_model_preset(model)
-    nano = get_model_preset("dpa4-nano-v20260901")
+    nano = get_model_preset("dpa4-nano-v20260911")
     shared = expanded["shared_dict"]
     # Referenced entries: preset region plus the run-specific keys; a shared
     # level suffix in the reference does not change the entry it names.
@@ -345,7 +357,7 @@ def test_shared_dict_entries_take_the_top_level_preset_as_base() -> None:
         "shared_dict": {"type_map": ["O"], "descriptor": {"rcut": 5.0}},
         "model_dict": {
             "a": {
-                "preset": "dpa4-nano-v20260901",
+                "preset": "dpa4-nano-v20260911",
                 "type_map": "type_map",
                 "descriptor": "descriptor",
             }
@@ -360,7 +372,7 @@ def test_explicit_alias_replaces_the_preset_canonical_key() -> None:
     """
     model = expand_model_preset(
         {
-            "preset": "dpa4-nano-v20260901",
+            "preset": "dpa4-nano-v20260911",
             "type_map": ["O", "H"],
             "descriptor": {"so2_layers": 5},
         }
@@ -375,7 +387,7 @@ def test_explicit_dict_for_a_whole_value_region_does_not_crash() -> None:
     dict there must not raise before the argument check reports it.
     """
     expanded = expand_model_preset(
-        {"preset": "dpa4-nano-v20260901", "type_map": {"O": 0, "H": 1}}
+        {"preset": "dpa4-nano-v20260911", "type_map": {"O": 0, "H": 1}}
     )
     assert expanded["type_map"] == {"O": 0, "H": 1}
     with pytest.raises(Exception):
@@ -386,9 +398,9 @@ def test_shared_dict_role_is_recognised_through_a_branch_default() -> None:
     """A `descriptor`/`fitting_net` reference inherited by a branch only
     through the top-level default must still be recognised as referenced.
     """
-    nano = get_model_preset("dpa4-nano-v20260901")
+    nano = get_model_preset("dpa4-nano-v20260911")
     model = {
-        "preset": "dpa4-nano-v20260901",
+        "preset": "dpa4-nano-v20260911",
         "descriptor": "desc",
         "shared_dict": {"type_map": ["O", "H"], "desc": {"seed": 42}},
         "model_dict": {"water_1": {"type_map": "type_map"}},
@@ -399,10 +411,10 @@ def test_shared_dict_role_is_recognised_through_a_branch_default() -> None:
 
 
 def test_malformed_multi_task_layout_is_left_to_argcheck() -> None:
-    malformed = {"preset": "dpa4-nano-v20260901", "model_dict": "water"}
+    malformed = {"preset": "dpa4-nano-v20260911", "model_dict": "water"}
     assert expand_model_preset(malformed) is malformed
     branch_not_mapping = {
-        "preset": "dpa4-nano-v20260901",
+        "preset": "dpa4-nano-v20260911",
         "model_dict": {"water": "not a mapping"},
     }
     assert expand_model_preset(branch_not_mapping)["model_dict"] == {
@@ -412,7 +424,7 @@ def test_malformed_multi_task_layout_is_left_to_argcheck() -> None:
 
 def test_update_deepmd_input_expands_presets() -> None:
     jdata = {
-        "model": {"preset": "dpa4c-mini-v20260901"},
+        "model": {"preset": "dpa4c-mini-v20260911"},
         "training": copy.deepcopy(TRAINING),
     }
     out = update_deepmd_input(jdata, warning=False)

From 0ad6b492e261ee18fbfaba583a70b88893133000 Mon Sep 17 00:00:00 2001
From: OutisLi <LTC201806070316@gmail.com>
Date: Sat, 12 Sep 2026 16:24:38 +0800
Subject: [PATCH 8/8] fix(dpa4): preserve empty-edge shapes and fp64 compile
 precision

---
 deepmd/dpmodel/descriptor/dpa4_nn/grid_net.py | 62 ++++++++++++-------
 .../pt/model/descriptor/sezm_nn/activation.py |  2 +-
 .../pt/model/descriptor/sezm_nn/grid_net.py   | 54 +++++++++-------
 deepmd/pt/model/descriptor/sezm_nn/so2.py     |  3 +-
 .../pt_expt/descriptor/dpa4_nn/activation.py  |  2 +-
 deepmd/pt_expt/descriptor/dpa4_nn/so2.py      |  3 +-
 .../kernels/triton/sezm/gated_activation.py   |  4 +-
 .../pt_expt/kernels/triton/sezm/radial_mix.py |  2 +-
 .../kernels/triton/sezm/so2_value_path.py     | 14 +++--
 source/tests/pt/model/test_descriptor_sezm.py | 32 +++++++---
 .../model/test_descriptor_sezm_train_paths.py | 14 +++--
 .../pt/model/test_descriptor_sezm_triton.py   | 37 +++++++++++
 source/tests/pt/model/test_embedding.py       | 14 ++++-
 source/tests/pt_expt/descriptor/test_dpa4.py  | 53 ++++++++++++++++
 .../descriptor/test_dpa4_train_paths.py       | 17 ++---
 15 files changed, 231 insertions(+), 82 deletions(-)

diff --git a/deepmd/dpmodel/descriptor/dpa4_nn/grid_net.py b/deepmd/dpmodel/descriptor/dpa4_nn/grid_net.py
index 597ac02261..22bf8569aa 100644
--- a/deepmd/dpmodel/descriptor/dpa4_nn/grid_net.py
+++ b/deepmd/dpmodel/descriptor/dpa4_nn/grid_net.py
@@ -148,9 +148,15 @@ def _project_frames(coeff: Any, proj: ChannelLinear, n_frames: int) -> Any:
     than the ``G``-point grid.
     """
     xp = array_api_compat.array_namespace(coeff)
-    n_batch, coeff_dim, n_focus, _ = coeff.shape
-    projected = proj(xp.reshape(coeff, (n_batch, coeff_dim, n_focus, n_frames, -1)))
-    return xp.reshape(projected, (n_batch, coeff_dim, n_focus, -1))
+    n_batch, coeff_dim, n_focus, n_channels = coeff.shape
+    projected = proj(
+        xp.reshape(
+            coeff, (n_batch, coeff_dim, n_focus, n_frames, n_channels // n_frames)
+        )
+    )
+    return xp.reshape(
+        projected, (n_batch, coeff_dim, n_focus, n_frames * projected.shape[-1])
+    )
 
 
 def _project_pair_in_one_transform(
@@ -162,14 +168,14 @@ def _project_pair_in_one_transform(
 ) -> tuple[Any, Any]:
     """Project two equally shaped coefficient operands in one linear transform."""
     xp = array_api_compat.array_namespace(left, right)
-    n_batch, coeff_dim, n_focus, _ = left.shape
-    frame_shape = (n_batch, coeff_dim, n_focus, n_frames, -1)
+    n_batch, coeff_dim, n_focus, n_channels = left.shape
+    frame_shape = (n_batch, coeff_dim, n_focus, n_frames, n_channels // n_frames)
     pair = xp.reshape(
         xp.concat(
             [xp.reshape(left, frame_shape), xp.reshape(right, frame_shape)],
             axis=-1,
         ),
-        (n_batch, coeff_dim, n_focus, -1),
+        (n_batch, coeff_dim, n_focus, 2 * n_channels),
     )
     pair_grid = to_grid(pair)
     split = pair_grid.shape[-1] // 2
@@ -354,10 +360,11 @@ def _project_operands(
         """Apply the two coefficient-space channel projections."""
         xp = array_api_compat.array_namespace(left)
         if self.mode == "self":
-            shape = (*left.shape[:-1], self.n_frames, -1)
+            n_channels = left.shape[-1]
+            shape = (*left.shape[:-1], self.n_frames, n_channels // self.n_frames)
             fused = xp.reshape(
                 xp.concat([xp.reshape(left, shape), xp.reshape(right, shape)], axis=-1),
-                (*left.shape[:-1], -1),
+                (*left.shape[:-1], 2 * n_channels),
             )  # per-frame concat -> (N, D, F, K*2C)
             left = _project_frames(fused, self.left_proj, self.n_frames)
             right = _project_frames(fused, self.right_proj, self.n_frames)
@@ -1254,11 +1261,14 @@ def _project_pair_in_one_transform(
         )
 
     def _to_grid(self, coeff: Any) -> Any:
-        # The per-frame channel width is inferred so the projector also serves
-        # widened operands (e.g. a branch hidden width ``n_branches * C``).
+        # Derive the per-frame width from the channel axis so empty batches
+        # and widened operands (e.g. ``n_branches * C``) are both valid.
         xp = array_api_compat.array_namespace(coeff)
-        n_batch, coeff_dim, n_focus, _ = coeff.shape
-        coeff_view = xp.reshape(coeff, (n_batch, coeff_dim, n_focus, self.n_frames, -1))
+        n_batch, coeff_dim, n_focus, n_channels = coeff.shape
+        n_channels //= self.n_frames
+        coeff_view = xp.reshape(
+            coeff, (n_batch, coeff_dim, n_focus, self.n_frames, n_channels)
+        )
         to_grid = xp_asarray_nodetach(
             xp, self.projector.to_grid_mat[...], device=array_api_compat.device(coeff)
         )
@@ -1274,7 +1284,6 @@ def _to_grid(self, coeff: Any) -> Any:
         # (`xp_einsum("gdk,ndfkc->ngfc")`) costs 3.6 ms. The same contraction
         # is faster there and slower here, so the choice belongs to the graph
         # around it rather than to the contraction itself.
-        n_channels = coeff_view.shape[-1]
         coeff_dk = xp.permute_dims(coeff_view, (0, 1, 3, 2, 4))  # (N, D, K, F, C)
         coeff_flat = xp.reshape(
             coeff_dk, (n_batch, coeff_dim * self.n_frames, n_focus * n_channels)
@@ -1285,7 +1294,7 @@ def _to_grid(self, coeff: Any) -> Any:
     def _from_grid(self, grid: Any) -> Any:
         # Channel width is inferred to match the (possibly widened) grid field.
         xp = array_api_compat.array_namespace(grid)
-        n_batch, _, n_focus, _ = grid.shape
+        n_batch, _, n_focus, n_channels = grid.shape
         coeff_dim = self.projector.coeff_dim // self.n_frames
         from_grid = xp_asarray_nodetach(
             xp, self.projector.from_grid_mat[...], device=array_api_compat.device(grid)
@@ -1294,7 +1303,6 @@ def _from_grid(self, grid: Any) -> Any:
         # einsum "dkg,ngfc->ndfkc" (with from_grid reshaped (D, K, G)) as a
         # broadcast batched matmul, then a reshape to (N, D, F, K*C). from_grid
         # is already stored as (D*K, G); the matmul output is reshaped/permuted.
-        n_channels = grid.shape[-1]
         grid_flat = xp.reshape(
             grid, (n_batch, self.projector.grid_size, n_focus * n_channels)
         )
@@ -1310,12 +1318,11 @@ def _from_grid(self, grid: Any) -> Any:
     def _from_grid_scalar(self, grid: Any) -> Any:
         """Project a grid field to the ``l=0`` coefficient only."""
         xp = array_api_compat.array_namespace(grid)
-        n_batch, _, n_focus, _ = grid.shape
+        n_batch, _, n_focus, n_channels = grid.shape
         from_grid = xp_asarray_nodetach(
             xp, self.projector.from_grid_mat[...], device=array_api_compat.device(grid)
         )
         from_grid = xp.astype(from_grid[: self.n_frames], grid.dtype)
-        n_channels = grid.shape[-1]
         grid_flat = xp.reshape(
             grid, (n_batch, self.projector.grid_size, n_focus * n_channels)
         )
@@ -1334,9 +1341,16 @@ def _scalar_so3_product(self, left: Any, right: Any) -> Any:
             xp, weight[...], device=array_api_compat.device(left)
         )
         weight = xp.astype(weight, left.dtype)
-        n_batch, coeff_dim, n_focus, _ = left.shape
-        left_view = xp.reshape(left, (n_batch, coeff_dim, n_focus, self.n_frames, -1))
-        right_view = xp.reshape(right, (n_batch, coeff_dim, n_focus, self.n_frames, -1))
+        n_batch, coeff_dim, n_focus, n_channels = left.shape
+        frame_shape = (
+            n_batch,
+            coeff_dim,
+            n_focus,
+            self.n_frames,
+            n_channels // self.n_frames,
+        )
+        left_view = xp.reshape(left, frame_shape)
+        right_view = xp.reshape(right, frame_shape)
         scalar = xp.sum(
             left_view * weight[None, :, None, :, None] * right_view,
             axis=(1, 3),
@@ -1355,9 +1369,11 @@ def _to_ndfc(self, value: Any) -> tuple[Any, tuple[int, ...]]:
             return xp.permute_dims(value, (0, 2, 1, 3)), tuple(value.shape)
         if self.layout == "fndc":
             return xp.permute_dims(value, (1, 2, 0, 3)), tuple(value.shape)
-        n_batch, coeff_dim, _ = value.shape
+        n_batch, coeff_dim, n_channels = value.shape
         return (
-            xp.reshape(value, (n_batch, coeff_dim, self.n_focus, -1)),
+            xp.reshape(
+                value, (n_batch, coeff_dim, self.n_focus, n_channels // self.n_focus)
+            ),
             tuple(value.shape),
         )
 
@@ -1377,7 +1393,7 @@ def _restore_layout(
             return xp.permute_dims(value, (2, 0, 1, 3))
         n_batch, input_coeff_dim, _ = shape_info
         coeff_dim = 1 if scalar_only else input_coeff_dim
-        return xp.reshape(value, (n_batch, coeff_dim, -1))
+        return xp.reshape(value, (n_batch, coeff_dim, value.shape[2] * value.shape[3]))
 
     def _slice_scalar_layout(self, value: Any) -> Any:
         """Select the degree axis from a restored full-layout tensor."""
diff --git a/deepmd/pt/model/descriptor/sezm_nn/activation.py b/deepmd/pt/model/descriptor/sezm_nn/activation.py
index fecd849fbf..d3e583f526 100644
--- a/deepmd/pt/model/descriptor/sezm_nn/activation.py
+++ b/deepmd/pt/model/descriptor/sezm_nn/activation.py
@@ -249,7 +249,7 @@ def forward(
             gw = weight.permute(1, 0, 2).contiguous()
             gwt = weight.permute(1, 2, 0).contiguous()
             out = self._fused_gated_act(
-                x.reshape(n_focus, n_edge, -1).contiguous(),
+                x.reshape(n_focus, n_edge, x.shape[2] * x.shape[3]).contiguous(),
                 gw,
                 gwt,
                 self.lmax,
diff --git a/deepmd/pt/model/descriptor/sezm_nn/grid_net.py b/deepmd/pt/model/descriptor/sezm_nn/grid_net.py
index c04673d298..4a8b73f0a0 100644
--- a/deepmd/pt/model/descriptor/sezm_nn/grid_net.py
+++ b/deepmd/pt/model/descriptor/sezm_nn/grid_net.py
@@ -136,9 +136,13 @@ def _project_frames(
     to applying it on the grid field while touching ``n_frames``-fold fewer rows
     than the ``G``-point grid.
     """
-    n_batch, coeff_dim, n_focus, _ = coeff.shape
-    projected = proj(coeff.reshape(n_batch, coeff_dim, n_focus, n_frames, -1))
-    return projected.reshape(n_batch, coeff_dim, n_focus, -1)
+    n_batch, coeff_dim, n_focus, n_channels = coeff.shape
+    projected = proj(
+        coeff.reshape(n_batch, coeff_dim, n_focus, n_frames, n_channels // n_frames)
+    )
+    return projected.reshape(
+        n_batch, coeff_dim, n_focus, n_frames * projected.shape[-1]
+    )
 
 
 def _project_pair_in_one_transform(
@@ -149,12 +153,12 @@ def _project_pair_in_one_transform(
     to_grid: Callable[[torch.Tensor], torch.Tensor],
 ) -> tuple[torch.Tensor, torch.Tensor]:
     """Project two equally shaped coefficient operands in one linear transform."""
-    n_batch, coeff_dim, n_focus, _ = left.shape
-    frame_shape = (n_batch, coeff_dim, n_focus, n_frames, -1)
+    n_batch, coeff_dim, n_focus, n_channels = left.shape
+    frame_shape = (n_batch, coeff_dim, n_focus, n_frames, n_channels // n_frames)
     pair = torch.cat(
         [left.reshape(frame_shape), right.reshape(frame_shape)],
         dim=-1,
-    ).reshape(n_batch, coeff_dim, n_focus, -1)
+    ).reshape(n_batch, coeff_dim, n_focus, 2 * n_channels)
     return torch.chunk(to_grid(pair), chunks=2, dim=-1)
 
 
@@ -347,10 +351,13 @@ def _project_operands(
     ) -> tuple[torch.Tensor, torch.Tensor]:
         """Apply the two coefficient-space channel projections."""
         if self.mode == "self":
-            shape = (*left.shape[:-1], self.n_frames, -1)
+            n_channels = left.shape[-1]
+            shape = (*left.shape[:-1], self.n_frames, n_channels // self.n_frames)
             fused = torch.cat(
                 [left.reshape(shape), right.reshape(shape)], dim=-1
-            ).reshape(*left.shape[:-1], -1)  # per-frame concat -> (N, D, F, K*2C)
+            ).reshape(
+                *left.shape[:-1], 2 * n_channels
+            )  # per-frame concat -> (N, D, F, K*2C)
             left = _project_frames(fused, self.left_proj, self.n_frames)
             right = _project_frames(fused, self.right_proj, self.n_frames)
         else:
@@ -1091,10 +1098,13 @@ def _project_pair_in_one_transform(
         )
 
     def _to_grid(self, coeff: torch.Tensor) -> torch.Tensor:
-        # The per-frame channel width is inferred so the projector also serves
-        # widened operands (e.g. a branch hidden width ``n_branches * C``).
-        n_batch, coeff_dim, n_focus, _ = coeff.shape
-        coeff_view = coeff.reshape(n_batch, coeff_dim, n_focus, self.n_frames, -1)
+        # Derive the per-frame width from the channel axis so empty batches
+        # and widened operands (e.g. ``n_branches * C``) are both valid.
+        n_batch, coeff_dim, n_focus, n_channels = coeff.shape
+        n_channels //= self.n_frames
+        coeff_view = coeff.reshape(
+            n_batch, coeff_dim, n_focus, self.n_frames, n_channels
+        )
         to_grid = self.projector.to_grid_mat.reshape(
             self.projector.grid_size,
             coeff_dim,
@@ -1104,7 +1114,7 @@ def _to_grid(self, coeff: torch.Tensor) -> torch.Tensor:
 
     def _from_grid(self, grid: torch.Tensor) -> torch.Tensor:
         # Channel width is inferred to match the (possibly widened) grid field.
-        n_batch, _, n_focus, _ = grid.shape
+        n_batch, _, n_focus, n_channels = grid.shape
         coeff_dim = self.projector.coeff_dim // self.n_frames
         from_grid = self.projector.from_grid_mat.reshape(
             coeff_dim,
@@ -1112,11 +1122,11 @@ def _from_grid(self, grid: torch.Tensor) -> torch.Tensor:
             self.projector.grid_size,
         )
         coeff = torch.einsum("dkg,ngfc->ndfkc", from_grid, grid)
-        return coeff.reshape(n_batch, coeff_dim, n_focus, -1)
+        return coeff.reshape(n_batch, coeff_dim, n_focus, self.n_frames * n_channels)
 
     def _from_grid_scalar(self, grid: torch.Tensor) -> torch.Tensor:
         """Project a grid field to the ``l=0`` coefficient only."""
-        n_batch, _, n_focus, _ = grid.shape
+        n_batch, _, n_focus, n_channels = grid.shape
         coeff_dim = self.projector.coeff_dim // self.n_frames
         from_grid = self.projector.from_grid_mat.reshape(
             coeff_dim,
@@ -1124,7 +1134,7 @@ def _from_grid_scalar(self, grid: torch.Tensor) -> torch.Tensor:
             self.projector.grid_size,
         )[0:1]
         coeff = torch.einsum("dkg,ngfc->ndfkc", from_grid, grid)
-        return coeff.reshape(n_batch, 1, n_focus, -1)
+        return coeff.reshape(n_batch, 1, n_focus, self.n_frames * n_channels)
 
     def _scalar_so3_product(
         self,
@@ -1135,8 +1145,10 @@ def _scalar_so3_product(
         weight = self._scalar_product_weight
         if weight is None:
             raise RuntimeError("SO(3) scalar product weights are unavailable")
-        n_batch, coeff_dim, n_focus, _ = left.shape
-        left_view = left.reshape(n_batch, coeff_dim, n_focus, self.n_frames, -1)
+        n_batch, coeff_dim, n_focus, n_channels = left.shape
+        left_view = left.reshape(
+            n_batch, coeff_dim, n_focus, self.n_frames, n_channels // self.n_frames
+        )
         right_view = right.reshape_as(left_view)
         # A weighted diagonal product-sum over (d, k); written out so that no
         # contraction-path search runs on the symbolic node count under export.
@@ -1156,9 +1168,9 @@ def _to_ndfc(self, value: torch.Tensor) -> tuple[torch.Tensor, tuple[int, ...]]:
             return value.transpose(1, 2), tuple(value.shape)
         if self.layout == "fndc":
             return value.permute(1, 2, 0, 3), tuple(value.shape)
-        n_batch, coeff_dim, _ = value.shape
+        n_batch, coeff_dim, n_channels = value.shape
         return (
-            value.reshape(n_batch, coeff_dim, self.n_focus, -1),
+            value.reshape(n_batch, coeff_dim, self.n_focus, n_channels // self.n_focus),
             tuple(value.shape),
         )
 
@@ -1177,7 +1189,7 @@ def _restore_layout(
             return value.permute(2, 0, 1, 3)
         n_batch, input_coeff_dim, _ = shape_info
         coeff_dim = 1 if scalar_only else input_coeff_dim
-        return value.reshape(n_batch, coeff_dim, -1)
+        return value.reshape(n_batch, coeff_dim, value.shape[2] * value.shape[3])
 
     def _slice_scalar_layout(self, value: torch.Tensor) -> torch.Tensor:
         """Select the degree axis from a restored full-layout tensor."""
diff --git a/deepmd/pt/model/descriptor/sezm_nn/so2.py b/deepmd/pt/model/descriptor/sezm_nn/so2.py
index a366babdd2..0a3418efcf 100644
--- a/deepmd/pt/model/descriptor/sezm_nn/so2.py
+++ b/deepmd/pt/model/descriptor/sezm_nn/so2.py
@@ -1739,10 +1739,11 @@ def __init__(
         # destination-sorted view with the flash aggregation; its backward and
         # hand-derived second order keep the force-loss trace from expanding
         # the chain into materialized surfaces and serialized scatters. The
-        # source-gated (SFPG) form keeps the reference path.
+        # source-gated (SFPG) form and fp64 compute keep the reference path.
         self._segment_softmax_fn = None
         if (
             max(self.triton_infer_level, self.triton_train_level) >= 1
+            and self.compute_dtype is torch.float32
             and self.attn_n_focus * self.n_atten_head <= 16
         ):
             from deepmd.pt_expt.kernels.triton.sezm.segment_softmax import (
diff --git a/deepmd/pt_expt/descriptor/dpa4_nn/activation.py b/deepmd/pt_expt/descriptor/dpa4_nn/activation.py
index 9a8d1a770a..d9d6be711e 100644
--- a/deepmd/pt_expt/descriptor/dpa4_nn/activation.py
+++ b/deepmd/pt_expt/descriptor/dpa4_nn/activation.py
@@ -90,7 +90,7 @@ def call(self, x: torch.Tensor, gate: torch.Tensor | None = None) -> torch.Tenso
                 self.channels, self.n_focus, self.lmax * self.channels
             )
             out = self._fused_gated_act(
-                x.reshape(n_focus, n_edge, -1).contiguous(),
+                x.reshape(n_focus, n_edge, x.shape[2] * x.shape[3]).contiguous(),
                 weight.permute(1, 0, 2).contiguous(),
                 weight.permute(1, 2, 0).contiguous(),
                 self.lmax,
diff --git a/deepmd/pt_expt/descriptor/dpa4_nn/so2.py b/deepmd/pt_expt/descriptor/dpa4_nn/so2.py
index b244f33e6a..87e3344e7c 100644
--- a/deepmd/pt_expt/descriptor/dpa4_nn/so2.py
+++ b/deepmd/pt_expt/descriptor/dpa4_nn/so2.py
@@ -373,10 +373,11 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
         # destination-sorted view with the flash aggregation; its backward and
         # hand-derived second order keep the force-loss trace from expanding
         # the chain into materialized surfaces and serialized scatters. The
-        # source-gated (SFPG) form keeps the reference path.
+        # source-gated (SFPG) form and fp64 compute keep the reference path.
         self._segment_softmax_fn = None
         if (
             max(self.triton_infer_level, self.triton_train_level) >= 1
+            and self.compute_precision == "float32"
             and self.attn_n_focus * self.n_atten_head <= 16
         ):
             from deepmd.pt_expt.kernels.triton.sezm.segment_softmax import (
diff --git a/deepmd/pt_expt/kernels/triton/sezm/gated_activation.py b/deepmd/pt_expt/kernels/triton/sezm/gated_activation.py
index 45399f132d..58d0d3be60 100644
--- a/deepmd/pt_expt/kernels/triton/sezm/gated_activation.py
+++ b/deepmd/pt_expt/kernels/triton/sezm/gated_activation.py
@@ -155,7 +155,7 @@ def gated_activation_second_order_reference(
 
     def fold(value: Tensor) -> Tensor:
         """Sum the two signed-``m`` halves that share a gate group."""
-        return value.view(n_focus, n_edge, 2, -1).sum(2)
+        return value.view(n_focus, n_edge, 2, value.shape[-1] // 2).sum(2)
 
     scalar = z[:, :, :focus_dim]
     grad_scalar = grad[:, :, :focus_dim]
@@ -512,7 +512,7 @@ def gated_activation_second_order(
         sig = torch.sigmoid(torch.bmm(z[:, :, :focus_dim], gw))
         grad_sig = grad[:, :, focus_dim:m0] * z[:, :, focus_dim:m0] + (
             grad[:, :, m0:] * z[:, :, m0:]
-        ).view(n_focus, n_edge, 2, -1).sum(2)
+        ).view(n_focus, n_edge, 2, lmax * focus_dim).sum(2)
         grad_wrt_gw = grad_wrt_gw + torch.bmm(
             grad_grad_z[:, :, :focus_dim].transpose(1, 2),
             grad_sig * sig * (1.0 - sig),
diff --git a/deepmd/pt_expt/kernels/triton/sezm/radial_mix.py b/deepmd/pt_expt/kernels/triton/sezm/radial_mix.py
index 2a95cb6727..2201f3b7d1 100644
--- a/deepmd/pt_expt/kernels/triton/sezm/radial_mix.py
+++ b/deepmd/pt_expt/kernels/triton/sezm/radial_mix.py
@@ -194,7 +194,7 @@ def _radial_mix_backward_reference(
         gk = torch.einsum("eoc,eic,rc->eoir", g_block, x_block, channel_basis)
         grad_compact[:, comp0 : comp0 + num_l * num_l, :] += gk.permute(
             0, 2, 1, 3
-        ).reshape(n_edge, num_l * num_l, -1)
+        ).reshape(n_edge, num_l * num_l, compact.shape[-1])
     return grad_compact, grad_x_local
 
 
diff --git a/deepmd/pt_expt/kernels/triton/sezm/so2_value_path.py b/deepmd/pt_expt/kernels/triton/sezm/so2_value_path.py
index c0932ba1c4..f42b546941 100644
--- a/deepmd/pt_expt/kernels/triton/sezm/so2_value_path.py
+++ b/deepmd/pt_expt/kernels/triton/sezm/so2_value_path.py
@@ -432,7 +432,7 @@ def _mixing_stack_backward_reference(
         gz1 = g_cur[:, :, m0:] * sig2
         g_sig = (g_cur[:, :, focus_dim:m0] * z0[:, :, focus_dim:]).view(*sig.shape) + (
             g_cur[:, :, m0:] * z1
-        ).view(sig.shape[0], sig.shape[1], 2, -1).sum(2)
+        ).view(sig.shape[0], sig.shape[1], 2, sig.shape[2]).sum(2)
         g_logit = g_sig * sig * (1.0 - sig)
         gz0 = torch.cat(
             [
@@ -2242,7 +2242,9 @@ def _stack_weight_gradients(
 
     u_flat = state.inputs.reshape(n_gated * n_focus, n_edge, row)
     gz_flat = state.grad_z.reshape(n_gated * n_focus, n_edge, row)
-    gq_flat = state.grad_logit.reshape(n_gated * n_focus, n_edge, -1)
+    gq_flat = state.grad_logit.reshape(
+        n_gated * n_focus, n_edge, state.grad_logit.shape[-1]
+    )
     z_flat = z_all.reshape(n_gated * n_focus, n_edge, row)
 
     gw0 = torch.empty(
@@ -2649,10 +2651,12 @@ def _gated_act_reference(
     focus_dim = int(focus_dim)
     m0 = (lmax + 1) * focus_dim
     scalar = z[:, :, :focus_dim]
-    sig = torch.sigmoid(torch.bmm(scalar.float(), gw.float())).to(z.dtype)
+    compute_dtype = torch.float64 if z.dtype is torch.float64 else torch.float32
+    scalar_compute = scalar.to(compute_dtype)
+    sig = torch.sigmoid(torch.bmm(scalar_compute, gw.to(compute_dtype))).to(z.dtype)
     return torch.cat(
         [
-            scalar * torch.sigmoid(scalar.float()).to(z.dtype),
+            scalar * torch.sigmoid(scalar_compute).to(z.dtype),
             z[:, :, focus_dim:m0] * sig,
             z[:, :, m0:] * sig.repeat(1, 1, 2),
         ],
@@ -2715,7 +2719,7 @@ def _stack_point_bwd_reference(
 
     grad_sig = (grad[:, :, focus_dim:m0] * z[:, :, focus_dim:m0]) + (
         grad[:, :, m0:] * z[:, :, m0:]
-    ).view(sig.shape[0], sig.shape[1], 2, -1).sum(2)
+    ).view(sig.shape[0], sig.shape[1], 2, sig.shape[2]).sum(2)
     grad_logit = grad_sig * sig * (1.0 - sig)
     if not fold_logit:
         gz_scalar = gz_scalar + torch.bmm(grad_logit, gw.transpose(1, 2))
diff --git a/source/tests/pt/model/test_descriptor_sezm.py b/source/tests/pt/model/test_descriptor_sezm.py
index a077c623f0..c3646da4a2 100644
--- a/source/tests/pt/model/test_descriptor_sezm.py
+++ b/source/tests/pt/model/test_descriptor_sezm.py
@@ -371,9 +371,6 @@ def test_so3_readout_empty_edge_shrinking_schedule(self) -> None:
 
     def test_edge_free_frame_continues_the_cutoff_limit(self) -> None:
         """A frame without edges is the limit of a frame whose last edge leaves the cutoff."""
-        model = DescrptSeZM(**_descriptor_kwargs(precision="float64", seed=5))
-        model = model.to(self.device).eval()
-        _perturb_parameters(model, seed=5)
         atype = torch.tensor([[0, 1]], dtype=torch.int32, device=self.device)
         empty_nlist = torch.full((1, 2, 2), -1, dtype=torch.int64, device=self.device)
         pair_nlist = torch.tensor(
@@ -388,13 +385,28 @@ def descriptor(distance: float, nlist: torch.Tensor) -> torch.Tensor:
             ).reshape(1, -1)
             return model(coord, atype, nlist, mapping=None, comm_dict=None)[0]
 
-        isolated = descriptor(10.0, empty_nlist)
-        torch.testing.assert_close(
-            descriptor(model.rcut - 1e-6, pair_nlist), isolated, rtol=0.0, atol=1e-12
-        )
-        self.assertFalse(
-            torch.allclose(descriptor(model.rcut - 0.5, pair_nlist), isolated)
-        )
+        for options in (
+            {},
+            {"node_wise_s2": True},
+            {"node_wise_so3": True},
+            {"s2_activation": [True, False]},
+        ):
+            with self.subTest(options=options):
+                model = DescrptSeZM(
+                    **_descriptor_kwargs(precision="float64", seed=5, **options)
+                )
+                model = model.to(self.device).eval()
+                _perturb_parameters(model, seed=5)
+                isolated = descriptor(10.0, empty_nlist)
+                torch.testing.assert_close(
+                    descriptor(model.rcut - 1e-6, pair_nlist),
+                    isolated,
+                    rtol=0.0,
+                    atol=1e-12,
+                )
+                self.assertFalse(
+                    torch.allclose(descriptor(model.rcut - 0.5, pair_nlist), isolated)
+                )
 
     def test_so3_readout_scalar_path_matches_full_output(self) -> None:
         """The scalar-specialized final FFN matches slicing its full output."""
diff --git a/source/tests/pt/model/test_descriptor_sezm_train_paths.py b/source/tests/pt/model/test_descriptor_sezm_train_paths.py
index 263e74ba08..cdf53b1d5e 100644
--- a/source/tests/pt/model/test_descriptor_sezm_train_paths.py
+++ b/source/tests/pt/model/test_descriptor_sezm_train_paths.py
@@ -83,6 +83,7 @@ def _make_descriptor(
     rcut: float,
     *,
     source_gated: bool = False,
+    precision: str = "float32",
 ) -> DescrptSeZM:
     """Build a small SeZM descriptor in the deployed layout."""
     return DescrptSeZM(
@@ -101,7 +102,7 @@ def _make_descriptor(
         grid_branch=[1, 1, 1],
         s2_activation=[False, True],
         random_gamma=False,
-        precision="float32",
+        precision=precision,
         seed=7,
         inner_clamp_r_inner=0.8 if source_gated else None,
         inner_clamp_r_outer=1.2 if source_gated else None,
@@ -120,8 +121,9 @@ def _clear_gates(monkeypatch) -> None:
     ids=("training", "inference"),
 )
 @pytest.mark.parametrize("enabled", [0, 1])
+@pytest.mark.parametrize("precision", ["float32", "float64"])
 def test_triton_mode_gate_binds_each_stage(
-    monkeypatch, gate_name: str, training: bool, enabled: int
+    monkeypatch, gate_name: str, training: bool, enabled: int, precision: str
 ) -> None:
     """Each Triton gate binds every supported stage for only its own mode."""
     _clear_gates(monkeypatch)
@@ -130,7 +132,7 @@ def test_triton_mode_gate_binds_each_stage(
     train_level = enabled if training else 0
     infer_level = enabled if not training else 0
 
-    descriptor = _make_descriptor(2, [20], 4.0)
+    descriptor = _make_descriptor(2, [20], 4.0, precision=precision)
     convolutions = [
         module for module in descriptor.modules() if isinstance(module, SO2Convolution)
     ]
@@ -145,10 +147,10 @@ def test_triton_mode_gate_binds_each_stage(
         assert (conv._rotate_back_fn is not None) is requested
         assert (conv._flash_atten_fn is not None) is requested
         assert conv._flash_atten_trains is (requested and training)
-        # Segment softmax has no wrapper-level fallback and binds only when its
-        # own Triton implementation is importable.
+        # Segment softmax uses fp32 accumulation and preserves fp64 compute
+        # through the descriptor's reference path.
         assert (conv._segment_softmax_fn is not None) is (
-            requested and SEGMENT_SOFTMAX_TRITON_AVAILABLE
+            requested and SEGMENT_SOFTMAX_TRITON_AVAILABLE and precision == "float32"
         )
         # The rotate-mix front end is bound by a profitability bound on the
         # hidden width, which this narrow block sits below.
diff --git a/source/tests/pt/model/test_descriptor_sezm_triton.py b/source/tests/pt/model/test_descriptor_sezm_triton.py
index 796d1c8178..daa6f17143 100644
--- a/source/tests/pt/model/test_descriptor_sezm_triton.py
+++ b/source/tests/pt/model/test_descriptor_sezm_triton.py
@@ -25,6 +25,9 @@
 import math
 import typing
 import unittest
+from unittest import (
+    mock,
+)
 
 import torch
 from torch.fx.experimental.proxy_tensor import (
@@ -659,7 +662,17 @@ class TestSeZMTritonValuePath(unittest.TestCase):
     N_NODE = 512
     N_EDGE = 20000
 
+    @mock.patch.dict(
+        "os.environ",
+        {
+            "DP_TRITON_INFER": "0",
+            "DP_CUDA_INFER": "0",
+            "DP_TRITON_TRAIN": "0",
+            "DP_CUDA_TRAIN": "0",
+        },
+    )
     def _build_conv(self, lmax, channels, n_focus, focus_dim, layers, mode, rank):
+        """Build the dense reference independently of ambient acceleration gates."""
         from deepmd.pt.model.descriptor.sezm_nn.so2 import (
             SO2Convolution,
         )
@@ -850,6 +863,30 @@ def test_float64_fallback_channel_basis_gradient_preserves_precision(self) -> No
         (grad_reference,) = torch.autograd.grad(reference, basis_reference, grad_out)
         torch.testing.assert_close(grad_fused, grad_reference, atol=1e-12, rtol=1e-12)
 
+    def test_float64_gated_activation_derivatives(self) -> None:
+        """The fp64 fallback preserves first and second finite-difference derivatives."""
+        from deepmd.pt_expt.kernels.triton.sezm.so2_value_path import (
+            fused_gated_activation,
+        )
+
+        generator = torch.Generator(device="cpu").manual_seed(41)
+        z = torch.randn(
+            1, 2, 14, dtype=torch.float64, device="cpu", generator=generator
+        ).requires_grad_(True)
+        weight = torch.randn(
+            1, 2, 4, dtype=torch.float64, device="cpu", generator=generator
+        ).requires_grad_(True)
+
+        def activate(z: torch.Tensor, weight: torch.Tensor) -> torch.Tensor:
+            return fused_gated_activation(
+                z, weight, weight.transpose(1, 2).contiguous(), 2, 2
+            )
+
+        for check in (torch.autograd.gradcheck, torch.autograd.gradgradcheck):
+            self.assertTrue(
+                check(activate, (z, weight), atol=1e-8, rtol=1e-6, fast_mode=True)
+            )
+
     @_GPU_KERNELS
     def test_competition_gradient_preserves_small_positive_scale(self) -> None:
         from deepmd.pt_expt.kernels.triton.sezm.so2_value_path import (
diff --git a/source/tests/pt/model/test_embedding.py b/source/tests/pt/model/test_embedding.py
index 2617a0afba..f49a362cc0 100644
--- a/source/tests/pt/model/test_embedding.py
+++ b/source/tests/pt/model/test_embedding.py
@@ -28,12 +28,20 @@
 from deepmd.pt.utils import (
     env,
 )
+from deepmd.pt.utils.compile_compat import (
+    SUPPORTED_COMPILE_TORCH,
+)
 
-# The SeZM compile path is validated on torch 2.11.x / 2.12.x only.
+# Keep the compile-test gate aligned with the runtime allowlist.
 _TORCH_VERSION = parse_version(torch.__version__)
-_SKIP_COMPILE = (_TORCH_VERSION.major, _TORCH_VERSION.minor) not in {(2, 11), (2, 12)}
+_SKIP_COMPILE = (
+    _TORCH_VERSION.major,
+    _TORCH_VERSION.minor,
+) not in SUPPORTED_COMPILE_TORCH
 _SKIP_COMPILE_REASON = (
-    "SeZM's torch.compile path is only supported on torch 2.11.x and 2.12.x."
+    "SeZM's torch.compile path is only supported on torch "
+    + ", ".join(f"{major}.{minor}.x" for major, minor in SUPPORTED_COMPILE_TORCH)
+    + f"; current torch is {torch.__version__}."
 )
 
 
diff --git a/source/tests/pt_expt/descriptor/test_dpa4.py b/source/tests/pt_expt/descriptor/test_dpa4.py
index 1504708050..ca9e2f34f1 100644
--- a/source/tests/pt_expt/descriptor/test_dpa4.py
+++ b/source/tests/pt_expt/descriptor/test_dpa4.py
@@ -16,6 +16,9 @@
 )
 
 from deepmd.dpmodel.descriptor.dpa4 import DescrptDPA4 as DPDescrptDPA4
+from deepmd.dpmodel.utils.neighbor_graph import (
+    NeighborGraph,
+)
 from deepmd.pt_expt.descriptor.dpa4 import (
     DescrptDPA4,
 )
@@ -56,6 +59,56 @@ def setup_method(self) -> None:
         TestCaseSingleFrameWithNlist.setUp(self)
         self.device = env.DEVICE
 
+    @pytest.mark.parametrize(
+        "options",
+        [
+            {"node_wise_s2": True},
+            {"node_wise_so3": True},
+            {"s2_activation": [True, False]},
+        ],
+    )
+    @pytest.mark.parametrize("training", [False, True])
+    def test_empty_edge_grid_paths(self, options: dict, training: bool) -> None:
+        """Empty graph edges and fully masked neighbors give the same descriptors."""
+        descriptor = make_descriptor(2, 2, 3.0, channels=4, **options).to(self.device)
+        descriptor.train(training)
+        graph_data = {
+            "n_node": np.array([2], dtype=np.int64),
+            "edge_index": np.empty((2, 0), dtype=np.int64),
+            "edge_vec": np.empty((0, 3), dtype=np.float64),
+            "edge_mask": np.empty((0,), dtype=bool),
+        }
+        graph = NeighborGraph(
+            **{
+                key: torch.as_tensor(value, device=self.device)
+                for key, value in graph_data.items()
+            }
+        )
+        graph.edge_vec.requires_grad_(True)
+        atype = torch.tensor([0, 1], dtype=torch.int64, device=self.device)
+        actual = descriptor.call_graph(graph, atype)[0]
+        reference = DPDescrptDPA4.deserialize(descriptor.serialize()).call_graph(
+            NeighborGraph(**graph_data), atype.cpu().numpy()
+        )[0]
+        np.testing.assert_allclose(
+            actual.detach().cpu().numpy(), reference, rtol=1e-10, atol=1e-12
+        )
+        coord = torch.zeros((1, 2, 3), dtype=torch.float64, device=self.device)
+        nlist = torch.full((1, 2, 2), -1, dtype=torch.int64, device=self.device)
+        padded = descriptor(coord, atype[None, :], nlist)[0]
+        torch.testing.assert_close(actual, padded[0], rtol=1e-10, atol=1e-12)
+        edge_grad = torch.autograd.grad(
+            actual.sum(), graph.edge_vec, create_graph=training
+        )[0]
+        assert edge_grad.shape == (0, 3)
+        if training:
+            (actual.square().sum() + edge_grad.square().sum()).backward()
+            assert all(
+                torch.isfinite(parameter.grad).all()
+                for parameter in descriptor.parameters()
+                if parameter.grad is not None
+            )
+
     @pytest.mark.parametrize("use_env_seed", [True, False])  # env seed feature
     @pytest.mark.parametrize("use_mapping", [True, False])  # pass mapping vs None
     def test_consistency(self, use_env_seed, use_mapping) -> None:
diff --git a/source/tests/pt_expt/descriptor/test_dpa4_train_paths.py b/source/tests/pt_expt/descriptor/test_dpa4_train_paths.py
index c29e89d622..397ccdc40a 100644
--- a/source/tests/pt_expt/descriptor/test_dpa4_train_paths.py
+++ b/source/tests/pt_expt/descriptor/test_dpa4_train_paths.py
@@ -77,7 +77,9 @@
 INFER_GATES = ("DP_TRITON_INFER", "DP_CUDA_INFER", "DP_CUTILE_INFER", "DP_CUTE_INFER")
 
 
-def _make_descriptor(ntypes: int, sel: list[int], rcut: float) -> DescrptDPA4:
+def _make_descriptor(
+    ntypes: int, sel: list[int], rcut: float, *, precision: str = "float32"
+) -> DescrptDPA4:
     """Build a small DPA4 descriptor in the deployed layout."""
     return DescrptDPA4(
         ntypes=ntypes,
@@ -95,7 +97,7 @@ def _make_descriptor(ntypes: int, sel: list[int], rcut: float) -> DescrptDPA4:
         grid_branch=[1, 1, 1],
         s2_activation=[False, True],
         random_gamma=False,
-        precision="float32",
+        precision=precision,
         seed=7,
     )
 
@@ -119,8 +121,9 @@ def test_cuda_train_gate_accepts_shared_truthy_values(monkeypatch, value: str) -
     ids=("training", "inference"),
 )
 @pytest.mark.parametrize("enabled", [0, 1])
+@pytest.mark.parametrize("precision", ["float32", "float64"])
 def test_triton_mode_gate_binds_each_stage(
-    monkeypatch, gate_name: str, training: bool, enabled: int
+    monkeypatch, gate_name: str, training: bool, enabled: int, precision: str
 ) -> None:
     """Each Triton gate binds every supported stage for only its own mode."""
     _clear_gates(monkeypatch)
@@ -129,7 +132,7 @@ def test_triton_mode_gate_binds_each_stage(
     train_level = enabled if training else 0
     infer_level = enabled if not training else 0
 
-    descriptor = _make_descriptor(2, [20], 4.0)
+    descriptor = _make_descriptor(2, [20], 4.0, precision=precision)
     convolutions = [
         module for module in descriptor.modules() if isinstance(module, SO2Convolution)
     ]
@@ -144,10 +147,10 @@ def test_triton_mode_gate_binds_each_stage(
         assert (conv._rotate_back_fn is not None) is requested
         assert (conv._flash_atten_fn is not None) is requested
         assert conv._flash_atten_trains is (requested and training)
-        # Segment softmax has no wrapper-level fallback and binds only when its
-        # own Triton implementation is importable.
+        # Segment softmax uses fp32 accumulation and preserves fp64 compute
+        # through the descriptor's reference path.
         assert (conv._segment_softmax_fn is not None) is (
-            requested and SEGMENT_SOFTMAX_TRITON_AVAILABLE
+            requested and SEGMENT_SOFTMAX_TRITON_AVAILABLE and precision == "float32"
         )
         # The rotate-mix front end is bound by a profitability bound on the
         # hidden width, which this narrow block sits below.