Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 19 additions & 2 deletions zig/src/cuda/roce.zig
Original file line number Diff line number Diff line change
Expand Up @@ -449,11 +449,24 @@ pub const Gather = struct {
for (0..g.n_hca) |h| try g.drain(h);
}

/// A 20 us nap; where std.c has no timespec (Windows) the nap is simply absent.
fn napTimespec() ?std.c.timespec {
const T = @TypeOf(std.c.timespec);
if (T == void) return null;
switch (@typeInfo(T)) {
.@"struct" => {},
else => return null,
}
const v: std.c.timespec = .{ .sec = 0, .nsec = 20_000 };
return v;
}

/// The proxy: spins on the doorbell and posts each sequence in order (a missed one from its slot's size).
fn proxy(g: *Gather) void {
const c = g.ctrl();
var idle: u64 = 0;
const nap: std.c.timespec = .{ .sec = 0, .nsec = 20_000 };
// std.c.timespec is void on Windows (no CRT nanosleep), so the nap is built only where it exists.
const nap: ?std.c.timespec = napTimespec();
while (g.running.load(.monotonic)) {
const seq = @atomicLoad(u32, &c[c_seq], .acquire);
if (seq == g.last_seq) {
Expand All @@ -462,7 +475,11 @@ pub const Gather = struct {
g.failed.store(true, .release);
return;
};
if (idle >= idle_spins) _ = std.c.nanosleep(&nap, null);
if (idle >= idle_spins) {
if (nap) |*t| {
if (@TypeOf(std.c.nanosleep) != void) _ = std.c.nanosleep(t, null);
}
}
continue;
}
idle = 0;
Expand Down
5 changes: 5 additions & 0 deletions zig/tests/cuda/exl3_tp2_tests.zig
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
//! EXL3 tensor parallel: each rank's shard of real RED-SNOW projections and experts vs the engine's, all-gathered.
const std = @import("std");
const builtin = @import("builtin");
const posix = std.posix;
const cuda = @import("cuda");
const check = @import("check.zig");
Expand Down Expand Up @@ -336,6 +337,8 @@ fn bf16r(v: f32) u16 {

/// Rank 0 binds host:port and sends its unique id to the one peer (70 s for rank 1's retries).
fn serve(addr_text: []const u8, bytes: []const u8) !void {
// Windows hands socket() back as SOCKET (usize), so the fd-shaped cast below cannot compile there.
if (builtin.os.tag == .windows) return error.SocketFailed;
const a = try parseAddr(addr_text);
var storage: std.Io.Threaded.PosixAddress = undefined;
const len = try posixLen(a.host, a.port, &storage);
Expand Down Expand Up @@ -375,6 +378,8 @@ fn serve(addr_text: []const u8, bytes: []const u8) !void {

/// Rank 1 connects to host:port, retrying for 60 s (rank 0 may not have bound yet), and reads the id.
fn receive(gpu: Gpu, addr_text: []const u8, bytes: []u8) !void {
// Windows hands socket() back as SOCKET (usize), so the fd-shaped cast below cannot compile there.
if (builtin.os.tag == .windows) return error.SocketFailed;
const a = try parseAddr(addr_text);
const ip = try std.Io.net.IpAddress.parseIp4(a.host, a.port);
var fd: posix.socket_t = -1;
Expand Down
12 changes: 10 additions & 2 deletions zig/tests/cuda/flashnext_rank_tests.zig
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
//! Real two-process/two-device NCCL sums and bounded disconnect/timeout behavior over the framed control link.
const std = @import("std");
const builtin = @import("builtin");
const cuda = @import("cuda");
const flash = @import("flashnext");
const check = @import("check.zig");
Expand Down Expand Up @@ -45,8 +46,7 @@ pub fn run(gpu: check.Gpu, rank: usize, port: u16, mode: []const u8) !void {
if (std.mem.eql(u8, mode, "eof")) {
_ = std.c.shutdown(link.fd, std.posix.SHUT.RDWR);
} else {
const delay: std.c.timespec = .{ .sec = 0, .nsec = 600_000_000 };
_ = std.c.nanosleep(&delay, null);
sleepSixTenths();
}
check.pass("rank 1: injected {s}", .{mode});
return;
Expand Down Expand Up @@ -119,3 +119,11 @@ fn completeSum(c: *flash.collective.Collective, stream: cuda.Stream, values: u64
try c.sum(stream, values, count);
try c.finish(stream);
}

/// The injected 0.6 s peer delay. Windows' CRT has neither nanosleep nor a real std.c.timespec,
/// and this test never runs there.
fn sleepSixTenths() void {
if (builtin.os.tag == .windows) return;
const delay: std.c.timespec = .{ .sec = 0, .nsec = 600_000_000 };
_ = std.c.nanosleep(&delay, null);
}
Loading