chore(pt): merge updated CuTe pull request

# Conflicts:
#	deepmd/pt/model/descriptor/sezm.py
#	deepmd/pt_expt/kernels/cute/sezm/so2/kernels/phase_a_radial_forward.py
#	deepmd/pt_expt/kernels/cute/sezm/so2/message_grid.py
#	doc/install/easy-install.md

Author: OutisLi

deepmd/backend/backend.py

@@ -153,6 +153,14 @@ class Feature(Flag):
     """The supported suffixes of the saved model.
 
     The first element is considered as the default suffix."""
+    preserves_lower_input_kind: ClassVar[bool] = False
+    """Whether the IO hook preserves lower-ABI metadata without materializing it.
+
+    Schema-neutral model containers retain ``lower_input_kind`` as provenance
+    even though their deserializer does not accept a concrete ``lower_kind``.
+    Executable backends instead materialize a lower ABI and must expose that
+    choice through their deserializer signature.
+    """
 
     @abstractmethod
     def is_available(self) -> bool:

deepmd/backend/dpmodel.py

@@ -42,6 +42,8 @@ class DPModelBackend(Backend):
     """The features of the backend."""
     suffixes: ClassVar[list[str]] = [".dp", ".yaml", ".yml"]
     """The suffixes of the backend."""
+    preserves_lower_input_kind: ClassVar[bool] = True
+    """DPModel files retain lower provenance without binding an execution ABI."""
 
     def is_available(self) -> bool:
         """Check if the backend is available.
@@ -106,10 +108,10 @@ def serialize_hook(self) -> Callable[[str], dict]:
             The serialize hook of the backend.
         """
         from deepmd.dpmodel.utils.serialization import (
-            load_dp_model,
+            serialize_from_file,
         )
 
-        return load_dp_model
+        return serialize_from_file
 
     @property
     def deserialize_hook(self) -> Callable[[str, dict], None]:

deepmd/dpmodel/utils/serialization.py

@@ -199,6 +199,29 @@ def convert_numpy_ndarray(x: Any) -> Any:
     return model_dict
 
 
+def serialize_from_file(filename: str) -> dict:
+    """Serialize a DPModel container for backend conversion.
+
+    DPModel files store model parameters rather than an executable lower ABI.
+    Concrete provenance written by an earlier conversion is retained; a native
+    file without provenance reports ``"auto"`` so the executable target selects
+    a compatible lower from the model capabilities.
+
+    Parameters
+    ----------
+    filename : str
+        The DPModel filename.
+
+    Returns
+    -------
+    dict
+        The serialized model data with declared lower-input semantics.
+    """
+    model_dict = load_dp_model(filename)
+    model_dict.setdefault("lower_input_kind", "auto")
+    return model_dict
+
+
 def format_big_number(x: int) -> str:
     """Format a big number with suffixes.
 

deepmd/entrypoints/convert_backend.py

@@ -30,6 +30,14 @@ def convert_backend(
         If True, export .pt2/.pte models with per-atom virial correction.
         This adds ~2.5x inference cost.  Default False.  Silently ignored
         (with a warning) for backends that don't support the flag.
+
+    Notes
+    -----
+    Backend conversion preserves an explicit ``lower_input_kind`` reported by
+    the source serializer. Sources without this metadata retain the target's
+    automatic lower selection for backward compatibility. A target backend
+    that cannot represent an explicit non-dense lower is rejected rather than
+    silently changing the model function.
     """
     inp_backend: Backend = Backend.detect_backend_by_model(INPUT)()
     out_backend: Backend = Backend.detect_backend_by_model(OUTPUT)()
