Skip to content

perf(pt_expt): fuse the mixing-stack per-layer backward chain (level-2 tuned table) #6060

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
long-yi-2019 wants to merge 11 commits into
base: master
Choose a base branch
from
Open

perf(pt_expt): fuse the mixing-stack per-layer backward chain (level-2 tuned table) #6060

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
11 commits
Select commit Hold shift + click to select a range
a3dcd97
perf(pt): fuse the mixing-stack per-layer backward chain (DP_STACK_BW…
Oct 8, 2026
4455e67
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Oct 8, 2026
ad9af59
test(pt): exercise the stack-backward fallback guards with the flag e…
Oct 16, 2026
633e87c
Merge branch 'master' into pr/stack-bwd-fused
OutisLi Oct 17, 2026
fa5e98e
feat(pt): select the fused mixing-stack backward by tuned tile entry
Oct 18, 2026
8476318
test(pt): retarget the fused stack-backward tests at the tuned selection
Oct 18, 2026
2e8e6f3
fix(pt): time fused sweep candidates under no_grad and assert the sel…
Oct 18, 2026
5df6813
fix(pt_expt): make stack_bwd_fused tuning trustworthy and dual-backend
Oct 21, 2026
ef2eaf7
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Oct 21, 2026
f524062
fix(pt): keep the freeze tuning seam import-safe and backend-neutral
Oct 22, 2026
8a61fbb
fix(pt_expt): give the fused stack backward a non-aliasing output buffer
Oct 25, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
50 changes: 14 additions & 36 deletions deepmd/pt/entrypoints/freeze_pt2.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -57,7 +57,6 @@
GeometricInitialEmbedding,
)
from deepmd.pt.model.descriptor.sezm_nn.so2 import (
SO2Convolution,
SO2Linear,
)
from deepmd.pt.model.model import (
Expand Down Expand Up @@ -331,55 +330,34 @@ def _tune_triton_configs(model: torch.nn.Module, target_device: torch.device) ->
width/degree) are swept here on the local GPU -- the exact hardware the
``.pt2`` will run on, since AOTInductor artifacts are not portable across
GPU models -- and registered for the current process before tracing.
Keys already covered cost nothing.

The fused value-path entries are then rebound: the mixing-stack operator
selection (fp32 versus fp16x3) is fixed at construction time, which
predates the registrations made here.
Keys already covered cost nothing. The fused value-path entries are
then rebound (the mixing-stack operator selection, fp32 versus fp16x3,
is fixed at construction time, which predates the registrations).

The whole flow lives in
:func:`deepmd.pt_expt.kernels.triton.sezm.sweep_tile_configs.autotune_sezm`
and is shared with the PT-expt export path, which reaches it through
the freeze-time autotuner registry. The early returns below must stay
ahead of the sweep-module import: that module imports Triton
unconditionally, and this helper runs on every freeze -- a CPU-only or
otherwise Triton-less installation must skip tuning (as it always did)
instead of failing the import.
"""
if triton_infer_level() < 2:
return
if target_device.type != "cuda" or not torch.cuda.is_available():
return
from deepmd.pt_expt.kernels.triton.sezm.so2_value_path import (
SO2_VALUE_PATH_TRITON_AVAILABLE,
make_triton_value_path,
)

if not SO2_VALUE_PATH_TRITON_AVAILABLE:
return
from deepmd.pt_expt.kernels.triton.sezm.sweep_tile_configs import (

@njzjz-bot njzjz-bot Oct 21, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[P2] Keep optional-Triton guards ahead of the sweep-module import

freeze_sezm_to_pt2 calls this helper unconditionally (including a CPU target), but the new helper imports sweep_tile_configs before checking the inference level/device/availability. That module imports triton unconditionally at module scope. On a supported CPU-only installation without Triton, a CPU freeze that previously skipped tuning now fails with ModuleNotFoundError, even with the default DP_TRITON_INFER=0; the guards inside autotune_sezm are never reached. Preserve an import-safe entry seam or the early guards before importing the Triton-dependent sweep. Add a missing-Triton CPU/default-level regression test for the actual freeze seam, not only a CUDA-mocked routing test. Exact old/new helper replay reproduces the changed import behavior; no GPU result is asserted.

collect_model_shape_keys,
tune_missing_configs,
autotune_sezm,
)

# The built-in tables and the sweep both resolve against the current
# device; pin it to the AOTI target so a freeze aimed at a secondary GPU
# tunes and looks up the right hardware (mixed-model hosts).
if target_device.index is not None:
torch.cuda.set_device(target_device)

shape_keys = collect_model_shape_keys(model)
registered = tune_missing_configs(
shape_keys, level=triton_infer_level(), device=target_device
)
if registered:
log.info(
"Registered freshly tuned Triton launch configurations: %s",
{family: sorted(entries) for family, entries in registered.items()},
)
else:
log.info(
"Triton launch tables already cover this checkpoint's shapes on %s; "
"no tuning needed.",
torch.cuda.get_device_name(target_device),
)
# Rebind unconditionally: the fp32-versus-fp16x3 stack selection was made
# at construction time, possibly against a different current device's
# tables, and must reflect the target device and any fresh registrations.
for module in model.modules():
if isinstance(module, SO2Convolution) and module.triton_infer_level >= 2:
module._triton_value_path = make_triton_value_path(module)
autotune_sezm(model, triton_infer_level(), target_device)


# The trace-time sendlist for the with-comm artifact embeds the address of a
Expand Down
Loading
Loading