fix(dpmodel,cc): strictly positive phantom floor; device-edge flat aparam; effective-layer pair capability

Round-3 review findings (OutisLi + njzjz-bot re-review):

1. The signed phantom denominator could cross zero for real logits
   below phantom_logit (the smooth envelope maps raw < -attnw_shift to
   l < -shift at sw > 0 -- reachable with finite, valid parameters),
   producing a pole / the where-guard's invalid weights-sum>1 fallback.
   segment_softmax now applies an exponential-tail floor on the signed
   path: D~ = D + delta*exp(-D/delta), delta = sel*ph_term/40. It is
   C-infinity, monotone toward deeper deficits, bounded below by delta
   (weights <= 40*n_real/sel), and in the in-design regime (all logits
   >= phantom_logit, so D >= sel*ph_term = 40*delta) the correction is
   <= exp(-40)*delta -- it underflows to EXACT identity in float64, so
   the dense term-for-term parity remains bit-preserved (verified: the
   repformer parity suites still pass at 1e-12). Regressions: both
   reviewers' repros (sel=1 logits [-21,-21]; three entries at -100),
   a resolution-scaling smoothness sweep across the former zero
   crossing, and an in-design bit-parity pin vs the exact fixed-width
   softmax.

2. compute_edges_gpu_impl (Kokkos device-edge route) still padded
   aparam to rank-3 (1, nnode, daparam) and fed it to run_model_graph,
   which the flat-ABI artifact rejects; it now routes through
   extend_graph_aparam (flat (N, daparam), width-validated, ghost-only
   synthesis). The phantom edge-schema branch keeps its rank-3 layout
   (run_model_edges_with_comm is a different ABI).

3. uses_compact_edge_pairs() read the CONFIGURED repformer args, but
   the last layer is built with update_chnnl_2=False which forces its
   g2/h2 updates off -- an nlayers=1 model never builds compact pairs
   yet was rejected on torch < 2.6. The capability now reads the
   EFFECTIVE per-layer flags (the same gate call_graph uses); the guard
   regression uses nlayers=2 for the positive cases and pins that
   nlayers=1 is accepted on torch 2.5.1.

Author: Han Wang

deepmd/dpmodel/descriptor/dpa2.py

@@ -757,13 +757,22 @@ def uses_compact_edge_pairs(self) -> bool:
         EquiVarApply).  ``check_graph_trace_torch_version`` keys its
         torch >= 2.6 requirement on this capability.
 
+        Derived from the EFFECTIVE per-layer flags, not the configured
+        arguments: the LAST repformer layer is built with
+        ``update_chnnl_2=False``, which forces its ``update_g2_has_attn``
+        and ``update_h2`` off -- so e.g. ``nlayers=1`` never runs
+        ``center_edge_pairs`` even when the arguments enable it, and old
+        torch stays usable.  This mirrors the gate
+        ``DescrptBlockRepformers.call_graph`` itself uses to decide whether
+        to build the pairs.
+
         Returns
         -------
         bool
             Whether tracing :meth:`call_graph` runs ``center_edge_pairs``.
         """
-        return bool(
-            self.repformer_args.update_g2_has_attn or self.repformer_args.update_h2
+        return any(
+            ll.update_g2_has_attn or ll.update_h2 for ll in self.repformers.layers
         )
 
     def disable_graph_lower(self) -> None:

deepmd/dpmodel/utils/neighbor_graph/segment.py

@@ -110,13 +110,17 @@ def segment_softmax(
     carry-all regime ``n_real > sel``, where a clamped (non-negative) count
     would drop the compensation and leave a finite ``exp(-attnw_shift)``
     denominator step at the crossing.  For ``n_real <= sel`` the scheme is
-    term-for-term equal to the dense softmax.  The denominator is positive
-    whenever every real logit is ``>= phantom_logit`` (guaranteed by the
-    smooth-envelope logit construction for sane pre-shift logits); the
-    existing non-positive-denominator guard below keeps the out-of-design
-    corner defined.
+    term-for-term equal to the dense softmax.  Real logits are NOT bounded
+    below by ``phantom_logit`` (the smooth envelope maps pre-shift logits
+    ``raw < -attnw_shift`` to ``l < phantom_logit`` at ``sw > 0``), so for
+    negative counts the raw signed denominator can cross zero; a smooth
+    strictly-positive floor (see the inline comment at the denominator)
+    keeps the normalization finite, positive and C-infinity for arbitrary
+    finite logits while perturbing the in-design regime by only
+    ~``exp(-2 * attnw_shift)`` relative.
     """
     xp = array_api_compat.array_namespace(data)