@@ -40,8 +48,21 @@ def convert_backend(
 
     sig = inspect.signature(out_hook)
     hook_kwargs: dict[str, Any] = {}
+    lower_input_kind = data.get("lower_input_kind")
     if "lower_kind" in sig.parameters:
-        hook_kwargs["lower_kind"] = "auto"
+        hook_kwargs["lower_kind"] = (
+            lower_input_kind if lower_input_kind is not None else "auto"
+        )
+    elif (
+        lower_input_kind not in (None, "auto", "nlist")
+        and not out_backend.preserves_lower_input_kind
+    ):
+        raise ValueError(
+            f"Cannot preserve lower_input_kind {lower_input_kind!r} when "
+            f"converting to output backend {out_backend.name!r}: its "
+            "deserializer does not accept a lower_kind. Retrain or freeze the "
+            "model with that backend instead of converting this artifact."
+        )
     if "do_atomic_virial" in sig.parameters:
         hook_kwargs["do_atomic_virial"] = atomic_virial
     elif atomic_virial:

deepmd/jax/utils/serialization.py

@@ -427,6 +427,7 @@ def restore_model(model_params: dict, model_state: dict) -> BaseModel:
             "jax_version": jax.__version__,
             "model": model_dict,
             "model_def_script": model_def_script,
+            "lower_input_kind": "nlist",
             "@variables": {},
         }
         if min_nbor_dist is not None:
@@ -436,6 +437,7 @@ def restore_model(model_params: dict, model_state: dict) -> BaseModel:
         data = load_dp_model(model_file)
         data.pop("constants")
         data["@variables"].pop("stablehlo")
+        data["lower_input_kind"] = "nlist"
         return data
     elif model_file.endswith(".savedmodel"):
         raise ValueError(

deepmd/pt/model/descriptor/sezm.py

@@ -1816,6 +1816,8 @@ def prepare_cute_infer_so2_metadata(
             Destination row pointers, source order, and source row pointers, or
             ``None`` when the exact CuTe contract is not satisfied.
         """
+        if torch.jit.is_scripting():
+            return None
         if (
             self.training
             or not edge_cache.destinations_sorted

deepmd/pt/model/model/model.py

@@ -26,6 +26,17 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
             "min_nbor_dist", torch.tensor(-1.0, dtype=torch.float64, device=env.DEVICE)
         )
 
+    def export_lower_input_kind(self) -> str:
+        """Return the lower-input ABI that preserves this model's semantics.
+
+        Returns
+        -------
+        str
+            ``"nlist"`` for the standard PyTorch model contract. Models with
+            a graph-native deployment ABI override this method.
+        """
+        return "nlist"
+
     def compute_or_load_stat(
         self,
         sampled_func: Any,

deepmd/pt/model/model/spin_model.py

@@ -454,6 +454,17 @@ def has_spin(self) -> bool:
         """Returns whether it has spin input and output."""
         return True
 
+    def export_lower_input_kind(self) -> str:
+        """Return the dense ABI used by the virtual-atom spin model.
+
+        Returns
+        -------
+        str
+            ``"nlist"``, because virtual atoms are expanded inside the
+            bounded neighbor-list contract.
+        """
+        return "nlist"
+
     @torch.jit.export
     def has_message_passing(self) -> bool:
         """Returns whether the model has message passing."""

deepmd/pt/utils/serialization.py

@@ -53,6 +53,7 @@ def serialize_from_file(model_file: str) -> dict:
         "pt_version": str(torch.__version__),
         "model": model_dict,
         "model_def_script": model_def_script,
+        "lower_input_kind": model.export_lower_input_kind(),
         "@variables": {},
     }
     if model.get_min_nbor_dist() is not None:

deepmd/pt_expt/kernels/cute/sezm/output_grid/kernels/readout_l0.py

@@ -15,6 +15,7 @@
 
 import cutlass
 import cutlass.cute as cute
+import cutlass.pipeline as pipeline
 from cuda.bindings.driver import (
     CUstream,
 )
@@ -32,7 +33,6 @@
     TILE_K,
     TILE_M,
     TILE_N,
-    TiledOutputGridProductBackward,
 )
 
 # CuTe JIT functions use DSL-inferred argument and return types.
@@ -116,11 +116,16 @@ def kernel(
         out[node, channel] = value.to(out.element_type)
 
 
-class TiledReadoutL0GramBackward(TiledOutputGridProductBackward):
+class TiledReadoutL0GramBackward:
     """Apply the frozen 48x48 Gram matrix to one 64-channel tile."""
 
     def __init__(self) -> None:
-        super().__init__(HIDDEN_CHANNELS)
+        self.cta_tiler = (TILE_M, TILE_N, TILE_K)
+        self.channel_tiles = HIDDEN_CHANNELS // TILE_N
+        self.cta_sync_barrier = pipeline.NamedBarrier(
+            barrier_id=1,
+            num_threads=THREADS,
+        )
 
     @cute.jit
     def __call__(