perf(dsv41): prompt-chunk partials over a reduce-scatter, the same bits from 3/4 of the bytes

src/tensorfold/cuda/comm.py

@@ -72,6 +72,8 @@ def __init__(self, rank: int, world: int, master: str, port: int, *, timeout_s:
                                       ctypes.c_void_p]
         lib.ncclAllReduce.argtypes = [ctypes.c_void_p, ctypes.c_void_p, ctypes.c_size_t, ctypes.c_int, ctypes.c_int,
                                       ctypes.c_void_p, ctypes.c_void_p]
+        lib.ncclReduceScatter.argtypes = [ctypes.c_void_p, ctypes.c_void_p, ctypes.c_size_t, ctypes.c_int, ctypes.c_int,
+                                          ctypes.c_void_p, ctypes.c_void_p]
         self.store = TCPStore(master, port, world, rank == 0, timeout=timedelta(seconds=timeout_s))
         uid = _UniqueId()
         if rank == 0:
@@ -108,6 +110,15 @@ def all_reduce(self, value: torch.Tensor) -> None:
         self._check(self.lib.ncclAllReduce(value.data_ptr(), value.data_ptr(), value.numel(), _DTYPES[value.dtype],
                                            NCCL_SUM, self.comm, stream))
 
+    def reduce_scatter(self, send: torch.Tensor, recv: torch.Tensor) -> None:
+        """recv [n] <- ncclSum over the ranks of their send's rank-th block (send [world * n])."""
+
+        if send.numel() != recv.numel() * self.world or send.dtype != recv.dtype:
+            raise ValueError("reduce_scatter: send must hold world x recv of the same dtype")
+        stream = torch.cuda.current_stream().cuda_stream
+        self._check(self.lib.ncclReduceScatter(send.data_ptr(), recv.data_ptr(), recv.numel(), _DTYPES[send.dtype],
+                                               NCCL_SUM, self.comm, stream))
+
     def ready(self, label: str, *, every: float = 60.0, timeout: float = 3600.0) -> None:
         """Every rank finishes ``label`` before any goes on; a rank missing after ``timeout`` s is named."""
 

src/tensorfold/families/deepseek_v41/cuda/model.py

@@ -21,6 +21,9 @@
 RAW = 64                            # per-position compressor inputs kept (a verify window rolls back by length alone)
 RING_EXTRA = 16                     # window ring slots beyond the 128-token window (verify windows never clobber)
 KV_QUANT = os.environ.get("TF_DS_KV", "native") != "bf16"
+# TF_DS_RS=1 (two ranks): a prompt chunk's attention and MoE partials reach the mHC post as bf16(p0 + p1), from an
+# fp32 reduce-scatter and a bf16 all-gather (Comm.reduced). The same bits as the fp32 gather, 3/4 of its bytes.
+RS = os.environ.get("TF_DS_RS", "0") == "1"
 KERNELS = os.environ.get("TF_DS_KERNELS", "1") != "0"      # 0: the plain-torch phase-1 path (A/B and debugging)
 TIMING = os.environ.get("TF_DS_TIMING", "0") == "1"        # per-section wall times (synchronizing; profiling only)
 # TF_DS_REPLAY_FLOOR=1: a prompt chunk of RING_EXTRA rows or fewer whose window reaches below the bounded replay's cut
@@ -115,6 +118,21 @@ def gather(self, x: torch.Tensor) -> torch.Tensor:
         self.nccl.all_gather(x.view(-1), out.view(-1))
         return out
 
