From 0e3fcc6c5aa2be12634e0a20f5b4d6212220a430 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Wed, 6 May 2026 18:33:02 +0800 Subject: [PATCH] fix(pt): replace in-place nlist masking to enable CUDA graph capture `nlist[nlist == -1] = 0` in DescrptBlockRepformers.forward and DescrptBlockRepflows.forward is lowered by TorchScript into an `index_put_` whose value is built by `torch.tensor(0, device=...)` on every forward pass. The per-call scalar allocation triggers a CPU->GPU sync, which is forbidden during CUDA stream capture (`cudaErrorStreamCaptureUnsupported`), making `forward_lower` of DPA-2 and DPA-3 models non-capturable from the LAMMPS C++ plugin. Switch to `nlist = torch.where(nlist == -1, 0, nlist)`, matching the pattern already used in se_atten.py and se_t_tebd.py. `aten::where` with a Python scalar takes the literal through the kernel without allocating a device tensor, so the resulting frozen IR is CUDA-graph capturable. `nlist`/`a_nlist` were already rebound to local copies earlier in both functions, so dropping the in-place mutation has no observable effect on callers. Closes #5432. --- deepmd/pt/model/descriptor/repflows.py | 4 ++-- deepmd/pt/model/descriptor/repformers.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/deepmd/pt/model/descriptor/repflows.py b/deepmd/pt/model/descriptor/repflows.py index a9d2326b93..031287a797 100644 --- a/deepmd/pt/model/descriptor/repflows.py +++ b/deepmd/pt/model/descriptor/repflows.py @@ -496,8 +496,8 @@ def forward( a_sw = a_sw.masked_fill(~a_nlist_mask, 0.0) # set all padding positions to index of 0 # if the a neighbor is real or not is indicated by nlist_mask - nlist[nlist == -1] = 0 - a_nlist[a_nlist == -1] = 0 + nlist = torch.where(nlist == -1, 0, nlist) + a_nlist = torch.where(a_nlist == -1, 0, a_nlist) # get node embedding # [nframes, nloc, tebd_dim] diff --git a/deepmd/pt/model/descriptor/repformers.py b/deepmd/pt/model/descriptor/repformers.py index 75e2f97576..ca8da0ac1b 100644 --- a/deepmd/pt/model/descriptor/repformers.py +++ b/deepmd/pt/model/descriptor/repformers.py @@ -457,7 +457,7 @@ def forward( # set all padding positions to index of 0 # if the a neighbor is real or not is indicated by nlist_mask - nlist[nlist == -1] = 0 + nlist = torch.where(nlist == -1, 0, nlist) # nb x nall x ng1 if comm_dict is None: assert mapping is not None