From e0e99ca71769753cd47817f5c2d70b9ab0d485e0 Mon Sep 17 00:00:00 2001
From: Hark <hark@local>
Date: Sun, 18 Oct 2026 17:46:12 +0000
Subject: [PATCH] =?UTF-8?q?fix(kin):=20gaps=20=E6=8C=89=E5=B9=B8=E5=AD=98?=
 =?UTF-8?q?=E8=83=B6=E5=9B=8A=E5=AF=B9=E6=95=B0=EF=BC=88PAIR=5FBUDGET=3D2e?=
 =?UTF-8?q?6=EF=BC=89=E8=87=AA=E5=8A=A8=E5=88=86=E5=9D=97=EF=BC=8C?=
 =?UTF-8?q?=E6=98=BE=E5=AD=98=E5=B3=B0=E5=80=BC=E4=B8=8E=E8=B0=83=E7=94=A8?=
 =?UTF-8?q?=E6=96=B9=20batch=20=E6=97=A0=E5=85=B3?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 ml/common/kin_torch.py | 29 +++++++++++++++++++++++++++++
 1 file changed, 29 insertions(+)

diff --git a/ml/common/kin_torch.py b/ml/common/kin_torch.py
index 91fb176..31587a1 100644
--- a/ml/common/kin_torch.py
+++ b/ml/common/kin_torch.py
@@ -118,12 +118,41 @@ class Kin:
         return a, b
 
     # ---------------- 间隙 ----------------
+    # 2026-10-19（用户实测 dagger/诊断吃满 24 GB 显存）：粗筛后幸存胶囊对数随 点数 × cap 暴涨，
+    # 一次几千个点、cap=40 mm 时中间张量可达十几 GB（训练里 eval 峰值 15.9 GiB 也是它）。
+    # 现在先按簇级粗筛估出每个点的幸存胶囊对数，按 PAIR_BUDGET 把 batch 切块逐块算，峰值与调用方的 batch 大小无关。
+    PAIR_BUDGET = 2_000_000
+
     def gaps(self, q, cap=40.0, xf=None, max_pairs=None):
         """返回 dict：
              gm  [B,49]：组对最小间隙 min(gap, cap)（gA*7+gB，gA<gB；其余为 cap）
              b, i, j, gap：幸存胶囊对（gap 可微）
         """
         xf = xf if xf is not None else self.cum_xf(q)
+        Rs, Ts = xf
+        B = q.shape[0]
+        with torch.no_grad():
+            cc = torch.einsum('bkij,kj->bki', Rs[:, self.cl_g], self.cl_c.to(q.dtype)) + Ts[:, self.cl_g]
+            d2 = (cc[:, self.cp_a] - cc[:, self.cp_b]).pow(2).sum(-1)
+            per = ((d2 <= (self.cp_lim + cap).pow(2)).long() * self.cp_cnt).sum(1)       # [B] 每点幸存胶囊对数
+            del cc, d2
+            if B <= 1 or int(per.sum()) <= self.PAIR_BUDGET:
+                bounds = [0, B]
+            else:
+                cs = torch.cumsum(per, 0).tolist(); bounds = [0]; base = 0
+                for k, c in enumerate(cs):
+                    if c - base > self.PAIR_BUDGET and k > bounds[-1]:
+                        bounds.append(k); base = cs[k - 1]
+                bounds.append(B)
+        if len(bounds) == 2:
+            return self._gaps_one(q, cap, xf)
+        parts = [self._gaps_one(q[s:e], cap, (Rs[s:e], Ts[s:e])) for s, e in zip(bounds[:-1], bounds[1:])]
+        return {'gm': torch.cat([p['gm'] for p in parts], 0),
+                'b': torch.cat([p['b'] + s for p, s in zip(parts, bounds[:-1])]),
+                'i': torch.cat([p['i'] for p in parts]), 'j': torch.cat([p['j'] for p in parts]),
+                'gap': torch.cat([p['gap'] for p in parts])}
+
+    def _gaps_one(self, q, cap, xf):
         Rs, Ts = xf
         B = q.shape[0]
         with torch.no_grad():