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

[ChenYuHo]

Summary:

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

[meta-codesync]

@ChenYuHo has exported this pull request. If you are a Meta employee, you can view the originating Diff in D107172780.