deepmd.kernels.cute.sezm

CuTe-DSL fused SO(2) value-path operator for SeZM / DPA4.

This package hosts a single bucketed CuTe operator that folds the entire per-edge value path of SO2Convolution (rotate_to_local -> radial degree mix -> the three-layer gated SO(2) mixing stack -> focus competition) into a fused forward kernel and a matching recompute backward kernel, keeping the per-edge intermediates on chip. It is an opt-in inference path enabled by DP_CUTE_INFER; the final local features are handed to the committed flash-attention aggregation for rotate-back and scatter. Kernel entry points are internal implementation details of the SeZM descriptor; the package-level API only exposes availability and the value-path factory.

Submodules

• deepmd.kernels.cute.sezm.backward
• deepmd.kernels.cute.sezm.forward
• deepmd.kernels.cute.sezm.operator

Attributes

SEZM_CUTE_AVAILABLE

Functions

make_cute_value_path(→ _CuteSO2ValuePath | None)

Build the fused CuTe value-path entry for a convolution block.

Package Contents

deepmd.kernels.cute.sezm.SEZM_CUTE_AVAILABLE = True

deepmd.kernels.cute.sezm.make_cute_value_path(conv: deepmd.pt.model.descriptor.sezm_nn.so2.SO2Convolution) → _CuteSO2ValuePath | None

Build the fused CuTe value-path entry for a convolution block.

Parameters:

convSO2Convolution

The convolution block to accelerate.

Returns:

_CuteSO2ValuePath or None

The entry callable when the CuTe backend is available and conv matches the validated configuration; otherwise None (the caller falls back to the reference path).