diff --git a/deepmd/pt_expt/kernels/triton/sezm/so2_edge_cohort.py b/deepmd/pt_expt/kernels/triton/sezm/so2_edge_cohort.py new file mode 100644 index 0000000000..c5a3435f82 --- /dev/null +++ b/deepmd/pt_expt/kernels/triton/sezm/so2_edge_cohort.py @@ -0,0 +1,251 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Experimental edge-cohort backward for frozen, narrow SO(2) inference. + +Keep native GEMM column parallelism and arithmetic, but traverse all layers +of one cohort before advancing to the next cohort. Scratch scales with the +cohort, not the entire graph. The saved activations and returned input gradient +still have full edge extent. This Python runtime launcher is an online-compile +prototype; standalone AOTInductor/C++ deployment is not validated. +""" + +from __future__ import annotations + +import torch +import triton +from torch import Tensor +from torch.library import wrap_triton + +from . import so2_value_path as sv + +COHORT_EDGES = 8192 + + +def supports_edge_cohort( + grad_out: Tensor, z_all: Tensor, lmax: int, focus_dim: int, apply_alpha: bool +) -> bool: + """Restrict opt-in dispatch to the experimentally checked Neo layout.""" + return ( + grad_out.is_cuda + and z_all.device == grad_out.device + and grad_out.dtype == torch.float32 + and z_all.dtype == torch.float32 + and grad_out.is_contiguous() + and z_all.is_contiguous() + and lmax == 3 + and focus_dim == 32 + and apply_alpha + and grad_out.ndim == 3 + and z_all.ndim == 4 + and grad_out.shape[1:] == (2, 320) + and z_all.shape[:2] == (2, 2) + and z_all.shape[2] == grad_out.shape[0] + and z_all.shape[3] == 320 + ) + + +def cohort_input_backward_impl( + grad_out: Tensor, + z_all: Tensor, + alpha: Tensor, + w0t: Tensor, + w1t: Tensor, + gw: Tensor, + gwt: Tensor, + cohort_edges: int, +) -> Tensor: + """Compute input gradients without retaining unrelated forward outputs. + + Kernels infer scratch focus strides from the cohort edge count. Independent + saved/output strides retain full-graph layout, including partial cohorts. + Only an internally owned residual is overwritten: matrix input ``gz`` and + GEMM output never alias, and caller-owned tensors remain read-only. + """ + if not supports_edge_cohort(grad_out, z_all, 3, 32, True): + raise ValueError("edge cohorts require contiguous FP32 Neo inference tensors") + if cohort_edges <= 0: + raise ValueError("cohort_edges must be positive") + if not all(t.is_contiguous() for t in (alpha, w0t, w1t, gw, gwt)): + raise ValueError("edge-cohort parameters must be contiguous") + n_gated, n_focus, n_edge, row = z_all.shape + result = grad_out.new_empty((n_focus, n_edge, row)) + if n_edge == 0: + return result + cap = min(n_edge, cohort_edges) + g = grad_out.new_empty((n_focus, cap, row)) + gz = torch.empty_like(g) + sig = grad_out.new_empty((n_focus, cap, 96)) + bm, bn, bk, warps, stages = sv.stack_fp32_configs(32, 3)[2] + n_tiles = triton.cdiv(128, bn) + triton.cdiv(192, bn) + inside, point_config, recompute_config = sv._point_backward_schedule( + 32, 3, train=False + ) + for start in range(0, n_edge, cap): + count = min(cap, n_edge - start) + go = grad_out[start:] + a = alpha[start:] + z = z_all.view(-1)[start * row :] + grid = (triton.cdiv(count, bm) * n_tiles, n_focus) + sv._stack_gemm_bwd_kernel[grid]( + go, + go, + w0t, + w1t, + a, + g, + count, + n_gated, + L=3, + CF=32, + G_EDGE_MAJOR=True, + FOLD_ALPHA=True, + RES_IS_GZ=True, + BLOCK_M=bm, + BLOCK_N=bn, + BLOCK_K=bk, + num_warps=warps, + num_stages=stages, + ) + for layer in range(n_gated - 1, -1, -1): + if not inside: + rbm, rw, rs = recompute_config + sv._stack_recompute_kernel[(triton.cdiv(count, rbm), n_focus)]( + z, + gw, + sig, + count, + layer, + L=3, + CF=32, + BLOCK_M=rbm, + saved_edges=n_edge, + num_warps=rw, + num_stages=rs, + ) + pbm, pw, ps = point_config + sv._stack_point_bwd_kernel[(triton.cdiv(count, pbm), n_focus)]( + g, + z, + sig, + gwt, + gz, + sig, + g, + g, + count, + layer, + L=3, + CF=32, + GLOGIT_OUT=False, + GLOGIT_STORE=False, + RECOMPUTE_SIG=inside, + RECOVER_INPUT=False, + BLOCK_M=pbm, + saved_edges=n_edge, + num_warps=pw, + num_stages=ps, + ) + # Residual elements are loaded before their matching output store. + # Distinct (edge, focus, column) tiles own disjoint residual elements. + target = result.view(-1)[start * row :] if layer == 0 else g + sv._stack_gemm_bwd_kernel[grid]( + gz, + g, + w0t, + w1t, + gz, + target, + count, + layer, + L=3, + CF=32, + G_EDGE_MAJOR=False, + FOLD_ALPHA=False, + RES_IS_GZ=False, + BLOCK_M=bm, + BLOCK_N=bn, + BLOCK_K=bk, + output_edges=n_edge if layer == 0 else count, + num_warps=warps, + num_stages=stages, + ) + return result + + +@torch.library.custom_op("sezm_triton::so2_cohort_input_bwd", mutates_args=()) +def _cohort_input_backward( + grad_out: Tensor, + z_all: Tensor, + alpha: Tensor, + w0t: Tensor, + w1t: Tensor, + gw: Tensor, + gwt: Tensor, + cohort_edges: int, +) -> Tensor: + """Keep the runtime edge loop opaque without retaining x_local/u_final.""" + return cohort_input_backward_impl( + grad_out, z_all, alpha, w0t, w1t, gw, gwt, cohort_edges + ) + + +@_cohort_input_backward.register_fake +def _cohort_input_backward_fake( + grad_out: Tensor, + z_all: Tensor, + alpha: Tensor, + w0t: Tensor, + w1t: Tensor, + gw: Tensor, + gwt: Tensor, + cohort_edges: int, +) -> Tensor: + """Preserve full input-gradient shape under FakeTensor tracing.""" + return grad_out.new_empty(z_all.shape[1:]) + + +@torch.library.triton_op("sezm_triton::so2_cohort_alpha_bwd", mutates_args=()) +def _cohort_alpha_backward(grad_out: Tensor, x_local: Tensor, alpha: Tensor) -> Tensor: + """Expose alpha's short-lived output dependency to the graph scheduler.""" + n_edge, n_focus = alpha.shape + result = torch.empty_like(alpha) + if sv._has_no_edges(n_edge): + return result + bm, warps, stages = sv.gate_config(32, 3) + wrap_triton(sv._stack_grad_alpha_kernel)[(triton.cdiv(n_edge, bm), n_focus)]( + grad_out, + x_local, + alpha, + result, + n_edge, + L=3, + CF=32, + BLOCK_M=bm, + num_warps=warps, + num_stages=stages, + ) + return result + + +def cohort_backward( + grad_out: Tensor, + x_local: Tensor, + z_all: Tensor, + u_final: Tensor, + alpha: Tensor, + w0t: Tensor, + w1t: Tensor, + gw: Tensor, + gwt: Tensor, + lmax: int, + focus_dim: int, + apply_alpha: bool, +) -> tuple[Tensor, Tensor]: + """First-order dispatcher; training/higher derivatives never enter here.""" + del u_final + if not supports_edge_cohort(grad_out, z_all, lmax, focus_dim, apply_alpha): + raise ValueError("unsupported edge-cohort layout") + grad_alpha = _cohort_alpha_backward(grad_out, x_local, alpha) + grad_u0 = _cohort_input_backward( + grad_out, z_all, alpha, w0t, w1t, gw, gwt, COHORT_EDGES + ) + return grad_u0, grad_alpha 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 7a8ac3fd43..b65d1793e5 100644 --- a/deepmd/pt_expt/kernels/triton/sezm/so2_value_path.py +++ b/deepmd/pt_expt/kernels/triton/sezm/so2_value_path.py @@ -91,6 +91,7 @@ annotations, ) +import os from typing import ( TYPE_CHECKING, NamedTuple, @@ -161,6 +162,10 @@ _ROTATE_MIX_BWD_CONFIG = (1, 2) # per-edge backward (warps, stages) +# Experimental, frozen first-order inference only. Read before compilation; +# changing the switch requires a new process/graph. Disabled by default. +_EDGE_COHORT_INFER = os.environ.get("DP_SO2_EDGE_COHORT_INFER", "0") == "1" + # ====================================================================== # Eager reference / fallback implementations @@ -1369,6 +1374,7 @@ def _stack_recompute_kernel( L: tl.constexpr, CF: tl.constexpr, BLOCK_M: tl.constexpr, + saved_edges=0, ): """Recompute the gate sigmoids from the saved pre-activation (backward).""" ROW: tl.constexpr = (3 * L + 1) * CF @@ -1386,7 +1392,8 @@ def _stack_recompute_kernel( cm = mm & (nc < CF)[None, :] wm = ((nc < CF)[:, None]) & ((nc < CF)[None, :]) - z_row = z_ptr + (layer * n_focus + fid) * n_edge * ROW + offs_m * ROW + z_stride = tl.where(saved_edges > 0, saved_edges, n_edge).to(tl.int64) + z_row = z_ptr + (layer * n_focus + fid) * z_stride * ROW + offs_m * ROW z_s = tl.load(z_row[:, None] + nc[None, :], mask=cm, other=0.0).to(tl.float32) sig_row = sig_ptr + (fid * n_edge + offs_m) * LG for g in tl.static_range(L): @@ -1420,6 +1427,7 @@ def _stack_point_bwd_kernel( RECOMPUTE_SIG: tl.constexpr, RECOVER_INPUT: tl.constexpr, BLOCK_M: tl.constexpr, + saved_edges=0, ): """Pointwise part of the gated-layer backward. @@ -1456,7 +1464,8 @@ def _stack_point_bwd_kernel( wm = ((nc < CF)[:, None]) & ((nc < CF)[None, :]) g_row = g_ptr + fid * n_edge * ROW + offs_m * ROW - z_row = z_ptr + (layer * n_focus + fid) * n_edge * ROW + offs_m * ROW + z_stride = tl.where(saved_edges > 0, saved_edges, n_edge).to(tl.int64) + z_row = z_ptr + (layer * n_focus + fid) * z_stride * ROW + offs_m * ROW gz_row = gz_ptr + fid * n_edge * ROW + offs_m * ROW sig_row = sig_ptr + (fid * n_edge + offs_m) * LG gl_row = gl_ptr + (fid * n_edge + offs_m) * LG @@ -2014,6 +2023,7 @@ def _stack_gemm_bwd_kernel( BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr, + output_edges=0, ): """Backward block GEMM ``g_u = residual + gz @ W^T`` over both blocks. @@ -2044,7 +2054,8 @@ def _stack_gemm_bwd_kernel( gz_row = gz_ptr + fid * ROW + offs_m * (n_focus * ROW) else: gz_row = gz_ptr + fid * n_edge * ROW + offs_m * ROW - gu_row = gu_ptr + fid * n_edge * ROW + offs_m * ROW + output_stride = tl.where(output_edges > 0, output_edges, n_edge).to(tl.int64) + gu_row = gu_ptr + fid * output_stride * ROW + offs_m * ROW if FOLD_ALPHA: alpha = tl.load(alpha_ptr + offs_m * n_focus + fid, mask=m_mask, other=0.0) @@ -4956,7 +4967,18 @@ def _mixing_stack_backward(ctx, grad_out, grad_z, grad_u): None, None, ) - grad_u0, grad_alpha = _mixing_stack_bwd_op( + backward_op = _mixing_stack_bwd_op + if _EDGE_COHORT_INFER: + from .so2_edge_cohort import ( + cohort_backward, + supports_edge_cohort, + ) + + if supports_edge_cohort( + grad_out, z_all, ctx.lmax, ctx.focus_dim, ctx.apply_alpha + ): + backward_op = cohort_backward + grad_u0, grad_alpha = backward_op( grad_out, x_local, z_all, diff --git a/source/tests/pt_expt/kernels/test_so2_edge_cohort.py b/source/tests/pt_expt/kernels/test_so2_edge_cohort.py new file mode 100644 index 0000000000..b7a7f1a8b9 --- /dev/null +++ b/source/tests/pt_expt/kernels/test_so2_edge_cohort.py @@ -0,0 +1,159 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Generated-input checks for the opt-in frozen Neo edge-cohort prototype.""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch + +pytest.importorskip("triton") + +from deepmd.pt_expt.kernels.triton.sezm import so2_edge_cohort as ecs +from deepmd.pt_expt.kernels.triton.sezm import so2_value_path as sv + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + + +def operands(edges: int, mode: str = "random") -> tuple: + """Generate well-conditioned operands without a downloaded checkpoint.""" + torch.manual_seed(719) + u = torch.randn(2, edges, 320, device="cuda") * 0.1 + if mode == "zero": + u.zero_() + alpha = torch.rand(edges, 2, device="cuda") * 0.8 + 0.1 + if mode == "tiny": + alpha.fill_(1e-5) + w0 = torch.randn(3, 2, 128, 128, device="cuda") * 0.01 + w1 = torch.randn(3, 2, 192, 192, device="cuda") * 0.01 + gw = torch.randn(2, 2, 32, 96, device="cuda") * 0.01 + out, z, uf = sv._mixing_stack_impl(u, alpha, w0, w1, gw, 3, 32, True) + go = torch.randn_like(out) * 0.1 + return ( + go, + out, + z, + uf, + alpha, + w0.transpose(2, 3).contiguous(), + w1.transpose(2, 3).contiguous(), + gw, + gw.transpose(2, 3).contiguous(), + 3, + 32, + True, + ) + + +@pytest.mark.parametrize("edges", [0, 1, 37, 129, 1025, 8193]) +@pytest.mark.parametrize("cap", [64, 8192]) +@pytest.mark.parametrize("mode", ["random", "zero", "tiny"]) +def test_cohort_matches_native(edges: int, cap: int, mode: str) -> None: + """Cover tail cohorts, two focus strides, legal tiny alpha and no mutation.""" + args = operands(edges, mode) + copies = [x.clone() for x in args[:9]] + expected_u, expected_a = sv._mixing_stack_bwd_impl(*args) + go, out, z, _, alpha, w0t, w1t, gw, gwt, *_ = args + actual_u = ecs.cohort_input_backward_impl(go, z, alpha, w0t, w1t, gw, gwt, cap) + actual_a = ecs._cohort_alpha_backward(go, out, alpha) + for actual, expected in ((actual_u, expected_u), (actual_a, expected_a)): + assert torch.isfinite(actual).all() + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + for before, after in zip(copies, args[:9], strict=True): + assert torch.equal(before, after) + + +def test_compiled_dynamic_cohort() -> None: + """The opaque runtime edge loop supports online dynamic graph execution.""" + compiled = torch.compile(ecs.cohort_backward, fullgraph=True, dynamic=True) + for edges in (37, 129): + args = operands(edges) + actual = compiled(*args) + expected = sv._mixing_stack_bwd_impl(*args) + for a, b in zip(actual, expected, strict=True): + torch.testing.assert_close(a, b, atol=0, rtol=0) + + +@pytest.mark.parametrize( + "lmax,cf,apply", [(2, 32, True), (3, 64, True), (3, 32, False)] +) +def test_unsupported_layout(lmax: int, cf: int, apply: bool) -> None: + """Untested layouts must not enter the experimental implementation.""" + args = operands(1) + assert not ecs.supports_edge_cohort(args[0], args[2], lmax, cf, apply) + assert not ecs.supports_edge_cohort(args[0].cpu(), args[2].cpu(), 3, 32, True) + + +@pytest.mark.parametrize("reason", ["weights", "saved_z", "saved_u", "create_graph"]) +def test_training_dispatch_preserved(monkeypatch, reason: str) -> None: + """Weight/state/higher-order requests bypass ECS even when opted in.""" + args = operands(1) + go, out, z, uf, alpha, w0t, w1t, gw, *_ = args + u = torch.zeros_like(uf) + w0 = w0t.transpose(2, 3).contiguous() + w1 = w1t.transpose(2, 3).contiguous() + ctx = SimpleNamespace( + saved_tensors=(u, alpha, out, z, uf, w0, w1, gw), + needs_input_grad=( + True, + True, + reason == "weights", + False, + False, + False, + False, + False, + ), + lmax=3, + focus_dim=32, + apply_alpha=True, + ) + calls = [] + + def training(*unused): + calls.append(True) + return torch.zeros_like(u), torch.zeros_like(alpha), w0, w1, gw + + def forbidden(*unused): + raise AssertionError("ECS entered a training/higher-order request") + + monkeypatch.setattr(sv, "_EDGE_COHORT_INFER", True) + monkeypatch.setattr(sv, "_mixing_stack_train_bwd_op", training) + monkeypatch.setattr(ecs, "cohort_backward", forbidden) + with torch.set_grad_enabled(reason == "create_graph"): + sv._mixing_stack_backward( + ctx, + go, + torch.zeros_like(z) if reason == "saved_z" else None, + torch.zeros_like(uf) if reason == "saved_u" else None, + ) + assert calls == [True] + + +def test_frozen_dispatch(monkeypatch) -> None: + """Enabling the switch actually changes the frozen inference dispatch.""" + args = operands(37) + go, out, z, uf, alpha, w0t, w1t, gw, *_ = args + ctx = SimpleNamespace( + saved_tensors=( + uf, + alpha, + out, + z, + uf, + w0t.transpose(2, 3).contiguous(), + w1t.transpose(2, 3).contiguous(), + gw, + ), + needs_input_grad=(True, True, False, False, False, False, False, False), + lmax=3, + focus_dim=32, + apply_alpha=True, + ) + monkeypatch.setattr(sv, "_EDGE_COHORT_INFER", True) + expected = sv._mixing_stack_bwd_impl(*args) + with torch.no_grad(): + actual = sv._mixing_stack_backward(ctx, go, None, None) + for a, b in zip(actual[:2], expected, strict=True): + torch.testing.assert_close(a, b, atol=0, rtol=0)