diff --git a/RUNBOOK.md b/RUNBOOK.md index 67d8ecdc2a..b604f2b308 100644 --- a/RUNBOOK.md +++ b/RUNBOOK.md @@ -168,6 +168,12 @@ The server logs `heat_wait_s`, a cancel ends the wait, and the reply's tokens do On three GB10s serving GLM-5.3-Flash NVFP4 at TP=3, a 1,035,295-token prompt without the bands kept one host's hottest zone at 94 C or more for most of the prompt (97.7 C at the top); `TF_HEAT_HIGH=92` and `TF_HEAT_LOW=88` held all three at 93.0 C or under, and the first token came 10% later. +`TF_GRAPH_CACHE=0` (or `off`, any case) runs every step eagerly in any family that installs the keyed CUDA graph cache, +without a rebuild. Set it alike on every tensor-parallel rank: an off rank never joins the cache's per-miss vote, so a +rank left on would wait at its first miss, and only the off rank logs the switch. Ranks that share a cache compare +`Settings.digest` of their effective settings at start, before building it, and refuse to start when they differ. +A tensor-parallel family passes its cache settings in the TP config; the bootstrap then compares them and refuses to +start on every rank, naming TF_GRAPH_CACHE among the settings that must match. Archives whose names contain `host-only` are CPU verification artifacts and cannot serve CUDA inference. ## Flash Next and paired Metal serving diff --git a/THIRD_PARTY_NOTICES.md b/THIRD_PARTY_NOTICES.md index db27613ce9..3706d0d4fe 100644 --- a/THIRD_PARTY_NOTICES.md +++ b/THIRD_PARTY_NOTICES.md @@ -183,6 +183,12 @@ the one-shot all-gather protocol of [b12x](https://github.com/local-inference-la sequence parity and a timeout that fails the transport. Copyright 2026 Luke Alonso and the b12x contributors, Apache License 2.0 (see [the license text](LICENSES/Apache-2.0.txt)); b12x ships no NOTICE file. +The Zig tensor-parallel fabric's RoCE path (`zig/src/fabric/tp/mailbox.cu`, `roce.zig`, `roce_verbs.zig`) adapts +the one-shot all-gather protocol of [b12x](https://github.com/local-inference-lab/b12x)'s RoCEnante +(`b12x/comm/roce/`, PRs #295 and #315): a device-resident epoch, sequence flags after each payload, two slots by +sequence parity and a timeout that fails the transport. Copyright 2026 Luke Alonso and the b12x contributors, +Apache License 2.0 (see [the license text](LICENSES/Apache-2.0.txt)); b12x ships no NOTICE file. + ## Vendored code and weights In the Python engine, frozen at 0.6.6, `src/tensorfold/drafters/vendor/z_lab_dflash/model_mlx.py` is the unmodified diff --git a/build.zig b/build.zig index 8361b06746..44a1c9dabe 100644 --- a/build.zig +++ b/build.zig @@ -19,6 +19,7 @@ const golden_build = @import("zig/build/golden.zig"); const contract_build = @import("zig/build/kernel_contract.zig"); const video_build = @import("zig/build/video.zig"); const sys_build = @import("zig/build/sys.zig"); +const tp_build = @import("zig/build/tp.zig"); comptime { const required = std.mem.trim(u8, @embedFile(".zig-version"), "\r\n"); @@ -202,6 +203,7 @@ pub fn build(b: *std.Build) void { b.step("test-learn-warning", "CPU production-log refusal, backoff and recovery test").dependOn(&warning_run.step); grammar_build.steps(b, target, xgrammar); contract_build.hostTests(b, gemma4_contract_host); + tp_build.add(b, target, optimize, test_step); sys_build.wire(b); @import("zig/build/checkpoint_config.zig").wire(b); } diff --git a/zig/build/tp.zig b/zig/build/tp.zig new file mode 100644 index 0000000000..e07230abbc --- /dev/null +++ b/zig/build/tp.zig @@ -0,0 +1,88 @@ +//! Tensor parallelism's build: Linux host tests, and on Linux the `tf-tp-test` GPU runner with the mailbox kernel. + +const std = @import("std"); + +/// The CUDA runtime without the engine's kernels (tp only needs the driver, memory, streams and NCCL). +fn runtime(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode) *std.Build.Module { + const options = b.addOptions(); + options.addOption(bool, "with_kernels", false); + const cuda = b.createModule(.{ .root_source_file = b.path("zig/src/cuda/root.zig"), .target = target, .optimize = optimize, .link_libc = true }); + cuda.addOptions("kernel_options", options); + cuda.addImport("stagger", b.createModule(.{ .root_source_file = b.path("zig/src/core/stagger.zig"), .target = target, .optimize = optimize })); + return cuda; +} + +/// The `tp` module over `cuda` (an engine passes its own runtime module); verbs (zig/src/fabric/verbs.zig) for RoCE. +pub fn module(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, cuda: *std.Build.Module) *std.Build.Module { + const m = b.createModule(.{ .root_source_file = b.path("zig/src/fabric/tp/tp.zig"), .target = target, .optimize = optimize, .link_libc = true }); + m.addImport("cuda", cuda); + m.addImport("verbs", b.createModule(.{ .root_source_file = b.path("zig/src/fabric/verbs.zig"), .target = target, .optimize = optimize, .link_libc = true })); + return m; +} + +/// mailbox.cu as a fatbin: sm_120 (RTX PRO 6000), sm_121 (GB10), and compute_90 PTX for anything else. +fn mailboxFatbin(b: *std.Build, nvcc: []const u8) std.Build.LazyPath { + const run = b.addSystemCommand(&.{ nvcc, "-fatbin", "-O3", "-std=c++17", "-gencode=arch=compute_120,code=sm_120", "-gencode=arch=compute_121,code=sm_121", "-gencode=arch=compute_90,code=compute_90", "-o" }); + const out = run.addOutputFileArg("tp_mailbox.fatbin"); + run.addFileArg(b.path("zig/src/fabric/tp/mailbox.cu")); + return out; +} + +/// `pub const mailbox: []const u8`, the fatbin's bytes or empty. +fn kernelModule(b: *std.Build, image: ?std.Build.LazyPath) *std.Build.Module { + const wf = b.addWriteFiles(); + if (image) |file| { + _ = wf.addCopyFile(file, "tp_mailbox.fatbin"); + // the driver takes 8-byte aligned images (Module.load refuses others); @embedFile alone promises nothing + return b.createModule(.{ .root_source_file = wf.add("tp_kernel.zig", "const bytes align(16) = @embedFile(\"tp_mailbox.fatbin\").*;\npub const mailbox: []const u8 = &bytes;\n") }); + } + return b.createModule(.{ .root_source_file = wf.add("tp_kernel.zig", "pub const mailbox: []const u8 = \"\";\n") }); +} + +pub fn add(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, test_step: *std.Build.Step) void { + const host = b.graph.host; + const cli_tests = b.addRunArtifact(b.addTest(.{ .root_module = b.createModule(.{ .root_source_file = b.path("zig/src/fabric/tp/gpu_cli.zig"), .target = host, .optimize = optimize }) })); + const cli_step = b.step("test-tp-cli", "TP runner command validation (no GPU)"); + cli_step.dependOn(&cli_tests.step); + test_step.dependOn(&cli_tests.step); + if (host.result.os.tag == .linux) { + const tests = b.addRunArtifact(b.addTest(.{ .root_module = module(b, host, .debug, runtime(b, host, .debug)) })); + const tp_step = b.step("test-tp", "Tensor parallelism's host tests: bootstrap, fate channel, in-process ranks (no GPU)"); + tp_step.dependOn(&tests.step); + tp_step.dependOn(cli_step); + test_step.dependOn(&tests.step); + } + if (target.result.os.tag != .linux) return; + const nvcc = b.option([]const u8, "tp-nvcc", "nvcc that builds the TP mailbox kernel"); + const prebuilt = b.option([]const u8, "tp-fatbin", "absolute path of a prebuilt tp_mailbox.fatbin to embed"); + const image: ?std.Build.LazyPath = if (prebuilt) |p| b.graph.cwdRelativePath(p) else if (nvcc) |tool| mailboxFatbin(b, tool) else null; + const cuda = runtime(b, target, optimize); + const runner = b.createModule(.{ .root_source_file = b.path("zig/src/fabric/tp/gpu_main.zig"), .target = target, .optimize = optimize, .link_libc = true }); + runner.addImport("cuda", cuda); + runner.addImport("tp", module(b, target, optimize, cuda)); + runner.addImport("tp_kernel", kernelModule(b, image)); + const exe = b.addExecutable(.{ .name = "tf-tp-test", .root_module = runner }); + if (host.result.os.tag == .linux and target.query.isNative()) { + const cli = @import("../src/fabric/tp/gpu_cli.zig"); + for (cli.invalid) |args| { + const check = b.addRunArtifact(exe); + check.addArgs(args); + check.setEnvironmentVariable("TF_TP_WORLD", "invalid"); + check.expectExitCode(2); + check.expectStdErrEqual(cli.usage); + cli_step.dependOn(&check.step); + test_step.dependOn(&check.step); + } + if (image == null) { + const kernel = b.addRunArtifact(exe); + kernel.addArg("kernel"); + kernel.setEnvironmentVariable("TF_TP_WORLD", "invalid"); + kernel.expectExitCode(1); + kernel.expectStdErrEqual("mailbox kernel 0 B\n"); + cli_step.dependOn(&kernel.step); + test_step.dependOn(&kernel.step); + } + } + b.installArtifact(exe); + b.step("tf-tp-test", "TP=2 on two GPUs: bootstrap, every collective bit-exact, graph capture, fail-fast, latency").dependOn(&b.addInstallArtifact(exe, .{}).step); +} diff --git a/zig/src/cuda/abi.zig b/zig/src/cuda/abi.zig index b784479964..fe21a66fc4 100644 --- a/zig/src/cuda/abi.zig +++ b/zig/src/cuda/abi.zig @@ -39,6 +39,8 @@ pub const DeviceAttribute = enum(c_int) { compute_capability_minor = 76, max_shared_memory_per_multiprocessor = 81, max_shared_memory_per_block_optin = 97, + virtual_memory_management_supported = 102, + gpu_direct_rdma_with_cuda_vmm_supported = 110, }; pub const FunctionAttribute = enum(c_int) { @@ -242,3 +244,10 @@ comptime { std.debug.assert(@sizeOf(KernelNodeParams) == 72 and @offsetOf(KernelNodeParams, "params") == 40); std.debug.assert(@sizeOf(ExecUpdateResultInfo) == 24); } + +test "CUDA device attribute matches the RDMA VMM capability ABI" { + // CUDA's CUdevice_attribute assigns 110 to RDMA VMM support; 126 is the memory sync domain count. + try @import("std").testing.expectEqual(@as(c_int, 110), @backingInt( + DeviceAttribute.gpu_direct_rdma_with_cuda_vmm_supported, + )); +} diff --git a/zig/src/cuda/arena.zig b/zig/src/cuda/arena.zig new file mode 100644 index 0000000000..de3da5718c --- /dev/null +++ b/zig/src/cuda/arena.zig @@ -0,0 +1,405 @@ +//! A device sub-allocator under DeviceBuffer.alloc: buffers carved from a few large chunks, not one cuMemAlloc each. + +const std = @import("std"); +const abi = @import("abi.zig"); +const Driver = @import("driver.zig").Driver; +const memory = @import("memory.zig"); + +pub const Error = error{ DriverUnavailable, MissingSymbol, Unsupported, CudaFailed, OutOfMemory }; + +/// Where chunks come from: `chunk` maps `len` bytes (a multiple of `granularity`), `release` gives a chunk back. +pub const Backend = struct { + ctx: *anyopaque, + granularity: usize, + chunk: *const fn (ctx: *anyopaque, len: usize) ?u64, + release: *const fn (ctx: *anyopaque, base: u64, len: usize) void, + name: []const u8, +}; + +pub const Options = struct { + small_max: usize = 4 << 20, // requests up to this many bytes share chunks + small_chunk: usize = 64 << 20, // a shared chunk's bytes, rounded up to the granularity + alignment: usize = 512, // every range starts at a multiple of this +}; + +pub const Stats = struct { small: u64 = 0, large: u64 = 0, reused: u64 = 0, fallback: u64 = 0, chunk_bytes: u64 = 0 }; + +/// Small requests bump 512-byte aligned through shared chunks and are reused by size; large ones get their own chunk. +pub const Arena = struct { + gpa: std.mem.Allocator, + be: Backend, + o: Options, + mutex: std.atomic.Mutex = .unlocked, + small: std.ArrayList([2]u64) = .empty, // shared chunks (base, len) + cursor: u64 = 0, // the current shared chunk's next free byte + limit: u64 = 0, + free_small: std.AutoHashMapUnmanaged(u64, std.ArrayList(u64)) = .empty, // freed small ranges by rounded size + live: std.AutoHashMapUnmanaged(u64, u64) = .empty, // live small ranges: address to rounded size + large: std.AutoHashMapUnmanaged(u64, u64) = .empty, // large chunks: address to mapped size + dead: [2]u32 = .{ 0, 0 }, // removals from `live` and `large` since each was last rehashed + stats: Stats = .{}, + + pub fn init(gpa: std.mem.Allocator, be: Backend, o: Options) Arena { + return .{ .gpa = gpa, .be = be, .o = o }; + } + + pub fn deinit(a: *Arena) void { + var it = a.large.iterator(); + while (it.next()) |e| a.be.release(a.be.ctx, e.key_ptr.*, e.value_ptr.*); + for (a.small.items) |c| a.be.release(a.be.ctx, c[0], c[1]); + var f = a.free_small.valueIterator(); + while (f.next()) |l| l.deinit(a.gpa); + a.free_small.deinit(a.gpa); + a.live.deinit(a.gpa); + a.large.deinit(a.gpa); + a.small.deinit(a.gpa); + a.* = undefined; + } + + fn lock(a: *Arena) void { + while (!a.mutex.tryLock()) std.Thread.yield() catch {}; + } + + fn roundUp(x: usize, m: usize) usize { + return std.mem.alignForward(usize, x, m); + } + + /// An address for `len` bytes, or null (the caller then allocates as before). + pub fn alloc(a: *Arena, len: usize) ?u64 { + if (len == 0) return null; + a.lock(); + defer a.mutex.unlock(); + return a.allocLocked(len) catch null; + } + + fn allocLocked(a: *Arena, len: usize) !?u64 { + if (len > a.o.small_max) { + const n = roundUp(len, a.be.granularity); + const p = a.be.chunk(a.be.ctx, n) orelse { + a.stats.fallback += 1; + return null; + }; + errdefer a.be.release(a.be.ctx, p, n); + try a.large.put(a.gpa, p, n); + a.stats.large += 1; + a.stats.chunk_bytes += n; + return p; + } + const n: u64 = roundUp(len, a.o.alignment); + try a.live.ensureUnusedCapacity(a.gpa, 1); // before a range is taken: a failure leaves free list and cursor alone + if (a.free_small.getPtr(n)) |l| if (l.pop()) |p| { + a.live.putAssumeCapacityNoClobber(p, n); + a.stats.reused += 1; + return p; + }; + if (a.limit - a.cursor < n) { + const cl = roundUp(@max(a.o.small_chunk, n), a.be.granularity); + const base = a.be.chunk(a.be.ctx, cl) orelse { + a.stats.fallback += 1; + return null; + }; + errdefer a.be.release(a.be.ctx, base, cl); + try a.small.append(a.gpa, .{ base, cl }); + a.cursor = base; + a.limit = base + cl; + a.stats.chunk_bytes += cl; + } + const p = a.cursor; + a.live.putAssumeCapacityNoClobber(p, n); + a.cursor += n; + a.stats.small += 1; + return p; + } + + /// True when `p` was the arena's (and is now free); false: not ours, the caller frees it as before. + pub fn free(a: *Arena, p: u64) bool { + a.lock(); + defer a.mutex.unlock(); + if (a.large.fetchRemove(p)) |e| { + removed(&a.large, &a.dead[1]); + a.be.release(a.be.ctx, p, e.value); + return true; + } + const e = a.live.fetchRemove(p) orelse return false; + removed(&a.live, &a.dead[0]); + const g = a.free_small.getOrPut(a.gpa, e.value) catch return true; // leaked range, still ours + if (!g.found_existing) g.value_ptr.* = .empty; + g.value_ptr.append(a.gpa, p) catch {}; + return true; + } + + /// std's tombstones never clear and keys fill 80 %: rehash at a tenth, so a tenth of the map always stays free. + fn removed(map: *std.AutoHashMapUnmanaged(u64, u64), dead: *u32) void { + dead.* += 1; + if (@as(u64, dead.*) * 10 < map.capacity()) return; + map.rehash(std.hash_map.AutoContext(u64){}); + dead.* = 0; + } + + // -- memory.zig's hook ---------------------------------------------------------------------------------------- + + fn hookAlloc(ctx: *anyopaque, len: usize) ?u64 { + const a: *Arena = @ptrCast(@alignCast(ctx)); + return a.alloc(len); + } + + fn hookFree(ctx: *anyopaque, p: u64) bool { + const a: *Arena = @ptrCast(@alignCast(ctx)); + return a.free(p); + } + + /// Every later DeviceBuffer.alloc comes from the arena, which must outlive its buffers; only addresses change. + pub fn install(a: *Arena) void { + memory.setHook(.{ .ctx = a, .alloc = hookAlloc, .free = hookFree }); + } +}; + +// -- backends --------------------------------------------------------------------------------------------------------- + +/// One cuMemAlloc a chunk (the driver's own 2 MiB-aligned allocations). +pub const Plain = struct { + d: *const Driver, + + pub fn backend(p: *Plain) Backend { + return .{ .ctx = p, .granularity = 2 << 20, .chunk = chunk, .release = release, .name = "plain" }; + } + + fn chunk(ctx: *anyopaque, len: usize) ?u64 { + const p: *Plain = @ptrCast(@alignCast(ctx)); + var q: abi.DevicePtr = 0; + if (p.d.api.cuMemAlloc_v2(&q, len) != abi.success) return null; + return q; + } + + fn release(ctx: *anyopaque, base: u64, _: usize) void { + const p: *Plain = @ptrCast(@alignCast(ctx)); + _ = p.d.api.cuMemFree_v2(base); + } +}; + +/// CUDA's virtual memory management: one reserved virtual range, chunks of the recommended granularity mapped in order. +pub const Vmm = struct { + lib: std.DynLib, + device: abi.Device, + prop: Prop, + granularity: usize, + va: u64 = 0, + va_len: usize = 0, + next: u64 = 0, + api: Api, + + const Api = struct { + reserve: *const fn (*u64, usize, usize, u64, c_ulonglong) callconv(.c) abi.Result, + addressFree: *const fn (u64, usize) callconv(.c) abi.Result, + create: *const fn (*u64, usize, *const Prop, c_ulonglong) callconv(.c) abi.Result, + release: *const fn (u64) callconv(.c) abi.Result, + map: *const fn (u64, usize, usize, u64, c_ulonglong) callconv(.c) abi.Result, + unmap: *const fn (u64, usize) callconv(.c) abi.Result, + setAccess: *const fn (u64, usize, *const AccessDesc, usize) callconv(.c) abi.Result, + granularity: *const fn (*usize, *const Prop, c_int) callconv(.c) abi.Result, + }; + + pub const Prop = extern struct { // CUmemAllocationProp + type: c_int = 1, // CU_MEM_ALLOCATION_TYPE_PINNED + handle_types: c_int = 0, + loc_type: c_int = 1, // CU_MEM_LOCATION_TYPE_DEVICE + loc_id: c_int = 0, + win32: ?*anyopaque = null, + compression: u8 = 0, + rdma: u8 = 0, + usage: u16 = 0, + reserved: [4]u8 = .{ 0, 0, 0, 0 }, + }; + const AccessDesc = extern struct { loc_type: c_int = 1, loc_id: c_int, flags: c_int = 3 }; // CUmemAccessDesc, read-write + + comptime { + std.debug.assert(@sizeOf(Prop) == 32 and @offsetOf(Prop, "compression") == 24 and @sizeOf(AccessDesc) == 12); + } + + /// `va_bytes` of virtual addresses reserved (cheap: no memory behind them until a chunk maps). + pub fn open(d: *const Driver, device: abi.Device, va_bytes: usize) Error!Vmm { + var vmm_ok: c_int = 0; + if (d.api.cuDeviceGetAttribute(&vmm_ok, .virtual_memory_management_supported, device) != abi.success or vmm_ok == 0) return error.Unsupported; // VIRTUAL_MEMORY_MANAGEMENT_SUPPORTED + var lib = std.DynLib.open("libcuda.so.1") catch return error.DriverUnavailable; + errdefer lib.close(); + var api: Api = undefined; + const names = .{ .{ "reserve", "cuMemAddressReserve" }, .{ "addressFree", "cuMemAddressFree" }, .{ "create", "cuMemCreate" }, .{ "release", "cuMemRelease" }, .{ "map", "cuMemMap" }, .{ "unmap", "cuMemUnmap" }, .{ "setAccess", "cuMemSetAccess" }, .{ "granularity", "cuMemGetAllocationGranularity" } }; + inline for (names) |n| @field(api, n[0]) = lib.lookup(@FieldType(Api, n[0]), n[1]) orelse return error.MissingSymbol; + var prop: Prop = .{ .loc_id = device }; + var rdma: c_int = 0; // GPU_DIRECT_RDMA_WITH_CUDA_VMM_SUPPORTED: keep the memory registrable as cuMemAlloc's is + if (d.api.cuDeviceGetAttribute(&rdma, .gpu_direct_rdma_with_cuda_vmm_supported, device) == abi.success and rdma != 0) prop.rdma = 1; + var g: usize = 0; + if (api.granularity(&g, &prop, 1) != abi.success or g == 0) return error.Unsupported; // RECOMMENDED + g = @max(g, 2 << 20); + const len = std.mem.alignForward(usize, va_bytes, g); + var va: u64 = 0; + if (api.reserve(&va, len, g, 0, 0) != abi.success) return error.CudaFailed; + return .{ .lib = lib, .device = device, .prop = prop, .granularity = g, .va = va, .va_len = len, .next = va, .api = api }; + } + + pub fn close(v: *Vmm) void { + _ = v.api.addressFree(v.va, v.va_len); + v.lib.close(); + } + + pub fn backend(v: *Vmm) Backend { + return .{ .ctx = v, .granularity = v.granularity, .chunk = chunk, .release = release, .name = "vmm" }; + } + + fn chunk(ctx: *anyopaque, len: usize) ?u64 { + const v: *Vmm = @ptrCast(@alignCast(ctx)); + if (v.va + v.va_len - v.next < len) return null; + var h: u64 = 0; + if (v.api.create(&h, len, &v.prop, 0) != abi.success) return null; + defer _ = v.api.release(h); // the mapping holds the memory until it is unmapped + const at = v.next; + if (v.api.map(at, len, 0, h, 0) != abi.success) return null; + const desc: AccessDesc = .{ .loc_id = v.device }; + if (v.api.setAccess(at, len, &desc, 1) != abi.success) { + _ = v.api.unmap(at, len); + return null; + } + v.next += len; + return at; + } + + fn release(ctx: *anyopaque, base: u64, len: usize) void { + const v: *Vmm = @ptrCast(@alignCast(ctx)); + _ = v.api.unmap(base, len); // the virtual range is not reused: the reservation is sized for the process + } +}; + +// -- tests (a fake backend: addresses only) --------------------------------------------------------------------------- + +const Fake = struct { + next: u64 = 1 << 40, + chunks: u32 = 0, + released: u32 = 0, + fail: bool = false, + + fn backend(f: *Fake, g: usize) Backend { + return .{ .ctx = f, .granularity = g, .chunk = chunk, .release = release, .name = "fake" }; + } + fn chunk(ctx: *anyopaque, len: usize) ?u64 { + const f: *Fake = @ptrCast(@alignCast(ctx)); + if (f.fail) return null; + const p = f.next; + f.next += len; + f.chunks += 1; + return p; + } + fn release(ctx: *anyopaque, _: u64, _: usize) void { + const f: *Fake = @ptrCast(@alignCast(ctx)); + f.released += 1; + } +}; + +test "arena: small requests pack 512-aligned into shared chunks, are reused by size, large ones get their own" { + const t = std.testing; + var f: Fake = .{}; + var a = Arena.init(t.allocator, f.backend(2 << 20), .{ .small_max = 1 << 20, .small_chunk = 4 << 20 }); + defer a.deinit(); + const p1 = a.alloc(20480).?; // a norm weight + const p2 = a.alloc(983040).?; // an mHC fn16 + const p3 = a.alloc(100).?; + try t.expectEqual(p1 + 20480, p2); + try t.expectEqual(p2 + 983040, p3); + for ([_]u64{ p1, p2, p3 }) |p| try t.expectEqual(@as(u64, 0), p % 512); + try t.expectEqual(@as(u32, 1), f.chunks); + // a freed range comes back for the same rounded size + try t.expect(a.free(p3)); + try t.expectEqual(p3, a.alloc(17).?); + try t.expectEqual(@as(u64, 1), a.stats.reused); + // three more 1 MiB ranges fit the 4 MiB chunk, the fourth opens a new shared chunk + for (0..3) |_| _ = a.alloc(1 << 20).?; + try t.expectEqual(@as(u32, 1), f.chunks); + _ = a.alloc(1 << 20).?; + try t.expectEqual(@as(u32, 2), f.chunks); + // large: its own chunk, rounded to the granularity, released when freed + const big = a.alloc(5 << 20).?; + try t.expectEqual(@as(u32, 3), f.chunks); + try t.expectEqual(@as(u64, 6 << 20), a.large.get(big).?); + try t.expect(a.free(big)); + try t.expectEqual(@as(u32, 1), f.released); + // not ours + try t.expect(!a.free(12345)); + // a backend that cannot map: null, the caller allocates as before + f.fail = true; + try t.expectEqual(@as(?u64, null), a.alloc(64 << 20)); + try t.expectEqual(@as(u64, 1), a.stats.fallback); +} + +test "arena: under DeviceBuffer.alloc through memory.setHook" { + const t = std.testing; + var f: Fake = .{}; + var a = Arena.init(t.allocator, f.backend(2 << 20), .{}); + defer a.deinit(); + a.install(); + defer memory.setHook(null); + const h = memory.currentHook().?; + const p = h.alloc(h.ctx, 4096).?; + try t.expect(h.free(h.ctx, p)); + try t.expect(!h.free(h.ctx, 7)); +} + +/// Slots of `m` that hold neither a key nor a tombstone. +fn freeSlots(m: *const std.AutoHashMapUnmanaged(u64, u64)) usize { + const meta = m.metadata orelse return 0; + var n: usize = 0; + for (meta[0..m.capacity()]) |x| n += @intFromBool(x.isFree()); + return n; +} + +test "arena: a long run of allocations and frees leaves both address maps free slots" { + const t = std.testing; + var f: Fake = .{}; + var a = Arena.init(t.allocator, f.backend(2 << 20), .{ .small_max = 1 << 20, .small_chunk = 4 << 20 }); + defer a.deinit(); + var held: std.ArrayList(u64) = .empty; + defer held.deinit(t.allocator); + var rng = std.Random.DefaultPrng.init(3); + for (0..20000) |_| { + // a serving process's buffers come and go (prompt workspaces, grown buffers): small and large alike + if (held.items.len > 0 and rng.random().boolean()) { + const k = rng.random().uintLessThan(usize, held.items.len); + try t.expect(a.free(held.swapRemove(k))); + } else { + const big = rng.random().uintLessThan(u32, 4) == 0; + try held.append(t.allocator, a.alloc(if (big) 3 << 20 else 512 * (1 + rng.random().uintLessThan(u64, 64))).?); + } + for ([_]*const std.AutoHashMapUnmanaged(u64, u64){ &a.live, &a.large }) |m| { + if (m.capacity() == 0) continue; + try t.expect(freeSlots(m) * 10 >= m.capacity()); + } + } +} + +test "arena: a reuse whose address-map insertion fails keeps its range free, and the retry gets it" { + const t = std.testing; + var f: Fake = .{}; + var fa = t.FailingAllocator.init(t.allocator, .{}); + { + var a = Arena.init(fa.allocator(), f.backend(2 << 20), .{ .small_max = 1 << 20, .small_chunk = 4 << 20 }); + defer a.deinit(); // testing.allocator underneath reports any leaked host bytes + const p = a.alloc(4096).?; + try t.expect(a.free(p)); + // other sizes take back the room the free gave: inserting the reused range must grow the map + while (a.live.available > 0) _ = a.alloc(512).?; + const live = a.live.count(); + const cursor = a.cursor; + const chunks = f.chunks; + fa.fail_index = fa.alloc_index; // the next allocation fails, and every one after it until healed + try t.expectEqual(@as(?u64, null), a.alloc(4096)); + try t.expectEqual(live, a.live.count()); + try t.expectEqual(cursor, a.cursor); + try t.expectEqual(chunks, f.chunks); + try t.expectEqualSlices(u64, &.{p}, a.free_small.get(4096).?.items); + try t.expectEqual(@as(u64, 0), a.stats.reused); + fa.fail_index = std.math.maxInt(usize); + try t.expectEqual(p, a.alloc(4096).?); + try t.expectEqual(@as(u64, 1), a.stats.reused); + try t.expectEqual(@as(usize, 0), a.free_small.get(4096).?.items.len); + try t.expect(a.free(p)); + } + try t.expectEqual(f.chunks, f.released); // every chunk went back to the backend +} diff --git a/zig/src/cuda/driver.zig b/zig/src/cuda/driver.zig index 7c7fb7f068..38a9b930d7 100644 --- a/zig/src/cuda/driver.zig +++ b/zig/src/cuda/driver.zig @@ -9,6 +9,10 @@ pub const Error = error{ CopySourceLimit, StageBusy, StageUnprepared, UnownedCap pub const Driver = struct { lib: platform.DynLib, api: abi.Api, + /// cuFuncGetName (CUDA 12.3+; optional: an older driver opens without it): a failed launch names its kernel + func_name: ?FuncGetName = null, + + pub const FuncGetName = *const fn (*?[*:0]const u8, abi.Function) callconv(.c) abi.Result; pub fn open() Error!Driver { return openPath(platform.cuda_driver); @@ -26,7 +30,7 @@ pub const Driver = struct { return error.MissingSymbol; }; } - const d: Driver = .{ .lib = lib, .api = api }; + const d: Driver = .{ .lib = lib, .api = api, .func_name = lib.lookup(FuncGetName, "cuFuncGetName") }; try d.check(api.cuInit(0), "cuInit"); return d; } @@ -43,18 +47,33 @@ pub const Driver = struct { pub fn check(self: *const Driver, res: abi.Result, what: []const u8) Error!void { if (res == abi.success) return; if (res == abi.error_not_ready) return error.NotReady; - std.log.err("{s}: {s} ({d}) {s}", .{ what, self.errorName(res), res, self.errorText(res) }); + std.log.err("{s}: {s} ({d}) {s}{s}", .{ what, self.errorName(res), res, self.errorText(res), shutdownNote(res) }); + return errorOf(res); + } + + /// The error a failed call returns (`check` without its log line, for callers that log more). + pub fn errorOf(res: abi.Result) Error { return switch (res) { 2 => error.OutOfDeviceMemory, abi.error_not_found => error.NotFound, + abi.error_not_ready => error.NotReady, else => error.CudaFailed, }; } + /// The driver's name for the code, else knownName's: during the exit's teardown cuGetErrorName fails too. pub fn errorName(self: *const Driver, res: abi.Result) []const u8 { var s: ?[*:0]const u8 = null; - if (self.api.cuGetErrorName(res, &s) != abi.success) return "CUDA_ERROR_UNKNOWN_CODE"; - return if (s) |p| std.mem.span(p) else "CUDA_ERROR_UNKNOWN_CODE"; + if (self.api.cuGetErrorName(res, &s) == abi.success) if (s) |p| return std.mem.span(p); + return knownName(res) orelse "CUDA_ERROR_UNKNOWN_CODE"; + } + + /// A kernel's symbol name (cuFuncGetName), or "?" when the driver cannot say. + pub fn functionName(self: *const Driver, f: abi.Function) []const u8 { + const get = self.func_name orelse return "?"; + var s: ?[*:0]const u8 = null; + if (get(&s, f) != abi.success) return "?"; + return if (s) |p| std.mem.span(p) else "?"; } pub fn errorText(self: *const Driver, res: abi.Result) []const u8 { @@ -96,3 +115,43 @@ pub const Driver = struct { return std.mem.sliceTo(buf, 0); } }; + +/// CUresult names (cuda.h) for when cuGetErrorName cannot answer. +pub fn knownName(res: abi.Result) ?[]const u8 { + return switch (res) { + 1 => "CUDA_ERROR_INVALID_VALUE", + 2 => "CUDA_ERROR_OUT_OF_MEMORY", + 3 => "CUDA_ERROR_NOT_INITIALIZED", + 4 => "CUDA_ERROR_DEINITIALIZED", + 101 => "CUDA_ERROR_INVALID_DEVICE", + 200 => "CUDA_ERROR_INVALID_IMAGE", + 201 => "CUDA_ERROR_INVALID_CONTEXT", + 209 => "CUDA_ERROR_NO_BINARY_FOR_GPU", + 400 => "CUDA_ERROR_INVALID_HANDLE", + 500 => "CUDA_ERROR_NOT_FOUND", + 600 => "CUDA_ERROR_NOT_READY", + 700 => "CUDA_ERROR_ILLEGAL_ADDRESS", + 701 => "CUDA_ERROR_LAUNCH_OUT_OF_RESOURCES", + 702 => "CUDA_ERROR_LAUNCH_TIMEOUT", + 709 => "CUDA_ERROR_CONTEXT_IS_DESTROYED", + 710 => "CUDA_ERROR_ASSERT", + 714 => "CUDA_ERROR_HARDWARE_STACK_ERROR", + 715 => "CUDA_ERROR_ILLEGAL_INSTRUCTION", + 716 => "CUDA_ERROR_MISALIGNED_ADDRESS", + 719 => "CUDA_ERROR_LAUNCH_FAILED", + 999 => "CUDA_ERROR_UNKNOWN", + else => null, + }; +} + +/// CUDA_ERROR_DEINITIALIZED: the process is exiting (a stop signal, then exit) while a thread still calls the driver. +pub fn shutdownNote(res: abi.Result) []const u8 { + return if (res == 4) " (the CUDA driver is shutting down: the process is exiting)" else ""; +} + +test "CUresult names without the driver: DEINITIALIZED is 4, unknown codes stay unknown" { + try std.testing.expectEqualStrings("CUDA_ERROR_DEINITIALIZED", knownName(4).?); + try std.testing.expectEqualStrings("CUDA_ERROR_ILLEGAL_ADDRESS", knownName(700).?); + try std.testing.expect(knownName(12345) == null); + try std.testing.expect(shutdownNote(4).len > 0 and shutdownNote(1).len == 0); +} diff --git a/zig/src/cuda/graph_cache.zig b/zig/src/cuda/graph_cache.zig new file mode 100644 index 0000000000..1916b22b57 --- /dev/null +++ b/zig/src/cuda/graph_cache.zig @@ -0,0 +1,448 @@ +//! Keyed CUDA graphs for any family: one captured executable per key, launched again when the key comes back. + +const std = @import("std"); +const graph = @import("graph.zig"); +const Driver = @import("driver.zig").Driver; +const Stream = @import("stream.zig").Stream; +const quiet = @import("builtin").is_test; // a test's stderr is the build runner's channel + +/// One step's launches on `stream`; it must not synchronize, allocate device memory or read the device on the host. +pub const Body = struct { + ctx: *anyopaque, + run: *const fn (ctx: *anyopaque, stream: Stream) anyerror!void, +}; + +/// Every rank's verdict AND-ed on a miss, never a hit; an off cache never votes, so ranks match `Settings.digest`. +pub const Agree = struct { + ctx: *anyopaque, + all: *const fn (ctx: *anyopaque, ok: bool) anyerror!bool, +}; + +/// Capture a body into an executable graph, launch it, free it: a fake engine runs the policy in host tests. +pub const Engine = struct { + ctx: *anyopaque, + capture: *const fn (ctx: *anyopaque, stream: Stream, body: Body) anyerror!*anyopaque, + launch: *const fn (ctx: *anyopaque, exec: *anyopaque, stream: Stream) anyerror!void, + free: *const fn (ctx: *anyopaque, exec: *anyopaque) void, +}; + +pub const Settings = struct { + on: bool = true, // off: every step runs eagerly + max: u32 = 48, // graphs held at most + min_keep: u32 = 8, // under the memory floor the cache evicts down to this many, then stops capturing + hold: bool = false, // under the floor a miss runs eagerly with no eviction: held graphs cost next to no memory + + /// These settings under the operator's TF_GRAPH_CACHE value: `0` or `off` runs every step eagerly; else unchanged. + pub fn operator(s: Settings, value: ?[]const u8) Settings { + var out = s; + const v = value orelse return out; + if (std.mem.eql(u8, v, "0") or std.ascii.eqlIgnoreCase(v, "off")) out.on = false; + return out; + } + + /// What a cache runs with: the operator's value, min_keep under max (0 with hold), and every off cache one value. + pub fn effective(s: Settings, value: ?[]const u8) Settings { + var out = s.operator(value); + if (!out.on or out.max == 0) return .{ .on = false, .max = 0, .min_keep = 0, .hold = false }; + out.min_keep = if (out.hold) 0 else @min(out.min_keep, out.max); + return out; + } + + /// `effective` under this process's TF_GRAPH_CACHE. + pub fn fromEnv(s: Settings) Settings { + return s.effective(if (std.c.getenv("TF_GRAPH_CACHE")) |v| std.mem.span(v) else null); + } + + /// The effective settings as bytes for ranks to compare at start (a family's bootstrap agreement), never per step. + pub fn digest(s: Settings) [10]u8 { + const e = s.effective(null); + var out: [10]u8 = undefined; + out[0] = @intFromBool(e.on); + out[1] = @intFromBool(e.hold); + std.mem.writeInt(u32, out[2..6], e.max, .little); + std.mem.writeInt(u32, out[6..10], e.min_keep, .little); + return out; + } +}; + +pub const Outcome = enum { replayed, captured, eager }; + +pub fn Cache(comptime Key: type) type { + return struct { + const Self = @This(); + const Entry = struct { exec: *anyopaque, used: u64 }; + + gpa: std.mem.Allocator, + engine: Engine, + settings: Settings, + agree: ?Agree = null, + room: ?*const fn () bool = null, // false when this rank is below its memory floor; null: always room + entries: std.AutoArrayHashMapUnmanaged(Key, Entry) = .empty, + cap: u32, + tick: u64 = 0, + fingerprint: ?u64 = null, + failed: bool = false, // a capture failed on some rank: eager from here on + stats: struct { replayed: u64 = 0, captured: u64 = 0, eager: u64 = 0, evicted: u64 = 0, dropped: u64 = 0, floor_hits: u64 = 0 } = .{}, + + /// The cache with `settings.fromEnv()`: every rank sharing an `agree` must hold an equal `Settings.digest`. + pub fn init(gpa: std.mem.Allocator, engine: Engine, settings: Settings) Self { + const s = settings.fromEnv(); + if (!quiet and settings.on and settings.max > 0 and !s.on) std.log.info("graph cache: off (TF_GRAPH_CACHE), every step eager", .{}); + return .{ .gpa = gpa, .engine = engine, .settings = s, .cap = s.max }; + } + + pub fn deinit(c: *Self) void { + c.dropAll(); + c.entries.deinit(c.gpa); + } + + pub fn count(c: *const Self) usize { + return c.entries.count(); + } + + pub fn has(c: *const Self, key: Key) bool { + return c.entries.contains(key); + } + + /// Every graph freed (a fingerprint change, a shutdown). + pub fn dropAll(c: *Self) void { + for (c.entries.values()) |e| c.engine.free(c.engine.ctx, e.exec); + c.stats.dropped += c.entries.count(); + c.entries.clearRetainingCapacity(); + } + + fn agreed(c: *Self, ok: bool) !bool { + const a = c.agree orelse return ok; + return a.all(a.ctx, ok); + } + + fn evictOldest(c: *Self) void { + var at: ?usize = null; + for (c.entries.values(), 0..) |e, i| { + if (at == null or e.used < c.entries.values()[at.?].used) at = i; + } + const i = at orelse return; + c.engine.free(c.engine.ctx, c.entries.values()[i].exec); + c.entries.swapRemoveAt(i); + c.stats.evicted += 1; + } + + fn eager(c: *Self, stream: Stream, body: Body) !Outcome { + try body.run(body.ctx, stream); + c.stats.eager += 1; + return .eager; + } + + /// One step: the key's graph, captured first on a miss; a new `fingerprint` (same step on all ranks) drops all. + pub fn run(c: *Self, key: Key, stream: Stream, fingerprint: u64, body: Body) !Outcome { + if (!c.settings.on or c.failed or c.settings.max == 0) return c.eager(stream, body); + if (c.fingerprint != null and c.fingerprint.? != fingerprint) c.dropAll(); + c.fingerprint = fingerprint; + c.tick += 1; + if (c.entries.getPtr(key)) |e| { + e.used = c.tick; + try c.engine.launch(c.engine.ctx, e.exec, stream); + c.stats.replayed += 1; + return .replayed; + } + // a miss: room on every rank (a host map that cannot grow is none), else shed a quarter and lower the cap + const reserved = if (c.entries.ensureUnusedCapacity(c.gpa, 1)) |_| true else |_| false; + const room_here = reserved and (if (c.room) |f| f() else true); + if (!try c.agreed(room_here)) { + if (c.settings.hold) { + if (!quiet and c.stats.floor_hits == 0) std.log.warn("graph cache: under the memory floor on {s} rank: this key eager, {d} graphs held (hold)", .{ if (room_here) "another" else "this", c.entries.count() }); + c.stats.floor_hits += 1; + return c.eager(stream, body); + } + if (c.entries.count() > c.settings.min_keep) { + const drop = @max(1, c.entries.count() / 4); + for (0..drop) |_| if (c.entries.count() > c.settings.min_keep) c.evictOldest(); + } + const was = c.cap; + c.cap = @intCast(@max(c.entries.count(), c.settings.min_keep)); // min_keep <= max (init) + c.stats.floor_hits += 1; + if (!quiet and c.cap != was) std.log.warn("graph cache: under the memory floor on {s} rank: cap {d} -> {d} ({d} held, {d} evicted, {d} eager so far)", .{ if (room_here) "another" else "this", was, c.cap, c.entries.count(), c.stats.evicted, c.stats.eager + 1 }); + return c.eager(stream, body); + } + if (c.cap < c.settings.max) { + if (!quiet) std.log.info("graph cache: room again on every rank: cap {d} -> {d}", .{ c.cap, c.settings.max }); + c.cap = c.settings.max; + } + while (c.entries.count() >= c.cap) c.evictOldest(); + const exec = c.engine.capture(c.engine.ctx, stream, body) catch |err| blk: { + // the only trace of why graphs went off (the step then runs eagerly, which may fail the same way) + if (!quiet) std.log.warn("graph capture failed: {s} (key {any}, {d} graphs held)", .{ @errorName(err), key, c.entries.count() }); + break :blk null; + }; + var owned = exec; // freed here until the map holds it, so a failed agreement cannot leak it + defer if (owned) |x| c.engine.free(c.engine.ctx, x); + if (!try c.agreed(exec != null)) { + c.failed = true; + c.dropAll(); + return c.eager(stream, body); + } + c.entries.putAssumeCapacityNoClobber(key, .{ .exec = exec.?, .used = c.tick }); + owned = null; + try c.engine.launch(c.engine.ctx, exec.?, stream); + c.stats.captured += 1; + return .captured; + } + }; +} + +/// The driver's engine: thread-local stream capture (other threads' CUDA calls stay out), instantiate, upload, launch. +pub const CudaEngine = struct { + d: *const Driver, + + pub fn engine(e: *CudaEngine) Engine { + return .{ .ctx = e, .capture = capture, .launch = launch, .free = free }; + } + + const Held = struct { exec: graph.Exec }; + + fn capture(ctx: *anyopaque, stream: Stream, body: Body) anyerror!*anyopaque { + const e: *CudaEngine = @ptrCast(@alignCast(ctx)); + try graph.beginCapture(stream, .thread_local); + body.run(body.ctx, stream) catch |err| { + // end the capture so the stream is usable again; the partial graph is dropped + if (graph.endCapture(stream)) |g| { + var gg = g; + gg.deinit(); + } else |_| {} + return err; + }; + var g = try graph.endCapture(stream); + defer g.deinit(); + const h = try std.heap.c_allocator.create(Held); + errdefer std.heap.c_allocator.destroy(h); + h.exec = try g.instantiate(); + errdefer h.exec.deinit(); + try h.exec.upload(stream); + _ = e; + return h; + } + + fn launch(_: *anyopaque, exec: *anyopaque, stream: Stream) anyerror!void { + const h: *Held = @ptrCast(@alignCast(exec)); + try h.exec.launchOn(stream); + } + + fn free(_: *anyopaque, exec: *anyopaque) void { + const h: *Held = @ptrCast(@alignCast(exec)); + h.exec.deinit(); + std.heap.c_allocator.destroy(h); + } +}; + +/// FNV-1a over the addresses a body captures (buffer bases, sizes): the fingerprint `run` compares, on one rank only. +pub fn fingerprintOf(words: []const u64) u64 { + var h: u64 = 0xcbf29ce484222325; + for (words) |w| { + var x = w; + for (0..8) |_| { + h = (h ^ (x & 0xff)) *% 0x100000001b3; + x >>= 8; + } + } + return h; +} + +// -- host tests: the policy with a fake engine -------------------------------------------------------------------- + +const Fake = struct { + captures: u32 = 0, + launches: u32 = 0, + live: u32 = 0, + fail_capture: bool = false, + slots: [64]u8 = undefined, + + fn engine(f: *Fake) Engine { + return .{ .ctx = f, .capture = cap, .launch = lau, .free = fre }; + } + fn cap(ctx: *anyopaque, stream: Stream, body: Body) anyerror!*anyopaque { + const f: *Fake = @ptrCast(@alignCast(ctx)); + if (f.fail_capture) return error.CaptureFailed; + _ = stream; + _ = body; // a capture records the body; nothing runs + f.captures += 1; + f.live += 1; + return &f.slots[f.captures % f.slots.len]; + } + fn lau(ctx: *anyopaque, _: *anyopaque, _: Stream) anyerror!void { + const f: *Fake = @ptrCast(@alignCast(ctx)); + f.launches += 1; + } + fn fre(ctx: *anyopaque, _: *anyopaque) void { + const f: *Fake = @ptrCast(@alignCast(ctx)); + f.live -= 1; + } +}; + +const Ran = struct { + n: u32 = 0, + fn body(r: *Ran) Body { + return .{ .ctx = r, .run = go }; + } + fn go(ctx: *anyopaque, _: Stream) anyerror!void { + const r: *Ran = @ptrCast(@alignCast(ctx)); + r.n += 1; + } +}; + +const no_stream: Stream = .{ .d = undefined, .handle = null }; + +test "graph cache: capture on first use, replay after, LRU at the cap" { + var f: Fake = .{}; + var r: Ran = .{}; + var c = Cache(u32).init(std.testing.allocator, f.engine(), .{ .max = 2 }); + defer c.deinit(); + try std.testing.expectEqual(Outcome.captured, try c.run(1, no_stream, 7, r.body())); + try std.testing.expectEqual(Outcome.replayed, try c.run(1, no_stream, 7, r.body())); + try std.testing.expectEqual(Outcome.captured, try c.run(2, no_stream, 7, r.body())); + _ = try c.run(1, no_stream, 7, r.body()); // 1 is now the most recent + try std.testing.expectEqual(Outcome.captured, try c.run(3, no_stream, 7, r.body())); // evicts 2 + try std.testing.expect(c.has(1) and c.has(3) and !c.has(2)); + try std.testing.expectEqual(@as(u32, 0), r.n); // the body never ran eagerly + try std.testing.expectEqual(@as(u32, 2), f.live); + try std.testing.expectEqual(@as(u32, 5), f.launches); +} + +test "graph cache: a new fingerprint drops every graph" { + var f: Fake = .{}; + var r: Ran = .{}; + var c = Cache(u32).init(std.testing.allocator, f.engine(), .{}); + defer c.deinit(); + _ = try c.run(1, no_stream, 7, r.body()); + _ = try c.run(2, no_stream, 7, r.body()); + try std.testing.expectEqual(Outcome.captured, try c.run(1, no_stream, 8, r.body())); + try std.testing.expectEqual(@as(usize, 1), c.count()); + try std.testing.expectEqual(@as(u32, 1), f.live); +} + +const Peer = struct { + theirs: []const bool, // the other rank's verdicts, in call order + at: usize = 0, + fn agree(p: *Peer) Agree { + return .{ .ctx = p, .all = all }; + } + fn all(ctx: *anyopaque, ok: bool) anyerror!bool { + const p: *Peer = @ptrCast(@alignCast(ctx)); + const t = p.theirs[p.at]; + p.at += 1; + return ok and t; + } +}; + +test "graph cache: a capture failing on another rank turns graphs off here too" { + var f: Fake = .{}; + var r: Ran = .{}; + // per miss: room, then the capture's verdict; the peer's capture fails on the second key + var peer: Peer = .{ .theirs = &.{ true, true, true, false } }; + var c = Cache(u32).init(std.testing.allocator, f.engine(), .{}); + c.agree = peer.agree(); + defer c.deinit(); + try std.testing.expectEqual(Outcome.captured, try c.run(1, no_stream, 7, r.body())); + try std.testing.expectEqual(Outcome.eager, try c.run(2, no_stream, 7, r.body())); + try std.testing.expect(c.failed); + try std.testing.expectEqual(@as(u32, 0), f.live); // this rank's good graphs are freed as well + try std.testing.expectEqual(Outcome.eager, try c.run(1, no_stream, 7, r.body())); + try std.testing.expectEqual(@as(u32, 2), r.n); +} + +var room_flag = true; +fn roomFn() bool { + return room_flag; +} + +test "graph cache: below the memory floor the least recent quarter goes, the cap follows, the step runs eagerly" { + var f: Fake = .{}; + var r: Ran = .{}; + var c = Cache(u32).init(std.testing.allocator, f.engine(), .{ .max = 48, .min_keep = 2 }); + c.room = roomFn; + defer c.deinit(); + room_flag = true; + for (0..8) |k| _ = try c.run(@intCast(k), no_stream, 7, r.body()); + room_flag = false; + try std.testing.expectEqual(Outcome.eager, try c.run(100, no_stream, 7, r.body())); + try std.testing.expectEqual(@as(usize, 6), c.count()); + try std.testing.expectEqual(@as(u32, 6), c.cap); + try std.testing.expect(!c.has(0) and !c.has(1) and c.has(7)); + room_flag = true; + // room again: the cap is back to max, the miss captures without evicting + try std.testing.expectEqual(Outcome.captured, try c.run(100, no_stream, 7, r.body())); + try std.testing.expectEqual(@as(usize, 7), c.count()); + try std.testing.expectEqual(@as(u32, 48), c.cap); + try std.testing.expect(c.has(7) and c.has(2)); + try std.testing.expectEqual(@as(u64, 1), c.stats.floor_hits); +} + +test "graph cache: a dip under the floor does not leave later keys churning once memory is back" { + var f: Fake = .{}; + var r: Ran = .{}; + var c = Cache(u32).init(std.testing.allocator, f.engine(), .{ .max = 16, .min_keep = 2 }); + c.room = roomFn; + defer c.deinit(); + room_flag = true; + for (0..12) |k| _ = try c.run(@intCast(k), no_stream, 7, r.body()); + room_flag = false; // a long prompt's workspace held + _ = try c.run(200, no_stream, 7, r.body()); + room_flag = true; // freed: the decode after it meets 12 new keys (a new context bucket), then replays them + const before = c.stats.evicted; + for (0..12) |k| _ = try c.run(@intCast(300 + k), no_stream, 7, r.body()); + for (0..12) |k| try std.testing.expectEqual(Outcome.replayed, try c.run(@intCast(300 + k), no_stream, 7, r.body())); + try std.testing.expectEqual(@as(u32, 16), c.cap); + // only what the cap of 16 forces: 9 left after the dip + 12 new = 21 -> 5 evictions, none during the replays + try std.testing.expectEqual(before + 5, c.stats.evicted); +} + +test "graph cache: hold - below the floor a miss runs eagerly, nothing is evicted, the held keys replay" { + var f: Fake = .{}; + var r: Ran = .{}; + var c = Cache(u32).init(std.testing.allocator, f.engine(), .{ .max = 48, .min_keep = 2, .hold = true }); + c.room = roomFn; + defer c.deinit(); + room_flag = true; + for (0..8) |k| _ = try c.run(@intCast(k), no_stream, 7, r.body()); + room_flag = false; + try std.testing.expectEqual(Outcome.eager, try c.run(100, no_stream, 7, r.body())); + try std.testing.expectEqual(Outcome.eager, try c.run(101, no_stream, 7, r.body())); + try std.testing.expectEqual(@as(usize, 8), c.count()); + try std.testing.expectEqual(@as(u32, 48), c.cap); + for (0..8) |k| try std.testing.expectEqual(Outcome.replayed, try c.run(@intCast(k), no_stream, 7, r.body())); + try std.testing.expectEqual(@as(u64, 0), c.stats.evicted); + try std.testing.expectEqual(@as(u64, 2), c.stats.floor_hits); + room_flag = true; + try std.testing.expectEqual(Outcome.captured, try c.run(100, no_stream, 7, r.body())); + try std.testing.expectEqual(@as(usize, 9), c.count()); +} + +test "graph cache: off runs every step eagerly" { + var f: Fake = .{}; + var r: Ran = .{}; + var c = Cache(u32).init(std.testing.allocator, f.engine(), .{ .on = false }); + defer c.deinit(); + try std.testing.expectEqual(Outcome.eager, try c.run(1, no_stream, 7, r.body())); + try std.testing.expectEqual(@as(u32, 0), f.captures); +} + +test "graph cache: TF_GRAPH_CACHE=0 or off turns a family's cache off without a rebuild; anything else leaves it" { + const on: Settings = .{}; + for ([_][]const u8{ "0", "off", "OFF" }) |v| try std.testing.expect(!on.operator(v).on); + for ([_]?[]const u8{ null, "1", "on", "" }) |v| try std.testing.expect(on.operator(v).on); + try std.testing.expect(!(Settings{ .on = false }).operator("1").on); // the operator only turns a cache off + var f: Fake = .{}; + var r: Ran = .{}; + var c = Cache(u32).init(std.testing.allocator, f.engine(), on.operator("off")); + defer c.deinit(); + for (0..3) |_| try std.testing.expectEqual(Outcome.eager, try c.run(1, no_stream, 7, r.body())); + try std.testing.expectEqual(@as(u32, 0), f.captures); +} + +test "graph cache: fingerprint" { + try std.testing.expect(fingerprintOf(&.{ 1, 2 }) != fingerprintOf(&.{ 2, 1 })); +} + +test { + _ = @import("graph_cache_regression_test.zig"); + _ = @import("graph_cache_ranks_test.zig"); +} diff --git a/zig/src/cuda/graph_cache_ranks_test.zig b/zig/src/cuda/graph_cache_ranks_test.zig new file mode 100644 index 0000000000..18b8a87e9d --- /dev/null +++ b/zig/src/cuda/graph_cache_ranks_test.zig @@ -0,0 +1,136 @@ +//! Two in-process ranks with fake engines: equal effective settings vote on misses only; a mismatch is a stall. + +const std = @import("std"); +const cache = @import("graph_cache.zig"); +const Stream = @import("stream.zig").Stream; +const no_stream: Stream = .{ .d = undefined, .handle = null }; +const equal = std.testing.expectEqual; +const Settings = cache.Settings; + +const Engine = struct { + captures: usize = 0, + live: usize = 0, + handles: [8]u8 = undefined, + + fn engine(e: *Engine) cache.Engine { + return .{ .ctx = e, .capture = capture, .launch = launch, .free = free }; + } + fn capture(p: *anyopaque, _: Stream, _: cache.Body) !*anyopaque { + const e: *Engine = @ptrCast(@alignCast(p)); + e.captures += 1; + e.live += 1; + return &e.handles[e.captures % e.handles.len]; + } + fn launch(_: *anyopaque, _: *anyopaque, _: Stream) !void {} + fn free(p: *anyopaque, _: *anyopaque) void { + const e: *Engine = @ptrCast(@alignCast(p)); + e.live -= 1; + } + fn body(e: *Engine) cache.Body { + return .{ .ctx = e, .run = run }; + } + fn run(_: *anyopaque, _: Stream) !void {} +}; + +/// The two ranks' votes in call order; a rank waits at most `wait_ns` for the other's vote of the same round. +const Board = struct { + votes: [2][16]bool = undefined, + cast: [2]std.atomic.Value(u32) = .{ .init(0), .init(0) }, + wait_ns: u64, + + fn vote(b: *Board, rank: u1, ok: bool) !bool { + const r = b.cast[rank].load(.monotonic); + if (r == b.votes[rank].len) return error.TooManyVotes; + b.votes[rank][r] = ok; + b.cast[rank].store(r + 1, .release); + const other = 1 - rank; + const deadline = now() + b.wait_ns; + while (b.cast[other].load(.acquire) <= r) { + if (now() > deadline) return error.PeerNeverVoted; + std.Thread.yield() catch {}; + } + return ok and b.votes[other][r]; + } + + fn now() u64 { + return @intCast(std.Io.Clock.awake.now(std.testing.io).toNanoseconds()); + } +}; + +const Rank = struct { + board: *Board, + id: u1, + engine: Engine = .{}, + outcomes: [8]cache.Outcome = undefined, + err: ?anyerror = null, + + fn all(ctx: *anyopaque, ok: bool) anyerror!bool { + const r: *Rank = @ptrCast(@alignCast(ctx)); + return r.board.vote(r.id, ok); + } + + /// One rank's steps over `keys` with `settings` exactly as its process resolved them. + fn steps(r: *Rank, settings: Settings, ks: []const u32) void { + var c = cache.Cache(u32).init(std.testing.allocator, r.engine.engine(), settings); + defer c.deinit(); + c.agree = .{ .ctx = r, .all = all }; + for (ks, 0..) |k, i| r.outcomes[i] = c.run(k, no_stream, 7, r.engine.body()) catch |e| { + r.err = e; + return; + }; + } +}; + +const keys = [_]u32{ 1, 2, 1, 2, 1 }; + +test "graph cache ranks: equal enabled settings capture together and vote on each miss, never on a hit" { + var board: Board = .{ .wait_ns = 5 * std.time.ns_per_s }; + var ranks = [2]Rank{ .{ .board = &board, .id = 0 }, .{ .board = &board, .id = 1 } }; + const s = (Settings{ .max = 4 }).effective(null); + const t = try std.Thread.spawn(.{}, Rank.steps, .{ &ranks[1], s, &keys }); + ranks[0].steps(s, &keys); + t.join(); + const want = [_]cache.Outcome{ .captured, .captured, .replayed, .replayed, .replayed }; + for (&ranks) |*r| { + try equal(@as(?anyerror, null), r.err); + try std.testing.expectEqualSlices(cache.Outcome, &want, r.outcomes[0..keys.len]); + try equal(@as(u32, 4), board.cast[r.id].load(.acquire)); // two misses: room and capture each, hits none + try equal(@as(usize, 0), r.engine.live); + } +} + +test "graph cache ranks: TF_GRAPH_CACHE 0, off and OFF are one effective setting, and an off cache never votes" { + const d: Settings = .{}; + const off = d.effective("0").digest(); + try std.testing.expectEqualSlices(u8, &off, &d.effective("off").digest()); + try std.testing.expectEqualSlices(u8, &off, &d.effective("OFF").digest()); + try std.testing.expectEqualSlices(u8, &off, &(Settings{ .max = 0, .hold = true }).digest()); + try std.testing.expect(!std.mem.eql(u8, &off, &d.effective(null).digest())); + const held: Settings = .{ .hold = true, .min_keep = 2 }; + try std.testing.expectEqualSlices(u8, &held.digest(), &(Settings{ .hold = true, .min_keep = 30 }).digest()); + try std.testing.expectEqualSlices(u8, &d.effective("1").digest(), &d.effective(null).digest()); + var board: Board = .{ .wait_ns = 0 }; + var ranks = [2]Rank{ .{ .board = &board, .id = 0 }, .{ .board = &board, .id = 1 } }; + ranks[0].steps(d.effective("0"), &keys); + ranks[1].steps(d.effective("OFF"), &keys); + for (&ranks) |*r| { + try equal(@as(?anyerror, null), r.err); + for (r.outcomes[0..keys.len]) |o| try equal(cache.Outcome.eager, o); + try equal(@as(u32, 0), board.cast[r.id].load(.acquire)); + try equal(@as(usize, 0), r.engine.captures); + } +} + +test "graph cache ranks: one rank switched off leaves the other waiting at its first miss; digests show it first" { + const d: Settings = .{}; + try std.testing.expect(!std.mem.eql(u8, &d.effective(null).digest(), &d.effective("off").digest())); + var board: Board = .{ .wait_ns = 50 * std.time.ns_per_ms }; + var ranks = [2]Rank{ .{ .board = &board, .id = 0 }, .{ .board = &board, .id = 1 } }; + ranks[1].steps(d.effective("off"), &keys); + try equal(@as(?anyerror, null), ranks[1].err); + try equal(@as(u32, 0), board.cast[1].load(.acquire)); // the off rank ran every step eagerly without a vote + ranks[0].steps(d.effective(null), &keys); + try equal(@as(?anyerror, error.PeerNeverVoted), ranks[0].err); // in production this wait has no deadline + try equal(@as(u32, 1), board.cast[0].load(.acquire)); + try equal(@as(usize, 0), ranks[0].engine.live); +} diff --git a/zig/src/cuda/graph_cache_regression_test.zig b/zig/src/cuda/graph_cache_regression_test.zig new file mode 100644 index 0000000000..7f2da1063a --- /dev/null +++ b/zig/src/cuda/graph_cache_regression_test.zig @@ -0,0 +1,181 @@ +//! Failure ownership, rank verdict counts and bounded-cap regressions for the graph cache. + +const std = @import("std"); +const cache = @import("graph_cache.zig"); +const Stream = @import("stream.zig").Stream; +const no_stream: Stream = .{ .d = undefined, .handle = null }; +const expect = std.testing.expect; +const equal = std.testing.expectEqual; + +const Engine = struct { + captures: usize = 0, + live: usize = 0, + launches: usize = 0, + bodies: usize = 0, + handles: [32]u8 = undefined, + + fn engine(e: *Engine) cache.Engine { + return .{ .ctx = e, .capture = capture, .launch = launch, .free = free }; + } + fn capture(p: *anyopaque, _: Stream, _: cache.Body) !*anyopaque { + const e: *Engine = @ptrCast(@alignCast(p)); + e.captures += 1; + e.live += 1; + return &e.handles[e.captures % e.handles.len]; + } + fn launch(p: *anyopaque, _: *anyopaque, _: Stream) !void { + const e: *Engine = @ptrCast(@alignCast(p)); + e.launches += 1; + } + fn free(p: *anyopaque, _: *anyopaque) void { + const e: *Engine = @ptrCast(@alignCast(p)); + std.debug.assert(e.live > 0); + e.live -= 1; + } + fn body(e: *Engine) cache.Body { + return .{ .ctx = e, .run = run }; + } + fn run(p: *anyopaque, _: Stream) !void { + const e: *Engine = @ptrCast(@alignCast(p)); + e.bodies += 1; + } +}; + +const Peer = struct { + verdicts: []const bool, + at: usize = 0, + throw_at: ?usize = null, + fn agree(p: *Peer) cache.Agree { + return .{ .ctx = p, .all = all }; + } + fn all(ctx: *anyopaque, ok: bool) !bool { + const p: *Peer = @ptrCast(@alignCast(ctx)); + const at = p.at; + p.at += 1; + if (p.throw_at == at) return error.CollectiveFailed; + if (at >= p.verdicts.len) return error.UnexpectedCollective; + return ok and p.verdicts[at]; + } +}; + +test "graph cache regression: capture agreement throw frees the not-yet-inserted executable" { + var e: Engine = .{}; + // per miss: room, then the capture's verdict + var p: Peer = .{ .verdicts = &.{ true, true }, .throw_at = 1 }; + var c = cache.Cache(u32).init(std.testing.allocator, e.engine(), .{}); + defer c.deinit(); + c.agree = p.agree(); + try std.testing.expectError(error.CollectiveFailed, c.run(1, no_stream, 7, e.body())); + try equal(@as(usize, 1), e.captures); + try equal(@as(usize, 0), e.live); + try equal(@as(usize, 0), c.count()); + try equal(@as(usize, 2), p.at); +} + +test "graph cache regression: local map OOM declines before any capture and can recover" { + var e: Engine = .{}; + var failing = std.testing.FailingAllocator.init(std.testing.allocator, .{ .fail_index = 0 }); + var c = cache.Cache(u32).init(failing.allocator(), e.engine(), .{}); + defer c.deinit(); + try equal(cache.Outcome.eager, try c.run(1, no_stream, 7, e.body())); + try expect(failing.has_induced_failure); + try equal(@as(usize, 0), e.captures); + try equal(@as(usize, 0), e.live); + try expect(!c.failed); + failing.fail_index = std.math.maxInt(usize); + try equal(cache.Outcome.captured, try c.run(1, no_stream, 7, e.body())); + try equal(@as(usize, 1), e.live); +} + +test "graph cache regression: peer map OOM makes a locally prepared rank eager before capture" { + var e: Engine = .{}; + var p: Peer = .{ .verdicts = &.{false} }; + var c = cache.Cache(u32).init(std.testing.allocator, e.engine(), .{}); + defer c.deinit(); + c.agree = p.agree(); + try equal(cache.Outcome.eager, try c.run(1, no_stream, 7, e.body())); + try equal(@as(usize, 1), p.at); + try equal(@as(usize, 0), e.captures); + try equal(@as(usize, 1), e.bodies); +} + +test "graph cache regression: local OOM still contributes the room verdict" { + var e: Engine = .{}; + var failing = std.testing.FailingAllocator.init(std.testing.allocator, .{ .fail_index = 0 }); + var p: Peer = .{ .verdicts = &.{true} }; + var c = cache.Cache(u32).init(failing.allocator(), e.engine(), .{}); + defer c.deinit(); + c.agree = p.agree(); + try equal(cache.Outcome.eager, try c.run(1, no_stream, 7, e.body())); + try equal(@as(usize, 1), p.at); + try equal(@as(usize, 0), e.captures); +} + +test "graph cache regression: hits make no collective; a fingerprint change shared by every rank recaptures" { + var e: Engine = .{}; + var p: Peer = .{ .verdicts = &.{ true, true, true, true, true, true } }; + var c = cache.Cache(u32).init(std.testing.allocator, e.engine(), .{}); + defer c.deinit(); + c.agree = p.agree(); + _ = try c.run(1, no_stream, 7, e.body()); + _ = try c.run(2, no_stream, 7, e.body()); + try equal(@as(usize, 4), p.at); + for (0..4) |_| try equal(cache.Outcome.replayed, try c.run(1, no_stream, 7, e.body())); + try equal(@as(usize, 4), p.at); + // a new layout generation on every rank: both graphs go, the next miss agrees twice as any other + try equal(cache.Outcome.captured, try c.run(1, no_stream, 8, e.body())); + try equal(@as(usize, 6), p.at); + try equal(@as(u64, 2), c.stats.dropped); + try equal(@as(usize, 1), c.count()); + try equal(@as(usize, 1), e.live); +} + +var room = true; +fn available() bool { + return room; +} + +test "graph cache regression: zero maximum remains eager across pressure recovery, with no collective" { + var e: Engine = .{}; + var p: Peer = .{ .verdicts = &.{} }; + var c = cache.Cache(u32).init(std.testing.allocator, e.engine(), .{ .max = 0 }); + defer c.deinit(); + c.room = available; + c.agree = p.agree(); + for ([_]bool{ false, true, false, true }) |ok| { + room = ok; + try equal(cache.Outcome.eager, try c.run(1, no_stream, 7, e.body())); + try equal(@as(u32, 0), c.cap); + try equal(@as(usize, 0), c.count()); + } + try equal(@as(usize, 0), e.captures); + try equal(@as(usize, 0), p.at); +} + +test "graph cache regression: default floor cannot enlarge a two-entry maximum" { + var e: Engine = .{}; + var c = cache.Cache(u32).init(std.testing.allocator, e.engine(), .{ .max = 2 }); + defer c.deinit(); + c.room = available; + try equal(@as(u32, 2), c.settings.min_keep); + room = false; + _ = try c.run(1, no_stream, 7, e.body()); + room = true; + for (0..8) |key| { + _ = try c.run(@intCast(key), no_stream, 7, e.body()); + try expect(c.count() <= 2 and c.cap <= 2 and e.live <= 2); + } +} + +test "graph cache regression: zero keep floor recovers after reducing the cap to zero" { + var e: Engine = .{}; + var c = cache.Cache(u32).init(std.testing.allocator, e.engine(), .{ .max = 2, .min_keep = 0 }); + defer c.deinit(); + c.room = available; + room = false; + try equal(cache.Outcome.eager, try c.run(1, no_stream, 7, e.body())); + try equal(@as(u32, 0), c.cap); + room = true; + try equal(cache.Outcome.captured, try c.run(1, no_stream, 7, e.body())); + try equal(@as(u32, 2), c.cap); +} diff --git a/zig/src/cuda/launch.zig b/zig/src/cuda/launch.zig index 56bba5272a..47cbf720b6 100644 --- a/zig/src/cuda/launch.zig +++ b/zig/src/cuda/launch.zig @@ -4,7 +4,8 @@ const std = @import("std"); const abi = @import("abi.zig"); const Function = @import("module.zig").Function; const Stream = @import("stream.zig").Stream; -const Error = @import("driver.zig").Error; +const driver = @import("driver.zig"); +const Error = driver.Error; pub const Dim3 = abi.Dim3; @@ -88,7 +89,7 @@ pub fn launch(f: Function, cfg: Config, stream: Stream, args: *Args) Error!void const g = cfg.grid; const b = cfg.block; if (cfg.plain()) { - return d.check(d.api.cuLaunchKernel(f.handle, g.x, g.y, g.z, b.x, b.y, b.z, cfg.shared, stream.handle, args.pointers(), null), "cuLaunchKernel"); + return launched(f, cfg, d.api.cuLaunchKernel(f.handle, g.x, g.y, g.z, b.x, b.y, b.z, cfg.shared, stream.handle, args.pointers(), null), "cuLaunchKernel"); } var attrs: [4]abi.LaunchAttribute = undefined; const n = fillAttributes(cfg, &attrs); @@ -104,7 +105,37 @@ pub fn launch(f: Function, cfg: Config, stream: Stream, args: *Args) Error!void .attrs = &attrs, .num_attrs = @intCast(n), }; - try d.check(d.api.cuLaunchKernelEx(&lc, f.handle, args.pointers(), null), "cuLaunchKernelEx"); + try launched(f, cfg, d.api.cuLaunchKernelEx(&lc, f.handle, args.pointers(), null), "cuLaunchKernelEx"); +} + +/// A launch's result; a failure names the kernel and its configuration with the code. +fn launched(f: Function, cfg: Config, res: abi.Result, what: []const u8) Error!void { + if (res == abi.success) return; + if (res == abi.error_not_ready) return error.NotReady; // as Driver.check: not a failure to log + const d = f.d; + var buf: [160]u8 = undefined; + std.log.err("{s}: {s} ({d}) {s}; {s}{s}", .{ what, d.errorName(res), res, d.errorText(res), describe(&buf, d.functionName(f.handle), cfg), driver.shutdownNote(res) }); + return driver.Driver.errorOf(res); +} + +/// "kernel , grid XxYxZ, block XxYxZ, shared N B" (+ cluster / PDL / cooperative when set). +pub fn describe(buf: []u8, name: []const u8, cfg: Config) []const u8 { + const g = cfg.grid; + const b = cfg.block; + var w = std.Io.Writer.fixed(buf); + w.print("kernel {s}, grid {d}x{d}x{d}, block {d}x{d}x{d}, shared {d} B", .{ name, g.x, g.y, g.z, b.x, b.y, b.z, cfg.shared }) catch return w.buffered(); + if (cfg.cluster) |c| w.print(", cluster {d}x{d}x{d}", .{ c.x, c.y, c.z }) catch return w.buffered(); + if (cfg.pdl) w.writeAll(", pdl") catch return w.buffered(); + if (cfg.cooperative) w.writeAll(", cooperative") catch return w.buffered(); + return w.buffered(); +} + +test "a failed launch's description: the kernel, its grid, block and shared bytes, bounded" { + var buf: [160]u8 = undefined; + const cfg: Config = .{ .grid = .{ .x = 70000, .y = 2, .z = 1 }, .block = .{ .x = 128, .y = 1, .z = 1 }, .shared = 49152, .pdl = true }; + try std.testing.expectEqualStrings("kernel _stream_pf, grid 70000x2x1, block 128x1x1, shared 49152 B, pdl", describe(&buf, "_stream_pf", cfg)); + var small: [24]u8 = undefined; + try std.testing.expect(describe(&small, "a_very_long_kernel_symbol_name", cfg).len <= small.len); } fn attribute(id: abi.LaunchAttributeId) abi.LaunchAttribute { diff --git a/zig/src/cuda/memory.zig b/zig/src/cuda/memory.zig index 28c5b6244f..0f8f407806 100644 --- a/zig/src/cuda/memory.zig +++ b/zig/src/cuda/memory.zig @@ -100,18 +100,39 @@ fn lockCounts() void { while (!counts_mutex.tryLock()) std.Thread.yield() catch {}; } +/// A sub-allocator under DeviceBuffer (arena.zig): `alloc` gives an address or null for cuMemAlloc, `free` owns it. +pub const Hook = struct { + ctx: *anyopaque, + alloc: *const fn (ctx: *anyopaque, len: usize) ?abi.DevicePtr, + free: *const fn (ctx: *anyopaque, p: abi.DevicePtr) bool, // true when the address was the hook's +}; + +var alloc_hook: ?Hook = null; + +/// Installs, or with null removes, the sub-allocator: once at boot, before the buffers it should hold. +pub fn setHook(h: ?Hook) void { + alloc_hook = h; +} + +pub fn currentHook() ?Hook { + return alloc_hook; +} + pub const DeviceBuffer = struct { d: *const Driver, ptr: abi.DevicePtr, len: usize, borrowed: bool = false, // a span of memory another value owns (a carveout's): free leaves it alone - /// `len` bytes, 256-byte aligned by the driver; zero bytes allocate nothing and hold address 0. + /// `len` bytes, 256-byte aligned (512 from an installed arena); zero bytes allocate nothing and hold address 0. pub fn alloc(d: *const Driver, len: usize) Error!DeviceBuffer { try reserveAllocation(len); errdefer finishAllocation(null, len); var p: abi.DevicePtr = 0; - if (len > 0) try d.check(d.api.cuMemAlloc_v2(&p, len), "cuMemAlloc"); + if (len > 0) { + if (alloc_hook) |h| p = h.alloc(h.ctx, len) orelse 0; + if (p == 0) try d.check(d.api.cuMemAlloc_v2(&p, len), "cuMemAlloc"); + } finishAllocation(&device_bytes, len); return .{ .d = d, .ptr = p, .len = len }; } @@ -129,7 +150,10 @@ pub const DeviceBuffer = struct { pub fn free(self: *DeviceBuffer) void { if (!self.borrowed) { - if (self.ptr != 0) _ = self.d.api.cuMemFree_v2(self.ptr); + if (self.ptr != 0) { + const mine = if (alloc_hook) |h| h.free(h.ctx, self.ptr) else false; + if (!mine) _ = self.d.api.cuMemFree_v2(self.ptr); + } freed(&device_bytes, self.len); } self.* = undefined; @@ -266,6 +290,28 @@ test "freeing a borrowed span leaves the driver and the counts alone" { try std.testing.expectEqual(before.device, usage(false).device); } +test "a buffer the hook owns is freed by the hook, never by the driver" { + const Owner = struct { + var frees: u32 = 0; + fn alloc(_: *anyopaque, _: usize) ?abi.DevicePtr { + return null; + } + fn free(_: *anyopaque, p: abi.DevicePtr) bool { + frees += 1; + return p == 0x2000; + } + }; + var ctx: u8 = 0; + setHook(.{ .ctx = &ctx, .alloc = Owner.alloc, .free = Owner.free }); + defer setHook(null); + held(&device_bytes, 4096); + var b: DeviceBuffer = .{ .d = undefined, .ptr = 0x2000, .len = 4096 }; + b.free(); + var c: DeviceBuffer = .{ .d = undefined, .ptr = 0x3000, .len = 64, .borrowed = true }; + c.free(); + try std.testing.expectEqual(@as(u32, 1), Owner.frees); +} + test "peak reset preserves an allocation between its read and store" { const before = usage(true); held(&device_bytes, 100); diff --git a/zig/src/cuda/nccl.zig b/zig/src/cuda/nccl.zig index 420b2d460a..66fc50c3f1 100644 --- a/zig/src/cuda/nccl.zig +++ b/zig/src/cuda/nccl.zig @@ -39,6 +39,8 @@ pub const Config = extern struct { pub const Error = error{ LibraryUnavailable, MissingSymbol, NcclFailed }; +pub const in_progress: Result = 7; // ncclInProgress: a non-blocking communicator's call is still running + const R = Result; const D = abi.DevicePtr; diff --git a/zig/src/cuda/root.zig b/zig/src/cuda/root.zig index bd3ffa6845..0a52db2a17 100644 --- a/zig/src/cuda/root.zig +++ b/zig/src/cuda/root.zig @@ -54,6 +54,8 @@ pub const tp_group = @import("tp_group.zig"); pub const tp_watch = @import("tp_watch.zig"); pub const ck = @import("ck.zig"); pub const modelopt = @import("modelopt.zig"); +pub const graph_cache = @import("graph_cache.zig"); +pub const arena = @import("arena.zig"); pub const tree_attention = @import("tree_attention.zig"); pub const block_attention = @import("block_attention.zig"); pub const window_attention = @import("window_attention.zig"); @@ -100,6 +102,8 @@ test { _ = @import("modelopt.zig"); _ = @import("ck.zig"); _ = @import("tp_watch.zig"); + _ = graph_cache; + _ = arena; _ = @import("shared_test.zig"); _ = @import("tree_attention.zig"); _ = @import("block_attention.zig"); diff --git a/zig/src/fabric/tp/README.md b/zig/src/fabric/tp/README.md new file mode 100644 index 0000000000..7ad00ef7d5 --- /dev/null +++ b/zig/src/fabric/tp/README.md @@ -0,0 +1,72 @@ +# Tensor parallelism across processes + +`zig/src/fabric/tp` runs tensor parallelism with one GPU per process, either on one host or on two hosts over RoCE. +`TF_COMM_BACKEND` picks the transport for small exchanges: `nccl` (the default), `mailbox` (two GPUs on one host) +or `roce` (two hosts). NCCL carries everything larger, so it must work on its own as well. + +## Building `tf-tp-test` + +The `mailbox` and `roce` backends need the mailbox kernel (`mailbox.cu`) built into the binary: + +```sh +zig build tf-tp-test -Doptimize=ReleaseSafe -Dtp-nvcc=/usr/local/cuda/bin/nvcc +# or embed a fatbin built elsewhere: +zig build tf-tp-test -Doptimize=ReleaseSafe -Dtp-fatbin=/abs/path/tp_mailbox.fatbin +``` + +A build without `-Dtp-nvcc` or `-Dtp-fatbin` still runs `nccl` and `roce-host`, which needs no GPU. +In that build the `mailbox` and `roce` suites fail at load with `NoKernel` (an engine falls back to NCCL instead). +The `INFO` line `tf-tp-test` prints at start shows the embedded kernel's size; `0 B` means none was built in. + +## Choosing the RDMA devices + +| Variable | Used by | Needed when | +| --- | --- | --- | +| `TF_TP_ROCE_HCA` | the RoCE wire | the host has more than one active RoCE port that is not cabled to the peer | +| `NCCL_IB_HCA` | NCCL | the same as above; NCCL otherwise picks its own devices | +| `NCCL_SOCKET_IFNAME` | NCCL bootstrap | the host has several IP interfaces and NCCL picks one the peer cannot reach | +| `TF_TP_MASTER` | the TCP bootstrap | always on two hosts: rank 0's address as the other rank reaches it | + +`TF_TP_ROCE_HCA` is `name[:gid],name[:gid]`: up to two devices, and optionally each one's RoCE v2 GID index. +Without it the rank picks the devices itself, at most `TF_TP_ROCE_HCAS` (default 2). +A host cabled to more than one neighbour should always set it explicitly. +The two hosts may list their devices in a different order and use different names. +Each device pairs with the peer device on the same IPv4 /24. + +`NCCL_IB_HCA` should name the same devices without the `:gid` suffix. +In NCCL, `:N` after a device name selects port `N`, not a GID index. + +`run_hosts.sh` sets `NCCL_IB_HCA` from `TF_TP_ROCE_HCA` when only that is set. +When neither is set, it first runs `tf-tp-test hca` on both ranks: the RoCE pick, made with the peer over the bootstrap. +The same rule picks the wire's devices, and each rank prints its own as `NCCL_IB_HCA=name,...`. +`hca` needs no GPU, and it exits 1 when no port shares a /24 with the peer, so no host falls back to its first ports. +The runner logs each rank's pick and stops with that reason when `hca` fails. +Both hosts must use the same mode: if one names its ports and the other names none, only the latter runs `hca`. +It finds no peer and stops after `TF_TP_PLAN_TIMEOUT` (default 300 s) with that reason. + +Two DGX Spark nodes, with the second ConnectX-7 port of each cabled to the other (both PCIe functions of that +port), and the first port cabled to a third node: + +```sh +# on both hosts; rank 0's link address is 192.0.2.1 +export TF_TP_ROCE_HCA=rocep1s0f1,roceP2p1s0f1 +export NCCL_IB_HCA=rocep1s0f1,roceP2p1s0f1 +export NCCL_SOCKET_IFNAME=enp1s0f1np1 +zig/src/fabric/tp/run_hosts.sh 0 192.0.2.1 ./zig-out/bin/tf-tp-test out-rank0 # host A +zig/src/fabric/tp/run_hosts.sh 1 192.0.2.1 ./zig-out/bin/tf-tp-test out-rank1 # host B +``` + +`ibv_devices` lists the RDMA devices; `ibv_devinfo -d NAME` shows a port's state and link layer. +`/sys/class/infiniband/NAME/ports/1/gids/N` and `gid_attrs/ndevs/N` show each GID and its netdev. +`ulimit -l` must allow the RoCE region to be pinned (`unlimited` is simplest). + +## RoCE limits and the failed-transport check + +No RDMA write is larger than 1 MiB, so on `TF_COMM_BACKEND=roce` a `TF_TP_MAILBOX_MAX_KB` above 1024 fails at boot. +`TF_TP_ROCE_CHUNKED_KB` (default 0, off) sends variable-length gathers past the slot over RoCE in slot-sized ops. +It needs a kernel built from this tree's `mailbox.cu`. +At most 2 MiB is posted and not yet completed at any time. +A rank whose transport fails moves its queue pairs to the error state and keeps its region until the process exits. + +`run_transport.sh MASTER_ADDR PEER_SSH OUTDIR` checks that failure path on two hosts with `tf-tp-test failfast V +transport` and stops at the first kernel log line or HCA error counter that grows; run it with no arguments for help. diff --git a/zig/src/fabric/tp/bootstrap.zig b/zig/src/fabric/tp/bootstrap.zig new file mode 100644 index 0000000000..e1a6911efb --- /dev/null +++ b/zig/src/fabric/tp/bootstrap.zig @@ -0,0 +1,311 @@ +//! Rank discovery over TCP (a star on rank 0): small host blobs all-gathered, settings agreed, sockets kept for fate. +const std = @import("std"); +const log = @import("log.zig"); +const sock = @import("sock.zig"); +const Config = @import("config.zig").Config; + +pub const max_ranks = 64; +/// Largest blob a rank contributes to one `allGather`. +pub const max_blob = 64 * 1024; + +const magic: u32 = 0x31425054; // "TPB1" +const version: u32 = 2; + +pub const Error = error{ Protocol, Mismatch, TooLarge } || sock.Error; + +/// `role`: 0 the control connection (bootstrap, then the fate channel), 1 the plan link. +const Hello = extern struct { magic: u32, version: u32, world: u32, rank: u32, role: u32 }; + +pub const Bootstrap = struct { + rank: u32, + world: u32, + /// Rank 0: the socket of rank i at index i (index 0 unused); other ranks: rank 0's socket at index 0. + fds: [max_ranks]sock.Fd = @splat(-1), + /// The plan link's sockets, laid out as `fds`. + plan_fds: [max_ranks]sock.Fd = @splat(-1), + /// Each call's limit for the peers' bytes (the connect timeout). + timeout_ns: u64, + + fn deadline(b: *const Bootstrap) u64 { + return sock.nowNs() + b.timeout_ns; + } + + /// Rank 0 listens on `link_port` and accepts every other rank; the others connect to `master` (connect timeout). + pub fn open(cfg: Config) Error!Bootstrap { + if (cfg.world > max_ranks or cfg.rank >= cfg.world) return error.Protocol; + var b: Bootstrap = .{ .rank = cfg.rank, .world = cfg.world, .timeout_ns = cfg.connect_timeout_ns }; + errdefer b.close(); + if (cfg.world == 1) return b; + if (cfg.rank == 0) { + const l = try sock.listen(cfg.link_port); + defer sock.close(l); + try b.acceptAll(l); + } else { + const until = b.deadline(); + for ([_]u32{ 0, 1 }) |role| { + const fd = try sock.connect(cfg.master, cfg.link_port, until); + if (role == 0) b.fds[0] = fd else b.plan_fds[0] = fd; + const h: Hello = .{ .magic = magic, .version = version, .world = cfg.world, .rank = cfg.rank, .role = role }; + try sock.sendAll(fd, std.mem.asBytes(&h)); + } + } + return b; + } + + /// Rank 0 over an already listening socket (tests: an ephemeral port). + pub fn acceptOn(l: sock.Fd, world: u32, timeout_ns: u64) Error!Bootstrap { + var b: Bootstrap = .{ .rank = 0, .world = world, .timeout_ns = timeout_ns }; + errdefer b.close(); + try b.acceptAll(l); + return b; + } + + fn acceptAll(b: *Bootstrap, l: sock.Fd) Error!void { + // every other rank's control and plan connections, in any order + var joined: u32 = 0; + const until = b.deadline(); + while (joined < 2 * (b.world - 1)) { + const fd = try sock.accept(l, until); + var h: Hello = undefined; + sock.recvExact(fd, std.mem.asBytes(&h), until) catch { + sock.close(fd); + continue; + }; + if (h.magic != magic or h.version != version) { + log.warn("tp bootstrap: a connection that is not a rank (magic {x}); ignored", .{h.magic}); + sock.close(fd); + continue; + } + const table = if (h.role == 0) &b.fds else &b.plan_fds; + if (h.world != b.world or h.rank == 0 or h.rank >= b.world or h.role > 1 or table[h.rank] >= 0) { + log.warn("tp bootstrap: rank {d} of {d} joined a world of {d}", .{ h.rank, h.world, b.world }); + sock.close(fd); + return error.Mismatch; + } + table[h.rank] = fd; + joined += 1; + } + } + + /// Takes the control sockets (the fate channel keeps them); `close` then leaves them open. + pub fn release(b: *Bootstrap) [max_ranks]sock.Fd { + const out = b.fds; + b.fds = @splat(-1); + return out; + } + + /// Takes the plan link's sockets. + pub fn releasePlan(b: *Bootstrap) [max_ranks]sock.Fd { + const out = b.plan_fds; + b.plan_fds = @splat(-1); + return out; + } + + pub fn close(b: *Bootstrap) void { + for ([_]*[max_ranks]sock.Fd{ &b.fds, &b.plan_fds }) |table| for (table) |*fd| if (fd.* >= 0) { + sock.close(fd.*); + fd.* = -1; + }; + } + + /// Every rank's `mine` in rank order into `out`, returned as slices of `out`; every rank calls it equally often. + pub fn allGather(b: *Bootstrap, mine: []const u8, out: []u8, blobs: [][]const u8) Error![]const []const u8 { + if (mine.len > max_blob) return error.TooLarge; + std.debug.assert(blobs.len >= b.world); + if (b.world == 1) { + if (out.len < mine.len) return error.TooLarge; + @memcpy(out[0..mine.len], mine); + blobs[0] = out[0..mine.len]; + return blobs[0..1]; + } + if (b.rank != 0) { + try sendFrame(b.fds[0], mine); + return b.readAll(b.fds[0], out, blobs); + } + // rank 0: everyone's blob, then the whole table to everyone + var used: usize = 0; + for (0..b.world) |i| { + const len: u32 = if (i == 0) @intCast(mine.len) else blk: { + var l: u32 = undefined; + try sock.recvExact(b.fds[i], std.mem.asBytes(&l), b.deadline()); + break :blk l; + }; + if (len > max_blob or used + len > out.len) return error.TooLarge; + if (i == 0) @memcpy(out[used..][0..len], mine) else try sock.recvExact(b.fds[i], out[used..][0..len], b.deadline()); + blobs[i] = out[used..][0..len]; + used += len; + } + for (1..b.world) |i| { + for (blobs[0..b.world]) |blob| try sendFrame(b.fds[i], blob); + } + return blobs[0..b.world]; + } + + fn readAll(b: *Bootstrap, fd: sock.Fd, out: []u8, blobs: [][]const u8) Error![]const []const u8 { + var used: usize = 0; + for (0..b.world) |i| { + var l: u32 = undefined; + try sock.recvExact(fd, std.mem.asBytes(&l), b.deadline()); + if (l > max_blob or used + l > out.len) return error.TooLarge; + try sock.recvExact(fd, out[used..][0..l], b.deadline()); + blobs[i] = out[used..][0..l]; + used += l; + } + return blobs[0..b.world]; + } + + /// Rank 0's `bytes` on every rank (into `out` on the others; `bytes` is ignored there). + pub fn broadcast(b: *Bootstrap, bytes: []const u8, out: []u8) Error!void { + var blobs: [max_ranks][]const u8 = undefined; + var scratch: [4096]u8 = undefined; + if (out.len * b.world > scratch.len) return error.TooLarge; + const all = try b.allGather(if (b.rank == 0) bytes else &.{}, &scratch, &blobs); + if (all[0].len != out.len) return error.Protocol; + @memcpy(out, all[0]); + } + + /// Every rank's `bytes` must be equal (settings that change bytes on the wire); `what` names them in the log. + pub fn agree(b: *Bootstrap, bytes: []const u8, what: []const u8) Error!void { + var blobs: [max_ranks][]const u8 = undefined; + var scratch: [8192]u8 = undefined; + const all = try b.allGather(bytes, &scratch, &blobs); + for (all, 0..) |blob, i| if (!std.mem.eql(u8, blob, all[0])) { + log.warn("tp bootstrap: rank {d}'s {s} differs from rank 0's; both ranks refuse to start", .{ i, what }); + return error.Mismatch; + }; + } +}; + +fn sendFrame(fd: sock.Fd, bytes: []const u8) sock.Error!void { + const l: u32 = @intCast(bytes.len); + try sock.sendAll(fd, std.mem.asBytes(&l)); + try sock.sendAll(fd, bytes); +} + +/// An ephemeral listening port for tests. +pub fn testListener() !struct { fd: sock.Fd, port: u16 } { + const l = try sock.listen(0); + var addr: std.posix.sockaddr.in = undefined; + var len: std.posix.socklen_t = @sizeOf(std.posix.sockaddr.in); + _ = std.c.getsockname(l, @ptrCast(&addr), &len); + return .{ .fd = l, .port = std.mem.bigToNative(u16, addr.port) }; +} + +fn testRank(cfg: Config, out: *[3][32]u8, err: *?anyerror) void { + var b = Bootstrap.open(cfg) catch |e| { + err.* = e; + return; + }; + defer b.close(); + var id: [32]u8 = undefined; + b.broadcast(&@as([32]u8, @splat(0xAB)), &id) catch |e| { + err.* = e; + return; + }; + var blobs: [max_ranks][]const u8 = undefined; + var buf: [256]u8 = undefined; + const mine: [3]u8 = @splat(@intCast(cfg.rank)); + const all = b.allGather(mine[0 .. cfg.rank + 1], &buf, &blobs) catch |e| { + err.* = e; + return; + }; + for (all, 0..) |blob, i| @memcpy(out[i][0..blob.len], blob); + b.agree("same", "settings") catch |e| { + err.* = e; + }; +} + +test "three ranks gather in rank order and agree" { + const l = try testListener(); + defer sock.close(l.fd); + var outs: [3][3][32]u8 = undefined; + var errs: [3]?anyerror = @splat(null); + var threads: [2]std.Thread = undefined; + for (0..2) |i| threads[i] = try std.Thread.spawn(.{}, testRank, .{ Config{ .world = 3, .rank = @intCast(i + 1), .link_port = l.port, .connect_timeout_ns = 5 * std.time.ns_per_s }, &outs[i + 1], &errs[i + 1] }); + var b = try Bootstrap.acceptOn(l.fd, 3, 5 * std.time.ns_per_s); + defer b.close(); + var id: [32]u8 = undefined; + try b.broadcast(&@as([32]u8, @splat(0xAB)), &id); + var blobs: [max_ranks][]const u8 = undefined; + var buf: [256]u8 = undefined; + const all = try b.allGather(&.{0}, &buf, &blobs); + try b.agree("same", "settings"); + for (threads) |t| t.join(); + for (errs) |e| try std.testing.expect(e == null); + try std.testing.expectEqual(@as(usize, 3), all.len); + try std.testing.expectEqualSlices(u8, &.{ 2, 2, 2 }, all[2]); + try std.testing.expectEqualSlices(u8, &.{ 1, 1 }, outs[2][1][0..2]); +} + +test "settings that differ stop every rank" { + const l = try testListener(); + defer sock.close(l.fd); + const T = struct { + fn go(port: u16, err: *?anyerror) void { + var b = Bootstrap.open(.{ .world = 2, .rank = 1, .link_port = port, .connect_timeout_ns = 5 * std.time.ns_per_s }) catch |e| { + err.* = e; + return; + }; + defer b.close(); + b.agree("knobs-b", "settings") catch |e| { + err.* = e; + }; + } + }; + var err: ?anyerror = null; + const t = try std.Thread.spawn(.{}, T.go, .{ l.port, &err }); + var b = try Bootstrap.acceptOn(l.fd, 2, 5 * std.time.ns_per_s); + defer b.close(); + try std.testing.expectError(error.Mismatch, b.agree("knobs-a", "settings")); + t.join(); + try std.testing.expectEqual(@as(?anyerror, error.Mismatch), err); +} + +/// Rank 1 agrees on `theirs.digest()` against rank 0's `ours.digest()`; returns both ranks' results. +fn agreeDigests(ours: Config, theirs: Config) !struct { ?anyerror, ?anyerror } { + const l = try testListener(); + defer sock.close(l.fd); + const T = struct { + fn go(port: u16, cfg: Config, err: *?anyerror) void { + var c = cfg; + c.world = 2; + c.rank = 1; + c.link_port = port; + c.connect_timeout_ns = 5 * std.time.ns_per_s; + var b = Bootstrap.open(c) catch |e| { + err.* = e; + return; + }; + defer b.close(); + b.agree(&c.digest(), "settings") catch |e| { + err.* = e; + }; + } + }; + var err1: ?anyerror = null; + const t = try std.Thread.spawn(.{}, T.go, .{ l.port, theirs, &err1 }); + var b = try Bootstrap.acceptOn(l.fd, 2, 5 * std.time.ns_per_s); + defer b.close(); + var err0: ?anyerror = null; + b.agree(&ours.digest(), "settings") catch |e| { + err0 = e; + }; + t.join(); + return .{ err0, err1 }; +} + +test "a graph cache switched off on one rank stops both at start; 0 and OFF on the two ranks start" { + const on: Config = .{ .graph_cache = .{} }; + const off: Config = .{ .graph_cache = .{}, .graph_cache_env = "off" }; + const r = try agreeDigests(on, off); + try std.testing.expectEqual(@as(?anyerror, error.Mismatch), r[0]); + try std.testing.expectEqual(@as(?anyerror, error.Mismatch), r[1]); + const none = try agreeDigests(.{}, .{ .graph_cache_env = "off" }); // a family without the cache is not refused + try std.testing.expectEqual(@as(?anyerror, null), none[0]); + try std.testing.expectEqual(@as(?anyerror, null), none[1]); + const same = try agreeDigests(.{ .graph_cache = .{}, .graph_cache_env = "0" }, .{ .graph_cache = .{}, .graph_cache_env = "OFF" }); + try std.testing.expectEqual(@as(?anyerror, null), same[0]); + try std.testing.expectEqual(@as(?anyerror, null), same[1]); + const family = try agreeDigests(.{ .graph_cache = .{ .max = 8 } }, .{ .graph_cache = .{ .max = 16 } }); + try std.testing.expectEqual(@as(?anyerror, error.Mismatch), family[0]); +} diff --git a/zig/src/fabric/tp/collective.zig b/zig/src/fabric/tp/collective.zig new file mode 100644 index 0000000000..d2334dc74a --- /dev/null +++ b/zig/src/fabric/tp/collective.zig @@ -0,0 +1,136 @@ +//! The tensor-parallel communicator every backend implements: stream-ordered, graph-capturable, rank-order sums. +const std = @import("std"); +const cuda = @import("cuda"); + +pub const DevicePtr = cuda.abi.DevicePtr; +pub const Stream = cuda.abi.Stream; + +pub const Error = error{ Aborted, Unsupported, Invalid, PeerFailed, Timeout, BackendFailed }; + +/// Element types the collectives move (`reduce` takes the float ones and the 32/64-bit integers). +pub const DType = enum(u8) { + u8, + i32, + i64, + f16, + bf16, + f32, + f64, + + pub fn size(t: DType) usize { + return switch (t) { + .u8 => 1, + .f16, .bf16 => 2, + .i32, .f32 => 4, + .i64, .f64 => 8, + }; + } +}; + +pub const Op = enum(u8) { sum, max, min }; + +/// Which implementation is behind a `Collective` (logs, bench labels). +pub const Kind = enum { nccl, mailbox, hybrid, roce, host }; + +pub const Collective = struct { + ptr: *anyopaque, + vtable: *const VTable, + + pub const VTable = struct { + kind: Kind, + rank: *const fn (ptr: *anyopaque) u32, + world: *const fn (ptr: *anyopaque) u32, + /// recv[count] <- the element-wise reduction of every rank's send[count], added in rank order. + all_reduce: *const fn (ptr: *anyopaque, send: DevicePtr, recv: DevicePtr, count: usize, t: DType, op: Op, stream: Stream) Error!void, + /// recv[world * count] <- every rank's send[count], rank 0's first (NCCL's layout). + all_gather: *const fn (ptr: *anyopaque, send: DevicePtr, recv: DevicePtr, count: usize, t: DType, stream: Stream) Error!void, + /// Two ranks trade: this rank's send[count] lands in `peer`'s recv[count] and the peer's in ours (one group). + exchange: *const fn (ptr: *anyopaque, send: DevicePtr, recv: DevicePtr, count: usize, t: DType, peer: u32, stream: Stream) Error!void, + send: *const fn (ptr: *anyopaque, buf: DevicePtr, count: usize, t: DType, peer: u32, stream: Stream) Error!void, + recv: *const fn (ptr: *anyopaque, buf: DevicePtr, count: usize, t: DType, peer: u32, stream: Stream) Error!void, + /// Host: returns once every rank has called it and this rank's earlier collectives have finished. + barrier: *const fn (ptr: *anyopaque) Error!void, + /// Host: a failure the transport recorded since the last check (a peer that timed out, an NCCL async error). + check: *const fn (ptr: *anyopaque) Error!void, + /// Any thread: end pending and later work at once; safe while another thread is inside a call. + abort: *const fn (ptr: *anyopaque) void, + /// recv[world][stride]: rank r's first lens[r] bytes (device int32 [world]) land at recv + r * stride. + all_gather_v: ?*const fn (ptr: *anyopaque, send: DevicePtr, recv: DevicePtr, stride: usize, lens: DevicePtr, t: DType, stream: Stream) Error!void = null, + /// Whether `all_gather_v` moves only the lengths for shards of `bytes` here (rank-local). Null: never. + var_fits: ?*const fn (ptr: *anyopaque, bytes: usize) bool = null, + /// Host: whether every rank's transport takes device lengths, agreed once per communicator. Null: false. + var_agreed: ?*const fn (ptr: *anyopaque) Error!bool = null, + }; + + pub fn kind(c: Collective) Kind { + return c.vtable.kind; + } + + pub fn rank(c: Collective) u32 { + return c.vtable.rank(c.ptr); + } + + pub fn world(c: Collective) u32 { + return c.vtable.world(c.ptr); + } + + pub fn allReduce(c: Collective, src: DevicePtr, dst: DevicePtr, count: usize, t: DType, op: Op, stream: Stream) Error!void { + return c.vtable.all_reduce(c.ptr, src, dst, count, t, op, stream); + } + + pub fn allGather(c: Collective, src: DevicePtr, dst: DevicePtr, count: usize, t: DType, stream: Stream) Error!void { + return c.vtable.all_gather(c.ptr, src, dst, count, t, stream); + } + + /// `allGather` of `stride` elements a rank where only rank r's first lens[r] bytes matter (`VTable.all_gather_v`). + pub fn allGatherV(c: Collective, src: DevicePtr, dst: DevicePtr, stride: usize, lens: DevicePtr, t: DType, stream: Stream) Error!void { + if (c.vtable.all_gather_v) |f| return f(c.ptr, src, dst, stride, lens, t, stream); + return c.vtable.all_gather(c.ptr, src, dst, stride, t, stream); + } + + /// Whether `allGatherV` of `bytes`-byte strides moves only the lengths here (callers agree first: `varAgreed`). + pub fn varFits(c: Collective, bytes: usize) bool { + const f = c.vtable.var_fits orelse return false; + return f(c.ptr, bytes); + } + + /// Every rank's transport takes device lengths: one agreement a communicator, called alike outside any graph. + pub fn varAgreed(c: Collective) Error!bool { + const f = c.vtable.var_agreed orelse return false; + return f(c.ptr); + } + + pub fn exchange(c: Collective, src: DevicePtr, dst: DevicePtr, count: usize, t: DType, peer: u32, stream: Stream) Error!void { + return c.vtable.exchange(c.ptr, src, dst, count, t, peer, stream); + } + + pub fn send(c: Collective, buf: DevicePtr, count: usize, t: DType, peer: u32, stream: Stream) Error!void { + return c.vtable.send(c.ptr, buf, count, t, peer, stream); + } + + pub fn recv(c: Collective, buf: DevicePtr, count: usize, t: DType, peer: u32, stream: Stream) Error!void { + return c.vtable.recv(c.ptr, buf, count, t, peer, stream); + } + + pub fn barrier(c: Collective) Error!void { + return c.vtable.barrier(c.ptr); + } + + pub fn check(c: Collective) Error!void { + return c.vtable.check(c.ptr); + } + + pub fn abort(c: Collective) void { + c.vtable.abort(c.ptr); + } +}; + +/// Refuses a peer that is this rank or out of range. +pub fn checkPeer(rank: u32, world: u32, peer: u32) Error!void { + if (peer == rank or peer >= world) return error.Invalid; +} + +test "dtype sizes" { + try std.testing.expectEqual(@as(usize, 2), DType.bf16.size()); + try std.testing.expectEqual(@as(usize, 8), DType.i64.size()); +} diff --git a/zig/src/fabric/tp/config.zig b/zig/src/fabric/tp/config.zig new file mode 100644 index 0000000000..6f96091c1c --- /dev/null +++ b/zig/src/fabric/tp/config.zig @@ -0,0 +1,288 @@ +//! A rank's TP settings from the environment: rank, master, port, backend, fail-fast knobs. +const std = @import("std"); +const log = @import("log.zig"); +const graph_cache = @import("cuda").graph_cache; + +pub const Backend = enum { + /// NCCL for everything (the default). + nccl, + /// Exchanges of at most `mailbox_max_bytes` a rank through the one-shot mailbox kernel, NCCL above that. + mailbox, + /// The same split between two hosts: the small exchanges over RoCE (doorbell + RDMA proxy), NCCL above. + roce, +}; + +/// What `roce` does when its setup or load-time probe fails on any rank (TF_TP_ROCE_FALLBACK). +pub const Fallback = enum { nccl, @"error" }; + +pub const Error = error{BadSetting}; + +/// RoCE's per-write cap (roce_guard.max_write): on the roce backend a slot past it fails boot, never clamped. +pub const max_rdma_write: usize = @import("roce_guard.zig").max_write; + +/// TF_TP_PLAN_PIN: the CPU the plan threads pin to at their first plan (auto: the highest cpu_capacity). +pub const PlanPin = union(enum) { auto, cpu: u32 }; + +pub const Config = struct { + world: u32 = 2, + rank: u32 = 0, + /// Rank 0's address as the other ranks reach it. + master: []const u8 = "localhost", + /// The rendezvous port (TF_TP_PORT); the link listens on `link_port`. + port: u16 = 29500, + /// Rank 0 listens here: TF_TP_PLAN_PORT, else `port` + 7. + link_port: u16 = 29507, + /// How long rank 0 waits for the others and they keep retrying the connect (TF_TP_PLAN_TIMEOUT). + connect_timeout_ns: u64 = 300 * std.time.ns_per_s, + /// CUDA device ordinal this rank uses. + device: u32 = 0, + backend: Backend = .nccl, + mailbox_max_bytes: usize = 256 * 1024, + /// TF_TP_ROCE_CHUNKED_KB: `allGatherV` strides up to this go RoCE in ops of at most a slot (0: off). + roce_chunked_bytes: usize = 0, + /// A mailbox wait for the peer gives up after this long and poisons the transport (fail-stop). + mailbox_timeout_ns: u64 = 120 * std.time.ns_per_s, + failfast: bool = true, + /// After a failure: time for the HTTP threads to answer before exit 70. + grace_ns: u64 = std.time.ns_per_s, + /// Rank 1's memory report period on the fate channel (also how fast it notices rank 0 hanging up). + report_ns: u64 = std.time.ns_per_s / 2, + nccl_lib: []const u8 = "libnccl.so.2", + /// `name[:gid],...` RDMA devices for `roce` (null: active RoCE v2 ports on an IPv4 /24 the peer has, roce_pick). + roce_hca: ?[]const u8 = null, + /// Stripe over at most this many HCAs (auto detection only). + roce_hcas: u8 = 2, + roce_tc: u8 = 0, + /// TF_TP_ROCE_FAST as the RoCE kernel's `opts` word: bit 0 on, bits 8-15 pollers, bits 16-31 stagger ns (local). + roce_fast: u32 = 0, + /// TF_TP_ROCE_BLOCKS: the largest RoCE grid, a power of two 1..16; launch shape only, the bytes are the same. + roce_blocks: u32 = 16, + /// Pin the proxy thread (it busy-spins one core between exchanges). + roce_cpu: ?u32 = null, + /// `shm` runs the RoCE protocol on one host without RDMA (tests); `verbs` is the real wire. + roce_wire: enum { verbs, shm } = .verbs, + roce_fallback: Fallback = .nccl, + /// TF_TP_PLAN_LINK=rdma: the follower polls the plan socket this long before blocking (TF_TP_PLAN_SPIN_US). + plan_spin_ns: u64 = 0, + plan_pin: ?PlanPin = null, + /// The keyed graph cache settings a family builds with, set before `Session.open` (null: the family has none). + graph_cache: ?graph_cache.Settings = null, + /// TF_GRAPH_CACHE as read (`Cache.init` applies the same value); only its effect on `graph_cache` is compared. + graph_cache_env: ?[]const u8 = null, + + /// Reads every setting from `get` (the environment in `fromEnv`); unset names keep the defaults. + pub fn parse(get: *const fn ([]const u8) ?[]const u8) Error!Config { + var c: Config = .{}; + c.world = try int(u32, get, "TF_TP_WORLD", c.world, 1, 64); + c.rank = try int(u32, get, "TF_TP_RANK", c.rank, 0, c.world - 1); + if (get("TF_TP_MASTER")) |m| c.master = m; + c.port = try int(u16, get, "TF_TP_PORT", c.port, 1, 65535 - 8); + c.link_port = try int(u16, get, "TF_TP_PLAN_PORT", c.port + 7, 1, 65535); + c.connect_timeout_ns = try seconds(get, "TF_TP_PLAN_TIMEOUT", 300, 0.1, 86400); + c.device = try int(u32, get, "TF_TP_DEVICE", c.device, 0, 63); + const backend = "TF_COMM_BACKEND"; + if (get(backend)) |b| c.backend = std.meta.stringToEnum(Backend, std.mem.trim(u8, b, " ")) orelse { + log.warn("{s}={s}: expected nccl, mailbox or roce", .{ backend, b }); + return error.BadSetting; + }; + c.mailbox_max_bytes = 1024 * try int(usize, get, "TF_TP_MAILBOX_MAX_KB", 256, 4, 16 * 1024); + if (c.backend == .roce and c.mailbox_max_bytes > max_rdma_write) { + log.warn("TF_TP_MAILBOX_MAX_KB={d}: a RoCE write is capped at {d} KiB; refusing to start", .{ c.mailbox_max_bytes / 1024, max_rdma_write / 1024 }); + return error.BadSetting; + } + c.roce_chunked_bytes = 1024 * try int(usize, get, "TF_TP_ROCE_CHUNKED_KB", 0, 0, 64 * 1024); + c.mailbox_timeout_ns = try seconds(get, "TF_TP_MAILBOX_TIMEOUT_S", 120, 0.001, 86400); + if (get("TF_TP_FAILFAST")) |v| c.failfast = !(std.mem.eql(u8, v, "0") or std.mem.eql(u8, v, "off")); + c.grace_ns = try seconds(get, "TF_TP_FAILFAST_GRACE_S", 1.0, 0, 60); + c.report_ns = try seconds(get, "TF_TP_PEER_MEM_S", 0.5, 0.01, 60); + if (get("TF_NCCL_LIB")) |p| c.nccl_lib = p; + c.roce_hca = get("TF_TP_ROCE_HCA"); + c.roce_hcas = try int(u8, get, "TF_TP_ROCE_HCAS", 2, 1, 2); + c.roce_tc = try int(u8, get, "TF_TP_ROCE_TC", try int(u8, get, "NCCL_IB_TC", 0, 0, 255), 0, 255); + if (try int(u8, get, "TF_TP_ROCE_FAST", 0, 0, 1) == 1) { + const pollers = try int(u32, get, "TF_TP_ROCE_POLLERS", 1, 1, 8); + const stagger = try int(u32, get, "TF_TP_ROCE_STAGGER_NS", 150, 0, 65535); + c.roce_fast = 1 | (pollers << 8) | (stagger << 16); + } + c.roce_blocks = try int(u32, get, "TF_TP_ROCE_BLOCKS", 16, 1, 16); + if (c.roce_blocks & (c.roce_blocks - 1) != 0) { + log.warn("TF_TP_ROCE_BLOCKS={d}: expected a power of two", .{c.roce_blocks}); + return error.BadSetting; + } + if (get("TF_TP_ROCE_CPU")) |_| c.roce_cpu = try int(u32, get, "TF_TP_ROCE_CPU", 0, 0, 4095); + if (get("TF_TP_ROCE_WIRE")) |w| c.roce_wire = std.meta.stringToEnum(@TypeOf(c.roce_wire), w) orelse { + log.warn("TF_TP_ROCE_WIRE={s}: expected verbs or shm", .{w}); + return error.BadSetting; + }; + const fallback = "TF_TP_ROCE_FALLBACK"; + if (get(fallback)) |f| c.roce_fallback = std.meta.stringToEnum(Fallback, f) orelse { + log.warn("{s}={s}: expected nccl or error", .{ fallback, f }); + return error.BadSetting; + }; + if (get("TF_TP_PLAN_LINK")) |l| { + const mode = std.mem.trim(u8, l, " "); + if (std.mem.eql(u8, mode, "rdma")) { + c.plan_spin_ns = 1000 * try int(u64, get, "TF_TP_PLAN_SPIN_US", 5000, 0, 10_000_000); + } else if (!std.mem.eql(u8, mode, "tcp") and !std.mem.eql(u8, mode, "nccl")) { + log.warn("TF_TP_PLAN_LINK={s}: expected nccl, tcp or rdma", .{l}); + return error.BadSetting; + } + } + if (get("TF_TP_PLAN_PIN")) |raw| { + const v = std.mem.trim(u8, raw, " "); + if (std.mem.eql(u8, v, "auto")) { + c.plan_pin = .auto; + } else if (!std.mem.eql(u8, v, "-1") and !std.mem.eql(u8, v, "off")) { + c.plan_pin = .{ .cpu = try int(u32, get, "TF_TP_PLAN_PIN", 0, 0, 4095) }; + } + } + c.graph_cache_env = get("TF_GRAPH_CACHE"); + return c; + } + + pub fn fromEnv() Error!Config { + return parse(&envGet); + } + + /// The settings both ranks must share (everything but rank, device and timeouts), as bytes to compare. + pub fn digest(c: Config) [32]u8 { + var h = std.crypto.hash.sha2.Sha256.init(.{}); + h.update(std.mem.asBytes(&c.world)); + h.update(std.mem.asBytes(&c.backend)); + h.update(std.mem.asBytes(&c.mailbox_max_bytes)); + h.update(std.mem.asBytes(&c.roce_chunked_bytes)); + h.update(std.mem.asBytes(&c.failfast)); + h.update(std.mem.asBytes(&c.roce_wire)); + h.update(std.mem.asBytes(&c.roce_fallback)); + if (c.graph_cache) |g| h.update(&g.effective(c.graph_cache_env).digest()); + var out: [32]u8 = undefined; + h.final(&out); + return out; + } +}; + +fn envGet(name: []const u8) ?[]const u8 { + var buf: [64]u8 = undefined; + if (name.len >= buf.len) return null; + @memcpy(buf[0..name.len], name); + buf[name.len] = 0; + const v = std.c.getenv(buf[0..name.len :0]) orelse return null; + const s = std.mem.span(v); + return if (s.len == 0) null else s; +} + +fn int(comptime T: type, get: *const fn ([]const u8) ?[]const u8, name: []const u8, default: T, lo: T, hi: T) Error!T { + const raw = get(name) orelse return default; + const v = std.fmt.parseInt(T, std.mem.trim(u8, raw, " "), 10) catch { + log.warn("{s}={s}: expected an integer", .{ name, raw }); + return error.BadSetting; + }; + if (v < lo or v > hi) { + log.warn("{s}={d}: expected {d}..{d}", .{ name, v, lo, hi }); + return error.BadSetting; + } + return v; +} + +fn seconds(get: *const fn ([]const u8) ?[]const u8, name: []const u8, default: f64, lo: f64, hi: f64) Error!u64 { + const v = if (get(name)) |raw| std.fmt.parseFloat(f64, std.mem.trim(u8, raw, " ")) catch { + log.warn("{s}={s}: expected seconds", .{ name, raw }); + return error.BadSetting; + } else default; + if (!(v >= lo and v <= hi)) { + log.warn("{s}={d}: expected {d}..{d} s", .{ name, v, lo, hi }); + return error.BadSetting; + } + return @intFromFloat(v * std.time.ns_per_s); +} + +const Fake = struct { + var pairs: []const [2][]const u8 = &.{}; + fn get(name: []const u8) ?[]const u8 { + for (pairs) |p| if (std.mem.eql(u8, p[0], name)) return p[1]; + return null; + } +}; + +test "defaults: the plan link 7 ports past the rendezvous, NCCL, fail-fast on" { + Fake.pairs = &.{ .{ "TF_TP_RANK", "1" }, .{ "TF_TP_PORT", "29600" } }; + const c = try Config.parse(&Fake.get); + try std.testing.expectEqual(@as(u32, 1), c.rank); + try std.testing.expectEqual(@as(u16, 29607), c.link_port); + try std.testing.expectEqual(Backend.nccl, c.backend); + try std.testing.expectEqual(std.time.ns_per_s, c.grace_ns); + try std.testing.expectEqual(std.time.ns_per_s / 2, c.report_ns); +} + +test "bad settings are refused and the digest ignores the rank" { + Fake.pairs = &.{.{ "TF_TP_RANK", "2" }}; + try std.testing.expectError(error.BadSetting, Config.parse(&Fake.get)); + Fake.pairs = &.{.{ "TF_COMM_BACKEND", "rdma" }}; + try std.testing.expectError(error.BadSetting, Config.parse(&Fake.get)); + Fake.pairs = &.{ .{ "TF_COMM_BACKEND", "roce" }, .{ "TF_TP_ROCE_HCA", "rocep1s0f1,roceP2p1s0f1" }, .{ "NCCL_IB_TC", "106" } }; + const r = try Config.parse(&Fake.get); + try std.testing.expectEqual(Backend.roce, r.backend); + try std.testing.expectEqual(@as(u8, 106), r.roce_tc); + try std.testing.expectEqualStrings("rocep1s0f1,roceP2p1s0f1", r.roce_hca.?); + Fake.pairs = &.{ .{ "TF_COMM_BACKEND", "mailbox" }, .{ "TF_TP_MAILBOX_MAX_KB", "64" }, .{ "TF_TP_FAILFAST_GRACE_S", "0.25" } }; + const a = try Config.parse(&Fake.get); + try std.testing.expectEqual(@as(usize, 64 * 1024), a.mailbox_max_bytes); + try std.testing.expectEqual(std.time.ns_per_s / 4, a.grace_ns); + var b = a; + b.rank = 1; + try std.testing.expectEqualSlices(u8, &a.digest(), &b.digest()); + b.mailbox_max_bytes = 1024; + try std.testing.expect(!std.mem.eql(u8, &a.digest(), &b.digest())); +} + +test "the plan link's spin and pin, the RoCE grid cap" { + Fake.pairs = &.{ .{ "TF_COMM_BACKEND", "roce" }, .{ "TF_TP_MAILBOX_MAX_KB", "1024" }, .{ "TF_TP_ROCE_FALLBACK", "error" }, .{ "TF_TP_PLAN_LINK", "rdma" }, .{ "TF_TP_PLAN_PIN", "auto" } }; + const p = try Config.parse(&Fake.get); + try std.testing.expectEqual(Backend.roce, p.backend); + try std.testing.expectEqual(@as(usize, 1024 * 1024), p.mailbox_max_bytes); + try std.testing.expectEqual(Fallback.@"error", p.roce_fallback); + try std.testing.expectEqual(@as(u64, 5000 * 1000), p.plan_spin_ns); + try std.testing.expectEqual(PlanPin.auto, p.plan_pin.?); + Fake.pairs = &.{ .{ "TF_TP_PLAN_LINK", "tcp" }, .{ "TF_TP_PLAN_PIN", "7" } }; // tcp blocks at once; an explicit CPU + const o = try Config.parse(&Fake.get); + try std.testing.expectEqual(@as(u64, 0), o.plan_spin_ns); + try std.testing.expectEqual(PlanPin{ .cpu = 7 }, o.plan_pin.?); + Fake.pairs = &.{ .{ "TF_TP_PLAN_LINK", "rdma" }, .{ "TF_TP_PLAN_SPIN_US", "200" }, .{ "TF_TP_PLAN_PIN", "-1" } }; + const s = try Config.parse(&Fake.get); + try std.testing.expectEqual(@as(u64, 200 * 1000), s.plan_spin_ns); + try std.testing.expect(s.plan_pin == null); + Fake.pairs = &.{.{ "TF_TP_PLAN_LINK", "udp" }}; + try std.testing.expectError(error.BadSetting, Config.parse(&Fake.get)); + Fake.pairs = &.{}; // the RoCE grid cap: 16 unset, powers of two only + try std.testing.expectEqual(@as(u32, 16), (try Config.parse(&Fake.get)).roce_blocks); + Fake.pairs = &.{.{ "TF_TP_ROCE_BLOCKS", "4" }}; + try std.testing.expectEqual(@as(u32, 4), (try Config.parse(&Fake.get)).roce_blocks); + Fake.pairs = &.{.{ "TF_TP_ROCE_BLOCKS", "12" }}; + try std.testing.expectError(error.BadSetting, Config.parse(&Fake.get)); + Fake.pairs = &.{}; +} + +test "the digest holds the graph cache's effective settings: 0, off and OFF agree, on and off differ" { + Fake.pairs = &.{.{ "TF_GRAPH_CACHE", "0" }}; + var zero = try Config.parse(&Fake.get); + Fake.pairs = &.{.{ "TF_GRAPH_CACHE", "OFF" }}; + var upper = try Config.parse(&Fake.get); + Fake.pairs = &.{.{ "TF_GRAPH_CACHE", "1" }}; + var one = try Config.parse(&Fake.get); + Fake.pairs = &.{}; + var unset = try Config.parse(&Fake.get); + try std.testing.expectEqualSlices(u8, &zero.digest(), &unset.digest()); // no cache installed: not compared + for ([_]*Config{ &zero, &upper, &one, &unset }) |c| c.graph_cache = .{}; + var max0 = unset; + max0.graph_cache = .{ .max = 0 }; + try std.testing.expectEqualSlices(u8, &zero.digest(), &upper.digest()); + try std.testing.expectEqualSlices(u8, &zero.digest(), &max0.digest()); + try std.testing.expectEqualSlices(u8, &one.digest(), &unset.digest()); + try std.testing.expect(!std.mem.eql(u8, &zero.digest(), &unset.digest())); + var family = unset; + family.graph_cache = .{ .max = 8, .hold = true }; + try std.testing.expect(!std.mem.eql(u8, &family.digest(), &unset.digest())); + var family_off = upper; + family_off.graph_cache = family.graph_cache; + try std.testing.expectEqualSlices(u8, &zero.digest(), &family_off.digest()); +} diff --git a/zig/src/fabric/tp/fate.zig b/zig/src/fabric/tp/fate.zig new file mode 100644 index 0000000000..5084c4648b --- /dev/null +++ b/zig/src/fabric/tp/fate.zig @@ -0,0 +1,261 @@ +//! Fail fast on every rank: a failure or a vanished rank ends all ranks with exit 70 about a second later. +const std = @import("std"); +const log = @import("log.zig"); +const sock = @import("sock.zig"); +const Config = @import("config.zig").Config; +const posix = std.posix; +/// Linux POLLRDHUP: the peer shut its side down (std has no name for it). +const pollrdhup: i16 = 0x2000; + +pub const exit_code: u8 = 70; +const magic: u32 = 0x46463431; +const kind_mem: u32 = 1; +const kind_fail: u32 = 2; +/// rank 0's clean stop: the hang-up that follows is no failure (sent before rank 0 closes its sockets) +const kind_stop: u32 = 3; +const Frame = extern struct { magic: u32, kind: u32, value: i64 }; + +pub const max_ranks = 64; +pub const max_hooks = 8; + +/// A callback with its context (an abort of a collective, rank 0's request failing); `reason` lives until exit. +pub const Hook = struct { ctx: *anyopaque, call: *const fn (ctx: *anyopaque, reason: []const u8) void }; + +pub const Fate = struct { + rank: u32, + world: u32, + fds: [max_ranks]sock.Fd, + grace_ns: u64, + report_ns: u64, + /// This rank's usable memory in bytes for the MEM report (default: MemAvailable). + usable: *const fn () u64 = memAvailable, + /// The process exit; tests replace it (it may return there). + exit_fn: *const fn (u8) void = exitNow, + hooks: [max_hooks]Hook = undefined, + n_hooks: std.atomic.Value(u32) = .init(0), + dead: std.atomic.Value(bool) = .init(false), + closing: std.atomic.Value(bool) = .init(false), + /// Set once the exit was asked for (tests read it; the process is gone in serving). + exited: std.atomic.Value(bool) = .init(false), + reason_buf: [512]u8 = undefined, + reason_len: usize = 0, + send_lock: std.atomic.Value(bool) = .init(false), + /// Rank 0: each rank's last report (bytes) and its monotonic time. + peer_mem: [max_ranks]std.atomic.Value(u64) = @splat(.init(0)), + peer_mem_at: [max_ranks]std.atomic.Value(u64) = @splat(.init(0)), + wake: [2]sock.Fd = .{ -1, -1 }, + thread: ?std.Thread = null, + + /// The channel over the bootstrap's sockets (`Bootstrap.release`); `start` runs its thread. + pub fn init(cfg: Config, fds: [max_ranks]sock.Fd) Fate { + return .{ .rank = cfg.rank, .world = cfg.world, .fds = fds, .grace_ns = cfg.grace_ns, .report_ns = cfg.report_ns }; + } + + pub fn start(f: *Fate) !void { + var p: [2]c_int = undefined; + if (std.c.pipe2(&p, .{ .CLOEXEC = true }) != 0) return error.SocketFailed; + f.wake = p; + f.thread = if (f.rank == 0) try std.Thread.spawn(.{}, watch0, .{f}) else try std.Thread.spawn(.{}, watchK, .{f}); + } + + /// Runs `hook` when the ranks go down; collectives register their abort here. + pub fn onFail(f: *Fate, hook: Hook) void { + const i = f.n_hooks.load(.acquire); + std.debug.assert(i < max_hooks); + f.hooks[i] = hook; + f.n_hooks.store(i + 1, .release); + } + + /// This rank failed (`where`: what it was doing): every rank goes down. Returns only in tests. + pub fn fatal(f: *Fate, where: []const u8, err: anyerror) void { + var buf: [256]u8 = undefined; + const msg = std.fmt.bufPrint(&buf, "rank {d} failed in {s}: {t}", .{ f.rank, where, err }) catch "rank failed"; + f.die(msg, true); + } + + /// This rank failed while booting: every rank goes down, and the thread is joined before the caller frees this. + pub fn failBoot(f: *Fate, where: []const u8, err: anyerror) void { + f.fatal(where, err); + f.close(); + } + + /// A normal stop is coming (rank 0 before it tells the others; the others when told): later hang-ups are fine. + pub fn expectStop(f: *Fate) void { + f.closing.store(true, .release); + } + + /// Rank 0's clean stop of every rank: a STOP frame to each ahead of its close, so followers end quietly. + pub fn announceStop(f: *Fate) void { + if (f.closing.swap(true, .acq_rel)) return; + if (f.rank != 0) return; + const h: Frame = .{ .magic = magic, .kind = kind_stop, .value = 0 }; + f.lock(); + defer f.unlock(); + for (f.fds[1..f.world]) |fd| if (fd >= 0) { + _ = sock.trySend(fd, std.mem.asBytes(&h)); + }; + } + + /// A normal stop: the thread ends and the sockets close. + pub fn close(f: *Fate) void { + f.closing.store(true, .release); + if (f.wake[1] >= 0) _ = std.c.write(f.wake[1], "x", 1); + if (f.thread) |t| t.join(); + f.thread = null; + for (&f.fds) |*fd| if (fd.* >= 0) { + sock.close(fd.*); + fd.* = -1; + }; + for (f.wake) |fd| if (fd >= 0) sock.close(fd); + f.wake = .{ -1, -1 }; + } + + /// Rank 0: rank `r`'s last reported usable bytes and the age of the report, or null before its first. + pub fn peerMem(f: *const Fate, r: u32) ?struct { bytes: u64, age_ns: u64 } { + const at = f.peer_mem_at[r].load(.acquire); + if (at == 0) return null; + return .{ .bytes = f.peer_mem[r].load(.acquire), .age_ns = sock.nowNs() -| at }; + } + + pub fn reason(f: *const Fate) []const u8 { + return f.reason_buf[0..f.reason_len]; + } + + fn lock(f: *Fate) void { + while (f.send_lock.swap(true, .acquire)) std.atomic.spinLoopHint(); + } + + fn unlock(f: *Fate) void { + f.send_lock.store(false, .release); + } + + fn die(f: *Fate, why: []const u8, own: bool) void { + if (f.dead.swap(true, .acq_rel)) { + // the first caller exits the process; a second (the fate thread or a round thread) waits for it + while (!f.exited.load(.acquire)) sock.sleepNs(10 * std.time.ns_per_ms); + return; + } + const n = @min(why.len, f.reason_buf.len); + @memcpy(f.reason_buf[0..n], why[0..n]); + f.reason_len = n; + const r = f.reason(); + log.warn("[tensorfold] tp fail-fast: rank {d}: {s}; exiting ({d}) in {d} ms so every rank restarts", .{ f.rank, r, exit_code, f.grace_ns / std.time.ns_per_ms }); + // aborts first: kernels waiting on the dead peer end now; a hook that blocks cannot hold the exit + if (std.Thread.spawn(.{}, runHooks, .{ f, r })) |t| t.detach() else |_| runHooks(f, r); + if (!f.closing.load(.acquire)) { + const head: Frame = .{ .magic = magic, .kind = kind_fail, .value = @intCast(r.len) }; + f.lock(); + for (f.fds[0..f.world]) |fd| if (fd >= 0) { + if (own or f.rank == 0) { + _ = sock.trySend(fd, std.mem.asBytes(&head)); + _ = sock.trySend(fd, r); + } + sock.hangup(fd); + }; + f.unlock(); + } + sock.sleepNs(f.grace_ns); + f.exited.store(true, .release); + f.exit_fn(exit_code); + } + + fn runHooks(f: *Fate, r: []const u8) void { + const n = f.n_hooks.load(.acquire); + for (f.hooks[0..n]) |h| h.call(h.ctx, r); + } + + /// Reads one frame from `fd`: a MEM report, a FAIL (the ranks go down), or the end (the peer is gone). + fn readFrame(f: *Fate, peer: u32, fd: sock.Fd) bool { + var h: Frame = undefined; + const deadline = sock.nowNs() + std.time.ns_per_s; + sock.recvExact(fd, std.mem.asBytes(&h), deadline) catch { + if (f.closing.load(.acquire)) return false; + var buf: [128]u8 = undefined; + f.die(std.fmt.bufPrint(&buf, "rank {d} closed the fate channel (it failed, exited or was killed)", .{peer}) catch "peer gone", false); + return false; + }; + if (h.magic != magic) { + f.die("a frame that is not the fate channel's", false); + return false; + } + if (h.kind == kind_stop) { + // rank 0 stops every rank: its hang-up is next, and no failure + f.closing.store(true, .release); + return false; + } + if (h.kind == kind_mem) { + f.peer_mem[peer].store(@intCast(@max(h.value, 0)), .release); + f.peer_mem_at[peer].store(sock.nowNs(), .release); + return true; + } + var text: [512]u8 = undefined; + const len: usize = @intCast(std.math.clamp(h.value, 0, text.len)); + sock.recvExact(fd, text[0..len], deadline) catch {}; + f.die(if (len > 0) text[0..len] else "a peer failed", false); + return false; + } + + fn watch0(f: *Fate) void { + var p: [max_ranks + 1]std.c.pollfd = undefined; + var who: [max_ranks + 1]u32 = undefined; + var n: usize = 0; + for (f.fds[1..f.world], 1..) |fd, r| if (fd >= 0) { + p[n] = .{ .fd = fd, .events = posix.POLL.IN | pollrdhup, .revents = 0 }; + who[n] = @intCast(r); + n += 1; + }; + p[n] = .{ .fd = f.wake[0], .events = posix.POLL.IN, .revents = 0 }; + while (!f.dead.load(.acquire) and !f.closing.load(.acquire)) { + if (std.c.poll(&p, @intCast(n + 1), -1) <= 0) continue; + for (p[0..n], who[0..n]) |*q, r| if (q.revents != 0) { + if (!f.readFrame(r, q.fd)) return; + }; + } + } + + fn watchK(f: *Fate) void { + const fd = f.fds[0]; + var p = [_]std.c.pollfd{ + .{ .fd = fd, .events = posix.POLL.IN | pollrdhup, .revents = 0 }, + .{ .fd = f.wake[0], .events = posix.POLL.IN, .revents = 0 }, + }; + const ms: c_int = @intCast(@max(f.report_ns / std.time.ns_per_ms, 1)); + f.report(fd); + while (!f.dead.load(.acquire) and !f.closing.load(.acquire)) { + const ready = std.c.poll(&p, 2, ms); + if (ready == 0) { + f.report(fd); + continue; + } + if (ready > 0 and p[0].revents != 0 and !f.readFrame(0, fd)) return; + } + } + + /// One MEM frame, never blocking (a full buffer drops it: the next comes one period later). + fn report(f: *Fate, fd: sock.Fd) void { + const h: Frame = .{ .magic = magic, .kind = kind_mem, .value = @intCast(f.usable()) }; + f.lock(); + defer f.unlock(); + _ = sock.trySend(fd, std.mem.asBytes(&h)); + } +}; + +fn exitNow(code: u8) void { + std.c._exit(code); +} + +/// MemAvailable from /proc/meminfo in bytes (0 when unreadable). +pub fn memAvailable() u64 { + const fd = std.c.open("/proc/meminfo", .{ .ACCMODE = .RDONLY }); + if (fd < 0) return 0; + defer _ = std.c.close(fd); + var buf: [4096]u8 = undefined; + const n = std.c.read(fd, &buf, buf.len); + if (n <= 0) return 0; + const text = buf[0..@intCast(n)]; + const key = "MemAvailable:"; + const at = std.mem.indexOf(u8, text, key) orelse return 0; + var it = std.mem.tokenizeAny(u8, text[at + key.len ..], " \n"); + const kb = std.fmt.parseInt(u64, it.next() orelse return 0, 10) catch return 0; + return kb * 1024; +} diff --git a/zig/src/fabric/tp/fate_test.zig b/zig/src/fabric/tp/fate_test.zig new file mode 100644 index 0000000000..728f4bfdf3 --- /dev/null +++ b/zig/src/fabric/tp/fate_test.zig @@ -0,0 +1,241 @@ +//! The fate channel between two ranks over real TCP: a failure or a vanished rank ends both, a stop ends neither. +const std = @import("std"); +const sock = @import("sock.zig"); +const fate = @import("fate.zig"); +const bootstrap = @import("bootstrap.zig"); +const Config = @import("config.zig").Config; +const Fate = fate.Fate; + +var exits: std.atomic.Value(u32) = .init(0); +var hooks: std.atomic.Value(u32) = .init(0); + +fn recordExit(code: u8) void { + std.debug.assert(code == fate.exit_code); + _ = exits.fetchAdd(1, .acq_rel); +} + +fn hook(_: *anyopaque, reason: []const u8) void { + std.debug.assert(reason.len > 0); + _ = hooks.fetchAdd(1, .acq_rel); +} + +const grace = 50 * std.time.ns_per_ms; + +/// Two connected ranks' fate channels (rank 0 first), started, with recorded exits and hooks. +fn pair(fates: *[2]Fate) !void { + return pairAt(.{ &fates[0], &fates[1] }); +} + +fn pairAt(fates: [2]*Fate) !void { + const l = try bootstrap.testListener(); + defer sock.close(l.fd); + const deadline = sock.nowNs() + 5 * std.time.ns_per_s; + const a = try sock.connect("localhost", l.port, deadline); + const b = try sock.accept(l.fd, deadline); + var fds0: [fate.max_ranks]sock.Fd = @splat(-1); + var fds1: [fate.max_ranks]sock.Fd = @splat(-1); + fds0[1] = b; + fds1[0] = a; + for (fates, [2][fate.max_ranks]sock.Fd{ fds0, fds1 }, 0..) |f, fds, r| { + f.* = Fate.init(.{ .world = 2, .rank = @intCast(r), .grace_ns = grace, .report_ns = 20 * std.time.ns_per_ms }, fds); + f.exit_fn = recordExit; + f.usable = fakeUsable; + f.onFail(.{ .ctx = f, .call = hook }); + } + for (fates) |f| try f.start(); +} + +fn fakeUsable() u64 { + return 123 << 20; +} + +fn waitExits(want: u32) !u64 { + const t0 = sock.nowNs(); + while (exits.load(.acquire) < want) { + if (sock.nowNs() - t0 > 5 * std.time.ns_per_s) return error.Timeout; + sock.sleepNs(std.time.ns_per_ms); + } + return sock.nowNs() - t0; +} + +fn reset() void { + exits.store(0, .release); + hooks.store(0, .release); +} + +test "rank 1's failure ends both ranks with its reason" { + reset(); + var fates: [2]Fate = undefined; + try pair(&fates); + defer for (&fates) |*f| f.close(); + sock.sleepNs(60 * std.time.ns_per_ms); + const mem = fates[0].peerMem(1) orelse return error.NoReport; + try std.testing.expectEqual(@as(u64, 123 << 20), mem.bytes); + fates[1].fatal("prefill", error.OutOfMemory); + const took = try waitExits(2); + try std.testing.expect(took < grace + 500 * std.time.ns_per_ms); + try std.testing.expectEqual(@as(u32, 2), hooks.load(.acquire)); + try std.testing.expect(std.mem.indexOf(u8, fates[0].reason(), "rank 1 failed in prefill: OutOfMemory") != null); +} + +test "rank 0's failure ends rank 1" { + reset(); + var fates: [2]Fate = undefined; + try pair(&fates); + defer for (&fates) |*f| f.close(); + fates[0].fatal("window", error.Unexpected); + _ = try waitExits(2); + try std.testing.expect(std.mem.indexOf(u8, fates[1].reason(), "rank 0 failed in window") != null); +} + +test "a rank that vanishes without a word ends the other" { + reset(); + var fates: [2]Fate = undefined; + try pair(&fates); + defer fates[0].close(); + // rank 1 killed: its socket closes with no FAIL frame (its own fate thread ends quietly first) + fates[1].expectStop(); + fates[1].close(); + _ = try waitExits(1); + try std.testing.expect(std.mem.indexOf(u8, fates[0].reason(), "rank 1 closed the fate channel") != null); +} + +test "a normal stop ends nobody" { + reset(); + var fates: [2]Fate = undefined; + try pair(&fates); + fates[0].expectStop(); + fates[1].expectStop(); + fates[1].close(); + fates[0].close(); + sock.sleepNs(2 * grace); + try std.testing.expectEqual(@as(u32, 0), exits.load(.acquire)); +} + +test "a rank failing at boot ends both with its reason, and its channel is down before its owner frees it" { + // a follower's boot failed after its session came up; its Fate is freed while the fate thread could still read it + reset(); + // rank 1's Fate where a Model keeps it: on the heap, its pages unmapped when freed (a stale read faults) + const pa = std.heap.page_allocator; + var f0: Fate = undefined; + const f1 = try pa.create(Fate); + try pairAt(.{ &f0, f1 }); + defer f0.close(); + sock.sleepNs(60 * std.time.ns_per_ms); + f1.failBoot("boot", error.RopeTooShort); + try std.testing.expect(f1.thread == null); + pa.destroy(f1); + _ = try waitExits(2); + // several report periods with rank 1's Fate gone: its thread would have reported from freed memory + sock.sleepNs(5 * 20 * std.time.ns_per_ms); + try std.testing.expect(std.mem.indexOf(u8, f0.reason(), "rank 1 failed in boot: RopeTooShort") != null); +} + +test "rank 0's clean stop ends every rank with no fail-fast exit, whatever order the ranks close in" { + // rank 0 ended its run and closed; the follower, still tearing down, must not see the hang-up as a failure + reset(); + var fates: [2]Fate = undefined; + try pair(&fates); + sock.sleepNs(60 * std.time.ns_per_ms); + fates[0].announceStop(); + fates[0].close(); // rank 0 gone before rank 1 even started its own stop + sock.sleepNs(2 * grace + 100 * std.time.ns_per_ms); + fates[1].close(); + try std.testing.expectEqual(@as(u32, 0), exits.load(.acquire)); + try std.testing.expectEqual(@as(u32, 0), hooks.load(.acquire)); +} + +// -- a server's drain on two ranks: the plan link and the fate channel in the order a family's stop uses them -- + +const planlink = @import("planlink.zig"); + +/// Plan words a family's rank 0 sends: a round's operation, and a stop word that ends a follower's loop. +const op_round: i64 = 1; +const op_stop: i64 = 4; + +fn linkPair(links: *[2]planlink.PlanLink) !void { + const l = try bootstrap.testListener(); + defer sock.close(l.fd); + const deadline = sock.nowNs() + 5 * std.time.ns_per_s; + const a = try sock.connect("localhost", l.port, deadline); + const b = try sock.accept(l.fd, deadline); + var f0: [planlink.max_ranks]sock.Fd = @splat(-1); + var f1: [planlink.max_ranks]sock.Fd = @splat(-1); + f0[1] = b; + f1[0] = a; + links[0] = planlink.PlanLink.init(0, 2, f0); + links[1] = planlink.PlanLink.init(1, 2, f1); +} + +/// Rank 1's loop: each round as it comes, until the stop word, or until the link ends. +const Follower = struct { + link: *planlink.PlanLink, + rounds: u32 = 0, + ended: enum { running, stopped, peer_closed, failed } = .running, + + fn run(f: *Follower) void { + var buf: [16]i64 = undefined; + while (true) { + const msg = f.link.recv(&buf) catch |e| { + f.ended = if (e == error.PeerClosed) .peer_closed else .failed; + return; + }; + if (msg.len == 0) continue; + if (msg[0] == op_stop) { + f.ended = .stopped; + return; + } + if (msg[0] == op_round) f.rounds += 1; + } + } +}; + +test "a drain on rank 0, then its clean stop at a round boundary: rank 1 ends its loop there and nobody exits 70" { + reset(); + var fates: [2]Fate = undefined; + try pair(&fates); + var links: [2]planlink.PlanLink = undefined; + try linkPair(&links); + var follower: Follower = .{ .link = &links[1] }; + const t = try std.Thread.spawn(.{}, Follower.run, .{&follower}); + // serving, then draining (rank 0 refuses new work, rounds go on; the fate channel is untouched), past a grace + var sent: u32 = 0; + for (0..30) |_| { + try links[0].send(&.{ op_round, sent }); + sent += 1; + sock.sleepNs(5 * std.time.ns_per_ms); + } + try std.testing.expect(!fates[0].closing.load(.acquire)); + // the stop's order: the stop word first, so the follower ends its loop before the fate STOP and the hang-up + try links[0].send(&.{op_stop}); + fates[0].announceStop(); + links[0].close(); + fates[0].close(); + t.join(); + // rank 1's loop returned: its model closes and the process exits 0 + fates[1].close(); + links[1].close(); + sock.sleepNs(2 * grace + 100 * std.time.ns_per_ms); + try std.testing.expectEqual(.stopped, follower.ended); + try std.testing.expectEqual(sent, follower.rounds); + try std.testing.expectEqual(@as(u32, 0), exits.load(.acquire)); + try std.testing.expectEqual(@as(u32, 0), hooks.load(.acquire)); +} + +test "a drain on rank 0 does not disarm fail-fast: rank 1 killed mid-drain still ends rank 0 (exit 70)" { + reset(); + var fates: [2]Fate = undefined; + try pair(&fates); + defer fates[0].close(); + var links: [2]planlink.PlanLink = undefined; + try linkPair(&links); + defer links[0].close(); + for (0..5) |i| try links[0].send(&.{ op_round, @intCast(i) }); + // rank 0 is draining (no stop announced); rank 1 is killed: its sockets close with no word + links[1].close(); + fates[1].expectStop(); + fates[1].close(); + _ = try waitExits(1); + try std.testing.expect(std.mem.indexOf(u8, fates[0].reason(), "rank 1 closed the fate channel") != null); + try std.testing.expectEqual(@as(u32, 1), exits.load(.acquire)); +} diff --git a/zig/src/fabric/tp/gpu_bench.zig b/zig/src/fabric/tp/gpu_bench.zig new file mode 100644 index 0000000000..99831e6d77 --- /dev/null +++ b/zig/src/fabric/tp/gpu_bench.zig @@ -0,0 +1,82 @@ +//! Latency and bandwidth of each backend's collectives, eager and graphed; RESULT lines are JSON, busbw as nccl-tests. +const std = @import("std"); +const cuda = @import("cuda"); +const tp = @import("tp"); +const Ctx = @import("gpu_checks.zig").Ctx; +const Collective = tp.Collective; + +pub const Kind = enum { all_reduce, all_gather }; + +fn issue(c: *Ctx, comm: Collective, kind: Kind, bytes: usize) !void { + const count = bytes / 2; + switch (kind) { + .all_reduce => try comm.allReduce(c.in.ptr, c.out.ptr, count, .bf16, .sum, c.stream.handle), + .all_gather => try comm.allGather(c.in.ptr, c.out.ptr, count, .bf16, c.stream.handle), + } +} + +/// GPU microseconds a call, `iters` calls issued back to back on the stream. +fn eager(c: *Ctx, comm: Collective, kind: Kind, bytes: usize, iters: u32) !f64 { + var e0 = try cuda.Event.init(c.d, true); + defer e0.deinit(); + var e1 = try cuda.Event.init(c.d, true); + defer e1.deinit(); + for (0..10) |_| try issue(c, comm, kind, bytes); + try c.stream.synchronize(); + try comm.barrier(); + try e0.record(c.stream); + for (0..iters) |_| try issue(c, comm, kind, bytes); + try e1.record(c.stream); + try e1.synchronize(); + return @as(f64, try cuda.Event.elapsedMs(e0, e1)) * 1000.0 / @as(f64, @floatFromInt(iters)); +} + +/// GPU microseconds a call inside one graph of `per_graph` calls, replayed `replays` times. +fn graphed(c: *Ctx, comm: Collective, kind: Kind, bytes: usize, per_graph: u32, replays: u32) !f64 { + try cuda.graph.beginCapture(c.stream, .thread_local); + const err: ?tp.collective.Error = blk: { + for (0..per_graph) |_| issue(c, comm, kind, bytes) catch |e| break :blk e; + break :blk null; + }; + var g = try cuda.graph.endCapture(c.stream); + defer g.deinit(); + if (err) |e| return e; + var exec = try g.instantiate(); + defer exec.deinit(); + try exec.upload(c.stream); + try exec.launchOn(c.stream); + try c.stream.synchronize(); + try comm.barrier(); + var e0 = try cuda.Event.init(c.d, true); + defer e0.deinit(); + var e1 = try cuda.Event.init(c.d, true); + defer e1.deinit(); + try e0.record(c.stream); + for (0..replays) |_| try exec.launchOn(c.stream); + try e1.record(c.stream); + try e1.synchronize(); + try comm.check(); + return @as(f64, try cuda.Event.elapsedMs(e0, e1)) * 1000.0 / @as(f64, @floatFromInt(per_graph * replays)); +} + +/// Every power of two from `lo` to `hi` bytes a rank, both kinds, eager and graphed. +pub fn sweep(c: *Ctx, comm: Collective, label: []const u8, lo: usize, hi: usize) !void { + for ([_]Kind{ .all_reduce, .all_gather }) |kind| { + var bytes = lo; + while (bytes <= hi) : (bytes *= 2) { + const per_graph: u32 = if (bytes >= 1 << 20) 20 else 100; + const e = try eager(c, comm, kind, bytes, if (bytes >= 1 << 20) 50 else 200); + const g = try graphed(c, comm, kind, bytes, per_graph, 5); + const moved: f64 = @floatFromInt(if (kind == .all_gather) 2 * bytes else bytes); + const factor: f64 = if (kind == .all_gather) 0.5 else 1.0; + const busbw = moved / (g * 1e-6) * factor / 1e9; + if (c.rank == 0) std.debug.print("RESULT {{\"backend\":\"{s}\",\"op\":\"{t}\",\"bytes\":{d},\"eager_us\":{d:.2},\"graph_us\":{d:.2},\"busbw_GBps\":{d:.2}}}\n", .{ label, kind, bytes, e, g, busbw }); + } + } +} + +/// One decode window's exchanges: `n` all-gathers of one row's bf16 partials (`row_bytes` a rank) in one graph. +pub fn decodeWindow(c: *Ctx, comm: Collective, label: []const u8, n: u32, row_bytes: usize) !void { + const per = try graphed(c, comm, .all_gather, row_bytes, n, 20); + if (c.rank == 0) std.debug.print("RESULT {{\"backend\":\"{s}\",\"op\":\"decode_window\",\"exchanges\":{d},\"bytes\":{d},\"per_exchange_us\":{d:.2},\"window_us\":{d:.1}}}\n", .{ label, n, row_bytes, per, per * @as(f64, @floatFromInt(n)) }); +} diff --git a/zig/src/fabric/tp/gpu_checks.zig b/zig/src/fabric/tp/gpu_checks.zig new file mode 100644 index 0000000000..5dc4ed6d4c --- /dev/null +++ b/zig/src/fabric/tp/gpu_checks.zig @@ -0,0 +1,188 @@ +//! Every collective on real GPUs against the host reference, eager and replayed from a graph. +const std = @import("std"); +const cuda = @import("cuda"); +const tp = @import("tp"); +const Collective = tp.Collective; +const DType = tp.DType; +const Op = tp.Op; +const reference = tp.reference; + +pub const Ctx = struct { + d: *const cuda.Driver, + gpa: std.mem.Allocator, + stream: cuda.Stream, + rank: u32, + /// Device buffers sized for the largest case: input, output (world x), a spare. + in: cuda.DeviceBuffer, + out: cuda.DeviceBuffer, + /// Host scratch: this rank's input, the peer's, the expected and the downloaded output. + mine: []u8, + peer: []u8, + want: []u8, + got: []u8, + failures: u32 = 0, + passes: u32 = 0, + + pub fn init(d: *const cuda.Driver, gpa: std.mem.Allocator, rank: u32, max_bytes: usize) !Ctx { + return .{ + .d = d, + .gpa = gpa, + .stream = try cuda.Stream.init(d, true), + .rank = rank, + .in = try cuda.DeviceBuffer.alloc(d, 2 * max_bytes + 256), + .out = try cuda.DeviceBuffer.alloc(d, 2 * max_bytes + 256), + .mine = try gpa.alloc(u8, max_bytes), + .peer = try gpa.alloc(u8, max_bytes), + .want = try gpa.alloc(u8, 2 * max_bytes), + .got = try gpa.alloc(u8, 2 * max_bytes), + }; + } + + pub fn deinit(c: *Ctx) void { + c.stream.deinit(); + c.in.free(); + c.out.free(); + for ([_][]u8{ c.mine, c.peer, c.want, c.got }) |b| c.gpa.free(b); + } + + fn verdict(c: *Ctx, ok: bool, comptime fmt: []const u8, args: anytype) void { + if (ok) c.passes += 1 else c.failures += 1; + std.debug.print("{s} " ++ fmt ++ "\n", .{if (ok) "PASS" else "FAIL"} ++ args); + } + + /// Both ranks' inputs for case `seed`: this rank's into `mine`, the peer's into `peer` (n bytes each). + fn inputs(c: *Ctx, n: usize, t: DType, seed: u64) void { + reference.fill(c.mine[0..n], t, c.rank, seed); + reference.fill(c.peer[0..n], t, 1 - c.rank, seed); + } + + fn rankInput(c: *Ctx, r: u32, n: usize) []const u8 { + return if (r == c.rank) c.mine[0..n] else c.peer[0..n]; + } + + fn compare(c: *Ctx, comm: Collective, label: []const u8, what: []const u8, t: DType, n: usize, out_off: usize, want: []const u8) !void { + try c.stream.synchronize(); + comm.check() catch |e| { + c.verdict(false, "{s} {s} {t} {d} B: transport error {t}", .{ label, what, t, n, e }); + return; + }; + try c.out.download(out_off, c.got[0..want.len]); + const ok = std.mem.eql(u8, c.got[0..want.len], want); + var first: usize = 0; + if (!ok) first = std.mem.indexOfDiff(u8, c.got[0..want.len], want) orelse 0; + c.verdict(ok, "{s} {s} {t} {d} B a rank{s}", .{ label, what, t, n, if (ok) "" else " (first difference at a byte shown below)" }); + if (!ok) std.debug.print(" byte {d}: got {x} want {x}\n", .{ first, c.got[first], want[first] }); + } + + /// all-reduce (out of place, in place, misaligned by `skew` bytes), all-gather (out of place, in place), exchange. + pub fn collectives(c: *Ctx, comm: Collective, label: []const u8, n_bytes: usize, seed: u64) !void { + const reduce_types = [_]DType{ .f32, .bf16, .f16, .i32 }; + for (reduce_types, 0..) |t, ti| { + const n = n_bytes - n_bytes % t.size(); + const count = n / t.size(); + c.inputs(n, t, seed + ti); + reference.reduce(c.want[0..n], &.{ c.rankInput(0, n), c.rankInput(1, n) }, t, .sum); + try c.in.uploadAsync(0, c.mine[0..n], c.stream.handle); + try comm.allReduce(c.in.ptr, c.out.ptr, count, t, .sum, c.stream.handle); + try c.compare(comm, label, "all-reduce sum", t, n, 0, c.want[0..n]); + // in place, at an offset of one element (not 16-byte aligned for 2-byte types) + const skew = t.size(); + try c.out.uploadAsync(skew, c.mine[0..n], c.stream.handle); + try comm.allReduce(c.out.ptr + skew, c.out.ptr + skew, count, t, .sum, c.stream.handle); + try c.compare(comm, label, "all-reduce sum in place, skewed", t, n, skew, c.want[0..n]); + } + for ([_]Op{ .max, .min }) |op| { + const n = n_bytes - n_bytes % 4; + c.inputs(n, .f32, seed + 11); + reference.reduce(c.want[0..n], &.{ c.rankInput(0, n), c.rankInput(1, n) }, .f32, op); + try c.in.uploadAsync(0, c.mine[0..n], c.stream.handle); + try comm.allReduce(c.in.ptr, c.out.ptr, n / 4, .f32, op, c.stream.handle); + try c.compare(comm, label, if (op == .max) "all-reduce max" else "all-reduce min", .f32, n, 0, c.want[0..n]); + } + for ([_]DType{ .u8, .bf16, .f32 }, 0..) |t, ti| { + const n = n_bytes - n_bytes % t.size(); + c.inputs(n, t, seed + 20 + ti); + @memcpy(c.want[0..n], c.rankInput(0, n)); + @memcpy(c.want[n..][0..n], c.rankInput(1, n)); + try c.in.uploadAsync(0, c.mine[0..n], c.stream.handle); + try comm.allGather(c.in.ptr, c.out.ptr, n / t.size(), t, c.stream.handle); + try c.compare(comm, label, "all-gather", t, n, 0, c.want[0 .. 2 * n]); + try c.out.fill8(0, c.stream.handle); + try c.out.uploadAsync(c.rank * n, c.mine[0..n], c.stream.handle); + try comm.allGather(c.out.ptr + c.rank * n, c.out.ptr, n / t.size(), t, c.stream.handle); + try c.compare(comm, label, "all-gather in place", t, n, 0, c.want[0 .. 2 * n]); + try c.in.uploadAsync(0, c.mine[0..n], c.stream.handle); + try comm.exchange(c.in.ptr, c.out.ptr, n / t.size(), t, 1 - c.rank, c.stream.handle); + try c.compare(comm, label, "exchange", t, n, 0, c.peer[0..n]); + } + } + + /// send then recv one way, then the other way (NCCL behind every backend). + pub fn pointToPoint(c: *Ctx, comm: Collective, label: []const u8, n_bytes: usize, seed: u64) !void { + const n = n_bytes - n_bytes % 4; + c.inputs(n, .f32, seed); + try c.in.uploadAsync(0, c.mine[0..n], c.stream.handle); + for (0..2) |dir| { + const sender: u32 = @intCast(dir); + if (c.rank == sender) { + try comm.send(c.in.ptr, n / 4, .f32, 1 - c.rank, c.stream.handle); + try c.stream.synchronize(); + } else { + try comm.recv(c.out.ptr, n / 4, .f32, sender, c.stream.handle); + try c.compare(comm, label, "send/recv", .f32, n, 0, c.peer[0..n]); + } + } + } + + /// A captured chain replayed with fresh inputs, eager calls between replays (the mailbox epoch spans both). + pub fn graphs(c: *Ctx, comm: Collective, label: []const u8, n_bytes: usize, replays: u32) !void { + const n = n_bytes - n_bytes % 16; + const half = n / 2; // bf16 elements + try cuda.graph.beginCapture(c.stream, .thread_local); + // x1 = sum(x0) [bf16], x2 = gather(x1), x3 = sum(x2) [f32 view], x4 = exchange(x3) + const x0 = c.in.ptr; + const x1 = c.out.ptr; + const x2 = c.in.ptr + n; + const x3 = c.out.ptr + n; + const x4 = c.out.ptr + 3 * n; + const ops: ?tp.collective.Error = blk: { + comm.allReduce(x0, x1, half, .bf16, .sum, c.stream.handle) catch |e| break :blk e; + comm.allGather(x1, x2, half, .bf16, c.stream.handle) catch |e| break :blk e; + comm.allReduce(x2, x3, n / 2, .f32, .sum, c.stream.handle) catch |e| break :blk e; + comm.exchange(x3, x4, n / 2, .f32, 1 - c.rank, c.stream.handle) catch |e| break :blk e; + break :blk null; + }; + var g = try cuda.graph.endCapture(c.stream); + defer g.deinit(); + if (ops) |e| return e; + var exec = try g.instantiate(); + defer exec.deinit(); + const nodes = try g.nodeCount(); + const want1 = try c.gpa.alloc(u8, 2 * n); + defer c.gpa.free(want1); + const want3 = try c.gpa.alloc(u8, 2 * n); + defer c.gpa.free(want3); + for (0..replays) |k| { + c.inputs(n, .bf16, 1000 + k); + try c.in.uploadAsync(0, c.mine[0..n], c.stream.handle); + reference.reduce(want1[0..n], &.{ c.rankInput(0, n), c.rankInput(1, n) }, .bf16, .sum); + @memcpy(want1[n..][0..n], want1[0..n]); // the gather of equal sums: twice the same + reference.reduce(want3[0 .. 2 * n], &.{ want1[0 .. 2 * n], want1[0 .. 2 * n] }, .f32, .sum); + try exec.launchOn(c.stream); + try c.compare(comm, label, "graph replay: sum, gather, sum, exchange", .bf16, n, n, want3[0 .. 2 * n]); + try c.compare(comm, label, "graph replay: exchange of the sum", .f32, 2 * n, 3 * n, want3[0 .. 2 * n]); + // an eager call between replays + try c.collectivesQuick(comm, label, k); + } + std.debug.print(" ({s}: {d} graph nodes, {d} replays)\n", .{ label, nodes, replays }); + } + + fn collectivesQuick(c: *Ctx, comm: Collective, label: []const u8, k: usize) !void { + const n = 4096; + c.inputs(n, .bf16, 5000 + k); + reference.reduce(c.want[0..n], &.{ c.rankInput(0, n), c.rankInput(1, n) }, .bf16, .sum); + try c.in.uploadAsync(5 * n, c.mine[0..n], c.stream.handle); + try comm.allReduce(c.in.ptr + 5 * n, c.out.ptr + 5 * n, n / 2, .bf16, .sum, c.stream.handle); + try c.compare(comm, label, "eager between replays", .bf16, n, 5 * n, c.want[0..n]); + } +}; diff --git a/zig/src/fabric/tp/gpu_cli.zig b/zig/src/fabric/tp/gpu_cli.zig new file mode 100644 index 0000000000..dc576fc03b --- /dev/null +++ b/zig/src/fabric/tp/gpu_cli.zig @@ -0,0 +1,66 @@ +//! Validate the runner's command before opening a driver, bootstrap or fate channel. +const std = @import("std"); + +pub const usage = "usage: tf-tp-test suite | failfast RANK fatal|kill|transport [N [HOLD_S]] | roce-host | hca | kernel\n"; +pub const Mode = enum { fatal, kill, transport }; +pub const Failure = struct { victim: u32, mode: Mode, n: u32 = 50, hold: ?u32 = null }; +pub const Command = union(enum) { suite, failfast: Failure, @"roce-host", hca, kernel }; + +pub fn parse(args: []const []const u8) error{InvalidArguments}!Command { + if (args.len == 0) return error.InvalidArguments; + const tag = std.meta.stringToEnum(std.meta.Tag(Command), args[0]) orelse return error.InvalidArguments; + switch (tag) { + .suite, .@"roce-host", .hca, .kernel => { + if (args.len != 1) return error.InvalidArguments; + return switch (tag) { + .suite => .suite, + .@"roce-host" => .@"roce-host", + .hca => .hca, + .kernel => .kernel, + else => unreachable, + }; + }, + .failfast => { + if (args.len < 3 or args.len > 5) return error.InvalidArguments; + const mode = std.meta.stringToEnum(Mode, args[2]) orelse return error.InvalidArguments; + if (mode != .transport and args.len != 3) return error.InvalidArguments; + return .{ .failfast = .{ + .victim = try number(args[1]), + .mode = mode, + .n = if (args.len > 3) try number(args[3]) else 50, + .hold = if (args.len > 4) try number(args[4]) else null, + } }; + }, + } +} + +fn number(s: []const u8) error{InvalidArguments}!u32 { + return std.fmt.parseInt(u32, s, 10) catch error.InvalidArguments; +} + +// Also executed against the Linux entry point with an invalid environment, before any CUDA work. +pub const invalid = [_][]const []const u8{ + &.{}, &.{"UNKNOWN"}, &.{ "suite", "extra" }, &.{ "kernel", "extra" }, + &.{ "hca", "extra" }, &.{ "roce-host", "extra" }, &.{"failfast"}, &.{ "failfast", "0" }, + &.{ "failfast", "0", "typo" }, &.{ "failfast", "bad", "fatal" }, &.{ "failfast", "-1", "kill" }, &.{ "failfast", "4294967296", "fatal" }, + &.{ "failfast", "0", "fatal", "50" }, &.{ "failfast", "0", "kill", "50" }, &.{ "failfast", "0", "transport", "bad" }, &.{ "failfast", "0", "transport", "50", "bad" }, + &.{ "failfast", "0", "transport", "50", "0", "extra" }, +}; + +test "invalid commands are refused before resources open" { + for (invalid) |args| try std.testing.expectError(error.InvalidArguments, parse(args)); +} + +test "documented commands and failfast defaults" { + inline for (.{ "suite", "roce-host", "hca", "kernel" }) |name| { + try std.testing.expectEqualStrings(name, @tagName(try parse(&.{name}))); + } + inline for (.{ Mode.fatal, Mode.kill, Mode.transport }) |mode| { + const f = (try parse(&.{ "failfast", "2", @tagName(mode) })).failfast; + try std.testing.expectEqualDeep(Failure{ .victim = 2, .mode = mode }, f); + } + const explicit = (try parse(&.{ "failfast", "63", "transport", "0", "0" })).failfast; + try std.testing.expectEqualDeep(Failure{ .victim = 63, .mode = .transport, .n = 0, .hold = 0 }, explicit); + const n = (try parse(&.{ "failfast", "1", "transport", "123" })).failfast; + try std.testing.expectEqualDeep(Failure{ .victim = 1, .mode = .transport, .n = 123 }, n); +} diff --git a/zig/src/fabric/tp/gpu_main.zig b/zig/src/fabric/tp/gpu_main.zig new file mode 100644 index 0000000000..1243ab794e --- /dev/null +++ b/zig/src/fabric/tp/gpu_main.zig @@ -0,0 +1,155 @@ +//! `tf-tp-test suite | failfast RANK fatal|kill|transport [N [HOLD_S]] | roce-host | hca | kernel`: one rank of TP=2. +const std = @import("std"); +const cuda = @import("cuda"); +const tp = @import("tp"); +const tp_kernel = @import("tp_kernel"); +const checks = @import("gpu_checks.zig"); +const bench = @import("gpu_bench.zig"); +const cli = @import("gpu_cli.zig"); + +const max_bytes = (4 << 20) + 4096; + +pub fn main(init: std.process.Init) !u8 { + const args = try init.minimal.args.toSlice(init.arena.allocator()); + const command = cli.parse(args[1..]) catch { + std.debug.print("{s}", .{cli.usage}); + return 2; + }; + // `kernel`: whether the build embeds the mailbox kernel (-Dtp-nvcc or -Dtp-fatbin), for scripts to check first + if (command == .kernel) { + std.debug.print("mailbox kernel {d} B\n", .{tp_kernel.mailbox.len}); + return if (tp_kernel.mailbox.len > 0) 0 else 1; + } + const cfg = try tp.Config.fromEnv(); + if (command == .@"roce-host") return @import("roce_host.zig").run(cfg, init.gpa); + if (command == .hca) return hca(cfg, init.io); + if (command == .failfast and (!cfg.failfast or command.failfast.hold != null) and !hostOnlyAllowed(cfg)) { + std.debug.print("REFUSED: failfast with TF_TP_FAILFAST=0 or a hold is host-only: TF_TP_ALLOW_HOST_ONLY=1 and no RDMA wire\n", .{}); + return 2; + } + if (cfg.backend != .nccl and tp_kernel.mailbox.len == 0) { + std.debug.print("FAIL TF_COMM_BACKEND={t} needs the mailbox kernel: rebuild with -Dtp-nvcc or -Dtp-fatbin\n", .{cfg.backend}); + return 2; + } + var driver = try cuda.Driver.open(); + defer driver.close(); + var ctx = try cuda.Context.init(&driver, @intCast(cfg.device)); + defer ctx.deinit(); + var name_buf: [128]u8 = undefined; + const t0 = tp.sock.nowNs(); + var s: tp.Session = undefined; + try s.open(&driver, cfg, tp_kernel.mailbox); + const up_ms = @as(f64, @floatFromInt(tp.sock.nowNs() - t0)) / 1e6; + std.debug.print("INFO rank {d}/{d} on {s} (sm_{d}), backend {t}, mailbox kernel {d} B, bootstrap + NCCL init {d:.0} ms\n", .{ cfg.rank, cfg.world, try ctx.name(&name_buf), try ctx.capability(), s.comm().kind(), tp_kernel.mailbox.len, up_ms }); + return switch (command) { + .suite => suite(&s, &driver, init.gpa, cfg), + .failfast => |f| if (f.mode == .transport) + @import("gpu_transport.zig").run(&s, &driver, init.gpa, cfg, f.victim, f.n, f.hold orelse 0) + else + failfast(&s, &driver, init.gpa, cfg, f.victim, @tagName(f.mode)), + else => unreachable, + }; +} + +/// `hca` (host only, no GPU): both ranks' RoCE pick over the bootstrap; this rank's `NCCL_IB_HCA=...` line on stdout. +fn hca(cfg: tp.Config, io: std.Io) !u8 { + var b = tp.Bootstrap.open(cfg) catch |e| { + const s = @as(f64, @floatFromInt(cfg.connect_timeout_ns)) / std.time.ns_per_s; + switch (e) { + // a peer that never runs `hca` (it names its ports, or is not up) leaves this rank waiting until here + error.Timeout => std.debug.print("FAIL hca rank {d}: no peer joined within TF_TP_PLAN_TIMEOUT ({d:.1} s): is its runner up, and leaving TF_TP_ROCE_HCA and NCCL_IB_HCA unset as this one does?\n", .{ cfg.rank, s }), + error.AddressInUse => std.debug.print("FAIL hca rank {d}: port {d} (TF_TP_PORT + 7) is in use on this host\n", .{ cfg.rank, cfg.link_port }), + error.Closed => std.debug.print("FAIL hca rank {d}: the peer closed the bootstrap connection\n", .{cfg.rank}), + else => std.debug.print("FAIL hca rank {d}: bootstrap: {t}\n", .{ cfg.rank, e }), + } + return 1; + }; + defer b.close(); + var specs: [tp.roce_verbs.max_hcas]tp.roce_verbs.HcaSpec = undefined; + const pick = tp.roce_pick.select(cfg, &b, &specs) catch |e| { + std.debug.print("FAIL hca rank {d}: {t} while picking the ports with the peer\n", .{ cfg.rank, e }); + return 1; + }; + var buf: [tp.roce_pick.list_bytes + 16]u8 = undefined; + const line = tp.roce_pick.hcaLine(pick, &specs, &buf) catch |e| { + const why = switch (e) { + error.LibraryUnavailable => "could not load libibverbs.so.1 (LibraryUnavailable); check the RDMA userspace installation", + error.MissingSymbol => "libibverbs.so.1 lacks required symbols (MissingSymbol); check the RDMA userspace installation", + error.NoRoceDevice => "no active Ethernet RDMA port with an IPv4 RoCE v2 GID on this host: name the peer's ports in TF_TP_ROCE_HCA or NCCL_IB_HCA", + error.NoSharedSubnet => "no RoCE port here shares an IPv4 /24 with the peer's: name the peer's ports in TF_TP_ROCE_HCA or NCCL_IB_HCA", + else => @errorName(e), + }; + std.debug.print("FAIL hca rank {d}: {s}\n", .{ cfg.rank, why }); + return 1; + }; + var names: [tp.roce_pick.list_bytes]u8 = undefined; + const how = if (pick.explicit) "TF_TP_ROCE_HCA" else "on a /24 the peer shares"; + std.debug.print("INFO hca rank {d}: {s} ({s})\n", .{ cfg.rank, try tp.roce_pick.joinNames(specs[0..pick.n], true, &names), how }); + try std.Io.File.stdout().writeStreamingAll(io, line); + return 0; +} + +fn suite(s: *tp.Session, d: *const cuda.Driver, gpa: std.mem.Allocator, cfg: tp.Config) !u8 { + var c = try checks.Ctx.init(d, gpa, cfg.rank, max_bytes); + defer c.deinit(); + var backends: [2]struct { comm: tp.Collective, label: []const u8 } = undefined; + var n_backends: usize = 1; + backends[0] = .{ .comm = s.nccl.iface(), .label = "nccl" }; + if (s.comm().kind() != .nccl) { + backends[1] = .{ .comm = s.comm(), .label = if (s.roce != null) "roce" else "mailbox" }; + n_backends = 2; + } + const sizes = [_]usize{ 6, 1000, 4096, 65536, 256 * 1024, 256 * 1024 + 16, 1 << 20, 4 << 20 }; + for (backends[0..n_backends]) |b| { + for (sizes, 0..) |n, i| try c.collectives(b.comm, b.label, n, 100 * i); + try c.pointToPoint(b.comm, b.label, 64 * 1024, 7); + for ([_]usize{ 4096, 65536, 1 << 20 }) |n| try c.graphs(b.comm, b.label, n, 4); + try b.comm.barrier(); + } + for (backends[0..n_backends]) |b| { + try bench.sweep(&c, b.comm, b.label, 4096, 4 << 20); + try bench.decodeWindow(&c, b.comm, b.label, 86, 8192); + } + // the stop: rank 0 marks it expected first, then both meet and close + if (s.fate) |*f| f.expectStop(); + try s.comm().barrier(); + std.debug.print("SUMMARY rank {d}: {d} passed, {d} failed\n", .{ cfg.rank, c.passes, c.failures }); + s.close(); + return if (c.failures == 0) 0 else 1; +} + +/// The failed-transport check's host-only forms (fail-fast off, a hold): TF_TP_ALLOW_HOST_ONLY=1 and no RDMA wire. +fn hostOnlyAllowed(cfg: tp.Config) bool { + const flag = std.mem.eql(u8, std.mem.span(std.c.getenv("TF_TP_ALLOW_HOST_ONLY") orelse ""), "1"); + return flag and (cfg.backend != .roce or cfg.roce_wire == .shm); +} + +pub fn realNow() f64 { + var ts: std.c.timespec = undefined; + _ = std.c.clock_gettime(.REALTIME, &ts); + return @as(f64, @floatFromInt(ts.sec)) + @as(f64, @floatFromInt(ts.nsec)) / 1e9; +} + +/// Exchanges until the victim fails; the survivor is left waiting inside one, which only the fate channel ends. +fn failfast(s: *tp.Session, d: *const cuda.Driver, gpa: std.mem.Allocator, cfg: tp.Config, victim: u32, mode: []const u8) !u8 { + var c = try checks.Ctx.init(d, gpa, cfg.rank, 1 << 20); + defer c.deinit(); + const comm = s.comm(); + var k: u32 = 0; + while (true) : (k += 1) { + if (k == 50 and cfg.rank == victim) { + std.debug.print("INJECT {s} rank {d} at {d:.6}\n", .{ mode, cfg.rank, realNow() }); + if (std.mem.eql(u8, mode, "kill")) _ = std.c.raise(.KILL); + s.fatal("window", error.Injected); + return 3; // not reached: the process exits 70 + } + try comm.allReduce(c.in.ptr, c.out.ptr, 2048, .bf16, .sum, c.stream.handle); + c.stream.synchronize() catch |e| { + std.debug.print("rank {d}: stream error {t} after the peer failed\n", .{ cfg.rank, e }); + }; + comm.check() catch |e| { + std.debug.print("rank {d}: transport reports {t} at exchange {d}; waiting for the fate channel\n", .{ cfg.rank, e, k }); + while (true) tp.sock.sleepNs(std.time.ns_per_s); + }; + } +} diff --git a/zig/src/fabric/tp/gpu_transport.zig b/zig/src/fabric/tp/gpu_transport.zig new file mode 100644 index 0000000000..49cd625844 --- /dev/null +++ b/zig/src/fabric/tp/gpu_transport.zig @@ -0,0 +1,71 @@ +//! `tf-tp-test failfast V transport [N [HOLD_S]]`: the failed-transport path on real HCAs, injected on an idle link. +const std = @import("std"); +const cuda = @import("cuda"); +const tp = @import("tp"); +const checks = @import("gpu_checks.zig"); +const realNow = @import("gpu_main.zig").realNow; + +/// N clean sums, an NCCL meet, the victim's `Roce.inject` once every link is idle, then its close (fail-fast: exit 70). +pub fn run(s: *tp.Session, d: *const cuda.Driver, gpa: std.mem.Allocator, cfg: tp.Config, victim: u32, n: u32, hold_s: u32) !u8 { + const r = s.roce orelse { + std.debug.print("failfast transport needs TF_COMM_BACKEND=roce (the session runs {t})\n", .{s.comm().kind()}); + return 2; + }; + var c = try checks.Ctx.init(d, gpa, cfg.rank, 4096); + defer c.deinit(); + const comm = s.comm(); + for (0..n) |_| { + try comm.allReduce(c.in.ptr, c.out.ptr, 2048, .bf16, .sum, c.stream.handle); + try c.stream.synchronize(); + try comm.check(); + } + try s.nccl.iface().barrier(); + std.debug.print("rank {d}: {d} clean exchanges; both ranks met over NCCL\n", .{ cfg.rank, n }); + // the inject only on an idle link: every rank's posts completed, else no rank injects + const t_idle = tp.sock.nowNs(); + while (!r.linkIdle() and tp.sock.nowNs() - t_idle < 2 * std.time.ns_per_s) tp.sock.sleepNs(std.time.ns_per_ms); + var line: [384]u8 = undefined; + std.debug.print("rank {d}: link {s} before the inject: {s}\n", .{ cfg.rank, if (r.linkIdle()) "idle" else "BUSY", r.describeLink(&line) }); + if (try idleRanks(s, &c, r.linkIdle()) != cfg.world) { + std.debug.print("REFUSED rank {d}: a link was not idle; no inject, a clean stop\n", .{cfg.rank}); + if (s.fate) |*f| f.expectStop(); + try s.nccl.iface().barrier(); + s.close(); + return 2; + } + if (cfg.rank == victim) { + std.debug.print("INJECT transport rank {d} at {d:.6}\n", .{ cfg.rank, realNow() }); + r.inject(); + comm.check() catch |e| std.debug.print("rank {d}: transport reports {t}\n", .{ cfg.rank, e }); + } else if (s.fate != null) { + while (true) tp.sock.sleepNs(std.time.ns_per_s); // the fate channel ends this rank (exit 70) + } + const kept0 = tp.roce.Roce.kept.load(.monotonic); + s.close(); // fail-fast on: exit 70 from here + const kept = tp.roce.Roce.kept.load(.monotonic) - kept0; + const want: u32 = @intFromBool(cfg.rank == victim); + const live = cfg.rank != victim or mapped(r.region); + std.debug.print("rank {d}: closed; regions kept {d} (want {d}){s}\n", .{ cfg.rank, kept, want, if (cfg.rank != victim) "" else if (live) ", region still mapped" else ", region NOT mapped" }); + if (hold_s > 0) { + std.debug.print("HOLD rank {d} pid {d} for {d} s\n", .{ cfg.rank, std.c.getpid(), hold_s }); + tp.sock.sleepNs(@as(u64, hold_s) * std.time.ns_per_s); + } + const ok = kept == want and live; + std.debug.print("{s} rank {d}\n", .{ if (ok) "PASS" else "FAIL", cfg.rank }); + return if (ok) 0 else 1; +} + +fn mapped(region: []const u8) bool { + var vec: [1]u8 = undefined; + return std.os.linux.errno(std.os.linux.mincore(@constCast(region.ptr), 4096, &vec)) == .SUCCESS; +} + +/// The ranks whose link was idle, summed over NCCL (one f32 a rank). +fn idleRanks(s: *tp.Session, c: *checks.Ctx, idle: bool) !u32 { + var v: [1]f32 = .{@floatFromInt(@intFromBool(idle))}; + try c.in.uploadAsync(0, std.mem.asBytes(&v), c.stream.handle); + try s.nccl.iface().allReduce(c.in.ptr, c.out.ptr, 1, .f32, .sum, c.stream.handle); + try c.stream.synchronize(); + try c.out.download(0, std.mem.asBytes(&v)); + return @intFromFloat(v[0]); +} diff --git a/zig/src/fabric/tp/host.zig b/zig/src/fabric/tp/host.zig new file mode 100644 index 0000000000..051daea537 --- /dev/null +++ b/zig/src/fabric/tp/host.zig @@ -0,0 +1,293 @@ +//! In-process ranks over host memory: the `Collective` contract without a GPU (desyncs refused, abort wakes waiters). +const std = @import("std"); +const collective = @import("collective.zig"); +const reference = @import("reference.zig"); +const DevicePtr = collective.DevicePtr; +const DType = collective.DType; +const Op = collective.Op; +const Error = collective.Error; + +pub const max_ranks = 8; + +const Call = enum(u8) { none, all_reduce, all_gather, all_gather_v, exchange, barrier }; + +const Post = struct { call: Call = .none, send: DevicePtr = 0, recv: DevicePtr = 0, count: usize = 0, t: DType = .u8, op: Op = .sum, peer: u32 = 0, lens: DevicePtr = 0 }; + +/// One-way message slot from rank `src` to rank `dst` (send / recv). +const Slot = struct { buf: std.atomic.Value(DevicePtr) = .init(0), bytes: usize = 0, taken: std.atomic.Value(bool) = .init(false) }; + +pub const Group = struct { + world: u32, + /// A wait longer than this is `error.Timeout` (a peer that never calls). + timeout_ns: u64 = 10 * std.time.ns_per_s, + posts: [max_ranks]Post = @splat(.{}), + arrived: std.atomic.Value(u32) = .init(0), + generation: std.atomic.Value(u32) = .init(0), + aborted: std.atomic.Value(bool) = .init(false), + /// Ranks (a bit each) whose transport refuses device lengths: their `all_gather_v` moves whole strides (tests). + var_refuse: u32 = 0, + slots: [max_ranks][max_ranks]Slot = @splat(@splat(.{})), + + pub fn init(world: u32) Group { + std.debug.assert(world >= 1 and world <= max_ranks); + return .{ .world = world }; + } + + pub fn rank(g: *Group, r: u32) Rank { + std.debug.assert(r < g.world); + return .{ .group = g, .me = r }; + } + + /// Sense-reversing barrier over every rank; `error.Aborted` once any rank aborts. + fn meet(g: *Group) Error!void { + const gen = g.generation.load(.acquire); + if (g.arrived.fetchAdd(1, .acq_rel) + 1 == g.world) { + g.arrived.store(0, .release); + _ = g.generation.fetchAdd(1, .acq_rel); + return if (g.aborted.load(.acquire)) error.Aborted else {}; + } + const t0 = now(); + while (g.generation.load(.acquire) == gen) { + if (g.aborted.load(.acquire)) return error.Aborted; + if (now() - t0 > g.timeout_ns) return error.Timeout; + std.Thread.yield() catch {}; + } + if (g.aborted.load(.acquire)) return error.Aborted; + } +}; + +fn now() u64 { + var ts: std.c.timespec = undefined; + _ = std.c.clock_gettime(.MONOTONIC, &ts); + return @as(u64, @intCast(ts.sec)) * std.time.ns_per_s + @as(u64, @intCast(ts.nsec)); +} + +fn bytesAt(p: DevicePtr, n: usize) []u8 { + if (n == 0) return &.{}; + return @as([*]u8, @ptrFromInt(p))[0..n]; +} + +pub const Rank = struct { + group: *Group, + me: u32, + /// `varAgreed`'s answer, once this communicator agreed. + var_agreed: ?bool = null, + + pub fn iface(self: *Rank) collective.Collective { + return .{ .ptr = self, .vtable = &vtable }; + } + + const vtable: collective.Collective.VTable = .{ + .kind = .host, + .rank = rankOf, + .world = worldOf, + .all_reduce = allReduce, + .all_gather = allGather, + .exchange = exchange, + .send = send, + .recv = recv, + .barrier = barrier, + .check = check, + .abort = abort, + .all_gather_v = allGatherV, + .var_fits = varFits, + .var_agreed = varAgreed, + }; + + fn self_(ptr: *anyopaque) *Rank { + return @ptrCast(@alignCast(ptr)); + } + + fn rankOf(ptr: *anyopaque) u32 { + return self_(ptr).me; + } + + fn worldOf(ptr: *anyopaque) u32 { + return self_(ptr).group.world; + } + + /// Posts this rank's call, waits for all, and refuses calls that do not match rank 0's. + fn enter(self: *Rank, p: Post) Error!void { + const g = self.group; + if (g.aborted.load(.acquire)) return error.Aborted; + g.posts[self.me] = p; + try g.meet(); + const r0 = g.posts[0]; + if (r0.call != p.call or r0.count != p.count or r0.t != p.t or r0.op != p.op) { + abort(self); + return error.Invalid; + } + } + + fn allReduce(ptr: *anyopaque, s: DevicePtr, r: DevicePtr, count: usize, t: DType, op: Op, _: collective.Stream) Error!void { + const self = self_(ptr); + const g = self.group; + try self.enter(.{ .call = .all_reduce, .send = s, .recv = r, .count = count, .t = t, .op = op }); + const n = count * t.size(); + var ins: [max_ranks][]const u8 = undefined; + for (0..g.world) |i| ins[i] = bytesAt(g.posts[i].send, n); + // the output may alias this rank's input: reduce into scratch first, then meet before writing + var scratch: [1 << 16]u8 = undefined; + if (n > scratch.len) return error.Unsupported; + if (n > 0) reference.reduce(scratch[0..n], ins[0..g.world], t, op); + try g.meet(); + @memcpy(bytesAt(r, n), scratch[0..n]); + try g.meet(); + } + + fn allGather(ptr: *anyopaque, s: DevicePtr, r: DevicePtr, count: usize, t: DType, _: collective.Stream) Error!void { + const self = self_(ptr); + const g = self.group; + try self.enter(.{ .call = .all_gather, .send = s, .recv = r, .count = count, .t = t }); + const n = count * t.size(); + const out = bytesAt(r, n * g.world); + for (0..g.world) |i| @memcpy(out[i * n ..][0..n], bytesAt(g.posts[i].send, n)); + try g.meet(); + } + + fn lensAt(p: DevicePtr, world: u32) []const i32 { + return @as([*]const i32, @ptrFromInt(p))[0..world]; + } + + /// Copies only each rank's lens[r] bytes (`var_refuse` ranks take whole strides); unequal lengths are refused. + fn allGatherV(ptr: *anyopaque, s: DevicePtr, r: DevicePtr, stride: usize, lens: DevicePtr, t: DType, _: collective.Stream) Error!void { + const self = self_(ptr); + const g = self.group; + try self.enter(.{ .call = .all_gather_v, .send = s, .recv = r, .count = stride, .t = t, .lens = lens }); + const n = stride * t.size(); + const mine = lensAt(lens, g.world); + for (0..g.world) |i| if (!std.mem.eql(i32, mine, lensAt(g.posts[i].lens, g.world))) { + abort(self); + return error.Invalid; + }; + const whole = g.var_refuse & (@as(u32, 1) << @intCast(self.me)) != 0; + const out = bytesAt(r, n * g.world); + for (0..g.world) |i| { + if (mine[i] < 0 or mine[i] > n) { + abort(self); + return error.Invalid; + } + const len: usize = if (whole) n else @intCast(mine[i]); + @memcpy(out[i * n ..][0..len], bytesAt(g.posts[i].send, len)); + } + try g.meet(); + } + + fn varFits(ptr: *anyopaque, bytes: usize) bool { + const self = self_(ptr); + return bytes > 0 and self.group.var_refuse & (@as(u32, 1) << @intCast(self.me)) == 0; + } + + /// One int all-gather of `varFits(16)` the first time, then the kept answer. + fn varAgreed(ptr: *anyopaque) Error!bool { + const self = self_(ptr); + if (self.var_agreed) |v| return v; + const mine: i32 = @intFromBool(varFits(ptr, 16)); + var all: [max_ranks]i32 = undefined; + try allGather(ptr, @intFromPtr(&mine), @intFromPtr(&all), 1, .i32, null); + self.var_agreed = std.mem.min(i32, all[0..self.group.world]) == 1; + return self.var_agreed.?; + } + + fn exchange(ptr: *anyopaque, s: DevicePtr, r: DevicePtr, count: usize, t: DType, peer: u32, _: collective.Stream) Error!void { + const self = self_(ptr); + const g = self.group; + try collective.checkPeer(self.me, g.world, peer); + try self.enter(.{ .call = .exchange, .send = s, .recv = r, .count = count, .t = t, .peer = peer }); + if (g.posts[peer].peer != self.me) { + abort(self); + return error.Invalid; + } + const n = count * t.size(); + @memcpy(bytesAt(r, n), bytesAt(g.posts[peer].send, n)); + try g.meet(); + } + + fn send(ptr: *anyopaque, buf: DevicePtr, count: usize, t: DType, peer: u32, _: collective.Stream) Error!void { + const self = self_(ptr); + const g = self.group; + try collective.checkPeer(self.me, g.world, peer); + const slot = &g.slots[self.me][peer]; + slot.bytes = count * t.size(); + slot.taken.store(false, .release); + slot.buf.store(if (buf == 0) 1 else buf, .release); + try waitFor(g, &slot.taken, true); + slot.buf.store(0, .release); + } + + fn recv(ptr: *anyopaque, buf: DevicePtr, count: usize, t: DType, peer: u32, _: collective.Stream) Error!void { + const self = self_(ptr); + const g = self.group; + try collective.checkPeer(self.me, g.world, peer); + const slot = &g.slots[peer][self.me]; + const t0 = now(); + var src = slot.buf.load(.acquire); + while (src == 0) : (src = slot.buf.load(.acquire)) { + if (g.aborted.load(.acquire)) return error.Aborted; + if (now() - t0 > g.timeout_ns) return error.Timeout; + std.Thread.yield() catch {}; + } + const n = count * t.size(); + if (slot.bytes != n) { + abort(self); + return error.Invalid; + } + @memcpy(bytesAt(buf, n), bytesAt(src, n)); + slot.taken.store(true, .release); + } + + fn waitFor(g: *Group, flag: *std.atomic.Value(bool), want: bool) Error!void { + const t0 = now(); + while (flag.load(.acquire) != want) { + if (g.aborted.load(.acquire)) return error.Aborted; + if (now() - t0 > g.timeout_ns) return error.Timeout; + std.Thread.yield() catch {}; + } + } + + fn barrier(ptr: *anyopaque) Error!void { + const self = self_(ptr); + try self.enter(.{ .call = .barrier }); + try self.group.meet(); + } + + fn check(ptr: *anyopaque) Error!void { + if (self_(ptr).group.aborted.load(.acquire)) return error.Aborted; + } + + fn abort(ptr: *anyopaque) void { + self_(ptr).group.aborted.store(true, .release); + } +}; + +/// Runs `body(rank_collective, rank, ctx)` on `world` threads at once and returns the first error any rank hit. +pub fn run(world: u32, ctx: anytype, comptime body: fn (collective.Collective, u32, @TypeOf(ctx)) anyerror!void) !void { + return runWith(world, .{}, ctx, body); +} + +pub const Options = struct { + /// `Group.var_refuse`: ranks whose transport takes no device lengths. + var_refuse: u32 = 0, +}; + +pub fn runWith(world: u32, opts: Options, ctx: anytype, comptime body: fn (collective.Collective, u32, @TypeOf(ctx)) anyerror!void) !void { + var group = Group.init(world); + group.var_refuse = opts.var_refuse; + var ranks: [max_ranks]Rank = undefined; + var errs: [max_ranks]?anyerror = @splat(null); + var threads: [max_ranks]std.Thread = undefined; + const Wrap = struct { + fn go(r: *Rank, c: @TypeOf(ctx), e: *?anyerror) void { + body(r.iface(), r.me, c) catch |err| { + e.* = err; + r.group.aborted.store(true, .release); + }; + } + }; + for (0..world) |i| { + ranks[i] = group.rank(@intCast(i)); + threads[i] = try std.Thread.spawn(.{}, Wrap.go, .{ &ranks[i], ctx, &errs[i] }); + } + for (threads[0..world]) |t| t.join(); + for (errs[0..world]) |e| if (e) |err| return err; +} diff --git a/zig/src/fabric/tp/host_test.zig b/zig/src/fabric/tp/host_test.zig new file mode 100644 index 0000000000..7de32008af --- /dev/null +++ b/zig/src/fabric/tp/host_test.zig @@ -0,0 +1,154 @@ +//! The `Collective` contract on the in-process ranks: every call's bits against the reference, desyncs and aborts. +const std = @import("std"); +const collective = @import("collective.zig"); +const reference = @import("reference.zig"); +const host = @import("host.zig"); +const Collective = collective.Collective; +const DType = collective.DType; + +fn addr(b: []u8) collective.DevicePtr { + return @intFromPtr(b.ptr); +} + +const n = 96; + +const Calls = struct { + fn go(c: Collective, rank: u32, _: void) !void { + const world = c.world(); + inline for (.{ DType.f32, DType.bf16, DType.i32, DType.f16 }) |t| { + const bytes = n * t.size(); + var mine: [n * 8]u8 = undefined; + var out: [n * 8 * host.max_ranks]u8 = undefined; + reference.fill(mine[0..bytes], t, rank, 7); + // all-reduce against the reference sum of every rank's input, in rank order + var ins: [host.max_ranks][n * 8]u8 = undefined; + var views: [host.max_ranks][]const u8 = undefined; + for (0..world) |r| { + reference.fill(ins[r][0..bytes], t, @intCast(r), 7); + views[r] = ins[r][0..bytes]; + } + var want: [n * 8]u8 = undefined; + reference.reduce(want[0..bytes], views[0..world], t, .sum); + try c.allReduce(addr(&mine), addr(&out), n, t, .sum, null); + try std.testing.expectEqualSlices(u8, want[0..bytes], out[0..bytes]); + // in place + try c.allReduce(addr(&mine), addr(&mine), n, t, .sum, null); + try std.testing.expectEqualSlices(u8, want[0..bytes], mine[0..bytes]); + // all-gather: rank order + reference.fill(mine[0..bytes], t, rank, 7); + try c.allGather(addr(&mine), addr(&out), n, t, null); + for (0..world) |r| try std.testing.expectEqualSlices(u8, views[r], out[r * bytes ..][0..bytes]); + } + try c.barrier(); + } +}; + +test "every call on two and three in-process ranks equals the reference" { + try host.run(2, {}, Calls.go); + try host.run(3, {}, Calls.go); +} + +const Pairwise = struct { + fn go(c: Collective, rank: u32, _: void) !void { + var mine: [64]u8 = @splat(@intCast(rank + 1)); + var got: [64]u8 = undefined; + const peer = 1 - rank; + try c.exchange(addr(&mine), addr(&got), 64, .u8, peer, null); + try std.testing.expectEqualSlices(u8, &@as([64]u8, @splat(@intCast(peer + 1))), &got); + // send / recv, ordered so neither waits forever + if (rank == 0) { + try c.send(addr(&mine), 16, .f32, 1, null); + } else { + var buf: [64]u8 = undefined; + try c.recv(addr(&buf), 16, .f32, 0, null); + try std.testing.expectEqualSlices(u8, &@as([64]u8, @splat(1)), &buf); + } + try std.testing.expectError(error.Invalid, c.exchange(addr(&mine), addr(&got), 64, .u8, rank, null)); + } +}; + +test "exchange, send and recv between two ranks" { + try host.run(2, {}, Pairwise.go); +} + +const Desync = struct { + fn go(c: Collective, rank: u32, _: void) !void { + var a: [64]u8 = @splat(0); + var b: [128]u8 = undefined; + // the ranks disagree on the count: both see it instead of hanging + const r = c.allGather(addr(&a), addr(&b), if (rank == 0) 16 else 32, .u8, null); + if (r) |_| return error.ShouldHaveFailed else |e| if (e != error.Invalid and e != error.Aborted) return e; + } +}; + +test "a desynchronised call is refused on every rank" { + try host.run(2, {}, Desync.go); +} + +const Aborts = struct { + fn go(c: Collective, rank: u32, _: void) !void { + var a: [16]u8 = @splat(0); + var b: [32]u8 = undefined; + if (rank == 1) { + // rank 1 never calls: it aborts after a while, which must free rank 0's waiting call + @import("sock.zig").sleepNs(50 * std.time.ns_per_ms); + c.abort(); + return; + } + try std.testing.expectError(error.Aborted, c.allGather(addr(&a), addr(&b), 16, .u8, null)); + try std.testing.expectError(error.Aborted, c.check()); + try std.testing.expectError(error.Aborted, c.barrier()); + } +}; + +test "abort frees a rank waiting in a collective" { + try host.run(2, {}, Aborts.go); +} + +const VarGather = struct { + fn go(c: Collective, rank: u32, refuse: u32) !void { + const world = c.world(); + const stride = 200; + const whole = refuse & (@as(u32, 1) << @intCast(rank)) != 0; + try std.testing.expectEqual(!whole, c.varFits(stride)); + try std.testing.expectEqual(refuse == 0, try c.varAgreed()); + try std.testing.expectEqual(refuse == 0, try c.varAgreed()); // kept: no second collective + var ins: [host.max_ranks][stride]u8 = undefined; + for (0..world) |r| reference.fill(&ins[r], .u8, @intCast(r), 11); + // lengths of 0, odd, 16 and the whole stride, in bytes; .f32 strides count elements + const cases = [_][3]i32{ .{ 0, 17, 200 }, .{ 200, 0, 3 }, .{ 16, 16, 16 }, .{ 199, 200, 1 } }; + inline for (.{ DType.u8, DType.f32 }) |t| for (cases) |lens| { + var out: [stride * host.max_ranks]u8 = @splat(0xEE); + try c.allGatherV(addr(&ins[rank]), addr(&out), stride / t.size(), @intFromPtr(&lens), t, null); + for (0..world) |r| { + const len: usize = @intCast(lens[r]); + try std.testing.expectEqualSlices(u8, ins[r][0..len], out[r * stride ..][0..len]); + // past a length: the transport's choice (untouched here, or the whole stride where lengths are refused) + const past = out[r * stride + len .. (r + 1) * stride]; + if (whole) try std.testing.expectEqualSlices(u8, ins[r][len..], past) else for (past) |x| try std.testing.expectEqual(@as(u8, 0xEE), x); + } + }; + try c.barrier(); + } +}; + +test "a variable-length all-gather moves each rank's length, and the ranks agree on device lengths once" { + for ([_]u32{ 2, 3 }) |w| { + try host.runWith(w, .{}, @as(u32, 0), VarGather.go); + try host.runWith(w, .{ .var_refuse = 2 }, @as(u32, 2), VarGather.go); + } +} + +const VarDesync = struct { + fn go(c: Collective, rank: u32, _: void) !void { + var a: [64]u8 = @splat(0); + var b: [128]u8 = undefined; + const lens = [2]i32{ 8, if (rank == 0) 8 else 16 }; + const r = c.allGatherV(addr(&a), addr(&b), 64, @intFromPtr(&lens), .u8, null); + if (r) |_| return error.ShouldHaveFailed else |e| if (e != error.Invalid and e != error.Aborted) return e; + } +}; + +test "lengths that differ between ranks are refused on every rank" { + try host.run(2, {}, VarDesync.go); +} diff --git a/zig/src/fabric/tp/hybrid.zig b/zig/src/fabric/tp/hybrid.zig new file mode 100644 index 0000000000..e5f6d72c31 --- /dev/null +++ b/zig/src/fabric/tp/hybrid.zig @@ -0,0 +1,183 @@ +//! Small exchanges through the mailbox (one host) or RoCE (two hosts), the rest through NCCL; ranks always agree. +const std = @import("std"); +const collective = @import("collective.zig"); +const Mailbox = @import("mailbox.zig").Mailbox; +const Roce = @import("roce.zig").Roce; +const Nccl = @import("nccl.zig").Nccl; +const DevicePtr = collective.DevicePtr; +const DType = collective.DType; +const Op = collective.Op; +const Error = collective.Error; + +/// The one-kernel path for small shards: both expose the same calls. +pub const Small = union(enum) { + mailbox: *Mailbox, + roce: *Roce, + + fn fits(s: Small, n: usize) bool { + return switch (s) { + inline else => |p| p.fits(n), + }; + } + + fn sum(s: Small, src: DevicePtr, dst: DevicePtr, count: usize, t: DType, stream: collective.Stream) Error!void { + return switch (s) { + inline else => |p| p.sum(src, dst, count, t, stream), + }; + } + + fn gather(s: Small, src: DevicePtr, dst: DevicePtr, n: usize, stream: collective.Stream) Error!void { + return switch (s) { + inline else => |p| p.allGatherBytes(src, dst, n, stream), + }; + } + + fn exchange(s: Small, src: DevicePtr, dst: DevicePtr, n: usize, stream: collective.Stream) Error!void { + return switch (s) { + inline else => |p| p.exchangeBytes(src, dst, n, stream), + }; + } + + fn check(s: Small) Error!void { + return switch (s) { + inline else => |p| p.check(), + }; + } + + fn abort(s: Small) void { + switch (s) { + inline else => |p| p.abort(), + } + } +}; + +/// Whether a variable-length gather of `n`-byte strides one op cannot take goes RoCE in slot-sized ops (needs the ABI). +pub fn chunkRoute(chunked_bytes: usize, n: usize, single: bool, abi: bool) bool { + return abi and !single and n > 0 and n <= chunked_bytes; +} + +pub const Hybrid = struct { + small: Small, + nccl: *Nccl, + /// Largest shard a rank sends through the small path. + max_bytes: usize, + /// Largest stride a variable-length gather takes on RoCE in slot-sized ops (TF_TP_ROCE_CHUNKED_KB; 0: none). + chunked_bytes: usize = 0, + /// Every rank's `varFits(16)`, agreed at `Session.open` (the bootstrap's all-gather): `varAgreed`'s answer. + var_agreed: bool = false, + + pub fn iface(h: *Hybrid) collective.Collective { + return .{ .ptr = h, .vtable = if (h.small == .roce) &roce_vtable else &mailbox_vtable }; + } + + fn vtable(comptime kind: collective.Kind) collective.Collective.VTable { + return .{ + .kind = kind, + .rank = rankOf, + .world = worldOf, + .all_reduce = allReduce, + .all_gather = allGather, + .exchange = exchange, + .send = send, + .recv = recv, + .barrier = barrier, + .check = check, + .abort = abort, + .all_gather_v = allGatherV, + .var_fits = varFits, + .var_agreed = varAgreed, + }; + } + + const mailbox_vtable = vtable(.hybrid); + const roce_vtable = vtable(.roce); + + fn self_(ptr: *anyopaque) *Hybrid { + return @ptrCast(@alignCast(ptr)); + } + + fn routed(h: *const Hybrid, n: usize) bool { + return n <= h.max_bytes and h.small.fits(n); + } + + fn rankOf(ptr: *anyopaque) u32 { + return self_(ptr).nccl.me; + } + + fn worldOf(ptr: *anyopaque) u32 { + return self_(ptr).nccl.size; + } + + fn allReduce(ptr: *anyopaque, s: DevicePtr, r: DevicePtr, count: usize, t: DType, op: Op, stream: collective.Stream) Error!void { + const h = self_(ptr); + if (op == .sum and (t == .f32 or t == .bf16) and h.routed(count * t.size())) return h.small.sum(s, r, count, t, stream); + return Nccl.allReduce(h.nccl, s, r, count, t, op, stream); + } + + fn allGather(ptr: *anyopaque, s: DevicePtr, r: DevicePtr, count: usize, t: DType, stream: collective.Stream) Error!void { + const h = self_(ptr); + if (h.routed(count * t.size())) return h.small.gather(s, r, count * t.size(), stream); + return Nccl.allGather(h.nccl, s, r, count, t, stream); + } + + /// RoCE moves each rank's device length; the mailbox and NCCL move whole strides (the same bytes where lens says). + fn allGatherV(ptr: *anyopaque, s: DevicePtr, r: DevicePtr, stride: usize, lens: DevicePtr, t: DType, stream: collective.Stream) Error!void { + const h = self_(ptr); + const n = stride * t.size(); + if (singleFits(h, n)) return h.small.roce.allGatherV(s, r, n, lens, stream); + if (varFits(ptr, n)) return h.small.roce.allGatherVChunked(s, r, n, lens, stream); + return allGather(ptr, s, r, stride, t, stream); + } + + fn singleFits(h: *const Hybrid, n: usize) bool { + return switch (h.small) { + .roce => |p| n > 0 and h.routed(n) and p.d != null, + .mailbox => false, + }; + } + + /// RoCE with its device side, in one op or in slot-sized ops (the route is the size's alone: every rank alike). + pub fn varFits(ptr: *anyopaque, n: usize) bool { + const h = self_(ptr); + const single = singleFits(h, n); + return single or switch (h.small) { + .roce => |p| chunkRoute(h.chunked_bytes, n, single, p.d != null and p.chunk_abi), + .mailbox => false, + }; + } + + fn varAgreed(ptr: *anyopaque) Error!bool { + return self_(ptr).var_agreed; + } + + fn exchange(ptr: *anyopaque, s: DevicePtr, r: DevicePtr, count: usize, t: DType, peer: u32, stream: collective.Stream) Error!void { + const h = self_(ptr); + try collective.checkPeer(h.nccl.me, h.nccl.size, peer); + if (h.routed(count * t.size())) return h.small.exchange(s, r, count * t.size(), stream); + return Nccl.exchange(h.nccl, s, r, count, t, peer, stream); + } + + fn send(ptr: *anyopaque, buf: DevicePtr, count: usize, t: DType, peer: u32, stream: collective.Stream) Error!void { + return Nccl.iface(self_(ptr).nccl).send(buf, count, t, peer, stream); + } + + fn recv(ptr: *anyopaque, buf: DevicePtr, count: usize, t: DType, peer: u32, stream: collective.Stream) Error!void { + return Nccl.iface(self_(ptr).nccl).recv(buf, count, t, peer, stream); + } + + fn barrier(ptr: *anyopaque) Error!void { + return Nccl.iface(self_(ptr).nccl).barrier(); + } + + fn check(ptr: *anyopaque) Error!void { + const h = self_(ptr); + try h.small.check(); + try Nccl.iface(h.nccl).check(); + } + + fn abort(ptr: *anyopaque) void { + const h = self_(ptr); + h.small.abort(); + Nccl.iface(h.nccl).abort(); + } +}; diff --git a/zig/src/fabric/tp/log.zig b/zig/src/fabric/tp/log.zig new file mode 100644 index 0000000000..4e3234ab49 --- /dev/null +++ b/zig/src/fabric/tp/log.zig @@ -0,0 +1,15 @@ +//! TP's log lines, silent under test: a test's stderr is the build runner's channel. +const std = @import("std"); +const quiet = @import("builtin").is_test; + +pub fn err(comptime format: []const u8, args: anytype) void { + if (!quiet) std.log.err(format, args); +} + +pub fn warn(comptime format: []const u8, args: anytype) void { + if (!quiet) std.log.warn(format, args); +} + +pub fn info(comptime format: []const u8, args: anytype) void { + if (!quiet) std.log.info(format, args); +} diff --git a/zig/src/fabric/tp/mailbox.cu b/zig/src/fabric/tp/mailbox.cu new file mode 100644 index 0000000000..843b4f1af2 --- /dev/null +++ b/zig/src/fabric/tp/mailbox.cu @@ -0,0 +1,564 @@ +// One-shot two-rank exchange through pinned host memory both GPUs map (the TP mailbox; mailbox.zig is its host side). + +#include +#include + +// Protocol adapted from b12x's RoCEnante (Apache-2.0): a device epoch, slots by seq parity, a seq flag after the data. +namespace tftp { + +constexpr int kMaxBlocks = 16; +constexpr int kFlagStride = 32; // words: one 128-byte line per flag +// ctrl words: [0] abort (any rank, any host thread); per rank r at 16 + 8 r: failed, seq, seen, block, why, waited us +constexpr int kCtrlAbort = 0; +constexpr int kCtrlRank = 16; + +enum Mode : uint32_t { kGather = 0, kSumF32 = 1, kSumBf16 = 2, kExchange = 3 }; + +struct MailArgs { + const uint8_t *in; // this rank's shard (device) + uint8_t *out; // gather [2 n], exchange / sums [n] + uint8_t *peer_slots; // the peer's inbox data, as this GPU maps it + uint32_t *peer_flags; // the peer's inbox flags + const uint8_t *my_slots; // this rank's inbox data + const uint32_t *my_flags; // this rank's inbox flags + uint32_t *state; // device: [0] epoch, [1] tail counter, [2] poison + uint32_t *ctrl; // shared control words + uint64_t nbytes; // bytes a rank + uint64_t slot_bytes; + uint64_t timeout_ns; + uint32_t rank; + uint32_t mode; +}; + +__device__ __forceinline__ uint32_t ld_relaxed_gpu(const uint32_t *p) { + uint32_t v; + asm volatile("ld.relaxed.gpu.global.u32 %0, [%1];" : "=r"(v) : "l"(p) : "memory"); + return v; +} + +__device__ __forceinline__ uint32_t ld_relaxed_sys(const uint32_t *p) { + uint32_t v; + asm volatile("ld.relaxed.sys.global.u32 %0, [%1];" : "=r"(v) : "l"(p) : "memory"); + return v; +} + +__device__ __forceinline__ uint32_t ld_acquire_sys(const uint32_t *p) { + uint32_t v; + asm volatile("ld.acquire.sys.global.u32 %0, [%1];" : "=r"(v) : "l"(p) : "memory"); + return v; +} + +__device__ __forceinline__ uint4 ld_relaxed_sys_v4(const void *p) { + uint4 v; + asm volatile("ld.relaxed.sys.global.v4.u32 {%0, %1, %2, %3}, [%4];" + : "=r"(v.x), "=r"(v.y), "=r"(v.z), "=r"(v.w) : "l"(p) : "memory"); + return v; +} + +__device__ __forceinline__ uint32_t ld_relaxed_sys_u8(const uint8_t *p) { + uint32_t v; + asm volatile("ld.relaxed.sys.global.u8 %0, [%1];" : "=r"(v) : "l"(p) : "memory"); + return v; +} + +__device__ __forceinline__ void st_relaxed_sys(uint32_t *p, uint32_t v) { + asm volatile("st.relaxed.sys.global.u32 [%0], %1;" ::"l"(p), "r"(v) : "memory"); +} + +__device__ __forceinline__ void st_release_sys(uint32_t *p, uint32_t v) { + asm volatile("st.release.sys.global.u32 [%0], %1;" ::"l"(p), "r"(v) : "memory"); +} + +__device__ __forceinline__ void st_release_gpu(uint32_t *p, uint32_t v) { + asm volatile("st.release.gpu.global.u32 [%0], %1;" ::"l"(p), "r"(v) : "memory"); +} + +__device__ __forceinline__ uint32_t atom_add_acq_rel_gpu(uint32_t *p, uint32_t v) { + uint32_t old; + asm volatile("atom.acq_rel.gpu.global.add.u32 %0, [%1], %2;" : "=r"(old) : "l"(p), "r"(v) : "memory"); + return old; +} + +__device__ __forceinline__ uint32_t atom_add_acq_rel_sys(uint32_t *p, uint32_t v) { + uint32_t old; + asm volatile("atom.acq_rel.sys.global.add.u32 %0, [%1], %2;" : "=r"(old) : "l"(p), "r"(v) : "memory"); + return old; +} + +__device__ __forceinline__ uint32_t atom_cas_gpu(uint32_t *p, uint32_t cmp, uint32_t v) { + uint32_t old; + asm volatile("atom.relaxed.gpu.global.cas.b32 %0, [%1], %2, %3;" : "=r"(old) : "l"(p), "r"(cmp), "r"(v) : "memory"); + return old; +} + +__device__ __forceinline__ void fence_sc_sys() { asm volatile("fence.sc.sys;" ::: "memory"); } + +__device__ __forceinline__ uint64_t globaltimer() { + uint64_t t; + asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t)); + return t; +} + +// dst[i] <- src[i] for i in [0, n), grid-stride, 16-byte vectors when aligned; SysSrc: system-scope loads. +template +__device__ __forceinline__ void copy_bytes(uint8_t *dst, const uint8_t *src, uint64_t n, uint64_t index, + uint64_t stride) { + const uint64_t a = reinterpret_cast(dst) | reinterpret_cast(src) | n; + if ((a & 15u) == 0u) { + const uint64_t packs = n >> 4; + for (uint64_t i = index; i < packs; i += stride) { + const uint4 v = SysSrc ? ld_relaxed_sys_v4(src + (i << 4)) : reinterpret_cast(src)[i]; + reinterpret_cast(dst)[i] = v; + } + } else { + for (uint64_t i = index; i < n; i += stride) { + dst[i] = static_cast(SysSrc ? ld_relaxed_sys_u8(src + i) : src[i]); + } + } +} + +// TF_TP_ROCE_FAST: copy_bytes with two grid strides' loads issued before their stores (same bytes, one round trip). +template +__device__ __forceinline__ void copy_bytes2(uint8_t *dst, const uint8_t *src, uint64_t n, uint64_t index, + uint64_t stride) { + const uint64_t a = reinterpret_cast(dst) | reinterpret_cast(src) | n; + if ((a & 15u) != 0u) { + copy_bytes(dst, src, n, index, stride); + return; + } + const uint64_t packs = n >> 4; + for (uint64_t i = index; i < packs; i += 2u * stride) { + const uint64_t j = i + stride; + uint4 v0, v1 = make_uint4(0u, 0u, 0u, 0u); + if (SysSrc) { + v0 = ld_relaxed_sys_v4(src + (i << 4)); + if (j < packs) v1 = ld_relaxed_sys_v4(src + (j << 4)); + } else { + v0 = reinterpret_cast(src)[i]; + if (j < packs) v1 = reinterpret_cast(src)[j]; + } + reinterpret_cast(dst)[i] = v0; + if (j < packs) reinterpret_cast(dst)[j] = v1; + } +} + +template +__device__ __forceinline__ void copy_any(bool fast, uint8_t *dst, const uint8_t *src, uint64_t n, uint64_t index, + uint64_t stride) { + if (fast) { + copy_bytes2(dst, src, n, index, stride); + } else { + copy_bytes(dst, src, n, index, stride); + } +} + +__device__ __forceinline__ uint32_t add_bf16x2(uint32_t x, uint32_t y) { + const float2 a = __bfloat1622float2(*reinterpret_cast(&x)); + const float2 b = __bfloat1622float2(*reinterpret_cast(&y)); + const __nv_bfloat162 r = __floats2bfloat162_rn(a.x + b.x, a.y + b.y); + return *reinterpret_cast(&r); +} + +__device__ __forceinline__ uint32_t add_f32(uint32_t x, uint32_t y) { + return __float_as_uint(__uint_as_float(x) + __uint_as_float(y)); +} + +__device__ __forceinline__ uint32_t ld_relaxed_sys_u16(const uint8_t *p) { + uint16_t v; + asm volatile("ld.relaxed.sys.global.u16 %0, [%1];" : "=h"(v) : "l"(p) : "memory"); + return v; +} + +__device__ __forceinline__ uint16_t add_bf16(uint16_t x, uint16_t y) { + const __nv_bfloat16 r = __float2bfloat16_rn(__bfloat162float(__ushort_as_bfloat16(x)) + + __bfloat162float(__ushort_as_bfloat16(y))); + return __bfloat16_as_ushort(r); +} + +// out <- mine + peer's, rank 0's first; IEEE addition commutes, so both ranks get the same bits. +template +__device__ __forceinline__ void sum_bytes(uint8_t *out, const uint8_t *mine, const uint8_t *peer, uint64_t n, + uint32_t rank, uint64_t index, uint64_t stride) { + const uint64_t al = reinterpret_cast(out) | reinterpret_cast(mine) | + reinterpret_cast(peer) | n; + if ((al & 15u) == 0u) { + const uint64_t packs = n >> 4; + for (uint64_t i = index; i < packs; i += stride) { + const uint4 m = reinterpret_cast(mine)[i]; + const uint4 p = ld_relaxed_sys_v4(peer + (i << 4)); + const uint4 a = rank == 0 ? m : p; + const uint4 b = rank == 0 ? p : m; + uint4 r; + r.x = Bf16 ? add_bf16x2(a.x, b.x) : add_f32(a.x, b.x); + r.y = Bf16 ? add_bf16x2(a.y, b.y) : add_f32(a.y, b.y); + r.z = Bf16 ? add_bf16x2(a.z, b.z) : add_f32(a.z, b.z); + r.w = Bf16 ? add_bf16x2(a.w, b.w) : add_f32(a.w, b.w); + reinterpret_cast(out)[i] = r; + } + } else if (Bf16) { + for (uint64_t i = index; i < (n >> 1); i += stride) { + const uint16_t m = reinterpret_cast(mine)[i]; + const uint16_t p = static_cast(ld_relaxed_sys_u16(peer + (i << 1))); + reinterpret_cast(out)[i] = rank == 0 ? add_bf16(m, p) : add_bf16(p, m); + } + } else { + for (uint64_t i = index; i < (n >> 2); i += stride) { + const uint32_t m = reinterpret_cast(mine)[i]; + const uint32_t p = ld_relaxed_sys(reinterpret_cast(peer) + i); + reinterpret_cast(out)[i] = rank == 0 ? add_f32(m, p) : add_f32(p, m); + } + } +} + +// Thread 0 of a block: spin until the peer block's flag holds seq; false on abort (why 1) or timeout (why 2). +__device__ __noinline__ bool wait_flag(const uint32_t *flag, uint32_t seq, const uint32_t *ctrl, uint64_t timeout_ns, + uint32_t *seen, uint32_t *why, uint32_t *waited_us) { + const uint64_t t0 = globaltimer(); + uint32_t polls = 0; + while (true) { + const uint32_t v = ld_acquire_sys(flag); + if (v == seq) { + *seen = v; + return true; + } + if ((++polls & 63u) == 0u) { + const uint64_t now = globaltimer(); + const bool aborted = ld_relaxed_sys(ctrl + kCtrlAbort) != 0u; + const bool late = now > t0 && now - t0 > timeout_ns; + if (aborted || late) { + *seen = ld_acquire_sys(flag); + *why = aborted ? 1u : 2u; + *waited_us = now > t0 ? static_cast((now - t0) / 1000u) : 0u; + return *seen == seq; + } + } + } +} + +// TF_TP_ROCE_FAST: wait_flag for one of several pollers of a flag; `done` (shared) ends the others once one saw it. +__device__ __noinline__ bool wait_flag_shared(const uint32_t *flag, uint32_t seq, const uint32_t *ctrl, + uint64_t timeout_ns, uint32_t *seen, uint32_t *why, uint32_t *waited_us, + volatile uint32_t *done) { + const uint64_t t0 = globaltimer(); + uint32_t polls = 0; + while (true) { + const uint32_t v = ld_acquire_sys(flag); + if (v == seq) { + *done = 1u; + *seen = v; + return true; + } + if (*done != 0u) { + *seen = seq; + return true; + } + if ((++polls & 63u) == 0u) { + const uint64_t now = globaltimer(); + const bool aborted = ld_relaxed_sys(ctrl + kCtrlAbort) != 0u; + const bool late = now > t0 && now - t0 > timeout_ns; + if (aborted || late) { + *seen = ld_acquire_sys(flag); + *why = aborted ? 1u : 2u; + *waited_us = now > t0 ? static_cast((now - t0) / 1000u) : 0u; + if (*seen == seq) { + *done = 1u; + return true; + } + return *done != 0u; + } + } + } +} + +// One failure record a rank (the first failing wait's): its fields, then its valid word, each fenced at system scope. +__device__ __forceinline__ void fail_record(uint32_t *poison, uint32_t *rec, uint32_t seq, uint32_t seen, uint32_t who, + uint32_t why, uint32_t waited) { + if (atom_cas_gpu(poison, 0u, 1u) == 0u) { + st_relaxed_sys(rec + 1, seq); + st_relaxed_sys(rec + 2, seen); + st_relaxed_sys(rec + 3, who); + st_relaxed_sys(rec + 4, why); + st_relaxed_sys(rec + 5, waited); + fence_sc_sys(); + st_relaxed_sys(rec, 1u); + fence_sc_sys(); + } +} + +} // namespace tftp + +using namespace tftp; + +extern "C" __global__ void __launch_bounds__(512) tp_mailbox(MailArgs a) { + __shared__ uint32_t s_head[2]; // poison, seq + __shared__ uint32_t s_bad; + const uint32_t tid = threadIdx.x; + const uint32_t gdim = gridDim.x; + const uint64_t index = static_cast(blockIdx.x) * blockDim.x + tid; + const uint64_t stride = static_cast(gdim) * blockDim.x; + + // a poisoned transport does nothing (the host raises on its next check); poison and epoch in one round trip + if (tid == 0) { + s_head[0] = ld_relaxed_gpu(a.state + 2); + s_head[1] = ld_relaxed_gpu(a.state) + 1u; + s_bad = 0u; + } + __syncthreads(); + if (s_head[0] != 0u) { + return; + } + const uint32_t seq = s_head[1]; + const uint64_t slot = seq & 1u; + const uint64_t flag_at = (slot * kMaxBlocks + blockIdx.x) * kFlagStride; + + // 1. this block's share of the shard into the peer's inbox, then its flag after it (system-scope release) + copy_bytes(a.peer_slots + slot * a.slot_bytes, a.in, a.nbytes, index, stride); + __syncthreads(); + if (tid == 0) { + fence_sc_sys(); + st_release_sys(a.peer_flags + flag_at, seq); + } + // the local shard while the peer's writes travel + if (a.mode == kGather) { + copy_bytes(a.out + static_cast(a.rank) * a.nbytes, a.in, a.nbytes, index, stride); + } + + // 2. the peer's block of the same index (it wrote exactly the bytes this block reads) + if (tid == 0) { + uint32_t seen = 0u, why = 0u, waited = 0u; + if (!wait_flag(a.my_flags + flag_at, seq, a.ctrl, a.timeout_ns, &seen, &why, &waited)) { + s_bad = 1u; + fail_record(a.state + 2, a.ctrl + kCtrlRank + 8u * a.rank, seq, seen, blockIdx.x, why, waited); + } + } + __syncthreads(); + + // 3. copy out, gated by this block's own wait + if (s_bad == 0u) { + const uint8_t *inbox = a.my_slots + slot * a.slot_bytes; + switch (a.mode) { + case kGather: + copy_bytes(a.out + static_cast(1u - a.rank) * a.nbytes, inbox, a.nbytes, index, + stride); + break; + case kExchange: + copy_bytes(a.out, inbox, a.nbytes, index, stride); + break; + case kSumF32: + sum_bytes(a.out, a.in, inbox, a.nbytes, a.rank, index, stride); + break; + default: + sum_bytes(a.out, a.in, inbox, a.nbytes, a.rank, index, stride); + break; + } + } + + // 4. the last block publishes the epoch, unless a wait of this launch failed (later launches stay no-ops) + __syncthreads(); + if (tid == 0) { + // the counter is back at 0 for the next launch (stream order), whatever its grid + const bool last = gdim == 1u || atom_add_acq_rel_gpu(a.state + 1, 1u) == gdim - 1u; + const bool ok = gdim == 1u ? s_bad == 0u : ld_relaxed_gpu(a.state + 2) == 0u; + if (last && gdim > 1u) { + st_release_gpu(a.state + 1, 0u); + } + if (last && ok) { + st_release_gpu(a.state, seq); + } + } +} + +// The RoCE variant (roce.zig): blocks stage into the send slot, the last rings the doorbell, a host proxy RDMA-writes. +namespace tftp { + +constexpr int kRoceHcas = 2; +// ctrl words shared with roce.zig's proxy: [0] doorbell seq, [4 + slot] padded bytes, [8] completed, [16..21] failure +constexpr int kRoceSeq = 0; +constexpr int kRoceSlotBytes = 4; +constexpr int kRoceCompleted = 8; +constexpr int kRoceFailed = 16; +constexpr int kRoceAbort = 12; // wait_flag reads ctrl[kCtrlAbort]: the RoCE abort word is passed as that ctrl + +struct RoceArgs { + const uint8_t *in; + uint8_t *out; + uint8_t *send_slots; // this rank's send slots (pinned host, read by the NIC) + const uint8_t *recv_slots; // this rank's receive slots (written by the peer's NIC) + const uint32_t *flags; // this rank's flags [2][kRoceHcas], written by the peer's NIC + uint32_t *ctrl; + uint32_t *state; // device: [0] epoch, [1] tail counter, [2] poison, [3] stage counter + uint64_t nbytes; + uint64_t slot_bytes; + uint64_t timeout_ns; + uint32_t rank; + uint32_t mode; + uint32_t n_hca; + uint32_t opts; // TF_TP_ROCE_FAST: bit 0 fast, bits 8-15 pollers a flag, bits 16-31 start stagger ns; 0: classic + // gather only (nullptr: every shard is nbytes): lens[r] = rank r's bytes, padded to 16 on the wire + const int32_t *lens; + uint64_t chunk; // chunked gathers: this op's offset into each shard (at most a slot of each from there; else 0) +}; + +// Rank r's bytes this op: ``nbytes``, or lens[r] capped at nbytes. +__device__ __forceinline__ uint64_t roce_shard(const RoceArgs &a, uint32_t r) { + if (a.lens == nullptr) { + return a.nbytes; + } + const int32_t v = __ldg(a.lens + r); + const uint64_t n = v > 0 ? static_cast(v) : 0u; + return n < a.nbytes ? n : a.nbytes; +} + +// The bytes a copy of an n-byte shard moves: n rounded up to 16 when that stays inside a 16-aligned stride, else n. +__device__ __forceinline__ uint64_t roce_span(const RoceArgs &a, uint64_t n) { + if (a.lens == nullptr || (a.nbytes & 15u) != 0u) { + return n; + } + return (n + 15u) & ~uint64_t(15u); +} + +// A shard's bytes past `off`, at most a slot (chunked gathers; off 0 and a shard within the slot: the shard). +__device__ __forceinline__ uint64_t roce_part(const RoceArgs &a, uint64_t shard, uint64_t off) { + const uint64_t rest = shard > off ? shard - off : 0u; + return rest < a.slot_bytes ? rest : a.slot_bytes; +} + +} // namespace tftp + +// Present in images whose RoceArgs carry ``chunk`` (roce.zig refuses TF_TP_ROCE_CHUNKED_KB without it); never launched. +extern "C" __global__ void tp_roce_chunk_abi() {} + +extern "C" __global__ void __launch_bounds__(512) tp_roce(RoceArgs a) { + __shared__ uint32_t s_head[2]; // poison, seq + __shared__ uint32_t s_bad; + __shared__ uint32_t s_done[kRoceHcas]; // fast: some poller of HCA h saw its flag + const bool fast = (a.opts & 1u) != 0u; + const uint32_t tid = threadIdx.x; + const uint32_t gdim = gridDim.x; + const uint64_t index = static_cast(blockIdx.x) * blockDim.x + tid; + const uint64_t stride = static_cast(gdim) * blockDim.x; + + if (tid == 0) { + s_head[0] = ld_relaxed_gpu(a.state + 2); + s_head[1] = ld_relaxed_gpu(a.state) + 1u; + s_bad = 0u; + } + if (tid < static_cast(kRoceHcas)) { + s_done[tid] = 0u; + } + __syncthreads(); + if (s_head[0] != 0u) { + return; + } + const uint32_t seq = s_head[1]; + const uint64_t slot = seq & 1u; + const bool var = a.mode == kGather && a.lens != nullptr; + const uint64_t off = var ? a.chunk : 0u; + const uint64_t mine = var ? roce_part(a, roce_shard(a, a.rank), off) : a.nbytes; + const uint64_t theirs = var ? roce_part(a, roce_shard(a, 1u - a.rank), off) : a.nbytes; + const uint64_t padded = var ? (mine < 16u ? 16u : (mine + 15u) & ~uint64_t(15u)) : (a.nbytes + 15u) & ~uint64_t(15u); + if (off != 0u && mine == 0u && theirs == 0u) { + return; // nothing past this offset on either rank: no op, the epoch stays (both ranks skip alike) + } + if (padded > a.slot_bytes) { // never staged past a slot: a refused wait (why 3), the transport poisoned + if (tid == 0 && blockIdx.x == 0) { + fail_record(a.state + 2, a.ctrl + kRoceFailed, seq, 0u, 0u, 3u, 0u); + } + return; + } + + // 1. stage this block's share into the send slot; the last block to finish rings the doorbell + copy_any(fast, a.send_slots + slot * a.slot_bytes, a.in + off, var ? roce_span(a, mine) : a.nbytes, index, + stride); + __syncthreads(); + if (tid == 0 && fast) { + // lean doorbell: release RMW per block arrival, the last acquires, then a release store (no SC system fences) + const bool last = gdim == 1u || atom_add_acq_rel_sys(a.state + 3, 1u) == gdim - 1u; + if (last) { + if (gdim > 1u) { + st_release_gpu(a.state + 3, 0u); + } + st_relaxed_sys(a.ctrl + kRoceSlotBytes + slot, static_cast(padded)); + st_release_sys(a.ctrl + kRoceSeq, seq); + } + } else if (tid == 0) { + fence_sc_sys(); + const bool last = gdim == 1u || atom_add_acq_rel_gpu(a.state + 3, 1u) == gdim - 1u; + if (last) { + if (gdim > 1u) { + st_release_gpu(a.state + 3, 0u); + fence_sc_sys(); // every block's staging (seen through the counter) before the doorbell + } + st_relaxed_sys(a.ctrl + kRoceSlotBytes + slot, static_cast(padded)); + st_release_sys(a.ctrl + kRoceSeq, seq); + } + } + if (a.mode == kGather) { + copy_any(fast, a.out + static_cast(a.rank) * a.nbytes + off, a.in + off, + var ? roce_span(a, mine) : a.nbytes, index, stride); + } + + // 2. every HCA's flag for this slot; fast: poller k of HCA h is thread 32 k + h, started k x stagger ns late + const uint32_t pollers = (a.opts >> 8) & 0xffu; + const uint32_t stagger_ns = a.opts >> 16; + if (fast) { + const uint32_t lane = tid & 31u; + const uint32_t k = tid >> 5; + if (lane < a.n_hca && k < (pollers == 0u ? 1u : pollers)) { + if (k > 0u && stagger_ns != 0u) { + __nanosleep(k * stagger_ns); + } + uint32_t seen = 0u, why = 0u, waited = 0u; + if (!wait_flag_shared(a.flags + (slot * kRoceHcas + lane) * kFlagStride, seq, a.ctrl + kRoceAbort, + a.timeout_ns, &seen, &why, &waited, &s_done[lane])) { + s_bad = 1u; + fail_record(a.state + 2, a.ctrl + kRoceFailed, seq, seen, lane, why, waited); + } + } + } else if (tid == 0) { + for (uint32_t h = 0; h < a.n_hca; h++) { + uint32_t seen = 0u, why = 0u, waited = 0u; + if (!wait_flag(a.flags + (slot * kRoceHcas + h) * kFlagStride, seq, a.ctrl + kRoceAbort, a.timeout_ns, &seen, + &why, &waited)) { + s_bad = 1u; + fail_record(a.state + 2, a.ctrl + kRoceFailed, seq, seen, h, why, waited); + break; + } + } + } + __syncthreads(); + + // 3. copy out of the receive slot the peer's NIC wrote + if (s_bad == 0u) { + const uint8_t *inbox = a.recv_slots + slot * a.slot_bytes; + switch (a.mode) { + case kGather: + copy_any(fast, a.out + static_cast(1u - a.rank) * a.nbytes + off, inbox, + var ? roce_span(a, theirs) : a.nbytes, index, stride); + break; + case kExchange: + copy_any(fast, a.out, inbox, a.nbytes, index, stride); + break; + case kSumF32: + sum_bytes(a.out, a.in, inbox, a.nbytes, a.rank, index, stride); + break; + default: + sum_bytes(a.out, a.in, inbox, a.nbytes, a.rank, index, stride); + break; + } + } + + // 4. the epoch, as tp_mailbox; the host also sees the last completed seq + __syncthreads(); + if (tid == 0) { + const bool last = gdim == 1u || atom_add_acq_rel_gpu(a.state + 1, 1u) == gdim - 1u; + const bool ok = gdim == 1u ? s_bad == 0u : ld_relaxed_gpu(a.state + 2) == 0u; + if (last && gdim > 1u) { + st_release_gpu(a.state + 1, 0u); + } + if (last && ok) { + st_release_gpu(a.state, seq); + st_relaxed_sys(a.ctrl + kRoceCompleted, seq); + } + } +} diff --git a/zig/src/fabric/tp/mailbox.zig b/zig/src/fabric/tp/mailbox.zig new file mode 100644 index 0000000000..902bb79393 --- /dev/null +++ b/zig/src/fabric/tp/mailbox.zig @@ -0,0 +1,215 @@ +//! The mailbox backend: two-rank exchanges in one kernel through pinned host memory both GPUs map (host side). +const std = @import("std"); +const log = @import("log.zig"); +const cuda = @import("cuda"); +const collective = @import("collective.zig"); +const Bootstrap = @import("bootstrap.zig").Bootstrap; +const sock = @import("sock.zig"); +const posix = std.posix; +const DevicePtr = collective.DevicePtr; +const DType = collective.DType; +const Error = collective.Error; + +pub const max_blocks = 16; +const threads = 512; +const flag_bytes = 2 * max_blocks * 128; +const ctrl_bytes = 4096; +const ctrl_abort = 0; +const ctrl_rank = 16; + +const Mode = enum(u32) { gather = 0, sum_f32 = 1, sum_bf16 = 2, exchange = 3 }; + +/// mailbox.cu's `MailArgs`, field for field. +const MailArgs = extern struct { + in: DevicePtr, + out: DevicePtr, + peer_slots: DevicePtr, + peer_flags: DevicePtr, + my_slots: DevicePtr, + my_flags: DevicePtr, + state: DevicePtr, + ctrl: DevicePtr, + nbytes: u64, + slot_bytes: u64, + timeout_ns: u64, + rank: u32, + mode: u32, +}; + +/// Driver calls the runtime's table does not carry (page-locking memory this process mapped itself). +pub const HostRegister = struct { + lib: std.DynLib, + register: *const fn (?*anyopaque, usize, c_uint) callconv(.c) cuda.abi.Result, + unregister: *const fn (?*anyopaque) callconv(.c) cuda.abi.Result, + + pub fn open() !HostRegister { + var lib = std.DynLib.open("libcuda.so.1") catch return error.DriverUnavailable; + errdefer lib.close(); + return .{ + .register = lib.lookup(@FieldType(HostRegister, "register"), "cuMemHostRegister_v2") orelse return error.MissingSymbol, + .unregister = lib.lookup(@FieldType(HostRegister, "unregister"), "cuMemHostUnregister") orelse return error.MissingSymbol, + .lib = lib, + }; + } +}; + +/// Grid for `n` bytes a rank: one block per 16 KiB, a power of two, at most `max_blocks` (both ranks pick the same). +pub fn blocksFor(n: usize) u32 { + const want = std.math.divCeil(usize, @max(n, 1), 16 * 1024) catch 1; + return @intCast(@min(std.math.ceilPowerOfTwoAssert(usize, want), max_blocks)); +} + +pub const Mailbox = struct { + d: *const cuda.Driver, + me: u32, + slot_bytes: usize, + timeout_ns: u64, + region: []align(std.heap.page_size_min) u8, + dev: DevicePtr, + hr: HostRegister, + state: cuda.DeviceBuffer, + module: cuda.Module, + func: cuda.Function, + aborted: std.atomic.Value(bool) = .init(false), + + /// Both ranks call it after the bootstrap (two ranks, one host); `image` is mailbox.cu's fatbin. + pub fn init(d: *const cuda.Driver, b: *Bootstrap, image: []const u8, max_bytes: usize, timeout_ns: u64) !Mailbox { + if (b.world != 2) return error.Unsupported; + if (image.len == 0) return error.NoKernel; + const slot = std.mem.alignForward(usize, max_bytes, 4096); + const len = ctrl_bytes + 2 * (flag_bytes + 2 * slot); + // rank 0 creates and sizes the file before anyone learns its name, the others open it, then rank 0 unlinks it + var name: [64]u8 = undefined; + var name_len: u8 = 0; + var fd: std.c.fd_t = -1; + defer if (fd >= 0) { + _ = std.c.close(fd); + }; + if (b.rank == 0) blk: { + const s = std.fmt.bufPrint(&name, "/dev/shm/tf-tp-{d}-{x}", .{ std.c.getpid(), sock.nowNs() }) catch break :blk; + name[s.len] = 0; + fd = std.c.open(name[0..s.len :0], .{ .ACCMODE = .RDWR, .CREAT = true, .EXCL = true, .CLOEXEC = true }, @as(std.c.mode_t, 0o600)); + if (fd >= 0 and std.c.ftruncate(fd, @intCast(len)) == 0) name_len = @intCast(s.len); + } + // an empty name tells the other rank that rank 0 failed (it must not wait for a name that never comes) + try b.broadcast(std.mem.asBytes(&name_len), std.mem.asBytes(&name_len)); + if (name_len == 0) return error.SharedMemoryUnavailable; + try b.broadcast(name[0..name_len], name[0..name_len]); + name[name_len] = 0; + const path = name[0..name_len :0]; + if (b.rank != 0) { + fd = std.c.open(path, .{ .ACCMODE = .RDWR, .CLOEXEC = true }, @as(std.c.mode_t, 0)); + if (fd < 0) return error.SharedMemoryUnavailable; + } + try b.agree("mapped", "mailbox file"); + const region = try posix.mmap(null, len, .{ .READ = true, .WRITE = true }, .{ .TYPE = .SHARED }, fd, 0); + errdefer posix.munmap(region); + try b.agree("open", "mailbox mapping"); + if (b.rank == 0) _ = std.c.unlink(path); + var hr = try HostRegister.open(); + errdefer hr.lib.close(); + try d.check(hr.register(region.ptr, len, 1 | 2), "cuMemHostRegister"); // PORTABLE | DEVICEMAP + errdefer _ = hr.unregister(region.ptr); + var dev: DevicePtr = 0; + try d.check(d.api.cuMemHostGetDevicePointer_v2(&dev, region.ptr, 0), "cuMemHostGetDevicePointer"); + var state = try cuda.DeviceBuffer.alloc(d, 16); + errdefer state.free(); + try state.fill8(0, null); + var module = try cuda.Module.load(d, image); + errdefer module.unload(); + const func = try module.function("tp_mailbox"); + try b.agree("ready", "mailbox setup"); + return .{ .d = d, .me = b.rank, .slot_bytes = slot, .timeout_ns = timeout_ns, .region = region, .dev = dev, .hr = hr, .state = state, .module = module, .func = func }; + } + + pub fn deinit(m: *Mailbox) void { + m.module.unload(); + m.state.free(); + _ = m.hr.unregister(m.region.ptr); + m.hr.lib.close(); + posix.munmap(m.region); + } + + fn inbox(m: *const Mailbox, r: u32) usize { + return ctrl_bytes + r * (flag_bytes + 2 * m.slot_bytes); + } + + fn ctrlWord(m: *const Mailbox, i: usize) *std.atomic.Value(u32) { + return @ptrCast(@alignCast(m.region[4 * i ..].ptr)); + } + + /// Whether `n` bytes a rank fit one slot. + pub fn fits(m: *const Mailbox, n: usize) bool { + return n <= m.slot_bytes; + } + + fn launch(m: *Mailbox, mode: Mode, in: DevicePtr, out: DevicePtr, n: usize, stream: collective.Stream) Error!void { + if (m.aborted.load(.acquire)) return error.Aborted; + if (n > m.slot_bytes) return error.Unsupported; + if (n == 0) return; + const peer = 1 - m.me; + var a: MailArgs = .{ + .in = in, + .out = out, + .peer_slots = m.dev + m.inbox(peer) + flag_bytes, + .peer_flags = m.dev + m.inbox(peer), + .my_slots = m.dev + m.inbox(m.me) + flag_bytes, + .my_flags = m.dev + m.inbox(m.me), + .state = m.state.ptr, + .ctrl = m.dev, + .nbytes = n, + .slot_bytes = m.slot_bytes, + .timeout_ns = m.timeout_ns, + .rank = m.me, + .mode = @backingInt(mode), + }; + var params = [1]?*anyopaque{@ptrCast(&a)}; + const r = m.d.api.cuLaunchKernel(m.func.handle, blocksFor(n), 1, 1, threads, 1, 1, 0, stream, ¶ms, null); + m.d.check(r, "cuLaunchKernel tp_mailbox") catch return error.BackendFailed; + } + + pub fn allGatherBytes(m: *Mailbox, s: DevicePtr, r: DevicePtr, n: usize, stream: collective.Stream) Error!void { + return m.launch(.gather, s, r, n, stream); + } + + pub fn exchangeBytes(m: *Mailbox, s: DevicePtr, r: DevicePtr, n: usize, stream: collective.Stream) Error!void { + return m.launch(.exchange, s, r, n, stream); + } + + /// Sums only, f32 or bf16 (what the model's partials are); the rest is `error.Unsupported`. + pub fn sum(m: *Mailbox, s: DevicePtr, r: DevicePtr, count: usize, t: DType, stream: collective.Stream) Error!void { + const mode: Mode = switch (t) { + .f32 => .sum_f32, + .bf16 => .sum_bf16, + else => return error.Unsupported, + }; + return m.launch(mode, s, r, count * t.size(), stream); + } + + /// Ends pending and later exchanges on both ranks (the abort word is shared); this rank's launches stop. + pub fn abort(m: *Mailbox) void { + m.aborted.store(true, .release); + m.ctrlWord(ctrl_abort).store(1, .release); + } + + /// A wait that timed out or was aborted on either rank. + pub fn check(m: *Mailbox) Error!void { + for (0..2) |r| { + const rec = ctrl_rank + 8 * r; + if (m.ctrlWord(rec).load(.acquire) != 0) { + log.warn("tp mailbox: rank {d}'s wait for seq {d} failed ({s} after {d} us; block {d}, flag held {d})", .{ r, m.ctrlWord(rec + 1).load(.acquire), if (m.ctrlWord(rec + 4).load(.acquire) == 1) "aborted" else "timed out", m.ctrlWord(rec + 5).load(.acquire), m.ctrlWord(rec + 3).load(.acquire), m.ctrlWord(rec + 2).load(.acquire) }); + return error.PeerFailed; + } + } + if (m.aborted.load(.acquire) or m.ctrlWord(ctrl_abort).load(.acquire) != 0) return error.Aborted; + } +}; + +test "grids grow with the shard and stay powers of two" { + try std.testing.expectEqual(@as(u32, 1), blocksFor(4096)); + try std.testing.expectEqual(@as(u32, 1), blocksFor(16 * 1024)); + try std.testing.expectEqual(@as(u32, 2), blocksFor(16 * 1024 + 1)); + try std.testing.expectEqual(@as(u32, 4), blocksFor(64 * 1024)); + try std.testing.expectEqual(@as(u32, 16), blocksFor(4 << 20)); + try std.testing.expectEqual(@as(usize, 96), @sizeOf(MailArgs)); +} diff --git a/zig/src/fabric/tp/nccl.zig b/zig/src/fabric/tp/nccl.zig new file mode 100644 index 0000000000..c9369369ef --- /dev/null +++ b/zig/src/fabric/tp/nccl.zig @@ -0,0 +1,164 @@ +//! The NCCL backend: one communicator a rank, every collective one enqueue on the caller's stream. +const std = @import("std"); +const cuda = @import("cuda"); +const collective = @import("collective.zig"); +const nccl = cuda.nccl; +const DevicePtr = collective.DevicePtr; +const DType = collective.DType; +const Op = collective.Op; +const Error = collective.Error; + +pub fn dataType(t: DType) nccl.DataType { + return switch (t) { + .u8 => .u8, + .i32 => .i32, + .i64 => .i64, + .f16 => .f16, + .bf16 => .bf16, + .f32 => .f32, + .f64 => .f64, + }; +} + +pub fn redOp(op: Op) nccl.RedOp { + return switch (op) { + .sum => .sum, + .max => .max, + .min => .min, + }; +} + +pub const Nccl = struct { + lib: *const nccl.Library, + d: *const cuda.Driver, + comm: nccl.Comm, + me: u32, + size: u32, + aborted: std.atomic.Value(bool) = .init(false), + /// The barrier's one-float buffers and its own stream (host path only). + scratch: cuda.DeviceBuffer, + stream: cuda.Stream, + + /// Joins the communicator `id` names (every rank calls it with the same id; it returns once all have). + pub fn init(lib: *const nccl.Library, d: *const cuda.Driver, id: nccl.UniqueId, rank: u32, world: u32) !Nccl { + var comm: nccl.Comm = null; + try lib.check(lib.api.ncclCommInitRank(&comm, @intCast(world), id, @intCast(rank)), "ncclCommInitRank"); + errdefer _ = lib.api.ncclCommDestroy(comm); + var scratch = try cuda.DeviceBuffer.alloc(d, 4 * (world + 1)); + errdefer scratch.free(); + try scratch.fill8(0, null); + return .{ .lib = lib, .d = d, .comm = comm, .me = rank, .size = world, .scratch = scratch, .stream = try cuda.Stream.init(d, true) }; + } + + pub fn deinit(n: *Nccl) void { + _ = n.stream.synchronize() catch {}; + n.stream.deinit(); + n.scratch.free(); + if (!n.aborted.load(.acquire)) _ = n.lib.api.ncclCommDestroy(n.comm); + } + + pub fn iface(n: *Nccl) collective.Collective { + return .{ .ptr = n, .vtable = &vtable }; + } + + const vtable: collective.Collective.VTable = .{ + .kind = .nccl, + .rank = rankOf, + .world = worldOf, + .all_reduce = allReduce, + .all_gather = allGather, + .exchange = exchange, + .send = send, + .recv = recv, + .barrier = barrier, + .check = check, + .abort = abort, + }; + + fn self_(ptr: *anyopaque) *Nccl { + return @ptrCast(@alignCast(ptr)); + } + + fn rankOf(ptr: *anyopaque) u32 { + return self_(ptr).me; + } + + fn worldOf(ptr: *anyopaque) u32 { + return self_(ptr).size; + } + + fn call(n: *Nccl, r: nccl.Result, what: []const u8) Error!void { + if (r == 0) return; + if (n.aborted.load(.acquire)) return error.Aborted; + n.lib.check(r, what) catch return error.BackendFailed; + } + + fn live(n: *Nccl) Error!void { + if (n.aborted.load(.acquire)) return error.Aborted; + } + + /// NCCL adds in its ring's order: the rank-order sum `all_reduce` promises holds at two ranks, not past them. + pub fn allReduce(ptr: *anyopaque, s: DevicePtr, r: DevicePtr, count: usize, t: DType, op: Op, stream: collective.Stream) Error!void { + const n = self_(ptr); + try n.live(); + try n.call(n.lib.api.ncclAllReduce(s, r, count, dataType(t), redOp(op), n.comm, stream), "ncclAllReduce"); + } + + pub fn allGather(ptr: *anyopaque, s: DevicePtr, r: DevicePtr, count: usize, t: DType, stream: collective.Stream) Error!void { + const n = self_(ptr); + try n.live(); + try n.call(n.lib.api.ncclAllGather(s, r, count, dataType(t), n.comm, stream), "ncclAllGather"); + } + + pub fn exchange(ptr: *anyopaque, s: DevicePtr, r: DevicePtr, count: usize, t: DType, peer: u32, stream: collective.Stream) Error!void { + const n = self_(ptr); + try n.live(); + try collective.checkPeer(n.me, n.size, peer); + try n.call(n.lib.api.ncclGroupStart(), "ncclGroupStart"); + const a = n.lib.api.ncclSend(s, count, dataType(t), @intCast(peer), n.comm, stream); + const b = n.lib.api.ncclRecv(r, count, dataType(t), @intCast(peer), n.comm, stream); + const e = n.lib.api.ncclGroupEnd(); + try n.call(a, "ncclSend"); + try n.call(b, "ncclRecv"); + try n.call(e, "ncclGroupEnd"); + } + + fn send(ptr: *anyopaque, buf: DevicePtr, count: usize, t: DType, peer: u32, stream: collective.Stream) Error!void { + const n = self_(ptr); + try n.live(); + try collective.checkPeer(n.me, n.size, peer); + try n.call(n.lib.api.ncclSend(buf, count, dataType(t), @intCast(peer), n.comm, stream), "ncclSend"); + } + + fn recv(ptr: *anyopaque, buf: DevicePtr, count: usize, t: DType, peer: u32, stream: collective.Stream) Error!void { + const n = self_(ptr); + try n.live(); + try collective.checkPeer(n.me, n.size, peer); + try n.call(n.lib.api.ncclRecv(buf, count, dataType(t), @intCast(peer), n.comm, stream), "ncclRecv"); + } + + /// One float all-gathered on the barrier's own stream, then a sync of it (host path, never captured). + fn barrier(ptr: *anyopaque) Error!void { + const n = self_(ptr); + try n.live(); + try n.call(n.lib.api.ncclAllGather(n.scratch.ptr, n.scratch.ptr + 4, 1, .f32, n.comm, n.stream.handle), "ncclAllGather"); + n.stream.synchronize() catch return if (n.aborted.load(.acquire)) error.Aborted else error.BackendFailed; + } + + fn check(ptr: *anyopaque) Error!void { + const n = self_(ptr); + try n.live(); + var r: nccl.Result = 0; + try n.call(n.lib.api.ncclCommGetAsyncError(n.comm, &r), "ncclCommGetAsyncError"); + if (r != 0 and r != nccl.in_progress) { + n.lib.check(r, "NCCL async error") catch {}; + return error.PeerFailed; + } + } + + fn abort(ptr: *anyopaque) void { + const n = self_(ptr); + if (n.aborted.swap(true, .acq_rel)) return; + _ = n.lib.api.ncclCommAbort(n.comm); + } +}; diff --git a/zig/src/fabric/tp/planlink.zig b/zig/src/fabric/tp/planlink.zig new file mode 100644 index 0000000000..44dbb61a0c --- /dev/null +++ b/zig/src/fabric/tp/planlink.zig @@ -0,0 +1,467 @@ +//! The round plan from rank 0 to the other ranks over its own TCP connection: a 4-byte length, then int64 words. +const std = @import("std"); +const log = @import("log.zig"); +const sock = @import("sock.zig"); +const builtin = @import("builtin"); +const config = @import("config.zig"); +const linux = std.os.linux; + +pub const max_ranks = 64; + +/// The CPU a plan thread pinned itself to (TF_TP_PLAN_PIN), -1 until one did; an engine's CPU isolation keeps off it. +pub var pinned_cpu: std.atomic.Value(i32) = .init(-1); +const more: u32 = 0x8000; // MSG_MORE: the length and the ints leave as one segment + +pub const Error = error{ PeerClosed, TooLarge, NotRankZero, NotFollower }; + +/// A batch frame's first word (a plan's op is >= 0): [batch, n, len 1, plan 1, ..., len n, plan n]. +pub const batch: i64 = std.math.minInt(i64); +/// Words rank 0 holds at most (more flushes the held plans first), and the most a batch frame carries. +pub const hold_words: usize = 1024; +pub const batch_words: usize = 2 * hold_words + 2; + +/// Largest plan either side accepts: 64M words (512 MiB), so a garbled length cannot become a 4 GiB allocation. +pub const max_frame_words: usize = 64 << 20; + +comptime { + std.debug.assert(builtin.cpu.arch.endian() == .little); // the ints go out as they are in memory +} + +pub const PlanLink = struct { + rank: u32, + world: u32, + /// Rank 0: rank i's socket at index i; the others: rank 0's at index 0 (the bootstrap's `releasePlan`). + fds: [max_ranks]sock.Fd, + sent: u64 = 0, + received: u64 = 0, + /// TF_TP_PLAN_LINK=rdma (config.plan_spin_ns): the follower polls the socket this long before it blocks. + spin_ns: u64 = 0, + /// TF_TP_PLAN_PIN: the sending / receiving thread pins itself at its first plan; `avoid`: the RoCE proxy's CPU. + pin: ?config.PlanPin = null, + avoid: ?u32 = null, + pinned: bool = false, + /// rank 0: plans held for the next send (`hold`), each as [len, plan...]; `held_n` of them + held: [hold_words]i64 = undefined, + held_len: usize = 0, + held_n: usize = 0, + /// the follower: a batch frame's plans not handed out yet, as [len, plan...] from `inbox_at` + inbox: [batch_words]i64 = undefined, + inbox_len: usize = 0, + inbox_at: usize = 0, + batches: u64 = 0, + + pub fn init(rank: u32, world: u32, fds: [max_ranks]sock.Fd) PlanLink { + return .{ .rank = rank, .world = world, .fds = fds }; + } + + /// Rank 0: one plan to every other rank (host only); held plans go first, in the same frame when they fit. + pub fn send(p: *PlanLink, ints: []const i64) Error!void { + if (p.rank != 0) return error.NotRankZero; + if (ints.len > max_frame_words) return error.TooLarge; + if (p.held_n > 0) { + if (2 + p.held_len + 1 + ints.len <= batch_words) { + var frame: [batch_words]i64 = undefined; + frame[0] = batch; + frame[1] = @intCast(p.held_n + 1); + @memcpy(frame[2..][0..p.held_len], p.held[0..p.held_len]); + frame[2 + p.held_len] = @intCast(ints.len); + @memcpy(frame[3 + p.held_len ..][0..ints.len], ints); + const n = 3 + p.held_len + ints.len; + p.held_len = 0; + p.held_n = 0; + p.batches += 1; + return p.sendFrame(frame[0..n]); + } + try p.flush(); + } + return p.sendFrame(ints); + } + + /// Rank 0: a plan the followers run before the next sent one, carried in that plan's frame (`send` or `flush`). + pub fn hold(p: *PlanLink, ints: []const i64) Error!void { + if (p.rank != 0) return error.NotRankZero; + if (1 + ints.len > hold_words) { + try p.flush(); + return p.sendFrame(ints); + } + if (p.held_len + 1 + ints.len > hold_words) try p.flush(); + p.held[p.held_len] = @intCast(ints.len); + @memcpy(p.held[p.held_len + 1 ..][0..ints.len], ints); + p.held_len += 1 + ints.len; + p.held_n += 1; + } + + /// Rank 0: the held plans now, as one batch frame (nothing held: nothing sent). + pub fn flush(p: *PlanLink) Error!void { + if (p.held_n == 0) return; + var frame: [batch_words]i64 = undefined; + frame[0] = batch; + frame[1] = @intCast(p.held_n); + @memcpy(frame[2..][0..p.held_len], p.held[0..p.held_len]); + const n = 2 + p.held_len; + p.held_len = 0; + p.held_n = 0; + p.batches += 1; + return p.sendFrame(frame[0..n]); + } + + fn sendFrame(p: *PlanLink, ints: []const i64) Error!void { + p.pinOnce(); + const bytes = std.mem.sliceAsBytes(ints); + const len: u32 = @intCast(bytes.len); + for (p.fds[1..p.world]) |fd| { + if (std.c.send(fd, &len, 4, std.posix.MSG.NOSIGNAL | more) != 4) return error.PeerClosed; + sock.sendAll(fd, bytes) catch return error.PeerClosed; + } + p.sent += 1; + } + + /// The other ranks: the next plan into `out`; `error.PeerClosed` once rank 0 is gone, `error.TooLarge` skips one. + pub fn recv(p: *PlanLink, out: []i64) Error![]i64 { + if (p.inboxNext()) |m| { + if (m.len > out.len) return error.TooLarge; + @memcpy(out[0..m.len], m); + return out[0..m.len]; + } + const n = try p.header(); + if (n > out.len) { + try p.drop(n); + return error.TooLarge; + } + const got = try p.body(out[0..n]); + if (got.len > 0 and got[0] == batch) { + try p.unpack(got); + return p.recv(out); + } + return got; + } + + /// The next plan into `buf`, grown to the frame; the buffer is kept, so steady plans allocate nothing. + pub fn recvGrow(p: *PlanLink, gpa: std.mem.Allocator, buf: *std.ArrayList(i64)) (Error || std.mem.Allocator.Error)![]i64 { + if (p.inboxNext()) |m| { + try buf.resize(gpa, m.len); + @memcpy(buf.items, m); + return buf.items; + } + const n = try p.header(); + if (n > max_frame_words) { + try p.drop(n); + return error.TooLarge; + } + try buf.resize(gpa, n); + const got = try p.body(buf.items); + if (got.len > 0 and got[0] == batch) { + try p.unpack(got); + return p.recvGrow(gpa, buf); + } + return got; + } + + /// The follower: a batch frame's plans into the inbox (checked whole first: a garbled batch ends the link). + fn unpack(p: *PlanLink, frame: []const i64) Error!void { + if (frame.len < 2 or frame.len > batch_words or frame[1] < 1) return error.PeerClosed; + var at: usize = 2; + var k: i64 = 0; + while (k < frame[1]) : (k += 1) { + if (at >= frame.len or frame[at] < 0) return error.PeerClosed; + at += 1 + @as(usize, @intCast(frame[at])); + if (at > frame.len) return error.PeerClosed; + } + if (at != frame.len) return error.PeerClosed; + @memcpy(p.inbox[0 .. frame.len - 2], frame[2..]); + p.inbox_len = frame.len - 2; + p.inbox_at = 0; + p.batches += 1; + } + + /// The follower: the next plan of the last batch frame, if any (a slice of the inbox, valid until the next recv). + fn inboxNext(p: *PlanLink) ?[]const i64 { + if (p.inbox_at >= p.inbox_len) return null; + const n: usize = @intCast(p.inbox[p.inbox_at]); + const m = p.inbox[p.inbox_at + 1 ..][0..n]; + p.inbox_at += 1 + n; + return m; + } + + /// The next frame's length in words. + fn header(p: *PlanLink) Error!usize { + if (p.rank == 0) return error.NotFollower; + p.pinOnce(); + if (p.spin_ns > 0) p.spin(); + var len: u32 = undefined; + sock.recvExact(p.fds[0], std.mem.asBytes(&len), null) catch return error.PeerClosed; + if (len % 8 != 0) return error.PeerClosed; // not a frame boundary: the stream is lost + return len / 8; + } + + fn body(p: *PlanLink, out: []i64) Error![]i64 { + sock.recvExact(p.fds[0], std.mem.sliceAsBytes(out), null) catch return error.PeerClosed; + p.received += 1; + return out; + } + + /// Reads and drops `n` words. + fn drop(p: *PlanLink, n: usize) Error!void { + var left: usize = n * 8; + var scratch: [4096]u8 = undefined; + while (left > 0) : (left -= @min(left, scratch.len)) sock.recvExact(p.fds[0], scratch[0..@min(left, scratch.len)], null) catch return error.PeerClosed; + } + + /// Polls rank 0's socket until it is readable (or hung up) or `spin_ns` passed; the read that follows blocks. + fn spin(p: *PlanLink) void { + if (comptime builtin.os.tag != .linux) return; + var fd = [1]linux.pollfd{.{ .fd = p.fds[0], .events = linux.POLL.IN, .revents = 0 }}; + const end = sock.nowNs() + p.spin_ns; + while (true) { + const rc = linux.poll(&fd, 1, 0); + if (linux.errno(rc) == .SUCCESS and rc > 0) return; + if (sock.nowNs() >= end) return; + std.atomic.spinLoopHint(); + } + } + + /// Pins the calling thread once: an explicit CPU, or auto's highest cpu_capacity. + fn pinOnce(p: *PlanLink) void { + if (p.pinned) return; + p.pinned = true; + if (comptime builtin.os.tag != .linux) return; + const want = p.pin orelse return; + const cpu = switch (want) { + .cpu => |c| c, + .auto => bestCpu(p.avoid) orelse { + log.info("[tensorfold] tp: TF_TP_PLAN_PIN=auto: no cpu_capacity differs among the allowed CPUs; the plan thread stays unpinned", .{}); + return; + }, + }; + pinTo(cpu) catch |e| { + log.warn("[tensorfold] tp: TF_TP_PLAN_PIN: CPU {d} refused ({t}); the plan thread stays unpinned", .{ cpu, e }); + return; + }; + pinned_cpu.store(@intCast(cpu), .release); + log.info("[tensorfold] tp: plan link rank {d} pinned to CPU {d}", .{ p.rank, cpu }); + } + + pub fn close(p: *PlanLink) void { + for (&p.fds) |*fd| if (fd.* >= 0) { + sock.close(fd.*); + fd.* = -1; + }; + } +}; + +/// The calling thread onto `cpu` alone. +pub fn pinTo(cpu: u32) !void { + if (cpu >= linux.CPU_SETSIZE) return error.BadCpu; + var set: linux.cpu_set_t = @splat(0); + set[cpu / @bitSizeOf(usize)] |= @as(usize, 1) << @intCast(cpu % @bitSizeOf(usize)); + try linux.sched_setaffinity(0, &set); +} + +/// Auto pin: the allowed CPU (not `avoid`) with the highest cpu_capacity, lowest on ties; null when all are equal. +pub fn bestCpu(avoid: ?u32) ?u32 { + var set: linux.cpu_set_t = undefined; + if (linux.errno(linux.sched_getaffinity(0, @sizeOf(linux.cpu_set_t), &set)) != .SUCCESS) return null; + var best: ?u32 = null; + var best_cap: u32 = 0; + var low_cap: u32 = std.math.maxInt(u32); + for (set, 0..) |word, wi| { + var bits = word; + while (bits != 0) : (bits &= bits - 1) { + const cpu: u32 = @intCast(wi * @bitSizeOf(usize) + @ctz(bits)); + if (avoid == cpu) continue; + const c = capacity(cpu) orelse continue; + low_cap = @min(low_cap, c); + if (best == null or c > best_cap) { + best = cpu; + best_cap = c; + } + } + } + return if (best != null and best_cap > low_cap) best else null; +} + +fn capacity(cpu: u32) ?u32 { + var pb: [64]u8 = undefined; + const path = std.fmt.bufPrintSentinel(&pb, "/sys/devices/system/cpu/cpu{d}/cpu_capacity", .{cpu}, 0) catch return null; + const rc = linux.open(path, .{ .ACCMODE = .RDONLY, .CLOEXEC = true }, 0); + if (linux.errno(rc) != .SUCCESS) return null; + const fd: i32 = @intCast(rc); + defer _ = linux.close(fd); + var b: [16]u8 = undefined; + const n = linux.read(fd, &b, b.len); + if (linux.errno(n) != .SUCCESS) return null; + return std.fmt.parseInt(u32, std.mem.trim(u8, b[0..n], " \n"), 10) catch null; +} + +test "plans arrive whole and in order; a closed rank 0 ends the follower" { + const bootstrap = @import("bootstrap.zig"); + const l = try bootstrap.testListener(); + defer sock.close(l.fd); + const deadline = sock.nowNs() + 5 * std.time.ns_per_s; + const a = try sock.connect("localhost", l.port, deadline); + const b = try sock.accept(l.fd, deadline); + var f0: [max_ranks]sock.Fd = @splat(-1); + var f1: [max_ranks]sock.Fd = @splat(-1); + f0[1] = b; + f1[0] = a; + var lead = PlanLink.init(0, 2, f0); + var follow = PlanLink.init(1, 2, f1); + defer follow.close(); + var big: [5000]i64 = undefined; + for (&big, 0..) |*v, i| v.* = @as(i64, @intCast(i)) * -3 + 0x7FFF_0000_0000; + try lead.send(&.{ 0x504C414E, 2, 1 }); + try lead.send(&big); + try lead.send(&.{}); + var buf: [8192]i64 = undefined; + try std.testing.expectEqualSlices(i64, &.{ 0x504C414E, 2, 1 }, try follow.recv(&buf)); + try std.testing.expectEqualSlices(i64, &big, try follow.recv(&buf)); + try std.testing.expectEqual(@as(usize, 0), (try follow.recv(&buf)).len); + try std.testing.expectError(error.NotRankZero, follow.send(&.{1})); + var small: [2]i64 = undefined; + try lead.send(&.{ 1, 2, 3 }); + try std.testing.expectError(error.TooLarge, follow.recv(&small)); + lead.close(); + try std.testing.expectError(error.PeerClosed, follow.recv(&buf)); +} + +test "an admission-sized plan (1.2M words) arrives whole through the growable path" { + const bootstrap = @import("bootstrap.zig"); + const l = try bootstrap.testListener(); + defer sock.close(l.fd); + const deadline = sock.nowNs() + 5 * std.time.ns_per_s; + const a = try sock.connect("localhost", l.port, deadline); + const b = try sock.accept(l.fd, deadline); + var f0: [max_ranks]sock.Fd = @splat(-1); + var f1: [max_ranks]sock.Fd = @splat(-1); + f0[1] = b; + f1[0] = a; + var lead = PlanLink.init(0, 2, f0); + defer lead.close(); + var follow = PlanLink.init(1, 2, f1); + defer follow.close(); + const gpa = std.testing.allocator; + const prompt = try gpa.alloc(i64, 1_200_000); + defer gpa.free(prompt); + for (prompt, 0..) |*v, i| v.* = @intCast((i *% 2654435761) % 163840); + // the socket buffer holds far less than 9.6 MB: rank 0 sends from its own thread, as it would + const Send = struct { + fn go(link: *PlanLink, ints: []const i64) void { + link.send(&.{ 0x504C414E, 3, 20 }) catch {}; + link.send(ints) catch {}; + link.send(&.{ 0x504C414E, 3, 21 }) catch {}; + } + }; + const t = try std.Thread.spawn(.{}, Send.go, .{ &lead, prompt }); + var buf: std.ArrayList(i64) = .empty; + defer buf.deinit(gpa); + try std.testing.expectEqualSlices(i64, &.{ 0x504C414E, 3, 20 }, try follow.recvGrow(gpa, &buf)); + try std.testing.expectEqualSlices(i64, prompt, try follow.recvGrow(gpa, &buf)); + try std.testing.expectEqualSlices(i64, &.{ 0x504C414E, 3, 21 }, try follow.recvGrow(gpa, &buf)); + t.join(); + try std.testing.expect(buf.capacity >= prompt.len); // kept for the next admission +} + +test "TF_TP_PLAN_LINK=rdma: a spinning follower gets every plan, late or early; TF_TP_PLAN_PIN pins the thread" { + if (builtin.os.tag != .linux) return error.SkipZigTest; + const bootstrap = @import("bootstrap.zig"); + const l = try bootstrap.testListener(); + defer sock.close(l.fd); + const deadline = sock.nowNs() + 5 * std.time.ns_per_s; + const a = try sock.connect("localhost", l.port, deadline); + const b = try sock.accept(l.fd, deadline); + var f0: [max_ranks]sock.Fd = @splat(-1); + var f1: [max_ranks]sock.Fd = @splat(-1); + f0[1] = b; + f1[0] = a; + var lead = PlanLink.init(0, 2, f0); + defer lead.close(); + var follow = PlanLink.init(1, 2, f1); + defer follow.close(); + follow.spin_ns = 2 * std.time.ns_per_ms; + // the CPU this thread may run on first: the pin keeps it allowed, so the test runner is not starved + var set: linux.cpu_set_t = undefined; + _ = linux.sched_getaffinity(0, @sizeOf(linux.cpu_set_t), &set); + var first: u32 = 0; + for (set, 0..) |w, i| if (w != 0) { + first = @intCast(i * @bitSizeOf(usize) + @ctz(w)); + break; + }; + const Follower = struct { + fn run(f: *PlanLink, cpu: u32, out: *[3]i64, ok: *bool) void { + f.pin = .{ .cpu = cpu }; + var buf: [8]i64 = undefined; + const x = (f.recv(&buf) catch return)[0]; // arrives within the spin + const y = (f.recv(&buf) catch return)[0]; // arrives after it: the blocking read + var now: linux.cpu_set_t = undefined; + _ = linux.sched_getaffinity(0, @sizeOf(linux.cpu_set_t), &now); + var count: usize = 0; + for (now) |w| count += @popCount(w); + out.* = .{ x, y, @intCast(count) }; + ok.* = f.pinned and (now[cpu / @bitSizeOf(usize)] >> @intCast(cpu % @bitSizeOf(usize))) & 1 == 1; + } + }; + var out: [3]i64 = .{ 0, 0, 0 }; + var ok = false; + const t = try std.Thread.spawn(.{}, Follower.run, .{ &follow, first, &out, &ok }); + try lead.send(&.{11}); + var ts: linux.timespec = .{ .sec = 0, .nsec = 20 * std.time.ns_per_ms }; + _ = linux.nanosleep(&ts, null); + try lead.send(&.{22}); + t.join(); + try std.testing.expectEqual([3]i64{ 11, 22, 1 }, out); + try std.testing.expect(ok); + if (bestCpu(null)) |c| try std.testing.expect((set[c / @bitSizeOf(usize)] >> @intCast(c % @bitSizeOf(usize))) & 1 == 1); +} + +test "held plans leave with the next plan, in order, as one frame; a large plan after them flushes them first" { + const bootstrap = @import("bootstrap.zig"); + const l = try bootstrap.testListener(); + defer sock.close(l.fd); + const deadline = sock.nowNs() + 5 * std.time.ns_per_s; + const a = try sock.connect("127.0.0.1", l.port, deadline); + const b = try sock.accept(l.fd, deadline); + var f0: [max_ranks]sock.Fd = @splat(-1); + var f1: [max_ranks]sock.Fd = @splat(-1); + f0[1] = b; + f1[0] = a; + const lead = try std.testing.allocator.create(PlanLink); + defer std.testing.allocator.destroy(lead); + lead.* = PlanLink.init(0, 2, f0); + const follow = try std.testing.allocator.create(PlanLink); + defer std.testing.allocator.destroy(follow); + follow.* = PlanLink.init(1, 2, f1); + defer follow.close(); + // a round: three keeps held, the window sent: one frame + try lead.hold(&.{ 22, 0, 3 }); + try lead.hold(&.{ 22, 1, 0 }); + try lead.hold(&.{ 22, 2, 5 }); + try lead.send(&.{ 21, 4, 9, 9, 9 }); + try std.testing.expectEqual(@as(u64, 1), lead.sent); + // a held plan, then one too large for a batch: two frames, the held one first + var big: [5000]i64 = undefined; + for (&big, 0..) |*v, i| v.* = @intCast(i); + try lead.hold(&.{ 22, 3, 1 }); + try lead.send(&big); + try std.testing.expectEqual(@as(u64, 3), lead.sent); + // flush alone, an empty plan held, a plain send after + try lead.hold(&.{}); + try lead.flush(); + try lead.flush(); // nothing held: nothing sent + try lead.send(&.{7}); + try std.testing.expectEqual(@as(u64, 5), lead.sent); + var buf: std.ArrayList(i64) = .empty; + defer buf.deinit(std.testing.allocator); + const gpa = std.testing.allocator; + try std.testing.expectEqualSlices(i64, &.{ 22, 0, 3 }, try follow.recvGrow(gpa, &buf)); + try std.testing.expectEqualSlices(i64, &.{ 22, 1, 0 }, try follow.recvGrow(gpa, &buf)); + try std.testing.expectEqualSlices(i64, &.{ 22, 2, 5 }, try follow.recvGrow(gpa, &buf)); + try std.testing.expectEqualSlices(i64, &.{ 21, 4, 9, 9, 9 }, try follow.recvGrow(gpa, &buf)); + var out: [8192]i64 = undefined; + try std.testing.expectEqualSlices(i64, &.{ 22, 3, 1 }, try follow.recv(&out)); + try std.testing.expectEqualSlices(i64, &big, try follow.recvGrow(gpa, &buf)); + try std.testing.expectEqual(@as(usize, 0), (try follow.recv(&out)).len); + try std.testing.expectEqualSlices(i64, &.{7}, try follow.recv(&out)); + try std.testing.expectEqual(@as(u64, 3), follow.batches); + try std.testing.expectError(error.NotRankZero, follow.hold(&.{1})); + lead.close(); +} diff --git a/zig/src/fabric/tp/reference.zig b/zig/src/fabric/tp/reference.zig new file mode 100644 index 0000000000..b98966fcc8 --- /dev/null +++ b/zig/src/fabric/tp/reference.zig @@ -0,0 +1,121 @@ +//! Host references for the collectives: reductions in rank order, bf16 / f16 added in fp32 and rounded once. +const std = @import("std"); +const collective = @import("collective.zig"); +const DType = collective.DType; +const Op = collective.Op; + +/// bf16 bits to fp32 (exact). +pub fn bf16ToF32(h: u16) f32 { + return @bitCast(@as(u32, h) << 16); +} + +/// fp32 to bf16 bits, round to nearest even; NaN stays a quiet NaN. +pub fn f32ToBf16(f: f32) u16 { + const u: u32 = @bitCast(f); + if (std.math.isNan(f)) return @intCast((u >> 16) | 0x40); + const round = 0x7FFF + ((u >> 16) & 1); + return @intCast((u +% round) >> 16); +} + +fn combine(comptime T: type, a: T, b: T, op: Op) T { + return switch (op) { + .sum => if (@typeInfo(T) == .int) a +% b else a + b, + .max => @max(a, b), + .min => @min(a, b), + }; +} + +fn reduceTyped(comptime T: type, acc: []u8, src: []const u8, op: Op) void { + const n = acc.len / @sizeOf(T); + for (0..n) |i| { + const a = std.mem.readInt(@Int(.unsigned, @bitSizeOf(T)), acc[i * @sizeOf(T) ..][0..@sizeOf(T)], .little); + const b = std.mem.readInt(@Int(.unsigned, @bitSizeOf(T)), src[i * @sizeOf(T) ..][0..@sizeOf(T)], .little); + const r: T = combine(T, @bitCast(a), @bitCast(b), op); + std.mem.writeInt(@Int(.unsigned, @bitSizeOf(T)), acc[i * @sizeOf(T) ..][0..@sizeOf(T)], @bitCast(r), .little); + } +} + +fn reduceHalf(acc: []u8, src: []const u8, op: Op, bf: bool) void { + const n = acc.len / 2; + for (0..n) |i| { + const a = std.mem.readInt(u16, acc[2 * i ..][0..2], .little); + const b = std.mem.readInt(u16, src[2 * i ..][0..2], .little); + const fa: f32 = if (bf) bf16ToF32(a) else @floatCast(@as(f16, @bitCast(a))); + const fb: f32 = if (bf) bf16ToF32(b) else @floatCast(@as(f16, @bitCast(b))); + const r = combine(f32, fa, fb, op); + const out: u16 = if (bf) f32ToBf16(r) else @bitCast(@as(f16, @floatCast(r))); + std.mem.writeInt(u16, acc[2 * i ..][0..2], out, .little); + } +} + +/// acc <- acc (op) src, element-wise; `acc` holds the lower ranks' result, `src` the next rank's input. +pub fn accumulate(acc: []u8, src: []const u8, t: DType, op: Op) void { + std.debug.assert(acc.len == src.len and acc.len % t.size() == 0); + switch (t) { + .u8 => reduceTyped(u8, acc, src, op), + .i32 => reduceTyped(i32, acc, src, op), + .i64 => reduceTyped(i64, acc, src, op), + .f32 => reduceTyped(f32, acc, src, op), + .f64 => reduceTyped(f64, acc, src, op), + .bf16 => reduceHalf(acc, src, op, true), + .f16 => reduceHalf(acc, src, op, false), + } +} + +/// out <- inputs[0] (op) inputs[1] (op) ..., in rank order. +pub fn reduce(out: []u8, inputs: []const []const u8, t: DType, op: Op) void { + @memcpy(out, inputs[0]); + for (inputs[1..]) |in| accumulate(out, in, t, op); +} + +/// Deterministic test data for `rank`'s input of `bytes`: finite values of type `t` that differ per rank and position. +pub fn fill(out: []u8, t: DType, rank: u32, seed: u64) void { + var rng = std.Random.DefaultPrng.init(seed ^ (@as(u64, rank) *% 0x9E3779B97F4A7C15)); + const r = rng.random(); + const n = out.len / t.size(); + for (0..n) |i| { + const x = r.float(f32) * 8.0 - 4.0; + switch (t) { + .u8 => out[i] = r.int(u8), + .i32 => std.mem.writeInt(i32, out[4 * i ..][0..4], r.int(i32), .little), + .i64 => std.mem.writeInt(i64, out[8 * i ..][0..8], r.int(i64), .little), + .f32 => std.mem.writeInt(u32, out[4 * i ..][0..4], @bitCast(x), .little), + .f64 => std.mem.writeInt(u64, out[8 * i ..][0..8], @bitCast(@as(f64, x) * 1.000000123), .little), + .bf16 => std.mem.writeInt(u16, out[2 * i ..][0..2], f32ToBf16(x), .little), + .f16 => std.mem.writeInt(u16, out[2 * i ..][0..2], @bitCast(@as(f16, @floatCast(x))), .little), + } + } +} + +test "bf16 rounding is nearest even" { + try std.testing.expectEqual(@as(u16, 0x3F80), f32ToBf16(1.0)); + // 1 + 2^-8 is exactly halfway between bf16 1.0 and 1.0078125: rounds to the even one (1.0) + try std.testing.expectEqual(@as(u16, 0x3F80), f32ToBf16(1.00390625)); + try std.testing.expectEqual(@as(u16, 0x3F82), f32ToBf16(1.01171875)); + try std.testing.expect(std.math.isNan(bf16ToF32(f32ToBf16(std.math.nan(f32))))); +} + +test "rank-order reduction of bf16 adds in fp32 and rounds once" { + var a: [4]u8 = undefined; + var b: [4]u8 = undefined; + std.mem.writeInt(u16, a[0..2], f32ToBf16(1.0), .little); + std.mem.writeInt(u16, b[0..2], f32ToBf16(0.00390625), .little); + std.mem.writeInt(u16, a[2..4], f32ToBf16(-2.5), .little); + std.mem.writeInt(u16, b[2..4], f32ToBf16(3.0), .little); + var out: [4]u8 = undefined; + reduce(&out, &.{ &a, &b }, .bf16, .sum); + try std.testing.expectEqual(f32ToBf16(1.0), std.mem.readInt(u16, out[0..2], .little)); + try std.testing.expectEqual(f32ToBf16(0.5), std.mem.readInt(u16, out[2..4], .little)); + reduce(&out, &.{ &a, &b }, .bf16, .max); + try std.testing.expectEqual(f32ToBf16(3.0), std.mem.readInt(u16, out[2..4], .little)); +} + +test "integer sums wrap" { + var a: [4]u8 = undefined; + var b: [4]u8 = undefined; + std.mem.writeInt(i32, &a, std.math.maxInt(i32), .little); + std.mem.writeInt(i32, &b, 1, .little); + var out: [4]u8 = undefined; + reduce(&out, &.{ &a, &b }, .i32, .sum); + try std.testing.expectEqual(@as(i32, std.math.minInt(i32)), std.mem.readInt(i32, &out, .little)); +} diff --git a/zig/src/fabric/tp/roce.zig b/zig/src/fabric/tp/roce.zig new file mode 100644 index 0000000000..c9bb72a2c3 --- /dev/null +++ b/zig/src/fabric/tp/roce.zig @@ -0,0 +1,525 @@ +//! The RoCE backend between two hosts: the tp_roce kernel stages and rings a doorbell, a proxy thread moves the slot. +const std = @import("std"); +const log = @import("log.zig"); +const cuda = @import("cuda"); +const collective = @import("collective.zig"); +const Bootstrap = @import("bootstrap.zig").Bootstrap; +const sock = @import("sock.zig"); +const rv = @import("roce_verbs.zig"); +const shm = @import("roce_shm.zig"); +const guard = @import("roce_guard.zig"); +const HostRegister = @import("mailbox.zig").HostRegister; +const posix = std.posix; +const DevicePtr = collective.DevicePtr; +const DType = collective.DType; +const Error = collective.Error; + +pub const max_blocks = 16; +const threads = 512; +const page = std.heap.page_size_min; +const flag_stride = 128; +// region: ctrl [4 KiB], flags [2 slots][2 HCAs] at 128 B, send [2][slot], recv [2][slot] (mailbox.cu's tp_roce) +const ctrl_off = 0; +const flag_off = 4096; +const send_off = 8192; +// ctrl words (mailbox.cu): doorbell, padded bytes per slot, completed, abort, proxy failure, a failed wait's record +const w_seq = 0; +const w_slot_bytes = 4; +const w_completed = 8; +const w_abort = 12; +const w_proxy_failed = 13; +const w_failed = 16; +/// Model graphs leave sub-millisecond gaps: the proxy stays hot this many idle polls, then naps 20 us at a time. +const idle_spins: u64 = 20_000_000; + +const Mode = enum(u32) { gather = 0, sum_f32 = 1, sum_bf16 = 2, exchange = 3 }; + +/// mailbox.cu's `RoceArgs`, field for field. +const RoceArgs = extern struct { + in: DevicePtr, + out: DevicePtr, + send_slots: DevicePtr, + recv_slots: DevicePtr, + flags: DevicePtr, + ctrl: DevicePtr, + state: DevicePtr, + nbytes: u64, + slot_bytes: u64, + timeout_ns: u64, + rank: u32, + mode: u32, + n_hca: u32, + /// TF_TP_ROCE_FAST (config.roce_fast; 0: the classic path) + opts: u32 = 0, + /// gather only: device int32 [2] of each rank's bytes (0: every shard is `nbytes`) + lens: DevicePtr = 0, + /// Chunked gathers: this op's byte offset into each rank's shard (it moves at most a slot from there; 0 otherwise). + chunk: u64 = 0, +}; + +pub const WireKind = enum { + /// RDMA writes over the HCAs. + verbs, + /// The proxy copies into the peer's region in shared memory: the same protocol on one host without RDMA (tests). + shm, +}; + +pub const Options = struct { + wire: WireKind = .verbs, + hcas: []const rv.HcaSpec = &.{}, + traffic_class: u8 = 0, + cpu: ?u32 = null, + max_bytes: usize = 256 * 1024, + timeout_ns: u64 = 120 * std.time.ns_per_s, + /// TF_TP_ROCE_FAST: tp_roce's `opts` word (config.roce_fast) + fast: u32 = 0, + /// TF_TP_ROCE_BLOCKS: the largest grid (config.roce_blocks) + blocks: u32 = 16, + /// The shm wire's stripes, as HCAs: its posts take the verbs wire's split and bounds. + shm_stripes: u32 = 1, + /// Tests: sees every shm post's stripes before the copy. + tap: ?*shm.PostTap = null, +}; + +const Wire = union(WireKind) { + verbs: rv.VerbsWire, + shm: shm.ShmWire, + + fn post(w: *Wire, slot: u32, seq: u32, padded: u32) !void { + switch (w.*) { + .verbs => |*v| try v.post(slot, seq, padded), + .shm => |*s| try s.post(slot, seq, padded), + } + } + + fn reap(w: *Wire) !void { + switch (w.*) { + .verbs => |*v| try v.reap(), + .shm => {}, + } + } + + fn hcas(w: *const Wire) u32 { + return switch (w.*) { + .verbs => |v| v.n_hca, + .shm => |s| s.stripes, + }; + } +}; + +pub const Roce = struct { + me: u32, + slot_bytes: usize, + timeout_ns: u64, + /// TF_TP_ROCE_FAST: the kernel's `opts` word + fast: u32 = 0, + /// TF_TP_ROCE_BLOCKS: the largest grid + max_grid: u32 = 16, + /// This rank's region (for shm: half of the mapping in `mapping`). + region: []align(page) u8, + mapping: []align(page) u8, + wire: Wire, + cpu: ?u32, + // device side (absent in host-only use: tests and `tf-tp-test roce-host`) + d: ?*const cuda.Driver = null, + hr: ?HostRegister = null, + dev: DevicePtr = 0, + state: ?cuda.DeviceBuffer = null, + module: ?cuda.Module = null, + func: ?cuda.Function = null, + /// The kernel image takes `RoceArgs.chunk` (it exports tp_roce_chunk_abi): TF_TP_ROCE_CHUNKED_KB needs it. + chunk_abi: bool = false, + // the proxy thread + thread: ?std.Thread = null, + running: std.atomic.Value(bool) = .init(false), + last_seq: u32 = 0, + err_buf: [160]u8 = undefined, + err_len: usize = 0, + // host emulation of the kernel's side (tests) + host_epoch: u32 = 0, + aborted: std.atomic.Value(bool) = .init(false), + + pub fn regionBytes(slot: usize) usize { + return send_off + 4 * slot; + } + + /// Both ranks call it after the bootstrap; any rank failing makes all return `error.RoceUnavailable` (collective). + pub fn init(d: ?*const cuda.Driver, b: *Bootstrap, image: []const u8, o: Options) !*Roce { + if (b.world != 2) return error.Unsupported; + if (o.max_bytes > guard.max_write) return error.WriteTooLarge; // a slot no post could fill in one write + const slot = std.mem.alignForward(usize, o.max_bytes, 4096); + const len = regionBytes(slot); + const r = try std.heap.page_allocator.create(Roce); + errdefer std.heap.page_allocator.destroy(r); + r.* = .{ .me = b.rank, .slot_bytes = slot, .timeout_ns = o.timeout_ns, .fast = o.fast, .max_grid = o.blocks, .region = undefined, .mapping = undefined, .wire = undefined, .cpu = o.cpu }; + // 1. the regions and the wire, each rank alone; every rank learns whether all did + const ok = r.openWire(b, o, len) catch |e| blk: { + log.warn("[tensorfold] tp roce: rank {d} setup failed: {t}", .{ b.rank, e }); + break :blk false; + }; + if (!try agreeAll(b, ok, "RoCE setup")) { + if (ok) r.closeWire(); + return error.RoceUnavailable; + } + errdefer r.closeWire(); + // 2. the verbs connection records, through the bootstrap + const connected = r.connectWire(b) catch |e| blk: { + log.warn("[tensorfold] tp roce: rank {d} connect failed: {t}", .{ b.rank, e }); + break :blk false; + }; + if (!try agreeAll(b, connected, "RoCE connect")) return error.RoceUnavailable; + // 3. the device side + if (d) |drv| { + const dev_ok = r.openDevice(drv, image) catch |e| blk: { + log.warn("[tensorfold] tp roce: rank {d} device setup failed: {t}", .{ b.rank, e }); + break :blk false; + }; + if (!try agreeAll(b, dev_ok, "RoCE device setup")) { + if (dev_ok) r.closeDevice(); + return error.RoceUnavailable; + } + } + r.running.store(true, .release); + r.thread = try std.Thread.spawn(.{}, proxyLoop, .{r}); + try b.agree("proxy", "RoCE proxy start"); + return r; + } + + fn agreeAll(b: *Bootstrap, ok: bool, _: []const u8) !bool { + var blobs: [2][]const u8 = undefined; + var buf: [2]u8 = undefined; + const all = try b.allGather(&.{@intFromBool(ok)}, &buf, &blobs); + return all[0][0] == 1 and all[1][0] == 1; + } + + fn openWire(r: *Roce, b: *Bootstrap, o: Options, len: usize) !bool { + switch (o.wire) { + .verbs => { + r.mapping = try posix.mmap(null, len, .{ .READ = true, .WRITE = true }, .{ .TYPE = .PRIVATE, .ANONYMOUS = true }, -1, 0); + r.region = r.mapping; + errdefer posix.munmap(r.mapping); + try rv.prepareRegion(r.mapping); + r.wire = .{ .verbs = try rv.VerbsWire.open(o.hcas, r.region, .{ .send = send_off, .recv = send_off + 2 * r.slot_bytes, .flags = flag_off, .slot = r.slot_bytes }, o.traffic_class) }; + }, + .shm => { + r.mapping = try shm.sharedPair(b, len); + r.region = @alignCast(r.mapping[r.me * len ..][0..len]); + const layout: guard.Layout = .{ .send = send_off, .recv = send_off + 2 * r.slot_bytes, .flags = flag_off, .slot = r.slot_bytes, .len = len }; + r.wire = .{ .shm = .{ .peer = r.mapping[(1 - r.me) * len ..][0..len], .layout = layout, .local = r.region, .stripes = o.shm_stripes, .tap = o.tap } }; + }, + } + return true; + } + + fn connectWire(r: *Roce, b: *Bootstrap) !bool { + switch (r.wire) { + .shm => return true, + .verbs => |*v| { + const mine = v.record(); + var blobs: [2][]const u8 = undefined; + var buf: [2 * @sizeOf(rv.Record)]u8 = undefined; + const all = try b.allGather(std.mem.asBytes(&mine), &buf, &blobs); + var peer: rv.Record = undefined; + if (all[1 - r.me].len != @sizeOf(rv.Record)) return error.BadRecord; + @memcpy(std.mem.asBytes(&peer), all[1 - r.me]); + try v.connect(peer); + var line: [512]u8 = undefined; + log.info("[tensorfold] tp roce: rank {d} registered {s}", .{ r.me, v.describe(&line) }); + return true; + }, + } + } + + fn openDevice(r: *Roce, d: *const cuda.Driver, image: []const u8) !bool { + if (image.len == 0) return error.NoKernel; + r.d = d; + r.hr = try HostRegister.open(); + try d.check(r.hr.?.register(r.region.ptr, r.region.len, 1 | 2), "cuMemHostRegister"); // PORTABLE | DEVICEMAP + try d.check(d.api.cuMemHostGetDevicePointer_v2(&r.dev, r.region.ptr, 0), "cuMemHostGetDevicePointer"); + r.state = try cuda.DeviceBuffer.alloc(d, 16); + try r.state.?.fill8(0, null); + r.module = try cuda.Module.load(d, image); + r.func = try r.module.?.function("tp_roce"); + r.chunk_abi = if (r.module.?.function("tp_roce_chunk_abi")) |_| true else |_| false; + return true; + } + + fn closeDevice(r: *Roce) void { + if (r.module) |*m| m.unload(); + if (r.state) |*s| s.free(); + if (r.hr) |*h| { + _ = h.unregister(r.region.ptr); + h.lib.close(); + } + r.module = null; + r.state = null; + r.hr = null; + } + + fn closeWire(r: *Roce) void { + switch (r.wire) { + .verbs => |*v| v.close(), + .shm => {}, + } + posix.munmap(r.mapping); + } + + /// Regions a failed transport left registered and mapped (the process exit is their only teardown). + pub var kept: std.atomic.Value(u32) = .init(0); + + /// HCAs and GPU quiet, then unregister and unmap; a failed transport's region is kept until exit. + pub fn deinit(r: *Roce) void { + const was_failed = r.failed(); + r.word(w_abort).store(1, .release); + r.running.store(false, .release); + if (r.thread) |t| t.join(); + const quiet = r.quiesce(); + if (was_failed or !quiet) { + const why = if (was_failed) "failed" else "did not quiesce"; + log.warn("[tensorfold] tp roce: rank {d}: the transport {s}; region 0x{x} ({d} B) kept registered until exit", .{ r.me, why, @intFromPtr(r.region.ptr), r.region.len }); + _ = kept.fetchAdd(1, .monotonic); + return; + } + r.closeDevice(); + r.closeWire(); + std.heap.page_allocator.destroy(r); + } + + /// The QPs to error with every posted request reaped, and the GPU's work on the region done. + fn quiesce(r: *Roce) bool { + const wire = switch (r.wire) { + .verbs => |*v| v.quiesce(rv.quiesce_ns), + .shm => true, + }; + const gpu = if (r.d) |d| d.api.cuCtxSynchronize() == cuda.abi.success else true; + var line: [384]u8 = undefined; + log.info("[tensorfold] tp roce: rank {d} quiesced ({s}): {s}; GPU sync {s}", .{ r.me, if (wire) "drained" else "NOT drained", r.describeLink(&line), if (gpu) "done" else "FAILED" }); + return wire and gpu; + } + + /// Each HCA's QP state, posts and completions, and what a drain reaped (tf-tp-test's idle check; shm: none). + pub fn describeLink(r: *const Roce, buf: []u8) []const u8 { + return switch (r.wire) { + .verbs => |*v| v.report(buf), + .shm => "shm wire", + }; + } + + fn word(r: *const Roce, i: usize) *std.atomic.Value(u32) { + return @ptrCast(@alignCast(r.region[ctrl_off + 4 * i ..].ptr)); + } + + fn fail(r: *Roce, comptime fmt: []const u8, args: anytype) void { + const s = std.fmt.bufPrint(&r.err_buf, fmt, args) catch "proxy failure"; + r.err_len = s.len; + log.err("[tensorfold] tp roce proxy: {s}", .{s}); + r.word(w_proxy_failed).store(1, .release); + r.word(w_abort).store(1, .release); // waits on this rank end now instead of at the timeout + r.halt(); + } + + /// The HCAs stop: nothing of this rank's goes out and none of the peer's writes lands while the ranks go down. + fn halt(r: *Roce) void { + switch (r.wire) { + .verbs => |*v| { + const ok = v.halt(); + var line: [384]u8 = undefined; + log.warn("[tensorfold] tp roce: rank {d} halted{s}: {s}", .{ r.me, if (ok) "" else " (a QP refused ERR)", v.report(&line) }); + }, + .shm => {}, + } + } + + /// Every post of this rank's completed (shm: nothing is ever in flight). + pub fn linkIdle(r: *const Roce) bool { + return switch (r.wire) { + .verbs => |*v| v.idle(), + .shm => true, + }; + } + + /// tf-tp-test's `failfast RANK transport`: the proxy's failure path (`fail`), here with no op in flight. + pub fn inject(r: *Roce) void { + r.fail("failure injected by tf-tp-test (no RDMA fault)", .{}); + } + + /// The transport failed or was aborted: a device may still write to its region. + pub fn failed(r: *const Roce) bool { + if (r.aborted.load(.acquire) or r.word(w_abort).load(.acquire) != 0) return true; + return r.word(w_proxy_failed).load(.acquire) != 0 or r.word(w_failed).load(.acquire) != 0; + } + + /// The proxy loop: spin on the doorbell, post each new sequence (at most two pending), reap while idle, then nap. + fn proxyLoop(r: *Roce) void { + if (r.cpu) |cpu| pin(cpu); + var idle: u64 = 0; + while (r.running.load(.monotonic)) { + if (r.word(w_abort).load(.monotonic) != 0) return; + const seq = r.word(w_seq).load(.acquire); + if (seq == r.last_seq) { + idle += 1; + if (idle % 64 == 0) r.wire.reap() catch |e| return r.fail("completion: {t}", .{e}); + if (idle >= idle_spins) sock.sleepNs(20 * std.time.ns_per_us) else std.atomic.spinLoopHint(); + continue; + } + idle = 0; + var pending = seq -% r.last_seq; + if (pending > 2) return r.fail("doorbell skipped {d} ops (posted {d}, now {d})", .{ pending, r.last_seq, seq }); + var s = r.last_seq +% 1; + while (pending > 0) : ({ + s +%= 1; + pending -= 1; + }) { + const padded = r.word(w_slot_bytes + (s & 1)).load(.acquire); + if (padded == 0 or padded % 16 != 0 or padded > r.slot_bytes) return r.fail("op {d} of {d} bytes", .{ s, padded }); + r.wire.post(s & 1, s, padded) catch |e| return r.fail("post of op {d}: {t}", .{ s, e }); + r.last_seq = s; + } + r.wire.reap() catch |e| return r.fail("completion: {t}", .{e}); + } + } + + pub fn fits(r: *const Roce, n: usize) bool { + return n <= r.slot_bytes; + } + + fn launch(r: *Roce, mode: Mode, in: DevicePtr, out: DevicePtr, n: usize, lens: DevicePtr, stream: collective.Stream) Error!void { + return r.launchAt(mode, in, out, n, lens, 0, stream); + } + + fn launchAt(r: *Roce, mode: Mode, in: DevicePtr, out: DevicePtr, n: usize, lens: DevicePtr, chunk: u64, stream: collective.Stream) Error!void { + if (r.aborted.load(.acquire)) return error.Aborted; + if (n > r.slot_bytes and (lens == 0 or !r.chunk_abi)) return error.Unsupported; + if (n == 0) return; + const d = r.d orelse return error.Unsupported; + var a: RoceArgs = .{ + .in = in, + .out = out, + .send_slots = r.dev + send_off, + .recv_slots = r.dev + send_off + 2 * r.slot_bytes, + .flags = r.dev + flag_off, + .ctrl = r.dev + ctrl_off, + .state = r.state.?.ptr, + .nbytes = n, + .slot_bytes = r.slot_bytes, + .timeout_ns = r.timeout_ns, + .rank = r.me, + .mode = @backingInt(mode), + .n_hca = r.wire.hcas(), + .opts = r.fast, + .lens = lens, + .chunk = chunk, + }; + var params = [1]?*anyopaque{@ptrCast(&a)}; + const blocks = @min(@import("mailbox.zig").blocksFor(@min(n, r.slot_bytes)), r.max_grid); + d.check(d.api.cuLaunchKernel(r.func.?.handle, blocks, 1, 1, threads, 1, 1, 0, stream, ¶ms, null), "cuLaunchKernel tp_roce") catch return error.BackendFailed; + } + + pub fn allGatherBytes(r: *Roce, s: DevicePtr, o: DevicePtr, n: usize, stream: collective.Stream) Error!void { + return r.launch(.gather, s, o, n, 0, stream); + } + + /// `allGatherBytes` where rank r posts only lens[r] bytes (device int32 [2], padded to 16, at least 16). + pub fn allGatherV(r: *Roce, s: DevicePtr, o: DevicePtr, n: usize, lens: DevicePtr, stream: collective.Stream) Error!void { + if (n > r.slot_bytes) return error.Unsupported; // one op: past a slot is allGatherVChunked's + if (lens == 0) return error.Invalid; + return r.launch(.gather, s, o, n, lens, stream); + } + + /// `allGatherV` past the slot: op c moves at most a slot of each shard from c x slot (empty on both ranks: no op). + pub fn allGatherVChunked(r: *Roce, s: DevicePtr, o: DevicePtr, n: usize, lens: DevicePtr, stream: collective.Stream) Error!void { + if (lens == 0) return error.Invalid; + if (!r.chunk_abi) return error.Unsupported; + var off: usize = 0; + while (off < n) : (off += r.slot_bytes) try r.launchAt(.gather, s, o, n, lens, off, stream); + } + + pub fn exchangeBytes(r: *Roce, s: DevicePtr, o: DevicePtr, n: usize, stream: collective.Stream) Error!void { + return r.launch(.exchange, s, o, n, 0, stream); + } + + pub fn sum(r: *Roce, s: DevicePtr, o: DevicePtr, count: usize, t: DType, stream: collective.Stream) Error!void { + const mode: Mode = switch (t) { + .f32 => .sum_f32, + .bf16 => .sum_bf16, + else => return error.Unsupported, + }; + return r.launch(mode, s, o, count * t.size(), 0, stream); + } + + /// Ends this rank's waits and the proxy (the peer learns through the fate channel and aborts its own). + pub fn abort(r: *Roce) void { + r.aborted.store(true, .release); + r.word(w_abort).store(1, .release); + r.halt(); + } + + pub fn check(r: *Roce) Error!void { + if (r.word(w_failed).load(.acquire) != 0) { + log.warn("tp roce: rank {d}'s wait for seq {d} failed ({s} after {d} us; HCA {d}, flag held {d}); doorbell {d}, posted {d}, completed {d}, abort {d}, proxy failed {d}, flags {d} {d}", .{ r.me, r.word(w_failed + 1).load(.acquire), if (r.word(w_failed + 4).load(.acquire) == 1) "aborted" else "timed out", r.word(w_failed + 5).load(.acquire), r.word(w_failed + 3).load(.acquire), r.word(w_failed + 2).load(.acquire), r.word(w_seq).load(.acquire), r.last_seq, r.word(w_completed).load(.acquire), r.word(w_abort).load(.acquire), r.word(w_proxy_failed).load(.acquire), r.flagWord(0), r.flagWord(1) }); + return error.PeerFailed; + } + if (r.word(w_proxy_failed).load(.acquire) != 0) return error.BackendFailed; + if (r.aborted.load(.acquire) or r.word(w_abort).load(.acquire) != 0) return error.Aborted; + } + + fn flagWord(r: *const Roce, slot: usize) u32 { + const p: *std.atomic.Value(u32) = @ptrCast(@alignCast(r.region[flag_off + slot * 2 * flag_stride ..].ptr)); + return p.load(.acquire); + } + + pub fn proxyError(r: *const Roce) []const u8 { + return r.err_buf[0..r.err_len]; + } + + /// The kernel's side of one op done by the host (tests, `roce-host`): a gather through the same slots and flags. + pub fn hostGather(r: *Roce, in: []const u8, out: []u8) Error!void { + return r.hostGatherV(in, out, null); + } + + /// `hostGather` with variable lengths (`allGatherV`): rank k's first lens[k] bytes land at out + k x n. + pub fn hostGatherV(r: *Roce, in: []const u8, out: []u8, lens: ?[2]u32) Error!void { + if (in.len > r.slot_bytes) return error.Invalid; + return r.hostGatherAt(in, out, lens, 0); + } + + /// `allGatherVChunked`'s ops done by the host (the kernel's side of each, as `hostGatherV`). + pub fn hostGatherVChunked(r: *Roce, in: []const u8, out: []u8, lens: [2]u32) Error!void { + var off: usize = 0; + while (off < in.len) : (off += r.slot_bytes) try r.hostGatherAt(in, out, lens, off); + } + + /// One op from shard offset `off`: each rank's bytes past it, at most a slot; an op empty on both ranks is skipped. + fn hostGatherAt(r: *Roce, in: []const u8, out: []u8, lens: ?[2]u32, off: usize) Error!void { + const n = in.len; + if (n == 0 or out.len < 2 * n) return error.Invalid; + if (lens) |l| if (l[0] > n or l[1] > n) return error.Invalid; + const full: [2]usize = if (lens) |l| .{ l[0], l[1] } else .{ n, n }; + const mine: usize = @min(full[r.me] -| off, r.slot_bytes); + const theirs: usize = @min(full[1 - r.me] -| off, r.slot_bytes); + if (off > 0 and mine == 0 and theirs == 0) return; + const seq = r.host_epoch +% 1; + const slot: usize = seq & 1; + @memcpy(r.region[send_off + slot * r.slot_bytes ..][0..mine], in[off..][0..mine]); + r.word(w_slot_bytes + slot).store(@intCast(@max(16, std.mem.alignForward(usize, mine, 16))), .release); + r.word(w_seq).store(seq, .release); + @memcpy(out[r.me * n + off ..][0..mine], in[off..][0..mine]); + const t0 = sock.nowNs(); + for (0..r.wire.hcas()) |h| { + const flag: *std.atomic.Value(u32) = @ptrCast(@alignCast(r.region[flag_off + (slot * 2 + h) * flag_stride ..].ptr)); + while (flag.load(.acquire) != seq) { + if (r.word(w_abort).load(.monotonic) != 0) return error.Aborted; + if (sock.nowNs() - t0 > r.timeout_ns) return error.Timeout; + std.atomic.spinLoopHint(); + } + } + @memcpy(out[(1 - r.me) * n + off ..][0..theirs], r.region[send_off + 2 * r.slot_bytes + slot * r.slot_bytes ..][0..theirs]); + r.host_epoch = seq; + } +}; + +fn pin(cpu: u32) void { + var set: std.os.linux.cpu_set_t = @splat(0); + set[cpu / @bitSizeOf(usize)] |= @as(usize, 1) << @intCast(cpu % @bitSizeOf(usize)); + std.os.linux.sched_setaffinity(0, &set) catch log.warn("tp roce: could not pin the proxy to CPU {d}", .{cpu}); +} diff --git a/zig/src/fabric/tp/roce_fake.zig b/zig/src/fabric/tp/roce_fake.zig new file mode 100644 index 0000000000..c316448ef8 --- /dev/null +++ b/zig/src/fabric/tp/roce_fake.zig @@ -0,0 +1,143 @@ +//! Fake verbs for host tests of VerbsWire: each QP's posted writes counted from the requests, completions that lag. +const std = @import("std"); +const verbs = @import("verbs"); +const abi = verbs.abi; +const rv = @import("roce_verbs.zig"); + +pub const slot: u64 = 1 << 20; +pub const send_off = 8192; + +/// What the HCAs hold (posted, not completed); every count comes from the work requests themselves. +pub const Fake = struct { + pub var region: []const u8 = &.{}; + /// The wire's slot (`Hw.wire`), every data write's bound with the send slots. + pub var slot_bytes: u64 = slot; + pub var lag: u32 = 1; + pub var polls: u64 = 0; + pub var queue: [2][512]u64 = undefined; + pub var head: [2]u32 = .{ 0, 0 }; + pub var len: [2]u32 = .{ 0, 0 }; + pub var inflight: u64 = 0; + pub var max_inflight: u64 = 0; + pub var max_write: u64 = 0; + pub var outside: u64 = 0; + pub var bytes: u64 = 0; + pub var posts: u64 = 0; + + pub fn reset(l: u32) void { + lag = l; + polls = 0; + head = .{ 0, 0 }; + len = .{ 0, 0 }; + inflight = 0; + max_inflight = 0; + max_write = 0; + outside = 0; + bytes = 0; + posts = 0; + } + + pub fn postSend(qp: *abi.Qp, wr: *abi.SendWr, _: *?*abi.SendWr) callconv(.c) c_int { + const q = qp.qp_num; + posts += 1; + var sum: u64 = 0; + var at: ?*abi.SendWr = wr; + while (at) |w| : (at = w.next) { + if (w.send_flags & abi.send_inline == 0) { + const sge = w.sg_list[0]; + const lo = @intFromPtr(region.ptr); + if (sge.addr < lo + send_off or sge.addr + sge.length > lo + send_off + 2 * slot_bytes) outside += 1; + max_write = @max(max_write, sge.length); + sum += sge.length; + } + if (w.send_flags & abi.send_signaled != 0) { + queue[q][(head[q] + len[q]) % 512] = sum; + len[q] += 1; + inflight += sum; + bytes += sum; + max_inflight = @max(max_inflight, inflight); + sum = 0; + } + } + return 0; + } + + /// One completion every `lag` polls of a QP's CQ: the HCA finishes writes far slower than they are posted. + pub fn pollCq(cq: *abi.Cq, _: c_int, wc: [*]abi.Wc) callconv(.c) c_int { + const q = cq.handle; + polls += 1; + if (polls % lag != 0 or len[q] == 0) return 0; + inflight -= queue[q][head[q]]; + head[q] = (head[q] + 1) % 512; + len[q] -= 1; + wc[0] = std.mem.zeroes(abi.Wc); + wc[0].status = abi.wc_success; + return 1; + } +}; + +/// QPs that complete nothing on their own: a QP moved to error flushes what it holds, as an HCA reports it. +pub const Halts = struct { + pub var halted: [2]bool = .{ false, false }; + pub var flushed_before_halt: u32 = 0; + pub var flushes: u32 = 0; + + pub fn reset() void { + halted = .{ false, false }; + flushed_before_halt = 0; + flushes = 0; + } + + pub fn modifyQp(qp: *abi.Qp, attr: *abi.QpAttr, mask: c_int) callconv(.c) c_int { + if (attr.qp_state == abi.qps_err and mask == abi.mask.state) halted[qp.qp_num] = true; + return 0; + } + + /// ibv_query_qp as the device answers it: ERR once halted, else RTS. + pub fn queryQp(qp: *abi.Qp, attr: *abi.QpAttr, _: c_int, _: *abi.QpInitAttr) callconv(.c) c_int { + attr.qp_state = if (halted[qp.qp_num]) abi.qps_err else abi.qps_rts; + return 0; + } + + pub fn pollCq(cq: *abi.Cq, _: c_int, wc: [*]abi.Wc) callconv(.c) c_int { + const q = cq.handle; + if (Fake.len[q] == 0 or !halted[q]) return 0; + if (!halted[0] or !halted[1]) flushed_before_halt += 1; + Fake.inflight -= Fake.queue[q][Fake.head[q]]; + Fake.head[q] = (Fake.head[q] + 1) % 512; + Fake.len[q] -= 1; + flushes += 1; + wc[0] = std.mem.zeroes(abi.Wc); + wc[0].status = 5; // IBV_WC_WR_FLUSH_ERR + return 1; + } + + /// Points a wire's QPs at these handlers. + pub fn install(hw: *Hw, w: *rv.VerbsWire) void { + w.v.api.ibv_modify_qp = modifyQp; + w.v.query_qp = queryQp; + for (&hw.ctx) |*c| c.ops.poll_cq = pollCq; + } +}; + +/// Two fake HCAs behind one VerbsWire, connected to a peer whose region has this rank's layout. +pub const Hw = struct { + ctx: [2]abi.Context = undefined, + qp: [2]abi.Qp = undefined, + cq: [2]abi.Cq = undefined, + mr: [2]abi.Mr = undefined, + + pub fn wire(hw: *Hw, region: []u8) rv.VerbsWire { + const s = Fake.slot_bytes; + var w: rv.VerbsWire = .{ .v = undefined, .traffic_class = 0, .recv_off = send_off + 2 * s, .flag_off = 4096, .send_off = send_off, .slot_bytes = s, .region = region }; + for (0..2) |i| { + hw.ctx[i].ops = .{ ._before_poll = @splat(null), .poll_cq = Fake.pollCq, ._before_post = @splat(null), .post_send = Fake.postSend, .post_recv = undefined, ._after = @splat(null) }; + hw.qp[i].qp_num = @intCast(i); + hw.cq[i].handle = @intCast(i); + hw.mr[i].lkey = 7; + w.hcas[i] = .{ .ctx = &hw.ctx[i], .pd = undefined, .mr = &hw.mr[i], .cq = &hw.cq[i], .qp = &hw.qp[i], .gid_index = 0, .gid = @splat(0), .mtu = 0, .lid = 0 }; + } + w.n_hca = 2; + return w; + } +}; diff --git a/zig/src/fabric/tp/roce_guard.zig b/zig/src/fabric/tp/roce_guard.zig new file mode 100644 index 0000000000..8f85247afd --- /dev/null +++ b/zig/src/fabric/tp/roce_guard.zig @@ -0,0 +1,73 @@ +//! Every RoCE post's bounds before the verb: each HCA's stripe inside its slot and the registered region. +const std = @import("std"); + +pub const max_hcas = 2; +pub const flag_stride = 128; + +pub const Error = error{ OutOfRegion, WriteTooLarge }; + +/// The largest RDMA write a post may make, whatever the slot (a 9.6 MB post faulted a DGX Spark's HCA); never raised. +pub const max_write: u32 = 1 << 20; + +/// Both ranks' region layout (roce.zig's): byte offsets and the registered length every post must stay inside. +pub const Layout = struct { send: u64, recv: u64, flags: u64, slot: u64, len: u64 }; + +/// One HCA's share of a post: its data stripe in this rank's send slot and the peer's receive slot, and its flag word. +pub const Stripe = struct { local: u64, remote: u64, len: u32, flag: u64 }; + +/// Slot `slot`'s first `padded` bytes striped over `n_hca` HCAs; error.OutOfRegion before any verb past a slot or MR. +pub fn stripes(l: Layout, n_hca: u32, slot: u32, padded: u32, out: *[max_hcas]Stripe) Error!void { + if (n_hca == 0 or n_hca > max_hcas or slot > 1) return error.OutOfRegion; + if (padded > max_write) return error.WriteTooLarge; + if (padded == 0 or padded % 16 != 0 or padded > l.slot) return error.OutOfRegion; + const packs = padded / 16; + var offset: u64 = 0; + for (out[0..n_hca], 0..) |*o, i| { + var stripe = packs / n_hca; + if (i < packs % n_hca) stripe += 1; + o.* = .{ .local = l.send + slot * l.slot + offset, .remote = l.recv + slot * l.slot + offset, .len = stripe * 16, .flag = l.flags + (@as(u64, slot) * max_hcas + i) * flag_stride }; + if (!within(o.local, o.len, l.send + slot * l.slot, l.slot) or !within(o.remote, o.len, l.recv + slot * l.slot, l.slot)) return error.OutOfRegion; + if (!within(o.flag, 4, l.flags, l.send - l.flags) or !within(o.local, o.len, 0, l.len) or !within(o.remote, o.len, 0, l.len)) return error.OutOfRegion; + offset += o.len; + } +} + +fn within(at: u64, len: u64, lo: u64, size: u64) bool { + return at >= lo and at + len <= lo + size; +} + +/// roce.zig's layout for a 64 KiB slot: ctrl 4 KiB, flags 4 KiB, send [2][slot], recv [2][slot]. +fn testLayout() Layout { + const slot: u64 = 64 << 10; + return .{ .send = 8192, .recv = 8192 + 2 * slot, .flags = 4096, .slot = slot, .len = 8192 + 4 * slot }; +} + +test "a post striped over two HCAs: the stripes tile the slot's first bytes, each flag on its own word" { + const l = testLayout(); + var s: [max_hcas]Stripe = undefined; + try stripes(l, 2, 1, 48, &s); + try std.testing.expectEqual(Stripe{ .local = l.send + l.slot, .remote = l.recv + l.slot, .len = 32, .flag = l.flags + 2 * flag_stride }, s[0]); + try std.testing.expectEqual(Stripe{ .local = l.send + l.slot + 32, .remote = l.recv + l.slot + 32, .len = 16, .flag = l.flags + 3 * flag_stride }, s[1]); + try stripes(l, 2, 1, @intCast(l.slot), &s); + try std.testing.expectEqual(l.recv + 2 * l.slot, s[1].remote + s[1].len); +} + +test "a post past its slot, a third slot, an unpadded or empty post, or no HCA is refused" { + const l = testLayout(); + var s: [max_hcas]Stripe = undefined; + try std.testing.expectError(error.OutOfRegion, stripes(l, 2, 0, @intCast(l.slot + 16), &s)); + try std.testing.expectError(error.OutOfRegion, stripes(l, 2, 2, 16, &s)); + try std.testing.expectError(error.OutOfRegion, stripes(l, 2, 0, 24, &s)); + try std.testing.expectError(error.OutOfRegion, stripes(l, 2, 0, 0, &s)); + try std.testing.expectError(error.OutOfRegion, stripes(l, 0, 0, 16, &s)); + try std.testing.expectError(error.OutOfRegion, stripes(l, 3, 0, 16, &s)); + var big = l; + big.slot = 10 << 20; + big.recv = big.send + 2 * big.slot; + big.len = big.send + 4 * big.slot; + try std.testing.expectError(error.WriteTooLarge, stripes(big, 2, 0, max_write + 16, &s)); + try stripes(big, 2, 0, max_write, &s); + var short = l; + short.len -= 16; + try std.testing.expectError(error.OutOfRegion, stripes(short, 2, 1, @intCast(l.slot), &s)); +} diff --git a/zig/src/fabric/tp/roce_host.zig b/zig/src/fabric/tp/roce_host.zig new file mode 100644 index 0000000000..1dec07e9ba --- /dev/null +++ b/zig/src/fabric/tp/roce_host.zig @@ -0,0 +1,54 @@ +//! `tf-tp-test roce-host`: the RoCE wire and proxy between two hosts without a GPU, checked bit for bit, then timed. +const std = @import("std"); +const tp = @import("tp"); +const roce = tp.roce; +const rv = tp.roce_verbs; + +const sizes = [_]usize{ 16, 4096, 8192, 16384, 65536, 262144 }; + +pub fn run(cfg: tp.Config, gpa: std.mem.Allocator) !u8 { + var b = try tp.Bootstrap.open(cfg); + defer b.close(); + var specs: [rv.max_hcas]rv.HcaSpec = undefined; + var n: usize = 0; + if (cfg.roce_wire == .verbs) n = try tp.roce_pick.resolve(cfg, &b, &specs); + const r = try roce.Roce.init(null, &b, "", .{ .wire = if (cfg.roce_wire == .shm) .shm else .verbs, .hcas = specs[0..n], .traffic_class = cfg.roce_tc, .cpu = cfg.roce_cpu, .max_bytes = 256 * 1024, .timeout_ns = 10 * std.time.ns_per_s }); + defer r.deinit(); + const in = try gpa.alloc(u8, 256 * 1024); + defer gpa.free(in); + const peer = try gpa.alloc(u8, 256 * 1024); + defer gpa.free(peer); + const out = try gpa.alloc(u8, 512 * 1024); + defer gpa.free(out); + // 1. every size from 1 byte to the slot, both slots, bit for bit + var checked: u32 = 0; + for (0..1000) |k| { + const len = 1 + (k * 7919) % (256 * 1024); + tp.reference.fill(in[0..len], .u8, cfg.rank, k); + tp.reference.fill(peer[0..len], .u8, 1 - cfg.rank, k); + try r.hostGather(in[0..len], out[0 .. 2 * len]); + const mine_at = cfg.rank * len; + const peer_at = (1 - cfg.rank) * len; + if (!std.mem.eql(u8, out[mine_at..][0..len], in[0..len]) or !std.mem.eql(u8, out[peer_at..][0..len], peer[0..len])) { + std.debug.print("FAIL roce-host gather {d} of {d} bytes\n", .{ k, len }); + return 1; + } + checked += 1; + } + try r.check(); + std.debug.print("PASS roce-host: {d} gathers of 1 B to 256 KiB bit-exact over {s}\n", .{ checked, @tagName(cfg.roce_wire) }); + // 2. one exchange's host-to-host time by size: median and p90 of 2,000 back-to-back ops + var times: [2000]u64 = undefined; + for (sizes) |len| { + for (×) |*t| { + const t0 = tp.sock.nowNs(); + try r.hostGather(in[0..len], out[0 .. 2 * len]); + t.* = tp.sock.nowNs() - t0; + } + std.mem.sort(u64, ×, {}, std.sort.asc(u64)); + if (cfg.rank == 0) std.debug.print("RESULT {{\"backend\":\"roce-host\",\"wire\":\"{s}\",\"bytes\":{d},\"p50_us\":{d:.2},\"p90_us\":{d:.2}}}\n", .{ @tagName(cfg.roce_wire), len, @as(f64, @floatFromInt(times[1000])) / 1e3, @as(f64, @floatFromInt(times[1800])) / 1e3 }); + } + try r.check(); + try b.agree("done", "roce-host end"); + return 0; +} diff --git a/zig/src/fabric/tp/roce_inflight_test.zig b/zig/src/fabric/tp/roce_inflight_test.zig new file mode 100644 index 0000000000..fd7be53262 --- /dev/null +++ b/zig/src/fabric/tp/roce_inflight_test.zig @@ -0,0 +1,110 @@ +//! VerbsWire on fake verbs whose completions lag: what it holds in flight and what a failed transport leaves behind. +const std = @import("std"); +const rv = @import("roce_verbs.zig"); +const fake = @import("roce_fake.zig"); +const Roce = @import("roce.zig").Roce; +const Fake = fake.Fake; +const Halts = fake.Halts; +const slot = fake.slot; + +test "chunked posts: the wire's own in-flight count and peak match the HCAs', and reach the 2 MiB cap" { + const region = try std.testing.allocator.alignedAlloc(u8, .fromByteUnits(4096), Roce.regionBytes(slot)); + defer std.testing.allocator.free(region); + Fake.region = region; + Fake.slot_bytes = slot; + Fake.reset(3); + var hw: fake.Hw = .{}; + var w = hw.wire(region); + for (1..9) |seq| { + try w.post(@intCast(seq & 1), @intCast(seq), slot); + try std.testing.expectEqual(Fake.inflight, w.inflight()); + try std.testing.expect(w.inflight() <= rv.max_inflight); + } + while (w.inflight() > 0) try w.reap(); + try std.testing.expectEqual(@as(u64, 0), Fake.inflight); + try std.testing.expectEqual(Fake.max_inflight, w.peak.load(.monotonic)); + try std.testing.expectEqual(@as(u64, 2 << 20), w.peak.load(.monotonic)); +} + +test "a failed transport: both QPs go to error before the first flush and every posted write is reaped" { + const region = try std.testing.allocator.alignedAlloc(u8, .fromByteUnits(4096), Roce.regionBytes(slot)); + defer std.testing.allocator.free(region); + @memset(region, 0); + Fake.region = region; + Fake.slot_bytes = slot; + Fake.reset(1); + Halts.reset(); + var hw: fake.Hw = .{}; + var w = hw.wire(region); + Halts.install(&hw, &w); + // two 1 MiB ops posted, neither completed: 2 MiB in flight at the HCAs, the link busy + try std.testing.expect(w.idle()); + try w.post(1, 1, slot); + try w.post(0, 2, slot); + try std.testing.expectEqual(@as(u64, 2 * slot), w.inflight()); + try std.testing.expect(!w.idle()); + try std.testing.expect(w.quiesce(std.time.ns_per_s)); + try std.testing.expect(Halts.halted[0] and Halts.halted[1]); + try std.testing.expectEqual(@as(u32, 0), Halts.flushed_before_halt); + try std.testing.expectEqual(@as(u32, 4), Halts.flushes); + try std.testing.expectEqual(@as(u64, 0), w.inflight()); + try std.testing.expectEqual(@as(u64, 0), Fake.inflight); + // the line tf-tp-test's transport run reads: the device's states and every reap, all of them flushes + var line: [256]u8 = undefined; + try std.testing.expect(w.idle()); + try std.testing.expectEqualStrings("HCA 0 QP ERR; posted 2, completed 2; HCA 1 QP ERR; posted 2, completed 2; reaped 4 (flushed 4), still posted 0", w.report(&line)); + // an aborted transport reports itself failed: its deinit keeps the region + Halts.reset(); + var r: Roce = .{ .me = 0, .slot_bytes = slot, .timeout_ns = 1, .region = region, .mapping = region, .wire = .{ .verbs = w }, .cpu = null }; + try std.testing.expect(!r.failed()); + r.abort(); + try std.testing.expect(Halts.halted[0] and Halts.halted[1]); + try std.testing.expect(r.failed()); +} + +test "a QP whose inline capacity cannot take the 4-byte flag post is refused" { + try std.testing.expect(!rv.inlineFits(.{ .max_send_wr = 1, .max_recv_wr = 1, .max_send_sge = 1, .max_recv_sge = 1, .max_inline_data = 3 })); + try std.testing.expect(rv.inlineFits(.{ .max_send_wr = 1, .max_recv_wr = 1, .max_send_sge = 1, .max_recv_sge = 1, .max_inline_data = 4 })); +} + +/// /proc/self/smaps' VmFlags of the mapping that contains `at`. +fn vmFlags(at: usize, buf: []u8) ![]const u8 { + const fd = std.c.open("/proc/self/smaps", .{}); + if (fd < 0) return error.NoSmaps; + defer _ = std.c.close(fd); + var n: usize = 0; + while (n < buf.len) { + const got = std.c.read(fd, buf[n..].ptr, buf.len - n); + if (got <= 0) break; + n += @intCast(got); + } + var lines = std.mem.splitScalar(u8, buf[0..n], '\n'); + var inside = false; + while (lines.next()) |l| { + const sp = std.mem.indexOfScalar(u8, l, ' ') orelse l.len; + if (std.mem.indexOfScalar(u8, l[0..sp], '-')) |dash| { + const lo = std.fmt.parseInt(usize, l[0..dash], 16) catch continue; + const hi = std.fmt.parseInt(usize, l[dash + 1 .. sp], 16) catch continue; + inside = lo <= at and at < hi; + } + if (inside and std.mem.startsWith(u8, l, "VmFlags:")) return l; + } + return error.NotFound; +} + +test "the region is DONTFORK and NOHUGEPAGE before registration" { + if (@import("builtin").os.tag != .linux) return error.SkipZigTest; + const len = Roce.regionBytes(slot); + const m = try std.posix.mmap(null, len, .{ .READ = true, .WRITE = true }, .{ .TYPE = .PRIVATE, .ANONYMOUS = true }, -1, 0); + defer std.posix.munmap(m); + const buf = try std.testing.allocator.alloc(u8, 4 << 20); + defer std.testing.allocator.free(buf); + // smaps: dc = MADV_DONTFORK, nh = MADV_NOHUGEPAGE + const before = try vmFlags(@intFromPtr(m.ptr), buf); + try std.testing.expect(std.mem.indexOf(u8, before, " dc") == null and std.mem.indexOf(u8, before, " nh") == null); + try rv.prepareRegion(m); + // An interior page stays inside the VMA even when a neighbouring mapping merges with it. + const flags = try vmFlags(@intFromPtr(m.ptr) + 4096, buf); + try std.testing.expect(std.mem.indexOf(u8, flags, " dc") != null); + try std.testing.expect(std.mem.indexOf(u8, flags, " nh") != null); +} diff --git a/zig/src/fabric/tp/roce_pick.zig b/zig/src/fabric/tp/roce_pick.zig new file mode 100644 index 0000000000..d1ed7a6529 --- /dev/null +++ b/zig/src/fabric/tp/roce_pick.zig @@ -0,0 +1,448 @@ +//! Which RoCE devices a rank stripes over: TF_TP_ROCE_HCA, else the usable ports on an IPv4 /24 the peer also has. +const std = @import("std"); +const log = @import("log.zig"); +const rv = @import("roce_verbs.zig"); +const bootstrap = @import("bootstrap.zig"); +const sock = @import("sock.zig"); +const Config = @import("config.zig").Config; +const Verbs = @import("verbs").Verbs; +const OpenFn = *const fn (?[:0]const u8) @import("verbs").Error!Verbs; + +/// Ports a rank lists to its peers (a DGX Spark's ConnectX-7 shows four RDMA devices). +pub const max_ports = 8; +const blob_bytes = 1 + 4 * max_ports; + +/// How many devices, and how: on a /24 the peer shares (`shared`), named by TF_TP_ROCE_HCA (`explicit`), or neither. +pub const Pick = struct { + n: usize, + shared: bool, + explicit: bool = false, + verbs_open_error: ?@import("verbs").Error = null, +}; + +fn usable(p: rv.Port) bool { + return p.active and p.ethernet and p.ipv4 != null; +} + +fn subnet(a: [4]u8) u32 { + return @as(u32, a[0]) << 16 | @as(u32, a[1]) << 8 | a[2]; +} + +fn countIn(list: []const [4]u8, net: u32) usize { + var k: usize = 0; + for (list) |a| k += @intFromBool(subnet(a) == net); + return k; +} + +/// Usable ports on a /24 the peers list, by subnet, each /24 as often as both sides have it; else the first usable. +pub fn choose(ports: []const rv.Port, peer: []const [4]u8, out: *[rv.max_hcas]rv.HcaSpec, limit: usize) Pick { + const cap = @min(limit, rv.max_hcas); + var order: [max_ports]usize = undefined; + var m: usize = 0; + for (ports[0..@min(ports.len, max_ports)], 0..) |p, i| if (usable(p)) { + order[m] = i; + m += 1; + }; + var n: usize = 0; + for (order[0..m]) |i| n += @intFromBool(countIn(peer, subnet(ports[i].ipv4.?)) > 0); + if (n == 0) { + for (order[0..@min(m, cap)], 0..) |i, k| out[k] = ports[i].spec; + return .{ .n = @min(m, cap), .shared = false }; + } + // stable insertion sort by subnet (at most max_ports entries) + for (1..m) |a| { + var j = a; + while (j > 0 and subnet(ports[order[j - 1]].ipv4.?) > subnet(ports[order[j]].ipv4.?)) : (j -= 1) { + std.mem.swap(usize, &order[j - 1], &order[j]); + } + } + n = 0; + var taken: [max_ports]u32 = undefined; + for (order[0..m]) |i| { + if (n == cap) break; + const net = subnet(ports[i].ipv4.?); + var mine: usize = 0; + for (taken[0..n]) |t| mine += @intFromBool(t == net); + if (mine >= countIn(peer, net)) continue; + taken[n] = net; + out[n] = ports[i].spec; + n += 1; + } + return .{ .n = n, .shared = true }; +} + +/// `[count][count x IPv4]` of the usable ports, the blob each rank sends. +fn encode(ports: []const rv.Port, out: *[blob_bytes]u8) []const u8 { + var n: u8 = 0; + for (ports[0..@min(ports.len, max_ports)]) |p| if (usable(p)) { + out[1 + 4 * @as(usize, n) ..][0..4].* = p.ipv4.?; + n += 1; + }; + out[0] = n; + return out[0 .. 1 + 4 * @as(usize, n)]; +} + +/// A peer's blob into `out`; a malformed one counts as no ports. +fn decode(blob: []const u8, out: [][4]u8) usize { + if (blob.len == 0 or blob[0] > max_ports or blob.len != 1 + 4 * @as(usize, blob[0]) or out.len < blob[0]) return 0; + for (0..blob[0]) |k| out[k] = blob[1 + 4 * k ..][0..4].*; + return blob[0]; +} + +/// The devices for the verbs wire. A collective: every rank calls it, including ranks that set TF_TP_ROCE_HCA. +pub fn resolve(cfg: Config, b: *bootstrap.Bootstrap, out: *[rv.max_hcas]rv.HcaSpec) !usize { + return (try select(cfg, b, out)).n; +} + +/// `resolve` with how the devices were chosen, over this host's ports (a collective, as `resolve`). +pub fn select(cfg: Config, b: *bootstrap.Bootstrap, out: *[rv.max_hcas]rv.HcaSpec) !Pick { + return selectWithOpen(cfg, b, out, Verbs.open); +} + +fn selectWithOpen(cfg: Config, b: *bootstrap.Bootstrap, out: *[rv.max_hcas]rv.HcaSpec, open: OpenFn) !Pick { + var ports: [max_ports]rv.Port = undefined; + var n_ports: usize = 0; + var open_error: ?@import("verbs").Error = null; + if (open(null)) |opened| { + var v = opened; + defer v.close(); + n_ports = rv.scanPorts(&v, &ports); + } else |e| { + open_error = e; + } + var pick = try selectFrom(cfg, b, ports[0..n_ports], out); + pick.verbs_open_error = open_error; + return pick; +} + +/// `select` over given ports (the scan's, or a test's fixture ports). +pub fn selectFrom(cfg: Config, b: *bootstrap.Bootstrap, ports: []const rv.Port, out: *[rv.max_hcas]rv.HcaSpec) !Pick { + var mine: [blob_bytes]u8 = undefined; + var buf: [bootstrap.max_ranks * blob_bytes]u8 = undefined; + var blobs: [bootstrap.max_ranks][]const u8 = undefined; + const all = try b.allGather(encode(ports, &mine), &buf, &blobs); + if (cfg.roce_hca) |text| return .{ .n = try rv.parseSpecs(text, out), .shared = false, .explicit = true }; + var peer: [bootstrap.max_ranks * max_ports][4]u8 = undefined; + var n_peer: usize = 0; + for (all, 0..) |blob, r| if (r != b.rank) { + n_peer += decode(blob, peer[n_peer..]); + }; + const pick = choose(ports, peer[0..n_peer], out, cfg.roce_hcas); + if (!pick.shared and pick.n > 0) log.warn("tp roce: no RoCE port shares an IPv4 /24 with the peer's; using the first {d} active (set TF_TP_ROCE_HCA)", .{pick.n}); + return pick; +} + +/// Device names joined by commas, with each `:gid` when `gid` (TF_TP_ROCE_HCA's form) or without (NCCL_IB_HCA's). +pub fn joinNames(specs: []const rv.HcaSpec, gid: bool, buf: []u8) error{WriteFailed}![]const u8 { + var w: std.Io.Writer = .fixed(buf); + for (specs, 0..) |s, i| { + try w.print("{s}{s}", .{ if (i == 0) "" else ",", s.name[0..s.name_len] }); + if (gid) if (s.gid) |g| try w.print(":{d}", .{g}); + } + return w.buffered(); +} + +/// Room for `joinNames` of every device a rank stripes over (a name is under 64 bytes, a GID index at most 3 digits). +pub const list_bytes = rv.max_hcas * (64 + 4); + +pub const HcaError = error{ + /// no active Ethernet port with an IPv4 RoCE v2 GID on this host + NoRoceDevice, + /// this host's ports share no IPv4 /24 with the peer's: NCCL would not reach the peer on them either + NoSharedSubnet, +} || @import("verbs").Error; + +/// `tf-tp-test hca`'s line for run_hosts.sh: `NCCL_IB_HCA=` the picked devices; never the first local ports alone. +pub fn hcaLine(pick: Pick, specs: []const rv.HcaSpec, buf: *[list_bytes + 16]u8) HcaError![]const u8 { + if (pick.n == 0) { + if (pick.verbs_open_error) |e| return e; + return error.NoRoceDevice; + } + if (!pick.shared and !pick.explicit) return error.NoSharedSubnet; + const prefix = "NCCL_IB_HCA="; + @memcpy(buf[0..prefix.len], prefix); + const joined = joinNames(specs[0..pick.n], false, buf[prefix.len .. buf.len - 1]) catch unreachable; // list_bytes holds them + buf[prefix.len + joined.len] = '\n'; + return buf[0 .. prefix.len + joined.len + 1]; +} + +fn port(name: []const u8, active: bool, ipv4: ?[4]u8) rv.Port { + var p: rv.Port = .{ .active = active, .ethernet = true, .ipv4 = ipv4 }; + p.spec.name_len = @intCast(name.len); + @memcpy(p.spec.name[0..name.len], name); + p.spec.gid = 3; + return p; +} + +fn names(specs: []const rv.HcaSpec, buf: []u8) []const u8 { + return joinNames(specs, false, buf) catch unreachable; +} + +// A host with two neighbours: port f0 cabled to one Spark (both functions on one /24), port f1 to the other. +const a_ports = [_]rv.Port{ + port("rocep1s0f0", true, .{ 203, 0, 113, 1 }), + port("rocep1s0f1", true, .{ 192, 0, 2, 1 }), + port("roceP2p1s0f0", true, .{ 203, 0, 113, 65 }), + port("roceP2p1s0f1", true, .{ 198, 51, 100, 1 }), +}; +// Its peer for this run: only port f1 is cabled, f0 is down. +const b_ports = [_]rv.Port{ + port("rocep1s0f0", false, null), + port("rocep1s0f1", true, .{ 192, 0, 2, 2 }), + port("roceP2p1s0f0", false, null), + port("roceP2p1s0f1", true, .{ 198, 51, 100, 2 }), +}; + +fn listed(ports: []const rv.Port, out: *[max_ports][4]u8) []const [4]u8 { + var blob: [blob_bytes]u8 = undefined; + return out[0..decode(encode(ports, &blob), out)]; +} + +test "a two-neighbour host picks the rails its peer shares, not its first two active ports" { + var specs: [rv.max_hcas]rv.HcaSpec = undefined; + var buf: [128]u8 = undefined; + // the old rule (no peer list): the first two active ports, one of them cabled to the other neighbour + const old = choose(&a_ports, &.{}, &specs, 2); + try std.testing.expect(!old.shared); + try std.testing.expectEqualStrings("rocep1s0f0,rocep1s0f1", names(specs[0..old.n], &buf)); + var peer_buf: [max_ports][4]u8 = undefined; + const b_addrs = listed(&b_ports, &peer_buf); + try std.testing.expectEqual(@as(usize, 0), countIn(b_addrs, subnet(.{ 203, 0, 113, 1 }))); + // with the peer's addresses: the two functions of the cabled port, on each side + const a = choose(&a_ports, b_addrs, &specs, 2); + try std.testing.expect(a.shared); + try std.testing.expectEqualStrings("rocep1s0f1,roceP2p1s0f1", names(specs[0..a.n], &buf)); + var a_buf: [max_ports][4]u8 = undefined; + const b = choose(&b_ports, listed(&a_ports, &a_buf), &specs, 2); + try std.testing.expectEqual(@as(usize, 2), b.n); + try std.testing.expectEqualStrings("rocep1s0f1,roceP2p1s0f1", names(specs[0..b.n], &buf)); + // one rail asked for: the lower /24 on both sides + try std.testing.expectEqualStrings("rocep1s0f1", names(specs[0..choose(&a_ports, b_addrs, &specs, 1).n], &buf)); +} + +test "two rails on one /24 count once per peer port; no shared /24 keeps the list order" { + var specs: [rv.max_hcas]rv.HcaSpec = undefined; + var buf: [128]u8 = undefined; + // the neighbour cabled to port f0 lists one address on that /24: one rail each, equal counts + const one = [_][4]u8{.{ 203, 0, 113, 2 }}; + const p = choose(&a_ports, &one, &specs, 2); + try std.testing.expectEqual(@as(usize, 1), p.n); + try std.testing.expectEqualStrings("rocep1s0f0", names(specs[0..p.n], &buf)); + const two = [_][4]u8{ .{ 203, 0, 113, 2 }, .{ 203, 0, 113, 66 } }; + try std.testing.expectEqualStrings("rocep1s0f0,roceP2p1s0f0", names(specs[0..choose(&a_ports, &two, &specs, 2).n], &buf)); + // inactive and non-Ethernet ports never count; a peer on other subnets falls back to the first usable ports + var mixed = b_ports; + mixed[1].ethernet = false; + const other = [_][4]u8{.{ 198, 51, 100, 9 }}; + const q = choose(&mixed, &.{.{ 203, 0, 113, 9 }}, &specs, 2); + try std.testing.expect(!q.shared); + try std.testing.expectEqualStrings("roceP2p1s0f1", names(specs[0..q.n], &buf)); + try std.testing.expectEqualStrings("roceP2p1s0f1", names(specs[0..choose(&mixed, &other, &specs, 2).n], &buf)); + // a malformed blob is no ports; the detected GID index travels with the pick + var out: [max_ports][4]u8 = undefined; + try std.testing.expectEqual(@as(usize, 0), decode(&.{ 2, 1, 2, 3, 4 }, &out)); + try std.testing.expectEqual(@as(?u8, 3), specs[0].gid); +} + +test "hca's line: the picked names for NCCL; no device or no shared /24 is refused, an explicit list is taken" { + var specs: [rv.max_hcas]rv.HcaSpec = undefined; + var line: [list_bytes + 16]u8 = undefined; + var peer_buf: [max_ports][4]u8 = undefined; + const shared = choose(&a_ports, listed(&b_ports, &peer_buf), &specs, 2); + try std.testing.expectEqualStrings("NCCL_IB_HCA=rocep1s0f1,roceP2p1s0f1\n", try hcaLine(shared, &specs, &line)); + var buf: [list_bytes]u8 = undefined; + try std.testing.expectEqualStrings("rocep1s0f1:3,roceP2p1s0f1:3", try joinNames(specs[0..shared.n], true, &buf)); + // the old fallback (first active ports, no /24 in common) is what NCCL failed on: refused, never printed + try std.testing.expectError(error.NoSharedSubnet, hcaLine(choose(&a_ports, &.{.{ 10, 9, 9, 9 }}, &specs, 2), &specs, &line)); + try std.testing.expectError(error.NoRoceDevice, hcaLine(choose(&.{}, &.{}, &specs, 2), &specs, &line)); + const n = try rv.parseSpecs("rocep1s0f1:3, roceP2p1s0f1", &specs); + try std.testing.expectEqualStrings("NCCL_IB_HCA=rocep1s0f1,roceP2p1s0f1\n", try hcaLine(.{ .n = n, .shared = false, .explicit = true }, &specs, &line)); +} + +/// One RDMA device for `rv.scanPorts` through a stub verbs table: port 1 up or down, Ethernet, mlx5's GIDs 0-3. +const FakeDev = struct { name: [:0]const u8, up: bool, ipv4: [4]u8 = @splat(0) }; + +const ScanFixture = struct { + const vabi = @import("verbs").abi; + var host: []const FakeDev = &.{}; + var list: [max_ports:null]?*vabi.Device = undefined; + + fn handle(i: usize) usize { + return 0x10000 * (i + 1); + } + fn index(p: *const anyopaque) usize { + return @intFromPtr(p) / 0x10000 - 1; + } + fn getList(num: *c_int) callconv(.c) ?[*:null]?*vabi.Device { + for (host, 0..) |_, i| list[i] = @ptrFromInt(handle(i)); + list[host.len] = null; + num.* = @intCast(host.len); + return &list; + } + fn freeList(_: [*:null]?*vabi.Device) callconv(.c) void {} + fn devName(d: *vabi.Device) callconv(.c) ?[*:0]const u8 { + return host[index(d)].name.ptr; + } + fn openDev(d: *vabi.Device) callconv(.c) ?*vabi.Context { + return @ptrFromInt(@intFromPtr(d)); + } + fn closeDev(_: *vabi.Context) callconv(.c) c_int { + return 0; + } + fn queryPort(c: *vabi.Context, _: u8, a: *vabi.PortAttr) callconv(.c) c_int { + a.* = std.mem.zeroes(vabi.PortAttr); + a.state = if (host[index(c)].up) vabi.port_active else 1; + a.link_layer = 2; // Ethernet + a.gid_tbl_len = 4; + return 0; + } + fn queryGid(c: *vabi.Context, _: u8, i: c_int, g: *vabi.Gid) callconv(.c) c_int { + g.* = .{ .raw = @splat(0) }; + if (i < 2) { + g.raw[0] = 0xfe; + g.raw[1] = 0x80; + g.raw[15] = @intCast(index(c) + 1); + } else { + g.raw[10] = 0xff; + g.raw[11] = 0xff; + g.raw[12..16].* = host[index(c)].ipv4; + } + return 0; + } + fn gidEx(_: *vabi.Context, _: u32, i: u32, e: *vabi.GidEntry, _: u32, _: usize) callconv(.c) c_int { + e.* = std.mem.zeroes(vabi.GidEntry); + e.gid_type = if (i % 2 == 1) 2 else 1; // odd indices RoCE v2 + return 0; + } + + /// `rv.scanPorts` over `devs` (names that no real /sys/class/infiniband has, so sysfs never answers for them). + fn scan(devs: []const FakeDev, out: *[max_ports]rv.Port) []const rv.Port { + host = devs; + var v: Verbs = undefined; + v.api.ibv_get_device_list = getList; + v.api.ibv_free_device_list = freeList; + v.api.ibv_get_device_name = devName; + v.api.ibv_open_device = openDev; + v.api.ibv_close_device = closeDev; + v.api.ibv_query_port = queryPort; + v.api.ibv_query_gid = queryGid; + v.gid_ex = gidEx; + return out[0..rv.scanPorts(&v, out)]; + } +}; + +/// #636's reported two GB10 hosts, modelled (no capture): device order and picks reported; addresses, states assumed. +const hosts = [2][4]FakeDev{ .{ + .{ .name = "fx_rocep1s0f0", .up = true, .ipv4 = .{ 192, 0, 2, 1 } }, + .{ .name = "fx_rocep1s0f1", .up = false }, + .{ .name = "fx_roceP2p1s0f0", .up = true, .ipv4 = .{ 198, 51, 100, 1 } }, + .{ .name = "fx_roceP2p1s0f1", .up = false }, +}, .{ + .{ .name = "fx_rocep1s0f0", .up = false }, + .{ .name = "fx_rocep1s0f1", .up = true, .ipv4 = .{ 192, 0, 2, 2 } }, + .{ .name = "fx_roceP2p1s0f0", .up = false }, + .{ .name = "fx_roceP2p1s0f1", .up = true, .ipv4 = .{ 198, 51, 100, 2 } }, +} }; + +test "the port scan reads a host's devices through verbs: state, Ethernet, the IPv4 RoCE v2 GID index" { + var out: [max_ports]rv.Port = undefined; + const a = ScanFixture.scan(&hosts[0], &out); + try std.testing.expectEqual(@as(usize, 4), a.len); + try std.testing.expectEqualStrings("fx_rocep1s0f0", a[0].spec.name[0..a[0].spec.name_len]); + try std.testing.expect(a[0].active and a[0].ethernet and !a[1].active); + try std.testing.expectEqual(@as(?[4]u8, .{ 192, 0, 2, 1 }), a[0].ipv4); + try std.testing.expectEqual(@as(?u8, 3), a[0].spec.gid); // the v2 IPv4 GID, not v1's index 2 + try std.testing.expectEqual(@as(?[4]u8, null), a[1].ipv4); +} + +/// One rank's `selectFrom` over a real bootstrap (rank 0 on `l`), its hca line or error into `out`. +fn hcaRank(cfg: Config, l: ?sock.Fd, ports: []const rv.Port, open: ?OpenFn, out: *[list_bytes + 16]u8, line: *HcaError![]const u8, err: *?anyerror) void { + var b = (if (l) |fd| bootstrap.Bootstrap.acceptOn(fd, 2, 5 * std.time.ns_per_s) else bootstrap.Bootstrap.open(cfg)) catch |e| { + err.* = e; + return; + }; + defer b.close(); + var specs: [rv.max_hcas]rv.HcaSpec = undefined; + const pick = (if (open) |f| selectWithOpen(cfg, &b, &specs, f) else selectFrom(cfg, &b, ports, &specs)) catch |e| { + err.* = e; + return; + }; + line.* = hcaLine(pick, &specs, out); +} + +fn hcaPair(cfgs: [2]Config, ports: [2][]const rv.Port, lines: *[2]HcaError![]const u8, bufs: *[2][list_bytes + 16]u8) !void { + return hcaPairWithOpen(cfgs, ports, .{ null, null }, lines, bufs); +} + +fn hcaPairWithOpen(cfgs: [2]Config, ports: [2][]const rv.Port, open: [2]?OpenFn, lines: *[2]HcaError![]const u8, bufs: *[2][list_bytes + 16]u8) !void { + const l = try bootstrap.testListener(); + defer sock.close(l.fd); + var errs: [2]?anyerror = @splat(null); + var c1 = cfgs[1]; + c1.link_port = l.port; + c1.connect_timeout_ns = 5 * std.time.ns_per_s; + const t = try std.Thread.spawn(.{}, hcaRank, .{ c1, null, ports[1], open[1], &bufs[1], &lines[1], &errs[1] }); + hcaRank(cfgs[0], l.fd, ports[0], open[0], &bufs[0], &lines[0], &errs[0]); + t.join(); + for (errs) |e| if (e) |x| return x; +} + +fn missingLibrary(_: ?[:0]const u8) @import("verbs").Error!Verbs { + return error.LibraryUnavailable; +} + +fn missingSymbol(_: ?[:0]const u8) @import("verbs").Error!Verbs { + return error.MissingSymbol; +} + +test "hca keeps a verbs load failure while completing the peer's port exchange" { + var lines: [2]HcaError![]const u8 = undefined; + var bufs: [2][list_bytes + 16]u8 = undefined; + var scanned: [max_ports]rv.Port = undefined; + const peer = ScanFixture.scan(&hosts[1], &scanned); + const cfgs: [2]Config = .{ .{ .rank = 0 }, .{ .rank = 1 } }; + try hcaPairWithOpen(cfgs, .{ &.{}, peer }, .{ missingLibrary, null }, &lines, &bufs); + try std.testing.expectError(error.LibraryUnavailable, lines[0]); + try std.testing.expectError(error.NoSharedSubnet, lines[1]); +} + +test "hca distinguishes missing verbs symbols from an empty successful scan" { + const cfg: Config = .{ .world = 1 }; + var b = try bootstrap.Bootstrap.open(cfg); + defer b.close(); + var specs: [rv.max_hcas]rv.HcaSpec = undefined; + var buf: [list_bytes + 16]u8 = undefined; + const pick = try selectWithOpen(cfg, &b, &specs, missingSymbol); + try std.testing.expectError(error.MissingSymbol, hcaLine(pick, &specs, &buf)); + try std.testing.expectError(error.NoRoceDevice, hcaLine(try selectFrom(cfg, &b, &.{}, &specs), &specs, &buf)); +} + +test "hca keeps an explicit list when the verbs scan cannot open its library" { + const cfg: Config = .{ .world = 1, .roce_hca = "fx_rocep1s0f0:3" }; + var b = try bootstrap.Bootstrap.open(cfg); + defer b.close(); + var specs: [rv.max_hcas]rv.HcaSpec = undefined; + var buf: [list_bytes + 16]u8 = undefined; + const pick = try selectWithOpen(cfg, &b, &specs, missingLibrary); + try std.testing.expectEqualStrings("NCCL_IB_HCA=fx_rocep1s0f0\n", try hcaLine(pick, &specs, &buf)); +} + +test "hca over a real bootstrap: both ranks print their peer-linked ports; an explicit rank keeps its list; no ports fail both" { + var lines: [2]HcaError![]const u8 = undefined; + var bufs: [2][list_bytes + 16]u8 = undefined; + var scanned: [2][max_ports]rv.Port = undefined; + const ports = [2][]const rv.Port{ ScanFixture.scan(&hosts[0], &scanned[0]), ScanFixture.scan(&hosts[1], &scanned[1]) }; + const auto: [2]Config = .{ .{ .rank = 0 }, .{ .rank = 1 } }; + try hcaPair(auto, ports, &lines, &bufs); + try std.testing.expectEqualStrings("NCCL_IB_HCA=fx_rocep1s0f0,fx_roceP2p1s0f0\n", try lines[0]); + try std.testing.expectEqualStrings("NCCL_IB_HCA=fx_rocep1s0f1,fx_roceP2p1s0f1\n", try lines[1]); + // TF_TP_ROCE_HCA on one rank: its list as given (no GID for NCCL); it still sends its ports to the other + const one: [2]Config = .{ .{ .rank = 0, .roce_hca = "fx_rocep1s0f0:3" }, .{ .rank = 1 } }; + try hcaPair(one, ports, &lines, &bufs); + try std.testing.expectEqualStrings("NCCL_IB_HCA=fx_rocep1s0f0\n", try lines[0]); + try std.testing.expectEqualStrings("NCCL_IB_HCA=fx_rocep1s0f1,fx_roceP2p1s0f1\n", try lines[1]); + // a host without usable ports: it has none, its peer shares no /24 with it, and both refuse + try hcaPair(auto, .{ &.{}, ports[1] }, &lines, &bufs); + try std.testing.expectError(error.NoRoceDevice, lines[0]); + try std.testing.expectError(error.NoSharedSubnet, lines[1]); +} diff --git a/zig/src/fabric/tp/roce_post_test.zig b/zig/src/fabric/tp/roce_post_test.zig new file mode 100644 index 0000000000..a284293000 --- /dev/null +++ b/zig/src/fabric/tp/roce_post_test.zig @@ -0,0 +1,120 @@ +//! VerbsWire.post on fake verbs: a post that would leave its slot or region is refused before any verb. +const std = @import("std"); +const fake = @import("roce_fake.zig"); +const Roce = @import("roce.zig").Roce; +const Config = @import("config.zig").Config; +const Fake = fake.Fake; + +test "a post past its slot is refused before any verb; a full slot posts inside the region" { + const small: u64 = 64 << 10; + const region = try std.testing.allocator.alignedAlloc(u8, .fromByteUnits(4096), Roce.regionBytes(small)); + defer std.testing.allocator.free(region); + Fake.region = region; + Fake.slot_bytes = small; + defer Fake.slot_bytes = fake.slot; + Fake.reset(1); + var hw: fake.Hw = .{}; + var w = hw.wire(region); + try std.testing.expectError(error.OutOfRegion, w.post(1, 1, small + 16)); + try std.testing.expectError(error.OutOfRegion, w.post(2, 1, 16)); + try std.testing.expectEqual(@as(u64, 0), Fake.posts); + try w.post(1, 1, small); + try std.testing.expectEqual(@as(u64, 2), Fake.posts); + try std.testing.expectEqual(@as(u64, 0), Fake.outside); +} + +test "no RDMA write past 1 MiB, whatever the slot: a 9.57 MB and a 1 MiB + 16 post refused in a 10 MiB slot" { + const big: u64 = 10 << 20; + const region = try std.testing.allocator.alignedAlloc(u8, .fromByteUnits(4096), Roce.regionBytes(big)); + defer std.testing.allocator.free(region); + Fake.region = region; + Fake.slot_bytes = big; + defer Fake.slot_bytes = fake.slot; + Fake.reset(1); + var hw: fake.Hw = .{}; + var w = hw.wire(region); + try std.testing.expectError(error.WriteTooLarge, w.post(0, 1, 9_568_256)); + try std.testing.expectError(error.WriteTooLarge, w.post(1, 1, (1 << 20) + 16)); + try std.testing.expectEqual(@as(u64, 0), Fake.posts); + try w.post(0, 1, 1 << 20); + try std.testing.expectEqual(@as(u64, 1 << 20), Fake.max_write * 2); +} + +const Env = struct { + var pairs: []const [2][]const u8 = &.{}; + + fn get(name: []const u8) ?[]const u8 { + for (pairs) |p| if (std.mem.eql(u8, p[0], name)) return p[1]; + return null; + } +}; + +test "boot refuses a RoCE slot past 1 MiB (TF_TP_MAILBOX_MAX_KB), never clamping it" { + for ([_][]const u8{ "1028", "2048", "16384" }) |kb| { + Env.pairs = &.{ .{ "TF_COMM_BACKEND", "roce" }, .{ "TF_TP_MAILBOX_MAX_KB", kb } }; + try std.testing.expectError(error.BadSetting, Config.parse(&Env.get)); + } + Env.pairs = &.{ .{ "TF_COMM_BACKEND", "roce" }, .{ "TF_TP_MAILBOX_MAX_KB", "1024" } }; + try std.testing.expectEqual(@as(usize, 1 << 20), (try Config.parse(&Env.get)).mailbox_max_bytes); + // the cap is RoCE's: the same-host mailbox keeps its larger slots + Env.pairs = &.{ .{ "TF_COMM_BACKEND", "mailbox" }, .{ "TF_TP_MAILBOX_MAX_KB", "4096" } }; + try std.testing.expectEqual(@as(usize, 4 << 20), (try Config.parse(&Env.get)).mailbox_max_bytes); +} + +test "chunked strides posted back to back: at most 2 MiB in flight over both HCAs and no write past 1 MiB" { + const region = try std.testing.allocator.alignedAlloc(u8, .fromByteUnits(4096), Roce.regionBytes(fake.slot)); + defer std.testing.allocator.free(region); + Fake.region = region; + Fake.slot_bytes = fake.slot; + // gathers of R x 512 x 584 B strides and 584 B rows up to a 128 MiB union, each in <= 1 MiB ops + var strides: [16]u64 = undefined; + var n: usize = 0; + for ([_]u64{ 1, 2, 3, 4, 6, 16, 32, 64 }) |r| { + strides[n] = r * 512 * 584; + n += 1; + } + for ([_]u64{ 1795, 1796, (128 << 20) / (3 * 584) }) |m| { + strides[n] = m * 584; + n += 1; + } + for ([_]u32{ 1, 3, 17 }) |lag| { + Fake.reset(lag); + var hw: fake.Hw = .{}; + var w = hw.wire(region); + var seq: u32 = 0; + var want: u64 = 0; + for (strides[0..n]) |stride| { + var off: u64 = 0; + while (off < stride) : (off += fake.slot) { + seq += 1; + const padded: u32 = @intCast(std.mem.alignForward(u64, @min(fake.slot, stride - off), 16)); + try w.post(seq & 1, seq, padded); + want += padded; + } + } + while (Fake.inflight > 0) try w.reap(); + try std.testing.expectEqual(want, Fake.bytes); + try std.testing.expectEqual(@as(u64, 0), Fake.outside); + // the cap's numbers, not roce_verbs.zig's constants: 1 MiB a write, 2 MiB in flight + try std.testing.expect(Fake.max_write <= 1 << 20); + try std.testing.expect(Fake.max_inflight <= 2 << 20); + // the bound is reached, not avoided: posts outran completions + try std.testing.expect(Fake.max_inflight > 1 << 20); + } +} + +test "Roce.abort (the fail-fast hook) moves both QPs to error: no send goes out and no peer write lands" { + const region = try std.testing.allocator.alignedAlloc(u8, .fromByteUnits(4096), Roce.regionBytes(fake.slot)); + defer std.testing.allocator.free(region); + @memset(region, 0); + Fake.region = region; + Fake.slot_bytes = fake.slot; + Fake.reset(1); + fake.Halts.reset(); + var hw: fake.Hw = .{}; + var w = hw.wire(region); + fake.Halts.install(&hw, &w); + var r: Roce = .{ .me = 0, .slot_bytes = fake.slot, .timeout_ns = 1, .region = region, .mapping = region, .wire = .{ .verbs = w }, .cpu = null }; + r.abort(); + try std.testing.expect(fake.Halts.halted[0] and fake.Halts.halted[1]); +} diff --git a/zig/src/fabric/tp/roce_shm.zig b/zig/src/fabric/tp/roce_shm.zig new file mode 100644 index 0000000000..3924be3e96 --- /dev/null +++ b/zig/src/fabric/tp/roce_shm.zig @@ -0,0 +1,63 @@ +//! The RoCE protocol's shm wire: both ranks' regions in one shared mapping, posts through the guard's stripes. +const std = @import("std"); +const Bootstrap = @import("bootstrap.zig").Bootstrap; +const sock = @import("sock.zig"); +const guard = @import("roce_guard.zig"); +const posix = std.posix; +const page = std.heap.page_size_min; + +/// A test's view of each post: slot, padded bytes and the stripes the wire is about to move. +pub const PostTap = struct { + ctx: *anyopaque, + seen: *const fn (ctx: *anyopaque, l: guard.Layout, slot: u32, padded: u32, s: []const guard.Stripe) void, +}; + +/// The peer's region is mapped here; a post copies each stripe, then stores its flag with release order. +pub const ShmWire = struct { + peer: []u8, + layout: guard.Layout, + local: []u8, + /// Stripes a post makes, as HCAs would: the verbs wire's split, flags and bounds on one host. + stripes: u32 = 1, + tap: ?*PostTap = null, + + pub fn post(w: *ShmWire, slot: u32, seq: u32, padded: u32) guard.Error!void { + var plan: [guard.max_hcas]guard.Stripe = undefined; + try guard.stripes(w.layout, w.stripes, slot, padded, &plan); + if (w.tap) |t| t.seen(t.ctx, w.layout, slot, padded, plan[0..w.stripes]); + for (plan[0..w.stripes]) |s| { + @memcpy(w.peer[s.remote..][0..s.len], w.local[s.local..][0..s.len]); + const flag: *std.atomic.Value(u32) = @ptrCast(@alignCast(w.peer[s.flag..].ptr)); + flag.store(seq, .release); + } + } +}; + +/// One shared mapping holding both ranks' regions (rank 0 creates and sizes the file before anyone learns its name). +pub fn sharedPair(b: *Bootstrap, len: usize) ![]align(page) u8 { + var name: [64]u8 = undefined; + var name_len: u8 = 0; + var fd: std.c.fd_t = -1; + defer if (fd >= 0) { + _ = std.c.close(fd); + }; + if (b.rank == 0) blk: { + const s = std.fmt.bufPrint(&name, "/dev/shm/tf-tp-roce-{d}-{x}", .{ std.c.getpid(), sock.nowNs() }) catch break :blk; + name[s.len] = 0; + fd = std.c.open(name[0..s.len :0], .{ .ACCMODE = .RDWR, .CREAT = true, .EXCL = true, .CLOEXEC = true }, @as(std.c.mode_t, 0o600)); + if (fd >= 0 and std.c.ftruncate(fd, @intCast(2 * len)) == 0) name_len = @intCast(s.len); + } + // an empty name tells the other rank that rank 0 failed (it must not wait for a name that never comes) + try b.broadcast(std.mem.asBytes(&name_len), std.mem.asBytes(&name_len)); + if (name_len == 0) return error.SharedMemoryUnavailable; + try b.broadcast(name[0..name_len], name[0..name_len]); + name[name_len] = 0; + if (b.rank != 0) { + fd = std.c.open(name[0..name_len :0], .{ .ACCMODE = .RDWR, .CLOEXEC = true }, @as(std.c.mode_t, 0)); + if (fd < 0) return error.SharedMemoryUnavailable; + } + const m = try posix.mmap(null, 2 * len, .{ .READ = true, .WRITE = true }, .{ .TYPE = .SHARED }, fd, 0); + try b.agree("mapped", "RoCE shm wire"); + if (b.rank == 0) _ = std.c.unlink(name[0..name_len :0]); + return m; +} diff --git a/zig/src/fabric/tp/roce_test.zig b/zig/src/fabric/tp/roce_test.zig new file mode 100644 index 0000000000..ec6bc4d65e --- /dev/null +++ b/zig/src/fabric/tp/roce_test.zig @@ -0,0 +1,173 @@ +//! The RoCE protocol end to end on host threads: two ranks, TCP bootstrap, proxy, shm or verbs (TF_TP_ROCE_TEST_HCA). +const std = @import("std"); +const roce = @import("roce.zig"); +const rv = @import("roce_verbs.zig"); +const shm = @import("roce_shm.zig"); +const guard = @import("roce_guard.zig"); +const bootstrap = @import("bootstrap.zig"); +const sock = @import("sock.zig"); +const reference = @import("reference.zig"); +const Bootstrap = bootstrap.Bootstrap; + +const Run = struct { + port: u16, + wire: roce.WireKind, + devs: [2][]const u8 = .{ "", "" }, + ops: u32, + timeout_ns: u64 = 10 * std.time.ns_per_s, + /// Rank 1 stops after this many ops (rank 0 must then fail its next wait, by timeout or abort). + stop_after: ?u32 = null, + abort_after_ms: ?u64 = null, + errs: [2]?anyerror = .{ null, null }, + done: [2]u32 = .{ 0, 0 }, + max_bytes: usize = 16 * 1024, + /// Variable lengths (`Roce.allGatherV`'s contract): each op's lens[r] bytes of rank r's n, 0 and n included. + var_lens: bool = false, + /// Strides up to four slots through `hostGatherVChunked` (variable lengths, slot-sized ops). + chunked: bool = false, + /// The shm wire's stripes (2: the verbs wire's split over two HCAs). + stripes: u32 = 1, + /// Each rank's shm posts as the tap saw them: how many, and how many left their slot or the region. + taps: [2]Tap = .{ .{}, .{} }, +}; + +/// Counts an shm wire's posts and any stripe outside its slot, its flag words or the region. +const Tap = struct { + posts: u32 = 0, + outside: u32 = 0, + tap: shm.PostTap = undefined, + + fn seen(ctx: *anyopaque, l: guard.Layout, slot: u32, padded: u32, s: []const guard.Stripe) void { + const t: *Tap = @ptrCast(@alignCast(ctx)); + t.posts += 1; + var sum: u64 = 0; + for (s) |x| { + sum += x.len; + const in_slot = x.local >= l.send + slot * l.slot and x.local + x.len <= l.send + (slot + 1) * l.slot; + const in_peer = x.remote >= l.recv + slot * l.slot and x.remote + x.len <= l.recv + (slot + 1) * l.slot; + if (!in_slot or !in_peer or x.remote + x.len > l.len or x.flag < l.flags or x.flag + 4 > l.send) t.outside += 1; + } + if (sum != padded) t.outside += 1; + } +}; + +fn rank(run: *Run, me: u32, lfd: sock.Fd) void { + run.errs[me] = body(run, me, lfd); +} + +fn body(run: *Run, me: u32, lfd: sock.Fd) ?anyerror { + var b = (if (me == 0) Bootstrap.acceptOn(lfd, 2, 5 * std.time.ns_per_s) else Bootstrap.open(.{ .world = 2, .rank = 1, .link_port = run.port, .connect_timeout_ns = 5 * std.time.ns_per_s })) catch |e| return e; + defer b.close(); + var spec: [1]rv.HcaSpec = .{.{}}; + if (run.wire == .verbs) _ = rv.parseSpecs(run.devs[me], @ptrCast(&spec)) catch |e| return e; + const mine = &run.taps[me]; + mine.tap = .{ .ctx = mine, .seen = Tap.seen }; + const r = roce.Roce.init(null, &b, "", .{ .wire = run.wire, .hcas = if (run.wire == .verbs) &spec else &.{}, .max_bytes = run.max_bytes, .timeout_ns = run.timeout_ns, .shm_stripes = run.stripes, .tap = &mine.tap }) catch |e| return e; + defer r.deinit(); + var abort_thread: ?std.Thread = null; + defer if (abort_thread) |t| t.join(); + var in: [64 * 1024]u8 = undefined; + var peer: [64 * 1024]u8 = undefined; + var out: [128 * 1024]u8 = undefined; + for (0..run.ops) |k| { + if (run.stop_after) |s| if (k == s) { + // rank 1 skips op s; rank 0 waits in it until its timeout, or until the abort armed now + if (me == 1) { + sock.sleepNs(300 * std.time.ns_per_ms); + return null; + } + if (run.abort_after_ms) |ms| abort_thread = std.Thread.spawn(.{}, abortLater, .{ r, ms }) catch null; + }; + // sizes from 1 byte to the slot, odd ones included, both slots and wrap-arounds + const n = 1 + (k * 7919) % (if (run.chunked) 4 * run.max_bytes else run.max_bytes); + reference.fill(in[0..n], .u8, me, k); + reference.fill(peer[0..n], .u8, 1 - me, k); + const lens: [2]u32 = if (k % 5 == 0) .{ 0, @intCast(n) } else .{ @intCast((k * 31) % (n + 1)), @intCast((k * 17 + 5) % (n + 1)) }; + const l: [2]usize = if (run.var_lens or run.chunked) .{ lens[0], lens[1] } else .{ n, n }; + if (run.chunked) { + r.hostGatherVChunked(in[0..n], out[0 .. 2 * n], lens) catch |e| return e; + } else r.hostGatherV(in[0..n], out[0 .. 2 * n], if (run.var_lens) lens else null) catch |e| return e; + const r0 = if (me == 0) in[0..n] else peer[0..n]; + const r1 = if (me == 1) in[0..n] else peer[0..n]; + if (!std.mem.eql(u8, out[0..l[0]], r0[0..l[0]]) or !std.mem.eql(u8, out[n..][0..l[1]], r1[0..l[1]])) return error.WrongBytes; + run.done[me] += 1; + } + r.check() catch |e| return e; + // both ranks finish before either tears down (the peer may still read a flag this rank's wire writes) + b.agree("done", "test end") catch |e| return e; + return null; +} + +fn abortLater(r: *roce.Roce, ms: u64) void { + sock.sleepNs(ms * std.time.ns_per_ms); + r.abort(); +} + +fn twoRanks(run: *Run) !void { + const l = try bootstrap.testListener(); + defer sock.close(l.fd); + run.port = l.port; + const t1 = try std.Thread.spawn(.{}, rank, .{ run, 1, l.fd }); + rank(run, 0, l.fd); + t1.join(); +} + +test "the RoCE protocol over the shm wire: 2,000 gathers of every size, bit for bit" { + var run: Run = .{ .port = 0, .wire = .shm, .ops = 2000 }; + try twoRanks(&run); + for (run.errs) |e| if (e) |err| return err; + try std.testing.expectEqual([2]u32{ 2000, 2000 }, run.done); +} + +test "the RoCE protocol with variable lengths over the shm wire: each rank posts its own padded count, empty shards included" { + var run: Run = .{ .port = 0, .wire = .shm, .ops = 2000, .var_lens = true }; + try twoRanks(&run); + for (run.errs) |e| if (e) |err| return err; + try std.testing.expectEqual([2]u32{ 2000, 2000 }, run.done); +} + +test "the shm wire striped as over two HCAs: every post inside its slot and the region, bytes and lengths exact" { + var run: Run = .{ .port = 0, .wire = .shm, .ops = 2000, .var_lens = true, .stripes = 2 }; + try twoRanks(&run); + for (run.errs) |e| if (e) |err| return err; + try std.testing.expectEqual([2]u32{ 2000, 2000 }, run.done); + for (run.taps) |t| { + try std.testing.expectEqual(@as(u32, 2000), t.posts); + try std.testing.expectEqual(@as(u32, 0), t.outside); + } +} + +test "chunked gathers over the shm wire: strides up to four slots in slot-sized ops, tails empty on both ranks skipped" { + var run: Run = .{ .port = 0, .wire = .shm, .ops = 2000, .chunked = true, .stripes = 2 }; + try twoRanks(&run); + for (run.errs) |e| if (e) |err| return err; + try std.testing.expectEqual([2]u32{ 2000, 2000 }, run.done); + for (run.taps) |t| { + try std.testing.expect(t.posts > 2000); + try std.testing.expectEqual(@as(u32, 0), t.outside); + } +} + +test "a peer that stops makes the waiting rank time out, and abort ends a wait at once" { + var run: Run = .{ .port = 0, .wire = .shm, .ops = 50, .stop_after = 20, .timeout_ns = 100 * std.time.ns_per_ms }; + try twoRanks(&run); + try std.testing.expectEqual(@as(?anyerror, error.Timeout), run.errs[0]); + try std.testing.expectEqual(@as(u32, 20), run.done[0]); + var run2: Run = .{ .port = 0, .wire = .shm, .ops = 50, .stop_after = 20, .abort_after_ms = 50 }; + const kept = roce.Roce.kept.load(.monotonic); + try twoRanks(&run2); + // the aborted rank keeps its region until exit, the other frees it + try std.testing.expectEqual(kept + 1, roce.Roce.kept.load(.monotonic)); + try std.testing.expectEqual(@as(?anyerror, error.Aborted), run2.errs[0]); + try std.testing.expectEqual(@as(?anyerror, null), run2.errs[1]); +} + +test "the RoCE protocol over verbs (TF_TP_ROCE_TEST_HCA=dev0,dev1: soft-RoCE or real HCAs)" { + const raw = std.c.getenv("TF_TP_ROCE_TEST_HCA") orelse return error.SkipZigTest; + var it = std.mem.splitScalar(u8, std.mem.span(raw), ','); + var run: Run = .{ .port = 0, .wire = .verbs, .ops = 2000 }; + run.devs = .{ it.next() orelse return error.SkipZigTest, it.next() orelse return error.SkipZigTest }; + try twoRanks(&run); + for (run.errs) |e| if (e) |err| return err; + try std.testing.expectEqual([2]u32{ 2000, 2000 }, run.done); +} diff --git a/zig/src/fabric/tp/roce_verbs.zig b/zig/src/fabric/tp/roce_verbs.zig new file mode 100644 index 0000000000..6917d292be --- /dev/null +++ b/zig/src/fabric/tp/roce_verbs.zig @@ -0,0 +1,455 @@ +//! The RoCE wire over verbs: one RC queue pair a HCA to the peer, a send slot striped over the HCAs with seq flags. +const std = @import("std"); +const log = @import("log.zig"); +const verbs = @import("verbs"); +const abi = verbs.abi; +const Verbs = verbs.Verbs; +const guard = @import("roce_guard.zig"); +const sock = @import("sock.zig"); + +pub const max_hcas = guard.max_hcas; +/// Work requests a queue pair holds; at most a quarter outstanding (two requests a stripe). +const send_depth = 256; +const link_ethernet = 2; +const record_magic: u32 = 0x5446_5244 + 2; // "TFRD", bumped for each change of this record's layout + +pub const Error = error{ NoHca, VerbsFailed, PostFailed, Completion, BadRecord } || guard.Error; + +/// Bytes posted and not yet completed, every HCA together: a post waits for completions past it; never raised. +pub const max_inflight: u64 = 2 * @as(u64, guard.max_write); +/// A post that waits on completions this long fails the proxy (a peer or HCA that stopped answering). +const inflight_wait_ns: u64 = 10 * std.time.ns_per_s; +const ring_len = send_depth / 4; +/// The flag post's bytes: a u32 on the stack under lkey 0, legal only while the QP sends it inline. +pub const flag_bytes: u32 = @sizeOf(u32); +/// The longest `quiesce` waits for the HCAs to flush what was posted. +pub const quiesce_ns: u64 = 5 * std.time.ns_per_s; + +/// Before ibv_reg_mr: no fork shares or copies the region (DONTFORK), its pages stay 4 KiB (NOHUGEPAGE); else no start. +pub fn prepareRegion(region: []align(std.heap.page_size_min) u8) Error!void { + if (comptime @import("builtin").os.tag != .linux) return; // the verbs wire runs on Linux only + const advice = [_]u32{ std.os.linux.MADV.DONTFORK, std.os.linux.MADV.NOHUGEPAGE }; + for (advice) |a| std.posix.madvise(region.ptr, region.len, a) catch |e| { + log.warn("tp roce: madvise({d}) on the {d} B region failed: {t}; refusing to start", .{ a, region.len, e }); + return error.VerbsFailed; + }; +} + +/// Whether a QP's actual inline capacity takes the flag post. +pub fn inlineFits(cap: abi.QpCap) bool { + return cap.max_inline_data >= flag_bytes; +} + +/// One HCA as configured: its device name and RoCE v2 GID index (null: detected). +pub const HcaSpec = struct { name: [64]u8 = @splat(0), name_len: u8 = 0, gid: ?u8 = null }; + +/// What a rank tells its peer to connect (both ranks build it the same way; sent through the bootstrap). +pub const Record = extern struct { + magic: u32 = record_magic, + n_hca: u32 = 0, + region: u64 = 0, + /// The region's slot and registered length: a peer whose layout differs is refused at connect. + slot_bytes: u64 = 0, + region_len: u64 = 0, + rkey: [max_hcas]u32 = @splat(0), + lid: [max_hcas]u32 = @splat(0), + mtu: [max_hcas]u32 = @splat(0), + qpn: [max_hcas]u32 = @splat(0), + gid: [max_hcas][16]u8 = @splat(@splat(0)), +}; + +const Hca = struct { + ctx: *abi.Context, + pd: *abi.Pd, + mr: *abi.Mr, + cq: *abi.Cq, + qp: *abi.Qp, + gid_index: u8, + gid: [16]u8, + mtu: u32, + lid: u16, + outstanding: u32 = 0, + /// Each outstanding post's stripe bytes, oldest at `head` (completions arrive in order). + ring: [ring_len]u32 = undefined, + head: u32 = 0, + inflight: u64 = 0, + /// Posts made and their completions reaped (the proxy writes them, others read): equal = nothing in flight. + posts: std.atomic.Value(u64) = .init(0), + done: std.atomic.Value(u64) = .init(0), +}; + +/// Parses `name[:gid],name[:gid]` (TF_TP_ROCE_HCA). +pub fn parseSpecs(text: []const u8, out: *[max_hcas]HcaSpec) !usize { + var n: usize = 0; + var it = std.mem.tokenizeScalar(u8, text, ','); + while (it.next()) |item| { + if (n == max_hcas) return error.TooManyHcas; + var parts = std.mem.splitScalar(u8, std.mem.trim(u8, item, " "), ':'); + const name = parts.next().?; + if (name.len == 0 or name.len >= 64) return error.BadHcaName; + out[n] = .{ .name_len = @intCast(name.len) }; + @memcpy(out[n].name[0..name.len], name); + if (parts.next()) |g| out[n].gid = try std.fmt.parseInt(u8, g, 10); + n += 1; + } + return n; +} + +/// RoCE v2 by sysfs (the kernel's own word for each GID); null when sysfs cannot say. +fn sysfsRoceV2(name: []const u8, index: u32) ?bool { + var path: [256]u8 = undefined; + const p = std.fmt.bufPrintSentinel(&path, "/sys/class/infiniband/{s}/ports/1/gid_attrs/types/{d}", .{ name, index }, 0) catch return null; + const fd = std.c.open(p, .{ .ACCMODE = .RDONLY }); + if (fd < 0) return null; + defer _ = std.c.close(fd); + var buf: [32]u8 = undefined; + const n = std.c.read(fd, &buf, buf.len); + if (n <= 0) return false; + return std.mem.startsWith(u8, buf[0..@intCast(n)], "RoCE v2"); +} + +/// An IPv4-mapped GID (::ffff:a.b.c.d): the address the HCA's netdev carries. +pub fn ipv4Mapped(g: [16]u8) bool { + return std.mem.allEqual(u8, g[0..10], 0) and g[10] == 0xff and g[11] == 0xff; +} + +/// The lowest RoCE v2 GID index that is IPv4-mapped (GID indices move on reboot: never hard-coded). +fn detectGid(v: *const Verbs, ctx: *abi.Context, name: []const u8, table: c_int) ?u8 { + var i: u32 = 0; + while (i < @min(@as(u32, @intCast(@max(table, 0))), 255)) : (i += 1) { + const g = v.gid(ctx, 1, i) catch continue; + if (!ipv4Mapped(g.raw)) continue; + const v2 = sysfsRoceV2(name, i) orelse (if (v.gidType(ctx, 1, i)) |t| t == 2 else true); + if (v2) return @intCast(i); + } + return null; +} + +/// A device's port 1 as `roce_pick.choose` sees it: up, Ethernet, and its IPv4 RoCE v2 GID (index and address). +pub const Port = struct { spec: HcaSpec = .{}, active: bool = false, ethernet: bool = false, ipv4: ?[4]u8 = null }; + +/// Port 1 of every device that opens, in the device list's order, at most `out.len`. +pub fn scanPorts(v: *const Verbs, out: []Port) usize { + var count: c_int = 0; + const list = v.api.ibv_get_device_list(&count) orelse return 0; + defer v.api.ibv_free_device_list(list); + var n: usize = 0; + var i: usize = 0; + while (list[i]) |dev| : (i += 1) { + if (n == out.len) break; + const name = std.mem.span(v.api.ibv_get_device_name(dev) orelse continue); + const ctx = v.api.ibv_open_device(dev) orelse continue; + defer _ = v.api.ibv_close_device(ctx); + const port = v.port(ctx, 1) catch continue; + var p: Port = .{ .active = port.state == abi.port_active, .ethernet = port.link_layer == link_ethernet }; + p.spec.name_len = @intCast(@min(name.len, 63)); + @memcpy(p.spec.name[0..p.spec.name_len], name[0..p.spec.name_len]); + if (p.active and p.ethernet) if (detectGid(v, ctx, name, port.gid_tbl_len)) |gi| { + if (v.gid(ctx, 1, gi)) |g| { + p.spec.gid = gi; + p.ipv4 = g.raw[12..16].*; + } else |_| {} + }; + out[n] = p; + n += 1; + } + return n; +} + +pub const VerbsWire = struct { + v: Verbs, + hcas: [max_hcas]Hca = undefined, + n_hca: u32 = 0, + traffic_class: u8, + peer_region: u64 = 0, + peer_rkey: [max_hcas]u32 = @splat(0), + /// Offsets in both ranks' regions (the same layout): receive slots and flags. + recv_off: usize, + flag_off: usize, + send_off: usize, + slot_bytes: usize, + region: []u8, + posted: u64 = 0, + /// The most stripe bytes in flight at once (posted, not completed), for tests and reports. + peak: std.atomic.Value(u64) = .init(0), + completed: u64 = 0, + /// Completions `quiesce` reaped, and those of them flushed (not sent: the QP was in error). + reaped: u64 = 0, + flushed: u64 = 0, + + /// Opens the HCAs, registers `region` on each and creates the queue pairs (INIT); `record` then goes to the peer. + pub fn open(specs: []const HcaSpec, region: []u8, layout: struct { send: usize, recv: usize, flags: usize, slot: usize }, traffic_class: u8) !VerbsWire { + if (specs.len == 0 or specs.len > max_hcas) return error.NoHca; + var w: VerbsWire = .{ .v = try Verbs.open(null), .traffic_class = traffic_class, .recv_off = layout.recv, .flag_off = layout.flags, .send_off = layout.send, .slot_bytes = layout.slot, .region = region }; + errdefer w.close(); + for (specs) |s| { + try w.openHca(s); + } + return w; + } + + fn openHca(w: *VerbsWire, s: HcaSpec) !void { + const api = w.v.api; + const name = s.name[0..s.name_len]; + const ctx = try w.v.openDevice(name); + errdefer _ = api.ibv_close_device(ctx); + const port = try w.v.port(ctx, 1); + if (port.state != abi.port_active) { + log.warn("tp roce: {s} port 1 is not active", .{name}); + return error.NoHca; + } + const gi = s.gid orelse detectGid(&w.v, ctx, name, port.gid_tbl_len) orelse { + log.warn("tp roce: {s} has no IPv4 RoCE v2 GID", .{name}); + return error.NoHca; + }; + const gid = try w.v.gid(ctx, 1, gi); + const pd = api.ibv_alloc_pd(ctx) orelse return error.VerbsFailed; + errdefer _ = api.ibv_dealloc_pd(pd); + const mr = api.ibv_reg_mr(pd, w.region.ptr, w.region.len, abi.access_local_write | abi.access_remote_write) orelse { + log.warn("tp roce: ibv_reg_mr on {s} failed (memlock limit? ulimit -l unlimited)", .{name}); + return error.VerbsFailed; + }; + errdefer _ = api.ibv_dereg_mr(mr); + const cq = api.ibv_create_cq(ctx, send_depth, null, null, 0) orelse return error.VerbsFailed; + errdefer _ = api.ibv_destroy_cq(cq); + var init: abi.QpInitAttr = .{ .send_cq = cq, .recv_cq = cq, .cap = .{ .max_send_wr = send_depth, .max_recv_wr = 1, .max_send_sge = 1, .max_recv_sge = 1, .max_inline_data = 16 }, .qp_type = abi.qpt_rc }; + const qp = api.ibv_create_qp(pd, &init) orelse return error.VerbsFailed; + errdefer _ = api.ibv_destroy_qp(qp); + if (!inlineFits(init.cap)) { + log.warn("tp roce: {s}'s QP sends {d} B inline, the flag post needs {d}; refusing to start", .{ name, init.cap.max_inline_data, flag_bytes }); + return error.VerbsFailed; + } + var a: abi.QpAttr = .{ .qp_state = abi.qps_init, .port_num = 1, .qp_access_flags = abi.access_remote_write }; + if (api.ibv_modify_qp(qp, &a, abi.mask.state | abi.mask.pkey_index | abi.mask.port | abi.mask.access_flags) != 0) return error.VerbsFailed; + w.hcas[w.n_hca] = .{ .ctx = ctx, .pd = pd, .mr = mr, .cq = cq, .qp = qp, .gid_index = gi, .gid = gid.raw, .mtu = @intCast(port.active_mtu), .lid = port.lid }; + w.n_hca += 1; + log.info("[tensorfold] tp roce: {s} GID {d} ({d}.{d}.{d}.{d}), MTU enum {d}", .{ name, gi, gid.raw[12], gid.raw[13], gid.raw[14], gid.raw[15], port.active_mtu }); + } + + pub fn close(w: *VerbsWire) void { + const api = w.v.api; + for (w.hcas[0..w.n_hca]) |h| { + _ = api.ibv_destroy_qp(h.qp); + _ = api.ibv_destroy_cq(h.cq); + _ = api.ibv_dereg_mr(h.mr); + _ = api.ibv_dealloc_pd(h.pd); + _ = api.ibv_close_device(h.ctx); + } + w.n_hca = 0; + w.v.close(); + } + + /// The region and each HCA's keys, this side's and the peer's (a protection error is checked against them). + pub fn describe(w: *const VerbsWire, buf: []u8) []const u8 { + var out: std.Io.Writer = .fixed(buf); + out.print("region 0x{x} len {d}", .{ @intFromPtr(w.region.ptr), w.region.len }) catch return out.buffered(); + for (w.hcas[0..w.n_hca], 0..) |h, i| { + out.print("; HCA {d} lkey 0x{x} rkey 0x{x}, peer region 0x{x} rkey 0x{x}", .{ i, h.mr.lkey, h.mr.rkey, w.peer_region, w.peer_rkey[i] }) catch break; + } + return out.buffered(); + } + + /// The QPs to error: the HCAs stop sending and refuse the peer's writes, and what was posted completes as flushed. + pub fn halt(w: *VerbsWire) bool { + var ok = true; + for (w.hcas[0..w.n_hca]) |h| { + var a: abi.QpAttr = .{ .qp_state = abi.qps_err }; + if (w.v.api.ibv_modify_qp(h.qp, &a, abi.mask.state) != 0) ok = false; + } + return ok; + } + + /// `halt`, then every posted request's completion reaped, flushed or not: no request of the region left at an HCA. + pub fn quiesce(w: *VerbsWire, timeout_ns: u64) bool { + if (!w.halt()) return false; + const t0 = sock.nowNs(); + for (w.hcas[0..w.n_hca]) |*h| while (h.outstanding > 0) { + var wc: [32]abi.Wc = undefined; + const n = h.ctx.ops.poll_cq(h.cq, 32, &wc); + if (n < 0) return false; + for (wc[0..@intCast(n)]) |c| w.flushed += @intFromBool(c.status != abi.wc_success); + w.reaped += @intCast(n); + if (n > 0) w.popped(h, @intCast(n)) else if (sock.nowNs() - t0 > timeout_ns) return false; + }; + return true; + } + + /// Every HCA's posts completed: no request of the region at an HCA (readable beside the proxy). + pub fn idle(w: *const VerbsWire) bool { + for (w.hcas[0..w.n_hca]) |*h| if (h.posts.load(.acquire) != h.done.load(.acquire)) return false; + return true; + } + + /// Each QP's state as its device reports it (ibv_query_qp), then what `quiesce` reaped and what is still posted. + pub fn report(w: *const VerbsWire, buf: []u8) []const u8 { + const names = [_][]const u8{ "RESET", "INIT", "RTR", "RTS", "SQD", "SQE", "ERR" }; + var out: std.Io.Writer = .fixed(buf); + for (w.hcas[0..w.n_hca], 0..) |*h, i| { + var a: abi.QpAttr = .{}; + var init: abi.QpInitAttr = undefined; + const q = w.v.query_qp orelse { + out.print("HCA {d} QP state unknown (no ibv_query_qp); ", .{i}) catch break; + continue; + }; + if (q(h.qp, &a, abi.mask.state, &init) != 0) { + out.print("HCA {d} QP query failed; ", .{i}) catch break; + } else if (a.qp_state >= 0 and a.qp_state < names.len) { + out.print("HCA {d} QP {s}; ", .{ i, names[@intCast(a.qp_state)] }) catch break; + } else out.print("HCA {d} QP state {d}; ", .{ i, a.qp_state }) catch break; + out.print("posted {d}, completed {d}; ", .{ h.posts.load(.acquire), h.done.load(.acquire) }) catch break; + } + var held: u64 = 0; + for (w.hcas[0..w.n_hca]) |h| held += h.outstanding; + out.print("reaped {d} (flushed {d}), still posted {d}", .{ w.reaped, w.flushed, held }) catch {}; + return out.buffered(); + } + + pub fn record(w: *const VerbsWire) Record { + var r: Record = .{ .n_hca = w.n_hca, .region = @intFromPtr(w.region.ptr), .slot_bytes = w.slot_bytes, .region_len = w.region.len }; + for (w.hcas[0..w.n_hca], 0..) |h, i| { + r.rkey[i] = h.mr.rkey; + r.lid[i] = h.lid; + r.mtu[i] = h.mtu; + r.qpn[i] = h.qp.qp_num; + r.gid[i] = h.gid; + } + return r; + } + + /// The peer HCA for each of ours: the same IPv4 /24 (functions pair by subnet, names may differ), else position. + pub fn pairing(mine: Record, peer: Record) ![max_hcas]u32 { + if (peer.magic != record_magic) return error.BadRecord; + if (peer.n_hca != mine.n_hca or peer.slot_bytes != mine.slot_bytes or peer.region_len != mine.region_len) return error.BadRecord; + var out: [max_hcas]u32 = undefined; + for (0..mine.n_hca) |i| { + out[i] = @intCast(i); + var hits: u32 = 0; + for (0..peer.n_hca) |j| if (std.mem.eql(u8, mine.gid[i][12..15], peer.gid[j][12..15])) { + out[i] = @intCast(j); + hits += 1; + }; + if (hits > 1) out[i] = @intCast(i); // one subnet for both: keep the order + } + return out; + } + + /// INIT to RTR to RTS against the peer's record. + pub fn connect(w: *VerbsWire, peer: Record) !void { + const pairs = try pairing(w.record(), peer); + const api = w.v.api; + const m = abi.mask; + w.peer_region = peer.region; + for (w.hcas[0..w.n_hca], 0..) |*h, i| { + const j = pairs[i]; + w.peer_rkey[i] = peer.rkey[j]; + var rtr: abi.QpAttr = .{ .qp_state = abi.qps_rtr, .path_mtu = @intCast(@min(peer.mtu[j], h.mtu)), .dest_qp_num = peer.qpn[j], .rq_psn = 0, .max_dest_rd_atomic = 1, .min_rnr_timer = 12 }; + rtr.ah_attr = .{ .grh = .{ .dgid = .{ .raw = peer.gid[j] }, .sgid_index = h.gid_index, .hop_limit = 64, .traffic_class = w.traffic_class }, .dlid = @intCast(peer.lid[j]), .is_global = 1, .port_num = 1 }; + if (api.ibv_modify_qp(h.qp, &rtr, m.state | m.av | m.path_mtu | m.dest_qpn | m.rq_psn | m.max_dest_rd_atomic | m.min_rnr_timer) != 0) return error.VerbsFailed; + var rts: abi.QpAttr = .{ .qp_state = abi.qps_rts, .timeout = 14, .retry_cnt = 7, .rnr_retry = 7, .sq_psn = 0, .max_rd_atomic = 1 }; + if (api.ibv_modify_qp(h.qp, &rts, m.state | m.timeout | m.retry_cnt | m.rnr_retry | m.sq_psn | m.max_qp_rd_atomic) != 0) return error.VerbsFailed; + } + } + + /// Slot `slot`'s first `padded` bytes to the peer, striped over the HCAs, each stripe's `seq` flag on the same QP. + pub fn post(w: *VerbsWire, slot: u32, seq: u32, padded: u32) !void { + var plan: [max_hcas]guard.Stripe = undefined; + try guard.stripes(w.regionLayout(), w.n_hca, slot, padded, &plan); + while (w.inflight() + padded > max_inflight) try w.reapWait(); + var seq_copy = seq; + for (w.hcas[0..w.n_hca], plan[0..w.n_hca], 0..) |*h, s, i| { + while (h.outstanding >= ring_len) try w.reapHca(h); + var flag_sge: abi.Sge = .{ .addr = @intFromPtr(&seq_copy), .length = 4, .lkey = 0 }; + var flag_wr: abi.SendWr = .{ .wr_id = seq, .sg_list = @ptrCast(&flag_sge), .num_sge = 1, .opcode = abi.wr_rdma_write, .send_flags = abi.send_signaled | abi.send_inline, .remote_addr = w.peer_region + s.flag, .rkey = w.peer_rkey[i] }; + var data_sge: abi.Sge = .{ .addr = @intFromPtr(w.region.ptr) + s.local, .length = s.len, .lkey = h.mr.lkey }; + var data_wr: abi.SendWr = .{ .wr_id = seq, .next = &flag_wr, .sg_list = @ptrCast(&data_sge), .num_sge = 1, .opcode = abi.wr_rdma_write, .send_flags = 0, .remote_addr = w.peer_region + s.remote, .rkey = w.peer_rkey[i] }; + var bad: ?*abi.SendWr = null; + const first: *abi.SendWr = if (s.len > 0) &data_wr else &flag_wr; + if (h.ctx.ops.post_send(h.qp, first, &bad) != 0) return error.PostFailed; + h.ring[(h.head + h.outstanding) % ring_len] = s.len; + h.outstanding += 1; + h.inflight += s.len; + _ = h.posts.fetchAdd(1, .release); + } + if (w.inflight() > w.peak.load(.monotonic)) w.peak.store(w.inflight(), .monotonic); + w.posted += 1; + } + + /// This rank's layout; the MR covers `region` exactly and the peer's record matched it at connect. + pub fn regionLayout(w: *const VerbsWire) guard.Layout { + return .{ .send = w.send_off, .recv = w.recv_off, .flags = w.flag_off, .slot = w.slot_bytes, .len = w.region.len }; + } + + fn reapHca(w: *VerbsWire, h: *Hca) !void { + var wc: [32]abi.Wc = undefined; + const n = h.ctx.ops.poll_cq(h.cq, 32, &wc); + if (n < 0) return error.Completion; + for (wc[0..@intCast(n)]) |c| if (c.status != abi.wc_success) { + log.err("tp roce: RDMA write of seq {d} failed: {s} (vendor 0x{x})", .{ c.wr_id, w.v.statusText(c.status), c.vendor_err }); + return error.Completion; + }; + w.popped(h, @intCast(n)); + } + + fn popped(w: *VerbsWire, h: *Hca, n: u32) void { + for (0..n) |_| { + h.inflight -= h.ring[h.head]; + h.head = (h.head + 1) % ring_len; + } + h.outstanding -= n; + w.completed += n; + _ = h.done.fetchAdd(n, .release); + } + + /// Bytes posted and not yet completed over every HCA. + pub fn inflight(w: *const VerbsWire) u64 { + var sum: u64 = 0; + for (w.hcas[0..w.n_hca]) |h| sum += h.inflight; + return sum; + } + + /// Reaps until some post completes (error.Completion when none does within `inflight_wait_ns`). + fn reapWait(w: *VerbsWire) !void { + const before = w.completed; + const t0 = sock.nowNs(); + while (w.completed == before) { + try w.reap(); + if (w.completed == before and sock.nowNs() - t0 > inflight_wait_ns) return error.Completion; + std.atomic.spinLoopHint(); + } + } + + /// Reap every HCA's completions (after each op, and in the proxy's idle loop). + pub fn reap(w: *VerbsWire) !void { + for (w.hcas[0..w.n_hca]) |*h| try w.reapHca(h); + } +}; + +test "HCA lists parse and subnets pair the functions" { + var specs: [max_hcas]HcaSpec = undefined; + const n = try parseSpecs("rocep1s0f1:3, roceP2p1s0f1", &specs); + try std.testing.expectEqual(@as(usize, 2), n); + try std.testing.expectEqualStrings("rocep1s0f1", specs[0].name[0..specs[0].name_len]); + try std.testing.expectEqual(@as(?u8, 3), specs[0].gid); + try std.testing.expectEqual(@as(?u8, null), specs[1].gid); + var a: Record = .{ .n_hca = 2 }; + var b: Record = .{ .n_hca = 2 }; + const sub100 = [16]u8{ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0xff, 0xff, 192, 168, 100, 10 }; + var sub101 = sub100; + sub101[14] = 101; + a.gid = .{ sub100, sub101 }; + b.gid = .{ sub101, sub100 }; // the peer lists its functions the other way round + b.gid[0][15] = 11; + b.gid[1][15] = 11; + const p = try VerbsWire.pairing(a, b); + try std.testing.expectEqual(@as(u32, 1), p[0]); + try std.testing.expectEqual(@as(u32, 0), p[1]); + try std.testing.expect(ipv4Mapped(sub100)); + b.n_hca = 1; + try std.testing.expectError(error.BadRecord, VerbsWire.pairing(a, b)); + b.n_hca = 2; + b.slot_bytes = 1 << 20; + try std.testing.expectError(error.BadRecord, VerbsWire.pairing(a, b)); + b.slot_bytes = 0; + b.region_len = 4096; + try std.testing.expectError(error.BadRecord, VerbsWire.pairing(a, b)); +} diff --git a/zig/src/fabric/tp/run_hosts.sh b/zig/src/fabric/tp/run_hosts.sh new file mode 100644 index 0000000000..75f0c7900b --- /dev/null +++ b/zig/src/fabric/tp/run_hosts.sh @@ -0,0 +1,51 @@ +#!/usr/bin/env bash +# One rank of the two-host TP=2 gate over RoCE: `run_hosts.sh RANK RANK0_ADDR ./tf-tp-test OUTDIR` on both hosts. +set -u +RANK=${1:?rank}; MASTER=${2:?master address}; BIN=${3:?tf-tp-test}; OUT=${4:?outdir} +mkdir -p "$OUT" +[ -n "${NCCL_SOCKET_IFNAME:-}" ] && export NCCL_SOCKET_IFNAME # else NCCL picks the interface +# NCCL uses the RoCE wire's devices unless NCCL_IB_HCA is set: the names without `:gid` (NCCL reads `:N` as a port) +[ -n "${TF_TP_ROCE_HCA:-}" ] && [ -z "${NCCL_IB_HCA:-}" ] && NCCL_IB_HCA=$(echo "$TF_TP_ROCE_HCA" | sed 's/:[0-9]*//g; s/ //g') +[ -n "${TF_TP_ROCE_HCA:-}" ] && export TF_TP_ROCE_HCA +"$BIN" kernel > /dev/null 2>&1 || { echo "error: $BIN has no mailbox kernel: rebuild with -Dtp-nvcc=PATH or -Dtp-fatbin=PATH" >&2; exit 2; } +ulimit -l unlimited 2> /dev/null || echo "warning: memlock is $(ulimit -l) KiB: ibv_reg_mr of the RoCE region may fail" >&2 +# neither named: NCCL takes the ports the RoCE picker chose with the peer (`tf-tp-test hca`, both ranks), not its own +if [ -z "${TF_TP_ROCE_HCA:-}${NCCL_IB_HCA:-}" ]; then + # hca gives up after TF_TP_PLAN_TIMEOUT (default 300 s) with its reason; the cap here only catches a hang past that + CAP=$(awk 'BEGIN { t = ARGV[1] + 0; if (t <= 0) t = 300; printf "%d", int(t) + 31 }' "${TF_TP_PLAN_TIMEOUT:-300}") + echo "hca: rank $RANK picking the ports with the peer (up to $CAP s)" + PICK=$(TF_TP_RANK=$RANK TF_TP_MASTER=$MASTER TF_TP_PORT=29700 timeout "$CAP" "$BIN" hca 2> "$OUT/hca.log") + rc=$? + if [ "$rc" -eq 124 ]; then + echo "error: $BIN hca still running after $CAP s: killed ($OUT/hca.log)" >&2 + exit 3 + fi + if [ "$rc" -ne 0 ]; then + echo "error: $BIN hca exited $rc: $(grep -h '^FAIL\|^error' "$OUT/hca.log" | tail -n 1 | grep . || tail -n 1 "$OUT/hca.log") ($OUT/hca.log)" >&2 + exit 3 + fi + if ! echo "$PICK" | grep -Eqx 'NCCL_IB_HCA=[A-Za-z0-9_.-]+(,[A-Za-z0-9_.-]+)*'; then + echo "error: $BIN hca printed '$PICK', not NCCL_IB_HCA=name,... ($OUT/hca.log)" >&2 + exit 3 + fi + NCCL_IB_HCA=${PICK#NCCL_IB_HCA=} + echo "hca: rank $RANK NCCL_IB_HCA=$NCCL_IB_HCA ($(grep -h '^INFO hca' "$OUT/hca.log"))" +fi +[ -n "${NCCL_IB_HCA:-}" ] && export NCCL_IB_HCA +{ hostname; nvidia-smi -L; ibv_devices 2>&1; ulimit -l; env | grep -E '^(TF_TP_ROCE_HCA|NCCL_IB_HCA|NCCL_SOCKET_IFNAME)='; } > "$OUT/host.txt" 2>&1 +PORT=29700 +step() { # NAME BACKEND ARGS...: this host's rank; log OUT/NAME.log ending with an EXIT line + local name=$1 backend=$2; shift 2 + PORT=$((PORT + 10)) + ( TF_TP_RANK=$RANK TF_TP_MASTER=$MASTER TF_TP_PORT=$PORT TF_COMM_BACKEND=$backend TF_TP_ROCE_FALLBACK=error \ + timeout 900 "$BIN" "$@"; echo "EXIT rank $RANK rc $? at $(date +%s.%N)" ) > "$OUT/$name.log" 2>&1 + echo "$name: $(grep -h '^SUMMARY\|^PASS roce-host\|^INJECT\|^EXIT' "$OUT/$name.log" | tr '\n' ' ')" +} +step roce-host roce roce-host +step suite-roce roce suite +step suite-nccl nccl suite +for backend in roce nccl; do + for mode in fatal kill; do + for victim in 0 1; do step "ff-$backend-$mode-$victim" "$backend" failfast "$victim" "$mode"; done + done +done diff --git a/zig/src/fabric/tp/run_pair.sh b/zig/src/fabric/tp/run_pair.sh new file mode 100644 index 0000000000..1cf73c6261 --- /dev/null +++ b/zig/src/fabric/tp/run_pair.sh @@ -0,0 +1,52 @@ +#!/usr/bin/env bash +# Two ranks of tf-tp-test on one two-GPU host: `[SUITES=..] [FAILFAST=..] run_pair.sh BIN OUTDIR`. +set -u +BIN=${1:?bin}; OUT=${2:?outdir} +mkdir -p "$OUT" +nvidia-smi > "$OUT/smi.txt" 2>&1 +nvidia-smi topo -m > "$OUT/topo.txt" 2>&1 +ldconfig -p | grep -E "libnccl|libcuda" > "$OUT/libs.txt" 2>&1 +PORT=29500 +pair() { # NAME BACKEND ARGS...: both ranks, logs OUT/NAME.rR.log ending with an EXIT line + local name=$1 backend=$2 r; shift 2 + PORT=$((PORT + 10)) + for r in 0 1; do + ( TF_TP_RANK=$r TF_TP_DEVICE=$r TF_TP_PORT=$PORT TF_COMM_BACKEND=$backend NCCL_DEBUG=${NCCL_DEBUG:-WARN} \ + timeout 600 "$BIN" "$@"; echo "EXIT rank $r rc $? at $(date +%s.%N)" ) > "$OUT/$name.r$r.log" 2>&1 & + done + wait +} +# ffsum OUT NAME VICTIM: a fail-fast pair's verdict line (exit codes, exit times after the injection) +ffsum() { + [ -f "$1/$2.r0.log" ] && [ -f "$1/$2.r1.log" ] || { echo "ffsum: $1/$2.r0.log or .r1.log missing" >&2; return 1; } + awk -v name="$2" -v victim="$3" ' + FNR == 1 { r = (FILENAME ~ /\.r1\.log$/) ? 1 : 0 } + !(r in ex) && match($0, /EXIT rank [0-9] rc [0-9]+ at [0-9.]+/) { + split(substr($0, RSTART, RLENGTH), f, " "); rc[r] = f[5]; at[r] = f[7]; ex[r] = 1 + } + r == victim && !inj && match($0, /INJECT [A-Za-z0-9_]+ rank [0-9]+ at [0-9.]+/) { + split(substr($0, RSTART, RLENGTH), g, " "); t0 = g[6]; inj = 1 + } + END { + if (!inj || !(0 in ex) || !(1 in ex)) { print name ": FAIL (no inject / exit line)"; exit } + for (i = 0; i < 2; i++) { want[i] = (i == victim && name ~ /kill/) ? 137 : 70; dt[i] = at[i] - t0 } + ok = rc[0] == want[0] && rc[1] == want[1] && dt[1 - victim] < 3.0 + printf "%s: %s rc [%d, %d] (want [%d, %d]), exit after injection: rank0 %.3f s, rank1 %.3f s\n", name, + ok ? "PASS" : "FAIL", rc[0], rc[1], want[0], want[1], dt[0], dt[1] + }' "$1/$2.r0.log" "$1/$2.r1.log" +} +for backend in ${SUITES:-mailbox nccl}; do + pair "suite-$backend" "$backend" suite +done +: > "$OUT/failfast.txt" +for backend in ${FAILFAST:-nccl mailbox}; do + for mode in fatal kill; do + for victim in 0 1; do + name=ff-$backend-$mode-$victim + pair "$name" "$backend" failfast "$victim" "$mode" + ffsum "$OUT" "$name" "$victim" >> "$OUT/failfast.txt" + done + done +done +grep -h "^SUMMARY\|^FAIL" "$OUT"/suite*.log +cat "$OUT/failfast.txt" diff --git a/zig/src/fabric/tp/run_transport.sh b/zig/src/fabric/tp/run_transport.sh new file mode 100755 index 0000000000..68ebcddd9a --- /dev/null +++ b/zig/src/fabric/tp/run_transport.sh @@ -0,0 +1,120 @@ +#!/usr/bin/env bash +# The failed-transport check on two hosts (tf-tp-test failfast V transport); `run_transport.sh` alone prints how. +set -u +usage() { + cat << 'TXT' +run_transport.sh MASTER_ADDR PEER_SSH OUTDIR (on rank 0's host, as root; rank 1 runs over ssh) + Two hosts with this script and tf-tp-test (built with the mailbox kernel) at the same path on both. Before the steps, + on both hosts: no GPU process and memlock unlimited. Steps: T0 (suite), T1 (victim 1), T3 (victim 0), fail-fast on; + tf-tp-test injects only when both links are idle. After each step, on both hosts: the kernel log since the step + began and every HCA *err* counter against the step's start. A matching line, a counter that grew, an unexpected rc + or a failure line prints STOP: and ends the run (exit 1), with no later step. +run_transport.sh t2 OUTDIR (one host: TF_TP_ALLOW_HOST_ONLY=1) + T2 (fail-fast off, a hold) runs both ranks on this host over the shm wire; tf-tp-test refuses it on an RDMA wire. +TXT +} +[ $# -ge 2 ] || { usage; exit 2; } +KIT=$(cd "$(dirname "$0")" && pwd) +BIN=$KIT/tf-tp-test +KLOG='EXT4-fs|smmu|iommu|mlx5|AER|nvme|Xid' +BAD='RDMA write of seq|local protection|NOT drained|GPU sync FAILED|REFUSED|refused ERR|panic' +"$BIN" kernel > /dev/null 2>&1 || { echo "STOP: $BIN has no mailbox kernel: rebuild with -Dtp-nvcc=PATH or -Dtp-fatbin=PATH" >&2; exit 2; } + +if [ "${1:-}" = t2 ]; then + OUT=${2:?outdir} + [ "${TF_TP_ALLOW_HOST_ONLY:-}" = 1 ] || { echo "STOP: T2 is host-only: set TF_TP_ALLOW_HOST_ONLY=1" >&2; exit 2; } + mkdir -p "$OUT" + for r in 0 1; do + TF_TP_RANK=$r TF_TP_MASTER=127.0.0.1 TF_TP_PORT=29990 TF_COMM_BACKEND=roce TF_TP_ROCE_WIRE=shm TF_TP_FAILFAST=0 \ + timeout 300 "$BIN" failfast 1 transport 200 5 > "$OUT/T2-r$r.log" 2>&1 & + done + wait + grep -h '^PASS\|^FAIL\|REFUSED\|regions kept' "$OUT"/T2-r*.log + exit 0 +fi + +MASTER=${1:?master address}; PEER=${2:?peer ssh target}; OUT=${3:?outdir} +mkdir -p "$OUT" +stop() { echo "STOP: $*" | tee -a "$OUT/verdict.txt"; exit 1; } + +precheck() { + [ -z "$(nvidia-smi --query-compute-apps=pid --format=csv,noheader 2> /dev/null)" ] || { echo "a GPU process runs"; return 1; } + [ "$(bash -c 'ulimit -l unlimited 2> /dev/null; ulimit -l')" = unlimited ] || { echo "memlock not unlimited"; return 1; } + echo "ok $(hostname) $(uname -r)" +} +counters() { + find -L /sys/class/infiniband/ -maxdepth 5 -path '*ports*counters*' -name '*err*' 2> /dev/null | sort | while read -r f; do + echo "$f $(cat "$f" 2> /dev/null)" + done +} +klog() { # kernel lines since epoch $1 matching $2 (journalctl must answer: an unreadable log is a STOP) + journalctl -k --since "@$1" --no-pager -o short-iso 2> /dev/null | grep -E "$2" + return "${PIPESTATUS[0]}" +} +probe() { # this host's tf-tp-test QPs and MRs once a second until it exits + local pid="" t + for _ in $(seq 200); do pid=$(pgrep -nx tf-tp-test) && break; sleep 0.1; done + [ -z "$pid" ] && return + for t in $(seq 180); do + kill -0 "$pid" 2> /dev/null || break + echo "== t=$t pid $pid" + rdma resource show qp pid "$pid" 2>&1 + rdma resource show mr pid "$pid" 2>&1 + sleep 1 + done +} +on() { # NODE CMD ARGS...: a command or one of the functions above, here (local) or on the peer + local node=$1 fn=$2 def + shift 2 + if [ "$node" = local ]; then "$fn" "$@"; return; fi + def=$(declare -f "$fn" 2> /dev/null) + ssh -o BatchMode=yes "$PEER" "${def:+$def +}$fn $(printf '%q ' "$@")" +} +grew() { # BEFORE AFTER: counters that went up + awk 'NR == FNR { b[$1] = $2; next } ($1 in b) && $2 > b[$1] { print $1 " " b[$1] " -> " $2 }' "$1" "$2" +} + +for n in local peer; do + v=$(on $n precheck) || stop "$n: $v" + echo "precheck $n: $v" | tee -a "$OUT/verdict.txt" +done +on peer test -x "$BIN" || stop "peer: no $BIN (the kit must sit at the same path on both hosts)" + +PORT=29900 +step() { # NAME WANT_RC -- ARGS... + local name=$1 want=$2 n rc rc0 rc1 hits g + shift 3 + PORT=$((PORT + 10)) + declare -A t + for n in local peer; do + t[$n]=$(on $n date +%s) + on $n counters > "$OUT/$name.$n.err0" + done + on local probe > "$OUT/$name.local.rdma.txt" 2>&1 & + on peer probe > "$OUT/$name.peer.rdma.txt" 2>&1 & + local envs=(TF_TP_MASTER="$MASTER" TF_TP_PORT=$PORT TF_COMM_BACKEND=roce TF_TP_ROCE_FALLBACK=error TF_TP_MAILBOX_MAX_KB=1024 + TF_TP_ROCE_FAST=1 TF_TP_FAILFAST_GRACE_S=15) + ssh -o BatchMode=yes "$PEER" "ulimit -l unlimited && env $(printf '%q ' "${envs[@]}") TF_TP_RANK=1 timeout 300 $(printf '%q ' "$BIN" "$@"); echo \"EXIT rank 1 rc \$? at \$(date +%s.%N)\"" > "$OUT/$name-r1.log" 2>&1 & + ( ulimit -l unlimited && env "${envs[@]}" TF_TP_RANK=0 timeout 300 "$BIN" "$@"; echo "EXIT rank 0 rc $? at $(date +%s.%N)" ) > "$OUT/$name-r0.log" 2>&1 + wait + rc0=$(grep -oP '^EXIT rank 0 rc \K[0-9]+' "$OUT/$name-r0.log") + rc1=$(grep -oP '^EXIT rank 1 rc \K[0-9]+' "$OUT/$name-r1.log") + echo "$name: rank 0 rc ${rc0:-?}, rank 1 rc ${rc1:-?} (want $want)" | tee -a "$OUT/verdict.txt" + for n in local peer; do + on $n counters > "$OUT/$name.$n.err1" + g=$(grew "$OUT/$name.$n.err0" "$OUT/$name.$n.err1") + [ -n "$g" ] && stop "$name $n: HCA error counters grew: $g" + hits=$(on $n klog "${t[$n]}" "$KLOG"); rc=$? + [ "$rc" -ne 0 ] && stop "$name $n: the kernel log is unreadable (journalctl rc $rc)" + [ -n "$hits" ] && { echo "$hits" > "$OUT/$name.$n.klog.txt"; stop "$name $n: kernel lines: $hits"; } + done + hits=$(grep -hE "$BAD" "$OUT/$name-r0.log" "$OUT/$name-r1.log") + [ -n "$hits" ] && stop "$name: $hits" + [ "$rc0" = "$want" ] && [ "$rc1" = "$want" ] || stop "$name: rc $rc0 / $rc1, want $want" + grep -h 'halted\|quiesced\|idle before\|regions kept' "$OUT/$name"-r*.log | sed "s/^/ /" +} +step T0 0 -- suite +step T1 70 -- failfast 1 transport 200 +step T3 70 -- failfast 0 transport 200 +echo "PASS transport check (T0, T1, T3)" | tee -a "$OUT/verdict.txt" diff --git a/zig/src/fabric/tp/sock.zig b/zig/src/fabric/tp/sock.zig new file mode 100644 index 0000000000..b382884034 --- /dev/null +++ b/zig/src/fabric/tp/sock.zig @@ -0,0 +1,169 @@ +//! TCP for the bootstrap and the fate channel: listen, accept and connect with deadlines, whole sends and receives. +const std = @import("std"); +const c = std.c; +const posix = std.posix; + +pub const Fd = c.fd_t; + +pub const Error = error{ SocketFailed, AddressUnresolved, AddressInUse, Timeout, Closed, Interrupted }; + +pub fn nowNs() u64 { + var ts: c.timespec = undefined; + _ = c.clock_gettime(.MONOTONIC, &ts); + return @as(u64, @intCast(ts.sec)) * std.time.ns_per_s + @as(u64, @intCast(ts.nsec)); +} + +pub fn sleepNs(ns: u64) void { + var ts: c.timespec = .{ .sec = @intCast(ns / std.time.ns_per_s), .nsec = @intCast(ns % std.time.ns_per_s) }; + while (c.nanosleep(&ts, &ts) != 0) {} +} + +/// TCP_NODELAY (frames are tiny and latency matters) and keepalive (a vanished host fails a read in ~90 s, not never). +fn tune(fd: Fd) void { + const one: c_int = 1; + _ = c.setsockopt(fd, posix.IPPROTO.TCP, posix.TCP.NODELAY, &one, @sizeOf(c_int)); + _ = c.setsockopt(fd, posix.SOL.SOCKET, posix.SO.KEEPALIVE, &one, @sizeOf(c_int)); + const idle: c_int = 30; + const intvl: c_int = 10; + const cnt: c_int = 6; + _ = c.setsockopt(fd, posix.IPPROTO.TCP, posix.TCP.KEEPIDLE, &idle, @sizeOf(c_int)); + _ = c.setsockopt(fd, posix.IPPROTO.TCP, posix.TCP.KEEPINTVL, &intvl, @sizeOf(c_int)); + _ = c.setsockopt(fd, posix.IPPROTO.TCP, posix.TCP.KEEPCNT, &cnt, @sizeOf(c_int)); +} + +/// Rank 0's listening socket on every IPv4 address at `port`. +pub fn listen(port: u16) Error!Fd { + const fd = c.socket(posix.AF.INET, posix.SOCK.STREAM | posix.SOCK.CLOEXEC, 0); + if (fd < 0) return error.SocketFailed; + errdefer _ = c.close(fd); + const one: c_int = 1; + _ = c.setsockopt(fd, posix.SOL.SOCKET, posix.SO.REUSEADDR, &one, @sizeOf(c_int)); + var addr: posix.sockaddr.in = .{ .port = std.mem.nativeToBig(u16, port), .addr = 0 }; + if (c.bind(fd, @ptrCast(&addr), @sizeOf(posix.sockaddr.in)) != 0) return error.AddressInUse; + if (c.listen(fd, 64) != 0) return error.SocketFailed; + return fd; +} + +/// One connection before `deadline_ns` (monotonic). +pub fn accept(fd: Fd, deadline_ns: u64) Error!Fd { + while (true) { + try waitReadable(fd, deadline_ns); + const conn = c.accept(fd, null, null); + if (conn >= 0) { + tune(conn); + return conn; + } + } +} + +/// Connects to `host:port`, retrying every 50 ms (the master may not listen yet) until `deadline_ns`. +pub fn connect(host: []const u8, port: u16, deadline_ns: u64) Error!Fd { + var name: [256]u8 = undefined; + if (host.len >= name.len) return error.AddressUnresolved; + @memcpy(name[0..host.len], host); + name[host.len] = 0; + var service: [8]u8 = undefined; + const svc = std.fmt.bufPrintSentinel(&service, "{d}", .{port}, 0) catch unreachable; + const hints: c.addrinfo = .{ .flags = .{}, .family = posix.AF.INET, .socktype = posix.SOCK.STREAM, .protocol = 0, .addrlen = 0, .addr = null, .canonname = null, .next = null }; + while (true) { + var res: ?*c.addrinfo = null; + if (@backingInt(c.getaddrinfo(name[0..host.len :0], svc, &hints, &res)) == 0) { + defer c.freeaddrinfo(res.?); + const fd = c.socket(posix.AF.INET, posix.SOCK.STREAM | posix.SOCK.CLOEXEC, 0); + if (fd < 0) return error.SocketFailed; + if (c.connect(fd, res.?.addr.?, res.?.addrlen) == 0) { + tune(fd); + return fd; + } + _ = c.close(fd); + } + if (nowNs() >= deadline_ns) return error.Timeout; + sleepNs(50 * std.time.ns_per_ms); + } +} + +fn waitReadable(fd: Fd, deadline_ns: u64) Error!void { + while (true) { + const t = nowNs(); + if (t >= deadline_ns) return error.Timeout; + const ms: c_int = @intCast(@min((deadline_ns - t) / std.time.ns_per_ms + 1, 1000)); + var p = [_]c.pollfd{.{ .fd = fd, .events = posix.POLL.IN, .revents = 0 }}; + const n = c.poll(&p, 1, ms); + if (n > 0) return; + } +} + +/// Every byte of `bytes`, or `error.Closed`. +pub fn sendAll(fd: Fd, bytes: []const u8) Error!void { + var off: usize = 0; + while (off < bytes.len) { + const n = c.send(fd, bytes[off..].ptr, bytes.len - off, posix.MSG.NOSIGNAL); + if (n <= 0) { + if (n < 0 and posix.errno(n) == .INTR) continue; + return error.Closed; + } + off += @intCast(n); + } +} + +/// Fills `out` before `deadline_ns` (null: no deadline), or `error.Closed` at the end of the stream. +pub fn recvExact(fd: Fd, out: []u8, deadline_ns: ?u64) Error!void { + var off: usize = 0; + while (off < out.len) { + if (deadline_ns) |d| try waitReadable(fd, d); + const n = c.recv(fd, out[off..].ptr, out.len - off, 0); + if (n <= 0) { + if (n < 0 and posix.errno(n) == .INTR) continue; + return error.Closed; + } + off += @intCast(n); + } +} + +/// A non-blocking send of the whole of `bytes` if the socket buffer has room (false: dropped, or the peer is gone). +pub fn trySend(fd: Fd, bytes: []const u8) bool { + const n = c.send(fd, bytes.ptr, bytes.len, posix.MSG.NOSIGNAL | posix.MSG.DONTWAIT); + return n == @as(isize, @intCast(bytes.len)); +} + +pub fn hangup(fd: Fd) void { + _ = c.shutdown(fd, posix.SHUT.RDWR); +} + +pub fn close(fd: Fd) void { + _ = c.close(fd); +} + +/// This host's identity (the kernel's boot id, else the host name): ranks with equal ids share memory. +pub fn hostId(out: *[64]u8) []const u8 { + @memset(out, 0); + const f = c.open("/proc/sys/kernel/random/boot_id", .{ .ACCMODE = .RDONLY }); + if (f >= 0) { + defer _ = c.close(f); + const n = c.read(f, out, out.len); + if (n > 0) return std.mem.trim(u8, out[0..@intCast(n)], " \n"); + } + if (c.gethostname(out, out.len) != 0) return ""; + return std.mem.sliceTo(out, 0); +} + +test "a listener and a client trade bytes and see the hang-up" { + const l = try listen(0); + defer close(l); + var addr: posix.sockaddr.in = undefined; + var len: posix.socklen_t = @sizeOf(posix.sockaddr.in); + try std.testing.expectEqual(@as(c_int, 0), c.getsockname(l, @ptrCast(&addr), &len)); + const port = std.mem.bigToNative(u16, addr.port); + const deadline = nowNs() + 5 * std.time.ns_per_s; + const a = try connect("localhost", port, deadline); + defer close(a); + const b = try accept(l, deadline); + defer close(b); + try sendAll(a, "hello"); + var got: [5]u8 = undefined; + try recvExact(b, &got, deadline); + try std.testing.expectEqualStrings("hello", &got); + try std.testing.expectError(error.Timeout, recvExact(b, &got, nowNs() + 20 * std.time.ns_per_ms)); + hangup(a); + try std.testing.expectError(error.Closed, recvExact(b, &got, deadline)); +} diff --git a/zig/src/fabric/tp/tp.zig b/zig/src/fabric/tp/tp.zig new file mode 100644 index 0000000000..2c9eeb85ec --- /dev/null +++ b/zig/src/fabric/tp/tp.zig @@ -0,0 +1,217 @@ +//! Tensor parallelism across processes, one GPU each: bootstrap, collectives, fail-fast; `Session.open` starts a rank. +const std = @import("std"); +const log = @import("log.zig"); +const cuda = @import("cuda"); + +pub const collective = @import("collective.zig"); +pub const reference = @import("reference.zig"); +pub const config = @import("config.zig"); +pub const sock = @import("sock.zig"); +pub const bootstrap = @import("bootstrap.zig"); +pub const fate = @import("fate.zig"); +pub const host = @import("host.zig"); +pub const nccl = @import("nccl.zig"); +pub const mailbox = @import("mailbox.zig"); +pub const hybrid = @import("hybrid.zig"); +pub const roce = @import("roce.zig"); +pub const roce_verbs = @import("roce_verbs.zig"); +pub const roce_pick = @import("roce_pick.zig"); +pub const planlink = @import("planlink.zig"); +pub const Verbs = @import("verbs").Verbs; + +pub const Collective = collective.Collective; +pub const DType = collective.DType; +pub const Op = collective.Op; +pub const Config = config.Config; +pub const Bootstrap = bootstrap.Bootstrap; +pub const Fate = fate.Fate; + +pub const Session = struct { + cfg: Config, + lib: cuda.nccl.Library, + nccl: nccl.Nccl, + mail: ?mailbox.Mailbox = null, + roce: ?*roce.Roce = null, + route: hybrid.Hybrid = undefined, + /// The round plan's link (rank 0 sends, the others receive), on its own TCP connection. + plan: planlink.PlanLink = undefined, + fate: ?Fate = null, + + /// Brings this rank up in place; `kernel_image` is mailbox.cu's fatbin, empty sends those backends to NCCL. + pub fn open(s: *Session, d: *const cuda.Driver, cfg: Config, kernel_image: []const u8) !void { + var b = try Bootstrap.open(cfg); + defer b.close(); + try b.agree(&cfg.digest(), "settings (TF_TP_WORLD, TF_COMM_BACKEND, TF_TP_MAILBOX_MAX_KB, TF_TP_ROCE_CHUNKED_KB, TF_TP_FAILFAST, TF_TP_ROCE_WIRE / _FALLBACK, TF_GRAPH_CACHE)"); + var lib = try cuda.nccl.Library.openPath(cfg.nccl_lib); + errdefer lib.close(); + const v = try lib.version(); + try b.agree(std.mem.asBytes(&v), "NCCL version"); + var id: cuda.nccl.UniqueId = undefined; + if (cfg.rank == 0) try lib.check(lib.api.ncclGetUniqueId(&id), "ncclGetUniqueId"); + try b.broadcast(&id.internal, &id.internal); + s.* = .{ .cfg = cfg, .lib = lib, .nccl = undefined }; + s.nccl = try nccl.Nccl.init(&s.lib, d, id, cfg.rank, cfg.world); + errdefer s.nccl.deinit(); + switch (cfg.backend) { + .nccl => {}, + .mailbox => if (try mailboxPossible(&b, kernel_image)) { + s.mail = try mailbox.Mailbox.init(d, &b, kernel_image, cfg.mailbox_max_bytes, cfg.mailbox_timeout_ns); + s.route = .{ .small = .{ .mailbox = &s.mail.? }, .nccl = &s.nccl, .max_bytes = cfg.mailbox_max_bytes }; + } else if (cfg.rank == 0) { + log.warn("[tensorfold] tp: TF_COMM_BACKEND=mailbox needs two ranks on one host and the kernel; every rank uses NCCL", .{}); + }, + .roce => try s.openRoce(d, &b, kernel_image), + } + try s.agreeVar(&b); + s.plan = planlink.PlanLink.init(cfg.rank, cfg.world, b.releasePlan()); + s.plan.spin_ns = cfg.plan_spin_ns; + s.plan.pin = cfg.plan_pin; + s.plan.avoid = cfg.roce_cpu; + if (cfg.failfast) { + s.fate = Fate.init(cfg, b.release()); + s.fate.?.onFail(.{ .ctx = s, .call = abortHook }); + try s.fate.?.start(); + } + log.info("[tensorfold] tp: rank {d} of {d} up ({t}, NCCL {d}, fail-fast {s}, roce fast {s}: {d} poller(s), stagger {d} ns)", .{ cfg.rank, cfg.world, s.comm().kind(), v, if (cfg.failfast) "on" else "off", if (cfg.roce_fast & 1 != 0) "on" else "off", (cfg.roce_fast >> 8) & 0xff, cfg.roce_fast >> 16 }); + } + + /// RoCE for the small exchanges, probed against NCCL at load; on failure every rank obeys TF_TP_ROCE_FALLBACK. + fn openRoce(s: *Session, d: *const cuda.Driver, b: *Bootstrap, image: []const u8) !void { + const cfg = s.cfg; + var specs: [roce_verbs.max_hcas]roce_verbs.HcaSpec = undefined; + var n_specs: usize = 0; + if (cfg.roce_wire == .verbs) n_specs = try roce_pick.resolve(cfg, b, &specs); + const opts: roce.Options = .{ .wire = if (cfg.roce_wire == .shm) .shm else .verbs, .hcas = specs[0..n_specs], .traffic_class = cfg.roce_tc, .cpu = cfg.roce_cpu, .max_bytes = cfg.mailbox_max_bytes, .timeout_ns = cfg.mailbox_timeout_ns, .fast = cfg.roce_fast, .blocks = cfg.roce_blocks }; + const r = roce.Roce.init(d, b, image, opts) catch |e| return s.roceFallback(e); + s.roce = r; + if (cfg.roce_chunked_bytes > 0 and !r.chunk_abi) { + log.warn("TF_TP_ROCE_CHUNKED_KB={d}: the RoCE kernel image has no chunked ABI (tp_roce_chunk_abi); refusing to start", .{cfg.roce_chunked_bytes / 1024}); + r.deinit(); + s.roce = null; + return error.ChunkAbiMissing; + } + s.route = .{ .small = .{ .roce = r }, .nccl = &s.nccl, .max_bytes = cfg.mailbox_max_bytes, .chunked_bytes = cfg.roce_chunked_bytes }; + const good = probe(s, d) catch false; + var blobs: [bootstrap.max_ranks][]const u8 = undefined; + var buf: [bootstrap.max_ranks]u8 = undefined; + const all = try b.allGather(&.{@intFromBool(good)}, &buf, &blobs); + for (all) |x| if (x[0] != 1) { + r.deinit(); + s.roce = null; + return s.roceFallback(error.ProbeMismatch); + }; + } + + /// `Collective.varAgreed`: whether every rank's transport moves device lengths, agreed once a session. + fn agreeVar(s: *Session, b: *Bootstrap) !void { + var blobs: [bootstrap.max_ranks][]const u8 = undefined; + var buf: [bootstrap.max_ranks]u8 = undefined; + const all = try b.allGather(&.{@intFromBool(s.comm().varFits(16))}, &buf, &blobs); + var ok = true; + for (all) |x| ok = ok and x[0] == 1; + if (s.mail != null or s.roce != null) s.route.var_agreed = ok; + } + + fn roceFallback(s: *Session, e: anyerror) !void { + if (s.cfg.roce_fallback == .@"error") return e; + if (s.cfg.rank == 0) log.warn("[tensorfold] tp: RoCE unavailable ({t}); every rank uses NCCL (TF_TP_ROCE_FALLBACK=error refuses instead)", .{e}); + } + + /// One 4 KiB all-gather and one bf16 sum through RoCE and through NCCL: the bytes must be equal. + fn probe(s: *Session, d: *const cuda.Driver) !bool { + const n = 4096; + // input [0, n), gathers [n, 3n) and [3n, 5n), sums [5n, 6n) and [6n, 7n) + var a = try cuda.DeviceBuffer.alloc(d, 7 * n); + defer a.free(); + var stream = try cuda.Stream.init(d, true); + defer stream.deinit(); + var bytes: [n]u8 = undefined; + reference.fill(&bytes, .bf16, s.cfg.rank, 0x7072_6f62); + try a.uploadAsync(0, &bytes, stream.handle); + const fast = s.route.iface(); + const slow = s.nccl.iface(); + try fast.allGather(a.ptr, a.ptr + n, n, .u8, stream.handle); + try slow.allGather(a.ptr, a.ptr + 3 * n, n, .u8, stream.handle); + try fast.allReduce(a.ptr, a.ptr + 5 * n, n / 2, .bf16, .sum, stream.handle); + try slow.allReduce(a.ptr, a.ptr + 6 * n, n / 2, .bf16, .sum, stream.handle); + try stream.synchronize(); + try fast.check(); + var got: [6 * n]u8 = undefined; + try a.download(n, &got); + const gathers = std.mem.eql(u8, got[0 .. 2 * n], got[2 * n .. 4 * n]); + const sums = std.mem.eql(u8, got[4 * n .. 5 * n], got[5 * n .. 6 * n]); + if (!gathers or !sums) log.warn("[tensorfold] tp roce probe: gather {s}, sum {s} against NCCL", .{ if (gathers) "equal" else "DIFFERS", if (sums) "equal" else "DIFFERS" }); + return gathers and sums; + } + + fn mailboxPossible(b: *Bootstrap, image: []const u8) !bool { + if (b.world != 2) return false; + var id_buf: [64]u8 = undefined; + const id = sock.hostId(&id_buf); + var mine: [65]u8 = @splat(0); + mine[0] = @intFromBool(image.len > 0); + @memcpy(mine[1..][0..id.len], id); + var blobs: [bootstrap.max_ranks][]const u8 = undefined; + var buf: [2 * 65]u8 = undefined; + const all = try b.allGather(&mine, &buf, &blobs); + return all[0][0] == 1 and std.mem.eql(u8, all[0], all[1]); + } + + fn abortHook(ctx: *anyopaque, _: []const u8) void { + const s: *Session = @ptrCast(@alignCast(ctx)); + s.comm().abort(); + } + + /// The communicator the model uses: NCCL, or the hybrid route over the mailbox / RoCE. + pub fn comm(s: *Session) Collective { + if (s.mail != null or s.roce != null) return s.route.iface(); + return s.nccl.iface(); + } + + /// This rank failed inside a round: every rank goes down (exit 70). Without fail-fast it only aborts. + pub fn fatal(s: *Session, where: []const u8, err: anyerror) void { + if (s.fate) |*f| return f.fatal(where, err); + s.comm().abort(); + } + + /// A clean stop announced by rank 0 before it closes: followers end with no fail-fast exit. + pub fn announceStop(s: *Session) void { + if (s.fate) |*f| f.announceStop(); + } + + /// This rank failed after `open`, before serving: the other ranks learn why and go down (fail-fast). + pub fn bootFailed(s: *Session, err: anyerror) void { + if (s.fate) |*f| f.failBoot("boot", err); + s.close(); + } + + /// A normal stop: every rank calls it (rank 0 first marks the stop as expected with `fate.expectStop`). + pub fn close(s: *Session) void { + // a failed transport: every rank goes down (the peers stop posting) before this one lets go of anything + if (s.roce) |r| if (r.failed()) if (s.fate) |*f| f.fatal("transport", error.BackendFailed); + if (s.fate) |*f| f.close(); + s.plan.close(); + if (s.roce) |r| r.deinit(); + if (s.mail) |*m| m.deinit(); + s.nccl.deinit(); + s.lib.close(); + } +}; + +test { + _ = collective; + _ = reference; + _ = config; + _ = sock; + _ = bootstrap; + _ = mailbox; + _ = roce_verbs; + _ = roce_pick; + _ = planlink; + _ = @import("roce_guard.zig"); + _ = @import("roce_test.zig"); + _ = @import("roce_post_test.zig"); + _ = @import("roce_inflight_test.zig"); + _ = @import("host_test.zig"); + _ = @import("fate_test.zig"); +} diff --git a/zig/src/fabric/verbs.zig b/zig/src/fabric/verbs.zig index 9b9353c554..3c32714d30 100644 --- a/zig/src/fabric/verbs.zig +++ b/zig/src/fabric/verbs.zig @@ -2,7 +2,7 @@ const std = @import("std"); const platform = @import("sys"); const builtin = @import("builtin"); -const abi = @import("verbs_abi.zig"); +pub const abi = @import("verbs_abi.zig"); pub const Error = error{ LibraryUnavailable, MissingSymbol, NoDeviceList, NoSuchDevice, VerbsFailed, BadDeviceList, NoUsableGid }; pub const max_links = 8; @@ -51,12 +51,15 @@ pub const Api = struct { /// Exported by rdma-core 34+ but not in every header; used only when present. const GidEx = *const fn (ctx: Ctx, port: u32, index: u32, entry: *abi.GidEntry, flags: u32, size: usize) callconv(.c) c_int; const InOrder = *const fn (qp: *abi.Qp, op: c_int, flags: u32) callconv(.c) c_int; +/// ibv_query_qp: the QP's state as the device holds it (not every library exports it; used only when present). +pub const QueryQp = *const fn (qp: *abi.Qp, attr: *abi.QpAttr, mask: c_int, init: *abi.QpInitAttr) callconv(.c) c_int; pub const Verbs = struct { lib: platform.DynLib, api: Api, gid_ex: ?GidEx = null, in_order: ?InOrder = null, + query_qp: ?QueryQp = null, pub fn open(path: ?[:0]const u8) Error!Verbs { var lib = platform.DynLib.openZ((path orelse default_path).ptr) catch return error.LibraryUnavailable; @@ -69,6 +72,7 @@ pub const Verbs = struct { var v: Verbs = .{ .lib = lib, .api = api }; v.gid_ex = v.lib.lookup(GidEx, "_ibv_query_gid_ex"); v.in_order = v.lib.lookup(InOrder, "ibv_query_qp_data_in_order"); + v.query_qp = v.lib.lookup(QueryQp, "ibv_query_qp"); return v; }