+    def reduced(self, x: torch.Tensor) -> torch.Tensor:
+        """The partials [n, d] the mHC post takes: ``gather(x)`` [world, n, d], or TF_DS_RS on two ranks [1, n, d]
+        bf16 = bf16(p0 + p1) from an fp32 reduce-scatter and a bf16 all-gather (one fp32 add of two operands, the same
+        bits in either order: the same y from 3/4 of the bytes)."""
+
+        rdma = getattr(self.nccl, "rdma", None)
+        if not (RS and self.world == 2 and x.shape[0] % 2 == 0 and x.dtype == F32) or (rdma and rdma.fits(x)):
+            return self.gather(x)
+        x = x.contiguous()
+        half = torch.empty((x.shape[0] // 2, *x.shape[1:]), dtype=F32, device=x.device)
+        self.nccl.reduce_scatter(x.view(-1), half.view(-1))
+        y = torch.empty((1, *x.shape), dtype=BF16, device=x.device)
+        self.nccl.all_gather(half.to(BF16).view(-1), y.view(-1))
+        return y
+
     def sum(self, x: torch.Tensor) -> torch.Tensor:
         """Sum rank partials; NCCL can modify x, so callers must treat x as consumed."""
 
@@ -1377,7 +1395,7 @@ def _forward_k(self, sc, ids, start, all_logits, taps, h, pre, hashes, shared, r
                 pa = self.attention_k(lay, x, sc, start, shared, pos, floor=floor,
                                       kv_done=kv_done and lay.idx == c.n_layers // 2)
             with _T("gather"):
-                g = self.comm.gather(pa)
+                g = self.comm.reduced(pa)
             with _T("hc"):
                 fn, scale, base = lay.hc_ffn
                 if hc_pf:
@@ -1394,7 +1412,7 @@ def _forward_k(self, sc, ids, start, all_logits, taps, h, pre, hashes, shared, r
             with _T("moe"):
                 pm = self.moe(lay, x, img=img)
             with _T("gather"):
-                g = self.comm.gather(pm)
+                g = self.comm.reduced(pm)
             with _T("hc"):
                 K.hc_post(g, h, post, comb, h)
             pre, pre_f = pre_f, pre

tests/test_dsv41_reduced.py

@@ -0,0 +1,101 @@
+"""Comm.reduced (TF_DS_RS): two ranks get bf16(p0 + p1), the y the mHC post kernels form from the fp32 gather, and
+the gather itself where the reduce-scatter does not apply."""
+
+from __future__ import annotations
+
+import threading
+
+import pytest
+
+torch = pytest.importorskip("torch")
+pytest.importorskip("triton")
+
+from tensorfold.families.deepseek_v41.cuda import model as M
+
+
+class Link:
+    """Two ranks' collectives on the host: each call waits for both ranks' buffers, as NCCL does."""
+
+    def __init__(self):
+        self.barrier, self.sends, self.calls = threading.Barrier(2), [None, None], []
+
+    def swap(self, rank, send):
+        self.sends[rank] = send.clone()
+        self.barrier.wait()
+        both = list(self.sends)
+        self.barrier.wait()
+        return both
+
+
+class Rank:
+    def __init__(self, link, rank, rdma=None):
+        self.link, self.rank, self.world = link, rank, 2
+        if rdma is not None:
+            self.rdma = rdma
+
+    def reduce_scatter(self, send, recv):
+        a, b = self.link.swap(self.rank, send)
+        self.link.calls.append(("reduce_scatter", self.rank))
+        recv.copy_((a + b).view(2, -1)[self.rank])                  # ncclSum of two operands
+
+    def all_gather(self, send, recv):
+        a, b = self.link.swap(self.rank, send)
+        self.link.calls.append(("all_gather", send.dtype))
+        recv.copy_(torch.cat([a, b]))
+
+
+class Fits:
+    def fits(self, x):
+        return True
+
+
+def both_ranks(p0, p1, rdma=None):
+    """Comm.reduced on two threads, one a rank: (rank 0's result, rank 1's result, the calls)."""
+
+    link, out = Link(), [None, None]
+
+    def run(rank, mine):
+        c = M.Comm.__new__(M.Comm)
+        c.nccl, c.world, c._buf = Rank(link, rank, rdma), 2, {}
+        out[rank] = c.reduced(mine)
+
+    ts = [threading.Thread(target=run, args=(r, p)) for r, p in ((0, p0), (1, p1))]
+    for t in ts:
+        t.start()
+    for t in ts:
+        t.join()
+    return out[0], out[1], link.calls
+
+
+def parts(n=8, d=64, seed=0):
+    g = torch.Generator().manual_seed(seed)
+    return torch.randn((n, d), generator=g) * 3, torch.randn((n, d), generator=g) * 1e-3
+
+
+def test_both_ranks_get_the_posts_y(monkeypatch):
+    monkeypatch.setattr(M, "RS", True)
+    p0, p1 = parts()
+    y0, y1, calls = both_ranks(p0, p1)
+    want = (p0 + p1).to(torch.bfloat16)[None]                    # what _hc_post forms from the gather, before the mix
+    for y in (y0, y1):
+        assert y.dtype == torch.bfloat16 and torch.equal(y.view(torch.int16), want.view(torch.int16))
+    assert sorted(c[0] for c in calls) == ["all_gather", "all_gather", "reduce_scatter", "reduce_scatter"]
+    assert all(c[1] == torch.bfloat16 for c in calls if c[0] == "all_gather")
+
+
+@pytest.mark.parametrize("case", ["off", "odd rows", "rdma holds it"])
+def test_falls_back_to_the_gather(monkeypatch, case):
+    monkeypatch.setattr(M, "RS", case != "off")
+    p0, p1 = parts(n=7 if case == "odd rows" else 8)
+    y0, y1, calls = both_ranks(p0, p1, Fits() if case == "rdma holds it" else None)
+    for y in (y0, y1):
+        assert y.dtype == torch.float32 and torch.equal(y, torch.stack([p0, p1]))
+    assert [c[0] for c in calls] == ["all_gather", "all_gather"]
+
+
+def test_one_rank_is_the_gather(monkeypatch):
+    monkeypatch.setattr(M, "RS", True)
+    c = M.Comm.__new__(M.Comm)
+    c.nccl, c.world, c._buf = None, 1, {}
+    p0, _ = parts()
+    assert torch.equal(c.reduced(p0), p0[None])