+    dev = array_api_compat.device(data)
     if mask is not None:
         # broadcast mask (n,) over any trailing feature dims of data (n, *f)
         mask_b = xp.reshape(mask, mask.shape + (1,) * (data.ndim - 1))
@@ -154,9 +158,46 @@ def segment_softmax(
         ex = ex * xp.astype(mask_b, ex.dtype)
     denom = segment_sum(ex, segment_ids, num_segments)
     if ph_b is not None:
-        denom = denom + xp.astype(ph_b, denom.dtype) * xp.exp(
-            xp.full_like(seg_max, phantom_logit) - seg_max
-        )
+        ph_term = xp.exp(xp.full_like(seg_max, phantom_logit) - seg_max)
+        denom = denom + xp.astype(ph_b, denom.dtype) * ph_term
+        # Smooth strictly-positive floor for the SIGNED compensation: real
+        # logits are not bounded below by ``phantom_logit`` (the smooth
+        # envelope maps ``raw < -shift`` to ``l < -shift`` at ``sw > 0``),
+        # so with a negative count the signed denominator D can cross zero
+        # -- a pole in the attention weights reachable with finite, valid
+        # model parameters.  Exponential-tail floor:
+        #
+        #     D~ = D + delta * exp(-D / delta),   delta = sel*ph_term / 40
+        #
+        # Properties: (i) C-infinity in the logits and monotone-decreasing
+        # towards larger deficits; (ii) D~ >= delta > 0 for ANY finite
+        # logits (minimum at D == 0), so the weights stay finite and are
+        # bounded by ``40 * n_real / sel``; (iii) in the in-design regime
+        # every real logit is >= phantom_logit, hence D >= sel * ph_term
+        # == 40 * delta and the correction is <= exp(-40) * delta -- more
+        # than 1e-17 RELATIVE below D, i.e. it underflows to EXACT identity
+        # in float64 and the dense term-for-term parity is bit-preserved;
+        # (iv) the boundary-edge cancellation survives (a smooth function
+        # of the already-continuous D).  The exponent is clamped at 700 to
+        # avoid inf*0 pathologies for astronomically negative D (there the
+        # floor is astronomically large and the weights are ~0, as a
+        # maximal phantom deficit should give).
+        if mask is not None:
+            n_real_seg = segment_sum(
+                xp.astype(mask, denom.dtype), segment_ids, num_segments
+            )
+        else:
+            n_real_seg = segment_sum(
+                xp.ones(data.shape[:1], dtype=denom.dtype, device=dev),
+                segment_ids,
+                num_segments,
+            )
+        sel_b = xp.reshape(
+            n_real_seg, n_real_seg.shape + (1,) * (denom.ndim - 1)
+        ) + xp.astype(ph_b, denom.dtype)
+        delta = xp.maximum(sel_b, xp.ones_like(sel_b)) * ph_term / 40.0
+        e_arg = xp.minimum(-denom / delta, xp.full_like(denom, 700.0))
+        denom = denom + delta * xp.exp(e_arg)
     denom_e = xp.take(denom, segment_ids, axis=0)
     safe = xp.where(denom_e > 0, denom_e, xp.ones_like(denom_e))
     return ex / safe

source/api_cc/src/DeepPotPTExpt.cc

@@ -2386,12 +2386,14 @@ void DeepPotPTExpt::compute_edges_gpu_impl(double* d_atom_energy,
           torch::full({1}, static_cast<std::int64_t>(nnode), opt_i64);
       const at::Tensor n_local =
           torch::full({1}, static_cast<std::int64_t>(nloc), opt_i64);
-      at::Tensor graph_aparam = aparam_tensor;
-      if (daparam > 0 && nnode > nloc) {
-        graph_aparam = torch::cat(
-            {aparam_tensor, torch::zeros({1, nnode - nloc, daparam}, opt_f64)},
-            1);
-      }
+      // Flat (N, daparam) graph ABI: validate the width, zero-pad ghost
+      // rows, and synthesize a zero tensor on a ghost-only subdomain --
+      // the rank-3 (1, nnode, daparam) pad this branch used before is
+      // rejected by the artifact (GeneralFitting.call_graph requires
+      // rank-2 aparam), so device-edge graph inference with
+      // numb_aparam > 0 failed at the artifact boundary.
+      at::Tensor graph_aparam =
+          extend_graph_aparam(aparam_tensor, nnode, nloc, daparam);
       GraphTensorPack graph_pack;
       graph_pack.atype = atype_t.reshape({nnode});
       graph_pack.n_node = n_node;

source/tests/common/dpmodel/test_segment_softmax.py

@@ -110,3 +110,108 @@ def test_masked_entry_larger_than_unmasked_max_no_nan() -> None:
     ref = np.exp([1.0, 2.0]) / np.exp([1.0, 2.0]).sum()
     np.testing.assert_allclose(out[:2], ref, rtol=1e-12)
     assert out[2] == 0.0
+
+
+class TestSignedPhantomStrictPositivity:
+    """The SIGNED phantom denominator must stay strictly positive for
+    arbitrary finite logits (OutisLi / njzjz-bot review).
+
+    Real logits are not bounded below by ``phantom_logit``: the smooth
+    envelope maps pre-shift logits ``raw < -attnw_shift`` to
+    ``l < phantom_logit`` whenever ``sw > 0``.  With a negative count the
+    raw signed denominator ``sum exp(l) + (sel - n) exp(ph)`` then crosses
+    zero (e.g. ``sel=1``, logits ``[-21, -21]``, ``ph=-20``), which used to
+    hit the ``denom > 0`` where-guard and return weights ``[1, 1]`` (sum 2)
+    on one side and a pole on the other.  The smooth floor
+    ``(D + sqrt(D^2 + delta^2)) / 2`` keeps the normalization finite,
+    positive and smooth for any finite logits.
+    """
+
+    def test_reviewer_repro_finite_positive(self) -> None:
+        # sel=1, two real entries below the phantom logit -> raw D < 0.
+        data = np.array([-21.0, -21.0])
+        seg = np.array([0, 0])
+        w = segment_softmax(
+            data,
+            seg,
+            1,
+            phantom_count=np.array([-1.0]),  # sel - n_real = 1 - 2
+            phantom_logit=-20.0,
+        )
+        assert np.all(np.isfinite(w))
+        assert np.all(w >= 0.0)
+        # a positive normalization cannot return the where-guard's [1, 1]
+        assert w.sum() < 2.0
+
+    def test_njzjz_repro_three_entries(self) -> None:
+        data = np.array([-100.0, -100.0, -100.0])
+        seg = np.array([0, 0, 0])
+        w = segment_softmax(
+            data,
+            seg,
+            1,
+            phantom_count=np.array([-2.0]),  # sel=1, n_real=3
+            phantom_logit=-20.0,
+        )
+        assert np.all(np.isfinite(w))
+        assert np.all(w >= 0.0)
+        # deep-suppressed entries with a huge phantom deficit -> near-zero
+        # weights, NOT the where-guard's [1, 1, 1].
+        assert w.sum() < 1.0
+
+    def test_smooth_across_former_zero_crossing(self) -> None:
+        """Sweep a logit through the raw denominator's zero crossing: the
+        weights must stay finite and BOUNDED (<= 40 * n_real / sel from the
+        floor scale), and their increments must SCALE with the sweep
+        resolution -- the smoothness signature a pole cannot fake (at a
+        pole, refining the grid does not shrink the largest step).
+        """
+
+        def _sweep(lo: float, hi: float, num: int) -> np.ndarray:
+            outs = []
+            for l2 in np.linspace(lo, hi, num):
+                w = segment_softmax(
+                    np.array([-21.0, float(l2)]),
+                    np.array([0, 0]),
+                    1,
+                    phantom_count=np.array([-1.0]),
+                    phantom_logit=-20.0,
+                )
+                assert np.all(np.isfinite(w))
+                assert np.all(w >= 0.0)
+                # floor-scale bound: 40 * n_real / sel = 80 here
+                assert np.all(w <= 80.0)
+                outs.append(w)
+            return np.stack(outs)
+
+        # D_raw(l) = exp(-21) + exp(l) - exp(-20) crosses zero at l = ln(
+        # exp(-20) - exp(-21)) ~= -20.4587 (the LocalAtten sw-sweep pole
+        # from the review, expressed directly in logit space).
+        crossing = np.log(np.exp(-20.0) - np.exp(-21.0))
+        coarse = _sweep(crossing - 0.5, crossing + 0.5, 201)
+        step_coarse = np.abs(np.diff(coarse, axis=0)).max()
+        # refine 10x: a smooth function's largest step shrinks ~10x; a pole
+        # (or the old where-guard plateau jump) would keep an O(1) step.
+        fine = _sweep(crossing - 0.5, crossing + 0.5, 2001)
+        step_fine = np.abs(np.diff(fine, axis=0)).max()
+        assert step_fine < 0.2 * step_coarse, (
+            f"steps do not shrink under refinement (coarse {step_coarse:.4f} "
+            f"-> fine {step_fine:.4f}): discontinuity or pole at the crossing"
+        )
+
+    def test_floor_does_not_disturb_dense_parity_regime(self) -> None:
+        """In the in-design regime (all logits >= phantom_logit, count >= 0)
+        the floor's correction is ~exp(-2*shift) relative -- invisible at
+        the 1e-12 level used by the dense-parity suites.
+        """
+        rng = np.random.default_rng(3)
+        data = rng.normal(size=7)  # normal-scale logits, shift 20 below
+        seg = np.zeros(7, dtype=np.int64)
+        w_floor = segment_softmax(
+            data, seg, 1, phantom_count=np.array([3.0]), phantom_logit=-20.0
+        )
+        # reference: the exact fixed-width softmax with 3 phantom slots
+        full = np.concatenate([data, np.full(3, -20.0)])
+        ref = np.exp(full - full.max())
+        ref = ref / ref.sum()
+        np.testing.assert_allclose(w_floor, ref[:7], rtol=1e-12, atol=1e-15)

source/tests/pt_expt/utils/test_graph_pt2_metadata.py

@@ -275,11 +275,20 @@ class _NoDesc:
 @pytest.mark.parametrize(
     ("repformer_overrides", "should_raise"),
     [
-        ({}, True),  # default dpa2: update_g2_has_attn=True -> compact pairs
-        ({"update_g2_has_attn": False, "update_h2": True}, True),  # h2 consumer
-        ({"update_g2_has_attn": False, "update_h2": False}, False),  # no pairs
+        # nlayers >= 2 so a non-last layer actually consumes compact pairs
+        # (the LAST layer is built with update_chnnl_2=False, which forces
+        # its g2/h2 updates off).
+        ({"nlayers": 2}, True),  # default update_g2_has_attn=True
+        ({"nlayers": 2, "update_g2_has_attn": False, "update_h2": True}, True),
+        (
+            {"nlayers": 2, "update_g2_has_attn": False, "update_h2": False},
+            False,
+        ),  # no pair consumers on any layer
+        # nlayers=1: the only layer is the last -> NO effective compact-pair
+        # consumer even with the arguments enabled; torch 2.5 stays usable.
+        ({"nlayers": 1}, False),
     ],
-    ids=["default_g2_attn", "update_h2", "no_pair_consumers"],
+    ids=["g2_attn_2layers", "update_h2_2layers", "no_pair_consumers", "single_layer"],
 )
 def test_graph_trace_version_guard_dpa2_compact_pairs(
     monkeypatch, repformer_overrides, should_raise
@@ -293,7 +302,10 @@ def test_graph_trace_version_guard_dpa2_compact_pairs(
     the descriptor capability ``uses_compact_edge_pairs()``: DPA2's
     ``update_g2_has_attn`` (default True) and ``update_h2`` both run the
     compact ``center_edge_pairs`` realization; with both off the lower
-    traces backed symbols only and old torch stays usable.
+    traces backed symbols only and old torch stays usable.  The capability
+    reads the EFFECTIVE per-layer flags: the last layer's g2/h2 updates
+    are structurally off (``update_chnnl_2=False``), so a single-layer
+    repformer never builds compact pairs and must NOT be rejected.
     """
     import torch