Composable send/recv interface (transport + hooks + minimal NVLink path) (meta-pytorch#2829)

Summary:
Pull Request resolved: meta-pytorch#2829

First, interface-only diff for the composable device send/recv framework that
fused compute/comm kernels (all-to-all, allgather, ...) will build on. No
production pipeline is wired; review overhead is intentionally low.

New package `comms/dsl/`:

* `ctx.py` — `Ctx`, the DSL-agnostic field-set spec of the hook contract,
  realized on-device per DSL (Triton `_aggregate` in `triton/ctx.py`).
* `ops.py` — the device transport-ops seam (`put` / `get` /
  `signal` / `wait`): the only transport-specific device primitives.
* `transport.py` — user-owned p2p transport objects:
  - `P2pTransport` Protocol + `PeerEndpoint` (host-resolved per-peer device state),
  - real, minimal `NvlTransport` + `nvl_rendezvous` (one collective rendezvous;
    no hidden cache — the user holds the object),
  - reserved `IbTransport` / `ib_rendezvous`, and `MeshTransport` which routes
    per peer via `link_kind` so a single collective can mix NVLink (intra-domain)
    and IB (inter-domain) later,
  - `check_transfer()` — fail-loud guard that numel fits the per-peer staging
    region and num_blocks fits the signal-pad slots (a violation would silently
    overrun the next peer's region on the remote rank).
* `triton/` — minimal, real (no-pipeline, single-shot) device kernels `send_tiles`/`recv_tiles` written
  against the ops seam; hooks take a single `Ctx` aggregate
  (`produce(ctx) -> regs`, `consume(ctx, regs)`) — `triton/ctx.py` realizes `Ctx`
  as a Triton `_aggregate`. `nvl_ops` (real, over the framework's own self-contained PTX in
  `triton/device_utils.py`, no `comms/pipes` dep); `ib_ops` (reserved); `copy_*` default hooks; `send`/`recv`/
  `sendrecv` launchers + a commented mixed-transport sketch.
* `cute/` — `send`/`recv` interface stubs reserved for the CuTe backend (made
  real in the next diff).

Design notes:

* The hook seam is **full-leg** (one `produce`/`consume` per direction) — minimal
  yet sufficient: it subsumes address-gather, value-transform, packed staging,
  accumulate, and compute-into-send. The `addr_fn`+`transform` form will be a
  composer on top (follow-up).
* Hooks take a single opaque `Ctx`, so the contract is churn-proof: future needs
  (pipeline `slot`/`step`, MoE `expert_counts`, FP8 `scales`, warp-spec role, ...)
  become `Ctx` fields and never change a hook signature. `Ctx` is also
  schedule-agnostic, so the same hooks serve a user-written schedule today and a
  library schedule skeleton later.
* Transport choice binds **per call-site** (ops as `constexpr`), which is what
  lets one kernel mix NVLink and IB once the IB ops/transport are filled in — no
  change to `send`/`recv`, hooks, or `Ctx`.
* **CUDA-graph-ready (architecture).** The collective rendezvous + symm-mem
  allocation run at setup (outside capture); the captured region is pure kernel
  launches over persistent, user-held buffers. Verified on 2x H100 for the
  **Triton** path: send/recv captures and a single replay is correct (after a
  host-side signal-pad reset; kernels must be warmed up before capture). NOTE:
  repeated replay is not yet gated — the minimal single-shot signal (`seq=1`) is
  not re-armed — so guaranteed repeat-replay support lands with the
  monotonic-counter pipelined follow-up. (CuTe graph capture is a separate
  follow-up — see D107283879.)

Follow-up stacks: (1) Triton — real pipelined `send`/`recv` + port
`all_to_all_single`; (2) CuTe — real CuTe `send`/`recv` on the same contract +
cross-DSL interop test; (3) IB + mixed transport.

Differential Revision: D107172780

Author: Elton Ho

comms/dsl/__init__.py

@@ -0,0 +1,46 @@
+# (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary.
+
+# pyre-strict
+
+"""Composable send/recv framework — DSL-agnostic contract.
+
+This package exposes the stable, backend-independent contract:
+
+* the hook contract :class:`Ctx` (+ ``ProduceFn`` / ``ConsumeFn`` aliases),
+* the user-owned p2p transport objects (:class:`NvlTransport` now,
+  :class:`IbTransport` / :class:`MeshTransport` reserved) and
+  :func:`nvl_rendezvous`.
+
+Device ``send``/``recv`` live in the per-DSL subpackages
+(``comms.dsl.triton``, ``comms.dsl.cute``).
+"""
+
+from __future__ import annotations
+
+from .ctx import ConsumeFn, Ctx, ProduceFn
+from .transport import (
+    check_transfer,
+    ib_rendezvous,
+    IbTransport,
+    LinkKind,
+    MeshTransport,
+    nvl_rendezvous,
+    NvlTransport,
+    P2pTransport,
+    PeerEndpoint,
+)
+
+__all__ = [
+    "Ctx",
+    "ProduceFn",
+    "ConsumeFn",
+    "LinkKind",
+    "P2pTransport",
+    "PeerEndpoint",
+    "NvlTransport",
+    "nvl_rendezvous",
+    "check_transfer",
+    "IbTransport",
+    "ib_rendezvous",
+    "MeshTransport",
+]

comms/dsl/ctx.py

@@ -0,0 +1,62 @@
+# (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary.
+
+# pyre-strict
+
+"""DSL-agnostic hook contract for the composable send/recv framework.
+
+``Ctx`` describes the *fields* a per-tile compute hook may read. It is the stable
+contract shared by every backend, realized on-device per DSL: the Triton backend
+as a ``@_aggregate`` struct (``triton/ctx.py``) and the CuTe backend as a
+``@dataclass`` (``cute/send_recv.py``). The field set stays identical so the same
+hook is portable across backends and across schedules (a user-written schedule
+today, a library schedule skeleton later).
+
+This dataclass is the DSL-agnostic *spec* of that field set (and is used by the
+GPU-free tests). The device hooks read the per-backend realization; adding a
+field here (and to each realization) never changes a hook signature.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from dataclasses import dataclass, field
+
+
+@dataclass
+class Ctx:
+    """The full-leg hook contract (field set a ``produce``/``consume`` may read)."""
+
+    # --- pointers ---
+    in_ptr: int = 0
+    out_ptr: int = 0
+    # Base pointer for the current tile's transport-addressable data region.
+    # On send paths derived from PeerEndpoint.send_dst; on recv from .recv_src.
+    region_ptr: int = 0
+    slot_base: int = 0
+    # --- pipeline position ---
+    tile_idx: int = 0
+    slot: int = 0
+    step: int = 0
+    tile_rows: int = 0
+    tile_cols: int = 0
+    # --- intra-block work split ---
+    block_id: int = 0
+    flat_tid: int = 0
+    num_blocks: int = 0
+    # --- schedule state (lets one hook serve multiple schedules) ---
+    peer: int = 0
+    recv_peer: int = 0
+    shard_idx: int = 0
+    # --- layout ---
+    in_strides: tuple[int, ...] = ()
+    out_strides: tuple[int, ...] = ()
+    # --- op-specific extension point (e.g. expert_counts, scales) ---
+    extra: dict[str, object] = field(default_factory=dict)
+
+
+# The two hooks. ``produce(ctx) -> regs``: HBM -> staging (reads the input leg).
+# ``consume(ctx, regs)``: staging -> HBM (writes the output leg with the loaded
+# tile payload ``regs``). These aliases document the contract; concrete device
+# hooks are DSL kernels reading the per-backend ``Ctx`` realization.
+ProduceFn = Callable[[Ctx], object]
+ConsumeFn = Callable[[Ctx, object], None]

comms/dsl/cute/__init__.py

@@ -0,0 +1,9 @@
+# (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary.
+
+# pyre-strict
+
+"""CuTe backend: send/recv interface stubs (reserved for the CuTe stack)."""
+
+from .send_recv import recv, send
+
+__all__ = ["send", "recv"]

comms/dsl/cute/send_recv.py

@@ -0,0 +1,25 @@
+# (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary.
+
+# pyre-strict
+
+"""CuTe send/recv — minimal interface stubs.
+
+Reserves the CuTe backend's device API so the contract is visible now; the real
+CuTe kernels (mirroring the Triton ``send``/``recv`` against the same transport +
+ops seam) land in the CuTe stack. Kept as plain stubs to avoid pulling the CuTe
+DSL dependency into this interface diff.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+_RESERVED = "framework CuTe send/recv is reserved; implemented in the CuTe stack"
+
+
+def send(transport: Any, send_buf: Any, peer: int, **kwargs: Any) -> None:
+    raise NotImplementedError(_RESERVED)
+
+
+def recv(transport: Any, recv_buf: Any, peer: int, **kwargs: Any) -> None:
+    raise NotImplementedError(_RESERVED)

comms/dsl/ops.py

@@ -0,0 +1,45 @@
+# (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary.
+
+# pyre-strict
+
+"""Device transport-ops seam.
+
+These four device primitives are the *only* transport-specific operations.
+``send``/``recv`` are written once against this seam and treat the per-peer
+buffers as opaque, so a single ``send``/``recv``:
+
+  * serves NVLink today and IB later, and
+  * serves **mixed-transport kernels** — a schedule may use NVLink for an
+    intra-domain peer and IB for an inter-domain peer by passing different ops
+    at each call site (the binding is per-call-site, not per-kernel).
+
+There is no runtime dispatch: in Triton the concrete ops are passed as
+``tl.constexpr`` and selected at compile time. Concrete implementations live per
+transport + DSL (e.g. ``framework/triton/nvl_ops.py``). This module documents
+the contract.
+
+Conceptual signatures (the concrete versions are ``@triton.jit`` / cute kernels):
+
+    put(region, idx, regs, mask) -> None   # produced tile -> peer region
+    get(region, idx, mask)        -> regs   # received tile  <- region
+    signal(sig_ptr, sig_idx, seq)          -> None   # publish "data ready"
+    wait(sig_ptr, sig_idx, seq)            -> None   # wait for peer's data
+"""
+
+from __future__ import annotations
+
+from typing import Protocol
+
+
+class TransportOps(Protocol):
+    """Structural description of the four device ops a transport must provide.
+
+    Used for documentation/typing only; the attributes are DSL kernels
+    (``@triton.jit`` functions for the Triton backend), not plain callables, so
+    the parameter types are intentionally left loose.
+    """
+
+    put: object
+    get: object
+    signal: object
+    wait: object

comms/dsl/tests/__init__.py

@@ -0,0 +1 @@
+# (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary.

comms/dsl/tests/test_interface.py

@@ -0,0 +1,190 @@
+# (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary.
+
+# pyre-unsafe
+
+"""GPU-free interface tests for the composable send/recv framework.
+
+Asserts the *contract* (types, signatures, transport abstraction, reserved
+stubs) without compiling or launching any kernel, so this runs in CI with no
+GPU. The real end-to-end path is exercised by ``test_minimal_sendrecv``.
+"""
+
+import dataclasses
+import inspect
+import unittest
+from typing import Any, cast
+
+from torch.utils._triton import has_triton
+
+
+class ContractTest(unittest.TestCase):
+    def test_ctx_fields(self) -> None:
+        from comms.dsl import Ctx
+
+        names = {f.name for f in dataclasses.fields(Ctx)}
+        for expected in (
+            "in_ptr",
+            "out_ptr",
+            "region_ptr",
+            "tile_idx",
+            "block_id",
+            "flat_tid",
+            "num_blocks",
+            "peer",
+            "recv_peer",
+            "shard_idx",
+            "in_strides",
+            "out_strides",
+            "extra",
+        ):
+            self.assertIn(expected, names)
+        # Construct + extension point works.
+        ctx = Ctx(in_ptr=1, extra={"scales": 7})
+        self.assertEqual(ctx.in_ptr, 1)
+        self.assertEqual(ctx.extra["scales"], 7)
+
+    def test_hook_aliases_importable(self) -> None:
+        from comms.dsl import ConsumeFn, ProduceFn
+
+        self.assertIsNotNone(ProduceFn)
+        self.assertIsNotNone(ConsumeFn)
+
+    def test_exports(self) -> None:
+        import comms.dsl as fw
+
+        for name in (
+            "Ctx",
+            "NvlTransport",
+            "nvl_rendezvous",
+            "MeshTransport",
+            "LinkKind",
+        ):
+            self.assertIn(name, fw.__all__)
+
+
+class TransportTest(unittest.TestCase):
+    def test_rendezvous_signature(self) -> None:
+        from comms.dsl import nvl_rendezvous
+
+        params = list(inspect.signature(nvl_rendezvous).parameters)
+        for expected in ("group", "device", "per_peer_bytes"):
+            self.assertIn(expected, params)
+
+    def test_nvl_link_kind(self) -> None:
+        from comms.dsl import LinkKind, NvlTransport
+
+        # handle is not touched by link_kind, so we can build this GPU-free.
+        t = NvlTransport(
+            handle=cast(Any, None), world_size=4, local_rank=0, per_peer_bytes=1024
+        )
+        self.assertIs(t.link_kind(1), LinkKind.NVLINK)
+
+    def test_check_transfer_guards(self) -> None:
+        import torch
+        from comms.dsl import check_transfer, NvlTransport
+
+        # per_peer_bytes=4096, fp32 -> 1024 elems per peer; 4 signal slots.
+        t = NvlTransport(
+            handle=cast(Any, None),
+            world_size=2,
+            local_rank=0,
+            per_peer_bytes=4096,
+            max_blocks_per_peer=4,
+        )
+        # Valid transfer: no raise.
+        check_transfer(t, numel=512, dtype=torch.float32, num_blocks=2)
+        # numel exceeds per-peer capacity -> raise (would corrupt the next peer).
+        with self.assertRaises(ValueError):
+            check_transfer(t, numel=2000, dtype=torch.float32, num_blocks=2)
+        # num_blocks exceeds signal slots -> raise (would OOB-write the pad).
+        with self.assertRaises(ValueError):
+            check_transfer(t, numel=512, dtype=torch.float32, num_blocks=8)
+        # num_blocks must be >= 1.
+        with self.assertRaises(ValueError):
+            check_transfer(t, numel=512, dtype=torch.float32, num_blocks=0)
+        # MeshTransport delegates max_blocks_per_peer to intra, so the
+        # num_blocks guard still fires through a mesh.
+        from comms.dsl import MeshTransport
+
+        mesh = MeshTransport(intra=t)
+        check_transfer(mesh, numel=512, dtype=torch.float32, num_blocks=2)
+        with self.assertRaises(ValueError):
+            check_transfer(mesh, numel=512, dtype=torch.float32, num_blocks=8)
+
+    def test_ib_reserved(self) -> None:
+        from comms.dsl import ib_rendezvous, IbTransport
+
+        ib = IbTransport(world_size=8, per_peer_bytes=1024)
+        with self.assertRaises(NotImplementedError):
+            ib.endpoint(1, dtype=cast(Any, None))
+        with self.assertRaises(NotImplementedError):
+            ib_rendezvous(cast(Any, None), cast(Any, None), 1024)
+
+    def test_mesh_routes_by_domain(self) -> None:
+        from comms.dsl import IbTransport, LinkKind, MeshTransport, NvlTransport
+
+        intra = NvlTransport(
+            handle=cast(Any, None), world_size=8, local_rank=0, per_peer_bytes=1024
+        )
+        # No inter transport yet -> everything is NVLINK.
+        mesh = MeshTransport(intra=intra)
+        self.assertIs(mesh.link_kind(3), LinkKind.NVLINK)
+
+        # With an inter transport + domain size 4: peers 0-3 NVLINK, 4-7 IB.
+        mesh2 = MeshTransport(
+            intra=intra,
+            inter=IbTransport(world_size=8, per_peer_bytes=1024),
+            local_domain_size=4,
+        )
+        self.assertIs(mesh2.link_kind(2), LinkKind.NVLINK)
+        self.assertIs(mesh2.link_kind(5), LinkKind.IB)
+
+
+@unittest.skipUnless(has_triton(), "Triton not available")
+class TritonInterfaceTest(unittest.TestCase):
+    def _arg_names(self, jit_fn: object) -> list[str]:
+        names = getattr(jit_fn, "arg_names", None)
+        if names is not None:
+            return list(names)
+        return list(inspect.signature(getattr(jit_fn, "fn", jit_fn)).parameters)
+
+    def test_send_recv_are_jit(self) -> None:
+        from comms.dsl.triton import recv_tiles as recv, send_tiles as send
+        from triton.runtime.jit import JITFunction
+
+        self.assertIsInstance(send, JITFunction)
+        self.assertIsInstance(recv, JITFunction)
+        for p in ("in_ptr", "produce", "put", "signal"):
+            self.assertIn(p, self._arg_names(send))
+        for p in ("out_ptr", "consume", "get", "wait"):
+            self.assertIn(p, self._arg_names(recv))
+
+    def test_nvl_and_ib_ops_present(self) -> None:
+        from comms.dsl.triton import ib_ops, nvl_ops
+        from triton.runtime.jit import JITFunction
+
+        for ops in (nvl_ops, ib_ops):
+            for name in ("put", "get", "signal", "wait"):
+                self.assertIsInstance(getattr(ops, name), JITFunction)
+
+    def test_hook_ctx_contract(self) -> None:
+        # Hooks take a single Ctx aggregate; consume also gets the tile payload.
+        from comms.dsl.triton import hooks
+        from comms.dsl.triton.ctx import Ctx
+        from triton.runtime.jit import JITFunction
+
+        self.assertIsInstance(hooks.copy_produce, JITFunction)
+        self.assertIsInstance(hooks.copy_consume, JITFunction)
+        self.assertEqual(self._arg_names(hooks.copy_produce), ["ctx"])
+        self.assertEqual(self._arg_names(hooks.copy_consume), ["ctx", "regs"])
+        self.assertIsNotNone(Ctx)
+
+
+class CuteInterfaceTest(unittest.TestCase):
+    def test_cute_stubs_reserved(self) -> None:
+        from comms.dsl.cute import recv, send
+
+        with self.assertRaises(NotImplementedError):
+            send(None, None, 1)
+        with self.assertRaises(NotImplementedError):
+            recv(None, None, 1)