From 2d49e2512ac3e23bcfcb48591952bb10a7a71965 Mon Sep 17 00:00:00 2001 From: Daigo Kimura <84138151+imD-5@users.noreply.github.com> Date: Sat, 24 Oct 2026 00:11:36 +0900 Subject: [PATCH] perf(dsv41): prompt-chunk partials over a reduce-scatter, the same bits from 3/4 of the bytes TF_DS_RS=1 (two ranks, off by default): 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, instead of the fp32 all-gather. The post kernels add the two ranks' rows and round to bf16; one fp32 add of two operands has the same bits in either order, so they form the same y. Odd row counts, rows the RDMA gather holds, and one rank keep the gather. --- src/tensorfold/cuda/comm.py | 11 ++ .../families/deepseek_v41/cuda/model.py | 22 +++- tests/test_dsv41_reduced.py | 101 ++++++++++++++++++ 3 files changed, 132 insertions(+), 2 deletions(-) create mode 100644 tests/test_dsv41_reduced.py diff --git a/src/tensorfold/cuda/comm.py b/src/tensorfold/cuda/comm.py index 9e9b16a1d6..bc322252d1 100644 --- a/src/tensorfold/cuda/comm.py +++ b/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.""" diff --git a/src/tensorfold/families/deepseek_v41/cuda/model.py b/src/tensorfold/families/deepseek_v41/cuda/model.py index c33ee49303..df2cc88fe3 100644 --- a/src/tensorfold/families/deepseek_v41/cuda/model.py +++ b/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 diff --git a/tests/test_dsv41_reduced.py b/tests/test_dsv41_reduced.py new file mode 100644 index 0000000000..ce712d8907 --- /dev/null +++ b/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])