deepmd.pt_expt.kernels.triton.sezm.radial_mix

deepmd.pt_expt.kernels.triton.sezm.radial_mix#

Fused Triton dynamic radial degree mixer for the SeZM/DPA4 descriptor.

This module provides a clean-room Triton implementation of the degree_channel branch of DynamicRadialDegreeMixer for the mmax == 1 reduced layout. The eager reference applies, per edge e and output coefficient o:

out[e, o, c] = sum_r channel_basis[r, c] * sum_i K_r[e, o, i] * x[e, i, c]

where K_r is the edge-conditioned degree kernel obtained by scattering the projected radial features compact into a (reduced_dim, reduced_dim) matrix. K_r is block-diagonal over the |m| groups, so for mmax == 1 only a (lmax+1) x (lmax+1) block (orders m = 0) and two identical lmax x lmax blocks (orders m = -1 and m = +1) are non-zero.

Design goals#

  1. Skip the structural zeros and the dense scratch. The eager path materializes the dense kernel (E, reduced_dim, reduced_dim, rank) via a scatter and then contracts it with a batched einsum/bmm whose matrices are tiny (reduced_dim <= 16), which is inefficient on cuBLAS and wastes roughly two thirds of the multiply-adds on off-block zeros. The kernel instead reads compact directly and contracts only the structural non-zeros, with the channel axis vectorized and one program per edge.

  2. Match eager fp32 accuracy. Accumulation is in fp32, matching the smooth potential-energy surface contract used throughout the SeZM descriptor.

  3. Compose with the SeZM ``make_fx`` lowering *and* the AOTInductor freeze. The forward and backward are functional torch.library.triton_op instances (mutates_args=()) with registered fake kernels and an autograd formula, so make_fx(tracing_mode="symbolic") captures the energy path together with the force autograd graph used by inference. triton_op + wrap_triton (vs custom_op) lets Inductor see through to the Triton kernel and bake the cubin into the AOTInductor .pt2, so the frozen package runs the fused mixer inside the LAMMPS C++ runtime without any Python op registration.

Inference-only contract#

The operator is opt-in through DP_TRITON_INFER >= 1 and is only used in evaluation, where the force is obtained from autograd.grad(energy, coord). The backward therefore returns gradients with respect to compact and x_local (both of which carry a path to the coordinates) and None for channel_basis, which is a parameter and never differentiated by the force computation.

Attributes#

RADIAL_MIX_TRITON_AVAILABLE

Functions#

radial_mix_reference(→ torch.Tensor)

Eager ground truth for radial_mix_block().

radial_mix_block(→ torch.Tensor)

Apply the block-diagonal dynamic radial degree mixer (mmax == 1).

Module Contents#

deepmd.pt_expt.kernels.triton.sezm.radial_mix.RADIAL_MIX_TRITON_AVAILABLE = True[source]#
deepmd.pt_expt.kernels.triton.sezm.radial_mix.radial_mix_reference(compact: torch.Tensor, x_local: torch.Tensor, channel_basis: torch.Tensor, lmax: int) → torch.Tensor[source]#

Eager ground truth for radial_mix_block().

Parameters:
compactTensor

Projected radial degree kernel with shape (E, degree_kernel_size, R).

x_localTensor

Edge-local reduced features with shape (E, reduced_dim, C).

channel_basisTensor

Per-rank channel basis with shape (R, C).

lmaxint

Maximum spherical-harmonic degree.

Returns:
Tensor

Mixed features with shape (E, reduced_dim, C).

deepmd.pt_expt.kernels.triton.sezm.radial_mix.radial_mix_block(compact: torch.Tensor, x_local: torch.Tensor, channel_basis: torch.Tensor, lmax: int) → torch.Tensor[source]#

Apply the block-diagonal dynamic radial degree mixer (mmax == 1).

Computes the same operation as radial_mix_reference() while avoiding the dense scattered kernel and the tiny batched matmul on CUDA.

Parameters:
compactTensor

Projected radial degree kernel with shape (E, degree_kernel_size, R).

x_localTensor

Edge-local reduced features with shape (E, reduced_dim, C).

channel_basisTensor

Per-rank channel basis with shape (R, C).

lmaxint

Maximum spherical-harmonic degree.

Returns:
Tensor

Mixed features with shape (E, reduced_dim, C).