[pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

Author: pre-commit-ci[bot]

deepmd/kernels/cute/neo/compile_cache.py

@@ -2,13 +2,25 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 """Device-aware caching for architecture-specific CuTe compilation."""
 
-from __future__ import annotations
-
-from collections.abc import Callable
-from contextlib import nullcontext
-from functools import lru_cache, wraps
-from typing import Any, TypeVar, cast
+from __future__ import (
+    annotations,
+)
 
+from collections.abc import (
+    Callable,
+)
+from contextlib import (
+    nullcontext,
+)
+from functools import (
+    lru_cache,
+    wraps,
+)
+from typing import (
+    Any,
+    TypeVar,
+    cast,
+)
 
 _T = TypeVar("_T", bound=Callable[..., Any])
 

deepmd/kernels/cute/neo/gie.py

@@ -15,23 +15,34 @@
 fusion by default; ``DP_CUTE_GIE=0`` disables it explicitly.
 """
 
-from __future__ import annotations
+from __future__ import (
+    annotations,
+)
 
 import threading
-from typing import TYPE_CHECKING, Any
+from typing import (
+    TYPE_CHECKING,
+    Any,
+)
 
 import torch
-from torch import Tensor
+from torch import (
+    Tensor,
+)
 
 if TYPE_CHECKING:
-    from collections.abc import Callable
+    from collections.abc import (
+        Callable,
+    )
 
 try:
     import cutlass
     import cutlass.cute as cute
     import cutlass.torch as cutlass_torch
     from cuda.bindings.driver import CUstream  # noqa: TC002
-    from cutlass.cute.runtime import from_dlpack
+    from cutlass.cute.runtime import (
+        from_dlpack,
+    )
 
     SEZM_CUTE_GIE_AVAILABLE = True
 except Exception:  # pragma: no cover - import guard for non-CuTe environments
@@ -40,7 +51,9 @@
 
 def is_cute_gie_enabled(device: torch.device | None = None) -> bool:
     """Return whether the architecture-selected GIE path is enabled."""
-    from .runtime_policy import is_gie_enabled
+    from .runtime_policy import (
+        is_gie_enabled,
+    )
 
     if device is None:
         if not torch.cuda.is_available():

deepmd/kernels/cute/neo/k1.py

@@ -2,21 +2,41 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 """Opt-in CuTe Neo K1 custom op for SeZM/DPA4 inference."""
 
-from __future__ import annotations
+from __future__ import (
+    annotations,
+)
 
 import threading
 import weakref
-from dataclasses import dataclass, field
-from functools import lru_cache
-from types import SimpleNamespace
-from typing import Any
+from dataclasses import (
+    dataclass,
+    field,
+)
+from functools import (
+    lru_cache,
+)
+from types import (
+    SimpleNamespace,
+)
+from typing import (
+    Any,
+)
 
 import torch
-from torch import Tensor
+from torch import (
+    Tensor,
+)
 
-from deepmd.pt.model.descriptor.sezm_nn.norm import EquivariantRMSNorm
-from . import runtime_policy
-from .k1_message_grid_packed import is_supported_message_grid
+from deepmd.pt.model.descriptor.sezm_nn.norm import (
+    EquivariantRMSNorm,
+)
+
+from . import (
+    runtime_policy,
+)
+from .k1_message_grid_packed import (
+    is_supported_message_grid,
+)
 
 
 @dataclass(frozen=True)
@@ -836,7 +856,9 @@ def _neo_so2_linear_backward_input_with_residual(
     inplace_residual: bool = False,
     out: Tensor | None = None,
 ) -> Tensor:
-    from .k1_so2linear import cached_neo_so2_linear_weights
+    from .k1_so2linear import (
+        cached_neo_so2_linear_weights,
+    )
 
     w0, wpair = cached_neo_so2_linear_weights(so2_linear)
     cache = getattr(so2_linear, "_deepmd_cute_neo_manual_weights_t", None)

deepmd/kernels/cute/neo/k1_gate_structural.py

@@ -7,12 +7,18 @@
 slice of ``grad_y`` through the GEMM beta epilogue.
 """
 
-from __future__ import annotations
+from __future__ import (
+    annotations,
+)
 
-from typing import Any
+from typing import (
+    Any,
+)
 
 import torch
-from torch import Tensor
+from torch import (
+    Tensor,
+)
 
 FOCUS_COUNT = 2
 CHANNELS = 32

deepmd/kernels/cute/neo/k1_kernels/cute_envelope_gated_softmax.py

@@ -8,21 +8,35 @@
 
 # ruff: noqa: ANN001, ANN201, ANN204, TC002
 
-from __future__ import annotations
+from __future__ import (
+    annotations,
+)
 
-from dataclasses import dataclass
-from typing import TYPE_CHECKING
+from dataclasses import (
+    dataclass,
+)
+from typing import (
+    TYPE_CHECKING,
+)
 
 import cutlass
 import cutlass.cute as cute
-from cuda.bindings.driver import CUstream
-from cutlass.cute.runtime import make_fake_compact_tensor, make_fake_stream
-
-from ..compile_cache import device_aware_lru_cache
-
+from cuda.bindings.driver import (
+    CUstream,
+)
+from cutlass.cute.runtime import (
+    make_fake_compact_tensor,
+    make_fake_stream,
+)
+
+from ..compile_cache import (
+    device_aware_lru_cache,
+)
 
 if TYPE_CHECKING:
-    from collections.abc import Callable
+    from collections.abc import (
+        Callable,
+    )
 
 
 @dataclass(frozen=True)

deepmd/kernels/cute/neo/k1_kernels/cute_neo_focus_src_backward.py

@@ -11,16 +11,26 @@
 The kernel is specialized to the Neo K1 shape `(E, F=2, C=32)`.
 """
 
-from __future__ import annotations
+from __future__ import (
+    annotations,
+)
 
-from functools import lru_cache
-from typing import Callable
+from functools import (
+    lru_cache,
+)
+from typing import (
+    Callable,
+)
 
 import cutlass
 import cutlass.cute as cute
-from cuda.bindings.driver import CUstream
-from cutlass.cute.runtime import make_fake_compact_tensor, make_fake_stream
-
+from cuda.bindings.driver import (
+    CUstream,
+)
+from cutlass.cute.runtime import (
+    make_fake_compact_tensor,
+    make_fake_stream,
+)
 
 FAKE_TENSOR_KW = {"assumed_align": 16, "use_32bit_stride": True}
 

deepmd/kernels/cute/neo/k1_kernels/cute_neo_gate_linear_residual_backward_fused.py

@@ -2,22 +2,36 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 """CuTe Neo gate-linear + gate/residual backward without saved logits."""
 
-from __future__ import annotations
+from __future__ import (
+    annotations,
+)
 
-from typing import TYPE_CHECKING
-
-# CuTe JIT functions use DSL-inferred argument and return types.
-# ruff: noqa: ANN001, ANN201, ANN202, TC002
+from typing import (
+    TYPE_CHECKING,
+)
 
 import cutlass
 import cutlass.cute as cute
-from cuda.bindings.driver import CUstream
-from cutlass.cute.runtime import make_fake_compact_tensor, make_fake_stream
+from cuda.bindings.driver import (
+    CUstream,
+)
+from cutlass.cute.runtime import (
+    make_fake_compact_tensor,
+    make_fake_stream,
+)
+
+from ..compile_cache import (
+    device_aware_lru_cache,
+)
+
+# CuTe JIT functions use DSL-inferred argument and return types.
+# ruff: noqa: ANN001, ANN201, ANN202, TC002
 
-from ..compile_cache import device_aware_lru_cache
 
 if TYPE_CHECKING:
-    from collections.abc import Callable
+    from collections.abc import (
+        Callable,
+    )
 
 
 FAKE_TENSOR_KW = {"assumed_align": 16, "use_32bit_stride": True}

deepmd/kernels/cute/neo/k1_kernels/cute_neo_gate_split_structural_vec4_sm80.py

@@ -9,18 +9,30 @@
 split-gate tensor contract and leave dense projections in PyTorch/cuBLAS.
 """
 
-from __future__ import annotations
+from __future__ import (
+    annotations,
+)
 
-from collections.abc import Callable
-from functools import lru_cache
+from collections.abc import (
+    Callable,
+)
+from functools import (
+    lru_cache,
+)
 
 import cutlass
 import cutlass.cute as cute
-from cuda.bindings.driver import CUstream
-from cutlass.cute.runtime import make_fake_compact_tensor, make_fake_stream
-
-from .. import runtime_policy
-
+from cuda.bindings.driver import (
+    CUstream,
+)
+from cutlass.cute.runtime import (
+    make_fake_compact_tensor,
+    make_fake_stream,
+)
+
+from .. import (
+    runtime_policy,
+)
 
 FOCUS_COUNT = 2
 REDUCED_COUNT = 10
@@ -39,7 +51,9 @@ def _guard_vec4_dispatch(
     kernel: Callable,
     tensor_names: tuple[str, ...],
 ) -> Callable:
-    from ..k1_gate_structural import _dispatch_aligned_vec4_kernel
+    from ..k1_gate_structural import (
+        _dispatch_aligned_vec4_kernel,
+    )
 
     def dispatch(*tensors: object):
         return _dispatch_aligned_vec4_kernel(kernel, tensor_names, *tensors)

deepmd/kernels/cute/neo/k1_kernels/cute_neo_message_grid_product.py

@@ -2,21 +2,34 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 """Packed strict-FP32 CuTe grid product for Neo's F=2 K1 branch."""
 
-from __future__ import annotations
-
-# CuTe JIT functions use DSL-inferred argument and return types.
-# ruff: noqa: ANN001, ANN201, TC003
-
-from collections.abc import Callable
-from functools import lru_cache
-from typing import Any
+from __future__ import (
+    annotations,
+)
+
+from collections.abc import (
+    Callable,
+)
+from functools import (
+    lru_cache,
+)
+from typing import (
+    Any,
+)
 
 import cutlass
 import cutlass.cute as cute
 import cutlass.pipeline as pipeline
-from cutlass.cute.runtime import make_fake_compact_tensor, make_fake_stream
+from cutlass.cute.runtime import (
+    make_fake_compact_tensor,
+    make_fake_stream,
+)
 
-from .. import runtime_policy
+from .. import (
+    runtime_policy,
+)
+
+# CuTe JIT functions use DSL-inferred argument and return types.
+# ruff: noqa: ANN001, ANN201, TC003
 
 
 PACKED_COEFF_DIM = 48