From 3688f3cec9b3086bb76a068729965d50b0693980 Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Sat, 10 Oct 2026 10:28:21 +0700 Subject: [PATCH 001/104] CUDA keyed graph cache and device arena, both off until a family installs them. graph_cache.Cache captures one executable graph per caller-named key and replays it when the key comes back; ranks agree on every capture, eviction and fallback through an Agree hook, and a memory floor turns misses eager. arena.Arena sub-allocates DeviceBuffer.alloc from a few large chunks (VMM or plain cuMemAlloc) through memory.setHook; borrowed spans and pointers it does not own are freed as before. --- zig/src/cuda/abi.zig | 2 + zig/src/cuda/arena.zig | 331 +++++++++++++++++++++++++++++ zig/src/cuda/graph_cache.zig | 392 +++++++++++++++++++++++++++++++++++ zig/src/cuda/memory.zig | 52 ++++- zig/src/cuda/root.zig | 5 + 5 files changed, 779 insertions(+), 3 deletions(-) create mode 100644 zig/src/cuda/arena.zig create mode 100644 zig/src/cuda/graph_cache.zig diff --git a/zig/src/cuda/abi.zig b/zig/src/cuda/abi.zig index a608fd6114..95c9291dd0 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 = 126, }; pub const FunctionAttribute = enum(c_int) { diff --git a/zig/src/cuda/arena.zig b/zig/src/cuda/arena.zig new file mode 100644 index 0000000000..2c80b7581b --- /dev/null +++ b/zig/src/cuda/arena.zig @@ -0,0 +1,331 @@ +//! 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 + 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); + if (a.free_small.getPtr(n)) |l| if (l.pop()) |p| { + try a.live.put(a.gpa, 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; + try a.live.put(a.gpa, 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| { + a.be.release(a.be.ctx, p, e.value); + return true; + } + const e = a.live.fetchRemove(p) orelse return false; + 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; + } + + // -- 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)); +} diff --git a/zig/src/cuda/graph_cache.zig b/zig/src/cuda/graph_cache.zig new file mode 100644 index 0000000000..a2fb8e18bb --- /dev/null +++ b/zig/src/cuda/graph_cache.zig @@ -0,0 +1,392 @@ +//! 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, so ranks capture, evict and fall back at the same steps; never called in a capture. +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 +}; + +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 } = .{}, + + pub fn init(gpa: std.mem.Allocator, engine: Engine, settings: Settings) Self { + return .{ .gpa = gpa, .engine = engine, .settings = settings, .cap = settings.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 launched, captured first on a miss; a new `fingerprint` drops every graph first. + pub fn run(c: *Self, key: Key, stream: Stream, fingerprint: u64, body: Body) !Outcome { + if (!c.settings.on or c.failed) 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, else shed the least recent quarter and lower the cap (once below it, eager) + const room_here = if (c.room) |f| f() else true; + if (!try c.agreed(room_here and c.cap > 0)) { + 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)); + 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; + }; + if (!try c.agreed(exec != null)) { + if (exec) |x| c.engine.free(c.engine.ctx, x); + c.failed = true; + c.dropAll(); + return c.eager(stream, body); + } + try c.entries.put(c.gpa, key, .{ .exec = exec.?, .used = c.tick }); + 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. +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: fingerprint" { + try std.testing.expect(fingerprintOf(&.{ 1, 2 }) != fingerprintOf(&.{ 2, 1 })); +} diff --git a/zig/src/cuda/memory.zig b/zig/src/cuda/memory.zig index 2c469bae2c..6d802980c7 100644 --- a/zig/src/cuda/memory.zig +++ b/zig/src/cuda/memory.zig @@ -49,16 +49,37 @@ 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 { 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"); + } held(&device_bytes, len); return .{ .d = d, .ptr = p, .len = len }; } @@ -73,7 +94,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; @@ -209,6 +233,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/root.zig b/zig/src/cuda/root.zig index ebbb12b935..cf7e74d559 100644 --- a/zig/src/cuda/root.zig +++ b/zig/src/cuda/root.zig @@ -6,6 +6,7 @@ pub const Error = @import("driver.zig").Error; pub const Context = @import("context.zig").Context; pub const DeviceBuffer = @import("memory.zig").DeviceBuffer; pub const HostBuffer = @import("memory.zig").HostBuffer; +pub const memory = @import("memory.zig"); pub const usage = @import("memory.zig").usage; pub const Usage = @import("memory.zig").Usage; pub const carveout = @import("carveout.zig"); @@ -32,6 +33,8 @@ pub const grouped = @import("grouped.zig"); pub const qlinear = @import("qlinear.zig"); pub const experts = @import("experts.zig"); pub const modelopt = @import("modelopt.zig"); +pub const graph_cache = @import("graph_cache.zig"); +pub const arena = @import("arena.zig"); test { _ = @import("memory.zig"); @@ -47,4 +50,6 @@ test { _ = carveout; _ = @import("nvfp4.zig"); _ = @import("modelopt.zig"); + _ = graph_cache; + _ = arena; } From b7be793d801490d243cb61c3cc37f2882ed20e9c Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Sat, 10 Oct 2026 14:24:18 +0700 Subject: [PATCH 002/104] A failed CUDA kernel launch names its kernel and launch configuration; CUresult names fall back to a table when the driver cannot answer. cuFuncGetName is optional (CUDA 12.3+), so an older driver still opens. During the exit's teardown cuGetErrorName fails too, which printed CUDA_ERROR_DEINITIALIZED as "CUDA_ERROR_UNKNOWN_CODE (4)"; code 4 now says the process is exiting. --- zig/src/cuda/driver.zig | 67 ++++++++++++++++++++++++++++++++++++++--- zig/src/cuda/launch.zig | 36 ++++++++++++++++++++-- 2 files changed, 96 insertions(+), 7 deletions(-) diff --git a/zig/src/cuda/driver.zig b/zig/src/cuda/driver.zig index 0a860eac92..88ee9050dc 100644 --- a/zig/src/cuda/driver.zig +++ b/zig/src/cuda/driver.zig @@ -8,6 +8,10 @@ pub const Error = error{ DriverUnavailable, MissingSymbol, CudaFailed, OutOfDevi pub const Driver = struct { lib: std.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("libcuda.so.1"); @@ -25,7 +29,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; } @@ -38,18 +42,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 { @@ -82,3 +101,43 @@ pub const Driver = struct { return @intCast(10 * major + minor); } }; + +/// 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/launch.zig b/zig/src/cuda/launch.zig index 2472042638..c80684f291 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; @@ -80,7 +81,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); @@ -96,7 +97,36 @@ 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; + 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 { From faf590149837e02436edcf19c83d3f7c29e06eb9 Mon Sep 17 00:00:00 2001 From: Jonathan Regan-Beasley <197533392+jregan-beasley-smc@users.noreply.github.com> Date: Sun, 11 Oct 2026 10:23:24 +1100 Subject: [PATCH 003/104] fix(cuda): correct RDMA VMM capability attribute Use CUDA device attribute 110 for GPUDirect RDMA with VMM support. Attribute 126 reports the memory sync domain count and can incorrectly enable the RDMA allocation flag. Add a regression for the external ABI contract. --- zig/src/cuda/abi.zig | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/zig/src/cuda/abi.zig b/zig/src/cuda/abi.zig index 95c9291dd0..dd1bbe1033 100644 --- a/zig/src/cuda/abi.zig +++ b/zig/src/cuda/abi.zig @@ -40,7 +40,7 @@ pub const DeviceAttribute = enum(c_int) { 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 = 126, + gpu_direct_rdma_with_cuda_vmm_supported = 110, }; pub const FunctionAttribute = enum(c_int) { @@ -222,3 +222,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, + )); +} From 3cc0397482e01bbf88f3dcfbbded98725e0eeca1 Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Sun, 11 Oct 2026 17:29:31 +0700 Subject: [PATCH 004/104] Graph cache: a capture is freed until the map holds it, the keep floor stays under the max, a max of 0 is eager - A capture whose agreement throws, or whose insert runs out of memory, no longer leaks its executable: the map's slot is reserved before the capture, and the executable is owned here until the insert, which cannot fail. - A host map that cannot grow counts as no room in the existing room vote, so every rank declines that miss together. - `min_keep` is clamped to `max`, so under the memory floor the cap never rises above the configured max. A max of 0 runs every step eagerly, with no collective. - A cap lowered to 0 under the floor goes back to max once every rank has room, instead of staying eager for good. - Hits make no collective: with `Agree` set, keys, settings and fingerprint changes must already match on every rank. The docs now say so: the fingerprint changes at the same step on every rank, and `fingerprintOf` is for one rank. Co-authored-by: Jonathan Regan-Beasley <197533392+jregan-beasley-smc@users.noreply.github.com> --- zig/src/cuda/graph_cache.zig | 31 ++-- zig/src/cuda/graph_cache_regression_test.zig | 181 +++++++++++++++++++ 2 files changed, 201 insertions(+), 11 deletions(-) create mode 100644 zig/src/cuda/graph_cache_regression_test.zig diff --git a/zig/src/cuda/graph_cache.zig b/zig/src/cuda/graph_cache.zig index a2fb8e18bb..693f8ce25a 100644 --- a/zig/src/cuda/graph_cache.zig +++ b/zig/src/cuda/graph_cache.zig @@ -12,7 +12,7 @@ pub const Body = struct { run: *const fn (ctx: *anyopaque, stream: Stream) anyerror!void, }; -/// Every rank's verdict AND-ed, so ranks capture, evict and fall back at the same steps; never called in a capture. +/// Every rank's verdict AND-ed on a miss (hits never call it), so ranks capture, evict and fall back together. pub const Agree = struct { ctx: *anyopaque, all: *const fn (ctx: *anyopaque, ok: bool) anyerror!bool, @@ -53,7 +53,9 @@ pub fn Cache(comptime Key: type) type { stats: struct { replayed: u64 = 0, captured: u64 = 0, eager: u64 = 0, evicted: u64 = 0, dropped: u64 = 0, floor_hits: u64 = 0 } = .{}, pub fn init(gpa: std.mem.Allocator, engine: Engine, settings: Settings) Self { - return .{ .gpa = gpa, .engine = engine, .settings = settings, .cap = settings.max }; + var s = settings; + s.min_keep = @min(s.min_keep, s.max); // the floor never holds more than the configured max + return .{ .gpa = gpa, .engine = engine, .settings = s, .cap = s.max }; } pub fn deinit(c: *Self) void { @@ -98,9 +100,9 @@ pub fn Cache(comptime Key: type) type { return .eager; } - /// One step: the key's graph launched, captured first on a miss; a new `fingerprint` drops every graph first. + /// 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) return c.eager(stream, body); + 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; @@ -110,9 +112,10 @@ pub fn Cache(comptime Key: type) type { c.stats.replayed += 1; return .replayed; } - // a miss: room on every rank, else shed the least recent quarter and lower the cap (once below it, eager) - const room_here = if (c.room) |f| f() else true; - if (!try c.agreed(room_here and c.cap > 0)) { + // 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; @@ -123,7 +126,7 @@ pub fn Cache(comptime Key: type) type { 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)); + 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); @@ -138,13 +141,15 @@ pub fn Cache(comptime Key: type) type { 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)) { - if (exec) |x| c.engine.free(c.engine.ctx, x); c.failed = true; c.dropAll(); return c.eager(stream, body); } - try c.entries.put(c.gpa, key, .{ .exec = exec.?, .used = c.tick }); + 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; @@ -196,7 +201,7 @@ pub const CudaEngine = struct { } }; -/// FNV-1a over the addresses a body captures (buffer bases, sizes): the fingerprint `run` compares. +/// 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| { @@ -390,3 +395,7 @@ test "graph cache: off runs every step eagerly" { test "graph cache: fingerprint" { try std.testing.expect(fingerprintOf(&.{ 1, 2 }) != fingerprintOf(&.{ 2, 1 })); } + +test { + _ = @import("graph_cache_regression_test.zig"); +} 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); +} From c58d83ca53086a0f60944ce050cab5167838f40a Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Sun, 11 Oct 2026 18:33:36 +0700 Subject: [PATCH 005/104] CUDA arena: the address maps are rehashed before tombstones fill them The arena's `live` (small ranges) and `large` (own chunks) maps drop an entry on every free and gain one on every allocation for the life of the process. std's tombstones never clear, so in time they take every free slot, and each allocation, free and miss then probes the whole map. Each map now counts its removals and is rehashed in place once they reach a tenth of its capacity. Test: 20,000 random allocations and frees (small ranges and large chunks) keep a tenth of each map free at every step; without the rehash it fails. --- zig/src/cuda/arena.zig | 43 ++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 43 insertions(+) diff --git a/zig/src/cuda/arena.zig b/zig/src/cuda/arena.zig index 2c80b7581b..3886c81403 100644 --- a/zig/src/cuda/arena.zig +++ b/zig/src/cuda/arena.zig @@ -36,6 +36,7 @@ pub const Arena = struct { 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 { @@ -114,16 +115,26 @@ pub const Arena = struct { 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 { @@ -329,3 +340,35 @@ test "arena: under DeviceBuffer.alloc through memory.setHook" { 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()); + } + } +} From 25c02d3865f323c4907b4949dac216909fe5d5e8 Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Sat, 10 Oct 2026 10:38:24 +0700 Subject: [PATCH 006/104] Tensor parallelism across CUDA processes: bootstrap, NCCL, mailbox and RoCE collectives, plan link, fail-fast. zig/src/fabric/tp brings a rank up from TF_TP_WORLD, TF_TP_RANK, TF_TP_MASTER and TF_TP_PORT, agrees the settings digest over every rank, and gives a family all-gather, all-reduce and broadcast over NCCL (TF_COMM_BACKEND=nccl, the default), a pinned host mailbox for ranks on one host, or RoCE verbs between two hosts (TF_TP_ROCE_*), with an NCCL fallback when RoCE setup fails. Rank 0 sends each round's plan to the others over its own link (TF_TP_PLAN_*). A failure or a vanished rank ends every rank shortly after (TF_TP_FAILFAST, TF_TP_FAILFAST_GRACE_S). No family uses it yet. zig build test-tp runs the host tests; zig build tf-tp-test builds the two-GPU runner (-Dtp-nvcc= or -Dtp-fatbin= for the mailbox kernel). The RoCE all-gather adapts b12x's RoCEnante protocol. --- THIRD_PARTY_NOTICES.md | 6 + build.zig | 2 + zig/build/tp.zig | 59 ++++ zig/src/cuda/nccl.zig | 7 + zig/src/fabric/tp/bootstrap.zig | 262 +++++++++++++++ zig/src/fabric/tp/collective.zig | 136 ++++++++ zig/src/fabric/tp/config.zig | 245 ++++++++++++++ zig/src/fabric/tp/fate.zig | 261 +++++++++++++++ zig/src/fabric/tp/fate_test.zig | 146 ++++++++ zig/src/fabric/tp/gpu_bench.zig | 82 +++++ zig/src/fabric/tp/gpu_checks.zig | 188 +++++++++++ zig/src/fabric/tp/gpu_main.zig | 92 +++++ zig/src/fabric/tp/host.zig | 293 ++++++++++++++++ zig/src/fabric/tp/host_test.zig | 154 +++++++++ zig/src/fabric/tp/hybrid.zig | 167 ++++++++++ zig/src/fabric/tp/log.zig | 15 + zig/src/fabric/tp/mailbox.cu | 556 +++++++++++++++++++++++++++++++ zig/src/fabric/tp/mailbox.zig | 215 ++++++++++++ zig/src/fabric/tp/nccl.zig | 163 +++++++++ zig/src/fabric/tp/planlink.zig | 467 ++++++++++++++++++++++++++ zig/src/fabric/tp/reference.zig | 121 +++++++ zig/src/fabric/tp/roce.zig | 458 +++++++++++++++++++++++++ zig/src/fabric/tp/roce_host.zig | 58 ++++ zig/src/fabric/tp/roce_test.zig | 116 +++++++ zig/src/fabric/tp/roce_verbs.zig | 297 +++++++++++++++++ zig/src/fabric/tp/run_hosts.sh | 24 ++ zig/src/fabric/tp/run_pair.sh | 52 +++ zig/src/fabric/tp/sock.zig | 169 ++++++++++ zig/src/fabric/tp/tp.zig | 212 ++++++++++++ zig/src/fabric/verbs.zig | 2 +- 30 files changed, 5024 insertions(+), 1 deletion(-) create mode 100644 zig/build/tp.zig create mode 100644 zig/src/fabric/tp/bootstrap.zig create mode 100644 zig/src/fabric/tp/collective.zig create mode 100644 zig/src/fabric/tp/config.zig create mode 100644 zig/src/fabric/tp/fate.zig create mode 100644 zig/src/fabric/tp/fate_test.zig create mode 100644 zig/src/fabric/tp/gpu_bench.zig create mode 100644 zig/src/fabric/tp/gpu_checks.zig create mode 100644 zig/src/fabric/tp/gpu_main.zig create mode 100644 zig/src/fabric/tp/host.zig create mode 100644 zig/src/fabric/tp/host_test.zig create mode 100644 zig/src/fabric/tp/hybrid.zig create mode 100644 zig/src/fabric/tp/log.zig create mode 100644 zig/src/fabric/tp/mailbox.cu create mode 100644 zig/src/fabric/tp/mailbox.zig create mode 100644 zig/src/fabric/tp/nccl.zig create mode 100644 zig/src/fabric/tp/planlink.zig create mode 100644 zig/src/fabric/tp/reference.zig create mode 100644 zig/src/fabric/tp/roce.zig create mode 100644 zig/src/fabric/tp/roce_host.zig create mode 100644 zig/src/fabric/tp/roce_test.zig create mode 100644 zig/src/fabric/tp/roce_verbs.zig create mode 100644 zig/src/fabric/tp/run_hosts.sh create mode 100644 zig/src/fabric/tp/run_pair.sh create mode 100644 zig/src/fabric/tp/sock.zig create mode 100644 zig/src/fabric/tp/tp.zig diff --git a/THIRD_PARTY_NOTICES.md b/THIRD_PARTY_NOTICES.md index cb54012303..15680d812f 100644 --- a/THIRD_PARTY_NOTICES.md +++ b/THIRD_PARTY_NOTICES.md @@ -76,6 +76,12 @@ checked bit for bit against ExLlamaV3's dequantization. Flash Next's optional int8 and int4 KV caches (`families/qwen4_exp/cuda/kvcache.py`) follow the cache quantization scheme of [ExLlamaV3](https://github.com/turboderp-org/exllamav3) `-cq 8` and `-cq 4` (MIT License, Copyright (c) 2025 Turboderp, text below): groups of 32, one fp16 absmax scale per group, the group rotated by a 32-point Hadamard, midpoint-grid codes, `compand_a == 0`. 8-bit stores each code as a signed int8 (`q - 128`). 4-bit stores two unsigned codes per byte, low nibble first (the same bits as ExLlamaV3's little-endian packing, a uint8 tensor rather than their uint32 words). Their dequantizer folds another `1/sqrt(32)` into the scale and applies the unnormalized butterfly on the way out; this cache applies the normalized H32 to the query and to the merged output instead, and leaves the stored codes rotated. Scales match their quantizer bit for bit. Reconstructed values agree within fp16/bf16 rounding (under 0.01 on random groups), not bit for bit. The quantizer and the attention dequant are written for TensorFold and checked against an independent reference of that arithmetic. +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 `src/tensorfold/drafters/vendor/z_lab_dflash/model_mlx.py` is the unmodified `dflash/model_mlx.py` from diff --git a/build.zig b/build.zig index 054e688280..98ff809634 100644 --- a/build.zig +++ b/build.zig @@ -5,6 +5,7 @@ const dist_build = @import("zig/build/dist.zig"); const qwen27_build = @import("zig/build/qwen27.zig"); const hip_build = @import("zig/build/hip.zig"); const grammar_build = @import("zig/build/grammar.zig"); +const tp_build = @import("zig/build/tp.zig"); comptime { const required = std.mem.trim(u8, @embedFile(".zig-version"), "\r\n"); @@ -175,6 +176,7 @@ pub fn build(b: *std.Build) void { test_step.dependOn(&warning_run.step); b.step("test-learn-warning", "CPU production-log refusal, backoff and recovery test").dependOn(&warning_run.step); grammar_build.steps(b, target, xgrammar); + tp_build.add(b, target, optimize, test_step); } /// `zig build native -Dcpu=apple_m1`: tensorfold-native with the Metal engines for the Python package's bundle (a native M5 build traps on M1-M4). diff --git a/zig/build/tp.zig b/zig/build/tp.zig new file mode 100644 index 0000000000..f5dacc46e3 --- /dev/null +++ b/zig/build/tp.zig @@ -0,0 +1,59 @@ +//! Tensor parallelism's build: host tests on any OS, 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 tests = b.addRunArtifact(b.addTest(.{ .root_module = module(b, host, .debug, runtime(b, host, .debug)) })); + b.step("test-tp", "Tensor parallelism's host tests: bootstrap, fate channel, in-process ranks (no GPU)").dependOn(&tests.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 }); + 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/nccl.zig b/zig/src/cuda/nccl.zig index f742d62549..85f88bace8 100644 --- a/zig/src/cuda/nccl.zig +++ b/zig/src/cuda/nccl.zig @@ -12,6 +12,8 @@ pub const RedOp = enum(c_int) { sum = 0, prod = 1, max = 2, min = 3, avg = 4 }; 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; @@ -25,6 +27,11 @@ pub const Api = struct { ncclAllReduce: *const fn (D, D, usize, DataType, RedOp, Comm, abi.Stream) callconv(.c) R, ncclAllGather: *const fn (D, D, usize, DataType, Comm, abi.Stream) callconv(.c) R, ncclBroadcast: *const fn (D, D, usize, DataType, c_int, Comm, abi.Stream) callconv(.c) R, + ncclSend: *const fn (D, usize, DataType, c_int, Comm, abi.Stream) callconv(.c) R, + ncclRecv: *const fn (D, usize, DataType, c_int, Comm, abi.Stream) callconv(.c) R, + /// Frees the communicator after ending its pending work; safe from another thread while a kernel waits. + ncclCommAbort: *const fn (Comm) callconv(.c) R, + ncclCommGetAsyncError: *const fn (Comm, *Result) callconv(.c) R, ncclGroupStart: *const fn () callconv(.c) R, ncclGroupEnd: *const fn () callconv(.c) R, }; diff --git a/zig/src/fabric/tp/bootstrap.zig b/zig/src/fabric/tp/bootstrap.zig new file mode 100644 index 0000000000..1486e3a8f7 --- /dev/null +++ b/zig/src/fabric/tp/bootstrap.zig @@ -0,0 +1,262 @@ +//! 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); +} 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..13647f5921 --- /dev/null +++ b/zig/src/fabric/tp/config.zig @@ -0,0 +1,245 @@ +//! A rank's TP settings from the environment: rank, master, port, backend, fail-fast knobs. +const std = @import("std"); +const log = @import("log.zig"); + +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}; + +/// 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, + /// 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: every active Ethernet port with an IPv4 RoCE v2 GID). + 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, + + /// 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); + 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) }; + } + } + 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.failfast)); + h.update(std.mem.asBytes(&c.roce_wire)); + h.update(std.mem.asBytes(&c.roce_fallback)); + 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 = &.{}; +} 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..6019055f50 --- /dev/null +++ b/zig/src/fabric/tp/fate_test.zig @@ -0,0 +1,146 @@ +//! 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)); +} 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_main.zig b/zig/src/fabric/tp/gpu_main.zig new file mode 100644 index 0000000000..c5668a7b06 --- /dev/null +++ b/zig/src/fabric/tp/gpu_main.zig @@ -0,0 +1,92 @@ +//! `tf-tp-test suite | failfast RANK fatal|kill | roce-host`: one rank of a two-process TP=2 test on real GPUs. +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 max_bytes = (4 << 20) + 4096; + +pub fn main(init: std.process.Init) !u8 { + const args = try init.minimal.args.toSlice(init.arena.allocator()); + if (args.len < 2) { + std.debug.print("usage: tf-tp-test suite | failfast RANK fatal|kill | roce-host\n", .{}); + return 2; + } + const cfg = try tp.Config.fromEnv(); + if (std.mem.eql(u8, args[1], "roce-host")) return @import("roce_host.zig").run(cfg, init.gpa); + 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 }); + if (std.mem.eql(u8, args[1], "suite")) return suite(&s, &driver, init.gpa, cfg); + if (std.mem.eql(u8, args[1], "failfast") and args.len == 4) return failfast(&s, &driver, init.gpa, cfg, try std.fmt.parseInt(u32, args[2], 10), args[3]); + std.debug.print("unknown command {s}\n", .{args[1]}); + return 2; +} + +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; +} + +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/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..2cdeb0e4a4 --- /dev/null +++ b/zig/src/fabric/tp/hybrid.zig @@ -0,0 +1,167 @@ +//! 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(), + } + } +}; + +pub const Hybrid = struct { + small: Small, + nccl: *Nccl, + /// Largest shard a rank sends through the small path. + max_bytes: usize, + /// 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 (varFits(ptr, n)) return h.small.roce.allGatherV(s, r, n, lens, stream); + return allGather(ptr, s, r, stride, t, stream); + } + + /// RoCE with its device side, for a stride the small path takes (the route is the size's alone: every rank alike). + pub fn varFits(ptr: *anyopaque, n: usize) bool { + const h = self_(ptr); + return switch (h.small) { + .roce => |p| n > 0 and h.routed(n) and p.d != null, + .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..63d53b2484 --- /dev/null +++ b/zig/src/fabric/tp/mailbox.cu @@ -0,0 +1,556 @@ +// 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; + } + } + } +} + +} // 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; + if (atom_cas_gpu(a.state + 2, 0u, 1u) == 0u) { // one failure record a rank + uint32_t *rec = a.ctrl + kCtrlRank + 8u * a.rank; + st_relaxed_sys(rec + 1, seq); + st_relaxed_sys(rec + 2, seen); + st_relaxed_sys(rec + 3, blockIdx.x); + st_relaxed_sys(rec + 4, why); + st_relaxed_sys(rec + 5, waited); + fence_sc_sys(); + st_relaxed_sys(rec, 1u); + fence_sc_sys(); + } + } + } + __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; +}; + +// 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); +} + +} // namespace tftp + +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 mine = var ? roce_shard(a, a.rank) : a.nbytes; + const uint64_t theirs = var ? roce_shard(a, 1u - a.rank) : a.nbytes; + const uint64_t padded = var ? (mine < 16u ? 16u : (mine + 15u) & ~uint64_t(15u)) : (a.nbytes + 15u) & ~uint64_t(15u); + + // 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, 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, a.in, + 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; + if (atom_cas_gpu(a.state + 2, 0u, 1u) == 0u) { + st_relaxed_sys(a.ctrl + kRoceFailed + 1, seq); + st_relaxed_sys(a.ctrl + kRoceFailed + 2, seen); + st_relaxed_sys(a.ctrl + kRoceFailed + 3, lane); + st_relaxed_sys(a.ctrl + kRoceFailed + 4, why); + st_relaxed_sys(a.ctrl + kRoceFailed + 5, waited); + fence_sc_sys(); + st_relaxed_sys(a.ctrl + kRoceFailed, 1u); + fence_sc_sys(); + } + } + } + } 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; + if (atom_cas_gpu(a.state + 2, 0u, 1u) == 0u) { + st_relaxed_sys(a.ctrl + kRoceFailed + 1, seq); + st_relaxed_sys(a.ctrl + kRoceFailed + 2, seen); + st_relaxed_sys(a.ctrl + kRoceFailed + 3, h); + st_relaxed_sys(a.ctrl + kRoceFailed + 4, why); + st_relaxed_sys(a.ctrl + kRoceFailed + 5, waited); + fence_sc_sys(); + st_relaxed_sys(a.ctrl + kRoceFailed, 1u); + fence_sc_sys(); + } + 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, 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..b641839aad --- /dev/null +++ b/zig/src/fabric/tp/nccl.zig @@ -0,0 +1,163 @@ +//! 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; + } + + 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..58f88e3441 --- /dev/null +++ b/zig/src/fabric/tp/roce.zig @@ -0,0 +1,458 @@ +//! 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 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, +}; + +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: the peer's region is mapped here; a post copies the slot, then stores the flag with release order. +const ShmWire = struct { + peer: []u8, + local: []u8, + slot_bytes: usize, + + fn post(w: *ShmWire, slot: u32, seq: u32, padded: u32) void { + const recv = send_off + 2 * w.slot_bytes; + @memcpy(w.peer[recv + slot * w.slot_bytes ..][0..padded], w.local[send_off + slot * w.slot_bytes ..][0..padded]); + const flag: *std.atomic.Value(u32) = @ptrCast(@alignCast(w.peer[flag_off + slot * 2 * flag_stride ..].ptr)); + flag.store(seq, .release); + } +}; + +const Wire = union(WireKind) { + verbs: rv.VerbsWire, + 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| 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 => 1, + }; + } +}; + +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 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; + 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); + 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 sharedPair(b, len); + r.region = @alignCast(r.mapping[r.me * len ..][0..len]); + r.wire = .{ .shm = .{ .peer = r.mapping[(1 - r.me) * len ..][0..len], .local = r.region, .slot_bytes = r.slot_bytes } }; + }, + } + 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); + 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"); + 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); + } + + pub fn deinit(r: *Roce) void { + r.running.store(false, .release); + if (r.thread) |t| t.join(); + r.closeDevice(); + r.closeWire(); + std.heap.page_allocator.destroy(r); + } + + 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 + } + + /// 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 { + if (r.aborted.load(.acquire)) return error.Aborted; + if (n > r.slot_bytes) 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, + }; + var params = [1]?*anyopaque{@ptrCast(&a)}; + const blocks = @min(@import("mailbox.zig").blocksFor(n), 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 (lens == 0) return error.Invalid; + return r.launch(.gather, s, o, n, lens, 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); + } + + 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 { + const n = in.len; + if (n == 0 or n > r.slot_bytes or out.len < 2 * n) return error.Invalid; + if (lens) |l| if (l[0] > n or l[1] > n) return error.Invalid; + const mine: usize = if (lens) |l| l[r.me] else n; + const theirs: usize = if (lens) |l| l[1 - r.me] else n; + const seq = r.host_epoch +% 1; + const slot: usize = seq & 1; + @memcpy(r.region[send_off + slot * r.slot_bytes ..][0..mine], in[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 ..][0..mine], in[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 ..][0..theirs], r.region[send_off + 2 * r.slot_bytes + slot * r.slot_bytes ..][0..theirs]); + r.host_epoch = seq; + } +}; + +/// One shared mapping holding both ranks' regions (rank 0 creates and sizes the file before anyone learns its name). +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; +} + +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_host.zig b/zig/src/fabric/tp/roce_host.zig new file mode 100644 index 0000000000..afb705d41b --- /dev/null +++ b/zig/src/fabric/tp/roce_host.zig @@ -0,0 +1,58 @@ +//! `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_hca) |text| n = try rv.parseSpecs(text, &specs) else { + var v = try tp.Verbs.open(null); + defer v.close(); + n = rv.autoSpecs(&v, &specs, cfg.roce_hcas); + } + 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_test.zig b/zig/src/fabric/tp/roce_test.zig new file mode 100644 index 0000000000..51e9d79133 --- /dev/null +++ b/zig/src/fabric/tp/roce_test.zig @@ -0,0 +1,116 @@ +//! 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 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, +}; + +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 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 }) 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) % 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) .{ lens[0], lens[1] } else .{ n, n }; + 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 "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 }; + try twoRanks(&run2); + 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..bc99681b5e --- /dev/null +++ b/zig/src/fabric/tp/roce_verbs.zig @@ -0,0 +1,297 @@ +//! 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; + +pub const max_hcas = 2; +/// Work requests a queue pair holds; at most a quarter outstanding (two requests a stripe). +const send_depth = 256; +const flag_stride = 128; +const link_ethernet = 2; +const record_magic: u32 = 0x5446_5244 + 1; // "TFRD", bumped once for this record's layout + +pub const Error = error{ NoHca, VerbsFailed, PostFailed, Completion, BadRecord }; + +/// 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, + 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, +}; + +/// 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; +} + +/// Every device whose port 1 is an active Ethernet port with an IPv4 RoCE v2 GID, by name order, at most `limit`. +pub fn autoSpecs(v: *const Verbs, out: *[max_hcas]HcaSpec, limit: usize) 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 == limit) 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; + if (port.state != abi.port_active or port.link_layer != link_ethernet) continue; + const g = detectGid(v, ctx, name, port.gid_tbl_len) orelse continue; + out[n] = .{ .name_len = @intCast(@min(name.len, 63)), .gid = g }; + @memcpy(out[n].name[0..out[n].name_len], name[0..out[n].name_len]); + 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, + completed: 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); + 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(); + } + + pub fn record(w: *const VerbsWire) Record { + var r: Record = .{ .n_hca = w.n_hca, .region = @intFromPtr(w.region.ptr) }; + 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) 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 { + const packs = padded / 16; + var offset: u64 = 0; + var seq_copy = seq; + for (w.hcas[0..w.n_hca], 0..) |*h, i| { + while (h.outstanding >= send_depth / 4) try w.reapHca(h); + var stripe = packs / w.n_hca; + if (i < packs % w.n_hca) stripe += 1; + const bytes: u32 = stripe * 16; + 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 + w.flag_off + (@as(u64, slot) * max_hcas + i) * flag_stride, .rkey = w.peer_rkey[i] }; + var data_sge: abi.Sge = .{ .addr = @intFromPtr(w.region.ptr) + w.send_off + slot * w.slot_bytes + offset, .length = bytes, .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 + w.recv_off + slot * w.slot_bytes + offset, .rkey = w.peer_rkey[i] }; + var bad: ?*abi.SendWr = null; + const first: *abi.SendWr = if (bytes > 0) &data_wr else &flag_wr; + if (h.ctx.ops.post_send(h.qp, first, &bad) != 0) return error.PostFailed; + h.outstanding += 1; + offset += bytes; + } + w.posted += 1; + } + + 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; + }; + h.outstanding -= @intCast(n); + w.completed += @intCast(n); + } + + /// 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)); +} diff --git a/zig/src/fabric/tp/run_hosts.sh b/zig/src/fabric/tp/run_hosts.sh new file mode 100644 index 0000000000..17a4ccbe91 --- /dev/null +++ b/zig/src/fabric/tp/run_hosts.sh @@ -0,0 +1,24 @@ +#!/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 +ulimit -l unlimited 2> /dev/null || echo "warning: memlock is $(ulimit -l) KiB: ibv_reg_mr of the RoCE region may fail" >&2 +{ hostname; nvidia-smi -L; ibv_devices 2>&1; ulimit -l; } > "$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/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..9094429c29 --- /dev/null +++ b/zig/src/fabric/tp/tp.zig @@ -0,0 +1,212 @@ +//! 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 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_FAILFAST, TF_TP_ROCE_WIRE / _FALLBACK)"); + 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) { + if (cfg.roce_hca) |text| { + n_specs = try roce_verbs.parseSpecs(text, &specs); + } else if (Verbs.open(null)) |vb| { + var v = vb; + defer v.close(); + n_specs = roce_verbs.autoSpecs(&v, &specs, cfg.roce_hcas); + } else |_| {} + } + 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; + s.route = .{ .small = .{ .roce = r }, .nccl = &s.nccl, .max_bytes = cfg.mailbox_max_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 { + 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; + _ = planlink; + _ = @import("roce_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 1da50d386d..f3af36587b 100644 --- a/zig/src/fabric/verbs.zig +++ b/zig/src/fabric/verbs.zig @@ -1,7 +1,7 @@ //! The system's verbs library by dlopen (Apple's librdma, Linux libibverbs): every entry point resolved by name. const std = @import("std"); 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; From 6b5404dcdefc6e2ef91997dabc72dd304646b27f Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Sat, 10 Oct 2026 14:24:56 +0700 Subject: [PATCH 007/104] TP tests: a drain on rank 0 then its clean stop ends rank 1's loop at the stop word with no fail-fast exit; a rank killed mid-drain still fails the other fast. --- zig/src/fabric/tp/fate_test.zig | 95 +++++++++++++++++++++++++++++++++ 1 file changed, 95 insertions(+) diff --git a/zig/src/fabric/tp/fate_test.zig b/zig/src/fabric/tp/fate_test.zig index 6019055f50..728f4bfdf3 100644 --- a/zig/src/fabric/tp/fate_test.zig +++ b/zig/src/fabric/tp/fate_test.zig @@ -144,3 +144,98 @@ test "rank 0's clean stop ends every rank with no fail-fast exit, whatever order 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)); +} From 06c4e75f97a7aee994f71abe2074d219ec65cff2 Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Sun, 11 Oct 2026 18:09:44 +0700 Subject: [PATCH 008/104] TP RoCE: with no TF_TP_ROCE_HCA, each rank picks the ports on an IPv4 /24 the peer also has The automatic pick took the first N active RoCE ports in device-list order and never checked that the peer can reach them. On a host with two neighbours, one picked port can sit on a subnet the peer does not have, and the queue pair's RTR step then fails, which reports RoCE unavailable. Every rank now sends its usable ports' IPv4 addresses over the bootstrap. Each rank then keeps the ports whose /24 the peer also lists, ordered by subnet and each /24 no more often than both sides have it, so both ranks pick the same number. With no shared /24 the old first-N order is kept, with a warning to set TF_TP_ROCE_HCA. The exchange runs on every rank, including ranks that set TF_TP_ROCE_HCA, so the bootstrap stays in step. --- zig/src/fabric/tp/config.zig | 2 +- zig/src/fabric/tp/roce_host.zig | 6 +- zig/src/fabric/tp/roce_pick.zig | 186 +++++++++++++++++++++++++++++++ zig/src/fabric/tp/roce_verbs.zig | 23 ++-- zig/src/fabric/tp/tp.zig | 12 +- 5 files changed, 207 insertions(+), 22 deletions(-) create mode 100644 zig/src/fabric/tp/roce_pick.zig diff --git a/zig/src/fabric/tp/config.zig b/zig/src/fabric/tp/config.zig index 13647f5921..ddc0c2a1d6 100644 --- a/zig/src/fabric/tp/config.zig +++ b/zig/src/fabric/tp/config.zig @@ -42,7 +42,7 @@ pub const Config = struct { /// 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: every active Ethernet port with an IPv4 RoCE v2 GID). + /// `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, diff --git a/zig/src/fabric/tp/roce_host.zig b/zig/src/fabric/tp/roce_host.zig index afb705d41b..1dec07e9ba 100644 --- a/zig/src/fabric/tp/roce_host.zig +++ b/zig/src/fabric/tp/roce_host.zig @@ -11,11 +11,7 @@ pub fn run(cfg: tp.Config, gpa: std.mem.Allocator) !u8 { defer b.close(); var specs: [rv.max_hcas]rv.HcaSpec = undefined; var n: usize = 0; - if (cfg.roce_hca) |text| n = try rv.parseSpecs(text, &specs) else { - var v = try tp.Verbs.open(null); - defer v.close(); - n = rv.autoSpecs(&v, &specs, cfg.roce_hcas); - } + 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); diff --git a/zig/src/fabric/tp/roce_pick.zig b/zig/src/fabric/tp/roce_pick.zig new file mode 100644 index 0000000000..627d0aefe1 --- /dev/null +++ b/zig/src/fabric/tp/roce_pick.zig @@ -0,0 +1,186 @@ +//! 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 Config = @import("config.zig").Config; +const Verbs = @import("verbs").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; + +pub const Pick = struct { n: usize, shared: bool }; + +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 { + var ports: [max_ports]rv.Port = undefined; + var n_ports: usize = 0; + if (Verbs.open(null)) |opened| { + var v = opened; + defer v.close(); + n_ports = rv.scanPorts(&v, &ports); + } else |_| {} + 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[0..n_ports], &mine), &buf, &blobs); + if (cfg.roce_hca) |text| return rv.parseSpecs(text, out); + 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[0..n_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.n; +} + +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 { + var w: std.Io.Writer = .fixed(buf); + for (specs, 0..) |s, i| w.print("{s}{s}", .{ if (i == 0) "" else ",", s.name[0..s.name_len] }) catch unreachable; + return w.buffered(); +} + +// 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); +} diff --git a/zig/src/fabric/tp/roce_verbs.zig b/zig/src/fabric/tp/roce_verbs.zig index bc99681b5e..a540a4ce02 100644 --- a/zig/src/fabric/tp/roce_verbs.zig +++ b/zig/src/fabric/tp/roce_verbs.zig @@ -89,23 +89,32 @@ fn detectGid(v: *const Verbs, ctx: *abi.Context, name: []const u8, table: c_int) return null; } -/// Every device whose port 1 is an active Ethernet port with an IPv4 RoCE v2 GID, by name order, at most `limit`. -pub fn autoSpecs(v: *const Verbs, out: *[max_hcas]HcaSpec, limit: usize) usize { +/// 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 == limit) break; + 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; - if (port.state != abi.port_active or port.link_layer != link_ethernet) continue; - const g = detectGid(v, ctx, name, port.gid_tbl_len) orelse continue; - out[n] = .{ .name_len = @intCast(@min(name.len, 63)), .gid = g }; - @memcpy(out[n].name[0..out[n].name_len], name[0..out[n].name_len]); + 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; diff --git a/zig/src/fabric/tp/tp.zig b/zig/src/fabric/tp/tp.zig index 9094429c29..3fe6d3fab5 100644 --- a/zig/src/fabric/tp/tp.zig +++ b/zig/src/fabric/tp/tp.zig @@ -15,6 +15,7 @@ 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; @@ -79,15 +80,7 @@ pub const Session = struct { const cfg = s.cfg; var specs: [roce_verbs.max_hcas]roce_verbs.HcaSpec = undefined; var n_specs: usize = 0; - if (cfg.roce_wire == .verbs) { - if (cfg.roce_hca) |text| { - n_specs = try roce_verbs.parseSpecs(text, &specs); - } else if (Verbs.open(null)) |vb| { - var v = vb; - defer v.close(); - n_specs = roce_verbs.autoSpecs(&v, &specs, cfg.roce_hcas); - } else |_| {} - } + 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; @@ -205,6 +198,7 @@ test { _ = bootstrap; _ = mailbox; _ = roce_verbs; + _ = roce_pick; _ = planlink; _ = @import("roce_test.zig"); _ = @import("host_test.zig"); From 062af8b9e0c8e8a8c0299b760920839c72d4fa25 Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Sun, 11 Oct 2026 18:09:44 +0700 Subject: [PATCH 009/104] TP README: when -Dtp-nvcc, TF_TP_ROCE_HCA, NCCL_IB_HCA, NCCL_SOCKET_IFNAME and TF_TP_MASTER are needed, with a two-host example --- zig/src/fabric/tp/README.md | 53 +++++++++++++++++++++++++++++++++++++ 1 file changed, 53 insertions(+) create mode 100644 zig/src/fabric/tp/README.md diff --git a/zig/src/fabric/tp/README.md b/zig/src/fabric/tp/README.md new file mode 100644 index 0000000000..5d5c57358f --- /dev/null +++ b/zig/src/fabric/tp/README.md @@ -0,0 +1,53 @@ +# 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. + +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). From d7510ef8834b9b988e14891ca20b6da81aa69862 Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Sun, 11 Oct 2026 18:09:44 +0700 Subject: [PATCH 010/104] TP test runner: a binary built without the mailbox kernel refuses the kernel backends up front; run_hosts.sh passes HCA picks - `tf-tp-test kernel` reports the embedded kernel and exits 1 without one. `suite` and `failfast` on the mailbox or RoCE backend stop at once with a message that names `-Dtp-nvcc`, instead of failing later on a missing kernel. - `run_hosts.sh` stops before any step when the kernel is missing. It exports TF_TP_ROCE_HCA and NCCL_IB_HCA (derived from TF_TP_ROCE_HCA when unset), warns when a host has more than one RDMA device and neither is set, and records all three variables in host.txt. --- zig/src/fabric/tp/gpu_main.zig | 13 +++++++++++-- zig/src/fabric/tp/run_hosts.sh | 11 ++++++++++- 2 files changed, 21 insertions(+), 3 deletions(-) diff --git a/zig/src/fabric/tp/gpu_main.zig b/zig/src/fabric/tp/gpu_main.zig index c5668a7b06..e3b9a420e7 100644 --- a/zig/src/fabric/tp/gpu_main.zig +++ b/zig/src/fabric/tp/gpu_main.zig @@ -1,4 +1,4 @@ -//! `tf-tp-test suite | failfast RANK fatal|kill | roce-host`: one rank of a two-process TP=2 test on real GPUs. +//! `tf-tp-test suite | failfast RANK fatal|kill | roce-host | kernel`: one rank of a two-process TP=2 test. const std = @import("std"); const cuda = @import("cuda"); const tp = @import("tp"); @@ -11,11 +11,20 @@ const max_bytes = (4 << 20) + 4096; pub fn main(init: std.process.Init) !u8 { const args = try init.minimal.args.toSlice(init.arena.allocator()); if (args.len < 2) { - std.debug.print("usage: tf-tp-test suite | failfast RANK fatal|kill | roce-host\n", .{}); + std.debug.print("usage: tf-tp-test suite | failfast RANK fatal|kill | roce-host | kernel\n", .{}); return 2; } + // `kernel`: whether the build embeds the mailbox kernel (-Dtp-nvcc or -Dtp-fatbin), for scripts to check first + if (std.mem.eql(u8, args[1], "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 (std.mem.eql(u8, args[1], "roce-host")) return @import("roce_host.zig").run(cfg, init.gpa); + 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)); diff --git a/zig/src/fabric/tp/run_hosts.sh b/zig/src/fabric/tp/run_hosts.sh index 17a4ccbe91..5f00168f83 100644 --- a/zig/src/fabric/tp/run_hosts.sh +++ b/zig/src/fabric/tp/run_hosts.sh @@ -4,8 +4,17 @@ 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 +[ -n "${NCCL_IB_HCA:-}" ] && export NCCL_IB_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 -{ hostname; nvidia-smi -L; ibv_devices 2>&1; ulimit -l; } > "$OUT/host.txt" 2>&1 +DEVICES=$(ibv_devices 2> /dev/null | tail -n +3 | grep -c .) +if [ "$DEVICES" -gt 1 ] && [ -z "${TF_TP_ROCE_HCA:-}${NCCL_IB_HCA:-}" ]; then + echo "warning: $DEVICES RDMA devices and no TF_TP_ROCE_HCA / NCCL_IB_HCA: name the peer's ports (README.md)" >&2 +fi +{ 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 From e16e331a087bb6e60d0d04854a4cdcc7c3f02ed0 Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Sun, 11 Oct 2026 18:39:14 +0700 Subject: [PATCH 011/104] TP RoCE: every RDMA post is bounds-checked before the verb, and peers with another region layout are refused roce_guard.stripes computes each HCA's share of a post (its send stripe, the peer's receive range and its flag word) and refuses with error.OutOfRegion any range past its slot, its flag area or the registered region, before post_send. The verbs wire and the shm wire both post through it. The shm wire moves to roce_shm.zig and can stripe like the verbs wire (Options.shm_stripes), with an optional tap that sees every post for tests. The connection record carries the slot and registered length, so a peer with another layout is refused at connect; its magic is bumped. Tests: stripes tile the slot over two HCAs and refuse a post past the slot, a third slot, an unpadded or empty post, no HCA and a short region. On fake verbs a post past its slot reaches no verb, and a full slot stays inside the region. Over the shm wire striped as two HCAs, 2,000 variable-length gathers stay inside slot and region, bit for bit. --- zig/src/fabric/tp/roce.zig | 60 ++++------------- zig/src/fabric/tp/roce_fake.zig | 96 ++++++++++++++++++++++++++++ zig/src/fabric/tp/roce_guard.zig | 63 ++++++++++++++++++ zig/src/fabric/tp/roce_post_test.zig | 20 ++++++ zig/src/fabric/tp/roce_shm.zig | 63 ++++++++++++++++++ zig/src/fabric/tp/roce_test.zig | 41 +++++++++++- zig/src/fabric/tp/roce_verbs.zig | 44 ++++++++----- zig/src/fabric/tp/tp.zig | 2 + 8 files changed, 323 insertions(+), 66 deletions(-) create mode 100644 zig/src/fabric/tp/roce_fake.zig create mode 100644 zig/src/fabric/tp/roce_guard.zig create mode 100644 zig/src/fabric/tp/roce_post_test.zig create mode 100644 zig/src/fabric/tp/roce_shm.zig diff --git a/zig/src/fabric/tp/roce.zig b/zig/src/fabric/tp/roce.zig index 58f88e3441..7806fd60d5 100644 --- a/zig/src/fabric/tp/roce.zig +++ b/zig/src/fabric/tp/roce.zig @@ -6,6 +6,8 @@ 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; @@ -71,30 +73,20 @@ pub const Options = struct { fast: u32 = 0, /// TF_TP_ROCE_BLOCKS: the largest grid (config.roce_blocks) blocks: u32 = 16, -}; - -/// The shm wire: the peer's region is mapped here; a post copies the slot, then stores the flag with release order. -const ShmWire = struct { - peer: []u8, - local: []u8, - slot_bytes: usize, - - fn post(w: *ShmWire, slot: u32, seq: u32, padded: u32) void { - const recv = send_off + 2 * w.slot_bytes; - @memcpy(w.peer[recv + slot * w.slot_bytes ..][0..padded], w.local[send_off + slot * w.slot_bytes ..][0..padded]); - const flag: *std.atomic.Value(u32) = @ptrCast(@alignCast(w.peer[flag_off + slot * 2 * flag_stride ..].ptr)); - flag.store(seq, .release); - } + /// 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: ShmWire, + 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| s.post(slot, seq, padded), + .shm => |*s| try s.post(slot, seq, padded), } } @@ -108,7 +100,7 @@ const Wire = union(WireKind) { fn hcas(w: *const Wire) u32 { return switch (w.*) { .verbs => |v| v.n_hca, - .shm => 1, + .shm => |s| s.stripes, }; } }; @@ -204,9 +196,10 @@ pub const Roce = struct { 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 sharedPair(b, len); + r.mapping = try shm.sharedPair(b, len); r.region = @alignCast(r.mapping[r.me * len ..][0..len]); - r.wire = .{ .shm = .{ .peer = r.mapping[(1 - r.me) * len ..][0..len], .local = r.region, .slot_bytes = r.slot_bytes } }; + 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; @@ -422,35 +415,6 @@ pub const Roce = struct { } }; -/// One shared mapping holding both ranks' regions (rank 0 creates and sizes the file before anyone learns its name). -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; -} - 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)); diff --git a/zig/src/fabric/tp/roce_fake.zig b/zig/src/fabric/tp/roce_fake.zig new file mode 100644 index 0000000000..e556c1990a --- /dev/null +++ b/zig/src/fabric/tp/roce_fake.zig @@ -0,0 +1,96 @@ +//! 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 = &.{}; + 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) 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; + } +}; + +/// 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 { + var w: rv.VerbsWire = .{ .v = undefined, .traffic_class = 0, .recv_off = send_off + 2 * slot, .flag_off = 4096, .send_off = send_off, .slot_bytes = slot, .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..6851f9a6a9 --- /dev/null +++ b/zig/src/fabric/tp/roce_guard.zig @@ -0,0 +1,63 @@ +//! 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}; + +/// 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 == 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 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_post_test.zig b/zig/src/fabric/tp/roce_post_test.zig new file mode 100644 index 0000000000..fe931e1e2d --- /dev/null +++ b/zig/src/fabric/tp/roce_post_test.zig @@ -0,0 +1,20 @@ +//! 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 Fake = fake.Fake; + +test "a post past its slot is refused before any verb; a full slot posts inside the region" { + const region = try std.testing.allocator.alignedAlloc(u8, .fromByteUnits(4096), Roce.regionBytes(fake.slot)); + defer std.testing.allocator.free(region); + Fake.region = region; + Fake.reset(1); + var hw: fake.Hw = .{}; + var w = hw.wire(region); + try std.testing.expectError(error.OutOfRegion, w.post(1, 1, fake.slot + 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, fake.slot); + try std.testing.expectEqual(@as(u64, 2), Fake.posts); + try std.testing.expectEqual(@as(u64, 0), Fake.outside); +} 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 index 51e9d79133..d8cf3fc03f 100644 --- a/zig/src/fabric/tp/roce_test.zig +++ b/zig/src/fabric/tp/roce_test.zig @@ -2,6 +2,8 @@ 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"); @@ -21,6 +23,30 @@ const Run = struct { 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, + /// 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 { @@ -32,7 +58,9 @@ fn body(run: *Run, me: u32, lfd: sock.Fd) ?anyerror { defer b.close(); var spec: [1]rv.HcaSpec = .{.{}}; if (run.wire == .verbs) _ = rv.parseSpecs(run.devs[me], @ptrCast(&spec)) catch |e| return e; - 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 }) 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(); @@ -94,6 +122,17 @@ test "the RoCE protocol with variable lengths over the shm wire: each rank posts 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 "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); diff --git a/zig/src/fabric/tp/roce_verbs.zig b/zig/src/fabric/tp/roce_verbs.zig index a540a4ce02..2502aa0757 100644 --- a/zig/src/fabric/tp/roce_verbs.zig +++ b/zig/src/fabric/tp/roce_verbs.zig @@ -4,15 +4,15 @@ const log = @import("log.zig"); const verbs = @import("verbs"); const abi = verbs.abi; const Verbs = verbs.Verbs; +const guard = @import("roce_guard.zig"); -pub const max_hcas = 2; +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 flag_stride = 128; const link_ethernet = 2; -const record_magic: u32 = 0x5446_5244 + 1; // "TFRD", bumped once for this record's layout +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 }; +pub const Error = error{ NoHca, VerbsFailed, PostFailed, Completion, BadRecord } || guard.Error; /// 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 }; @@ -22,6 +22,9 @@ 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), @@ -195,7 +198,7 @@ pub const VerbsWire = struct { } pub fn record(w: *const VerbsWire) Record { - var r: Record = .{ .n_hca = w.n_hca, .region = @intFromPtr(w.region.ptr) }; + 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; @@ -209,7 +212,7 @@ pub const VerbsWire = struct { /// 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) 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); @@ -242,27 +245,28 @@ pub const VerbsWire = struct { /// 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 { - const packs = padded / 16; - var offset: u64 = 0; + var plan: [max_hcas]guard.Stripe = undefined; + try guard.stripes(w.regionLayout(), w.n_hca, slot, padded, &plan); var seq_copy = seq; - for (w.hcas[0..w.n_hca], 0..) |*h, i| { + for (w.hcas[0..w.n_hca], plan[0..w.n_hca], 0..) |*h, s, i| { while (h.outstanding >= send_depth / 4) try w.reapHca(h); - var stripe = packs / w.n_hca; - if (i < packs % w.n_hca) stripe += 1; - const bytes: u32 = stripe * 16; 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 + w.flag_off + (@as(u64, slot) * max_hcas + i) * flag_stride, .rkey = w.peer_rkey[i] }; - var data_sge: abi.Sge = .{ .addr = @intFromPtr(w.region.ptr) + w.send_off + slot * w.slot_bytes + offset, .length = bytes, .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 + w.recv_off + slot * w.slot_bytes + offset, .rkey = w.peer_rkey[i] }; + 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 (bytes > 0) &data_wr else &flag_wr; + 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.outstanding += 1; - offset += bytes; } 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); @@ -303,4 +307,10 @@ test "HCA lists parse and subnets pair the functions" { 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/tp.zig b/zig/src/fabric/tp/tp.zig index 3fe6d3fab5..8e5d5b0d34 100644 --- a/zig/src/fabric/tp/tp.zig +++ b/zig/src/fabric/tp/tp.zig @@ -200,7 +200,9 @@ test { _ = roce_verbs; _ = roce_pick; _ = planlink; + _ = @import("roce_guard.zig"); _ = @import("roce_test.zig"); + _ = @import("roce_post_test.zig"); _ = @import("host_test.zig"); _ = @import("fate_test.zig"); } From 649e78800076019d521213b00048d80175c401b7 Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Sun, 11 Oct 2026 19:03:59 +0700 Subject: [PATCH 012/104] TP RoCE: no RDMA write past 1 MiB, and boot refuses any setting that would ask for one A post of about 9.6 MB faulted a DGX Spark's HCA. roce_guard.max_write (1 MiB, a constant) is checked in stripes() before the slot bounds, so the verbs wire, the shm wire and the proxy refuse a larger post with error.WriteTooLarge before any verb, whatever the slot size. Roce.init refuses a slot past it. Config.parse fails boot with a clear message when TF_TP_MAILBOX_MAX_KB would make the RoCE slot pass 1 MiB; the value is never clamped. Tests: a 9.57 MB and a 1 MiB + 16 post are refused in a 10 MiB slot and reach no verb, a 1 MiB post goes out; boot refuses TF_TP_MAILBOX_MAX_KB 1028, 2048 and 16384 and keeps 1024. --- zig/src/fabric/tp/config.zig | 7 ++++ zig/src/fabric/tp/roce.zig | 1 + zig/src/fabric/tp/roce_fake.zig | 7 ++-- zig/src/fabric/tp/roce_guard.zig | 12 ++++++- zig/src/fabric/tp/roce_post_test.zig | 48 ++++++++++++++++++++++++++-- 5 files changed, 69 insertions(+), 6 deletions(-) diff --git a/zig/src/fabric/tp/config.zig b/zig/src/fabric/tp/config.zig index ddc0c2a1d6..d823297743 100644 --- a/zig/src/fabric/tp/config.zig +++ b/zig/src/fabric/tp/config.zig @@ -16,6 +16,9 @@ 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 }; @@ -76,6 +79,10 @@ pub const Config = struct { 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.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); diff --git a/zig/src/fabric/tp/roce.zig b/zig/src/fabric/tp/roce.zig index 7806fd60d5..e3695929f1 100644 --- a/zig/src/fabric/tp/roce.zig +++ b/zig/src/fabric/tp/roce.zig @@ -142,6 +142,7 @@ pub const Roce = struct { /// 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); diff --git a/zig/src/fabric/tp/roce_fake.zig b/zig/src/fabric/tp/roce_fake.zig index e556c1990a..a6b02db77d 100644 --- a/zig/src/fabric/tp/roce_fake.zig +++ b/zig/src/fabric/tp/roce_fake.zig @@ -10,6 +10,8 @@ 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; @@ -44,7 +46,7 @@ pub const Fake = struct { 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) outside += 1; + 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; } @@ -82,7 +84,8 @@ pub const Hw = struct { mr: [2]abi.Mr = undefined, pub fn wire(hw: *Hw, region: []u8) rv.VerbsWire { - var w: rv.VerbsWire = .{ .v = undefined, .traffic_class = 0, .recv_off = send_off + 2 * slot, .flag_off = 4096, .send_off = send_off, .slot_bytes = slot, .region = region }; + 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); diff --git a/zig/src/fabric/tp/roce_guard.zig b/zig/src/fabric/tp/roce_guard.zig index 6851f9a6a9..8f85247afd 100644 --- a/zig/src/fabric/tp/roce_guard.zig +++ b/zig/src/fabric/tp/roce_guard.zig @@ -4,7 +4,10 @@ const std = @import("std"); pub const max_hcas = 2; pub const flag_stride = 128; -pub const Error = error{OutOfRegion}; +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 }; @@ -15,6 +18,7 @@ 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; @@ -57,6 +61,12 @@ test "a post past its slot, a third slot, an unpadded or empty post, or no HCA i 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_post_test.zig b/zig/src/fabric/tp/roce_post_test.zig index fe931e1e2d..3e815c2c8f 100644 --- a/zig/src/fabric/tp/roce_post_test.zig +++ b/zig/src/fabric/tp/roce_post_test.zig @@ -2,19 +2,61 @@ 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 region = try std.testing.allocator.alignedAlloc(u8, .fromByteUnits(4096), Roce.regionBytes(fake.slot)); + 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, fake.slot + 16)); + 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, fake.slot); + 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); +} From 664a8e2d57549ab9d040a2c4f9e8ce53e274760c Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Sun, 11 Oct 2026 19:04:52 +0700 Subject: [PATCH 013/104] TP RoCE: TF_TP_ROCE_CHUNKED_KB gathers strides past the slot in slot-sized ops, at most 2 MiB in flight Default off. allGatherV strides past the slot, up to TF_TP_ROCE_CHUNKED_KB, go RoCE as ceil(n / slot) launches of the tp_roce kernel: op c moves at most a slot of each shard from c x slot (RoceArgs.chunk), and an op empty on both ranks rings no doorbell and leaves the epoch. Every RDMA write keeps the 1 MiB cap. The kernel poisons the transport rather than stage past a slot, and its four failure-record blocks fold into fail_record (same stores and fences). VerbsWire.post waits for completions while the bytes posted and not completed over both HCAs would pass max_inflight (2 MiB, a constant in the post path); a wait that sees no completion for 10 s fails the proxy. Boot refuses TF_TP_ROCE_CHUNKED_KB when the kernel image lacks the tp_roce_chunk_abi marker. The setting is part of the digest both ranks agree on. Tests: chunked strides from 300 KB to a 128 MiB union posted back to back on fake verbs whose completions lag keep at most 2 MiB in flight and no write past 1 MiB, and reach the cap; the wire's own in-flight count and peak match the fake HCAs'. Over the shm wire, 2,000 chunked variable-length gathers of up to four slots arrive bit for bit. --- zig/src/fabric/tp/config.zig | 4 ++ zig/src/fabric/tp/hybrid.zig | 24 +++++-- zig/src/fabric/tp/mailbox.cu | 80 +++++++++++++----------- zig/src/fabric/tp/roce.zig | 48 +++++++++++--- zig/src/fabric/tp/roce_inflight_test.zig | 26 ++++++++ zig/src/fabric/tp/roce_post_test.zig | 42 +++++++++++++ zig/src/fabric/tp/roce_test.zig | 21 ++++++- zig/src/fabric/tp/roce_verbs.zig | 41 +++++++++++- zig/src/fabric/tp/tp.zig | 11 +++- 9 files changed, 243 insertions(+), 54 deletions(-) create mode 100644 zig/src/fabric/tp/roce_inflight_test.zig diff --git a/zig/src/fabric/tp/config.zig b/zig/src/fabric/tp/config.zig index d823297743..442f82ed0e 100644 --- a/zig/src/fabric/tp/config.zig +++ b/zig/src/fabric/tp/config.zig @@ -37,6 +37,8 @@ pub const Config = struct { 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, @@ -83,6 +85,7 @@ pub const Config = struct { 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); @@ -141,6 +144,7 @@ pub const Config = struct { 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)); diff --git a/zig/src/fabric/tp/hybrid.zig b/zig/src/fabric/tp/hybrid.zig index 2cdeb0e4a4..e5f6d72c31 100644 --- a/zig/src/fabric/tp/hybrid.zig +++ b/zig/src/fabric/tp/hybrid.zig @@ -51,11 +51,18 @@ pub const Small = union(enum) { } }; +/// 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, @@ -117,19 +124,28 @@ pub const Hybrid = struct { 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 (varFits(ptr, n)) return h.small.roce.allGatherV(s, r, n, lens, stream); + 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); } - /// RoCE with its device side, for a stride the small path takes (the route is the size's alone: every rank alike). - pub fn varFits(ptr: *anyopaque, n: usize) bool { - const h = self_(ptr); + 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; } diff --git a/zig/src/fabric/tp/mailbox.cu b/zig/src/fabric/tp/mailbox.cu index 63d53b2484..843b4f1af2 100644 --- a/zig/src/fabric/tp/mailbox.cu +++ b/zig/src/fabric/tp/mailbox.cu @@ -270,6 +270,21 @@ __device__ __noinline__ bool wait_flag_shared(const uint32_t *flag, uint32_t seq } } +// 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; @@ -313,17 +328,7 @@ extern "C" __global__ void __launch_bounds__(512) tp_mailbox(MailArgs a) { 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; - if (atom_cas_gpu(a.state + 2, 0u, 1u) == 0u) { // one failure record a rank - uint32_t *rec = a.ctrl + kCtrlRank + 8u * a.rank; - st_relaxed_sys(rec + 1, seq); - st_relaxed_sys(rec + 2, seen); - st_relaxed_sys(rec + 3, blockIdx.x); - st_relaxed_sys(rec + 4, why); - st_relaxed_sys(rec + 5, waited); - fence_sc_sys(); - st_relaxed_sys(rec, 1u); - fence_sc_sys(); - } + fail_record(a.state + 2, a.ctrl + kCtrlRank + 8u * a.rank, seq, seen, blockIdx.x, why, waited); } } __syncthreads(); @@ -391,6 +396,7 @@ struct RoceArgs { 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. @@ -411,8 +417,17 @@ __device__ __forceinline__ uint64_t roce_span(const RoceArgs &a, uint64_t 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; @@ -438,12 +453,23 @@ extern "C" __global__ void __launch_bounds__(512) tp_roce(RoceArgs a) { 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 mine = var ? roce_shard(a, a.rank) : a.nbytes; - const uint64_t theirs = var ? roce_shard(a, 1u - a.rank) : a.nbytes; + 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, var ? roce_span(a, mine) : a.nbytes, index, stride); + 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) @@ -468,7 +494,7 @@ extern "C" __global__ void __launch_bounds__(512) tp_roce(RoceArgs a) { } } if (a.mode == kGather) { - copy_any(fast, a.out + static_cast(a.rank) * a.nbytes, a.in, + copy_any(fast, a.out + static_cast(a.rank) * a.nbytes + off, a.in + off, var ? roce_span(a, mine) : a.nbytes, index, stride); } @@ -486,16 +512,7 @@ extern "C" __global__ void __launch_bounds__(512) tp_roce(RoceArgs a) { 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; - if (atom_cas_gpu(a.state + 2, 0u, 1u) == 0u) { - st_relaxed_sys(a.ctrl + kRoceFailed + 1, seq); - st_relaxed_sys(a.ctrl + kRoceFailed + 2, seen); - st_relaxed_sys(a.ctrl + kRoceFailed + 3, lane); - st_relaxed_sys(a.ctrl + kRoceFailed + 4, why); - st_relaxed_sys(a.ctrl + kRoceFailed + 5, waited); - fence_sc_sys(); - st_relaxed_sys(a.ctrl + kRoceFailed, 1u); - fence_sc_sys(); - } + fail_record(a.state + 2, a.ctrl + kRoceFailed, seq, seen, lane, why, waited); } } } else if (tid == 0) { @@ -504,16 +521,7 @@ extern "C" __global__ void __launch_bounds__(512) tp_roce(RoceArgs a) { if (!wait_flag(a.flags + (slot * kRoceHcas + h) * kFlagStride, seq, a.ctrl + kRoceAbort, a.timeout_ns, &seen, &why, &waited)) { s_bad = 1u; - if (atom_cas_gpu(a.state + 2, 0u, 1u) == 0u) { - st_relaxed_sys(a.ctrl + kRoceFailed + 1, seq); - st_relaxed_sys(a.ctrl + kRoceFailed + 2, seen); - st_relaxed_sys(a.ctrl + kRoceFailed + 3, h); - st_relaxed_sys(a.ctrl + kRoceFailed + 4, why); - st_relaxed_sys(a.ctrl + kRoceFailed + 5, waited); - fence_sc_sys(); - st_relaxed_sys(a.ctrl + kRoceFailed, 1u); - fence_sc_sys(); - } + fail_record(a.state + 2, a.ctrl + kRoceFailed, seq, seen, h, why, waited); break; } } @@ -525,7 +533,7 @@ extern "C" __global__ void __launch_bounds__(512) tp_roce(RoceArgs a) { 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, inbox, + 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: diff --git a/zig/src/fabric/tp/roce.zig b/zig/src/fabric/tp/roce.zig index e3695929f1..3e7c1e644a 100644 --- a/zig/src/fabric/tp/roce.zig +++ b/zig/src/fabric/tp/roce.zig @@ -53,6 +53,8 @@ const RoceArgs = extern struct { 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 { @@ -125,6 +127,8 @@ pub const Roce = struct { 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), @@ -233,6 +237,7 @@ pub const Roce = struct { 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; } @@ -311,8 +316,12 @@ pub const Roce = struct { } 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) return error.Unsupported; + 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 = .{ @@ -331,9 +340,10 @@ pub const Roce = struct { .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(n), r.max_grid); + 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; } @@ -347,6 +357,14 @@ pub const Roce = struct { 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); } @@ -391,17 +409,31 @@ pub const Roce = struct { /// `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 n > r.slot_bytes or out.len < 2 * n) return error.Invalid; + 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 mine: usize = if (lens) |l| l[r.me] else n; - const theirs: usize = if (lens) |l| l[1 - r.me] else n; + 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[0..mine]); + @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 ..][0..mine], in[0..mine]); + @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)); @@ -411,7 +443,7 @@ pub const Roce = struct { std.atomic.spinLoopHint(); } } - @memcpy(out[(1 - r.me) * n ..][0..theirs], r.region[send_off + 2 * r.slot_bytes + slot * r.slot_bytes ..][0..theirs]); + @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; } }; 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..a8b83a1564 --- /dev/null +++ b/zig/src/fabric/tp/roce_inflight_test.zig @@ -0,0 +1,26 @@ +//! 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 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)); +} diff --git a/zig/src/fabric/tp/roce_post_test.zig b/zig/src/fabric/tp/roce_post_test.zig index 3e815c2c8f..a4f7971fed 100644 --- a/zig/src/fabric/tp/roce_post_test.zig +++ b/zig/src/fabric/tp/roce_post_test.zig @@ -60,3 +60,45 @@ test "boot refuses a RoCE slot past 1 MiB (TF_TP_MAILBOX_MAX_KB), never clamping 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); + } +} diff --git a/zig/src/fabric/tp/roce_test.zig b/zig/src/fabric/tp/roce_test.zig index d8cf3fc03f..dff1ab7eb9 100644 --- a/zig/src/fabric/tp/roce_test.zig +++ b/zig/src/fabric/tp/roce_test.zig @@ -23,6 +23,8 @@ const Run = struct { 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. @@ -77,12 +79,14 @@ fn body(run: *Run, me: u32, lfd: sock.Fd) ?anyerror { 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) % run.max_bytes; + 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) .{ lens[0], lens[1] } else .{ n, n }; - r.hostGatherV(in[0..n], out[0 .. 2 * n], if (run.var_lens) lens else null) catch |e| return e; + 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; @@ -133,6 +137,17 @@ test "the shm wire striped as over two HCAs: every post inside its slot and the } } +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); diff --git a/zig/src/fabric/tp/roce_verbs.zig b/zig/src/fabric/tp/roce_verbs.zig index 2502aa0757..cfcb2d143d 100644 --- a/zig/src/fabric/tp/roce_verbs.zig +++ b/zig/src/fabric/tp/roce_verbs.zig @@ -5,6 +5,7 @@ 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). @@ -14,6 +15,12 @@ const record_magic: u32 = 0x5446_5244 + 2; // "TFRD", bumped for each change of 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; + /// 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 }; @@ -43,6 +50,10 @@ const Hca = struct { 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, }; /// Parses `name[:gid],name[:gid]` (TF_TP_ROCE_HCA). @@ -137,6 +148,8 @@ pub const VerbsWire = struct { 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, /// Opens the HCAs, registers `region` on each and creates the queue pairs (INIT); `record` then goes to the peer. @@ -247,9 +260,10 @@ pub const VerbsWire = struct { 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 >= send_depth / 4) try w.reapHca(h); + 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 }; @@ -257,8 +271,11 @@ pub const VerbsWire = struct { 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; } + if (w.inflight() > w.peak.load(.monotonic)) w.peak.store(w.inflight(), .monotonic); w.posted += 1; } @@ -275,10 +292,32 @@ pub const VerbsWire = struct { 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; }; + for (0..@intCast(n)) |_| { + h.inflight -= h.ring[h.head]; + h.head = (h.head + 1) % ring_len; + } h.outstanding -= @intCast(n); w.completed += @intCast(n); } + /// 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); diff --git a/zig/src/fabric/tp/tp.zig b/zig/src/fabric/tp/tp.zig index 8e5d5b0d34..901544b58b 100644 --- a/zig/src/fabric/tp/tp.zig +++ b/zig/src/fabric/tp/tp.zig @@ -41,7 +41,7 @@ pub const Session = struct { 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_FAILFAST, TF_TP_ROCE_WIRE / _FALLBACK)"); + 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)"); var lib = try cuda.nccl.Library.openPath(cfg.nccl_lib); errdefer lib.close(); const v = try lib.version(); @@ -84,7 +84,13 @@ pub const Session = struct { 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; - s.route = .{ .small = .{ .roce = r }, .nccl = &s.nccl, .max_bytes = cfg.mailbox_max_bytes }; + 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; @@ -203,6 +209,7 @@ test { _ = @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"); } From 4ad5029960fb26959b8787833aa04c39a131afaf Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Sun, 11 Oct 2026 19:05:28 +0700 Subject: [PATCH 014/104] TP RoCE: a failed transport halts its QPs, drains them and keeps its region; the region is DONTFORK and NOHUGEPAGE A local protection error can arrive while the IOMMU invalidates its mappings lazily. On such an error the old path left the other HCA's QP ready to send (still taking the peer's writes), hung up the fate channel without telling the peer to stop, and freed the region at once: unregistered from CUDA with no GPU sync, then QP destroy, MR deregistration and munmap. A device write after the unmap could then reach a reused page through a stale IOTLB entry. - Roce.fail and Roce.abort (the fail-fast hook) move every QP to the error state: no send goes out and no peer write lands. - Roce.deinit quiesces first: QPs to error, every posted request reaped (flushed or not), cuCtxSynchronize. A failed transport, or one that does not quiesce, keeps its region registered and mapped; the process exit tears it down. - Session.close on a failed transport goes fail-fast before closing, so the peer aborts and halts its HCAs first. - The region is MADV_DONTFORK and MADV_NOHUGEPAGE before ibv_reg_mr; a refusal fails the start. - A QP whose actual max_inline_data is under 4 B is refused (the flag post uses lkey 0 on a stack word). - One log line a rank: region base and length, each HCA's lkey and rkey, the peer's region and rkeys. Tests: on fake verbs Roce.abort halts both QPs; quiesce halts both before the first flush and reaps every posted write; an aborted transport reports itself failed, and over the shm wire the aborted rank keeps its region while the other frees it; smaps shows dc and nh on a prepared region; the inline bound refuses 3 B and takes 4. --- zig/src/fabric/tp/roce.zig | 42 ++++++++++++++ zig/src/fabric/tp/roce_fake.zig | 37 ++++++++++++ zig/src/fabric/tp/roce_inflight_test.zig | 74 ++++++++++++++++++++++++ zig/src/fabric/tp/roce_post_test.zig | 16 +++++ zig/src/fabric/tp/roce_test.zig | 3 + zig/src/fabric/tp/roce_verbs.zig | 66 ++++++++++++++++++++- zig/src/fabric/tp/tp.zig | 2 + 7 files changed, 237 insertions(+), 3 deletions(-) diff --git a/zig/src/fabric/tp/roce.zig b/zig/src/fabric/tp/roce.zig index 3e7c1e644a..c8e8010810 100644 --- a/zig/src/fabric/tp/roce.zig +++ b/zig/src/fabric/tp/roce.zig @@ -198,6 +198,7 @@ pub const Roce = struct { 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 => { @@ -222,6 +223,8 @@ pub const Roce = struct { 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; }, } @@ -261,14 +264,37 @@ pub const Roce = struct { 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; + return wire and gpu; + } + fn word(r: *const Roce, i: usize) *std.atomic.Value(u32) { return @ptrCast(@alignCast(r.region[ctrl_off + 4 * i ..].ptr)); } @@ -279,6 +305,21 @@ pub const Roce = struct { 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| _ = v.halt(), + .shm => {}, + } + } + + /// 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. @@ -382,6 +423,7 @@ pub const Roce = struct { 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 { diff --git a/zig/src/fabric/tp/roce_fake.zig b/zig/src/fabric/tp/roce_fake.zig index a6b02db77d..fbb55891f6 100644 --- a/zig/src/fabric/tp/roce_fake.zig +++ b/zig/src/fabric/tp/roce_fake.zig @@ -76,6 +76,43 @@ pub const Fake = struct { } }; +/// 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; + } + + 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; + 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, diff --git a/zig/src/fabric/tp/roce_inflight_test.zig b/zig/src/fabric/tp/roce_inflight_test.zig index a8b83a1564..130f1e1eab 100644 --- a/zig/src/fabric/tp/roce_inflight_test.zig +++ b/zig/src/fabric/tp/roce_inflight_test.zig @@ -4,6 +4,7 @@ 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" { @@ -24,3 +25,76 @@ test "chunked posts: the wire's own in-flight count and peak match the HCAs', an 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 + 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.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); + // 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 starts at `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 head: [32]u8 = undefined; + const want = try std.fmt.bufPrint(&head, "{x}-", .{at}); + var lines = std.mem.splitScalar(u8, buf[0..n], '\n'); + var inside = false; + while (lines.next()) |l| { + if (std.mem.startsWith(u8, l, want)) inside = true; + 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); + const flags = try vmFlags(@intFromPtr(m.ptr), 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_post_test.zig b/zig/src/fabric/tp/roce_post_test.zig index a4f7971fed..a284293000 100644 --- a/zig/src/fabric/tp/roce_post_test.zig +++ b/zig/src/fabric/tp/roce_post_test.zig @@ -102,3 +102,19 @@ test "chunked strides posted back to back: at most 2 MiB in flight over both HCA 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_test.zig b/zig/src/fabric/tp/roce_test.zig index dff1ab7eb9..ec6bc4d65e 100644 --- a/zig/src/fabric/tp/roce_test.zig +++ b/zig/src/fabric/tp/roce_test.zig @@ -154,7 +154,10 @@ test "a peer that stops makes the waiting rank time out, and abort ends a wait a 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]); } diff --git a/zig/src/fabric/tp/roce_verbs.zig b/zig/src/fabric/tp/roce_verbs.zig index cfcb2d143d..8ed5e3eca4 100644 --- a/zig/src/fabric/tp/roce_verbs.zig +++ b/zig/src/fabric/tp/roce_verbs.zig @@ -20,6 +20,25 @@ 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 }; @@ -190,6 +209,10 @@ pub const VerbsWire = struct { 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 }; @@ -210,6 +233,39 @@ pub const VerbsWire = struct { 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; + if (n > 0) w.popped(h, @intCast(n)) else if (sock.nowNs() - t0 > timeout_ns) return false; + }; + return true; + } + 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| { @@ -292,12 +348,16 @@ pub const VerbsWire = struct { 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; }; - for (0..@intCast(n)) |_| { + 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 -= @intCast(n); - w.completed += @intCast(n); + h.outstanding -= n; + w.completed += n; } /// Bytes posted and not yet completed over every HCA. diff --git a/zig/src/fabric/tp/tp.zig b/zig/src/fabric/tp/tp.zig index 901544b58b..6cc5ceb6ec 100644 --- a/zig/src/fabric/tp/tp.zig +++ b/zig/src/fabric/tp/tp.zig @@ -187,6 +187,8 @@ pub const Session = struct { /// 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(); From 9bbc21aca301d9d0a32025b12bdd424b12c68711 Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Sun, 11 Oct 2026 19:06:03 +0700 Subject: [PATCH 015/104] TP RoCE: tf-tp-test failfast V transport checks the failed-transport path on real HCAs, only on an idle link - `tf-tp-test failfast V transport [N [HOLD_S]]`: N clean 4 KiB sums, the ranks meet over NCCL, each waits for its link to go idle (every HCA's posts completed), both agree over NCCL, then the victim takes the proxy's failure path (`Roce.inject`, which is `Roce.fail`) and closes its session: both ranks halt their QPs and exit 70, and the victim's region stays registered. When either rank's link is not idle there is no inject, a clean stop, REFUSED and rc 2. - Its host-only forms (TF_TP_FAILFAST=0, or a hold) are refused unless TF_TP_ALLOW_HOST_ONLY=1 and no RDMA wire is in use (the shm wire, or no RoCE), so they never run on an HCA. - `run_transport.sh MASTER PEER_SSH OUT`, driven from rank 0's host (rank 1 over ssh): prechecks on both hosts (no GPU process, memlock unlimited), then T0 (suite), T1 and T3 (each rank as the victim). After each step it reads both hosts' kernel log since the step and every HCA *err* counter; a hit, a counter that grew, an unexpected rc or a failure line prints STOP: and exits 1. `run_transport.sh t2 OUT` runs the host-only form on one host. - The `halted` and `quiesced` lines carry each QP's state read back with ibv_query_qp (an optional symbol), each HCA's posts and completions, the completions the drain reaped and how many of them were flushes. - `Roce.allGatherV` refuses a stride past the slot (one op; the route never sends one). - The README states the RoCE limits and how to run the check. Tests: on fake verbs the link reads busy with two posts out and idle after the drain, and the report line names both QPs in the error state with every completion a flush. --- zig/src/fabric/tp/README.md | 11 +++ zig/src/fabric/tp/gpu_main.zig | 21 +++- zig/src/fabric/tp/gpu_transport.zig | 71 ++++++++++++++ zig/src/fabric/tp/roce.zig | 30 +++++- zig/src/fabric/tp/roce_fake.zig | 7 ++ zig/src/fabric/tp/roce_inflight_test.zig | 8 +- zig/src/fabric/tp/roce_verbs.zig | 40 ++++++++ zig/src/fabric/tp/run_transport.sh | 120 +++++++++++++++++++++++ zig/src/fabric/verbs.zig | 4 + 9 files changed, 307 insertions(+), 5 deletions(-) create mode 100644 zig/src/fabric/tp/gpu_transport.zig create mode 100755 zig/src/fabric/tp/run_transport.sh diff --git a/zig/src/fabric/tp/README.md b/zig/src/fabric/tp/README.md index 5d5c57358f..a0b4af5342 100644 --- a/zig/src/fabric/tp/README.md +++ b/zig/src/fabric/tp/README.md @@ -51,3 +51,14 @@ zig/src/fabric/tp/run_hosts.sh 1 192.0.2.1 ./zig-out/bin/tf-tp-test out-rank1 `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/gpu_main.zig b/zig/src/fabric/tp/gpu_main.zig index e3b9a420e7..07d7c7afb9 100644 --- a/zig/src/fabric/tp/gpu_main.zig +++ b/zig/src/fabric/tp/gpu_main.zig @@ -1,4 +1,4 @@ -//! `tf-tp-test suite | failfast RANK fatal|kill | roce-host | kernel`: one rank of a two-process TP=2 test. +//! `tf-tp-test suite | failfast RANK fatal|kill|transport [N [HOLD_S]] | roce-host | kernel`: one rank of a TP=2 test. const std = @import("std"); const cuda = @import("cuda"); const tp = @import("tp"); @@ -11,7 +11,7 @@ const max_bytes = (4 << 20) + 4096; pub fn main(init: std.process.Init) !u8 { const args = try init.minimal.args.toSlice(init.arena.allocator()); if (args.len < 2) { - std.debug.print("usage: tf-tp-test suite | failfast RANK fatal|kill | roce-host | kernel\n", .{}); + std.debug.print("usage: tf-tp-test suite | failfast RANK fatal|kill|transport [N [HOLD_S]] | roce-host | kernel\n", .{}); return 2; } // `kernel`: whether the build embeds the mailbox kernel (-Dtp-nvcc or -Dtp-fatbin), for scripts to check first @@ -21,6 +21,10 @@ pub fn main(init: std.process.Init) !u8 { } const cfg = try tp.Config.fromEnv(); if (std.mem.eql(u8, args[1], "roce-host")) return @import("roce_host.zig").run(cfg, init.gpa); + if (std.mem.eql(u8, args[1], "failfast") and (!cfg.failfast or args.len > 5) 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; @@ -36,6 +40,11 @@ pub fn main(init: std.process.Init) !u8 { 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 }); if (std.mem.eql(u8, args[1], "suite")) return suite(&s, &driver, init.gpa, cfg); + if (std.mem.eql(u8, args[1], "failfast") and args.len >= 4 and std.mem.eql(u8, args[3], "transport")) { + const n = if (args.len > 4) try std.fmt.parseInt(u32, args[4], 10) else 50; + const hold = if (args.len > 5) try std.fmt.parseInt(u32, args[5], 10) else 0; + return @import("gpu_transport.zig").run(&s, &driver, init.gpa, cfg, try std.fmt.parseInt(u32, args[2], 10), n, hold); + } if (std.mem.eql(u8, args[1], "failfast") and args.len == 4) return failfast(&s, &driver, init.gpa, cfg, try std.fmt.parseInt(u32, args[2], 10), args[3]); std.debug.print("unknown command {s}\n", .{args[1]}); return 2; @@ -70,7 +79,13 @@ fn suite(s: *tp.Session, d: *const cuda.Driver, gpa: std.mem.Allocator, cfg: tp. return if (c.failures == 0) 0 else 1; } -fn realNow() f64 { +/// 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; 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/roce.zig b/zig/src/fabric/tp/roce.zig index c8e8010810..c9bb72a2c3 100644 --- a/zig/src/fabric/tp/roce.zig +++ b/zig/src/fabric/tp/roce.zig @@ -292,9 +292,19 @@ pub const Roce = struct { .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)); } @@ -311,11 +321,28 @@ pub const Roce = struct { /// 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| _ = v.halt(), + .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; @@ -394,6 +421,7 @@ pub const Roce = struct { /// `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); } diff --git a/zig/src/fabric/tp/roce_fake.zig b/zig/src/fabric/tp/roce_fake.zig index fbb55891f6..c316448ef8 100644 --- a/zig/src/fabric/tp/roce_fake.zig +++ b/zig/src/fabric/tp/roce_fake.zig @@ -93,6 +93,12 @@ pub const Halts = struct { 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; @@ -109,6 +115,7 @@ pub const Halts = struct { /// 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; } }; diff --git a/zig/src/fabric/tp/roce_inflight_test.zig b/zig/src/fabric/tp/roce_inflight_test.zig index 130f1e1eab..6441431794 100644 --- a/zig/src/fabric/tp/roce_inflight_test.zig +++ b/zig/src/fabric/tp/roce_inflight_test.zig @@ -37,16 +37,22 @@ test "a failed transport: both QPs go to error before the first flush and every 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 + // 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 }; diff --git a/zig/src/fabric/tp/roce_verbs.zig b/zig/src/fabric/tp/roce_verbs.zig index 8ed5e3eca4..6917d292be 100644 --- a/zig/src/fabric/tp/roce_verbs.zig +++ b/zig/src/fabric/tp/roce_verbs.zig @@ -73,6 +73,9 @@ const Hca = struct { 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). @@ -170,6 +173,9 @@ pub const VerbsWire = struct { /// 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 { @@ -261,11 +267,43 @@ pub const VerbsWire = struct { 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| { @@ -330,6 +368,7 @@ pub const VerbsWire = struct { 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; @@ -358,6 +397,7 @@ pub const VerbsWire = struct { } h.outstanding -= n; w.completed += n; + _ = h.done.fetchAdd(n, .release); } /// Bytes posted and not yet completed over every HCA. 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/verbs.zig b/zig/src/fabric/verbs.zig index f3af36587b..8ae58532ef 100644 --- a/zig/src/fabric/verbs.zig +++ b/zig/src/fabric/verbs.zig @@ -50,12 +50,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: std.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 = std.DynLib.openZ((path orelse default_path).ptr) catch return error.LibraryUnavailable; @@ -68,6 +71,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; } From cf19042f9d92031d937ea8305990edd17400052a Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Sun, 11 Oct 2026 21:41:15 +0700 Subject: [PATCH 016/104] Fix RoCE test smaps lookup for merged mappings --- zig/src/fabric/tp/roce_inflight_test.zig | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/zig/src/fabric/tp/roce_inflight_test.zig b/zig/src/fabric/tp/roce_inflight_test.zig index 6441431794..fd7be53262 100644 --- a/zig/src/fabric/tp/roce_inflight_test.zig +++ b/zig/src/fabric/tp/roce_inflight_test.zig @@ -67,7 +67,7 @@ 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 = 4 })); } -/// /proc/self/smaps' VmFlags of the mapping that starts at `at`. +/// /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; @@ -78,12 +78,15 @@ fn vmFlags(at: usize, buf: []u8) ![]const u8 { if (got <= 0) break; n += @intCast(got); } - var head: [32]u8 = undefined; - const want = try std.fmt.bufPrint(&head, "{x}-", .{at}); var lines = std.mem.splitScalar(u8, buf[0..n], '\n'); var inside = false; while (lines.next()) |l| { - if (std.mem.startsWith(u8, l, want)) inside = true; + 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; @@ -100,7 +103,8 @@ test "the region is DONTFORK and NOHUGEPAGE before registration" { 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); - const flags = try vmFlags(@intFromPtr(m.ptr), buf); + // 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); } From 989ca7e6578309da169d94f55b9ffd17541b1bf1 Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Mon, 12 Oct 2026 06:30:24 +0700 Subject: [PATCH 017/104] TP RoCE: run_hosts.sh gives NCCL the RoCE picker's ports when neither TF_TP_ROCE_HCA nor NCCL_IB_HCA is set - `tf-tp-test hca` (host only, no GPU): both ranks run the RoCE pick over the bootstrap, the same rule the wire uses, and each prints its own devices as `NCCL_IB_HCA=name,...` (no GIDs). It exits 1 with the reason when the host has no usable port or no port on a /24 the peer shares, so no rank falls back to its first ports for NCCL. - run_hosts.sh runs it when neither variable is set, exports the names for every step and logs each rank's pick; a failed or malformed `hca` stops the runner (exit 3) with the reason; `hca` is capped at TF_TP_PLAN_TIMEOUT + 31 s, so its own reason (no peer joined in time, port in use, the peer closed, a bootstrap error) comes first. An explicit NCCL_IB_HCA, an explicit TF_TP_ROCE_HCA (NCCL_IB_HCA derived from it) or both keep their behaviour; the many-devices warning is gone. Both hosts use the same mode: with one naming its ports and the other not, the latter's `hca` finds no peer and stops after TF_TP_PLAN_TIMEOUT with that reason (README). - roce_pick: `select` / `selectFrom` return how the devices were chosen (`resolve` is unchanged for the wire), `joinNames` and `hcaLine` give the NCCL list. Host tests: the port scan over a stub verbs table (state, Ethernet, the IPv4 RoCE v2 GID index), then both ranks through a real bootstrap on the two hosts #636's receipt reports (modelled, documentation addresses), an explicit rank, and a host without ports. --- zig/src/fabric/tp/README.md | 8 ++ zig/src/fabric/tp/gpu_main.zig | 40 +++++- zig/src/fabric/tp/roce_pick.zig | 216 ++++++++++++++++++++++++++++++-- zig/src/fabric/tp/run_hosts.sh | 26 +++- 4 files changed, 276 insertions(+), 14 deletions(-) diff --git a/zig/src/fabric/tp/README.md b/zig/src/fabric/tp/README.md index a0b4af5342..7ad00ef7d5 100644 --- a/zig/src/fabric/tp/README.md +++ b/zig/src/fabric/tp/README.md @@ -36,6 +36,14 @@ 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: diff --git a/zig/src/fabric/tp/gpu_main.zig b/zig/src/fabric/tp/gpu_main.zig index 07d7c7afb9..82fc1edd5b 100644 --- a/zig/src/fabric/tp/gpu_main.zig +++ b/zig/src/fabric/tp/gpu_main.zig @@ -1,4 +1,4 @@ -//! `tf-tp-test suite | failfast RANK fatal|kill|transport [N [HOLD_S]] | roce-host | kernel`: one rank of a TP=2 test. +//! `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"); @@ -11,7 +11,7 @@ const max_bytes = (4 << 20) + 4096; pub fn main(init: std.process.Init) !u8 { const args = try init.minimal.args.toSlice(init.arena.allocator()); if (args.len < 2) { - std.debug.print("usage: tf-tp-test suite | failfast RANK fatal|kill|transport [N [HOLD_S]] | roce-host | kernel\n", .{}); + std.debug.print("usage: tf-tp-test suite | failfast RANK fatal|kill|transport [N [HOLD_S]] | roce-host | hca | kernel\n", .{}); return 2; } // `kernel`: whether the build embeds the mailbox kernel (-Dtp-nvcc or -Dtp-fatbin), for scripts to check first @@ -21,6 +21,7 @@ pub fn main(init: std.process.Init) !u8 { } const cfg = try tp.Config.fromEnv(); if (std.mem.eql(u8, args[1], "roce-host")) return @import("roce_host.zig").run(cfg, init.gpa); + if (std.mem.eql(u8, args[1], "hca")) return hca(cfg, init.io); if (std.mem.eql(u8, args[1], "failfast") and (!cfg.failfast or args.len > 5) 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; @@ -50,6 +51,41 @@ pub fn main(init: std.process.Init) !u8 { return 2; } +/// `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.NoRoceDevice => "no active Ethernet RDMA port with an IPv4 RoCE v2 GID on this host", + error.NoSharedSubnet => "no RoCE port here shares an IPv4 /24 with the peer's", + }; + std.debug.print("FAIL hca rank {d}: {s}: name the peer's ports in TF_TP_ROCE_HCA or NCCL_IB_HCA\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(); diff --git a/zig/src/fabric/tp/roce_pick.zig b/zig/src/fabric/tp/roce_pick.zig index 627d0aefe1..b9355c401a 100644 --- a/zig/src/fabric/tp/roce_pick.zig +++ b/zig/src/fabric/tp/roce_pick.zig @@ -3,6 +3,7 @@ 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; @@ -10,7 +11,8 @@ const Verbs = @import("verbs").Verbs; pub const max_ports = 8; const blob_bytes = 1 + 4 * max_ports; -pub const Pick = struct { n: usize, shared: bool }; +/// 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 }; fn usable(p: rv.Port) bool { return p.active and p.ethernet and p.ipv4 != null; @@ -83,6 +85,11 @@ fn decode(blob: []const u8, out: [][4]u8) usize { /// 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 { var ports: [max_ports]rv.Port = undefined; var n_ports: usize = 0; if (Verbs.open(null)) |opened| { @@ -90,19 +97,55 @@ pub fn resolve(cfg: Config, b: *bootstrap.Bootstrap, out: *[rv.max_hcas]rv.HcaSp defer v.close(); n_ports = rv.scanPorts(&v, &ports); } else |_| {} + return selectFrom(cfg, b, ports[0..n_ports], out); +} + +/// `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[0..n_ports], &mine), &buf, &blobs); - if (cfg.roce_hca) |text| return rv.parseSpecs(text, out); + 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[0..n_ports], peer[0..n_peer], out, cfg.roce_hcas); + 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.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, +}; + +/// `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) 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 { @@ -114,9 +157,7 @@ fn port(name: []const u8, active: bool, ipv4: ?[4]u8) rv.Port { } fn names(specs: []const rv.HcaSpec, buf: []u8) []const u8 { - var w: std.Io.Writer = .fixed(buf); - for (specs, 0..) |s, i| w.print("{s}{s}", .{ if (i == 0) "" else ",", s.name[0..s.name_len] }) catch unreachable; - return w.buffered(); + 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. @@ -184,3 +225,162 @@ test "two rails on one /24 count once per peer port; no shared /24 keeps the lis 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, 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 = 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 { + 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], &bufs[1], &lines[1], &errs[1] }); + hcaRank(cfgs[0], l.fd, ports[0], &bufs[0], &lines[0], &errs[0]); + t.join(); + for (errs) |e| if (e) |x| return x; +} + +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/run_hosts.sh b/zig/src/fabric/tp/run_hosts.sh index 5f00168f83..75f0c7900b 100644 --- a/zig/src/fabric/tp/run_hosts.sh +++ b/zig/src/fabric/tp/run_hosts.sh @@ -7,13 +7,31 @@ mkdir -p "$OUT" # 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 -[ -n "${NCCL_IB_HCA:-}" ] && export NCCL_IB_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 -DEVICES=$(ibv_devices 2> /dev/null | tail -n +3 | grep -c .) -if [ "$DEVICES" -gt 1 ] && [ -z "${TF_TP_ROCE_HCA:-}${NCCL_IB_HCA:-}" ]; then - echo "warning: $DEVICES RDMA devices and no TF_TP_ROCE_HCA / NCCL_IB_HCA: name the peer's ports (README.md)" >&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 From 05c5213a6934a6f11249a738b0903cec58c7bbb8 Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Sat, 10 Oct 2026 10:42:56 +0700 Subject: [PATCH 018/104] Session memory for paged-KV families: page pool, prefix trie, RAM residency budget, streamed NVMe tier. zig/src/sessions keeps a chat's KV pages after its request ends: entries are found by a trie over chained page digests, held in RAM within a residency budget (chats and pool bytes), and the least recently used idle chat is parked to an NVMe tier and streamed back on resume. Ranks that split KV run every store call in the same order, so their books stay identical. It complements core/prompt_cache, which snapshots a whole family state; a family that uses it reports the resident bytes as Memory.cache and its budget as Info.prompt_cache_plan. No family uses it yet; zig build test-sessions runs the host tests, the NVMe tier on real files. --- build.zig | 2 + zig/build/sessions.zig | 13 + zig/src/sessions/disk.zig | 587 ++++++++++++++++++++++++++++++++ zig/src/sessions/disk_test.zig | 125 +++++++ zig/src/sessions/format.zig | 197 +++++++++++ zig/src/sessions/freeset.zig | 114 +++++++ zig/src/sessions/jobs.zig | 150 ++++++++ zig/src/sessions/lanes.zig | 336 ++++++++++++++++++ zig/src/sessions/pagestore.zig | 129 +++++++ zig/src/sessions/pool.zig | 475 ++++++++++++++++++++++++++ zig/src/sessions/prefix.zig | 457 +++++++++++++++++++++++++ zig/src/sessions/sessions.zig | 33 ++ zig/src/sessions/store.zig | 570 +++++++++++++++++++++++++++++++ zig/src/sessions/store_test.zig | 216 ++++++++++++ 14 files changed, 3404 insertions(+) create mode 100644 zig/build/sessions.zig create mode 100644 zig/src/sessions/disk.zig create mode 100644 zig/src/sessions/disk_test.zig create mode 100644 zig/src/sessions/format.zig create mode 100644 zig/src/sessions/freeset.zig create mode 100644 zig/src/sessions/jobs.zig create mode 100644 zig/src/sessions/lanes.zig create mode 100644 zig/src/sessions/pagestore.zig create mode 100644 zig/src/sessions/pool.zig create mode 100644 zig/src/sessions/prefix.zig create mode 100644 zig/src/sessions/sessions.zig create mode 100644 zig/src/sessions/store.zig create mode 100644 zig/src/sessions/store_test.zig diff --git a/build.zig b/build.zig index 98ff809634..bd0476b368 100644 --- a/build.zig +++ b/build.zig @@ -6,6 +6,7 @@ const qwen27_build = @import("zig/build/qwen27.zig"); const hip_build = @import("zig/build/hip.zig"); const grammar_build = @import("zig/build/grammar.zig"); const tp_build = @import("zig/build/tp.zig"); +const sessions_build = @import("zig/build/sessions.zig"); comptime { const required = std.mem.trim(u8, @embedFile(".zig-version"), "\r\n"); @@ -177,6 +178,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); tp_build.add(b, target, optimize, test_step); + sessions_build.add(b, test_step); } /// `zig build native -Dcpu=apple_m1`: tensorfold-native with the Metal engines for the Python package's bundle (a native M5 build traps on M1-M4). diff --git a/zig/build/sessions.zig b/zig/build/sessions.zig new file mode 100644 index 0000000000..a5d8ea6f3c --- /dev/null +++ b/zig/build/sessions.zig @@ -0,0 +1,13 @@ +//! Session memory's build: the family-neutral `sessions` module (zig/src/sessions) and its host tests on any OS. + +const std = @import("std"); + +pub fn module(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode) *std.Build.Module { + return b.createModule(.{ .root_source_file = b.path("zig/src/sessions/sessions.zig"), .target = target, .optimize = optimize, .link_libc = true }); +} + +pub fn add(b: *std.Build, test_step: *std.Build.Step) void { + const run = b.addRunArtifact(b.addTest(.{ .root_module = module(b, b.graph.host, .debug) })); + b.step("test-sessions", "Session memory's host tests: pool, prefix trie, RAM tier, NVMe tier on real files (no GPU)").dependOn(&run.step); + test_step.dependOn(&run.step); +} diff --git a/zig/src/sessions/disk.zig b/zig/src/sessions/disk.zig new file mode 100644 index 0000000000..f5ab58f06f --- /dev/null +++ b/zig/src/sessions/disk.zig @@ -0,0 +1,587 @@ +//! The NVMe tier: a file an entry under //rank/, written via .tmp + rename, streamed via the lanes. + +const std = @import("std"); +const linux = std.os.linux; +const Io = std.Io; +const lanes_mod = @import("lanes.zig"); +const Lanes = lanes_mod.Lanes; +const Staging = lanes_mod.Staging; +const Op = lanes_mod.Op; +const up = lanes_mod.up; +const format = @import("format.zig"); +const Layout = format.Layout; +const Head = format.Head; +const Seg = format.Seg; +const PageStore = @import("pagestore.zig").PageStore; +const pool_mod = @import("pool.zig"); +const Family = pool_mod.Family; +const Digest = @import("prefix.zig").Digest; + +pub const none = std.math.maxInt(u32); + +pub const Options = struct { + root: []const u8, + compat: Digest, + rank: u32 = 0, + world: u32 = 1, + /// bytes of files the tier keeps (least recently used files without dependents go first) + budget: u64 = 64 << 30, + /// entries shorter than this are dropped, not parked + min_tokens: u64 = 1024, + /// O_DIRECT where the filesystem takes it (buffered with the page cache dropped otherwise) + direct: bool = true, + /// staging buffers and their size: the host transient of any park or restore + staging: u32 = 16, + chunk: u32 = 8 << 20, + lanes: u32 = 4, + /// delta files a chain may stack before a park writes the entry whole + max_chain: u32 = 8, +}; + +/// A piece of bounded state (rings, carries, taps): small, host bytes. +pub const Blob = struct { + name: []const u8, + bytes: []const u8, +}; + +/// What a park writes: the entry, its tokens from first_page on, and this rank's pages per family, in logical order. +pub const WriteSrc = struct { + key: Digest, + /// the file holding pages [0, first_page) (a delta), or null (a whole entry: first_page 0) + base: ?Digest = null, + tag: u32, + kind: u32 = 0, + page: u32, + pos: u64, + first_page: u32 = 0, + npages: u32, + /// tokens [first_page x page, pos) + ids: []const i32, + store: PageStore, + fence: u64 = 0, + families: []const Family, + /// per family: family-tensor pages of this rank's logical pages [first_page, npages), in logical order + pages: []const []const u32, + blobs: []const Blob = &.{}, +}; + +/// One file of a restore: read logical pages [lo, hi) of it into the slot. +pub const Extent = struct { + key: Digest, + first_page: u32, + lo: u32, + hi: u32, +}; + +/// Where a restore puts rows: per family, the family-tensor page of each of the slot's logical pages. +pub const ReadSink = struct { + store: PageStore, + families: []const Family, + page: u32, + world: u32, + rank: u32, + /// per family: family-tensor pages for logical pages [0, npages) + pages: []const []const u32, +}; + +/// The bounded state a restore returns (owned). +pub const Restored = struct { + blobs: []Blob, + + pub fn deinit(r: *Restored, gpa: std.mem.Allocator) void { + for (r.blobs) |b| { + gpa.free(b.name); + gpa.free(b.bytes); + } + gpa.free(r.blobs); + } +}; + +/// The index record of one file. +pub const Rec = struct { + key: Digest, + base: u32 = none, + tag: u32, + kind: u32, + pos: u64, + first_page: u32, + npages: u32, + size: u64, + /// files stacked on this one + refs: u32 = 0, + /// files under it, this one included + depth: u32 = 1, + older: u32 = none, + newer: u32 = none, + live: bool = true, +}; + +pub const Stats = struct { + writes: u64 = 0, + reads: u64 = 0, + bytes_written: u64 = 0, + bytes_read: u64 = 0, + bad: u64 = 0, + evicted: u64 = 0, +}; + +pub const Disk = struct { + gpa: std.mem.Allocator, + io: Io, + opts: Options, + dir: Io.Dir, + staging: Staging, + lanes: *Lanes, + recs: std.ArrayList(Rec) = .empty, + free_recs: std.ArrayList(u32) = .empty, + by_key: std.AutoHashMapUnmanaged(Digest, u32) = .empty, + /// least recently used first + oldest: u32 = none, + newest: u32 = none, + used: u64 = 0, + stats: Stats = .{}, + + pub fn open(gpa: std.mem.Allocator, io: Io, opts: Options) !*Disk { + const d = try gpa.create(Disk); + errdefer gpa.destroy(d); + var name_buf: [64]u8 = undefined; + const sub = try std.fmt.bufPrint(&name_buf, "{x}/rank{d}", .{ opts.compat[0..8], opts.rank }); + var root = try Io.Dir.cwd().createDirPathOpen(io, opts.root, .{}); + defer root.close(io); + const dir = try root.createDirPathOpen(io, sub, .{ .open_options = .{ .iterate = true } }); + errdefer dir.close(io); + d.* = .{ .gpa = gpa, .io = io, .opts = opts, .dir = dir, .staging = try Staging.init(gpa, io, @max(2, opts.staging), up(opts.chunk)), .lanes = undefined }; + errdefer d.staging.deinit(); + d.lanes = try Lanes.start(gpa, io, &d.staging, @max(1, opts.lanes)); + return d; + } + + pub fn close(d: *Disk) void { + d.lanes.shutdown(); + d.staging.deinit(); + d.dir.close(d.io); + d.recs.deinit(d.gpa); + d.free_recs.deinit(d.gpa); + d.by_key.deinit(d.gpa); + const gpa = d.gpa; + gpa.destroy(d); + } + + pub fn accepts(d: *const Disk, pos: u64) bool { + return pos >= d.opts.min_tokens; + } + + pub fn get(d: *Disk, key: Digest) ?*Rec { + const i = d.by_key.get(key) orelse return null; + return &d.recs.items[i]; + } + + fn fileName(key: Digest, ext: []const u8, buf: *[48]u8) [:0]const u8 { + const n = (std.fmt.bufPrint(buf[0..47], "{x}.{s}", .{ &key, ext }) catch unreachable).len; + buf[n] = 0; + return buf[0..n :0]; + } + + fn openFile(d: *Disk, name: [:0]const u8, write: bool) !struct { fd: linux.fd_t, direct: bool } { + var flags: linux.O = .{ .ACCMODE = if (write) .WRONLY else .RDONLY, .CLOEXEC = true }; + if (write) { + flags.CREAT = true; + flags.TRUNC = true; + } + if (d.opts.direct) { + var f2 = flags; + f2.DIRECT = true; + const rc = linux.openat(d.dir.handle, name, f2, 0o644); + if (linux.errno(rc) == .SUCCESS) return .{ .fd = @intCast(rc), .direct = true }; + } + const rc = linux.openat(d.dir.handle, name, flags, 0o644); + if (linux.errno(rc) != .SUCCESS) return if (linux.errno(rc) == .NOENT) error.FileNotFound else error.OpenFailed; + return .{ .fd = @intCast(rc), .direct = false }; + } + + /// The layout a WriteSrc gets: segments, chunk sizes (whole pages a chunk for pool segments). + pub fn plan(d: *Disk, src: *const WriteSrc) !Layout { + var segs: std.ArrayList(Seg) = .empty; + defer segs.deinit(d.gpa); + const cap = d.opts.chunk; + try segs.append(d.gpa, .{ .name = format.segName("ids"), .kind = .ids, .family = 0, .nbytes = src.ids.len * 4, .chunk = cap }); + for (src.families, 0..) |f, i| { + const pb: u32 = @intCast(f.pageBytes(src.page)); + const k = @max(1, cap / pb); + if (@as(u64, k) * pb > d.staging.size) return error.PageTooLarge; + try segs.append(d.gpa, .{ .name = format.segName(f.name), .kind = .pool, .family = @intCast(i), .nbytes = @as(u64, src.pages[i].len) * pb, .chunk = k * pb }); + } + for (src.blobs) |b| try segs.append(d.gpa, .{ .name = format.segName(b.name), .kind = .blob, .family = 0, .nbytes = b.bytes.len, .chunk = cap }); + const head: Head = .{ .key = src.key, .base = src.base orelse @splat(0), .compat = d.opts.compat, .tag = src.tag, .kind = src.kind, .page = src.page, .world = d.opts.world, .rank = d.opts.rank, .first_page = src.first_page, .npages = src.npages, .nseg = 0, .ndigest = 0, .pos = src.pos }; + return Layout.build(d.gpa, head, segs.items); + } + + /// Writes an entry's file streamed (the pump copies pages to staging, lanes hash and write); returns its bytes. + pub fn writeFile(d: *Disk, src: *const WriteSrc) !u64 { + var lay = try d.plan(src); + defer lay.deinit(d.gpa); + var tb: [48]u8 = undefined; + var fb: [48]u8 = undefined; + const tmp = fileName(src.key, "tmp", &tb); + const f = try d.openFile(tmp, true); + var ok = false; + defer { + _ = linux.close(f.fd); + if (!ok) _ = linux.unlinkat(d.dir.handle, tmp, 0); + } + try src.store.waitFence(src.fence); + var op = try Op.init(d.gpa, d.io, &d.staging); + defer op.deinit(); + errdefer op.wait() catch {}; + for (lay.segs) |*s| { + for (0..s.nchunks) |ci| { + if (op.failed()) |e| return e; + const c = s.chunkAt(@intCast(ci)); + const sb = try d.staging.get(); + errdefer d.staging.put(sb); + const buf = d.staging.buf(sb)[0..c.data]; + const at = @as(u64, ci) * s.chunk; + switch (s.kind) { + .ids => @memcpy(buf, std.mem.sliceAsBytes(src.ids)[at..][0..c.data]), + .blob => @memcpy(buf, blobOf(src, s).bytes[at..][0..c.data]), + .pool => { + const pb: u32 = @intCast(src.families[s.family].pageBytes(src.page)); + const p0: usize = @intCast(at / pb); + try src.store.read(s.family, src.pages[s.family][p0..][0 .. c.data / pb], buf); + }, + _ => unreachable, + } + d.lanes.submit(.{ .kind = .write, .fd = f.fd, .staging = sb, .data = c.data, .span = c.span, .off = c.off, .digest = &lay.digests[s.first_digest + ci], .op = &op }); + } + } + try op.wait(); + const hb = try d.staging.get(); + defer d.staging.put(hb); + if (lay.head_span > d.staging.size) return error.HeaderTooLarge; + lay.encode(d.staging.buf(hb)); + try lanes_mod.pwriteAll(f.fd, d.staging.buf(hb)[0..lay.head_span], 0); + if (linux.errno(linux.fdatasync(f.fd)) != .SUCCESS) return error.SyncFailed; + if (!f.direct) _ = linux.fadvise(f.fd, 0, 0, linux.POSIX_FADV.DONTNEED); + if (linux.errno(linux.renameat(d.dir.handle, tmp, d.dir.handle, fileName(src.key, "tfs", &fb))) != .SUCCESS) return error.RenameFailed; + ok = true; + return lay.head.total; + } + + fn blobOf(src: *const WriteSrc, s: *const Seg) Blob { + for (src.blobs) |b| if (std.mem.eql(u8, b.name, s.nameOf())) return b; + unreachable; + } + + /// A file's header, checked. + pub fn readHead(d: *Disk, fd: linux.fd_t) !Layout { + const hb = try d.staging.get(); + defer d.staging.put(hb); + const buf = d.staging.buf(hb); + try lanes_mod.preadAll(fd, buf[0..4096], 0); + const hn = std.mem.readInt(u32, buf[12..16], .little); + if (up(hn) > buf.len) return error.BadHeader; + if (hn > 4096) try lanes_mod.preadAll(fd, buf[4096..up(hn)], 4096); + return Layout.decode(d.gpa, buf[0..up(hn)]); + } + + /// Reads pages [lo, hi) of one restore file into the sink, each chunk checked first; blobs too when want_blobs. + pub fn readFile(d: *Disk, ext: Extent, sink: *const ReadSink, want_blobs: bool) !?Restored { + var nb: [48]u8 = undefined; + const f = try d.openFile(fileName(ext.key, "tfs", &nb), false); + defer _ = linux.close(f.fd); + var lay = try d.readHead(f.fd); + defer lay.deinit(d.gpa); + if (!std.mem.eql(u8, &lay.head.key, &ext.key) or lay.head.first_page != ext.first_page or lay.head.page != sink.page) return error.BadHeader; + var st: linux.Statx = undefined; + if (linux.errno(linux.statx(f.fd, "", linux.AT.EMPTY_PATH, .{ .SIZE = true }, &st)) != .SUCCESS or st.size < lay.head.total) return error.Truncated; + var out: std.ArrayList(Blob) = .empty; + errdefer { + for (out.items) |b| { + d.gpa.free(b.name); + d.gpa.free(b.bytes); + } + out.deinit(d.gpa); + } + var op = try Op.init(d.gpa, d.io, &d.staging); + defer op.deinit(); + errdefer op.wait() catch {}; + // (segment, chunk) jobs; the pump keeps at most staging - 1 in flight (one buffer for the header path) + var jobs: std.ArrayList(struct { seg: u32, chunk: u32 }) = .empty; + defer jobs.deinit(d.gpa); + for (lay.segs, 0..) |*s, si| switch (s.kind) { + .pool => { + if (s.family >= sink.families.len) return error.BadHeader; + const fam = sink.families[s.family]; + const pb: u32 = @intCast(fam.pageBytes(sink.page)); + if (s.chunk % pb != 0 or !std.mem.eql(u8, s.nameOf(), fam.name)) return error.BadHeader; + const need = ownedIn(fam, sink, ext.first_page, ext.hi) - ownedIn(fam, sink, ext.first_page, ext.lo); + const skip = ownedIn(fam, sink, ext.first_page, ext.lo); + if (s.nbytes < (@as(u64, skip) + need) * pb) return error.BadHeader; + const kpp = s.chunk / pb; + const c0 = skip / kpp; + const c1 = if (need == 0) c0 else (skip + need + kpp - 1) / kpp; + for (c0..c1) |c| try jobs.append(d.gpa, .{ .seg = @intCast(si), .chunk = @intCast(c) }); + }, + .blob => if (want_blobs) { + const buf = try d.gpa.alloc(u8, s.nbytes); + errdefer d.gpa.free(buf); + const name = try d.gpa.dupe(u8, s.nameOf()); + errdefer d.gpa.free(name); + try out.append(d.gpa, .{ .name = name, .bytes = buf }); + for (0..s.nchunks) |c| try jobs.append(d.gpa, .{ .seg = @intCast(si), .chunk = @intCast(c) }); + }, + else => {}, + }; + const window = d.staging.count - 1; + var sent: usize = 0; + var seen: usize = 0; + var bytes: u64 = 0; + while (seen < jobs.items.len) { + while (sent < jobs.items.len and sent - seen < window and op.failed() == null) : (sent += 1) { + const jb = jobs.items[sent]; + const s = &lay.segs[jb.seg]; + const c = s.chunkAt(jb.chunk); + d.lanes.submit(.{ .kind = .read, .fd = f.fd, .staging = try d.staging.get(), .data = c.data, .span = c.span, .off = c.off, .digest = &lay.digests[s.first_digest + jb.chunk], .op = &op, .tag = @intCast(sent) }); + bytes += c.span; + } + const j = op.next() orelse break; + seen += 1; + defer d.staging.put(j.staging); + if (op.failed() != null) continue; + const jb = jobs.items[j.tag]; + const s = &lay.segs[jb.seg]; + const data = d.staging.buf(j.staging)[0..j.data]; + if (s.kind == .blob) { + for (out.items) |b| if (std.mem.eql(u8, b.name, s.nameOf())) @memcpy(@constCast(b.bytes)[@as(u64, jb.chunk) * s.chunk ..][0..data.len], data); + continue; + } + try d.drain(s, jb.chunk, data, ext, sink); + } + try op.wait(); + if (!f.direct) _ = linux.fadvise(f.fd, 0, 0, linux.POSIX_FADV.DONTNEED); + d.stats.bytes_read += bytes; + if (!want_blobs) { + out.deinit(d.gpa); + return null; + } + return .{ .blobs = try out.toOwnedSlice(d.gpa) }; + } + + /// A checked pool chunk's pages that fall in [lo, hi) into the sink's pages. + fn drain(d: *Disk, s: *const Seg, chunk: u32, data: []const u8, ext: Extent, sink: *const ReadSink) !void { + _ = d; + const fam = sink.families[s.family]; + const pb: u32 = @intCast(fam.pageBytes(sink.page)); + const kpp = s.chunk / pb; + const skip = ownedIn(fam, sink, ext.first_page, ext.lo); + const end = ownedIn(fam, sink, ext.first_page, ext.hi); + const a = @max(chunk * kpp, skip); + const b = @min(chunk * kpp + @as(u32, @intCast(data.len / pb)), end); + if (a >= b) return; + var list: [512]u32 = undefined; + var n: u32 = 0; + var j = a; + while (j < b) : (j += 1) { + list[n] = sink.pages[s.family][logicalOf(fam, sink, ext.first_page, j)]; + n += 1; + if (n == list.len or j + 1 == b) { + const first = j + 1 - n; + try sink.store.write(s.family, list[0..n], data[(first - chunk * kpp) * pb ..][0 .. n * pb]); + n = 0; + } + } + } + + /// Records a written file in the index (round thread, at the plan's point); its base gains a dependent. + pub fn register(d: *Disk, src: *const WriteSrc, size: u64) !u32 { + const id = try d.add(src.key, src.base, src.tag, src.kind, src.pos, src.first_page, src.npages, size); + d.stats.writes += 1; + d.stats.bytes_written += size; + return id; + } + + fn add(d: *Disk, key: Digest, base: ?Digest, tag: u32, kind: u32, pos: u64, first_page: u32, npages: u32, size: u64) !u32 { + var r: Rec = .{ .key = key, .tag = tag, .kind = kind, .pos = pos, .first_page = first_page, .npages = npages, .size = size }; + if (base) |b| { + const bi = d.by_key.get(b) orelse return error.NoBase; + r.base = bi; + r.depth = d.recs.items[bi].depth + 1; + } + const id: u32 = if (d.free_recs.pop()) |i| i else blk: { + try d.recs.append(d.gpa, undefined); + break :blk @intCast(d.recs.items.len - 1); + }; + try d.by_key.put(d.gpa, key, id); + if (r.base != none) d.recs.items[r.base].refs += 1; + d.recs.items[id] = r; + d.link(id); + d.used += size; + return id; + } + + /// A valid file found on restart: its header, size and verified ids (tokens [first_page x page, pos), owned). + pub const Found = struct { + head: Head, + size: u64, + mtime: i64, + ids: []i32, + + fn older(_: void, a: Found, b: Found) bool { + return a.mtime < b.mtime or (a.mtime == b.mtime and a.head.first_page < b.head.first_page); + } + }; + + /// Indexes the directory on restart: drops .tmp and bad files, checks ids, registers bases before their deltas. + pub fn scan(d: *Disk, out: *std.ArrayList(Found)) !void { + var it = d.dir.iterate(); + var nb: [48]u8 = undefined; + while (try it.next(d.io)) |ent| { + if (ent.name.len >= nb.len) continue; + @memcpy(nb[0..ent.name.len], ent.name); + nb[ent.name.len] = 0; + const name = nb[0..ent.name.len :0]; + if (std.mem.endsWith(u8, name, ".tmp")) { + _ = linux.unlinkat(d.dir.handle, name, 0); + continue; + } + if (!std.mem.endsWith(u8, name, ".tfs")) continue; + const f = d.found(name) catch { + d.stats.bad += 1; + _ = linux.unlinkat(d.dir.handle, name, 0); + continue; + }; + try out.append(d.gpa, f); + } + std.mem.sort(Found, out.items, {}, Found.older); + var i: usize = 0; + while (i < out.items.len) { + const f = out.items[i]; + const base: ?Digest = if (f.head.hasBase()) f.head.base else null; + if (base != null and d.by_key.get(base.?) == null) { + d.stats.bad += 1; + var kb: [48]u8 = undefined; + _ = linux.unlinkat(d.dir.handle, fileName(f.head.key, "tfs", &kb), 0); + d.gpa.free(f.ids); + _ = out.orderedRemove(i); + continue; + } + _ = try d.add(f.head.key, base, f.head.tag, f.head.kind, f.head.pos, f.head.first_page, f.head.npages, f.size); + i += 1; + } + } + + fn found(d: *Disk, name: [:0]const u8) !Found { + const f = try d.openFile(name, false); + defer _ = linux.close(f.fd); + var lay = try d.readHead(f.fd); + defer lay.deinit(d.gpa); + const h = lay.head; + var kb: [48]u8 = undefined; + if (!std.mem.eql(u8, fileName(h.key, "tfs", &kb), name)) return error.BadHeader; + if (!std.mem.eql(u8, &h.compat, &d.opts.compat) or h.rank != d.opts.rank or h.world != d.opts.world) return error.Foreign; + var st: linux.Statx = undefined; + if (linux.errno(linux.statx(f.fd, "", linux.AT.EMPTY_PATH, .{ .SIZE = true, .MTIME = true }, &st)) != .SUCCESS or st.size < h.total) return error.Truncated; + const seg = lay.find(.ids, "ids") orelse return error.BadHeader; + if (seg.nbytes % 4 != 0 or seg.nbytes / 4 != h.pos - @as(u64, h.first_page) * h.page) return error.BadHeader; + const ids = try d.gpa.alloc(i32, seg.nbytes / 4); + errdefer d.gpa.free(ids); + const sb = try d.staging.get(); + defer d.staging.put(sb); + for (0..seg.nchunks) |ci| { + const c = seg.chunkAt(@intCast(ci)); + const buf = d.staging.buf(sb); + try lanes_mod.preadAll(f.fd, buf[0..c.span], c.off); + var dg: [32]u8 = undefined; + std.crypto.hash.sha2.Sha256.hash(buf[0..c.data], &dg, .{}); + if (!std.mem.eql(u8, &dg, &lay.digests[seg.first_digest + ci])) return error.ChecksumMismatch; + @memcpy(std.mem.sliceAsBytes(ids)[@as(usize, ci) * seg.chunk ..][0..c.data], buf[0..c.data]); + } + return .{ .head = h, .size = st.size, .mtime = st.mtime.sec * std.time.ns_per_s + st.mtime.nsec, .ids = ids }; + } + + fn link(d: *Disk, id: u32) void { + const r = &d.recs.items[id]; + r.older = d.newest; + r.newer = none; + if (d.newest != none) d.recs.items[d.newest].newer = id else d.oldest = id; + d.newest = id; + } + + fn unlink(d: *Disk, id: u32) void { + const r = &d.recs.items[id]; + if (r.older != none) d.recs.items[r.older].newer = r.newer else d.oldest = r.newer; + if (r.newer != none) d.recs.items[r.newer].older = r.older else d.newest = r.older; + } + + /// Marks a file (and its bases) used now: O(chain). + pub fn touch(d: *Disk, key: Digest) void { + var id = d.by_key.get(key) orelse return; + while (id != none) : (id = d.recs.items[id].base) { + d.unlink(id); + d.link(id); + } + // bases end up newest: walk order keeps the entry itself older than its bases, so trims reach dependents first + } + + /// The files a restore of key reads, base first, with the logical pages each gives. + pub fn extents(d: *Disk, key: Digest, out: *std.ArrayList(Extent)) !void { + out.clearRetainingCapacity(); + var id = d.by_key.get(key) orelse return error.FileNotFound; + var hi = d.recs.items[id].npages; + while (id != none) { + const r = &d.recs.items[id]; + try out.append(d.gpa, .{ .key = r.key, .first_page = r.first_page, .lo = r.first_page, .hi = hi }); + hi = r.first_page; + id = r.base; + } + std.mem.reverse(Extent, out.items); + } + + /// Deletes a file without dependents (error when one stacks on it). + pub fn remove(d: *Disk, key: Digest) !void { + const id = d.by_key.get(key) orelse return; + const r = &d.recs.items[id]; + if (r.refs > 0) return error.HasDependents; + if (r.base != none) d.recs.items[r.base].refs -= 1; + d.unlink(id); + _ = d.by_key.remove(key); + d.used -= r.size; + r.live = false; + var nb: [48]u8 = undefined; + _ = linux.unlinkat(d.dir.handle, fileName(key, "tfs", &nb), 0); + try d.free_recs.append(d.gpa, id); + } + + /// Deletes least recently used files without dependents while over budget; their keys go to out. + pub fn trim(d: *Disk, out: *std.ArrayList(Digest)) !void { + var id = d.oldest; + while (d.used > d.opts.budget and id != none) { + const next = d.recs.items[id].newer; + if (d.recs.items[id].refs == 0) { + const key = d.recs.items[id].key; + try d.remove(key); + try out.append(d.gpa, key); + d.stats.evicted += 1; + id = d.oldest; // a removed file may free its base + continue; + } + id = next; + } + } +}; + +/// Owned (this rank's) logical pages of family f in [first, k): the file's index of logical page k. +pub fn ownedIn(f: Family, sink: *const ReadSink, first: u32, k: u32) u32 { + if (!(f.split and sink.world > 1)) return k - first; + var by: [pool_mod.max_world]u32 = undefined; + pool_mod.span(first, k, sink.world, &by); + return by[sink.rank]; +} + +/// The logical page of a family's j-th stored page of a file that starts at first. +pub fn logicalOf(f: Family, sink: *const ReadSink, first: u32, j: u32) u32 { + if (!(f.split and sink.world > 1)) return first + j; + const w = sink.world; + const r0 = first + (sink.rank + w - first % w) % w; // the first owned logical page at or after first + return r0 + j * w; +} diff --git a/zig/src/sessions/disk_test.zig b/zig/src/sessions/disk_test.zig new file mode 100644 index 0000000000..6ac0525e45 --- /dev/null +++ b/zig/src/sessions/disk_test.zig @@ -0,0 +1,125 @@ +//! The NVMe tier on real files: whole and delta files round-trip bit for bit, split ranks, damage. + +const std = @import("std"); +const testing = std.testing; +const pool_mod = @import("pool.zig"); +const Pool = pool_mod.Pool; +const Family = pool_mod.Family; +const HostPool = @import("pagestore.zig").HostPool; +const disk = @import("disk.zig"); + +const fams = [_]Family{ + .{ .name = "comp.2", .ratio = 2, .row_bytes = 584, .split = true }, + .{ .name = "index_k.2", .ratio = 2, .row_bytes = 132 }, + .{ .name = "comp.20", .ratio = 1, .row_bytes = 584, .split = true }, +}; + +/// Fills a page of every family with bytes that name (family, logical page, byte). +fn fill(h: *HostPool, p: *const Pool, logical: u32, pg: u32, salt: u8) void { + for (fams, 0..) |f, fi| { + if (f.split and !p.owns(pg)) continue; + const v = h.pageOf(@intCast(fi), p.familyPage(f, pg)); + for (v, 0..) |*b, i| b.* = @truncate(i *% 131 +% logical *% 7 +% fi *% 29 +% salt); + } +} + +fn familyPages(gpa: std.mem.Allocator, p: *const Pool, pages: []const u32, first: u32, owned_only: bool) ![][]u32 { + const out = try gpa.alloc([]u32, fams.len); + for (fams, 0..) |f, fi| { + var l: std.ArrayList(u32) = .empty; + for (pages[first..], first..) |pg, k| { + if (f.split and p.world > 1 and owned_only and k % p.world != p.rank) continue; + try l.append(gpa, p.familyPage(f, pg)); + } + out[fi] = try l.toOwnedSlice(gpa); + } + return out; +} + +fn freePages(gpa: std.mem.Allocator, x: [][]u32) void { + for (x) |l| gpa.free(l); + gpa.free(x); +} + +fn roundTrip(world: u32, rank: u32) !void { + const gpa = testing.allocator; + const io = testing.io; + var tmp = testing.tmpDir(.{}); + defer tmp.cleanup(); + const root = try tmp.dir.realPathFileAlloc(io, ".", gpa); + defer gpa.free(root); + var p = try Pool.init(gpa, .{ .families = &fams, .page = 256 }, 64 * 256, world, rank); + defer p.deinit(); + var h = try HostPool.init(gpa, &p); + defer h.deinit(); + // small chunks so pool segments span many chunks; few staging buffers so the pump waits on the lanes + const d = try disk.Disk.open(gpa, io, .{ .root = root, .compat = @splat(5), .rank = rank, .world = world, .chunk = 3 * 74_752, .staging = 3, .lanes = 3, .min_tokens = 1 }); + defer d.close(); + const a = try p.newSlot(32 * 256); + const pos: u64 = 9 * 256 + 100; + try a.ensure(pos); + for (a.mapped(), 0..) |pg, k| fill(&h, &p, @intCast(k), pg, 1); + const ids = try gpa.alloc(i32, pos); + defer gpa.free(ids); + for (ids, 0..) |*t, i| t.* = @intCast(i * 3 + 1); + var blob: [5000]u8 = undefined; + for (&blob, 0..) |*x, i| x.* = @truncate(9 - i % 5); + const fp = try familyPages(gpa, &p, a.mapped(), 0, true); + defer freePages(gpa, fp); + const src: disk.WriteSrc = .{ .key = @splat(1), .tag = 4, .page = 256, .pos = pos, .npages = a.len, .ids = ids, .store = h.store(), .families = &fams, .pages = fp, .blobs = &.{.{ .name = "rings", .bytes = &blob }} }; + const size = try d.writeFile(&src); + _ = try d.register(&src, size); + // a newer turn stacks a delta on it: pages [9, 14) only (page 9 was partial in the base) + const pos2: u64 = 13 * 256 + 7; + try a.ensure(pos2); + for (a.mapped()[9..], 9..) |pg, k| fill(&h, &p, @intCast(k), pg, 2); + const ids2 = try gpa.alloc(i32, pos2); + defer gpa.free(ids2); + for (ids2, 0..) |*t, i| t.* = @intCast(i * 3 + 1); + const fp2 = try familyPages(gpa, &p, a.mapped(), 9, true); + defer freePages(gpa, fp2); + const src2: disk.WriteSrc = .{ .key = @splat(2), .base = @splat(1), .tag = 4, .page = 256, .pos = pos2, .first_page = 9, .npages = a.len, .ids = ids2[9 * 256 ..], .store = h.store(), .families = &fams, .pages = fp2, .blobs = &.{.{ .name = "rings", .bytes = blob[0..77] }} }; + const size2 = try d.writeFile(&src2); + _ = try d.register(&src2, size2); + try testing.expect(size2 < size); + try testing.expectError(error.HasDependents, d.remove(@splat(1))); + // restore the delta onto other pages through its chain + const b = try p.newSlot(32 * 256); + try b.ensure(pos2); + var ext: std.ArrayList(disk.Extent) = .empty; + defer ext.deinit(gpa); + try d.extents(@splat(2), &ext); + try testing.expectEqual(@as(usize, 2), ext.items.len); + try testing.expectEqual(@as(u32, 9), ext.items[0].hi); + const tp = try familyPages(gpa, &p, b.mapped(), 0, false); + defer freePages(gpa, tp); + const sink: disk.ReadSink = .{ .store = h.store(), .families = &fams, .page = 256, .world = world, .rank = rank, .pages = tp }; + try testing.expect(try d.readFile(ext.items[0], &sink, false) == null); + var got = (try d.readFile(ext.items[1], &sink, true)).?; + defer got.deinit(gpa); + try testing.expectEqualSlices(u8, blob[0..77], got.blobs[0].bytes); + for (a.mapped(), b.mapped()) |pa, pb| for (fams, 0..) |f, fi| { + if (f.split and !p.owns(pa)) continue; + try testing.expectEqualSlices(u8, h.pageOf(@intCast(fi), p.familyPage(f, pa)), h.pageOf(@intCast(fi), p.familyPage(f, pb))); + }; + // damage one chunk: the read refuses it before or as its rows land + var nb: [48]u8 = undefined; + const key2: [16]u8 = @splat(2); + const name = try std.fmt.bufPrint(&nb, "{x}.tfs", .{&key2}); + const f = try d.dir.openFile(io, name, .{ .mode = .read_write }); + var one = [_]u8{0}; + _ = try f.readPositional(io, &.{&one}, 3 * 4096 + 100); + one[0] ^= 0xff; + try f.writePositionalAll(io, &one, 3 * 4096 + 100); + f.close(io); + try testing.expectError(error.ChecksumMismatch, d.readFile(ext.items[1], &sink, true)); +} + +test "whole and delta files round-trip bit for bit (replicated)" { + try roundTrip(1, 0); +} + +test "split ranks write and restore their own pages (rank 0 and 1 of 2)" { + try roundTrip(2, 0); + try roundTrip(2, 1); +} diff --git a/zig/src/sessions/format.zig b/zig/src/sessions/format.zig new file mode 100644 index 0000000000..2ac391d156 --- /dev/null +++ b/zig/src/sessions/format.zig @@ -0,0 +1,197 @@ +//! A session file (format 2): a header, then segments of 4 KiB-aligned chunks, each with its own SHA-256. + +const std = @import("std"); +const Sha256 = std.crypto.hash.sha2.Sha256; +const up = @import("lanes.zig").up; + +pub const magic = "TFSESS02".*; +pub const version: u32 = 2; + +pub const SegKind = enum(u32) { ids = 1, pool = 2, blob = 3, _ }; + +pub const Head = extern struct { + magic: [8]u8 = magic, + version: u32 = version, + /// bytes of Head + segments + digests (before the pad) + head_bytes: u32 = 0, + /// SHA-256 of those bytes with this field zero + head_sha: [32]u8 = @splat(0), + key: [16]u8, + /// the entry whose files hold logical pages [0, first_page); zero: none + base: [16]u8 = @splat(0), + /// what makes these bytes mean the same thing to a build (family, knobs, layout, split) + compat: [16]u8, + tag: u32, + kind: u32, + page: u32, + world: u32, + rank: u32, + first_page: u32, + npages: u32, + nseg: u32, + ndigest: u32, + _pad: u32 = 0, + pos: u64, + total: u64 = 0, + + pub fn hasBase(h: *const Head) bool { + return !std.mem.allEqual(u8, &h.base, 0); + } +}; + +pub const Seg = extern struct { + name: [40]u8, + kind: SegKind, + /// pool: the family's index in the layout + family: u32, + nbytes: u64, + offset: u64 = 0, + /// bytes a chunk (the last may be shorter) and file bytes a chunk (4 KiB multiple) + chunk: u32, + stride: u32 = 0, + first_digest: u32 = 0, + nchunks: u32 = 0, + + pub fn nameOf(s: *const Seg) []const u8 { + return std.mem.sliceTo(&s.name, 0); + } + + /// Chunk i's data bytes and file span. + pub fn chunkAt(s: *const Seg, i: u32) struct { off: u64, data: u32, span: u32 } { + const data: u32 = @intCast(@min(@as(u64, s.chunk), s.nbytes - @as(u64, i) * s.chunk)); + return .{ .off = s.offset + @as(u64, i) * s.stride, .data = data, .span = @intCast(up(data)) }; + } +}; + +pub const Digest32 = [32]u8; + +pub fn segName(name: []const u8) [40]u8 { + var n: [40]u8 = @splat(0); + @memcpy(n[0..@min(name.len, 39)], name[0..@min(name.len, 39)]); + return n; +} + +/// A file's layout: offsets, chunks and the digest slots, before or after its chunks are written. +pub const Layout = struct { + head: Head, + segs: []Seg, + digests: []Digest32, + /// file bytes of the padded header + head_span: u64, + + /// Places segments (each with name, kind, family, nbytes, chunk) after the header. + pub fn build(gpa: std.mem.Allocator, head: Head, segs_in: []const Seg) !Layout { + const segs = try gpa.dupe(Seg, segs_in); + errdefer gpa.free(segs); + var nd: u32 = 0; + for (segs) |*s| { + if (s.chunk == 0) return error.BadSegment; + s.nchunks = @intCast((s.nbytes + s.chunk - 1) / s.chunk); + s.stride = @intCast(up(s.chunk)); + s.first_digest = nd; + nd += s.nchunks; + } + var h = head; + h.nseg = @intCast(segs.len); + h.ndigest = nd; + h.head_bytes = @intCast(@sizeOf(Head) + segs.len * @sizeOf(Seg) + nd * @sizeOf(Digest32)); + const span = up(h.head_bytes); + var off = span; + for (segs) |*s| { + s.offset = off; + if (s.nchunks > 0) off += @as(u64, s.nchunks - 1) * s.stride + up(s.nbytes - @as(u64, s.nchunks - 1) * s.chunk); + } + h.total = off; + const d = try gpa.alloc(Digest32, nd); + @memset(d, @splat(0)); + return .{ .head = h, .segs = segs, .digests = d, .head_span = span }; + } + + pub fn deinit(l: *Layout, gpa: std.mem.Allocator) void { + gpa.free(l.segs); + gpa.free(l.digests); + l.* = undefined; + } + + /// The header's bytes with its checksum into out (len >= head_span; the pad is zeroed). + pub fn encode(l: *Layout, out: []u8) void { + @memset(out[0..l.head_span], 0); + var h = l.head; + h.head_sha = @splat(0); + const hb = std.mem.asBytes(&h); + const sb = std.mem.sliceAsBytes(l.segs); + const db = std.mem.sliceAsBytes(l.digests); + @memcpy(out[0..hb.len], hb); + @memcpy(out[hb.len..][0..sb.len], sb); + @memcpy(out[hb.len + sb.len ..][0..db.len], db); + var sha: [32]u8 = undefined; + Sha256.hash(out[0..h.head_bytes], &sha, .{}); + @memcpy(out[@offsetOf(Head, "head_sha")..][0..32], &sha); + l.head.head_sha = sha; + } + + /// The header read back from bytes (at least head_bytes of them), every field and offset checked. + pub fn decode(gpa: std.mem.Allocator, bytes: []const u8) !Layout { + if (bytes.len < @sizeOf(Head)) return error.BadHeader; + var h: Head = undefined; + @memcpy(std.mem.asBytes(&h), bytes[0..@sizeOf(Head)]); + if (!std.mem.eql(u8, &h.magic, &magic) or h.version != version) return error.BadHeader; + const want = @sizeOf(Head) + @as(u64, h.nseg) * @sizeOf(Seg) + @as(u64, h.ndigest) * @sizeOf(Digest32); + if (h.head_bytes != want or bytes.len < want) return error.BadHeader; + var check: [32]u8 = undefined; + var sha = Sha256.init(.{}); + sha.update(bytes[0..@offsetOf(Head, "head_sha")]); + sha.update(&@as([32]u8, @splat(0))); + sha.update(bytes[@offsetOf(Head, "head_sha") + 32 .. want]); + sha.final(&check); + if (!std.mem.eql(u8, &check, &h.head_sha)) return error.ChecksumMismatch; + const segs = try gpa.alloc(Seg, h.nseg); + errdefer gpa.free(segs); + @memcpy(std.mem.sliceAsBytes(segs), bytes[@sizeOf(Head)..][0 .. segs.len * @sizeOf(Seg)]); + const d = try gpa.alloc(Digest32, h.ndigest); + errdefer gpa.free(d); + @memcpy(std.mem.sliceAsBytes(d), bytes[@sizeOf(Head) + segs.len * @sizeOf(Seg) ..][0 .. d.len * 32]); + var nd: u64 = 0; + for (segs) |s| { + if (s.chunk == 0 or s.stride != up(s.chunk) or s.first_digest != nd or s.offset % 4096 != 0) return error.BadHeader; + if (s.nchunks != (s.nbytes + s.chunk - 1) / s.chunk) return error.BadHeader; + if (s.nchunks > 0 and s.offset + @as(u64, s.nchunks - 1) * s.stride + up(s.nbytes - @as(u64, s.nchunks - 1) * s.chunk) > h.total) return error.BadHeader; + nd += s.nchunks; + } + if (nd != h.ndigest) return error.BadHeader; + return .{ .head = h, .segs = segs, .digests = d, .head_span = up(h.head_bytes) }; + } + + pub fn find(l: *const Layout, kind: SegKind, name: []const u8) ?*const Seg { + for (l.segs) |*s| if (s.kind == kind and std.mem.eql(u8, s.nameOf(), name)) return s; + return null; + } +}; + +test "a layout encodes, decodes, and refuses a damaged header" { + const gpa = std.testing.allocator; + const head: Head = .{ .key = @splat(1), .compat = @splat(2), .tag = 3, .kind = 0, .page = 256, .world = 2, .rank = 1, .first_page = 4, .npages = 9, .nseg = 0, .ndigest = 0, .pos = 2200 }; + var l = try Layout.build(gpa, head, &.{ + .{ .name = segName("ids"), .kind = .ids, .family = 0, .nbytes = 4 * (2200 - 1024), .chunk = 1 << 20 }, + .{ .name = segName("comp.20"), .kind = .pool, .family = 0, .nbytes = 3 * 149_504, .chunk = 2 * 149_504 }, + .{ .name = segName("taps"), .kind = .blob, .family = 0, .nbytes = 0, .chunk = 4096 }, + }); + defer l.deinit(gpa); + try std.testing.expectEqual(@as(u32, 1 + 2 + 0), l.head.ndigest); + try std.testing.expectEqual(@as(u64, 4096), l.segs[0].offset); + try std.testing.expectEqual(@as(u64, 4096 + up(4 * 1176)), l.segs[1].offset); + const c1 = l.segs[1].chunkAt(1); + try std.testing.expectEqual(@as(u32, 149_504), c1.data); + try std.testing.expectEqual(l.segs[1].offset + up(2 * 149_504), c1.off); + try std.testing.expectEqual(l.segs[2].offset, l.head.total); + l.digests[1][5] = 9; + var buf: [8192]u8 = undefined; + l.encode(&buf); + var back = try Layout.decode(gpa, &buf); + defer back.deinit(gpa); + try std.testing.expectEqual(l.head.total, back.head.total); + try std.testing.expectEqual(@as(u8, 9), back.digests[1][5]); + try std.testing.expectEqualStrings("comp.20", back.find(.pool, "comp.20").?.nameOf()); + buf[200] ^= 0x40; + try std.testing.expectError(error.ChecksumMismatch, Layout.decode(gpa, &buf)); +} diff --git a/zig/src/sessions/freeset.zig b/zig/src/sessions/freeset.zig new file mode 100644 index 0000000000..b8cce323af --- /dev/null +++ b/zig/src/sessions/freeset.zig @@ -0,0 +1,114 @@ +//! A set of free page indices, lowest first in O(log64 n): a bit a page, then a bit a non-empty word a level up. + +const std = @import("std"); + +const max_levels = 6; // 64^6 pages: more than any pool + +pub const FreeSet = struct { + /// levels[0] holds a bit a page, levels[k] a bit a non-empty word of levels[k - 1] + levels: [max_levels][]u64 = @splat(&.{}), + depth: u8 = 0, + n: u32 = 0, + free: u32 = 0, + + /// All of 0 .. n - 1 free. + pub fn init(gpa: std.mem.Allocator, n: u32) !FreeSet { + var s: FreeSet = .{ .n = n }; + errdefer s.deinit(gpa); + var bits: usize = @max(n, 1); + while (true) { + const words = (bits + 63) / 64; + s.levels[s.depth] = try gpa.alloc(u64, words); + @memset(s.levels[s.depth], 0); + s.depth += 1; + if (words == 1) break; + bits = words; + } + for (0..n) |i| s.setBit(@intCast(i)); + s.free = n; + return s; + } + + pub fn deinit(s: *FreeSet, gpa: std.mem.Allocator) void { + for (s.levels[0..s.depth]) |l| gpa.free(l); + s.* = .{}; + } + + pub fn count(s: *const FreeSet) u32 { + return s.free; + } + + pub fn contains(s: *const FreeSet, i: u32) bool { + return s.levels[0][i / 64] & (@as(u64, 1) << @intCast(i % 64)) != 0; + } + + /// The lowest free index, taken; null when none is free. + pub fn takeLowest(s: *FreeSet) ?u32 { + if (s.free == 0) return null; + var idx: usize = 0; + var lv: usize = s.depth; + while (lv > 0) { + lv -= 1; + const w = s.levels[lv][idx]; + idx = idx * 64 + @ctz(w); + } + const i: u32 = @intCast(idx); + s.clearBit(i); + s.free -= 1; + return i; + } + + /// Takes a given free index (error when it is not free). + pub fn take(s: *FreeSet, i: u32) error{NotFree}!void { + if (i >= s.n or !s.contains(i)) return error.NotFree; + s.clearBit(i); + s.free -= 1; + } + + /// Gives an index back (error when it is already free: a double free). + pub fn give(s: *FreeSet, i: u32) error{DoubleFree}!void { + if (i >= s.n or s.contains(i)) return error.DoubleFree; + s.setBit(i); + s.free += 1; + } + + fn setBit(s: *FreeSet, i: u32) void { + var idx: usize = i; + for (s.levels[0..s.depth]) |l| { + const was = l[idx / 64]; + l[idx / 64] = was | (@as(u64, 1) << @intCast(idx % 64)); + if (was != 0) return; // the word was non-empty: the levels above already say so + idx /= 64; + } + } + + fn clearBit(s: *FreeSet, i: u32) void { + var idx: usize = i; + for (s.levels[0..s.depth]) |l| { + l[idx / 64] &= ~(@as(u64, 1) << @intCast(idx % 64)); + if (l[idx / 64] != 0) return; // the word still has a free bit: the levels above stay set + idx /= 64; + } + } +}; + +test "lowest first, give back, double free refused, across levels" { + const gpa = std.testing.allocator; + var s = try FreeSet.init(gpa, 70_000); // three levels + defer s.deinit(gpa); + try std.testing.expectEqual(@as(u8, 3), s.depth); + for (0..5000) |i| try std.testing.expectEqual(@as(?u32, @intCast(i)), s.takeLowest()); + try s.give(4097); + try s.give(17); + try std.testing.expectError(error.DoubleFree, s.give(17)); + try std.testing.expectEqual(@as(?u32, 17), s.takeLowest()); + try std.testing.expectEqual(@as(?u32, 4097), s.takeLowest()); + try std.testing.expectEqual(@as(?u32, 5000), s.takeLowest()); + try s.take(69_999); + try std.testing.expectError(error.NotFree, s.take(69_999)); + try std.testing.expectEqual(@as(u32, 70_000 - 5001 - 1), s.count()); + while (s.takeLowest()) |_| {} + try std.testing.expectEqual(@as(u32, 0), s.count()); + try s.give(69_999); + try std.testing.expectEqual(@as(?u32, 69_999), s.takeLowest()); +} diff --git a/zig/src/sessions/jobs.zig b/zig/src/sessions/jobs.zig new file mode 100644 index 0000000000..5b445f20e7 --- /dev/null +++ b/zig/src/sessions/jobs.zig @@ -0,0 +1,150 @@ +//! Parks and restores off the round: a tier thread runs each job's file I/O until the plan settles it (Store.settle). + +const std = @import("std"); +const Io = std.Io; +const disk_mod = @import("disk.zig"); +const Disk = disk_mod.Disk; + +pub const Kind = enum { park, restore }; +pub const State = enum(u8) { queued, running, done, failed }; + +pub const Job = struct { + kind: Kind, + /// the store's entry + entry: u32, + state: std.atomic.Value(State) = .init(.queued), + err: ?anyerror = null, + started_ns: i96 = 0, + elapsed_ns: u64 = 0, + /// park: the file to write (owned buffers below) and its size once written + src: disk_mod.WriteSrc = undefined, + size: u64 = 0, + /// restore: files base first, the sink, the bounded state read + extents: []disk_mod.Extent = &.{}, + sink: disk_mod.ReadSink = undefined, + restored: ?disk_mod.Restored = null, + /// restore: the slot it fills + slot: u32 = 0, + /// pool pages the job holds until settled + held: []u32 = &.{}, + /// owned buffers + ids: []i32 = &.{}, + lists: [][]u32 = &.{}, + blobs: []disk_mod.Blob = &.{}, + + pub fn done(j: *const Job) bool { + const s = j.state.load(.acquire); + return s == .done or s == .failed; + } + + pub fn deinit(j: *Job, gpa: std.mem.Allocator) void { + gpa.free(j.held); + gpa.free(j.ids); + for (j.lists) |l| gpa.free(l); + gpa.free(j.lists); + for (j.blobs) |b| { + gpa.free(b.name); + gpa.free(b.bytes); + } + gpa.free(j.blobs); + gpa.free(j.extents); + if (j.restored) |*r| r.deinit(gpa); + gpa.destroy(j); + } + + fn run(j: *Job, d: *Disk, io: Io) void { + j.state.store(.running, .release); + const t0 = Io.Clock.awake.now(io).toNanoseconds(); + j.err = switch (j.kind) { + .park => if (d.writeFile(&j.src)) |n| blk: { + j.size = n; + break :blk null; + } else |e| e, + .restore => j.restore(d), + }; + j.elapsed_ns = @intCast(Io.Clock.awake.now(io).toNanoseconds() - t0); + j.state.store(if (j.err == null) .done else .failed, .release); + } + + fn restore(j: *Job, d: *Disk) ?anyerror { + for (j.extents, 0..) |e, i| { + const last = i + 1 == j.extents.len; + const r = d.readFile(e, &j.sink, last) catch |err| return err; + if (last) j.restored = r; + } + j.sink.store.publish() catch |err| return err; + return null; + } +}; + +/// The tier thread and its queue (null thread: jobs run inline). +pub const Runner = struct { + gpa: std.mem.Allocator, + io: Io, + disk: *Disk, + thread: ?std.Thread = null, + queue: std.ArrayList(*Job) = .empty, + head: usize = 0, + mutex: Io.Mutex = .init, + cond: Io.Condition = .init, + stop: bool = false, + + pub fn init(gpa: std.mem.Allocator, io: Io, d: *Disk) Runner { + return .{ .gpa = gpa, .io = io, .disk = d }; + } + + /// Starts the tier thread: jobs then run off the caller's thread. + pub fn startThread(r: *Runner) !void { + if (r.thread != null) return; + r.thread = try std.Thread.spawn(.{ .stack_size = 1 << 20 }, loop, .{r}); + } + + pub fn deinit(r: *Runner) void { + if (r.thread) |t| { + r.mutex.lockUncancelable(r.io); + r.stop = true; + r.cond.broadcast(r.io); + r.mutex.unlock(r.io); + t.join(); + } + r.queue.deinit(r.gpa); + } + + pub fn submit(r: *Runner, j: *Job) !void { + if (r.thread == null) return j.run(r.disk, r.io); + r.mutex.lockUncancelable(r.io); + defer r.mutex.unlock(r.io); + if (r.head == r.queue.items.len) { + r.queue.clearRetainingCapacity(); + r.head = 0; + } + try r.queue.append(r.gpa, j); + r.cond.signal(r.io); + } + + /// Waits for a job (the plan settles it next). + pub fn wait(r: *Runner, j: *Job) void { + while (!j.done()) { + r.mutex.lockUncancelable(r.io); + defer r.mutex.unlock(r.io); + if (!j.done()) r.cond.waitTimeout(r.io, &r.mutex, .{ .duration = .{ .raw = .fromMilliseconds(5), .clock = .awake } }) catch {}; + } + } + + fn loop(r: *Runner) void { + while (true) { + const j = blk: { + r.mutex.lockUncancelable(r.io); + defer r.mutex.unlock(r.io); + while (r.head == r.queue.items.len and !r.stop) r.cond.waitUncancelable(r.io, &r.mutex); + if (r.head == r.queue.items.len) return; + r.head += 1; + break :blk r.queue.items[r.head - 1]; + }; + j.run(r.disk, r.io); + r.mutex.lockUncancelable(r.io); + r.cond.broadcast(r.io); + r.mutex.unlock(r.io); + } + } +}; diff --git a/zig/src/sessions/lanes.zig b/zig/src/sessions/lanes.zig new file mode 100644 index 0000000000..38fad8cd89 --- /dev/null +++ b/zig/src/sessions/lanes.zig @@ -0,0 +1,336 @@ +//! The NVMe tier's lane threads, which hash and write (or read and check) 4 KiB chunks, and their staging buffers. + +const std = @import("std"); +const linux = std.os.linux; +const Io = std.Io; +const Sha256 = std.crypto.hash.sha2.Sha256; + +pub const block = 4096; + +pub fn up(n: u64) u64 { + return (n + block - 1) / block * block; +} + +/// Bounded, aligned staging buffers (allocated once; `release` gives the memory back while idle). +pub const Staging = struct { + io: Io, + size: usize, + count: u32, + bufs: [][]align(block) u8, + free_list: []u32, + free_n: u32, + mutex: Io.Mutex = .init, + cond: Io.Condition = .init, + gpa: std.mem.Allocator, + + pub fn init(gpa: std.mem.Allocator, io: Io, count: u32, size: usize) !Staging { + if (count == 0 or size % block != 0) return error.BadStaging; + const bufs = try gpa.alloc([]align(block) u8, count); + errdefer gpa.free(bufs); + const fl = try gpa.alloc(u32, count); + errdefer gpa.free(fl); + for (bufs, 0..) |*b, i| { + b.* = &.{}; + fl[i] = @intCast(i); + } + return .{ .io = io, .size = size, .count = count, .bufs = bufs, .free_list = fl, .free_n = count, .gpa = gpa }; + } + + pub fn deinit(s: *Staging) void { + s.release(); + s.gpa.free(s.bufs); + s.gpa.free(s.free_list); + s.* = undefined; + } + + /// A free buffer's index, waiting for one; its memory is mapped on first use. + pub fn get(s: *Staging) !u32 { + s.mutex.lockUncancelable(s.io); + defer s.mutex.unlock(s.io); + while (s.free_n == 0) s.cond.waitUncancelable(s.io, &s.mutex); + s.free_n -= 1; + const i = s.free_list[s.free_n]; + if (s.bufs[i].len == 0) s.bufs[i] = try s.gpa.alignedAlloc(u8, .fromByteUnits(block), s.size); + return i; + } + + pub fn put(s: *Staging, i: u32) void { + s.mutex.lockUncancelable(s.io); + defer s.mutex.unlock(s.io); + s.free_list[s.free_n] = i; + s.free_n += 1; + s.cond.signal(s.io); + } + + pub fn buf(s: *Staging, i: u32) []align(block) u8 { + return s.bufs[i]; + } + + /// Unmaps every buffer (all must be free): the memory goes back between operations. + pub fn release(s: *Staging) void { + std.debug.assert(s.free_n == s.count); + for (s.bufs) |*b| if (b.len != 0) { + s.gpa.free(b.*); + b.* = &.{}; + }; + } + + pub fn bytes(s: *const Staging) u64 { + return @as(u64, s.count) * s.size; + } +}; + +pub const Kind = enum { write, read }; + +/// One chunk's I/O: buf[0 .. span] at file offset off (span: the chunk padded to 4 KiB; data: the chunk's own bytes). +pub const Job = struct { + kind: Kind, + fd: linux.fd_t, + staging: u32, + data: u32, + span: u32, + off: u64, + /// write: the chunk's digest goes here; read: the digest it must have + digest: *[32]u8, + op: *Op, + /// the pump's tag for a finished read + tag: u32 = 0, +}; + +/// A park's or restore's chunks in flight: how many are pending, the first error, and finished reads for the pump. +pub const Op = struct { + io: Io, + mutex: Io.Mutex = .init, + cond: Io.Condition = .init, + pending: u32 = 0, + err: ?anyerror = null, + /// finished reads (their jobs), drained by the pump + done: std.ArrayList(Job) = .empty, + gpa: std.mem.Allocator, + /// write jobs hand their staging back themselves + staging: *Staging, + + pub fn init(gpa: std.mem.Allocator, io: Io, staging: *Staging) !Op { + var o: Op = .{ .io = io, .gpa = gpa, .staging = staging }; + try o.done.ensureTotalCapacity(gpa, staging.count); + return o; + } + + pub fn deinit(o: *Op) void { + o.done.deinit(o.gpa); + } + + fn finish(o: *Op, j: Job, err: ?anyerror) void { + o.mutex.lockUncancelable(o.io); + defer o.mutex.unlock(o.io); + o.pending -= 1; + if (err) |e| { + if (o.err == null) o.err = e; + } + if (j.kind == .read) o.done.appendAssumeCapacity(j) else o.staging.put(j.staging); + o.cond.broadcast(o.io); + } + + /// The next finished read, waiting for one; null when none is pending. + pub fn next(o: *Op) ?Job { + o.mutex.lockUncancelable(o.io); + defer o.mutex.unlock(o.io); + while (o.done.items.len == 0 and o.pending > 0) o.cond.waitUncancelable(o.io, &o.mutex); + return o.done.pop(); + } + + /// Waits for every pending job; returns the first error. + pub fn wait(o: *Op) !void { + o.mutex.lockUncancelable(o.io); + defer o.mutex.unlock(o.io); + while (o.pending > 0) o.cond.waitUncancelable(o.io, &o.mutex); + if (o.err) |e| return e; + } + + pub fn failed(o: *Op) ?anyerror { + o.mutex.lockUncancelable(o.io); + defer o.mutex.unlock(o.io); + return o.err; + } +}; + +/// The lane threads and their queue. +pub const Lanes = struct { + io: Io, + gpa: std.mem.Allocator, + staging: *Staging, + threads: []std.Thread, + /// queued jobs: a ring as long as the staging set (every queued job holds a buffer) + ring: []Job, + head: usize = 0, + len: usize = 0, + mutex: Io.Mutex = .init, + cond: Io.Condition = .init, + stop: bool = false, + /// seconds the lanes spent hashing and in I/O (summed over lanes) + hash_ns: std.atomic.Value(u64) = .init(0), + io_ns: std.atomic.Value(u64) = .init(0), + + pub fn start(gpa: std.mem.Allocator, io: Io, staging: *Staging, n: u32) !*Lanes { + const l = try gpa.create(Lanes); + errdefer gpa.destroy(l); + const ring = try gpa.alloc(Job, staging.count); + errdefer gpa.free(ring); + l.* = .{ .io = io, .gpa = gpa, .staging = staging, .threads = try gpa.alloc(std.Thread, n), .ring = ring }; + errdefer gpa.free(l.threads); + var made: usize = 0; + errdefer { + l.halt(); + for (l.threads[0..made]) |t| t.join(); + } + for (l.threads) |*t| { + t.* = try std.Thread.spawn(.{ .stack_size = 1 << 20 }, run, .{l}); + made += 1; + } + return l; + } + + fn halt(l: *Lanes) void { + l.mutex.lockUncancelable(l.io); + l.stop = true; + l.cond.broadcast(l.io); + l.mutex.unlock(l.io); + } + + pub fn shutdown(l: *Lanes) void { + l.halt(); + for (l.threads) |t| t.join(); + l.gpa.free(l.ring); + l.gpa.free(l.threads); + l.gpa.destroy(l); + } + + /// Queues a job (each holds a staging buffer, so the ring never overflows). + pub fn submit(l: *Lanes, j: Job) void { + { + j.op.mutex.lockUncancelable(l.io); + defer j.op.mutex.unlock(l.io); + j.op.pending += 1; + } + l.mutex.lockUncancelable(l.io); + defer l.mutex.unlock(l.io); + std.debug.assert(l.len < l.ring.len); + l.ring[(l.head + l.len) % l.ring.len] = j; + l.len += 1; + l.cond.signal(l.io); + } + + fn run(l: *Lanes) void { + while (true) { + const j = blk: { + l.mutex.lockUncancelable(l.io); + defer l.mutex.unlock(l.io); + while (l.len == 0 and !l.stop) l.cond.waitUncancelable(l.io, &l.mutex); + if (l.len == 0) return; + const j = l.ring[l.head]; + l.head = (l.head + 1) % l.ring.len; + l.len -= 1; + break :blk j; + }; + j.op.finish(j, l.work(j)); + } + } + + fn work(l: *Lanes, j: Job) ?anyerror { + const b = l.staging.buf(j.staging); + var t0 = std.Io.Clock.awake.now(l.io).toNanoseconds(); + switch (j.kind) { + .write => { + Sha256.hash(b[0..j.data], j.digest, .{}); + const t1 = std.Io.Clock.awake.now(l.io).toNanoseconds(); + _ = l.hash_ns.fetchAdd(@intCast(t1 - t0), .monotonic); + t0 = t1; + @memset(b[j.data..j.span], 0); + pwriteAll(j.fd, b[0..j.span], j.off) catch |e| return e; + _ = l.io_ns.fetchAdd(@intCast(std.Io.Clock.awake.now(l.io).toNanoseconds() - t0), .monotonic); + }, + .read => { + preadAll(j.fd, b[0..j.span], j.off) catch |e| return e; + const t1 = std.Io.Clock.awake.now(l.io).toNanoseconds(); + _ = l.io_ns.fetchAdd(@intCast(t1 - t0), .monotonic); + var d: [32]u8 = undefined; + Sha256.hash(b[0..j.data], &d, .{}); + _ = l.hash_ns.fetchAdd(@intCast(std.Io.Clock.awake.now(l.io).toNanoseconds() - t1), .monotonic); + if (!std.mem.eql(u8, &d, j.digest)) return error.ChecksumMismatch; + }, + } + return null; + } +}; + +pub fn pwriteAll(fd: linux.fd_t, data: []const u8, off: u64) !void { + var done: usize = 0; + while (done < data.len) { + const rc = linux.pwrite(fd, data[done..].ptr, data.len - done, @intCast(off + done)); + switch (linux.errno(rc)) { + .SUCCESS => {}, + .INTR => continue, + .NOSPC => return error.NoSpaceLeft, + else => return error.WriteFailed, + } + if (rc == 0) return error.WriteFailed; + done += rc; + } +} + +pub fn preadAll(fd: linux.fd_t, out: []u8, off: u64) !void { + var done: usize = 0; + while (done < out.len) { + const rc = linux.pread(fd, out[done..].ptr, out.len - done, @intCast(off + done)); + switch (linux.errno(rc)) { + .SUCCESS => {}, + .INTR => continue, + else => return error.ReadFailed, + } + if (rc == 0) return error.ShortRead; + done += rc; + } +} + +test "lanes write, read back and check chunks through bounded staging" { + const gpa = std.testing.allocator; + const io = std.testing.io; + var st = try Staging.init(gpa, io, 3, 2 * block); + defer st.deinit(); + const l = try Lanes.start(gpa, io, &st, 2); + defer l.shutdown(); + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + const f = try tmp.dir.createFile(io, "lanes.bin", .{ .read = true }); + defer f.close(io); + const fd = f.handle; + var digests: [7][32]u8 = undefined; + var op = try Op.init(gpa, io, &st); + defer op.deinit(); + for (0..7) |i| { + const s = try st.get(); + const n: u32 = if (i == 6) 1000 else 2 * block - 7; + for (st.buf(s)[0..n], 0..) |*b, k| b.* = @truncate(k * 31 + i); + l.submit(.{ .kind = .write, .fd = fd, .staging = s, .data = n, .span = @intCast(up(n)), .off = i * 2 * block, .digest = &digests[i], .op = &op }); + } + try op.wait(); + var rop = try Op.init(gpa, io, &st); + defer rop.deinit(); + var seen: u32 = 0; + var sent: u32 = 0; + while (seen < 7) { + while (sent < 7 and sent - seen < 3) : (sent += 1) { + const n: u32 = if (sent == 6) 1000 else 2 * block - 7; + l.submit(.{ .kind = .read, .fd = fd, .staging = try st.get(), .data = n, .span = @intCast(up(n)), .off = sent * 2 * block, .digest = &digests[sent], .op = &rop, .tag = sent }); + } + const j = rop.next().?; + for (st.buf(j.staging)[0..j.data], 0..) |b, k| try std.testing.expectEqual(@as(u8, @truncate(k * 31 + j.tag)), b); + st.put(j.staging); + seen += 1; + } + try rop.wait(); + digests[2][0] ^= 1; // a damaged chunk is refused + l.submit(.{ .kind = .read, .fd = fd, .staging = try st.get(), .data = 2 * block - 7, .span = 2 * block, .off = 2 * 2 * block, .digest = &digests[2], .op = &rop }); + st.put(rop.next().?.staging); + try std.testing.expectError(error.ChecksumMismatch, rop.wait()); +} diff --git a/zig/src/sessions/pagestore.zig b/zig/src/sessions/pagestore.zig new file mode 100644 index 0000000000..0e4792699b --- /dev/null +++ b/zig/src/sessions/pagestore.zig @@ -0,0 +1,129 @@ +//! The bytes behind the pool's pages: a family's tensor of pages, read and written a page run at a time. + +const std = @import("std"); +const pool_mod = @import("pool.zig"); +const Pool = pool_mod.Pool; +const Family = pool_mod.Family; + +pub const PageStore = struct { + ptr: *anyopaque, + vtable: *const VTable, + + pub const VTable = struct { + /// family fam's pages, back to back into dst (pages.len x a page's bytes) + read: *const fn (ptr: *anyopaque, fam: u32, pages: []const u32, dst: []u8) anyerror!void, + /// src (pages.len x a page's bytes) into family fam's pages + write: *const fn (ptr: *anyopaque, fam: u32, pages: []const u32, src: []const u8) anyerror!void, + /// one page's rows of family fam into another page (copy on write of a RAM resume's partial page) + copyPage: *const fn (ptr: *anyopaque, fam: u32, src: u32, dst: u32) anyerror!void, + /// round thread: a point after every row the compute stream has written so far + fence: *const fn (ptr: *anyopaque) anyerror!u64, + /// before reads: wait until the fence's rows are visible + waitFence: *const fn (ptr: *anyopaque, token: u64) anyerror!void, + /// after writes: make them visible to the compute stream + publish: *const fn (ptr: *anyopaque) anyerror!void, + }; + + pub fn read(s: PageStore, fam: u32, pages: []const u32, dst: []u8) !void { + return s.vtable.read(s.ptr, fam, pages, dst); + } + pub fn write(s: PageStore, fam: u32, pages: []const u32, src: []const u8) !void { + return s.vtable.write(s.ptr, fam, pages, src); + } + pub fn copyPage(s: PageStore, fam: u32, src: u32, dst: u32) !void { + return s.vtable.copyPage(s.ptr, fam, src, dst); + } + pub fn fence(s: PageStore) !u64 { + return s.vtable.fence(s.ptr); + } + pub fn waitFence(s: PageStore, token: u64) !void { + return s.vtable.waitFence(s.ptr, token); + } + pub fn publish(s: PageStore) !void { + return s.vtable.publish(s.ptr); + } +}; + +/// Every family's pages in host memory: [familyTensorPages x page bytes] a family, zeroed (the null page stays zero). +pub const HostPool = struct { + gpa: std.mem.Allocator, + page: u32, + families: []const Family, + tensors: [][]u8, + + pub fn init(gpa: std.mem.Allocator, p: *const Pool) !HostPool { + const t = try gpa.alloc([]u8, p.layout.families.len); + errdefer gpa.free(t); + var made: usize = 0; + errdefer for (t[0..made]) |x| gpa.free(x); + for (p.layout.families, 0..) |f, i| { + t[i] = try gpa.alloc(u8, @as(usize, p.familyTensorPages(f)) * f.pageBytes(p.page)); + @memset(t[i], 0); + made += 1; + } + return .{ .gpa = gpa, .page = p.page, .families = p.layout.families, .tensors = t }; + } + + pub fn deinit(h: *HostPool) void { + for (h.tensors) |x| h.gpa.free(x); + h.gpa.free(h.tensors); + h.* = undefined; + } + + pub fn store(h: *HostPool) PageStore { + return .{ .ptr = h, .vtable = &vtable }; + } + + /// One page of family fam (a view). + pub fn pageOf(h: *HostPool, fam: u32, pg: u32) []u8 { + const n: usize = @intCast(h.families[fam].pageBytes(h.page)); + return h.tensors[fam][pg * n ..][0..n]; + } + + const vtable: PageStore.VTable = .{ .read = read, .write = write, .copyPage = copyPage, .fence = fence, .waitFence = waitFence, .publish = publish }; + + fn read(ptr: *anyopaque, fam: u32, pages: []const u32, dst: []u8) anyerror!void { + const h: *HostPool = @ptrCast(@alignCast(ptr)); + const n: usize = @intCast(h.families[fam].pageBytes(h.page)); + if (dst.len != pages.len * n) return error.BadLength; + for (pages, 0..) |pg, i| @memcpy(dst[i * n ..][0..n], h.pageOf(fam, pg)); + } + + fn write(ptr: *anyopaque, fam: u32, pages: []const u32, src: []const u8) anyerror!void { + const h: *HostPool = @ptrCast(@alignCast(ptr)); + const n: usize = @intCast(h.families[fam].pageBytes(h.page)); + if (src.len != pages.len * n) return error.BadLength; + for (pages, 0..) |pg, i| @memcpy(h.pageOf(fam, pg), src[i * n ..][0..n]); + } + + fn copyPage(ptr: *anyopaque, fam: u32, src: u32, dst: u32) anyerror!void { + const h: *HostPool = @ptrCast(@alignCast(ptr)); + @memcpy(h.pageOf(fam, dst), h.pageOf(fam, src)); + } + + fn fence(_: *anyopaque) anyerror!u64 { + return 0; + } + fn waitFence(_: *anyopaque, _: u64) anyerror!void {} + fn publish(_: *anyopaque) anyerror!void {} +}; + +test "host pool pages round-trip through the interface" { + const gpa = std.testing.allocator; + const fams = [_]Family{ .{ .name = "comp", .ratio = 1, .row_bytes = 8 }, .{ .name = "ik", .ratio = 2, .row_bytes = 4 } }; + var p = try Pool.init(gpa, .{ .families = &fams, .page = 16 }, 64, 1, 0); + defer p.deinit(); + var h = try HostPool.init(gpa, &p); + defer h.deinit(); + const s = h.store(); + var src: [2 * 128]u8 = undefined; + for (&src, 0..) |*b, i| b.* = @truncate(i * 7 + 3); + try s.write(0, &.{ 3, 1 }, &src); + var dst: [2 * 128]u8 = undefined; + try s.read(0, &.{ 3, 1 }, &dst); + try std.testing.expectEqualSlices(u8, &src, &dst); + try s.copyPage(0, 3, 0); + try std.testing.expectEqualSlices(u8, src[0..128], h.pageOf(0, 0)); + try std.testing.expect(std.mem.allEqual(u8, h.pageOf(0, p.null_page), 0)); + try std.testing.expectError(error.BadLength, s.read(1, &.{0}, &dst)); +} diff --git a/zig/src/sessions/pool.zig b/zig/src/sessions/pool.zig new file mode 100644 index 0000000000..d3e67516af --- /dev/null +++ b/zig/src/sessions/pool.zig @@ -0,0 +1,475 @@ +//! The paged KV pool's host book: lowest-free pages per residue mod W, holder counts, slot page tables, quotas. + +const std = @import("std"); +const FreeSet = @import("freeset.zig").FreeSet; + +pub const Error = error{ PoolExhausted, PastCapacity, SlotNotEmpty, DoubleFree, BadPool, ResidueMismatch } || std.mem.Allocator.Error; + +/// A cache family: rows of `row_bytes` every `ratio` tokens; `split` families shard by logical page over the ranks. +pub const Family = struct { + name: []const u8, + ratio: u32, + row_bytes: u32, + split: bool = false, + + pub fn pageRows(f: Family, page: u32) u32 { + return page / f.ratio; + } + + /// Bytes of one page of this family (a whole page's rows). + pub fn pageBytes(f: Family, page: u32) u64 { + return @as(u64, f.pageRows(page)) * f.row_bytes; + } +}; + +/// The pool's families and page size: what a token costs and how pages are laid out. +pub const Layout = struct { + families: []const Family, + /// tokens a page (a power of two; 256 for every family we serve) + page: u32 = 256, + + /// A rank's bytes a token: the split families count 1 / W. + pub fn bytesPerToken(l: Layout, world: u32) f64 { + var b: f64 = 0; + for (l.families) |f| b += @as(f64, @floatFromInt(f.row_bytes)) / @as(f64, @floatFromInt(f.ratio)) / @as(f64, @floatFromInt(if (f.split and world > 1) world else 1)); + return b; + } + + /// A rank's average bytes a page (split families 1 / W): what the residency budget counts. + pub fn pageBytes(l: Layout, world: u32) u64 { + var b: u64 = 0; + for (l.families) |f| b += f.pageBytes(l.page) / (if (f.split and world > 1) world else 1); + return b; + } +}; + +pub fn pagesFor(tokens: u64, page: u32) u32 { + return @intCast((tokens + page - 1) / page); +} + +/// Logical pages [k0, k1) counted by residue mod `w` (O(w)). +pub fn span(k0: u32, k1: u32, w: u32, out: []u32) void { + for (out[0..w], 0..) |*o, r| { + const rr: u32 = @intCast(r); + const first = k0 + (rr + w - k0 % w) % w; + o.* = if (first >= k1) 0 else (k1 - 1 - first) / w + 1; + } +} + +pub const max_world = 8; + +/// A file whose logical pages [0, pages) match a slot's or entry's rows: a later park writes only past it. +pub const Origin = struct { + key: [16]u8, + pages: u32, +}; + +/// One slot's page table: logical page i -> physical page; the device tables follow through the dirty range. +pub const Slot = struct { + pool: *Pool, + id: u32, + capacity: u64, + max_pages: u32, + /// physical page of each mapped logical page; len is the mapped count + pages: []u32, + len: u32 = 0, + /// pages the admission reserved for this slot's request (null: no admission control) + quota: ?u32 = null, + /// pages taken past the quota, from pages nobody had reserved + over: u32 = 0, + /// table entries changed since the device side last synced them: [lo, hi) + dirty_lo: u32 = std.math.maxInt(u32), + dirty_hi: u32 = 0, + /// the file the leading pages came from (truncation shortens it, release clears it) + origin: ?Origin = null, + + pub fn mapped(s: *const Slot) []const u32 { + return s.pages[0..s.len]; + } + + pub fn tokens(s: *const Slot) u64 { + return @as(u64, s.len) * s.pool.page; + } + + fn mark(s: *Slot, lo: u32, hi: u32) void { + s.dirty_lo = @min(s.dirty_lo, lo); + s.dirty_hi = @max(s.dirty_hi, hi); + } + + /// The table range to upload since the last call (empty when nothing changed): entries past len are the null page. + pub fn takeDirty(s: *Slot) ?struct { lo: u32, hi: u32 } { + if (s.dirty_lo >= s.dirty_hi) return null; + defer { + s.dirty_lo = std.math.maxInt(u32); + s.dirty_hi = 0; + } + return .{ .lo = s.dirty_lo, .hi = s.dirty_hi }; + } + + /// The device table's value at logical page i (the page, or the null page past len). + pub fn tableAt(s: *const Slot, i: u32) u32 { + return if (i < s.len) s.pages[i] else s.pool.null_page; + } + + /// The split families' device table at logical page i: the local page this rank stores, else the discard page. + pub fn localTableAt(s: *const Slot, i: u32) u32 { + return s.pool.localPage(s.tableAt(i)); + } + + fn reservedSpan(s: *const Slot, out: []u32) void { + const q = s.quota orelse 0; + span(s.len, @max(s.len, q), s.pool.world, out); + } + + /// Maps every page holding positions < end (a no-op when they are: the common case). + pub fn ensure(s: *Slot, end: u64) Error!void { + const p = s.pool; + const need = pagesFor(end, p.page); + if (need <= s.len) return; + if (need > s.max_pages) return error.PastCapacity; + const w = p.world; + var before: [max_world]u32 = undefined; + s.reservedSpan(&before); + if (s.quota) |q| if (need > q) { + // past the reservation: only pages nobody else reserved, residue by residue + var extra: [max_world]u32 = undefined; + span(@max(q, s.len), need, w, &extra); + for (0..w) |r| if (extra[r] > p.free[r].count() -| (p.reserved[r] - before[r])) return error.PoolExhausted; + s.over += need - @max(q, s.len); + }; + var want: [max_world]u32 = undefined; + span(s.len, need, w, &want); + for (0..w) |r| if (want[r] > p.free[r].count()) return error.PoolExhausted; + const lo = s.len; + for (lo..need) |k| { + const pg = p.takeAt(@intCast(k)) catch unreachable; // counted above + p.hold(pg, 0, 1); // the taker's hold becomes the slot's + s.pages[k] = pg; + } + s.len = need; + p.unreserve(&before); + var after: [max_world]u32 = undefined; + s.reservedSpan(&after); + p.addReserved(&after); + s.mark(lo, need); + } + + /// Maps shared pages as the first pages of an empty slot (a RAM resume: read only, the slot writes past them). + pub fn adopt(s: *Slot, shared: []const u32) Error!void { + if (s.len != 0) return error.SlotNotEmpty; + if (shared.len > s.max_pages) return error.PastCapacity; + const p = s.pool; + var before: [max_world]u32 = undefined; + s.reservedSpan(&before); + for (shared, 0..) |pg, k| if (pg >= p.npages or pg % p.world != k % p.world) return error.ResidueMismatch; + for (shared, 0..) |pg, k| { + p.hold(pg, 1, 1); + s.pages[k] = pg; + } + s.len = @intCast(shared.len); + p.unreserve(&before); + var after: [max_world]u32 = undefined; + s.reservedSpan(&after); + p.addReserved(&after); + s.mark(0, s.len); + } + + /// Appends one page taken for logical page len (the copy-on-write page of a RAM resume). + pub fn push(s: *Slot, pg: u32) Error!void { + if (s.len >= s.max_pages) return error.PastCapacity; + const p = s.pool; + if (p.world > 1 and pg % p.world != s.len % p.world) return error.ResidueMismatch; + var before: [max_world]u32 = undefined; + s.reservedSpan(&before); + p.hold(pg, 0, 1); // the taker's hold becomes the slot's + s.pages[s.len] = pg; + s.len += 1; + p.unreserve(&before); + var after: [max_world]u32 = undefined; + s.reservedSpan(&after); + p.addReserved(&after); + s.mark(s.len - 1, s.len); + } + + /// Lets go of the pages past position keep (rows nothing reads again); returns how many. + pub fn truncate(s: *Slot, keep: u64) Error!u32 { + const p = s.pool; + // positions >= keep get written again: a page holding one no longer equals the origin file's + if (s.origin) |*o| o.pages = @min(o.pages, @as(u32, @intCast(keep / p.page))); + const k = pagesFor(keep, p.page); + if (k >= s.len) return 0; + var before: [max_world]u32 = undefined; + s.reservedSpan(&before); + const n = s.len - k; + for (s.pages[k..s.len]) |pg| try p.release(pg, 1); + const old = s.len; + s.len = k; + p.unreserve(&before); + var after: [max_world]u32 = undefined; + s.reservedSpan(&after); + p.addReserved(&after); + s.mark(k, old); + return n; + } + + /// Sets the admission's reservation (null: none). + pub fn reserve(s: *Slot, quota: ?u32) void { + var before: [max_world]u32 = undefined; + s.reservedSpan(&before); + s.quota = quota; + s.pool.unreserve(&before); + var after: [max_world]u32 = undefined; + s.reservedSpan(&after); + s.pool.addReserved(&after); + } + + /// Frees the slot's pages and its reservation. + pub fn releaseAll(s: *Slot) Error!u32 { + s.reserve(null); + s.over = 0; + s.origin = null; + return s.truncate(0); + } +}; + +/// The pool's book: free sets by residue, holders a page, the slots and their reservations. +pub const Pool = struct { + gpa: std.mem.Allocator, + layout: Layout, + npages: u32, + page: u32, + /// log2 of page tokens + shift: u5, + /// ranks the split families are sharded over (1: every rank stores every row) + world: u32, + rank: u32, + /// the page unmapped table entries point at (all zeros: a stray read reads zeros) + null_page: u32, + /// a split family's local page count (npages / W), its local null page and its discard page (other ranks' rows) + local_pages: u32, + local_null: u32, + discard: u32, + free: [max_world]FreeSet = @splat(.{}), + /// pages each residue's slots reserved and have not mapped yet + reserved: [max_world]u32 = @splat(0), + /// holders a page (slots + entries + in-flight jobs) and how many are slots + refs: []u16, + slot_refs: []u16, + /// pages held by something and by no slot (session entries alone): what parking can free + entry_only: u32 = 0, + slots: std.ArrayList(*Slot) = .empty, + + pub fn init(gpa: std.mem.Allocator, layout: Layout, tokens: u64, world: u32, rank: u32) Error!Pool { + if (layout.page == 0 or layout.page & (layout.page - 1) != 0) return error.BadPool; + if (tokens % layout.page != 0 or tokens == 0) return error.BadPool; + if (world == 0 or world > max_world or rank >= world) return error.BadPool; + const npages: u32 = @intCast(tokens / layout.page); + if (npages % world != 0) return error.BadPool; + for (layout.families) |f| if (f.ratio == 0 or layout.page % f.ratio != 0) return error.BadPool; + var p: Pool = .{ + .gpa = gpa, + .layout = layout, + .npages = npages, + .page = layout.page, + .shift = @intCast(std.math.log2_int(u32, layout.page)), + .world = world, + .rank = rank, + .null_page = npages, + .local_pages = npages / world, + .local_null = npages / world, + .discard = npages / world + 1, + .refs = &.{}, + .slot_refs = &.{}, + }; + errdefer p.deinit(); + for (0..world) |r| p.free[r] = try FreeSet.init(gpa, npages / world); + p.refs = try gpa.alloc(u16, npages); + @memset(p.refs, 0); + p.slot_refs = try gpa.alloc(u16, npages); + @memset(p.slot_refs, 0); + return p; + } + + pub fn deinit(p: *Pool) void { + for (p.slots.items) |s| { + p.gpa.free(s.pages); + p.gpa.destroy(s); + } + p.slots.deinit(p.gpa); + for (0..p.world) |r| p.free[r].deinit(p.gpa); + p.gpa.free(p.refs); + p.gpa.free(p.slot_refs); + p.* = undefined; + } + + pub fn newSlot(p: *Pool, capacity: u64) Error!*Slot { + const s = try p.gpa.create(Slot); + errdefer p.gpa.destroy(s); + const mp = pagesFor(capacity, p.page); + s.* = .{ .pool = p, .id = @intCast(p.slots.items.len), .capacity = capacity, .max_pages = mp, .pages = try p.gpa.alloc(u32, mp) }; + errdefer p.gpa.free(s.pages); + try p.slots.append(p.gpa, s); + return s; + } + + /// A free physical page for logical page k (its residue when split), held once by the caller. + pub fn takeAt(p: *Pool, k: u32) Error!u32 { + const r = k % p.world; + const local = p.free[r].takeLowest() orelse return error.PoolExhausted; + const pg = local * p.world + r; + p.hold(pg, 1, 0); // the caller's hold (an entry-like holder until a slot maps it) + return pg; + } + + /// Holders change on page pg: dref total, dslot of them slots; keeps entry_only and gives the page back at zero. + fn hold(p: *Pool, pg: u32, dref: i32, dslot: i32) void { + const r0 = p.refs[pg]; + const s0 = p.slot_refs[pg]; + const r1: u16 = @intCast(@as(i32, r0) + dref); + const s1: u16 = @intCast(@as(i32, s0) + dslot); + p.refs[pg] = r1; + p.slot_refs[pg] = s1; + const was = r0 > 0 and s0 == 0; + const now = r1 > 0 and s1 == 0; + if (was and !now) p.entry_only -= 1; + if (!was and now) p.entry_only += 1; + } + + /// One holder lets go (slot: from a slot's table); the page returns to its free set when nobody holds it. + fn release(p: *Pool, pg: u32, slot: u1) Error!void { + if (p.refs[pg] == 0) return error.DoubleFree; + p.hold(pg, -1, -@as(i32, slot)); + if (p.refs[pg] == 0) p.free[pg % p.world].give(pg / p.world) catch return error.DoubleFree; + } + + /// One more holder (a session entry or an in-flight job) of each page. + pub fn share(p: *Pool, pages: []const u32) void { + for (pages) |pg| p.hold(pg, 1, 0); + } + + /// Those holders let go. + pub fn drop(p: *Pool, pages: []const u32) Error!void { + for (pages) |pg| try p.release(pg, 0); + } + + pub fn holders(p: *const Pool, pg: u32) u16 { + return p.refs[pg]; + } + + fn unreserve(p: *Pool, by: []const u32) void { + for (0..p.world) |r| p.reserved[r] -= by[r]; + } + + fn addReserved(p: *Pool, by: []const u32) void { + for (0..p.world) |r| p.reserved[r] += by[r]; + } + + /// Free pages no admission reserved (split: W x the scarcer residue, so any logical run fits). + pub fn available(p: *const Pool) u32 { + var low: u32 = std.math.maxInt(u32); + for (0..p.world) |r| low = @min(low, p.free[r].count() -| p.reserved[r]); + return low * p.world; + } + + pub fn freePages(p: *const Pool) u32 { + var n: u32 = 0; + for (0..p.world) |r| n += p.free[r].count(); + return n; + } + + /// Pages an admission needs for a request (prompt + reply + slack, at most the slot's capacity). + pub fn needPages(p: *const Pool, prompt: u64, max_tokens: u64, capacity: u64, slack: u64) u32 { + return pagesFor(@min(prompt + max_tokens + slack, capacity), p.page); + } + + /// This rank stores the split families' rows of page pg. + pub fn owns(p: *const Pool, pg: u32) bool { + return pg != p.null_page and pg % p.world == p.rank; + } + + /// A split family's local page of pg: its own when owned, the local null for the null page, else the discard page. + pub fn localPage(p: *const Pool, pg: u32) u32 { + if (pg == p.null_page) return p.local_null; + return if (pg % p.world == p.rank) pg / p.world else p.discard; + } + + /// The page a family's tensor indexes for physical page pg (split families: localPage). + pub fn familyPage(p: *const Pool, f: Family, pg: u32) u32 { + return if (f.split and p.world > 1) p.localPage(pg) else pg; + } + + /// Pages a family's device tensor holds: npages + null, or (split) local pages + local null + discard page. + pub fn familyTensorPages(p: *const Pool, f: Family) u32 { + return if (f.split and p.world > 1) p.local_pages + 2 else p.npages + 1; + } + + /// Bytes this rank stores of family f for logical pages [first, first + n) of a slot or entry. + pub fn familyBytes(p: *const Pool, f: Family, first: u32, n: u32) u64 { + if (!(f.split and p.world > 1)) return @as(u64, n) * f.pageBytes(p.page); + var by: [max_world]u32 = undefined; + span(first, first + n, p.world, &by); + return @as(u64, by[p.rank]) * f.pageBytes(p.page); + } + + /// Bytes of every family's device tensor on this rank. + pub fn tensorBytes(p: *const Pool) u64 { + var b: u64 = 0; + for (p.layout.families) |f| b += @as(u64, p.familyTensorPages(f)) * f.pageBytes(p.page); + return b; + } + + /// Pages held by session entries alone and their bytes on this rank (the residency budget's measure): O(1). + pub fn entryOnlyBytes(p: *const Pool) u64 { + return @as(u64, p.entry_only) * p.layout.pageBytes(p.world); + } +}; + +test "residues, holders, quotas and the device tables" { + const gpa = std.testing.allocator; + const fams = [_]Family{ .{ .name = "comp", .ratio = 2, .row_bytes = 584, .split = true }, .{ .name = "index_k", .ratio = 2, .row_bytes = 132 } }; + var p = try Pool.init(gpa, .{ .families = &fams }, 16 * 256, 2, 1); + defer p.deinit(); + try std.testing.expectEqual(@as(u32, 8 + 2), p.familyTensorPages(fams[0])); + try std.testing.expectEqual(@as(u32, 16 + 1), p.familyTensorPages(fams[1])); + try std.testing.expectApproxEqAbs(@as(f64, 146 + 66), p.layout.bytesPerToken(2), 1e-9); + const a = try p.newSlot(8 * 256); + const b = try p.newSlot(8 * 256); + try a.ensure(3 * 256 + 1); // 4 pages: residues 0,1,0,1 + try std.testing.expectEqualSlices(u32, &.{ 0, 1, 2, 3 }, a.mapped()); + try b.ensure(256); + try std.testing.expectEqualSlices(u32, &.{4}, b.mapped()); + try std.testing.expectEqual(p.discard, a.localTableAt(0)); // page 0 is rank 0's + try std.testing.expectEqual(@as(u32, 0), a.localTableAt(1)); // page 1 -> rank 1's local page 0 + try std.testing.expectEqual(p.local_null, a.localTableAt(7)); + const d = a.takeDirty().?; + try std.testing.expectEqual(@as(u32, 0), d.lo); + try std.testing.expectEqual(@as(u32, 4), d.hi); + try std.testing.expect(a.takeDirty() == null); + // an entry shares a's pages; a lets go; the pages stay held by the entry alone + p.share(a.mapped()[0..3]); + try std.testing.expectEqual(@as(u32, 0), p.entry_only); + _ = try a.releaseAll(); + try std.testing.expectEqual(@as(u32, 3), p.entry_only); + try std.testing.expectEqual(@as(u32, 16 - 1 - 3), p.freePages()); + try p.drop(&.{ 0, 1, 2 }); + try std.testing.expectEqual(@as(u32, 0), p.entry_only); + try std.testing.expectError(error.DoubleFree, p.drop(&.{0})); + // quotas: b reserves 6 pages; available counts W x the scarcer residue + b.reserve(6); + var by: [2]u32 = undefined; + span(1, 6, 2, &by); + try std.testing.expectEqual(@as(u32, 2), by[0]); + try std.testing.expectEqual(@as(u32, 3), by[1]); + try std.testing.expectEqual(@as(u32, 2 * @min(7 - 2, 8 - 3)), p.available()); + try b.ensure(6 * 256); + try std.testing.expectEqual(@as(u32, 0), p.reserved[0] + p.reserved[1]); + for (b.mapped(), 0..) |pg, k| try std.testing.expectEqual(k % 2, pg % 2); + // RAM resume: adopt full pages of the right residues, refuse the wrong ones + const c = try p.newSlot(8 * 256); + try std.testing.expectError(error.ResidueMismatch, c.adopt(&.{b.pages[1]})); + p.share(b.mapped()[0..2]); + try c.adopt(b.mapped()[0..2]); + try c.push(try p.takeAt(2)); + try std.testing.expectEqual(@as(u16, 3), p.holders(b.pages[0])); + try std.testing.expectEqual(@as(u32, 0), c.pages[2] % 2); +} diff --git a/zig/src/sessions/prefix.zig b/zig/src/sessions/prefix.zig new file mode 100644 index 0000000000..28c639fbd7 --- /dev/null +++ b/zig/src/sessions/prefix.zig @@ -0,0 +1,457 @@ +//! Session entries by prefix: a trie over chained page digests, one node per distinct run of full pages a tag. + +const std = @import("std"); + +pub const Digest = [16]u8; +pub const none = std.math.maxInt(u32); +const Blake = std.crypto.hash.blake2.Blake2b128; + +/// Chained digests of the full pages of ids under a tag: out[j] determines tag and ids[0 .. page (j + 1)]. +pub fn chain(tag: u32, ids: []const i32, page: u32, out: []Digest) void { + var prev = root(tag); + for (out, 0..) |*o, j| { + o.* = step(prev, ids[j * page ..][0..page]); + prev = o.*; + } +} + +/// The digest of a tag's empty prefix. +pub fn root(tag: u32) Digest { + var h = Blake.init(.{}); + h.update("tf-sess-root"); + h.update(std.mem.asBytes(&tag)); + return h.finalResult(); +} + +/// The next page's digest. +pub fn step(prev: Digest, page_ids: []const i32) Digest { + var h = Blake.init(.{}); + h.update(&prev); + h.update(std.mem.sliceAsBytes(page_ids)); + return h.finalResult(); +} + +/// An entry's identity: its node's digest (tag and full pages), its length and tail. +pub fn entryKey(node: Digest, pos: u64, tail: []const i32) Digest { + var h = Blake.init(.{}); + h.update("tf-sess-entry"); + h.update(&node); + h.update(std.mem.asBytes(&pos)); + h.update(std.mem.sliceAsBytes(tail)); + return h.finalResult(); +} + +pub const Tier = enum { ram, any }; + +pub fn Index(comptime Extra: type) type { + return struct { + const Self = @This(); + + pub const Node = struct { + digest: Digest, + parent: u32, + depth: u32, + first_child: u32 = none, + next_sib: u32 = none, + prev_sib: u32 = none, + first_entry: u32 = none, + /// entries in the subtree (this node's included): all, and those in RAM + sub_all: u32 = 0, + sub_ram: u32 = 0, + /// RAM entries at this node itself + own_ram: u32 = 0, + /// the tokens of the page leading into this node, while sub_ram > 0 + tokens: ?[]i32 = null, + }; + + pub const Entry = struct { + key: Digest, + tag: u32, + pos: u64, + node: u32, + tail: []i32, + ram: bool, + next_at: u32 = none, + prev_at: u32 = none, + live: bool = true, + extra: Extra, + }; + + /// An entry inserted, or the one already there with the same key. + pub const Placed = struct { id: u32, new: bool }; + + gpa: std.mem.Allocator, + page: u32, + nodes: std.ArrayList(Node) = .empty, + free_nodes: std.ArrayList(u32) = .empty, + by_digest: std.AutoHashMapUnmanaged(Digest, u32) = .empty, + entries: std.ArrayList(Entry) = .empty, + free_entries: std.ArrayList(u32) = .empty, + by_key: std.AutoHashMapUnmanaged(Digest, u32) = .empty, + /// page-sized token blocks, reused + blocks: std.ArrayList([]i32) = .empty, + + pub fn init(gpa: std.mem.Allocator, page: u32) Self { + return .{ .gpa = gpa, .page = page }; + } + + pub fn deinit(x: *Self) void { + for (x.nodes.items) |n| if (n.tokens) |t| x.gpa.free(t); + for (x.blocks.items) |b| x.gpa.free(b); + for (x.entries.items) |e| if (e.live) x.gpa.free(e.tail); + x.nodes.deinit(x.gpa); + x.free_nodes.deinit(x.gpa); + x.by_digest.deinit(x.gpa); + x.entries.deinit(x.gpa); + x.free_entries.deinit(x.gpa); + x.by_key.deinit(x.gpa); + x.blocks.deinit(x.gpa); + x.* = undefined; + } + + pub fn get(x: *Self, e: u32) *Entry { + return &x.entries.items[e]; + } + + pub fn node(x: *Self, n: u32) *Node { + return &x.nodes.items[n]; + } + + pub fn lookup(x: *const Self, key: Digest) ?u32 { + return x.by_key.get(key); + } + + fn newNode(x: *Self, digest: Digest, parent: u32, depth: u32) !u32 { + const id: u32 = if (x.free_nodes.pop()) |i| i else blk: { + try x.nodes.append(x.gpa, undefined); + break :blk @intCast(x.nodes.items.len - 1); + }; + x.nodes.items[id] = .{ .digest = digest, .parent = parent, .depth = depth }; + try x.by_digest.put(x.gpa, digest, id); + if (parent != none) { + const p = &x.nodes.items[parent]; + x.nodes.items[id].next_sib = p.first_child; + if (p.first_child != none) x.nodes.items[p.first_child].prev_sib = id; + p.first_child = id; + } + return id; + } + + fn takeBlock(x: *Self) ![]i32 { + if (x.blocks.pop()) |b| return b; + return x.gpa.alloc(i32, x.page); + } + + /// Adds an entry for ids (chain: its full pages' digests), or returns the one with the same key. + pub fn insert(x: *Self, tag: u32, ids: []const i32, digests: []const Digest, ram: bool, extra: Extra) !Placed { + const full: u32 = @intCast(ids.len / x.page); + std.debug.assert(digests.len >= full); + const r = root(tag); + var n = x.by_digest.get(r) orelse try x.newNode(r, none, 0); + for (0..full) |j| { + const d = digests[j]; + n = x.by_digest.get(d) orelse try x.newNode(d, n, @intCast(j + 1)); + } + const r2 = try x.place(tag, n, ids.len, ids[full * x.page ..], extra); + if (r2.new and ram) try x.setRam(r2.id, true, ids); + return r2; + } + + /// Adds a disk entry continuing from node anchor (none: the tag's root); suffix: tokens [depth x page, pos). + pub fn insertFrom(x: *Self, tag: u32, anchor: u32, suffix: []const i32, extra: Extra) !Placed { + var n = if (anchor != none) anchor else blk: { + const r = root(tag); + break :blk x.by_digest.get(r) orelse try x.newNode(r, none, 0); + }; + const d0 = x.nodes.items[n].depth; + const full: u32 = @intCast(suffix.len / x.page); + for (0..full) |j| { + const d = step(x.nodes.items[n].digest, suffix[j * x.page ..][0..x.page]); + n = x.by_digest.get(d) orelse try x.newNode(d, n, d0 + @as(u32, @intCast(j)) + 1); + } + return x.place(tag, n, @as(u64, d0) * x.page + suffix.len, suffix[full * x.page ..], extra); + } + + /// The node at depth on an entry's path (its first `depth` full pages). + pub fn nodeAt(x: *Self, id: u32, depth: u32) ?u32 { + var n = x.entries.items[id].node; + while (n != none and x.nodes.items[n].depth > depth) n = x.nodes.items[n].parent; + return if (n != none and x.nodes.items[n].depth == depth) n else null; + } + + fn place(x: *Self, tag: u32, n: u32, pos: u64, tail_ids: []const i32, extra: Extra) !Placed { + const key = entryKey(x.nodes.items[n].digest, pos, tail_ids); + if (x.by_key.get(key)) |e| return .{ .id = e, .new = false }; + const tail = try x.gpa.dupe(i32, tail_ids); + errdefer x.gpa.free(tail); + const id: u32 = if (x.free_entries.pop()) |i| i else blk: { + try x.entries.append(x.gpa, undefined); + break :blk @intCast(x.entries.items.len - 1); + }; + try x.by_key.put(x.gpa, key, id); + const nd = &x.nodes.items[n]; + x.entries.items[id] = .{ .key = key, .tag = tag, .pos = pos, .node = n, .tail = tail, .ram = false, .next_at = nd.first_entry, .extra = extra }; + if (nd.first_entry != none) x.entries.items[nd.first_entry].prev_at = id; + nd.first_entry = id; + x.addPath(n, 1, 0); + return .{ .id = id, .new = true }; + } + + fn addPath(x: *Self, from: u32, d_all: i32, d_ram: i32) void { + var n = from; + while (n != none) { + const nd = &x.nodes.items[n]; + nd.sub_all = @intCast(@as(i64, nd.sub_all) + d_all); + nd.sub_ram = @intCast(@as(i64, nd.sub_ram) + d_ram); + n = nd.parent; + } + } + + /// Moves an entry in or out of RAM; going in needs its ids (to keep its path's tokens). + pub fn setRam(x: *Self, id: u32, ram: bool, ids: ?[]const i32) !void { + const e = &x.entries.items[id]; + if (e.ram == ram) return; + if (ram) { + const all = ids orelse return error.IdsNeeded; + var n = e.node; + while (n != none and x.nodes.items[n].depth > 0) : (n = x.nodes.items[n].parent) { + const nd = &x.nodes.items[n]; + if (nd.tokens != null) continue; + const b = try x.takeBlock(); + @memcpy(b, all[(nd.depth - 1) * x.page ..][0..x.page]); + nd.tokens = b; + } + } + e.ram = ram; + x.nodes.items[e.node].own_ram = if (ram) x.nodes.items[e.node].own_ram + 1 else x.nodes.items[e.node].own_ram - 1; + x.addPath(e.node, 0, if (ram) 1 else -1); + if (!ram) x.dropTokens(e.node); + } + + fn dropTokens(x: *Self, from: u32) void { + var n = from; + while (n != none) : (n = x.nodes.items[n].parent) { + const nd = &x.nodes.items[n]; + if (nd.sub_ram > 0) return; + if (nd.tokens) |t| { + x.blocks.append(x.gpa, t) catch x.gpa.free(t); + nd.tokens = null; + } + } + } + + /// Removes an entry and the nodes no entry needs any more. + pub fn remove(x: *Self, id: u32) void { + x.setRam(id, false, null) catch unreachable; + const e = &x.entries.items[id]; + const n = e.node; + if (e.prev_at != none) x.entries.items[e.prev_at].next_at = e.next_at else x.nodes.items[n].first_entry = e.next_at; + if (e.next_at != none) x.entries.items[e.next_at].prev_at = e.prev_at; + _ = x.by_key.remove(e.key); + x.gpa.free(e.tail); + e.live = false; + x.free_entries.append(x.gpa, id) catch {}; + x.addPath(n, -1, 0); + var m = n; + while (m != none and x.nodes.items[m].sub_all == 0) { + const nd = x.nodes.items[m]; + if (nd.parent != none) { + if (nd.prev_sib != none) x.nodes.items[nd.prev_sib].next_sib = nd.next_sib else x.nodes.items[nd.parent].first_child = nd.next_sib; + if (nd.next_sib != none) x.nodes.items[nd.next_sib].prev_sib = nd.prev_sib; + } + _ = x.by_digest.remove(nd.digest); + x.free_nodes.append(x.gpa, m) catch {}; + m = nd.parent; + } + } + + /// The entry's tail matches the prompt there and the entry is shorter (its full pages matched by digest). + fn tailOf(e: *const Entry, prompt: []const i32) bool { + const at: usize = @intCast(e.pos - e.tail.len); + return e.pos < prompt.len and std.mem.eql(i32, e.tail, prompt[at..][0..e.tail.len]); + } + + /// The longest entry of tag whose ids strictly prefix prompt (RAM first on a tie); digests: the prompt's chain. + pub fn find(x: *Self, tag: u32, prompt: []const i32, digests: []const Digest, tier: Tier) ?u32 { + const full: u32 = @intCast(prompt.len / x.page); + const r = x.by_digest.get(root(tag)) orelse return null; + // the deepest node on the prompt's path: each prefix of an entry's path has a node: binary search + var lo: u32 = 0; + var hi: u32 = full; + while (lo < hi) { + const mid = lo + (hi - lo + 1) / 2; + if (x.by_digest.contains(digests[mid - 1])) lo = mid else hi = mid - 1; + } + var n: u32 = if (lo == 0) r else x.by_digest.get(digests[lo - 1]).?; + while (n != none) : (n = x.nodes.items[n].parent) { + const nd = &x.nodes.items[n]; + if (nd.sub_all == 0 or (tier == .ram and nd.own_ram == 0)) continue; + var best: u32 = none; + var e = nd.first_entry; + while (e != none) : (e = x.entries.items[e].next_at) { + const en = &x.entries.items[e]; + if (tier == .ram and !en.ram) continue; + if (!tailOf(en, prompt)) continue; + if (best == none or en.pos > x.entries.items[best].pos or (en.pos == x.entries.items[best].pos and en.ram)) best = e; + } + if (best != none) return best; + } + return null; + } + + /// A newer entry of id's chain holds every position it does (tier .ram: in RAM), so id need not be written. + pub fn covered(x: *Self, id: u32, tier: Tier) bool { + const e = &x.entries.items[id]; + const nd = &x.nodes.items[e.node]; + var f = nd.first_entry; + while (f != none) : (f = x.entries.items[f].next_at) { + const o = &x.entries.items[f]; + if (f == id or (tier == .ram and !o.ram)) continue; + if (o.pos > e.pos and std.mem.startsWith(i32, o.tail, e.tail)) return true; + } + var c = nd.first_child; + while (c != none) : (c = x.nodes.items[c].next_sib) { + const ch = &x.nodes.items[c]; + const n = if (tier == .ram) ch.sub_ram else ch.sub_all; + if (n == 0) continue; + if (e.tail.len == 0) return true; + if (ch.tokens) |t| if (std.mem.startsWith(i32, t, e.tail)) return true; + } + return false; + } + + /// How many RAM entries extend a RAM entry (longer, same tag, its ids their prefix): branch points have more. + pub fn extensions(x: *Self, id: u32) u32 { + const e = &x.entries.items[id]; + const nd = &x.nodes.items[e.node]; + var n: u32 = 0; + var f = nd.first_entry; + while (f != none) : (f = x.entries.items[f].next_at) { + const o = &x.entries.items[f]; + if (f != id and o.ram and o.pos > e.pos and std.mem.startsWith(i32, o.tail, e.tail)) n += 1; + } + var c = nd.first_child; + while (c != none) : (c = x.nodes.items[c].next_sib) { + const ch = &x.nodes.items[c]; + if (ch.sub_ram == 0) continue; + if (std.mem.startsWith(i32, ch.tokens.?, e.tail)) n += ch.sub_ram; + } + return n; + } + + /// A RAM entry no other RAM entry extends: a chat's newest state. + pub fn isLeaf(x: *Self, id: u32) bool { + return x.entries.items[id].ram and !x.covered(id, .ram); + } + + /// A RAM leaf's older turns (the RAM entries it extends) into out; returns the chat's recency (newest `used`). + pub fn members(x: *Self, leaf: u32, out: *std.ArrayList(u32), used: *const fn (*const Extra) u64) !u64 { + const l = &x.entries.items[leaf]; + var recency = used(&l.extra); + var below: u32 = none; // the path's node under n: its tokens are the page an entry at n must prefix + var n = l.node; + while (n != none) : ({ + below = n; + n = x.nodes.items[n].parent; + }) { + const nd = &x.nodes.items[n]; + if (nd.own_ram == 0) continue; + const page_ids: []const i32 = if (below == none) l.tail else x.nodes.items[below].tokens.?; + var f = nd.first_entry; + while (f != none) : (f = x.entries.items[f].next_at) { + const o = &x.entries.items[f]; + if (f == leaf or !o.ram or o.pos >= l.pos) continue; + if (!std.mem.startsWith(i32, page_ids, o.tail)) continue; + try out.append(x.gpa, f); + recency = @max(recency, used(&o.extra)); + } + } + return recency; + } + + /// A RAM entry's ids, rebuilt from its path's pages and tail, into out (len = pos). + pub fn idsOf(x: *Self, id: u32, out: []i32) !void { + const e = &x.entries.items[id]; + if (out.len != e.pos) return error.BadLength; + const at = e.pos - e.tail.len; + @memcpy(out[at..], e.tail); + var n = e.node; + while (n != none and x.nodes.items[n].depth > 0) : (n = x.nodes.items[n].parent) { + const nd = &x.nodes.items[n]; + const t = nd.tokens orelse return error.NotInRam; + @memcpy(out[(nd.depth - 1) * x.page ..][0..x.page], t); + } + } + }; +} + +const T = Index(u64); + +fn usedOf(e: *const u64) u64 { + return e.*; +} + +fn seq(gpa: std.mem.Allocator, n: usize, salt: i32) ![]i32 { + const v = try gpa.alloc(i32, n); + for (v, 0..) |*t, i| t.* = @as(i32, @intCast(i % 997)) * 3 + salt; + return v; +} + +fn add(x: *T, tag: u32, ids: []const i32, ram: bool, used: u64) !u32 { + var d: [64]Digest = undefined; + chain(tag, ids, x.page, d[0 .. ids.len / x.page]); + return (try x.insert(tag, ids, d[0 .. ids.len / x.page], ram, used)).id; +} + +test "longest strict prefix, ties, tags and coverage" { + const gpa = std.testing.allocator; + var x = T.init(gpa, 16); + defer x.deinit(); + const chat = try seq(gpa, 200, 0); + defer gpa.free(chat); + const t1 = try add(&x, 7, chat[0..40], true, 1); // a prompt's replay point + const t2 = try add(&x, 7, chat[0..72], true, 2); // a turn's end + const t3 = try add(&x, 7, chat[0..128], false, 3); // on disk only, page aligned + _ = try add(&x, 8, chat[0..100], true, 4); // another tag + var d: [16]Digest = undefined; + chain(7, chat[0..40], 16, d[0..2]); + const again = try x.insert(7, chat[0..40], d[0..2], true, 9); + try std.testing.expect(again.id == t1 and !again.new); + chain(7, chat[0..150], 16, d[0..9]); + try std.testing.expectEqual(@as(?u32, t3), x.find(7, chat[0..150], d[0..9], .any)); + try std.testing.expectEqual(@as(?u32, t2), x.find(7, chat[0..150], d[0..9], .ram)); + try std.testing.expectEqual(@as(?u32, t2), x.find(7, chat[0..128], d[0..8], .any)); // strict: 128 is not a prefix of 128 + chain(7, chat[0..72], 16, d[0..4]); + try std.testing.expectEqual(@as(?u32, t1), x.find(7, chat[0..72], d[0..4], .any)); + var other = try gpa.dupe(i32, chat[0..150]); + defer gpa.free(other); + other[50] = -1; // diverges inside t2's tail page: t1 still prefixes it + chain(7, other, 16, d[0..9]); + try std.testing.expectEqual(@as(?u32, t1), x.find(7, other, d[0..9], .any)); + // t3 extends t2 on disk only: its page's tokens are gone, so a tail there is undecided (the store marks covered_by) + try std.testing.expect(x.covered(t1, .ram) and !x.covered(t2, .any) and !x.covered(t2, .ram) and !x.covered(t3, .any)); + const t4 = try add(&x, 7, chat[0..64], false, 5); // page aligned: any deeper entry covers it + try std.testing.expect(x.covered(t4, .any) and x.covered(t4, .ram)); + x.remove(t4); + try std.testing.expect(x.isLeaf(t2) and !x.isLeaf(t1)); + try std.testing.expectEqual(@as(u32, 1), x.extensions(t1)); + try std.testing.expectEqual(@as(u32, 0), x.extensions(t2)); + var m: std.ArrayList(u32) = .empty; + defer m.deinit(gpa); + try std.testing.expectEqual(@as(u64, 2), try x.members(t2, &m, usedOf)); + try std.testing.expectEqualSlices(u32, &.{t1}, m.items); + const ids = try gpa.alloc(i32, 72); + defer gpa.free(ids); + try x.idsOf(t2, ids); + try std.testing.expectEqualSlices(i32, chat[0..72], ids); + // going to disk drops the path's tokens once no RAM entry needs them; removal prunes the nodes + try x.setRam(t2, false, null); + try x.setRam(t1, false, null); + try std.testing.expectError(error.NotInRam, x.idsOf(t2, ids)); + x.remove(t3); + x.remove(t2); + x.remove(t1); + chain(7, chat[0..150], 16, d[0..9]); + try std.testing.expectEqual(@as(?u32, null), x.find(7, chat[0..150], d[0..9], .any)); + try std.testing.expectEqual(@as(u32, 1 + 6), x.by_digest.count()); // tag 8's root + its 6 full pages +} diff --git a/zig/src/sessions/sessions.zig b/zig/src/sessions/sessions.zig new file mode 100644 index 0000000000..7bdc46b386 --- /dev/null +++ b/zig/src/sessions/sessions.zig @@ -0,0 +1,33 @@ +//! Session memory, family-neutral: the KV pool book, page bytes, prefix entries, the RAM tier and the NVMe tier. + +pub const freeset = @import("freeset.zig"); +pub const pool = @import("pool.zig"); +pub const Pool = pool.Pool; +pub const Slot = pool.Slot; +pub const Family = pool.Family; +pub const Layout = pool.Layout; +pub const pagestore = @import("pagestore.zig"); +pub const PageStore = pagestore.PageStore; +pub const HostPool = pagestore.HostPool; +pub const prefix = @import("prefix.zig"); +pub const lanes = @import("lanes.zig"); +pub const format = @import("format.zig"); +pub const disk = @import("disk.zig"); +pub const Disk = disk.Disk; +pub const jobs = @import("jobs.zig"); +pub const store = @import("store.zig"); +pub const Store = store.Store; + +test { + _ = freeset; + _ = pool; + _ = pagestore; + _ = prefix; + _ = lanes; + _ = format; + _ = disk; + _ = jobs; + _ = store; + _ = @import("disk_test.zig"); + _ = @import("store_test.zig"); +} diff --git a/zig/src/sessions/store.zig b/zig/src/sessions/store.zig new file mode 100644 index 0000000000..851ce509aa --- /dev/null +++ b/zig/src/sessions/store.zig @@ -0,0 +1,570 @@ +//! The session store: RAM entries holding pool pages and bounded state, a residency budget, and the NVMe tier. + +const std = @import("std"); +const Io = std.Io; +const pool_mod = @import("pool.zig"); +const Pool = pool_mod.Pool; +const Slot = pool_mod.Slot; +const PageStore = @import("pagestore.zig").PageStore; +const prefix = @import("prefix.zig"); +const Digest = prefix.Digest; +const disk_mod = @import("disk.zig"); +const Disk = disk_mod.Disk; +const jobs = @import("jobs.zig"); +const Job = jobs.Job; + +pub const none = prefix.none; + +/// A slot's bounded state (rings, carries, lookback, taps) as the family snapshots it: opaque to the store. +pub const Bounded = struct { + ptr: *anyopaque, + vtable: *const VTable, + + pub const VTable = struct { + nbytes: *const fn (ptr: *anyopaque) u64, + /// host copies of its pieces for a park (owned by the caller: name and bytes allocated with gpa) + blobs: *const fn (ptr: *anyopaque, gpa: std.mem.Allocator) anyerror![]disk_mod.Blob, + release: *const fn (ptr: *anyopaque) void, + }; +}; + +pub const Policy = struct { + /// bounded state the RAM tier keeps (least recently used entries go past it) + ram_bytes: u64 = 512 << 20, + /// idle chats resident in the pool (0: no limit) and pool bytes held by entries alone (0: no limit) + chats: u32 = 0, + resident_bytes: u64 = 0, + /// park an evicted chat's older turns too (default: drop them, the leaf's file covers them) + park_prefixes: bool = false, + /// drop, never write, an entry a newer entry of its chain covers + skip_covered: bool = true, + /// write a newer turn as a delta on the file its leading pages came from + delta: bool = true, + /// parks and restores on the tier thread (settled by the plan) + async_io: bool = false, +}; + +pub const Extra = struct { + /// RAM: the pool pages it holds, logical [0, npages) + pages: []u32 = &.{}, + bounded: ?Bounded = null, + /// a file with this key is registered on the NVMe tier + file: bool = false, + used: u64 = 0, + origin: ?pool_mod.Origin = null, + /// a newer entry of its chain was parked from RAM: writing this one again is redundant + covered_by: bool = false, + kind: u32 = 0, + /// an unsettled park of it + job: ?*Job = null, + older: u32 = none, + newer: u32 = none, +}; + +const Index = prefix.Index(Extra); + +pub const Hit = struct { id: u32, ram: bool, pos: u64 }; + +pub const Stats = struct { + saves: u64 = 0, + dups: u64 = 0, + hits_ram: u64 = 0, + hits_disk: u64 = 0, + parks: u64 = 0, + delta_parks: u64 = 0, + drops: u64 = 0, + covered_drops: u64 = 0, + chat_evictions: u64 = 0, + park_ns: u64 = 0, + restore_ns: u64 = 0, +}; + +pub const Store = struct { + gpa: std.mem.Allocator, + io: Io, + pool: *Pool, + pages: PageStore, + index: Index, + disk: ?*Disk, + /// heap-held: the tier thread keeps its address + runner: ?*jobs.Runner = null, + policy: Policy, + clock: u64 = 0, + bounded_total: u64 = 0, + /// RAM entries, least recently used first + oldest: u32 = none, + newest: u32 = none, + ram_count: u32 = 0, + pending: std.ArrayList(*Job) = .empty, + scratch: std.ArrayList(Digest) = .empty, + stats: Stats = .{}, + + pub fn init(gpa: std.mem.Allocator, io: Io, pool: *Pool, pages: PageStore, d: ?*Disk, policy: Policy) !Store { + var s: Store = .{ .gpa = gpa, .io = io, .pool = pool, .pages = pages, .index = Index.init(gpa, pool.page), .disk = d, .policy = policy }; + if (d) |dd| { + const r = try gpa.create(jobs.Runner); + r.* = jobs.Runner.init(gpa, io, dd); + s.runner = r; + if (policy.async_io) r.startThread() catch |e| { + gpa.destroy(r); + return e; + }; + } + return s; + } + + pub fn deinit(s: *Store) void { + if (s.runner) |r| { + for (s.pending.items) |j| r.wait(j); + r.deinit(); + s.gpa.destroy(r); + } + for (s.pending.items) |j| j.deinit(s.gpa); + s.pending.deinit(s.gpa); + for (s.index.entries.items) |*e| if (e.live) { + if (e.extra.bounded) |b| b.vtable.release(b.ptr); + s.pool.drop(e.extra.pages) catch {}; + s.gpa.free(e.extra.pages); + }; + s.index.deinit(); + s.scratch.deinit(s.gpa); + s.* = undefined; + } + + fn digests(s: *Store, tag: u32, ids: []const i32) ![]const Digest { + const n = ids.len / s.pool.page; + try s.scratch.resize(s.gpa, n); + prefix.chain(tag, ids, s.pool.page, s.scratch.items); + return s.scratch.items; + } + + pub fn entry(s: *Store, id: u32) *Index.Entry { + return s.index.get(id); + } + + /// The longest stored entry (RAM or NVMe) whose ids strictly prefix prompt, with the same tag. + pub fn find(s: *Store, tag: u32, prompt: []const i32) !?Hit { + const d = try s.digests(tag, prompt); + const id = s.index.find(tag, prompt, d, .any) orelse return null; + const e = s.index.get(id); + return .{ .id = id, .ram = e.ram, .pos = e.pos }; + } + + fn link(s: *Store, id: u32) void { + const x = &s.index.get(id).extra; + x.older = s.newest; + x.newer = none; + if (s.newest != none) s.index.get(s.newest).extra.newer = id else s.oldest = id; + s.newest = id; + } + + fn unlink(s: *Store, id: u32) void { + const x = &s.index.get(id).extra; + if (x.older != none) s.index.get(x.older).extra.newer = x.newer else s.oldest = x.newer; + if (x.newer != none) s.index.get(x.newer).extra.older = x.older else s.newest = x.older; + } + + fn touch(s: *Store, id: u32) void { + s.clock += 1; + const e = s.index.get(id); + e.extra.used = s.clock; + if (e.ram) { + s.unlink(id); + s.link(id); + } + } + + /// An entry for a slot's state at len(ids), sharing its pages; null for a duplicate (bounded is released then). + pub fn save(s: *Store, slot: *Slot, bounded: Bounded, ids: []const i32, tag: u32, kind: u32) !?u32 { + if (ids.len == 0) return error.EmptyEntry; + const need = pool_mod.pagesFor(ids.len, s.pool.page); + if (slot.len < need) return error.SlotTooShort; + const d = try s.digests(tag, ids); + const ins = try s.index.insert(tag, ids, d, false, .{ .kind = kind }); + errdefer if (ins.new) s.index.remove(ins.id); + const e = s.index.get(ins.id); + if (e.ram) { + bounded.vtable.release(bounded.ptr); + s.touch(ins.id); + s.stats.dups += 1; + return null; + } + const pages = try s.gpa.dupe(u32, slot.mapped()[0..need]); + errdefer s.gpa.free(pages); + try s.index.setRam(ins.id, true, ids); + s.pool.share(pages); + const x = &s.index.get(ins.id).extra; + x.pages = pages; + x.bounded = bounded; + x.kind = kind; + if (slot.origin) |o| x.origin = .{ .key = o.key, .pages = @min(o.pages, @as(u32, @intCast(ids.len / s.pool.page))) }; + s.bounded_total += bounded.vtable.nbytes(bounded.ptr); + s.ram_count += 1; + s.link(ins.id); + s.touch(ins.id); + s.stats.saves += 1; + try s.trimRam(); + _ = try s.enforce(); + return ins.id; + } + + /// A RAM entry into an empty slot: full pages shared, the partial last page copied; returns its bounded state. + pub fn restoreRam(s: *Store, id: u32, slot: *Slot) !Bounded { + const e = s.index.get(id); + if (!e.ram) return error.NotInRam; + const full: u32 = @intCast(e.pos / s.pool.page); + try slot.adopt(e.extra.pages[0..full]); + errdefer _ = slot.releaseAll() catch {}; + if (e.pos % s.pool.page != 0) { + const src = e.extra.pages[full]; + const dst = try s.pool.takeAt(full); + for (s.pool.layout.families, 0..) |f, fi| { + if (f.split and !s.pool.owns(src)) continue; + try s.pages.copyPage(@intCast(fi), s.pool.familyPage(f, src), s.pool.familyPage(f, dst)); + } + try slot.push(dst); + } + if (e.extra.origin) |o| slot.origin = .{ .key = o.key, .pages = @min(o.pages, full) }; + s.touch(id); + s.stats.hits_ram += 1; + return e.extra.bounded.?; + } + + /// Starts an NVMe resume into an empty slot (pages mapped now, filled by the job); settle it before the slot runs. + pub fn beginRestore(s: *Store, id: u32, slot: *Slot) !*Job { + const d = s.disk orelse return error.NoDisk; + const e = s.index.get(id); + if (!e.extra.file) return error.NotOnDisk; + if (slot.len != 0) return error.SlotNotEmpty; + try slot.ensure(e.pos); + const j = try s.gpa.create(Job); + j.* = .{ .kind = .restore, .entry = id, .slot = slot.id }; + errdefer j.deinit(s.gpa); + var ext: std.ArrayList(disk_mod.Extent) = .empty; + errdefer ext.deinit(s.gpa); + try d.extents(e.key, &ext); + j.extents = try ext.toOwnedSlice(s.gpa); + const np = pool_mod.pagesFor(e.pos, s.pool.page); + j.lists = try s.familyLists(slot.mapped()[0..np], 0, false); + j.sink = .{ .store = s.pages, .families = s.pool.layout.families, .page = s.pool.page, .world = s.pool.world, .rank = s.pool.rank, .pages = j.lists }; + try s.pending.append(s.gpa, j); + try s.runner.?.submit(j); + return j; + } + + /// Per family: family-tensor pages of logical [first, len(pages)) (owned_only: this rank's of split families). + fn familyLists(s: *Store, pages: []const u32, first: u32, owned_only: bool) ![][]u32 { + const fams = s.pool.layout.families; + const out = try s.gpa.alloc([]u32, fams.len); + var made: usize = 0; + errdefer { + for (out[0..made]) |l| s.gpa.free(l); + s.gpa.free(out); + } + for (fams, 0..) |f, fi| { + const w = s.pool.world; + const sharded = f.split and w > 1 and owned_only; + var n: usize = pages.len - first; + if (sharded) { + var by: [pool_mod.max_world]u32 = undefined; + pool_mod.span(first, @intCast(pages.len), w, &by); + n = by[s.pool.rank]; + } + const l = try s.gpa.alloc(u32, n); + var i: usize = 0; + var k: usize = first; + if (sharded) k = first + (s.pool.rank + w - first % w) % w; + while (k < pages.len) : (k += if (sharded) w else 1) { + l[i] = s.pool.familyPage(f, pages[k]); + i += 1; + } + out[fi] = l; + made += 1; + } + return out; + } + + /// Starts a park of a RAM entry (null: nothing to write, the entry is already settled out of RAM). + pub fn beginPark(s: *Store, id: u32) !?*Job { + const d = s.disk orelse return error.NoDisk; + const e = s.index.get(id); + if (!e.ram or e.extra.job != null) return error.NotParkable; + if (e.extra.file and d.get(e.key) != null) { + try s.leaveRam(id); + return null; + } + const page = s.pool.page; + const npages = pool_mod.pagesFor(e.pos, page); + var base: ?Digest = null; + var first: u32 = 0; + if (s.policy.delta) if (e.extra.origin) |o| if (o.pages > 0) { + var r = d.get(o.key); + while (r) |rec| : (r = if (rec.base == disk_mod.none) null else &d.recs.items[rec.base]) { + if (rec.first_page < o.pages) break; + } + if (r) |rec| if (rec.depth < d.opts.max_chain) { + base = rec.key; + first = o.pages; + }; + }; + const j = try s.gpa.create(Job); + j.* = .{ .kind = .park, .entry = id }; + errdefer j.deinit(s.gpa); + j.ids = try s.gpa.alloc(i32, e.pos); + try s.index.idsOf(id, j.ids); + j.lists = try s.familyLists(e.extra.pages, first, true); + j.blobs = try e.extra.bounded.?.vtable.blobs(e.extra.bounded.?.ptr, s.gpa); + j.src = .{ .key = e.key, .base = base, .tag = e.tag, .kind = e.extra.kind, .page = page, .pos = e.pos, .first_page = first, .npages = npages, .ids = j.ids[first * page ..], .store = s.pages, .fence = try s.pages.fence(), .families = s.pool.layout.families, .pages = j.lists, .blobs = j.blobs }; + if (base) |b| d.get(b).?.refs += 1; // held while the job runs: a trim cannot take the base + e.extra.job = j; + try s.pending.append(s.gpa, j); + try s.runner.?.submit(j); + return j; + } + + /// Settles a finished job at the plan's point: a park registers its file, a restore returns the bounded state. + pub fn settle(s: *Store, j: *Job) !?disk_mod.Restored { + s.runner.?.wait(j); + defer { + for (s.pending.items, 0..) |p, i| if (p == j) { + _ = s.pending.orderedRemove(i); + break; + }; + j.deinit(s.gpa); + } + const d = s.disk.?; + switch (j.kind) { + .park => { + if (j.src.base) |b| d.get(b).?.refs -= 1; + const e = s.index.get(j.entry); + e.extra.job = null; + if (j.err) |err| return err; // the entry stays in RAM; the .tmp is gone + _ = try d.register(&j.src, j.size); + e.extra.file = true; + s.stats.parks += 1; + if (j.src.base != null) s.stats.delta_parks += 1; + s.stats.park_ns += j.elapsed_ns; + if (s.policy.skip_covered) { + var m: std.ArrayList(u32) = .empty; + defer m.deinit(s.gpa); + _ = try s.index.members(j.entry, &m, usedOf); + for (m.items) |o| s.index.get(o).extra.covered_by = true; + } + try s.leaveRam(j.entry); + try s.trimDisk(); + return null; + }, + .restore => { + if (j.err) |err| { + // the slot as before the restore: empty, its reservation kept + _ = try s.pool.slots.items[j.slot].truncate(0); + return err; + } + const e = s.index.get(j.entry); + d.touch(e.key); + s.touch(j.entry); + s.pool.slots.items[j.slot].origin = .{ .key = e.key, .pages = @intCast(e.pos / s.pool.page) }; + s.stats.hits_disk += 1; + s.stats.restore_ns += j.elapsed_ns; + const r = j.restored; + j.restored = null; + return r; + }, + } + } + + /// Settles every pending job in submission order (one rank, or the plan's "settle all"). + pub fn settleAll(s: *Store) !void { + while (s.pending.items.len > 0) { + if (try s.settle(s.pending.items[0])) |r| { + var rr = r; + rr.deinit(s.gpa); + } + } + } + + fn usedOf(x: *const Extra) u64 { + return x.used; + } + + /// A RAM entry lets go of its pages and bounded state; it stays indexed if a file holds it. + fn leaveRam(s: *Store, id: u32) !void { + const e = s.index.get(id); + if (!e.ram) return; + try s.pool.drop(e.extra.pages); + s.gpa.free(e.extra.pages); + e.extra.pages = &.{}; + if (e.extra.bounded) |b| { + s.bounded_total -= b.vtable.nbytes(b.ptr); + b.vtable.release(b.ptr); + } + e.extra.bounded = null; + s.unlink(id); + s.ram_count -= 1; + try s.index.setRam(id, false, null); + if (!e.extra.file) s.index.remove(id); + } + + /// An entry whose file could not be read back: the file deleted, the entry gone unless RAM holds it (every rank). + pub fn forget(s: *Store, id: u32) !void { + const e = s.index.get(id); + if (e.extra.job != null) return error.Busy; + if (e.extra.file) if (s.disk) |d| { + d.stats.bad += 1; + // a delta's base stays while a newer file stacks on it (an opt-in delta mode) + d.remove(e.key) catch |err| return if (err == error.HasDependents) {} else err; + e.extra.file = false; + }; + if (!e.ram) s.index.remove(id); + } + + /// Drops an entry from RAM without writing it. + pub fn drop(s: *Store, id: u32) !void { + try s.leaveRam(id); + s.stats.drops += 1; + } + + /// Lets go of a RAM entry: dropped when covered, parked if the tier takes it, else dropped; returns a started park. + pub fn evict(s: *Store, id: u32) !?*Job { + const e = s.index.get(id); + if (e.extra.job != null) return null; + if (s.policy.skip_covered and (e.extra.covered_by or s.index.covered(id, .any))) { + try s.drop(id); + s.stats.covered_drops += 1; + return null; + } + if (s.disk) |d| if (d.accepts(e.pos)) { + const j = try s.beginPark(id) orelse return null; + if (!s.policy.async_io) _ = try s.settle(j); + return if (s.policy.async_io) j else null; + }; + try s.drop(id); + return null; + } + + /// Evicts least recently used RAM entries past the bounded-state budget. + fn trimRam(s: *Store) !void { + var id = s.oldest; + while (s.bounded_total > s.policy.ram_bytes and s.ram_count > 1 and id != none) { + const next = s.index.get(id).extra.newer; + if (s.index.get(id).extra.job == null) _ = try s.evict(id); + id = next; + } + } + + /// Files the NVMe budget deleted leave the index (a RAM entry with that key stays, unwritten again). + fn trimDisk(s: *Store) !void { + const d = s.disk orelse return; + var gone: std.ArrayList(Digest) = .empty; + defer gone.deinit(s.gpa); + try d.trim(&gone); + for (gone.items) |k| if (s.index.lookup(k)) |id| { + s.index.get(id).extra.file = false; + if (!s.index.get(id).ram) s.index.remove(id); + }; + } + + const Chat = struct { leaf: u32, recency: u64 }; + + fn chatLess(_: void, a: Chat, b: Chat) bool { + return a.recency < b.recency or (a.recency == b.recency and a.leaf < b.leaf); + } + + /// A chat is live while a slot maps its leaf's first page. + fn idleLeaf(s: *Store, id: u32) bool { + const e = s.index.get(id); + return e.extra.job == null and s.index.isLeaf(id) and (e.extra.pages.len == 0 or s.pool.slot_refs[e.extra.pages[0]] == 0); + } + + /// Evicts LRU idle chats past the residency budget (leaf parked, unshared older turns dropped); returns how many. + pub fn enforce(s: *Store) !u32 { + const pol = s.policy; + if (pol.chats == 0 and pol.resident_bytes == 0) return 0; + var chats: std.ArrayList(Chat) = .empty; + defer chats.deinit(s.gpa); + var members: std.ArrayList(u32) = .empty; + defer members.deinit(s.gpa); + var id = s.oldest; + while (id != none) : (id = s.index.get(id).extra.newer) { + if (!s.idleLeaf(id)) continue; + members.clearRetainingCapacity(); + try chats.append(s.gpa, .{ .leaf = id, .recency = try s.index.members(id, &members, usedOf) }); + } + std.mem.sort(Chat, chats.items, {}, chatLess); + var idle: u32 = @intCast(chats.items.len); + var gone: u32 = 0; + for (chats.items) |c| { + const over_n = pol.chats > 0 and idle > pol.chats; + const over_b = pol.resident_bytes > 0 and s.pool.entryOnlyBytes() > pol.resident_bytes; + if (!over_n and !over_b) break; + if (!s.index.get(c.leaf).ram or !s.idleLeaf(c.leaf)) continue; + members.clearRetainingCapacity(); + _ = try s.index.members(c.leaf, &members, usedOf); + _ = try s.evict(c.leaf); + // deepest first: a member more RAM entries extend than this chat's leaf is a branch point, as is all above + const own: u32 = if (s.index.get(c.leaf).ram) 1 else 0; + for (members.items) |m| { + if (s.index.get(m).extra.job != null) continue; + if (s.index.extensions(m) > own) break; + if (pol.park_prefixes) _ = try s.evict(m) else try s.drop(m); + } + idle -= 1; + gone += 1; + s.stats.chat_evictions += 1; + } + return gone; + } + + /// RAM entries to let go of (coldest first, never keep) so need pages are free: [] if they are, null if impossible. + pub fn planSpill(s: *Store, need: u32, keep: []const u32, out: *std.ArrayList(u32)) !bool { + out.clearRetainingCapacity(); + var avail = s.pool.available(); + if (avail >= need) return true; + var left: std.AutoHashMapUnmanaged(u32, u16) = .empty; + defer left.deinit(s.gpa); + var id = s.oldest; + while (id != none and avail < need) : (id = s.index.get(id).extra.newer) { + if (std.mem.indexOfScalar(u32, keep, id) != null or s.index.get(id).extra.job != null) continue; + try out.append(s.gpa, id); + for (s.index.get(id).extra.pages) |pg| { + const g = try left.getOrPut(s.gpa, pg); + if (!g.found_existing) g.value_ptr.* = s.pool.refs[pg]; + g.value_ptr.* -= 1; + if (g.value_ptr.* == 0) avail += 1; // split: a page of the scarcer residue may count for less; the check below decides + } + } + return avail >= need; + } + + /// Rebuilds NVMe entries on restart (ranks then intersect them), deltas on their bases' paths; returns how many. + pub fn reconcile(s: *Store) !u32 { + const d = s.disk orelse return 0; + var found: std.ArrayList(Disk.Found) = .empty; + defer { + for (found.items) |f| s.gpa.free(f.ids); + found.deinit(s.gpa); + } + try d.scan(&found); + for (found.items) |f| { + var anchor: u32 = none; + if (f.head.hasBase()) { + const b = s.index.lookup(f.head.base) orelse return error.NoBase; + anchor = s.index.nodeAt(b, f.head.first_page) orelse return error.BadChain; + } + const r = try s.index.insertFrom(f.head.tag, anchor, f.ids, .{ .file = true, .kind = f.head.kind }); + s.clock += 1; + s.index.get(r.id).extra.used = s.clock; + } + return @intCast(found.items.len); + } + + /// Idle resident chats and the pool bytes held by entries alone. + pub fn resident(s: *Store) struct { chats: u32, bytes: u64 } { + var n: u32 = 0; + var id = s.oldest; + while (id != none) : (id = s.index.get(id).extra.newer) if (s.idleLeaf(id)) { + n += 1; + }; + return .{ .chats = n, .bytes = s.pool.entryOnlyBytes() }; + } +}; diff --git a/zig/src/sessions/store_test.zig b/zig/src/sessions/store_test.zig new file mode 100644 index 0000000000..86b58a4f59 --- /dev/null +++ b/zig/src/sessions/store_test.zig @@ -0,0 +1,216 @@ +//! The session store on a host pool and a real NVMe tier: chats, evictions, resumes bit for bit, deltas, two ranks. + +const std = @import("std"); +const testing = std.testing; +const pool_mod = @import("pool.zig"); +const Pool = pool_mod.Pool; +const Slot = pool_mod.Slot; +const Family = pool_mod.Family; +const HostPool = @import("pagestore.zig").HostPool; +const disk_mod = @import("disk.zig"); +const store_mod = @import("store.zig"); +const Store = store_mod.Store; +const Bounded = store_mod.Bounded; + +const page = 64; +const fams = [_]Family{ + .{ .name = "comp.2", .ratio = 2, .row_bytes = 584, .split = true }, + .{ .name = "index_k.2", .ratio = 2, .row_bytes = 132 }, + .{ .name = "comp.20", .ratio = 1, .row_bytes = 584, .split = true }, +}; + +/// Bounded state in host memory: one blob naming its chat and length. +const HostBounded = struct { + gpa: std.mem.Allocator, + bytes: []u8, + + fn make(gpa: std.mem.Allocator, chat: u32, pos: u64) !Bounded { + const h = try gpa.create(HostBounded); + h.* = .{ .gpa = gpa, .bytes = try gpa.alloc(u8, 3000) }; + for (h.bytes, 0..) |*b, i| b.* = @truncate(i +% chat *% 17 +% pos); + return .{ .ptr = h, .vtable = &vt }; + } + const vt: Bounded.VTable = .{ .nbytes = nbytes, .blobs = blobs, .release = release }; + fn nbytes(p: *anyopaque) u64 { + const h: *HostBounded = @ptrCast(@alignCast(p)); + return h.bytes.len; + } + fn blobs(p: *anyopaque, gpa: std.mem.Allocator) anyerror![]disk_mod.Blob { + const h: *HostBounded = @ptrCast(@alignCast(p)); + const out = try gpa.alloc(disk_mod.Blob, 1); + out[0] = .{ .name = try gpa.dupe(u8, "rings"), .bytes = try gpa.dupe(u8, h.bytes) }; + return out; + } + fn release(p: *anyopaque) void { + const h: *HostBounded = @ptrCast(@alignCast(p)); + h.gpa.free(h.bytes); + h.gpa.destroy(h); + } +}; + +/// One rank: its pool, host pages and store, plus what each chat's slot wrote (the reference rows by logical page). +const Rank = struct { + pool: Pool, + host: HostPool, + disk: *disk_mod.Disk, + store: Store, + + fn init(r: *Rank, gpa: std.mem.Allocator, root: []const u8, world: u32, rank: u32, policy: store_mod.Policy) !void { + r.pool = try Pool.init(gpa, .{ .families = &fams, .page = page }, 96 * page, world, rank); + r.host = try HostPool.init(gpa, &r.pool); + r.disk = try disk_mod.Disk.open(gpa, testing.io, .{ .root = root, .compat = @splat(7), .rank = rank, .world = world, .chunk = 64 << 10, .staging = 4, .lanes = 2, .min_tokens = 1 }); + r.store = try Store.init(gpa, testing.io, &r.pool, r.host.store(), r.disk, policy); + } + + fn deinit(r: *Rank) void { + r.store.deinit(); + r.disk.close(); + r.host.deinit(); + r.pool.deinit(); + } + + /// The rows a chat's slot writes at a logical page: a function of (chat, logical page, family) alone. + fn write(r: *Rank, s: *const Slot, chat: u32, from_page: u32) void { + for (s.mapped()[from_page..], from_page..) |pg, k| for (fams, 0..) |f, fi| { + if (f.split and !r.pool.owns(pg)) continue; + const v = r.host.pageOf(@intCast(fi), r.pool.familyPage(f, pg)); + for (v, 0..) |*b, i| b.* = @truncate(i *% 13 +% k *% 101 +% chat *% 7 +% fi); + }; + } + + /// Every stored row of the slot's pages below pos equals what write() put there. + fn check(r: *Rank, s: *const Slot, chat: u32, pos: u64) !void { + const n = pool_mod.pagesFor(pos, page); + for (s.mapped()[0..n], 0..) |pg, k| for (fams, 0..) |f, fi| { + if (f.split and !r.pool.owns(pg)) continue; + const v = r.host.pageOf(@intCast(fi), r.pool.familyPage(f, pg)); + const rows: usize = if (k + 1 == n and pos % page != 0) (pos % page) / f.ratio * f.row_bytes else v.len; + for (v[0..rows], 0..) |b, i| try testing.expectEqual(@as(u8, @truncate(i *% 13 +% k *% 101 +% chat *% 7 +% fi)), b); + }; + } +}; + +fn chatIds(gpa: std.mem.Allocator, chat: u32, n: usize) ![]i32 { + const v = try gpa.alloc(i32, n); + for (v, 0..) |*t, i| t.* = @intCast(chat * 100_000 + i); + return v; +} + +/// Prefill or resume a chat to len(ids) in slot s, save it, free the slot, enforce residency; returns the resume point. +fn turn(gpa: std.mem.Allocator, r: *Rank, s: *Slot, chat: u32, ids: []const i32) ![]const u8 { + var tier: []const u8 = "fresh"; + var from: u32 = 0; + if (try r.store.find(1, ids)) |hit| { + if (hit.ram) { + const b = try r.store.restoreRam(hit.id, s); + try testing.expectEqual(@as(u64, 3000), b.vtable.nbytes(b.ptr)); + tier = "ram"; + } else { + const j = try r.store.beginRestore(hit.id, s); + var got = (try r.store.settle(j)).?; + defer got.deinit(gpa); + try testing.expectEqual(@as(u8, @truncate(chat *% 17 +% hit.pos)), got.blobs[0].bytes[0]); + tier = "disk"; + } + try r.check(s, chat, hit.pos); + from = @intCast(hit.pos / page); + } + try s.ensure(ids.len); + r.write(s, chat, from); + _ = try r.store.save(s, try HostBounded.make(gpa, chat, ids.len), ids, 1, 0); + _ = try s.releaseAll(); + _ = try r.store.enforce(); + try r.store.settleAll(); // the plan's settle point (async parks finish here) + return tier; +} + +/// The scenario on one rank; returns a trace of its decisions (stats and page ids) to compare ranks. +fn scenario(gpa: std.mem.Allocator, root: []const u8, world: u32, rank: u32, async_io: bool, trace: *std.ArrayList(u64)) !void { + var r: Rank = undefined; + try r.init(gpa, root, world, rank, .{ .chats = 1, .async_io = async_io }); + defer r.deinit(); + const s = try r.pool.newSlot(64 * page); + const a = try chatIds(gpa, 1, 20 * page); + defer gpa.free(a); + const b = try chatIds(gpa, 2, 20 * page); + defer gpa.free(b); + const st = &r.store.stats; + try testing.expectEqualStrings("fresh", try turn(gpa, &r, s, 1, a[0 .. 5 * page + 10])); + try testing.expectEqualStrings("ram", try turn(gpa, &r, s, 1, a[0 .. 8 * page + 3])); + try testing.expectEqual(@as(u64, 0), st.parks); + // a second idle chat: chat 1 (least recent) parked whole, its older turn dropped (not written) + try testing.expectEqualStrings("fresh", try turn(gpa, &r, s, 2, b[0 .. 6 * page])); + try testing.expectEqual(@as(u64, 1), st.parks); + try testing.expectEqual(@as(u64, 1), st.drops); + try testing.expectEqual(@as(u32, 1), r.store.resident().chats); + const whole = r.disk.used; + // chat 1 resumes from NVMe bit for bit; chat 2 goes + try testing.expectEqualStrings("disk", try turn(gpa, &r, s, 1, a[0 .. 12 * page + 40])); + try testing.expectEqual(@as(u64, 2), st.parks); + const before = r.disk.used; + // chat 2 resumes; chat 1 goes again, written as a delta on its first file (pages 8..12 of 13) + try testing.expectEqualStrings("disk", try turn(gpa, &r, s, 2, b[0 .. 9 * page])); + try testing.expectEqual(@as(u64, 3), st.parks); + try testing.expectEqual(@as(u64, 1), st.delta_parks); + // the delta holds 5 of the entry's 13 pages: less than chat 1's first whole file of 9 pages + try testing.expect(r.disk.used - before < whole); + // chat 1's newest turn restores through its two files, bit for bit + try testing.expectEqualStrings("disk", try turn(gpa, &r, s, 1, a[0 .. 15 * page])); + try testing.expectEqual(@as(u64, 3), st.hits_disk); + try trace.appendSlice(gpa, &.{ st.parks, st.delta_parks, st.drops, st.covered_drops, st.hits_ram, st.hits_disk, r.pool.freePages(), r.pool.entry_only }); + var id = r.store.oldest; + while (id != store_mod.none) : (id = r.store.entry(id).extra.newer) for (r.store.entry(id).extra.pages) |pg| try trace.append(gpa, pg); +} + +/// After the scenario: a new process on the same tier finds every file and resumes chat 1 through its delta chain. +fn restart(gpa: std.mem.Allocator, root: []const u8, world: u32, rank: u32) !void { + var r: Rank = undefined; + try r.init(gpa, root, world, rank, .{ .chats = 1 }); + defer r.deinit(); + try testing.expectEqual(@as(u32, 4), try r.store.reconcile()); + const s = try r.pool.newSlot(64 * page); + const a = try chatIds(gpa, 1, 20 * page); + defer gpa.free(a); + try testing.expectEqualStrings("disk", try turn(gpa, &r, s, 1, a[0 .. 16 * page])); + try testing.expectEqual(@as(u64, 1), r.store.stats.hits_disk); +} + +fn tmpRoot(gpa: std.mem.Allocator, tmp: *std.testing.TmpDir, sub: []const u8) ![:0]u8 { + try tmp.dir.createDirPath(testing.io, sub); + return tmp.dir.realPathFileAlloc(testing.io, sub, gpa); +} + +test "chats over turns: residency to NVMe, bit-identical resumes, deltas, older turns dropped (sync and async)" { + const gpa = testing.allocator; + var tmp = testing.tmpDir(.{}); + defer tmp.cleanup(); + var t0: std.ArrayList(u64) = .empty; + defer t0.deinit(gpa); + var t1: std.ArrayList(u64) = .empty; + defer t1.deinit(gpa); + const r0 = try tmpRoot(gpa, &tmp, "sync"); + defer gpa.free(r0); + try scenario(gpa, r0, 1, 0, false, &t0); + try restart(gpa, r0, 1, 0); + const r1 = try tmpRoot(gpa, &tmp, "async"); + defer gpa.free(r1); + try scenario(gpa, r1, 1, 0, true, &t1); + try testing.expectEqualSlices(u64, t0.items, t1.items); // the tier thread changes when bytes move, not what is decided +} + +test "two split ranks make the same decisions on the same calls and each restores its own rows" { + const gpa = testing.allocator; + var tmp = testing.tmpDir(.{}); + defer tmp.cleanup(); + var t0: std.ArrayList(u64) = .empty; + defer t0.deinit(gpa); + var t1: std.ArrayList(u64) = .empty; + defer t1.deinit(gpa); + const root = try tmpRoot(gpa, &tmp, "tier"); + defer gpa.free(root); + try scenario(gpa, root, 2, 0, false, &t0); + try scenario(gpa, root, 2, 1, true, &t1); + try restart(gpa, root, 2, 0); + try restart(gpa, root, 2, 1); + try testing.expectEqualSlices(u64, t0.items, t1.items); +} From 0ba7d30557c5a4bd6f0646dd7fd125bb58105704 Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Sun, 11 Oct 2026 18:33:54 +0700 Subject: [PATCH 019/104] Sessions: the page map is rehashed before tombstones fill it The prefix index keeps one node per distinct full page of every saved entry in `by_digest`. Each saved turn adds its new pages, and each entry dropped or evicted removes the pages no other entry passes through, for the life of the server. std's tombstones never clear, so in time they take every free slot, and each save's new pages and each find's walk past the cached prefix then probe the whole map. The map now counts its removals and is rehashed in place once they reach a tenth of its capacity. Test: 3,000 turns of up to 60 fresh pages, the oldest dropped past eight, keep a tenth of the map free after every turn; without the rehash it fails. --- zig/src/sessions/prefix.zig | 31 ++++++++++++++++++++++++++++++- 1 file changed, 30 insertions(+), 1 deletion(-) diff --git a/zig/src/sessions/prefix.zig b/zig/src/sessions/prefix.zig index 28c639fbd7..98634bbfd2 100644 --- a/zig/src/sessions/prefix.zig +++ b/zig/src/sessions/prefix.zig @@ -85,6 +85,7 @@ pub fn Index(comptime Extra: type) type { nodes: std.ArrayList(Node) = .empty, free_nodes: std.ArrayList(u32) = .empty, by_digest: std.AutoHashMapUnmanaged(Digest, u32) = .empty, + dead: u32 = 0, // nodes dropped from `by_digest` since it was last rehashed entries: std.ArrayList(Entry) = .empty, free_entries: std.ArrayList(u32) = .empty, by_key: std.AutoHashMapUnmanaged(Digest, u32) = .empty, @@ -259,12 +260,20 @@ pub fn Index(comptime Extra: type) type { if (nd.prev_sib != none) x.nodes.items[nd.prev_sib].next_sib = nd.next_sib else x.nodes.items[nd.parent].first_child = nd.next_sib; if (nd.next_sib != none) x.nodes.items[nd.next_sib].prev_sib = nd.prev_sib; } - _ = x.by_digest.remove(nd.digest); + if (x.by_digest.remove(nd.digest)) x.dropped(); x.free_nodes.append(x.gpa, m) catch {}; m = nd.parent; } } + /// std's tombstones never clear and keys fill 80 %: rehash at a tenth, so a tenth of the map always stays free. + fn dropped(x: *Self) void { + x.dead += 1; + if (@as(u64, x.dead) * 10 < x.by_digest.capacity()) return; + x.by_digest.rehash(std.hash_map.AutoContext(Digest){}); + x.dead = 0; + } + /// The entry's tail matches the prompt there and the entry is shorter (its full pages matched by digest). fn tailOf(e: *const Entry, prompt: []const i32) bool { const at: usize = @intCast(e.pos - e.tail.len); @@ -455,3 +464,23 @@ test "longest strict prefix, ties, tags and coverage" { try std.testing.expectEqual(@as(?u32, null), x.find(7, chat[0..150], d[0..9], .any)); try std.testing.expectEqual(@as(u32, 1 + 6), x.by_digest.count()); // tag 8's root + its 6 full pages } + +test "a long run of turns saved and dropped leaves the page map free slots" { + const gpa = std.testing.allocator; + var x = T.init(gpa, 16); + defer x.deinit(); + var live: std.ArrayList(u32) = .empty; + defer live.deinit(gpa); + var rng = std.Random.DefaultPrng.init(5); + for (0..3000) |i| { + // a chat's turn (its own pages), the oldest dropped once eight are kept: nodes come and go + const ids = try seq(gpa, 16 * 60, @intCast(i)); + defer gpa.free(ids); + try live.append(gpa, try add(&x, 7, ids[0 .. 16 * (1 + rng.random().uintLessThan(usize, 60))], true, i)); + if (live.items.len > 8) x.remove(live.orderedRemove(0)); + const meta = x.by_digest.metadata orelse continue; + var free: usize = 0; + for (meta[0..x.by_digest.capacity()]) |m| free += @intFromBool(m.isFree()); + try std.testing.expect(free * 10 >= x.by_digest.capacity()); + } +} From 0918809379cadc6ad118ab2ab780e7641b0ae2b6 Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Sat, 10 Oct 2026 11:10:50 +0700 Subject: [PATCH 020/104] Let a family's round planner drive the lane host with prompt pieces A family can now hand the lane host a round planner (api.Rounds) that admits requests, splits prompts into pieces and runs them in the same rounds as the other streams' decode rows. The planner is opt-in: the host's own loop, with a backend's prefill_step a chunk a loop between rounds, stays the default, and a host with a planner runs the planner's loop instead (LaneHost.run picks one). The backend gains optional entries for this: begin, piece, pieces (several streams' pieces in one call), finish and join, so a round's finished prompts get their first drafts in one backend call. A planner's prompts open through lanes/open.zig like the whole and the chunked pass (open.openedWith; their first drafts sized by the head's reach under node probabilities), and a prompt cancelled at its final releases its lane as a cancelled chunk does. The lane core also takes a family row policy (lanes/trim.zig) in place of the built-in allocation, proposers that price their own drafts, and a timed switch for backends without round timers. Under a planner the host still steps the learner while no request waits or runs, and a failed piece or piece run fails only its own requests. A planner over a backend without prompt pieces (such as the external drafter wrapper) is refused at start with error.PiecesUnsupported. The round loop lives in lane_rounds.zig, a request's stream setup in lane_stream.zig and the pieces' steps in lanes/pieces.zig, keeping the host and the engine under the line limit. --- zig/src/core/engine_api.zig | 3 + zig/src/core/lane_host.zig | 114 ++++----- zig/src/core/lane_host_test.zig | 14 +- zig/src/core/lane_rounds.zig | 279 +++++++++++++++++++++ zig/src/core/lane_rounds_test.zig | 387 +++++++++++++++++++++++++++++ zig/src/core/lane_stream.zig | 60 +++++ zig/src/core/lanes/backend.zig | 20 ++ zig/src/core/lanes/config.zig | 17 ++ zig/src/core/lanes/engine.zig | 33 ++- zig/src/core/lanes/engine_test.zig | 121 +++++++++ zig/src/core/lanes/fake.zig | 69 +++++ zig/src/core/lanes/lanes.zig | 2 + zig/src/core/lanes/open.zig | 17 +- zig/src/core/lanes/pieces.zig | 144 +++++++++++ zig/src/core/lanes/proposer.zig | 8 + zig/src/core/lanes/trim.zig | 72 ++++++ zig/src/core/lanes/windows.zig | 34 +++ 17 files changed, 1314 insertions(+), 80 deletions(-) create mode 100644 zig/src/core/lane_rounds.zig create mode 100644 zig/src/core/lane_rounds_test.zig create mode 100644 zig/src/core/lane_stream.zig create mode 100644 zig/src/core/lanes/pieces.zig create mode 100644 zig/src/core/lanes/trim.zig diff --git a/zig/src/core/engine_api.zig b/zig/src/core/engine_api.zig index f72c13b227..de9b9da1af 100644 --- a/zig/src/core/engine_api.zig +++ b/zig/src/core/engine_api.zig @@ -381,6 +381,8 @@ pub const LoneHooks = struct { /// The lane core served to the HTTP threads (lane_host.zig). pub const serial_host = @import("serial_host.zig"); pub const LaneHost = @import("lane_host.zig").LaneHost; +/// An opt-in family round planner for the lane host, in place of its `prefill_step` chunks (lane_rounds.zig). +pub const Rounds = @import("lane_host.zig").Rounds; /// Exact prompt reuse between requests, for any family (prompt_cache.zig). pub const prompt_cache = @import("prompt_cache.zig"); @@ -402,6 +404,7 @@ test { _ = @import("serial_host_test.zig"); _ = @import("lane_host.zig"); _ = @import("lane_host_test.zig"); + _ = @import("lane_rounds_test.zig"); _ = @import("lane_host_reuse_test.zig"); _ = @import("lane_host_warm_test.zig"); _ = @import("prompt_cache.zig"); diff --git a/zig/src/core/lane_host.zig b/zig/src/core/lane_host.zig index 09ff0fa2de..585f5cea3f 100644 --- a/zig/src/core/lane_host.zig +++ b/zig/src/core/lane_host.zig @@ -17,12 +17,19 @@ const Reason = api.Reason; const Stats = api.Stats; const SubmitError = api.SubmitError; +pub const Rounds = rounds.Rounds; +pub const Proposers = rounds.Proposers; +const rounds = @import("lane_rounds.zig"); +const lane_stream = @import("lane_stream.zig"); + pub const LaneHost = struct { gpa: Allocator, io: std.Io, core: *lanes.Engine, info_: Info, min_match: i64 = 4, + /// A family's own copy proposer for each request (null: lanes.SuffixLookup at `min_match`) + proposers: ?Proposers = null, mutex: std.Io.Mutex = .init, wake: std.Io.Condition = .init, queued: std.ArrayList(*Job) = .empty, @@ -47,20 +54,31 @@ pub const LaneHost = struct { learner: ?api.Learner = null, // the family's Sliding Weights learner; null: learn requests are refused lessons: std.ArrayList(Lesson) = .empty, // learn requests waiting for an idle engine lesson_open: bool = false, // the learner holds a begun job (engine thread only) + /// Runs first on the engine thread: a backend whose API state is per thread (a CUDA context) binds it there. + on_thread: ?ThreadInit = null, + /// opt-in: a family round planner (prompt pieces, decode rows in one pass); null: whole or `prefill_step` prompts + rounds: ?Rounds = null, + pieces: std.ArrayList(Rounds.Piece) = .empty, + finals: std.ArrayList(usize) = .empty, + + pub const ThreadInit = struct { ctx: *anyopaque, run: *const fn (ctx: *anyopaque) anyerror!void }; const Lesson = struct { request: *const api.LearnRequest, sink: api.LearnSink }; - const Job = struct { + pub const Job = struct { host: *LaneHost, id: Id, request: *const Request, sink: Sink, stream: lanes.Stream = undefined, proposer: lanes.SuffixLookup = undefined, + /// the family's proposer (LaneHost.proposers), in place of `proposer` + custom: ?lanes.proposer.Proposer = null, delivered: usize = 0, started: bool = false, prefill_sent: bool = false, // a lone driver's prefilled event went out fill_began: bool = false, // its first prompt chunk ran + submitted: i96 = 0, began: i96 = 0, prefilled: ?i96 = null, entry: ?*pc.Entry = null, // the kept state the backend restores, until its prompt pass reports @@ -68,7 +86,7 @@ pub const LaneHost = struct { kept0: u64 = 0, // the store's kept count when the job looked it up /// The backend's prompt pass stands at a mark: the cache keeps the stream's state there. - fn kept(ptr: *anyopaque, s: *lanes.Stream, at: u32) void { + pub fn kept(ptr: *anyopaque, s: *lanes.Stream, at: u32) void { const job: *Job = @ptrCast(@alignCast(ptr)); job.reported(); // before a keep can evict the entry the pass restored if (job.host.cache) |store| _ = store.keep(job.request.prompt, at, s, job.request.chunks, job.host.spans(job.request)); @@ -92,6 +110,7 @@ pub const LaneHost = struct { } pub fn start(h: *LaneHost) !void { + if (h.rounds != null and !h.core.backend.pieces()) return error.PiecesUnsupported; // a drafter wrapper runs whole prompts h.thread = try std.Thread.spawn(.{ .stack_size = 16 << 20 }, run, .{h}); } @@ -112,6 +131,8 @@ pub const LaneHost = struct { h.cancels.deinit(h.gpa); h.decoded.deinit(h.gpa); h.live_tokens.deinit(h.gpa); + h.pieces.deinit(h.gpa); + h.finals.deinit(h.gpa); } pub fn engine(h: *LaneHost) Engine { @@ -134,7 +155,7 @@ pub const LaneHost = struct { fn submitFn(ctx: *anyopaque, id: Id, request: *const Request, sink: Sink) SubmitError!void { const h = self(ctx); const job = h.gpa.create(Job) catch return error.Busy; - job.* = .{ .host = h, .id = id, .request = request, .sink = sink }; + job.* = .{ .host = h, .id = id, .request = request, .sink = sink, .submitted = std.Io.Clock.awake.now(h.io).toNanoseconds() }; h.mutex.lockUncancelable(h.io); defer h.mutex.unlock(h.io); if (h.closing) { @@ -171,7 +192,7 @@ pub const LaneHost = struct { } /// A job's cancel hook for its prompt pass: its id is in `cancels`, or the host is closing (read under the lock). - fn cancelled(ctx: *anyopaque) bool { + pub fn cancelled(ctx: *anyopaque) bool { const job: *Job = @ptrCast(@alignCast(ctx)); job.host.lock(); defer job.host.unlock(); @@ -179,7 +200,7 @@ pub const LaneHost = struct { } /// A background pass's yield hook: a foreground request waits or a stream decodes (read under the lock). - fn yieldsTo(ctx: *anyopaque) bool { + pub fn yieldsTo(ctx: *anyopaque) bool { const job: *Job = @ptrCast(@alignCast(ctx)); job.host.lock(); defer job.host.unlock(); @@ -203,7 +224,7 @@ pub const LaneHost = struct { } /// Sends a stream's new tokens; true once it has finished (its job freed). - fn deliver(h: *LaneHost, job: *Job) bool { + pub fn deliver(h: *LaneHost, job: *Job) bool { h.send(job); if (!job.stream.finished) return false; const reason: Reason = switch (job.stream.reason) { @@ -216,7 +237,8 @@ pub const LaneHost = struct { return true; } - fn finish(h: *LaneHost, job: *Job, reason: Reason, message: []const u8) void { + pub fn finish(h: *LaneHost, job: *Job, reason: Reason, message: []const u8) void { + if (h.rounds) |rd| rd.vtable.left(rd.ctx, @intFromPtr(job)); job.reported(); h.gpa.free(job.marks); const s = &job.stream; @@ -225,13 +247,13 @@ pub const LaneHost = struct { if (job.started) { if (s.grammar) |g| g.deinit(h.gpa); s.deinit(h.gpa); - job.proposer.deinit(); + lane_stream.freeProposer(h, job); } h.gpa.destroy(job); } /// Cancels queued and admitted jobs named since the last round. - fn takeCancels(h: *LaneHost) void { + pub fn takeCancels(h: *LaneHost) void { h.mutex.lockUncancelable(h.io); const ids = h.gpa.dupe(Id, h.cancels.items) catch &.{}; h.cancels.clearRetainingCapacity(); @@ -259,11 +281,11 @@ pub const LaneHost = struct { h.gpa.free(ids); } - fn lock(h: *LaneHost) void { + pub fn lock(h: *LaneHost) void { h.mutex.lockUncancelable(h.io); } - fn unlock(h: *LaneHost) void { + pub fn unlock(h: *LaneHost) void { h.mutex.unlock(h.io); } @@ -288,50 +310,9 @@ pub const LaneHost = struct { return true; }; h.unlock(); - const r = job.request; - var reuse: lanes.stream.Reuse = .{}; - // the entry stays alive until the backend restores it: nothing keeps between here and this stream's own pass - if (h.cache) |store| if (store.lookupRewind(h.gpa, r.prompt, r.history_len, r.rewind_len, r.shared_prefixes, r.chunks, h.spans(r))) |l| { - job.entry = l.entry; - job.kept0 = store.counts.kept; - job.marks = l.marks; - reuse = .{ .saved = if (l.entry) |e| e.saved else null, .at = if (l.entry) |e| e.at else 0, .marks = l.marks, .hook = .{ .ptr = job, .at = Job.kept } }; - } else |_| {}; - job.proposer = lanes.SuffixLookup.init(h.gpa, .{ .min_match = h.min_match }) catch return h.drop(job, "the drafter could not start"); - job.stream = lanes.Stream.init(h.gpa, .{ - .id = "request", - .prompt = r.prompt, - .max_new = r.max_tokens, - .eos = r.eos, - .sampling = r.sampling, - .drafts = r.drafts, - .proposer = job.proposer.proposer(), - .stop_check = if (r.stop) |s| .{ .ptr = s.ctx, .check = s.check } else null, - .cancel_check = .{ .ptr = job, .check = cancelled }, - .yield_check = if (r.background and r.max_tokens == 0) .{ .ptr = job, .check = yieldsTo } else null, - .think_budget = r.think_budget, - .think_close = r.think_close, - .think_end = if (r.think_end) |t| t else -1, - .logprobs = r.logprobs, - .loop_guard = r.loop_guard, - .chunks = r.chunks, - .history_len = r.history_len, - .shared_prefixes = r.shared_prefixes, - .reuse = reuse, - .grammar = if (r.structure) |st| h.constraint(st) catch |e| { - job.proposer.deinit(); - return h.drop(job, if (e == error.OutOfMemory) "out of memory" else "the reply's grammar could not start (xgrammar)"); - } else null, - .decode_spans = h.spans(r), - .decide = r.decide, - }) catch { - job.proposer.deinit(); - return h.drop(job, "out of memory"); - }; - job.started = true; - const began = std.Io.Clock.awake.now(h.io).toNanoseconds(); - job.began = began; - if (r.max_tokens == 0 and h.info_.warm_turns) return h.warmPass(job, began); + if (!lane_stream.open(h, job)) return true; + const began = job.began; + if (job.request.max_tokens == 0 and h.info_.warm_turns) return h.warmPass(job, began); if (h.loneFits(job)) return h.runLone(job, began); if (h.core.fills()) { h.filling.append(h.gpa, job) catch return h.drop(job, "out of memory"); @@ -378,12 +359,12 @@ pub const LaneHost = struct { } /// A reply's grammar from the server's compiler (Info.structures engines get requests that carry one). - fn constraint(h: *LaneHost, st: api.Structure) !*lanes.grammar.Constraint { + pub fn constraint(h: *LaneHost, st: api.Structure) !*lanes.grammar.Constraint { const c = st.compiler orelse return error.NoGrammarCompiler; return c.constraint(h.gpa, h.io, @fromBackingInt(@intCast(@backingInt(st.kind))), st.text, st.after); } - fn prefilled(h: *LaneHost, job: *Job, began: i96) void { + pub fn prefilled(h: *LaneHost, job: *Job, began: i96) void { const done = std.Io.Clock.awake.now(h.io).toNanoseconds(); h.lock(); if (done > began) h.prefill_rate = @as(f64, @floatFromInt(job.request.prompt.len - job.stream.cached)) / (@as(f64, @floatFromInt(done - began)) / 1e9); @@ -396,7 +377,7 @@ pub const LaneHost = struct { } /// A request's decoded prompt rows, for a backend that keeps their bits (`decoded_rows`); otherwise none. - fn spans(h: *const LaneHost, r: *const Request) []const [2]u32 { + pub fn spans(h: *const LaneHost, r: *const Request) []const [2]u32 { return if (h.decoded_rows) r.decode_spans else &.{}; } @@ -473,7 +454,7 @@ pub const LaneHost = struct { return x.text(x.ctx, e) orelse @errorName(e); } - fn drop(h: *LaneHost, job: *Job, message: []const u8) bool { + pub fn drop(h: *LaneHost, job: *Job, message: []const u8) bool { h.remove(job); if (job.started and !job.stream.finished) h.core.discard(&job.stream); h.finish(job, .failed, message); @@ -483,7 +464,7 @@ pub const LaneHost = struct { /// No memory for another lane yet: the job waits at the queue's head, and admission holds until a lane frees. fn requeue(h: *LaneHost, job: *Job) bool { job.stream.deinit(h.gpa); - job.proposer.deinit(); + lane_stream.freeProposer(h, job); job.started = false; job.fill_began = false; job.entry = null; // looked up again at its next admission @@ -506,13 +487,13 @@ pub const LaneHost = struct { } /// A job cancelled in its prompt pass, its lane already released. - fn cancel(h: *LaneHost, job: *Job) bool { + pub fn cancel(h: *LaneHost, job: *Job) bool { h.remove(job); h.finish(job, .cancelled, ""); return true; } - fn remove(h: *LaneHost, job: *Job) void { + pub fn remove(h: *LaneHost, job: *Job) void { h.unfill(job); h.lock(); defer h.unlock(); @@ -523,6 +504,9 @@ pub const LaneHost = struct { } fn run(h: *LaneHost) void { + if (h.on_thread) |t| t.run(t.ctx) catch |e| std.log.err("lane host: the engine thread's init failed ({t})", .{e}); + // a family planner: prompt pieces and decode rows in one pass; else whole or `prefill_step` chunked prompts + if (h.rounds) |rd| return rounds.run(h, rd); while (true) { h.takeCancels(); while (h.admitOne()) {} @@ -573,7 +557,8 @@ pub const LaneHost = struct { } /// A failed round ends every stream it held, with the backend's error. - fn failAll(h: *LaneHost, message: []const u8) void { + pub fn failAll(h: *LaneHost, message: []const u8) void { + std.log.err("lane host: a round failed ({s}): every admitted stream fails", .{message}); h.lock(); const jobs = h.gpa.dupe(*Job, h.admitted.items) catch &.{}; h.admitted.clearRetainingCapacity(); @@ -586,7 +571,7 @@ pub const LaneHost = struct { h.gpa.free(jobs); } - fn closeAll(h: *LaneHost) void { + pub fn closeAll(h: *LaneHost) void { h.lock(); for (h.queued.items) |job| h.cancels.append(h.gpa, job.id) catch {}; for (h.admitted.items) |job| h.cancels.append(h.gpa, job.id) catch {}; @@ -597,4 +582,5 @@ pub const LaneHost = struct { test { _ = @import("lane_host_test.zig"); + _ = @import("lane_rounds_test.zig"); } diff --git a/zig/src/core/lane_host_test.zig b/zig/src/core/lane_host_test.zig index d02118e137..4ef15db8f9 100644 --- a/zig/src/core/lane_host_test.zig +++ b/zig/src/core/lane_host_test.zig @@ -257,7 +257,7 @@ test "a lane host serves the core's own tokens, in order, and cancels between ro } /// A learner that takes `total` steps over a lesson, then reports it learned in that many. -const FakeLearner = struct { +pub const FakeLearner = struct { sink: ?api.LearnSink = null, steps: u32 = 0, total: u32 = 3, @@ -266,11 +266,11 @@ const FakeLearner = struct { fn emit(l: *FakeLearner, event: api.LearnEvent) void { l.sink.?.event(l.sink.?.ctx, &event); } - fn begin(ctx: *anyopaque, request: *const api.LearnRequest, sink: api.LearnSink) anyerror!void { + pub fn begin(ctx: *anyopaque, request: *const api.LearnRequest, sink: api.LearnSink) anyerror!void { const l: *FakeLearner = @ptrCast(@alignCast(ctx)); l.* = .{ .sink = sink, .total = l.total, .examples = request.train.len }; } - fn step(ctx: *anyopaque) api.Learner.Step { + pub fn step(ctx: *anyopaque) api.Learner.Step { const l: *FakeLearner = @ptrCast(@alignCast(ctx)); l.steps += 1; if (l.steps < l.total) return .{ .done = false, .changed = false }; @@ -278,20 +278,20 @@ const FakeLearner = struct { l.emit(.{ .done = .{} }); return .{ .done = true, .changed = true }; } - fn abort(ctx: *anyopaque) void { + pub fn abort(ctx: *anyopaque) void { const l: *FakeLearner = @ptrCast(@alignCast(ctx)); l.emit(.{ .done = .{ .message = "aborted" } }); } }; /// Learn events as tags, read on the test's thread once `done` arrives. -const LearnBox = struct { +pub const LearnBox = struct { mutex: std.Io.Mutex = .init, tags: std.ArrayList(std.meta.Tag(api.LearnEvent)) = .empty, steps: u32 = 0, done: bool = false, - fn sink(b: *LearnBox) api.LearnSink { + pub fn sink(b: *LearnBox) api.LearnSink { return .{ .ctx = b, .event = event }; } fn event(ctx: *anyopaque, e: *const api.LearnEvent) void { @@ -302,7 +302,7 @@ const LearnBox = struct { if (e.* == .learned) b.steps = e.learned.steps; if (e.* == .done) b.done = true; } - fn wait(b: *LearnBox) void { + pub fn wait(b: *LearnBox) void { while (true) { b.mutex.lockUncancelable(std.testing.io); const d = b.done; diff --git a/zig/src/core/lane_rounds.zig b/zig/src/core/lane_rounds.zig new file mode 100644 index 0000000000..6db9191cca --- /dev/null +++ b/zig/src/core/lane_rounds.zig @@ -0,0 +1,279 @@ +//! A family's round planner under the lane host: prompts in pieces between the other streams' decode rounds. +const std = @import("std"); +const lanes = @import("lanes"); +const api = @import("engine_api.zig"); +const host = @import("lane_host.zig"); +const lane_stream = @import("lane_stream.zig"); +const learn = @import("lane_host_learn.zig"); +const status = @import("lane_host_status.zig"); +const LaneHost = host.LaneHost; +const Job = LaneHost.Job; +const Allocator = std.mem.Allocator; + +/// A family's round planner: admissions, then `admitted`, begins, prompt pieces, a decode round, finals, `after`. +pub const Rounds = struct { + ctx: *anyopaque, + vtable: *const VTable, + + pub const Verdict = enum { admit, wait, skip, refuse }; + pub const Piece = struct { key: usize, start: u64, end: u64, save: bool, run: bool = true }; + pub const Ask = struct { key: usize, prompt: []const u32, max_tokens: u32, submitted_s: f64 }; + + pub const VTable = struct { + /// An admission pass starts (a lane is free and requests wait). + begin_admit: *const fn (ctx: *anyopaque) void, + /// The next queued request: start it, wait (it and those behind it), skip past it, or refuse it (`words`). + admit: *const fn (ctx: *anyopaque, ask: Ask, words: *[]const u8) Verdict, + /// The pass is decided: what runs before the admitted streams begin. + admitted: *const fn (ctx: *anyopaque) anyerror!void, + /// Around an admitted request's stream begin (`damaged`: its saved prefix could not be restored). + arm: *const fn (ctx: *anyopaque, key: usize) void, + began: *const fn (ctx: *anyopaque, key: usize, damaged: bool) void, + /// This round's pieces, in order, and the requests whose prompts are in. + plan: *const fn (ctx: *anyopaque, pieces: *std.ArrayList(Piece), finals: *std.ArrayList(usize)) anyerror!void, + /// A request left: finished, cancelled, failed or refused, started or not. + left: *const fn (ctx: *anyopaque, key: usize) void, + /// After the round: the pieces' seconds and the decode round's (finals included). + after: *const fn (ctx: *anyopaque, piece_s: f64, decode_s: f64) void, + }; +}; + +/// Makes and frees a request's copy proposer (`prompt` and `eos` outlive it). +pub const Proposers = struct { + ctx: *anyopaque, + make: *const fn (ctx: *anyopaque, gpa: Allocator, prompt: []const u32, eos: []const u32) anyerror!lanes.proposer.Proposer, + free: *const fn (ctx: *anyopaque, p: lanes.proposer.Proposer) void, +}; + +/// One round under the family's planner; false when it did nothing (nothing to admit, prefill or decode). +fn round(h: *LaneHost, rd: Rounds) bool { + var did = false; + // the admission pass, decided whole before any of it runs + var admits: std.ArrayList(*Job) = .empty; + defer admits.deinit(h.gpa); + var refused: std.ArrayList(struct { job: *Job, why: []const u8 }) = .empty; + defer refused.deinit(h.gpa); + h.lock(); + if (h.queued.items.len > 0 and h.admitted.items.len < h.info_.lanes) { + rd.vtable.begin_admit(rd.ctx); + var free = h.info_.lanes - h.admitted.items.len; + var i: usize = 0; + while (i < h.queued.items.len and free > 0) { + const job = h.queued.items[i]; + var why: []const u8 = ""; + const v = rd.vtable.admit(rd.ctx, .{ .key = @intFromPtr(job), .prompt = job.request.prompt, .max_tokens = job.request.max_tokens, .submitted_s = seconds(job.submitted) }, &why); + switch (v) { + .admit => { + if (h.admitted.append(h.gpa, job)) |_| {} else |_| break; + if (admits.append(h.gpa, job)) |_| {} else |_| { + _ = h.admitted.pop(); + break; + } + _ = h.queued.orderedRemove(i); + free -= 1; + }, + .wait => break, + .skip => i += 1, + .refuse => { + if (refused.append(h.gpa, .{ .job = job, .why = why })) |_| {} else |_| break; + _ = h.queued.orderedRemove(i); + }, + } + } + } + h.unlock(); + for (refused.items) |x| h.finish(x.job, .failed, x.why); + rd.vtable.admitted(rd.ctx) catch |e| { + h.failAll(h.words(e)); + return true; + }; + for (admits.items) |job| { + did = true; + if (!lane_stream.open(h, job)) continue; + rd.vtable.arm(rd.ctx, @intFromPtr(job)); + const b = h.core.beginStream(&job.stream) catch |e| { + _ = h.drop(job, h.words(e)); + continue; + }; + rd.vtable.began(rd.ctx, @intFromPtr(job), b.damaged); + } + // this round's pieces, a decode round of the streams decoding, then the prompts whose rows are in + rd.vtable.plan(rd.ctx, &h.pieces, &h.finals) catch |e| { + h.failAll(h.words(e)); + return true; + }; + const t0 = std.Io.Clock.awake.now(h.io).toNanoseconds(); + if (h.core.cfg.piece_runs and h.core.backend.vtable.pieces != null) { + // one call for the round's pieces (the backend runs several slots' rows in one forward); a failure fails each + var list: [16]lanes.backend.Backend.Piece = undefined; + var jobs: [16]*Job = undefined; + var n: usize = 0; + for (h.pieces.items) |piece| { + if (!piece.run) continue; + did = true; + const job: *Job = @ptrFromInt(piece.key); + if (job.stream.finished) continue; + if (n == list.len) { + _ = h.drop(job, "too many prompt pieces in one round"); + continue; + } + list[n] = .{ .stream = &job.stream, .start = piece.start, .end = piece.end, .save = piece.save }; + jobs[n] = job; + n += 1; + } + if (n > 0) h.core.pieces(list[0..n]) catch |e| { + for (jobs[0..n]) |job| dropPiece(h, job, h.words(e)); + }; + } else for (h.pieces.items) |piece| { + if (!piece.run) continue; + did = true; + const job: *Job = @ptrFromInt(piece.key); + if (job.stream.finished) continue; + h.core.piece(&job.stream, piece.start, piece.end, piece.save) catch |e| { + dropPiece(h, job, h.words(e)); + }; + } + const t1 = std.Io.Clock.awake.now(h.io).toNanoseconds(); + // Config.join_tail: a joined prompt's last token decodes in this round's window; else its own runs after + const join = h.core.cfg.join_tail; + if (join and h.finals.items.len > 0) { + did = true; + runFinals(h); + } + if (h.core.activeCount() > 0) { + did = true; + h.core.step() catch |e| { + h.failAll(@errorName(e)); + return true; + }; + } + if (!join and h.finals.items.len > 0) { + did = true; + runFinals(h); + } + const t2 = std.Io.Clock.awake.now(h.io).toNanoseconds(); + var i: usize = 0; + while (i < h.admitted.items.len) { + const job = h.admitted.items[i]; + if (job.started and h.deliver(job)) { + h.lock(); + _ = h.admitted.orderedRemove(i); + h.unlock(); + } else i += 1; + } + rd.vtable.after(rd.ctx, seconds(t1 - t0), seconds(t2 - t1)); + return did; +} + +/// A job whose piece failed: dropped and taken out of this round's finals (`drop` frees the job). +fn dropPiece(h: *LaneHost, job: *Job, message: []const u8) void { + const key = @intFromPtr(job); + var i: usize = 0; + while (i < h.finals.items.len) { + if (h.finals.items[i] == key) _ = h.finals.orderedRemove(i) else i += 1; + } + _ = h.drop(job, message); +} + +/// The prompts whose rows are in: each stream's last row (Engine.finishStream), then it decodes with the rounds. +fn runFinals(h: *LaneHost) void { + if (h.core.cfg.replay_runs) { + // the finals' shared work first (one forward for all); a failure fails each + var ss: [16]*lanes.Stream = undefined; + var jobs: [16]*Job = undefined; + var n: usize = 0; + for (h.finals.items) |key| { + const job: *Job = @ptrFromInt(key); + if (job.stream.finished or n == ss.len) continue; + ss[n] = &job.stream; + jobs[n] = job; + n += 1; + } + h.core.prepareFinals(ss[0..n]) catch |e| { + for (jobs[0..n]) |job| _ = h.drop(job, h.words(e)); + return; + }; + } + if (h.core.cfg.join_tail and h.core.cfg.join_drafts) { + // the round's finals together: the joined prompts' first drafts in one pass (Engine.finishStreams) + var ss: [16]*lanes.Stream = undefined; + var jobs: [16]*Job = undefined; + var errs: [16]?anyerror = undefined; + var at: usize = 0; + while (at < h.finals.items.len) { + var n: usize = 0; + while (at < h.finals.items.len and n < ss.len) : (at += 1) { + const job: *Job = @ptrFromInt(h.finals.items[at]); + if (job.stream.finished) continue; + ss[n] = &job.stream; + jobs[n] = job; + n += 1; + } + const joined = h.core.finishStreams(ss[0..n], errs[0..n]) catch |e| blk: { + for (errs[0..n]) |*x| x.* = e; + break :blk 0; + }; + // a round's line when its finished prompts' first drafts joined in one pass + if (joined > 0) std.log.info("join drafts: {d} of {d} finished prompts drafted in one pass", .{ joined, n }); + for (jobs[0..n], errs[0..n]) |job, x| { + if (x) |e| { + _ = if (e == error.Cancelled) h.cancel(job) else h.drop(job, h.words(e)); + continue; + } + h.prefilled(job, job.began); + } + } + return; + } + for (h.finals.items) |key| { + const job: *Job = @ptrFromInt(key); + if (job.stream.finished) continue; + h.core.finishStream(&job.stream) catch |e| { + _ = if (e == error.Cancelled) h.cancel(job) else h.drop(job, h.words(e)); + continue; + }; + h.prefilled(job, job.began); + } +} + +/// The engine thread under a round planner; an idle round waits for an arrival, a cancel or a short timeout. +pub fn run(h: *LaneHost, rd: Rounds) void { + while (true) { + h.takeCancels(); + h.lock(); + if (h.closing) { + const left = h.queued.items.len + h.admitted.items.len; + h.unlock(); + if (h.lesson_open) { + h.learner.?.abort(h.learner.?.ctx); + h.lesson_open = false; + } + if (left == 0) return; + h.closeAll(); + continue; + } + const work = h.admitted.items.len > 0 or (h.queued.items.len > 0 and h.admitted.items.len < h.info_.lanes); + if (!work and learn.ready(h)) { + h.unlock(); + learn.step(h); + continue; + } + if (!work) { + if (h.cancels.items.len == 0) h.wake.waitTimeout(h.io, &h.mutex, .{ .duration = .{ .raw = .fromMilliseconds(100), .clock = .awake } }) catch {}; + h.unlock(); + continue; + } + h.unlock(); + const did = round(h, rd); + status.noteLive(h); + if (!did) { + h.lock(); + if (h.cancels.items.len == 0) h.wake.waitTimeout(h.io, &h.mutex, .{ .duration = .{ .raw = .fromMilliseconds(20), .clock = .awake } }) catch {}; + h.unlock(); + } + } +} + +fn seconds(ns: i96) f64 { + return @as(f64, @floatFromInt(ns)) / 1e9; +} diff --git a/zig/src/core/lane_rounds_test.zig b/zig/src/core/lane_rounds_test.zig new file mode 100644 index 0000000000..029f6d0e3e --- /dev/null +++ b/zig/src/core/lane_rounds_test.zig @@ -0,0 +1,387 @@ +//! The lane host under a family round planner: prompt pieces, failed piece runs, learning, refusal. +const std = @import("std"); +const lanes = @import("lanes"); +const api = @import("engine_api.zig"); +const LaneHost = @import("lane_host.zig").LaneHost; +const Memory = api.Memory; +const Reason = api.Reason; +const Id = api.Id; +const Event = api.Event; +const Request = api.Request; +const Engine = api.Engine; +const FakeLearner = @import("lane_host_test.zig").FakeLearner; +const LearnBox = @import("lane_host_test.zig").LearnBox; + +test "a round planner: prompts in pieces between the other streams' decode rounds, queued admissions, the same tokens" { + const gpa = std.testing.allocator; + var cfg = try lanes.Config.init(gpa, .{ .exact_width = 8, .gpu_tokens = true, .hidden_rows = true }, 8, 7); + defer cfg.deinit(gpa); + var target: lanes.fake.Fake = .{ .gpa = gpa }; + defer target.deinit(); + var clock: lanes.fake.FixedClock = .{}; + var core = lanes.Engine.init(gpa, &cfg, target.pieceBackend(), clock.clock()); + defer core.deinit(); + var host = LaneHost.init(gpa, std.testing.io, &core, .{ .lanes = 2 }); + // one prompt prefilling at a time (the next waits queued), 3 rows a round, a final when its rows are in + const Plan = struct { + const Seq = struct { key: usize, n: u64, done: u64 = 0, decoding: bool = false }; + seqs: std.ArrayList(Seq) = .empty, + waits: u32 = 0, + rounds: u32 = 0, + mixed: u32 = 0, // rounds that ran a piece while another stream decoded + fn self(ctx: *anyopaque) *@This() { + return @ptrCast(@alignCast(ctx)); + } + fn beginAdmit(_: *anyopaque) void {} + fn admit(ctx: *anyopaque, ask: api.Rounds.Ask, _: *[]const u8) api.Rounds.Verdict { + const p = self(ctx); + for (p.seqs.items) |x| if (!x.decoding) { + p.waits += 1; + return .wait; + }; + p.seqs.append(gpa, .{ .key = ask.key, .n = ask.prompt.len }) catch return .refuse; + return .admit; + } + fn admitted(_: *anyopaque) anyerror!void {} + fn arm(_: *anyopaque, _: usize) void {} + fn began(_: *anyopaque, _: usize, _: bool) void {} + fn plan(ctx: *anyopaque, pieces: *std.ArrayList(api.Rounds.Piece), finals: *std.ArrayList(usize)) anyerror!void { + const p = self(ctx); + pieces.clearRetainingCapacity(); + finals.clearRetainingCapacity(); + var decoding = false; + for (p.seqs.items) |x| decoding = decoding or x.decoding; + for (p.seqs.items) |*x| if (!x.decoding) { + if (x.done < x.n - 1) { + const end = @min(x.done + 3, x.n - 1); + try pieces.append(gpa, .{ .key = x.key, .start = x.done, .end = end, .save = false }); + x.done = end; + if (decoding) p.mixed += 1; + } + if (x.done == x.n - 1) { + try finals.append(gpa, x.key); + x.decoding = true; + } + }; + p.rounds += 1; + } + fn left(ctx: *anyopaque, key: usize) void { + const p = self(ctx); + for (p.seqs.items, 0..) |x, i| if (x.key == key) { + _ = p.seqs.orderedRemove(i); + return; + }; + } + fn after(_: *anyopaque, _: f64, _: f64) void {} + }; + var pl: Plan = .{}; + defer pl.seqs.deinit(gpa); + host.rounds = .{ .ctx = &pl, .vtable = &.{ .begin_admit = Plan.beginAdmit, .admit = Plan.admit, .admitted = Plan.admitted, .arm = Plan.arm, .began = Plan.began, .plan = Plan.plan, .left = Plan.left, .after = Plan.after } }; + const Box = struct { + mutex: std.Io.Mutex = .init, + tokens: std.ArrayList(u32) = .empty, + done: ?Reason = null, + fn event(ctx: *anyopaque, _: Id, e: *const Event) void { + const b: *@This() = @ptrCast(@alignCast(ctx)); + b.mutex.lockUncancelable(std.testing.io); + defer b.mutex.unlock(std.testing.io); + switch (e.*) { + .tokens => |t| b.tokens.appendSlice(gpa, t) catch {}, + .finished => |f| b.done = f.reason, + else => {}, + } + } + fn wait(b: *@This()) Reason { + while (true) { + b.mutex.lockUncancelable(std.testing.io); + const d = b.done; + b.mutex.unlock(std.testing.io); + if (d) |r| return r; + std.Io.sleep(std.testing.io, .fromMilliseconds(1), .awake) catch {}; + } + } + }; + const p1 = [_]u32{ 3, 1, 4, 1, 5, 9, 2, 6, 5, 3 }; + const p2 = [_]u32{ 2, 7, 1, 8, 2, 8, 1, 8, 2, 8, 4, 5, 9 }; + var b1: Box = .{}; + var b2: Box = .{}; + defer b1.tokens.deinit(gpa); + defer b2.tokens.deinit(gpa); + const r1: Request = .{ .prompt = &p1, .max_tokens = 40 }; + const r2: Request = .{ .prompt = &p2, .max_tokens = 12 }; + const e = host.engine(); + try host.start(); + try e.submit(1, &r1, .{ .ctx = &b1, .event = Box.event }); + try e.submit(2, &r2, .{ .ctx = &b2, .event = Box.event }); + try std.testing.expectEqual(Reason.length, b1.wait()); + try std.testing.expectEqual(Reason.length, b2.wait()); + host.stop(); + // the replies are the whole-prompt path's tokens + for ([_]struct { p: []const u32, b: *Box }{ .{ .p = &p1, .b = &b1 }, .{ .p = &p2, .b = &b2 } }) |c| { + var history: std.ArrayList(u32) = .empty; + defer history.deinit(gpa); + try history.appendSlice(gpa, c.p); + for (c.b.tokens.items) |t| { + try std.testing.expectEqual(lanes.fake.next(history.items, null, history.items.len), t); + try history.append(gpa, t); + } + } + try std.testing.expectEqual(@as(usize, 40), b1.tokens.items.len); + try std.testing.expectEqual(@as(usize, 12), b2.tokens.items.len); + // the second prompt waited queued, then ran its pieces while the first decoded + try std.testing.expect(pl.waits > 0 and pl.mixed >= 3); + try std.testing.expectEqual(@as(usize, 3 + 4), target.pieces_run); + try std.testing.expectEqual(@as(usize, 0), pl.seqs.items.len); +} + +test "a failed piece run fails its requests alone: their finals of the same round do not run, the host keeps serving" { + // a run's error drops its jobs, whose finals are listed in the same round; runFinals must not read them + const gpa = std.testing.allocator; + var costs: [16]lanes.config.Cost = undefined; + for (&costs, 1..) |*c, w| c.* = .{ .width = @intCast(w), .ms = 5.0 + 0.8 * @as(f64, @floatFromInt(w)) }; + var cfg = try lanes.Config.init(gpa, .{ .exact_width = 16, .gpu_tokens = true, .mtp = true, .speculate = true, .speculate_early = false, .drafts = 4, .window_costs = &costs, .mtp_step_ms = 0.5, .hidden_rows = true, .batch_rows = 32, .max_streams = 8, .draft_streams = true, .join_tail = true, .piece_runs = true }, 16, 15); + defer cfg.deinit(gpa); + var target: lanes.fake.Fake = .{ .gpa = gpa, .fail_runs = 1 }; + defer target.deinit(); + var clock: lanes.fake.FixedClock = .{}; + var core = lanes.Engine.init(gpa, &cfg, target.runBackend(), clock.clock()); + defer core.deinit(); + // jobs on the page allocator: a freed job is unmapped, so reading one faults (the testing allocator hides it) + var host = LaneHost.init(std.heap.page_allocator, std.testing.io, &core, .{ .lanes = 4 }); + // every admitted prompt's rows in one piece and its final in the same round (as for short prompts) + const Plan = struct { + const Seq = struct { key: usize, n: u64, done: bool = false }; + seqs: std.ArrayList(Seq) = .empty, + fn self(ctx: *anyopaque) *@This() { + return @ptrCast(@alignCast(ctx)); + } + fn beginAdmit(_: *anyopaque) void {} + fn admit(ctx: *anyopaque, ask: api.Rounds.Ask, _: *[]const u8) api.Rounds.Verdict { + self(ctx).seqs.append(gpa, .{ .key = ask.key, .n = ask.prompt.len }) catch return .refuse; + return .admit; + } + fn admitted(_: *anyopaque) anyerror!void {} + fn arm(_: *anyopaque, _: usize) void {} + fn began(_: *anyopaque, _: usize, _: bool) void {} + fn plan(ctx: *anyopaque, pieces: *std.ArrayList(api.Rounds.Piece), finals: *std.ArrayList(usize)) anyerror!void { + pieces.clearRetainingCapacity(); + finals.clearRetainingCapacity(); + for (self(ctx).seqs.items) |*x| if (!x.done) { + // the host's lists: its allocator (it frees them) + try pieces.append(std.heap.page_allocator, .{ .key = x.key, .start = 0, .end = x.n - 1, .save = false }); + try finals.append(std.heap.page_allocator, x.key); + x.done = true; + }; + } + fn left(ctx: *anyopaque, key: usize) void { + const p = self(ctx); + for (p.seqs.items, 0..) |x, i| if (x.key == key) { + _ = p.seqs.orderedRemove(i); + return; + }; + } + fn after(_: *anyopaque, _: f64, _: f64) void {} + }; + var pl: Plan = .{}; + defer pl.seqs.deinit(gpa); + host.rounds = .{ .ctx = &pl, .vtable = &.{ .begin_admit = Plan.beginAdmit, .admit = Plan.admit, .admitted = Plan.admitted, .arm = Plan.arm, .began = Plan.began, .plan = Plan.plan, .left = Plan.left, .after = Plan.after } }; + const Box = struct { + mutex: std.Io.Mutex = .init, + tokens: usize = 0, + done: ?Reason = null, + fn event(ctx: *anyopaque, _: Id, e: *const Event) void { + const b: *@This() = @ptrCast(@alignCast(ctx)); + b.mutex.lockUncancelable(std.testing.io); + defer b.mutex.unlock(std.testing.io); + switch (e.*) { + .tokens => |t| b.tokens += t.len, + .finished => |f| b.done = f.reason, + else => {}, + } + } + fn wait(b: *@This()) Reason { + while (true) { + b.mutex.lockUncancelable(std.testing.io); + const d = b.done; + b.mutex.unlock(std.testing.io); + if (d) |r| return r; + std.Io.sleep(std.testing.io, .fromMilliseconds(1), .awake) catch {}; + } + } + }; + const p1 = [_]u32{ 3, 1, 4, 1, 5, 9, 2, 6 }; + const p2 = [_]u32{ 2, 7, 1, 8, 2, 8 }; + var b1: Box = .{}; + var b2: Box = .{}; + var b3: Box = .{}; + const r1: Request = .{ .prompt = &p1, .max_tokens = 12 }; + const r2: Request = .{ .prompt = &p2, .max_tokens = 12 }; + const e = host.engine(); + // both admitted before the host's first round: one run of their pieces, which fails + try e.submit(1, &r1, .{ .ctx = &b1, .event = Box.event }); + try e.submit(2, &r2, .{ .ctx = &b2, .event = Box.event }); + try host.start(); + defer host.stop(); + try std.testing.expectEqual(Reason.failed, b1.wait()); + try std.testing.expectEqual(Reason.failed, b2.wait()); + try std.testing.expectEqual(@as(usize, 0), target.joins); // no final of theirs ran + // the host serves the next request + try e.submit(3, &r1, .{ .ctx = &b3, .event = Box.event }); + try std.testing.expectEqual(Reason.length, b3.wait()); + try std.testing.expectEqual(@as(usize, 12), b3.tokens); + try std.testing.expectEqual(@as(usize, 1), target.runs); +} + +test "under a round planner a learn request still steps while no request waits or runs" { + const gpa = std.testing.allocator; + var cfg = try lanes.Config.init(gpa, .{ .exact_width = 8, .gpu_tokens = true, .hidden_rows = true }, 8, 7); + defer cfg.deinit(gpa); + var target: lanes.fake.Fake = .{ .gpa = gpa }; + defer target.deinit(); + var clock: lanes.fake.FixedClock = .{}; + var core = lanes.Engine.init(gpa, &cfg, target.pieceBackend(), clock.clock()); + defer core.deinit(); + const Idle = struct { + fn none(_: *anyopaque) void {} + fn admit(_: *anyopaque, _: api.Rounds.Ask, _: *[]const u8) api.Rounds.Verdict { + return .wait; + } + fn admitted(_: *anyopaque) anyerror!void {} + fn key(_: *anyopaque, _: usize) void {} + fn began(_: *anyopaque, _: usize, _: bool) void {} + fn plan(_: *anyopaque, _: *std.ArrayList(api.Rounds.Piece), _: *std.ArrayList(usize)) anyerror!void {} + fn after(_: *anyopaque, _: f64, _: f64) void {} + }; + var ctx: u8 = 0; + var host = LaneHost.init(gpa, std.testing.io, &core, .{ .lanes = 1 }); + host.rounds = .{ .ctx = &ctx, .vtable = &.{ .begin_admit = Idle.none, .admit = Idle.admit, .admitted = Idle.admitted, .arm = Idle.key, .began = Idle.began, .plan = Idle.plan, .left = Idle.key, .after = Idle.after } }; + var fake: FakeLearner = .{}; + host.learner = .{ .ctx = &fake, .begin = FakeLearner.begin, .step = FakeLearner.step, .abort = FakeLearner.abort }; + try host.start(); + defer host.stop(); + const example: api.Example = .{ .ids = &.{ 1, 2, 3 }, .start = 2 }; + const learn: api.LearnRequest = .{ .train = &.{example} }; + var box: LearnBox = .{}; + defer box.tags.deinit(gpa); + try host.engine().learn(&learn, box.sink()); + box.wait(); + try std.testing.expectEqualSlices(std.meta.Tag(api.LearnEvent), &.{ .learned, .done }, box.tags.items); + try std.testing.expectEqual(@as(u32, 3), box.steps); +} + +test "a round planner over a backend without prompt pieces is refused at start" { + const gpa = std.testing.allocator; + var cfg = try lanes.Config.init(gpa, .{ .exact_width = 8, .gpu_tokens = true, .hidden_rows = true }, 8, 7); + defer cfg.deinit(gpa); + var target: lanes.fake.Fake = .{ .gpa = gpa }; + defer target.deinit(); + var whole = target.backend(); + var vt = whole.vtable.*; + vt.begin = null; + whole.vtable = &vt; + var clock: lanes.fake.FixedClock = .{}; + var core = lanes.Engine.init(gpa, &cfg, whole, clock.clock()); + defer core.deinit(); + var host = LaneHost.init(gpa, std.testing.io, &core, .{ .lanes = 1 }); + var ctx: u8 = 0; + host.rounds = .{ .ctx = &ctx, .vtable = undefined }; + try std.testing.expectError(error.PiecesUnsupported, host.start()); +} + +test "a failed piece fails its request alone: its final of the same round does not run, the host keeps serving" { + // a job dropped for a failed piece must leave the round's finals: `drop` frees it, runFinals must not read it + const gpa = std.testing.allocator; + var costs: [16]lanes.config.Cost = undefined; + for (&costs, 1..) |*c, w| c.* = .{ .width = @intCast(w), .ms = 5.0 + 0.8 * @as(f64, @floatFromInt(w)) }; + var cfg = try lanes.Config.init(gpa, .{ .exact_width = 16, .gpu_tokens = true, .mtp = true, .speculate = true, .speculate_early = false, .drafts = 4, .window_costs = &costs, .mtp_step_ms = 0.5, .hidden_rows = true, .batch_rows = 32, .max_streams = 8, .draft_streams = true, .join_tail = true }, 16, 15); + defer cfg.deinit(gpa); + var target: lanes.fake.Fake = .{ .gpa = gpa, .fail_pieces = 1 }; + defer target.deinit(); + var clock: lanes.fake.FixedClock = .{}; + var core = lanes.Engine.init(gpa, &cfg, target.joinBackend(), clock.clock()); + defer core.deinit(); + // jobs on the page allocator: a freed job is unmapped, so reading one faults (the testing allocator hides it) + const pa = std.heap.page_allocator; + var host = LaneHost.init(pa, std.testing.io, &core, .{ .lanes = 4 }); + // every admitted prompt's rows in one piece and its final in the same round (as for short prompts) + const Plan = struct { + const Seq = struct { key: usize, n: u64, done: bool = false }; + seqs: std.ArrayList(Seq) = .empty, + fn self(ctx: *anyopaque) *@This() { + return @ptrCast(@alignCast(ctx)); + } + fn beginAdmit(_: *anyopaque) void {} + fn admit(ctx: *anyopaque, ask: api.Rounds.Ask, _: *[]const u8) api.Rounds.Verdict { + self(ctx).seqs.append(gpa, .{ .key = ask.key, .n = ask.prompt.len }) catch return .refuse; + return .admit; + } + fn admitted(_: *anyopaque) anyerror!void {} + fn arm(_: *anyopaque, _: usize) void {} + fn began(_: *anyopaque, _: usize, _: bool) void {} + fn plan(ctx: *anyopaque, pieces: *std.ArrayList(api.Rounds.Piece), finals: *std.ArrayList(usize)) anyerror!void { + pieces.clearRetainingCapacity(); + finals.clearRetainingCapacity(); + for (self(ctx).seqs.items) |*x| if (!x.done) { + // the host's lists: its allocator (it frees them) + try pieces.append(pa, .{ .key = x.key, .start = 0, .end = x.n - 1, .save = false }); + try finals.append(pa, x.key); + x.done = true; + }; + } + fn left(ctx: *anyopaque, key: usize) void { + const p = self(ctx); + for (p.seqs.items, 0..) |x, i| if (x.key == key) { + _ = p.seqs.orderedRemove(i); + return; + }; + } + fn after(_: *anyopaque, _: f64, _: f64) void {} + }; + var pl: Plan = .{}; + defer pl.seqs.deinit(gpa); + host.rounds = .{ .ctx = &pl, .vtable = &.{ .begin_admit = Plan.beginAdmit, .admit = Plan.admit, .admitted = Plan.admitted, .arm = Plan.arm, .began = Plan.began, .plan = Plan.plan, .left = Plan.left, .after = Plan.after } }; + const Box = struct { + mutex: std.Io.Mutex = .init, + tokens: usize = 0, + done: ?Reason = null, + fn event(ctx: *anyopaque, _: Id, e: *const Event) void { + const b: *@This() = @ptrCast(@alignCast(ctx)); + b.mutex.lockUncancelable(std.testing.io); + defer b.mutex.unlock(std.testing.io); + switch (e.*) { + .tokens => |t| b.tokens += t.len, + .finished => |f| b.done = f.reason, + else => {}, + } + } + fn wait(b: *@This()) Reason { + while (true) { + b.mutex.lockUncancelable(std.testing.io); + const d = b.done; + b.mutex.unlock(std.testing.io); + if (d) |r| return r; + std.Io.sleep(std.testing.io, .fromMilliseconds(1), .awake) catch {}; + } + } + }; + const p1 = [_]u32{ 3, 1, 4, 1, 5, 9, 2, 6 }; + const p2 = [_]u32{ 2, 7, 1, 8, 2, 8 }; + var b1: Box = .{}; + var b2: Box = .{}; + var b3: Box = .{}; + const r1: Request = .{ .prompt = &p1, .max_tokens = 12 }; + const r2: Request = .{ .prompt = &p2, .max_tokens = 12 }; + const e = host.engine(); + // both admitted before the host's first round: the first piece fails, the second's request is served + try e.submit(1, &r1, .{ .ctx = &b1, .event = Box.event }); + try e.submit(2, &r2, .{ .ctx = &b2, .event = Box.event }); + try host.start(); + defer host.stop(); + try std.testing.expectEqual(Reason.failed, b1.wait()); + try std.testing.expectEqual(Reason.length, b2.wait()); + try std.testing.expectEqual(@as(usize, 12), b2.tokens); + // the host serves the next request + try e.submit(3, &r1, .{ .ctx = &b3, .event = Box.event }); + try std.testing.expectEqual(Reason.length, b3.wait()); + try std.testing.expectEqual(@as(usize, 12), b3.tokens); +} diff --git a/zig/src/core/lane_stream.zig b/zig/src/core/lane_stream.zig new file mode 100644 index 0000000000..b213336edb --- /dev/null +++ b/zig/src/core/lane_stream.zig @@ -0,0 +1,60 @@ +//! A request's stream on the lane host: its kept prefix looked up, its drafter made, the stream set up. +const std = @import("std"); +const lanes = @import("lanes"); +const LaneHost = @import("lane_host.zig").LaneHost; +const Job = LaneHost.Job; + +/// A job's drafter freed: the family's (LaneHost.proposers) or the built-in one. +pub fn freeProposer(h: *LaneHost, job: *Job) void { + if (job.custom) |c| h.proposers.?.free(h.proposers.?.ctx, c) else job.proposer.deinit(); + job.custom = null; +} + +/// An admitted job's stream (its kept prefix looked up, its drafter): false when it failed (the job is finished). +pub fn open(h: *LaneHost, job: *Job) bool { + const r = job.request; + var reuse: lanes.stream.Reuse = .{}; + // the entry stays alive until the backend restores it: nothing keeps between here and this stream's own pass + if (h.cache) |store| if (store.lookupRewind(h.gpa, r.prompt, r.history_len, r.rewind_len, r.shared_prefixes, r.chunks, h.spans(r))) |l| { + job.entry = l.entry; + job.kept0 = store.counts.kept; + job.marks = l.marks; + reuse = .{ .saved = if (l.entry) |e| e.saved else null, .at = if (l.entry) |e| e.at else 0, .marks = l.marks, .hook = .{ .ptr = job, .at = Job.kept } }; + } else |_| {}; + if (h.proposers) |ps| { + job.custom = ps.make(ps.ctx, h.gpa, r.prompt, r.eos) catch return !h.drop(job, "the drafter could not start"); + } else job.proposer = lanes.SuffixLookup.init(h.gpa, .{ .min_match = h.min_match }) catch return !h.drop(job, "the drafter could not start"); + job.stream = lanes.Stream.init(h.gpa, .{ + .id = "request", + .prompt = r.prompt, + .max_new = r.max_tokens, + .eos = r.eos, + .sampling = r.sampling, + .drafts = r.drafts, + .proposer = job.custom orelse job.proposer.proposer(), + .stop_check = if (r.stop) |s| .{ .ptr = s.ctx, .check = s.check } else null, + .cancel_check = .{ .ptr = job, .check = LaneHost.cancelled }, + .yield_check = if (r.background and r.max_tokens == 0) .{ .ptr = job, .check = LaneHost.yieldsTo } else null, + .think_budget = r.think_budget, + .think_close = r.think_close, + .think_end = if (r.think_end) |t| t else -1, + .logprobs = r.logprobs, + .loop_guard = r.loop_guard, + .chunks = r.chunks, + .history_len = r.history_len, + .shared_prefixes = r.shared_prefixes, + .reuse = reuse, + .grammar = if (r.structure) |st| h.constraint(st) catch |e| { + freeProposer(h, job); + return !h.drop(job, if (e == error.OutOfMemory) "out of memory" else "the reply's grammar could not start (xgrammar)"); + } else null, + .decode_spans = h.spans(r), + .decide = r.decide, + }) catch { + freeProposer(h, job); + return !h.drop(job, "out of memory"); + }; + job.started = true; + job.began = std.Io.Clock.awake.now(h.io).toNanoseconds(); + return true; +} diff --git a/zig/src/core/lanes/backend.zig b/zig/src/core/lanes/backend.zig index 9d8d9d337f..d8d3451f41 100644 --- a/zig/src/core/lanes/backend.zig +++ b/zig/src/core/lanes/backend.zig @@ -100,8 +100,28 @@ pub const Backend = struct { held: ?*const fn (ptr: *anyopaque, s: *Stream, out: []u32) anyerror!void = null, /// The stream left the rounds: free its caches and held drafts. release: *const fn (ptr: *anyopaque, s: *Stream) void, + /// A prompt in pieces between rounds (`Rounds`) instead of `prefill`: begin, pieces, finish; null: whole only. + begin: ?*const fn (ptr: *anyopaque, s: *Stream) anyerror!Begun = null, + piece: ?*const fn (ptr: *anyopaque, s: *Stream, start: u64, end: u64, save: bool) anyerror!void = null, + finish: ?*const fn (ptr: *anyopaque, s: *Stream) anyerror!void = null, + /// `finish` minus the last row (Config.join_tail): true when that token is left for the first round. + join: ?*const fn (ptr: *anyopaque, s: *Stream) anyerror!bool = null, + /// A round's pieces in one call (Config.piece_runs), maybe one forward; a failure fails them all. + pieces: ?*const fn (ptr: *anyopaque, list: []const Piece) anyerror!void = null, + /// Before a round's `finish` / `join` calls (Config.replay_runs): their shared work in one go. + finals: ?*const fn (ptr: *anyopaque, streams: []const *Stream) anyerror!void = null, }; + /// One stream's prompt rows [start, end) of a round (`save`: the prompt's snapshot right after them). + pub const Piece = struct { stream: *Stream, start: u64, end: u64, save: bool }; + + /// `begin`'s outcome: restored prefix tokens, or a damaged saved entry dropped (the prompt runs from 0). + pub const Begun = struct { cached: u64 = 0, damaged: bool = false }; + + pub fn pieces(b: Backend) bool { + return b.vtable.begin != null and b.vtable.piece != null and b.vtable.finish != null; + } + pub fn prefill(b: Backend, s: *Stream) !void { return b.vtable.prefill(b.ptr, s); } diff --git a/zig/src/core/lanes/config.zig b/zig/src/core/lanes/config.zig index 85361bd839..c215c87f86 100644 --- a/zig/src/core/lanes/config.zig +++ b/zig/src/core/lanes/config.zig @@ -32,6 +32,10 @@ pub const Model = struct { draft_probabilities: bool = false, // the head gives each drafted node's chance of landing draft_streams: bool = false, // the head drafts every stream of a shared round in one batch head_trees: bool = false, // the head drafts trees of lanes (DraftRequest.lanes) and holds them + join_tail: bool = false, // a prompt in pieces leaves its last token for the next shared round's window (Backend.join) + piece_runs: bool = false, // a round's pieces go to the backend in one call (Backend.pieces: one forward over them) + join_drafts: bool = false, // with join_tail: a round's joined prompts' first drafts in one call (Engine.finishStreams) + replay_runs: bool = false, // a round's finals go to the backend together first (Backend.finals) lane_costs: []const Cost = &.{}, // a lone stream's window ms by rows, up to its widest tree (timed at load) lone_copies: bool = true, // a lone stream proposes copies (false: only streams in shared rounds do) }; @@ -55,9 +59,18 @@ pub const Config = struct { fill_floor: f64 = 0.35, // the share of grafted rows a stream must keep to go on reading its chain back every round draft_streams: bool, head_trees: bool, + /// a prompt in pieces leaves its last token as the pending row of its first shared round's window + join_tail: bool = false, + /// a round's pieces of several streams in one backend call (Backend.pieces), which may run them as one forward + piece_runs: bool = false, + /// with join_tail: a round's joined prompts get their first drafts in one backend call + join_drafts: bool = false, + /// a round's finished prompts to the backend together before their finishes (Backend.finals) + replay_runs: bool = false, head_lanes: u32 = 0, // a lone stream's lanes at most, the planner picking each round's (0: the depth rule's chain) head_depth: u32 = 16, // and their depth at most own_prices: bool = true, // the planner prices each pick from the stream's own measured rounds (seeded from the table) + timed: bool = true, // false: rounds priced from the cost tables alone, no wall clock (a request drafts the same whatever the load or the requests before it) depth_prior: []f64, family_costs: Table, shared_costs: Table, @@ -115,6 +128,10 @@ pub const Config = struct { .node_probabilities = m.draft_probabilities, .draft_streams = m.draft_streams, .head_trees = m.head_trees, + .join_tail = m.join_tail, + .piece_runs = m.piece_runs, + .join_drafts = m.join_drafts, + .replay_runs = m.replay_runs, .depth_prior = depth_prior, .family_costs = family_costs, .shared_costs = try alloc.extendCosts(gpa, timed.items, batch_rows), diff --git a/zig/src/core/lanes/engine.zig b/zig/src/core/lanes/engine.zig index 4339af40aa..41c8dfcfec 100644 --- a/zig/src/core/lanes/engine.zig +++ b/zig/src/core/lanes/engine.zig @@ -13,6 +13,7 @@ const plan_lanes = @import("plan_lanes.zig"); const fill = @import("fill.zig"); const trail = @import("trail.zig"); const open = @import("open.zig"); +const trim = @import("trim.zig"); const Stream = sm.Stream; const Feed = be.Feed; const LogRow = @import("logprob.zig").Row; @@ -26,11 +27,20 @@ const Result = struct { rows: u32, keep: u32, got: []const u32 }; const Unread = struct { event: usize, handle: u64 }; pub const Engine = struct { + const pieces_ = @import("pieces.zig"); + pub const beginStream = pieces_.beginStream; + pub const piece = pieces_.piece; + pub const pieces = pieces_.pieces; + pub const prepareFinals = pieces_.prepareFinals; + pub const finishStream = pieces_.finishStream; + pub const finishStreams = pieces_.finishStreams; + gpa: Allocator, cfg: *const Config, backend: be.Backend, clock: be.Clock, log: ?*ev.Log = null, + trim: ?trim.Trim = null, // a family's row policy in place of the built-in allocation (lanes/trim.zig) rule: depth.Rule, live: std.ArrayList(*Stream) = .empty, // admitted streams in admission order (Python `_live`) drafted: u64 = 0, @@ -218,7 +228,7 @@ pub const Engine = struct { s.priced = null; const a = e.arena.allocator(); var plans = [_]Plan{try win.plan(e, s, copied)}; - if (plans[0].kind == .head and e.cfg.node_probabilities) try win.allocate(e, &plans); + if (plans[0].kind == .head and (e.cfg.node_probabilities or e.trim != null)) try win.allocate(e, &plans); try win.constrain(e, &plans[0]); const plan = plans[0]; const early = e.cfg.family_mtp and s.drafts and plan.kind != .forced and e.cfg.speculate_early and plan.parents == null and s.grammar == null; @@ -240,14 +250,14 @@ pub const Engine = struct { if (early) if (e.backend.vtable.unspeculate) |undo| try undo(e.backend.ptr, s); try e.draftLate(s, plan.position, o.follow, o.path, null); } - if (plan.kind == .head or (plan.kind == .none and e.cfg.plain_guard)) { + if (e.cfg.timed and (plan.kind == .head or (plan.kind == .none and e.cfg.plain_guard))) { const ms = e.clock.elapsedMs(.cost); const first = s.rounds == 1; // a stream's first round carries the prefill-to-decode switch try trail.event(e, &.{ f("ev", str("cost")), f("stream", str(s.id)), f("drafts", int(rows - 1)), f("init", .{ .bool = first }), f("ms", ev.bits(ms)) }); try e.rule.observeCost(@intCast(rows - 1), ms, first, win.who(s)); } } - if (s.graft_room != null and s.rounds > 1) { + if (e.cfg.timed and s.graft_room != null and s.rounds > 1) { // the stream's round time beyond its window and head levels, for the lane planner if (e.cfg.lane_costs.get(@intCast(rows))) |w| { const over = e.clock.elapsedMs(.cost) - w - e.cfg.mtp_step_ms * @as(f64, @floatFromInt(held_levels)); @@ -257,7 +267,7 @@ pub const Engine = struct { // the pick that shaped this round, priced from its own measured rounds next time if (priced) |p| s.measured[p.choice].add(e.clock.elapsedMs(.cost) - p.ms); } - if (plan.kind == .head) { + if (e.cfg.timed and plan.kind == .head) { const ms = e.clock.elapsedMs(.overhead); try trail.event(e, &.{ f("ev", str("overhead")), f("streams", int(1)), f("rows", int(rows)), f("ms", ev.bits(ms)) }); try e.rule.observeOverhead(1, rows, ms); @@ -314,9 +324,11 @@ pub const Engine = struct { for (heads.items, budgets) |i, d| try e.draftLate(plans[i].stream, plans[i].position, outcomes[i].follow, outcomes[i].path, d); } try e.backend.keep(windows, paths); - const ms = e.clock.elapsedMs(.overhead); - try trail.event(e, &.{ f("ev", str("overhead")), f("streams", int(plans.len)), f("rows", int(total)), f("ms", ev.bits(ms)) }); - try e.rule.observeOverhead(@intCast(plans.len), total, ms); + if (e.cfg.timed) { + const ms = e.clock.elapsedMs(.overhead); + try trail.event(e, &.{ f("ev", str("overhead")), f("streams", int(plans.len)), f("rows", int(total)), f("ms", ev.bits(ms)) }); + try e.rule.observeOverhead(@intCast(plans.len), total, ms); + } } /// Commit the target-sampled path (after the thinking budget's cut), each token with its path row's logprob row. @@ -334,6 +346,7 @@ pub const Engine = struct { path = try accept.acceptPath(a, tokens, parents, sampled); for (path[1..]) |r| try committed.append(a, tokens[r]); try committed.append(a, sampled[path[path.len - 1]]); + if (s.proposer) |pr| pr.round(@intCast(rows), @intCast(path.len - 1)); if (rows > 1) { const kept = path.len - 1; e.drafted += rows - 1; @@ -371,6 +384,10 @@ pub const Engine = struct { } } } + if (e.trim) |t| if (p.kind != .forced) { + const chain = parents.len == 0 or accept.isChain(parents); + try t.commit(.{ .stream = s, .drafted = p.kind == .head, .rows = @intCast(rows), .parents = if (chain) null else parents, .path = path, .bonus = sampled[path[path.len - 1]] }); + }; const budget_cut = s.thinkCut(committed.items); const loop_cut = sm.loopCut(s, committed.items); const loop_wins = if (loop_cut) |c| budget_cut == null or c < budget_cut.? else false; @@ -427,7 +444,7 @@ pub const Engine = struct { const s = requests[0].stream; if (s.odds == null) s.odds = shape.Odds.init(shape.start_odds, shape.start_weight); const cap: usize = @intCast(@max(1, @min(@min(shape.max_depth, e.cfg.head_depth), s.draftRoom() - 1))); - const costs = plan_lanes.Costs{ .window = &e.cfg.lane_costs, .head_ms = e.cfg.mtp_step_ms, .over_ms = s.round_over orelse 0, .tree_over_ms = s.tree_over, .measured = if (e.cfg.own_prices) &s.measured else null }; + const costs = plan_lanes.Costs{ .window = &e.cfg.lane_costs, .head_ms = e.cfg.mtp_step_ms, .over_ms = s.round_over orelse 0, .tree_over_ms = s.tree_over, .measured = if (e.cfg.own_prices and e.cfg.timed) &s.measured else null }; // every eighth round one size wider than the best, so the odds of deeper lanes stay measured const p = try plan_lanes.pick(e.gpa, s.odds.?, costs, e.cfg.head_lanes, cap, try e.graftProspect(s), s.rounds % 8 == 4); s.picks[p.choice] += 1; diff --git a/zig/src/core/lanes/engine_test.zig b/zig/src/core/lanes/engine_test.zig index 2100cc3762..5c9043d336 100644 --- a/zig/src/core/lanes/engine_test.zig +++ b/zig/src/core/lanes/engine_test.zig @@ -460,3 +460,124 @@ test "a family that keeps copies to shared rounds proposes none for a lone strea try std.testing.expectEqual(lone_copies, (try win.copyProposal(&engine, &s)).len >= 2); } } + +/// The cases' prompts in pieces, case i opened before round 2 i, others decoding; `join`: Config.join_tail. +fn runPieces(cases: []const Case, join: bool, joins: *usize) ![][]u32 { + var cfg = try model(); + defer cfg.deinit(gpa); + cfg.join_tail = join; + var target: fake.Fake = .{ .gpa = gpa }; + defer target.deinit(); + var clock: fake.FixedClock = .{}; + var engine = Engine.init(gpa, &cfg, target.joinBackend(), clock.clock()); + defer engine.deinit(); + const streams = try gpa.alloc(sm.Stream, cases.len); + defer gpa.free(streams); + for (cases, streams) |c, *s| s.* = try sm.Stream.init(gpa, .{ .id = "s", .prompt = c.prompt, .max_new = c.max_new, .eos = &.{96}, .sampling = c.sampling, .drafts = c.drafts, .think_budget = c.think_budget, .think_close = &.{ 90, 91, 92 }, .think_end = 91 }); + defer for (streams) |*s| s.deinit(gpa); + var opened: usize = 0; + var round: usize = 0; + while (opened < streams.len or engine.activeCount() > 0) : (round += 1) { + if (opened < streams.len and round == 2 * opened) { + const s = &streams[opened]; + _ = try engine.beginStream(s); + if (s.prompt_len > 1) try engine.piece(s, 0, s.prompt_len - 1, false); + try engine.finishStream(s); + opened += 1; + } + if (engine.activeCount() > 0) try engine.step(); + } + joins.* = target.joins; + const out = try gpa.alloc([]u32, cases.len); + for (out, streams) |*o, *s| o.* = try gpa.dupe(u32, s.emitted()); + return out; +} + +test "a prompt's last token joins its first round's shared window: the same tokens, greedy and sampled" { + const cases = [_]Case{ + .{ .prompt = &.{ 3, 1, 4, 1, 5, 9, 2, 6 }, .max_new = 37 }, + .{ .prompt = &.{ 2, 7, 1, 8, 2, 8 }, .max_new = 29, .sampling = .{ .seed = 5, .temperature = 0.8, .top_k = 7 } }, + .{ .prompt = &.{4}, .max_new = 21 }, + .{ .prompt = &.{ 1, 1, 2, 3, 5 }, .max_new = 25, .drafts = false }, + .{ .prompt = &.{ 6, 6, 6, 1 }, .max_new = 30, .think_budget = 1 }, + }; + var j0: usize = 0; + var j1: usize = 0; + const off = try runPieces(&cases, false, &j0); + defer free(off); + const on = try runPieces(&cases, true, &j1); + defer free(on); + try std.testing.expectEqual(@as(usize, 0), j0); + try std.testing.expectEqual(@as(usize, 3), j1); // not the undrafted stream nor the budget cutting the first token + for (off, on) |a, b| try std.testing.expectEqualSlices(u32, a, b); + // and the whole prompts' tokens + const whole = try run(&cases); + defer free(whole); + for (whole, on) |a, b| try std.testing.expectEqualSlices(u32, a, b); +} + +/// Every prompt in one piece, all finished in one round (`together`: Engine.finishStreams), then rounds until done. +fn runFinals(cases: []const Case, together: bool, calls: *usize, most: *usize) ![][]u32 { + var cfg = try model(); + defer cfg.deinit(gpa); + cfg.join_tail = true; + cfg.join_drafts = together; + var target: fake.Fake = .{ .gpa = gpa }; + defer target.deinit(); + var clock: fake.FixedClock = .{}; + var engine = Engine.init(gpa, &cfg, target.joinBackend(), clock.clock()); + defer engine.deinit(); + const streams = try gpa.alloc(sm.Stream, cases.len); + defer gpa.free(streams); + for (cases, streams) |c, *s| s.* = try sm.Stream.init(gpa, .{ .id = "s", .prompt = c.prompt, .max_new = c.max_new, .eos = &.{96}, .sampling = c.sampling, .drafts = c.drafts, .think_budget = c.think_budget, .think_close = &.{ 90, 91, 92 }, .think_end = 91 }); + defer for (streams) |*s| s.deinit(gpa); + const ptrs = try gpa.alloc(*sm.Stream, cases.len); + defer gpa.free(ptrs); + for (streams, ptrs) |*s, *p| { + _ = try engine.beginStream(s); + if (s.prompt_len > 1) try engine.piece(s, 0, s.prompt_len - 1, false); + p.* = s; + } + // the finish phase alone: its draft calls and their largest + target.drafts_calls = 0; + target.drafts_most = 0; + if (together) { + const errs = try gpa.alloc(?anyerror, cases.len); + defer gpa.free(errs); + _ = try engine.finishStreams(ptrs, errs); + for (errs) |x| if (x) |e| return e; + } else for (ptrs) |s| try engine.finishStream(s); + calls.* = target.drafts_calls; + most.* = target.drafts_most; + while (engine.activeCount() > 0) try engine.step(); + const out = try gpa.alloc([]u32, cases.len); + for (out, streams) |*o, *s| o.* = try gpa.dupe(u32, s.emitted()); + return out; +} + +test "a round's finished prompts get their first drafts in one call (Python's _windows), the same tokens" { + const cases = [_]Case{ + .{ .prompt = &.{ 3, 1, 4, 1, 5, 9, 2, 6 }, .max_new = 37 }, + .{ .prompt = &.{ 2, 7, 1, 8, 2, 8 }, .max_new = 29, .sampling = .{ .seed = 5, .temperature = 0.8, .top_k = 7 } }, + .{ .prompt = &.{4}, .max_new = 21 }, + .{ .prompt = &.{ 1, 1, 2, 3, 5 }, .max_new = 25, .drafts = false }, + .{ .prompt = &.{ 6, 6, 6, 1 }, .max_new = 30, .think_budget = 1 }, + }; + var c1: usize = 0; + var m1: usize = 0; + var c0: usize = 0; + var m0: usize = 0; + const together = try runFinals(&cases, true, &c1, &m1); + defer free(together); + const apart = try runFinals(&cases, false, &c0, &m0); + defer free(apart); + for (apart, together) |a, b| try std.testing.expectEqualSlices(u32, a, b); + const whole = try run(&cases); + defer free(whole); + for (whole, together) |a, b| try std.testing.expectEqualSlices(u32, a, b); + // three joined prompts in one call, the budget-cut one its own, the undrafted none: 2 calls against 4 + try std.testing.expectEqual(@as(usize, 2), c1); + try std.testing.expectEqual(@as(usize, 3), m1); + try std.testing.expectEqual(@as(usize, 4), c0); + try std.testing.expectEqual(@as(usize, 1), m0); +} diff --git a/zig/src/core/lanes/fake.zig b/zig/src/core/lanes/fake.zig index dd60ca7f64..50290a9702 100644 --- a/zig/src/core/lanes/fake.zig +++ b/zig/src/core/lanes/fake.zig @@ -53,11 +53,18 @@ pub const Fake = struct { answer_cycles: bool = false, prefill_chunks: usize = 0, prefill_count: usize = 0, + fail_pieces: usize = 0, // the next this many `piece` calls fail with error.Unbound prefill_hook: ?*const fn (ctx: *anyopaque, s: *Stream, chunk: usize) void = null, prefill_hook_ctx: ?*anyopaque = null, refuse_sampled: bool = false, // prefill refuses a sampled stream with error.SamplingRefused ranked: bool = false, // the head gives its held chain back (alternatives) memory_lanes: ?usize = null, // prefill past this many live lanes fails with error.LaneMemory (no room yet) + pieces_run: usize = 0, // prompt pieces run (`pieceBackend`) + joins: usize = 0, // prompts whose last token joined a round (`joinBackend`) + drafts_calls: usize = 0, // `draft` calls, and the most requests one took + drafts_most: usize = 0, + fail_runs: usize = 0, // the next this many `pieces` calls fail with error.Unbound (`runBackend`) + runs: usize = 0, // `pieces` calls that ran pub fn deinit(x: *Fake) void { var it = x.lanes.valueIterator(); @@ -81,6 +88,66 @@ pub const Fake = struct { return .{ .ptr = x, .vtable = &.{ .prefill = prefill, .prefill_step = prefillStep, .first = first, .queue = queue, .read = read, .verify = verify, .keep = keep, .draft = draft, .features = features, .release = release, .first_row = firstRow } }; } + /// The backend with prompts in pieces (begin / piece / finish) beside whole prompts. + pub fn pieceBackend(x: *Fake) be.Backend { + return .{ .ptr = x, .vtable = &.{ .prefill = prefill, .first = first, .queue = queue, .read = read, .verify = verify, .keep = keep, .draft = draft, .release = release, .begin = begin, .piece = piece, .finish = finish } }; + } + + /// `pieceBackend` whose prompts may leave their last token to the stream's first round (Config.join_tail). + pub fn joinBackend(x: *Fake) be.Backend { + return .{ .ptr = x, .vtable = &.{ .prefill = prefill, .first = first, .queue = queue, .read = read, .verify = verify, .keep = keep, .draft = draft, .release = release, .begin = begin, .piece = piece, .finish = finish, .join = join } }; + } + + /// `joinBackend` that takes a round's pieces in one call (Config.piece_runs): a piece at a time, or a failure. + pub fn runBackend(x: *Fake) be.Backend { + return .{ .ptr = x, .vtable = &.{ .prefill = prefill, .first = first, .queue = queue, .read = read, .verify = verify, .keep = keep, .draft = draft, .release = release, .begin = begin, .piece = piece, .finish = finish, .join = join, .pieces = pieces } }; + } + + fn pieces(ptr: *anyopaque, list: []const be.Backend.Piece) anyerror!void { + const x = self(ptr); + if (x.fail_runs > 0) { + x.fail_runs -= 1; + return error.Unbound; + } + x.runs += 1; + for (list) |p| try piece(ptr, p.stream, p.start, p.end, p.save); + } + + fn join(ptr: *anyopaque, s: *Stream) anyerror!bool { + const x = self(ptr); + if (x.lane(s).history.items.len + 1 != s.prompt_len) return error.PositionMismatch; + x.joins += 1; + return true; + } + + fn begin(ptr: *anyopaque, s: *Stream) anyerror!be.Backend.Begun { + const x = self(ptr); + const got = try x.lanes.getOrPut(x.gpa, s); + if (got.found_existing) freeLane(x.gpa, got.value_ptr); + got.value_ptr.* = .{}; + return .{}; + } + + fn piece(ptr: *anyopaque, s: *Stream, start: u64, end: u64, snapshot: bool) anyerror!void { + _ = snapshot; + const x = self(ptr); + if (x.fail_pieces > 0) { + x.fail_pieces -= 1; + return error.Unbound; + } + const h = &x.lane(s).history; + if (h.items.len != start or end + 1 > s.prompt_len) return error.PositionMismatch; + try h.appendSlice(x.gpa, s.prompt()[start..end]); + x.pieces_run += 1; + } + + fn finish(ptr: *anyopaque, s: *Stream) anyerror!void { + const x = self(ptr); + const h = &x.lane(s).history; + if (h.items.len + 1 != s.prompt_len) return error.PositionMismatch; + try h.append(x.gpa, s.prompt()[s.prompt_len - 1]); + } + fn self(ptr: *anyopaque) *Fake { return @ptrCast(@alignCast(ptr)); } @@ -285,6 +352,8 @@ pub const Fake = struct { fn draft(ptr: *anyopaque, requests: []const be.DraftRequest) anyerror!void { const x = self(ptr); + x.drafts_calls += 1; + x.drafts_most = @max(x.drafts_most, requests.len); for (requests) |r| { const l = x.lane(r.stream); // the head reads the last verify's kept rows (a shared round drafts before its rollback) diff --git a/zig/src/core/lanes/lanes.zig b/zig/src/core/lanes/lanes.zig index 910f501675..8a1f7213e9 100644 --- a/zig/src/core/lanes/lanes.zig +++ b/zig/src/core/lanes/lanes.zig @@ -14,6 +14,7 @@ pub const cost_rule = @import("cost_rule.zig"); pub const gpu_rule = @import("gpu_rule.zig"); pub const gpu_full = @import("gpu_full.zig"); pub const depth = @import("depth.zig"); +pub const trim = @import("trim.zig"); pub const config = @import("config.zig"); pub const stream = @import("stream.zig"); pub const backend = @import("backend.zig"); @@ -34,6 +35,7 @@ pub const Stream = stream.Stream; pub const Sampling = sampling.Sampling; pub const SuffixLookup = proposer.SuffixLookup; pub const LogprobRow = logprob.Row; +pub const Trim = trim.Trim; test { std.testing.refAllDecls(@This()); diff --git a/zig/src/core/lanes/open.zig b/zig/src/core/lanes/open.zig index 7cddf9a064..cd832256ab 100644 --- a/zig/src/core/lanes/open.zig +++ b/zig/src/core/lanes/open.zig @@ -23,8 +23,16 @@ pub fn add(e: *Engine, s: *Stream) !void { try opened(e, s); } +/// How a stream's first drafts are sized: the depth rule, or the head's whole reach under node probabilities. +pub const Reach = enum { rule, head }; + /// After the prompt pass: the first token drawn and committed, the first drafts asked; rounds from the next. pub fn opened(e: *Engine, s: *Stream) !void { + return openedWith(e, s, .rule); +} + +/// `opened` with its first drafts sized by `reach` (a family planner's prompts in pieces: `.head`). +pub fn openedWith(e: *Engine, s: *Stream, reach: Reach) !void { if (s.isCancelled()) { // cancelled in its last chunk: no first token e.backend.release(s); return error.Cancelled; @@ -40,7 +48,7 @@ pub fn opened(e: *Engine, s: *Stream) !void { var asked: ?u32 = null; if (e.cfg.family_mtp and s.drafts) { // the head reads the prompt's last row and the first token, and drafts the one after it - const d: u32 = @intCast(try e.rule.depth(win.who(s))); + const d = try firstDepth(e, s, reach); asked = d; try e.backend.draft(&.{.{ .stream = s, .follow = &.{}, .first = feed, .rows = null, .start = s.prompt_len, .position = position + 1, .depth = d }}); s.dropHeld(e.gpa); @@ -70,6 +78,13 @@ pub fn opened(e: *Engine, s: *Stream) !void { try e.live.append(e.gpa, s); } +/// The first drafts' depth for `reach`. +pub fn firstDepth(e: *Engine, s: *Stream, reach: Reach) !u32 { + var who = win.who(s); + if (reach == .head and e.cfg.node_probabilities) return @intCast(try e.rule.headDepth(&who, null, win.probe(e))); + return @intCast(try e.rule.depth(who)); +} + /// Whether the backend can run a prompt pass a chunk at a time between rounds (`fillStream`). pub fn fills(e: *const Engine) bool { return e.backend.vtable.prefill_step != null; diff --git a/zig/src/core/lanes/pieces.zig b/zig/src/core/lanes/pieces.zig new file mode 100644 index 0000000000..cd64bb1eb0 --- /dev/null +++ b/zig/src/core/lanes/pieces.zig @@ -0,0 +1,144 @@ +//! Prompts in pieces on the lane core: a stream begun, its prompt rows run in pieces, then its first token. +const std = @import("std"); +const be = @import("backend.zig"); +const sm = @import("stream.zig"); +const ev = @import("events.zig"); +const trail = @import("trail.zig"); +const open = @import("open.zig"); +const Engine = @import("engine.zig").Engine; +const Stream = sm.Stream; +const f = ev.f; +const str = trail.str; +const int = trail.int; + +/// A stream whose prompt runs in pieces between rounds: `piece`s, then `finishStream`; a failed begin holds nothing. +pub fn beginStream(e: *Engine, s: *Stream) !be.Backend.Begun { + try trail.event(e, &.{ f("ev", str("add")), f("stream", str(s.id)) }); + const begin = e.backend.vtable.begin orelse return error.NoPieces; + const b = try begin(e.backend.ptr, s); + s.cached = @intCast(b.cached); + return b; +} + +/// Prompt rows [start, end) of a begun stream (`save`: the prompt's snapshot right after them). +pub fn piece(e: *Engine, s: *Stream, start: u64, end: u64, save: bool) !void { + const run = e.backend.vtable.piece orelse return error.NoPieces; + return run(e.backend.ptr, s, start, end, save); +} + +/// A round's pieces of several streams in one backend call (Config.piece_runs, Backend.pieces); else one at a time. +pub fn pieces(e: *Engine, list: []const be.Backend.Piece) !void { + if (e.cfg.piece_runs) if (e.backend.vtable.pieces) |run| return run(e.backend.ptr, list); + for (list) |p| try piece(e, p.stream, p.start, p.end, p.save); +} + +/// A round's finished prompts before their `finishStream` calls (Config.replay_runs, Backend.finals). +pub fn prepareFinals(e: *Engine, streams: []const *Stream) !void { + if (!e.cfg.replay_runs or streams.len < 2) return; + const run = e.backend.vtable.finals orelse return; + return run(e.backend.ptr, streams); +} + +/// A begun stream's prompt is in: first token and drafts; a failure is the host's to discard (Cancelled: released). +pub fn finishStream(e: *Engine, s: *Stream) !void { + _ = e.arena.reset(.retain_capacity); + // Config.join_tail: the last token waits for the next round (a forced or budget-cut first token keeps its row) + if (e.cfg.join_tail and e.cfg.family_mtp and s.drafts and s.prompt_len > 0 and s.force.items.len == 0 and !s.cutsNext()) { + if (e.backend.vtable.join) |join| { + if (join(e.backend.ptr, s) catch |x| return released(e, s, x)) return joined(e, s); + return open.openedWith(e, s, .head); + } + } + const fin = e.backend.vtable.finish orelse return error.NoPieces; + fin(e.backend.ptr, s) catch |x| return released(e, s, x); + return open.openedWith(e, s, .head); +} + +/// A round's finished prompts as `finishStream`, joined ones' first drafts in one call; returns how many joined. +pub fn finishStreams(e: *Engine, streams: []const *Stream, errs: []?anyerror) !usize { + _ = e.arena.reset(.retain_capacity); + const a = e.arena.allocator(); + var reqs: std.ArrayList(be.DraftRequest) = .empty; + var joins: std.ArrayList(usize) = .empty; + for (streams, errs, 0..) |s, *err, i| { + err.* = null; + const can = e.cfg.join_tail and e.cfg.family_mtp and s.drafts and s.prompt_len > 0 and s.force.items.len == 0 and !s.cutsNext(); + const join = if (can) e.backend.vtable.join else null; + if (join) |j| { + const ok = j(e.backend.ptr, s) catch |x| { + err.* = released(e, s, x); + continue; + }; + if (ok) { + try reqs.append(a, joinRequest(e, s) catch |x| { + err.* = x; + continue; + }); + try joins.append(a, i); + continue; + } + } else { + const fin = e.backend.vtable.finish orelse return error.NoPieces; + fin(e.backend.ptr, s) catch |x| { + err.* = released(e, s, x); + continue; + }; + } + open.openedWith(e, s, .head) catch |x| { + err.* = x; + }; + } + if (reqs.items.len == 0) return 0; + e.backend.draft(reqs.items) catch |x| { + for (joins.items) |i| errs[i] = x; + return 0; + }; + for (reqs.items) |r| joinedDrafted(e, r) catch |x| { + for (joins.items) |i| if (streams[i] == r.stream) { + errs[i] = x; + }; + }; + return reqs.items.len; +} + +/// A cancelled stream's lane is released here: the host finishes it without the core. +fn released(e: *Engine, s: *Stream, err: anyerror) anyerror { + if (err == error.Cancelled) e.backend.release(s); + return err; +} + +/// A joined prompt (Backend.join): its last token pending, drafts asked now, verified in its first shared round. +fn joined(e: *Engine, s: *Stream) !void { + s.context.shrinkRetainingCapacity(s.prompt_len); + const last = s.context.items[s.prompt_len - 1]; + s.cache_len = s.prompt_len - 1; + s.pending = last; + const position: u64 = s.prompt_len; + const d = try open.firstDepth(e, s, .head); + try e.backend.draft(&.{.{ .stream = s, .follow = &.{}, .first = .{ .value = last }, .rows = null, .start = s.cache_len, .position = position, .depth = d }}); + s.dropHeld(e.gpa); + s.next = .{ .count = d }; + try trail.event(e, &.{ f("ev", str("draft")), f("stream", str(s.id)), f("depth", int(d)), f("position", int(position)), f("follow", .{ .u32s = &.{last} }), f("rows", .null) }); + try trail.resolve(e); + try e.live.append(e.gpa, s); +} + +/// After its first draft: the joined stream holds it and takes part from the next round. +fn joinedDrafted(e: *Engine, r: be.DraftRequest) !void { + const s = r.stream; + s.dropHeld(e.gpa); + s.next = .{ .count = r.depth }; + try trail.event(e, &.{ f("ev", str("draft")), f("stream", str(s.id)), f("depth", int(r.depth)), f("position", int(r.position)), f("follow", .{ .u32s = &.{s.pending.?} }), f("rows", .null) }); + try trail.resolve(e); + try e.live.append(e.gpa, s); +} + +/// A joined prompt's state and its first draft request (the last token pending at the prompt's last position). +fn joinRequest(e: *Engine, s: *Stream) !be.DraftRequest { + s.context.shrinkRetainingCapacity(s.prompt_len); + const last = s.context.items[s.prompt_len - 1]; + s.cache_len = s.prompt_len - 1; + s.pending = last; + const d = try open.firstDepth(e, s, .head); + return .{ .stream = s, .follow = &.{}, .first = .{ .value = last }, .rows = null, .start = s.cache_len, .position = s.prompt_len, .depth = d }; +} diff --git a/zig/src/core/lanes/proposer.zig b/zig/src/core/lanes/proposer.zig index f49626e74a..3df9a73393 100644 --- a/zig/src/core/lanes/proposer.zig +++ b/zig/src/core/lanes/proposer.zig @@ -11,6 +11,10 @@ pub const Proposer = struct { propose: *const fn (ptr: *anyopaque, context: []const u32, max_draft: i64) anyerror![]const u32, last_match: *const fn (ptr: *anyopaque) i64, observe: ?*const fn (ptr: *anyopaque, proposed: i64, accepted: i64) void = null, + /// The proposer prices its own drafts: `propose`'s result is taken as is, no width, length or match gate. + priced: bool = false, + /// Every non-forced round of the stream, whatever drafted it: the window's rows and its drafts kept. + round: ?*const fn (ptr: *anyopaque, rows: u32, kept: u32) void = null, }; /// Tokens valid until the next `propose`. @@ -25,6 +29,10 @@ pub const Proposer = struct { pub fn observe(p: Proposer, proposed: i64, accepted: i64) void { if (p.vtable.observe) |f| f(p.ptr, proposed, accepted); } + + pub fn round(p: Proposer, rows: u32, kept: u32) void { + if (p.vtable.round) |f| f(p.ptr, rows, kept); + } }; pub const max_ngram = 8; diff --git a/zig/src/core/lanes/trim.zig b/zig/src/core/lanes/trim.zig new file mode 100644 index 0000000000..426577030f --- /dev/null +++ b/zig/src/core/lanes/trim.zig @@ -0,0 +1,72 @@ +//! An optional row policy in place of windows.allocate: a family's own depth rule trims drafts. +const std = @import("std"); +const Allocator = std.mem.Allocator; +const Stream = @import("stream.zig").Stream; + +/// One planned window as the policy sees it. +pub const Window = struct { + stream: *Stream, + drafted: bool, // head drafts the policy may trim; every other window runs whole + rows: u32, // rows as planned, the pending row included + parents: ?[]const i32, // the drafts' parents for a tree (null: a chain) +}; + +/// What a window keeps: its first `count` drafts, or the listed drafts (ascending, each one's parent listed). +pub const Choice = struct { count: u32, nodes: ?[]const u32 = null }; + +/// A committed window: its rows, the accepted path (row 0 first) and the target's choice after the path's last row. +pub const Commit = struct { + stream: *Stream, + drafted: bool, + rows: u32, + parents: ?[]const i32, // each row's parent row (row 0: -1) for a tree; null for a chain + path: []const u32, + bonus: u32, +}; + +pub const Trim = struct { + ptr: *anyopaque, + vtable: *const VTable, + + pub const VTable = struct { + /// Each window's drafts for this round; `alone`: the round's only window. Slices in `out` live in `arena`. + choose: *const fn (ptr: *anyopaque, windows: []const Window, alone: bool, arena: Allocator, out: []Choice) anyerror!void, + /// Every window's commit, in plan order (a shared round's windows one after the other). + commit: *const fn (ptr: *anyopaque, c: Commit) anyerror!void, + }; + + pub fn choose(t: Trim, windows: []const Window, alone: bool, arena: Allocator, out: []Choice) !void { + return t.vtable.choose(t.ptr, windows, alone, arena, out); + } + pub fn commit(t: Trim, c: Commit) !void { + return t.vtable.commit(t.ptr, c); + } +}; + +pub const Sub = struct { tokens: []u32, parents: []i32 }; + +/// The listed drafts of a tree re-indexed among themselves (parents -1: the pending row); arena-owned. +pub fn select(arena: Allocator, tokens: []const u32, parents: []const i32, nodes: []const u32) !Sub { + const at = try arena.alloc(i32, tokens.len); + @memset(at, -1); + const out_t = try arena.alloc(u32, nodes.len); + const out_p = try arena.alloc(i32, nodes.len); + for (nodes, out_t, out_p, 0..) |n, *t, *p, i| { + const q = parents[n]; + if (q >= 0 and at[@intCast(q)] < 0) return error.ParentNotKept; + at[n] = @intCast(i); + t.* = tokens[n]; + p.* = if (q < 0) -1 else at[@intCast(q)]; + } + return .{ .tokens = out_t, .parents = out_p }; +} + +test "select keeps listed drafts with their parents re-indexed" { + var arena = std.heap.ArenaAllocator.init(std.testing.allocator); + defer arena.deinit(); + // chain 10 11 12, sibling 20 (on the pending row) with child 21 + const s = try select(arena.allocator(), &.{ 10, 11, 12, 20, 21 }, &.{ -1, 0, 1, -1, 3 }, &.{ 0, 3, 4 }); + try std.testing.expectEqualSlices(u32, &.{ 10, 20, 21 }, s.tokens); + try std.testing.expectEqualSlices(i32, &.{ -1, -1, 1 }, s.parents); + try std.testing.expectError(error.ParentNotKept, select(arena.allocator(), &.{ 10, 11 }, &.{ -1, 0 }, &.{1})); +} diff --git a/zig/src/core/lanes/windows.zig b/zig/src/core/lanes/windows.zig index 901cae7b8d..20e401e05e 100644 --- a/zig/src/core/lanes/windows.zig +++ b/zig/src/core/lanes/windows.zig @@ -4,6 +4,7 @@ const accept = @import("accept.zig"); const alloc = @import("allocate.zig"); const be = @import("backend.zig"); const depth = @import("depth.zig"); +const trim = @import("trim.zig"); const fill = @import("fill.zig"); const shape = @import("shape.zig"); const ev = @import("events.zig"); @@ -137,6 +138,7 @@ pub fn copyProposal(e: *Engine, s: *Stream) ![]const u32 { const p = s.proposer orelse return &.{}; if (e.cfg.max_copy == 0 or s.force.items.len > 0) return &.{}; if (e.alone and !e.cfg.lone_copies) return &.{}; + if (p.vtable.priced) return p.propose(s.context.items, @min(@as(i64, e.cfg.max_copy), s.draftRoom() - 1)) catch &.{}; const width: u32 = if (e.alone) (s.copy_width orelse e.cfg.first_copy) else e.cfg.first_copy; const copied = p.propose(s.context.items, @min(@as(i64, width), s.draftRoom() - 1)) catch return &.{}; if (copied.len < 2 or p.lastMatch() < e.cfg.enter_match) return &.{}; @@ -160,6 +162,7 @@ pub fn who(s: *Stream) depth.Who { /// Trim draft prefixes by landing probability to the shared width with the most expected tokens a ms. pub fn allocate(e: *Engine, plans: []Plan) !void { + if (e.trim) |t| return trimmed(e, plans, t); const a = e.arena.allocator(); const fixed = try a.alloc(u32, plans.len); const probs = try a.alloc([]const f64, plans.len); @@ -203,6 +206,37 @@ pub fn allocate(e: *Engine, plans: []Plan) !void { try trail.event(e, &.{ f("ev", trail.str("alloc")), f("plans", .{ .list = rows }) }); } +/// A family's row policy decides each head window's drafts (lanes/trim.zig); other windows run whole. +fn trimmed(e: *Engine, plans: []Plan, t: trim.Trim) !void { + const a = e.arena.allocator(); + const asks = try a.alloc(trim.Window, plans.len); + for (asks, plans) |*w, p| w.* = .{ .stream = p.stream, .drafted = p.kind == .head, .rows = @intCast(1 + p.count()), .parents = p.parents }; + const out = try a.alloc(trim.Choice, plans.len); + for (out, plans) |*o, p| o.* = .{ .count = @intCast(p.count()) }; + try t.choose(asks, e.alone and plans.len == 1, a, out); + for (plans, out) |*p, c| { + if (p.kind != .head) continue; + if (c.nodes) |nodes| { + if (p.parents == null or p.held > 0) return error.TrimNodesNeedTree; + const sub = try trim.select(a, p.tokens, p.parents.?, nodes); + p.tokens = sub.tokens; + p.parents = if (accept.isChain(try accept.rowParents(a, 1 + sub.tokens.len, sub.parents))) null else sub.parents; + } else if (c.count == 0) { + p.kind = .none; + p.held = 0; + p.tokens = &.{}; + p.parents = null; + } else if (c.count < p.count()) { + if (p.parents) |parents| { + const tree = try accept.sanitizeTree(a, p.tokens, parents, c.count); + p.tokens = tree.tokens; + p.parents = tree.parents; + } else if (p.held > 0) p.held = c.count else p.tokens = p.tokens[0..c.count]; + } + if (p.count() == 0) p.kind = .none; + } +} + /// The head's chance for each held draft, else chain chances from the stream's per-depth acceptance. pub fn nodeChances(e: *Engine, s: *Stream, n: usize) ![]const f64 { const a = e.arena.allocator(); From 361f06340a9c9aafb529ee8b172aeb8eb84dde70 Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Sun, 11 Oct 2026 03:39:17 +0700 Subject: [PATCH 021/104] Lane host halt: a stop's drain deadline ends the rounds at the next boundary, and every request queued, filling or running fails with the reason. Opened.halt (optional) lets a server stop an engine between rounds instead of under a running one; LaneHost.halt is that hook for lane-host families. A request still filling its prompt a chunk a round fails too, with its lane freed, and a prompt pass the halt cancels mid-chunk reports the halt's reason rather than a cancel. --- zig/src/core/engine_api.zig | 9 +- zig/src/core/lane_halt.zig | 34 ++++++ zig/src/core/lane_halt_test.zig | 210 ++++++++++++++++++++++++++++++++ zig/src/core/lane_host.zig | 17 ++- 4 files changed, 266 insertions(+), 4 deletions(-) create mode 100644 zig/src/core/lane_halt.zig create mode 100644 zig/src/core/lane_halt_test.zig diff --git a/zig/src/core/engine_api.zig b/zig/src/core/engine_api.zig index de9b9da1af..2187c24416 100644 --- a/zig/src/core/engine_api.zig +++ b/zig/src/core/engine_api.zig @@ -216,7 +216,13 @@ pub const Open = struct { }; /// An opened engine; ``close`` stops its thread and frees its backend. -pub const Opened = struct { engine: Engine, close: *const fn (ctx: *anyopaque) void, ctx: *anyopaque }; +pub const Opened = struct { + engine: Engine, + close: *const fn (ctx: *anyopaque) void, + ctx: *anyopaque, + /// A drain's deadline: rounds end at the next boundary, what is left fails with `reason`; null: `close` alone. + halt: ?*const fn (ctx: *anyopaque, reason: []const u8) void = null, +}; pub const Memory = struct { active: u64 = 0, cache: u64 = 0, peak: u64 = 0, pinned: u64 = 0 }; @@ -405,6 +411,7 @@ test { _ = @import("lane_host.zig"); _ = @import("lane_host_test.zig"); _ = @import("lane_rounds_test.zig"); + _ = @import("lane_halt_test.zig"); _ = @import("lane_host_reuse_test.zig"); _ = @import("lane_host_warm_test.zig"); _ = @import("prompt_cache.zig"); diff --git a/zig/src/core/lane_halt.zig b/zig/src/core/lane_halt.zig new file mode 100644 index 0000000000..466de561d8 --- /dev/null +++ b/zig/src/core/lane_halt.zig @@ -0,0 +1,34 @@ +//! A stop's drain deadline on the lane host: the rounds end at a boundary and what is left fails with a reason. +const std = @import("std"); +const LaneHost = @import("lane_host.zig").LaneHost; +const Job = LaneHost.Job; + +/// `LaneHost.halt`: the engine thread stops at the next round boundary; what is left fails with `reason`. +pub fn halt(h: *LaneHost, reason: []const u8) void { + h.mutex.lockUncancelable(h.io); + h.closing = true; + h.halt_reason = reason; + h.wake.broadcast(h.io); + h.mutex.unlock(h.io); + if (h.thread) |t| t.join(); + h.thread = null; +} + +/// `halt`: between rounds, every request queued, filling or running fails with `reason` (its stream discarded). +pub fn haltAll(h: *LaneHost, reason: []const u8) void { + h.lock(); + const queued = h.gpa.dupe(*Job, h.queued.items) catch &.{}; + h.queued.clearRetainingCapacity(); + const jobs = h.gpa.dupe(*Job, h.admitted.items) catch &.{}; + h.admitted.clearRetainingCapacity(); + h.filling.clearRetainingCapacity(); // a job still filling is admitted: it fails below, its lane freed + h.unlock(); + if (queued.len + jobs.len > 0) std.log.warn("lane host: stopping at a round boundary: {d} running and {d} queued request(s) end ({s})", .{ jobs.len, queued.len, reason }); + for (jobs) |job| { + if (job.started and !job.stream.finished) h.core.discard(&job.stream); + h.finish(job, .failed, reason); + } + for (queued) |job| h.finish(job, .failed, reason); + h.gpa.free(queued); + h.gpa.free(jobs); +} diff --git a/zig/src/core/lane_halt_test.zig b/zig/src/core/lane_halt_test.zig new file mode 100644 index 0000000000..87ba6710ab --- /dev/null +++ b/zig/src/core/lane_halt_test.zig @@ -0,0 +1,210 @@ +//! The lane host's halt: a stop's drain deadline ends the rounds at a boundary, what is left fails with the reason. +const std = @import("std"); +const lanes = @import("lanes"); +const api = @import("engine_api.zig"); +const LaneHost = @import("lane_host.zig").LaneHost; +const Reason = api.Reason; +const Id = api.Id; +const Event = api.Event; +const Request = api.Request; + +test "halt (a stop's drain deadline): the rounds end at a round boundary, every request queued or running fails with the reason" { + const gpa = std.testing.allocator; + var cfg = try lanes.Config.init(gpa, .{ .exact_width = 8, .gpu_tokens = true, .hidden_rows = true }, 8, 7); + defer cfg.deinit(gpa); + var target: lanes.fake.Fake = .{ .gpa = gpa }; + defer target.deinit(); + var clock: lanes.fake.FixedClock = .{}; + var core = lanes.Engine.init(gpa, &cfg, target.pieceBackend(), clock.clock()); + defer core.deinit(); + var host = LaneHost.init(gpa, std.testing.io, &core, .{ .lanes = 2 }); + // one prompt prefilling at a time (the next waits queued), 3 rows a round + const Plan = struct { + const Seq = struct { key: usize, n: u64, done: u64 = 0, decoding: bool = false }; + seqs: std.ArrayList(Seq) = .empty, + rounds: u32 = 0, + fn self(ctx: *anyopaque) *@This() { + return @ptrCast(@alignCast(ctx)); + } + fn beginAdmit(_: *anyopaque) void {} + fn admit(ctx: *anyopaque, ask: api.Rounds.Ask, _: *[]const u8) api.Rounds.Verdict { + const p = self(ctx); + for (p.seqs.items) |x| if (!x.decoding) return .wait; + p.seqs.append(gpa, .{ .key = ask.key, .n = ask.prompt.len }) catch return .refuse; + return .admit; + } + fn admitted(_: *anyopaque) anyerror!void {} + fn arm(_: *anyopaque, _: usize) void {} + fn began(_: *anyopaque, _: usize, _: bool) void {} + fn plan(ctx: *anyopaque, pieces: *std.ArrayList(api.Rounds.Piece), finals: *std.ArrayList(usize)) anyerror!void { + const p = self(ctx); + pieces.clearRetainingCapacity(); + finals.clearRetainingCapacity(); + for (p.seqs.items) |*x| if (!x.decoding) { + if (x.done < x.n - 1) { + const end = @min(x.done + 3, x.n - 1); + try pieces.append(gpa, .{ .key = x.key, .start = x.done, .end = end, .save = false }); + x.done = end; + } + if (x.done == x.n - 1) { + try finals.append(gpa, x.key); + x.decoding = true; + } + }; + p.rounds += 1; + } + fn left(ctx: *anyopaque, key: usize) void { + const p = self(ctx); + for (p.seqs.items, 0..) |x, i| if (x.key == key) { + _ = p.seqs.orderedRemove(i); + return; + }; + } + fn after(_: *anyopaque, _: f64, _: f64) void {} + }; + var pl: Plan = .{}; + defer pl.seqs.deinit(gpa); + host.rounds = .{ .ctx = &pl, .vtable = &.{ .begin_admit = Plan.beginAdmit, .admit = Plan.admit, .admitted = Plan.admitted, .arm = Plan.arm, .began = Plan.began, .plan = Plan.plan, .left = Plan.left, .after = Plan.after } }; + const Box = struct { + mutex: std.Io.Mutex = .init, + tokens: std.ArrayList(u32) = .empty, + done: ?Reason = null, + message: []const u8 = "", + /// the planner's rounds when the finish arrived (on the engine thread, between rounds) + rounds_at_end: u32 = 0, + plan: *Plan, + fn event(ctx: *anyopaque, _: Id, e: *const Event) void { + const b: *@This() = @ptrCast(@alignCast(ctx)); + b.mutex.lockUncancelable(std.testing.io); + defer b.mutex.unlock(std.testing.io); + switch (e.*) { + .tokens => |t| b.tokens.appendSlice(gpa, t) catch {}, + .finished => |f| { + b.done = f.reason; + b.message = f.message; + b.rounds_at_end = b.plan.rounds; + }, + else => {}, + } + } + fn count(b: *@This()) usize { + b.mutex.lockUncancelable(std.testing.io); + defer b.mutex.unlock(std.testing.io); + return b.tokens.items.len; + } + }; + var long_prompt: [600]u32 = undefined; + for (&long_prompt, 0..) |*t, i| t.* = @intCast(1 + i % 13); + const p1 = [_]u32{ 3, 1, 4, 1, 5, 9, 2, 6, 5, 3 }; + var boxes: [3]Box = .{ .{ .plan = &pl }, .{ .plan = &pl }, .{ .plan = &pl } }; + defer for (&boxes) |*b| b.tokens.deinit(gpa); + // 1: decoding (no end of its own), 2: prefilling a long prompt in pieces, 3: queued behind 2 + const r1: Request = .{ .prompt = &p1, .max_tokens = 1 << 30 }; + const r2: Request = .{ .prompt = &long_prompt, .max_tokens = 1 << 30 }; + const r3: Request = .{ .prompt = &p1, .max_tokens = 8 }; + const e = host.engine(); + try host.start(); + try e.submit(1, &r1, .{ .ctx = &boxes[0], .event = Box.event }); + while (boxes[0].count() < 4) std.Io.sleep(std.testing.io, .fromMilliseconds(1), .awake) catch {}; + try e.submit(2, &r2, .{ .ctx = &boxes[1], .event = Box.event }); + try e.submit(3, &r3, .{ .ctx = &boxes[2], .event = Box.event }); + while (boxes[0].count() < 20) std.Io.sleep(std.testing.io, .fromMilliseconds(1), .awake) catch {}; + host.halt("server restarting: retry shortly"); + const rounds = pl.rounds; + for (&boxes) |*b| { + try std.testing.expectEqual(@as(?Reason, .failed), b.done); + try std.testing.expectEqualStrings("server restarting: retry shortly", b.message); + // no round was planned after the finishes: the loop stopped at the boundary they were sent at + try std.testing.expectEqual(rounds, b.rounds_at_end); + } + // the decoding stream's tokens are whole rounds' (the fake's own sequence), the prefilling one had none yet + var history: std.ArrayList(u32) = .empty; + defer history.deinit(gpa); + try history.appendSlice(gpa, &p1); + for (boxes[0].tokens.items) |t| { + try std.testing.expectEqual(lanes.fake.next(history.items, null, history.items.len), t); + try history.append(gpa, t); + } + try std.testing.expectEqual(@as(usize, 0), boxes[1].tokens.items.len); + try std.testing.expectEqual(@as(usize, 0), boxes[2].tokens.items.len); + try std.testing.expectEqual(@as(usize, 0), pl.seqs.items.len); + // nothing runs after it; the stop that follows (a close) finds the thread gone + std.Io.sleep(std.testing.io, .fromMilliseconds(20), .awake) catch {}; + try std.testing.expectEqual(rounds, pl.rounds); + host.stop(); +} + +test "halt fails a request still filling its prompt with the reason and frees its lane; a decoding one fails too" { + const gpa = std.testing.allocator; + var cfg = try lanes.Config.init(gpa, .{ .exact_width = 8, .gpu_tokens = true, .hidden_rows = true }, 8, 7); + defer cfg.deinit(gpa); + var target: lanes.fake.Fake = .{ .gpa = gpa, .prefill_chunks = 2 }; + defer target.deinit(); + var clock: lanes.fake.FixedClock = .{}; + var core = lanes.Engine.init(gpa, &cfg, target.stepped(), clock.clock()); + defer core.deinit(); + var host = LaneHost.init(gpa, std.testing.io, &core, .{ .lanes = 2 }); + const Chunks = struct { + n: std.atomic.Value(usize) = .init(0), // chunks the second request's pass ran + fn call(ctx: *anyopaque, _: *lanes.Stream, _: usize) void { + const c: *@This() = @ptrCast(@alignCast(ctx)); + _ = c.n.fetchAdd(1, .monotonic); + } + }; + var chunks: Chunks = .{}; + const Box = struct { + mutex: std.Io.Mutex = .init, + tokens: std.ArrayList(u32) = .empty, + done: ?Reason = null, + message: []const u8 = "", + fn event(ctx: *anyopaque, _: Id, e: *const Event) void { + const b: *@This() = @ptrCast(@alignCast(ctx)); + b.mutex.lockUncancelable(std.testing.io); + defer b.mutex.unlock(std.testing.io); + switch (e.*) { + .tokens => |t| b.tokens.appendSlice(gpa, t) catch {}, + .finished => |f| { + b.done = f.reason; + b.message = f.message; + }, + else => {}, + } + } + fn count(b: *@This()) usize { + b.mutex.lockUncancelable(std.testing.io); + defer b.mutex.unlock(std.testing.io); + return b.tokens.items.len; + } + }; + const p1 = [_]u32{ 3, 1, 4, 1, 5, 9, 2, 6 }; + const p2 = [_]u32{ 2, 7, 1, 8, 2, 8 }; + var boxes: [3]Box = .{ .{}, .{}, .{} }; + defer for (&boxes) |*b| b.tokens.deinit(gpa); + // 1: decoding (no end of its own), 2: filling a prompt that never ends, 3: queued behind both lanes + const r1: Request = .{ .prompt = &p1, .max_tokens = 1 << 30 }; + const r2: Request = .{ .prompt = &p2, .max_tokens = 8 }; + const r3: Request = .{ .prompt = &p2, .max_tokens = 8 }; + const e = host.engine(); + try host.start(); + try e.submit(1, &r1, .{ .ctx = &boxes[0], .event = Box.event }); + while (boxes[0].count() < 4) std.Io.sleep(std.testing.io, .fromMilliseconds(1), .awake) catch {}; + // nothing fills now: the next pass takes chunks until the halt + target.prefill_chunks = std.math.maxInt(usize); + target.prefill_hook = Chunks.call; + target.prefill_hook_ctx = &chunks; + try e.submit(2, &r2, .{ .ctx = &boxes[1], .event = Box.event }); + try e.submit(3, &r3, .{ .ctx = &boxes[2], .event = Box.event }); + while (chunks.n.load(.monotonic) < 3) std.Io.sleep(std.testing.io, .fromMilliseconds(1), .awake) catch {}; + host.halt("server restarting: retry shortly"); + for (&boxes) |*b| { + try std.testing.expectEqual(@as(?Reason, .failed), b.done); + try std.testing.expectEqualStrings("server restarting: retry shortly", b.message); + } + try std.testing.expect(boxes[0].tokens.items.len >= 4); + try std.testing.expectEqual(@as(usize, 0), boxes[1].tokens.items.len); + try std.testing.expectEqual(@as(usize, 0), boxes[2].tokens.items.len); + // every lane freed: the decoding stream's and the filling one's + try std.testing.expectEqual(@as(usize, 0), target.lanes.count()); + try std.testing.expectEqual(@as(usize, 0), host.filling.items.len); + host.stop(); +} diff --git a/zig/src/core/lane_host.zig b/zig/src/core/lane_host.zig index 585f5cea3f..248d4230f2 100644 --- a/zig/src/core/lane_host.zig +++ b/zig/src/core/lane_host.zig @@ -21,6 +21,7 @@ pub const Rounds = rounds.Rounds; pub const Proposers = rounds.Proposers; const rounds = @import("lane_rounds.zig"); const lane_stream = @import("lane_stream.zig"); +const lane_halt = @import("lane_halt.zig"); pub const LaneHost = struct { gpa: Allocator, @@ -37,6 +38,8 @@ pub const LaneHost = struct { filling: std.ArrayList(*Job) = .empty, // admitted jobs whose prompt is still going in, a chunk a round, oldest first cancels: std.ArrayList(Id) = .empty, closing: bool = false, + /// `halt`'s reason: what is left finishes failed with it (null: a stop cancels what is left) + halt_reason: ?[]const u8 = null, thread: ?std.Thread = null, decoded: std.ArrayList(status.Mark) = .empty, // tokens a round landed, for the 2 s decode rate prefill_rate: f64 = 0, @@ -114,7 +117,7 @@ pub const LaneHost = struct { h.thread = try std.Thread.spawn(.{ .stack_size = 16 << 20 }, run, .{h}); } - /// Stops admitting, cancels what is left and joins the engine thread. + /// Stops admitting, cancels what is left and joins the engine thread (after `halt`: the thread is gone already). pub fn stop(h: *LaneHost) void { h.mutex.lockUncancelable(h.io); h.closing = true; @@ -135,6 +138,9 @@ pub const LaneHost = struct { h.finals.deinit(h.gpa); } + /// A drain's deadline: the round in flight completes on every rank, so the ranks stop at the same boundary. + pub const halt = lane_halt.halt; + pub fn engine(h: *LaneHost) Engine { return .{ .ctx = h, .vtable = &.{ .info = infoFn, .submit = submitFn, .cancel = cancelFn, .status = status.read, .memory = memoryFn, .keepalive = keepaliveFn, .learn = learn.submit } }; } @@ -486,10 +492,13 @@ pub const LaneHost = struct { return false; } - /// A job cancelled in its prompt pass, its lane already released. + /// A job cancelled in its prompt pass, its lane already released (by a `halt`: failed with its reason). pub fn cancel(h: *LaneHost, job: *Job) bool { h.remove(job); - h.finish(job, .cancelled, ""); + h.lock(); + const why = h.halt_reason; + h.unlock(); + h.finish(job, if (why == null) .cancelled else .failed, why orelse ""); return true; } @@ -572,6 +581,7 @@ pub const LaneHost = struct { } pub fn closeAll(h: *LaneHost) void { + if (h.halt_reason) |reason| return lane_halt.haltAll(h, reason); h.lock(); for (h.queued.items) |job| h.cancels.append(h.gpa, job.id) catch {}; for (h.admitted.items) |job| h.cancels.append(h.gpa, job.id) catch {}; @@ -583,4 +593,5 @@ pub const LaneHost = struct { test { _ = @import("lane_host_test.zig"); _ = @import("lane_rounds_test.zig"); + _ = @import("lane_halt_test.zig"); } From e9f9210be2ed986edc61cbe57e5d431570215ce4 Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Sun, 11 Oct 2026 06:52:10 +0700 Subject: [PATCH 022/104] Round planner test: both prompts queue before the host starts, and its waits time out The second prompt was submitted after start(), so a loaded scheduler could let the first finish its pieces before the second arrived: no queued admission, and the test's waits > 0 check failed. Both are queued now before the first round. Each wait gives up after 10 s, and a timed-out wait stops the host thread before the test's frame is gone. --- zig/src/core/lane_rounds_test.zig | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/zig/src/core/lane_rounds_test.zig b/zig/src/core/lane_rounds_test.zig index 029f6d0e3e..f14f685302 100644 --- a/zig/src/core/lane_rounds_test.zig +++ b/zig/src/core/lane_rounds_test.zig @@ -91,14 +91,15 @@ test "a round planner: prompts in pieces between the other streams' decode round else => {}, } } - fn wait(b: *@This()) Reason { - while (true) { + fn wait(b: *@This()) !Reason { + for (0..10_000) |_| { // 10 s b.mutex.lockUncancelable(std.testing.io); const d = b.done; b.mutex.unlock(std.testing.io); if (d) |r| return r; std.Io.sleep(std.testing.io, .fromMilliseconds(1), .awake) catch {}; } + return error.Timeout; } }; const p1 = [_]u32{ 3, 1, 4, 1, 5, 9, 2, 6, 5, 3 }; @@ -110,11 +111,15 @@ test "a round planner: prompts in pieces between the other streams' decode round const r1: Request = .{ .prompt = &p1, .max_tokens = 40 }; const r2: Request = .{ .prompt = &p2, .max_tokens = 12 }; const e = host.engine(); - try host.start(); + // both queued before the first round, so the second waits whatever the scheduler does try e.submit(1, &r1, .{ .ctx = &b1, .event = Box.event }); try e.submit(2, &r2, .{ .ctx = &b2, .event = Box.event }); - try std.testing.expectEqual(Reason.length, b1.wait()); - try std.testing.expectEqual(Reason.length, b2.wait()); + try host.start(); + { + errdefer host.stop(); // a timed-out wait stops the host thread before this frame goes + try std.testing.expectEqual(Reason.length, try b1.wait()); + try std.testing.expectEqual(Reason.length, try b2.wait()); + } host.stop(); // the replies are the whole-prompt path's tokens for ([_]struct { p: []const u32, b: *Box }{ .{ .p = &p1, .b = &b1 }, .{ .p = &p2, .b = &b2 } }) |c| { From 96ae8658371704822c70bbc05d50c51e35b7e378 Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Sat, 10 Oct 2026 13:57:31 +0700 Subject: [PATCH 023/104] Server hooks for a model family's own chat rules. family.zig is an optional vtable a family installs in place of the chat template: prompt rendering, thinking and effort fields, a streaming reply reader that parses its own call markup, and the turn it remembers. Without a family the server runs as before. chat.zig's reply collection moves to chat_generation.zig and openai.zig's run to openai_run.zig (pure moves), so neither file passes 600 lines with the hooks in. --- zig/src/server/anthropic.zig | 3 +- zig/src/server/chat.zig | 349 +++-------------------- zig/src/server/chat_generation.zig | 436 +++++++++++++++++++++++++++++ zig/src/server/family.zig | 135 +++++++++ zig/src/server/openai.zig | 325 ++------------------- zig/src/server/openai_run.zig | 328 ++++++++++++++++++++++ zig/src/server/prompt.zig | 19 +- zig/src/server/root.zig | 1 + zig/src/server/serve.zig | 16 ++ zig/src/server/server.zig | 13 +- zig/src/server/tokens.zig | 22 +- 11 files changed, 1036 insertions(+), 611 deletions(-) create mode 100644 zig/src/server/chat_generation.zig create mode 100644 zig/src/server/family.zig create mode 100644 zig/src/server/openai_run.zig diff --git a/zig/src/server/anthropic.zig b/zig/src/server/anthropic.zig index 4821e0f3eb..79113e569c 100644 --- a/zig/src/server/anthropic.zig +++ b/zig/src/server/anthropic.zig @@ -202,7 +202,8 @@ const Reply = struct { fn text(r: *Reply, kind: []const u8, value: []const u8) E!void { if (value.len == 0) return; if (r.active != null) { - r.pending.append(r.a, .{ .kind = kind, .text = value }) catch return error.Closed; + // held until the open tool_use block closes: copied, ``value`` is a slice of a text that grows by then + r.pending.append(r.a, .{ .kind = kind, .text = r.a.dupe(u8, value) catch return error.Closed }) catch return error.Closed; return; } const key = if (std.mem.eql(u8, kind, "thinking")) "thinking" else "text"; diff --git a/zig/src/server/chat.zig b/zig/src/server/chat.zig index 6315e32991..b039cff1f9 100644 --- a/zig/src/server/chat.zig +++ b/zig/src/server/chat.zig @@ -12,7 +12,11 @@ const log = @import("log.zig"); const ids = @import("ids.zig"); const clock = @import("clock.zig"); const Server = @import("server.zig").Server; +const family_mod = @import("family.zig"); const chunk_plan = @import("chunk_plan.zig"); +const chat_generation = @import("chat_generation.zig"); +const Shown = chat_generation.Shown; +const Generation = chat_generation.Generation; const Value = json.Value; const Cx = errors.Cx; const Allocator = std.mem.Allocator; @@ -32,6 +36,10 @@ pub const Input = struct { id: []const u8 = "", /// ``logprobs: true`` with this many ``top_logprobs``; null: not asked. logprobs: ?u8 = null, + /// The request body (a family reads fields of its own from it, such as ``response_format``). + body: Value = .null, + /// ``parallel_tool_calls: false``: the reply keeps its first call (a family applies it while streaming). + single_call: bool = false, }; /// A streamed piece: content text (a string) or a delta object (reasoning or tool calls). @@ -57,12 +65,14 @@ pub const Reply = struct { effort: ?[]const u8 = null, /// ``choices[0].logprobs`` (``{"content": [...]}``) for a request that asked; null otherwise. logprobs: ?Value = null, + /// A family's parsed calls (OpenAI tool-call objects); null: the route parses the content itself. + calls: ?[]Value = null, }; pub const Failure = error{ Refused, Cancelled, Failed, OutOfMemory }; /// Events the engine thread hands this request, read by the request's own thread. -const Mailbox = struct { +pub const Mailbox = struct { io: std.Io, gpa: Allocator, mutex: std.Io.Mutex = .init, @@ -140,6 +150,9 @@ pub const Prepared = struct { drafts: bool, stops: fields_mod.Stops, preparing: bool, + images_omitted: u32 = 0, + /// the prompt's prepared images (a family's), held until the reply ends + images: ?family_mod.Images = null, }; /// Render ``input`` and run every check that can refuse it, before anything reaches the client or the engine. @@ -159,8 +172,10 @@ pub fn prepare(srv: *Server, cx: *Cx, input: Input, gone: anytype) Failure!Prepa if (gone.check()) return error.Cancelled; var thinking = flag(f, "enable_thinking") orelse srv.config.enable_thinking; if (input.prompt != null) thinking = false; - const effort = srv.effortFor(if (f.get("reasoning_effort")) |e| (if (e == .string) e.string else null) else null); + const asked: ?[]const u8 = if (f.get("reasoning_effort")) |e| (if (e == .string) e.string else null) else null; + const effort: ?[]const u8 = if (srv.family) |fam| asked orelse srv.config.reasoning_effort orelse fam.defaultEffort() else srv.effortFor(asked); const rendered = try prompt_mod.prepare(srv, cx, input, thinking, effort); + errdefer if (rendered.images) |im| im.release(); if (gone.check()) return error.Cancelled; if (rendered.ids.len == 0) return cx.refuse("rendered prompt is empty"); const window: i64 = srv.info.context_window; @@ -206,7 +221,7 @@ pub fn prepare(srv: *Server, cx: *Cx, input: Input, gone: anytype) Failure!Prepa } } if (input.logprobs) |count| request.logprobs = try @import("logprobs.zig").admit(srv, cx, thinking, request.think_budget, count); - return .{ .input = input, .request = request, .prompt_len = rendered.ids.len, .received = received, .thinking = thinking, .effort = effort, .sampled = sampling != null, .drafts = drafts, .stops = stops_opt, .preparing = preparing }; + return .{ .input = input, .request = request, .prompt_len = rendered.ids.len, .received = received, .thinking = thinking, .effort = effort, .sampled = sampling != null, .drafts = drafts, .stops = stops_opt, .preparing = preparing, .images_omitted = rendered.images_omitted, .images = rendered.images }; } /// Submit a prepared request and collect its reply; ``sink`` hears the stream (null: not streamed). @@ -221,6 +236,7 @@ pub fn generate(srv: *Server, cx: *Cx, prepared: Prepared, sink: ?Sink, gone: an const stops_opt = prepared.stops; var preparing = prepared.preparing; defer release(srv, preparing); + defer if (prepared.images) |im| im.release(); // after the engine's finished event (the loop below waits for it) if (request.background) { while (true) { const now = nowNs(io); @@ -232,16 +248,19 @@ pub fn generate(srv: *Server, cx: *Cx, prepared: Prepared, sink: ?Sink, gone: an } if (gone.check()) return error.Cancelled; var stop_hook: StopHook = .{ .srv = srv, .stops = .{ .strings = stops_opt.strings } }; - if (stops_opt.strings.len > 0) request.stop = .{ .ctx = &stop_hook, .check = StopHook.check }; + // a family matches stop strings in the answer (``push``); the generic path in the raw text on the engine's thread + if (stops_opt.strings.len > 0 and srv.family == null) request.stop = .{ .ctx = &stop_hook, .check = StopHook.check }; var box: Mailbox = .{ .io = io, .gpa = srv.gpa }; defer box.deinit(); const id = srv.next_id.fetchAdd(1, .monotonic); if (srv.keepalive) |k| k.begin(); // the GPU is busy: the idle ticker holds its commits defer if (srv.keepalive) |k| k.end(); - srv.engine.submit(id, &request, .{ .ctx = &box, .event = Mailbox.onEvent }) catch |e| return switch (e) { - error.Busy => cx.fail(.capacity, "the engine is busy; retry shortly", .{}), - error.Closed => cx.fail(.other, "the scheduler is closed", .{}), - error.InvalidSpans => cx.fail(.request, "the prompt span layout is invalid", .{}), + srv.engine.submit(id, &request, .{ .ctx = &box, .event = Mailbox.onEvent }) catch |e| { + return switch (e) { + error.Busy => cx.fail(.capacity, "the engine is busy; retry shortly", .{}), + error.Closed => cx.fail(.other, "the scheduler is closed", .{}), + error.InvalidSpans => cx.fail(.request, "the prompt span layout is invalid", .{}), + }; }; release(srv, preparing); // a background request waits only while a foreground one prepares preparing = false; @@ -249,9 +268,11 @@ pub fn generate(srv: *Server, cx: *Cx, prepared: Prepared, sink: ?Sink, gone: an defer srv.noteRequest(prepared.prompt_len, box.cached orelse 0, gen.collected.items.len, box.stats.drafted, box.stats.accepted, box.stats.rounds, received, gen.first_ns, gen.last_ns, box.stats.prefill_seconds); errdefer if (!gen.engine_done) gen.cancel(); // the engine writes to the mailbox until it says finished const result: Failure!Reply = blk: { - if (sink != null and input.tools.len > 0) gen.calls = tool_stream.Streamer.init(a, input.tools) catch |e| break :blk e; + if (srv.family) |fam| { + gen.reader = fam.reader(a, thinking, input.tools, input.single_call) catch |e| break :blk e; + } else if (sink != null and input.tools.len > 0) gen.calls = tool_stream.Streamer.init(a, input.tools) catch |e| break :blk e; gen.loop(gone) catch |e| break :blk e; - break :blk gen.finish(cx, prepared.prompt_len, received, thinking, effort, prepared.sampled, prepared.drafts); + break :blk gen.finish(cx, prepared.prompt_len, received, thinking, effort, prepared.sampled, prepared.drafts, prepared.images_omitted); }; if (result) |_| {} else |e| logEnded(input.id, e, cx.message, prepared.prompt_len, gen.collected.items.len, seconds(nowNs(io) - received)); return result; @@ -279,298 +300,6 @@ pub fn release(srv: *Server, preparing: bool) void { if (preparing) _ = srv.preparing.fetchSub(1, .acq_rel); } -/// The text a stream has sent: Python's ``streamed = visible``, grown by its delta while it only grows. -const Shown = struct { - text: std.ArrayList(u8) = .empty, - chars: usize = 0, - at: ?[*]const u8 = null, // where the last shown text lay; the decode buffer only grows, so that prefix holds - extends: bool = false, - - /// Python's ``now[len(shown):]``; ``stable``: ``now`` lies in the append-only decode buffer. - fn after(s: *Shown, now: []const u8, stable: bool) []const u8 { - const sent = s.text.items; - s.extends = (stable and s.at == now.ptr and now.len >= sent.len) or std.mem.startsWith(u8, now, sent); - if (!s.extends) return reply_text.afterChars(now, s.chars); - if (stable) s.at = now.ptr; // the buffer moved but kept its bytes: the next check is free again - return now[sent.len..]; - } - - fn set(s: *Shown, a: Allocator, now: []const u8, delta: []const u8) Allocator.Error!void { - if (s.extends) { - try s.text.appendSlice(a, delta); - s.chars += reply_text.charCount(delta); - } else { - s.text = .empty; - try s.text.appendSlice(a, now); - s.chars = reply_text.charCount(now); - } - s.at = now.ptr; - } -}; - -/// The token loop and the text it streams. -const Generation = struct { - srv: *Server, - a: Allocator, - box: *Mailbox, - id: api.Id, - reply_id: []const u8, // the id the client gets, which the done line prints - sink: ?Sink, - thinking: bool, - stops: reply_text.Stops, - ignore_eos: bool, - max_tokens: u32, - tools: []const Value, - calls: ?tool_stream.Streamer = null, - collected: std.ArrayList(u32) = .empty, - visible: reply_text.Incremental = .{}, - hidden: std.ArrayList(u8) = .empty, // reused for the answer without its call blocks - streamed: Shown = .{}, // what content streamed, kept apart from the decode buffer that grows under slices - streamed_reasoning: Shown = .{}, - streaming_done: bool = false, - first_ns: ?i96 = null, - last_ns: ?i96 = null, - reason: ?[]const u8 = null, // the server ended the reply (a stop string, the length) before the engine said so - engine_done: bool = false, - consumed: usize = 0, - chunk_index: usize = 0, - logprobs: bool = false, // the request asked for logprobs: the mailbox holds a row a token - - fn eos(g: *const Generation, t: u32) bool { - return !g.ignore_eos and std.mem.indexOfScalar(u32, g.srv.eos, t) != null; - } - - /// What closes a call the model's end token left open (``tool_parse.closeCall``); nothing when a stop string, the length or ``ignore_eos`` ended the reply instead. - fn closeCall(g: *const Generation, text: []const u8) Allocator.Error![]const u8 { - const t = g.collected.items; - if (g.tools.len == 0 or t.len == 0 or !g.eos(t[t.len - 1])) return ""; - return tool_parse.closeCall(g.a, text, g.tools); - } - - /// The engine's stop check, here: the newest tokens' text holds a stop string. - fn stopHit(g: *Generation) Allocator.Error!bool { - if (g.stops.strings.len == 0) return false; - const t = g.collected.items; - const tail = t[t.len -| g.stops.tail()..]; - const text = try g.srv.text.decode(g.a, tail); - for (g.stops.strings) |s| if (std.mem.indexOf(u8, text, s) != null) return true; - return false; - } - - /// Commits a round's tokens as ``LaneStream.commit`` would, then streams what they add. - fn commit(g: *Generation, chunk: []const u32) Failure!void { - var landed: usize = 0; - for (chunk) |t| { - if (g.reason != null) break; - try g.collected.append(g.a, t); - landed += 1; - if (g.eos(t)) { - g.reason = "stop"; - } else if (try g.stopHit()) { - g.reason = "stop"; - g.srv.engine.cancel(g.id); // the engine ends EOS and length itself; a stop string is the server's - } else if (g.collected.items.len >= g.max_tokens) g.reason = "length"; - } - if (landed == 0) return; - const arrived = nowNs(g.srv.io); - if (g.first_ns == null) g.first_ns = arrived; - g.last_ns = arrived; - const sink = g.sink orelse return; - if (g.streaming_done) return; - var fresh: std.ArrayList(u32) = .empty; - for (chunk[0..landed]) |t| { - if (g.eos(t)) { - g.streaming_done = true; - break; - } - try fresh.append(g.a, t); - } - const all = try g.visible.extend(g.a, g.srv.text, decodeText, fresh.items); - const text = g.stops.visible(all, true); - if (g.visible.pending()) return; // a character still split across tokens - var answer = text; - if (g.thinking) { - const split = try reply_text.splitThinking(g.a, text, false, g.srv.markers); - const piece = g.streamed_reasoning.after(split.reasoning, true); - if (piece.len > 0) { - try g.streamed_reasoning.set(g.a, split.reasoning, piece); - try g.emit(sink, try deltaOf(g.a, "reasoning_content", piece)); - } - answer = split.answer; - } - const shown = if (g.calls != null) try reply_text.hideInto(g.a, &g.hidden, answer, false) else answer; - const vis = reply_text.heldBack(reply_text.streamingVisible(shown), g.calls != null); - const copied = @intFromPtr(vis.ptr) >= @intFromPtr(g.hidden.items.ptr) and @intFromPtr(vis.ptr) <= @intFromPtr(g.hidden.items.ptr) + g.hidden.items.len; - const delta = g.streamed.after(vis, !copied); - if (delta.len > 0) { - try g.streamed.set(g.a, vis, delta); - try g.emit(sink, .{ .string = delta }); - } - if (g.calls) |*c| { - var out: std.ArrayList(Value) = .empty; - try c.feed(answer, &out); // never the reasoning: a call it mentions is not made - for (out.items) |d| try g.emit(sink, d); - } - } - - fn emit(g: *Generation, sink: Sink, delta: Value) Failure!void { - sink.call(sink.ctx, delta) catch { - g.cancel(); - return error.Cancelled; - }; - } - - /// Ends the request and waits for the engine; one it ended unfinished counts as a disconnect, as the Mac scheduler counts it. - fn cancel(g: *Generation) void { - if (!g.engine_done) g.srv.engine.cancel(g.id); - g.drain(); - if (g.box.reason == .cancelled and g.reason == null) g.srv.metrics.disconnected(g.srv.io); - } - - /// Waits for the engine's own end, so the request it holds may be freed. - fn drain(g: *Generation) void { - const m = g.box; - m.mutex.lockUncancelable(m.io); - defer m.mutex.unlock(m.io); - while (!m.finished) m.cond.waitUncancelable(m.io, &m.mutex); - g.engine_done = true; - } - - /// Takes rounds until the reply ends; a client that leaves cancels it. - fn loop(g: *Generation, gone: anytype) Failure!void { - const m = g.box; - while (true) { - m.mutex.lockUncancelable(m.io); - var chunk: ?[]u32 = null; - var ended = false; - if (g.chunk_index < m.chunks.items.len) { - const end = m.chunks.items[g.chunk_index]; - g.chunk_index += 1; - chunk = g.a.dupe(u32, m.tokens.items[g.consumed..end]) catch null; - g.consumed = end; - } else if (m.finished) { - ended = true; - } else { - m.cond.waitTimeout(m.io, &m.mutex, .{ .duration = .{ .raw = .fromMilliseconds(50), .clock = .awake } }) catch {}; - } - m.mutex.unlock(m.io); - if (ended) { - g.engine_done = true; - return; - } - if (gone.check()) { - g.cancel(); - return error.Cancelled; - } - const tokens = chunk orelse continue; - if (g.reason == null) try g.commit(tokens); - } - } - - fn finish(g: *Generation, cx: *Cx, prompt_len: usize, received: i96, thinking: bool, effort: ?[]const u8, exact: bool, drafts: bool) Failure!Reply { - const a = g.a; - const m = g.box; - if (!g.engine_done) g.drain(); - const reason: []const u8 = g.reason orelse switch (m.reason) { - .stop => "stop", - .length => "length", - .cancelled => "cancelled", - .failed => return cx.other(if (m.message.len > 0) try a.dupe(u8, m.message) else "the reply failed"), - }; - const finished_ns = nowNs(g.srv.io); - const content_tokens = if (g.ignore_eos) g.collected.items else reply_text.stripTrailing(g.collected.items, g.srv.eos); - const raw = try g.srv.text.decode(a, content_tokens); - const visible = g.stops.visible(raw, false); - // closed before the think split, so a call ending an unclosed think block is the answer - const close = try g.closeCall(visible); - const text = if (close.len > 0) try std.mem.concat(a, u8, &.{ visible, close }) else visible; - var content: []const u8 = text; - var reasoning: ?[]const u8 = null; - if (g.thinking) { - const split = try reply_text.splitThinking(a, text, true, g.srv.markers); - const r = reply_text.pyStrip(split.reasoning); - reasoning = if (r.len > 0) r else null; - content = split.answer; - } else { - const h = reply_text.parseHarmony(text); - content = h.content; - reasoning = h.reasoning; - } - if (g.sink) |sink| { - const thought = g.streamed_reasoning.text.items; - if (reasoning) |r| if (std.mem.startsWith(u8, r, thought) and r.len > thought.len) - try g.emit(sink, try deltaOf(a, "reasoning_content", r[thought.len..])); - const shown = if (g.calls != null) try reply_text.finishedProse(a, content) else content; - const sent = g.streamed.text.items; - if (std.mem.startsWith(u8, shown, sent) and shown.len > sent.len) try g.emit(sink, .{ .string = shown[sent.len..] }); - if (g.calls) |*c| if (close.len > 0) { - // the streamer reads the closers as markup the model wrote, so the call ends with the same deltas - var out: std.ArrayList(Value) = .empty; - try c.feed(if (g.thinking) content else text, &out); - for (out.items) |d| try g.emit(sink, d); - }; - } - const prefilled = m.prefilled_ns; - const decode_s: f64 = if (prefilled) |p| @max(0, seconds(finished_ns - p)) else 0; - const decode_tokens = g.collected.items.len -| 1; - const runtime = try json.newObject(a); - try runtime.put(a, "enable_thinking", .{ .bool = thinking }); - try runtime.put(a, "reasoning_effort", if (!thinking) .{ .string = "none" } else if (effort) |e| .{ .string = e } else .null); - try runtime.put(a, "engine", .{ .string = g.srv.info.name }); - try runtime.put(a, "tokens_per_second", .{ .float = if (decode_s > 0) @as(f64, @floatFromInt(decode_tokens)) / decode_s else 0 }); - try runtime.put(a, "seconds", .{ .float = @max(0, decode_s) }); // the reply's decode seconds, as tokens_per_second divides them - try runtime.put(a, "prefill_seconds", if (m.stats.prefill_seconds) |p| .{ .float = p } else .null); - const widths = try a.alloc(Value, m.widths.len); - for (m.widths, widths) |w, *slot| slot.* = try json.intValue(a, w); - try runtime.put(a, "prefill_widths", .{ .array = widths }); - const raised = try a.alloc(Value, m.raised.len); - for (m.raised, raised) |r, *slot| slot.* = .{ .bool = r }; - try runtime.put(a, "prefill_raised", .{ .array = raised }); - try runtime.put(a, "time_to_first_token", if (g.first_ns) |t| .{ .float = seconds(t - received) } else .null); - try runtime.put(a, "sampling", .{ .string = if (exact) "exact" else "greedy" }); - try runtime.put(a, "drafts", .{ .bool = drafts }); - const sha = @import("tokens.zig").tokenSha(g.collected.items); - try runtime.put(a, "token_sha", .{ .string = try a.dupe(u8, &sha) }); - try runtime.put(a, "min_rows", try json.intValue(a, m.stats.min_rows)); - if (m.stats.loop_period) |period| { - const loop_field = try json.newObject(a); - try loop_field.put(a, "period", try json.intValue(a, period)); - try runtime.put(a, "loop", .{ .object = loop_field }); - } - const s = m.stats; - const spec = try json.newObject(a); - try spec.put(a, "rounds", try json.intValue(a, s.rounds)); - try spec.put(a, "drafted", try json.intValue(a, s.drafted)); - try spec.put(a, "accepted", try json.intValue(a, s.accepted)); - try spec.put(a, "acceptance_rate", .{ .float = if (s.drafted > 0) @as(f64, @floatFromInt(s.accepted)) / @as(f64, @floatFromInt(s.drafted)) else 0 }); - try spec.put(a, "tokens_per_round", .{ .float = if (s.rounds > 0) @as(f64, @floatFromInt(g.collected.items.len)) / @as(f64, @floatFromInt(s.rounds)) else 0 }); - if (m.telemetry.len > 0) if ((try json.parse(a, m.telemetry)) == .ok) try spec.put(a, "proposer", (try json.parse(a, m.telemetry)).ok); - const think_end: ?u32 = if (thinking) g.srv.text.tokenId(g.srv.markers.close) else null; - const reply: Reply = .{ - .content = content, - .stop_sequence = g.stops.matched(raw), - .reasoning = reasoning, - .tool_calls_streamed = g.calls != null and g.calls.?.streamed, - .finish_reason = reason, - .prompt_tokens = prompt_len, - .cached_tokens = m.cached orelse 0, - .completion_tokens = g.collected.items.len, - .reasoning_tokens = reply_text.reasoningCount(g.collected.items, think_end), - .runtime = .{ .object = runtime }, - .speculative = .{ .object = spec }, - .thinking = thinking, - .effort = effort, - .logprobs = if (g.logprobs) try @import("logprobs.zig").value(g.srv, cx, a, g.thinking, content_tokens, m.rows.items) else null, - }; - if (std.mem.eql(u8, reason, "length") and thinking and reply_text.pyStrip(content).len == 0) - log.line("warning: a reply reached max_tokens while still thinking, so its content is empty and its text is all in reasoning_content; raise max_tokens, or send chat_template_kwargs {{\"enable_thinking\": false}} (server: --no-thinking)", .{}); - var cycle_text: [32]u8 = undefined; - const cycle = if (s.loop_period) |period| std.fmt.bufPrint(&cycle_text, " loop=period:{d}", .{period}) catch "" else ""; - log.line("done {s} prompt={d} cached={d} thinking={s} effort={s} tokens={d} sha={s} finish={s}{s} rounds={d} accepted={d}/{d}", .{ g.reply_id, prompt_len, reply.cached_tokens, if (thinking) "True" else "False", if (thinking) effort orelse "none" else "none", g.collected.items.len, sha, reason, cycle, s.rounds, s.accepted, s.drafted }); - return reply; - } -}; - /// The request's stop strings as the engine checks them after each token: the newest tokens' text holds one. const StopHook = struct { srv: *Server, @@ -587,7 +316,21 @@ const StopHook = struct { } }; -fn decodeText(t: anytype, a: Allocator, tokens: []const u32) ![]u8 { +/// What closes a call the model's end token left open (``tool_parse.closeCall``); nothing when a stop string, the length or ``ignore_eos`` ended the reply instead. +pub fn closeCall(g: *const Generation, text: []const u8) Allocator.Error![]const u8 { + const t = g.collected.items; + if (g.tools.len == 0 or t.len == 0 or !g.eos(t[t.len - 1])) return ""; + return tool_parse.closeCall(g.a, text, g.tools); +} + +/// Ends the request and waits for the engine; one it ended unfinished counts as a disconnect, as the Mac scheduler counts it. +pub fn cancel(g: *Generation) void { + if (!g.engine_done) g.srv.engine.cancel(g.id); + g.drain(); + if (g.box.reason == .cancelled and g.reason == null) g.srv.metrics.disconnected(g.srv.io); +} + +pub fn decodeText(t: anytype, a: Allocator, tokens: []const u32) ![]u8 { return t.decode(a, tokens); } diff --git a/zig/src/server/chat_generation.zig b/zig/src/server/chat_generation.zig new file mode 100644 index 0000000000..635fc7ac97 --- /dev/null +++ b/zig/src/server/chat_generation.zig @@ -0,0 +1,436 @@ +//! A submitted reply's collection: the engine's events in, streamed deltas and the finished ``Reply`` out. +const std = @import("std"); +const api = @import("engine_api"); +const json = @import("json.zig"); +const errors = @import("errors.zig"); +const reply_text = @import("reply_text.zig"); +const tool_stream = @import("tool_stream.zig"); +const tool_parse = @import("tool_parse.zig"); +const log = @import("log.zig"); +const clock = @import("clock.zig"); +const Server = @import("server.zig").Server; +const family_mod = @import("family.zig"); +const chunk_plan = @import("chunk_plan.zig"); +const Value = json.Value; +const Cx = errors.Cx; +const Allocator = std.mem.Allocator; +const chat = @import("chat.zig"); +const Sink = chat.Sink; +const Reply = chat.Reply; +const Failure = chat.Failure; +const Mailbox = chat.Mailbox; +const decodeText = chat.decodeText; +const deltaOf = chat.deltaOf; +const nowNs = clock.nowNs; +const seconds = clock.seconds; + +/// The text a stream has sent: Python's ``streamed = visible``, grown by its delta while it only grows. +pub const Shown = struct { + text: std.ArrayList(u8) = .empty, + chars: usize = 0, + at: ?[*]const u8 = null, // where the last shown text lay; the decode buffer only grows, so that prefix holds + extends: bool = false, + + /// Python's ``now[len(shown):]``; ``stable``: ``now`` lies in the append-only decode buffer. + fn after(s: *Shown, now: []const u8, stable: bool) []const u8 { + const sent = s.text.items; + s.extends = (stable and s.at == now.ptr and now.len >= sent.len) or std.mem.startsWith(u8, now, sent); + if (!s.extends) return reply_text.afterChars(now, s.chars); + if (stable) s.at = now.ptr; // the buffer moved but kept its bytes: the next check is free again + return now[sent.len..]; + } + + fn set(s: *Shown, a: Allocator, now: []const u8, delta: []const u8) Allocator.Error!void { + if (s.extends) { + try s.text.appendSlice(a, delta); + s.chars += reply_text.charCount(delta); + } else { + s.text = .empty; + try s.text.appendSlice(a, now); + s.chars = reply_text.charCount(now); + } + s.at = now.ptr; + } +}; + +/// The token loop and the text it streams. +pub const Generation = struct { + srv: *Server, + a: Allocator, + box: *Mailbox, + id: api.Id, + reply_id: []const u8, // the id the client gets, which the done line prints + sink: ?Sink, + thinking: bool, + stops: reply_text.Stops, + ignore_eos: bool, + max_tokens: u32, + tools: []const Value, + calls: ?tool_stream.Streamer = null, + collected: std.ArrayList(u32) = .empty, + visible: reply_text.Incremental = .{}, + hidden: std.ArrayList(u8) = .empty, // reused for the answer without its call blocks + streamed: Shown = .{}, // what content streamed, kept apart from the decode buffer that grows under slices + streamed_reasoning: Shown = .{}, + streaming_done: bool = false, + first_ns: ?i96 = null, + last_ns: ?i96 = null, + reason: ?[]const u8 = null, // the server ended the reply (a stop string, the length) before the engine said so + engine_done: bool = false, + consumed: usize = 0, + chunk_index: usize = 0, + logprobs: bool = false, // the request asked for logprobs: the mailbox holds a row a token + // a family's reply: its reader, the answer so far, how much of it streamed, and a stop string's hit + reader: ?family_mod.Reader = null, + answer: std.ArrayList(u8) = .empty, + answer_sent: usize = 0, + deltas: std.ArrayList(Value) = .empty, + stop_hit: bool = false, + stop_at: usize = 0, + + pub fn eos(g: *const Generation, t: u32) bool { + return !g.ignore_eos and std.mem.indexOfScalar(u32, g.srv.eos, t) != null; + } + + const closeCall = chat.closeCall; // in chat.zig with its doc + + /// The engine's stop check, here: the newest tokens' text holds a stop string. + fn stopHit(g: *Generation) Allocator.Error!bool { + if (g.stops.strings.len == 0) return false; + const t = g.collected.items; + const tail = t[t.len -| g.stops.tail()..]; + const text = try g.srv.text.decode(g.a, tail); + for (g.stops.strings) |s| if (std.mem.indexOf(u8, text, s) != null) return true; + return false; + } + + /// Commits a round's tokens as ``LaneStream.commit`` would, then streams what they add. + fn commit(g: *Generation, chunk: []const u32) Failure!void { + if (g.reader) |r| return g.commitFamily(r, chunk); + var landed: usize = 0; + for (chunk) |t| { + if (g.reason != null) break; + try g.collected.append(g.a, t); + landed += 1; + if (g.eos(t)) { + g.reason = "stop"; + } else if (try g.stopHit()) { + g.reason = "stop"; + g.srv.engine.cancel(g.id); // the engine ends EOS and length itself; a stop string is the server's + } else if (g.collected.items.len >= g.max_tokens) g.reason = "length"; + } + if (landed == 0) return; + const arrived = nowNs(g.srv.io); + if (g.first_ns == null) g.first_ns = arrived; + g.last_ns = arrived; + const sink = g.sink orelse return; + if (g.streaming_done) return; + var fresh: std.ArrayList(u32) = .empty; + for (chunk[0..landed]) |t| { + if (g.eos(t)) { + g.streaming_done = true; + break; + } + try fresh.append(g.a, t); + } + const all = try g.visible.extend(g.a, g.srv.text, decodeText, fresh.items); + const text = g.stops.visible(all, true); + if (g.visible.pending()) return; // a character still split across tokens + var answer = text; + if (g.thinking) { + const split = try reply_text.splitThinking(g.a, text, false, g.srv.markers); + const piece = g.streamed_reasoning.after(split.reasoning, true); + if (piece.len > 0) { + try g.streamed_reasoning.set(g.a, split.reasoning, piece); + try g.emit(sink, try deltaOf(g.a, "reasoning_content", piece)); + } + answer = split.answer; + } + const shown = if (g.calls != null) try reply_text.hideInto(g.a, &g.hidden, answer, false) else answer; + const vis = reply_text.heldBack(reply_text.streamingVisible(shown), g.calls != null); + const copied = @intFromPtr(vis.ptr) >= @intFromPtr(g.hidden.items.ptr) and @intFromPtr(vis.ptr) <= @intFromPtr(g.hidden.items.ptr) + g.hidden.items.len; + const delta = g.streamed.after(vis, !copied); + if (delta.len > 0) { + try g.streamed.set(g.a, vis, delta); + try g.emit(sink, .{ .string = delta }); + } + if (g.calls) |*c| { + var out: std.ArrayList(Value) = .empty; + try c.feed(answer, &out); // never the reasoning: a call it mentions is not made + for (out.items) |d| try g.emit(sink, d); + } + } + + fn emit(g: *Generation, sink: Sink, delta: Value) Failure!void { + sink.call(sink.ctx, delta) catch { + g.cancel(); + return error.Cancelled; + }; + } + + pub const cancel = chat.cancel; // in chat.zig with its doc + + /// Waits for the engine's own end, so the request it holds may be freed. + pub fn drain(g: *Generation) void { + const m = g.box; + m.mutex.lockUncancelable(m.io); + defer m.mutex.unlock(m.io); + while (!m.finished) m.cond.waitUncancelable(m.io, &m.mutex); + g.engine_done = true; + } + + /// Takes rounds until the reply ends; a client that leaves cancels it. + pub fn loop(g: *Generation, gone: anytype) Failure!void { + const m = g.box; + while (true) { + m.mutex.lockUncancelable(m.io); + var chunk: ?[]u32 = null; + var ended = false; + if (g.chunk_index < m.chunks.items.len) { + const end = m.chunks.items[g.chunk_index]; + g.chunk_index += 1; + chunk = g.a.dupe(u32, m.tokens.items[g.consumed..end]) catch null; + g.consumed = end; + } else if (m.finished) { + ended = true; + } else { + m.cond.waitTimeout(m.io, &m.mutex, .{ .duration = .{ .raw = .fromMilliseconds(50), .clock = .awake } }) catch {}; + } + m.mutex.unlock(m.io); + if (ended) { + g.engine_done = true; + return; + } + if (gone.check()) { + g.cancel(); + return error.Cancelled; + } + const tokens = chunk orelse continue; + if (g.reason == null) try g.commit(tokens); + } + } + + /// A family's round: committed tokens through the family's reader, its deltas pushed (an EOS id ends, unshown). + fn commitFamily(g: *Generation, r: family_mod.Reader, chunk: []const u32) Failure!void { + var landed: usize = 0; + const start = g.collected.items.len; + for (chunk) |t| { + if (g.reason != null) break; + try g.collected.append(g.a, t); + landed += 1; + if (g.eos(t)) { + g.reason = "stop"; + } else if (g.collected.items.len >= g.max_tokens) g.reason = "length"; + } + if (landed == 0) return; + const arrived = nowNs(g.srv.io); + if (g.first_ns == null) g.first_ns = arrived; + g.last_ns = arrived; + for (g.collected.items[start..]) |t| if (!g.isEos(t)) try r.push(&.{t}); + if (r.pending()) return; // a character still split across tokens + g.deltas.clearRetainingCapacity(); + try r.feed(false, &g.deltas); + try g.pushFamily(g.deltas.items, false); + } + + /// Any end-of-sentence id, ``ignore_eos`` or not: never text. + fn isEos(g: *const Generation, t: u32) bool { + return std.mem.indexOfScalar(u32, g.srv.eos, t) != null; + } + + /// ``push``: reasoning and calls go out as they come; the answer is held where a stop may begin, cut where it does. + fn pushFamily(g: *Generation, deltas: []const Value, finished: bool) Failure!void { + for (deltas) |d| switch (d) { + .string => |t| try g.answer.appendSlice(g.a, t), + else => if (g.sink) |sink| if (!g.streaming_done or finished) try g.emit(sink, d), + }; + var full = g.answer.items; + const stops = g.stops.strings; + if (stops.len > 0 and !g.stop_hit) { + var longest: usize = 0; + for (stops) |st| longest = @max(longest, st.len); + const from = (g.answer_sent + 1) -| longest; + var cut: ?usize = null; + for (stops) |st| if (std.mem.indexOfPos(u8, full, @min(from, full.len), st)) |at| { + if (cut == null or at < cut.?) cut = at; + }; + if (cut) |c| { + g.answer.shrinkRetainingCapacity(c); + full = g.answer.items; + g.stop_hit = true; + g.stop_at = g.collected.items.len; + if (g.reason == null) { + g.reason = "stop"; + g.srv.engine.cancel(g.id); // the engine ends EOS and length itself; a stop string is the server's + } + } + } + var upto = full.len; + if (!finished and !g.stop_hit) { + var hold: usize = 0; + for (stops) |st| { + var k = @min(st.len - 1, full.len); + while (k > 0) : (k -= 1) if (std.mem.endsWith(u8, full, st[0..k])) break; + hold = @max(hold, k); + } + upto -= hold; + } + if (upto > g.answer_sent) { + if (g.sink) |sink| try g.emit(sink, .{ .string = full[g.answer_sent..upto] }); + g.answer_sent = upto; + } + } + + /// A family's finished reply: the stream's last deltas, then the parts (``app._run``'s end). + fn finishFamily(g: *Generation, r: family_mod.Reader) Failure!struct { content: []const u8, reasoning: ?[]const u8, calls: []Value, reason: []const u8, used: usize } { + const streaming = g.sink != null; + g.deltas.clearRetainingCapacity(); + try r.feed(true, &g.deltas); + try g.pushFamily(g.deltas.items, true); + const parsed = try r.parse(); + var content: []const u8 = undefined; + var calls: []Value = undefined; + if (streaming) { + calls = try r.calls(); + content = g.answer.items[0..g.answer_sent]; + } else { + calls = parsed.calls; + content = parsed.content; + for (g.stops.strings) |st| if (std.mem.indexOf(u8, content, st)) |at| { + content = content[0..at]; + if (!g.stop_hit) { + g.stop_hit = true; + g.stop_at = g.collected.items.len; + } + }; + } + const t = g.collected.items; + const reason: []const u8 = if (calls.len > 0) "tool_calls" else if (g.stop_hit or (!g.ignore_eos and t.len > 0 and g.isEos(t[t.len - 1]))) "stop" else "length"; + if (g.srv.family) |fam| fam.remember(parsed.reasoning, content, calls); + // tools offered, no call, but tool markup: log up to 300 bytes of what the parser saw (UTF-8 cut, escaped) + if (calls.len == 0 and g.tools.len > 0) if (r.markup() catch null) |m| + std.log.warn("tool markup with no call: request {s}, finish {s}: {s}; reply bytes {d}..: {f}", .{ g.reply_id, reason, m.reason, m.at, std.zig.fmtString(m.window) }); + return .{ .content = content, .reasoning = if (parsed.reasoning.len > 0) parsed.reasoning else null, .calls = calls, .reason = reason, .used = if (g.stop_hit and g.stop_at > 0) g.stop_at else t.len }; + } + + pub fn finish(g: *Generation, cx: *Cx, prompt_len: usize, received: i96, thinking: bool, effort: ?[]const u8, exact: bool, drafts: bool, images_omitted: u32) Failure!Reply { + const a = g.a; + const m = g.box; + if (!g.engine_done) g.drain(); + const reason_generic: []const u8 = g.reason orelse switch (m.reason) { + .stop => "stop", + .length => "length", + .cancelled => "cancelled", + .failed => return cx.other(if (m.message.len > 0) try a.dupe(u8, m.message) else "the reply failed"), + }; + var reason = reason_generic; + var raw: []const u8 = ""; + var content: []const u8 = undefined; + var reasoning: ?[]const u8 = null; + var calls: ?[]Value = null; + var used = g.collected.items.len; + var logprobs: ?Value = null; + if (g.reader) |r| { + const f = try g.finishFamily(r); + content = f.content; + reasoning = f.reasoning; + calls = f.calls; + used = f.used; + if (!(m.reason == .cancelled and g.reason == null)) reason = f.reason; + } else { + const content_tokens = if (g.ignore_eos) g.collected.items else reply_text.stripTrailing(g.collected.items, g.srv.eos); + raw = try g.srv.text.decode(a, content_tokens); + if (g.logprobs) logprobs = try @import("logprobs.zig").value(g.srv, cx, a, g.thinking, content_tokens, m.rows.items); + const visible = g.stops.visible(raw, false); + // closed before the think split, so a call ending an unclosed think block is the answer + const close = try g.closeCall(visible); + const text = if (close.len > 0) try std.mem.concat(a, u8, &.{ visible, close }) else visible; + content = text; + if (g.thinking) { + const split = try reply_text.splitThinking(a, text, true, g.srv.markers); + const r = reply_text.pyStrip(split.reasoning); + reasoning = if (r.len > 0) r else null; + content = split.answer; + } else { + const h = reply_text.parseHarmony(text); + content = h.content; + reasoning = h.reasoning; + } + if (g.sink) |sink| { + const thought = g.streamed_reasoning.text.items; + if (reasoning) |r| if (std.mem.startsWith(u8, r, thought) and r.len > thought.len) + try g.emit(sink, try deltaOf(a, "reasoning_content", r[thought.len..])); + const shown = if (g.calls != null) try reply_text.finishedProse(a, content) else content; + const sent = g.streamed.text.items; + if (std.mem.startsWith(u8, shown, sent) and shown.len > sent.len) try g.emit(sink, .{ .string = shown[sent.len..] }); + if (g.calls) |*c| if (close.len > 0) { + // the streamer reads the closers as markup the model wrote, so the call ends with the same deltas + var out: std.ArrayList(Value) = .empty; + try c.feed(if (g.thinking) content else text, &out); + for (out.items) |d| try g.emit(sink, d); + }; + } + } + const finished_ns = nowNs(g.srv.io); + const prefilled = m.prefilled_ns; + const decode_s: f64 = if (prefilled) |p| @max(0, seconds(finished_ns - p)) else 0; + const decode_tokens = g.collected.items.len -| 1; + const runtime = try json.newObject(a); + try runtime.put(a, "enable_thinking", .{ .bool = thinking }); + try runtime.put(a, "reasoning_effort", if (!thinking) .{ .string = "none" } else if (effort) |e| .{ .string = e } else .null); + try runtime.put(a, "engine", .{ .string = g.srv.info.name }); + try runtime.put(a, "tokens_per_second", .{ .float = if (decode_s > 0) @as(f64, @floatFromInt(decode_tokens)) / decode_s else 0 }); + try runtime.put(a, "seconds", .{ .float = @max(0, decode_s) }); // the reply's decode seconds, as tokens_per_second divides them + try runtime.put(a, "prefill_seconds", if (m.stats.prefill_seconds) |p| .{ .float = p } else .null); + const widths = try a.alloc(Value, m.widths.len); + for (m.widths, widths) |w, *slot| slot.* = try json.intValue(a, w); + try runtime.put(a, "prefill_widths", .{ .array = widths }); + const raised = try a.alloc(Value, m.raised.len); + for (m.raised, raised) |r, *slot| slot.* = .{ .bool = r }; + try runtime.put(a, "prefill_raised", .{ .array = raised }); + try runtime.put(a, "time_to_first_token", if (g.first_ns) |t| .{ .float = seconds(t - received) } else .null); + try runtime.put(a, "sampling", .{ .string = if (exact) "exact" else "greedy" }); + try runtime.put(a, "drafts", .{ .bool = drafts }); + const sha = @import("tokens.zig").tokenSha(g.collected.items); + try runtime.put(a, "token_sha", .{ .string = try a.dupe(u8, &sha) }); + if (images_omitted > 0) try runtime.put(a, "images_omitted", try json.intValue(a, images_omitted)); + try runtime.put(a, "min_rows", try json.intValue(a, m.stats.min_rows)); + if (m.stats.loop_period) |period| { + const loop_field = try json.newObject(a); + try loop_field.put(a, "period", try json.intValue(a, period)); + try runtime.put(a, "loop", .{ .object = loop_field }); + } + const s = m.stats; + const spec = try json.newObject(a); + try spec.put(a, "rounds", try json.intValue(a, s.rounds)); + try spec.put(a, "drafted", try json.intValue(a, s.drafted)); + try spec.put(a, "accepted", try json.intValue(a, s.accepted)); + try spec.put(a, "acceptance_rate", .{ .float = if (s.drafted > 0) @as(f64, @floatFromInt(s.accepted)) / @as(f64, @floatFromInt(s.drafted)) else 0 }); + try spec.put(a, "tokens_per_round", .{ .float = if (s.rounds > 0) @as(f64, @floatFromInt(g.collected.items.len)) / @as(f64, @floatFromInt(s.rounds)) else 0 }); + if (m.telemetry.len > 0) if ((try json.parse(a, m.telemetry)) == .ok) try spec.put(a, "proposer", (try json.parse(a, m.telemetry)).ok); + const think_end: ?u32 = if (thinking) g.srv.text.tokenId(g.srv.markers.close) else null; + const reply: Reply = .{ + .content = content, + .stop_sequence = g.stops.matched(raw), + .reasoning = reasoning, + .tool_calls_streamed = g.calls != null and g.calls.?.streamed, + .finish_reason = reason, + .prompt_tokens = prompt_len, + .cached_tokens = m.cached orelse 0, + .completion_tokens = used, + .reasoning_tokens = reply_text.reasoningCount(g.collected.items, think_end), + .runtime = .{ .object = runtime }, + .speculative = .{ .object = spec }, + .thinking = thinking, + .effort = effort, + .logprobs = logprobs, + .calls = calls, + }; + if (std.mem.eql(u8, reason, "length") and thinking and reply_text.pyStrip(content).len == 0) + log.line("warning: a reply reached max_tokens while still thinking, so its content is empty and its text is all in reasoning_content; raise max_tokens, or send chat_template_kwargs {{\"enable_thinking\": false}} (server: --no-thinking)", .{}); + var cycle_text: [32]u8 = undefined; + const cycle = if (s.loop_period) |period| std.fmt.bufPrint(&cycle_text, " loop=period:{d}", .{period}) catch "" else ""; + log.line("done {s} prompt={d} cached={d} thinking={s} effort={s} tokens={d} sha={s} finish={s}{s} rounds={d} accepted={d}/{d}", .{ g.reply_id, prompt_len, reply.cached_tokens, if (thinking) "True" else "False", if (thinking) effort orelse "none" else "none", g.collected.items.len, sha, reason, cycle, s.rounds, s.accepted, s.drafted }); + return reply; + } +}; diff --git a/zig/src/server/family.zig b/zig/src/server/family.zig new file mode 100644 index 0000000000..9e24c8dab4 --- /dev/null +++ b/zig/src/server/family.zig @@ -0,0 +1,135 @@ +//! A model family's own chat rules (prompt encoding, effort budgets, call markup) in place of the chat template. +const std = @import("std"); +const json = @import("json.zig"); +const errors = @import("errors.zig"); +const Value = json.Value; +const Cx = errors.Cx; +const Allocator = std.mem.Allocator; + +/// A request's thinking switch and effort as it asked for them (null: the server's default). +pub const Thinking = struct { enable: ?bool = null, effort: ?[]const u8 = null }; + +/// What a chat prompt is rendered from: the request's own messages and tools, as the client sent them. +pub const Prompt = struct { + messages: Value, + tools: []const Value, + thinking: bool, + /// The effort label (``Family.thinking``'s, else the default). + effort: []const u8, + generation: bool = true, + /// The request body, for fields the family reads itself (``response_format``). + body: Value = .null, +}; + +pub const Rendered = struct { ids: []const u32, images_omitted: u32 = 0, images: ?Images = null }; + +/// A prompt's prepared images: ``held`` lives until ``release``; ``token`` for /tokenize. +pub const Images = struct { + held: *const anyopaque, + token: u32, + ctx: *anyopaque, + release_fn: *const fn (ctx: *anyopaque) void, + + pub fn release(im: Images) void { + im.release_fn(im.ctx); + } +}; + +/// A finished reply's parts, calls as OpenAI tool-call objects. +pub const Parsed = struct { reasoning: []const u8, content: []const u8, calls: []Value }; + +/// One reply as its tokens arrive: the text so far, and the deltas it adds. +pub const Reader = struct { + ctx: *anyopaque, + vt: *const VTable, + + pub const VTable = struct { + /// Tokens committed (end-of-sentence ids already left out). + push: *const fn (ctx: *anyopaque, ids: []const u32) Allocator.Error!void, + /// A character is split across tokens: nothing new is readable yet. + pending: *const fn (ctx: *anyopaque) bool, + /// Deltas since the last feed (content a string, reasoning and calls delta objects); ``finished`` flushes. + feed: *const fn (ctx: *anyopaque, finished: bool, out: *std.ArrayList(Value)) Allocator.Error!void, + /// The whole reply parsed (after the last ``push``; a split character at the end reads as U+FFFD). + parse: *const fn (ctx: *anyopaque) Allocator.Error!Parsed, + /// The calls the stream sent, in order, as OpenAI tool-call objects. + calls: *const fn (ctx: *anyopaque) Allocator.Error![]Value, + /// The finished reply's tool markup that yielded no call (a reason, at most 300 bytes) for the log; null: none + markup: ?*const fn (ctx: *anyopaque) Allocator.Error!?Markup = null, + }; + + pub fn push(r: Reader, ids: []const u32) Allocator.Error!void { + return r.vt.push(r.ctx, ids); + } + pub fn pending(r: Reader) bool { + return r.vt.pending(r.ctx); + } + pub fn feed(r: Reader, finished: bool, out: *std.ArrayList(Value)) Allocator.Error!void { + return r.vt.feed(r.ctx, finished, out); + } + pub fn parse(r: Reader) Allocator.Error!Parsed { + return r.vt.parse(r.ctx); + } + pub fn calls(r: Reader) Allocator.Error![]Value { + return r.vt.calls(r.ctx); + } + pub fn markup(r: Reader) Allocator.Error!?Markup { + const f = r.vt.markup orelse return null; + return f(r.ctx); + } +}; + +/// A reply's tool markup that the parser turned into no call (``Reader.markup``). +pub const Markup = struct { at: usize, reason: []const u8, window: []const u8 }; + +/// Where a reply's reasoning goes: ``reasoning_content`` (OpenAI's), ``reasoning`` (vLLM's), or both. +pub const ReasoningFields = enum { both, reasoning_content, reasoning }; + +pub const Family = struct { + ctx: *anyopaque, + vt: *const VTable, + + pub const VTable = struct { + name: []const u8, + /// ``chat_template_kwargs`` and ``reasoning_effort`` read; a bad value refuses the request. + thinking: *const fn (ctx: *anyopaque, cx: *Cx, body: Value) errors.Refused!Thinking, + /// The effort a request that names none gets. + default_effort: *const fn (ctx: *anyopaque) []const u8, + /// The tools the prompt offers: ``tool_choice`` none offers none; otherwise all of them. + tools: *const fn (ctx: *anyopaque, cx: *Cx, body: Value) errors.Refused![]const Value, + render: *const fn (ctx: *anyopaque, cx: *Cx, p: Prompt) errors.Refused!Rendered, + reader: *const fn (ctx: *anyopaque, a: Allocator, thinking: bool, tools: []const Value, single: bool) Allocator.Error!Reader, + /// A finished reply's reasoning, kept for a later request whose history comes back without it. + remember: *const fn (ctx: *anyopaque, reasoning: []const u8, content: []const u8, calls: []const Value) void, + reasoning_fields: ReasoningFields = .both, + /// Whether the request's ``min_p`` reaches the sampler (some families ignore it). + reads_min_p: bool = true, + }; + + pub fn thinking(f: Family, cx: *Cx, body: Value) errors.Refused!Thinking { + return f.vt.thinking(f.ctx, cx, body); + } + pub fn defaultEffort(f: Family) []const u8 { + return f.vt.default_effort(f.ctx); + } + pub fn tools(f: Family, cx: *Cx, body: Value) errors.Refused![]const Value { + return f.vt.tools(f.ctx, cx, body); + } + pub fn render(f: Family, cx: *Cx, p: Prompt) errors.Refused!Rendered { + return f.vt.render(f.ctx, cx, p); + } + pub fn reader(f: Family, a: Allocator, think: bool, offered: []const Value, single: bool) Allocator.Error!Reader { + return f.vt.reader(f.ctx, a, think, offered, single); + } + pub fn remember(f: Family, reasoning: []const u8, content: []const u8, calls: []const Value) void { + f.vt.remember(f.ctx, reasoning, content, calls); + } + + /// A reasoning delta under the family's field names. + pub fn reasoningDelta(f: Family, a: Allocator, text: []const u8) Allocator.Error!Value { + const o = try json.newObject(a); + if (f.vt.reasoning_fields != .reasoning) try o.put(a, "reasoning_content", .{ .string = text }); + if (f.vt.reasoning_fields != .reasoning_content) try o.put(a, "reasoning", .{ .string = text }); + return .{ .object = o }; + } +}; diff --git a/zig/src/server/openai.zig b/zig/src/server/openai.zig index 91d02cceb8..4f0f463018 100644 --- a/zig/src/server/openai.zig +++ b/zig/src/server/openai.zig @@ -19,6 +19,8 @@ const routes = @import("routes.zig"); const warm = @import("warm.zig"); const Server = @import("server.zig").Server; const Conn = @import("http_conn.zig").Conn; +const openai_run = @import("openai_run.zig"); +const Run = openai_run.Run; const Value = json.Value; const Cx = errors.Cx; const Allocator = std.mem.Allocator; @@ -107,7 +109,7 @@ pub fn replyModel(srv: *const Server, body: Value) []const u8 { return srv.config.served_name; } -const Plan = struct { +pub const Plan = struct { input: chat.Input, stream: bool, separate_usage: bool, @@ -123,11 +125,18 @@ fn plan(srv: *Server, cx: *Cx, is_chat: bool, raw: Value) errors.Refused!Plan { const exact_bytes = json.truthyField(body, "logprobs") and srv.info.logprobs and (srv.text.tokenBytes(cx.a, 0) catch return error.OutOfMemory) != null; const logprobs = try fields.probabilityOptions(cx, body, exact_bytes); if (srv.config.request_log) |path| request_log.append(cx.a, path, body); - var input: chat.Input = .{ .fields = undefined }; - if (is_chat) { + var input: chat.Input = .{ .fields = undefined, .body = body }; + if (srv.family) |fam| if (is_chat) { + // a family renders the client's messages and tools itself (its encoding reads OpenAI's shapes) + const m = body.get("messages") orelse Value.null; + if (m != .array or m.array.len == 0) return cx.refuse("messages must be a non-empty list"); + input.messages = m; + input.tools = try fam.tools(cx, body); + }; + if (is_chat and srv.family == null) { input.messages = try messages.normalize(cx, body.get("messages"), "system", srv.needs_user_after_tool); input.tools = try tool_specs.active(cx, body.get("tools"), body.get("tool_choice")); - } else if (body.get("messages")) |m| if (m == .array and m.array.len > 0) { + } else if (!is_chat) if (body.get("messages")) |m| if (m == .array and m.array.len > 0) { input.messages = try messages.normalize(cx, m, "system", srv.needs_user_after_tool); }; if (!is_chat and input.messages.array.len == 0) input.prompt = try legacyPrompt(srv, cx, body.get("prompt")); @@ -145,9 +154,12 @@ fn plan(srv: *Server, cx: *Cx, is_chat: bool, raw: Value) errors.Refused!Plan { const sampling = try json.newObject(cx.a); for ([_][]const u8{ "temperature", "top_p", "top_k", "min_p", "seed", "priority", "draft", "thinking_budget", "ignore_eos", "stop" } ++ grammar.fields) |k| if (body.get(k)) |v| try sampling.put(cx.a, k, v); try grammar.refusal(cx, body); - if (input.tools.len > 0 and try tool_specs.choiceRequiresCall(cx.a, body.get("tool_choice"))) try sampling.put(cx.a, "tool_call_required", .{ .bool = true }); - if (body.field("parallel_tool_calls")) |p| if (p == .bool and !p.bool) try sampling.put(cx.a, "tool_call_single", .{ .bool = true }); - const thinking = try fields.thinkingFields(cx, body, srv.effort_levels); + if (srv.family == null and input.tools.len > 0 and try tool_specs.choiceRequiresCall(cx.a, body.get("tool_choice"))) try sampling.put(cx.a, "tool_call_required", .{ .bool = true }); + if (srv.family == null) if (body.field("parallel_tool_calls")) |p| if (p == .bool and !p.bool) try sampling.put(cx.a, "tool_call_single", .{ .bool = true }); + const thinking: fields.Thinking = if (srv.family) |fam| blk: { + const t = try fam.thinking(cx, body); + break :blk .{ .enable = t.enable, .effort = t.effort }; + } else try fields.thinkingFields(cx, body, srv.effort_levels); if (thinking.effort) |e| try sampling.put(cx.a, "reasoning_effort", .{ .string = e }); if (thinking.enable) |on| try sampling.put(cx.a, "enable_thinking", .{ .bool = on }); input.fields = .{ .object = sampling }; @@ -165,6 +177,7 @@ fn plan(srv: *Server, cx: *Cx, is_chat: bool, raw: Value) errors.Refused!Plan { if (body.field("parallel_tool_calls")) |p| { if (p != .bool) return cx.refuse("parallel_tool_calls must be a boolean"); policy.single = !p.bool; + input.single_call = policy.single; } const options = body.get("stream_options"); return .{ @@ -237,307 +250,13 @@ pub fn run(srv: *Server, a: Allocator, out: Out, gone: Gone, is_chat: bool, raw: } /// Why a request got an error reply, logged before it: its access line shows only the status. -fn logRefused(id: []const u8, message: []const u8) void { +pub fn logRefused(id: []const u8, message: []const u8) void { var buf: [1400]u8 = undefined; log.line("{s}", .{log.refused(&buf, id, message)}); } -fn errorOther(a: Allocator, message: []const u8) Allocator.Error!Value { +pub fn errorOther(a: Allocator, message: []const u8) Allocator.Error!Value { const o = try json.newObject(a); try o.put(a, "message", .{ .string = message }); return .{ .object = o }; } - -const Run = struct { - srv: *Server, - a: Allocator, - out: Out, - is_chat: bool, - plan: Plan, - id: []const u8, - created: i64, - streamed_prose: bool = false, - prose_sent: std.ArrayList(u8) = .empty, // the content a tool stream sent - - fn chunk(r: *Run, delta: ?Value, finish: ?[]const u8) Allocator.Error!Value { - const a = r.a; - const o = try json.newObject(a); - try o.put(a, "id", .{ .string = r.id }); - try o.put(a, "object", .{ .string = if (r.is_chat) "chat.completion.chunk" else "text_completion" }); - try o.put(a, "created", try json.intValue(a, r.created)); - try o.put(a, "model", .{ .string = r.plan.named }); - const choice = try json.newObject(a); - try choice.put(a, "index", .{ .int = "0" }); - if (r.is_chat) { - const d: Value = if (delta) |v| (if (v == .object) v else if (v.string.len > 0) try chat.deltaOf(a, "content", v.string) else .{ .object = try json.newObject(a) }) else .{ .object = try json.newObject(a) }; - try choice.put(a, "delta", d); - try choice.put(a, "finish_reason", if (finish) |f| .{ .string = f } else .null); - } else { - try choice.put(a, "text", delta orelse .{ .string = "" }); - try choice.put(a, "finish_reason", if (finish) |f| .{ .string = f } else .null); - try choice.put(a, "logprobs", .null); - } - const list = try a.alloc(Value, 1); - list[0] = .{ .object = choice }; - try o.put(a, "choices", .{ .array = list }); - return .{ .object = o }; - } - - fn emit(r: *Run, payload: Value) error{Closed}!void { - return r.out.vt.event(r.out.ctx, payload); - } - - fn usage(r: *Run, reply: *const chat.Reply) Allocator.Error!Value { - const a = r.a; - const o = try json.newObject(a); - try o.put(a, "prompt_tokens", try json.intValue(a, reply.prompt_tokens)); - try o.put(a, "completion_tokens", try json.intValue(a, reply.completion_tokens)); - try o.put(a, "total_tokens", try json.intValue(a, reply.prompt_tokens + reply.completion_tokens)); - const cached = try json.newObject(a); - try cached.put(a, "cached_tokens", try json.intValue(a, reply.cached_tokens)); - try o.put(a, "prompt_tokens_details", .{ .object = cached }); - const thought = try json.newObject(a); - try thought.put(a, "reasoning_tokens", try json.intValue(a, reply.reasoning_tokens)); - try o.put(a, "completion_tokens_details", .{ .object = thought }); - return .{ .object = o }; - } - - /// ``response_extras``: exact_mode, the matched stop, the runtime and the draft counters. - fn extras(r: *Run, o: *json.Object, reply: *const chat.Reply) Allocator.Error!void { - try o.put(r.a, "exact_mode", .{ .string = "exact" }); - if (reply.stop_sequence) |s| try o.put(r.a, "stop_sequence", .{ .string = s }); - try o.put(r.a, "tensorfold", reply.runtime); - try o.put(r.a, "speculative", reply.speculative); - } - - /// ``ToolCallPolicy.finish``: the reply's calls parsed out of its content. - fn attachCalls(r: *Run, reply: *chat.Reply) Allocator.Error!?[]Value { - const tools = r.plan.input.tools; - if (tools.len == 0) return null; - const parsed = try tool_parse.parse(r.a, reply.content, tools, r.plan.policy.maxCalls()); - const calls = parsed.calls orelse { - if (r.plan.policy.single) reply.content = try r.plan.policy.parsedContent(r.a, parsed.content); - return null; - }; - reply.content = try r.plan.policy.parsedContent(r.a, parsed.content); - reply.finish_reason = "tool_calls"; - if (r.plan.policy.single) reply.tool_calls_streamed = false; - return if (r.plan.policy.single and calls.len > 1) calls[0..1] else calls; - } - - fn fail(r: *Run, cx: *const Cx, field: []const u8) void { - const a = r.a; - const kind_ok = cx.kind != .other and cx.kind != .server; - const body = (if (kind_ok) errorBody(a, cx, field) else errorOther(a, cx.message)) catch return; - r.out.vt.reply(r.out.ctx, if (kind_ok) 400 else 500, wrapError(a, body) catch return); - } - - /// A reply that ended before anything was sent: a refusal's 400, a failure's 500, nothing when the client left. - fn unsent(r: *Run, cx: *Cx, e: chat.Failure, field: []const u8) void { - switch (e) { - error.Cancelled => {}, - error.OutOfMemory => r.fail(&.{ .a = r.a, .kind = .server, .message = "out of memory" }, field), - error.Failed => r.fail(&.{ .a = r.a, .kind = .server, .message = "the reply failed" }, field), - error.Refused => { - if (cx.kind == .other) cx.kind = .server; - if (cx.kind != .server) logRefused(r.id, cx.message); // a 500 is a failure, not a refusal - r.fail(cx, field); - }, - } - } - - /// A context-window refusal on a stream: the role chunk when there are no tools, then the error event and [DONE]. - fn streamedContext(r: *Run, cx: *const Cx, tools: bool, field: []const u8) void { - const a = r.a; - r.out.vt.open(r.out.ctx) catch return; - if (!tools and r.is_chat) r.emit(r.chunk(roleDelta(a) catch return, null) catch return) catch return; - logRefused(r.id, cx.message); - const body = errorBody(a, cx, field) catch return; - r.emit(wrapError(a, body) catch return) catch return; - r.out.vt.event(r.out.ctx, null) catch {}; - } - - fn whole(r: *Run, gone: Gone, field: []const u8) void { - const a = r.a; - var cx: Cx = .{ .a = a }; - var reply = chat.run(r.srv, &cx, r.plan.input, null, gone) catch |e| return r.unsent(&cx, e, field); - const calls = r.attachCalls(&reply) catch return; - const o = json.newObject(a) catch return; - r.wholeBody(o, &reply, calls) catch return; - r.out.vt.reply(r.out.ctx, 200, .{ .object = o }); - r.warmNext(&reply, calls); - } - - /// The next turn's prompt prefilled in the background, from the reply as its client got it (warm.zig). - fn warmNext(r: *Run, reply: *const chat.Reply, calls: ?[]Value) void { - if (!r.is_chat) return; - const message = json.newObject(r.a) catch return; - message.put(r.a, "role", .{ .string = "assistant" }) catch return; - message.put(r.a, "content", if (calls != null) .null else .{ .string = reply.content }) catch return; - if (calls) |c| message.put(r.a, "tool_calls", .{ .array = c }) catch return; - warm.reply(r.srv, r.plan.input, .{ .object = message }, reply.thinking, reply.effort); - } - - fn wholeBody(r: *Run, o: *json.Object, reply: *const chat.Reply, calls: ?[]Value) Allocator.Error!void { - const a = r.a; - try o.put(a, "id", .{ .string = r.id }); - try o.put(a, "object", .{ .string = if (r.is_chat) "chat.completion" else "text_completion" }); - try o.put(a, "created", try json.intValue(a, r.created)); - try o.put(a, "model", .{ .string = r.plan.named }); - const choice = try json.newObject(a); - try choice.put(a, "index", .{ .int = "0" }); - if (r.is_chat) { - const message = try json.newObject(a); - try message.put(a, "role", .{ .string = "assistant" }); - try message.put(a, "content", if (calls != null) .null else .{ .string = reply.content }); - if (reply.reasoning) |t| if (t.len > 0) try message.put(a, "reasoning_content", .{ .string = t }); - if (calls) |c| try message.put(a, "tool_calls", .{ .array = c }); - try choice.put(a, "message", .{ .object = message }); - try choice.put(a, "finish_reason", .{ .string = reply.finish_reason }); - if (reply.logprobs) |lp| try choice.put(a, "logprobs", lp); - } else { - try choice.put(a, "text", .{ .string = reply.content }); - try choice.put(a, "finish_reason", .{ .string = reply.finish_reason }); - try choice.put(a, "logprobs", .null); - } - const list = try a.alloc(Value, 1); - list[0] = .{ .object = choice }; - try o.put(a, "choices", .{ .array = list }); - try o.put(a, "usage", try r.usage(reply)); - try r.extras(o, reply); - } - - const StreamSink = struct { - run: *Run, - tools: bool, - - fn call(ctx: *anyopaque, delta: Value) error{Closed}!void { - const s: *StreamSink = @ptrCast(@alignCast(ctx)); - const r = s.run; - if (!s.tools) { - if (!r.is_chat and delta != .string) return; // completions stream text only - return r.emit(r.chunk(delta, null) catch return error.Closed); - } - const filtered = r.policyDelta(delta) catch return error.Closed; - const d = filtered orelse return; - try r.prose(); - r.prose_sent.appendSlice(r.a, if (d == .string) d.string else d.strField("content") orelse "") catch return error.Closed; - return r.emit(r.chunk(d, null) catch return error.Closed); - } - }; - - /// The one-call policy on a streamed delta: text filtered, tool_calls dropped; null when nothing is left. - fn policyDelta(r: *Run, delta: Value) Allocator.Error!?Value { - if (!r.plan.policy.single) return if (delta.truthy()) delta else null; - if (delta == .string) { - const t = try r.plan.policy.text(r.a, delta.string); - return if (t.len > 0) Value{ .string = t } else null; - } - const o = try json.newObject(r.a); - var content: ?[]const u8 = null; - for (delta.object.keys(), delta.object.values()) |k, v| { - if (std.mem.eql(u8, k, "tool_calls")) continue; - if (std.mem.eql(u8, k, "content") and v == .string) { - content = v.string; - continue; - } - try o.put(r.a, k, v); - } - if (content) |c| { - const t = try r.plan.policy.text(r.a, c); - if (t.len > 0) try o.put(r.a, "content", .{ .string = t }); - } - return if (o.count() > 0) Value{ .object = o } else null; - } - - /// The assistant role before the first prose of a tool-using stream. - fn prose(r: *Run) error{Closed}!void { - if (r.streamed_prose) return; - r.streamed_prose = true; - const role = roleDelta(r.a) catch return error.Closed; - return r.emit(r.chunk(role, null) catch return error.Closed); - } - - fn stream(r: *Run, gone: Gone, field: []const u8) void { - const a = r.a; - const tools = r.plan.input.tools.len > 0; - var cx: Cx = .{ .a = a }; - // a context-window refusal is reported inside the stream; every other refusal is still a 400 before it opens - var compaction: ?compact.Stamp = null; - const prepared = if (r.srv.config.compact_at == null) chat.prepare(r.srv, &cx, r.plan.input, gone) catch |e| { - if (e == error.Refused and cx.kind == .context_length) return r.streamedContext(&cx, tools, field); - return r.unsent(&cx, e, field); - } else blk: { - const ready = compact.prepare(r.srv, &cx, r.plan.input, gone) catch |e| { - if (e == error.Refused and cx.kind == .context_length) return r.streamedContext(&cx, tools, field); - return r.unsent(&cx, e, field); - }; - compaction = ready.stamp; - break :blk ready.prepared; - }; - var handed = false; // generate gives the preparing count back from here on - defer if (!handed) chat.release(r.srv, prepared.preparing); - r.out.vt.open(r.out.ctx) catch return; - if (!tools and r.is_chat) r.emit(r.chunk(roleDelta(a) catch return, null) catch return) catch return; - var sink_state: StreamSink = .{ .run = r, .tools = tools }; - handed = true; - var reply = chat.generate(r.srv, &cx, prepared, .{ .ctx = &sink_state, .call = StreamSink.call }, gone) catch |e| { - switch (e) { - error.Cancelled => return, - error.Refused => if (cx.kind != .other and cx.kind != .server) { - logRefused(r.id, cx.message); - const body = errorBody(a, &cx, field) catch return; - r.emit(wrapError(a, body) catch return) catch return; - r.out.vt.event(r.out.ctx, null) catch {}; - return; - }, - else => {}, - } - const message = if (e == error.Refused) cx.message else if (e == error.OutOfMemory) "out of memory" else "the reply failed"; - log.line("stream error: {s}", .{message}); - const err = json.newObject(a) catch return; - err.put(a, "message", .{ .string = message }) catch return; - err.put(a, "type", .{ .string = "server_error" }) catch return; - r.emit(wrapError(a, .{ .object = err }) catch return) catch return; - r.out.vt.event(r.out.ctx, null) catch {}; - return; - }; - var calls: ?[]Value = null; - if (tools) { - calls = r.attachCalls(&reply) catch return; - const tail = r.plan.policy.flush(); - if (tail.len > 0) { - r.prose() catch return; - r.prose_sent.appendSlice(a, tail) catch return; - r.emit(r.chunk(.{ .string = tail }, null) catch return) catch return; - } - if (calls != null and !reply.tool_calls_streamed) { - r.emit(r.chunk(roleDelta(a) catch return, null) catch return) catch return; - const deltas = tool_parse.deltas(a, calls.?) catch return; - for (deltas) |d| r.emit(r.chunk(d, null) catch return) catch return; - } else if (reply.content.len > 0 and !r.streamed_prose) { - r.emit(r.chunk(.{ .string = reply.content }, null) catch return) catch return; - } else if (r.plan.policy.kept(reply.content, r.prose_sent.items)) |rest| { - r.emit(r.chunk(.{ .string = rest }, null) catch return) catch return; - } - } - const last = r.chunk(.{ .string = "" }, if (reply.finish_reason.len > 0) reply.finish_reason else "length") catch return; - if (compaction) |s| compact.stamp(&cx, &reply, s) catch return; - r.extras(last.object, &reply) catch return; - const use = r.usage(&reply) catch return; - if (!r.plan.separate_usage) last.object.put(a, "usage", use) catch return; - r.emit(last) catch return; - if (r.plan.separate_usage) { - const u = r.chunk(null, null) catch return; - u.object.put(a, "choices", .{ .array = &.{} }) catch return; - u.object.put(a, "usage", use) catch return; - r.emit(u) catch return; - } - r.out.vt.event(r.out.ctx, null) catch {}; - r.warmNext(&reply, calls); - } -}; - -fn roleDelta(a: Allocator) Allocator.Error!Value { - return chat.deltaOf(a, "role", "assistant"); -} diff --git a/zig/src/server/openai_run.zig b/zig/src/server/openai_run.zig new file mode 100644 index 0000000000..3f9978259d --- /dev/null +++ b/zig/src/server/openai_run.zig @@ -0,0 +1,328 @@ +//! One OpenAI request's run: the reply generated, then sent whole or streamed as chunks. +const std = @import("std"); +const json = @import("json.zig"); +const errors = @import("errors.zig"); +const chat = @import("chat.zig"); +const compact = @import("compact.zig"); +const tool_parse = @import("tool_parse.zig"); +const tool_stream = @import("tool_stream.zig"); +const ids = @import("ids.zig"); +const log = @import("log.zig"); +const request_log = @import("request_log.zig"); +const sse = @import("sse.zig"); +const warm = @import("warm.zig"); +const Server = @import("server.zig").Server; +const Value = json.Value; +const Cx = errors.Cx; +const Allocator = std.mem.Allocator; +const openai = @import("openai.zig"); +const Out = openai.Out; +const Gone = openai.Gone; +const errorBody = openai.errorBody; +const wrapError = openai.wrapError; +const replyModel = openai.replyModel; +const Plan = openai.Plan; +const logRefused = openai.logRefused; +const errorOther = openai.errorOther; + +pub const Run = struct { + srv: *Server, + a: Allocator, + out: Out, + is_chat: bool, + plan: Plan, + id: []const u8, + created: i64, + streamed_prose: bool = false, + prose_sent: std.ArrayList(u8) = .empty, // the content a tool stream sent + + fn chunk(r: *Run, delta: ?Value, finish: ?[]const u8) Allocator.Error!Value { + const a = r.a; + const o = try json.newObject(a); + try o.put(a, "id", .{ .string = r.id }); + try o.put(a, "object", .{ .string = if (r.is_chat) "chat.completion.chunk" else "text_completion" }); + try o.put(a, "created", try json.intValue(a, r.created)); + try o.put(a, "model", .{ .string = r.plan.named }); + const choice = try json.newObject(a); + try choice.put(a, "index", .{ .int = "0" }); + if (r.is_chat) { + const d: Value = if (delta) |v| (if (v == .object) v else if (v.string.len > 0) try chat.deltaOf(a, "content", v.string) else .{ .object = try json.newObject(a) }) else .{ .object = try json.newObject(a) }; + try choice.put(a, "delta", d); + try choice.put(a, "finish_reason", if (finish) |f| .{ .string = f } else .null); + } else { + try choice.put(a, "text", delta orelse .{ .string = "" }); + try choice.put(a, "finish_reason", if (finish) |f| .{ .string = f } else .null); + try choice.put(a, "logprobs", .null); + } + const list = try a.alloc(Value, 1); + list[0] = .{ .object = choice }; + try o.put(a, "choices", .{ .array = list }); + return .{ .object = o }; + } + + fn emit(r: *Run, payload: Value) error{Closed}!void { + return r.out.vt.event(r.out.ctx, payload); + } + + fn usage(r: *Run, reply: *const chat.Reply) Allocator.Error!Value { + const a = r.a; + const o = try json.newObject(a); + try o.put(a, "prompt_tokens", try json.intValue(a, reply.prompt_tokens)); + try o.put(a, "completion_tokens", try json.intValue(a, reply.completion_tokens)); + try o.put(a, "total_tokens", try json.intValue(a, reply.prompt_tokens + reply.completion_tokens)); + const cached = try json.newObject(a); + try cached.put(a, "cached_tokens", try json.intValue(a, reply.cached_tokens)); + try o.put(a, "prompt_tokens_details", .{ .object = cached }); + const thought = try json.newObject(a); + try thought.put(a, "reasoning_tokens", try json.intValue(a, reply.reasoning_tokens)); + try o.put(a, "completion_tokens_details", .{ .object = thought }); + return .{ .object = o }; + } + + /// ``response_extras``: exact_mode, the matched stop, the runtime and the draft counters. + fn extras(r: *Run, o: *json.Object, reply: *const chat.Reply) Allocator.Error!void { + try o.put(r.a, "exact_mode", .{ .string = "exact" }); + if (reply.stop_sequence) |s| try o.put(r.a, "stop_sequence", .{ .string = s }); + try o.put(r.a, "tensorfold", reply.runtime); + try o.put(r.a, "speculative", reply.speculative); + } + + /// ``ToolCallPolicy.finish``: the reply's calls parsed out of its content. + fn attachCalls(r: *Run, reply: *chat.Reply) Allocator.Error!?[]Value { + if (reply.calls) |calls| return if (calls.len > 0) calls else null; // the family parsed them (and applied the one-call policy) + const tools = r.plan.input.tools; + if (tools.len == 0) return null; + const parsed = try tool_parse.parse(r.a, reply.content, tools, r.plan.policy.maxCalls()); + const calls = parsed.calls orelse { + if (r.plan.policy.single) reply.content = try r.plan.policy.parsedContent(r.a, parsed.content); + return null; + }; + reply.content = try r.plan.policy.parsedContent(r.a, parsed.content); + reply.finish_reason = "tool_calls"; + if (r.plan.policy.single) reply.tool_calls_streamed = false; + return if (r.plan.policy.single and calls.len > 1) calls[0..1] else calls; + } + + fn fail(r: *Run, cx: *const Cx, field: []const u8) void { + const a = r.a; + const kind_ok = cx.kind != .other and cx.kind != .server; + const body = (if (kind_ok) errorBody(a, cx, field) else errorOther(a, cx.message)) catch return; + r.out.vt.reply(r.out.ctx, if (kind_ok) 400 else 500, wrapError(a, body) catch return); + } + + /// A reply that ended before anything was sent: a refusal's 400, a failure's 500, nothing when the client left. + fn unsent(r: *Run, cx: *Cx, e: chat.Failure, field: []const u8) void { + switch (e) { + error.Cancelled => {}, + error.OutOfMemory => r.fail(&.{ .a = r.a, .kind = .server, .message = "out of memory" }, field), + error.Failed => r.fail(&.{ .a = r.a, .kind = .server, .message = "the reply failed" }, field), + error.Refused => { + if (cx.kind == .other) cx.kind = .server; + if (cx.kind != .server) logRefused(r.id, cx.message); // a 500 is a failure, not a refusal + r.fail(cx, field); + }, + } + } + + /// A context-window refusal on a stream: the role chunk when there are no tools, then the error event and [DONE]. + fn streamedContext(r: *Run, cx: *const Cx, tools: bool, field: []const u8) void { + const a = r.a; + r.out.vt.open(r.out.ctx) catch return; + if (!tools and r.is_chat) r.emit(r.chunk(roleDelta(a) catch return, null) catch return) catch return; + logRefused(r.id, cx.message); + const body = errorBody(a, cx, field) catch return; + r.emit(wrapError(a, body) catch return) catch return; + r.out.vt.event(r.out.ctx, null) catch {}; + } + + pub fn whole(r: *Run, gone: Gone, field: []const u8) void { + const a = r.a; + var cx: Cx = .{ .a = a }; + var reply = chat.run(r.srv, &cx, r.plan.input, null, gone) catch |e| return r.unsent(&cx, e, field); + const calls = r.attachCalls(&reply) catch return; + const o = json.newObject(a) catch return; + r.wholeBody(o, &reply, calls) catch return; + r.out.vt.reply(r.out.ctx, 200, .{ .object = o }); + r.warmNext(&reply, calls); + } + + /// The next turn's prompt prefilled in the background, from the reply as its client got it (warm.zig). + fn warmNext(r: *Run, reply: *const chat.Reply, calls: ?[]Value) void { + if (!r.is_chat) return; + const message = json.newObject(r.a) catch return; + message.put(r.a, "role", .{ .string = "assistant" }) catch return; + message.put(r.a, "content", if (calls != null) .null else .{ .string = reply.content }) catch return; + if (calls) |c| message.put(r.a, "tool_calls", .{ .array = c }) catch return; + warm.reply(r.srv, r.plan.input, .{ .object = message }, reply.thinking, reply.effort); + } + + fn wholeBody(r: *Run, o: *json.Object, reply: *const chat.Reply, calls: ?[]Value) Allocator.Error!void { + const a = r.a; + try o.put(a, "id", .{ .string = r.id }); + try o.put(a, "object", .{ .string = if (r.is_chat) "chat.completion" else "text_completion" }); + try o.put(a, "created", try json.intValue(a, r.created)); + try o.put(a, "model", .{ .string = r.plan.named }); + const choice = try json.newObject(a); + try choice.put(a, "index", .{ .int = "0" }); + if (r.is_chat) { + const message = try json.newObject(a); + try message.put(a, "role", .{ .string = "assistant" }); + const empty_null = r.srv.family != null and reply.content.len == 0; // a family's reply: ``content or None`` + try message.put(a, "content", if (calls != null or empty_null) .null else .{ .string = reply.content }); + if (reply.reasoning) |t| if (t.len > 0) { + if (r.srv.family) |fam| { + const d = try fam.reasoningDelta(a, t); + for (d.object.keys(), d.object.values()) |k, v| try message.put(a, k, v); + } else try message.put(a, "reasoning_content", .{ .string = t }); + }; + if (calls) |c| try message.put(a, "tool_calls", .{ .array = c }); + try choice.put(a, "message", .{ .object = message }); + try choice.put(a, "finish_reason", .{ .string = reply.finish_reason }); + if (reply.logprobs) |lp| try choice.put(a, "logprobs", lp); + } else { + try choice.put(a, "text", .{ .string = reply.content }); + try choice.put(a, "finish_reason", .{ .string = reply.finish_reason }); + try choice.put(a, "logprobs", .null); + } + const list = try a.alloc(Value, 1); + list[0] = .{ .object = choice }; + try o.put(a, "choices", .{ .array = list }); + try o.put(a, "usage", try r.usage(reply)); + try r.extras(o, reply); + } + + const StreamSink = struct { + run: *Run, + tools: bool, + + fn call(ctx: *anyopaque, delta: Value) error{Closed}!void { + const s: *StreamSink = @ptrCast(@alignCast(ctx)); + const r = s.run; + if (!s.tools) { + if (!r.is_chat and delta != .string) return; // completions stream text only + return r.emit(r.chunk(delta, null) catch return error.Closed); + } + const filtered = r.policyDelta(delta) catch return error.Closed; + const d = filtered orelse return; + try r.prose(); + r.prose_sent.appendSlice(r.a, if (d == .string) d.string else d.strField("content") orelse "") catch return error.Closed; + return r.emit(r.chunk(d, null) catch return error.Closed); + } + }; + + /// The one-call policy on a streamed delta: text filtered, tool_calls dropped; null when nothing is left. + fn policyDelta(r: *Run, delta: Value) Allocator.Error!?Value { + if (!r.plan.policy.single) return if (delta.truthy()) delta else null; + if (delta == .string) { + const t = try r.plan.policy.text(r.a, delta.string); + return if (t.len > 0) Value{ .string = t } else null; + } + const o = try json.newObject(r.a); + var content: ?[]const u8 = null; + for (delta.object.keys(), delta.object.values()) |k, v| { + if (std.mem.eql(u8, k, "tool_calls")) continue; + if (std.mem.eql(u8, k, "content") and v == .string) { + content = v.string; + continue; + } + try o.put(r.a, k, v); + } + if (content) |c| { + const t = try r.plan.policy.text(r.a, c); + if (t.len > 0) try o.put(r.a, "content", .{ .string = t }); + } + return if (o.count() > 0) Value{ .object = o } else null; + } + + /// The assistant role before the first prose of a tool-using stream. + fn prose(r: *Run) error{Closed}!void { + if (r.streamed_prose) return; + r.streamed_prose = true; + const role = roleDelta(r.a) catch return error.Closed; + return r.emit(r.chunk(role, null) catch return error.Closed); + } + + pub fn stream(r: *Run, gone: Gone, field: []const u8) void { + const a = r.a; + // a family streams calls and prose itself, so its stream is the plain one + const tools = r.plan.input.tools.len > 0 and r.srv.family == null; + var cx: Cx = .{ .a = a }; + // a context-window refusal is reported inside the stream + var compaction: ?compact.Stamp = null; + const prepared = if (r.srv.config.compact_at == null) chat.prepare(r.srv, &cx, r.plan.input, gone) catch |e| { + if (e == error.Refused and cx.kind == .context_length) return r.streamedContext(&cx, tools, field); + return r.unsent(&cx, e, field); + } else blk: { + const ready = compact.prepare(r.srv, &cx, r.plan.input, gone) catch |e| { + if (e == error.Refused and cx.kind == .context_length) return r.streamedContext(&cx, tools, field); + return r.unsent(&cx, e, field); + }; + compaction = ready.stamp; + break :blk ready.prepared; + }; + var handed = false; // generate gives the preparing count back from here on + defer if (!handed) chat.release(r.srv, prepared.preparing); + r.out.vt.open(r.out.ctx) catch return; + if (!tools and r.is_chat) r.emit(r.chunk(roleDelta(a) catch return, null) catch return) catch return; + var sink_state: StreamSink = .{ .run = r, .tools = tools }; + handed = true; + var reply = chat.generate(r.srv, &cx, prepared, .{ .ctx = &sink_state, .call = StreamSink.call }, gone) catch |e| { + switch (e) { + error.Cancelled => return, + error.Refused => if (cx.kind != .other and cx.kind != .server) { + logRefused(r.id, cx.message); + const body = errorBody(a, &cx, field) catch return; + r.emit(wrapError(a, body) catch return) catch return; + r.out.vt.event(r.out.ctx, null) catch {}; + return; + }, + else => {}, + } + const message = if (e == error.Refused) cx.message else if (e == error.OutOfMemory) "out of memory" else "the reply failed"; + log.line("stream error: {s}", .{message}); + const err = json.newObject(a) catch return; + err.put(a, "message", .{ .string = message }) catch return; + err.put(a, "type", .{ .string = "server_error" }) catch return; + r.emit(wrapError(a, .{ .object = err }) catch return) catch return; + r.out.vt.event(r.out.ctx, null) catch {}; + return; + }; + var calls: ?[]Value = null; + if (tools) { + calls = r.attachCalls(&reply) catch return; + const tail = r.plan.policy.flush(); + if (tail.len > 0) { + r.prose() catch return; + r.prose_sent.appendSlice(a, tail) catch return; + r.emit(r.chunk(.{ .string = tail }, null) catch return) catch return; + } + if (calls != null and !reply.tool_calls_streamed) { + r.emit(r.chunk(roleDelta(a) catch return, null) catch return) catch return; + const deltas = tool_parse.deltas(a, calls.?) catch return; + for (deltas) |d| r.emit(r.chunk(d, null) catch return) catch return; + } else if (reply.content.len > 0 and !r.streamed_prose) { + r.emit(r.chunk(.{ .string = reply.content }, null) catch return) catch return; + } else if (r.plan.policy.kept(reply.content, r.prose_sent.items)) |rest| { + r.emit(r.chunk(.{ .string = rest }, null) catch return) catch return; + } + } + const last = r.chunk(.{ .string = "" }, if (reply.finish_reason.len > 0) reply.finish_reason else "length") catch return; + if (compaction) |s| compact.stamp(&cx, &reply, s) catch return; + r.extras(last.object, &reply) catch return; + const use = r.usage(&reply) catch return; + if (!r.plan.separate_usage) last.object.put(a, "usage", use) catch return; + r.emit(last) catch return; + if (r.plan.separate_usage) { + const u = r.chunk(null, null) catch return; + u.object.put(a, "choices", .{ .array = &.{} }) catch return; + u.object.put(a, "usage", use) catch return; + r.emit(u) catch return; + } + r.out.vt.event(r.out.ctx, null) catch {}; + r.warmNext(&reply, calls); + } +}; + +fn roleDelta(a: Allocator) Allocator.Error!Value { + return chat.deltaOf(a, "role", "assistant"); +} diff --git a/zig/src/server/prompt.zig b/zig/src/server/prompt.zig index e68363bfe7..f3fa2c9931 100644 --- a/zig/src/server/prompt.zig +++ b/zig/src/server/prompt.zig @@ -6,10 +6,11 @@ const messages_mod = @import("messages.zig"); const model_text = @import("model_text.zig"); const chat = @import("chat.zig"); const Server = @import("server.zig").Server; +const family_mod = @import("family.zig"); const Value = json.Value; const Cx = errors.Cx; -pub const Rendered = struct { ids: []const u32, history_len: usize = 0, rewind_len: usize = 0 }; +pub const Rendered = struct { ids: []const u32, history_len: usize = 0, rewind_len: usize = 0, images_omitted: u32 = 0, images: ?family_mod.Images = null }; /// Where lessons decide to act on a reply: the row before the generation prompt's ``, else before its last. pub fn decideOf(srv: *const Server, ids: []const u32) u32 { @@ -27,6 +28,11 @@ pub const isTitle = messages_mod.isTitleRequest; /// ``render_prompt_ids``: messages normalized for this template, rendered, and a dangling ```` closed. pub fn renderIds(srv: *Server, cx: *Cx, messages: Value, tools: []const Value, thinking: bool, effort: ?[]const u8, generation: bool) errors.Refused![]const u32 { + if (srv.family) |fam| { + const got = try fam.render(cx, .{ .messages = messages, .tools = tools, .thinking = thinking, .effort = effort orelse fam.defaultEffort(), .generation = generation }); + if (got.images) |im| im.release(); // ids only: the images are not held past the render + return got.ids; + } const normalized = try messages_mod.toolArguments(cx, try messages_mod.normalize(cx, messages, srv.late_system, srv.needs_user_after_tool)); var problem: []const u8 = ""; const options: model_text.RenderOptions = .{ @@ -96,6 +102,17 @@ pub fn prepare(srv: *Server, cx: *Cx, input: chat.Input, thinking: bool, effort: } }, .ids => |ids| return .{ .ids = ids }, }; + if (srv.family) |fam| { + const p: family_mod.Prompt = .{ .messages = input.messages, .tools = input.tools, .thinking = thinking, .effort = effort orelse fam.defaultEffort(), .body = input.body }; + const got = try fam.render(cx, p); + var history = p; + history.generation = false; + errdefer if (got.images) |im| im.release(); + const before = try fam.render(cx, history); + if (before.images) |im| im.release(); + const n = before.ids.len; + return .{ .ids = got.ids, .history_len = if (n > 0 and n < got.ids.len and std.mem.eql(u32, got.ids[0..n], before.ids)) n else 0, .images_omitted = got.images_omitted, .images = got.images }; + } const prompt = try renderIds(srv, cx, input.messages, input.tools, thinking, effort, true); const history = try renderIds(srv, cx, input.messages, input.tools, thinking, effort, false); var history_len: usize = 0; diff --git a/zig/src/server/root.zig b/zig/src/server/root.zig index ecc81c575a..441d11a144 100644 --- a/zig/src/server/root.zig +++ b/zig/src/server/root.zig @@ -12,6 +12,7 @@ pub const server = @import("server.zig"); pub const serve = @import("serve.zig"); pub const hub = @import("hub.zig"); pub const tool_specs = @import("tool_specs.zig"); +pub const family = @import("family.zig"); pub const Server = server.Server; pub const Config = server.Config; diff --git a/zig/src/server/serve.zig b/zig/src/server/serve.zig index 967c0979b1..a689aa3f88 100644 --- a/zig/src/server/serve.zig +++ b/zig/src/server/serve.zig @@ -73,6 +73,8 @@ pub const Setup = struct { started: i96 = 0, /// Called once the socket listens, with its port (tests read it). on_listen: ?*const fn (port: u16) void = null, + /// The checkpoint family's own chat rules (family.zig), or null. + family: ?@import("family.zig").Family = null, }; fn env(s: Setup, name: []const u8) ?[]const u8 { @@ -119,6 +121,7 @@ pub fn run(gpa: Allocator, io: std.Io, args: cli.Args, s: Setup) u8 { .seed_salt = salt, .request_log = env(s, "TENSORFOLD_REQUEST_LOG"), .dashboard = args.dashboard, + .family = s.family, .compact_at = if (args.compact_auto) .{ .auto = {} } else if (args.compact_fraction) |f| .{ .fraction = f } else null, .compact_keep = args.compact_keep, .compact_memory = args.compact_memory, @@ -154,6 +157,8 @@ pub fn run(gpa: Allocator, io: std.Io, args: cli.Args, s: Setup) u8 { s.served, args.host, port, shownSampling(a, s.sampling), if (args.no_drafts or srv.info.plain_only) "off" else "on", if (window > 0) std.fmt.bufPrint(&window_text, "{d}", .{window}) catch "?" else "unlimited", loaded, }); + if (greedy(s.sampling)) + log.line("WARNING: default sampling is greedy (temperature <= 0) - a client that omits `temperature` decodes greedily, which on a long agentic turn can loop in reasoning and return empty content", .{}); if (args.thinking and std.mem.indexOf(u8, s.text.templateSource(), "enable_thinking") != null) log.line("thinking on (the chat template's default): replies reason in reasoning_content before the answer in content, and max_tokens counts both. --no-thinking turns it off; a request can send chat_template_kwargs {{\"enable_thinking\": false}}", .{}); live.columns_env = env(s, "COLUMNS"); @@ -177,6 +182,17 @@ fn resolve(io: std.Io, host: []const u8, port: u16) ?std.Io.net.IpAddress { return std.Io.net.IpAddress.resolve(io, host, port) catch null; } +/// Whether requests that name no temperature decode greedily (the startup warning). +fn greedy(sampling: ?json.Value) bool { + const v = sampling orelse return true; + const t = v.get("temperature") orelse return true; + return switch (t) { + .float => |f| f <= 0, + .int => |i| (std.fmt.parseFloat(f64, i) catch 0) <= 0, + else => true, + }; +} + /// The serving line's sampling: greedy, or each default as ``name value``. fn shownSampling(a: Allocator, sampling: ?json.Value) []const u8 { const v = sampling orelse return "greedy"; diff --git a/zig/src/server/server.zig b/zig/src/server/server.zig index 255b02a2b6..9d6eaff3f7 100644 --- a/zig/src/server/server.zig +++ b/zig/src/server/server.zig @@ -22,6 +22,7 @@ const log = @import("log.zig"); const live = @import("live.zig"); const slide_graph = @import("slide_graph.zig"); const slide_lesson = @import("slide_lesson.zig"); +const family_mod = @import("family.zig"); const Value = json.Value; const Cx = errors.Cx; const Allocator = std.mem.Allocator; @@ -57,6 +58,8 @@ pub const Config = struct { /// --slide-graph: the file holding the Sliding Weights fact graph; null keeps the graph in memory. slide_graph: ?[]const u8 = null, timeouts: http_conn.Timeouts = .{}, + /// The checkpoint family's own chat rules (family.zig); null: the chat template path. + family: ?family_mod.Family = null, }; pub const Server = struct { @@ -77,6 +80,7 @@ pub const Server = struct { think_close_end: ?u32 = null, metrics: metrics.Metrics, store: responses.Store, + family: ?family_mod.Family = null, next_id: std.atomic.Value(u64) = .init(1), /// Connections open now; a stop waits for them before freeing what they read. open_connections: std.atomic.Value(u32) = .init(0), @@ -95,7 +99,7 @@ pub const Server = struct { /// Reads what the template and tokenizer decide once: late system role, think markers, efforts, the forced close. pub fn init(gpa: Allocator, io: std.Io, engine: api.Engine, text: model_text.Text, config: Config, keys: ?*auth.Store) !*Server { const srv = try gpa.create(Server); - srv.* = .{ .gpa = gpa, .io = io, .engine = engine, .info = engine.info(), .text = text, .config = config, .keys = keys, .metrics = .{ .gpa = gpa }, .store = .{ .gpa = gpa }, .arena = .init(gpa), .slide = .init(gpa, io, config.slide_graph), .teacher = .init(gpa) }; + srv.* = .{ .gpa = gpa, .io = io, .engine = engine, .info = engine.info(), .text = text, .config = config, .keys = keys, .metrics = .{ .gpa = gpa }, .store = .{ .gpa = gpa }, .arena = .init(gpa), .slide = .init(gpa, io, config.slide_graph), .teacher = .init(gpa), .family = config.family }; if (config.keep_warm_s > 0) if (engine.keepaliveTarget()) |target| { srv.keepalive = api.keepalive.Keepalive.start(gpa, io, target, @as(i64, config.keep_warm_s) * std.time.ns_per_s) catch |e| blk: { log.line("idle keepalive off: {s}", .{@errorName(e)}); @@ -145,7 +149,12 @@ pub const Server = struct { const options = try json.newObject(cx.a); const defaults = if (srv.config.default_sampling) |d| try fields.parseNumbers(cx, d) else null; if (defaults) |d| try merge(options, cx, d); - try merge(options, cx, try fields.parseNumbers(cx, f)); + const asked = try fields.parseNumbers(cx, f); + // a family whose Python server ignores the request's min_p keeps the default's + if (srv.family) |fam| if (!fam.vt.reads_min_p) { + _ = asked.object.orderedRemove("min_p"); + }; + try merge(options, cx, asked); const temp = number(options.get("temperature")) orelse 0; if (temp <= 0) return null; if (srv.info.greedy_only) return cx.refuse(if ((if (defaults) |d| number(d.get("temperature")) else null) orelse 0 > 0) diff --git a/zig/src/server/tokens.zig b/zig/src/server/tokens.zig index 3a2e351921..52bcb87092 100644 --- a/zig/src/server/tokens.zig +++ b/zig/src/server/tokens.zig @@ -27,7 +27,27 @@ fn tokenize(srv: *Server, cx: *Cx, body: Value) errors.Refused!Value { if (body != .object) return cx.refuse("the request body must be a JSON object"); const strings = try flag(cx, body, "return_token_strs", false); var ids: []const u32 = undefined; - if (body.has("messages")) { + if (body.has("messages") and srv.family != null) { + // rendered exactly as a chat request with the same body + const fam = srv.family.?; + const generation = try flag(cx, body, "add_generation_prompt", true); + const m = body.get("messages").?; + if (m != .array or m.array.len == 0) return cx.refuse("messages must be a non-empty list"); + const t = try fam.thinking(cx, body); + const rendered = fam.render(cx, .{ .messages = m, .tools = try fam.tools(cx, body), .thinking = t.enable orelse srv.config.enable_thinking, .effort = t.effort orelse srv.config.reasoning_effort orelse fam.defaultEffort(), .generation = generation, .body = body }) catch |e| switch (e) { + error.Refused => return cx.refuse(cx.message), + else => |x| return x, + }; + ids = rendered.ids; + if (rendered.images) |im| { // /tokenize shows the image token at each image position (expand(virtual=False)) + defer im.release(); + const out = try cx.a.dupe(u32, rendered.ids); + for (out) |*id| if (id.* >= 1 << 24) { + id.* = im.token; + }; + ids = out; + } + } else if (body.has("messages")) { const generation = try flag(cx, body, "add_generation_prompt", true); const msgs = try messages.normalize(cx, body.get("messages"), "system", srv.needs_user_after_tool); const tools = tool_specs.active(cx, body.get("tools"), body.get("tool_choice")) catch |e| switch (e) { From df480e0a8cbee6cbc49727e3df5a6b9765b923bb Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Sat, 10 Oct 2026 14:25:57 +0700 Subject: [PATCH 024/104] Server: a stop logs its signal and the open connections when the accept loop ends, and again if it exits with connections open. --- zig/src/server/serve.zig | 18 ++++++++++++++++-- 1 file changed, 16 insertions(+), 2 deletions(-) diff --git a/zig/src/server/serve.zig b/zig/src/server/serve.zig index a689aa3f88..4293da22ee 100644 --- a/zig/src/server/serve.zig +++ b/zig/src/server/serve.zig @@ -13,11 +13,20 @@ const Allocator = std.mem.Allocator; const posix = std.posix; var stop_requested: std.atomic.Value(bool) = .init(false); +/// The signal that asked for the stop (0: none yet), for the shutdown line +var stop_signal: std.atomic.Value(i32) = .init(0); -fn onStop(_: posix.SIG) callconv(.c) void { +fn onStop(sig: posix.SIG) callconv(.c) void { + stop_signal.store(@intCast(@backingInt(sig)), .release); stop_requested.store(true, .release); } +fn signalName(sig: i32) []const u8 { + if (sig == @backingInt(posix.SIG.TERM)) return "SIGTERM"; + if (sig == @backingInt(posix.SIG.INT)) return "SIGINT"; + return "a stop signal"; +} + /// SIGTERM and SIGINT stop the server cleanly; SIGHUP rereads the key file; SIGPIPE is ignored, as Python ignores it. pub fn installSignals(keys: bool) void { const stop: posix.Sigaction = .{ .handler = .{ .handler = onStop }, .mask = posix.sigemptyset(), .flags = 0 }; @@ -168,10 +177,15 @@ pub fn run(gpa: Allocator, io: std.Io, args: cli.Args, s: Setup) u8 { const accept = std.Thread.spawn(.{}, server_mod.Server.serve, .{ srv, lis, &stop_requested }) catch return 1; accept.join(); lis.close(); + // before the replies still running fail (their rounds' CUDA calls fail once the exit tears the driver down) + log.line("shutting down ({s}): no new requests; {d} connection(s) open, waiting up to 2 s before exiting", .{ signalName(stop_signal.load(.acquire)), srv.open_connections.load(.acquire) }); if (drawing) ticker.finish(); var waited: u32 = 0; while (srv.open_connections.load(.acquire) > 0 and waited < 40) : (waited += 1) std.Io.sleep(io, .fromMilliseconds(50), .awake) catch {}; - if (srv.open_connections.load(.acquire) > 0) std.process.exit(0); // replies still open: end without freeing what they read + if (srv.open_connections.load(.acquire) > 0) { + log.line("shutting down: exiting with {d} connection(s) still open; their requests fail", .{srv.open_connections.load(.acquire)}); + std.process.exit(0); // replies still open: end without freeing what they read + } return 0; } From 2c4d837bd01cba152ac66aae74c5eb8d635c6f07 Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Sat, 10 Oct 2026 14:28:48 +0700 Subject: [PATCH 025/104] Server: SIGTERM / SIGINT drains before stopping. New generation requests get 503 with Retry-After (code server_restarting; Anthropic overloaded_error) and /health answers "draining" without asking the engine. Requests in progress get TF_DRAIN_S seconds (default 20; a second signal cuts it short), then the engine's optional halt ends its rounds at a boundary and what is left fails with "server restarting: retry shortly". The engine closes before the server is freed, once. --- zig/src/server/drain_test.zig | 302 +++++++++++++++++++++++++++++++ zig/src/server/main.zig | 25 ++- zig/src/server/root.zig | 1 + zig/src/server/routes.zig | 46 ++++- zig/src/server/serve.zig | 65 ++++++- zig/src/server/server.zig | 4 + zig/src/server/status_routes.zig | 2 + 7 files changed, 434 insertions(+), 11 deletions(-) create mode 100644 zig/src/server/drain_test.zig diff --git a/zig/src/server/drain_test.zig b/zig/src/server/drain_test.zig new file mode 100644 index 0000000000..75106abcb9 --- /dev/null +++ b/zig/src/server/drain_test.zig @@ -0,0 +1,302 @@ +//! A stop's drain over HTTP: serve.run with a real SIGTERM, a slow stub engine, bytes as tokens. +const std = @import("std"); +const api = @import("engine_api"); +const json = @import("json.zig"); +const model_text = @import("model_text.zig"); +const serve = @import("serve.zig"); +const cli = @import("cli.zig"); +const routes = @import("routes.zig"); +const Allocator = std.mem.Allocator; +const testing = std.testing; + +/// Slow enough to stop mid-reply: "@short" ends after a dozen tokens, "@long" only when halted or cancelled. +const Slow = struct { + gpa: Allocator, + io: std.Io, + halted: std.atomic.Value(bool) = .init(false), + reason: []const u8 = "", + active: std.atomic.Value(u32) = .init(0), + halts: u32 = 0, + closes: u32 = 0, + /// halts before the first close (the halt must come first) + halts_at_close: u32 = 0, + cancelled: std.atomic.Value(u64) = .init(0), + + fn engine(s: *Slow) api.Engine { + return .{ .ctx = s, .vtable = &.{ .info = info, .submit = submit, .cancel = cancel, .status = status, .memory = memory } }; + } + + fn self(ctx: *anyopaque) *Slow { + return @ptrCast(@alignCast(ctx)); + } + + fn info(_: *anyopaque) api.Info { + return .{ .context_window = 4096, .name = "slow" }; + } + + fn cancel(ctx: *anyopaque, id: api.Id) void { + self(ctx).cancelled.store(id, .release); + } + + fn status(ctx: *anyopaque, out: *api.Status, _: []u32) void { + out.* = .{ .running = self(ctx).active.load(.acquire) }; + } + + fn memory(_: *anyopaque, _: bool) ?api.Memory { + return null; + } + + fn submit(ctx: *anyopaque, id: api.Id, r: *const api.Request, sink: api.Sink) api.SubmitError!void { + const s = self(ctx); + _ = s.active.fetchAdd(1, .acq_rel); + const t = std.Thread.spawn(.{}, run, .{ s, id, r, sink }) catch { + _ = s.active.fetchSub(1, .acq_rel); + return error.Busy; + }; + t.detach(); + } + + fn run(s: *Slow, id: api.Id, r: *const api.Request, sink: api.Sink) void { + defer _ = s.active.fetchSub(1, .acq_rel); + var arena = std.heap.ArenaAllocator.init(s.gpa); + defer arena.deinit(); + const a = arena.allocator(); + const prompt = a.alloc(u8, r.prompt.len) catch return sink.event(sink.ctx, id, &.{ .finished = .{ .reason = .failed, .message = "out of memory" } }); + for (r.prompt, prompt) |t, *b| b.* = @truncate(t); + const long = std.mem.indexOf(u8, prompt, "@long") != null; + const word: []const u32 = if (long) &.{ ' ', 'l', 'a' } else &.{ ' ', 'o', 'k' }; + sink.event(sink.ctx, id, &.{ .prefilled = 0 }); + var n: usize = 0; + while (true) : (n += 1) { + if (s.halted.load(.acquire)) return sink.event(sink.ctx, id, &.{ .finished = .{ .reason = .failed, .message = s.reason } }); + if (s.cancelled.load(.acquire) == id) return sink.event(sink.ctx, id, &.{ .finished = .{ .reason = .cancelled } }); + if (!long and n == 12) return sink.event(sink.ctx, id, &.{ .finished = .{ .reason = .stop } }); + sink.event(sink.ctx, id, &.{ .tokens = word }); + std.Io.sleep(s.io, .fromMilliseconds(15), .awake) catch {}; + } + } + + fn halt(ctx: *anyopaque, reason: []const u8) void { + const s = self(ctx); + s.reason = reason; + s.halts += 1; + s.halted.store(true, .release); + // a round boundary: every reply has ended before halt returns + while (s.active.load(.acquire) > 0) std.Io.sleep(s.io, .fromMilliseconds(1), .awake) catch {}; + } + + fn close(ctx: *anyopaque) void { + const s = self(ctx); + if (s.closes == 0) s.halts_at_close = s.halts; + s.closes += 1; + } +}; + +/// Bytes as tokens; a prompt renders as its messages' JSON, so the engine sees the client's marker. +const TestText = struct { + fn text(t: *@This()) model_text.Text { + return .{ .ctx = t, .vtable = &.{ .encode = encode, .decode = decode, .token_id = tokenId, .token_string = tokenString, .vocab_size = vocabSize, .eos_ids = eosIds, .render = render, .template_source = templateSource } }; + } + + fn encode(_: *anyopaque, a: Allocator, input: []const u8, _: bool) model_text.Error![]u32 { + const ids = try a.alloc(u32, input.len); + for (input, ids) |byte, *id| id.* = byte; + return ids; + } + + fn decode(_: *anyopaque, a: Allocator, ids: []const u32) Allocator.Error![]u8 { + const decoded = try a.alloc(u8, ids.len); + for (ids, decoded) |id, *byte| byte.* = @intCast(id); + return decoded; + } + + fn tokenId(_: *anyopaque, _: []const u8) ?u32 { + return null; + } + + fn tokenString(_: *anyopaque, a: Allocator, id: u32) Allocator.Error![]u8 { + return std.fmt.allocPrint(a, "{d}", .{id}); + } + + fn vocabSize(_: *anyopaque) u32 { + return 256; + } + + fn eosIds(_: *anyopaque) []const u32 { + return &.{}; + } + + fn render(_: *anyopaque, a: Allocator, messages: json.Value, _: model_text.RenderOptions, _: *[]const u8) model_text.Error![]u8 { + return json.stringify(a, messages, .{}) catch error.OutOfMemory; + } + + fn templateSource(_: *anyopaque) []const u8 { + return ""; + } +}; + +var listen_port: std.atomic.Value(u16) = .init(0); + +fn onListen(port: u16) void { + listen_port.store(port, .release); +} + +const Reply = struct { status: u16, head: []const u8, body: []const u8 }; + +/// One request on its own connection (closed by us after it), the whole response read. +fn request(a: Allocator, io: std.Io, port: u16, method: []const u8, path: []const u8, body: []const u8) !Reply { + const addr: std.Io.net.IpAddress = .{ .ip4 = .loopback(port) }; + const stream = try addr.connect(io, .{ .mode = .stream }); + defer stream.close(io); + const head = try std.fmt.allocPrint(a, "{s} {s} HTTP/1.1\r\nHost: x\r\nContent-Type: application/json\r\nContent-Length: {d}\r\nConnection: close\r\n\r\n{s}", .{ method, path, body.len, body }); + var wbuf: [4096]u8 = undefined; + var w = std.Io.net.Stream.Writer.init(stream, io, &wbuf); + try w.interface.writeAll(head); + try w.interface.flush(); + var rbuf: [16384]u8 = undefined; + var r = std.Io.net.Stream.Reader.init(stream, io, &rbuf); + const text = try r.interface.allocRemaining(a, .unlimited); + const split = std.mem.indexOf(u8, text, "\r\n\r\n") orelse return error.Head; + var payload: []const u8 = text[split + 4 ..]; + if (std.mem.indexOf(u8, text[0..split], "chunked") != null) payload = try unchunk(a, payload); + return .{ .status = try std.fmt.parseInt(u16, text[9..12], 10), .head = text[0..split], .body = payload }; +} + +fn unchunk(a: Allocator, s: []const u8) ![]const u8 { + var out: std.ArrayList(u8) = .empty; + var i: usize = 0; + while (i < s.len) { + const eol = std.mem.indexOfPos(u8, s, i, "\r\n") orelse break; + const n = try std.fmt.parseInt(usize, std.mem.trim(u8, s[i..eol], " "), 16); + if (n == 0) break; + try out.appendSlice(a, s[eol + 2 .. eol + 2 + n]); + i = eol + 2 + n + 2; + } + return out.items; +} + +const Background = struct { + a: Allocator, + io: std.Io, + port: u16, + body: []const u8, + reply: ?Reply = null, + failed: bool = false, + + fn run(b: *Background) void { + b.reply = request(b.a, b.io, b.port, "POST", "/v1/chat/completions", b.body) catch { + b.failed = true; + return; + }; + } +}; + +fn chatBody(a: Allocator, marker: []const u8, stream: bool) ![]const u8 { + return std.fmt.allocPrint(a, "{{\"messages\":[{{\"role\":\"user\",\"content\":\"go {s}\"}}],\"stream\":{},\"max_tokens\":4000}}", .{ marker, stream }); +} + +const Served = struct { + code: u8 = 255, + fn run(x: *Served, gpa: Allocator, io: std.Io, setup: serve.Setup) void { + x.code = serve.run(gpa, io, .{ .host = "127.0.0.1", .port = 0 }, setup); + } +}; + +/// serve.run on its own thread with `drain_s`, the engine `slow`; the port once it listens. +fn start(gpa: Allocator, io: std.Io, env: *std.process.Environ.Map, slow: *Slow, tt: *TestText, x: *Served) !std.Thread { + listen_port.store(0, .release); + const t = try std.Thread.spawn(.{}, Served.run, .{ x, gpa, io, serve.Setup{ + .engine = slow.engine(), + .text = tt.text(), + .served = "slow", + .environ = env, + .on_listen = onListen, + .stop = .{ .ctx = slow, .halt = Slow.halt, .close = Slow.close }, + } }); + while (listen_port.load(.acquire) == 0) std.Io.sleep(io, .fromMilliseconds(2), .awake) catch {}; + return t; +} + +fn waitActive(io: std.Io, slow: *Slow, n: u32) void { + while (slow.active.load(.acquire) < n) std.Io.sleep(io, .fromMilliseconds(2), .awake) catch {}; +} + +test "SIGTERM drains: 503 + Retry-After and /health draining while a reply finishes; at the deadline the rest end \"server restarting\" at a round boundary, then the engine closes" { + const gpa = testing.allocator; + const io = testing.io; + var tt: TestText = .{}; + var arena = std.heap.ArenaAllocator.init(gpa); + defer arena.deinit(); + const a = arena.allocator(); + var env = std.process.Environ.Map.init(gpa); + defer env.deinit(); + try env.put("TENSORFOLD_NO_LIVE", "1"); + + // 1. the reply in progress finishes within the drain: no halt, a whole reply; new work refused meanwhile + { + try env.put("TF_DRAIN_S", "10"); + var slow: Slow = .{ .gpa = gpa, .io = io }; + var x: Served = .{}; + const t = try start(gpa, io, &env, &slow, &tt, &x); + const port = listen_port.load(.acquire); + var bg: Background = .{ .a = a, .io = io, .port = port, .body = try chatBody(a, "@short", true) }; + const bt = try std.Thread.spawn(.{}, Background.run, .{&bg}); + waitActive(io, &slow, 1); + try std.posix.raise(.TERM); + // the drain has begun once a new request is refused + var refused: Reply = undefined; + while (true) { + refused = try request(a, io, port, "POST", "/v1/chat/completions", try chatBody(a, "@short", false)); + if (refused.status == 503) break; + try testing.expectEqual(@as(u16, 200), refused.status); // only before the signal lands + } + try testing.expect(std.mem.indexOf(u8, refused.head, "Retry-After: 30") != null); + try testing.expect(std.mem.indexOf(u8, refused.body, "\"code\": \"server_restarting\"") != null); + const anth = try request(a, io, port, "POST", "/v1/messages", "{\"model\":\"x\",\"max_tokens\":5,\"messages\":[{\"role\":\"user\",\"content\":\"hi\"}]}"); + try testing.expectEqual(@as(u16, 503), anth.status); + try testing.expect(std.mem.indexOf(u8, anth.body, "overloaded_error") != null); + const health = try request(a, io, port, "GET", "/health", ""); + try testing.expectEqual(@as(u16, 503), health.status); + try testing.expect(std.mem.indexOf(u8, health.head, "Retry-After: 30") != null); + try testing.expect(std.mem.indexOf(u8, health.body, "\"status\": \"draining\"") != null); + bt.join(); + t.join(); + try testing.expect(!bg.failed); + try testing.expectEqual(@as(u16, 200), bg.reply.?.status); + try testing.expect(std.mem.indexOf(u8, bg.reply.?.body, "\"finish_reason\": \"stop\"") != null); + try testing.expect(std.mem.indexOf(u8, bg.reply.?.body, "restarting") == null); + try testing.expect(std.mem.endsWith(u8, std.mem.trimEnd(u8, bg.reply.?.body, "\r\n"), "data: [DONE]")); + try testing.expectEqual(@as(u8, 0), x.code); + try testing.expectEqual(@as(u32, 0), slow.halts); + try testing.expectEqual(@as(u32, 1), slow.closes); + } + // 2. replies that outlast the drain end "server restarting"; the engine is halted, then closed once + { + try env.put("TF_DRAIN_S", "0.3"); + var slow: Slow = .{ .gpa = gpa, .io = io }; + var x: Served = .{}; + const t = try start(gpa, io, &env, &slow, &tt, &x); + const port = listen_port.load(.acquire); + var streamed: Background = .{ .a = a, .io = io, .port = port, .body = try chatBody(a, "@long", true) }; + var whole: Background = .{ .a = a, .io = io, .port = port, .body = try chatBody(a, "@long", false) }; + const st = try std.Thread.spawn(.{}, Background.run, .{&streamed}); + const wt = try std.Thread.spawn(.{}, Background.run, .{&whole}); + waitActive(io, &slow, 2); + try std.posix.raise(.TERM); + st.join(); + wt.join(); + t.join(); + try testing.expectEqual(@as(u8, 0), x.code); + try testing.expectEqual(@as(u32, 1), slow.halts); + try testing.expectEqual(@as(u32, 1), slow.closes); + try testing.expectEqual(@as(u32, 1), slow.halts_at_close); + try testing.expectEqualStrings(routes.restarting, slow.reason); + const sb = streamed.reply.?.body; + try testing.expectEqual(@as(u16, 200), streamed.reply.?.status); + try testing.expect(std.mem.indexOf(u8, sb, "\"type\": \"server_error\"") != null); + try testing.expect(std.mem.indexOf(u8, sb, routes.restarting) != null); + try testing.expect(std.mem.indexOf(u8, sb, "data: [DONE]") != null); + try testing.expect(whole.reply.?.status >= 500); + try testing.expect(std.mem.indexOf(u8, whole.reply.?.body, routes.restarting) != null); + } +} diff --git a/zig/src/server/main.zig b/zig/src/server/main.zig index 5ca1fd32f6..8ea644c331 100644 --- a/zig/src/server/main.zig +++ b/zig/src/server/main.zig @@ -75,8 +75,10 @@ pub fn main(init: std.process.Init) !u8 { return e; } orelse return fail(problem); defer up.text.deinit(); - defer up.engine.close(up.engine.ctx); + var closer: Closer = .{ .opened = up.engine }; + defer closer.closeOnce(); return serve.run(gpa, io, args, .{ + .stop = .{ .ctx = &closer, .halt = if (up.engine.halt != null) Closer.halt else null, .close = Closer.close }, .engine = up.engine.engine, .text = up.text.text(), .served = hub.servedName(args.name, args.model, dir), @@ -87,6 +89,27 @@ pub fn main(init: std.process.Init) !u8 { }); } +/// Closes once: serve.run closes the engine before freeing the server, main's defer covers paths that end before it. +const Closer = struct { + opened: engines.Opened, + closed: bool = false, + + fn closeOnce(c: *Closer) void { + if (c.closed) return; + c.closed = true; + c.opened.close(c.opened.ctx); + } + + fn close(ctx: *anyopaque) void { + closeOnce(@ptrCast(@alignCast(ctx))); + } + + fn halt(ctx: *anyopaque, reason: []const u8) void { + const c: *Closer = @ptrCast(@alignCast(ctx)); + if (c.opened.halt) |h| h(c.opened.ctx, reason); + } +}; + /// HfText.load with every failure named in `problem`, so main tells the text's failure from the engine's. fn loadText(gpa: std.mem.Allocator, io: std.Io, dir: []const u8, template: ?[]const u8, pa: std.mem.Allocator, problem: *[]const u8) !*hf_text.HfText { return hf_text.HfText.load(gpa, io, dir, template, pa, problem) catch |e| { diff --git a/zig/src/server/root.zig b/zig/src/server/root.zig index 441d11a144..301d627386 100644 --- a/zig/src/server/root.zig +++ b/zig/src/server/root.zig @@ -41,6 +41,7 @@ test { _ = @import("messages.zig"); _ = @import("stream_preflight_test.zig"); _ = @import("logprobs_test.zig"); + _ = @import("drain_test.zig"); _ = @import("late_system_test.zig"); _ = @import("compact_test.zig"); _ = @import("completion_prompt_test.zig"); diff --git a/zig/src/server/routes.zig b/zig/src/server/routes.zig index 44981cdda9..8fe21318e1 100644 --- a/zig/src/server/routes.zig +++ b/zig/src/server/routes.zig @@ -63,15 +63,51 @@ fn discardBody(conn: *Conn, a: Allocator) void { fn post(srv: *Server, conn: *Conn, a: Allocator) void { const route = auth.routePath(conn.path); - if (std.mem.eql(u8, route, "/v1/slide/learn")) return slide.learn(srv, conn, a); - if (std.mem.endsWith(u8, route, "/decisions")) return decisions.post(srv, conn, a); - if (anthropic.route(conn.path)) return anthropic.post(srv, conn, a); - if (responses.route(route)) |rid| if (rid.len == 0) return responses.post(srv, conn, a); for (tokens.paths) |r| if (std.mem.eql(u8, route, r)) return tokens.post(srv, conn, a, std.mem.endsWith(u8, route, "/detokenize")); + const is_anthropic = anthropic.route(conn.path); + const is_responses = if (responses.route(route)) |rid| rid.len == 0 else false; + const is_decisions = std.mem.endsWith(u8, route, "/decisions"); + const is_slide = std.mem.eql(u8, route, "/v1/slide/learn"); const is_chat = std.mem.endsWith(u8, route, "/chat/completions"); - if (!is_chat and !std.mem.endsWith(u8, route, "/completions")) { + if (!is_slide and !is_anthropic and !is_responses and !is_decisions and !is_chat and !std.mem.endsWith(u8, route, "/completions")) { discardBody(conn, a); return unknown(conn, a); } + // counted before the drain flag is read: a stop that sees no request in progress sees every later one refused + _ = srv.generating.fetchAdd(1, .seq_cst); + defer _ = srv.generating.fetchSub(1, .seq_cst); + if (srv.draining.load(.seq_cst)) { + discardBody(conn, a); + return draining(conn, is_anthropic); + } + if (is_slide) return slide.learn(srv, conn, a); + if (is_decisions) return decisions.post(srv, conn, a); + if (is_anthropic) return anthropic.post(srv, conn, a); + if (is_responses) return responses.post(srv, conn, a); openai.post(srv, conn, a, is_chat); } + +/// Seconds a refused client is told to wait while the server drains for a restart (Retry-After). +pub const retry_after_s = 30; +pub const restarting = "server restarting: retry shortly"; + +/// A request that arrives while a stop drains: 503, Retry-After, the API's own error shape; the connection closes. +pub fn draining(conn: *Conn, anthropic_shape: bool) void { + const body = if (anthropic_shape) + "{\"type\": \"error\", \"error\": {\"type\": \"overloaded_error\", \"message\": \"" ++ restarting ++ "\"}}" + else + "{\"error\": {\"message\": \"" ++ restarting ++ "\", \"type\": \"server_error\", \"param\": null, \"code\": \"server_restarting\"}}"; + sendDraining(conn, body); +} + +/// 503 with Retry-After and Connection: close, `body` as JSON. +pub fn sendDraining(conn: *Conn, body: []const u8) void { + conn.close = true; + conn.startResponse(503, null) catch return; + conn.addHeader("Content-Type", "application/json") catch return; + conn.addHeader("Retry-After", std.fmt.comptimePrint("{d}", .{retry_after_s})) catch return; + var len: [24]u8 = undefined; + conn.addHeader("Content-Length", std.fmt.bufPrint(&len, "{d}", .{body.len}) catch unreachable) catch return; + conn.addHeader("Connection", "close") catch return; + conn.finish(body) catch {}; +} diff --git a/zig/src/server/serve.zig b/zig/src/server/serve.zig index 4293da22ee..52aca924d4 100644 --- a/zig/src/server/serve.zig +++ b/zig/src/server/serve.zig @@ -9,15 +9,19 @@ const live = @import("live.zig"); const listener_mod = @import("listener.zig"); const model_text = @import("model_text.zig"); const server_mod = @import("server.zig"); +const routes = @import("routes.zig"); const Allocator = std.mem.Allocator; const posix = std.posix; var stop_requested: std.atomic.Value(bool) = .init(false); /// The signal that asked for the stop (0: none yet), for the shutdown line var stop_signal: std.atomic.Value(i32) = .init(0); +/// Stop signals so far: a second one cuts the drain short +var stop_count: std.atomic.Value(u32) = .init(0); fn onStop(sig: posix.SIG) callconv(.c) void { stop_signal.store(@intCast(@backingInt(sig)), .release); + _ = stop_count.fetchAdd(1, .acq_rel); stop_requested.store(true, .release); } @@ -84,6 +88,15 @@ pub const Setup = struct { on_listen: ?*const fn (port: u16) void = null, /// The checkpoint family's own chat rules (family.zig), or null. family: ?@import("family.zig").Family = null, + /// Run by `run` so the engine stops before the server is freed; null: the caller closes it after `run`. + stop: ?EngineStop = null, +}; + +/// `halt` (optional) ends the rounds at a boundary at the drain's deadline, so `close` never meets a running round. +pub const EngineStop = struct { + ctx: *anyopaque, + halt: ?*const fn (ctx: *anyopaque, reason: []const u8) void = null, + close: *const fn (ctx: *anyopaque) void, }; fn env(s: Setup, name: []const u8) ?[]const u8 { @@ -93,6 +106,9 @@ fn env(s: Setup, name: []const u8) ?[]const u8 { /// Serves until a stop signal; the process exit status. pub fn run(gpa: Allocator, io: std.Io, args: cli.Args, s: Setup) u8 { + stop_requested.store(false, .release); + stop_count.store(0, .release); + stop_signal.store(0, .release); var arena: std.heap.ArenaAllocator = .init(gpa); defer arena.deinit(); const a = arena.allocator(); @@ -115,6 +131,10 @@ pub fn run(gpa: Allocator, io: std.Io, args: cli.Args, s: Setup) u8 { std.debug.print("tensorfold: TENSORFOLD_SEED_SALT={s}: an integer\n", .{t}); return 1; } else 0; + const drain_s = seconds(env(s, "TF_DRAIN_S"), 20) orelse { + std.debug.print("tensorfold: TF_DRAIN_S: seconds, a number >= 0\n", .{}); + return 1; + }; const config: server_mod.Config = .{ .served_name = s.served, .model_dir = s.model_dir, @@ -174,21 +194,56 @@ pub fn run(gpa: Allocator, io: std.Io, args: cli.Args, s: Setup) u8 { var ticker: live.Ticker = .{ .engine = s.engine }; const drawing = live.wanted(io, if (env(s, "TENSORFOLD_NO_LIVE")) |v| std.mem.eql(u8, v, "1") else false); if (drawing) ticker.start(); - const accept = std.Thread.spawn(.{}, server_mod.Server.serve, .{ srv, lis, &stop_requested }) catch return 1; + // the listener stays open through the drain: new requests get 503 + Retry-After, /health "draining" + var accept_stop: std.atomic.Value(bool) = .init(false); + const accept = std.Thread.spawn(.{}, server_mod.Server.serve, .{ srv, lis, &accept_stop }) catch return 1; + while (!stop_requested.load(.acquire)) std.Io.sleep(io, .fromMilliseconds(50), .awake) catch {}; + drain(srv, io, s.stop, drain_s); + accept_stop.store(true, .release); accept.join(); lis.close(); - // before the replies still running fail (their rounds' CUDA calls fail once the exit tears the driver down) - log.line("shutting down ({s}): no new requests; {d} connection(s) open, waiting up to 2 s before exiting", .{ signalName(stop_signal.load(.acquire)), srv.open_connections.load(.acquire) }); if (drawing) ticker.finish(); + // the engine stops while nothing runs on it: no round meets a CUDA driver torn down by the exit + if (s.stop) |st| st.close(st.ctx); var waited: u32 = 0; while (srv.open_connections.load(.acquire) > 0 and waited < 40) : (waited += 1) std.Io.sleep(io, .fromMilliseconds(50), .awake) catch {}; if (srv.open_connections.load(.acquire) > 0) { - log.line("shutting down: exiting with {d} connection(s) still open; their requests fail", .{srv.open_connections.load(.acquire)}); - std.process.exit(0); // replies still open: end without freeing what they read + log.line("shutting down: exiting with {d} idle connection(s) still open", .{srv.open_connections.load(.acquire)}); + std.process.exit(0); // connections still open: end without freeing what they read } + log.line("shutting down: stopped", .{}); return 0; } +/// Work in progress gets `drain_s` to finish (new work 503s) before the rest is halted; a second signal cuts it short. +fn drain(srv: *server_mod.Server, io: std.Io, stop: ?EngineStop, drain_s: f64) void { + srv.draining.store(true, .seq_cst); + const signals = stop_count.load(.acquire); + log.line("shutting down ({s}): draining: new requests get 503; {d} request(s) in progress may finish within {d:.1} s", .{ signalName(stop_signal.load(.acquire)), srv.generating.load(.seq_cst), drain_s }); + const t0 = std.Io.Clock.awake.now(io).toNanoseconds(); + const deadline = t0 + @as(i96, @intFromFloat(drain_s * 1e9)); + while (srv.generating.load(.seq_cst) > 0 and std.Io.Clock.awake.now(io).toNanoseconds() < deadline and stop_count.load(.acquire) == signals) + std.Io.sleep(io, .fromMilliseconds(20), .awake) catch {}; + const left = srv.generating.load(.seq_cst); + if (left == 0) { + log.line("shutting down: drained", .{}); + return; + } + const cut = if (stop_count.load(.acquire) != signals) "a second stop signal" else "the drain deadline (TF_DRAIN_S)"; + log.line("shutting down: {s}: {d} request(s) still in progress end \"{s}\" at the next round boundary", .{ cut, left, routes.restarting }); + if (stop) |st| if (st.halt) |h| h(st.ctx, routes.restarting); + var waited: u32 = 0; + while (srv.generating.load(.seq_cst) > 0 and waited < 100) : (waited += 1) std.Io.sleep(io, .fromMilliseconds(20), .awake) catch {}; +} + +/// A seconds setting (TF_DRAIN_S): unset or blank keeps the default; null: not a number >= 0. +fn seconds(raw: ?[]const u8, default: f64) ?f64 { + const t = std.mem.trim(u8, raw orelse "", " \t"); + if (t.len == 0) return default; + const x = std.fmt.parseFloat(f64, t) catch return null; + return if (x >= 0) x else null; +} + /// ``--host`` as an address: an IP literal, else a name the resolver knows (localhost first). fn resolve(io: std.Io, host: []const u8, port: u16) ?std.Io.net.IpAddress { if (std.Io.net.IpAddress.parse(host, port)) |ip| return ip else |_| {} diff --git a/zig/src/server/server.zig b/zig/src/server/server.zig index 9d6eaff3f7..6129bb9602 100644 --- a/zig/src/server/server.zig +++ b/zig/src/server/server.zig @@ -84,6 +84,10 @@ pub const Server = struct { next_id: std.atomic.Value(u64) = .init(1), /// Connections open now; a stop waits for them before freeing what they read. open_connections: std.atomic.Value(u32) = .init(0), + /// A stop is draining (serve.zig): new generation requests get 503 with Retry-After, /health answers "draining" + draining: std.atomic.Value(bool) = .init(false), + /// Generation requests in progress (routes.post), what a drain waits for; counted before `draining` is read + generating: std.atomic.Value(u32) = .init(0), /// The idle keepalive on the engine's queue, armed by --keep-warm; null when off or not Metal. keepalive: ?*api.keepalive.Keepalive = null, preparing: std.atomic.Value(i64) = .init(0), diff --git a/zig/src/server/status_routes.zig b/zig/src/server/status_routes.zig index 89ee5c8488..9c1f93a07e 100644 --- a/zig/src/server/status_routes.zig +++ b/zig/src/server/status_routes.zig @@ -11,6 +11,8 @@ const Allocator = std.mem.Allocator; /// With keys on, health names nothing about the model. pub fn health(srv: *Server, conn: *Conn, a: Allocator) !void { + // draining: the engine is not asked, it may be stopping + if (srv.draining.load(.seq_cst)) return routes.sendDraining(conn, try std.fmt.allocPrint(a, "{{\"status\": \"draining\", \"requests_in_progress\": {d}}}", .{srv.generating.load(.seq_cst)})); if (srv.keys) |k| if (k.enabled()) return conn.sendJson(200, "{\"status\": \"ok\"}"); const o = try json.newObject(a); try o.put(a, "status", .{ .string = "ok" }); From 55bf0a787df02e42174d7bd837f13c0d1db1b02d Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Sat, 10 Oct 2026 14:29:01 +0700 Subject: [PATCH 026/104] Server tests: a stop that lands between a request's count and its drain check refuses the request (503) and the drain waits for it. A tests-only seam (routes.counted_hook) delivers SIGTERM in that gap; with the drain check moved before the count the test fails. --- zig/src/server/drain_test.zig | 62 +++++++++++++++++++++++++++++++++++ zig/src/server/routes.zig | 4 +++ 2 files changed, 66 insertions(+) diff --git a/zig/src/server/drain_test.zig b/zig/src/server/drain_test.zig index 75106abcb9..2bdc9f4a59 100644 --- a/zig/src/server/drain_test.zig +++ b/zig/src/server/drain_test.zig @@ -300,3 +300,65 @@ test "SIGTERM drains: 503 + Retry-After and /health draining while a reply finis try testing.expect(std.mem.indexOf(u8, whole.reply.?.body, routes.restarting) != null); } } + +/// The race seam's state: the stop delivered between a request's count and its drain check. +const Gap = struct { + var io: std.Io = undefined; + var port: u16 = 0; + var fired: std.atomic.Value(bool) = .init(false); + var saw_draining: bool = false; + var drain_waited: bool = false; + var served: *Served = undefined; + + fn hook() void { + if (fired.swap(true, .acq_rel)) return; + std.posix.raise(.TERM) catch return; + var arena = std.heap.ArenaAllocator.init(testing.allocator); + defer arena.deinit(); + // the drain has begun (it reads the count after it sets the flag): /health says so + var tries: u32 = 0; + while (tries < 500) : (tries += 1) { + const h = request(arena.allocator(), io, port, "GET", "/health", "") catch break; + if (h.status == 503 and std.mem.indexOf(u8, h.body, "draining") != null) { + saw_draining = true; + break; + } + std.Io.sleep(io, .fromMilliseconds(5), .awake) catch {}; + } + // this request is counted: the drain may not finish while it is between its count and its check + std.Io.sleep(io, .fromMilliseconds(300), .awake) catch {}; + drain_waited = served.code == 255; + } +}; + +test "a stop that lands between a request's count and its drain check: the request is refused (503) and the drain waited for it, never lost" { + const gpa = testing.allocator; + const io = testing.io; + var tt: TestText = .{}; + var arena = std.heap.ArenaAllocator.init(gpa); + defer arena.deinit(); + const a = arena.allocator(); + var env = std.process.Environ.Map.init(gpa); + defer env.deinit(); + try env.put("TENSORFOLD_NO_LIVE", "1"); + try env.put("TF_DRAIN_S", "10"); + var slow: Slow = .{ .gpa = gpa, .io = io }; + var x: Served = .{}; + const t = try start(gpa, io, &env, &slow, &tt, &x); + Gap.io = io; + Gap.port = listen_port.load(.acquire); + Gap.served = &x; + Gap.fired.store(false, .release); + routes.counted_hook = Gap.hook; + defer routes.counted_hook = null; + const r = try request(a, io, Gap.port, "POST", "/v1/chat/completions", try chatBody(a, "@short", false)); + t.join(); + try testing.expect(Gap.fired.load(.acquire)); + try testing.expect(Gap.saw_draining); + try testing.expect(Gap.drain_waited); + try testing.expectEqual(@as(u16, 503), r.status); + try testing.expect(std.mem.indexOf(u8, r.body, "server_restarting") != null); + try testing.expectEqual(@as(u32, 0), slow.closes -| 1); // closed once, after + try testing.expectEqual(@as(u32, 0), slow.active.load(.acquire)); // nothing reached the engine + try testing.expectEqual(@as(u8, 0), x.code); +} diff --git a/zig/src/server/routes.zig b/zig/src/server/routes.zig index 8fe21318e1..75b3b18a18 100644 --- a/zig/src/server/routes.zig +++ b/zig/src/server/routes.zig @@ -76,6 +76,7 @@ fn post(srv: *Server, conn: *Conn, a: Allocator) void { // counted before the drain flag is read: a stop that sees no request in progress sees every later one refused _ = srv.generating.fetchAdd(1, .seq_cst); defer _ = srv.generating.fetchSub(1, .seq_cst); + if (@import("builtin").is_test) if (counted_hook) |h| h(); if (srv.draining.load(.seq_cst)) { discardBody(conn, a); return draining(conn, is_anthropic); @@ -87,6 +88,9 @@ fn post(srv: *Server, conn: *Conn, a: Allocator) void { openai.post(srv, conn, a, is_chat); } +/// Tests only: the gap between a request's count and its drain check, where a stop must still see it counted. +pub var counted_hook: ?*const fn () void = null; + /// Seconds a refused client is told to wait while the server drains for a restart (Retry-After). pub const retry_after_s = 30; pub const restarting = "server restarting: retry shortly"; From 9fbc2b3edc2e772d28124dfea2da0a0be554de86 Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Sat, 10 Oct 2026 14:53:36 +0700 Subject: [PATCH 027/104] Server: a stream that has not sent its first token gets keepalives every TF_SSE_KEEPALIVE_S seconds (default 15). Queueing and a long prefill send nothing after the head, and clients or proxies that time out a silent body leave and retry from zero. Chat, completions and responses streams get an SSE comment (": keepalive"); Anthropic streams get its ping event. A keepalive that cannot be written cancels the request as a client that left. --- zig/src/server/anthropic.zig | 8 +++- zig/src/server/chat_generation.zig | 13 ++++++ zig/src/server/drain_test.zig | 29 ++++++++---- zig/src/server/keepalive_test.zig | 75 ++++++++++++++++++++++++++++++ zig/src/server/openai.zig | 21 ++++++++- zig/src/server/openai_run.zig | 5 +- zig/src/server/responses.zig | 7 ++- zig/src/server/root.zig | 1 + zig/src/server/serve.zig | 7 ++- zig/src/server/server.zig | 2 + zig/src/server/sse.zig | 5 ++ 11 files changed, 158 insertions(+), 15 deletions(-) create mode 100644 zig/src/server/keepalive_test.zig diff --git a/zig/src/server/anthropic.zig b/zig/src/server/anthropic.zig index 79113e569c..1dfc527985 100644 --- a/zig/src/server/anthropic.zig +++ b/zig/src/server/anthropic.zig @@ -105,7 +105,7 @@ const Wire = struct { payload: Value = .null, fn out(w: *Wire) openai.Out { - return .{ .ctx = w, .vt = &.{ .open = open, .event = event, .reply = whole } }; + return .{ .ctx = w, .vt = &.{ .open = open, .event = event, .reply = whole, .keepalive = ping } }; } fn open(ctx: *anyopaque) error{Closed}!void { @@ -125,6 +125,12 @@ const Wire = struct { w.status = status; w.payload = payload; } + + /// Anthropic's own keepalive event. + fn ping(ctx: *anyopaque) error{Closed}!void { + const w: *Wire = @ptrCast(@alignCast(ctx)); + return w.reply.send("ping", try w.reply.obj("ping")); + } }; const Call = struct { id: []const u8 = "", name: []const u8 = "", arguments: []const u8 = "", sent: usize = 0, done: bool = false }; diff --git a/zig/src/server/chat_generation.zig b/zig/src/server/chat_generation.zig index 635fc7ac97..905dcc489d 100644 --- a/zig/src/server/chat_generation.zig +++ b/zig/src/server/chat_generation.zig @@ -75,6 +75,8 @@ pub const Generation = struct { streaming_done: bool = false, first_ns: ?i96 = null, last_ns: ?i96 = null, + /// when the stream last carried bytes (a keepalive goes out after a silence) + wrote_ns: ?i96 = null, reason: ?[]const u8 = null, // the server ended the reply (a stop string, the length) before the engine said so engine_done: bool = false, consumed: usize = 0, @@ -166,6 +168,7 @@ pub const Generation = struct { g.cancel(); return error.Cancelled; }; + g.wrote_ns = nowNs(g.srv.io); } pub const cancel = chat.cancel; // in chat.zig with its doc @@ -205,6 +208,16 @@ pub const Generation = struct { g.cancel(); return error.Cancelled; } + if (chunk == null and @hasDecl(@TypeOf(gone), "beat")) { + // queueing and a long prefill send nothing: a client that times out a silent stream would leave + const now = nowNs(g.srv.io); + const since = g.wrote_ns orelse now; + if (g.wrote_ns == null) g.wrote_ns = now; + if (gone.beat(now - since) catch { + g.cancel(); + return error.Cancelled; + }) g.wrote_ns = now; + } const tokens = chunk orelse continue; if (g.reason == null) try g.commit(tokens); } diff --git a/zig/src/server/drain_test.zig b/zig/src/server/drain_test.zig index 2bdc9f4a59..411b651cf3 100644 --- a/zig/src/server/drain_test.zig +++ b/zig/src/server/drain_test.zig @@ -10,7 +10,7 @@ const Allocator = std.mem.Allocator; const testing = std.testing; /// Slow enough to stop mid-reply: "@short" ends after a dozen tokens, "@long" only when halted or cancelled. -const Slow = struct { +pub const Slow = struct { gpa: Allocator, io: std.Io, halted: std.atomic.Value(bool) = .init(false), @@ -65,6 +65,15 @@ const Slow = struct { for (r.prompt, prompt) |t, *b| b.* = @truncate(t); const long = std.mem.indexOf(u8, prompt, "@long") != null; const word: []const u32 = if (long) &.{ ' ', 'l', 'a' } else &.{ ' ', 'o', 'k' }; + const hang = std.mem.indexOf(u8, prompt, "@hang") != null; + const wait = std.mem.indexOf(u8, prompt, "@wait") != null; + // a long prefill: nothing for 800 ms ("@wait"), or nothing until stopped ("@hang") + var quiet: u32 = 0; + while (hang or (wait and quiet < 80)) : (quiet += 1) { + if (s.halted.load(.acquire)) return sink.event(sink.ctx, id, &.{ .finished = .{ .reason = .failed, .message = s.reason } }); + if (s.cancelled.load(.acquire) == id) return sink.event(sink.ctx, id, &.{ .finished = .{ .reason = .cancelled } }); + std.Io.sleep(s.io, .fromMilliseconds(10), .awake) catch {}; + } sink.event(sink.ctx, id, &.{ .prefilled = 0 }); var n: usize = 0; while (true) : (n += 1) { @@ -93,7 +102,7 @@ const Slow = struct { }; /// Bytes as tokens; a prompt renders as its messages' JSON, so the engine sees the client's marker. -const TestText = struct { +pub const TestText = struct { fn text(t: *@This()) model_text.Text { return .{ .ctx = t, .vtable = &.{ .encode = encode, .decode = decode, .token_id = tokenId, .token_string = tokenString, .vocab_size = vocabSize, .eos_ids = eosIds, .render = render, .template_source = templateSource } }; } @@ -135,16 +144,16 @@ const TestText = struct { } }; -var listen_port: std.atomic.Value(u16) = .init(0); +pub var listen_port: std.atomic.Value(u16) = .init(0); fn onListen(port: u16) void { listen_port.store(port, .release); } -const Reply = struct { status: u16, head: []const u8, body: []const u8 }; +pub const Reply = struct { status: u16, head: []const u8, body: []const u8 }; /// One request on its own connection (closed by us after it), the whole response read. -fn request(a: Allocator, io: std.Io, port: u16, method: []const u8, path: []const u8, body: []const u8) !Reply { +pub fn request(a: Allocator, io: std.Io, port: u16, method: []const u8, path: []const u8, body: []const u8) !Reply { const addr: std.Io.net.IpAddress = .{ .ip4 = .loopback(port) }; const stream = try addr.connect(io, .{ .mode = .stream }); defer stream.close(io); @@ -175,7 +184,7 @@ fn unchunk(a: Allocator, s: []const u8) ![]const u8 { return out.items; } -const Background = struct { +pub const Background = struct { a: Allocator, io: std.Io, port: u16, @@ -191,11 +200,11 @@ const Background = struct { } }; -fn chatBody(a: Allocator, marker: []const u8, stream: bool) ![]const u8 { +pub fn chatBody(a: Allocator, marker: []const u8, stream: bool) ![]const u8 { return std.fmt.allocPrint(a, "{{\"messages\":[{{\"role\":\"user\",\"content\":\"go {s}\"}}],\"stream\":{},\"max_tokens\":4000}}", .{ marker, stream }); } -const Served = struct { +pub const Served = struct { code: u8 = 255, fn run(x: *Served, gpa: Allocator, io: std.Io, setup: serve.Setup) void { x.code = serve.run(gpa, io, .{ .host = "127.0.0.1", .port = 0 }, setup); @@ -203,7 +212,7 @@ const Served = struct { }; /// serve.run on its own thread with `drain_s`, the engine `slow`; the port once it listens. -fn start(gpa: Allocator, io: std.Io, env: *std.process.Environ.Map, slow: *Slow, tt: *TestText, x: *Served) !std.Thread { +pub fn start(gpa: Allocator, io: std.Io, env: *std.process.Environ.Map, slow: *Slow, tt: *TestText, x: *Served) !std.Thread { listen_port.store(0, .release); const t = try std.Thread.spawn(.{}, Served.run, .{ x, gpa, io, serve.Setup{ .engine = slow.engine(), @@ -217,7 +226,7 @@ fn start(gpa: Allocator, io: std.Io, env: *std.process.Environ.Map, slow: *Slow, return t; } -fn waitActive(io: std.Io, slow: *Slow, n: u32) void { +pub fn waitActive(io: std.Io, slow: *Slow, n: u32) void { while (slow.active.load(.acquire) < n) std.Io.sleep(io, .fromMilliseconds(2), .awake) catch {}; } diff --git a/zig/src/server/keepalive_test.zig b/zig/src/server/keepalive_test.zig new file mode 100644 index 0000000000..43435c1d04 --- /dev/null +++ b/zig/src/server/keepalive_test.zig @@ -0,0 +1,75 @@ +//! A silent stream keeps its client: keepalives on both routes before the first token; a failed one ends the request. +const std = @import("std"); +const server = @import("server.zig"); +const d = @import("drain_test.zig"); +const testing = std.testing; + +/// A raw request whose response is read until the connection ends. +fn messagesBody(a: std.mem.Allocator, marker: []const u8) ![]const u8 { + return std.fmt.allocPrint(a, "{{\"model\":\"x\",\"max_tokens\":50,\"stream\":true,\"messages\":[{{\"role\":\"user\",\"content\":\"go {s}\"}}]}}", .{marker}); +} + +/// Sends `body` streamed and hangs up once the response head arrived and the request reached `slow`. +fn leave(a: std.mem.Allocator, io: std.Io, port: u16, body: []const u8, slow: *d.Slow) !void { + const addr: std.Io.net.IpAddress = .{ .ip4 = .loopback(port) }; + const stream = try addr.connect(io, .{ .mode = .stream }); + defer stream.close(io); + const head = try std.fmt.allocPrint(a, "POST /v1/chat/completions HTTP/1.1\r\nHost: x\r\nContent-Type: application/json\r\nContent-Length: {d}\r\n\r\n{s}", .{ body.len, body }); + var wbuf: [4096]u8 = undefined; + var w = std.Io.net.Stream.Writer.init(stream, io, &wbuf); + try w.interface.writeAll(head); + try w.interface.flush(); + var rbuf: [64]u8 = undefined; + var r = std.Io.net.Stream.Reader.init(stream, io, &rbuf); + _ = try r.interface.peek(1); + // the head goes out before the engine has the request: a hang-up before that leaves nothing to cancel + d.waitActive(io, slow, 1); +} + +test "a silent engine's stream gets keepalives before its first token (chat: SSE comments, messages: ping events)" { + const period = comptime std.meta.fieldInfo(server.Config, .sse_keepalive_s).attrs.defaultValue(f64).?; + try testing.expect(period > 0 and period <= 15); + const gpa = testing.allocator; + const io = testing.io; + var tt: d.TestText = .{}; + var arena = std.heap.ArenaAllocator.init(gpa); + defer arena.deinit(); + const a = arena.allocator(); + var env = std.process.Environ.Map.init(gpa); + defer env.deinit(); + try env.put("TENSORFOLD_NO_LIVE", "1"); + try env.put("TF_DRAIN_S", "5"); + try env.put("TF_SSE_KEEPALIVE_S", "0.1"); + var slow: d.Slow = .{ .gpa = gpa, .io = io }; + var x: d.Served = .{}; + const t = try d.start(gpa, io, &env, &slow, &tt, &x); + const port = d.listen_port.load(.acquire); + + const chat = try d.request(a, io, port, "POST", "/v1/chat/completions", try d.chatBody(a, "@wait", true)); + try testing.expectEqual(@as(u16, 200), chat.status); + const first_token = std.mem.indexOf(u8, chat.body, "\" ok\"") orelse { + std.debug.print("{s}\n", .{chat.body[0..@min(chat.body.len, 1500)]}); + return error.NoToken; + }; + try testing.expect(std.mem.count(u8, chat.body[0..first_token], ": keepalive\n\n") >= 3); + try testing.expect(std.mem.indexOf(u8, chat.body, "data: [DONE]") != null); + + const msg = try d.request(a, io, port, "POST", "/v1/messages", try messagesBody(a, "@wait")); + try testing.expectEqual(@as(u16, 200), msg.status); + const start = std.mem.indexOf(u8, msg.body, "event: message_start") orelse return error.NoStart; + const ping = std.mem.indexOf(u8, msg.body, "event: ping\ndata: {\"type\": \"ping\"}") orelse return error.NoPing; + const delta = std.mem.indexOf(u8, msg.body, "event: content_block_delta") orelse return error.NoDelta; + try testing.expect(start < ping and ping < delta); + + // a client that leaves a silent stream: the request is cancelled (the disconnect poll or a failed keepalive) + var idle: u32 = 0; + while (slow.active.load(.acquire) > 0 and idle < 300) : (idle += 1) std.Io.sleep(io, .fromMilliseconds(10), .awake) catch {}; + try leave(a, io, port, try d.chatBody(a, "@hang", true), &slow); + var waited: u32 = 0; + while (slow.cancelled.load(.acquire) == 0 and waited < 300) : (waited += 1) std.Io.sleep(io, .fromMilliseconds(10), .awake) catch {}; + try testing.expect(slow.cancelled.load(.acquire) != 0); + + try std.posix.raise(.TERM); + t.join(); + try testing.expectEqual(@as(u8, 0), x.code); +} diff --git a/zig/src/server/openai.zig b/zig/src/server/openai.zig index 4f0f463018..bfdc5658ea 100644 --- a/zig/src/server/openai.zig +++ b/zig/src/server/openai.zig @@ -37,6 +37,8 @@ pub const Out = struct { event: *const fn (ctx: *anyopaque, payload: ?Value) error{Closed}!void, /// A whole JSON reply. reply: *const fn (ctx: *anyopaque, status: u16, payload: Value) void, + /// Bytes that keep a silent stream alive (an SSE comment, or the API's ping); null: none. + keepalive: ?*const fn (ctx: *anyopaque) error{Closed}!void = null, }; }; @@ -57,7 +59,7 @@ const ClientOut = struct { a: Allocator, fn out(c: *ClientOut) Out { - return .{ .ctx = c, .vt = &.{ .open = open, .event = event, .reply = reply } }; + return .{ .ctx = c, .vt = &.{ .open = open, .event = event, .reply = reply, .keepalive = keepalive } }; } fn open(ctx: *anyopaque) error{Closed}!void { @@ -74,14 +76,31 @@ const ClientOut = struct { const c: *ClientOut = @ptrCast(@alignCast(ctx)); routes.sendValue(c.conn, c.a, status, payload); } + + fn keepalive(ctx: *anyopaque) error{Closed}!void { + const c: *ClientOut = @ptrCast(@alignCast(ctx)); + return sse.comment(c.conn); + } }; /// Whether the client has gone, read between rounds (Python's socket_cancellation). pub const Gone = struct { conn: *Conn, + /// The stream a keepalive goes to once it has been silent `every_ns` (0: never) + out: ?Out = null, + every_ns: i96 = 0, pub fn check(g: Gone) bool { return g.conn.peerGone(); } + + /// A keepalive if the stream has been silent long enough; a failed write means the client left. + pub fn beat(g: Gone, silent_ns: i96) error{Closed}!bool { + const o = g.out orelse return false; + const send = o.vt.keepalive orelse return false; + if (g.every_ns <= 0 or silent_ns < g.every_ns) return false; + try send(o.ctx); + return true; + } }; /// ``error_body``: OpenAI's error object, with param and code where clients key on them. diff --git a/zig/src/server/openai_run.zig b/zig/src/server/openai_run.zig index 3f9978259d..15d0785cf3 100644 --- a/zig/src/server/openai_run.zig +++ b/zig/src/server/openai_run.zig @@ -266,7 +266,10 @@ pub const Run = struct { if (!tools and r.is_chat) r.emit(r.chunk(roleDelta(a) catch return, null) catch return) catch return; var sink_state: StreamSink = .{ .run = r, .tools = tools }; handed = true; - var reply = chat.generate(r.srv, &cx, prepared, .{ .ctx = &sink_state, .call = StreamSink.call }, gone) catch |e| { + var live = gone; + live.out = r.out; + live.every_ns = @intFromFloat(r.srv.config.sse_keepalive_s * std.time.ns_per_s); + var reply = chat.generate(r.srv, &cx, prepared, .{ .ctx = &sink_state, .call = StreamSink.call }, live) catch |e| { switch (e) { error.Cancelled => return, error.Refused => if (cx.kind != .other and cx.kind != .server) { diff --git a/zig/src/server/responses.zig b/zig/src/server/responses.zig index 309aade433..1bca9c72fa 100644 --- a/zig/src/server/responses.zig +++ b/zig/src/server/responses.zig @@ -203,7 +203,7 @@ const Wire = struct { payload: Value = .null, fn out(w: *Wire) openai.Out { - return .{ .ctx = w, .vt = &.{ .open = open, .event = event, .reply = whole } }; + return .{ .ctx = w, .vt = &.{ .open = open, .event = event, .reply = whole, .keepalive = keepalive } }; } fn open(ctx: *anyopaque) error{Closed}!void { @@ -223,6 +223,11 @@ const Wire = struct { w.status = status; w.payload = payload; } + + fn keepalive(ctx: *anyopaque) error{Closed}!void { + const w: *Wire = @ptrCast(@alignCast(ctx)); + return sse.comment(w.reply.conn); + } }; /// A Response built from chat-completion deltas, in order: reasoning, text, calls; streamed as typed events. diff --git a/zig/src/server/root.zig b/zig/src/server/root.zig index 301d627386..0c2f28c41e 100644 --- a/zig/src/server/root.zig +++ b/zig/src/server/root.zig @@ -42,6 +42,7 @@ test { _ = @import("stream_preflight_test.zig"); _ = @import("logprobs_test.zig"); _ = @import("drain_test.zig"); + _ = @import("keepalive_test.zig"); _ = @import("late_system_test.zig"); _ = @import("compact_test.zig"); _ = @import("completion_prompt_test.zig"); diff --git a/zig/src/server/serve.zig b/zig/src/server/serve.zig index 52aca924d4..1574ad697b 100644 --- a/zig/src/server/serve.zig +++ b/zig/src/server/serve.zig @@ -135,6 +135,10 @@ pub fn run(gpa: Allocator, io: std.Io, args: cli.Args, s: Setup) u8 { std.debug.print("tensorfold: TF_DRAIN_S: seconds, a number >= 0\n", .{}); return 1; }; + const keepalive_s = seconds(env(s, "TF_SSE_KEEPALIVE_S"), 15) orelse { + std.debug.print("tensorfold: TF_SSE_KEEPALIVE_S: seconds, a number >= 0\n", .{}); + return 1; + }; const config: server_mod.Config = .{ .served_name = s.served, .model_dir = s.model_dir, @@ -148,6 +152,7 @@ pub fn run(gpa: Allocator, io: std.Io, args: cli.Args, s: Setup) u8 { .default_sampling = s.sampling, .use_drafts = !args.no_drafts, .seed_salt = salt, + .sse_keepalive_s = keepalive_s, .request_log = env(s, "TENSORFOLD_REQUEST_LOG"), .dashboard = args.dashboard, .family = s.family, @@ -236,7 +241,7 @@ fn drain(srv: *server_mod.Server, io: std.Io, stop: ?EngineStop, drain_s: f64) v while (srv.generating.load(.seq_cst) > 0 and waited < 100) : (waited += 1) std.Io.sleep(io, .fromMilliseconds(20), .awake) catch {}; } -/// A seconds setting (TF_DRAIN_S): unset or blank keeps the default; null: not a number >= 0. +/// A seconds setting (TF_DRAIN_S, TF_SSE_KEEPALIVE_S): unset or blank keeps the default; null: not a number >= 0. fn seconds(raw: ?[]const u8, default: f64) ?f64 { const t = std.mem.trim(u8, raw orelse "", " \t"); if (t.len == 0) return default; diff --git a/zig/src/server/server.zig b/zig/src/server/server.zig index 6129bb9602..0467c4ee0c 100644 --- a/zig/src/server/server.zig +++ b/zig/src/server/server.zig @@ -58,6 +58,8 @@ pub const Config = struct { /// --slide-graph: the file holding the Sliding Weights fact graph; null keeps the graph in memory. slide_graph: ?[]const u8 = null, timeouts: http_conn.Timeouts = .{}, + /// TF_SSE_KEEPALIVE_S: a silent stream's keepalive period before its first token (serve.zig; 0: none) + sse_keepalive_s: f64 = 15, /// The checkpoint family's own chat rules (family.zig); null: the chat template path. family: ?family_mod.Family = null, }; diff --git a/zig/src/server/sse.zig b/zig/src/server/sse.zig index de30ac1212..bcdcf56f4c 100644 --- a/zig/src/server/sse.zig +++ b/zig/src/server/sse.zig @@ -19,6 +19,11 @@ pub fn data(conn: *Conn, a: Allocator, payload: json.Value) error{Closed}!void { return conn.writeAll(std.mem.concat(a, u8, &.{ "data: ", text, "\n\n" }) catch return error.Closed); } +/// An SSE comment line, which clients skip: it keeps a stream alive while nothing else is sent. +pub fn comment(conn: *Conn) error{Closed}!void { + return conn.writeAll(": keepalive\n\n"); +} + pub fn done(conn: *Conn) error{Closed}!void { return conn.writeAll("data: [DONE]\n\n"); } From bb53ff4e95f76c632d5b9cafc08fec7a5fd2d4c6 Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Sat, 10 Oct 2026 16:10:45 +0700 Subject: [PATCH 028/104] Server tests: the drain test's reply in progress is held until its drain checks are done The listener stays open only while a generation runs. Case 1 streamed a "@short" reply (about 180 ms) and then sent a messages request and a /health request expecting 503; when the reply ended first, the drain finished and those requests found the listener closed. The stub engine's "@hold" reply now waits for the test's release. --- zig/src/server/drain_test.zig | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/zig/src/server/drain_test.zig b/zig/src/server/drain_test.zig index 411b651cf3..09873e92f5 100644 --- a/zig/src/server/drain_test.zig +++ b/zig/src/server/drain_test.zig @@ -9,7 +9,7 @@ const routes = @import("routes.zig"); const Allocator = std.mem.Allocator; const testing = std.testing; -/// Slow enough to stop mid-reply: "@short" ends after a dozen tokens, "@long" only when halted or cancelled. +/// Stops mid-reply: "@short" ends after 12 tokens, "@long" when halted or cancelled, "@hold" waits for `release`. pub const Slow = struct { gpa: Allocator, io: std.Io, @@ -21,6 +21,8 @@ pub const Slow = struct { /// halts before the first close (the halt must come first) halts_at_close: u32 = 0, cancelled: std.atomic.Value(u64) = .init(0), + /// "@hold" replies stay silent until this is set + release: std.atomic.Value(bool) = .init(false), fn engine(s: *Slow) api.Engine { return .{ .ctx = s, .vtable = &.{ .info = info, .submit = submit, .cancel = cancel, .status = status, .memory = memory } }; @@ -67,9 +69,10 @@ pub const Slow = struct { const word: []const u32 = if (long) &.{ ' ', 'l', 'a' } else &.{ ' ', 'o', 'k' }; const hang = std.mem.indexOf(u8, prompt, "@hang") != null; const wait = std.mem.indexOf(u8, prompt, "@wait") != null; + const hold = std.mem.indexOf(u8, prompt, "@hold") != null; // a long prefill: nothing for 800 ms ("@wait"), or nothing until stopped ("@hang") var quiet: u32 = 0; - while (hang or (wait and quiet < 80)) : (quiet += 1) { + while (hang or (wait and quiet < 80) or (hold and !s.release.load(.acquire))) : (quiet += 1) { if (s.halted.load(.acquire)) return sink.event(sink.ctx, id, &.{ .finished = .{ .reason = .failed, .message = s.reason } }); if (s.cancelled.load(.acquire) == id) return sink.event(sink.ctx, id, &.{ .finished = .{ .reason = .cancelled } }); std.Io.sleep(s.io, .fromMilliseconds(10), .awake) catch {}; @@ -248,8 +251,10 @@ test "SIGTERM drains: 503 + Retry-After and /health draining while a reply finis var x: Served = .{}; const t = try start(gpa, io, &env, &slow, &tt, &x); const port = listen_port.load(.acquire); - var bg: Background = .{ .a = a, .io = io, .port = port, .body = try chatBody(a, "@short", true) }; + // held through the checks below: a reply that ended first would end the drain and close the listener + var bg: Background = .{ .a = a, .io = io, .port = port, .body = try chatBody(a, "@hold", true) }; const bt = try std.Thread.spawn(.{}, Background.run, .{&bg}); + errdefer slow.release.store(true, .release); waitActive(io, &slow, 1); try std.posix.raise(.TERM); // the drain has begun once a new request is refused @@ -268,6 +273,7 @@ test "SIGTERM drains: 503 + Retry-After and /health draining while a reply finis try testing.expectEqual(@as(u16, 503), health.status); try testing.expect(std.mem.indexOf(u8, health.head, "Retry-After: 30") != null); try testing.expect(std.mem.indexOf(u8, health.body, "\"status\": \"draining\"") != null); + slow.release.store(true, .release); bt.join(); t.join(); try testing.expect(!bg.failed); From 9eb765d490c4537d2ab8ca64d2801d0425d68bd0 Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Sun, 11 Oct 2026 19:31:50 +0700 Subject: [PATCH 029/104] DeepSeek-V4.1 host side: release config, EXL3 pack headers, a rank's TP weight plan and its named tensors. zig/src/families/deepseek_v41 starts with the parts that need no GPU: the text config (config.json's text_config, with the per-layer CSA2 topology), EXL3 group geometry and a rank's TP slices, the pack's tensor names and shapes, the pack folder read from its safetensors headers alone, the rank's weight plan, its host images and the weights as named tensors in the reference loader's forms. zig build test-dsv41 runs the host tests on the release config, a synthetic pack and a 2.8 bpw pack's expert widths; TF_DSV41_HEADERS adds real header dumps. --- build.zig | 2 + zig/build/dsv41.zig | 24 ++ zig/src/families/deepseek_v41/config.zig | 340 ++++++++++++++++++ .../families/deepseek_v41/deepseek_v41.zig | 23 ++ zig/src/families/deepseek_v41/exl3.zig | 254 +++++++++++++ .../deepseek_v41/fixtures/config.json | 184 ++++++++++ .../deepseek_v41/fixtures/widths-exl3-2p8.txt | 44 +++ zig/src/families/deepseek_v41/host_test.zig | 330 +++++++++++++++++ zig/src/families/deepseek_v41/named.zig | 261 ++++++++++++++ zig/src/families/deepseek_v41/names.zig | 236 ++++++++++++ zig/src/families/deepseek_v41/pack.zig | 306 ++++++++++++++++ zig/src/families/deepseek_v41/plan.zig | 295 +++++++++++++++ zig/src/families/deepseek_v41/stage.zig | 148 ++++++++ 13 files changed, 2447 insertions(+) create mode 100644 zig/build/dsv41.zig create mode 100644 zig/src/families/deepseek_v41/config.zig create mode 100644 zig/src/families/deepseek_v41/deepseek_v41.zig create mode 100644 zig/src/families/deepseek_v41/exl3.zig create mode 100644 zig/src/families/deepseek_v41/fixtures/config.json create mode 100644 zig/src/families/deepseek_v41/fixtures/widths-exl3-2p8.txt create mode 100644 zig/src/families/deepseek_v41/host_test.zig create mode 100644 zig/src/families/deepseek_v41/named.zig create mode 100644 zig/src/families/deepseek_v41/names.zig create mode 100644 zig/src/families/deepseek_v41/pack.zig create mode 100644 zig/src/families/deepseek_v41/plan.zig create mode 100644 zig/src/families/deepseek_v41/stage.zig diff --git a/build.zig b/build.zig index bd0476b368..8f8ff87e93 100644 --- a/build.zig +++ b/build.zig @@ -7,6 +7,7 @@ const hip_build = @import("zig/build/hip.zig"); const grammar_build = @import("zig/build/grammar.zig"); const tp_build = @import("zig/build/tp.zig"); const sessions_build = @import("zig/build/sessions.zig"); +const dsv41_build = @import("zig/build/dsv41.zig"); comptime { const required = std.mem.trim(u8, @embedFile(".zig-version"), "\r\n"); @@ -179,6 +180,7 @@ pub fn build(b: *std.Build) void { grammar_build.steps(b, target, xgrammar); tp_build.add(b, target, optimize, test_step); sessions_build.add(b, test_step); + dsv41_build.add(b, test_step); } /// `zig build native -Dcpu=apple_m1`: tensorfold-native with the Metal engines for the Python package's bundle (a native M5 build traps on M1-M4). diff --git a/zig/build/dsv41.zig b/zig/build/dsv41.zig new file mode 100644 index 0000000000..bdedccec4c --- /dev/null +++ b/zig/build/dsv41.zig @@ -0,0 +1,24 @@ +//! DeepSeek-V4.1's build: the family's host side (zig/src/families/deepseek_v41) and its host tests on any OS. + +const std = @import("std"); + +/// The backend-neutral core (the safetensors reader the pack headers use) for the family's host side. +fn core(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode) *std.Build.Module { + const tokenizer = b.createModule(.{ .root_source_file = b.path("zig/src/core/tokenizer/tokenizer.zig"), .target = target, .optimize = optimize, .link_libc = true }); + const m = b.createModule(.{ .root_source_file = b.path("zig/src/core/root.zig"), .target = target, .optimize = optimize, .link_libc = true }); + m.addImport("tokenizer", tokenizer); + return m; +} + +/// The family's host side (config, pack headers, a rank's weight plan) as the `deepseek_v41` module. +pub fn module(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode) *std.Build.Module { + const m = b.createModule(.{ .root_source_file = b.path("zig/src/families/deepseek_v41/deepseek_v41.zig"), .target = target, .optimize = optimize, .link_libc = true }); + m.addImport("core", core(b, target, optimize)); + return m; +} + +pub fn add(b: *std.Build, test_step: *std.Build.Step) void { + const run = b.addRunArtifact(b.addTest(.{ .root_module = module(b, b.graph.host, .debug) })); + b.step("test-dsv41", "DeepSeek-V4.1's host-side unit tests (no GPU work)").dependOn(&run.step); + test_step.dependOn(&run.step); +} diff --git a/zig/src/families/deepseek_v41/config.zig b/zig/src/families/deepseek_v41/config.zig new file mode 100644 index 0000000000..85e1401978 --- /dev/null +++ b/zig/src/families/deepseek_v41/config.zig @@ -0,0 +1,340 @@ +//! DeepSeek-V4.1-Flash text config (config.json -> text_config) and the per-layer CSA2 topology. + +const std = @import("std"); + +/// Fixed-capacity id lists of the config (layers, sources, Engram tables); the release's longest has 43 entries. +pub fn List(comptime T: type, comptime cap: usize) type { + return struct { + const Self = @This(); + buf: [cap]T = undefined, + len: usize = 0, + + pub fn of(values: []const T) Self { + var s: Self = .{}; + for (values) |v| s.push(v) catch unreachable; + return s; + } + + pub fn push(s: *Self, v: T) !void { + if (s.len == cap) return error.ConfigListTooLong; + s.buf[s.len] = v; + s.len += 1; + } + + pub fn items(s: *const Self) []const T { + return s.buf[0..s.len]; + } + + pub fn contains(s: *const Self, v: T) bool { + return std.mem.indexOfScalar(T, s.items(), v) != null; + } + }; +} + +pub const Ids = List(u32, 64); + +pub const Rope = struct { + theta: f64 = 10000.0, + compress_theta: f64 = 160000.0, + factor: f64 = 16.0, + original_max_positions: u64 = 65536, + beta_fast: f64 = 32.0, + beta_slow: f64 = 1.0, + max_positions: u64 = 1048576, + yarn: bool = true, +}; + +/// The attention a layer runs: sliding window only, KV source + indexer, own indexer on a source's keys, or reuse. +pub const Mode = enum { swa, full, reindex, reuse }; + +pub const Config = struct { + vocab: u32 = 129280, + hidden: u32 = 5120, + layers: u32 = 40, + heads: u32 = 64, + head_dim: u32 = 512, + rope_dim: u32 = 64, + q_lora: u32 = 1280, + o_lora: u32 = 1024, + o_groups: u32 = 8, + eps: f32 = 1e-20, + window: u32 = 128, + experts: u32 = 384, + shared_experts: u32 = 1, + topk: u32 = 6, + expert_width: u32 = 2304, + routed_scale: f32 = 1.5, + norm_topk: bool = true, + swiglu_limit: f32 = 10.0, + compress_ratios: Ids = Ids.of(&.{ 0, 0, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0 }), + kv_sources: Ids = Ids.of(&.{ 2, 8, 14, 20 }), + index_sources: Ids = Ids.of(&.{ 2, 8, 14, 20, 24, 28, 32, 36 }), + index_heads: u32 = 32, + index_dim: u32 = 128, + index_topk: u32 = 512, + candidate_source: i64 = 20, + candidate_blocks: u32 = 2048, + candidate_block_size: u32 = 8, + hc_mult: u32 = 4, + hc_sinkhorn_iters: u32 = 20, + hc_eps: f32 = 1e-6, + hc_post_alpha: f32 = 2.0, + engram_layers: Ids = Ids.of(&.{ 1, 14 }), + engram_rows: List(u64, 8) = List(u64, 8).of(&.{ 384006168, 384016682 }), + engram_max_ngram: u32 = 4, + engram_vocab: u64 = 16000000, + engram_heads: u32 = 8, + engram_head_dim: u32 = 256, + engram_pad: u32 = 2, + engram_compressed_vocab: u32 = 99092, + engram_gate_clamp: f32 = 1e-6, + mtp_layers: u32 = 3, + dspark_block: u32 = 5, + dspark_noise_token: u32 = 128799, + dspark_targets: Ids = Ids.of(&.{ 37, 38, 39 }), + dspark_markov_rank: u32 = 256, + dspark_experts: u32 = 128, + dspark_topk: u32 = 3, + rope: Rope = .{}, + bos: u32 = 0, + eos: u32 = 1, + + /// config.json in `dir`. + pub fn read(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) !Config { + const path = try std.fs.path.join(gpa, &.{ dir, "config.json" }); + defer gpa.free(path); + const bytes = try std.Io.Dir.cwd().readFileAlloc(io, path, gpa, .limited(1 << 24)); + defer gpa.free(bytes); + return parse(gpa, bytes); + } + + /// Checkpoint JSON (text_config, else top level); absent keys keep release values, unknown routing is refused. + pub fn parse(gpa: std.mem.Allocator, json: []const u8) !Config { + const parsed = try std.json.parseFromSlice(std.json.Value, gpa, json, .{}); + defer parsed.deinit(); + if (parsed.value != .object) return error.BadConfig; + const top = parsed.value.object; + const text = if (top.get("text_config")) |t| (if (t == .object) t.object else return error.BadConfig) else top; + var c: Config = .{}; + const ints = .{ + .{ "vocab", "vocab_size" }, .{ "hidden", "hidden_size" }, + .{ "layers", "num_hidden_layers" }, .{ "heads", "num_attention_heads" }, + .{ "head_dim", "head_dim" }, .{ "rope_dim", "qk_rope_head_dim" }, + .{ "q_lora", "q_lora_rank" }, .{ "o_lora", "o_lora_rank" }, + .{ "o_groups", "o_groups" }, .{ "window", "sliding_window" }, + .{ "experts", "n_routed_experts" }, .{ "shared_experts", "n_shared_experts" }, + .{ "topk", "num_experts_per_tok" }, .{ "expert_width", "moe_intermediate_size" }, + .{ "index_heads", "index_n_heads" }, .{ "index_dim", "index_head_dim" }, + .{ "index_topk", "index_topk" }, .{ "candidate_source", "candidate_source_layer_id" }, + .{ "candidate_blocks", "candidate_topk_blocks" }, .{ "candidate_block_size", "candidate_block_size" }, + .{ "hc_mult", "hc_mult" }, .{ "hc_sinkhorn_iters", "hc_sinkhorn_iters" }, + .{ "engram_max_ngram", "engram_max_ngram_size" }, .{ "engram_vocab", "engram_vocab_size" }, + .{ "engram_heads", "engram_n_heads" }, .{ "engram_head_dim", "engram_head_dim" }, + .{ "engram_pad", "engram_pad_token_id" }, .{ "engram_compressed_vocab", "engram_compressed_vocab_size" }, + .{ "mtp_layers", "num_nextn_predict_layers" }, .{ "dspark_block", "dspark_block_size" }, + .{ "dspark_noise_token", "dspark_noise_token_id" }, .{ "dspark_markov_rank", "dspark_markov_rank" }, + .{ "dspark_experts", "dspark_n_routed_experts" }, .{ "dspark_topk", "dspark_num_experts_per_tok" }, + }; + inline for (ints) |p| if (text.get(p[1])) |v| { + @field(c, p[0]) = try intAs(@TypeOf(@field(c, p[0])), v); + }; + const floats = .{ + .{ "eps", "rms_norm_eps" }, .{ "routed_scale", "routed_scaling_factor" }, .{ "swiglu_limit", "swiglu_limit" }, + .{ "hc_eps", "hc_eps" }, .{ "hc_post_alpha", "hc_post_alpha" }, .{ "engram_gate_clamp", "engram_gate_clamp" }, + }; + inline for (floats) |p| if (text.get(p[1])) |v| { + @field(c, p[0]) = @floatCast(try float(v)); + }; + if (text.get("norm_topk_prob")) |v| c.norm_topk = if (v == .bool) v.bool else return error.BadConfig; + const lists = .{ + .{ "compress_ratios", "compress_ratios" }, .{ "kv_sources", "kv_source_layer_ids" }, + .{ "index_sources", "index_source_layer_ids" }, .{ "engram_layers", "engram_layer_ids" }, + .{ "dspark_targets", "dspark_target_layer_ids" }, .{ "engram_rows", "engram_num_embeddings" }, + }; + inline for (lists) |p| if (text.get(p[1])) |v| { + @field(c, p[0]) = try list(@TypeOf(@field(c, p[0])), v); + }; + // routing the kernels implement: sqrt(softplus) scores, bias-corrected top-k (vLLM noaux_tc) + if (text.get("scoring_func")) |v| if (v != .string or !std.mem.eql(u8, v.string, "sqrtsoftplus")) return error.UnsupportedScoring; + if (text.get("topk_method")) |v| if (v != .string or !std.mem.eql(u8, v.string, "noaux_tc")) return error.UnsupportedScoring; + c.rope = try rope(text); + if (top.get("bos_token_id")) |v| c.bos = try intAs(u32, v); + if (top.get("eos_token_id")) |v| c.eos = try intAs(u32, v); + try c.validate(); + return c; + } + + /// Shapes the kernels and the loader assume; a release that breaks one is refused before any allocation. + pub fn validate(c: *const Config) !void { + if (c.hidden % 128 != 0 or c.expert_width % 128 != 0 or c.vocab % 128 != 0) return error.BadConfig; + if (c.heads % c.o_groups != 0 or c.rope_dim > c.head_dim or c.topk > c.experts) return error.BadConfig; + if (c.dspark_topk > c.dspark_experts and c.dspark_experts != 0) return error.BadConfig; + if (c.compress_ratios.len < c.layers) return error.BadConfig; + for (c.compress_ratios.items()) |r| if (r > 2) return error.BadConfig; + for (c.kv_sources.items()) |s| if (s >= c.layers or c.compressRatio(s) == 0) return error.BadConfig; + for (c.index_sources.items()) |s| if (s >= c.layers or c.compressRatio(s) == 0) return error.BadConfig; + if (c.engram_layers.len != c.engram_rows.len) return error.BadConfig; + // every compressed layer needs a KV source and an index source at or before it + var l: u32 = 0; + while (l < c.layers) : (l += 1) if (c.compressRatio(l) > 0 and (c.kvSource(l) == null or c.indexSource(l) == null)) return error.BadConfig; + } + + /// Layers the forward runs: the backbone, then the DSpark blocks (checkpoint mtp.*). + pub fn allLayers(c: *const Config) u32 { + return c.layers + c.mtp_layers; + } + + pub fn compressRatio(c: *const Config, layer: u32) u32 { + return if (layer < c.compress_ratios.len) c.compress_ratios.buf[layer] else 0; + } + + pub fn isBackbone(c: *const Config, layer: u32) bool { + return layer < c.layers; + } + + /// The latest source at or before `layer` (null for sliding-window layers). + fn latest(c: *const Config, sources: *const Ids, layer: u32) ?u32 { + if (c.compressRatio(layer) == 0) return null; + var best: ?u32 = null; + for (sources.items()) |s| if (s <= layer and (best == null or s > best.?)) { + best = s; + }; + return best; + } + + pub fn kvSource(c: *const Config, layer: u32) ?u32 { + return c.latest(&c.kv_sources, layer); + } + + pub fn indexSource(c: *const Config, layer: u32) ?u32 { + return c.latest(&c.index_sources, layer); + } + + pub fn isKvSource(c: *const Config, layer: u32) bool { + return c.isBackbone(layer) and c.kv_sources.contains(layer); + } + + pub fn isIndexSource(c: *const Config, layer: u32) bool { + return c.isBackbone(layer) and c.index_sources.contains(layer); + } + + pub fn isCandidateSource(c: *const Config, layer: u32) bool { + return c.candidate_source == layer and c.candidate_blocks > 0; + } + + pub fn usesCandidates(c: *const Config, layer: u32) bool { + return c.isIndexSource(layer) and c.candidate_blocks > 0 and c.candidate_source >= 0 and c.candidate_source < layer; + } + + pub fn mode(c: *const Config, layer: u32) Mode { + if (c.compressRatio(layer) == 0) return .swa; + if (c.isKvSource(layer)) return .full; + if (c.isIndexSource(layer)) return .reindex; + return .reuse; + } + + pub fn isEngram(c: *const Config, layer: u32) bool { + return c.engram_layers.contains(layer); + } + + /// (routed experts, experts a token) of a backbone layer or a DSpark block. + pub fn expertsOf(c: *const Config, layer: u32) struct { count: u32, topk: u32 } { + if (c.isBackbone(layer)) return .{ .count = c.experts, .topk = c.topk }; + return .{ .count = if (c.dspark_experts != 0) c.dspark_experts else c.experts, .topk = if (c.dspark_topk != 0) c.dspark_topk else c.topk }; + } + + /// Engram's hash columns: (max n-gram - 1) x heads. + pub fn hashCols(c: *const Config) u32 { + return (c.engram_max_ngram - 1) * c.engram_heads; + } +}; + +fn intAs(comptime T: type, v: std.json.Value) !T { + return switch (v) { + .integer => |i| std.math.cast(T, i) orelse error.BadConfig, + .float => |f| if (@floor(f) == f) std.math.cast(T, @as(i64, @intFromFloat(f))) orelse error.BadConfig else error.BadConfig, + else => error.BadConfig, + }; +} + +fn float(v: std.json.Value) !f64 { + return switch (v) { + .integer => |i| @floatFromInt(i), + .float => |f| f, + else => error.BadConfig, + }; +} + +fn list(comptime L: type, v: std.json.Value) !L { + if (v != .array) return error.BadConfig; + var out: L = .{}; + const T = @TypeOf(out.buf[0]); + for (v.array.items) |x| try out.push(try intAs(T, x)); + return out; +} + +/// rope_theta, compress_rope_theta and rope_scaling (or rope_parameters), as RopeConfig.from_dict. +fn rope(text: std.json.ObjectMap) !Rope { + var r: Rope = .{}; + if (text.get("rope_theta")) |v| r.theta = try float(v); + if (text.get("compress_rope_theta")) |v| r.compress_theta = try float(v); + if (text.get("max_position_embeddings")) |v| r.max_positions = try intAs(u64, v); + const scaling = text.get("rope_scaling") orelse text.get("rope_parameters"); + r.factor = 1.0; + r.original_max_positions = r.max_positions; + r.yarn = false; + if (scaling) |s| if (s == .object) { + const o = s.object; + if (o.get("factor")) |v| r.factor = try float(v); + if (o.get("original_max_position_embeddings")) |v| r.original_max_positions = try intAs(u64, v); + if (o.get("beta_fast")) |v| r.beta_fast = try float(v); + if (o.get("beta_slow")) |v| r.beta_slow = try float(v); + const kind = o.get("rope_type") orelse o.get("type"); + r.yarn = if (kind) |k| (k == .string and !std.mem.eql(u8, k.string, "default")) else false; + }; + return r; +} + +/// A small config with every V4.1 mechanism; split widths are 256 so TP=2 keeps 128-value blocks. +pub fn tiny() Config { + return .{ + .vocab = 256, + .hidden = 256, + .layers = 8, + .heads = 4, + .head_dim = 64, + .rope_dim = 8, + .q_lora = 32, + .o_lora = 128, + .o_groups = 2, + .eps = 1e-6, + .window = 4, + .experts = 8, + .topk = 2, + .expert_width = 256, + .compress_ratios = Ids.of(&.{ 0, 0, 2, 2, 1, 1, 1, 1, 0, 0 }), + .kv_sources = Ids.of(&.{ 2, 4 }), + .index_sources = Ids.of(&.{ 2, 4, 6 }), + .index_heads = 2, + .index_dim = 16, + .index_topk = 3, + .candidate_source = 4, + .candidate_blocks = 2, + .candidate_block_size = 2, + .engram_layers = Ids.of(&.{1}), + .engram_rows = List(u64, 8).of(&.{0}), + .engram_max_ngram = 3, + .engram_vocab = 31, + .engram_heads = 2, + .engram_head_dim = 8, + .engram_compressed_vocab = 50, + .mtp_layers = 2, + .dspark_block = 3, + .dspark_noise_token = 96, + .dspark_targets = Ids.of(&.{ 5, 6, 7 }), + .dspark_markov_rank = 8, + .dspark_experts = 4, + .dspark_topk = 2, + .rope = .{ .factor = 4.0, .original_max_positions = 64, .max_positions = 256 }, + }; +} diff --git a/zig/src/families/deepseek_v41/deepseek_v41.zig b/zig/src/families/deepseek_v41/deepseek_v41.zig new file mode 100644 index 0000000000..fcca76853d --- /dev/null +++ b/zig/src/families/deepseek_v41/deepseek_v41.zig @@ -0,0 +1,23 @@ +//! DeepSeek-V4.1-Flash on the Zig engine: release config, EXL3 pack headers, a rank's TP weight plan (host side). + +pub const config = @import("config.zig"); +pub const Config = config.Config; +pub const exl3 = @import("exl3.zig"); +pub const names = @import("names.zig"); +pub const pack = @import("pack.zig"); +pub const Pack = pack.Pack; +pub const plan = @import("plan.zig"); +pub const Plan = plan.Plan; +pub const stage = @import("stage.zig"); +pub const named = @import("named.zig"); + +test { + _ = config; + _ = exl3; + _ = names; + _ = pack; + _ = plan; + _ = stage; + _ = named; + _ = @import("host_test.zig"); +} diff --git a/zig/src/families/deepseek_v41/exl3.zig b/zig/src/families/deepseek_v41/exl3.zig new file mode 100644 index 0000000000..68050611a4 --- /dev/null +++ b/zig/src/families/deepseek_v41/exl3.zig @@ -0,0 +1,254 @@ +//! EXL3 groups as the pack stores them: trellis geometry, codebooks, a rank's TP slice as byte spans, expert layouts. + +const std = @import("std"); + +/// The codebook a group's trellis decodes with: the `.mul1` / `.mcg` marker's value, else the 3-instruction default. +pub const Codebook = enum(u8) { + inst3, + mul1, + mcg, + + pub const marker_mul1: u32 = 0x83DCD12D; + pub const marker_mcg: u32 = 0xCBAC1FED; + + /// A marker tensor's value (the low 32 bits of its first element); unknown values are refused. + pub fn fromMarker(v: u32) !Codebook { + return switch (v) { + marker_mul1 => .mul1, + marker_mcg => .mcg, + else => error.UnknownCodebookMarker, + }; + } +}; + +/// How a rank holds a group under TP (every boundary a multiple of 128, so Hadamard blocks stay whole). +pub const Split = enum { + /// all of it, on every rank + whole, + /// this rank's output columns: N tiles and `svh` + col, + /// this rank's input rows: K tiles and `suh` (a partial sum the ranks exchange) + row, +}; + +/// One EXL3 matrix y = x @ W, W [K, N]: trellis int16 [K/16, N/16, words], suh fp16 [K], svh fp16 [N]. +pub const Group = struct { + k_tiles: u32, + n_tiles: u32, + /// int16 words a 16x16 tile: 16 x bits (the trellis's last dimension) + words: u32, + + /// From the trellis entry's shape; a tile must hold whole K2 bytes and 1 to 8 bits a value. + pub fn fromShape(shape: []const usize) !Group { + if (shape.len != 3) return error.BadTrellis; + const w = shape[2]; + if (w == 0 or w % 8 != 0 or w > 128 or shape[0] == 0 or shape[1] == 0) return error.BadTrellis; + return .{ .k_tiles = @intCast(shape[0]), .n_tiles = @intCast(shape[1]), .words = @intCast(w) }; + } + + pub fn k(g: Group) u32 { + return 16 * g.k_tiles; + } + + pub fn n(g: Group) u32 { + return 16 * g.n_tiles; + } + + /// Half-bits a value (the kernels' K2): 4 for 2 bits, 16 for 8 bits. + pub fn k2(g: Group) u32 { + return g.words / 8; + } + + /// Bits a value, in halves exact (k2 / 2). + pub fn bits(g: Group) f32 { + return @as(f32, @floatFromInt(g.words)) / 16.0; + } + + pub fn tileBytes(g: Group) u64 { + return @as(u64, g.words) * 2; + } + + pub fn trellisBytes(g: Group) u64 { + return @as(u64, g.k_tiles) * g.n_tiles * g.tileBytes(); + } + + /// The rank's part of the group; a split that does not fall on 128-value blocks is refused. + pub fn part(g: Group, split: Split, rank: u32, world: u32) !Part { + std.debug.assert(rank < world); + var p: Part = .{ .group = g, .trellis = .{ .offset = 0, .len = g.trellisBytes() }, .suh = .{ .offset = 0, .len = 2 * @as(u64, g.k()) }, .svh = .{ .offset = 0, .len = 2 * @as(u64, g.n()) } }; + switch (split) { + .whole => {}, + .row => { + if (g.k() % (world * 128) != 0) return error.SplitNotOnBlocks; + const per = g.k_tiles / world; + p.group.k_tiles = per; + const row = @as(u64, g.n_tiles) * g.tileBytes(); + p.trellis = .{ .offset = rank * per * row, .len = per * row }; + p.suh = .{ .offset = 2 * 16 * @as(u64, rank * per), .len = 2 * 16 * @as(u64, per) }; + }, + .col => { + if (g.n() % (world * 128) != 0) return error.SplitNotOnBlocks; + const per = g.n_tiles / world; + p.group.n_tiles = per; + // one K-tile row at a time: this rank's N tiles are a run inside each row + p.trellis = .{ .offset = rank * per * g.tileBytes(), .len = per * g.tileBytes(), .stride = @as(u64, g.n_tiles) * g.tileBytes(), .count = g.k_tiles }; + p.svh = .{ .offset = 2 * 16 * @as(u64, rank * per), .len = 2 * 16 * @as(u64, per) }; + }, + } + return p; + } +}; + +/// Bytes of a stored tensor to read: `count` runs of `len` bytes, `stride` apart, from `offset` into the tensor. +pub const Span = struct { + offset: u64, + len: u64, + stride: u64 = 0, + count: u64 = 1, + + pub fn bytes(s: Span) u64 { + return s.len * s.count; + } + + /// The `i`th run's start. + pub fn at(s: Span, i: u64) u64 { + return s.offset + i * s.stride; + } +}; + +/// A rank's slice of a group: its geometry and where its trellis, suh and svh bytes sit in the stored tensors. +pub const Part = struct { + group: Group, + trellis: Span, + suh: Span, + svh: Span, + + pub fn bytes(p: Part) u64 { + return p.trellis.bytes() + p.suh.bytes() + p.svh.bytes(); + } +}; + +/// Routed experts of one projection on a rank: one width (X3Stack) or a width per expert (X3Ragged, E2). +pub const Layout = struct { + kind: enum { stack, ragged }, + k_tiles: u32, + n_tiles: u32, + /// each expert's int16 words a tile + words: []u16, + /// each expert's first int16 element in the projection's trellis buffer (back to back, as `ragged()` packs them) + offsets: []u64, + /// int16 elements of every expert's trellis + elements: u64, + + /// The layout of the rank's expert parts; parts must share K and N (the kernels' tables index one shape). + pub fn of(gpa: std.mem.Allocator, parts: []const Part) !Layout { + if (parts.len == 0) return error.NoExperts; + const words = try gpa.alloc(u16, parts.len); + errdefer gpa.free(words); + const offsets = try gpa.alloc(u64, parts.len); + errdefer gpa.free(offsets); + const kt = parts[0].group.k_tiles; + const nt = parts[0].group.n_tiles; + var at: u64 = 0; + var same = true; + for (parts, words, offsets) |p, *w, *o| { + if (p.group.k_tiles != kt or p.group.n_tiles != nt) return error.ExpertShapesDiffer; + w.* = @intCast(p.group.words); + o.* = at; + at += @as(u64, kt) * nt * p.group.words; + same = same and p.group.words == parts[0].group.words; + } + return .{ .kind = if (same) .stack else .ragged, .k_tiles = kt, .n_tiles = nt, .words = words, .offsets = offsets, .elements = at }; + } + + pub fn deinit(l: *Layout, gpa: std.mem.Allocator) void { + gpa.free(l.words); + gpa.free(l.offsets); + l.* = undefined; + } + + pub fn count(l: *const Layout) usize { + return l.words.len; + } + + pub fn trellisBytes(l: *const Layout) u64 { + return 2 * l.elements; + } + + /// Byte offset of expert `e`'s trellis in the projection buffer: 16-byte aligned (a tile is 16 K2 bytes). + pub fn byteOffset(l: *const Layout, e: usize) u64 { + return 2 * l.offsets[e]; + } + + /// Experts at each K2 (index = K2, 1..16): the widths x3gm launches once each. + pub fn histogram(l: *const Layout) [17]u32 { + var h: [17]u32 = @splat(0); + for (l.words) |w| h[w / 8] += 1; + return h; + } + + /// Distinct widths: x3gm's launches for this projection. + pub fn widths(l: *const Layout) u32 { + var n: u32 = 0; + for (l.histogram()) |c| n += @intFromBool(c > 0); + return n; + } +}; + +const testing = std.testing; + +test "group geometry and K2" { + const g = try Group.fromShape(&.{ 320, 144, 32 }); + try testing.expectEqual(@as(u32, 5120), g.k()); + try testing.expectEqual(@as(u32, 2304), g.n()); + try testing.expectEqual(@as(u32, 4), g.k2()); + try testing.expectEqual(@as(f32, 2.0), g.bits()); + try testing.expectError(error.BadTrellis, Group.fromShape(&.{ 320, 144, 30 })); + try testing.expectError(error.BadTrellis, Group.fromShape(&.{ 320, 144 })); + try testing.expectEqual(Codebook.mcg, try Codebook.fromMarker(0xCBAC1FED)); + try testing.expectError(error.UnknownCodebookMarker, Codebook.fromMarker(7)); +} + +test "TP parts: rows contiguous, columns a run a K-tile row" { + const w2 = try Group.fromShape(&.{ 144, 320, 48 }); // a 3-bit routed down projection: 2304 -> 5120 + const r1 = try w2.part(.row, 1, 2); + try testing.expectEqual(@as(u32, 72), r1.group.k_tiles); + try testing.expectEqual(@as(u64, 72 * 320 * 96), r1.trellis.len); + try testing.expectEqual(r1.trellis.len, r1.trellis.offset); + try testing.expectEqual(Span{ .offset = 2 * 1152, .len = 2 * 1152 }, r1.suh); + try testing.expectEqual(@as(u64, 2 * 5120), r1.svh.len); + const w1 = try Group.fromShape(&.{ 320, 144, 32 }); + const c1 = try w1.part(.col, 1, 2); + try testing.expectEqual(@as(u32, 72), c1.group.n_tiles); + try testing.expectEqual(@as(u64, 320), c1.trellis.count); + try testing.expectEqual(@as(u64, 144 * 64), c1.trellis.stride); + try testing.expectEqual(@as(u64, 72 * 64), c1.trellis.offset); + try testing.expectEqual(c1.group.trellisBytes(), c1.trellis.bytes()); + // the 8-bit indexer key projection (512 -> 128) cannot split its 128 outputs over two ranks + const wk = try Group.fromShape(&.{ 32, 8, 128 }); + try testing.expectError(error.SplitNotOnBlocks, wk.part(.col, 0, 2)); + try testing.expectEqual(wk.trellisBytes(), (try wk.part(.whole, 1, 2)).trellis.bytes()); +} + +test "expert layouts: one width stacks, mixed widths go ragged with aligned offsets" { + const a = testing.allocator; + const two = try (try Group.fromShape(&.{ 144, 320, 32 })).part(.row, 0, 2); + const three = try (try Group.fromShape(&.{ 144, 320, 48 })).part(.row, 0, 2); + var stack = try Layout.of(a, &.{ two, two, two }); + defer stack.deinit(a); + try testing.expect(stack.kind == .stack); + try testing.expectEqual(@as(u32, 1), stack.widths()); + // a mixed-bit pack's layer 2 down projection on rank 0: 383 experts at 2 bits and one at 3 bits + var parts: [384]Part = @splat(two); + parts[200] = three; + var rag = try Layout.of(a, &parts); + defer rag.deinit(a); + try testing.expect(rag.kind == .ragged); + try testing.expectEqual(@as(u64, 566968320), rag.trellisBytes()); + try testing.expectEqual(@as(u32, 383), rag.histogram()[4]); + try testing.expectEqual(@as(u32, 1), rag.histogram()[6]); + try testing.expectEqual(@as(u64, 0), rag.byteOffset(201) % 16); + try testing.expectEqual(rag.byteOffset(200) + three.trellis.bytes(), rag.byteOffset(201)); + const odd = try (try Group.fromShape(&.{ 144, 160, 32 })).part(.row, 0, 2); + try testing.expectError(error.ExpertShapesDiffer, Layout.of(a, &.{ two, odd })); +} diff --git a/zig/src/families/deepseek_v41/fixtures/config.json b/zig/src/families/deepseek_v41/fixtures/config.json new file mode 100644 index 0000000000..eeec94126b --- /dev/null +++ b/zig/src/families/deepseek_v41/fixtures/config.json @@ -0,0 +1,184 @@ +{ + "architectures": [ + "DeepseekV41ForCausalLM" + ], + "model_type": "deepseek_v41", + "dtype": "bfloat16", + "transformers_version": "5.6.0", + "bos_token_id": 0, + "eos_token_id": 1, + "pad_token_id": 2, + "image_token_id": 129264, + "quantization_config": { + "quant_method": "exl3", + "version": "1.4.2", + "avg_bits": 2.9, + "head_bits": 6, + "calibration": { + "rows": 250, + "cols": 2048 + }, + "out_scales": "always", + "codebook": "mul1", + "original_quantization_config": { + "quant_method": "fp8", + "activation_scheme": "dynamic", + "weight_block_size": [ + 32, + 32 + ], + "scale_fmt": "ue8m0", + "expert_dtype": "fp4" + }, + "mtp_bits": 4 + }, + "text_config": { + "model_type": "deepseek_v41_text", + "vocab_size": 129280, + "hidden_size": 5120, + "moe_intermediate_size": 2304, + "num_hidden_layers": 40, + "num_attention_heads": 64, + "num_key_value_heads": 1, + "head_dim": 512, + "qk_rope_head_dim": 64, + "q_lora_rank": 1280, + "o_lora_rank": 1024, + "o_groups": 8, + "hidden_act": "silu", + "swiglu_limit": 10.0, + "rms_norm_eps": 1e-20, + "attention_bias": false, + "attention_dropout": 0.0, + "initializer_range": 0.02, + "use_cache": true, + "tie_word_embeddings": false, + "max_position_embeddings": 1048576, + "rope_theta": 10000, + "rope_scaling": { + "rope_type": "yarn", + "factor": 16, + "beta_fast": 32, + "beta_slow": 1, + "original_max_position_embeddings": 65536 + }, + "n_routed_experts": 384, + "n_shared_experts": 1, + "num_experts_per_tok": 6, + "scoring_func": "sqrtsoftplus", + "topk_method": "noaux_tc", + "norm_topk_prob": true, + "routed_scaling_factor": 1.5, + "sliding_window": 128, + "compress_ratios": [ + 0, + 0, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 0, + 0, + 0 + ], + "compress_rope_theta": 160000, + "kv_source_layer_ids": [ + 2, + 8, + 14, + 20 + ], + "index_source_layer_ids": [ + 2, + 8, + 14, + 20, + 24, + 28, + 32, + 36 + ], + "index_n_heads": 32, + "index_head_dim": 128, + "index_topk": 512, + "candidate_source_layer_id": 20, + "candidate_topk_blocks": 2048, + "candidate_block_size": 8, + "hc_mult": 4, + "hc_sinkhorn_iters": 20, + "hc_eps": 1e-06, + "engram_layer_ids": [ + 1, + 14 + ], + "engram_num_embeddings": [ + 384006168, + 384016682 + ], + "engram_max_ngram_size": 4, + "engram_vocab_size": 16000000, + "engram_n_heads": 8, + "engram_head_dim": 256, + "engram_pad_token_id": 2, + "engram_compressed_vocab_size": 99092, + "num_nextn_predict_layers": 3, + "dspark_block_size": 5, + "dspark_noise_token_id": 128799, + "dspark_target_layer_ids": [ + 37, + 38, + 39 + ], + "dspark_markov_rank": 256, + "dspark_n_routed_experts": 128, + "dspark_num_experts_per_tok": 3 + }, + "vision_config": { + "model_type": "deepseek_v41_vision", + "num_hidden_layers": 32, + "hidden_size": 1024, + "num_attention_heads": 16, + "intermediate_size": 2816, + "patch_size": 14, + "rope_theta": 10000, + "downsample_ratio": 3, + "max_image_tokens": 1024, + "min_pixels": 295936, + "max_wh_ratio": null + }, + "engram_table_dir": null +} \ No newline at end of file diff --git a/zig/src/families/deepseek_v41/fixtures/widths-exl3-2p8.txt b/zig/src/families/deepseek_v41/fixtures/widths-exl3-2p8.txt new file mode 100644 index 0000000000..317ae6ed73 --- /dev/null +++ b/zig/src/families/deepseek_v41/fixtures/widths-exl3-2p8.txt @@ -0,0 +1,44 @@ +# expert widths of a 2.8 bpw EXL3 pack, world 2 rank 0: layer w1 w2 w3 shared_k2(w1,w2,w3) rank_bytes +0 4:384 4:384 4:384 10,10,10 1709752320 +1 4:384 4:384 4:384 10,10,10 1709752320 +2 4:384 4:383,6:1 4:384 10,10,10 1710489600 +3 4:384 4:384 4:384 10,10,10 1709752320 +4 4:384 4:383,6:1 4:384 10,10,10 1710489600 +5 4:384 4:384 4:384 10,10,10 1709752320 +6 4:384 4:384 4:384 8,8,8 1707540480 +7 4:384 4:382,6:2 4:384 10,10,10 1711226880 +8 4:384 4:376,6:8 4:384 12,12,12 1717862400 +9 4:382,6:2 4:369,6:14,8:1 4:382,6:2 12,12,12 1726709760 +10 4:376,6:8 4:360,6:23,8:1 4:376,6:8 12,12,12 1742192640 +11 4:382,6:2 4:364,6:20 4:382,6:2 12,12,12 1729658880 +12 4:382,6:2 4:373,6:10,8:1 4:382,6:2 12,12,12 1723760640 +13 4:384 4:380,6:4 4:384 12,12,12 1714913280 +14 4:384 4:384 4:384 10,10,10 1709752320 +15 4:384 4:384 4:384 10,10,10 1709752320 +16 4:384 4:384 4:384 10,10,10 1709752320 +17 4:384 4:384 4:384 8,8,8 1707540480 +18 4:384 4:384 4:384 6,6,6 1705328640 +19 4:71,6:222,8:90,10:1 4:34,6:126,8:217,10:7 4:71,6:222,8:90,10:1 12,12,12 2737520640 +20 4:41,6:162,8:166,10:15 4:18,6:69,8:221,10:76 4:41,6:162,8:166,10:15 12,12,12 3051601920 +21 4:18,6:104,8:216,10:46 4:5,6:38,8:182,10:159 4:18,6:104,8:216,10:46 12,12,12 3353886720 +22 4:13,6:69,8:226,10:76 4:3,6:25,8:145,10:211 4:13,6:69,8:226,10:76 12,12,12 3515351040 +23 4:24,6:55,8:220,10:85 4:21,6:31,8:139,10:193 4:24,6:55,8:220,10:85 12,12,12 3472588800 +24 4:14,6:117,8:221,10:32 4:6,6:34,8:213,10:131 4:14,6:117,8:221,10:32 12,12,12 3306700800 +25 4:28,6:175,8:167,10:14 4:5,6:66,8:237,10:76 4:28,6:175,8:167,10:14 12,12,12 3090677760 +26 4:31,6:239,8:108,10:6 4:8,6:116,8:228,10:32 4:31,6:239,8:108,10:6 12,12,12 2901934080 +27 4:54,6:199,8:126,10:5 4:20,6:113,8:219,10:32 4:54,6:199,8:126,10:5 12,12,12 2876129280 +28 4:109,6:213,8:59,10:3 4:39,6:178,8:146,10:21 4:109,6:213,8:59,10:3 12,12,12 2606284800 +29 4:116,6:219,8:47,10:2 4:47,6:189,8:136,10:12 4:116,6:219,8:47,10:2 12,12,12 2548776960 +30 4:164,6:192,8:28 4:76,6:217,8:88,10:3 4:164,6:192,8:28 12,12,12 2374041600 +31 4:180,6:177,8:26,10:1 4:84,6:213,8:84,10:3 4:180,6:177,8:26,10:1 12,12,12 2341601280 +32 4:188,6:172,8:24 4:84,6:205,8:92,10:3 4:188,6:172,8:24 12,12,12 2329804800 +33 4:187,6:175,8:22 4:92,6:217,8:72,10:3 4:187,6:175,8:22 12,12,12 2307686400 +34 4:208,6:157,8:19 4:122,6:207,8:53,10:2 4:208,6:157,8:19 12,12,12 2234695680 +35 4:65,6:186,8:121,10:12 4:36,6:97,8:207,10:44 4:65,6:186,8:121,10:12 12,12,12 2870231040 +36 6:147,8:205,10:32 6:73,8:207,10:104 6:147,8:205,10:32 12,12,12 3263938560 +37 6:91,8:222,10:71 6:47,8:158,10:179 6:91,8:222,10:71 12,12,12 3478487040 +38 6:50,8:129,10:205 6:32,8:61,10:291 6:50,8:129,10:205 12,12,12 3830169600 +39 6:48,8:103,10:233 6:38,8:42,10:304 6:48,8:103,10:233 12,12,12 3879567360 +40 8:128 8:128 8:128 8,8,8 1141309440 +41 8:128 8:128 8:128 8,8,8 1141309440 +42 8:128 8:128 8:128 8,8,8 1141309440 diff --git a/zig/src/families/deepseek_v41/host_test.zig b/zig/src/families/deepseek_v41/host_test.zig new file mode 100644 index 0000000000..1212ed3518 --- /dev/null +++ b/zig/src/families/deepseek_v41/host_test.zig @@ -0,0 +1,330 @@ +//! Host tests of the loader side: release config, synthetic pack, 2.8 bpw widths, header dumps (TF_DSV41_HEADERS). + +const std = @import("std"); +const testing = std.testing; +const config = @import("config.zig"); +const Config = config.Config; +const exl3 = @import("exl3.zig"); +const Pack = @import("pack.zig").Pack; +const plan = @import("plan.zig"); +const named = @import("named.zig"); + +test "the release config.json: topology, routing and rope" { + const c = try Config.parse(testing.allocator, @embedFile("fixtures/config.json")); + try testing.expectEqual(@as(u32, 129280), c.vocab); + try testing.expectEqual(@as(u32, 43), c.allLayers()); + try testing.expectEqual(@as(f32, 1e-20), c.eps); + try testing.expect(c.rope.yarn and c.rope.factor == 16 and c.rope.original_max_positions == 65536); + try testing.expectEqual(@as(f64, 160000), c.rope.compress_theta); + try testing.expectEqual(config.Mode.swa, c.mode(1)); + try testing.expectEqual(config.Mode.full, c.mode(2)); + try testing.expectEqual(config.Mode.reuse, c.mode(3)); + try testing.expectEqual(config.Mode.full, c.mode(20)); + try testing.expectEqual(config.Mode.reindex, c.mode(24)); + try testing.expectEqual(config.Mode.swa, c.mode(41)); + try testing.expectEqual(@as(?u32, 14), c.kvSource(19)); + try testing.expectEqual(@as(?u32, 20), c.kvSource(39)); + try testing.expectEqual(@as(?u32, 36), c.indexSource(39)); + try testing.expect(c.isCandidateSource(20) and c.usesCandidates(24) and !c.usesCandidates(20)); + try testing.expectEqual(@as(u32, 128), c.expertsOf(40).count); + try testing.expectEqual(@as(u32, 3), c.expertsOf(40).topk); + try testing.expectEqual(@as(u32, 24), c.hashCols()); + try testing.expect(c.isEngram(14) and !c.isEngram(2)); + try testing.expectEqual(@as(u32, 1), c.eos); + // the release's values equal the Config defaults + const d: Config = .{}; + try testing.expectEqualSlices(u32, d.compress_ratios.items(), c.compress_ratios.items()); + try testing.expectEqualSlices(u64, d.engram_rows.items(), c.engram_rows.items()); + try testing.expectError(error.UnsupportedScoring, Config.parse(testing.allocator, "{\"scoring_func\": \"sigmoid\"}")); + try testing.expectError(error.BadConfig, Config.parse(testing.allocator, "{\"compress_ratios\": [0, 3]}")); + try config.tiny().validate(); +} + +/// Header JSON of a pack for `c`, named as the release packs are (their index, pattern by pattern), header only. +fn tinyPack(a: std.mem.Allocator, c: *const Config, extra: []const u8, skip: []const u8) ![]u8 { + var n: u64 = 0; + return tinyPackLen(a, c, extra, skip, &n); +} + +/// `tinyPack`, and the data bytes its header describes. +fn tinyPackLen(a: std.mem.Allocator, c: *const Config, extra: []const u8, skip: []const u8, data_len: *u64) ![]u8 { + var w: std.Io.Writer.Allocating = .init(a); + errdefer w.deinit(); + var at: u64 = 0; + var first = true; + const T = struct { + fn put(wr: *std.Io.Writer, pos: *u64, f: *bool, sk: []const u8, name: []const u8, dtype: []const u8, shape: []const u64) !void { + if (sk.len > 0 and std.mem.eql(u8, name, sk)) return; + var n: u64 = if (std.mem.eql(u8, dtype, "F32") or std.mem.eql(u8, dtype, "I32")) 4 else 2; + for (shape) |d| n *= d; + try wr.print("{s}\"{s}\":{{\"dtype\":\"{s}\",\"shape\":[", .{ if (f.*) "" else ",", name, dtype }); + for (shape, 0..) |d, i| try wr.print("{s}{d}", .{ if (i == 0) "" else ",", d }); + try wr.print("],\"data_offsets\":[{d},{d}]}}", .{ pos.*, pos.* + n }); + pos.* += n; + f.* = false; + } + fn group(wr: *std.Io.Writer, pos: *u64, f: *bool, sk: []const u8, p: []const u8, k: u64, n: u64, words: u64) !void { + var b: [160]u8 = undefined; + try put(wr, pos, f, sk, try std.fmt.bufPrint(&b, "{s}.trellis", .{p}), "I16", &.{ k / 16, n / 16, words }); + try put(wr, pos, f, sk, try std.fmt.bufPrint(&b, "{s}.suh", .{p}), "F16", &.{k}); + try put(wr, pos, f, sk, try std.fmt.bufPrint(&b, "{s}.svh", .{p}), "F16", &.{n}); + try put(wr, pos, f, sk, try std.fmt.bufPrint(&b, "{s}.mul1", .{p}), "I32", &.{}); + } + }; + const wr = &w.writer; + try wr.writeAll("{"); + const D: u64 = c.hidden; + var nb: [160]u8 = undefined; + var l: u32 = 0; + while (l < c.allLayers()) : (l += 1) { + var pb: [32]u8 = undefined; + const p = if (l < c.layers) try std.fmt.bufPrint(&pb, "layers.{d}", .{l}) else try std.fmt.bufPrint(&pb, "mtp.{d}", .{l - c.layers}); + const E: u64 = if (l < c.layers) c.experts else c.dspark_experts; + inline for (.{ "attn_norm.weight", "ffn_norm.weight" }) |s| try T.put(wr, &at, &first, skip, try std.fmt.bufPrint(&nb, "{s}." ++ s, .{p}), "BF16", &.{D}); + inline for (.{ "attn", "ffn" }) |s| { + try T.put(wr, &at, &first, skip, try std.fmt.bufPrint(&nb, "{s}.hc_" ++ s ++ "_fn", .{p}), "F32", &.{ 24, 4 * D }); + try T.put(wr, &at, &first, skip, try std.fmt.bufPrint(&nb, "{s}.hc_" ++ s ++ "_base", .{p}), "F32", &.{24}); + try T.put(wr, &at, &first, skip, try std.fmt.bufPrint(&nb, "{s}.hc_" ++ s ++ "_scale", .{p}), "F32", &.{3}); + } + try T.put(wr, &at, &first, skip, try std.fmt.bufPrint(&nb, "{s}.attn.q_norm.weight", .{p}), "BF16", &.{c.q_lora}); + try T.put(wr, &at, &first, skip, try std.fmt.bufPrint(&nb, "{s}.attn.kv_norm.weight", .{p}), "BF16", &.{c.head_dim}); + try T.put(wr, &at, &first, skip, try std.fmt.bufPrint(&nb, "{s}.attn.attn_sink", .{p}), "F32", &.{c.heads}); + try T.group(wr, &at, &first, skip, try std.fmt.bufPrint(&nb, "{s}.attn.wq_a", .{p}), D, c.q_lora, 64); + try T.group(wr, &at, &first, skip, try std.fmt.bufPrint(&nb, "{s}.attn.wkv", .{p}), D, c.head_dim, 80); + try T.group(wr, &at, &first, skip, try std.fmt.bufPrint(&nb, "{s}.attn.wq_b", .{p}), c.q_lora, c.heads * c.head_dim, 64); + var g: u32 = 0; + while (g < c.o_groups) : (g += 1) try T.group(wr, &at, &first, skip, try std.fmt.bufPrint(&nb, "{s}.attn.wo_a.slice.{d}", .{ p, g }), c.heads / c.o_groups * c.head_dim, c.o_lora, 64); + try T.group(wr, &at, &first, skip, try std.fmt.bufPrint(&nb, "{s}.attn.wo_b", .{p}), c.o_groups * c.o_lora, D, 64); + if (l < c.layers and c.kv_sources.contains(l)) { + try T.group(wr, &at, &first, skip, try std.fmt.bufPrint(&nb, "{s}.attn.compressor.wkv", .{p}), D, c.head_dim, 80); + try T.put(wr, &at, &first, skip, try std.fmt.bufPrint(&nb, "{s}.attn.compressor.wkv.weight", .{p}), "BF16", &.{ c.head_dim, D }); + if (c.compressRatio(l) == 2) { + try T.group(wr, &at, &first, skip, try std.fmt.bufPrint(&nb, "{s}.attn.compressor.wgate", .{p}), D, c.head_dim, 80); + try T.put(wr, &at, &first, skip, try std.fmt.bufPrint(&nb, "{s}.attn.compressor.wgate.weight", .{p}), "BF16", &.{ c.head_dim, D }); + } + try T.put(wr, &at, &first, skip, try std.fmt.bufPrint(&nb, "{s}.attn.compressor.norm.weight", .{p}), "BF16", &.{c.head_dim}); + try T.group(wr, &at, &first, skip, try std.fmt.bufPrint(&nb, "{s}.attn.indexer.wk", .{p}), c.head_dim, c.index_dim, 128); + try T.put(wr, &at, &first, skip, try std.fmt.bufPrint(&nb, "{s}.attn.indexer.wk.weight", .{p}), "BF16", &.{ c.index_dim, c.head_dim }); + try T.put(wr, &at, &first, skip, try std.fmt.bufPrint(&nb, "{s}.attn.indexer.k_norm.weight", .{p}), "BF16", &.{c.index_dim}); + } + if (l < c.layers and c.index_sources.contains(l)) { + try T.group(wr, &at, &first, skip, try std.fmt.bufPrint(&nb, "{s}.attn.indexer.wq_b", .{p}), c.q_lora, c.index_heads * c.index_dim, 80); + try T.put(wr, &at, &first, skip, try std.fmt.bufPrint(&nb, "{s}.attn.indexer.weights_proj.weight", .{p}), "BF16", &.{ c.index_heads, D }); + } + if (c.engram_layers.contains(l)) { + try T.group(wr, &at, &first, skip, try std.fmt.bufPrint(&nb, "{s}.engram.wkv", .{p}), c.hashCols() * c.engram_head_dim, (c.engram_max_ngram + 1) * D, 64); + try T.put(wr, &at, &first, skip, try std.fmt.bufPrint(&nb, "{s}.engram.q_weight", .{p}), "BF16", &.{ c.hc_mult, D }); + try T.put(wr, &at, &first, skip, try std.fmt.bufPrint(&nb, "{s}.engram.k_weight", .{p}), "BF16", &.{ c.hc_mult, D }); + } + try T.put(wr, &at, &first, skip, try std.fmt.bufPrint(&nb, "{s}.ffn.gate.weight", .{p}), "BF16", &.{ E, D }); + try T.put(wr, &at, &first, skip, try std.fmt.bufPrint(&nb, "{s}.ffn.gate.bias", .{p}), "F32", &.{E}); + try T.put(wr, &at, &first, skip, try std.fmt.bufPrint(&nb, "{s}.ffn.gate.bias_vl", .{p}), "F32", &.{E}); + inline for (.{ "w1", "w2", "w3" }) |m| { + const kn: [2]u64 = if (comptime std.mem.eql(u8, m, "w2")) .{ c.expert_width, D } else .{ D, c.expert_width }; + try T.group(wr, &at, &first, skip, try std.fmt.bufPrint(&nb, "{s}.ffn.shared_experts." ++ m, .{p}), kn[0], kn[1], 80); + var e: u64 = 0; + while (e < E) : (e += 1) { + // block 2: expert 1's down projection at 3 bits, expert 3's gate at 3 bits (a gate/up width mismatch) + const words: u64 = if (l == 2 and ((e == 1 and comptime std.mem.eql(u8, m, "w2")) or (e == 3 and comptime std.mem.eql(u8, m, "w1")))) 48 else 32; + try T.group(wr, &at, &first, skip, try std.fmt.bufPrint(&nb, "{s}.ffn.experts.{d}." ++ m, .{ p, e }), kn[0], kn[1], words); + } + } + } + try T.put(wr, &at, &first, skip, "embed.weight", "BF16", &.{ c.vocab, D }); + try T.put(wr, &at, &first, skip, "norm.weight", "BF16", &.{D}); + try T.group(wr, &at, &first, skip, "head", D, c.vocab, 96); + const last = c.mtp_layers - 1; + try T.group(wr, &at, &first, skip, try std.fmt.bufPrint(&nb, "mtp.{d}.main_proj", .{last}), 3 * D, D, 64); + inline for (.{ "main_norm", "norm" }) |s| try T.put(wr, &at, &first, skip, try std.fmt.bufPrint(&nb, "mtp.{d}." ++ s ++ ".weight", .{last}), "BF16", &.{D}); + inline for (.{ "embed", "head" }) |s| try T.put(wr, &at, &first, skip, try std.fmt.bufPrint(&nb, "mtp.{d}.markov_head." ++ s ++ ".weight", .{last}), "BF16", &.{ c.vocab, c.dspark_markov_rank }); + try T.put(wr, &at, &first, skip, try std.fmt.bufPrint(&nb, "mtp.{d}.confidence_head.proj.weight", .{last}), "F32", &.{ 1, D + c.dspark_markov_rank }); + try T.put(wr, &at, &first, skip, "vision.norm.weight", "BF16", &.{D}); + try T.put(wr, &at, &first, skip, "image_start", "BF16", &.{D}); + if (extra.len > 0) try T.put(wr, &at, &first, skip, extra, "BF16", &.{4}); + try wr.writeAll("}"); + data_len.* = at; + return w.toOwnedSlice(); +} + +fn tinyPlan(a: std.mem.Allocator, c: *const Config, extra: []const u8, skip: []const u8, rank: u32) !plan.Plan { + return tinyPlanOf(a, c, c, extra, skip, rank); +} + +/// A pack written for `stored`, planned with `c`. +fn tinyPlanOf(a: std.mem.Allocator, c: *const Config, stored: *const Config, extra: []const u8, skip: []const u8, rank: u32) !plan.Plan { + const json = try tinyPack(a, stored, extra, skip); + defer a.free(json); + var p = Pack.init(a); + defer p.deinit(); + try p.addHeader("model.safetensors", json, std.math.maxInt(u32)); + return plan.build(a, c, &p, .{ .rank = rank, .world = 2 }); +} + +test "a synthetic pack: both ranks plan every block, the parts tile the full tensors, nothing is left" { + const a = testing.allocator; + const c = config.tiny(); + var r0 = try tinyPlan(a, &c, "", "", 0); + defer r0.deinit(); + var r1 = try tinyPlan(a, &c, "", "", 1); + defer r1.deinit(); + try testing.expectEqual(@as(usize, 0), r0.leftovers.len); + try testing.expectEqual(@as(usize, 10), r0.layers.len); + const l2 = &r0.layers[2]; + try testing.expectEqual(config.Mode.full, l2.mode); + try testing.expect(l2.experts[1].layout.kind == .ragged and l2.experts[0].layout.kind == .ragged and l2.experts[2].layout.kind == .stack); + try testing.expectEqual(@as(u32, 1), l2.gate_up_mismatch); + try testing.expectEqual(@as(u32, 2), l2.experts[1].layout.widths()); + try testing.expect(r0.layers[3].experts[1].layout.kind == .stack); + try testing.expectEqual(@as(usize, 4), r0.layers[8].experts[0].layout.count()); + // every split part on the two ranks adds up to the stored tensor + for (r0.layers, r1.layers) |x, y| for (x.groups, y.groups) |g0, g1| { + if (g0.split == .whole) { + try testing.expectEqual(g0.full.trellisBytes(), g0.part.trellis.bytes()); + } else try testing.expectEqual(g0.full.trellisBytes(), g0.part.trellis.bytes() + g1.part.trellis.bytes()); + }; + try testing.expectEqual(r0.head.?.full.trellisBytes(), r0.head.?.part.trellis.bytes() + r1.head.?.part.trellis.bytes()); + try testing.expectEqualStrings("mtp.1.main_proj", r0.dspark.?.main_proj.prefix); + try testing.expect(r0.dspark.?.confidence != null); + try testing.expectEqual(r0.bytes(), r1.bytes()); + // the embedding: half the vocabulary rows on each rank + try testing.expectEqual(@as(u64, 128 * 256 * 2), r0.top[0].bytes); +} + +test "a synthetic pack: unknown names are leftovers, missing ones and wrong shapes are refused" { + const a = testing.allocator; + const c = config.tiny(); + var p = try tinyPlan(a, &c, "layers.3.attn.mystery", "", 0); + defer p.deinit(); + try testing.expectEqual(@as(usize, 1), p.leftovers.len); + try testing.expectEqualStrings("layers.3.attn.mystery", p.leftovers[0]); + try testing.expectError(error.MissingTensor, tinyPlan(a, &c, "", "layers.5.ffn.experts.6.w3.svh", 0)); + try testing.expectError(error.MissingTensor, tinyPlan(a, &c, "", "layers.4.attn.indexer.k_norm.weight", 1)); + var wide = c; + wide.q_lora = 48; + try testing.expectError(error.UnexpectedShape, tinyPlanOf(a, &wide, &c, "", "", 0)); +} + +test "a synthetic pack on disk: every planned tensor named in the reference loader's form" { + const a = testing.allocator; + const io = testing.io; + const c = config.tiny(); + var n: u64 = 0; + const json = try tinyPackLen(a, &c, "", "", &n); + defer a.free(json); + var tmp = testing.tmpDir(.{}); + defer tmp.cleanup(); + { + var f = try tmp.dir.createFile(io, "model.safetensors", .{}); + defer f.close(io); + var len: [8]u8 = undefined; + std.mem.writeInt(u64, &len, json.len, .little); + try f.writePositionalAll(io, &len, 0); + try f.writePositionalAll(io, json, 8); + try f.setLength(io, 8 + json.len + n); + } + const dir = try tmp.dir.realPathFileAlloc(io, ".", a); + defer a.free(dir); + var p = try Pack.open(a, io, dir); + defer p.deinit(); + var pl = try plan.build(a, &c, &p, .{ .rank = 1, .world = 2 }); + defer pl.deinit(); + var b: named.Builder = .{ .gpa = a, .io = io, .pack = &p }; + defer b.deinit(); + try b.all(&pl, c.o_groups); + try testing.expectEqual(named.Kind.i16, b.find("L2.moe.w2.data").?.kind); + try testing.expect(b.find("L3.moe.w2.trellis") != null and b.find("L3.moe.w2.data") == null); + try testing.expectEqualSlices(usize, &.{ 8, 8, 16, 32 }, b.find("L3.moe.w2.trellis").?.shape[0..4]); + const qk = b.find("L1.engram.qk").?; + try testing.expect(qk.kind == .f32 and qk.bytes.len == 4 * c.hc_mult * c.hidden); + try testing.expectEqual(named.Kind.bf16, b.find("L0.moe.gate").?.kind); + try testing.expectEqual(named.Kind.f32, b.find("L2.attn.ix_wp").?.kind); + // rank 1 of 2: half the heads' sinks, half the vocabulary rows, wo_a's one local group + try testing.expectEqual(@as(usize, 4 * c.heads / 2), b.find("L4.attn.sink").?.bytes.len); + try testing.expectEqual(@as(usize, 2 * c.vocab / 2 * c.hidden), b.find("embed").?.bytes.len); + try testing.expect(b.find("L4.attn.wo_a.0.svh") != null and b.find("L4.attn.wo_a.1.svh") == null); + try testing.expect(b.find("head.trellis") != null and b.find("norm") != null); +} + +/// One fixture histogram ("4:383,6:1") as rank-0 parts of a TP=2 projection with full shape `kn`. +fn parts(a: std.mem.Allocator, text: []const u8, kn: [2]u32, split: exl3.Split) ![]exl3.Part { + var out: std.ArrayList(exl3.Part) = .empty; + errdefer out.deinit(a); + var it = std.mem.tokenizeScalar(u8, text, ','); + while (it.next()) |item| { + const colon = std.mem.indexOfScalar(u8, item, ':') orelse return error.BadFixture; + const k2 = try std.fmt.parseInt(u32, item[0..colon], 10); + const n = try std.fmt.parseInt(u32, item[colon + 1 ..], 10); + const g = try exl3.Group.fromShape(&.{ kn[0] / 16, kn[1] / 16, 8 * k2 }); + try out.appendNTimes(a, try g.part(split, 0, 2), n); + } + return out.toOwnedSlice(a); +} + +test "the 2.8 bpw EXL3 pack: rank 0's routed + shared trellis bytes of every block match the reference loader's listing" { + const a = testing.allocator; + const c: Config = .{}; + var lines = std.mem.tokenizeScalar(u8, @embedFile("fixtures/widths-exl3-2p8.txt"), '\n'); + var blocks: u32 = 0; + var ragged: u32 = 0; + while (lines.next()) |line| { + if (line[0] == '#') continue; + var f = std.mem.tokenizeScalar(u8, line, ' '); + const layer = try std.fmt.parseInt(u32, f.next().?, 10); + var total: u64 = 0; + const projs = [_]struct { kn: [2]u32, split: exl3.Split }{ .{ .kn = .{ c.hidden, c.expert_width }, .split = .col }, .{ .kn = .{ c.expert_width, c.hidden }, .split = .row }, .{ .kn = .{ c.hidden, c.expert_width }, .split = .col } }; + for (projs) |pj| { + const ps = try parts(a, f.next().?, pj.kn, pj.split); + defer a.free(ps); + try testing.expectEqual(@as(usize, c.expertsOf(layer).count), ps.len); + var lay = try exl3.Layout.of(a, ps); + defer lay.deinit(a); + ragged += @intFromBool(lay.kind == .ragged); + total += lay.trellisBytes(); + } + var sh = std.mem.tokenizeScalar(u8, f.next().?, ','); + for (projs) |pj| { + const g = try exl3.Group.fromShape(&.{ pj.kn[0] / 16, pj.kn[1] / 16, 8 * try std.fmt.parseInt(u32, sh.next().?, 10) }); + total += (try g.part(pj.split, 0, 2)).trellis.bytes(); + } + try testing.expectEqual(try std.fmt.parseInt(u64, f.next().?, 10), total); + blocks += 1; + } + try testing.expectEqual(@as(u32, 43), blocks); + try testing.expect(ragged > 40); +} + +test "real pack headers (TF_DSV41_HEADERS: model.safetensors.index.json and .hdr dumps): planned blocks" { + const dir = testing.environ.getPosix("TF_DSV41_HEADERS") orelse return error.SkipZigTest; + const a = testing.allocator; + const io = testing.io; + var p = Pack.init(a); + defer p.deinit(); + const cwd = std.Io.Dir.cwd(); + const index = try std.fs.path.join(a, &.{ dir, "model.safetensors.index.json" }); + defer a.free(index); + const text = try cwd.readFileAlloc(io, index, a, .limited(1 << 28)); + defer a.free(text); + try p.useIndex(text); + var d = try cwd.openDir(io, dir, .{ .iterate = true }); + defer d.close(io); + var it = d.iterate(); + var blocks: std.ArrayList(u32) = .empty; + defer blocks.deinit(a); + while (try it.next(io)) |e| { + if (!std.mem.endsWith(u8, e.name, ".hdr")) continue; + const bytes = try d.readFileAlloc(io, e.name, a, .limited(1 << 28)); + defer a.free(bytes); + const shard = e.name[0 .. e.name.len - 4]; + try p.addHeaderDump(shard, bytes); + if (std.mem.startsWith(u8, shard, "layer-")) try blocks.append(a, try std.fmt.parseInt(u32, shard[6..8], 10)); + } + const c: Config = .{}; + for ([_]u32{ 0, 1 }) |rank| { + var pl = try plan.build(a, &c, &p, .{ .rank = rank, .world = 2, .blocks = blocks.items }); + defer pl.deinit(); + for (pl.layers) |*l| { + const h = l.experts[1].layout.histogram(); + std.debug.print("rank {d} {s} {t}: down {t} K2 4/6/8 = {d}/{d}/{d}, gate/up mismatches {d}, experts {d} B, block {d} B\n", .{ rank, l.prefix, l.mode, l.experts[1].layout.kind, h[4], h[6], h[8], l.gate_up_mismatch, l.expertBytes(), l.bytes() }); + } + } +} diff --git a/zig/src/families/deepseek_v41/named.zig b/zig/src/families/deepseek_v41/named.zig new file mode 100644 index 0000000000..d439d03508 --- /dev/null +++ b/zig/src/families/deepseek_v41/named.zig @@ -0,0 +1,261 @@ +//! A rank's weights as named host tensors in the reference loader's forms, named as the reference capture names them. + +const std = @import("std"); +const plan = @import("plan.zig"); +const names = @import("names.zig"); +const stage = @import("stage.zig"); +const Pack = @import("pack.zig").Pack; +const DType = @import("pack.zig").DType; +const Io = std.Io; + +pub const Kind = enum { f32, bf16, f16, i16 }; + +/// One tensor the forward holds: its name (`L2.attn.wq_a.trellis`, `embed`), element kind, shape and bytes. +pub const Named = struct { + name: []const u8, + kind: Kind, + shape: [4]usize, + rank: u8, + bytes: []align(16) u8, +}; + +/// What a native tensor becomes (Loader.native and its callers). +const Out = enum { + /// fp32 values (`native(name)`: bf16 or fp32 storage widened) + f32, + /// bf16 values (`native(name, bf16, bf16_param=True)`: the router gate) + bf16, + /// fp32 holding bf16-rounded values (`native(name, bf16_param=True)`: the indexer's weights_proj) + f32_via_bf16, + /// Engram's q / k: combined into `engram.qk` = bf16(q) * bf16(k) in fp32 + qk, +}; + +const Rule = struct { name: []const u8, out: Out }; + +/// Native suffixes (names.zig blockNatives) -> the capture's field names and forms. +const rules = std.StaticStringMap(Rule).initComptime(.{ + .{ "attn_norm.weight", Rule{ .name = "attn_norm", .out = .f32 } }, + .{ "ffn_norm.weight", Rule{ .name = "ffn_norm", .out = .f32 } }, + .{ "hc_attn_fn", Rule{ .name = "hc_attn.fn", .out = .f32 } }, + .{ "hc_attn_base", Rule{ .name = "hc_attn.base", .out = .f32 } }, + .{ "hc_attn_scale", Rule{ .name = "hc_attn.scale", .out = .f32 } }, + .{ "hc_ffn_fn", Rule{ .name = "hc_ffn.fn", .out = .f32 } }, + .{ "hc_ffn_base", Rule{ .name = "hc_ffn.base", .out = .f32 } }, + .{ "hc_ffn_scale", Rule{ .name = "hc_ffn.scale", .out = .f32 } }, + .{ "attn.q_norm.weight", Rule{ .name = "attn.q_norm", .out = .f32 } }, + .{ "attn.kv_norm.weight", Rule{ .name = "attn.kv_norm", .out = .f32 } }, + .{ "attn.attn_sink", Rule{ .name = "attn.sink", .out = .f32 } }, + .{ "attn.compressor.norm.weight", Rule{ .name = "attn.comp_norm", .out = .f32 } }, + .{ "attn.indexer.k_norm.weight", Rule{ .name = "attn.ix_knorm", .out = .f32 } }, + .{ "attn.indexer.weights_proj.weight", Rule{ .name = "attn.ix_wp", .out = .f32_via_bf16 } }, + .{ "ffn.gate.weight", Rule{ .name = "moe.gate", .out = .bf16 } }, + .{ "ffn.gate.bias", Rule{ .name = "moe.bias", .out = .f32 } }, + .{ "engram.q_weight", Rule{ .name = "engram.qk", .out = .qk } }, + .{ "engram.k_weight", Rule{ .name = "engram.qk", .out = .qk } }, +}); + +/// The capture's name of a block group (AttnW / MoEW / EngramW fields). +fn groupField(g: plan.GroupPlan, gl: u32, buf: []u8) ![]const u8 { + return switch (g.proj) { + .wq_a => "attn.wq_a", + .wkv => "attn.wkv", + .wq_b => "attn.wq_b", + .wo_a => std.fmt.bufPrint(buf, "attn.wo_a.{d}", .{g.slice % gl}), + .wo_b => "attn.wo_b", + .comp_wkv => "attn.comp_wkv", + .comp_wgate => "attn.comp_wgate", + .ix_wq_b => "attn.ix_wq_b", + .ix_wk => "attn.ix_wk", + .shared_w1 => "moe.shared.0.w1", + .shared_w2 => "moe.shared.0.w2", + .shared_w3 => "moe.shared.0.w3", + .engram_wkv => "engram.wkv", + }; +} + +/// fp32 -> bf16, round to nearest even (torch's .to(bfloat16)); NaN stays a quiet NaN. +pub fn bf16Round(x: f32) u16 { + const b: u32 = @bitCast(x); + if (std.math.isNan(x)) return @intCast((b >> 16) | 0x40); + return @intCast((b + 0x7fff + ((b >> 16) & 1)) >> 16); +} + +pub fn bf16Widen(h: u16) f32 { + return @bitCast(@as(u32, h) << 16); +} + +/// One stored element as fp32 (bf16, fp16 and fp32 storage). +fn load(dtype: DType, src: []const u8, i: usize) !f32 { + return switch (dtype) { + .f32 => @bitCast(std.mem.readInt(u32, src[4 * i ..][0..4], .little)), + .bf16 => bf16Widen(std.mem.readInt(u16, src[2 * i ..][0..2], .little)), + .f16 => @floatCast(@as(f16, @bitCast(std.mem.readInt(u16, src[2 * i ..][0..2], .little)))), + else => error.UnexpectedTensor, + }; +} + +pub const Builder = struct { + gpa: std.mem.Allocator, + io: Io, + pack: *const Pack, + out: std.ArrayList(Named) = .empty, + + pub fn deinit(b: *Builder) void { + for (b.out.items) |n| { + b.gpa.free(n.name); + b.gpa.free(n.bytes); + } + b.out.deinit(b.gpa); + } + + fn put(b: *Builder, name: []const u8, kind: Kind, shape: []const usize, bytes: []align(16) u8) !void { + var n: Named = .{ .name = try b.gpa.dupe(u8, name), .kind = kind, .shape = @splat(1), .rank = @intCast(shape.len), .bytes = bytes }; + @memcpy(n.shape[0..shape.len], shape); + try b.out.append(b.gpa, n); + } + + fn copy(b: *Builder, src: []const u8) ![]align(16) u8 { + const d = try b.gpa.alignedAlloc(u8, .@"16", src.len); + @memcpy(d, src); + return d; + } + + /// A group's part as `.trellis`, `.suh`, `.svh`. + fn group(b: *Builder, base: []const u8, g: plan.GroupPlan) !void { + var img = try stage.group(b.gpa, b.io, b.pack, g); + defer img.deinit(b.gpa); + const p = g.part.group; + var nb: [160]u8 = undefined; + try b.put(try std.fmt.bufPrint(&nb, "{s}.trellis", .{base}), .i16, &.{ p.k_tiles, p.n_tiles, p.words }, try b.copy(img.trellis)); + try b.put(try std.fmt.bufPrint(&nb, "{s}.suh", .{base}), .f16, &.{p.k()}, try b.copy(img.suh)); + try b.put(try std.fmt.bufPrint(&nb, "{s}.svh", .{base}), .f16, &.{p.n()}, try b.copy(img.svh)); + } + + /// A routed projection: `.trellis` [E, kt, nt, words] (stack) or `.data` (ragged), `.suh` [E, K], `.svh` [E, N]. + fn experts(b: *Builder, base: []const u8, x: *const plan.ExpertsPlan) !void { + var img = try stage.experts(b.gpa, b.io, b.pack, x); + defer img.deinit(b.gpa); + const l = &x.layout; + const e = l.count(); + var nb: [160]u8 = undefined; + if (l.kind == .stack) { + try b.put(try std.fmt.bufPrint(&nb, "{s}.trellis", .{base}), .i16, &.{ e, l.k_tiles, l.n_tiles, l.words[0] }, try b.copy(img.trellis)); + } else try b.put(try std.fmt.bufPrint(&nb, "{s}.data", .{base}), .i16, &.{@intCast(l.elements)}, try b.copy(img.trellis)); + try b.put(try std.fmt.bufPrint(&nb, "{s}.suh", .{base}), .f16, &.{ e, 16 * @as(usize, l.k_tiles) }, try b.copy(img.suh)); + try b.put(try std.fmt.bufPrint(&nb, "{s}.svh", .{base}), .f16, &.{ e, 16 * @as(usize, l.n_tiles) }, try b.copy(img.svh)); + } + + /// The rank's slice of a native tensor, as stored. + fn raw(b: *Builder, v: plan.NativePlan, rank: u32, world: u32) ![]u8 { + const whole = v.slice == .whole; + const len = if (whole) v.info.nbytes else v.info.nbytes / world; + const bytes = try b.gpa.alloc(u8, @intCast(len)); + errdefer b.gpa.free(bytes); + try b.pack.read(b.io, v.info, .{ .offset = if (whole) 0 else rank * len, .len = len }, bytes); + return bytes; + } + + /// A native tensor in its forward form; Engram's q waits for its k (`pending`). + fn native(b: *Builder, name: []const u8, rule: Rule, v: plan.NativePlan, rank: u32, world: u32, pending: *?[]f32) !void { + const bytes = try b.raw(v, rank, world); + defer b.gpa.free(bytes); + const count = bytes.len / v.info.dtype.size(); + var shape: [4]usize = @splat(1); + const dims = v.info.dims(); + @memcpy(shape[0..dims.len], dims); + if (v.slice != .whole) shape[0] /= world; + switch (rule.out) { + .bf16 => { + const out = try b.gpa.alignedAlloc(u8, .@"16", 2 * count); + for (0..count) |i| std.mem.writeInt(u16, out[2 * i ..][0..2], bf16Round(try load(v.info.dtype, bytes, i)), .little); + try b.put(name, .bf16, shape[0..dims.len], out); + }, + .f32, .f32_via_bf16 => { + const out = try b.gpa.alignedAlloc(u8, .@"16", 4 * count); + for (0..count) |i| { + const x = try load(v.info.dtype, bytes, i); + const y = if (rule.out == .f32) x else bf16Widen(bf16Round(x)); + std.mem.writeInt(u32, out[4 * i ..][0..4], @bitCast(y), .little); + } + try b.put(name, .f32, shape[0..dims.len], out); + }, + .qk => { + const vals = try b.gpa.alloc(f32, count); + for (vals, 0..) |*x, i| x.* = bf16Widen(bf16Round(try load(v.info.dtype, bytes, i))); + if (pending.*) |q| { + defer b.gpa.free(q); + defer b.gpa.free(vals); + if (q.len != vals.len) return error.UnexpectedTensor; + const out = try b.gpa.alignedAlloc(u8, .@"16", 4 * count); + for (q, vals, 0..) |x, y, i| std.mem.writeInt(u32, out[4 * i ..][0..4], @bitCast(x * y), .little); + pending.* = null; + try b.put(name, .f32, shape[0..dims.len], out); + } else pending.* = vals; + }, + } + } + + pub fn layer(b: *Builder, l: *const plan.LayerPlan, o_groups: u32, rank: u32, world: u32) !void { + var nb: [160]u8 = undefined; + var fb: [32]u8 = undefined; + const gl = o_groups / world; + for (l.groups) |g| try b.group(try std.fmt.bufPrint(&nb, "L{d}.{s}", .{ l.index, try groupField(g, gl, &fb) }), g); + var pending: ?[]f32 = null; + defer if (pending) |q| b.gpa.free(q); + for (l.natives) |v| { + const suffix = v.name[l.prefix.len + 1 ..]; + const rule = rules.get(suffix) orelse return error.UnnamedNative; + try b.native(try std.fmt.bufPrint(&nb, "L{d}.{s}", .{ l.index, rule.name }), rule, v, rank, world, &pending); + } + if (pending != null) return error.EngramHalf; + for (&l.experts, [_][]const u8{ "w1", "w2", "w3" }) |*x, w| try b.experts(try std.fmt.bufPrint(&nb, "L{d}.moe.{s}", .{ l.index, w }), x); + } + + /// The vocabulary rows (bf16), the final norm (fp32) and the head's part. + pub fn top(b: *Builder, p: *const plan.Plan) !void { + var none: ?[]f32 = null; + for (p.top) |v| { + const is_embed = std.mem.eql(u8, v.name, "embed.weight"); + const rule: Rule = if (is_embed) .{ .name = "embed", .out = .bf16 } else .{ .name = "norm", .out = .f32 }; + try b.native(rule.name, rule, v, p.rank, p.world, &none); + } + if (p.head) |h| try b.group("head", h); + } + + /// DSpark's heads: main_proj (whole), main_norm and norm fp32, Markov embed / head bf16, confidence head fp32. + pub fn dspark(b: *Builder, p: *const plan.Plan) !void { + const d = p.dspark orelse return; + var none: ?[]f32 = null; + try b.group("dspark.main_proj", d.main_proj); + const out_names = [_][]const u8{ "dspark.main_norm", "dspark.norm", "dspark.markov.w1", "dspark.markov.w2" }; + const outs = [_]Out{ .f32, .f32, .bf16, .bf16 }; + for (d.natives, out_names, outs) |v, nm, o| try b.native(nm, .{ .name = nm, .out = o }, v, p.rank, p.world, &none); + if (d.confidence) |c| try b.native("dspark.conf", .{ .name = "dspark.conf", .out = .f32 }, c, p.rank, p.world, &none); + } + + /// Every planned block, then the model-level tensors (a whole-model plan only). + pub fn all(b: *Builder, p: *const plan.Plan, o_groups: u32) !void { + for (p.layers) |*l| try b.layer(l, o_groups, p.rank, p.world); + try b.top(p); + try b.dspark(p); + } + + pub fn find(b: *const Builder, name: []const u8) ?*const Named { + for (b.out.items) |*n| if (std.mem.eql(u8, n.name, name)) return n; + return null; + } +}; + +const testing = std.testing; + +test "the loader's conversions: bf16 round to nearest even, widening, fp16 storage" { + try testing.expectEqual(@as(u16, 0x3f80), bf16Round(1.0)); + // 1 + 2^-8 is a tie between 1.0 and 1 + 2^-7: even keeps 1.0; 1 + 3 * 2^-8 rounds up + try testing.expectEqual(@as(u16, 0x3f80), bf16Round(1.00390625)); + try testing.expectEqual(@as(u16, 0x3f82), bf16Round(1.01171875)); + try testing.expectEqual(@as(f32, -2.5), bf16Widen(bf16Round(-2.5))); + var h: [2]u8 = undefined; + std.mem.writeInt(u16, &h, @bitCast(@as(f16, 0.333251953125)), .little); + try testing.expectEqual(@as(f32, 0.333251953125), try load(.f16, &h, 0)); +} diff --git a/zig/src/families/deepseek_v41/names.zig b/zig/src/families/deepseek_v41/names.zig new file mode 100644 index 0000000000..16da179483 --- /dev/null +++ b/zig/src/families/deepseek_v41/names.zig @@ -0,0 +1,236 @@ +//! The pack's tensor names, each one's TP split and expected shape, plus the names the text engine leaves unread. + +const std = @import("std"); +const Config = @import("config.zig").Config; +const List = @import("config.zig").List; +const Split = @import("exl3.zig").Split; +const DType = @import("core").safetensors.DType; + +/// "layers.L" for the backbone, "mtp.i" for DSpark block i (layer = layers + i). +pub fn blockPrefix(c: *const Config, layer: u32, buf: []u8) ![]const u8 { + return if (c.isBackbone(layer)) std.fmt.bufPrint(buf, "layers.{d}", .{layer}) else std.fmt.bufPrint(buf, "mtp.{d}", .{layer - c.layers}); +} + +/// An EXL3 projection of a block, by the forward's name for it. +pub const Proj = enum { wq_a, wkv, wq_b, wo_a, wo_b, comp_wkv, comp_wgate, ix_wq_b, ix_wk, shared_w1, shared_w2, shared_w3, engram_wkv }; + +/// A block's EXL3 group: its name under the block prefix, how ranks split it, and its K x N (0: not checked). +pub const GroupSpec = struct { + proj: Proj, + suffix: []const u8, + /// wo_a: the global group index this slice is + slice: u32 = 0, + split: Split, + k: u32 = 0, + n: u32 = 0, +}; + +pub const Groups = List(GroupSpec, 32); + +/// The block's EXL3 groups this rank loads, in the loader's order (wo_a: only this rank's groups). +pub fn blockGroups(c: *const Config, layer: u32, rank: u32, world: u32, has_shared: bool) !Groups { + var g: Groups = .{}; + const per_group_k = c.heads / c.o_groups * c.head_dim; + try g.push(.{ .proj = .wq_a, .suffix = "attn.wq_a", .split = .whole, .k = c.hidden, .n = c.q_lora }); + try g.push(.{ .proj = .wkv, .suffix = "attn.wkv", .split = .whole, .k = c.hidden, .n = c.head_dim }); + try g.push(.{ .proj = .wq_b, .suffix = "attn.wq_b", .split = .col, .k = c.q_lora, .n = c.heads * c.head_dim }); + const gl = c.o_groups / world; + var s = rank * gl; + while (s < (rank + 1) * gl) : (s += 1) try g.push(.{ .proj = .wo_a, .suffix = "attn.wo_a.slice", .slice = s, .split = .whole, .k = per_group_k, .n = c.o_lora }); + try g.push(.{ .proj = .wo_b, .suffix = "attn.wo_b", .split = .row, .k = c.o_groups * c.o_lora, .n = c.hidden }); + if (c.isBackbone(layer) and c.isKvSource(layer)) { + try g.push(.{ .proj = .comp_wkv, .suffix = "attn.compressor.wkv", .split = .whole, .k = c.hidden, .n = c.head_dim }); + if (c.compressRatio(layer) > 1) try g.push(.{ .proj = .comp_wgate, .suffix = "attn.compressor.wgate", .split = .whole, .k = c.hidden, .n = c.head_dim }); + try g.push(.{ .proj = .ix_wk, .suffix = "attn.indexer.wk", .split = .whole, .k = c.head_dim, .n = c.index_dim }); + } + if (c.isBackbone(layer) and c.isIndexSource(layer)) + try g.push(.{ .proj = .ix_wq_b, .suffix = "attn.indexer.wq_b", .split = .whole, .k = c.q_lora, .n = c.index_heads * c.index_dim }); + if (has_shared) { + try g.push(.{ .proj = .shared_w1, .suffix = "ffn.shared_experts.w1", .split = .col, .k = c.hidden, .n = c.expert_width }); + try g.push(.{ .proj = .shared_w2, .suffix = "ffn.shared_experts.w2", .split = .row, .k = c.expert_width, .n = c.hidden }); + try g.push(.{ .proj = .shared_w3, .suffix = "ffn.shared_experts.w3", .split = .col, .k = c.hidden, .n = c.expert_width }); + } + if (c.isEngram(layer)) try g.push(.{ .proj = .engram_wkv, .suffix = "engram.wkv", .split = .whole, .k = c.hashCols() * c.engram_head_dim, .n = (c.engram_max_ngram + 1) * c.hidden }); + return g; +} + +/// The group's full name under `block`. +pub fn groupName(spec: GroupSpec, block: []const u8, buf: []u8) ![]const u8 { + return if (spec.proj == .wo_a) std.fmt.bufPrint(buf, "{s}.{s}.{d}", .{ block, spec.suffix, spec.slice }) else std.fmt.bufPrint(buf, "{s}.{s}", .{ block, spec.suffix }); +} + +/// Routed experts' projections: gate (w1) and up (w3) split by columns, down (w2) by rows. +pub const ExpertProj = enum { + w1, + w2, + w3, + + pub fn split(p: ExpertProj) Split { + return if (p == .w2) .row else .col; + } + + /// (K, N) of one expert's full matrix. + pub fn kn(p: ExpertProj, c: *const Config) [2]u32 { + return if (p == .w2) .{ c.expert_width, c.hidden } else .{ c.hidden, c.expert_width }; + } +}; + +pub fn expertName(block: []const u8, e: u32, p: ExpertProj, buf: []u8) ![]const u8 { + return std.fmt.bufPrint(buf, "{s}.ffn.experts.{d}.{t}", .{ block, e, p }); +} + +/// Which dtypes a native tensor may be stored in; the loader converts to what the forward holds. +pub const Store = enum { + /// norms: bf16 in the pack (fp32 in the forward), fp32 accepted + norm, + /// fp32 only (sinks, hc_*, router bias) + f32, + /// fp16 or bf16, rounded to bf16 once (router gate, indexer weights_proj, Markov heads) + half, + + pub fn accepts(s: Store, d: DType) bool { + return switch (s) { + .norm => d == .bf16 or d == .f32, + .f32 => d == .f32, + .half => d == .bf16 or d == .f16, + }; + } +}; + +/// How a rank holds a native tensor. +pub const Slice = enum { whole, heads, vocab_rows }; + +pub const NativeSpec = struct { + suffix: []const u8, + store: Store, + shape: [2]u32, + rank: u8, + slice: Slice = .whole, + + fn v(suffix: []const u8, store: Store, n: u32) NativeSpec { + return .{ .suffix = suffix, .store = store, .shape = .{ n, 0 }, .rank = 1 }; + } + + fn m(suffix: []const u8, store: Store, r: u32, cols: u32) NativeSpec { + return .{ .suffix = suffix, .store = store, .shape = .{ r, cols }, .rank = 2 }; + } + + pub fn dims(s: *const NativeSpec) []const u32 { + return s.shape[0..s.rank]; + } +}; + +pub const Natives = List(NativeSpec, 24); + +/// A block's native tensors: norms, mHC mixers, the sink, compressor / indexer pieces, the router, Engram's q.k. +pub fn blockNatives(c: *const Config, layer: u32) !Natives { + var n: Natives = .{}; + const hc = c.hc_mult; + const mix = (2 + hc) * hc; // pre (hc), post (hc) and the hc x hc residual mix + try n.push(.v("attn_norm.weight", .norm, c.hidden)); + try n.push(.v("ffn_norm.weight", .norm, c.hidden)); + inline for (.{ "attn", "ffn" }) |w| { + try n.push(.m("hc_" ++ w ++ "_fn", .f32, mix, hc * c.hidden)); + try n.push(.v("hc_" ++ w ++ "_base", .f32, mix)); + try n.push(.v("hc_" ++ w ++ "_scale", .f32, 3)); + } + try n.push(.v("attn.q_norm.weight", .norm, c.q_lora)); + try n.push(.v("attn.kv_norm.weight", .norm, c.head_dim)); + var sink = NativeSpec.v("attn.attn_sink", .f32, c.heads); + sink.slice = .heads; + try n.push(sink); + if (c.isBackbone(layer) and c.isKvSource(layer)) { + try n.push(.v("attn.compressor.norm.weight", .norm, c.head_dim)); + try n.push(.v("attn.indexer.k_norm.weight", .norm, c.index_dim)); + } + if (c.isBackbone(layer) and c.isIndexSource(layer)) try n.push(.m("attn.indexer.weights_proj.weight", .half, c.index_heads, c.hidden)); + const ex = c.expertsOf(layer); + try n.push(.m("ffn.gate.weight", .half, ex.count, c.hidden)); + try n.push(.v("ffn.gate.bias", .f32, ex.count)); + if (c.isEngram(layer)) { + try n.push(.m("engram.q_weight", .norm, hc, c.hidden)); + try n.push(.m("engram.k_weight", .norm, hc, c.hidden)); + } + return n; +} + +/// Model-level tensors outside the blocks: this rank's vocabulary rows, the final norm (the head is an EXL3 group). +pub fn topNatives(c: *const Config) !Natives { + var n: Natives = .{}; + var embed = NativeSpec.m("embed.weight", .half, c.vocab, c.hidden); + embed.slice = .vocab_rows; + try n.push(embed); + try n.push(.v("norm.weight", .norm, c.hidden)); + return n; +} + +pub const head = GroupSpec{ .proj = .wq_a, .suffix = "head", .split = .col }; + +/// DSpark's heads beside its blocks (the latest mtp.i that has each). +pub fn dsparkNatives(c: *const Config) !Natives { + var n: Natives = .{}; + try n.push(.v("main_norm.weight", .norm, c.hidden)); + try n.push(.v("norm.weight", .norm, c.hidden)); + try n.push(.m("markov_head.embed.weight", .half, c.vocab, c.dspark_markov_rank)); + try n.push(.m("markov_head.head.weight", .half, c.vocab, c.dspark_markov_rank)); + return n; +} + +/// The optional confidence head: fp32 values [D + r], any shape with that many. +pub const dspark_confidence = "confidence_head.proj.weight"; +pub const dspark_main_proj = "main_proj"; + +/// Names the text engine never reads: vision tower, image markers, image routing bias, bf16 copies of EXL3 groups. +pub fn ignored(name: []const u8, has: anytype) bool { + for ([_][]const u8{ "vision.", "aligner.", "image_" }) |p| if (std.mem.startsWith(u8, name, p)) return true; + if (std.mem.endsWith(u8, name, ".ffn.gate.bias_vl")) return true; + if (std.mem.endsWith(u8, name, ".weight")) { + var buf: [256]u8 = undefined; + const base = name[0 .. name.len - ".weight".len]; + const sib = std.fmt.bufPrint(&buf, "{s}.trellis", .{base}) catch return false; + return has.has(sib); + } + return false; +} + +const testing = std.testing; + +test "the release's block 2: a ratio-2 KV and index source with its compressor gate" { + const c: Config = .{}; + var buf: [64]u8 = undefined; + try testing.expectEqualStrings("mtp.1", try blockPrefix(&c, 41, &buf)); + const g = try blockGroups(&c, 2, 1, 2, true); + var names: [32][]const u8 = undefined; + var nbuf: [32][96]u8 = undefined; + for (g.items(), 0..) |s, i| names[i] = try groupName(s, "layers.2", &nbuf[i]); + // wo_a: rank 1 of 2 holds groups 4..7 + try testing.expectEqualStrings("layers.2.attn.wo_a.slice.4", names[3]); + try testing.expectEqual(@as(usize, 3 + 4 + 1 + 3 + 1 + 3), g.len); + try testing.expectEqual(@as(u32, 4096), g.items()[3].k); + try testing.expectEqual(@as(u32, 8192), g.items()[7].k); + // layer 20: ratio 1, no compressor gate; layer 24: reindex (own wq_b only); layer 30: reuse; DSpark blocks: neither + const g20 = try blockGroups(&c, 20, 0, 2, true); + for (g20.items()) |s| try testing.expect(s.proj != .comp_wgate); + try testing.expectEqual(@as(usize, 3 + 4 + 1 + 1), (try blockGroups(&c, 24, 0, 2, false)).len); + try testing.expectEqual(@as(usize, 3 + 4 + 1), (try blockGroups(&c, 30, 0, 2, false)).len); + try testing.expectEqual(@as(usize, 3 + 4 + 1 + 3), (try blockGroups(&c, 41, 0, 2, true)).len); + try testing.expect((try blockGroups(&c, 14, 0, 2, true)).items()[15].proj == .engram_wkv); + try testing.expectEqualStrings("layers.3.ffn.experts.17.w2", try expertName("layers.3", 17, .w2, &buf)); +} + +test "natives and the ignore list" { + const c: Config = .{}; + const n = try blockNatives(&c, 41); + for (n.items()) |s| if (std.mem.eql(u8, s.suffix, "ffn.gate.weight")) try testing.expectEqual(@as(u32, 128), s.shape[0]); + try testing.expectEqual(@as(u32, 24), (try blockNatives(&c, 0)).items()[2].shape[0]); + const Has = struct { + pub fn has(_: @This(), name: []const u8) bool { + return std.mem.eql(u8, name, "layers.2.attn.indexer.wk.trellis"); + } + }; + try testing.expect(ignored("layers.2.attn.indexer.wk.weight", Has{})); + try testing.expect(!ignored("layers.2.attn.kv_norm.weight", Has{})); + try testing.expect(ignored("vision.blocks.3.mlp.w1.weight", Has{})); + try testing.expect(ignored("layers.9.ffn.gate.bias_vl", Has{})); + try testing.expect(Store.half.accepts(.f16) and !Store.f32.accepts(.bf16)); +} diff --git a/zig/src/families/deepseek_v41/pack.zig b/zig/src/families/deepseek_v41/pack.zig new file mode 100644 index 0000000000..131978bcd3 --- /dev/null +++ b/zig/src/families/deepseek_v41/pack.zig @@ -0,0 +1,306 @@ +//! A DeepSeek-V4.1 pack folder by its safetensors headers alone, read with pread and no mmap. + +const std = @import("std"); +const quiet = @import("builtin").is_test; // a test's stderr is the build runner's channel +const log = std.log.scoped(.dsv41); +const core = @import("core"); +const exl3 = @import("exl3.zig"); +const st = core.safetensors; +const Io = std.Io; + +pub const DType = st.DType; + +/// One stored tensor: its shard, dtype, shape and where its bytes start in that file. +pub const Info = struct { + file: u32, + dtype: DType, + rank: u8, + shape: [st.max_rank]usize, + /// absolute byte offset of the first byte in the shard file + start: u64, + nbytes: u64, + + pub fn dims(i: *const Info) []const usize { + return i.shape[0..i.rank]; + } + + pub fn is(i: *const Info, dtype: DType, shape: []const usize) bool { + return i.dtype == dtype and std.mem.eql(usize, i.dims(), shape); + } +}; + +pub const Pack = struct { + arena: std.heap.ArenaAllocator, + /// the folder, else "" for header-only packs (tests, header dumps) + dir: []const u8 = "", + /// shards the index lists that the folder does not hold + missing: u32 = 0, + files: std.ArrayList([]const u8) = .empty, + names: std.StringArrayHashMapUnmanaged(Info) = .empty, + /// the index's weight map (name -> shard); when set, names another shard owns are skipped + owner: std.StringHashMapUnmanaged([]const u8) = .empty, + + pub fn init(gpa: std.mem.Allocator) Pack { + return .{ .arena = std.heap.ArenaAllocator.init(gpa) }; + } + + pub fn deinit(p: *Pack) void { + p.arena.deinit(); + p.* = undefined; + } + + /// Every shard's header in `dir`: the index's weight map if present, else every *.safetensors file. + pub fn open(gpa: std.mem.Allocator, io: Io, dir: []const u8) !Pack { + var p = init(gpa); + errdefer p.deinit(); + const a = p.arena.allocator(); + p.dir = try a.dupe(u8, dir); + const cwd = Io.Dir.cwd(); + var shards: std.ArrayList([]const u8) = .empty; + const index = try std.fs.path.join(a, &.{ dir, "model.safetensors.index.json" }); + if (cwd.readFileAlloc(io, index, a, .limited(1 << 28))) |text| { + try p.useIndex(text); + var it = p.owner.valueIterator(); + while (it.next()) |f| { + for (shards.items) |s| { + if (std.mem.eql(u8, s, f.*)) break; + } else try shards.append(a, f.*); + } + } else |err| switch (err) { + error.FileNotFound => { + var d = try cwd.openDir(io, dir, .{ .iterate = true }); + defer d.close(io); + var it = d.iterate(); + while (try it.next(io)) |e| if (std.mem.endsWith(u8, e.name, ".safetensors")) try shards.append(a, try a.dupe(u8, e.name)); + }, + else => return err, + } + if (shards.items.len == 0) return error.NoSafetensors; + std.mem.sort([]const u8, shards.items, {}, lessStr); + for (shards.items) |name| { + const path = try std.fs.path.join(a, &.{ dir, name }); + // a shard the index lists but the folder lacks: its tensors stay unknown (a plan needing one fails) + const file = cwd.openFile(io, path, .{}) catch |e| switch (e) { + error.FileNotFound => { + p.missing += 1; + continue; + }, + else => return e, + }; + defer file.close(io); + var head: [8]u8 = undefined; + if (try file.readPositionalAll(io, &head, 0) != 8) return error.BadSafetensors; + const n = std.mem.readInt(u64, &head, .little); + const len = try file.length(io); + if (n > len - 8 or n > 1 << 30) return error.BadSafetensors; + const json = try a.alloc(u8, @intCast(n)); + if (try file.readPositionalAll(io, json, 8) != json.len) return error.BadSafetensors; + try p.addHeader(name, json, len - 8 - n); + } + return p; + } + + /// model.safetensors.index.json's weight map, before any header is added. + pub fn useIndex(p: *Pack, json: []const u8) !void { + const a = p.arena.allocator(); + const parsed = try std.json.parseFromSliceLeaky(std.json.Value, a, json, .{}); + if (parsed != .object) return error.BadIndex; + const map = (parsed.object.get("weight_map") orelse return error.BadIndex); + if (map != .object) return error.BadIndex; + try p.owner.ensureTotalCapacity(a, @intCast(map.object.count())); + var it = map.object.iterator(); + while (it.next()) |e| { + if (e.value_ptr.* != .string) return error.BadIndex; + p.owner.putAssumeCapacity(e.key_ptr.*, e.value_ptr.string); + } + } + + /// A shard's header JSON (`data_len` bytes follow; maxInt(u64) for dumps); a name already indexed is refused. + pub fn addHeader(p: *Pack, file_name: []const u8, json: []const u8, data_len: u64) !void { + const a = p.arena.allocator(); + const id: u32 = @intCast(p.files.items.len); + try p.files.append(a, try a.dupe(u8, file_name)); + const header = try st.parseHeader(a, json, @intCast(data_len)); + try p.names.ensureUnusedCapacity(a, header.count()); + var it = header.iterator(); + while (it.next()) |e| { + const v = e.value_ptr.*; + if (p.owner.count() > 0) { + const f = p.owner.get(e.key_ptr.*) orelse continue; + if (!std.mem.eql(u8, f, file_name)) continue; + } + const gop = p.names.getOrPutAssumeCapacity(e.key_ptr.*); + if (gop.found_existing) return error.DuplicateTensor; + gop.value_ptr.* = .{ .file = id, .dtype = v.dtype, .rank = v.rank, .shape = v.shape, .start = 8 + json.len + v.begin, .nbytes = v.end - v.begin }; + } + } + + /// A header dump (the 8-byte length, then the JSON; nothing after) as shard `file_name`. + pub fn addHeaderDump(p: *Pack, file_name: []const u8, bytes: []const u8) !void { + if (bytes.len < 8) return error.BadSafetensors; + const n = std.mem.readInt(u64, bytes[0..8], .little); + if (n > bytes.len - 8) return error.BadSafetensors; + try p.addHeader(file_name, bytes[8..][0..@intCast(n)], std.math.maxInt(u64)); + } + + pub fn get(p: *const Pack, name: []const u8) ?Info { + return p.names.get(name); + } + + pub fn has(p: *const Pack, name: []const u8) bool { + return p.names.contains(name); + } + + /// A required tensor; a missing name is logged and refused. + pub fn need(p: *const Pack, name: []const u8) !Info { + return p.get(name) orelse { + if (!quiet) log.warn("pack has no tensor {s}", .{name}); + return error.MissingTensor; + }; + } + + /// The EXL3 group at `prefix` (its trellis entry), checked int16 [K/16, N/16, words] with fp16 suh [K] and svh [N]. + pub fn group(p: *const Pack, prefix: []const u8, buf: []u8) !exl3.Group { + const t = try p.need(try std.fmt.bufPrint(buf, "{s}.trellis", .{prefix})); + if (t.dtype != .i16) return error.BadTrellis; + const g = try exl3.Group.fromShape(t.dims()); + const suh = try p.need(try std.fmt.bufPrint(buf, "{s}.suh", .{prefix})); + const svh = try p.need(try std.fmt.bufPrint(buf, "{s}.svh", .{prefix})); + if (!suh.is(.f16, &.{g.k()}) or !svh.is(.f16, &.{g.n()})) { + if (!quiet) log.warn("{s}: suh/svh do not match trellis K={d} N={d}", .{ prefix, g.k(), g.n() }); + return error.BadTrellis; + } + return g; + } + + /// The group's marker tensor name (mul1 or mcg) if it has one; its value picks the codebook at load time. + pub fn marker(p: *const Pack, prefix: []const u8, buf: []u8) !?[]const u8 { + for ([_][]const u8{ "mul1", "mcg" }) |m| { + const name = try std.fmt.bufPrint(buf, "{s}.{s}", .{ prefix, m }); + if (p.names.getKey(name)) |k| return k; + } + return null; + } + + /// `span` of tensor `info`'s bytes into `out` (span.bytes() long), run by run. + pub fn read(p: *const Pack, io: Io, info: Info, span: exl3.Span, out: []u8) !void { + if (out.len != span.bytes()) return error.BadSpan; + if (span.count > 0 and span.at(span.count - 1) + span.len > info.nbytes) return error.BadSpan; + if (p.dir.len == 0) return error.HeaderOnlyPack; + var path_buf: [Io.Dir.max_path_bytes]u8 = undefined; + const path = try std.fmt.bufPrint(&path_buf, "{s}/{s}", .{ p.dir, p.files.items[info.file] }); + const file = try Io.Dir.cwd().openFile(io, path, .{}); + defer file.close(io); + if (span.count <= 1) { + if (try file.readPositionalAll(io, out, info.start + span.offset) != out.len) return error.ShortRead; + dropCache(file, info.start + span.offset, out.len); + return; + } + // a strided part (a column split): one read covering every run, then the runs gathered + const first = span.at(0); + const cover: usize = @intCast(span.at(span.count - 1) + span.len - first); + const tmp = try std.heap.page_allocator.alloc(u8, cover); + defer std.heap.page_allocator.free(tmp); + if (try file.readPositionalAll(io, tmp, info.start + first) != cover) return error.ShortRead; + dropCache(file, info.start + first, cover); + const len: usize = @intCast(span.len); + const stride: usize = @intCast(span.stride); + for (0..@intCast(span.count)) |i| @memcpy(out[i * len ..][0..len], tmp[i * stride ..][0..len]); + } + + /// Drops the page cache of bytes just read (POSIX_FADV_DONTNEED): on unified memory it competes with the device. + fn dropCache(file: Io.File, offset: u64, len: usize) void { + if (@import("builtin").os.tag != .linux) return; + const linux = std.os.linux; + _ = linux.fadvise(file.handle, @intCast(offset), @intCast(len), linux.POSIX_FADV.DONTNEED); + } + + /// The codebook of the group at `prefix`: its marker's value, else the 3-instruction default. + pub fn codebook(p: *const Pack, io: Io, prefix: []const u8) !exl3.Codebook { + var buf: [256]u8 = undefined; + const name = (try p.marker(prefix, &buf)) orelse return .inst3; + const info = p.get(name).?; + if (info.nbytes < 4) return error.BadMarker; + var v: [4]u8 = undefined; + try p.read(io, info, .{ .offset = 0, .len = 4 }, &v); + return exl3.Codebook.fromMarker(std.mem.readInt(u32, &v, .little)); + } +}; + +fn lessStr(_: void, x: []const u8, y: []const u8) bool { + return std.mem.order(u8, x, y) == .lt; +} + +const testing = std.testing; + +/// One tensor of an in-memory image. +pub const Spec = struct { name: []const u8, dtype: []const u8, shape: []const usize, fill: u8 = 0 }; + +/// A safetensors image in memory: tensors in order, each filled with `fill`. +pub fn image(a: std.mem.Allocator, specs: []const Spec) ![]u8 { + var json: std.Io.Writer.Allocating = .init(a); + defer json.deinit(); + try json.writer.writeAll("{\"__metadata__\":{\"format\":\"pt\"}"); + var at: usize = 0; + for (specs) |s| { + var n = (DType.parse(s.dtype) orelse return error.UnsupportedDType).size(); + for (s.shape) |d| n *= d; + try json.writer.print(",\"{s}\":{{\"dtype\":\"{s}\",\"shape\":[", .{ s.name, s.dtype }); + for (s.shape, 0..) |d, i| try json.writer.print("{s}{d}", .{ if (i == 0) "" else ",", d }); + try json.writer.print("],\"data_offsets\":[{d},{d}]}}", .{ at, at + n }); + at += n; + } + try json.writer.writeAll("}"); + const head = json.written(); + const out = try a.alloc(u8, 8 + head.len + at); + std.mem.writeInt(u64, out[0..8], head.len, .little); + @memcpy(out[8..][0..head.len], head); + var pos = 8 + head.len; + for (specs) |s| { + var n = DType.parse(s.dtype).?.size(); + for (s.shape) |d| n *= d; + @memset(out[pos..][0..n], s.fill); + pos += n; + } + return out; +} + +test "a folder by its headers: index order, groups checked, strided reads, codebook markers" { + const io = testing.io; + var tmp = testing.tmpDir(.{}); + defer tmp.cleanup(); + const a = testing.allocator; + const one = try image(a, &.{ + .{ .name = "x.trellis", .dtype = "I16", .shape = &.{ 2, 16, 16 }, .fill = 0 }, + .{ .name = "x.suh", .dtype = "F16", .shape = &.{32} }, + .{ .name = "x.svh", .dtype = "F16", .shape = &.{256} }, + .{ .name = "x.mcg", .dtype = "I32", .shape = &.{}, .fill = 0 }, + }); + defer a.free(one); + // tile (kt, nt) holds bytes of value 16 kt + nt, so a column part's runs are checkable + const hl = std.mem.readInt(u64, one[0..8], .little); + const data = one[8 + hl ..]; + for (0..2) |kt| for (0..16) |nt| @memset(data[(kt * 16 + nt) * 32 ..][0..32], @intCast(16 * kt + nt)); + std.mem.writeInt(u32, data[2 * 16 * 32 + 64 + 512 ..][0..4], exl3.Codebook.marker_mcg, .little); + const two = try image(a, &.{.{ .name = "y.weight", .dtype = "BF16", .shape = &.{ 3, 2 }, .fill = 7 }}); + defer a.free(two); + try tmp.dir.writeFile(io, .{ .sub_path = "b.safetensors", .data = one }); + try tmp.dir.writeFile(io, .{ .sub_path = "a.safetensors", .data = two }); + const dir = try tmp.dir.realPathFileAlloc(io, ".", a); + defer a.free(dir); + var p = try Pack.open(a, io, dir); + defer p.deinit(); + try testing.expectEqualStrings("a.safetensors", p.files.items[0]); + var buf: [128]u8 = undefined; + const g = try p.group("x", &buf); + try testing.expectEqual(@as(u32, 256), g.n()); + const part = try g.part(.col, 1, 2); + const out = try a.alloc(u8, @intCast(part.trellis.bytes())); + defer a.free(out); + try p.read(io, p.get("x.trellis").?, part.trellis, out); + for (0..2) |kt| for (0..8) |j| try testing.expectEqual(@as(u8, @intCast(16 * kt + 8 + j)), out[(kt * 8 + j) * 32]); + try testing.expectEqual(exl3.Codebook.mcg, try p.codebook(io, "x")); + try testing.expectEqual(exl3.Codebook.inst3, try p.codebook(io, "y")); + try testing.expectError(error.MissingTensor, p.group("y", &buf)); + try testing.expectError(error.BadSpan, p.read(io, p.get("y.weight").?, .{ .offset = 8, .len = 8 }, out[0..8])); +} diff --git a/zig/src/families/deepseek_v41/plan.zig b/zig/src/families/deepseek_v41/plan.zig new file mode 100644 index 0000000000..e730a4ee2d --- /dev/null +++ b/zig/src/families/deepseek_v41/plan.zig @@ -0,0 +1,295 @@ +//! A rank's weights planned from the pack's headers alone: group TP parts, routed layouts, byte totals, coverage. + +const std = @import("std"); +const quiet = @import("builtin").is_test; // a test's stderr is the build runner's channel +const log = std.log.scoped(.dsv41); +const Config = @import("config.zig").Config; +const exl3 = @import("exl3.zig"); +const names = @import("names.zig"); +const Pack = @import("pack.zig").Pack; +const Info = @import("pack.zig").Info; + +/// One EXL3 group of the rank: its stored name and the rank's part. +pub const GroupPlan = struct { + proj: names.Proj, + prefix: []const u8, + split: exl3.Split, + full: exl3.Group, + part: exl3.Part, + /// marker tensor (mul1 / mcg) whose value picks the codebook at load; null: the 3-instruction default + marker: ?[]const u8, + /// wo_a: the global group index this slice is + slice: u32 = 0, +}; + +pub const NativePlan = struct { + name: []const u8, + info: Info, + slice: names.Slice, + /// bytes of the rank's slice as stored + bytes: u64, +}; + +/// One routed projection: each expert's group part, their layout on the device, and the stored prefixes. +pub const ExpertsPlan = struct { + proj: names.ExpertProj, + parts: []exl3.Part, + prefixes: [][]const u8, + layout: exl3.Layout, +}; + +pub const LayerPlan = struct { + index: u32, + prefix: []const u8, + mode: @import("config.zig").Mode, + groups: []GroupPlan, + natives: []NativePlan, + experts: [3]ExpertsPlan, + /// experts whose gate (w1) and up (w3) widths differ: x3gm cannot fuse them, the block goes to x3pf / grouped + gate_up_mismatch: u32, + + /// Routed trellis bytes plus the shared expert's: the quant repo's `rank_bytes`. + pub fn expertBytes(l: *const LayerPlan) u64 { + var n: u64 = 0; + for (l.experts) |e| n += e.layout.trellisBytes(); + for (l.groups) |g| switch (g.proj) { + .shared_w1, .shared_w2, .shared_w3 => n += g.part.trellis.bytes(), + else => {}, + }; + return n; + } + + /// Everything the rank holds for the block as stored: trellises, suh/svh, natives. + pub fn bytes(l: *const LayerPlan) u64 { + var n: u64 = 0; + for (l.groups) |g| n += g.part.bytes(); + for (l.natives) |v| n += v.bytes; + for (l.experts) |e| { + n += e.layout.trellisBytes(); + for (e.parts) |p| n += p.suh.bytes() + p.svh.bytes(); + } + return n; + } +}; + +pub const Dspark = struct { + main_proj: GroupPlan, + natives: []NativePlan, + confidence: ?NativePlan, +}; + +pub const Plan = struct { + arena: std.heap.ArenaAllocator, + rank: u32, + world: u32, + layers: []LayerPlan, + top: []NativePlan, + /// null when only some blocks were planned + head: ?GroupPlan, + dspark: ?Dspark, + /// stored names neither planned nor on the ignore list (the loader refuses a pack with any) + leftovers: [][]const u8, + + pub fn deinit(p: *Plan) void { + p.arena.deinit(); + p.* = undefined; + } + + pub fn bytes(p: *const Plan) u64 { + var n: u64 = if (p.head) |h| h.part.bytes() else 0; + for (p.top) |v| n += v.bytes; + for (p.layers) |*l| n += l.bytes(); + if (p.dspark) |d| { + n += d.main_proj.part.bytes(); + for (d.natives) |v| n += v.bytes; + if (d.confidence) |c| n += c.bytes; + } + return n; + } +}; + +pub const Options = struct { + rank: u32, + world: u32, + /// plan only these blocks (single-layer runs, header dumps): no top tensors, head, DSpark or coverage; null: all + blocks: ?[]const u32 = null, + /// with `blocks`: also the vocabulary rows, the final norm and the head (a window's finish) + top: bool = false, + /// with `blocks` and `top`: also DSpark's heads (main_proj, its norms, the Markov and confidence heads) + dspark: bool = false, +}; + +const Builder = struct { + a: std.mem.Allocator, + c: *const Config, + pack: *const Pack, + o: Options, + used: std.StringHashMapUnmanaged(void) = .empty, + buf: [256]u8 = undefined, + + fn mark(b: *Builder, name: []const u8) !void { + try b.used.put(b.a, b.pack.names.getKey(name) orelse return error.MissingTensor, {}); + } + + fn dupe(b: *Builder, s: []const u8) ![]const u8 { + return b.a.dupe(u8, s); + } + + fn group(b: *Builder, proj: names.Proj, prefix: []const u8, split: exl3.Split, k: u32, n: u32) !GroupPlan { + const g = try b.pack.group(prefix, &b.buf); + if ((k != 0 and g.k() != k) or (n != 0 and g.n() != n)) { + if (!quiet) log.warn("{s}: K={d} N={d}, expected K={d} N={d}", .{ prefix, g.k(), g.n(), k, n }); + return error.UnexpectedShape; + } + inline for (.{ "trellis", "suh", "svh" }) |part| try b.mark(try std.fmt.bufPrint(&b.buf, "{s}." ++ part, .{prefix})); + const marker = try b.pack.marker(prefix, &b.buf); + if (marker) |m| try b.mark(m); + return .{ .proj = proj, .prefix = try b.dupe(prefix), .split = split, .full = g, .part = try g.part(split, b.o.rank, b.o.world), .marker = if (marker) |m| try b.dupe(m) else null }; + } + + fn native(b: *Builder, name: []const u8, spec: names.NativeSpec) !NativePlan { + const info = try b.pack.need(name); + const want = spec.dims(); + var ok = spec.store.accepts(info.dtype) and info.rank == want.len; + if (ok) for (want, info.dims()) |w, d| { + ok = ok and w == d; + }; + if (!ok) { + if (!quiet) log.warn("{s}: {t} {any}, expected {t} storage of {any}", .{ name, info.dtype, info.dims(), spec.store, want }); + return error.UnexpectedTensor; + } + try b.mark(name); + const bytes = switch (spec.slice) { + .whole => info.nbytes, + .heads, .vocab_rows => blk: { + if (info.shape[0] % b.o.world != 0) return error.SplitNotEven; + break :blk info.nbytes / b.o.world; + }, + }; + return .{ .name = try b.dupe(name), .info = info, .slice = spec.slice, .bytes = bytes }; + } + + fn experts(b: *Builder, block: []const u8, index: u32, proj: names.ExpertProj) !ExpertsPlan { + const count = b.c.expertsOf(index).count; + const parts = try b.a.alloc(exl3.Part, count); + const prefixes = try b.a.alloc([]const u8, count); + const kn = proj.kn(b.c); + var codebook_marker: ?bool = null; + for (parts, prefixes, 0..) |*p, *pre, e| { + var nb: [128]u8 = undefined; + const name = try names.expertName(block, @intCast(e), proj, &nb); + const g = try b.group(.wq_a, name, proj.split(), kn[0], kn[1]); + // the codebook is one a projection; here the marker's presence, its value at load + const has = g.marker != null; + if (codebook_marker) |m| if (m != has) return error.CodebooksDiffer; + codebook_marker = has; + p.* = g.part; + pre.* = g.prefix; + } + return .{ .proj = proj, .parts = parts, .prefixes = prefixes, .layout = try exl3.Layout.of(b.a, parts) }; + } + + fn layer(b: *Builder, index: u32) !LayerPlan { + var pb: [32]u8 = undefined; + const block = try b.dupe(try names.blockPrefix(b.c, index, &pb)); + const has_shared = b.pack.has(try std.fmt.bufPrint(&b.buf, "{s}.ffn.shared_experts.w1.trellis", .{block})); + const specs = try names.blockGroups(b.c, index, b.o.rank, b.o.world, has_shared); + const groups = try b.a.alloc(GroupPlan, specs.len); + for (specs.items(), groups) |s, *g| { + var nb: [128]u8 = undefined; + g.* = try b.group(s.proj, try names.groupName(s, block, &nb), s.split, s.k, s.n); + g.slice = s.slice; + } + const nspecs = try names.blockNatives(b.c, index); + const natives = try b.a.alloc(NativePlan, nspecs.len); + for (nspecs.items(), natives) |s, *v| { + var nb: [128]u8 = undefined; + v.* = try b.native(try std.fmt.bufPrint(&nb, "{s}.{s}", .{ block, s.suffix }), s); + } + // the other ranks' wo_a groups are theirs to read: mark them so coverage holds on every rank + var g: u32 = 0; + while (g < b.c.o_groups) : (g += 1) { + var nb: [128]u8 = undefined; + const pre = try std.fmt.bufPrint(&nb, "{s}.attn.wo_a.slice.{d}", .{ block, g }); + inline for (.{ "trellis", "suh", "svh" }) |part| try b.mark(try std.fmt.bufPrint(&b.buf, "{s}." ++ part, .{pre})); + if (try b.pack.marker(pre, &b.buf)) |m| try b.mark(m); + } + var out: LayerPlan = .{ .index = index, .prefix = block, .mode = b.c.mode(index), .groups = groups, .natives = natives, .experts = undefined, .gate_up_mismatch = 0 }; + inline for (.{ .w1, .w2, .w3 }, 0..) |p, i| out.experts[i] = try b.experts(block, index, p); + for (out.experts[0].layout.words, out.experts[2].layout.words) |w1, w3| out.gate_up_mismatch += @intFromBool(w1 != w3); + return out; + } + + /// The latest mtp.i holding `suffix` (as a tensor, a group or a .weight). + fn findDspark(b: *Builder, suffix: []const u8, buf: []u8) !?[]const u8 { + var i = b.c.mtp_layers; + while (i > 0) { + i -= 1; + const base = try std.fmt.bufPrint(buf, "mtp.{d}.{s}", .{ i, suffix }); + for ([_][]const u8{ "", ".trellis", ".weight" }) |ext| { + var nb: [160]u8 = undefined; + if (b.pack.has(try std.fmt.bufPrint(&nb, "{s}{s}", .{ base, ext }))) return base; + } + } + return null; + } + + fn dspark(b: *Builder) !?Dspark { + if (b.c.mtp_layers == 0) return null; + var fb: [160]u8 = undefined; + const mp = (try b.findDspark(names.dspark_main_proj, &fb)) orelse return null; + const main_proj = try b.group(.wq_a, mp, .whole, 3 * b.c.hidden, b.c.hidden); + const specs = try names.dsparkNatives(b.c); + const natives = try b.a.alloc(NativePlan, specs.len); + for (specs.items(), natives) |s, *v| { + const stem = s.suffix[0 .. s.suffix.len - ".weight".len]; + const base = (try b.findDspark(stem, &fb)) orelse return error.MissingTensor; + var nb: [160]u8 = undefined; + v.* = try b.native(try std.fmt.bufPrint(&nb, "{s}.weight", .{base}), s); + } + var confidence: ?NativePlan = null; + if (try b.findDspark("confidence_head.proj", &fb)) |base| { + var nb: [160]u8 = undefined; + const name = try std.fmt.bufPrint(&nb, "{s}.weight", .{base}); + const info = try b.pack.need(name); + var n: u64 = 1; + for (info.dims()) |d| n *= d; + if (info.dtype != .f32 and info.dtype != .bf16 or n != b.c.hidden + b.c.dspark_markov_rank) return error.UnexpectedTensor; + try b.mark(name); + confidence = .{ .name = try b.dupe(name), .info = info, .slice = .whole, .bytes = info.nbytes }; + } + return .{ .main_proj = main_proj, .natives = natives, .confidence = confidence }; + } +}; + +/// The rank's plan; heads, groups and the vocabulary must split over `world` (blocks of 128). +pub fn build(gpa: std.mem.Allocator, c: *const Config, pack: *const Pack, o: Options) !Plan { + if (o.world == 0 or o.rank >= o.world or c.heads % o.world != 0 or c.o_groups % o.world != 0 or c.vocab % (128 * o.world) != 0) return error.BadWorld; + var arena = std.heap.ArenaAllocator.init(gpa); + errdefer arena.deinit(); + var b: Builder = .{ .a = arena.allocator(), .c = c, .pack = pack, .o = o }; + const full = o.blocks == null; + const layers = try b.a.alloc(LayerPlan, if (o.blocks) |bl| bl.len else c.allLayers()); + for (layers, 0..) |*l, i| { + const index: u32 = if (o.blocks) |bl| bl[i] else @intCast(i); + if (index >= c.allLayers()) return error.LayerOutOfBounds; + l.* = try b.layer(index); + } + var top: []NativePlan = &.{}; + var head: ?GroupPlan = null; + var ds: ?Dspark = null; + if (full or o.top) { + const tspecs = try names.topNatives(c); + top = try b.a.alloc(NativePlan, tspecs.len); + for (tspecs.items(), top) |s, *v| v.* = try b.native(s.suffix, s); + head = try b.group(.wq_a, names.head.suffix, .col, c.hidden, c.vocab); + if (full or o.dspark) ds = try b.dspark(); + } + // coverage: every stored name was planned or is knowingly left (vision, image routing, bf16 copies of groups) + var left: std.ArrayList([]const u8) = .empty; + if (full) for (pack.names.keys()) |k| { + if (b.used.contains(k) or names.ignored(k, pack)) continue; + try left.append(b.a, try b.dupe(k)); + }; + return .{ .arena = arena, .rank = o.rank, .world = o.world, .layers = layers, .top = top, .head = head, .dspark = ds, .leftovers = left.items }; +} diff --git a/zig/src/families/deepseek_v41/stage.zig b/zig/src/families/deepseek_v41/stage.zig new file mode 100644 index 0000000000..726618d9b1 --- /dev/null +++ b/zig/src/families/deepseek_v41/stage.zig @@ -0,0 +1,148 @@ +//! A rank's EXL3 bytes read into host images the device buffers copy as is: group parts and ragged routed experts. + +const std = @import("std"); +const exl3 = @import("exl3.zig"); +const plan = @import("plan.zig"); +const Pack = @import("pack.zig").Pack; +const Io = std.Io; + +/// One group's part: trellis int16 [kt, nt, words], suh fp16 [K part], svh fp16 [N part]. +pub const GroupImage = struct { + trellis: []u8, + suh: []u8, + svh: []u8, + + pub fn deinit(g: *GroupImage, gpa: std.mem.Allocator) void { + gpa.free(g.trellis); + gpa.free(g.suh); + gpa.free(g.svh); + g.* = undefined; + } +}; + +pub fn group(gpa: std.mem.Allocator, io: Io, pack: *const Pack, g: plan.GroupPlan) !GroupImage { + var out: GroupImage = .{ .trellis = &.{}, .suh = &.{}, .svh = &.{} }; + errdefer out.deinit(gpa); + var buf: [256]u8 = undefined; + out.trellis = try gpa.alloc(u8, @intCast(g.part.trellis.bytes())); + try pack.read(io, try pack.need(try std.fmt.bufPrint(&buf, "{s}.trellis", .{g.prefix})), g.part.trellis, out.trellis); + out.suh = try gpa.alloc(u8, @intCast(g.part.suh.bytes())); + try pack.read(io, try pack.need(try std.fmt.bufPrint(&buf, "{s}.suh", .{g.prefix})), g.part.suh, out.suh); + out.svh = try gpa.alloc(u8, @intCast(g.part.svh.bytes())); + try pack.read(io, try pack.need(try std.fmt.bufPrint(&buf, "{s}.svh", .{g.prefix})), g.part.svh, out.svh); + return out; +} + +/// A routed projection: experts' trellises back to back, suh [E, K part] and svh [E, N part] stacked, K2 per expert. +pub const ExpertsImage = struct { + trellis: []align(16) u8, + suh: []u8, + svh: []u8, + /// the kernels' int32 K2 table + k2: []i32, + /// byte offset of each expert's trellis in `trellis` (the int64 address table is the device base plus these) + offsets: []u64, + + pub fn deinit(e: *ExpertsImage, gpa: std.mem.Allocator) void { + gpa.free(e.trellis); + gpa.free(e.suh); + gpa.free(e.svh); + gpa.free(e.k2); + gpa.free(e.offsets); + e.* = undefined; + } +}; + +pub fn experts(gpa: std.mem.Allocator, io: Io, pack: *const Pack, x: *const plan.ExpertsPlan) !ExpertsImage { + const n = x.layout.count(); + const kpart = 2 * @as(u64, x.layout.k_tiles) * 16; + const npart = 2 * @as(u64, x.layout.n_tiles) * 16; + var out: ExpertsImage = .{ .trellis = &.{}, .suh = &.{}, .svh = &.{}, .k2 = &.{}, .offsets = &.{} }; + errdefer out.deinit(gpa); + out.trellis = try gpa.alignedAlloc(u8, .@"16", @intCast(x.layout.trellisBytes())); + out.suh = try gpa.alloc(u8, @intCast(n * kpart)); + out.svh = try gpa.alloc(u8, @intCast(n * npart)); + out.k2 = try gpa.alloc(i32, n); + out.offsets = try gpa.alloc(u64, n); + var buf: [256]u8 = undefined; + for (x.parts, x.prefixes, 0..) |p, pre, e| { + const at = x.layout.byteOffset(e); + out.offsets[e] = at; + out.k2[e] = @intCast(p.group.k2()); + const t = out.trellis[@intCast(at)..][0..@intCast(p.trellis.bytes())]; + try pack.read(io, try pack.need(try std.fmt.bufPrint(&buf, "{s}.trellis", .{pre})), p.trellis, t); + try pack.read(io, try pack.need(try std.fmt.bufPrint(&buf, "{s}.suh", .{pre})), p.suh, out.suh[@intCast(e * kpart)..][0..@intCast(kpart)]); + try pack.read(io, try pack.need(try std.fmt.bufPrint(&buf, "{s}.svh", .{pre})), p.svh, out.svh[@intCast(e * npart)..][0..@intCast(npart)]); + } + return out; +} + +const testing = std.testing; +const image = @import("pack.zig").image; + +test "staged images: a column part's runs, ragged experts back to back, stacked scales" { + const a = testing.allocator; + const io = testing.io; + var tmp = testing.tmpDir(.{}); + defer tmp.cleanup(); + // three experts of a gate projection (K 32, N 256): expert 1 at 3 bits, the others at 2 + const words = [_]usize{ 32, 48, 32 }; + var specs: [12]@import("pack.zig").Spec = undefined; + var nb: [12][48]u8 = undefined; + const shapes = [3][3]usize{ .{ 2, 16, words[0] }, .{ 2, 16, words[1] }, .{ 2, 16, words[2] } }; + for (0..3) |e| { + const base = e * 4; + specs[base] = .{ .name = try std.fmt.bufPrint(&nb[base], "layers.0.ffn.experts.{d}.w1.trellis", .{e}), .dtype = "I16", .shape = &shapes[e] }; + specs[base + 1] = .{ .name = try std.fmt.bufPrint(&nb[base + 1], "layers.0.ffn.experts.{d}.w1.suh", .{e}), .dtype = "F16", .shape = &.{32} }; + specs[base + 2] = .{ .name = try std.fmt.bufPrint(&nb[base + 2], "layers.0.ffn.experts.{d}.w1.svh", .{e}), .dtype = "F16", .shape = &.{256} }; + specs[base + 3] = .{ .name = try std.fmt.bufPrint(&nb[base + 3], "layers.0.ffn.experts.{d}.w1.mul1", .{e}), .dtype = "I32", .shape = &.{} }; + } + const file = try image(a, &specs); + defer a.free(file); + // every byte of expert e's tensors: trellis tile (kt, nt) = 64 e + 16 kt + nt; svh element j's low byte = j + const hl = std.mem.readInt(u64, file[0..8], .little); + var pos: usize = 8 + hl; + for (0..3) |e| { + const tile = 2 * words[e]; + for (0..2) |kt| for (0..16) |nt| @memset(file[pos + (kt * 16 + nt) * tile ..][0..tile], @intCast(64 * e + 16 * kt + nt)); + pos += 2 * 16 * tile; + @memset(file[pos..][0..64], @intCast(200 + e)); + pos += 64; + for (0..256) |j| file[pos + 2 * j] = @intCast(j); + pos += 512 + 4; + } + try tmp.dir.writeFile(io, .{ .sub_path = "x.safetensors", .data = file }); + const dir = try tmp.dir.realPathFileAlloc(io, ".", a); + defer a.free(dir); + var pack = try Pack.open(a, io, dir); + defer pack.deinit(); + var arena = std.heap.ArenaAllocator.init(a); + defer arena.deinit(); + var parts: [3]exl3.Part = undefined; + var prefixes: [3][]const u8 = undefined; + for (0..3) |e| { + var buf: [128]u8 = undefined; + prefixes[e] = try std.fmt.bufPrint(&nb[e], "layers.0.ffn.experts.{d}.w1", .{e}); + parts[e] = try (try pack.group(prefixes[e], &buf)).part(.col, 1, 2); + } + const x: plan.ExpertsPlan = .{ .proj = .w1, .parts = &parts, .prefixes = &prefixes, .layout = try exl3.Layout.of(arena.allocator(), &parts) }; + var img = try experts(a, io, &pack, &x); + defer img.deinit(a); + try testing.expectEqualSlices(i32, &.{ 4, 6, 4 }, img.k2); + try testing.expectEqualSlices(u64, &.{ 0, 2 * 8 * 64, 2 * 8 * 64 + 2 * 8 * 96 }, img.offsets); + // rank 1's N tiles are 8..15 of each K-tile row + for (0..3) |e| for (0..2) |kt| for (0..8) |j| { + const tile = 2 * words[e]; + try testing.expectEqual(@as(u8, @intCast(64 * e + 16 * kt + 8 + j)), img.trellis[@intCast(img.offsets[e] + (kt * 8 + j) * tile)]); + }; + // suh whole (a column split keeps K), svh this rank's 128 columns + try testing.expectEqual(@as(u8, 201), img.suh[64]); + try testing.expectEqual(@as(u8, 128), img.svh[0]); + try testing.expectEqual(@as(u8, 128), img.svh[256 + 0]); + try testing.expectEqual(@as(usize, 3 * 256), img.svh.len); + var buf: [128]u8 = undefined; + const gp: plan.GroupPlan = .{ .proj = .wq_a, .prefix = prefixes[1], .split = .row, .full = try pack.group(prefixes[1], &buf), .part = try (try pack.group(prefixes[1], &buf)).part(.whole, 0, 2), .marker = null }; + var gi = try group(a, io, &pack, gp); + defer gi.deinit(a); + try testing.expectEqual(@as(u8, 64 + 16 + 15), gi.trellis[(16 + 15) * 96]); +} From 865335ec1c74ef16f872935b5f922a08b4ba967f Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Sun, 11 Oct 2026 19:31:50 +0700 Subject: [PATCH 030/104] DeepSeek-V4.1 weight forms built at load, and prepared weight folders cached on disk for a fast boot. prepare.zig builds the run-time forms from a rank's named host tensors: linear strips and the routed and shared experts' tables; its test checks the dense lanes layout at every width against a reference (fixtures/lanes-ref.bin). prepared.zig writes and reads built layer images under TF_DSV41_PREPARED (_WRITE builds them, _VERIFY refuses a damaged payload), keyed by the loader sources that build them, so a changed loader never reads an older folder. --- .../families/deepseek_v41/deepseek_v41.zig | 4 + .../deepseek_v41/fixtures/lanes-ref.bin | Bin 0 -> 114760 bytes zig/src/families/deepseek_v41/prepare.zig | 144 +++++++++ zig/src/families/deepseek_v41/prepared.zig | 281 ++++++++++++++++++ 4 files changed, 429 insertions(+) create mode 100644 zig/src/families/deepseek_v41/fixtures/lanes-ref.bin create mode 100644 zig/src/families/deepseek_v41/prepare.zig create mode 100644 zig/src/families/deepseek_v41/prepared.zig diff --git a/zig/src/families/deepseek_v41/deepseek_v41.zig b/zig/src/families/deepseek_v41/deepseek_v41.zig index fcca76853d..dd59ec7fe4 100644 --- a/zig/src/families/deepseek_v41/deepseek_v41.zig +++ b/zig/src/families/deepseek_v41/deepseek_v41.zig @@ -10,6 +10,7 @@ pub const plan = @import("plan.zig"); pub const Plan = plan.Plan; pub const stage = @import("stage.zig"); pub const named = @import("named.zig"); +pub const prepare = @import("prepare.zig"); test { _ = config; @@ -19,5 +20,8 @@ test { _ = plan; _ = stage; _ = named; + _ = prepare; _ = @import("host_test.zig"); + // prepared weight folders (TF_DSV41_PREPARED) + _ = @import("prepared.zig"); } diff --git a/zig/src/families/deepseek_v41/fixtures/lanes-ref.bin b/zig/src/families/deepseek_v41/fixtures/lanes-ref.bin new file mode 100644 index 0000000000000000000000000000000000000000..231acd0531953e107b9c64d1c53816ca15d00d2a GIT binary patch literal 114760 zcmYh>Q$Qf?!^Yw4db4fYwr$(4&25`)+na6IX4}}X*|z@koqY$-`5etN_w^eH5D*YB z;6MPU?Hi(!Qb_Efrk;K&jx~LVib|x#Ay0^y$QV_m&#z@Mg~3N7d0crsZgSsXuOcLcl21$yLl^58unb~)lQtAg{IwB-d!!<&6n#cORa z(JZ`!WmlaV9rxSR6%0|zh&S${?uSjPELTVMcWl@wqpVNWl%_eowKg(7w*cQI8+J(? zS{I8BHI*~*3YqzembVWzbZXJzu59r(f|xch(&dH&sNE5b*wtXNB=P3%8S_W99_lZO za>%gX=A)z&koP%by2w#qR#GiTcltZ%@2)&@zfAv2j3|1ho~^^ujW^O;rTi7sl4QKs z(o!t4w!v%bAkIOWh0X_B3u>4mdRjfLSEy<(| ztPwNfaw3X*^8_?)rxEzq3Dw(ALs#xiNv2fGo4HVOMt%tdE!jG8e#>~~JNWFNIK$8MoGoscyN1gS`v ze|~4UiAa}{o+>(pFWG3-E0~xMef2CMSR2o)OS2eO{YxwfuGW~U4e$!SqWRtU_s5)7Y*Q%p6AfrfP&}4N~HAu&O)#&Wd@W%RVsos(}`t+Ob z@(bN2bA)Akg)Zc^6{SriJshRVagi=^wOv5TwY&>)36cKEUd>H0GBt`lF-9aV`D&B6 z?MB|K0G0j%N`+|3y}aJhkufKefjyK(F z{VHBqH_kYUwXnm9m!jeQA~rcn%NJLkhF_ukM>+}l{`ao!_f^V>lnd*k4T^H@aMI!W|D~+*kn!V*CCOQIq)|dx z{ekL7Nn%bFgPi&79k+q}JftR!_0qEJk`zQ9?^~3 z2W!vwus^QZM8gjEs})^QxEo1Z)X`Qj6b`&`6w{}R-OFMQ_SZBQMijMhMuK~L=a*dG zfa3wnHZ$AcVrL^sy>!o}oV!R{`W>joADbbV0fKBHgOmrkpWj6*shsVkl@B0LGbV(P zjsJw1$BMrkM7LNdNlfyGooVDe=}-@=up7K1)5w-X50`VNgL*ENE-bwD;MapAbAG$# zojgza5*@-Jev1t_Ahi7D7unXycZwPALPJVhtLK!q@&BT&2};VZmx60Irw=3#Y{PO; zq$jT(Jesb%@&i2^!8@qYKYdAgLZ~Ik<>m>Uq9Luz-^2{ZV{`fs&*73GF2ayLUyVwl zaN4iW$9^Vf68yWjQ@Jzj_MUZYD%mlXYL
$CLALcYuMwpY9$o#nzcxC9 z$};UfW%cOBcve(Ak{n+4O2qkfUQY1L>Dj@ZQ2K-0}Nt_m474RyIQke59Kgew^d51UQ> zd7h;!i4{4TKX|OpX>T&*3WN4tABW}*6A9&`PHy`(R-hsmTW&$Ali5SoH>cGc-OT#0 z?}_Sx54#zT(7L4QER&)HmueptheBnPt)(4Xl?=VDBymA#Y^}RcNCRS!k?jX;KlA`a zK>A7?+}?Xi0{ZM3QOR{n!&~LZ&7~zgpZcj4C`M!8NnBwyL`ITate+`uYaO+XP}&Mp z`*RK{%gCp*XAx%pp)y7QPC45fm7s_1X{>M-dFZ0umv7*hYEg@>)xRcm*KG{fklF~6%G7s5djIrLZQZuv+5Vfv-l#4?1}7<-#A5uXH@iFXpwRcby=K20{8Z2E7m zH7JR!?(j8?{eaZZ89_5d9*S%M?ZK^jA%9O-`h$7rZcxauA4UYt3LmwYyD2+VD@=A( zUa{^8=RR*nWGA(oi84PBV4g?jtFKnFek#)~AVbXNYNIjn=5w(8n@S>L7lh5fjv8r( za|H*HDJMaJu*t9#M^re1dI|oNoaBuy+LUF8r{c#iuVzldk_KbH?kwIBD9A-F? zsFdeK{CBvZ+f6^}-Yx#ayesJS)Rk&#D7P(ZRi&6oq(=|}&PJNx9*5IH2HbwXoWstR zfaiAf^fo-NMe~gD6h`;(F?Tsc!bUz-%#q`Jt%RYLVDzlesO2NOfVJm}Tr_y|VCVjv zhBEjhTfK|9Yn(?G8+2%93xA~Ft~r=udRJ)ky-X#SiFtwV(se6tHn@jw>7=jjPE&#T z*qT05MJ|EVjG%kXXZ!Xd8=0}L`V>N+55s^}WDTSZZwGI=o1yfvW8PWR60YImFcL!^ ze1e80a*H!H$20sL{kQr1djG$44`qob$4Qg@6=s+wj;H@j6zEMiAu(yRf5|&z&s~Y= z9v1pF=yTsVcEB8e{*Ds06L=Rf)URlTxl5j=}M6jBjDk?`!pce2}~`wIigz zt}d>$CzFB=$NpU{0%nFk{RJaaFlgR5Z_%(FuQ&fjefeJ%nsuO4_^*n*&gSB}#ZG&e zA7>@CPImRqgIG98Tp_)4jo7PL0}wmOq0k$jguhz+$n zl$=k!SB29vDDaTgl96A33oHcBEj+e9z*R9nUP*9S*KtwS(hP&ryv65v4qRj%{ShLV zS^EwBDrSiqyO6N2t98gX)5uHt@STR<34gFc)15M5G$@Z;5HyFDLLuC@$<3cRaIE9@ zGyi^phVf2-_5?q+#c@}wAWX>X?w23F27N>3sXB_^{)g%n!(#kkQqR8(QsYn-v0~Cd zI_FJJ&tq>}LhyEx+CWYeEn$Twq$u&Gh(<(DAIjr?J!kBx<&Q}JXFO~S*N1d#EbqY% zCUi&l))?fRYx0hkBj^e0Xe5OLgpn=GV}!2e=yf)o{neUe#{j7`AK|$ff3s=y7kC$$ zYWfRRxvR!g!enO?7{RT%O}l9)rbk%*KhC+<1V)(+ zwS2FiuP7FXgrA&J$M<286MVip1)Y*%3aDpvLc$4w-l=R6s3hP4F7;ehNt?rRn5au0D zAw7x{!Se<1tXzrJm{z^341~T}F;HpL1BN-*7C5`2<;W_pJosM;HJ~?#D-paAs8Qvy zc_MJDN?UzVhnSs^th%Tiv5cHmS7)9$0>J2}PTGB?5MR5aEZ`wXZoUPOS%h7|M&*#b z?Y&Tjkj0ZD6KI)``lLI~h_f>vt?Wq6S{rn3scafH4-)DOT%Vvg-J(nJSr9grEwBf( zc^|`081O8!X$Ix*Jx~9F)_%HO4QVf%;`a=UxV0Iv7QyS7C?#1%&k=R8WxL+nTFULp zS;LT2lO3zM4!+xp1rO11nplwNLvSOXjv}zH8!pieszT8}LK#&a{2saPnHL*oNK*ef zN{>3+vbfl?t)P>v)+uMTkRo=_G90`Y#V&#xwxn-q$&itQ6TK~`dhj&Lq;}D5V}H`I zGa$5s3{i=U{~oaN684)f5bvW&q3zpRd-C8QvmCNa$m6I-Q|y87GUD<|jdbDwcW+Pe zo>GWgAa+^zr}-#)vI{s_0R2))9Ov~*OU_tbK{kWgZB0EJc2he)e4vRHJePCYaDUt*0p;xd|apC5E8%dT7DFVmzdBj#^A zAC2PeZ_}2RXKEQq_hy}p6oy?{blil8h$5=|L=<05{dH|#M5(kbEM2f=+RnN1$xQkCJy;W6;uM^FFTL`2}Bi;>+cq4`lBk)6qXq_V(o7AGX zn3l_#pNOVM?1rFrbw8B+D1R~IP`$};SHBJ>cT@|R?6uS|9gAT-a#@RA$#UsPCK+4Y zF-YKNoCj_0hK<_qi|p0HHwrh4v$RPPf5(1~wG=x~r8#3)!K{2l`&4zaSgP1yJq33S|AFFsBzKaaUd2|L?*D8rMk=bxM>TF{>UF`J^!oSA2B}g z;N*m^YHHdShS)UaGOs?gXsxw(!G!B=3KCEt##lyaNA^LTG%P#yEYu61m7x%0A0qBx zMH`!h(h(d)_FzZxaqPCX4|22P`0>NSUYhk-xfbTVk7ymfng5G>p-rM@yc4pFvl3G*9e^Ht5?k&dtw``zwVM4JC;D#T`hVpyMe3@-&b z!ws7xf82Otl95Vdc_{eeRFPaEnB1e6^i{q<)qX_mS7r3`w1+UyOPK3$Wcd38tFgi`6GjJbj}lMhWW zUiEFV+GO279UYix%-*n9PM*+(5cI!y@rP|32ZW%`NY>b1z59{xWl~UiAn|jqTueo2 z)q1Z|5IR)t(`H?P4a;l0!f*TCOk(Guu0zBm2Y)!qlpwnRyn}{fnHytXcoRu$YeuBP z=oou_FO|^=~^v0Y%3bCm13vpSH4|>vU`Zrj5yM)U(S(u-g>KKthdY1G! zqaV+vRcl@vKlX&9Dh*pTycr!MTmPR=5xIb%Ms-z9NhM=(dRVoB)R zv}Kzg-N?PMJl#?>ln_S$bcT^RCT@3Voc!X%bELtb_JR=OdwH)H=SF_9MI z-2;f1#Y@T|oCBO%bhuQbI%FdkMlQDl!daUbN6OGFvp*UJMMn=oege-sfXCK_HkfvX4+b}Gv&HfyhZ!C-Tp&^@85B$yJM7G|jpwjduPfuP zab|x=ZIbVJ;Pf9-n*JaBd?xpBX-~Sr#kdF?|3kfF!>Mjx9U3`iwo3+1Y4ZxV|4{p; z?W?d!>hslVg3|TK3?jN`ZBY+98x85dXg@^*Z22DC3nJ4tQShc)Gw zac$8aZasJ*1Nff7qZMKs>ux08GYfBmYyuYY+9Cc%C&AF0{Ky=+pB)@mev9iE5RSI4 zHtWGD4-75Sb-38cp^c+$Zh0pwsmib;c;SH{L2%%g-oIVHi~P4i&zggYd%D=;x_yaA zD{bTSPdke$>M(tT&YmYZ+Z#6y0QOI@yL_aqt`6R0b#XBQ_WW32f8*ekZ+FC%Y(MP$8aVX%36d#(}n{H!J8Z)1FD7a@`!SqP0bJ*wn zwXdQlCpGgpFqpx^prfYLe;2XI`TD^;vFz*`a?yXPEKj@k+fSd9)_`1(Lz7{mGaTi0 zX_-BS3R(O6t}o}Wb&G@&4yLW3F%52o=LJLxjq!%nyx0ikY71ze-lkw`KqR-Jzdq-1 zu=snPF*2MK3h_~s$T79^=GP0gXa)|m{U}h8e;oNtB08AYJ(TAhWc)QI-AxMYVORQ7 z7>P>gfpQD_r6Nq&d{D^cki!Moez{oxVP{WvLFewrT*QCrE15lI2+9~E%|-&C6nD*A=T6w(Z4>J3ENW$ z8>$t8(NovAO6+E8zR*R5qgeNk?6$0yh6g^;)hu#x#40bNmZ98BOY0I%Hs8SDPoKA* zt74{+8#qEtKD#fdmxS{I)>Ret(@YBb@1hpo*pu4eO+g5FzNRyq>@iizNh30HQIQ^u zMyh=%iGNM93^G6N&!q{hKud7B)^x3kHa8OwcHnH{3Mo{=*S0bE@t;S6PdF_CPhfGf zy3t+ZsSsdjDDQDHl3LmWQfn8Olde{}nC3D5b301g+E#+dBanUFxqyrf zO27lsHqu{dPow>KHP=YzoI)Y8wuN6^42fCA9MVe~IERL}983Wl?u>0+hc)1NyvdEX zW7kY3^-ErY3ww2wQ@F46QE0PTL`1iiCOMQ4eBij7VR{$)_uLg1CB%-mEq|u7s7fyW zv|j_w+b3ad;LrSy`S&s~v%|t+Q0R?z5goUNiv>QJf*VKAf&VVGglQNp;h31eVD!1; z^MpM)QSKKollWL4xUG~Ef%u=LPv$BWv{tTC)d zG`2W`Ps|2eH2jVdB7n{oZ^Dn|Ozo?CqZToI74x?8p)3fC+*&%(n-_Th4b6{>=((#E z{;t4Wh1OylyKo|==Y01o$Ma6}@<06wOVm~L<9;D8^LP48r*Lg9T`h-_KJy;>X9_)x8{LchJ-+A!H)M#*Sp*4oo znB%GX)UN(PGG$^R*wOu8hUumOYZ4dzbwSVr+%>sau>w6G1{RJTp*-?qD{@XKi>=)e zG{XGYm6hEP^SoFQyevUS3ld9UED1D{*r)5~N4GPiuo!SCYC?5}5o1wZkbG zEUB3j(lCRL$@VP52zb8y7L^T_5k2Z0(aw;3TrTe z)To7ZL)r;_@ldAuX{rv<40xN+!OUk-+^U#hQVY!~J4hx*_)HSpL1)ig7sPsWiu-s~+e zmLhXi$c$Z#t{`|We{V+|g@P7E#e6Bj#d@fi52)zbOs{JuBEm3>bEIKV@wIzu?S3_b0iL6=z{UZRj0fe|e*Jih zzC^=YHhT)uTqCZ4mpPM)#usDZXH_KLDW!4IplxSGD5*6H8W0YR+H7QkARp5@1}cN& z2wPh>hSmqo7Y}|kRQit|Owsv)<*aQF-%9IcjVO-Sd48N3gVa9lsjZeZsJKdj7J;Pt zAC;f zI>LCx9vY@>Wf|Z~fFU&9|HSVHMdgFMJGmiUuo#RT&#>|w&mWF03aEBW{F$r2f69Js zG#HOen;?L8YmN2y79W*th9@C!a5RDAOlMR34`&!{w?(bt$G{!oY;NOTvpC7O-&gq~ zPu3Uii|bAuuQHr@N9h|p{c{HEY~o9V=6a|Sk}TD}b3j`XEQg_Uqn-rKJK4SXOkpte}W0G0gZJ(c}w#B<44y`D&QOozswUWIERk9vIK{hJd_#h9IO`CTb0NhBKL%RuOs4Hksa&w zhRE?mgu1TS&%3a9K@l~1pNe{>PrjGOs*pP(>Kt45Ske7_yzOhnTi%N({ni{m>Q`cZl z%Q-A@h}{&ZL+!1C`Yy~!B$ir7?mzS(P-9Nb7@j~Yc|XJtJts;`Eiyb+QbhT+OV%u4 z{bnGYC`wjy_1%VHO4S+6&79rjkt)WS2eY?xG{FZD#OtuNd!KWq#*wELXzSKYUJpq;VrILYxgraXr{j< z$icRpKVLcnn;}xJn?qCzm3(lNoVhDZCB#3=HN)HWJ%aZZKJM1={q7_0iRGXZoO`v~ybr@8&-Q5sA8Olvf)jf@f8nHYo832b;{p#nVn%Q=j{Z_WrP*?q0ikfXGV_bW-?csvtT_J1o*e)3ryP3fPaG{^W6~ zZVDhqtm@WEg?WiU+eJg%a%+8lJNL0DL1y5bS)-v-Uk~nH{~3TWI^nKgmiKgbncTry z9@J7z=@_oW!L&_qJ@dgFwm%U?;^e=wNZociE7Q&&M6T=EEsvJ~>%-uzLnxEaqwJ1) zpPAkboB2(F<52*T*Fmcv8z-2Pj#`x*C@T;#HG&-(f&S~WuCj7TP^$mc*%e;!s*ym> z8kLu9Ty~SXM{i-x2Y(4>Ge-SW073d8>x-u#;74?QR%&ZqVA=#(k8?XTt4jQj9e9|P zwjvgW|2$|O+?;UU{;5+&jcYkGi8By0>Yxj_<@Q$LkA%cOe|;IEG^z0}%pHeoHuvs^ zFbR`cM$F{1uC*jp-^Z{_ zPhPaIUPSON{e7RIxnJWFxLp0xI@}+BqCk3cE&v@Za7oIRLWvj;DtnH5g7fxEjm@aa zQu)h=f-}dl7#5{UQfX0^Z!qlq5{0g?n+Afu=)0uj<5?W+b!AanAFrrj`{FSEvK7>g zPzk;j8Jtjq!MJ~`fxkaf@`CnweX!fD=$d_gm=@N0q#56!H+EQu?LTq@N#130fj*Gn zv<=(8$og$xylW)I%`nHKo++oVbOyx?3!l9`*A>d~Eq1=iQ~ebnqIWWhEM>;v+i*cq zhlSGlRAqS7|LCzrW696Uef|Z~hyASJ9mv}R#i*t*-1eF!>ZEYM><`HzNT0W3>Ln=- zlczK-cPjFx{Q4^r5^EpR$(b{I=x0rW1Gk&Hv+r;+Ibml+TF&e(Q}LKqV_!$xJonSX z3dCtDI}?;uwqg=-{rv0a&*To{E0G=u3E}q^a4v+Uamu(wWl%1ih`A=Ff)1Y=Dyx6s zn64{)H`L^d0Vb78%zRigos1eAe~ERp;FbxJ??pVOVeHTT<8JI9;hvS$RAg`GTKBkE z-#?hcuvfdv9f!ILp&aBhK5-MvJ=3cDdq#ZfQkobNj{%W_cx|d`0gaXASG3AoC6+Ja zQRta3@yg}tM1Cq^Kd3rd%-b-AScQ?l;V(-=qhY7R7okf1~;2a@~9hlqYe; zz;j=9Y%6A9hrIjN_>$*6REIlS^h4O_VJD0m!?+Y$X&8!j;}01;b>kGAzbux$p5IB2 zl#EQzt;HlXId!?-pThgZrqpBwS;oQ@1W8MD)4z-33NnWnbeV2$N661As)9zYQ^<$ zEmY6ATjqic_1jjRbNld2s!_TD^l07XWDVq;kVl>n5o0d7AkaJQMUYl8J7@GPW1-iHG9r7h$gqk+$O_MgjmQ0n3{c_14Y>gPv?XBy9O!=eGx9QZ9Bpa;y zFL0l{BWv(YMAQW>*W{K>Y`O(Xg5;uel1|P$3Fn9{AyA$+IT;d*@eip;l(RXA#8dLj zn&mgP_2>}7OuqiCNCxwns4Pmnnd0N@dQkVgYfp@NPw6AnrMdNHYql-9MH4c#4`UHv zyKIWdV{I<@^TyH@#MKhVU^>7ulA5b0$Kn?em7 zP`=<$ryruX7!oyqG7^)mQFP--L^6ZQAM-enhp zT_XWhQfG<%vI>Uuw3kTqoRmTd5m89wf;R*I{bn#Fud5AGIAU5E5R-~GD%!W3y|2b^ zu+d5Wv3ILdC0n>i6q80H8vg$pQ|R@4#E71KM@F@WY+We4PLlc;AMXTLk5W96IMV;{ z2Zr+|{OBuL{<#^uK2MWwCx#wy)I&+KlbNC55QB+zb_hpVy@(xBRoYC5%(89tNmcXU4l)(rb)kXYfjZiNUk^T1i^m)?}$KNy-pBV-QuhJHa+_Ix)5 znH7n`Ig~T**I5`Cd7+7ieqJC+!QnpbGwX5ByR!7!@!wd?jCtkG^M84#!J_Zfs}41E z;na_~reo=|$hz&O;WU+k)1oFRdz<~UG)~F4#|V&D4C{vf?VDPhx{6&w9M__0J4>^& zyJmld&LYJWd&@hs|I9Y7Z&p%%d-3H$T44)4nYwu6K(>1+e3^;O=~YY(<(&zdJ!%Sb z^Jsa>zXEUO8ghEwqX@+j+C!LRVS^5rz87_m;R>figfk*G5I8GYQr9=vD%(V{6 zdCeVy%gAwIJ&eI-843&VYLlwm&Tm}TgUmouH0IdCT`>61X&=gtXno;tQ3I5F;LyRg zn>3*+l+fK+TuvhbSSI0@0dJ{-?d10uL4rB9O+n76?T1FdWTI)AF=czolpepR2EmHb zXvpz=z1N<2mX~gToI>+zPESe18cYD}m9>jxIw_MgB>R?;s6jlZ^lwwRivsGZx|c(V zes_|)+v8tiVe6%EpjA59ide{S%;2@@+rA9+y&pNg0lw2V*-a&WT=@N!ufk*;LG$x> z)Y&~@;k%H%N~2~0P+So^&!T8YI<9?T$-ll4KEv?XTj5BbjeN)NZ7j4@%`3xjN=M0c zVnY%%bUnEChtwtuMoz^_qGt24Rmq!JcT$XI27M0>)w#(*xS<VvjKJ<1x%FPi(;W-EDWnrd+ai_P`i&yB3JB4?Mn!pu#45@n&vw zVJ!!hgClW+3;s!Fr{$(9NRcz*O=a6s&zH9q3~;t;hQlNxjhbA}*2Ls-doIn{?yzOT zuXrm^RN7A|cK91CD#u(1h*UDc*B0FyuTj~KhmP}G{6|%g@YYrSUmLHCqxLXIc|_Ou zW0xQ$^m4Ey=WKtI7qabsZs|gDT3imIrrw}cHdJP&96EQ{nj zKwhjd<*#-WD`iD;{5x;p%u{qmOqt+F>a@Wr(Y!F$im8A}cvLrugpx>>v$#n^s z{gG8njqm=gEi(;A#ZNO|BhKJ!sSQXH#lpGO$d&IXauK=6Na%= zcp$AlNz|d-{d-ya5a&5Gmw&>%5S2`Ltcs%DJY+MtBzsv;HbP_V;JsiJ89)N6ps8Rq?uJ+FE*TT6#fXd|9^S=|76 zjxYlpbg4kKOo5YWLcg#M$vgHL!;`sEBhrqn8o?4Hzw^*vNa-!gtY;w3Btyq{-9*TZHeobpV4rOGUyqeJGDjo;?t~re%3J|3N zM5zu?>XslC-HDeeiz~+cs@0^PzZ5IoVP1K!y!^~;=LZd`2L^I7`}@5=f=0Zn# z8YKE%X8Dz|E=?solH#SL?daJ8C5$Ish|LhZ*+8$0JRK$l>WZMY1~D&xesY$gxEcC< z+12cZ?Y{hK=3Mp2yRP>dPw5r7@>oKn-Rz@HNzdjz)suO)$1*s}Ohk2XPKF}9$|bYo z-)S<`oN)S=Q!0R%?h;~KO!&6C$Hpip+0Wb7WWMY)Vy);;tY}SSvEt$=5-GbRA8{N8 z>=MB!5i%4=LxTt1;qyuKUs^A#?sjP_!H$ThXEcInuy6f|I2^mxYEg>opC2z7LgsA} zQ|e}f&p6eGCiE*K`F;YIw+1rnd`Chs7Pkj`qraV% z@voeK#_0xqXI%h_63AW_2Vn@|)J_raL^hc_Atkf9KyLYBreQr77Nb&aKE^pbXUakO z(Snt08mt&|#4UcmXXAEcsp!(al|6}}`^D~Z#DfnhyEs}(Tp3Q@50MBLT-_FN@%9X{ z!{5TuC(iCdnZ>YEWGL;l+9}}Dl#m#5%fEgkF5>p|kMf6jr-H2bkB-!^?U?tWiG@dS zFh{^AP2C<5GdAKH3r0VVx5ZqcJ1OEU2p6F6JQkR?LpB0e z@s1vw;Y$=^{8f}*erXiPkCl4cDoaN)H)MN-YkgOu??V6Z00+oTDWQyo%%5U;_WCWv zn(7hk15>)yLYtQl<+$%uj;r>31(&LqEW~S4d3XN0R$5yiyEgK@m&yeqw(=QiN?AwZ zC;pHv?OO`?rmHx1$Wz$?$60kB4RZbPGHvsSx-qLd=TIQ;r{j*=cUVLsu6%UIhgG8| zeFQQK!Z+E$bE`8+q)TH4=hIT1e~e)HTSGtxr=J$bVH4X{KOZK956qy+*}XFmJJY=3 zyGCceR6983(Q1`={u>~l#x606;)8iig=2$B7?+sM7v?}i^Zlk$P4sYezo7<)g*bEV zB5*v1$_p>zY`gQJEQVxbmMbcXsa$~j@BMv$?6;32${&cRnC7H_CKxr6NvaobMkyj( z+J(W;SH_$@(YLop6msG759HO}74{r5)}h?dD4}74qAut3roJYb5SDLvghjFJ1IQ@? zwG?ZEmN!eRh|C`vzv;Q(wOf3uG}abxuI`b%Z9j0P#sYwqK}%d zc?{m8;O6dRGkMs(ahBFZ`SxUNQ{_vNnS+_Wj_Nhm>4Jr<3_H4N8h2BODC~_Rr<9!W z7UULUn{K$M+B<)o{T4B@V`2#GFz?@Nl1038`e%y<+mS8h?`gfX2FKOfM-6Y&narc2 z>(uV@fzGN}QC; zWEBcy57@vMjC1oZAQTv|HzfidZkY+7btTXY@MuLT0UB>31pIoXE?_rlmE;;S*R}h4 zg;)!Qf6j|hRMRfx;c?Fbgqr#ihzr&kjWd9upNIp{%@1@2Ol6-`fYVR=0rQ^Sa6{L7 zs6)WuX;vW^_5lH-o{*Bd;cqq&qqju?f`pgm78&5~=Yar@69H`kTNo1-Ao$=1 zz{y9tZV*4-Y*VOjD(?qQih?rKO})b%ph6QP&Zdj!O$}8xx+0SM&g4 z2?$*oA)T(*W3?3A7e6u4z0VDX3)w?%y}K(KkgU7ZWmTp3E_ltBfzX z&&*)T$~d{pVnU!RuA~M;nK19`0iZqTF#v%e!WM54Aok-~z`=)(j!9Y#KA`m6wHf~A zHX06B`24hW^Mb-);PcofA)M{X+NITeiJk>4{-GEEa}Q-xTU2@`B7o)5#sE9FT6AC) z zJ=OtlkBFDji`pLGLFF^ma4UVQtMhebvR-afyXgnzx6Lr`Xhm+Ac%45Q1I(Wz1AzG< z^0sz^RFVvU|7V;5qwl&>Y}mr-*Uzvbol;*9WyKa*jO^1+R345UkDH~6&}2W4Kwg82 zw*WdpOD2H%popr`EdsY}Kr#?$09?PQ$E0GsGS}#5Zfx=(hu{7(FXc;8{Tv0mW8yL+^ti8&PJqM zI4LO^VbFHljg5npk83`aw`g>87u4}K=-&d2jYa;MBx{} z`<%&u!E2gK!2#5qORByvxZYdka9)Rf*p_Y4(Ldn9Sf5dcYVSm79`(~h0_Xrsv;fwN z6-^caCfx@l0SA7-?HSi~H~^kDd!uwvsyp;ZbmtSB>3IK&6&BbO>P|)O7tJn8tU)IV z1loZ^6#)EF-krYz$NmB`!gVuX{r72%t^jF=BQjDsvL51RBy)QfL0Erm%fR>HwIr6j z$MOZ`OT(Z|1HD0S6^H_H-f>=leR~1vL1+r-0}V!MYyg92+3%TmS)u$&y?n)vJ9-pv zrKE|`d>2qsMBMaqhFWD91{%e#7!VRbw%Tq1FGm7$#S=Wx6!yOfcK|w%!B8aqZX~O* zE`q{%R{W6Eki+8dc)>^%-6Gl5*jz}$1-i!A2@oId&SsneV-EllhdnOPAz;B7uK{i^ zU@2oxpU-YGEqcpge+D;HH4qwvs>0@Gj&d?pwQTKv1X_sPI}j|8$smXUmnQm$2K z635lMD*=Rpi6{_xz{B@b0kqBpn!(B^;usD zQ7BsB_6`Mz@_5w(fj^>k%FYU|T!0|K;QW&ZaCe^<0LKx5Hj&R0KnxJPOd8!h(5z{w=#DzN-6j!<%I;`YJI?K}=n?Vf5kSO9?}m8+tUl8R0QD`PXHejvAp&gw zeFfP1ZajiuUgq-~!+-D~#ntySt;04b9t9J; zf#jq!a01jXxCMATvO+*{U>u!#)VNc!!_wain$M|G)T;G0adpb6?)wPdZXlKjM63q` zO+YphfcsIPn?yzgU;udln-MVlpZMbPbtC#^Zn+Gf30jA1?*A|a(0^S{jGzk$-cqc7 z-2s86aw+c#=m?c10QQ#yt>x9ZgA+&!{Re>iCsUNcUP)C`qBMmFq7}FxwU9ZkyTWq> z&3zeKccXxgDgjZ3K+Z87Xb+8tKp;5M91jABy%h>@a2VCx@v*QcI4H(vy{jvcSM_Mh z!T5&Q=QiFZ{}AC0<1;_EyE<%dKfvM-PXI7~KjJfnB(KT>SpG%}uya74&(KFuaG3pH zFT}p)A%{PF_%PS5k=->*ty*B-}kFq$G6lXxRYGUGp#!8Xm9(D%I-v}Cjd1hr*3-wHe zU4Z|2Rso~GZuKiVn_KjS#1c`;x}pA;i6WLSr|${k{3{Z^%4CJ_@Wqpt8Fonr=mb!O z0OnUZ%6QCm1$zO>V4DkYeV#;#SF!B+cn5CI@BB{2%6($;(@>5V3*VZAYA0f>nH2+7 z&rKdh1)wdQ4*~&TY)iskL~r~B$P57uKno!6F6Nc`La0p38Acwx7t!GvsKPjz3=S&p zjA%f#CrG|EmK|g81>XXCgqc4O4MKJDH|nrzFn}~cVF&bs4`CXOstXPLdG4_%8t&P% zu;>Ual+KrZh>o^LKKA|*qlZSFqyGIFpkWwJ0iglIxPMKL`fvcN+hyFW=4@!p@b9}5g|&MsjSFeh0Jt+&fjr-kK_I?uIqd}->;95 zS;&8aCR6Z%w|AmXVEx0m$AVfeR=b3L&5cXOn>EV3 zf1*G*z&2iDtY4iKx=U=*tONjxhN_VxB&K;8FE5zu=19wOhyi<04*LE>iOC zrBGadQRUOVaPU+SgMoD?f3k*^=2xmeCU4K+Qy;!2Ut2Ok(@)C+I=8h46t>N7Z$S8Jt05 zzv;yEbH2c!4V(AR-@80Jaij8dpVuSa_hCI-wvw~?gi-ANNI*hS#+|#wUfxY`Mc!k= z6t3QEwj`$WZ-UOkUe24l7K^%h*35Q5mlOH7JcIDUE7iNWbYnV^EMNsEk*$S|DUzQD6V&d(hqw42mHR2*H zWMiHH3pd-20AiOfC0Imth_H(WL&j)gqF3%X?ium+R1zm>^TC7__>pPL72Exi->8w zwU5}ZnNxn<)Ne8nk4W#!BVXSg5+^(Ia>J=4Mf0Q>BmIg4O+SGPG~rVm#D?yCg&4Yd z31_eqx|~JK=_OBz&6{H1HS>fsq4Yh&(j~Jp4Zl0$W_1VGps?D3aZ6;(Wvlastf z4#kG7&^L>+R=S7Er$5&FxpG}v^~vCF0!wncSHlSts9hz7dlkY>`aWl=6Fd;DM{M{B zulgg@zA@&q`62WfB=@jv2`255q6IeW4||HEra(+ z^$AkoenIU1Wl7WiQ35K=L6MOXSE2FVVU9*ChTEP z0|5m5Aqm68#6E0B?BJxL7tYbG3Av@07!5Caau;w861%k$UETiXJ2VbA)ax={^BIYY z#s4C<_|_(3Fdx0fmRDszCrZrnm9fNje!Dd~IG<|dZ0oEkA9l`j4Cc8 z$Ux#bvFpo}%TnHUU-k1I)^8yB{P{^-rrv#4aY+Nq&&8UO)zU$xp>yuuFOTdbY#~sX z0Dz`5>MajSCfW#Q7|10oAnN|j)RZR-7prpv$@QA-Z+~j%JlkjWaXqj%T{BSj~Mlc%az)}Hrg>}dbiTV z?;6P+(f`IQ_@6i&GtP<%!&P;EmGyjR8keLo+vv)xW|AyffG+3-fko2{b5&PgD()1*hx&y z;VlZMJMtp-u9aW5x!kU>9FW-(^6zu6^K8imSHJ!(DPzI`Sf_}w-s*(hLt>NWBS?V4 zj@a!L%u>9Fd3r@&EB2$nlOOH2Y6I!b&Xk9x>0iFpZu2I!^g1|_VU)s6*g=LjG4Qtp z?z~Lw*vSMV+;Skc{;Q2WuZT%|T~XlC51j$J%*VlU8^pRTp|SirQC%%a_Q$wV%XNOK zSA;jD%@HV|(=?(;Y~P^-J!t+Ve4r?s(Vm#WWt>N3(#qv(`+~j(#nMYzjmLX4yo}PZ z*FD3bwKuI%KASL#0~ZKLQ2a2KPweGF1Xt9X5T?LlXlX-C=Z&kT;Z#;)u0Q_9-VQk2 zZRRv!7&}|}N|r>9{GQ0CuQ8s4YwYwO@WF&ayn)!*xe4MZN+ldZqSn}vnA>aa6)ANM zMUFI>1|;hk`>i|bTVcEm2o|pAlRB2tg?Z z^@Ml)X?k#%nDu@2pR7K~EjpXeC{bR@92g#4n8LVtRgIxDU?Qb%jn0vpFc;b-0!Hp^ zUFjnB|1|9X^Bm``h?(1^i`b$~55ATjF}`p=&oSc29r?a)hC7bWZc7vxn17!!ZxSe2 zULl|$H>4|y6?5#hnK0Dc;NWcfr*64F$-4eR^=Dl8zd zq&MlE4`BiaEQ#S>lW-HRBLr@2yP_aMZ1}0okAzQlne=}lyS<&5^L-4`q6Gs4Bqt{5 zjKif=^Cy=HU^!_%{fck|R#9T?rzWh0geF;yAO$-QV)qxEl;hK3kxrj|J;ifuCr?q3 zLNpa&8IztND+gS-I-9-`DD!$i(1)-GY83(qu9-4Z6BGLxW?~1=ix41pCF5iupQz?J zk@3j0*y@Id1jBOaq&=!n{$8XYA-4FSy~JQ%;278W=b6Y)VwSfFCbn}1 zmnH?PU;V;^#*e?&F;5@4!b{TA&-7}-fxToVarEf}-T5LP;Wyq+#D4$l9Wk;KGj)@d zh?62Q<;T|&d;2BwL5F`5D@^OL=T{`D)gGmUhK7wA{R^Y6mDW^azmoJ#la?vAmm`wc z{4K+XVgBBA>BF?A-rphSf5S0iqqlz1_G_`PP<7=2wFVzs^&#DI!mK(z4`hyHTje z*5DWWQhl%dQ#U4J^?$^;+gZteKVb`JG6?{1?d?!j`F&|W!3>cOgaynREoWxP&Jg0D z7$m0>mhk21$p;77oz+jWevR{bao=NBJX&HdZ$~`g5p916G`Ro!*}21Gvn~WpjO{1< z!0fy8ezvh(*JpvpVy@&UAJ?%=zb2L4LNa`?a_FI_+E0?9XzM8_J8i-+Y)c4em^J7o zcPB80;EXvd!W?u`Y6h1bruzpUrs}0LzET=fS&b-39g*hU-{x<-y^S|XKSJ`=nNbG9 zJt*b~oR|pM5kt)T54ed9J~V~On3$UDd)Z#}ls0|bVbHwl!$hrpERFrFtDp?+YcA@& ztpPurwFn1r<|M}YQt?BY#3ucTAc3nP#BRUWyW=e}Pj{ZI3$zxl=hK;ANgpn+^(y&5;Pg4mj6|f!}6Bb)DF;0$nY*z#Wyl;XTVc!%A^FP+4!uVPQxn|g{YqPZexr_Ys= z`!`dRN68lh8Q+MBw#^N6JCqXUa-5BTk-Bfro5cR_K=uFL0&KIVFSh>DwJ$T_N2pSx z>Eq66?L(g2GORn}v%7PQH3l^XR`>X5`r4%55&d&$kJ~36uDg-0N1u!{)h7%#uHA^t z@9H|ePhvdk*^N`aOo>e;?driMm)<$oNKTBNtrAi^>1b;b^RxHno5PiCpXo=P{Nhi{ zEDt5c2gsh~`&)K7-C5AdP{L~PtbH#}xU59ir~J1bu9tr?YIuAr@0q6Wj z19ycwMNe2!H~Wa*YuwEzEJHsl`}elmAir98l=E?)InIfveVdK46`WO4hGk3nPl3w1VUD0`3;8zJfU%E@hSol%ZSwK1kjEIG)P^{3R%?9>%c0VDSlgCnP7i)~Df zGEaBRJoQu-_cQ45i0^zZyL$NTUYgn`LL8p_cR!J3mi31Yq_?^LBx}8@y#BRS?oX$* zu^Hqv(^@N@;4->YWsKrqJt=P{gYT~o?BSjF zznJqVMfbSx(y-rns29&JgESeCnaM0UwzYsKLKj~gbz;aPXss?ztI z{%~wi!@B-T(0Q_FOP97waJKSw9yH!NCtIksGdWZ;+a-V^<&Z#Dg}UpHq@&6^GhVO+ z`#jX73|kY{w0?Dv)9pyrWaq9K2UD)cBpH0gWO8!MTq)aXPn+AXrR}WGto8W5{ORWG z{l={Y$Md7N3hQecpZUG|6@U6lpNy~uW6OwU(xFR&KOYrK+D(Tfv;5tDhLc=jVcs^^ z&G9Sy*Jcy<+(FhO%jTn$Kf=55Gy)$CIeEnnr2uOj?Ow;?q!z zH}#FAmEJy*`b7L`V(H23rc6rPq8a%GrQIc@v=IwGuAX)AWjp@F`l8tRV2V7G-WvXO zW(%c~sl4_~-o^k@n?JvMC7wMZ{scvPKkzghd01?5m09fR_Me~Y5}GVyk}Dp)=G=AB zhWUtE>MjwRA19g*6}oYh+x`g36PSNVdh@=Y5@)D7+m1`xWZR#X?96z-5oEkKDt2SV z>J#6?i4XdtW`T9HTawnfnWgnUo9}G*bJAresT%vRdYf?uZ3-@UoL`A3EVX&-IQU1& zf1<`)GiJ}`$|LW1xu1falj?q@P7!<)KS=&*oiXY->0ejS@6@%cwB?nFQA;2Don*E( za%ub2C~Y3D0OFlMz0v3gsn8=D&$4y*WO7hDJ2j38v_w9P&WJD{JX?4xGa~3hNbsso z19j`Gf6v6XYAzV>II0`Qpz>_-$+*YmYm$*^hVP7R$DbPJ9G3lJ(HWCKOT@tSq&hz| zWlFwxg>U+DTHmhk73rF7%`XMSd!V0 zC%;KDd04H#U#i|GR`j>;z^3ZpGsjMzf@ZM9q>%Je8`+Zy7 zta8*%z-9U$jSThGUeCu9o7=qlx>TQ?;F}GP7aG%c@?q2ZS~9(Cl`ipFPRzb9Atg z$|u#gQkC5CeSYeVnD#Nlj=_xA$K|)Q#taYLO`pk?9d3R9C3|hh%gMBBfsCJ?oxjGG zdt|I;a-8>|zzv(YU?#D$`5bxefVwI>Kb~i0zPs~{tY1@DX1$ei(r8^={kRy_{IT5S z-dmb}^Tan112!|?r2M95mc;24?hXs7^ZgT(kH2>2&Mt<*q|ZNP{qAh3Hs!UuEls^v z-QiF--DF$qcCoF#(U4#CWhNQbuOX)%Hk#qD0vzj)Jp7Z;>bpCkK+xpvAyv}atg_+Nonz=m2dNvkf4J&We?BAr{q3yR``&wRUGxS^}kNjL=vG5CSv(h^ZA1OB; zeE)GSd5?-?zP;_$3nU)A^?8?c3SYl%mOi?f$X-qAm`Ro?OX{J!)3LDP&Opu^fv$&* zs}GqSWYnGv-aJ*kMgHc;X4e#}4B5!|t)~LJsoWA;gzI^eeuvEGYFSxS@3c58(qwr# zbXQOL-4oU=BiFcchaR|GDtxfQZdCm;b#!xpwES3s+Q>6Y!yhrh7n}|qt|x2bY4&GH zJH2lDo!>4cV&hM$C2J)4CO4%>bvE<%H=38+;h~NxYe`P-Jei=u!^_Wqrx(zOwsXEI-@(Suv-ae9 zNa=SCgECG2C@uc@b;Ics#)^*l$?<=ajZfHIBF|Z}{O;y^W`!x}lq$&+_R)`P-eOON zUfGn-njfhM7PK0uP2kYRZafH$Uf)3mnOE0$unlSC!%eD+p|sU`X9w}-Zb;o^KIDIj(1D)F;A}> zOMbd)Mk>gaIKTAwD|;Ff%hCI$g}FN$I!D}|pL|g}*Y<4c+~BLlrR)BTyY-z^WzSeB zJY_j0VK21bSn8IOibTTF@9HM`&-zCWy&eC>{!(YqFw~T?E{eUYb9wXPC1Z-(#?4u| zO2Y}4D>;7UDc{)0tWU{PGaWjb5%Tg(ZKR93&3V75Wf%7nlg4f_`?~nxnXFmU{OY>~ zX36YrZnwLi1>H?m=pcO|d~{%K^+A@8@vMW8ikWmb{3yluD&u1K+-z{_kZQioLe& z9x*eRGifkN5`FvNH|3|H2a|Gl%=$OoM62IOJ?ga0>2Np`dNkCllJ4e5i+0634PoVF zNdxZ3&Oa|)*;+WY@l%&7X-9~Hu5{4&8|Ivl{I>Mh6fxgt`+AbMUCvDxGF3mi(_geE zw5R=1P>57vtk+&`d2Pm}_1J@@=D(^fSX5Z)dr!X`K0nf3_xDtN+}F|>HO^mwq)-r4u&t3`{TYNeTd3+I~? zW1>rk4g@NR7RoQ=J{KFD-1Dub;edNre-_^|RYE!C0cWN$nxwkX?FtSjn8_N|f~l<7 z{8h(JOBkrkST<|9-S;Ya--m!)Z8YbXspzK z)hjdQ8*L?@&Az*#G{>H3Kkd`#lwi;FCRv-JSmE zH@eFo;jhv$;WZOHeCzkAV;$#pW*nlZ)v5cHMMXZ>gdVE&Kd67bu2c3+iA&Z@%~r?0 zWJj^Z^I_Bq4wuT>XmnRBewyE0wW{e^dL?}D;%BD*_O{5q?1veGBzx8Z5}04^?HHQ5 zLKCxmGI{t%i3*L(tuF)IGvAk+*$1fXW#5^p&wkDMzLKa)bNriYqW+fOfqn`_wdW(z zip);un(A-n8FPH~=3b27Zg6G()A0#%ql2W{Mh0SHZSLfY`!|O_J=i9wecfo_0VDO7 zo@S}D?S52_^>?1z*qhiW(NA zt36~*xniD5DtxExT31B+@bO1wC1mHyy397q@_kO$^pSYCyEq-Tjt%a-(0!1tb*6ZG z+2*G~o=|qjrS%N`!+b)+PI>ibBfEL6BsdvJ?A)Jk?miMFw*R@XQPFwZDbJy>J?2>Nulx&SCEg7!PdIbi2lrpBu%b%n^ktU3L+Q#q z@PN&sd#tHMu4Z>CSDCNH6-#CmL!^>KUTSyRE9#Xh zk;Kj&M~;8~HFlZLG~&k>TmNOsn%SE5cHcsy%d_W%zp-8`_BcCc{3M={^V29-{9WG@ z9RKDXZKR}bI$qxRck|rC$m=p+$q)V#VcgQ%aLy?Ibxhi_i+U=YV>?Fz(h0 z49sWNu77KyJEfp%L22?RNc)}Y;)m=zucklpHXkt@GheGeW?!}~V< z7o}fGcx;PLY?sLid6v!kllR!{BM#5sM`Sb()d$_Ay1gtX%G>?@?v{e$J~4T@-~Y~- z)qXArznbsHqo;A^Z{9vulZ+D&`qj(tJuvS~=Wz|RhztBk(;!;N|Jw9x>V-(7JZ`%r82|qnk0X6StZt;EtN5h!z zG_psNxv#IiuA{HWs%FvT+;Kq6S|EM3G)E&^kLB$)vIkU(?S2E5`9D)xTbpGzza86F zJZyJvDaKnEXGjFhB8~5{NReImi5cNLYKP- z^&X#FKIZ%6&$z*bup{p}{a#762XP-`YO;=Ao!|4Lw&PlaXh1>(=Q69L&95v~CKi?3 za%HYxHS3(mtGqmR=|?n>zWhmA&e|R;PLZz^=xJ(~aLhT zpSM)EyF62`T{ulnT{acyyLMRrfXNR!w!3PM@(asf>1ixFyF0Xf=4Ri0lJys93#pF# zS#rxpCZ7BG_^}?zz2gCc0h4dC{@v@najniNSF(pSt*44UwRz``hpjqi82>F!K5r4Bfi$%mo> zu2^+jI)A*i;`)Z~x@t<*Bw63TsEo?l6KB+wKI#~;N5A)0>});EZ!+{cD6>^9Dx1M* z_>|DcL7}4_Bl~`T(h3**e4+H=c?GHP=+8}=YVLwUHp_*_#@A(ZeO(?aE>%~Hv2T3M z)3 zRjwxSNoA$x{iTJ8s^j8aAGR+TTmN*h@{?)VJi1XKDYHD3vQYmrq-A!Rsnm9Vte|E2 z^B4o=t@P&e9&evr2)NdBZMxz@~G+57SaH*(mbg3~I@(l=Y8dUidTWd4%*WzWJc&lvH|H+(l)lB_r17I$-& z_;saOdR*B-+D|2uTyLM1rX}>NXgke$?t7lcI4OMH)_8?{z4{nZVt#&$skwi+a_`!e zvSS_hnRkkOh#)U(iKj2hD%lrSXM0_@@@RKDe_NsJ&8NM)*o$J$g-z7w+z{@|A2Q?l z;(z`{!5?{!R|mZxU$^gT3oDv1KGG_nG29lDsIj3CFpzwFY3#b=VW)bb4pB;-`(HaU z$TzQs-w_;TX=fU?aC~>ZxlyUEf&A08d!b3ElrEYGXbsPFS|sY3i5kRt*9TP&sonT? zLiBpmv!bEtc12;z;GHT>rXq)Rt}fS5>ho;hOM2dlreI~|>JKXA#?#gZxigpkXogU~ zd$~#~mVbfwuvJ;d9Am9I_d1_f@nnrBRcK&(d00tS5&4 zi+LE@|2>a(jwx!GpVvQ|`I)@&wt2NLa=l({Dy=yxrE7XWk6+`x(am~;_YLcLhlRy) zG8%~s7ZsZ9>k^(DP{&>BX1SVW7veg*W%;ACaKDS#p)S#iqH8KL??nu&N|=Ouxzj;g{s!CE~fsl%kN~z;=f+7rTv*B;At;AAHF|gdulyG$|yqO#`q-t-N`E5F{VEwhmV`bh2$@r-Kk|Q(Z zNe`pc5~3aK`M3VY_3lf)$Q>`IaoK{!XNWOb%vEu1j7O%ZiSl*HlyTh4#|A&HesCVU z`c&GSBed}i(~*Hnezoi(*QdX(8@g>%HafAhXiatKLL+s_MGLAYpF_p$ttbJd|MoIR;p((OM+~{Io)3E;ear=k%x2YEdG6~1Zk&Tst+8bunp8K%*(-YX%vYQ% zzSr9tX7@vQ^mBgh1!w-M`(e?i{6{p{br*(er|s>xXDPopbN)tK{{ke;=o35{1Lo~T^yd}hUM$ni0FegyU2+Af(JNm0V)QfYh zZPt(y@3=Ixr^~^W?eS=^kz0o4&}78Lr6$#d#{#079Iq1Dt_T=+a9#47H9VNM=4Wg( zcX@toEk|=-`N)mo`6t6x!I|XiJ~uvz)o;7Jrz3dK8#2UlU&^seb*zgEA;Z8umj)X_@{oH-LdZyQ*0@Znadc9f<|3+=<B)RfdIU!lKFmd%=ue?BrD35U6K=ySR2aWR+ zk|g0piHH8mby0V*p!THUkjUWFW9_2?Cec3$2kA&LU|VTRo1=9Doj?cXVhkj zNS@dK>Tlhfq{nb{IKFBtEl>DVWB=C{^7Zwb)};G(-Dxd8{Oi`|=kvK-2c&Ej=*-OA zsj_u^6ZJXSFML;i@z0BUXq&iIueQd7%Cg!uy|<>)#{q)@C`di*mI`Bu%0)8{am7KyuVK}BhC?T5>;9>WSls1rdc-wX`IY-V1=V!% zy-QYoIfwOQUn`%^P9CXgUrd}XJ~hRYoo*_w>DZzWz8WwQ;5&Jm<=@$N!B1IolI2~U ze>a@(q58*m|53@`bG{~LcXpn6wdyBx;Zx~A-oDw_GNPrA^A24dyMOES#$b)S`I9%J z;pPSx!v!dN&8`lF$2%STZC~+u>1{>#xhPl5o+!oe=zFOnducMFpFNx2Q7#fRH(RMm zT)nkdt~uPvZuzXvU#+t%?03XKj_f8yb^fg4!d=sDo3v?-Yf`U9mr8{wDeji)$a-fc z##J&Uj|52vQTsh*;TDe*YB}6p;_u#e*lt*@?7Ch`WG{yV(-r-TlH1HCQ!uv z59YOp#SE#(vP9lK_SO5v4GC>JS!imTLzKHdj#kP+f-69_k zlFzjq7rz;;NaIu|zPhUV%y!Xfv;XcQNk!mZ-G|fupXcjHZ#dPZ{|V54sxfi&=NIP# zniTVc^LI+zWqiW-(j|YOAK>1S?tdq!%UU@te0g49Ce&ki-u{5u#2LoS?1Te~%6V2Z zOs<-h^;hb)?X!*(Ryg6W7gKz;|KV-v#1oQX&n%e-id(a;TBzW2tm ztWM{YR9|cg1qHp=kQojuvi}*+M)JqBcGpkuj0r=X;#M%}7V95me=K5^XQ zUyI)I>VE#I-I3*=DH$|h@_r$o{sorEm=lF7}SUC~VM^~Naq|751*Ui6S~Epn5&@4V&_dfu>_=Fo@D z{M^eA{iXPwMP-$mdau^3i@8TQo_ssgP9q!rIB|S#D6n`V;@Qal{g>CIovQAt`RcKN;t=C4L6_Pr3UkR1yWG9yv`sD;c2MYh&E$5rL1()>sRnM@i?1Wd*M{~UEB|?Y_oh*YQR{8N&wExbQzb@C zx(nZU6QH-XG;_Uur+=T_t{(@fKPHD&ZI-ews!1FW*)5XfFcn$y`t?p*?EvR>+UL{z z@=JEHxL0|yNbx2%r0!vq9`u}~eX)MvW?6*8U&Xft%vpRn6&z;ixHJ4DdLp>s@IU(;vBSrShfidcBm%lmCP5mwi^H%=8|Bj9P zed;m!u(U@HMMF|L9_yyW7cM^EMJZ7;skJ+sTSD><&!HR@x)HH@5&>aqRda2p9Z{=i zk9Ugv{WtMu=ldI5^}%eP&6Z-P)UQStNhK`T?$7x6&{vW5K}F%*q@L88M%->KKN_}+ z?WH+7F$s2InIBY)*0$T_hieAO&+oG_68o1nb=Ra_Mrdeh@Zz9eOM35U=xofPG=?FA zXs&FY>ng(q)qzj`$ePj?+8gRdrhjs6i|4p8(*AbKt+pVgR`*s|xWwJn*!Dqg#rK7( zH@_v4%ZGbLdx{t`y4pUpWzvod|C0Xm+c%rqgMEze-YM}V6cImj_RIg&`PXYQ61!-% z_sLS{iKNk$m%i%ZlYLsdcm3!Tk*Z*Os+}^YFuupwnpILqxJx$g-0odIZPQm3>m^vG z8m3#nX|AXjWweSNJuQ18SM`}=nQ)JlMp)$XgS2UZU(p2&62N!=v|CC#0;W*PoqeP}e%n=vS-hGBP=M?>$3!i0@GJhNkW_IS^lTf)*c{2h2 z4}*t;!=ljwz~f=!aA|Y_uzC16j2evqoE}yVuSPEbvxl3*uF(#_?_ub0Y;**$e0Vxc z8%+URAGQwPMqdEqhqJ@F(Hg+}VeW8mbO*41_&W?74FWmuds*ggCl{Q?>SI6f>LEdzQ2m_A$`T?5(z_&$ssjRQIZSU9<2m=1vCS2dvp_M z7tjyD@X=79V?av)&qq&zrU6|6Y#(g}`UW%xaDH?aXdTcS!2Ho%pnE`j@Sp#~;Njq~ zXtV(Ec$hd`8eITv9zG7EMk4^Hhn2&t(F?%r;pVVwv;**a7&;sq9RVyKo(|JSQvlb8 zt;4s`7r^-8?67XM2Jn8EJKP)H0qh_C4g*JnfDQl_4-ZF=fF=Mg4;x3DfIa|54<|>b zfK~ur4>L!zfNlVG4?jo0fQA5$4@*bOfSv%R4_8OmfVKd>4`WB;fX)Ec4{t~BfaU=1 z4|_-Zfc^jt030421X=|205Ewp5$F=o2EgahN1#zaCjhHQD}i1C%>djU-2~bN^aC(_ zG!*C<&=SD&(NmylKvw|UM_Ylu0gVBiADsnS2lNIoe>4~99?%~A=l?KxI5;dCEdV?o zCJvWI7XX`wkHe_Z2*BxK}v#v!==#$ zz~vU z(Lta^Ko0WxpmjiR0P{z4f$jnA!GHb_gNK8|qR|4t<6+`(X>Ku z?r?8(2e5zmI}98R0y+R#JUkpd0-6B0JZv0o0{Q?LJ)9h!0$KrhJzodB#Jtps`nGy`yZbQ5S7&=0`y(NLgcKuZA6M^AyK0bK!X zA8iHt1~dk6esmUS9nc%V{Lx&Xdq8{epZ~+);oz`nv;gpUm^fS-T>xwzJ`STsBLJs| zmBXvi3&8B*=CEtD1MqtoIvg7v0W2S$4%0?c0N01D!?)2F!1&?pux_*l@P3#(+#B5i z>>vIP14n~^4geMp4@ZxHCIBuE8%LXfJ^)4!Cr77%RsdcPGe@(4ZUA-lrNMgUF^D~DI3 z7l7Hr&0*JQ2jKTGbT~FT0$4sg9j1+@0Im;Phi{`VfbqlGVclpA;QcUnxHq~3*gyOo z295>+9RMsI9*!OXO#oaTHjXv{eE^IePL56itpL0pW{ze7-2m(!evWmpW_Kx-e{Q(*PI6OKCv%-RJ+vp2m{BU+yH(CRDKg=EOjqU*U4}XV&qd`Ci z0E>r*qenm!0GEf2qfI~`0HcSKqfCwdAr4d@DB`)DiBH=r?q^P{ss z>ww+>=8xtA-2>W#|NI{Y4+n=uqXmG+!^Gjz=mKE#@NpP58UZ*xtQ=mAUI1nfH-}xL z9f04%(BatV2w?f}beJ}p0=PbG9lnje0LBkzhjpVhfcL}P;oj&DVE^!U7&saPbO5k; zcsP0lGy!mV*f`n*^Z_t>I5|26v;y#Ym^qpSbOW$^_&NFoGz4&bSUOq;^aL<{xH`H9 zv<2{e7&{sVbOx|~csqIrGzV~h*gM(>^ap4F;PB`m&?2A*fXSnYK$n0v06vdC0*wMX z0a!g+3G@nR2H^JSCeSXRAAsSbp+LugmH?iQo&rq+x&qif+6wdyXbj-|=q%7Wpf`Z| zqq#u$fcD@&|A)cD!C}#80pRg4akw_0gNBc4(mp10Plyn!@bcR!2aRyFmN;o=m22x@No19 zXaeB!uyM2r=mTK%aB_4CXa(T)Fmp5u=muc-@N@JFXb9l=uynKx=m}u@aCLMIXba%` zFm^N!=nP=}@OJbLXb#~1uy?c%=nv2Uz~RwBphZ9r0Fy@(fi3}U0DK;O1R4c&0I;8;&5qn0kC=aIE)&N0Gu9H4zETp0JDdi!>-W|!0%z`aBOr0 zuzYwrOdCxBTpzX$-$q{m0JuDC z9Bl&n02n=-9GwDM0eC&k9L)l{0oXnK9Q^_s0ysV_9W4WT0+>Er9bE(30{A|R9gPDz z16V)29lZmZ1Gqoz9qj}912h0|cythG5zqs`wUXcy2A!0^#fpkqKw0MAEHfu;dn0c;;_1^Naw25^3K7HA#N8^HY0T%dbEd+?wC z!{FiIuxPXZ@OYRwTpC>fY#u%iqedeDr-zlptI-Rw- z?||k2?hkuM`+)uc4FDV-9Rykg^Z+n{OA8LcsMvL z8Z7`k9wrW#Mi&5^hmXUk(FnllVdd~@^a3z@xH;?^?Ew59h7QL@M*z!*r^B?-6u|Xi z>+ou%X%JFFY60lXjP4);cP0Q-l(|IfgQ1_2!aEFK;n6{$ML-V#lmA~6L39ad1K{)MBhVhm^M4pT92^#n z762X(6NgKq3xLhT$6?fH1mN_ra(Fd*0hm499CnR%0Dccchhw87faSx}VcKX4;QFw2 z_%`|i7(bjH){WKx-Vbw!d!svm{lnj3;Ajxg0l?zn;ph?21i1Y|y6TtN0>gXEK7Qpvm>}VX&8NmAC?dToQ9Kiiy z?`R*;AD{t%!=r;hi+~;gCXXfpT>{zw_&oXuGz#bhVD)Gv&?}%BfZL;+K)Zl`0EUl- z0v!Wd0(d@p3N#Jq3Sj$aE6_KfF@W=y0<^tUV+Jpc69|jKxhee|WfXBnc z;nL^=VDs>C7&RIJI6bT!UX5M=W)C-qU85a<-^0-1*yspg`S5g@HktysK5QMnjlKZJ z4`+vUqcwo{!`$KC=ni22@OKzE8U%Czuy}YldIU59aCz7`+643gFnTySIt8=>@Oqdz zngw(NuzUD9`UNxuaC}%gS_bq4Fnzc>x(2ib@O>CN8V7U+uzq+udIvNIaDUi4+6VLp zXaL~w=pfJ{pa+1-qlrM5fHnX=k3Is80y+U$Jz5F$3TOu4_UIFM#pG*U z(Lta^Ko0WxpmjiR0P{z4f$jnA0sNol|J?%ce5x74r%QhB{a|v}tC|+3 zlSfJ?@7W8+{dAAKN&CWx;u!C!2_reVi!qrO+q}m?Ew2@}?^YY`*Q*83KUx~^w)B7P zpq!Pg{!N(k5}Dh{;LEp8QWQ>UM-~5#71tV0J1BThzNOmAId}SLwANPA;;e`19}E1!zrCt4#ctbXj#86 zulg%5$85&mCHBqC%^@X1DJV(Mn@qQDizVbxnoiHGZ&;7Swb zs#l$w`pBXx5?2x@;$@`~n<>PW8nO(7>jfqN?@=S9FUBU5L+#^SLR3$5P*>myHp&|lJ#TlYJd?D^xp zx8~`~lJ-My4vYp;vqnEkko?ebMa5UgMLPUQcIS29r&}6VeJ`#bFi(we7FB+DHu&Cl zpRJ^Nw##B4y^HRdL^Zz6I7EN$2lEl(`0<*b^&C{RKOY>ZidQvY;4D8x#mcYL+0l3D zsZE#V&aCQ(KNw{-SWgCrK36R`&JmGX8Rtl@^}g}x(nf)W#kA5?u5sl=hmzi9M+Nh4 zdzW(SVr_?-PaTJi?q&6z?0v_cs>2;Y_c(Bkwa+>4+xU);Y6CJXyKnUkubk9dm+HEu zMISQbvnN-+Kiyz_wIty5HS1lY69U{q&Pf`9oKZaC!?Oz9g)f~uK7CJ0tz?)>J`rs; zC_2}+95I+l607?4VX{lTNk(0{ac8yFX~pug)~fUQpAJRzors!Sw+%D$4W~`_^&h0) za;OV2H+@NyydY>|6#was?oUD9;TneFmi<$A@>}g+C2>vpp1q{^^x=2CRJmiCZJ))A z*Y{AVQKs@74!7z4=XqVp$urw}dOK?yYpf7;u};hPvQEmIX1w=0*>3Dg@_d!|X1B$? zYieAz#ZEM|3|9xIRGY+=_H&=H61Xe!)BO2zB7l7`E2LV!##E6Yqn3O z(i1(L#NryXXP72G*l@W&{9%w|-P+`>l}#qn=;QgUKDVCP-z|z$JuSTDz&VY)NPfD` z)Y#F22E(ks+C5YUgmhp0^DD7)Y^reRGBl7p{@lkU> zH#Hajy8Ng@!B_r6aihJNzw&PLoo}lPMi!68?0L>6!$AJ%lJN4cfZMuaBZ4+OTEk&V zIl*nb=8fl89-ruru`J-%)Bn_Znx)rk*la6HSm#B2ZK3Y}x*OD`ZafWqfOq0XKw`u?Vk(Li$ zVGWa1XYLk#OUlPh@jGAayL*vUsRw7gs@v?AYkK}%5F_c8Yv*fMO{zKtJkA7PonXDF zBEvMt)uZ>FY|&)xarO_Il$`0418NN2u4|e3Q=^|tj|3)0aYbM#2j>~H3&ET#M}mgIPu57r#@I?uUsVCcyFA=%Nv^G56*4<%j5{z%^{RGunPuF_uA zh!{Hm-RHOJy2y2}7YtWtb_*u^N-#fnifVnDo?WFPU+w6%S#e%*E>@PE;>-t+Z`WQ4 z&&ck`PwOgS_TLbd*~8VNB^iA^A>i!Y2+A2A+DqrJi92)s<5+l4I{(H0M94sBc&xNw z>5Dtwo~q7&=>kF)og6q)Pm@eAWgN9TWB=33{4~W021O61ceh-;2TuM>b=etfP;_kn z_3-^%9~iWw8n}abDJpg?9x8s-n}6>*i?!>~LPewJLnhu1Is^O5JVY&=$2Bu5$BHjn zN?dt-HCRf7Ioyc*b9y1;mGpSMKr=7O+*20|A_}D%_?fO;S4(&oBuB>lB`&3En^^n% zUFmH$J}2#>%0B*0KNlFEYRuk!W5o9C%M+)6U6eYQ983RBa*=VM;}(C&p;WuV_r`aR zPi0fI{rXfDJ15JMe2!*5g;YXP$-Y8l?9bt^pWL^di1?z%qZzU&K&q!Q!`_8YbCznc&9Zd4u&fHz|vEj`7FRlq1TFz2)y$+*} zOj~^o*U8@2dx~wmO4IrHZ$5idm)=O_Yu@A<@m24?N5jkFLkw~>{y4jR>)FBoT<}Xw z&9=I;>P}GwZKXkEL5#*ltS4WdjUc(*XmFUT_uTx;_}>7R5qGkI)TjH9dx zo#evNmXBuxALG=^A1^++)3xwS$!#^SUO}rvgjQS-?P{{^u#q`4*2-#;z~8Aba;9u{qx@Bw zfFa%c{)b1ZGo0JI%c_g+FhA=m^_FQ7-u{aF#8H*s9my@Lc;@{Q{Ov=_hyAN4gH*ORDZpiY#D|qj>N$ zGw@X;ZFS)5$u%pRB+)~@S6aoaPMNv43lz)s7z>(>ha657RBI2NTN0SF@+%^#@Hc$1 zxwPL?=<`1zuW036AFkZ%`W$$A+-wp?{-_^{-lqRYO^N&+$5+}9<$R;Bzt)D@(0({G67cV%_2 znxEWaQ$2L_?r^+!MNOn1XZzn;CFAC<=y29Q>J+Q|bwe)G!|xNTTed8Bu_S5R?0Wa3 z)>Nj>T5_9*VR_SB=uM?n0g|F$bB!u8kM;|Ap5Fd{Ox;ymkX^S0V7f!;Zs~522I-O% zknWZe5D=tP5J?FMDN$ZRL=b7EL#4Y*Lb{}sJ)eW`y7o_4r(@l7jOUIz5Aluk zoOA1c8xtdX6~gI+{S#9rZg`X&Ee^c$q?eDo%YHQfeEsODfe)k9ZzCdng*NCSqd8K$`t712mmj`E5qZYEr8ciST8l*Vc2I*oNM|onL zG%vFc^$oTfdd=m@*?ZYBd0bAN=n%Qk-5PxAqD6BcC)U3A}Q_Q9T!)|{AJtC_H@&UrCUp8C#D-Eo%Gm`OX{J@+d*hl5;ZlP%Qr znbQm36OV0+_GA%V;Qr41xr$Ch`YY6{Y&%B^qsj5BM?Te7G$*Nv$P-$MnvyPIBSxk{ zR+6_Qn*j_af2(f%q>>iQ8$W&4)iHV{_t7hLT(X%L=EDx33a$Wv&)r)Z9lknQ<2SuwZ=$aA{f5XJ?s!xY4<{3)PJnQF3Te)D6@et)< zCmSN8Red08KD8;@$2UNukIrYMRqP+#*U?R_W_a5;B#?P!Doc<>l1y!-Ln=Nr|cTNo4X9 z;pn`U+fTZ)MQV`4V2i;x219)%4X&Fwz+P8VO$!F zQOoghzQNXab!%V06kX7>j@l{FuyE&>=4`mM%c|RjeiG@sM;P9f-G2{ou%_K#u8S9D z-M5sM)mI!=)xLcF$#}EX^dcU^fwbYlf;FEzhxk=1E@?(|wr0T5!|S%QVOA9tlzX?O zU9BQ=SX56?%9G0d)n0c@47{2dcKY|>mlb)DZb+Zxw`tk++=#-HG1I);CVmzH=3^X7 z|LIc;KW`e%s?rT8$}F~fiN>%kk|X3wS8=!Odza=RYX0%2&>*K%QE$RNqtNeT^Z-Q# zUFBPYOD{?O=CyKgE z>G;{2kkICO)w2Ve`e)Lj2NKl6)}f_+^QM-6@KQz-hc3aK+CwXa@J zaSgbBbT1EITDpTJsd`2-oxr zQXM9Wnr7B`G8X2AyrvZunmeMrAF6yaHkoZWB?rbyO1A_O`^w7QJd^0?b5Si{aJ*yrfOKb~kc-)( z!j(Hq=;T52BZki1tv~LS6K*qT!9jd2%t_fji^n>NC{?JAA{8mUPsAD3u8Eb^t@-*I zJ28!m)8$q1C+fcS>=NrS9rLf9DDAvLv153s^`5#WXD>B|n4GA@YI#`)cc`apS2lfg zR3rX9NC)fuX|x>+rhG4*JdtqDP+M&h=3Fg>aXK?jS4sXC z5=m(x>E{YPhqoUVDK)5t5v>u%NO?XGvcHCK%P^*=k%V^(-EjiVA7wK{~dyf%7ANs;fn zS^KT|3S(yjLW3QP2(b|&D+1=URTa+OscS}xA6Rw=6RfeFG}`>?b&>j=ZPz0|)gH7x zqm>^p4HCrDEJfSgAIbH_C-zQt$;x_ev~^WI(GzRE`c|9YYrymPQ}Oi8h4f%1f0-DL z^Y0!7{jV>qzec*BdFKC{y3_okT*jzKE%Ds2PFDIKlSP%R?#}(S#?}6nb4eFYyOeaM zMK0V}DXBL>pDo3S=9*4@Ik5-Y`Z6!a&uu1unOF<{$k3qq_30_EPBS}Ms}|NIZcnoM zd6SIY?N@agI88om&%rO~?j;Y3UCFF!&0bM8i!DsINLkY-v9W)dY1}6(y zbL;PK(FafS;=Ht+;PwMcy2)4|n1iY;Y0pgi@9ZEB`DoVlJ! zI3I&O=R9zqz9OWXe+G-NmZ!3ug`cj8G2|ykMgB{a`zQzbJC!y1YMuurcTl<}11-PT zSGR@U>XY93i|c+OCmF{m(?1$Yew4w0S0~dK?N`j}w<&|tZl(QM{=g_ftd=S<_ugxd|% z{iA@J{q23Vp~xW-4R&j*cOuberH~h05B@6Rz9Blt$|_kdb`{`wMiX>L<#R!0Z$nV~ za3LoH@1OWo!~XJrToNBp1fu0ajnF8*?{Ukk)-Cvs95*zb_Z-QNamBOnS`5|Q8PYB7 zNXPe^!zmC~VtkvqoH264R*+-keb3(|k+lvAcgI&l`-&<&1M)(TrxR>P4smIxdQ(2h zD)pO2_L>o_$`3x?*Z-zFua-|uKx}8M?O$E|@3&shByYh{v>W5X!W66BaKyTDft9&^ zwaD?VcGKcDUeSRrtqQkusyZ%ze*(`xs^~hEn*vnQvaj}^lIL`(dQwl@)zxfIoCK8P zoci3YACFP=H(Gh>TtD>`?JH*}vAq&S5Wz8r$#sE>oWb_j$rAinykOB`MkXH1+Bc#M zX5-a;0-ZvG`S-r0D!O<2+{Ai+K+Yf2nc9RVygGRJhWc-OZx;Sr)srpnzvin7-`xkk zWt;_cXQpd#IqkLx;p>nX$+_#~OQ4hfwyOPAK!MYfD*fm(O&51N;q_u$@dI@K^m?z&t~Uv@ zgvEFEUQ2b2n2e1c(fkY+?XL3kl9)kFM!V(w-yeOosPF$g5?BTN27*uL3nv>`NWy}` zS8j?tdmgyj5=55w%9qXLTG`pGoRRq+x&MC-bjh3Q94#2nmRGTH-r^)lP;;8iJW#pp zdB;|r<&2Sj=;hMH`_6i|A&$O;Dq)h6QX{MP=^$13h|7&$#fkZ3`{y2yV+MGCK1)2X zIkAy&|L>RJz4r<><=#(S(@pMK-p?+;6g;<}e`WhNiv(|Qm^UXzY2hti-ff})KhKo2L5w1vpV$?li6>vF}s`BW*Ktf#mBaB1V)R+GM}V77^< z_4Tim_0WasK`o7Qg2r9cFKU4}zeqYaQw_I*$-DGKYaY>d8cj4o$&^C)u!#y3FiehwXhX7{3Vm)Z6}>l-)yPs#+J!Y zRtjavM~upb53b%J``i;t8|EF&HC_+#rfHY7Uq7Lg>SeG}-d{P$48@CCVM38>n~h8( zV&Ca0!t81**q@;idQi&JT%#2pXI|K`$HtM=w^z^ePU~yfq`4>WU4@w4Pl>ZG*=VgM z8Q#ASDHj)u!^Ld`ut>h8*(IVxyq!*&lhiu#%xuRwPyzE z3$qC{;vcM+@8$)D(a(q`L|U`cy?x1|S8Nr8iGLR_g~U@ky`v`kCbwE>kE>c9rl%RF z=nD(Ixrp0CDyxNJxZ~3lJQvJiRUHSVDGF@1={jsX!qv0gC>&12QSMH0)%nA|3AVR` zbQFjFQIAZtYF$oeSOg~-(LEqUMGw=jeIHSbVQ2e zVbyqhhDLdt(RsW-`!d5HdvrM!ffMVn2WK71Aw0Jqk@5=5_4x%TC&vXXM$kT(m^*f| zWeYs9AxSLWJuEYOgA&`iQW}pU$4k-ZK+tk%P7vzyIKbk1J>~6OtSj~Bg+fez!$rN` z2cKN;@;ba$3uVK99`eplKC%tPI4w;=Q}ow1xwitV$UOUoUgqqTCNaw)HLpd-qH}%M z{=&#{$#H*S5Zgbdo`LIAw>$}jd+WY)fAqgxh_T zL%%whqv7BB;^8EY&Ln>JBrjatU4X^N8GH%jdGYH=+!uQ)CF|gZX)0S(<*M%Bpe~DYbmovTTzR|JK+0!z8zermq|J z*RFmsNmOl>^#_R)GhaSn1iz}2t?SqDt5P*!i>9K|!;hE0>sJX(F}Z80z24Pn`E?D1-Olaai=vsmbAcvz zi(7HPb8llhe!V8JlTnh5#|<|jx59l`LVr9~GL$%D>4gZc$?0%F-b%Ol(Hm^~6tb`% zI}RpvhUL8n+F@9WWfPZ!4!zk0sLCZIn}Mi-I|h1Rp7=1NqcO4?YY|)1R0{pEnkox; z6ZsK!I`@0Og_-rep7Mlr6{1nwi9f@4@^CJ{VR2NAG?s53aehDd#W1HcVm7kL!oQ=# z-fx~1AI^!=T9Q&s)P`O+xOmMbL%C0#)c^@ zdu8nhckJ;m_havOJ=DAzI6PS_+spiCdG^1|Yhg<3yzh;8Ec$GV@s3CGKR0f$O#LJ^ zDQUki?RMoR7PjmYiD!^g5dB8pn`yRAjx~yRL|gH=$?}BG&8PM56TY;hsO9c?lf*mG zVjg)B!g4H9q>jW2MZWKO#PfRF3dki!$HW$`@dqT_NZg2D2rfzvl)2Aw_`15pw;XgY zSbxmnJ7?9$_uTgXlvXR1oVnbfPb&0v%~fFx9Br7bv96sm0^6mc_-BnTq=t5u<(w|$8~2@T>}kTU{9n)1H&Wn6 zs-R|sjJ0WW5YsMckS@7KQEKSh@ZLK#WqcUDqhV9b_6wzRU+Y2JX0YVFzHg36bu#i{ zOqz4Y28<6WZbn{FEKej5$Ob$qy3R>%Dk1-sYkoTx!^OG(AZtRJ%k>XMuYG2YOUAvV zd}dzUc>alses9V->~fcqpN|55jPR=l8)R5ovD&^8b@<~i{_pup{MSI1p6w(;rF?p& z$<~W%519rr25Jmpy1p{5n#IOF%}! zC*GcB!?=&(v>+a3eUv&8DL~V?Y~fpE8~GSt>IsMGEB>tAP3xw1HMDxmkN-wk%}%<} z8w>38-&*~Vy!8s@?xLdVsu^BG#{SQJmoTq0(PIzvuwR>gN8(;(7|Fk}3DwCB9CR+s zIGm=e7-AjGzGP~C)>V^;m(HQrzZF$UnQ|9<iF8~h&JeF_4%OaY zM=J58f(9FPq+9)_EEs=rS_=4mnIacnj@&ewr1^?Rz`-;*%VV}GkxY#BWXa}(z`N!V z{-w%?(Z4c&6I0qqJVle);_beD{I$Ng%P-rL`=xKk@4WKh!oq$+yR_!9RY@hDWBpY1 zxgUfFr{|Q|REoX*rfDBi&mFMYE9whm7{&kg30#LS@zROWiG2-CTHE2JX~N@}GGS87 zo$f=W+>)9SiM8-}_eJS~?@sqy`IpD!Lwn@U@g5gYaAIm*MK}975~M{-xNzUSSthC+ z{g>Ba>e5ubxc~dil>xycy<{{eclZ zq8V`S@R4sOqlyo;J!zEck(CkO^w)j$Hw=F)!M&v{ZK2=x>U+(aw@UMv-QNo0pLFT&2-k); z%yO~(X84g#ggs@fhofyWaTH8XY#cqChTX!(ygqBW7MpQ9U9(|3nx=!LWVnWcD>GQO+;Jxehs3(JDaXYO`$rn5 zS9N|oo~V`PJkx&Hbw>qle$VFE2*v2Rfar((nAE{2yQ{m}{*v@+RrRab4_z?T`oCXx zB{98n5h;GI3a2tI@+O87)XR!IJXEGs+YVy;g0UV+eGpAzAAballI|u|s*y zEL&5qs`YQ@p=AjD1P7{Pc=C)=ovH_cpvKOXV=hnisgl=>jzvP>>iq?3-W9)G_l$=2 z5e=fN=6i$Q{x}5lBgxmtbPd@2H;&9$8GE!-3S%RP>z-7Fi3xc+6s<;RP?tWDuB0SD z@mp=4Mm57hTe^dnzH#y6*4=S&l}$t8p@M8_Mn$8L_*@nKO|#iGPrt=P_fV%_f}PZ2d^;Q$qooz<5!^? z(V4KAmk=0b?R8;Rl-;VC(kn5bZ`IuXXUjY^d2;cOwuQ<`&L_mT8fPHyg7I-VrF6m{ z%+eaBRR1mx#rjBJMLI(6pV*#q&#!w1I|>dF{|Cg*vGB`hPTc+gf>FSL;O#O#OG!(; z*DyotCGA@nZTj2dG;Mn82lzJNfnQ9$$+e^(3(P12c0;TnaJ(nIO+gNocS3`)01kr9 zGe|7^{?Fe5pM2E_zcw2~hqW^=hW`RA98f~p^90J9i^- zH_ZYv1jk41LxZYE9kt2i6B_u3>4w?mk^X7;PjGX<4lS74LYjr*v$?YX(IEyAn4f}i zVzQEivWf5vfJ88Vcl73X0@s;X%{(y*GJ30%8JV4B?zDli1ER>CJZR+YN}xxI58w{* zia-LZO1Vru`L{v_GypGx4-|T0K{$0Qzk`sM<8M-T?z`-x)Xv>u@OZ!(muA~2v!qTw zZBzozLu@1P!GCO4rCzE^nXY~S9l;XruuwdxiFU#?K+#V>dhv|?l#GwVN(t;9&_@$0 zDcp^JYB*Z-0QnFH3DodeZTJn+`w{wR|N z=Lbv@-AuuUqVms9k0GGH5ipa$5^FrQJ1XpXmx;dtpai4%Y#|E#e)46rxmoVJH+|5%}{I@FU_2f$&+p z_w>76rb4Av1K1Ee-{+HqTdBYyKf1iLkt4TUYKqk^AacKzfLS18O#n$?5pwyG?(Y0-!i>XLtf&3W4pF2>ACjOiouG&H^w*3?DE# z;k#4XsHNkpWnlpN5HklX?VBH$PJu^v&ps^xIxuRQ#NZEs_~G5o%{_~~`z`w(U_`|G z0jY~U72DjU3w9Wt0{n;gIv{M)qd$#BPC~=p{{n0Wp6xgStRiqg&MfMuM2BO1Ckg-; z5gQ2n&S2JutyOO=pFIztAY$==yGeK7;||Q1V-8LQmNB6bjX9(&5q7z?u1?>4sq5`xu>(gueK)bTeh z&xOFHjBeK$piIO)0xir7{#n85ra0D$zfNZ_1W*pKXoQ{V&JXWhZ7L00biD4p`f39CLDMlux8gu%EL-q`K=Q2Thm(jGTiyu-0v<$s9uT-Aeo5~N4+Dw4HNbj^hg)#*%F!wm zStUvenm+&>N9VU2iP%Emd#{sri?Fsl(0ug)8X}etxZGfqrfuts#cBmsfPILaQ;pl8 ziVi+JB01_{J_ShUGOH+F4-OHiAJ(8xJ&(WAE)zeXM8y38trO(PSe`BaxJ|YQ2oQ00 z!sNxjKlWEXa?qADxC`jc`Nh0)5ST?^fks`mswROsV_x$B7!e}~46oNk?$x_rbXWDl z00t40mjpq$`47Kc|L7`rM*wI4Xxp?cGHuseMGW|5@$kE7u=2PGI(c z#oa^kN{?c6z&4==z=s$)AckAsD{L+@ntad20L+0m{N8;>RCDR&f%>G-Vfnw~h$b;I z92_6;eL(1ZjJ5AbhBTSCJp$~9csk&S_G2v!>4|;llw=#Dtmi zpMl-Wkb;oos2v8-NHArw7dOKykX?%qW=HAQ{jem`C=_ z^BY!Ie{YvGEY(~t_SvuSjqKeU2V)3~PeSlR$o}bb41HAqhluF|5?Mbp?nS4@;pXiD zzz=+6@~w~S-)^*^5@sK7+j>76bDvC%x=jNf5y;;_om`>wr*DP{n1B=i=l_7%EemE@ z?8c|BjGhAq1aD{5JT}@)c0(%tlL*`S$IVOs__AdSG4O4`13SLeBvRvJF&)?gc0;Tn za6I|O&BU6M^yY{xfP-N3K3G`ZN?m&7dmcx_`^L+8LYj|%lV|}f98kio>R2H*_Wp8P zT7YqgD+IbHGH%GHEE-cKWe>;@9N($wI^XCw4rLBWY_noZY}Y=)*H**K;O2lGUj6Ft z9tjx`w-o?HhZsa)ey#Okwp9Gb6)4vL62bh**Ir$IUg8?MOO9z-d$(q%h#{bcy&jAm z5Jj=M!Y@g_*93hb0PYa42qf_5P(-r?6=v3`3E)NWfsg*Ik$=FdskCvrL0u|d@$xep z$ETog;PHSn)}tScDj&b^<^Byg53!BF2g`RaQm8a9nbn2>I)WuUAHt_n?7Dm?Po*AZ zDOcGYUL$5pE+9q$4CA!ZU-V$Wjxf@f3i!=EPrD8VSY zU!D%w^}dbB1x2}gwX{p%JW)o;2__IoCF;uyy+$KHvH$)85M%=&h!{LzZY9RKf{ukv zIJ2(+vM04T1Aj)1K@+zPJcVE(3Uo*goKM z-Q3PqxWtzsKX?H8fu&;xu7*&=>{#B-$i;g4iM-NPRm3F%{kvdDQS4S-Rk26~q=+~` zpmv;DWJ!B1@wZdZ0R@7)voL*Jr3q{?i0f>;UcCINwL!|(L&Puw6U>N`Iq5qI%IrS_ z(1@5pV0oeJLrGs3Jx=dr03ZaTr#9GXRSv}Q7GGpJXKiGizSmrFJK`OI6q-6JwF8P#yvd%D2Z4^;C|ZQL&j4T#o1BD049R{a~r916ahN9%PGKskts32GbwKH2yv7^ z6}Pc(ZcX^aemb-UREfAnpn;69Z^yRJtv93(0%8OwD6?|RG6Kv}i+qiet6;@~{JwUn zDq=2yH5NSem^qwPHs_i`*)3cJq=z^-pk~i7!^iKXg{%_G0g3~6mWY@_V0&JidXy|< zYWl+m01Oes2TV@0C##S%p{kMa3_u@Z=76Ow@NKZ&Q_B-ftO7s>Mop_0{2>s(jTMdu z%Q-Zp!^waV5$^}2j=C*dP}04)*A2wbp7BH}>|-9#0G7r;iu4g$}sx7AZ#FIc;1jsr*t zR! zO|4d3V#iTq3W%+vyKH*^3?MMOI1|6ZlUs&T81n!E5pxHuEe}_b^y{Xk0_z|EJjA%o z%s!!AB6Fb1@%h523*as8MNNqsctId}F`31qS@BiEM$Ui-5uXPHE+B`5VbIf~UnLQ+ z9^&Cfsj6-`J)dUV`Bg`)2{>+!l}fr0Y$5PH!~Em2@8%hKRA!{XI_$8a;>reWM_JHm@0t|0mf>{I>sBI$O+oHfG%VP|H5ix?m z@Ju^3g_jpFX*W0k3?e4a>%4Z_{n1OqR$Oa7J^*^yT}7)`;2VJuipw!7Dv}SOYmx$X zL_8soy_n_H#lOR^QI$~u7b1S|@IZ1#B{4Nz+Lcu$4KQDIQa{ISu#mtF?~QW?D(Kd) zTv!2&MC>8(zSqF}21nYEiO_EB#j#0NcO?j!9|@H861vWU|JS>eOT2zx{eD zzpWV@9#FTMBOgkOJFdlwWPtJz_XcdR$=|JHgZCDeOa>r0Foablw5RUA6ni&qI6qVs zNXr+EY24Uh2D1k&Zc~vM)hcgj#H@l;~P;WZMU=M=t140*&A5*!>YMnqO1K1DobifgR<8>%~X*(mkCInmuHZefm zculD1x2^iWR1%%4<7c`rYd6=I!2$xe1Fa}DIK%$zAQoUCV()-1j!2_eKRev~yF>vH z4;*9I(kW9}X2&09ec_?s?L?%(M9X`!?BE80=4}r>+G`co@~rIxM2I*&V2-8Fxz5AM zxjBvl0PTT!tPbOA{Yt7wN*!iQ(|0D`MPnT9oP8gRAuzturs*YVp}#My-vT&9OdpWQ z;8{BE^vdrIOWy(bfse!x!J%!h6ppK17`D-POpB>?NyR_<3p^r_Km4`B?9BTw@dG&l zC;rd>0kP}t``c6EG0=G83m6c*T>!DYZ2If+E7j4yocWfnU9<${SFb;8rK=<^gaX@~pD!iBhWC)ILI5KJ@&}Df|#phx{ko?h>;}FX4 z=PBUkfE^}U;ObYmj}|H~0-{3+#W?d>b@KxIT&rVqd!;uV1ex=mf5SN_1_>*kW zbJ7eN-P7n6y94>&r{M8`Gx7=NU}i_g$`dsM&O>Y?@WG0+KKZG} zjAKVDKXOf+()-0YU9Ig|Nq!yJJ)n<_(n4LzXagjYO#t~22MN?rekAvWG`aWYejT7l zaEBB4#Ea$wSy{fj#H#q!giJlDE@?Dr;QW9|etQu%uNV8gCSe`WA7UnfB{IC#?yXw$ zv|gnKKnX_iz1P6C6JPN(|6b}tt#)FXoa0aO1m<7@fmGJSk{idw+|+OQ06_2pfFNS< zfVqWpINa=sbT8<~2apHmEeHLa{Se5YbuW3MKJTVRcLASE2;v8U=oPp*wFa%>VFnli zCPcg*khnti9-%tve`<+*fcL=19qIOWSphGs+n!18uB{2mJjYE^N30=mJ_b69qRB50 zSF1Sz9wN36_#CT}Nla2Y)4AkVfPP@<>^r1JX@DZyvCJDaqU72VS~U~fAubWM4d3uHrGT2C=lEoy)?%8Ah1Qw=MG8OQXW6z1YEU?5yJ>f zkR#17_3Te|5;p=G5ivSxOm=WbXk?1p6naJcShsbkdCNwDI?2 z4&H9IDcWMhQ36%;&lklPWDnORo&!{gxJ96WaiH&*i;P)+r2t|CCs@ic7X=gt?koxg zOd+tn`@{9`Zz|!^3hn?fL<}D=Io!zJ6z)ma7Gwmj5$NCnCqr*T(eZ9!Eg(n45dzg4rd3(wb+rvCZ3a|`xILh8C||gP9&4_r zc7y`r11Bf<7YrmY!ySt60?`6I|Io4lL?Y%8Sf7#XtGc~M)EIcP01y%52aN8#d9k@< zd{}^e48R|lI$1ODl0XupSa^Gm0;cWa-vFLOd?FA)u2Y|Fi~3)r*f#(xA|4RP-CWH~ z^FXl5RfZDaK=60s$6zagFJ@*g;}mc{b$%NR&=RqXzy%rYm7ZpBbmyEq18hX>An-hq zeJnzHt1o6P?f?nF>QTgj!vyLWJ+!~pfv+O>stHgg;vRt(ItGh=Rb=8hlrILwCk0Rr zv1qk#rg~qGS6iBl-Ptn%n8slCNp=KWAJDn3``J%jzx4H+jso&S934=#d+we3W;5?o z6>kBmL)_X;ta&vw{urM>`l!h{fY=VmsR{lA0|?B{dExEg@L{K&u^NCt#M}XEi}@SR zTYdWW4Rr$mJjA$V;YYhoZh7~H19@bEcHw{r5uXPH zEhaX4C30n$<6sN7ryhX~Yf*+g`!DMgw} z`YoVD#Qg!S6TJRNz2D>cb|ecBAmZ$%%Ij*Xw^mKl58tt&0=jdsn>fJ$vj{BEIS_}S z@%y)O94`Qjh!F&a_vK5<%3UhHwC9un1`(6T=ukgJNs!#o`!rz^My7tsw4;5^0^bOP zP-8JdzU{+&DDfa*N5m5X*-JEkSwLaGF5(;mxDfGs1j7c~kHrrTem;BdfdQDW*mkmn z7c3-j!-48Exwtpzv3Q99BN2NDyzdftpZDA8PW;X@l!#QoYGD1+XF2Suw@%GdTIMU5 z654cMS|Uf_omd(q{)0)XJa5Q@;x zgzT$WvxkK}K$qKa^(Ib;oLI63vj;5hq`8KhDTvOz;|u^FV&s4rlHJ97?bWA^iJuH$ z4!q%g5>KgK<(ll2#R$KE{U?;V=^a+eLE!s<(8XnjGQ7#VxpF`Y*bnh^z!AThan9nA z=wW&u0x0{MDV#kz@gg z5OI3I9Gjhp&Fc%}P&gI=?SXl0b9cI{U^wg>_23gGf1)RENOk`1r(Q6I!1(mC(PG*| z5@uao0URQx4@jh+Ti0UiyMiH;HUNI$BSY#k!^I=l9u=ge`N@+k7OOSI2w)k5M+EYR zVwY(1V_p7Rngrm)|M@>4cC$ga)=wOY@Ud6`1A@0B_V-m#GvsIR=C!!eoRE{F)=>Hv zc>ul*c;I1ltWAd~35pFCU^m1H0>|qoV!ctKJlCTB58xo!JhqQ)(ahdYg+akl zhCQ+Pg!6%5;eZk{y|R=!rJLsG(E}KVxI&rVXYR#N@D{ynb4zE8<+E8pr`FaNcqC*TKFu$qB)Glu8lIqe~0EuA!A|8%O zqLw?|_;_`o*YTiV4E-^M%DM^|J0OZ|GhejlrUVlwYysRMUJ*#(S!jL#hbl}%LIf{@ z4-6+Ttx_$g_F82zdXh(_;U~ad%NETE9uGL9JAQnLeg%-rtZb85d&_5{L4LXe!wKjaz>(mhq_B(>H5n%C&07&cjU!q6@6zzaa6sCkHJCskm1OC7#SLWL ze-#-42$lm7L<}A#$^D6VDzGtCMj@%I3AfK zrX^+6H_l&FtDqv@5lCTon?~fFzJ`#oB*2e|F9gEpShMKb!HL4&?+e%vJYUhfUmE(r zA>+mKh!qZ_D0%h`dG8@M68Pc!htW2@ut~r;3{VoWh`{{@*vL9}a?->fx&cfC`$zI~ zC9Ut>;riT>3EX~hZyB8=782RiBBd-v139KS)vSqwJbzWdTY4Hyyeen9HBj>-8jQnC3bXaN5qz77c6 zh(O6tZwudvNE*O);MpdZ!72g=>?Gj0Y2?U>iB|-05wU^5@4iM$L^EYMWLP!<6htf@ zaJMCE@qe~x^}Czz0Oo;x6Ho)!2z0QWACw;UU)~W*K#qtb1gh8k*i)g>t~H7951>NC z?E#HLv+|Ike}$s|wi^&1I62;1U?71RN`9@E+0vY0KE4Ye5;2Fs`fjh2aTV_$mT{c{ zKtzlmFuHOtv8+Z9beFJx0DoZW3Xj1{0!b82YX<#SC}9@{coOl6KmZ9pILVJ;?|Ww( z09HgiAdovk@_=_Z7iNU;BRuy5{N1?@*h=7w4vM`8r{Py8R2%>;5z7c%5Idx_t^ad8 z_f|Z>M#K&R&nu9Sc&|mTudDVKAR$;i@f>iNKpmy8!(|wsnF$Nr0+flkN1%m$Uj7$L zZ-#m4zGKWxcmtF}ESl^OJdP^O0M4OAR>w(zY0)DlO!eUUfX*dqxNY3$dxV0lIrHiPq`>HpX9!IeNV&j{>)140n5SX1N?nH29n7aH+ zI{<-*xdYa=)v87@l(*NH*$n^>F>d4jT30JOEnlApv(x?n@Rl{!>^A{k5J=viTIuQ4 z93O(TSAYi*p9ci)@tj7q3p%6!U>{&T#KRS!*$^<`?tHHNUcUVtaGasLg0H8 zAM%Y*bjgpDYylb~mJhg`z;i$M86kf!oIQYjh@De-?Xjs&)4hGG^tU}WK)U~;5hSUT z1=Oz@6I5M*^*liyphU#|0j+yzsf8L|tF`f!3lJdU>@d5O%w`{qiZlEOZ+HmkZuHZ5 z_-`lV`0XS7eOa_!PT^o;U!2 z9&bTSV;cBIAcV6MH1gArcJ(aT06QX{5Xc_Kj3i%cJ$~Z2J>Wvb@9An$r*U_c=Y>cJ zkBS54E1CTAN*ydDaKkN6Z1a;36_bMo03#852)yq*@IKNRC32Hmg7D;Ez-nOq22?$I ze?Qr|s*6+9^Ii19JQUQB!N~zz2Yl_v+e=ZL8OJE8Qh;`dWdjO0l+cIj`5(28#SXwW zaDf^37bGnS{c=|sur1f@AbaW8Kk$)PjuQWN=%N^3Vte{3(yMCu zgtA(O`~;5(^z+!#yg@;nm?== z_McQd2;6-d|K~p@F%q3-J10U)Cv6sVVOJD|vr&F@;b-S=w6W3ACMsVKMSn!r*IMT= z8x`E{F(#bC*?Nl=VxxBW_ht9o?X;{HPwxmG;7zmrRZn?c+bMc_=5ojOU0@&? zmJ-p6w%CwtR66?fPU-eVF3EsmPHhrptKcGqZzG>2OB_7T-qe<>F*W&q*hv$s*Y|Pp zZ7F8>uA`r*_eOrTI&V+8*o#r=(t*F|fN-%8A{$NxS^e2i8aNUnrqsN}9sn{=D} z&qE&25trxRaebAyq}Zract7jnODWtJMF3YJ?H*?yQm)x z%duc7nCRb!nr;Lz|D2X~vbNxH!N+;$0r@73xHf(S_g$+F#>;6$WE)u|O zqfWy)8L#|*GI3Z~!rc#bCI+t!|5eC~U=B-MjY(-8Px=@~4^E_g`59x$Y>GZXOYmW& z=fQQf@P^etr?>%jN*StqejE=&){P8y%p%(=W3sI)W{QGOsk$N}U7}0Szvf9y{Fi=b zvBr6)1?$BpCjMalaIBka6%09ZjMI-S=Dyi6Z_w?U1*_cIm33F9*7#*8^h-Dw@2JFS zjpf2tq?C^+1?`$o+UacjZDc;*(ec1Y5^aUq6M0*O`Z7LNB3WQTljEm+S~HW+LxXw!(&v@Jt6&0#qqUi8O3?i+?(t{W7js{jUe{_GrkG4PaSIe_`rU<`+KIg`o9*o= zV|&4las7YQi`}AMq?xltOthMPIPojZm=33cadYl9mZIkt5<&gpDS^4TE2>M>h-u=+ zpU!oslR=ewe?Dp#^4;nDK7F|jmi@814=Vwdn*P_rho7!ThB|gkE-dfAi_=$`)2BUF zY*RU&$1IZ_nds46&Z4NA=0Y)&|V{n*Qg0Z}=m z++lBp$v#)i$8pDZ(X%Ejq|fY-s(g!K{NSxBrfXR|lg6vBfM!sB!eM;#fw1-NR)|qq zc*W$3iP+=1|6>1*VR)`HSl_qW`T5RLZT!tlh)sDh(NC#!SVRBP%Bz9Y;pH%DTV>+Zj{czn-}QqjsvXCjT!r{8sdH91C3L^)?NDAC+R-Pc^IUd zu3oclrsw3jM}N3Nnc29+&|>d7{3O6ST0kw9ihiD4R63Tpv4G><+hmPt`%{TDQuEB8 zCrr--%bLf}0LxGyq8!Kll}k15QLj!_m)NlE;*0HnD^OCReK7D(fD{FHl8u^5`GI<{ z;Tx~(3ZbSe_E46w2j~LTG`-jxnKG&hMjukg$bZI(mJx0LV~RQapxanfQx(2*68?=| zc`q&Kd&W7PWV#~CKYA(rO&(&kmM9^~cf~(nc$vhO7H_{N#xu~0UMVt1Bbg+6B=hpX z=2+_6KXm$J_sf>^nVoRSXI-)AW2DqYCeZ`f)Qx2VZc92`t(m=7*;Sg~2){2x42p@+qP|Ut!3L>Ubb7VW%u5{=epkiW?jqgS4Q zyWJatNcX;6r~lafvu(u#d$V_V`b$7M}~%1n5X&4OdT`K2MQGJ79_U!+3 znvcR;(pisLt4TSxVEyz3JcXX8Ru2eYx2^b3|5S1AaXw^hKvVIWqPg-PU7r$Z)?7H^ z^1+gtZ{kSaF&eBi(ZootbVLWeRScP%$okpTvL_?tHLC zL`{%5HeOw54M#J*{)7!ReHN_nPULTL$DmXjntDP}s4=JeIK2Ac$>m=2yML8Ot8#P2 zTVs6(gVS}%yrYekl!IdR&2$gQ5?PCtUaV`JW>atQic@KuSJvH|ayjO*q(OL<1R8T& zISN*I8c<>RI8qU7(E8`)CdL0mLnrb4H^n_$4kleKKWx6yaT!n}tTMfutb!*dKesh( z>Tpc5yu&V2=U}SgnQefa*0pbPT8tuU{o^E^t9)Wdk7+ML-1LXnq+2yHn!Nt6g~VS0 z-0Ptn*oxktMk$kdolV;I1{0SSe*RpkfR*)UmM(5SF__cT+GdGEF zt&?uKk7TW?q5{H?z6kQ(Ave;`_UuP3NEza>m^;j8X{~T=i$)_pPuz5*bAs>gZ#%f0_jxVFE_LKZ-cF+$Fq72Hty89sSxBqku ziTJJy#2YdX*{#6i3-&qMW9~-5r>-8ih@F0N%R1i}>3OVm&vKlx z*}R317d-Q(`f?Rk4hdJ|c)icMd0(7v?B85?+0bG#v^7T7#YCyIJo@5~!bES-?6Vgv zwhPY@R3f!rtZ%#mqeq_m7o{?$U-VFu7ZdqzHZY~v{h#`4?EM;5e18c;iTPzgSv!B{ zTNimMDVXp$Q$`J${Tbtyr&U@1H)E9+eQtQJMqGeE%JsF3^?9DPawTk}BATQ9p%5S$v;SZL}Mivir-F_RcMfwbBZs&OHn5mM5jv zrKyEBk&xveLLqW7QDYUhtq0#5im33Qx^ADOAwB)L<)G34t(oarKkI1M+h2C6&nTam(!J-=a5X;8{}A6$vN>W#Rwy*8)vZ~* zp-Z`0B5Q$j>CgQS*IZ_sLo_;2ix(0HP`=pTC=%XgkEEo4D>gG6_8y@S&=Dc@WCKGWayRp2N3 z+z*pzR$n@lG&)ym8^@NM|FOfwFK#61WGup&(D&0LST`|qMRNzBB$QRE0R+M>h%m}VL4k>B0WRqw|@K8@0X#K86{?x z*ji69B{+GmYsDIF&wAzW>%W(6dy~sDE+3ThlQ*7CiiATHZhuFRnz`#%{FM^Xzu zW{UHEiCH=oQP(K{jQUj6QU}d)2D60!sPHr+--_#R-r^a+FReuNPbkznszf)GWhmwJ)Mg(X>^fxT(Lu;FT{U;ZiR7kc5t*7yeN`UktxqFo5R6{lF^ zX;nAsf|G^{D^iTVmHyo*xh4NpE9%9@7Cp4s4wClgCFbNi1BA|Ug+~Vt*uCmB*HgBR zn^+N03Gl}P1N?c7`F(jV9K@D!dCb?a6L+@CgQ_}6i)?L}zFZ;<{2o<+Ltxr^y1%#7 zg4Z{swQc=ogFU(VPw*=(>6zZztjX9xg-8MKX=A1 zVm@Z~45}!7;VFD2yWE`;-9RZewGm@!MyLE$m=aoBsk@X`A@Me%>A$_KxL5+8)|<9D zU0$>6J2N|#?sjRue0&(XaMK#zPb$!Q%x@UEM=LV1uB44v-OsxF^vmi?r|m>?%T3dI zSVKs4-k)?0EM$Y^t<)H}@z^s6Tm`EPOo5$a=!jsyNNlv4h~UQ|XQq*j8R{44cj!Ih z#sjGV*tY&|i18}H`-0~tdtobMmzkY(Yisu&S4o9Izh3R^un0QVr5tS0axE{EWs!f< zSun7kX|y%h@^c%ndHz+dVo-%v@ipAg>RWGBRianJjJV@9KVe*09ta+$4&rK6PGI#U z5Aa@@ar@|+@(aecA2H4kkCJ_%bt*f&g7BahXXRTQ28T3r`*$OnkXWZf(=7pZN@v^3 zK>szkZ^RCyO>O!|Ca9&&L@=Qp`cdT7VkgieVd|Mv@@VT^@$%vB4*Thbn{lE0ki)gh z#6MS^u73#0^5zLFAfSGV7maa3POI*VOJ}ClyN}^jW6l#LvS5RM-Zr;3)D{t``7H_=_s6u;V?U1u=He?Yvx_Ln0fMeUh7wxd=ZGtnNdJi$PJvFU z%cmWd$!qBP0gV@kd3u!_S-3J4i57Rg^u_bczud^4t*8lRJ!e}yD>4Buq(A4&hpHS zJz2hZ9NS|rM*o|Y8w0C{Oly5vDB3lY99%q96x*?F@v-DU0{qh`_-xx$#5y~9M%4pp6 zulg1NogXdYi?x(#Uff&3WWHr9`D703V7+48eDrmNR7q7#qubh3!@sIWB zP@=m%l-O8Qh>*FN&fzi9f#d#}&Q<87O#9DaG^*2}3etlPfo85WTo2EWL6@8NqAtY7 zKk=!Du49M#Ul^B5&q7y8wU^H>5QdR+*`~@e;JJ1$&^>df{25+!S@LJO>XCvha7dL5 zOXNikZ+RH%!o0V~uSD-DD*~QZBER7g=5NSrUbrSC?TW-0b3L=UR--C2?S>|!$#VC0 zA&!->lZI-gk>+qrE)Z#-pI7?dN$#E`?M7-bx%k;7SVU9kzfidjx)3Bw(9F&}G2pZw z?=3FCm}5xrYN+{Rmp2;25F}Mj+!xY}X9`z`=gpWV-A-`3_}>k=AMBzh)}F4Y#0Xz& z(2It3nRL78(r^204k6bN6J(X>*Q{FD`1#q}gLov-e9WQs`IS3FnO9Is0M~zaEaRf8q|P{yiiJO?v9w_e zE9kX%28%G`loru3yo!w{7g!f%jxu>wa#|ucjC#08{$Qu|ZYBYK(RJzBX5hN2OwU;w za>vcVIfi?|8sXO;w!jn9F!za|f^+r{5IhAm_1?K&Jx7`9mD$McLLV@4WwKi-14@FW zobqeQPg{d8-6*&YG5^)pxG$bMP z=Cde>WZI!xM)SIJ99`~y^^Vg~NO-Cy&mdEw5*o%*J|_Ecb2B zwmM(l))km3)Fx~ej{R?PHj2SsREvujy>jrI90C+aYh)?Wb9{%d`&Y*gbj_I~vWnw` zjoN)0Ohug3QAlqLt z6)eednI;iDr*dO5Uz!`s`0OfIUk>~27F4)1e&qB%zF1whgko4A1z(0jN;!3 zqV>@i-C>Q2{8m_bBx>-}3>4Z5W^7!mNtqvCTMA}BXZ#iEwfi9#^#IKi)8mPf2}bDW zxfx>XL`nY`$B2{e1@DPK#LN`Z4s-0UQ5%ZEy7p|SrB)nQHNg+#&WAy_V;{3|I!pILknsN4aeRmfh zUHsi0#F&+Jf=O=`qrCZDCapm)^22h;tu~IG@-G%<5f%sXHJjjHr+T=-niLV$M&UhC z+|7Q2_iWofs<-FPn}2QG=G4-tcuvG2F7fEBBQC5f(6lH*80HggHxcT&LnUM#q}JxN zK0p(}tvGkDIYz^MNG2=}+xk1#!2g>mpF<sbIm`Z|a_`i)5uX)iPOpNQ(Z@2UC`@ zCdu{dws@VJC_#CFv3kh(C(HFTP1A*$L&CA0sCJJv&rFWoenuA26kS*@=vf-FDu@yH z;##;9F^@SonQ;;Stj@r0ju>-k9UNi{dcrv|gD^&P`4qVSX6pnP^5>KvZArNlDC}nG zqAu^R-g^@S-6A@5wE1q(%oq#LkAog!a(he zL?=#Zf$o|H>$cVb@yh>}|N1-XkE<;dQ(%-xtS=WILY@8(L;5JWP4!CwFZN-323DK% zqt3=ZV~2y`LFGLd9%@W19$6m;e-OvH@Ee5omK zl5DlpRZaKxG??;r^{aU&C4wJrG(W7`WE(^Bl3=FsT(s$z{S|GNA4`qfJ$|5a;car} z{XV48pL`R}ALR(y!&}cP{>LZY8u~nfAD{otvTAp7ciZfBW%_D+6<0wXq{=u-Jfu+2 zvmZfSy?=UR%{MVt*AQa#&M4ZqTWR*ifUh-bj}otgJQzKolP2k|O11AP-Xh6q_V0!o zm#O#_<^EA$iF3)6FEE}EC&Lu;6VZ~-&{dBjl$h}WW8!f!_aZ6}YHoniV*=)=$gnZB ztxhX)jR~80Ef!*5$bZLT6$}xpZjL*>uA`NRxEn@241EsAaEtATvJO};#14_iDt5vg zoL_qkw}hO=2qjacL zLc$MaaKwa6H;kPCKdVtJNj%MBZ0YN+ z3(8ekk1ttDFqG5Px%|5N=VjT^meb7r5IL)d)3EVSTA9s%-9nF6Ndd@tslvCC=5fza zk1gCW0djj}f2=KVlSc;~MX=i5F_y)mEMkp6&PGKp&9~^lmL}8q=&n+jXXJMkjS^bp zxpRexAb$6j6>ZN6$G`W+Q7GsoljBpP6T?pNPfv0hi|x?O((Ss&MEB}9x%M#&9wx>4 z?1|2SYM;H++RlkqevH`s#zbMKj)cefBQB)S*=)Ri!S!uvplqV4>h9Bg9&Raj$<~x$ z$;b1nBE6bGA5$@MELdw>`y$NBI*pz#H`HrkxK&Y-+Cj@Z9stJ!rXPSzF4r9Z~Q6Q zdba$%6w7>%&M-DT_!AOM9ve|lH9}EO&!mGkikQi_r6icM$g;=YgJ0~TsM8X@A6E7o z(Ynhc7?9u9qaRvwXVIF+{I1=9jFsEI`AFb@Y-@$`j1}VwyWwqTL1Q{IpwPD@oxq8L z#Xu^RaB5i?r-1(QrBfH-7$wQw@4hHm@GzYE*#rH@zJCH98(meE>9?oDXboko1x-28yPoc@-`&$Xy)kGj=tTdvrKHC);W$E(F3SopbS z{ehfg%W;?quZQQIp@yD3CTC;5ko#@YOJ>54SszL}IjPAwMu(d<*Oz&!f ztm`_#mMWpoM@6^t&&;OY)EeD+o)|ZA!HlAeJtJx?%CWBll%tSqF%*pZ!Fr8yG>s+seq6^TY$F!kcirG9_NZE46y-oQURfFR% z8hR||rMA1piWezHu#?edj{}kAb0s3cbX(2GV$E_+r4O^_U(|8s9k?n%h~#ANf=k6t*@bpI&n|W=Ypij%$_13PUy&qYs`f^)a zeeLB?UF%v!!N~@`ocC+-%4|Szz?pYV4>ZyWZxeZlse7X)?d=rn&WS3&LyqHAT?pRg zVOpVM``%aLv^M9jC6K?TMBrE8*}cC@@jG~=7qQikk0MxAg@pBh$EB zgBoqHn1pXP5~5)Bua6(^?Oz-oRXo+_4hlK^p*wVJ94tZU=LeW zik|gPb|*8bcfT>%E1H;}-cFP02@;E{u}OVugX;3}Ay(!6?A1q?^V_*0yw|;!f7|~` z>O}h5YfYB;VSRDeqm15lc{1Xzq*eLm1-j1In*1}7)W|J6c4NHdA@OO(;&cY;FZvh_ znTvOd(0iixi@4EIrCPZh+xFy}G1n;$Py~yxdEd)&f6f*D_Ggxo7P0KWnJ_2O9B!#l zm_mMv(ZrVU`WSz$hS0Uyt1MBSShm3b^)~@8m$RF1-{# ztb+!j`gygAY)Wbl3QEYyi&Q+8u=j(mH2IxJulla&TZDrfj24A66rA6>Dtn=4ZMdVS z+9Sqm+RAFj47%NBe#;wjMPIcP$$FptnUd9N5rs)z3F34^RT#fAbS7I}_jeox%A zFE$*WeAdQW_CtqI9C0x_$ofLSSHtD@0lrgm0IvF)TCnUV?bkvR>4zm@oF!SQa}PV! z+xCVE&MP=PmjuNPe%NSHT?|<2;%t?05=n$bJ7G=Z@*h3}O99081+FW0?&wx%4I=+c zHzkCq`PZ0r3yCjGZKmV(aN^IPSvMA>|4rAZw-__0DZ{TR5%yZzkl?cjAF|(>ABza5 zyI4NOl8FckHl-oIIw}!=d~K+YyXa>n`b;tXHSiQRZyY|r;2FIWdHmP1g;^72Y%kil zfXnPX9cpy@H4(ZG{fP`_u)Oj@BXcwtg2_G7?pkVM+G?oWkC;PZc(>`I_G5M)U+)200VP%crI4 zlfjc)ygtturyJdV!GNAzRMJLt2abj9ZUm#P_4O}8*vzPf9x?m}v$_Z>E(nQ#uPNGvHhaslI!e+5IeMY^)>;!3`jO0*OZAkX}c zkcVLR%jwRGdQG+FJ;TrX1~*z2g^W(i|8t$gnJJq^bnOG2aw6lJ8x~`T#q<`LIcxDC z(i@V-!j=tuH{LU8h>kDB=2n7 z6Nh|1+gBWyLtCaxuU#Ujk6}J|Ne#?$+;EWnv*nKsv)hXacKz#ptuoN1DA%3N{3(3r zM7_=AhA9%6@(RijCqtHBdDe{JtL5rqDVYm9@;kl;+DJKyeWl4A@mTz~3!ZRoCPVmJ zH|v)#YI`H%6%airv+CG1(hu2W@Hr8WDIEnKjqD*hq@LX2J?_xKo32tjY!6h{{)XZI zAU03EhTQs4OoF&-PXN=~0=9L^-%N8CC9d=>xNaG zkxi<+wtN`2`CG9)6C~qf+K}(1%RlEw{KE)uFxPulqi@O!h^Y_f`;bCaEbA!tgqQ!` zML`*v)Ci-|S0g%C8hwa0(NyyCrW0l}<6g24n{eu;FeNG2{G#zQizA6zHNwHLVZPU4 z$yCOGy??rRjJmXf=W~s5-s#t4R4bk+->MNS?g#O1`4?rII)tUdHx%v}rOm6aOvRQE z^iMhgMx&=9SHoDr+GvRP>@!N}B2@}=?B4 zg{w6viqAX$)9I^~S>m1T_L9bsy*^Cc-X$7Lz9Jr*H7RH}TUH!T6;rZ^+oSoEYwoTA zbJhMu);2rjU2qoHKQhG6*+kl6#-8i82k~*^)ZI86C<$Q;&bm7MRB`yX^X`GfT*eJ8 zak%K$e#bg7d*OLeoTw(3ZZ6BH<<{g&TElmh+HQ`xyK9w2@Rx2|FoJl3bv}@8_;{vK zy_t=+ft7Y|@(ddG_w1XB(FLv}rzo@<2#ux{*5kkljKfF)`;~(z?j5J z^K-?v>>mkihogSt4lT<7Nu+Muz$?Q@4_{MorLRiAW!K$!vicog)t(?j!$aMzt=R482U~!>>B1M* z+Y)s0$Bcrv1@zx0vQJJAPG8K<=});sd!}YsORQ3~zdV$QDtKQ=Y;@%0FX4Se$3@FO zq$S5x-?}C;)hIR9?23dH5sW|imlw9$hwrB*L2e-Lex?Ndrr2?9K%#N4ZJV-1^CH$+ ztS?3#+_C_h{*m$+p^rP=KkYWr7~MKa2v!h}Dop;vunw14n+?16M7Hv=;l&(MNC zto}HrSM3JkMLV0`kXh9uJ^q$3C0ATdZc>dP{t1l{t#=E7x#TkrTIY*T>h&c& zRfL-jDb-|Wpa6pC=QODP`lxBaDvcvbs`!_kOQBR!V`r@a)*Zd=n@BO zUB{#6``7946CMBH9ajk?*0969>vD#;E@^fVGKq9gB;oJ3A_vFj@Ab>3$j83i=dtH6 zO0jPdUz?P^|GQiF%^5c6%cJj18W&m3x;6sjC_kxyJxodcS`C^9VmNxw;S%gDn#l|@ z?I*Si+pzBhi=*G?`!UDzuKuY90D zBY8E#I-XB-yb{S89?E0W;(iVd!k4K>&XwLkD*1D^SfqR8DvfQa5k5b3`$?~*VO7r-ykT~s^8Z`Gm8OPUG}M>V zE$!TZAkSy1cXA(;ocoAvi^IwnvwtT$d^ix^bo}kN&{=1@*`>9lJQ|;KEX|xUuKjp` zZFK1yKgLQ;ti#L>9$9Letl5!ahur%_9kH|fv`&lphY7C@o*#2#jMD6x(s}XHPR47J zMh4mxfm*1XwrpW4S3bFmxYv|e8fSDStp2SuB&`P!g~ic5G7YHk_uGd$BIJc(;dZ@0s*s^qgwywbr80bd0J1#e1jN;%eDP@MC>8D!Y+^U zB-G!rPm@pVMbmeTn!DUEpOrXfp)ui;IS|+n?#SGamr*@h?jI?Y!^I&(v36l6jUD9# zU6=7B#*mSNr&N^KYuzuF3kW%S7vL}&Weatr`Z-*v^N@|}rz=fM3=#4fysinyya^*H z=(mN>|Rb(Z{4gfJE-2(m*ycQfH^%Ko+!}-19MdeycH)%0dw2*_ z%U(jMR<9zEH|ag7wHV#&=~0wEeX80F=-(rg;nf*Kh!J7^zJ>G6ev%mH;h8~Y?9 z)aecj2_GVmFiotzTBvGvSub_eh;SH>IJ!q+_N1Q=^W8Cj4x}gh;^1rL)OKZTmFs(`wdk2_39r4%b@2vJ!^ap{Ua#9KOqe z<3Iu5tk(xoBoN0I#t%)TDfj%I^k{ymC&#tXMUGHPDDIu@1Ri=0#%g=v*r$?@3}fm) zlQ|Docd^e)Ay<65p--tGovCu_ourN3D3IUc1Ev#rRfBkyDYU}&zcQEO(fhKR3d4N0 z*k;vr{UbP@yMoUbM7;{#UjeHz=><#fqt+3yMYP;Lr_Wx7uLz0VC(`6CWTw_zXN>X4 zMYh=ULvL)H?m z$w`nCUhm9Nc)#GuYE@!8*SE!iPRdmO@I>im`P#LxJ9)y4KT@%E-}}c!qieW1_0HDz zjRLn|jO#aprF|Ui&+Z$mi+ZhCrm={EOgN{EYB5cL{~ucOze2QtokY+CP`(H;W;7rz1D_sk%5Wo-pePTFeGq2f8t4R!erUqO(V9yFvu^ueHQ?%M z`USXqa#IpV&<;fP-Y8~5ilG9>PL@}t9^#ou1uVYn8TS?F2{!3=mWpWk`hacM&>l+Q zarTG#R&dMZ^V!!1nVOo=by1p*AOHw#Y_!7g%{G zme$L2aRNi&{jpo}w?MP-4`tvs-;lWw_<5wO{LK^nB{a|gpvJMEgO1_k#A&hCTebq2 zde^i1KdsQB|3DWInmq;!+6LY@3H+^z*x4bT)UOu@l{F5<<`APeee z#yl|k`TQ*nu|8yE6fgB2}|W-IDTw3G2r&4FA!`e6oAWk5@B9DB$}?$yIFL5qr$h zxcFv%R@HI@UOqHnf6?>rQzJ0{D=?3rN7e3eG7XN|3f(J%fuWaR^!{0~GTaS1fafDJ z@AtPb=|LrpdpXng!?YI^n*n zf?fcPU6I)+XeJnaPm3mN%@+>%`;5Ig&kEv0LeLQOD}KB55ftPxCUXow(_ZTVCf~@0 zrtW}{{up!x)>ijFT@d%?#qly3)K~>AfzuCZ!_k%9cb*2VfmDtVRc3KOq1m3oFC+G= z5@7cqH4Q(LKcY~A{$Q?98}F^5d7g=HXW*n5>J51QLOH`PX7{*yph*b-n3w9PLW>^~ z%%=7}Dh?AE|KizS1_QWOZJ<-ww#^?P>FZ^&ct5Z3zGhq-xPOjjv2HSg)(6lo*b#s1 z;v=JX&WC0j`q9jb23mkh?Ve{&e-{PNGnAK)#FJj?AWJ2{{;{s7Mgx67c7Z5FATReb zXdIlF1ej1+TPP6{rWT4YumeCdU>*o-aZ#4k47vw|EN=f9&1)D2d`EF$@a2{8%rBF| zR6s|toae-J;*5O;O#9W|e{c5xxqKV1xuy$h`ZdrNIDa9(;E=4^0nWXWvQ-pV{ZQ4n z8Ena^Owb#gYFyUIC{FhS`+oKCR0{mQy1&F)1(fSo&>&okMkrr9DiZ(?uPeg-1DO6~ zMZp;fCz*B7CCF{V{t_@-7zRe(=rhYEaQ+;A9I`K|NhP3FIJ}}BnD5AS0B$~1H3tRQ zf2jg)hMU?XIM6Qy6sPQHD3q@OORs!KvJZNIgzEj;H_ff$1`=V4RGx9x%#K9 z?O&vU#g{9WqXs=eE&c@WK^im*uu@^Q30H0W-Tv@!1BW@vaW?6O}#Px zz~BmRxlsi+zGr_^uK8$00{DJeQtR%hg&cFvt-xh{f$INx`N^`a99X})Rbc+{G&JT$ zxD~7jU;6|3Pgf6sp^t4{gWk4^bp;(jxrtQ7ej46?;Ni8NDV9dWNu|7Yzl>9}ex1baX)fUvsxMqFV)EqQm}xYD$V1N=R{4EZ^u zfebom2%0+JiIvrC4TIi_4b4qDqJhaTgHq?H7&1lyT>;N5cA0xMQw>610AvsqdopnP z;=4|Z%_`OypfzayW8dTg5k-Xby?qJSXo(Qmee833ObB`9MbIA*H=Wp_J=%*Tj#WMs zT2m7M&%e_jO+H@DkpY?n3klUka?AR^P}9sX3h$n^A0i;>}AP;L(N0pnCg=VN`(*PwC0ukEnr`dnRGNvR)S-#pp~ngLEx>!d-k zH(by?9LW;lez?5l{vv2q1q^*Ai!C(h2rPVjmM@ab=7DLKGNhjcE}tQ+3JY@% zcLcNrT#f3>Vf{HIz`0A#>j?v^FToe5JHMcj4SIuIc=JA&Z?S)Yeb-k*X9j*hl$4P* zj2deQGzd~>Xln}gi;2L){~QcG0;XTJ@eO*)?)wbr5-wPTiR?z$XMvHI|Lq@Cs0gSYfcf0XP4XoLC&#f0rg1(?r6QR?sh?CSR1c(0X1T+HGb;3{$>4gWNbI5}7{k2ab<_FAu?#AN_ z=myZ5pGx$zDh)vU;BmGt!t`7)0UUmvuUiCY32>d=^56=!I6x0k|MzFiT#i;Hu=%`( zUV)%5keZ)s=`#ED0W=a4EOW#A@t2{%>o?HuXn^J*(bti0sU!jwbQ5a1x;QW+vg5$; ztN-*f1RX+Qg~eS(c!DZuDe5ijXDKF)M1bp;EP^QlZ9-|-Fz*OD)eYz?b`DdYQC(p4 zf%P|VvUUNz0=I4JkcTE#F=#FlQzyJ66P$jDl>UrD3`$L`O)S@K&S=Mq?aDR{~S&=WATC?R`_%_sre?sc+F5?ebg z3w-`S5`+e53~IO;G?c$zD+2G{Tx)yoS&x+f%zi)Q*KeRZD9dQ=R?hED1_s_BetAe1DP2WBxCgd$jeG z=C&r*I!)l^7h1FKwvH%7LqMOTLN3DMYg`fUoJaHfwhPGW|d*X)Y=Vt zft_Pq1!4zG0R*Nc=KFgEIN4@?Az5|!Wp4f-L z?k^LsjO{!mv4j4glWs5U=jvtV!a{-c+*~pz@cau1g+aV73(25Ku(P{LLjCmXm5(_C zn<}b{7#RP!bXzp{j7wrtRF4*X&R0pcyzMURdUI`KJoHhhgdyMu9&A!JgznT)^PlT@d>lRB6tD zj=-EvJ5(My4;q+u+kUr?z~w)8rGDMZ!{Y{R!PKsQ>Jo!^ByjHG6_o72>ce8B>I%5F zErH&E0sA{MDvVPZuMkeCR20u2IjQSyH~RC)!#!?RPyIswzq`Jd^> zHo5!(x&#E+?Ej2>>?|In8)rr>~9AuB>el?~x20J!-MebiE5|A_+H z=9+dW-9W$47SUDyH+^jgSo$gX6KK!_e0}*bO}9~*4Vs329%Nk1bEik(>kT*M!$BiZ z*XkmAx@W-#ItK(@TIP@tk^^Awovw##K{t@%6LH0Yf5{En2mgt@&!?^`k-*`{i{=S{ zmSAY@X{T!q_BXZ{6-zrd*aFX#|FBF?r7ER=6ROCb|r))U#nvjSZI z&*dOX&?X?k_lm6k>xKY*#XJ2n;w*i-3t;_GsD{-+uYji{+JTk8-Jij3*n)N-X_r%<$z}HhIQC19c%4Wsexwm|L$&?h9_8vk_hpD_hie(2hTc!C?B8hC#Qo{1pPEF{*Jo~gv_q60rqs?E53 zx)Qny8i2~|TrbcubeZI`>Hb#)08^g+cNDOEXOgnVfXkIDvqRUmwfP;tn?{0eMxS$kRe!-=*=~kQ_tdn9T_2gSe z5n$tQ$I`SMCm08S?{}iIQ~H>7;3V-3E9lY5t^>S$9ckjO>L;%~VE(b#nGA_9AwM_P z53jm!!(;+OkIne|ktGWo4|D)f8Y=Um?+3J2U`kG#uz6d7tC!=uL774=kppeOD`rJo z6cJ>*V(`0}#YH(AVC{on5j?%965v5E5R=I?nf*GID?m}EN0ze50sMUze&y=XnzJ-$ z2>R{#?9laI4zG@_5t=h!_kqdBH0XVNp!@X)bOo0!KXSY;U#ztI3eJl-2^)da&m4jp z(qLl$09pfMfdEXjTCAic-<~6hWH%3B_bVysi*!sf)j@wSDuUKmK7WVx5mhGJYb~S} zcz%y;Kh02^R!Yz$Yy}fF(3S)WF~yy>uAW;^0^=XYKfZ>e5qt?c1@`ll`IytC;PH3v zi%m&eufYAQ+2o?Wzf*kx?Sd!f?uVY712q3f$S2JNNM_Ii(BS+$|5wz^1bT+n2^b-( zU^B!s%hBH4P2nZb2UL^kUW9rv@q)%dAY#M`i<0-G3;n-NLx~v%Gy@AwzefM-9*{uy zFj&VNYhq(V;i}f24GjM8ZsCfotkFf#5q#Nl&uE*9onZO|Jq!)3fE|62I~?0Y`%*EHbwgMS4E;RdhS z{I5ayI(>MmSM>-yyqx8A5HS5cQ7Zj@gGS+?OR%}`#^-K!fB;6mSs`Nof6m{Ty~)N+ z^%k@W5a#E>n}5Bcftzp7p@0VVUmU`CgwC4h1oR946#%}Mcq@d!(jy7HNPr%IB4nH~ z(H2b$G!6LNE@lk#Z1%v{qroqgfkr^q0U~2gD^vn>4*A?crBZ7`QNY~45=m-+ZlHS& zk`L~~iZf^*{#Ie;J4Poh0EeIXSt1s+1al2O^-TF&e?bosLB=1;|GSbD*nB_!)gaIp zIQxCdBI_zx0F8vI`wrgJJ;gBa`XNcj%%C|CAJYGM=wq1-x`~3My7yA_kE+1%@1o|Y zgAU>Aw(7c@Xp$>vDUx5F)x_%@+kxvBWML);ZNm5`fu1x?^$^fkbp2NPG_Lad5k%5_x3$> z(x4p(Di<^GP#X*fjvZZ|aIOay~Au#YXGl))0nZ6sq z@uTuECW010HGRazV9HJxxcJiTd46Bdf&yUsjj10KK%XFjf2}<~{7wp3`3}vpf8j=z zO~Csj_SKkxW+A~WbFOHUdEaJg^H5HVQr%D@w$K{D(3gB1Z3$l1nj&lQrDA1+2Xr>VJzlQbSju7mzd$ zK)x0hzP0Q-E|y3~{hz<*F^dTZW#L`~4FTDo7VbF}xxIXD%+ag5dIn(f?;;&Ve6B@? zL09mOadeO%0;}yvnoYbUTJ#J!{ol`zv$o~6SfDjnn#USXi;0Sq42W3GYJqSnhi&4u( zJ~002-@@@P&xjpBr+}l-p3{+b?|rM%V;FDX83f#aWF(!~v&@zdXczp``P81Juv0{z zN3|NdajHQJz&Vhmn-V<01$qWd^bTEGi>~WB`>-uzR}mY~2M7rV4^REo*#V8iXHOrq ztP)b3MMA_B6F=k%&|$HkfxO&tlb85n%max8OVG&DTW5g5jP zFZHqBqXVXW!8N21xO`Y4w*M}@4stklpyf?+{(Y$htpm5L4zv(84f@y?SbErF2@22y@Dh64O^X?&fu_N926n@W(1Hf|`pfn@3eX7j7jvSZ zp}h|Rodcc`K90~Lg;naTjqqKsDz~-yVcA$g4fOI+mZ*$X-6EqT4Rr~&kPi%C+>u;7CWP;{^ zpNrfT`PMgU#0B^Sz}RQBvEqSF z;Mh0!+)*eN1(LjVU)8! zjWK8i$U}OdLGO@09dV5MMnnnhJ%0^F-xY7@F6akT`Dt%K10lRQ)u4C}1qD35XvyRN zuDzTNXbK7)>q0;m(G>gMaJ`ll78w2MXQYC&&P`6x8T_6~YXhx>4Qo)IE0B)i!Dvj%_w<(73T}+jbhKv2EM7?WD2o#*OX9Y_c<-G2SuuFIeB^ znmCW^2JC@HIg|*iMC=|!4*(lbvAo^?&{>zL zm`+4$^f-x@ zZG9DO-esh7d)?g^XrR@XbVd(;)i3xCc!TAZa|zgYX$7EtU*)GX0s4JbH@N+RA&yDF zAS9+4KT#cSqyjyBbM=)6(DZj<5kl`rY()bup{O(F#KP8X4rt_~1zi$=&L1)}DXmYa z5+AS%M;)Hnbs--T6#kJs#3rMtOx!S z3p5d_155+;1gz&ts}2Ov*I)l|M+J<)51k4zFbU;gz&ZTd(yr$rx+?*i`@cW9Hvl)# zL9;LR@cPFFun!zj=4W4PE*60fe_0$W9 zz!#XG!AXG`w><+!0=9q4E=gK>3h4D8&A;LT=0GNqj;g2S9vg5IU#hai?-NpvfQFx! z*To8O2*b2CZTFiHO@O6Hzy|lL_oC(ly8fm(M|r>|klRKXm7Ct(0KNh${BXEAB;65c z{UP&INCB_Ve-rQ*7Qf6GFcTmJv+|I_JA@%GN8dVv-_1ls1COgQoq(9_51 zciRFc0R9_&iUAgFB+#b^Oi9Dtf?-Jkjs4hmM=Rh2@JxELSIT~y0?m4fiKc+nAGcGw`0pEn z8_(Gai9p9+2Ww~ySOkxj9bzq_t^%NoFCMm-V-7KA1loQo{#awcCj@W*gx+Iv2M1dD zH>BkqSFv6kp!fgr)jkX`3)di;Y+fS<;XptCd97yHzIcoP3_yfKc^2Rp8qP-vCwPO{ zfTmts6+VB6=z1A&0fTc6tAK4-+LPZRX9igTI{OBfjkl9ZwIjd^T$ogI0N!Ec&R1pZ z^=$)a@6W0dN6a|vs{udoEF9$v7>Ktt7*(UH&O@NbfA-qAX>poF2TTDX`_D$eMX28r zU$;LpiU5s%Y}h@E6tOKg;0#{1OlScsvB2Sd*91q<3UvEDa;`@S(27HVJwS#PxLuHc zp{bR9Vg@?+VsyI?RDpa2pygMy53FGbmZrH(x@#5{U+n?f_+R|4V@$+!QZ2gR zU-kuJNcqx3%EDq0(94^g_`<+~3LOB=Kd6=>83MnXwsP~QTW}T&2O4^b!n}x22k^Ne zKhR(%IYO-oPqokaJNjN6O=1=3>PPs_HhaCQ3IH3>azQByabDp>Fl^OGIJvtIwDyc% zxj&O{Ne=)oaPS=@UrY|iS8l4FU*5uW4e0On*x{s8g^;NMLoj`69XvTQuIjF)F5jt} zYYsH|l%GrQ5Vv05fGbcLlap+nkIKnnn|6X#K?nvq{V)}8!E54PQotIRsY}~gpXK>) z+Oj0YO&NIt?S2aeEckIexf0+HV12#UOLtaYXS*X>YLk_+=0d)uJ@l0n~_l-;1t?4jc{!Pp)0gvYRB@T{LX;xUjwvb=^)&i z9k2_g(5Ak6holZ`NHD7;t8OlU1=#1&zpAn+`~W;dXQh8*E3FlkBX+S>O_8P>-~-xP zN0awAK&S!Z(2rG&T>9XMgu^?E0qO?93zz|j%Xa%lpI}YEJ=mwlc|$`-Pad}JCIAh- z&WI)M3w0+G;0W+C&i|O!D%k@~d(A1oGtlMd;%6kCt1_d7vC7gIK89@84GVD45^!q-kPS4O}-TwfCV0tC_dUxZb3-s_` zReq!cO@DXb0;8Y`*$m(kQZ)N!d0B;^Ba1P2Md13M3xPCx$Z}b(P7;posdf*=w_Mt0)eYn%V zly{K-R0N+L{^{VZBwz`Ov3Dpaiqdoe57Aw0n{qDwPY7u9Zw&%O0biis3Yq>RWe6WI z5@#^&%S}TJ-+*2}zo#q%FbCxM4`7?xYI=a1$g*g2o66s<1RDMa!!L1wLjYZ3G`JPI z{smZyn&@eNb!e|9pzGJ8W-9}1f;?>``&t~~6W}Y@IF7|f32Eqn*1x`>EFSO*LSyJ& zb&Z<+fVqJEt2Z?k*|7D${+~OMaM2EmNE~R<(N3zVWJsnIfu26R_3#le0sX%oeZG)Z z=Ky_rFQ2FP+Ipx1(AbMgUv>jdU@1|TQXmWC3~1K51*-)dZ_grt?moblViB+dZ=p!w z5=|QHK*w%tUEmso3w8yx_wqVizXUJQK7}g*+IDS6G4z!p2OFTzPt-C!1&qPc z$hnqIyIMHVyYIQqg+Ivn2LsLi7Du-Q;0~G*?y7V&Tn~W;zDNAzV1S$^4(RyBoy3^| zi*P6l$Nxm;j0JS@0>|l|zKE6&K-RqB1K|`T!?X|9=n3@rVV3Ze-VlFn0aFnDgNPAu5l}1WlTzB^@=0dBb5N!_yY{_Mbe!c8=&uhx}3T<8G1s$3sS%~PaE6@ zdig=>m`6+}`zfIL=UOZ{&`qE1a=ELIjffQ@1sZzDF?L`3DTf8X0nFWc5cG+#2+e>Q zbji$cC;(kO`=T1}(WDd#U;|9&sMUB2Fi6z*m!4JRAHM^wy|GQECOgEm7vKd>a9hJg z$6!8}=`+P2iJ_B${{8|}kbwC(FCH)ibhT7C+Jn_RwK2;uO)pIQ8< z=Zn@V{}o^XtTF8{n!=;~0M8&J!Nsm#zPvzTvM}0S!Js1VPp=BjXRi5s<^e z^O;!uYXq8h*yF+-pvzy!4YMczGnEV2f|GphQ?4K)6`*rRETt3$T754_Uo7P=wHd%0 zynRbgCh^g!2HN-Q`fEO*-`8VqNRR{VTm%e)I6mfM1uhFE(8JR)fe8Xlzp1i?ixHYP zAK((IApb4&H;TRijeNRK?jq3n%UvUNbmT}+0an51*U#*{C7C&(n>Rk?Isn@Ls#0?( z=D#9tfM1|r4^|Y?`c)0I^!`K||9JpN82<1``;IukG+;Ssb=B$YvjBblcwTl6U<8)g zi+CLP)(ipX5XL-A28#oUlI_+P;HU{iNU(Oqf6VB2q(BV@T!I=S; zpf-^&&TQrU58xrDaqz`D*$`}jHs4pb*8%VaiQ?~}PE@1$A{XSN* zE5ICx#qMJcPG!skZUS`u_)nwdRTH{7Sa{I~LB&`35WZ{0jWK|lWLxXuEu}esJORW&K>= z@+LNd#39-5EZ;RIo@Bq_;;1{_EO@tlU?65v*ppY?U}H_59=L*NkM7}NPTxaT~? zK-C4%YeiPzlAnw1i0FHr0c?t9@A5M@0? z(WWhL1-bDmx~!jRA~DhnwyyWCvpnRNO%o%E@mTw+ z2msXnW8S7L&ExSF{FlDQ~iGnJ*o_Mr1&eMiq@p> zO8$9lE4Kco`8+E{qRu^x;9*Zj+w1G&;Zk3o!2rggtjC~8<3PzCYs-S()hL2zpIPBy zC)tO(z>f1mE@@OI+A$m4W;Rg(zbAU;5ct2gKTUp0j<1W0!713`tanRlTyp~JbmmP~ zs_Ti(`NS!#3=q7pwyt3QTezOtjHmKKZ|Goq4EXOM-C#tI_E*DXRCN+Tx;9wd#o+Ym zKImU6mMpD4hH%7LB`b0GA|Wb#ufX|mr%@G|9n77Y&rE-$tx?)e3x{C6?3jcTLz}Y- zfsNuS+F-TvqG3L~zIK#ZYgLQ1g{G$0bJh;xB|wHYaQ;y24cVP0fx$?y`vSI8H-Sh>3hGGijHXM{{`W^GcFpFUm1CZ}GfnZSXDlVu$cJ074roD?d}0CDZ-?pT zi#Cx9Hf!)=YPmyu4U2c&n0YgX2oN4)1B+EDQ|gLNDY%u_`{x#K_bHSeV$%esU)?=S zpGU*eIG%t-LQ2v*%ncmg=lYw6>dB+Q42>@RYpL2qBmenq>ZHWSTev1WDx;in29f@P zvce|wwKA`9cXaPC*pbT9KVq*CY#B{f4&;EXjhMrNN8tnY11LJ zuFSmprCa6eG24dX7kN}WuOYRm#3$vFPNU{K6zwR_5fWeQ>iExBhuz)?3VnFOpu(!> zaFx2kQHS>^A1eC0`@C@~USEl1)D_M&{gl@E<438Ri26*j9hv2&S~hmXx1$W2e_0~* zzJ?h8Jft>z!{CbGeep>H?PmsUFl)6SJzK-JdnPx znmn5!0an!+Bwat*suyTKn)P8?nu{-LCcypQ%ezN01*Kzm#THi{Qhr*^)bp~vgl*k( zn-&DdOKzud9e1B1I&^?O*vMLh#6$KQK%~c9=b;`B6M(n&%7G?i!Z_|vz=}ayblC;) zVep*XeH&yJ#QnJ%@s|-+%t#E+F;0+^yW7{4v~n25OA=uddQy-1IAYKWAJ)S7Yunop z?chZ}%NJOjpd!w^fp|+R{&pQwyR;gN*3gSx`ivHF@B@h2XEdgQO2KE=DW}PDYBR8T zO^pf87~V&IVUZ#leyVpf_%+kYYh7BTl(o#?mWZmUU_11%d7tYqlUAz1-@zYW2uCIQ zhU(#uE#jUG6!n>@BzZ@O?Ydg0FOz?#fu}T{yr~mWf=XP`FV~E)A>HaRt;KkqXA>M& z>WWMUOIqQwEv&nlBH5OI7@@MLRJ%x{e)#xsPSOzXO?{L(BY4>!r$z3R!#}XOx^dit ztJDn;>eZ{-TI!5wRTbI6Wu&jl(`q9!c?nPTZ-Tge+Hn=(qWZwqFo!O&K)${CkJ-n8ku- zhq5g~V5$nsrhkh$&nmqA%M?dom_p}(Js<6xRT!?5fCeL$9j{oMl5@%}+i3d+IZgdr z12x8I9*T(XT>4}p^o0vC_cZAeZPQi~Jpv3$Itq2tK33(P%L*TdA=qMo9m0?mjU#^eCk9zX%emJ_nY(R`}q?2aK_?U4;sbxK3rEn-Rr8v1+zhcxMg-Hl;N|zeuF^YCl zI|Zov%H`C~7WwbMc=rWSHaDn7`9=QxCK=>-L<3J)^?4Qy#EDDU%B$VlgvUwZ8vQG( zb00XTe=_I@J83~Z`Dr<>7wD2#F(pg?wXL~dfcwAA2Hb?I8DM@^t z3Ur?J=l%-mfvDSihr)-&87H7bq$}tQ8~(-qQa3NYHHE==guFPEkg5DVRz$kYU;3DtekfqgYChx9>z=k>Sj&0KG`KNvPO&qOPab_yI`GB&Ftj_9K__G+5s>7kr&fVDtWf&dCtbHp%s@PW)j zRn88}=PxcxNixGO$kyt?=Yl!C2;vDcm>|W1wP;ZJ z&ZCkMj2E$vmhq4a!DW4@B307&q#}8gAM7%e>cUS*ww^QmH0uEETP*$ErsMIuRQ9Al zCZ@3+BGU2zgXMGkU7rRg4@05hU9}d~P{XTVavLdyc{Rzka@O%n@5RqV0^Sxx%rGv; zGs$BMLnq#IwXzvS#7EOiWisTv21^SDD{s0AF{4ymzOFn))Y!UkEbwt>2Ao17!gjwWd3 z&<~)GGh4Trp2udkD6d*&BS%h@8C3AhcE?JOcm>`;(?U^eC}@He2ld$zO`KAy9!kO( zJL(^MuzH~;w||bUced2}DI=GXx`z6~#!x^syp^uKgO)#^|C2ZlxPtl$Q7DCA!j(U7 zGn6uZHtqcE(Okc)6AQ4)L9l3bTkVh9hPu3vSzp;8(j{9SC+fWbOPI z6Zqv_(Cu#e=FT=<`uNDCzsuVz2_@g*q0K1R+U*$_QD(wOG;|ol>r#0+Q`vpY$by4v z9mPazPxYenxOcCmDuq^d&^Oh%nsJ#RpEa~s2+K+*b%^|Li+P;NLyxViPh+q&0*d&V zu$?PfPjBq@e;#RI%Bo5f((@p>3;Rj)-`L}nW3*4O4K1uy^zal0kc2gwf?!C-gxk}? zdO;b;tvEco=m^DKNu~)kEmAi(JN%a(C6%a|sC-vNaiedyyeMBP?=>UNT*RK3UKj1e z4{1Y;o;B*MHto9uRQzK@Xe4QbWYQv#2+v-t_Uq`*NpHO?3w)DFdufTU6RFLE5qu9S z=!C$&lY_gpFE%xYa>o2yRsR*8AJ0)Y%d$CJIZM<#U_EA$GN(4-qw%<_c%)Gewv zC}A{WJFq`SwDL^*HV@Tu_Y-*-A-)8u;oyJ?`t&{RK32aeM4PN2t^_bCs; zlUnoPhe?ccUZm_VxR(NoU@nr0e;gJx;ZZo0)Kx{K9~upOv*ALvdgT#lGz>`r^CfV27qIOUm?@ z4L0{4ta|6{BGhJ4Xf98z-Ph>Ge=Au5JUu2XVQAtVyLNiVMS_X#4T_&zKN z8pyE;*0y0VMt{G(Yd5XQg-rJ7fvo;wwZ-IhQu(%cbi}R|{wm^K-yKyO)}Yi!>-29w z&++^gy?8|oGpgWGI`ZD8i!Z*>|Ab!m^Bc|O^c?eZ`V8+8753mIf=~Se5k#-TW3k70 z>g02?6XRAt`H-dmga9>i;Co`b{U=dta*I)rkA%&%WJXStdmF!M?h5M+aE_el4Rd4k zPgso1B9*ceTX#Z@25I7dGz}jmeG9+TJKi}^ZNCv#)eYa=Rf;nxG5mITKyL-lCfdv* zRl*BJw(8EI+V~wUq8xoCWgU7J*Obd+j-xUC#Cf+UV|po=c%5@Evy8UysaSJHoqcKb zU!#q3ZC;)D9~PWvx>EV_^>eh|4G&Bl)}r- z*`jQ3mkLZ4L*8m<#|M!J3&x5cq4{$EJGF*VuE{J`;H6O084{O`hwmGsN`TNUwEfz6 zTKTEy4eaH#3Wkw2@)oXGcqa%n({g7m9xfk8U>x{zk!`93o0b?HeGZ?gp^L9s8}9q| zlxg$aZ6~^^?%JQGK_@C=KgejWFLx+MV32ysm*w>vB)+Sx(rl{f1om)0dPWWY%1eg3 zXi0e^Y8uRKNReqY5(wcZKl#Po$;2vN1^7f#*p4%Ha)7~>oHHrf)cnhH;>mdQtK1L| zEc0){lB>sCQ*6Br?A5_tJNL}zC=lO=t7M$W68}q$N6Y_tl2_x3>oEFn7xq>h)_t3f zesggRZ(9!{g;`Xg!@1~P8;ZXnnGE(8CKW`JK?jsZHV=|xxm^&?ark94hH<&hPNEfu zF0=8Kd3eFGU%EYhQ0GR)dDNv10LP zlu_*ta3MA^2xDm91Zy8lW$BeQ#PtdFPAIWt+xb|0@Ip3NNIpVk+_?_lby+Vu|9&}4 zfZKvCMn^zx)eHZ(bUDI}EgErr{121IX~Za=-+(-14qXdZtel?MfxD=qa`kL^EmhBgzy++?MA z4bs%!V>71$Fe=o^aD%16e7izrz2OVA2QxjxYUb@_rsS7(sXh%QD+u4+edd~ppkM>a zgVY;s7Ytovw4Q`1SrEb|i(|C7-d1?7Q|?7DC46^TjwQm>Tj$h#TXt^Bl(Tq6_Y zh44i^zGq=>3u873|H$NtLQu7uIO3g%4&8uSOGZCut+uc!$?ZjpIAmQq{^r*a$q$=2 zC~PgPsQ~TOvzX;)l+JZfE!w0*AA08(Ugtc8qNN%%$#^)~J4|LifymRM_KUr-dh}wkm zeAey@I;V7W)O@G;R1nZ`yw~l1{yWCIZztHMq4W>D=3z1`T0gC)JtCA9l5JBOzHpSw z4lkp7bi2=*V?1`f_q&okN8Q|;$Ih}S?<(7o&$V%6Jsk?}u-2j8k}bmGgS3^_wQCpL zb??Ocbu7@R>6V@E_Uosy1y9Cx+U|j_h>>x8?jTh5tvp}v`XD}=b3`bak%gW2*@3V> z;+RqXkpEH465o}VNTZU5^o4LySUNY7d`zv^rt&a>^Azo-Ex zc)IN)FSBvDJW8j0tmT+|cNjD;!F7r6mn17q`lP*wdI;BW5*H%c(B(z+o^qr6i2%E8 zj{dP+#Rq1s)P8fj5Tld${YaUdnkG*Q70WVr72>pYQX1y3-d0k?+vv*H!^=rQF|R`= zJ?UqM_@PU8R7_|KBofuct<;w>@zxMif@+mX*FO_WLH8-Tt=jtDm^!6Z89v{quLG@< zkLHU$X4}oFWVI^zWLXVUFFWGux)V)e;ox?2QQS>CI^e7;Q(=$8g;#r}l9q7DJ24cfa zc)RQ%#VDV)p~&r=oI%D*ZaRXaSAe=#1ZD}3TvSBRAJpa8+fGkpDB29@SDV?WV`3S@ zOg=S(eqEQQ!H#WT<9wqdrWc?lt+pEsf_I+o8Z+?{s5t5O8PN2n@U-%O zi5an)4wb-h}JetQhUuO=IGI7s8_wUg!e>f>Hi-b>=U zWQEGqUi9WbbC}U8uJaA+Z(|=T%Og5z=r?=Ekv~{Z5OYZ8dzC|f+J$D#*N1#R$(JWM zrIx0_lj^TY4t&9V{|U9LPOK4IKb7&sE2d(Q`b={Cfm7S>&Xo+Iv_JZWQK9<6Q)=uo zwXCFaDWS7w_4BaZTH^~!ROK-H)8b9Ff(OD8YqThQw{tHox^oN`bOk@ah|!O=c$XRj z-CEBzM^gJ7ArWZI?6h!T#l#60?%qN#HAAtC20k~Cbd`h*>vd;UCxVyjNy>W@kU7UPqx}} zc{7qd`=v$}d-@l*U>5@e(j3W9VKxRr5q40k;4#f)cP85x)g2d+I4P=?>T&4HzA0l; z+X-|^O5ttJ!(>f8I<-r5AB~SOzJ#5&U%YayIAeGvN(Aa*#tWL+Buf^D);OvIoF=QO zL_a3JeNft<7JX6nHIGWMDU_ScoKMnDOR7i8nvCw@lG=cO;T8T2f&6kd+^3#7H%!0b z)7$S%MdIsMjYr&k4)ic;w=DsqY{PWeUjL!O!AZ&8b&ww+(xOEbF93 zbq-Kh^xO`Jim#_!+J}mizZFB6W*%cvh2bSG;&_TTvTwV%w|6=B`sbxm#78>)=nF}c zJV=rTE=pF9`n^+p6e+CKVR4>vwR^JPuQRJOY~{@*VSd;w5R}#=@(eS?=G-uXWp~uG ztD8>g2@}6g$e;XL)rbF!cd_(@qX(Q8`+P4$Oq51Br?B4Tp8 zL3jH*R(%IEvh=p5YV{`SeWH5F2~%ZeVm~$5O4+1cj)V>xSQtE{aga3w%M9O=Ve1=F z$l&_CJKbbL#4ivFm?O&{?{%LIf4^T4muwdMm1^-W<7G|>o6UiI7I;UYb1w3TudTAF z1-`0Q*@2fk!z1XC6)YyU3B5cnYm86RPBJ>3lr5w$v?>oOO8{bDgB1kv_p_!p= zqAz??DVJjsxAl5qcUkljgwlb4Isd6pNKZ_%uV5u?kf!ae@nLQu_w_7U-uebeTSC`T z%6E-N0t}fgJ}ETCrs}qMeHlaEmWBql>SnU*;+S80vJfG0(5C%)tQD4+ zl#0Nb-LlmSrtXJ^Sc{Vw>!=8mpQ|tF7ptEt7QJ=^`8yDKAm`}_5{ z+M)wO!lhE*_fAMa=V-jD37&@g(!vDCU;KiuCf!*b9J25pG>t1LbTxx3LQJ)eq^&~N}B-G0}3{AQA$XEYeR){j;vEh(qoo)q`eja z1}UQBgLz+aBM#%EZad<5>p1J`JSGu(unXGGiaZu`!#P(q=KfsrLm9H|meaMfP#P#h z_Iqv2x<1H-QOiKwz{H9t3k>;@(GhGCL{|=A7A7y=$gT&Cj9E3x#kXpOcv8K61-iJ_^S!fn8#vVHHliY7Z@Cu4>L? z%$yqL#hKtsfz`JpX$h`|c?WQ>1cNayt_Ae~BA?9_RtPW4Q$P_%=RW(ueMdG_!GueXD z36|J}|5UF|-bGxIh9F#*|04J5w<9CM5YwN$*w|zdFL=#=hq@Gm5&U4DZk>K0Jh3io zc#V#&s@eWem70X3Wg{i5%Nz^e(_Fh_e9#00)ktv7X(RqvfxxfK&Bax~K1C2aO=fum zrhwnRL0Q1t3YA-HKcRCOPpxY%eU5)Krv-^Oz62*6q&SRF4goZ@+uqsV!4*m<&p>&b zpD8V9>!V{9LC<@vytJ3~d%pU3E9XMlg2RaWBb+XhYAlQ9(gkSK;*^8mf_B6jxO78TIW0DycB`r!b6^ijHfPn?pO4m&Fuu!nr9ng6*P48 z9%%B3mhda_a;;6>{3aN=eBRH$U~NiWm`D&g+l*0vDzl<0e=yym!4jgoi;+Drbh_bX z{QNrCxGST?`8R>WQ|F?oQe#n1d9Bs@NW1Ltkh|~EbS5%mqj00dUc;3MK>roj4*>y(}#2Xw(>5Y2aOaeIyRH-p9Nw*<>(-r5T?$uH?# z>a(SpGD`SzjoeG(BkrkLaU?{d+akZr9jpGq8+A5ZVV(tPk08v{JdF-~GbJ{w32Lh& z`F1y6WMH1(p(=?`U5h08Pp8rE)v>gbR;1NB$?03F?GM3QqZ!?+$ah(DqUmKvQT)QM zd2xJW1Re zo@7DRcs}xnuDpu1e||`9!O@*JxN$t7e+$LA;501feAO2CaAAley;k`$+;qIy_?Ipw z+LP93Pyb}{nXzIMoA=-A2ZN0LV%iP45kt`(ta@W{5r`%~Sdw!UrHKrL;(OPF_V;WN zP#EQnDizeIu3BpF7_GHF0&BXA-u%e7S@(pCF`R=zYcDERYZk?#a(fIi9rt$iY!B4SSQn z5f9gR^AXZWFObzAH#x60Xe0fw^TyCw{C)L9xY2_eeLdp^X$8)SEBX0|Vq+;ip$r9Q zKMd4@+6-}`_hrje?UqRKxg8eiFJqw5C_HwIPy%}X?oP}hXjxGFgx4h0>MaWR^-QN+ zfvH^?>e$=Pti3TpL>sGUk*_TFg)q-~Rs43qQsrP#CSlUW9zi4^6BNqjO#Vays=U?6lZTV}<>ooBA61I}&mL>nz|ZH!kSG;hll#F}Hx~+z1H;c^HKDu%%`qFEPzRQH|d|0%g%nM(G_t zIY*WLzQ0K4@FMw3R*uD3HYuw1Bmd3==gpu6F^KsAPQJe(U&Q#TlbmH)*#5+W(P-;+m{3zO_l|VSOvN(@ET{4 z^3Wwx334=0q^P28gI>HK2*wU>|M;!A*O?!yj#oq(J6wYd6f^$$rU`hxcq8t=Ifdp% zzvKC}MMfRBvrkvQ$!_Wp(ol*;mex`TTzsDV&=PpHycOs}V;oZTHHyib4~jtxw{ZC- zn#1sHh{k&};IPb`vwb+E>M7N*sH*>e_-y@otaeii$Mwwl%m zrs8^?Vf@wbp@5I7jrfj!?p_&BiDE(nHzW0+2G)H{EuNHbqg&+dL-rwB=S^Uec9}Ba z{U94v){DIP>y1Wk6Izl?BmM+86ECB>&iXY%Cd!a?jrHu?Gl_Q8SkT3sUcQsrFAIm< z#h|rc%8(~7F6PuMy@3;pp}0j4v)X^?O9V$p2{QfBy={*q=e3%zz68BGZ|Ax~Rw^DB z=FKA>e>*}+@)SnA(0jG#8~I&$Nbq%hB?Uz^XDv!88n0?45EI(G3MDOcszMlvQvStt zzJ}Gm{St5INyeiD6K!t$ohwU15+@|biGzn_EhW(}@j7J5Mk}-=oQhH>26dadS~UYR zS>&&2Aj(8BBtK;YBD;cU;G016VoRDY9HDphU#k#f7eyk}N2d*@%JmeyJs0`cKzJkb zl=gOF$q@f<;VWZRORq-LP5Oa6KVd`p`74bWHIbvK`#2R-E8ccRq(PLi-}!6~{0?V` zG>)`3=ihAjCYXCk3>+Bd(hqDpbGFE?+&;cUR)z~u7+bzb-V)hF)|&%aNLFA`J&IMM#>>-Mmbyz><64jA{DD^lnl`=% zMcE#@-TCNk+Z$92in_d6(&DAPn!kV2HdiMb$1B{EMp=_rw^7PjU!+`%Y5MRc)&Z~d z@V-<)l&Tw6Tea~`uolLtY_DxlpnoYF+94+>rP!!t6~~YezdTf*6Qr&jstZzVpD(^w zG}eHe@i2A-f$)cP)lK>%(iB+>f_^%|pt`&Y4*q8EQ#vALaq=q8y+OZ$^n~?hRfu@* zkF20B_gXazwo1~z2zd@WJCwzQw8s?bFxhE|V;8ob?3)#oQm&K_+0=e+wl=TvU*d$< z_TOwqkXZC1`I=@*Oq0NrO_xZ?#D#8({)|^{`0Wmsc9RVaZ4Kw5ki{X;9tLp3`-L3o z%v1>|u*$)a+vmdn32W6650`B#JKH)2yJ0=!Nb)Wbe{?@~Y@$H!@`Z2N7+jcJ*ilY0 zO2aTeU+Tx%Iqlg(${tkVO)m0w*U+UbRJ-*!B(-mlYK_o^8!cGeCD76?WBKI^F56kA z-~P{=$nobe{kGrq-@ntQ$TuVzE$>{`2-rVR ztXv7cYwToB;^r4B6S?A?$hH@{-X@eQ2v?&?$z-zsfXjEOdcwIz*8ARpMB-LWWWh*CMiC_Z+5<8Hf5c_BmgLrV4lHhn*yc$$#H z>g74|MaK^vPWOusjSuqi{?-g}TxPq3_3>AjPNsYDa4Nws=33_FiR4`MNY3ohP=ECg zdYp#9&4X^uczm-efBlf^)tSFPA9;k632>+w+1LLL;ng?phxS+O9>fnhY*g_fKNYOn zX8ACfEIHtV>L=xJBhF^a?o_va;f+W~jTJ-;m2!62)(n=}tdu?c*6blVJl}%Mwy0N1 zwU-};dZ2MCT4cl=(a35;Z4_fl&HiQs9-_;kUPw>J8$p>RaV(V*|H0*b-ZKNj;MA z9iP!AuS)pMjjkk0kTWf!BKHtCEJx8_kFhFBwv}CoNr;69g{<96JmLL*E`uQmCOzf4 zUZ&m0M17df9E^bn(_Mjk&`qDO`aA8^XG5>QWH(V1Qw=C|pO>n!VBHCN^7U%ZAYO*m zeegc@Ks%)4eU7q+l=Mpn)gfc<>NMX@ibULkZvoQHtjipdxWDFYolp#XzgJz_BV|^= z%BS(-fn%0K(^|bk>vEClvGOS9J|IBuvaS(N{qZ5j>lAPngHldy^z2Y6tMP8O$`-^l zmojn}I--av-;0j&7VYj*6r!gpjso8kv+Sm0@5{%VNop+_ai+bO_Kc>PUlUREK)Lz1 zvzm8crc+>SD?xZrmMXbLQC@{Eto_) z99b z;E%hRKjTBOV2jB*xR80F{Z;P6*h?luz7DJOg^m3>ct#SUpn z{(;7BJE&%hua$e$vyEBCru4Ngh^+Jj!o-&)I_u`F1DpV5pIrSAMd&8 zuC+RtkDNo^enmztHU-31x~tCF$O;V$o2W7m!q450v}x`RuA^XP-WmQ*@i118{(|8p z0y6Cp1$lxoPLDR`?_M<)GHSTB|G3EDY`-$p+RZjQ8idpS;-Q$7`tw&yd>hH>9}QVZ z^}?)#{)kiwUYre!9p(kLmY8x0b-$$mRHs7e!4-+;hQw@7_48LzFS{Tu-nvIx zG^;pAZ@%yq(8QdHl5FH@jYraKz6SV-JGoZM%z1ZUWFApd_|&|wNgF;Y3|R+SfYp)LCVW~j=YtdW|aRxJ?o)dq=Ek?Daeg= zuvqB>hZ%cS%_wh|kfsu&7ylanSfei(VZ!+D1h&TS>Q_gwIuvv?zDJhutB4Qfu{9Bx zOqzj9koM1Sgx2~)%fD_r2C8mZHy~#AC=L^DDjif}tDiI8JghZ9HgXVR6^v>(q-+B& zYX#?&^-h#4gta9Wv6+Vj3(s{1`;jq6x+HhpX&q%PF2l&|qK&=FR&%owyz;vX zLtaxQ3=QW5zoG5(oK~zg+U5G_?Pc}<88WB$VO-`lY{OMaCHkko|1wUm91%}v6rO|M z3B-$FjKe(rw4uOR^T?ar!%cHDr}=!Okx1wT5w0Zr0xo zu)5*#Ni+ch@5;{&Z3P}xhmZwUjIAo*-~ka}J81JdQ#LFfudX$4!?5zZ6JS=QXkZ9u z5#`}%xjj9lU#8lMWq%*ldzd-Jh4`meH8 z9bMB#S^uX{DBBZt=}+D6qNmGx+5r_Te|TTU4q1GgZr~cuVMj zlN+0wA-&zpx=VYHGj)*tF&Jv~=DffN$xgjzs?Gt3P?OlsOg7IYLKb|hQULFjJa6}E{0|3B z4i;82E(KY*H4Vr7_!}ymBKSL#G=~T@YgH0=pM;#@P=aAaEVQ+UYuw2yO6=GYcPuWq zi9wXGc&UXYCv)?LhWH4zz^%pVKFUK$V2Ul(hI8A3+iI{ovq4E#H@sqh`1}s}9{2kL ze)hPQT>3tTGSiEP_4QT-!ZzLYT+88|S(bcoDl+^z@0S;|frf10ZX=Ct`k7H(+t&%? zbS#81q4J7}ZsdIW#C<8Mf3U?kOFITV_@hLpo)x(QKB3Lc`T-2&p_mR%`_nT*1c@#= zv<3u*dyUdBW{<+%7qt`~>EWa|*@v=sK@gy5%@EYb2yLxfJt8d<4a0dDDf)zia@ zbXoX@_m)J=;u3a;P3f%<$z)Z*fED0c-b^8GMPh_vy~o>3H`J#j zF7dq=Q(`~H3S2z8zQp)roYHr1ELWNkIcZ%~jB%59$gx zjd);1q|t|zyKq$}O{xcmu}H-4ah|e~(KGz=?zUGUb;&-&>naC^H@;cZ?W`FLd`?YZ z8!k;yFTI4|4*P2lU=LF);E?KP?XnqT`oB)(!mODGAT*~Mm8MisrB>3L(}MEEvRc4` zGv{rX)^6uW2buZax}zC?G;Mz+e}=#Rz=%ZrdELZC4BofII!m`!n_ap~;fER_4Oan6 ze;BS%r;wC^N94(9K$f4e-5ue8e(7(Wc#*_M)Duyk)_;GBRj$cd4*&T(Y|Lo!jQeZ>{5QTP}CbZ+!2&a2NN%)AQd zgEVfrS1xXQduD4&6gFnSi|q1mrw0eLG#=7;#Oh}A9I;Wmaax?C@j@0Bx^5}K_puy*6a zvWI`$7_@Uw*{T1zkh}8BGhRz;Sa93;tgDXnZxJ-D`P`X33g6xuQ6PF$nVUnFeb6E; zYg|Z!dubP27yY_t|2zAeelRg$v$g5bXLzG+@iBGRr**6zo}T)~rueUBw%on3aeB^iBY)YwbyeHYmS@V( zNj|Z;(DrAZocVFBN;Olf#rNq@cV_gO8Xfx#i>ST#l|#WDvJae}XV#LgDKAtGDH>5b z{Kqx@&!v?vwd08jcXuAo=)V3^oman_+4}Xs=?NPyz5n}@M@rxMPsik}X$N`-ygfH& z=83#9dzwCz&?T%()74vJ&i>f!P`(_E{t6k~yz=5pvpYrHuHCs_>6@L7&90i1dbD)# z7wK0H+^ll_T%J9FyB7^S_;Z!LV|PVQsx$W7?Z?sQyid1!r^K+d!cE)m zPws!de8+V;&L158_dn6ok7j;8^8D4FxrXMr_E&Jsksj|(Sohe}bH&Ex`($0`uR4_M zp5?8ErK=S?QTUa?TQ>~t@=~qGM(^nqoHr%>gU!+VTJ-Mz<;*PZ)p!n%MB zpOttgKIZK?$%Xo?_@Q$5;$Q#oRNRMcn|zZzC(E@OAs+Q0t{qD0lKTa*s z^}V2{pI&u1xLC74-{?MMZ@(rjZ?)dre#qNDcK^Ccr{Mj+U9aDy-H1H9s%Av@nVxO+ zyXUKXQnTT$PxI|)STaxN+Xmk(kPm(C`VTLXIBLaf zVQV^uZwl=B=c%(fI(&KbyFc^Y>l3uB^1wy~_CHgs!-30@i^Fox`{q#h7SYASzaPD0 z^p6iJcRSbgQsU3^w)gsW=FN#o#Ygv@UN=6ubNaCr&0fm?%JIpIswA!%*L^`sw()Ou zN_*#Ex!@NY?LQFPDtgnZYhN`Vw5IQ<`COCjzZ^K8U}-x&x#>y!pj#sF>Kn))85NIar&sp zyWtHt9ayusX5`%SaX(#ab#dO^=|Szsq?KNFDt+XJ$jysR-x{<#DA&q1mjh~yX!q>C zq@*c9!CQAHRA0QWdhm|fU2Z2|%^vvrs5MDTZdELs*sgZo5*bHtJrQtW!rgZr5FTH6XUszWN^qRvr`?@$>G)%Ta+33ncb9 zvu0J5Tk(-A`aeFiUGJE?&GWRnSacTwQdhQu2Y)NwaF4OI&|0 zeN@t)RZB)J$$qz0lPcwl9{uvQgSA&5ely2k69X>X+g{H40BU);L?m z0eL-F`-g=u|Jftu!Js@- zmd01#aDQ*wveHlP>ae6jmBbuJO7A>gvE)a8rJNjfdC`@A_c!+%(0JL00pBj!xv^2H z(#Lyj|76qpP3@nq`q}i+n`>;1+@02|?dAIul5a#dosnaI@6#_fcrNV5!rr%=g>8Jk z$?A21TPh6LkYmhOADllJzVJd)$@53IhF>08GjQ6|eJa1Vzv!dSeo1}2M_9)RN#{?z z*L`iVKaO?2^6e)XuV)FW_xs{u6_Uz?OrMc|)(b;sAHE&Y>C2uG&s@yE=gm9K7hS)b ze(XxlleOoUT(c-(S&Zb86fzY`gr zvr<5f*|jUBb}e(KTk|pn8^8Q$;_3D4#~jSkGdz7zMw=JH;@)prW8%E(4dP3E^yS#H zH_xTb9FSbA`_pwpJKR3rBlce0y_$Jmi%Lug3Gey+q)YQ^25!y!b=caVeOV^e8J9Ss z_oPad%N1|jyZ8JP;m3}j+1;u^rXjSttJ99#IA6J3W_Yj(10YL-j&7Sx^aze6u%D6*6OX=9@mPe%kcJgY|YaS$XF_ zZKljG`fRDQi2=z~GEP4`qv5?n{c3*n{g&7<6Kk#vEYYg>$~M#1A4tftwOhlA70+DS zb-H%HxE6ChD3uzWcWFY7Ass@84Jg&5NWBG-HAhCq<$iC?_V+>x=Z^k${5y5WSG%6J zq2Jc;3-9hwKfdajrP~(Q{IbV|svQc370g&xXv#|&lPf-qy54c%;0C!4{&ecdks7&Q z=(K(LPfz3?ziM;VKFiy`I_AT(UtC;tA?mr(vzDCrti;bt-{0PNNa{OhRt(GW!HFK@ zI@MlQYU{x*U5c&UQ*Y?2LpJ{O^NP~R&(&*PXzJP?MN2LpHS)hTug{1WaDPhN-YT2Y ziZJUQj^(dM&yq?zPNe5 z5kHJBQRLTqLyq-*Sm3G4gL)QvvTWP-CC41v`}i~Ewyik$MzvkrHys(Awa=n~UvJzK zyQfHjMsJ26Z?$FJAH#dRGx@#mrpET#QX{#2oi))JNmn+U?(SSku zM`f?OclE8f$h%oaz7a7ndhoVYy9&jxJRMSU_Q~J7)LI+3Ao433{O0_Q-s_j$OZl~T?Ay(I4NeU^H6>)k z>;CkK7xzABR-i!3%OfiGUGkr(y$xztEZ#mfe8>}9UU;h3Ya3b&*%Gj_R)sgF-TQCx za+g*-x#PDg%NHM?(dMrXgXW)H6q~WSV(^+4KNn3}Tr{fwzVK=3la9}+98m52x>DO$ z7r(hRCHK6RL;FqMfA!Xz+j2};k?_})KaY>CKJ3{cJr)*d8{8)Asf|$uA4CnE(Wg}0 z*g}xlMrv(0RB9E0RDdV0$n4W!g{b%30@aco#2KV0G)by=N z@0QH@?Y7;wK07?)w}YY4zZdURqs%AQ_vH=z@28iq2Tf~H>cW8<(9j@WGPMBNcLOi9a%T^n{CDw~xxr|8wv<06e@N02i+dz{l$YaPm3< zyu4ljH?JGO&+7+p^g04Oy`BJ9uPeaU>kDx9Is?4D-T-&6JHX%T4{-Q81U$YT0hg~! zz~}1|aQZq0yuMxmx362k@9P(E{5l3azn%ftuWP{f>l<+XItRSJ-U0Wodyv`xbMQI< zJiHzN7q1Jz$Lj-d@;U*$yj}n|uN%P6>j(Vn=>O>m@br2DT)nOUU#~B~+3O7O_Id-{ zz3u>iuRp-y>k#nxdIVg)E&-pfPr&Kx6!7|b1>C-F0l%+b!13!C@ceoPT)(aX->+}L z`Rg3;{(1-8zwSY1|Ifke0Pyg709?E-03WXpz{%?b@bY>A+`Mi8Kd&Fa(d!8C^m+nZ zy{-UXuP?ya>kRPrdIQ|O?f`$UKfvMZ5b*eV1YEu@0iUl=!0G?nDfq|h>lJYOx&{2c zegVg?W5Dz48F2l&27JH10q3uC!29bRaR0gonf*TpuLHot>j7}_x&VBk)AIx&(Z_ zJ^`n%Q^4!%6>$5y1^@c}fBFR+zm5UVuV=vZ>l*O=`Uaf8&H?YQcfkGY9%T0a9J~$y z53dKn#p?p_@%jLqyiNcwuNT10>jv=i`T-ohjsQ=uC&1O~3h?#%0-U|h0B^51z}@Q( z@b~%y9KH?#kFQ6-lX0)`UM=njsefFXTbIA8u0!42Asdn z0q?JO!2RnUWcL3Yybb^luLr=z>jLoc`T(4~P5>{j7r@Qy2JrLx0UW)K08g(cz}4#t z@b&rvoW0HfZ?8AN-RlnU_xb}Iz77G8uSdY;>k{z!`UIT5P64m4SHSJ-7V!J}1suPQ z0ne{z!1e1I@csG*oWITi@2_{j{p%iN_WvBb4ge3Y2f)Sa0`T$r0Gzx|057i>z|H^L z4fw~;>j!Z3Is!bso&Z;`E5O(53vl*21H8T70C%rDz~Ac+aQHd|JiZlbkRItDzyo&nddYryyG8*u(Q2fV-D0r#(aklFup@HzlIydD4- zuM5D(>jQA|Isv@AUH~_*8^F)&2XOQ{0zAE*09UUoz}M>waP~R_yuIE4cdt9Z-|G)> z_&NkUz8(RWuS>w^>l1MLIt9GGUIDkSTkx;n|EFKT@#`4y{CWmlzpeq_uW!Kl>m2a@ zdI#LU?m=e%&%x^e@bG#7T)Zv-AFmI<$?F90@_GT>ylwzLuOGnC>j?1ldIDU%t^i-J zFTmOB4Dj}P1Khpt0DrGPz~So<@c4QJT)r*=pRZ5A>FX5m`g#T2zHR}(uV29N>lpC- zdIns-t^wb#Z@~HM9Ps{n2i(8zL1zEY!Rr9<@Ol7Tyejd!fdI8+LZU8^8 zAHdP;2=MfJ0$jbW0AH^!z}f2z@b-EG+`aAqf3H8l;p-6a_<96fzAgcuuTQ|~>lE<% zdIj9RZUMiqU%>I}81Vdh23)_c0pG80!1?PO@cw!S+`sNYX8+H@>j3cZdH`I!E&v~| z55URm1n}~D0o=T906(uEz|rdn@br2DT)nOUU#~B~+3O7O_Id-{z3u>iuRp-y>k#nx zdIVg)E&-pfPr&Kx6!7|b1>C-F0l%+b!13!C@ceoPT)(aX->+}*uk-(>bHMxS9dQ4; z2bujp2d@Lb!|MTX@wxzfygmRYuM@z_>jiN0x&i#WegH?WBf!(^32^ng0(`x`0B5f= zz}xE$aQC_c{Js7Fhp$7xlJYOx&{2cegVg?W5Dz48F2l& z27JH10q3uC!29bRaR0gonf*TpuLHot>j7}_x&VBk9Do`U0H2&H!((H^ANN4)FK-1022%0gtaoz~$=_@cH@#oW4#0udi3Y z?duls`}zeOzm5UVuV=vZ>l*O=`Uaf8&H?YQcfkGY9%T0a9J~$y53dKn#p?p_@%jLq zyiNcwuNT10>jv=i`T-ohjsQ=uC&1O~3h?#%0-U|h0B^51z}@Q(@b~%y9KH?#kFQ6- zlX0)`UM=njsefFXTbIA8u0!42Asdn0q?JO!2RnUWcL3Y zybb^luLr=z>jLoc`T(4~P5>{j7r@Qy2JrLx0UW)K08g(cz}4#t@b&rvoW0HfZ?8AN z-RlnU_xb}Iz77G8uSdY;>k{z!`UIT5P64m4SHSJ-7V!J}1suPQ0ne{z!1e1I@csG* zoWITi@2_{j{p%iN_WvBb4ge3Y2f)Sa0`T$r0Gzx|057i>z|HFh@bmfs9KDVJPp>Dy z)$0oI_4)#wz0Lq{uQ$Nm>kjbu`U4!k4grs^N5JLl67c!@1f0H30k5xD!0qc6@ca4& z9KVhM&#z~|_3Ik&{rU!+zs>>guXn)x>mFqG{~Wvy01vMRz{TqV@bUTpoV-o|FRvHC z&Fcp6^ZEfCy^a7+uP4CO>k9Do`U0H2&H!((H^ANN4)FK-1022%0gtaoz~$=_@cH@# zoW4#0udi3Y?duls`}zeOzm5UVuV=vZ>l*O=`Uaf8&H?YQcfkGY9%T0a9J~$y53dKn z#p?p_@%jLqyiNcwuNT10>jv=i`T-ohjsQ=uC&1O~3h?#%0-U|h0B^51z}@Q(@b~%y z9KH?#kFQ6-lX0)`UM=njsefFXTbIA8u0!42Asdn0q?JO z!2RnUWcL3Yybb^luLr=z>jLoc`T(4~P5>{j7r@Qy2JrLx0UW)K08g(cz}4#t@b&rv zoW0HfZ?8AN-RlnU_xb}Iz77G8uSdY;>k{z!`UIT5P64m4SHSJ-7V!J}1suPQ0ne{z z!1e1I@csG*oWITi@2_{j{p%iN_WvBb4ge3Y2f)Sa0`T$r0Gzx|057i>z|HFh@bmfs z9KDVJPp>Dy)$0oI_4)#wz0Lq{uQ$Nm>kjbu`U4!k4grs^N5JLl67c!@1f0H30k5xD z!0qc6@ca4&9KVhM&#z~|_3Ik&{rU!+zs>>guXn)x>mFqG{~Wvy01vMRz{TqV@bUTp zoV-o|FRvHC&Fcp6^ZEfCy^a7+uP4CO>k9Do`U0H2&H!((H^ANN4)FK-1022%0gtao zz~$=_@cH@#oW4#0udi3Y?duls`}zeOzm5UVuV=vZ>l*O=`Uaf8&H?YQcfkGY9%T0a z9J~$y53dKn#p?p_@%jLqyiNcwuNT10>jv=i`T-ohjsQ=uC&1O~3h?#%0-U|h0B^51 zz}@Q(@b~%y9KH?#kFQ6-lX0)`UM=njsefFXTbIA8u0!4 z2Asdn0q?JO!2RnUWcL3Yybb^luLr=z>jLoc`T(4~P5>{j7r@Qy2JrLx0UW)K08g(c zz}4#t@b&rvoW0HfZ?8AN-RlnU_xb}Iz77G8uSdY;>k{z!`UIT5P64m4SHSJ-7V!J} z1suPQ0ne{z!1e1I@csG*oWITi@2_{j{p%iN_WvBb4ge3Y2f)Sa0`T$r0Gzx|057i> zz|HFh@bmfs9KDVJPp>Dy)$0oI_4)#wz0Lq{uQ$Nm>kjbu`U4!k4grs^N5JLl67c!@ z1f0H30k5xD!0qc6@ca4&9KVhM&#z~|_3Ik&{rU!+zs>>guXn)x>mFqG{~Wvy01vMR zz{TqV@bUTpoV-o|FRvHC&Fcp6^ZEfCy^a7+uP4CO>k9Do`U0H2&H!((H^ANN4)FK- z1022%0gtaoz~$=_@cH@#oW4#0udi3Y?duls`}zeOzm5UVuV=vZ>l*O=`Uaf8&H?YQ zcfkGY9%T0a9J~$y53dKn#p?p_@%jLqyiNcwuNT10>jv=i`T-ohjsQ=uC&1O~3h?#% z0-U|h0B^51z}@Q(@b~%y9KH?#kFQ6-lX0)`UM=njsefF zXTbIA8u0!42Asdn0q?JO!2RnU{6GKykAv3%;NkTExOiOvK3*Sylh+C0<@Ex%dEEeh zUO#}N*Ad|9^#r(jT>-vcUx2gM8Q|^p2Dp3O0sdZpfWy}z;PLecxO`m#K3|`J)7L5B z_4Nw4ecb|nU%!Ck*D>Jv^$fUvT?4*f-+=SiIpF>E4!D2agUtS)gVzDz;q?HxcwGQK zULSyy*9qX|^#Ztg-2i@GKY*jx5#Z_d1h{%#0lr>efV0;b;O+GWxO?3J{$788!`C6; z@%0F}d|d)QU!Q=}*D2uj^$NIs-2#4JzkuV{G2r?247h$>1HND1fb-Wm;QjRuxPRS) z%>JK)*8$++^#HhdT>w5_AApnB3E<`R0=RkI0DfLSfTPzD;OX@QxO!ayzFuE|v)38m z?ezw@d))#4UVnhY*CF8X^$56pT>?H|pMcZXDd6?>3b=jU0)Ah=faBLO;Q93oxPDy& zzF*&f^Vd1x{q+vGf8B%3{-1-_0pQ{F0JwNv06tzHfRon=;N|rKxOv?GeqKL-qt_AO z>GcG-dR+m&USEK-*BRjL^#-_m-2whye}Kc+A>i@#2)KM*0zO}#fYaA0;Pv$ixP9FM zeqX-vcUx2gM8Q|^p2Dp3O0sdZpfWy}z;PLecxO`m# zK3|`J)7L5B_4Nw4ecb|nU%!Ck*D>Jv^$fUvT?4*f-+=SiIpF>E4!D2agUtS)gVzDz z;q?HxcwGQKULSyy*9qX|^#Ztg-2i@GKY*jx5#Z_d1h{%#0lr>efV0;b;O+GWxO?3J z{$788!`C6;@%0F}d|d)QU!Q=}*D2uj^$NIs-2#4JzkuV{G2r?247h$>1HND1fb-Wm P;QjRuxPRRP{{Q~~!`Ie> literal 0 HcmV?d00001 diff --git a/zig/src/families/deepseek_v41/prepare.zig b/zig/src/families/deepseek_v41/prepare.zig new file mode 100644 index 0000000000..0869e34320 --- /dev/null +++ b/zig/src/families/deepseek_v41/prepare.zig @@ -0,0 +1,144 @@ +//! The run-time weight forms built at load from named host tensors: linear strips, routed + shared experts' tables. + +const std = @import("std"); + +/// Trellis int16 [kt, nt, words] -> int32 [nt/8, kt, 8, words/2]: each 128-column block's tiles in K order. +pub fn strips(gpa: std.mem.Allocator, src: []const u8, kt: usize, nt: usize, words: usize) ![]align(16) u8 { + if (nt % 8 != 0 or src.len != kt * nt * words * 2) return error.BadTrellis; + const run = 8 * words * 2; // 8 tiles of one K row inside one 128-column block + const out = try gpa.alignedAlloc(u8, .@"16", src.len); + const blocks = nt / 8; + for (0..blocks * kt) |i| { + const nb = i / kt; + const k = i % kt; + @memcpy(out[i * run ..][0..run], src[(k * nt + nb * 8) * words * 2 ..][0..run]); + } + return out; +} + +/// Strips words -> dense3's lanes layout (dense3.cu to_lanes_kernel): per strip, lane L's bits of each G-step group. +pub fn lanes(gpa: std.mem.Allocator, strip_words: []const u8, kt: usize, k2: usize) ![]align(16) u8 { + const G: usize = if (k2 == 16) 1 else 2; + const nb = k2 / 2; // bytes of a lane's run of one tile (4 K2 bits) + const strip = 4 * kt * 32 * k2; // bytes a strip + if (k2 < 4 or k2 % 2 != 0 or kt % G != 0 or strip_words.len % strip != 0) return error.BadStrips; + const out = try gpa.alignedAlloc(u8, .@"16", strip_words.len); + const tile = 16 * k2; // bytes a tile (4 K2 words) + for (0..strip_words.len / strip) |s| { + const src = strip_words[s * strip ..][0..strip]; + const dst = out[s * strip ..][0..strip]; + for (0..kt / G) |gi| { + const g0 = gi * G * 8 * tile; // the group's first byte (src and dst) + for (0..32) |L| for (0..8 * G) |tt| for (0..nb) |j| { + const sb = g0 + tt * tile + L * nb + j; // stream byte in the source + const lb = tt * nb + j; // stream byte in the lane's run + const w = lb / 4; // the lane's word + const db = g0 + 4 * (((w / 4) * 32 + L) * 4 + w % 4) + (lb % 4); // stream byte in the destination + dst[4 * (db / 4) + 3 - db % 4] = src[4 * (sb / 4) + 3 - sb % 4]; + }; + } + } + return out; +} + +/// One expert of a projection: its trellis's device address and K2, its scales as stored (fp16). +pub const Expert = struct { trellis: u64, k2: i32, suh: []const u8, svh: []const u8 }; + +/// A projection's tables: int64 addresses, int32 K2s, suh [E, K] and svh [E, N] fp16. +pub const Tables = struct { + ptrs: []i64, + k2s: []i32, + suh: []align(16) u8, + svh: []align(16) u8, + k2_lo: i32, + k2_hi: i32, + + pub fn deinit(t: *Tables, gpa: std.mem.Allocator) void { + gpa.free(t.ptrs); + gpa.free(t.k2s); + gpa.free(t.suh); + gpa.free(t.svh); + t.* = undefined; + } +}; + +/// The tables of `experts` in order, the shared expert last; every suh / svh must have the projection's K and N. +pub fn tables(gpa: std.mem.Allocator, experts: []const Expert) !Tables { + if (experts.len == 0) return error.NoExperts; + const kb = experts[0].suh.len; + const nb = experts[0].svh.len; + var t: Tables = .{ + .ptrs = try gpa.alloc(i64, experts.len), + .k2s = try gpa.alloc(i32, experts.len), + .suh = try gpa.alignedAlloc(u8, .@"16", experts.len * kb), + .svh = try gpa.alignedAlloc(u8, .@"16", experts.len * nb), + .k2_lo = std.math.maxInt(i32), + .k2_hi = 0, + }; + errdefer t.deinit(gpa); + for (experts, 0..) |e, i| { + if (e.suh.len != kb or e.svh.len != nb) return error.ExpertShapesDiffer; + t.ptrs[i] = @bitCast(e.trellis); + t.k2s[i] = e.k2; + @memcpy(t.suh[i * kb ..][0..kb], e.suh); + @memcpy(t.svh[i * nb ..][0..nb], e.svh); + t.k2_lo = @min(t.k2_lo, e.k2); + t.k2_hi = @max(t.k2_hi, e.k2); + } + return t; +} + +const testing = std.testing; + +test "strips: each 128-column block's tiles in K order" { + const a = testing.allocator; + const kt = 3; + const nt = 16; + const words = 2; // one int32 a tile, to keep the check readable + var src: [kt * nt * words * 2]u8 = undefined; + // tile (k, n) holds the byte 16 k + n + for (0..kt) |k| for (0..nt) |n| @memset(src[(k * nt + n) * words * 2 ..][0 .. words * 2], @intCast(16 * k + n)); + const out = try strips(a, &src, kt, nt, words); + defer a.free(out); + // out[nb][k][j] = tile (k, 8 nb + j) + for (0..2) |nb| for (0..kt) |k| for (0..8) |j| { + try testing.expectEqual(@as(u8, @intCast(16 * k + 8 * nb + j)), out[((nb * kt + k) * 8 + j) * words * 2]); + }; + try testing.expectError(error.BadTrellis, strips(a, src[0..8], kt, nt, words)); +} + +test "lanes: the numpy reference's layout (to_lanes) at every dense3 width" { + // fixtures/lanes-ref.bin: for K2 4, 6, 8, 10, 12, 16: u32 K2, u32 NB, u32 KT, strips words, lanes words + const a = testing.allocator; + var rest: []const u8 = @embedFile("fixtures/lanes-ref.bin"); + var cases: usize = 0; + while (rest.len > 0) : (cases += 1) { + const k2 = std.mem.readInt(u32, rest[0..4], .little); + const nb = std.mem.readInt(u32, rest[4..8], .little); + const kt = std.mem.readInt(u32, rest[8..12], .little); + const bytes = 4 * nb * kt * 32 * k2; + const src = rest[12..][0..bytes]; + const want = rest[12 + bytes ..][0..bytes]; + const got = try lanes(a, src, kt, k2); + defer a.free(got); + try testing.expectEqualSlices(u8, want, got); + rest = rest[12 + 2 * bytes ..]; + } + try testing.expectEqual(@as(usize, 6), cases); +} + +test "expert tables: addresses, K2s and stacked scales in order, the width range" { + const a = testing.allocator; + const s1 = [_]u8{ 1, 2, 3, 4 }; + const s2 = [_]u8{ 5, 6, 7, 8 }; + const v1 = [_]u8{ 9, 9 }; + const v2 = [_]u8{ 7, 7 }; + var t = try tables(a, &.{ .{ .trellis = 0x7f00_0000_1000, .k2 = 4, .suh = &s1, .svh = &v1 }, .{ .trellis = 0x7f00_0002_0000, .k2 = 10, .suh = &s2, .svh = &v2 } }); + defer t.deinit(a); + try testing.expectEqualSlices(i64, &.{ 0x7f00_0000_1000, 0x7f00_0002_0000 }, t.ptrs); + try testing.expectEqualSlices(i32, &.{ 4, 10 }, t.k2s); + try testing.expectEqualSlices(u8, &.{ 1, 2, 3, 4, 5, 6, 7, 8 }, t.suh); + try testing.expectEqualSlices(u8, &.{ 9, 9, 7, 7 }, t.svh); + try testing.expect(t.k2_lo == 4 and t.k2_hi == 10); + try testing.expectError(error.ExpertShapesDiffer, tables(a, &.{ .{ .trellis = 1, .k2 = 4, .suh = &s1, .svh = &v1 }, .{ .trellis = 2, .k2 = 4, .suh = s2[0..2], .svh = &v2 } })); +} diff --git a/zig/src/families/deepseek_v41/prepared.zig b/zig/src/families/deepseek_v41/prepared.zig new file mode 100644 index 0000000000..4f0e0f911f --- /dev/null +++ b/zig/src/families/deepseek_v41/prepared.zig @@ -0,0 +1,281 @@ +//! Prepared weight folders (TF_DSV41_PREPARED, _WRITE, _VERIFY): built layer images cached on disk for a fast boot. +const std = @import("std"); +const quiet = @import("builtin").is_test; // a test's stderr is the build runner's channel +const named = @import("named.zig"); +const pack_mod = @import("pack.zig"); +const Io = std.Io; +const Allocator = std.mem.Allocator; + +const magic = "TFDSV41P"; +const version: u32 = 1; +const page: u64 = 4096; + +/// What the images depend on besides the pack's bytes. +pub const Key = struct { + rank: u32, + world: u32, + o_groups: u32, + blocks: []const u32, + dspark: bool, +}; + +/// The builders' source, so a change to how images are made never reads an old folder. +const sources = [_][]const u8{ @embedFile("named.zig"), @embedFile("plan.zig"), @embedFile("stage.zig"), @embedFile("pack.zig"), @embedFile("exl3.zig") }; + +pub const Stats = struct { hits: u32 = 0, misses: u32 = 0, written: u32 = 0, read_bytes: u64 = 0, read_ns: u64 = 0, write_ns: u64 = 0 }; + +pub const Cache = struct { + gpa: Allocator, + io: Io, + /// ``//rank`` + root: []u8, + write: bool, + verify: bool, + stats: Stats = .{}, + + /// The cache TF_DSV41_PREPARED names, or null when unset. + pub fn fromEnv(gpa: Allocator, io: Io, p: *const pack_mod.Pack, key: Key) !?Cache { + const dir = std.mem.span(std.c.getenv("TF_DSV41_PREPARED") orelse return null); + if (std.mem.trim(u8, dir, " ").len == 0) return null; + const flag = struct { + fn on(name: [*:0]const u8) bool { + const v = std.c.getenv(name) orelse return false; + return std.mem.eql(u8, std.mem.span(v), "1"); + } + }.on; + return try open(gpa, io, dir, try packDigest(io, p), key, flag("TF_DSV41_PREPARED_WRITE"), flag("TF_DSV41_PREPARED_VERIFY")); + } + + pub fn open(gpa: Allocator, io: Io, dir: []const u8, pack_digest: [32]u8, key: Key, write: bool, verify: bool) !Cache { + var h = std.crypto.hash.sha2.Sha256.init(.{}); + h.update(magic); + h.update(std.mem.asBytes(&version)); + for (sources) |s| h.update(s); + h.update(&pack_digest); + h.update(std.mem.asBytes(&[_]u32{ key.rank, key.world, key.o_groups, @intFromBool(key.dspark), @intCast(key.blocks.len) })); + h.update(std.mem.sliceAsBytes(key.blocks)); + var d: [32]u8 = undefined; + h.final(&d); + const hex = std.fmt.bytesToHex(d[0..8].*, .lower); + const root = try std.fmt.allocPrint(gpa, "{s}/{s}/rank{d}", .{ dir, &hex, key.rank }); + errdefer gpa.free(root); + if (write) try Io.Dir.cwd().createDirPath(io, root); + return .{ .gpa = gpa, .io = io, .root = root, .write = write, .verify = verify }; + } + + pub fn deinit(c: *Cache) void { + c.gpa.free(c.root); + } + + fn path(c: *const Cache, buf: []u8, tag: []const u8, suffix: []const u8) ![]const u8 { + return std.fmt.bufPrint(buf, "{s}/{s}.bin{s}", .{ c.root, tag, suffix }); + } + + /// Asks the kernel to read `.bin` ahead (POSIX_FADV_WILLNEED); a missing file is ignored, nothing waits. + pub fn willNeed(c: *const Cache, tag: []const u8) void { + var pbuf: [std.fs.max_path_bytes]u8 = undefined; + const p = c.path(pbuf[0 .. pbuf.len - 1], tag, "") catch return; + pbuf[p.len] = 0; + const linux = std.os.linux; + const rc = linux.openat(linux.AT.FDCWD, pbuf[0..p.len :0], .{ .ACCMODE = .RDONLY, .CLOEXEC = true }, 0); + if (linux.errno(rc) != .SUCCESS) return; + const fd: linux.fd_t = @intCast(rc); + defer _ = linux.close(fd); + _ = linux.fadvise(fd, 0, 0, linux.POSIX_FADV.WILLNEED); + } + + /// Fills `b` from ``.bin``: true when read, false when there is no usable file (`b` is then untouched). + pub fn get(c: *Cache, b: *named.Builder, tag: []const u8) !bool { + var pbuf: [std.fs.max_path_bytes]u8 = undefined; + const t0 = Io.Clock.awake.now(c.io).toNanoseconds(); + const file = Io.Dir.cwd().openFile(c.io, try c.path(&pbuf, tag, ""), .{}) catch { + c.stats.misses += 1; + return false; + }; + defer file.close(c.io); + const got = readInto(c.gpa, c.io, file, b, c.verify) catch |e| { + if (!quiet) std.log.warn("dsv41 prepared: {s} unusable ({t}); building it from the pack", .{ pbuf[0 .. c.root.len + tag.len + 5], e }); + c.stats.misses += 1; + return false; + }; + c.stats.hits += 1; + c.stats.read_bytes += got; + c.stats.read_ns += @intCast(Io.Clock.awake.now(c.io).toNanoseconds() - t0); + return true; + } + + /// Writes ``.bin`` from `b` when writing is on and the file is missing. + pub fn put(c: *Cache, b: *const named.Builder, tag: []const u8) !void { + if (!c.write) return; + var pbuf: [std.fs.max_path_bytes]u8 = undefined; + var tbuf: [std.fs.max_path_bytes]u8 = undefined; + const final_path = try c.path(&pbuf, tag, ""); + const tmp_path = try c.path(&tbuf, tag, ".tmp"); + const t0 = Io.Clock.awake.now(c.io).toNanoseconds(); + const dir = Io.Dir.cwd(); + { + const file = try dir.createFile(c.io, tmp_path, .{ .truncate = true }); + defer file.close(c.io); + try writeOut(c.gpa, c.io, file, b); + } + try dir.rename(tmp_path, dir, final_path, c.io); + c.stats.written += 1; + c.stats.write_ns += @intCast(Io.Clock.awake.now(c.io).toNanoseconds() - t0); + } +}; + +/// The pack's shards as they are on disk: each file's name, size and modification time. +pub fn packDigest(io: Io, p: *const pack_mod.Pack) ![32]u8 { + var h = std.crypto.hash.sha2.Sha256.init(.{}); + var buf: [std.fs.max_path_bytes]u8 = undefined; + for (p.files.items) |f| { + h.update(f); + const st = Io.Dir.cwd().statFile(io, try std.fmt.bufPrint(&buf, "{s}/{s}", .{ p.dir, f }), .{}) catch |e| switch (e) { + error.FileNotFound => { + h.update("missing"); + continue; + }, + else => return e, + }; + h.update(std.mem.asBytes(&[_]i128{ @intCast(st.size), st.mtime.toNanoseconds() })); + } + var d: [32]u8 = undefined; + h.final(&d); + return d; +} + +const Entry = extern struct { + kind: u8, + rank: u8, + name_len: u16, + pad: u32 = 0, + shape: [4]u64, + offset: u64, + len: u64, + sha256: [32]u8, +}; + +fn writeOut(gpa: Allocator, io: Io, file: Io.File, b: *const named.Builder) !void { + var head: std.ArrayList(u8) = .empty; + defer head.deinit(gpa); + try head.appendSlice(gpa, magic); + try head.appendSlice(gpa, std.mem.asBytes(&version)); + try head.appendSlice(gpa, std.mem.asBytes(&@as(u32, @intCast(b.out.items.len)))); + var size: u64 = 16; + for (b.out.items) |n| size += @sizeOf(Entry) + n.name.len; + var at = std.mem.alignForward(u64, size, page); + for (b.out.items) |n| { + var e: Entry = .{ .kind = @backingInt(n.kind), .rank = n.rank, .name_len = @intCast(n.name.len), .shape = undefined, .offset = at, .len = n.bytes.len, .sha256 = undefined }; + for (&e.shape, n.shape) |*d, s| d.* = s; + std.crypto.hash.sha2.Sha256.hash(n.bytes, &e.sha256, .{}); + try head.appendSlice(gpa, std.mem.asBytes(&e)); + try head.appendSlice(gpa, n.name); + at = std.mem.alignForward(u64, at + n.bytes.len, page); + } + try file.writePositionalAll(io, head.items, 0); + var cursor = std.mem.alignForward(u64, size, page); + for (b.out.items) |n| { + try file.writePositionalAll(io, n.bytes, cursor); + cursor = std.mem.alignForward(u64, cursor + n.bytes.len, page); + } + try file.setLength(io, cursor); +} + +/// Reads every tensor of a file into `b` (each into its own buffer); returns payload bytes. On error `b` gets nothing. +fn readInto(gpa: Allocator, io: Io, file: Io.File, b: *named.Builder, verify: bool) !u64 { + var fixed: [16]u8 = undefined; + if (try file.readPositionalAll(io, &fixed, 0) != 16) return error.Truncated; + if (!std.mem.eql(u8, fixed[0..8], magic)) return error.BadMagic; + if (std.mem.readInt(u32, fixed[8..12], .little) != version) return error.BadVersion; + const count = std.mem.readInt(u32, fixed[12..16], .little); + const total = try file.length(io); + var got: std.ArrayList(named.Named) = .empty; + errdefer { + for (got.items) |n| { + gpa.free(n.name); + gpa.free(n.bytes); + } + got.deinit(gpa); + } + var off: u64 = 16; + var bytes: u64 = 0; + for (0..count) |_| { + var e: Entry = undefined; + if (try file.readPositionalAll(io, std.mem.asBytes(&e), off) != @sizeOf(Entry)) return error.Truncated; + off += @sizeOf(Entry); + if (e.offset + e.len > total or e.rank > 4 or e.kind > @backingInt(named.Kind.i16)) return error.BadEntry; + const name = try gpa.alloc(u8, e.name_len); + errdefer gpa.free(name); + if (try file.readPositionalAll(io, name, off) != name.len) return error.Truncated; + off += e.name_len; + const data = try gpa.alignedAlloc(u8, .@"16", e.len); + errdefer gpa.free(data); + if (try file.readPositionalAll(io, data, e.offset) != data.len) return error.Truncated; + if (verify) { + var sha: [32]u8 = undefined; + std.crypto.hash.sha2.Sha256.hash(data, &sha, .{}); + if (!std.mem.eql(u8, &sha, &e.sha256)) return error.DigestMismatch; + } + var n: named.Named = .{ .name = name, .kind = @fromBackingInt(@intCast(e.kind)), .shape = undefined, .rank = e.rank, .bytes = data }; + for (&n.shape, e.shape) |*d, s| d.* = @intCast(s); + try got.append(gpa, n); + bytes += e.len; + } + try b.out.appendSlice(b.gpa, got.items); + got.deinit(gpa); + return bytes; +} + +test "a layer's images round-trip through its prepared file, and another key misses" { + const gpa = std.testing.allocator; + const io = std.testing.io; + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + const root = try std.fmt.allocPrint(gpa, ".zig-cache/tmp/{s}", .{&tmp.sub_path}); + defer gpa.free(root); + var p = pack_mod.Pack.init(gpa); + defer p.deinit(); + const key: Key = .{ .rank = 1, .world = 2, .o_groups = 8, .blocks = &.{ 0, 1, 2 }, .dspark = true }; + var c = try Cache.open(gpa, io, root, try packDigest(io, &p), key, true, true); + defer c.deinit(); + var b: named.Builder = .{ .gpa = gpa, .io = io, .pack = &p }; + defer b.deinit(); + const sizes = [_]usize{ 4096, 5000, 16, 1 }; + for (sizes, 0..) |n, i| { + const bytes = try gpa.alignedAlloc(u8, .@"16", n); + for (bytes, 0..) |*x, j| x.* = @truncate(j *% 31 +% i); + try b.out.append(gpa, .{ .name = try std.fmt.allocPrint(gpa, "L1.t{d}.trellis", .{i}), .kind = @fromBackingInt(@intCast(i % 4)), .shape = .{ n, 1, 1, 1 }, .rank = @intCast(1 + i % 4), .bytes = bytes }); + } + var miss: named.Builder = .{ .gpa = gpa, .io = io, .pack = &p }; + defer miss.deinit(); + try std.testing.expect(!try c.get(&miss, "L1")); + try c.put(&b, "L1"); + var back: named.Builder = .{ .gpa = gpa, .io = io, .pack = &p }; + defer back.deinit(); + try std.testing.expect(try c.get(&back, "L1")); + try std.testing.expectEqual(b.out.items.len, back.out.items.len); + for (b.out.items, back.out.items) |x, y| { + try std.testing.expectEqualStrings(x.name, y.name); + try std.testing.expectEqual(x.kind, y.kind); + try std.testing.expectEqual(x.rank, y.rank); + try std.testing.expectEqualSlices(usize, &x.shape, &y.shape); + try std.testing.expectEqualSlices(u8, x.bytes, y.bytes); + try std.testing.expect(std.mem.isAligned(@intFromPtr(y.bytes.ptr), 16)); + } + try std.testing.expectEqual(@as(u32, 1), c.stats.hits); + // another rank (or blocks, pack, builder source) is another folder + var other = try Cache.open(gpa, io, root, try packDigest(io, &p), .{ .rank = 0, .world = 2, .o_groups = 8, .blocks = &.{ 0, 1, 2 }, .dspark = true }, false, false); + defer other.deinit(); + var none: named.Builder = .{ .gpa = gpa, .io = io, .pack = &p }; + defer none.deinit(); + try std.testing.expect(!try other.get(&none, "L1")); + // a damaged payload is refused with VERIFY and the builder stays empty + var pbuf: [std.fs.max_path_bytes]u8 = undefined; + const f = try Io.Dir.cwd().openFile(io, try c.path(&pbuf, "L1", ""), .{ .mode = .read_write }); + try f.writePositionalAll(io, "x", page); + f.close(io); + var bad: named.Builder = .{ .gpa = gpa, .io = io, .pack = &p }; + defer bad.deinit(); + try std.testing.expect(!try c.get(&bad, "L1")); + try std.testing.expectEqual(@as(usize, 0), bad.out.items.len); +} From 7ccfc5cb0e15e70f06e8bdf305095da75859223c Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Sun, 11 Oct 2026 19:32:16 +0700 Subject: [PATCH 031/104] DeepSeek-V4.1 Engram host side: table layout, n-gram hashes of a window's tokens, a rank's heads. engram_host.zig lays out the Engram tables of the config, hashes each window's token n-grams into table rows as the reference hasher does (checked on the release config and tokenizer, fixtures/engram-hash-ref.bin), splits the heads over the TP ranks and builds the synthetic index rows the tests use. --- .../families/deepseek_v41/deepseek_v41.zig | 2 + zig/src/families/deepseek_v41/engram_host.zig | 307 ++++++++++++++++++ .../deepseek_v41/fixtures/engram-hash-ref.bin | Bin 0 -> 27596 bytes 3 files changed, 309 insertions(+) create mode 100644 zig/src/families/deepseek_v41/engram_host.zig create mode 100644 zig/src/families/deepseek_v41/fixtures/engram-hash-ref.bin diff --git a/zig/src/families/deepseek_v41/deepseek_v41.zig b/zig/src/families/deepseek_v41/deepseek_v41.zig index 13660d9cc2..9e024fd039 100644 --- a/zig/src/families/deepseek_v41/deepseek_v41.zig +++ b/zig/src/families/deepseek_v41/deepseek_v41.zig @@ -11,6 +11,7 @@ pub const Plan = plan.Plan; pub const stage = @import("stage.zig"); pub const named = @import("named.zig"); pub const prepare = @import("prepare.zig"); +pub const engram_host = @import("engram_host.zig"); test { _ = config; @@ -21,6 +22,7 @@ test { _ = stage; _ = named; _ = prepare; + _ = engram_host; _ = @import("host_test.zig"); _ = @import("phases.zig"); // the phase timers // prepared weight folders (TF_DSV41_PREPARED) diff --git a/zig/src/families/deepseek_v41/engram_host.zig b/zig/src/families/deepseek_v41/engram_host.zig new file mode 100644 index 0000000000..cb40006758 --- /dev/null +++ b/zig/src/families/deepseek_v41/engram_host.zig @@ -0,0 +1,307 @@ +//! Engram's host side: table layout, n-gram hashes of a window's tokens, a rank's heads, and the synthetic index rows. + +const std = @import("std"); +const quiet = @import("builtin").is_test; // a test's stderr is the build runner's channel + +pub const dead_id: i64 = -1; +pub const max_layers = 4; +pub const max_ngram = 8; +pub const max_cols = 64; + +/// Where a window's table rows come from (local NVMe shards or a remote service); every source gives the same bytes. +pub const Source = struct { + ptr: *anyopaque, + vtable: *const VTable, + + pub const VTable = struct { + /// Starts fetching Engram layer index `li`'s rows `idx` (this rank's heads' global rows); returns at once. + issue: *const fn (ptr: *anyopaque, li: usize, idx: []const i64) anyerror!void, + /// Layer index `li`'s rows `idx` as bf16 bits [idx.len, dim] (waits for any still in flight). + rows: *const fn (ptr: *anyopaque, li: usize, idx: []const i64, dim: usize, out: []u16) anyerror!void, + }; +}; + +pub const Host = struct { + layer_ids: []const u32, + max_ngram: u32, + heads: u32, + /// token id -> compressed id + token_map: []const u32, + pad_id: i64, + mult: [max_layers][max_ngram]i64 = undefined, + primes: [max_layers][max_cols]i64 = undefined, + offsets: [max_layers][max_cols]i64 = undefined, + + pub fn cols(h: *const Host) u32 { + return (h.max_ngram - 1) * h.heads; + } + + /// Prime bucket sizes: for each layer and order the next unused primes above `vocab_size - 1` (reference order). + pub fn layout(h: *Host, vocab_size: u64) void { + var seen: [max_layers * max_cols]i64 = undefined; + var nseen: usize = 0; + for (0..h.layer_ids.len) |li| { + var col: usize = 0; + for (0..h.max_ngram - 1) |_| { + var cur: i64 = @intCast(vocab_size - 1); + for (0..h.heads) |_| { + cur += 1; + while (!isPrime(@intCast(cur)) or std.mem.indexOfScalar(i64, seen[0..nseen], cur) != null) cur += 1; + seen[nseen] = cur; + nseen += 1; + h.primes[li][col] = cur; + col += 1; + } + } + var acc: i64 = 0; + for (0..col) |c| { + h.offsets[li][c] = acc; + acc += h.primes[li][c]; + } + } + } + + /// Table rows [t][layers][cols] of `ids` after `lookback` (oldest first; short: sequence start); past map: dead. + pub fn hash(h: *const Host, ids: []const u32, lookback: []const u32, out: []i64) void { + const nb = lookback.len; + const L = h.layer_ids.len; + const C = h.cols(); + std.debug.assert(out.len == ids.len * L * C); + for (0..ids.len) |t| for (0..L) |li| { + var blocked = false; + var rolling: i64 = 0; + for (0..h.max_ngram) |s| { + const j = @as(i64, @intCast(t + nb)) - @as(i64, @intCast(s)); + const before = j < 0; + const v: i64 = if (before) h.pad_id else h.src(if (j < @as(i64, @intCast(nb))) lookback[@intCast(j)] else ids[@as(usize, @intCast(j)) - nb]); + blocked = blocked or before or v == dead_id; + const value = if (blocked) h.pad_id else v; + rolling ^= value *% h.mult[li][s]; + if (s == 0) continue; + for (0..h.heads) |hh| { + const c = (s - 1) * h.heads + hh; + out[(t * L + li) * C + c] = @mod(rolling, h.primes[li][c]) + h.offsets[li][c]; + } + } + }; + } + + fn src(h: *const Host, id: u32) i64 { + return if (id >= h.token_map.len) dead_id else h.token_map[id]; + } + + /// (first head, heads) rank `rank` of `world` owns: complete heads, TP-major (head_shard). + pub fn headShard(h: *const Host, rank: u32, world: u32) [2]u32 { + const n = h.cols(); + const part = (n + world - 1) / world; + const h0 = @min(rank * part, n); + return .{ h0, @min(h0 + part, n) - h0 }; + } + + /// This rank's heads' global rows of Engram layer `layer` for a window's `ids` after `lookback`: idx [n][heads]. + pub fn shardIndex(h: *const Host, a: std.mem.Allocator, ids: []const u32, lookback: []const u32, layer: u32, rank: u32, world: u32) ![]i64 { + const li = std.mem.indexOfScalar(u32, h.layer_ids, layer) orelse return error.NotEngram; + const L = h.layer_ids.len; + const C = h.cols(); + const sh = h.headShard(rank, world); + const hashes = try a.alloc(i64, ids.len * L * C); + defer a.free(hashes); + h.hash(ids, lookback, hashes); + const idx = try a.alloc(i64, ids.len * sh[1]); + for (0..ids.len) |t| @memcpy(idx[t * sh[1] ..][0..sh[1]], hashes[(t * L + li) * C + sh[0] ..][0..sh[1]]); + return idx; + } + + /// Engram layer `layer`'s window rows, this rank's heads, bf16 [ids.len, heads x dim]: from `source` or index rows. + pub fn windowRows(h: *const Host, a: std.mem.Allocator, source: ?Source, ids: []const u32, lookback: []const u32, layer: u32, rank: u32, world: u32, dim: usize, out: []u16) !void { + const idx = try h.shardIndex(a, ids, lookback, layer, rank, world); + defer a.free(idx); + const s = source orelse return indexRows(idx, dim, out); + const li = std.mem.indexOfScalar(u32, h.layer_ids, layer).?; + return s.vtable.rows(s.ptr, li, idx, dim, out); + } + + /// Starts reading a coming window's rows on every Engram layer; a wrong guess costs a read; errors only logged. + pub fn prefetch(h: *const Host, a: std.mem.Allocator, source: ?Source, ids: []const u32, lookback: []const u32, rank: u32, world: u32) void { + const s = source orelse return; + for (h.layer_ids, 0..) |layer, li| { + const idx = h.shardIndex(a, ids, lookback, layer, rank, world) catch return; + defer a.free(idx); + s.vtable.issue(s.ptr, li, idx) catch |e| if (!quiet) std.log.scoped(.dsv41).warn("engram: prefetch of layer {d} failed ({t})", .{ layer, e }); + } + } + + /// Loads recipe/dsv41_engram_host.py's file (the map borrowed from `bytes`, which must outlive the host). + pub fn load(bytes: []const u8, heads: u32, pad_token: u32, vocab_size: u64, layer_buf: []u32) !Host { + if (bytes.len < 24 or !std.mem.eql(u8, bytes[0..8], "DSV41EH1")) return error.BadEngramHost; + const L = std.mem.readInt(u32, bytes[8..12], .little); + const M = std.mem.readInt(u32, bytes[12..16], .little); + const V = std.mem.readInt(u32, bytes[16..20], .little); + if (L > max_layers or M > max_ngram or (M - 1) * heads > max_cols or L > layer_buf.len) return error.BadEngramHost; + var at: usize = 24; + for (0..L) |i| layer_buf[i] = std.mem.readInt(u32, bytes[at + 4 * i ..][0..4], .little); + at += 4 * L; + var h: Host = .{ .layer_ids = layer_buf[0..L], .max_ngram = M, .heads = heads, .token_map = &.{}, .pad_id = 0 }; + for (0..L) |li| for (0..M) |s| { + h.mult[li][s] = std.mem.readInt(i64, bytes[at..][0..8], .little); + at += 8; + }; + if (bytes.len - at < 4 * @as(usize, V)) return error.BadEngramHost; + h.token_map = @alignCast(std.mem.bytesAsSlice(u32, bytes[at..][0 .. 4 * V])); + h.pad_id = h.token_map[pad_token]; + h.layout(vocab_size); + return h; + } +}; + +/// The committed ids before `start` (last `keep_n`): the tail, then kept leading rows of `pending`; null: unreachable. +pub fn lookbackAt(tail: []const u32, pos: u64, pending: ?[]const u32, start: u64, keep_n: usize, buf: []u32) ?[]const u32 { + if (start < pos) return null; + const k: usize = @intCast(start - pos); + const extra: []const u32 = if (pending) |p| (if (k <= p.len) p[0..k] else return null) else if (k == 0) &.{} else return null; + if (tail.len + extra.len > buf.len) return null; + @memcpy(buf[0..tail.len], tail); + @memcpy(buf[tail.len..][0..extra.len], extra); + const m = tail.len + extra.len; + const n = @min(m, keep_n); + return buf[m - n .. m]; +} + +/// A slot's Engram tail after keeping `ids`: the last `keep_n` of tail + ids, into `tail`; returns its new length. +pub fn tailAfter(tail: []u32, len: usize, keep_n: usize, ids: []const u32) usize { + var all: [2 * max_ngram + 64]u32 = undefined; + const old = tail[0..len]; + const take = ids[ids.len -| keep_n..]; + @memcpy(all[0..old.len], old); + @memcpy(all[old.len..][0..take.len], take); + const m = old.len + take.len; + const k = @min(m, keep_n); + @memcpy(tail[0..k], all[m - k .. m]); + return k; +} + +test "an Engram warm's lookback at a draft before its window's keep is the window's tail after it" { + // TF_DSV41_ENGRAM_WARM: warmEngram asks at pos + kept, pending window in flight; stageEngram reads after the keep + var prng = std.Random.DefaultPrng.init(56); + const rnd = prng.random(); + for (0..500) |_| { + const keep_n = 3; + var tail: [max_ngram]u32 = undefined; + const len = rnd.uintAtMost(usize, keep_n); + for (tail[0..len]) |*t| t.* = rnd.int(u32); + var pend: [6]u32 = undefined; + const pn = 1 + rnd.uintLessThan(usize, pend.len); + for (pend[0..pn]) |*t| t.* = rnd.int(u32); + const pos: u64 = 100 + rnd.uintAtMost(u64, 50); + const kept = rnd.uintLessThan(usize, pn); // drafts kept: rows 0..kept, the next window starts past them + const start = pos + kept + 1; + var buf: [2 * max_ngram + 64]u32 = undefined; + const before = lookbackAt(tail[0..len], pos, pend[0..pn], start, keep_n, &buf).?; + var after = tail; + const alen = tailAfter(&after, len, keep_n, pend[0 .. kept + 1]); + try std.testing.expectEqualSlices(u32, after[0..alen], before); + // a keep already done (the commit before the draft): no pending window, the same tail + var buf2: [2 * max_ngram + 64]u32 = undefined; + try std.testing.expectEqualSlices(u32, after[0..alen], lookbackAt(after[0..alen], start, null, start, keep_n, &buf2).?); + } +} + +/// Small primes, then Miller-Rabin with bases 2, 7, 61 (exact below 4.7e9). +pub fn isPrime(n: u64) bool { + if (n < 2) return false; + for ([_]u64{ 2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37 }) |p| if (n % p == 0) return n == p; + var d = n - 1; + var r: u32 = 0; + while (d % 2 == 0) : (r += 1) d /= 2; + for ([_]u64{ 2, 7, 61 }) |a| { + var x = powMod(a, d, n); + if (x == 1 or x == n - 1) continue; + var ok = false; + var i: u32 = 1; + while (i < r) : (i += 1) { + x = @intCast(@as(u128, x) * x % n); + if (x == n - 1) { + ok = true; + break; + } + } + if (!ok) return false; + } + return true; +} + +fn powMod(b0: u64, e0: u64, m: u64) u64 { + var b: u128 = b0 % m; + var e = e0; + var r: u128 = 1; + while (e > 0) : (e >>= 1) { + if (e & 1 == 1) r = r * b % m; + b = b * b % m; + } + return @intCast(r); +} + +/// Synthetic rows: row i, value j = (((i * 2654435761 + j * 40503) % 511) - 255) / 1024, exact in bf16, into `out`. +pub fn indexRows(rows: []const i64, dim: usize, out: []u16) void { + for (rows, 0..) |i, r| for (0..dim) |j| { + const v = @mod(i *% 2654435761 +% @as(i64, @intCast(j)) * 40503, 511) - 255; + const f: f32 = @as(f32, @floatFromInt(v)) / 1024.0; + out[r * dim + j] = @truncate(@as(u32, @bitCast(f)) >> 16); + }; +} + +const testing = std.testing; + +test "the n-gram hashes: the reference NgramHasher on the release config and tokenizer (fixtures/engram-hash-ref.bin)" { + const a = testing.allocator; + const fx: []const u8 = @embedFile("fixtures/engram-hash-ref.bin"); + const L = std.mem.readInt(u32, fx[8..12], .little); + const M = std.mem.readInt(u32, fx[12..16], .little); + const V = std.mem.readInt(u32, fx[16..20], .little); + var at: usize = 24 + 4 * L + 8 * L * M; + // a full-size map holding only the entries the fixture's ids use + const map = try a.alloc(u32, V); + defer a.free(map); + @memset(map, 0); + const used = std.mem.readInt(u32, fx[at..][0..4], .little); + at += 4; + for (0..used) |_| { + map[std.mem.readInt(u32, fx[at..][0..4], .little)] = std.mem.readInt(u32, fx[at + 4 ..][0..4], .little); + at += 8; + } + var file: std.ArrayList(u8) = .empty; + defer file.deinit(a); + try file.appendSlice(a, fx[0 .. 24 + 4 * L + 8 * L * M]); + try file.appendSlice(a, std.mem.sliceAsBytes(map)); + const bytes = try a.alignedAlloc(u8, .@"4", file.items.len); + defer a.free(bytes); + @memcpy(bytes, file.items); + var ids_buf: [max_layers]u32 = undefined; + const h = try Host.load(bytes, 8, 2, 16_000_000, &ids_buf); + // the reference layout: the first and last primes, the last head's offset + try testing.expectEqual(@as(i64, 16_000_057), h.primes[0][0]); + try testing.expectEqual(@as(i64, 16_000_889), h.primes[1][23]); + try testing.expectEqual(@as(i64, 368_005_705), h.offsets[0][23]); + const cases = std.mem.readInt(u32, fx[at..][0..4], .little); + at += 4; + for (0..cases) |_| { + const n = std.mem.readInt(u32, fx[at..][0..4], .little); + const nb = std.mem.readInt(u32, fx[at + 4 ..][0..4], .little); + at += 8; + const ids_c = try a.alloc(u32, n); + defer a.free(ids_c); + for (ids_c, 0..) |*x, i| x.* = std.mem.readInt(u32, fx[at + 4 * i ..][0..4], .little); + at += 4 * n; + const lb_c = try a.alloc(u32, nb); + defer a.free(lb_c); + for (lb_c, 0..) |*x, i| x.* = std.mem.readInt(u32, fx[at + 4 * i ..][0..4], .little); + at += 4 * nb; + const got = try a.alloc(i64, n * L * h.cols()); + defer a.free(got); + h.hash(ids_c, lb_c, got); + for (got, 0..) |g, i| try testing.expectEqual(std.mem.readInt(i64, fx[at + 8 * i ..][0..8], .little), g); + at += 8 * got.len; + } + try testing.expectEqual(fx.len, at); + try testing.expectEqual([2]u32{ 12, 12 }, h.headShard(1, 2)); +} diff --git a/zig/src/families/deepseek_v41/fixtures/engram-hash-ref.bin b/zig/src/families/deepseek_v41/fixtures/engram-hash-ref.bin new file mode 100644 index 0000000000000000000000000000000000000000..2e1b20b38c62cd97ad58ca49a4ce2870dc29116e GIT binary patch literal 27596 zcmeI)c{EmS!}sxhN=PE1(xgFXjx>i#8fcPA1C107DkTjlg$$KOL!?NP2926X6x9`? z@uH#(rCAy@sn_{E?mynY-hZBVz3;Qub6TtQyREy+^*Mg~IJdozy^m{{!(?+)n~|nU zk|ZhD{gHmiB;AcNNmlniHS7Ky{QY$F85>FRcWquWukJ5+o$%XuuB{|J-cb`^Yb{9@ zn#(swSV)r20!O>T?vfPm_hyMfD@k%(esGkbog^97{S5ifmj5wrBS~B9e(0nbl4RLH zlGZMiq^DgaDRrGB>A6VKCx1!0=qgD=0_sv*l5)07(u8JpKPy_2rh3=?b9*Ie@ft~b zm?=qX>VWxoC29D6Nir^wB)^W5^s=c;YSuxL>N zNuT=2q(yNu>DVBd#J!v4bm?ZqW70I=sJ{Gr%Yq`Z z+|}<=@ZI>!&Tx4?YW<$n$LO79Dw3nC0rqQ+KnQo;)?ma+}JtUd$nRNY=-i ze+A7Qc|Mu_Z;=b{wnTq^-S=d7SM=z$6Q}?#TofDl+u5?dTm;b#o zKYQF)8;@!v|GxTx8>y!n%Q$kI9mPL+sR3?ZZ4a_fNDBU-bJ6%;8I8f8wz(_*8TqOB zSCpKU|5#~UR3rQ!*XH7{N_s=;tJTNfx!FAY^Diaf_v^C~_x0U+^xwv&0d6PjLAY<0 z9UxtT;^<%hrVZ|q1_N*(F$^GUv&`syXVU<3LGQQtTRNHG-yA<2zpmtiKmV{Z{#!Q( z<1g}Ei{C}X6n|=8AKZp1d+Far+Zumv;5%~8ib?qQs#)My{?Zfw2YnO#8xAePT^{m? zd=PRQe_XCR#~~`*2zMW~0@8FyPuzF5Tp(jD5y;ZXo7Q&LaFd!@rTic9r-IE=TgzZU&>^Dn?(tTGn&rtRIB_sen)lf}-{8ON0DX^i91<&OAYHySJYx4wdZ~Po{oR2Q#<@!hS}pk9Ik?Y zL4SSve-$wu{~nj0iopq%i{a=j{=O5*= z1pI}bdbo3XHpbmbw-xS}1J{x*XOwar@|V8Af9ZA}e%+L7_#Ijt#lPTJE&hT%W%zrx zSHpks@(;3LkFWP?XH>3^Pv z8g5Cqnylv^gx`D3Gy2a}O2)tKM;`v)-!b{-|jg z_($e-#lJDr9KXYzx{LI>uL9l8_+5)f;4jr^j@xD8Q?k@>Ab#Dj*7SbOqyg?ozr&L*Qb`_=X!9V7SPkM%0V zZ*H7}|K-mj{1(69NghT4xN{<>(7)03Tlh!Kd`5b>@5eu~WgFa)<}tXJtvN*6syxI$ z#7TwTT?54VN6E>Har6yojeB`V@%wjrpZoN$dEXnqzR4u~+QqBrzox6Wet1*68-M51 zGx0CT66eoat6z*`;1mA3R63X2T4_h#KloSW2jRZ&Tuo{ge_^0MtB>Mu9Cnh_ZRm@; z+*h1Gb30q%A2qKH?o+=Gl4^%*=)cP39pm_MwmI$%@n1==&i~?Hcl`*U-|Mf@e}jQ9 z@mm*p;9tGe68C~9>*&8>%K+w2Y@=rMe_dG{|G9JG{^kDmAl!BtV*cEhDd7(F8cpv7 z+Sd5vbRGHJHW*(Y|AG;B$%@fG@$d3Hh2Qp<8om3cP9&9=Y{Wlh!!-Q<4F=+073Ylq z&Xz#@gW_M}PntCx|Cqnx{zdcUVEo5C=itAQc@cl|+CSu1>!bK*)}CPg_+~laUvn*+ z-1YV|HZiV|!R32&B%oev<<}xxQ<^}$VQfu7mFG9$tQ{(Z+$GG9QI@}QVP?fjj zXl>E|{QJxC=dD^z|8DOd;s3aG7JjWlf))^oeo!yFX$aS1#f^rcqDAlSuP=~`CN;%<+^ILc`__*lyAGU-zqPiQ|0%mG@t^7}=6{2b zUie!iE8#vHbDa!sCVu}088pFvbWSh)Blm2=e`$6Y?!pRF+^4Uc$NlO4bo$SHD?YzJ zQwteKK>2q1uly+d&$`XQUF`dZoZ2OX{*NBN$~Xe%cgKC;*?V&BouBj{6C>6i9OtOw zexfhtPya6UaG!9hA>+er=)Y}4P5e3rZSkkuoW;M(jHMlEuggIE>;6o^ztGbU|Ce6B z$(H@P(ZAhj4gBX-T62z!=+UGJ#1#-{~7j;3`l}0wO%03KQo80Tqre55dX)>askz^Cv0YkAZ1kI6(i8 zddA{cPAerf7Y5-Dv$;razYxK3*!R8wf6?oU_*HBYI1c}q7BUX+)=h8^^EgAMgkHv9 z`FSyZKi6RVdgY@3J^@Aaf2zY!{4uRG@S9GUj6X{~0DmVvYy8g)V)2(ojK+UUeE|N1 z-DRYcrv>xgtckdO=+__vfB3*Dxc#Gt;(p&|18Kha0RGvD^YNegD1LvLS7_sYaZG%E zgSYtNul7pDpZNt7 zFpe8tt??IR_~WyC!4p$e{<|<`cGYU4S$8k8T^%=Gx7flT#nx+ zTCD#>WT@iSIo^Q&9WPdr&0dN7-`SH734cyL{(jrt@w+%_;y*OX4EKec`nb!RHp5-3 zT|#Dl{ffU);u8ED<81L8P8o<_a<#;NvE~u}9m-<<>zL-^pLORZ{*dEh{W-khPsR~4 zz8C$U3z(0;K>IbB6#E%}+E!!SlQYHn)8~98ep&z7_}`oOk#8-M7>DMwf%xat7x#~S zCy4p;QFaOckim}lCu;=Z-!SGAes9Bc{BJgh_0NS>-t<4XEFXWFLoWUizGD6DUYAw$ zKWFb@{4;m`LmIRlgWGk^Wim1G9rI_5pZuwhBsHG)ioD=2);~|*6X*Xxb1(c(N<|!p zRT*Oaq~9q$+{x!1NS`%U_!swe#lQDb34Z@3V*R{_o>)I0v8Vw5xU8f2TQA&&|4rAQ z_)pB=ihE@zvHslrlrir1hF8g(XO_~x(whyqA9fqXd1vbOcrw1ChWbkhg%1@KNCi=KNMtvvCJ59&$*~*>%ujub({`{V&j{j{94C(PA@%#6J zi}?Je2h7HA;?NMk;ew;gn?5O9_y?ZsXhrWwBSzu>_htT&pW7v+(ERNLYo=d>F4`{Qfr-)FWL?$UcQdhb-Si7eal zjd7?st!EtH?Zox}9i1Kc-CdgE4)N5*9hc-m`c1lm|9oT|{<3X7@CQ1I^@j_E{52`w zcs=u`qFUVlj*AV&pE|9AyyTzFI0`2mW5CW2ipkjtcNxcrnD2~3>q9Q%*sH&eab&xR z`**d{llToD=Hj;vuH-l@GYi45d$~S-J+r~2L5;ZnUp&S>%!~bfZ4gLT9 z)r|RX`pyJ@scsVe2gQ!WeQfqV+(|A`^uFosH&Ui+jyv32te>wNU50;inpl7P_WlI^ zkAu~4U;Q)zcjLYBWVNLX_x9f6{Qu^g3-doMOmFX#eXVVO8NdC!8nS9*4dc*sH^6@=zZ>q2g-P`8sjh+l^yT`vXIdN6dz-_` zxbq9≦SBGJn>{#Q7&(Upzl@-zL_tu6rdhj`_zM()*{C<4FBy8}MuWXZ=}otTXe+ z-)Iv34{qg#|Iqy9__r>Zgj=cCM%-r4XVCv4>w9F}H5L4cs;cKY2Z-zElZkKe&$Oz0pj`K*JnmooGgM=7zYF|9CVh3szk0SI{;rPua8G(E z?!Tg5U*9Zn}UDweoOoZvc}MVQdkK7<|Xe*qv^+)|6Af~NlTqb`fsJa z2lwf#C&(lBgYggEm`)b1cEG*aMqK~64HM_j_$*KS`Uiyns#;TeS4rMM9=g8*|G>?q z_`kOe#=pa;A^w4HcHpGuZLT&k1oBZrt0Fi9c{zB>7{X*d@!*Xe{IDy`hVP}g87rbZZQ6z>SFz5 z(SM%*+BY!39~wIs|0<;s_#2J0{`Z^HE$s zT#PfL_p$rh;-9l+1^)5HuJ~)+l<|)|V2xYrhc~&f#|YeA)LY@cv2r14Q#uoW<}o#T z@AuH1Ogtv8pSv%aj(ZY|+*?&{lKXG@;BRMJpWaJ@O>k$q zMUeSTZ!-VCO)tiO{)>42?!3nqYLQQf08RJ6# zv0lOWyDb&(&n!3>{i z0OR_8V*YeX70(}S@?7xSeJLYb z-;cmwHBYP`CeL@F|J%n0h~7)`$eJ(n=>K`&!{o{y;`eV^)2;YreeCJI>sxKywnMg& zvzCebuUl&3{&$~c3I4fzD)84>MBx9J6OVsWQ*qzv9YZ>I38DY%J$5sW zS6eOUy_chS|6%aG1m=(7`dj#Gzli5|_Mvw8EoSQA@AtzWf9F|B_|0Z~#(%QE9)5e1 z*7)0H7~sDf+W>c@mjiCc=HmW0q`V#Oj7CG5H@6m;;kI6C!@PO7%@y|%-DbEGLL1@^ zx-gR7M?Y1?Z|Z!TJea=@zxUUF@$X7YqW@{LhT#uiZ-w8pP#3?Y$saOI#Ss5hle(Aa z>b_dkTY`W4KH)#~K&&5nKK{Ts!shQ|{+L|ojo;O`gxunJgmF{`_QU`7P*eOm1>*X1 z=*!cLmB~b&6IG@_K7CHH?P9Kqf=}AUh^B^kGL@xx0b~=`p$=uF0M+jcbmnbFG3pN*r$`hU+s zTN#IO?LEdZ@$V7*+f`f9d!de4Ki_Zh5dXfBQQ|zh_Z#CFqaTicp4KD$XP@1o|1Fy& z{11DG_h(8B-{N;oY>mHF>}33*G7bEPiqGNp+0lUhonDLQzX#Jj@W1MoOg^kQ#W+5$ z7v~?Rpgs7n9e;#@=g1h^Xx8&DsE!=tIN0MzGw8DS<`qTgDlgIG){;I(#iWx`HJ#qar!)G*p`zLC+kNyoLWlrH7hwClH_46C8VEWez_<+A=qIiDfo#cjp zV3iiV&vUOJyX@$V-)f6D{+Q9)_`gJ)#NXe#?uFgDugAp}^#AblH0FOuwz&Rrn~;aU z;*wbZ(YiGlcZl64`ft8hJU_M26z_i;9TV%H8Doy(pWSUM{>Z-8xSiV%#jSiam9$F> zq<@vp<@jGOZAbr)AMU`fGe+D$c5ihZe~!`w{O{f3@b{Y13;*HqY52Du5%Z^usx$6m zf#!UEKlSTQ|C^V8A=UH7;1ASa!u+{DOgw+rF%ZwM=T(dM-)if{;J@9Zn)$Ol?k3}S zq%7XQwYEKjKe1IUe$T~+@DI!Bh`;6Td-PxI@(llk=3DUJZrlO?lsFswb9{E;k94h1 z|Ca5qX)V z@iYF_OPuLHe4QiypP}OY2gxxMe}0HH{sEK4`osBv=J@+s1mnLk)enD_wphOneIJ8A zZ=zT~EPwHkaXcLDf;+t`j^6vYrI8bCZ!(Vz2d~B-bmJfLvC(1t2|bQ6e`aVT(0_V6 z@%%1ICkTJH3N76El_vDAG(oICER7N8pTGt$@Hg;ZME}vLxAFh`R~3JFZV3L}E9~jt zS7wC&RqRUKEvw7O{OObNN3IdiuUGF9?=M|_?8E%oA6-apiuAx8klq>h>W>X^N1Kc1 zFZb;R;kKGyM7CGkNB@fxjxdgm3oRK(J5_Q2-nmx1f93n#AOFCtX#9^(*yHyZF%ka| z8%z9up6tfIQ!>M?bLuzQKufIOj`EyE@8LT>lEWK|_YYGXI?=oSvOf5qjO~rzG@*

mb;h0V7Dc8UImvNY zxVAl9$;*=LK-@lLUS#InW!_k+OyJU?~%#c`-OJd6H6d5QOLuiA+D^P+J9{jVC` z48Q)VApEg|bMZUgiNN2)UYvgd>ep@BsrxEFQ$;R2^OXJvjT7r9)-FqM@0%>%|Jl=9 zod1nH*W>PSO}sxQtGY!0Q@Xvw|MEZU&s%n!rT?;?0gS`Asypsi!#0r_ek#nHzB5*m zw=ZSkfBth8e#2m~ep2S%8vl;jrTA@>J?X!VbPIRbjP|%2O+U^2&%5ix{P#5U!a(kmv=FgCe(fEHHxF9|UAH(T?{DYJD4<%1# z{=6y_&))}KT#Vn}dMP2CIV@jmB$B%qj%UI_Z#+)8HY-M_>BAaY)S9Wuf8W=^$()|nluN-;ce&091Hf6;r5v= z-Z37brH1?DC0+8)TCvBc$To{P7Mv4>|G&CN$unI%qff39{SWh>O8;LRqR2z@pVR-g ze!KCX$`WgIlU2n#S|MB4;lI2to#W8Ms2Af)=Tz*I zQtXqGe;yS3q!jz4*NDfUSz_DRX-fnuMOVxN?J z9w_!nDfUUp=Ye9MlwzNhd>$zFNh$V8$>)J$pOj*slzbj2_DL!BNy+DdVxN>^pOkza zDE3Jy_DRX-fnuMOVxN?J9w_!nDfUUp=Ye9MlwzNhd>$zFNh$V8$>)J$pVa@|eNz7m Dj3d4g literal 0 HcmV?d00001 From b775815b404f19f488c5615ff61e011f5d63f493 Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Sun, 11 Oct 2026 19:31:50 +0700 Subject: [PATCH 032/104] DeepSeek-V4.1 phase timers: where a serving round's time goes, by named phase (TF_DSV41_PHASES). TF_DSV41_PHASES=1 times each named phase of a round on the host clock, =sync also waits for the stream at each boundary, and TF_DSV41_PHASES_OUT= writes the totals as JSON. Off by default; nothing reports into it yet. --- .../families/deepseek_v41/deepseek_v41.zig | 1 + zig/src/families/deepseek_v41/phases.zig | 206 ++++++++++++++++++ 2 files changed, 207 insertions(+) create mode 100644 zig/src/families/deepseek_v41/phases.zig diff --git a/zig/src/families/deepseek_v41/deepseek_v41.zig b/zig/src/families/deepseek_v41/deepseek_v41.zig index dd59ec7fe4..13660d9cc2 100644 --- a/zig/src/families/deepseek_v41/deepseek_v41.zig +++ b/zig/src/families/deepseek_v41/deepseek_v41.zig @@ -22,6 +22,7 @@ test { _ = named; _ = prepare; _ = @import("host_test.zig"); + _ = @import("phases.zig"); // the phase timers // prepared weight folders (TF_DSV41_PREPARED) _ = @import("prepared.zig"); } diff --git a/zig/src/families/deepseek_v41/phases.zig b/zig/src/families/deepseek_v41/phases.zig new file mode 100644 index 0000000000..bfd86f1caf --- /dev/null +++ b/zig/src/families/deepseek_v41/phases.zig @@ -0,0 +1,206 @@ +//! TF_DSV41_PHASES (0 | 1 | sync, _OUT=): where every millisecond of a serving round goes, by named phase. + +const std = @import("std"); + +pub const Name = enum { + prefill, + draft, + @"draft.pass", + window, + forward, + @"graph.stage", + @"graph.replay", + @"graph.capture", + @"engram.wait", + prefetch, + candidates, + commit, + /// the speculative pass's launch after a window (TF_DSV41_SPEC_DRAFT; its collect is `draft.pass`) + spec, + /// host time between top-level phases (the lanes' own work, token events, plan link) + host, + /// inside `draft.pass` (TF_DSV41_DRAFT_GRAPHS): the host waiting for the device pass and its copy + @"draft.wait", + /// inside `draft.pass`: the host's chain (Markov bias, choice, confidences) over the gathered candidates + @"draft.chain", + /// inside `draft.pass`, the host step (DRAFT_GRAPHS off): the logits' download, the host candidates, their gather + @"draft.host", +}; + +/// A gap between top-level phases this long or longer is idle (no live stream), not a round's host time. +pub const idle_ns: u64 = 50 * std.time.ns_per_ms; + +const count = std.enums.values(Name).len; + +pub const Mode = enum { off, wall, sync }; + +/// The stream the `sync` mode waits for. +pub const Sync = struct { ctx: *anyopaque, run: *const fn (ctx: *anyopaque) anyerror!void }; + +/// Counters beside the phases (an object without "ms", so phase filters skip it): written into the report. +pub const Extra = struct { ctx: *anyopaque, write: *const fn (ctx: *anyopaque, w: *std.Io.Writer) anyerror!void }; + +pub const Phases = struct { + mode: Mode, + sync: ?Sync = null, + extra: ?Extra = null, + ns: [count]u64 = @splat(0), + n: [count]u64 = @splat(0), + rounds: u64 = 0, + tokens: u64 = 0, + t0: u64, + path: ?[]const u8 = null, + reported: u64 = 0, + /// phases open now (a top-level phase starts at depth 0), and when the last top-level phase ended (0: none yet) + depth: u32 = 0, + last_end: u64 = 0, + + pub fn parseMode(v: []const u8) !Mode { + if (std.mem.eql(u8, v, "0") or v.len == 0) return .off; + if (std.mem.eql(u8, v, "1")) return .wall; + if (std.mem.eql(u8, v, "sync")) return .sync; + return error.BadPhasesMode; + } + + /// From TF_DSV41_PHASES / TF_DSV41_PHASES_OUT; null when off. + pub fn fromEnv(gpa: std.mem.Allocator) !?*Phases { + const mode = try parseMode(if (std.c.getenv("TF_DSV41_PHASES")) |v| std.mem.span(v) else "0"); + if (mode == .off) return null; + const p = try gpa.create(Phases); + p.* = .{ .mode = mode, .t0 = nowNs(), .path = if (std.c.getenv("TF_DSV41_PHASES_OUT")) |v| std.mem.span(v) else null }; + return p; + } + + pub fn reset(p: *Phases) void { + p.ns = @splat(0); + p.n = @splat(0); + p.rounds = 0; + p.tokens = 0; + p.t0 = nowNs(); + p.depth = 0; + p.last_end = 0; + } + + fn settle(p: *Phases) void { + if (p.mode != .sync) return; + const s = p.sync orelse return; + s.run(s.ctx) catch {}; + } + + pub fn add(p: *Phases, name: Name, ns: u64) void { + p.ns[@backingInt(name)] += ns; + p.n[@backingInt(name)] += 1; + } + + /// The report: {name: {ms, n, ms_avg}} plus wall_ms, mode, rounds, tokens. + pub fn write(p: *const Phases, w: *std.Io.Writer) !void { + const wall = @as(f64, @floatFromInt(nowNs() - p.t0)) / 1e6; + try w.print("{{\"mode\": \"{s}\", \"wall_ms\": {d:.2}, \"rounds\": {d}, \"tokens\": {d}", .{ if (p.mode == .sync) "sync" else "1", wall, p.rounds, p.tokens }); + for (0..count) |i| { + if (p.n[i] == 0) continue; + const ms = @as(f64, @floatFromInt(p.ns[i])) / 1e6; + try w.print(", \"{s}\": {{\"ms\": {d:.2}, \"n\": {d}, \"ms_avg\": {d:.3}}}", .{ @tagName(@as(Name, @fromBackingInt(@intCast(i)))), ms, p.n[i], ms / @as(f64, @floatFromInt(p.n[i])) }); + } + if (p.extra) |x| try x.write(x.ctx, w); + try w.writeAll("}\n"); + } + + /// The report into TF_DSV41_PHASES_OUT (when set). + pub fn flush(p: *Phases, gpa: std.mem.Allocator, io: std.Io) void { + const path = p.path orelse return; + var w: std.Io.Writer.Allocating = .init(gpa); + defer w.deinit(); + p.write(&w.writer) catch return; + std.Io.Dir.cwd().writeFile(io, .{ .sub_path = path, .data = w.written() }) catch |e| std.log.err("phases: {s} not written ({t})", .{ path, e }); + p.reported = p.rounds; + } + + /// A verify window's round finished with `kept` tokens committed; the report is rewritten every 64 rounds. + pub fn round(p: *Phases, gpa: std.mem.Allocator, io: std.Io, kept: usize) void { + p.rounds += 1; + p.tokens += kept; + if (p.rounds - p.reported >= 64) p.flush(gpa, io); + } +}; + +/// One phase's timer: `const t = phases.start(p, .window); defer t.stop();` (p null: nothing). +pub const Timer = struct { + p: ?*Phases, + name: Name, + t0: u64, + + pub fn stop(t: Timer) void { + const p = t.p orelse return; + p.settle(); + const now = nowNs(); + p.add(t.name, now -| t.t0); + p.depth -|= 1; + if (p.depth == 0) p.last_end = now; + } +}; + +pub fn start(p: ?*Phases, name: Name) Timer { + const x = p orelse return .{ .p = null, .name = name, .t0 = 0 }; + if (x.depth == 0 and x.last_end != 0) { + const gap = nowNs() -| x.last_end; + if (gap < idle_ns) x.add(.host, gap); + } + x.depth += 1; + x.settle(); + return .{ .p = x, .name = name, .t0 = nowNs() }; +} + +pub fn nowNs() 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)); +} + +test "phases: the reference engine's names and report shape; sync mode syncs at both ends" { + const Count = struct { + n: u32 = 0, + fn run(ctx: *anyopaque) anyerror!void { + const c: *@This() = @ptrCast(@alignCast(ctx)); + c.n += 1; + } + }; + var c: Count = .{}; + var p: Phases = .{ .mode = .sync, .t0 = nowNs(), .sync = .{ .ctx = &c, .run = Count.run } }; + { + const t = start(&p, .@"draft.pass"); + defer t.stop(); + } + const off = start(null, .window); + off.stop(); + try std.testing.expectEqual(@as(u32, 2), c.n); + { + // a second top-level phase: one `host` gap before it; the nested phase opens none + const t = start(&p, .window); + defer t.stop(); + const n = start(&p, .forward); + n.stop(); + } + try std.testing.expectEqual(@as(u64, 1), p.n[@backingInt(Name.host)]); + p.last_end = nowNs() - idle_ns; // an idle gap: not counted + start(&p, .commit).stop(); + try std.testing.expectEqual(@as(u64, 1), p.n[@backingInt(Name.host)]); + p.add(.@"engram.wait", 1_500_000); + p.add(.@"engram.wait", 500_000); + p.rounds = 2; + p.tokens = 7; + var w: std.Io.Writer.Allocating = .init(std.testing.allocator); + defer w.deinit(); + try p.write(&w.writer); + const parsed = try std.json.parseFromSlice(std.json.Value, std.testing.allocator, w.written(), .{}); + defer parsed.deinit(); + const o = parsed.value.object; + try std.testing.expectEqualStrings("sync", o.get("mode").?.string); + try std.testing.expectEqual(@as(i64, 7), o.get("tokens").?.integer); + const ew = o.get("engram.wait").?.object; + try std.testing.expectEqual(@as(i64, 2), ew.get("n").?.integer); + try std.testing.expectEqual(@as(f64, 2.0), ew.get("ms").?.float); + try std.testing.expect(o.get("draft.pass") != null and o.get("window") != null); + try std.testing.expect(o.get("host") != null); // the gap between the two top-level phases above + try std.testing.expectEqual(Mode.off, try Phases.parseMode("0")); + try std.testing.expectError(error.BadPhasesMode, Phases.parseMode("nvtx")); +} From def4859f921f009284c395a0ff30c6675dc2821d Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Sun, 11 Oct 2026 19:32:16 +0700 Subject: [PATCH 033/104] DeepSeek-V4.1 Engram rows from packed per-rank shards on local NVMe, LRU-cached, read ahead of the window. engram_rows.zig serves a window's Engram rows from the rank's packed shards: io_uring where it works, else Linux native AIO (TF_DSV41_ENGRAM_AIO=1, TF_DSV41_ENGRAM_POLL) or a pool of reader threads (TF_DSV41_ENGRAM_QD, 0 reads serially); TF_DSV41_ENGRAM_URING=0 picks the threads even where io_uring works. Reads can be issued ahead and waited for later, bulk reads bypass the LRU cache, and eviction keeps the bits. The tests compare the bf16 table and a packed shard with the reference reader's rows (fixtures/engram-rows-ref.txt) on every read path. --- .../families/deepseek_v41/deepseek_v41.zig | 1 + zig/src/families/deepseek_v41/engram_io.zig | 534 ++++++++++++++++++ zig/src/families/deepseek_v41/engram_rows.zig | 306 ++++++++++ .../deepseek_v41/engram_rows_test.zig | 383 +++++++++++++ .../deepseek_v41/fixtures/engram-rows-ref.txt | 5 + 5 files changed, 1229 insertions(+) create mode 100644 zig/src/families/deepseek_v41/engram_io.zig create mode 100644 zig/src/families/deepseek_v41/engram_rows.zig create mode 100644 zig/src/families/deepseek_v41/engram_rows_test.zig create mode 100644 zig/src/families/deepseek_v41/fixtures/engram-rows-ref.txt diff --git a/zig/src/families/deepseek_v41/deepseek_v41.zig b/zig/src/families/deepseek_v41/deepseek_v41.zig index 9e024fd039..1bacbd5cdc 100644 --- a/zig/src/families/deepseek_v41/deepseek_v41.zig +++ b/zig/src/families/deepseek_v41/deepseek_v41.zig @@ -24,6 +24,7 @@ test { _ = prepare; _ = engram_host; _ = @import("host_test.zig"); + _ = @import("engram_rows_test.zig"); // Engram's shard rows _ = @import("phases.zig"); // the phase timers // prepared weight folders (TF_DSV41_PREPARED) _ = @import("prepared.zig"); diff --git a/zig/src/families/deepseek_v41/engram_io.zig b/zig/src/families/deepseek_v41/engram_io.zig new file mode 100644 index 0000000000..60be56ea5a --- /dev/null +++ b/zig/src/families/deepseek_v41/engram_io.zig @@ -0,0 +1,534 @@ +//! Engram shard reads: the batches, LRU tables and reader paths (io_uring, native AIO, a thread pool, serial). + +const std = @import("std"); +const linux = std.os.linux; +const engram_rows = @import("engram_rows.zig"); +const Rows = engram_rows.Rows; +const Header = engram_rows.Header; +const header_bytes = engram_rows.header_bytes; +const sector = engram_rows.sector; +const row_bytes = engram_rows.row_bytes; +const merge_gap = engram_rows.merge_gap; +const page = engram_rows.page; + +/// One extent of a batch: file bytes [off, off + len) into the batch buffer at `at`, its rows `first .. + count`. +const Extent = struct { off: u64, len: u64, at: usize, first: u32, count: u32 }; + +pub const Batch = struct { + table: u32, + /// sorted unique global rows + rows: []u64, + ext: []Extent, + buf: []align(page) u8, + left: u32, + failed: bool = false, + in_use: bool = true, + /// a `rows` call took at least half its rows (a bulk batch is read through once read); `taken`: this call's count + read: bool = false, + taken: u32 = 0, + unread: u8 = 0, // `rows` calls that settled it unread: a TF_DSV41_PREFETCH_AHEAD batch waits through one + + /// The record of `rows[k]` (k found by the caller). + pub fn record(b: *const Batch, k: usize, lo: u64) []const u8 { + var a: usize = 0; + var z: usize = b.ext.len; + while (z - a > 1) { + const m = (a + z) / 2; + if (b.ext[m].first <= k) a = m else z = m; + } + const e = b.ext[a]; + const at = e.at + @as(usize, @intCast(header_bytes + (b.rows[k] - lo) * row_bytes - e.off)); + return b.buf[at..][0..row_bytes]; + } +}; + +const Entry = union(enum) { ready: u32, pending: u32 }; +const none: u32 = std.math.maxInt(u32); + +/// One layer's shard and record cache. +pub const Table = struct { + fd: linux.fd_t, + direct: bool, + h: Header, + size: u64, + map: std.AutoHashMapUnmanaged(u64, Entry) = .empty, + /// cap records, their ids, and the LRU links (head: most recent) + recs: []u8, + ids: []u64, + prev: []u32, + next: []u32, + used: u32 = 0, + head: u32 = none, + tail: u32 = none, + /// keys removed since the map was last rehashed: each left a tombstone that a missing key's probe walks past + dead: u32 = 0, + + pub fn unlink(t: *Table, s: u32) void { + if (t.prev[s] != none) t.next[t.prev[s]] = t.next[s] else t.head = t.next[s]; + if (t.next[s] != none) t.prev[t.next[s]] = t.prev[s] else t.tail = t.prev[s]; + } + + pub fn front(t: *Table, s: u32) void { + t.prev[s] = none; + t.next[s] = t.head; + if (t.head != none) t.prev[t.head] = s; + t.head = s; + if (t.tail == none) t.tail = s; + } + + /// A slot for a new record: an unused one, else the least recent (its id leaves the map). + pub fn take(t: *Table) u32 { + if (t.used < t.ids.len) { + t.used += 1; + return t.used - 1; + } + const s = t.tail; + t.unlink(s); + t.drop(t.ids[s]); + return s; + } + + /// std's tombstones never clear and keys fill 80 %: rehash at a tenth, so a tenth of the map always stays free. + pub fn drop(t: *Table, x: u64) void { + if (!t.map.remove(x)) return; + t.dead += 1; + if (@as(u64, t.dead) * 10 < t.map.capacity()) return; + t.map.rehash(std.hash_map.AutoContext(u64){}); + t.dead = 0; + } +}; + +/// One extent for a pool thread: file bytes [off, off + buf.len) into `buf` (the batch's own buffer), batch `bi`. +const Job = struct { bi: u32, fd: linux.fd_t, size: u64, direct: bool, buf: []u8, off: u64 }; + +/// The reader threads without io_uring; `mu` guards the queue, every batch's `left` / `failed`, and the batch list. +pub const Pool = struct { + mu: std.c.pthread_mutex_t = std.c.PTHREAD_MUTEX_INITIALIZER, + work: std.c.pthread_cond_t = .{}, + done: std.c.pthread_cond_t = .{}, + jobs: std.ArrayList(Job) = .empty, + head: usize = 0, + stop: bool = false, + threads: []std.Thread = &.{}, + + pub fn lock(p: *Pool) void { + _ = std.c.pthread_mutex_lock(&p.mu); + } + + pub fn unlock(p: *Pool) void { + _ = std.c.pthread_mutex_unlock(&p.mu); + } +}; + +/// Linux native AIO (TF_DSV41_ENGRAM_AIO=1): one io_submit a batch, reaped by the waiter; works without io_uring. +pub const Aio = struct { + ctx: u64 = 0, + in_flight: u32 = 0, + poll: bool = false, // TF_DSV41_ENGRAM_POLL=1: a wait spins on io_getevents up to `poll_ns`, then sleeps + + pub const poll_ns: u64 = 20 * std.time.ns_per_ms; + + pub const depth: u32 = 1024; + + const Iocb = extern struct { + data: u64, + key: u32 = 0, + rw_flags: i32 = 0, + opcode: u16 = 0, // IOCB_CMD_PREAD + reqprio: i16 = 0, + fildes: u32, + buf: u64, + nbytes: u64, + offset: i64, + reserved2: u64 = 0, + flags: u32 = 0, + resfd: u32 = 0, + }; + + const Event = extern struct { data: u64, obj: u64, res: i64, res2: i64 }; + + pub fn init() !Aio { + var a: Aio = .{ .poll = pollEnv() }; + const rc = linux.syscall2(.io_setup, depth, @intFromPtr(&a.ctx)); + if (linux.errno(rc) != .SUCCESS) return error.AioRefused; + return a; + } + + pub fn deinit(a: *Aio) void { + _ = linux.syscall1(.io_destroy, a.ctx); + } +}; + +/// TF_DSV41_ENGRAM_POLL: 0 (default) | 1 (with TF_DSV41_ENGRAM_AIO=1). +pub fn pollEnv() bool { + const v = std.c.getenv("TF_DSV41_ENGRAM_POLL") orelse return false; + return std.mem.eql(u8, std.mem.span(v), "1"); +} + +/// TF_DSV41_ENGRAM_AIO: 0 (default) | 1. +pub fn aioEnv() bool { + const v = std.c.getenv("TF_DSV41_ENGRAM_AIO") orelse return false; + return std.mem.eql(u8, std.mem.span(v), "1"); +} + +/// TF_DSV41_ENGRAM_QD: the reader threads without io_uring (default 64; 0: none, reads serial). +pub fn poolThreads() u32 { + const v = std.c.getenv("TF_DSV41_ENGRAM_QD") orelse return 64; + return std.fmt.parseInt(u32, std.mem.span(v), 10) catch 64; +} + +/// Native AIO in place of io_uring or the pool (tests, the probe): every read async through one context. +pub fn startAio(r: *Rows) !void { + if (r.ring) |*g| g.deinit(); + r.ring = null; + r.stopPool(); + if (r.aio == null) r.aio = try Aio.init(); +} + +/// The reader threads (no io_uring): `n` workers waiting for extents. +pub fn startPool(r: *Rows, n: u32) !void { + if (r.pool != null or r.ring != null or n == 0) return error.EngramPool; + const p = try r.gpa.create(Pool); + errdefer r.gpa.destroy(p); + p.* = .{}; + p.threads = try r.gpa.alloc(std.Thread, n); + var made: usize = 0; + errdefer { + p.lock(); + p.stop = true; + p.unlock(); + _ = std.c.pthread_cond_broadcast(&p.work); + for (p.threads[0..made]) |t| t.join(); + r.gpa.free(p.threads); + } + r.pool = p; + errdefer r.pool = null; + for (p.threads) |*t| { + t.* = try std.Thread.spawn(.{ .stack_size = 1 << 20 }, worker, .{r}); + made += 1; + } +} + +/// The reader threads stopped (their queue drained); reads are serial afterwards. +pub fn stopPool(r: *Rows) void { + const p = r.pool orelse return; + p.lock(); + p.stop = true; + p.unlock(); + _ = std.c.pthread_cond_broadcast(&p.work); + for (p.threads) |t| t.join(); + r.gpa.free(p.threads); + p.jobs.deinit(r.gpa); + r.gpa.destroy(p); + r.pool = null; +} + +/// A reader thread: extents off the queue until `stop` (the queue drained first). +pub fn worker(r: *Rows) void { + const p = r.pool.?; + while (true) { + p.lock(); + while (p.head == p.jobs.items.len and !p.stop) _ = std.c.pthread_cond_wait(&p.work, &p.mu); + if (p.head == p.jobs.items.len) { + p.unlock(); + return; + } + const j = p.jobs.items[p.head]; + p.head += 1; + if (p.head == p.jobs.items.len) { + p.jobs.clearRetainingCapacity(); + p.head = 0; + } + p.unlock(); + const ok = if (readAt(j.fd, j.size, j.direct, j.buf, j.off)) true else |_| false; + p.lock(); + const b = &r.batches.items[j.bi]; + if (!ok) b.failed = true; + b.left -= 1; + const last = b.left == 0; + p.unlock(); + if (last) _ = std.c.pthread_cond_broadcast(&p.done); + } +} + +pub fn openTable(gpa: std.mem.Allocator, path: [:0]const u8, cap: u32) !Table { + var direct = true; + var rc = linux.openat(linux.AT.FDCWD, path, .{ .ACCMODE = .RDONLY, .CLOEXEC = true, .DIRECT = true }, 0); + if (linux.errno(rc) != .SUCCESS) { + direct = false; + rc = linux.openat(linux.AT.FDCWD, path, .{ .ACCMODE = .RDONLY, .CLOEXEC = true }, 0); + if (linux.errno(rc) != .SUCCESS) { + std.log.scoped(.dsv41).err("engram: cannot open {s}", .{path}); + return error.FileNotFound; + } + } + const fd: linux.fd_t = @intCast(rc); + errdefer _ = linux.close(fd); + var st: linux.Statx = undefined; + if (linux.errno(linux.statx(fd, "", linux.AT.EMPTY_PATH, .{ .SIZE = true }, &st)) != .SUCCESS) return error.StatFailed; + const hb = try gpa.alignedAlloc(u8, .fromByteUnits(page), page); + defer gpa.free(hb); + try preadAll(fd, hb, 0); + if (!direct) _ = linux.fadvise(fd, 0, page, linux.POSIX_FADV.DONTNEED); + const h = try Header.parse(hb[0..48], st.size); + const t: Table = .{ + .fd = fd, + .direct = direct, + .h = h, + .size = st.size, + .recs = try gpa.alloc(u8, cap * row_bytes), + .ids = try gpa.alloc(u64, cap), + .prev = try gpa.alloc(u32, cap), + .next = try gpa.alloc(u32, cap), + }; + return t; +} + +pub fn closeTable(r: *Rows, t: *Table) void { + _ = linux.close(t.fd); + t.map.deinit(r.gpa); + r.gpa.free(t.recs); + r.gpa.free(t.ids); + r.gpa.free(t.prev); + r.gpa.free(t.next); +} + +/// A batch for sorted unique `rows` of table `ti`, its extents submitted; `urgent`: runs before queued read-ahead. +pub fn start(r: *Rows, ti: u32, rows_in: []const u64, urgent: bool) !u32 { + const t = &r.tables[ti]; + const gpa = r.gpa; + var ext: std.ArrayList(Extent) = .empty; + errdefer ext.deinit(gpa); + var total: usize = 0; + for (rows_in, 0..) |x, k| { + const s = header_bytes + (x - t.h.lo) * row_bytes; + const s0 = s / sector * sector; + const s1 = (s + row_bytes + sector - 1) / sector * sector; + if (ext.items.len > 0) { + const e = &ext.items[ext.items.len - 1]; + if (s0 <= e.off + e.len + merge_gap) { + const end = @max(e.off + e.len, s1); + total += @intCast(end - e.off - e.len); + e.len = end - e.off; + e.count += 1; + continue; + } + } + total = std.mem.alignForward(usize, total, page); + try ext.append(gpa, .{ .off = s0, .len = s1 - s0, .at = total, .first = @intCast(k), .count = 1 }); + total += @intCast(s1 - s0); + } + // extents keep their buffer offsets page aligned (O_DIRECT) + const buf = try gpa.alignedAlloc(u8, .fromByteUnits(page), std.mem.alignForward(usize, total, page)); + errdefer gpa.free(buf); + const owned = try gpa.dupe(u64, rows_in); + errdefer gpa.free(owned); + const b: Batch = .{ .table = ti, .rows = owned, .ext = try ext.toOwnedSlice(gpa), .buf = buf, .left = 0 }; + if (r.pool) |p| { + // the batch's slot and its extents queued under the pool's lock (a worker settles batches.items[bi]) + try t.map.ensureUnusedCapacity(gpa, @intCast(owned.len)); + p.lock(); + const queued = r.queue(p, t, b, urgent); + p.unlock(); + const bi = try queued; + _ = std.c.pthread_cond_broadcast(&p.work); + for (owned) |x| t.map.putAssumeCapacity(x, .{ .pending = bi }); + return bi; + } + const bi = try r.slotFor(b); + for (owned) |x| try t.map.put(gpa, x, .{ .pending = bi }); + const bp = &r.batches.items[bi]; + r.stats.reads += bp.ext.len; + for (bp.ext) |e| r.stats.bytes += e.len; + if (r.ring) |*g| { + for (bp.ext, 0..) |e, k| { + const sqe = g.read(@as(u64, bi) << 32 | k, t.fd, .{ .buffer = bp.buf[e.at..][0..@intCast(e.len)] }, e.off) catch blk: { + // a full submission queue: submit what is queued, take finished reads, then the slot is free + _ = try g.submit(); + try r.reap(0); + break :blk try g.read(@as(u64, bi) << 32 | k, t.fd, .{ .buffer = bp.buf[e.at..][0..@intCast(e.len)] }, e.off); + }; + _ = sqe; + bp.left += 1; + } + _ = try g.submit(); + } else if (r.aio != null) { + try r.aioSubmit(bi, t); + } else { + for (bp.ext) |e| try r.readExtent(t, bp.buf[e.at..][0..@intCast(e.len)], e.off); + } + return bi; +} + +pub fn slotFor(r: *Rows, b: Batch) !u32 { + for (r.batches.items, 0..) |*x, i| if (!x.in_use) { + x.* = b; + return @intCast(i); + }; + try r.batches.append(r.gpa, b); + return @intCast(r.batches.items.len - 1); +} + +/// A batch `b` in a free slot, its extents queued on `p` (locked): at the end, or `urgent` at the head. +pub fn queue(r: *Rows, p: *Pool, t: *const Table, b: Batch, urgent: bool) !u32 { + try p.jobs.ensureUnusedCapacity(r.gpa, b.ext.len); + const bi = try r.slotFor(b); + const bp = &r.batches.items[bi]; + r.stats.reads += bp.ext.len; + const n = bp.ext.len; + const at = if (urgent) p.head else p.jobs.items.len; + if (urgent) { + // one shift of the queued jobs past the head, then this batch's in their place + const old = p.jobs.items.len; + p.jobs.items.len += n; + std.mem.copyBackwards(Job, p.jobs.items[at + n ..], p.jobs.items[at..old]); + } else p.jobs.items.len += n; + for (bp.ext, p.jobs.items[at..][0..n]) |e, *j| { + r.stats.bytes += e.len; + j.* = .{ .bi = bi, .fd = t.fd, .size = t.size, .direct = t.direct, .buf = bp.buf[e.at..][0..@intCast(e.len)], .off = e.off }; + } + bp.left = @intCast(bp.ext.len); + return bi; +} + +pub fn readExtent(r: *Rows, t: *const Table, buf: []u8, off: u64) !void { + _ = r; + return readAt(t.fd, t.size, t.direct, buf, off); +} + +/// Batch `bi`'s extents on the AIO context (all of them, reaping finished reads while the context is full). +pub fn aioSubmit(r: *Rows, bi: u32, t: *const Table) !void { + const a = &r.aio.?; + const bp = &r.batches.items[bi]; + var cbs: [64]Aio.Iocb = undefined; + var ptrs: [64]*Aio.Iocb = undefined; + var k: usize = 0; + while (k < bp.ext.len) { + if (a.in_flight == Aio.depth) try r.reap(1); + const n = @min(bp.ext.len - k, cbs.len, Aio.depth - a.in_flight); + for (cbs[0..n], ptrs[0..n], bp.ext[k..][0..n], 0..) |*cb, *pp, e, j| { + cb.* = .{ .data = @as(u64, bi) << 32 | (k + j), .fildes = @intCast(t.fd), .buf = @intFromPtr(bp.buf[e.at..].ptr), .nbytes = e.len, .offset = @intCast(e.off) }; + pp.* = cb; + } + const rc = linux.syscall3(.io_submit, a.ctx, n, @intFromPtr(&ptrs)); + switch (linux.errno(rc)) { + .SUCCESS => {}, + .AGAIN, .INTR => { + try r.reap(if (a.in_flight > 0) 1 else 0); + continue; + }, + else => return error.EngramReadFailed, + } + bp.left += @intCast(rc); + a.in_flight += @intCast(rc); + k += rc; + } +} + +/// Takes finished AIO reads (waiting for at least `at_least`). +pub fn aioReap(r: *Rows, at_least: u32) !void { + const a = &r.aio.?; + var evs: [64]Aio.Event = undefined; + var need = @min(at_least, a.in_flight); + const zero: linux.timespec = .{ .sec = 0, .nsec = 0 }; + const t0 = if (a.poll and need > 0) nowNs() else 0; + var polling = a.poll and need > 0; + while (true) { + // polling: whatever has landed, at once; else the kernel's wait for `need` + if (polling and nowNs() - t0 > Aio.poll_ns) polling = false; + const rc = if (polling) + linux.syscall5(.io_getevents, a.ctx, 0, evs.len, @intFromPtr(&evs), @intFromPtr(&zero)) + else + linux.syscall5(.io_getevents, a.ctx, need, evs.len, @intFromPtr(&evs), 0); + switch (linux.errno(rc)) { + .SUCCESS => {}, + .INTR => continue, + else => return error.EngramReadFailed, + } + for (evs[0..rc]) |ev| { + const b = &r.batches.items[@intCast(ev.data >> 32)]; + const e = b.ext[@intCast(ev.data & 0xFFFF_FFFF)]; + const t = &r.tables[b.table]; + const want = @min(e.len, t.size - e.off); + if (ev.res < 0 or @as(u64, @intCast(ev.res)) < want) b.failed = true; + if (!t.direct) _ = linux.fadvise(t.fd, @intCast(e.off), @intCast(e.len), linux.POSIX_FADV.DONTNEED); + b.left -= 1; + a.in_flight -= 1; + } + need -|= @intCast(rc); + if (need == 0) return; + if (polling and rc == 0) std.atomic.spinLoopHint(); + } +} + +/// Takes finished reads off the ring (waiting for at least `at_least`). +pub fn reap(r: *Rows, at_least: u32) !void { + if (r.aio != null) return r.aioReap(at_least); + const g = &(r.ring orelse return); + var cqes: [64]linux.io_uring_cqe = undefined; + var need = at_least; + while (true) { + const n = try g.copy_cqes(&cqes, need); + for (cqes[0..n]) |c| { + const b = &r.batches.items[@intCast(c.user_data >> 32)]; + const e = b.ext[@intCast(c.user_data & 0xFFFF_FFFF)]; + const t = &r.tables[b.table]; + const want = @min(e.len, t.size - e.off); + if (c.res < 0 or @as(u64, @intCast(c.res)) < want) b.failed = true; + if (!t.direct) _ = linux.fadvise(t.fd, @intCast(e.off), @intCast(e.len), linux.POSIX_FADV.DONTNEED); + b.left -= 1; + } + need -|= n; + if (need == 0) return; + } +} + +/// Blocks until batch `bi`'s reads are done. +pub fn wait(r: *Rows, bi: u32) !void { + if (r.pool) |p| { + p.lock(); + defer p.unlock(); + while (r.batches.items[bi].left > 0) _ = std.c.pthread_cond_wait(&p.done, &p.mu); + if (r.batches.items[bi].failed) return error.EngramReadFailed; + return; + } + while (r.batches.items[bi].left > 0) try r.reap(1); + if (r.batches.items[bi].failed) return error.EngramReadFailed; +} + +/// An extent's bytes [off, off + buf.len) of a shard of `size` bytes (stops at its end), pages dropped if buffered. +pub fn readAt(fd: linux.fd_t, size: u64, direct: bool, buf: []u8, off: u64) !void { + const want = @min(buf.len, size - off); + var done: usize = 0; + while (done < want) { + const rc = linux.pread(fd, buf[done..].ptr, buf.len - done, @intCast(off + done)); + switch (linux.errno(rc)) { + .SUCCESS => {}, + .INTR => continue, + else => return error.EngramReadFailed, + } + if (rc == 0) return error.EngramShortRead; + done += rc; + } + if (!direct) _ = linux.fadvise(fd, @intCast(off), @intCast(buf.len), linux.POSIX_FADV.DONTNEED); +} + +pub fn preadAll(fd: linux.fd_t, out: []u8, off: u64) !void { + var done: usize = 0; + while (done < out.len) { + const rc = linux.pread(fd, out[done..].ptr, out.len - done, @intCast(off + done)); + switch (linux.errno(rc)) { + .SUCCESS => {}, + .INTR => continue, + else => return error.EngramReadFailed, + } + if (rc == 0) return error.EngramShortRead; + done += rc; + } +} + +pub fn nowNs() 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)); +} diff --git a/zig/src/families/deepseek_v41/engram_rows.zig b/zig/src/families/deepseek_v41/engram_rows.zig new file mode 100644 index 0000000000..7854df9ec0 --- /dev/null +++ b/zig/src/families/deepseek_v41/engram_rows.zig @@ -0,0 +1,306 @@ +//! Engram table rows from packed per-rank shards on local NVMe, through io_uring, AIO or reader threads, LRU-cached. + +const std = @import("std"); +const quiet = @import("builtin").is_test; // a test's stderr is the build runner's channel +const linux = std.os.linux; +const eh = @import("engram_host.zig"); + +pub const magic: u64 = 0x31344E4531565344; // "DSV41EN1"; file `

/engram-l-rof.bin` +pub const header_bytes: u64 = 4096; // 6 little-endian u64: magic, layer, lo, hi, total rows, row bytes +pub const sector: u64 = 512; +pub const row_bytes: u64 = 264; // 256 e4m3 values and 8 UE8M0 scales (one a 32) +pub const values = 256; +pub const merge_gap: u64 = 4096; // holes up to this merged into one extent +pub const page = 4096; +const engram_io = @import("engram_io.zig"); + +pub const Header = struct { + layer: u64, + lo: u64, + hi: u64, + total: u64, + row_bytes: u64, + + /// The header's six words, checked against the file's size. + pub fn parse(b: *const [48]u8, size: u64) !Header { + var w: [6]u64 = undefined; + for (&w, 0..) |*x, i| x.* = std.mem.readInt(u64, b[8 * i ..][0..8], .little); + if (w[0] != magic) return error.NotEngramShard; + const h: Header = .{ .layer = w[1], .lo = w[2], .hi = w[3], .total = w[4], .row_bytes = w[5] }; + if (h.row_bytes != row_bytes or h.hi < h.lo or size != header_bytes + (h.hi - h.lo) * h.row_bytes) return error.BadEngramShard; + return h; + } +}; + +/// torch's float8_e4m3fn -> float32 (c10 fp8e4m3fn_to_fp32_value, bit for bit). +pub fn e4m3(b: u8) f32 { + const w: u32 = @as(u32, b) << 24; + const sign = w & 0x8000_0000; + const nonsign = w & 0x7FFF_FFFF; + var shift: u32 = @clz(nonsign); + shift = if (shift > 4) shift - 4 else 0; + const inf_nan: u32 = @bitCast((@as(i32, @bitCast(nonsign +% 0x0100_0000)) >> 8) & 0x7F80_0000); + const zero: u32 = @bitCast(@as(i32, @bitCast(nonsign -% 1)) >> 31); + const r = sign | ((((nonsign << @intCast(shift)) >> 4) +% ((0x78 -% shift) << 23) | inf_nan) & ~zero); + return @bitCast(r); +} + +/// float32 -> bfloat16 as torch's CPU vector path: round to nearest even, NaN -> 0xFFFF (no packed table holds a NaN). +pub fn bf16(x: f32) u16 { + if (std.math.isNan(x)) return 0xFFFF; + const u: u32 = @bitCast(x); + return @intCast((u +% (0x7FFF + ((u >> 16) & 1))) >> 16); +} + +/// The table [scale << 8 | byte] -> bf16 bits of e4m3(byte) x 2^(scale - 127) (scale 0: x 0). +pub fn fillTable(t: *[65536]u16) void { + for (0..256) |s| { + const mul: f32 = @bitCast(@as(u32, @intCast(s)) << 23); + for (0..256) |b| t[s << 8 | b] = bf16(e4m3(@intCast(b)) * mul); + } +} + +const Batch = engram_io.Batch; +const Table = engram_io.Table; +const Pool = engram_io.Pool; +const Aio = engram_io.Aio; +pub const pollEnv = engram_io.pollEnv; +pub const aioEnv = engram_io.aioEnv; +pub const poolThreads = engram_io.poolThreads; +const nowNs = engram_io.nowNs; + +pub const Stats = struct { hits: u64 = 0, waited: u64 = 0, misses: u64 = 0, issued: u64 = 0, reads: u64 = 0, bytes: u64 = 0, wait_ns: u64 = 0 }; + +pub const Rows = struct { + gpa: std.mem.Allocator, + tables: []Table, + ring: ?linux.IoUring, + /// no io_uring: the reader threads (null: serial reads on the engine thread) + pool: ?*Pool = null, + /// no io_uring, TF_DSV41_ENGRAM_AIO=1: native AIO (in place of the pool) + aio: ?Aio = null, + batches: std.ArrayList(Batch) = .empty, + lut: *[65536]u16, + cap: u32, + stats: Stats = .{}, + /// scratch reused by `rows` (sorted unique ids, their records) + uniq: std.ArrayList(u64) = .empty, + recs: std.ArrayList([*]const u8) = .empty, + want: std.ArrayList(u32) = .empty, + + /// Opens rank `rank` of `world`'s shard of every Engram layer in `layer_ids` under `dir`, `cap` records a layer. + pub fn open(gpa: std.mem.Allocator, dir: []const u8, rank: u32, world: u32, layer_ids: []const u32, cap: u32) !*Rows { + if (cap < 4096) return error.EngramCacheTooSmall; + const r = try gpa.create(Rows); + errdefer gpa.destroy(r); + r.* = .{ .gpa = gpa, .tables = try gpa.alloc(Table, layer_ids.len), .ring = null, .lut = try gpa.create([65536]u16), .cap = cap }; + fillTable(r.lut); + var made: usize = 0; + errdefer { + for (r.tables[0..made]) |*t| r.closeTable(t); + gpa.free(r.tables); + gpa.destroy(r.lut); + } + for (layer_ids, r.tables) |l, *t| { + var nb: [4096]u8 = undefined; + const p = try std.fmt.bufPrint(nb[0 .. nb.len - 1], "{s}/engram-l{d}-r{d}of{d}.bin", .{ dir, l, rank, world }); + nb[p.len] = 0; + const path = nb[0..p.len :0]; + t.* = try openTable(gpa, path, cap); + made += 1; + if (t.h.layer != l) return error.BadEngramShard; + } + // TF_DSV41_ENGRAM_URING=0: the reader threads even where io_uring works (a container that allows it, A/B) + const uring_off = if (std.c.getenv("TF_DSV41_ENGRAM_URING")) |v| std.mem.eql(u8, std.mem.span(v), "0") else false; + r.ring = (if (uring_off) error.PermissionDenied else linux.IoUring.init(256, 0)) catch |e| blk: { + if (aioEnv()) { + if (Aio.init()) |a| { + r.aio = a; + std.log.scoped(.dsv41).info("engram: no io_uring ({t}); rows read by native AIO (TF_DSV41_ENGRAM_AIO, {d} in flight), waits {s}", .{ e, Aio.depth, if (a.poll) "polled (TF_DSV41_ENGRAM_POLL)" else "asleep in the kernel" }); + break :blk null; + } else |ae| if (!quiet) std.log.scoped(.dsv41).warn("engram: TF_DSV41_ENGRAM_AIO=1 but io_setup failed ({t}): the reader threads", .{ae}); + } + const n = poolThreads(); + if (n > 0) { + try r.startPool(n); + std.log.scoped(.dsv41).info("engram: no io_uring ({t}); rows read by {d} threads (TF_DSV41_ENGRAM_QD)", .{ e, n }); + } else if (!quiet) std.log.scoped(.dsv41).warn("engram: no io_uring ({t}) and TF_DSV41_ENGRAM_QD=0; rows are read when a window needs them", .{e}); + break :blk null; + }; + return r; + } + + pub const startAio = engram_io.startAio; + pub const startPool = engram_io.startPool; + pub const stopPool = engram_io.stopPool; + const worker = engram_io.worker; + const openTable = engram_io.openTable; + const closeTable = engram_io.closeTable; + + pub fn close(r: *Rows) void { + // every read in flight lands before its buffer goes + for (r.batches.items, 0..) |*b, i| if (b.in_use) { + r.wait(@intCast(i)) catch {}; + b.read = true; + }; + r.settle(null); + r.stopPool(); + if (r.ring) |*g| g.deinit(); + if (r.aio) |*a| a.deinit(); + for (r.tables) |*t| r.closeTable(t); + r.gpa.free(r.tables); + r.batches.deinit(r.gpa); + r.uniq.deinit(r.gpa); + r.recs.deinit(r.gpa); + r.want.deinit(r.gpa); + r.gpa.destroy(r.lut); + r.gpa.destroy(r); + } + + /// engram_host's interface to these rows. + pub fn source(r: *Rows) eh.Source { + return .{ .ptr = r, .vtable = &.{ .issue = issueFn, .rows = rowsFn } }; + } + + fn issueFn(p: *anyopaque, li: usize, idx: []const i64) anyerror!void { + const r: *Rows = @ptrCast(@alignCast(p)); + return r.issue(li, idx); + } + + fn rowsFn(p: *anyopaque, li: usize, idx: []const i64, dim: usize, out: []u16) anyerror!void { + const r: *Rows = @ptrCast(@alignCast(p)); + return r.rows(li, idx, dim, out); + } + + /// Sorted unique `idx` into r.uniq (checked against the shard's range). + fn unique(r: *Rows, t: *const Table, idx: []const i64) !void { + r.uniq.clearRetainingCapacity(); + try r.uniq.ensureTotalCapacity(r.gpa, idx.len); + for (idx) |i| { + if (i < 0 or @as(u64, @intCast(i)) < t.h.lo or @as(u64, @intCast(i)) >= t.h.hi) return error.EngramRowOutsideShard; + r.uniq.appendAssumeCapacity(@intCast(i)); + } + std.mem.sort(u64, r.uniq.items, {}, std.sort.asc(u64)); + var m: usize = 0; + for (r.uniq.items) |x| if (m == 0 or r.uniq.items[m - 1] != x) { + r.uniq.items[m] = x; + m += 1; + }; + r.uniq.shrinkRetainingCapacity(m); + } + + /// Starts reading layer `li`'s rows `idx` that are neither cached nor in flight; returns at once. + pub fn issue(r: *Rows, li: usize, idx: []const i64) !void { + const t = &r.tables[li]; + try r.unique(t, idx); + var m: usize = 0; + for (r.uniq.items) |x| if (!t.map.contains(x)) { + r.uniq.items[m] = x; + m += 1; + }; + if (m == 0) return; + r.stats.issued += m; + _ = try r.start(@intCast(li), r.uniq.items[0..m], false); + } + const start = engram_io.start; + pub const slotFor = engram_io.slotFor; + pub const queue = engram_io.queue; + pub const readExtent = engram_io.readExtent; + pub const aioSubmit = engram_io.aioSubmit; + pub const aioReap = engram_io.aioReap; + pub const reap = engram_io.reap; + const wait = engram_io.wait; + + /// Finished batches into the caches (bulk read through); `li`: the table whose `rows` call settles (null: none). + fn settle(r: *Rows, li: ?usize) void { + // the pool's workers write `left` / `failed`: settled under its lock (the batches' list stays put) + if (r.pool) |p| p.lock(); + defer if (r.pool) |p| p.unlock(); + for (r.batches.items) |*b| { + if (!b.in_use or b.left > 0) continue; + const t = &r.tables[b.table]; + const keep = !b.failed and b.rows.len <= r.cap / 4; + // a settle on another table never drops it: segments read layer 1 then 14, so the next 14 batch sees both + if (!keep and !b.failed and !b.read and (b.unread < 1 or li != b.table)) { + if (li == b.table) b.unread += 1; + continue; + } + for (b.rows, 0..) |x, k| { + if (!keep) { + t.drop(x); + continue; + } + const s = t.take(); + @memcpy(t.recs[s * row_bytes ..][0..row_bytes], b.record(k, t.h.lo)); + t.ids[s] = x; + t.front(s); + // the entry is this batch's (a row is in one batch at most; evictions only remove ready entries) + t.map.getPtr(x).?.* = .{ .ready = s }; + } + r.gpa.free(b.rows); + r.gpa.free(b.ext); + r.gpa.free(b.buf); + b.in_use = false; + } + } + + /// Layer `li`'s rows `idx` (global row ids this rank owns, any order, duplicates allowed) as bf16 [idx.len, dim]. + pub fn rows(r: *Rows, li: usize, idx: []const i64, dim: usize, out: []u16) !void { + if (dim > values or dim % 32 != 0 or out.len != idx.len * dim) return error.BadEngramRows; + const t = &r.tables[li]; + try r.unique(t, idx); + const u = r.uniq.items; + try r.recs.resize(r.gpa, u.len); + r.want.clearRetainingCapacity(); + // cached records now; in-flight ones after their batches; the rest read as one batch + var miss: std.ArrayList(u64) = .empty; + defer miss.deinit(r.gpa); + for (u, 0..) |x, k| { + if (t.map.get(x)) |e| switch (e) { + .ready => |s| { + r.recs.items[k] = t.recs[s * row_bytes ..].ptr; + t.unlink(s); + t.front(s); + r.stats.hits += 1; + }, + .pending => |bi| { + r.batches.items[bi].taken += 1; + if (std.mem.indexOfScalar(u32, r.want.items, bi) == null) try r.want.append(r.gpa, bi); + r.stats.waited += 1; + }, + } else try miss.append(r.gpa, x); + } + if (miss.items.len > 0) { + r.stats.misses += miss.items.len; + const bi = try r.start(@intCast(li), miss.items, true); + r.batches.items[bi].read = true; + try r.want.append(r.gpa, bi); + } + const t0 = nowNs(); + for (r.want.items) |bi| { + try r.wait(bi); + const b = &r.batches.items[bi]; + if (2 * @as(usize, b.taken) >= b.rows.len) b.read = true; + b.taken = 0; + } + r.stats.wait_ns += nowNs() - t0; + for (u, 0..) |x, k| switch (t.map.get(x).?) { + .ready => {}, + .pending => |bi| { + const b = &r.batches.items[bi]; + r.recs.items[k] = b.record(std.sort.binarySearch(u64, b.rows, x, orderU64).?, t.h.lo).ptr; + }, + }; + // every position: its record's values through the table (scale byte of each 32 values, then the e4m3 byte) + for (idx, 0..) |i, p| { + const rec = r.recs.items[std.sort.binarySearch(u64, u, @as(u64, @intCast(i)), orderU64).?]; + const o = out[p * dim ..][0..dim]; + for (o, 0..) |*v, j| v.* = r.lut[@as(usize, rec[values + j / 32]) << 8 | rec[j]]; + } + r.settle(li); + } +}; + +fn orderU64(a: u64, b: u64) std.math.Order { + return std.math.order(a, b); +} diff --git a/zig/src/families/deepseek_v41/engram_rows_test.zig b/zig/src/families/deepseek_v41/engram_rows_test.zig new file mode 100644 index 0000000000..b750d6143a --- /dev/null +++ b/zig/src/families/deepseek_v41/engram_rows_test.zig @@ -0,0 +1,383 @@ +//! engram_rows.zig against reference Engram host fixtures: bf16 table, packed shard via sync, threads and io_uring. + +const std = @import("std"); +const testing = std.testing; +const er = @import("engram_rows.zig"); +const eh = @import("engram_host.zig"); + +const ref = @embedFile("fixtures/engram-rows-ref.txt"); + +/// engram_fixture.record: byte j of global row r. +fn recordByte(r: u64, j: u64) u8 { + const h = (r * 2654435761 + j * 40503) % (1 << 32); + return if (j < 256) @truncate(h >> 7) else @intCast(100 + (h >> 9) % 40); +} + +const Fx = struct { lo: u64, rows: u64, total: u64, sets: [3]struct { idx: []i64, sha: [32]u8 } }; + +fn parseFx(a: std.mem.Allocator, lut_sha: *[32]u8) !Fx { + var fx: Fx = undefined; + var set: usize = 0; + var it = std.mem.tokenizeScalar(u8, ref, '\n'); + while (it.next()) |line| { + var w = std.mem.tokenizeScalar(u8, line, ' '); + const kind = w.next().?; + if (std.mem.eql(u8, kind, "lut")) { + _ = try std.fmt.hexToBytes(lut_sha, w.next().?); + } else if (std.mem.eql(u8, kind, "shard")) { + fx.lo = try std.fmt.parseInt(u64, w.next().?, 10); + fx.rows = try std.fmt.parseInt(u64, w.next().?, 10); + fx.total = try std.fmt.parseInt(u64, w.next().?, 10); + } else { + var ids: std.ArrayList(i64) = .empty; + var c = std.mem.tokenizeScalar(u8, w.next().?, ','); + while (c.next()) |x| try ids.append(a, try std.fmt.parseInt(i64, x, 10)); + fx.sets[set].idx = try ids.toOwnedSlice(a); + _ = try std.fmt.hexToBytes(&fx.sets[set].sha, w.next().?); + set += 1; + } + } + return fx; +} + +/// Writes the fixture's shard (the packer's layout) as engram-l1-r0of2.bin under `dir`. +fn writeShard(a: std.mem.Allocator, io: std.Io, dir: std.Io.Dir, fx: Fx) !void { + return writeShardAs(a, io, dir, fx, 1); +} + +fn writeShardAs(a: std.mem.Allocator, io: std.Io, dir: std.Io.Dir, fx: Fx, layer: u64) !void { + const n: usize = @intCast(er.header_bytes + fx.rows * er.row_bytes); + const b = try a.alloc(u8, n); + defer a.free(b); + @memset(b, 0); + for ([_]u64{ er.magic, layer, fx.lo, fx.lo + fx.rows, fx.total, er.row_bytes }, 0..) |v, i| std.mem.writeInt(u64, b[8 * i ..][0..8], v, .little); + for (0..fx.rows) |k| for (0..er.row_bytes) |j| { + b[er.header_bytes + k * er.row_bytes + j] = recordByte(fx.lo + k, j); + }; + var nb: [32]u8 = undefined; + try dir.writeFile(io, .{ .sub_path = try std.fmt.bufPrint(&nb, "engram-l{d}-r0of2.bin", .{layer}), .data = b }); +} + +fn sha(out: []const u16) [32]u8 { + var d: [32]u8 = undefined; + std.crypto.hash.sha2.Sha256.hash(std.mem.sliceAsBytes(out), &d, .{}); + return d; +} + +test "engram rows: the bf16 table equals torch's dequant(...).to(bfloat16)" { + var want: [32]u8 = undefined; + var arena = std.heap.ArenaAllocator.init(testing.allocator); + defer arena.deinit(); + _ = try parseFx(arena.allocator(), &want); + const t = try testing.allocator.create([65536]u16); + defer testing.allocator.destroy(t); + er.fillTable(t); + try testing.expectEqual(want, sha(t)); + // spot checks: 1.0 x 2^0, -0.0, the smallest subnormal x 2^-6, scale byte 0 (x 0) + try testing.expectEqual(@as(u16, 0x3F80), t[127 << 8 | 0x38]); + try testing.expectEqual(@as(u16, 0x8000), t[127 << 8 | 0x80]); + try testing.expectEqual(@as(u16, 0x3800), t[121 << 8 | 0x01]); // 2^-9 x 2^-6 + try testing.expectEqual(@as(u16, 0x0000), t[0 << 8 | 0x38]); +} + +/// How a test reads: serially on the calling thread, by the reader threads (no io_uring), or by io_uring. +const Mode = enum { sync, pool, uring, aio, aio_poll }; +const modes = [_]Mode{ .sync, .pool, .uring, .aio, .aio_poll }; + +fn openRows(a: std.mem.Allocator, io: std.Io, tmp: *testing.TmpDir, mode: Mode) !*er.Rows { + const root = try tmp.dir.realPathFileAlloc(io, ".", a); + defer a.free(root); + const r = try er.Rows.open(a, root, 0, 2, &.{1}, 4096); + if (mode == .aio or mode == .aio_poll) { + try r.startAio(); // TF_DSV41_ENGRAM_AIO's reader: native AIO, no threads + r.aio.?.poll = mode == .aio_poll; // + TF_DSV41_ENGRAM_POLL: waits polled + return r; + } + if (r.aio) |*x| { // TF_DSV41_ENGRAM_AIO=1 in the environment: the other modes without it + _ = x; + r.aio.?.deinit(); + r.aio = null; + } + if (mode != .uring) if (r.ring) |*g| { + g.deinit(); + r.ring = null; + }; + if (mode == .pool and r.pool == null) try r.startPool(8); + if (mode == .sync) r.stopPool(); // no io_uring here: open started the pool + return r; +} + +test "engram rows: a packed shard read as the reference ShardRows.rows(...).to(bfloat16), synchronously, by the reader threads and by io_uring" { + const a = testing.allocator; + const io = testing.io; + var arena = std.heap.ArenaAllocator.init(a); + defer arena.deinit(); + var lut_sha: [32]u8 = undefined; + const fx = try parseFx(arena.allocator(), &lut_sha); + var tmp = testing.tmpDir(.{}); + defer tmp.cleanup(); + try writeShard(a, io, tmp.dir, fx); + for (modes) |mode| { + const r = try openRows(a, io, &tmp, mode); + defer r.close(); + try testing.expectEqual(fx.lo, r.tables[0].h.lo); + for (fx.sets) |s| { + const out = try a.alloc(u16, s.idx.len * 256); + defer a.free(out); + try r.rows(0, s.idx, 256, out); + try testing.expectEqual(s.sha, sha(out)); + // again: every record from the cache now, the same bits + const hits = r.stats.hits; + try r.rows(0, s.idx, 256, out); + try testing.expectEqual(s.sha, sha(out)); + try testing.expect(r.stats.hits > hits); + } + try testing.expectError(error.EngramRowOutsideShard, r.rows(0, &.{@intCast(fx.lo - 1)}, 256, try arena.allocator().alloc(u16, 256))); + } +} + +test "engram rows: issued ahead, then waited for; bulk reads past the cache; eviction keeps the bits" { + const a = testing.allocator; + const io = testing.io; + var arena = std.heap.ArenaAllocator.init(a); + defer arena.deinit(); + var lut_sha: [32]u8 = undefined; + const fx = try parseFx(arena.allocator(), &lut_sha); + var tmp = testing.tmpDir(.{}); + defer tmp.cleanup(); + try writeShard(a, io, tmp.dir, fx); + for (modes) |mode| { + const r = try openRows(a, io, &tmp, mode); + defer r.close(); + const s = fx.sets[1]; + try r.issue(0, s.idx); + try testing.expectEqual(@as(u64, s.idx.len), r.stats.issued); + try r.issue(0, s.idx); // nothing new: every row is in flight or cached + try testing.expectEqual(@as(u64, s.idx.len), r.stats.issued); + const out = try a.alloc(u16, s.idx.len * 256); + defer a.free(out); + try r.rows(0, s.idx, 256, out); + try testing.expectEqual(s.sha, sha(out)); + try testing.expectEqual(@as(u64, 0), r.stats.misses); + // every row of the shard once (a prompt's bulk read: past cap / 4, read through), then the sets again + const all = try a.alloc(i64, fx.rows); + defer a.free(all); + for (all, 0..) |*x, k| x.* = @intCast(fx.lo + fx.rows - 1 - k); + const big = try a.alloc(u16, all.len * 256); + defer a.free(big); + try r.rows(0, all, 256, big); + for (all, 0..) |x, k| for (0..256) |j| { + const row: u64 = @intCast(x); + const v = recordByte(row, j); + const sc = recordByte(row, 256 + j / 32); + try testing.expectEqual(r.lut[@as(usize, sc) << 8 | v], big[k * 256 + j]); + }; + for (fx.sets) |q| { + const o = try a.alloc(u16, q.idx.len * 256); + defer a.free(o); + try r.rows(0, q.idx, 256, o); + try testing.expectEqual(q.sha, sha(o)); + } + try testing.expect(r.tables[0].used <= 4096); + } +} + +test "engram rows: a prompt segment's bulk reads issued ahead wait through one unrelated read, then serve it" { + // TF_DSV41_PREFETCH_AHEAD: segment k + 1's read issued before k reads; k must not drop it, k + 1 finds every row + const a = testing.allocator; + const io = testing.io; + var arena = std.heap.ArenaAllocator.init(a); + defer arena.deinit(); + var lut_sha: [32]u8 = undefined; + const fx = try parseFx(arena.allocator(), &lut_sha); + var tmp = testing.tmpDir(.{}); + defer tmp.cleanup(); + try writeShard(a, io, tmp.dir, fx); + const all = try a.alloc(i64, fx.rows); + defer a.free(all); + for (all, 0..) |*x, k| x.* = @intCast(fx.lo + k); + try testing.expect(all.len > 4096 / 4); // a bulk batch: read through, not cached + const big = try a.alloc(u16, all.len * 256); + defer a.free(big); + for (modes) |mode| { + const r = try openRows(a, io, &tmp, mode); + defer r.close(); + try r.issue(0, all); + const s = fx.sets[0]; + const o = try a.alloc(u16, s.idx.len * 256); + defer a.free(o); + try r.rows(0, s.idx, 256, o); // segment k: waits for the rows it shares with the batch, the batch stays + try testing.expectEqual(s.sha, sha(o)); + const reads = r.stats.reads; + try r.rows(0, all, 256, big); // segment k + 1: every row from the batch issued ahead + try testing.expectEqual(@as(u64, 0), r.stats.misses); + try testing.expectEqual(reads, r.stats.reads); + for (all, 0..) |x, k| for (0..256) |j| { + const row: u64 = @intCast(x); + try testing.expectEqual(r.lut[@as(usize, recordByte(row, 256 + j / 32)) << 8 | recordByte(row, j)], big[k * 256 + j]); + }; + // a wrong guess: issued, then two unrelated reads of its table, then it is gone (read again on demand) + try r.issue(0, all); + try r.rows(0, s.idx, 256, o); + try r.rows(0, s.idx, 256, o); + const misses = r.stats.misses; + try r.rows(0, all, 256, big); + try testing.expect(r.stats.misses > misses); + for (r.batches.items) |b| try testing.expect(!b.in_use); + } +} + +test "engram rows: a prompt segment ahead on two tables survives the segment's layer-1 then layer-14 reads" { + // both layers' batches of k + 1 issued, k reads 1 then 14, k + 1 reads 1 (settle) then 14: no miss on either table + const a = testing.allocator; + const io = testing.io; + var arena = std.heap.ArenaAllocator.init(a); + defer arena.deinit(); + var lut_sha: [32]u8 = undefined; + const fx = try parseFx(arena.allocator(), &lut_sha); + var tmp = testing.tmpDir(.{}); + defer tmp.cleanup(); + try writeShardAs(a, io, tmp.dir, fx, 1); + try writeShardAs(a, io, tmp.dir, fx, 14); + const all = try a.alloc(i64, fx.rows); + defer a.free(all); + for (all, 0..) |*x, k| x.* = @intCast(fx.lo + k); + const big = try a.alloc(u16, all.len * 256); + defer a.free(big); + const s = fx.sets[0]; + const o = try a.alloc(u16, s.idx.len * 256); + defer a.free(o); + for (modes) |mode| { + const root = try tmp.dir.realPathFileAlloc(io, ".", a); + defer a.free(root); + const r = try er.Rows.open(a, root, 0, 2, &.{ 1, 14 }, 4096); + defer r.close(); + if (mode != .uring) if (r.ring) |*g| { + g.deinit(); + r.ring = null; + }; + if (mode == .pool and r.pool == null) try r.startPool(8); + if (mode == .sync) r.stopPool(); + try r.issue(0, all); + try r.issue(1, all); + for (0..2) |li| { + try r.rows(li, s.idx, 256, o); // segment k: a few shared rows of each table (waits for the batch) + try testing.expectEqual(s.sha, sha(o)); + } + const reads = r.stats.reads; + for (0..2) |li| { + try r.rows(li, all, 256, big); // segment k + 1: layer 1, then layer 14, all from the batches issued ahead + try testing.expectEqual(@as(u64, 0), r.stats.misses); + try testing.expectEqual(reads, r.stats.reads); + for (all, 0..) |x, k| for (0..256) |j| { + const row: u64 = @intCast(x); + try testing.expectEqual(r.lut[@as(usize, recordByte(row, 256 + j / 32)) << 8 | recordByte(row, j)], big[k * 256 + j]); + }; + } + for (r.batches.items) |b| try testing.expect(!b.in_use); + } +} + +test "engram rows: the reader threads with many batches in flight (issued, missed, waited for out of order)" { + const a = testing.allocator; + const io = testing.io; + var arena = std.heap.ArenaAllocator.init(a); + defer arena.deinit(); + var lut_sha: [32]u8 = undefined; + const fx = try parseFx(arena.allocator(), &lut_sha); + var tmp = testing.tmpDir(.{}); + defer tmp.cleanup(); + try writeShard(a, io, tmp.dir, fx); + const r = try openRows(a, io, &tmp, .pool); + defer r.close(); + var rng = std.Random.DefaultPrng.init(4101); + const rand = rng.random(); + const sets = try a.alloc([]i64, 48); + defer { + for (sets) |q| a.free(q); + a.free(sets); + } + for (sets) |*q| { + q.* = try a.alloc(i64, 1 + rand.uintLessThan(usize, 300)); + for (q.*) |*x| x.* = @intCast(fx.lo + rand.uintLessThan(u64, fx.rows)); + } + // every set issued ahead, then read in reverse (later batches waited for first), each read checked value by value + for (sets) |q| try r.issue(0, q); + var k = sets.len; + while (k > 0) { + k -= 1; + const q = sets[k]; + const out = try a.alloc(u16, q.len * 256); + defer a.free(out); + try r.rows(0, q, 256, out); + for (q, 0..) |x, i| for (0..256) |j| { + const row: u64 = @intCast(x); + try testing.expectEqual(r.lut[@as(usize, recordByte(row, 256 + j / 32)) << 8 | recordByte(row, j)], out[i * 256 + j]); + }; + } + try testing.expect(r.stats.issued > 0); +} + +test "engram rows: shard headers" { + var b: [48]u8 = undefined; + for ([_]u64{ er.magic, 14, 10, 20, 100, er.row_bytes }, 0..) |v, i| std.mem.writeInt(u64, b[8 * i ..][0..8], v, .little); + const h = try er.Header.parse(&b, er.header_bytes + 10 * er.row_bytes); + try testing.expectEqual(@as(u64, 14), h.layer); + try testing.expectError(error.BadEngramShard, er.Header.parse(&b, er.header_bytes + 9 * er.row_bytes)); + b[0] ^= 1; + try testing.expectError(error.NotEngramShard, er.Header.parse(&b, er.header_bytes + 10 * er.row_bytes)); +} + +test "engram lookback: the tail, then the kept rows of the window in flight" { + var buf: [16]u32 = undefined; + // nothing in flight: the next window starts at the slot's position + try testing.expectEqualSlices(u32, &.{ 7, 8, 9 }, eh.lookbackAt(&.{ 7, 8, 9 }, 50, null, 50, 3, &buf).?); + try testing.expect(eh.lookbackAt(&.{ 7, 8, 9 }, 50, null, 51, 3, &buf) == null); + // a window [a, d1, d2] at 50 not kept yet; the next starts at 52: a and d1 were kept + try testing.expectEqualSlices(u32, &.{ 9, 100, 101 }, eh.lookbackAt(&.{ 7, 8, 9 }, 50, &.{ 100, 101, 102 }, 52, 3, &buf).?); + try testing.expectEqualSlices(u32, &.{ 100, 101, 102 }, eh.lookbackAt(&.{ 7, 8, 9 }, 50, &.{ 100, 101, 102 }, 53, 3, &buf).?); + try testing.expect(eh.lookbackAt(&.{ 7, 8, 9 }, 50, &.{ 100, 101, 102 }, 54, 3, &buf) == null); + // the sequence start: fewer than keep_n + try testing.expectEqualSlices(u32, &.{5}, eh.lookbackAt(&.{}, 0, &.{ 5, 6 }, 1, 3, &buf).?); +} + +/// Slots of the record map that hold neither a key nor a tombstone: a lookup that misses probes until it finds one. +fn freeSlots(r: *const er.Rows) usize { + const m = &r.tables[0].map; + const meta = m.metadata orelse return 0; + var n: usize = 0; + for (meta[0..m.capacity()]) |x| n += @intFromBool(x.isFree()); + return n; +} + +test "engram rows: a long run of bulk reads and evictions leaves the record map free slots (no probe of the whole map)" { + // bulk reads past cap / 4 (inserted, then dropped) and evictions, as a long uncached prefill does + const a = testing.allocator; + const io = testing.io; + var tmp = testing.tmpDir(.{}); + defer tmp.cleanup(); + const fx: Fx = .{ .lo = 1000, .rows = 20000, .total = 40000, .sets = undefined }; + try writeShard(a, io, tmp.dir, fx); + const root = try tmp.dir.realPathFileAlloc(io, ".", a); + defer a.free(root); + const r = try er.Rows.open(a, root, 0, 2, &.{1}, 4096); + defer r.close(); + const ids = try a.alloc(i64, 1500); + defer a.free(ids); + const out = try a.alloc(u16, ids.len * 256); + defer a.free(out); + var rng = std.Random.DefaultPrng.init(7); + // the worst sample: free slots against the map's slots (it grows meanwhile) + var worst: [2]usize = .{ 1, 1 }; + for (0..300) |i| { + // a prompt segment's bulk read (past cap / 4: read through, its keys dropped), then a decode window's few + const n: usize = if (i % 2 == 0) ids.len else 64; + for (ids[0..n]) |*x| x.* = @intCast(fx.lo + rng.random().uintLessThan(u64, fx.rows)); + try r.rows(0, ids[0..n], 256, out[0 .. n * 256]); + const all = r.tables[0].map.capacity(); + const free = freeSlots(r); + if (free * worst[1] < worst[0] * all) worst = .{ free, all }; + } + errdefer std.debug.print("record map: {d} free of {d} slots at the worst\n", .{ worst[0], worst[1] }); + try testing.expect(worst[0] * 10 >= worst[1]); +} diff --git a/zig/src/families/deepseek_v41/fixtures/engram-rows-ref.txt b/zig/src/families/deepseek_v41/fixtures/engram-rows-ref.txt new file mode 100644 index 0000000000..d9ff54b326 --- /dev/null +++ b/zig/src/families/deepseek_v41/fixtures/engram-rows-ref.txt @@ -0,0 +1,5 @@ +lut ac4091794caec6f553a226c3857da2756d0b29c4c7b24e02ce61274a3df93185 +shard 1000 3000 9000 +rows 1000,1001,1001,3999,1007,2500,1007 82b5fc83bc4eb5195c2d36925f57ef8af91663ed60af344641924535b8dc25b4 +rows 1000,1977,2954,3931,1908,2885,3862,1839,2816,3793,1770,2747,3724,1701,2678,3655,1632,2609,3586,1563,2540,3517,1494,2471,3448,1425,2402,3379,1356,2333,3310,1287,2264,3241,1218,2195,3172,1149,2126,3103,1080,2057,3034,1011,1988,2965,3942,1919,2896,3873,1850,2827,3804,1781,2758,3735,1712,2689,3666,1643,2620,3597,1574,2551,3528,1505,2482,3459,1436,2413,3390,1367,2344,3321,1298,2275,3252,1229,2206,3183,1160,2137,3114,1091,2068,3045,1022,1999,2976,3953,1930,2907,3884,1861,2838,3815 6a0439b98aad93e03a744b515286f8041b7567bd09e15a2f1e358e00b4950199 +rows 1259,1258,1257,1256,1255,1254,1253,1252,1251,1250,1249,1248,1247,1246,1245,1244,1243,1242,1241,1240,1239,1238,1237,1236,1235,1234,1233,1232,1231,1230,1229,1228,1227,1226,1225,1224,1223,1222,1221,1220,1219,1218,1217,1216,1215,1214,1213,1212,1211,1210,1209,1208,1207,1206,1205,1204,1203,1202,1201,1200 f8ef91955c85ec1ba1caa46f4116ddc94cea8204c72f97670292cbb9d8ce82b5 From c7d29eabe3a5003d2f90e53e9706e64d91f692d7 Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Sun, 11 Oct 2026 22:00:59 +0700 Subject: [PATCH 034/104] DeepSeek-V4.1 CUDA kernel sources: CSA2 attention and top-k, mHC, router, EXL3 experts and linears, glue. zig/kernels/cuda/deepseek_v41 holds the decode kernels as .cu sources. The reference extensions' kernels are copied unchanged (each names its source and blob): CSA2 attention with its split kernel, the top-k, mHC with its parallel tail, the router GEMV with its narrow kernel, the EXL3 routed experts, x3ld, x3gm with its v2 kernel and the EXL3 linear with its decode headers. Ours, each bit-exact with the torch ops it replaces: the window glue and the prefill MoE's pointwise ops (--fmad=false), the decode router glue, x3gm's pass tables and the logits' row top-k. ptxas.sh compiles them in a CUDA image with -Xptxas -v to report spills and stack. Nothing builds them yet. --- zig/kernels/cuda/deepseek_v41/attn_cuda.cu | 541 ++++++++++++++++++ zig/kernels/cuda/deepseek_v41/attn_split.cuh | 538 +++++++++++++++++ zig/kernels/cuda/deepseek_v41/decode.cuh | 166 ++++++ zig/kernels/cuda/deepseek_v41/exl3_experts.cu | 283 +++++++++ .../cuda/deepseek_v41/experts_grouped.cuh | 311 ++++++++++ zig/kernels/cuda/deepseek_v41/glue.cu | 403 +++++++++++++ zig/kernels/cuda/deepseek_v41/linear.cu | 363 ++++++++++++ zig/kernels/cuda/deepseek_v41/mhc_cuda.cu | 539 +++++++++++++++++ zig/kernels/cuda/deepseek_v41/mhc_tail.cuh | 94 +++ zig/kernels/cuda/deepseek_v41/pointwise.cu | 40 ++ zig/kernels/cuda/deepseek_v41/ptxas.sh | 27 + zig/kernels/cuda/deepseek_v41/router_gemv.cu | 521 +++++++++++++++++ zig/kernels/cuda/deepseek_v41/router_glue.cu | 115 ++++ .../cuda/deepseek_v41/router_narrow.cuh | 96 ++++ zig/kernels/cuda/deepseek_v41/topk_cuda.cu | 296 ++++++++++ zig/kernels/cuda/deepseek_v41/topk_keys.cu | 94 +++ zig/kernels/cuda/deepseek_v41/x3gm.cu | 477 +++++++++++++++ zig/kernels/cuda/deepseek_v41/x3gm2.cuh | 256 +++++++++ zig/kernels/cuda/deepseek_v41/x3gm_plan.cu | 104 ++++ zig/kernels/cuda/deepseek_v41/x3ld.cu | 249 ++++++++ 20 files changed, 5513 insertions(+) create mode 100644 zig/kernels/cuda/deepseek_v41/attn_cuda.cu create mode 100644 zig/kernels/cuda/deepseek_v41/attn_split.cuh create mode 100644 zig/kernels/cuda/deepseek_v41/decode.cuh create mode 100644 zig/kernels/cuda/deepseek_v41/exl3_experts.cu create mode 100644 zig/kernels/cuda/deepseek_v41/experts_grouped.cuh create mode 100644 zig/kernels/cuda/deepseek_v41/glue.cu create mode 100644 zig/kernels/cuda/deepseek_v41/linear.cu create mode 100644 zig/kernels/cuda/deepseek_v41/mhc_cuda.cu create mode 100644 zig/kernels/cuda/deepseek_v41/mhc_tail.cuh create mode 100644 zig/kernels/cuda/deepseek_v41/pointwise.cu create mode 100755 zig/kernels/cuda/deepseek_v41/ptxas.sh create mode 100644 zig/kernels/cuda/deepseek_v41/router_gemv.cu create mode 100644 zig/kernels/cuda/deepseek_v41/router_glue.cu create mode 100644 zig/kernels/cuda/deepseek_v41/router_narrow.cuh create mode 100644 zig/kernels/cuda/deepseek_v41/topk_cuda.cu create mode 100644 zig/kernels/cuda/deepseek_v41/topk_keys.cu create mode 100644 zig/kernels/cuda/deepseek_v41/x3gm.cu create mode 100644 zig/kernels/cuda/deepseek_v41/x3gm2.cuh create mode 100644 zig/kernels/cuda/deepseek_v41/x3gm_plan.cu create mode 100644 zig/kernels/cuda/deepseek_v41/x3ld.cu diff --git a/zig/kernels/cuda/deepseek_v41/attn_cuda.cu b/zig/kernels/cuda/deepseek_v41/attn_cuda.cu new file mode 100644 index 0000000000..ec43c81ecb --- /dev/null +++ b/zig/kernels/cuda/deepseek_v41/attn_cuda.cu @@ -0,0 +1,541 @@ +// Copy of src/tensorfold/families/deepseek_v41/cuda/csa2/attn_cuda.cu (blob f4be064c7368): sass_check.sh +#define DSV41_ATTN_NO_TORCH + + + +// CSA2 decode attention (chunks + merge) in ONE launch, TF_DSV41_ATTN_CUDA=1: the Triton pair's ops, same bits +#include +#include +#include +#include + +namespace dsv41_attn { + +constexpr int D = 512, CH = 128, NCOMP = 4, NCH = 5, BMQ = 16, WINDOW = 128; // tiles of 32 keys +constexpr int NOPE = 448, VB = 576, SB = 8, NT = 256; +constexpr int SM_Q = 0; // q A fragments [32 steps][32 lanes] uint4 16384 +constexpr int SM_KV = 16384; // 128 rows x 576 B, 8-byte pieces swizzled by key & 7 73728 +constexpr int SM_SC = SM_KV + CH * VB; // 128 x 8 scale bytes 1024 +constexpr int SM_P = SM_SC + CH * SB; // P A fragments [4 tiles][2 steps][32 lanes] uint4 4096 +constexpr int SM_TMAX = SM_P + 4096; // [4 tiles][4 n8][16 rows] fp32 1024 +constexpr int SM_WSUM = SM_TMAX + 1024; // same shape: the warps' row sums 1024 +constexpr int SM_ROW = SM_WSUM + 1024; // [4 tiles][16 rows] nm, alpha, active 768 +constexpr int SM_PHYS = SM_ROW + 768; // 128 physical rows (int32, -1: masked) 512 +constexpr int SM_MISC = SM_PHYS + 512; // final m, l [16] each + the ticket flag 256 +constexpr int SMEM = SM_MISC + 256; // 98,816 B (sm_121: <= 101,376 a block) +static_assert(SMEM <= 101376, "shared memory"); +constexpr unsigned FULL = 0xffffffffu; + +struct Args { + const uint16_t* q; // [R, H, 512] bf16 (RoPE'd) + const uint8_t* cv; const uint8_t* csc; long long cvs, css; // compressed rows (values, scales), row strides + const int* tok; long long ts; const int* cnt; // selection [R, ts] int32 (-1 pad), counts [R] + const uint8_t* sv; const uint8_t* ssc; long long svs, sss; // SWA rings + const int* lo; const int* hi; // [R] window starts; window ends (HAS_HI) or null + const void* pos; // int32 [1] (POS + r) or int64 [R] (ROWS) + const long long* sl; // ROWS: the rows' slots [R] + const int* pt; long long pts; int psh; // page table of the compressed rows (psh 0: contiguous) + const float* sink; const float* cs; long long csst; // [H] fp32, RoPE table fp32 [max_pos, 64] + float* po; float* pm; float* pl; // partials (attn.Scratch layout) + uint16_t* out; // [R, H, 512] bf16 + int* ticket; // [R, H / 16] (self-cleaning) + int R, H, ring, has_comp; + int qrope; // split_kernel: q arrives un-rotated, RoPE'd as it loads +}; + +__device__ __forceinline__ float ex2(float x) { + float y; + asm("ex2.approx.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} +__device__ __forceinline__ float div_full(float a, float b) { + float y; + asm("div.full.f32 %0, %1, %2;" : "=f"(y) : "f"(a), "f"(b)); + return y; +} +__device__ __forceinline__ float texp(float x) { return ex2(__fmul_rn(x, __int_as_float(0x3FB8AA3B))); } +__device__ __forceinline__ float pow2b(uint32_t s) { return __int_as_float(static_cast(s) << 23); } + +__device__ __forceinline__ void mma(float (&d)[4], const uint4& a, uint32_t b0, uint32_t b1) { + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, " + "{%0,%1,%2,%3};" + : "+f"(d[0]), "+f"(d[1]), "+f"(d[2]), "+f"(d[3]) + : "r"(a.x), "r"(a.y), "r"(a.z), "r"(a.w), "r"(b0), "r"(b1)); +} + +// two e4m3 bytes (low byte = the lower k) with their own scales -> bf16x2 (exact: e4m3 x 2^e is a bf16) +__device__ __forceinline__ uint32_t deq2(uint32_t two, float s0, float s1) { + uint32_t h2; + asm("cvt.rn.f16x2.e4m3x2 %0, %1;" : "=r"(h2) : "h"(static_cast(two))); + const __half2 h = *reinterpret_cast(&h2); + const __nv_bfloat162 b = __floats2bfloat162_rn(__fmul_rn(__low2float(h), s0), __fmul_rn(__high2float(h), s1)); + return *reinterpret_cast(&b); +} + +// a byte offset x of row `key` in the staged rows: 8-byte pieces swizzled within groups of 8 (bank-conflict free) +__device__ __forceinline__ int swz(int key, int x) { + const int ch = x >> 3; + return key * VB + ((((ch & ~7) | ((ch & 7) ^ (key & 7)))) << 3) + (x & 7); +} + +// Triton's dot-operand k order: step s (16 k) of a 32-k block takes k = klo(s), klo + 1 and khi(s), khi + 1 +template __device__ __forceinline__ int klo(int s, int t) { + return KW == 4 ? 32 * (s >> 1) + 4 * t + 2 * (s & 1) : 16 * s + 2 * t; +} +template __device__ __forceinline__ int khi(int s, int t) { return klo(s, t) + (KW == 4 ? 16 : 8); } +// where element k (< 32) of a row sits in a 2-step A fragment: step, quad lane, high half (+16 / +8) +template __device__ __forceinline__ void frag_of(int k, int& step, int& tq, int& hi) { + if (KW == 4) { step = (k >> 1) & 1; hi = (k >> 4) & 1; tq = (k & 15) >> 2; } + else { step = k >> 4; hi = (k >> 3) & 1; tq = (k & 7) >> 1; } +} + +// the bf16 pair (k, k + 1) of staged row `key` (S's B operand: k along the row) +__device__ __forceinline__ uint32_t pair_k(const uint8_t* sm, int key, int k) { + if (k < NOPE) { + const uint32_t two = *reinterpret_cast(sm + SM_KV + swz(key, k)); + const float s = pow2b(sm[SM_SC + key * SB + (k >> 6)]); + return deq2(two, s, s); + } + return *reinterpret_cast(sm + SM_KV + swz(key, NOPE + 2 * (k - NOPE))); +} + +// value (key, n) of the staged rows as a bf16 bit pattern / its e4m3 byte (PV's B operand: k along the keys) +__device__ __forceinline__ uint32_t pair_n(const uint8_t* sm, int k0, int k1, int n) { + if (n < NOPE) { + const uint32_t two = sm[SM_KV + swz(k0, n)] | (static_cast(sm[SM_KV + swz(k1, n)]) << 8); + return deq2(two, pow2b(sm[SM_SC + k0 * SB + (n >> 6)]), pow2b(sm[SM_SC + k1 * SB + (n >> 6)])); + } + const int x = NOPE + 2 * (n - NOPE); + return *reinterpret_cast(sm + SM_KV + swz(k0, x)) | + (static_cast(*reinterpret_cast(sm + SM_KV + swz(k1, x))) << 16); +} + +__device__ __forceinline__ void cp8(uint32_t dst, const void* src, bool ok) { + asm volatile("cp.async.ca.shared.global [%0], [%1], 8, %2;" :: "r"(dst), "l"(src), "r"(ok ? 8 : 0)); +} + +template __device__ __forceinline__ int row_pos(const Args& a, int r) { + return ROWS ? static_cast(static_cast(a.pos)[r]) : static_cast(a.pos)[0] + r; +} + +// how many chunks of row r have keys (the merge's ticket target; the others are Triton's empty partial) +__device__ __forceinline__ bool chunk_works(const Args& a, int r, int c) { + return c == NCOMP || (a.has_comp && c * CH < a.cnt[r]); +} + +template +__device__ void merge(const Args& a, uint8_t* sm, int r, int hb) { + const int tid = threadIdx.x, H = a.H, R = a.R; + float* ca = reinterpret_cast(sm + SM_TMAX); // [5][16] a, [5][16] b, [16] af, [16] den + float* cb = ca + NCH * BMQ; + float* fa = cb + NCH * BMQ; + float* fd = fa + BMQ; + if (tid < BMQ) { // a head's scalar chain (Triton's _merge, unrolled) + const int h = hb * BMQ + tid; + float m = -INFINITY, l = 0.f; + for (int c = 0; c < NCH; ++c) { + const bool w = chunk_works(a, r, c); + const long long base = (static_cast(c) * R + r) * H + h; + const float cm = w ? __ldcg(a.pm + base) : -INFINITY, cl = w ? __ldcg(a.pl + base) : 0.f; + const bool act = cl > 0.f; + float aa, bb; + if (c == 0) { + const float nm = act ? fmaxf(cm, -INFINITY) : -INFINITY; + bb = act ? texp(__fsub_rn(cm, nm)) : 0.f; + aa = 0.f; + l = __fmaf_rn(bb, cl, 0.f); + m = nm; + } else { + const float nm = act ? fmaxf(m, cm) : m; + aa = act ? (m == -INFINITY ? 0.f : texp(__fsub_rn(m, nm))) : 1.f; + bb = act ? texp(__fsub_rn(cm, nm)) : 0.f; + l = __fmaf_rn(l, aa, __fmul_rn(bb, cl)); + m = nm; + } + ca[c * BMQ + tid] = aa; + cb[c * BMQ + tid] = bb; + } + const float sink = a.sink[h]; + const float top = fmaxf(m, sink); + const float af = m == -INFINITY ? 0.f : texp(__fsub_rn(m, top)); + const float e = sink == -INFINITY ? 0.f : texp(__fsub_rn(sink, top)); + fa[tid] = af; + fd[tid] = __fmaf_rn(l, af, e); + } + __syncthreads(); + const int at = row_pos(a, r); + const float* csr = a.cs + static_cast(at) * a.csst; + for (int u = tid; u < BMQ * (D / 4); u += NT) { // float4 units: head hl, columns d .. d + 3 + const int hl = u / (D / 4), d = (u % (D / 4)) * 4, h = hb * BMQ + hl; + float o[4]; + for (int c = 0; c < NCH; ++c) { + float4 v = make_float4(0.f, 0.f, 0.f, 0.f); + if (chunk_works(a, r, c)) + v = __ldcg(reinterpret_cast(a.po + ((static_cast(c) * R + r) * H + h) * D + + d)); + const float co[4] = {v.x, v.y, v.z, v.w}; + const float aa = ca[c * BMQ + hl], bb = cb[c * BMQ + hl]; + for (int i = 0; i < 4; ++i) + o[i] = c == 0 ? __fmaf_rn(bb, co[i], 0.f) : __fmaf_rn(o[i], aa, __fmul_rn(co[i], bb)); + } + float x[4]; + for (int i = 0; i < 4; ++i) x[i] = div_full(__fmul_rn(o[i], fa[hl]), fd[hl]); + uint16_t ob[4]; + for (int i = 0; i < 4; i += 2) { // pair (d + i, d + i + 1): GPT-J, inverse + const int p = (d + i) >> 1; + const float c = p >= 224 ? csr[p - 224] : 1.f, s = p >= 224 ? csr[32 + p - 224] : 0.f; + const float ne = __fmaf_rn(x[i], c, __fmul_rn(x[i + 1], s)); + const float no = __fmaf_rn(x[i + 1], c, -__fmul_rn(x[i], s)); + const __nv_bfloat16 b0 = __float2bfloat16_rn(ne), b1 = __float2bfloat16_rn(no); + ob[i] = *reinterpret_cast(&b0); + ob[i + 1] = *reinterpret_cast(&b1); + } + uint2 w; + w.x = ob[0] | (static_cast(ob[1]) << 16); + w.y = ob[2] | (static_cast(ob[3]) << 16); + *reinterpret_cast(a.out + (static_cast(r) * H + h) * D + d) = w; + } +} + +template +__global__ void __launch_bounds__(NT, 1) attn_kernel(Args a) { + extern __shared__ __align__(16) uint8_t sm[]; + const int c = blockIdx.x, hb = blockIdx.y, r = blockIdx.z; + const int tid = threadIdx.x, lane = tid & 31, warp = tid >> 5, g = lane >> 2, t = lane & 3; + const int H = a.H, R = a.R; + if (!chunk_works(a, r, c)) return; + int* phys = reinterpret_cast(sm + SM_PHYS); + const long long sl = ROWS ? a.sl[r] : 0; + if (tid < CH) { // the chunk's 128 keys -> physical rows (-1: masked) + int row = -1; + if (c < NCOMP) { + const int idx = c * CH + tid; + if (idx < a.cnt[r]) { + const int tk = a.tok[static_cast(r) * a.ts + idx]; + if (tk >= 0) { + row = tk; + if (a.psh > 0) { + const int* ptr = a.pt + (ROWS ? sl * a.pts : 0); + row = (ptr[tk >> a.psh] << a.psh) + (tk & ((1 << a.psh) - 1)); + } + } + } + } else { + const int hi = HAS_HI ? a.hi[r] : row_pos(a, r); + const int lo = max(a.lo[r], hi - (WINDOW - 1)); + const int p = hi - (WINDOW - 1) + tid; + if (p >= lo && p >= 0) row = (p % a.ring) + (ROWS ? static_cast(sl) * a.ring : 0); + } + phys[tid] = row; + } + { // q -> A fragments (Triton's k order) + const uint16_t* qb = a.q + (static_cast(r) * H + hb * BMQ) * D; + for (int i = tid; i < 32 * 32; i += NT) { + const int s = i >> 5, ln = i & 31, gg = ln >> 2, tt = ln & 3; + const int k0 = klo(s, tt), k1 = khi(s, tt); + uint4 f; + f.x = *reinterpret_cast(qb + gg * D + k0); + f.y = *reinterpret_cast(qb + (gg + 8) * D + k0); + f.z = *reinterpret_cast(qb + gg * D + k1); + f.w = *reinterpret_cast(qb + (gg + 8) * D + k1); + reinterpret_cast(sm + SM_Q)[i] = f; + } + } + __syncthreads(); + { // the rows: 72 value pieces + 1 scale piece a key + const bool comp = c < NCOMP; + const uint8_t* V = comp ? a.cv : a.sv; + const uint8_t* S = comp ? a.csc : a.ssc; + const long long vs = comp ? a.cvs : a.svs, ss = comp ? a.css : a.sss; + const uint32_t base = static_cast(__cvta_generic_to_shared(sm)); + for (int i = tid; i < CH * 73; i += NT) { + const int key = i / 73, pc = i % 73, row = phys[key]; + const bool ok = row >= 0; + const long long rr = ok ? row : 0; + if (pc < 72) cp8(base + SM_KV + swz(key, pc * 8), V + rr * vs + pc * 8, ok); + else cp8(base + SM_SC + key * SB, S + rr * ss, ok); + } + asm volatile("cp.async.commit_group;\ncp.async.wait_group 0;" ::: "memory"); + } + __syncthreads(); + const float NINF = -INFINITY; + const int n8 = warp & 3, ta = warp >> 2; // this warp: S of tiles ta and ta + 2, block n8 + float sv[2][4] = {{0.f, 0.f, 0.f, 0.f}, {0.f, 0.f, 0.f, 0.f}}; + { + const int kA = 32 * ta + 8 * n8 + g, kB = kA + 64; + const uint4* qf = reinterpret_cast(sm + SM_Q); +#pragma unroll 4 + for (int s = 0; s < 32; ++s) { + const uint4 af = qf[s * 32 + lane]; + const int k0 = klo(s, t), k1 = khi(s, t); + mma(sv[0], af, pair_k(sm, kA, k0), pair_k(sm, kA, k1)); + mma(sv[1], af, pair_k(sm, kB, k0), pair_k(sm, kB, k1)); + } + } + float* tmax = reinterpret_cast(sm + SM_TMAX); + float* wsum = reinterpret_cast(sm + SM_WSUM); + float* rowv = reinterpret_cast(sm + SM_ROW); // [tile][row]: nm, alpha, active at 0 / 64 / 128 + bool val[2][2]; + for (int j = 0; j < 2; ++j) { + const int tile = ta + 2 * j, kc = 32 * tile + 8 * n8 + 2 * t; + val[j][0] = phys[kc] >= 0; + val[j][1] = phys[kc + 1] >= 0; + for (int e = 0; e < 4; ++e) + sv[j][e] = val[j][e & 1] ? __fmul_rn(sv[j][e], __int_as_float(0x3D3504F3)) : NINF; + float m0 = fmaxf(sv[j][0], sv[j][1]), m1 = fmaxf(sv[j][2], sv[j][3]); + m0 = fmaxf(m0, __shfl_xor_sync(FULL, m0, 2)); + m0 = fmaxf(m0, __shfl_xor_sync(FULL, m0, 1)); + m1 = fmaxf(m1, __shfl_xor_sync(FULL, m1, 2)); + m1 = fmaxf(m1, __shfl_xor_sync(FULL, m1, 1)); + if (t == 0) { + tmax[(tile * 4 + n8) * 16 + g] = m0; + tmax[(tile * 4 + n8) * 16 + g + 8] = m1; + } + } + __syncthreads(); + if (tid < BMQ) { // the running max over the tiles, a row a thread + float m = NINF; + for (int tile = 0; tile < 4; ++tile) { + const float* tm4 = tmax + tile * 64 + tid; + const float tm = fmaxf(fmaxf(tm4[0], tm4[16]), fmaxf(tm4[32], tm4[48])); + const bool act = tm != NINF; + const float nm = act ? fmaxf(m, tm) : m; + const float al = act ? (m == NINF ? 0.f : texp(__fsub_rn(m, nm))) : 1.f; + rowv[tile * 16 + tid] = nm; + rowv[64 + tile * 16 + tid] = al; + rowv[128 + tile * 16 + tid] = act ? 1.f : 0.f; + m = nm; + } + reinterpret_cast(sm + SM_MISC)[tid] = m; + } + __syncthreads(); + uint32_t* pf = reinterpret_cast(sm + SM_P); // [tile][step][lane][4] + for (int j = 0; j < 2; ++j) { // p, the thread / quad sums, bf16(p) -> fragments + const int tile = ta + 2 * j; + float p[4]; + for (int e = 0; e < 4; ++e) { + const int row = g + (e >> 1) * 8; + const bool act = rowv[128 + tile * 16 + row] != 0.f; + p[e] = (val[j][e & 1] && act) ? texp(__fsub_rn(sv[j][e], rowv[tile * 16 + row])) : 0.f; + } + float a0 = __fadd_rn(p[0], p[1]), a1 = __fadd_rn(p[2], p[3]); + a0 = __fadd_rn(a0, __shfl_xor_sync(FULL, a0, 2)); + a0 = __fadd_rn(a0, __shfl_xor_sync(FULL, a0, 1)); + a1 = __fadd_rn(a1, __shfl_xor_sync(FULL, a1, 2)); + a1 = __fadd_rn(a1, __shfl_xor_sync(FULL, a1, 1)); + if (t == 0) { + wsum[(tile * 4 + n8) * 16 + g] = a0; + wsum[(tile * 4 + n8) * 16 + g + 8] = a1; + } + int step, tq, hi; + frag_of(8 * n8 + 2 * t, step, tq, hi); + const __nv_bfloat162 lo2 = __floats2bfloat162_rn(p[0], p[1]), hi2 = __floats2bfloat162_rn(p[2], p[3]); + uint32_t* f = pf + ((tile * 2 + step) * 32 + g * 4 + tq) * 4; + f[2 * hi] = *reinterpret_cast(&lo2); // reg 0 / 2: row g; 1 / 3: row g + 8 + f[2 * hi + 1] = *reinterpret_cast(&hi2); + } + __syncthreads(); + if (tid < BMQ) { // l: fma(l, alpha, (W0 + W2) + (W1 + W3)) + float l = 0.f; + for (int tile = 0; tile < 4; ++tile) { + const float* w4 = wsum + tile * 64 + tid; + const float sum = __fadd_rn(__fadd_rn(w4[0], w4[32]), __fadd_rn(w4[16], w4[48])); + l = __fmaf_rn(l, rowv[64 + tile * 16 + tid], sum); + } + reinterpret_cast(sm + SM_MISC)[16 + tid] = l; + } + float acc[8][4]; + for (int j = 0; j < 8; ++j) + for (int e = 0; e < 4; ++e) acc[j][e] = 0.f; + for (int tile = 0; tile < 4; ++tile) { // o = mma(P, kv, rn(o * alpha)), 2 k steps a tile + const float al0 = rowv[64 + tile * 16 + g], al1 = rowv[64 + tile * 16 + g + 8]; + for (int j = 0; j < 8; ++j) { + acc[j][0] = __fmul_rn(acc[j][0], al0); + acc[j][1] = __fmul_rn(acc[j][1], al0); + acc[j][2] = __fmul_rn(acc[j][2], al1); + acc[j][3] = __fmul_rn(acc[j][3], al1); + } + for (int step = 0; step < 2; ++step) { + const uint4 af = reinterpret_cast(pf)[(tile * 2 + step) * 32 + lane]; + const int k0 = 32 * tile + klo(step, t), k1 = 32 * tile + khi(step, t); +#pragma unroll + for (int j = 0; j < 8; ++j) { + const int n = 64 * j + 8 * warp + g; + mma(acc[j], af, pair_n(sm, k0, k0 + 1, n), pair_n(sm, k1, k1 + 1, n)); + } + } + } + const long long pbase = (static_cast(c) * R + r) * H + hb * BMQ; + for (int j = 0; j < 8; ++j) { + const int n = 64 * j + 8 * warp + 2 * t; + *reinterpret_cast(a.po + (pbase + g) * D + n) = make_float2(acc[j][0], acc[j][1]); + *reinterpret_cast(a.po + (pbase + g + 8) * D + n) = make_float2(acc[j][2], acc[j][3]); + } + if (tid < BMQ) { + a.pm[pbase + tid] = reinterpret_cast(sm + SM_MISC)[tid]; + a.pl[pbase + tid] = reinterpret_cast(sm + SM_MISC)[16 + tid]; + } + int* flag = reinterpret_cast(sm + SM_MISC) + 32; + __threadfence(); + __syncthreads(); + if (tid == 0) { + int need = 0; + for (int cc = 0; cc < NCH; ++cc) need += chunk_works(a, r, cc) ? 1 : 0; + *flag = atomicAdd(a.ticket + r * (H / BMQ) + hb, 1) == need - 1; + } + __syncthreads(); + if (!*flag) return; + __threadfence(); + if (tid == 0) atomicExch(a.ticket + r * (H / BMQ) + hb, 0); // self-cleaning: the next launch / replay + merge(a, sm, r, hb); +} + +template +cudaError_t launch_one(const Args& a, cudaStream_t stream) { + static bool init = false; + if (!init) { + const cudaError_t e = cudaFuncSetAttribute(attn_kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, SMEM); + if (e != cudaSuccess) return e; + init = true; + } + attn_kernel<<>>(a); + return cudaGetLastError(); +} + +#include "attn_split.cuh" + +template +cudaError_t launch_split(const Args& a, cudaStream_t stream) { + using G = Split

; + static bool init = false; + if (!init) { + const cudaError_t e = cudaFuncSetAttribute(split_kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, G::SMEM); + if (e != cudaSuccess) return e; + init = true; + } + cudaLaunchConfig_t cfg = {}; + cfg.gridDim = dim3(NCH * P, a.H / BMQ, a.R); + cfg.blockDim = dim3(NT); + cfg.dynamicSmemBytes = G::SMEM; + cfg.stream = stream; + cudaLaunchAttribute at[1]; + at[0].id = cudaLaunchAttributeClusterDimension; + at[0].val.clusterDim.x = P; + at[0].val.clusterDim.y = 1; + at[0].val.clusterDim.z = 1; + cfg.attrs = at; + cfg.numAttrs = P > 1 ? 1 : 0; + return cudaLaunchKernelEx(&cfg, split_kernel, a); +} + +template inline cudaError_t dispatch_split(const Args& a, int kw, bool rows, bool has_hi, cudaStream_t s) { + switch ((kw == 4 ? 4 : 0) + (rows ? 2 : 0) + (has_hi ? 1 : 0)) { + case 0: return launch_split<2, false, false, P>(a, s); + case 1: return launch_split<2, false, true, P>(a, s); + case 2: return launch_split<2, true, false, P>(a, s); + case 3: return launch_split<2, true, true, P>(a, s); + case 4: return launch_split<4, false, false, P>(a, s); + case 5: return launch_split<4, false, true, P>(a, s); + case 6: return launch_split<4, true, false, P>(a, s); + default: return launch_split<4, true, true, P>(a, s); + } +} + +// split: 0 attn_kernel (the default launch), 1 / 4 split_kernel with that many CTAs a chunk (the same bits) +inline cudaError_t dispatch(const Args& a, int kw, bool rows, bool has_hi, cudaStream_t stream, int split = 0) { + if (a.R < 1 || a.H % BMQ || a.ring < WINDOW || (a.ring & (a.ring - 1))) return cudaErrorInvalidValue; + if (split == 1) return dispatch_split<1>(a, kw, rows, has_hi, stream); + if (split == 4) return dispatch_split<4>(a, kw, rows, has_hi, stream); + if (split != 0) return cudaErrorInvalidValue; + switch ((kw == 4 ? 4 : 0) + (rows ? 2 : 0) + (has_hi ? 1 : 0)) { + case 0: return launch_one<2, false, false>(a, stream); + case 1: return launch_one<2, false, true>(a, stream); + case 2: return launch_one<2, true, false>(a, stream); + case 3: return launch_one<2, true, true>(a, stream); + case 4: return launch_one<4, false, false>(a, stream); + case 5: return launch_one<4, false, true>(a, stream); + case 6: return launch_one<4, true, false>(a, stream); + default: return launch_one<4, true, true>(a, stream); + } +} + +} // namespace dsv41_attn + +#ifndef DSV41_ATTN_NO_TORCH // the compile test builds the kernels alone (nvcc, no torch headers) +#include +#include +#include + +namespace { +template T* tp(const at::Tensor& t) { return t.numel() ? static_cast(t.data_ptr()) : nullptr; } +} + +void dsv41_attn_run_cuda(const at::Tensor& q, const at::Tensor& cv, const at::Tensor& csc, const at::Tensor& tok, + const at::Tensor& cnt, const at::Tensor& sv, const at::Tensor& ssc, const at::Tensor& lo, + const at::Tensor& hi, const at::Tensor& pos, const at::Tensor& sl, const at::Tensor& pt, + const at::Tensor& sink, const at::Tensor& cs, at::Tensor& po, at::Tensor& pm, + at::Tensor& pl, at::Tensor& out, at::Tensor& ticket, int64_t ring, int64_t psh, + int64_t pts, int64_t kw, int64_t split, int64_t qrope) { + using namespace dsv41_attn; + Args a{}; + a.q = tp(q); + a.has_comp = cv.numel() ? 1 : 0; + a.cv = tp(cv); a.csc = tp(csc); + a.cvs = cv.numel() ? cv.stride(0) : 0; a.css = csc.numel() ? csc.stride(0) : 0; + a.tok = tp(tok); a.ts = tok.numel() ? tok.stride(0) : 0; a.cnt = tp(cnt); + a.sv = tp(sv); a.ssc = tp(ssc); a.svs = sv.stride(0); a.sss = ssc.stride(0); + a.lo = tp(lo); a.hi = tp(hi); + a.pos = pos.data_ptr(); + a.sl = tp(sl); + a.pt = tp(pt); a.pts = pts; a.psh = static_cast(psh); + a.sink = tp(sink); a.cs = tp(cs); a.csst = cs.stride(0); + a.po = tp(po); a.pm = tp(pm); a.pl = tp(pl); + a.out = tp(out); a.ticket = tp(ticket); + a.R = static_cast(q.size(0)); a.H = static_cast(q.size(1)); a.ring = static_cast(ring); + a.qrope = qrope && split > 0 ? 1 : 0; + TORCH_CHECK(!qrope || split > 0, "attn_cuda: q's RoPE folds into the split kernel only"); + C10_CUDA_CHECK(dispatch(a, static_cast(kw), sl.numel() > 0, hi.numel() > 0, + at::cuda::getCurrentCUDAStream(), static_cast(split))); +} + +std::vector dsv41_attn_info_cuda(int64_t device) { // GB10's limits for the two launches + int sms = 0, optin = 0, cl = 0, blocks = 0; + cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, static_cast(device)); + cudaDeviceGetAttribute(&optin, cudaDevAttrMaxSharedMemoryPerBlockOptin, static_cast(device)); + cudaDeviceGetAttribute(&cl, cudaDevAttrClusterLaunch, static_cast(device)); + cudaFuncSetAttribute(dsv41_attn::attn_kernel<4, true, false>, cudaFuncAttributeMaxDynamicSharedMemorySize, + dsv41_attn::SMEM); + cudaOccupancyMaxActiveBlocksPerMultiprocessor(&blocks, dsv41_attn::attn_kernel<4, true, false>, dsv41_attn::NT, + dsv41_attn::SMEM); + return {sms, optin, cl, blocks, dsv41_attn::SMEM}; +} +#endif + +// The instantiations dispatch() launches: attn_kernel. +template __global__ void dsv41_attn::attn_kernel<2, false, false>(dsv41_attn::Args); +template __global__ void dsv41_attn::attn_kernel<2, false, true>(dsv41_attn::Args); +template __global__ void dsv41_attn::attn_kernel<2, true, false>(dsv41_attn::Args); +template __global__ void dsv41_attn::attn_kernel<2, true, true>(dsv41_attn::Args); +template __global__ void dsv41_attn::attn_kernel<4, false, false>(dsv41_attn::Args); +template __global__ void dsv41_attn::attn_kernel<4, false, true>(dsv41_attn::Args); +template __global__ void dsv41_attn::attn_kernel<4, true, false>(dsv41_attn::Args); +template __global__ void dsv41_attn::attn_kernel<4, true, true>(dsv41_attn::Args); +// dispatch_split's split_kernel at P 1 and 4 (a cluster of 4) +template __global__ void dsv41_attn::split_kernel<2, false, false, 1>(dsv41_attn::Args); +template __global__ void dsv41_attn::split_kernel<2, false, true, 1>(dsv41_attn::Args); +template __global__ void dsv41_attn::split_kernel<2, true, false, 1>(dsv41_attn::Args); +template __global__ void dsv41_attn::split_kernel<2, true, true, 1>(dsv41_attn::Args); +template __global__ void dsv41_attn::split_kernel<4, false, false, 1>(dsv41_attn::Args); +template __global__ void dsv41_attn::split_kernel<4, false, true, 1>(dsv41_attn::Args); +template __global__ void dsv41_attn::split_kernel<4, true, false, 1>(dsv41_attn::Args); +template __global__ void dsv41_attn::split_kernel<4, true, true, 1>(dsv41_attn::Args); +template __global__ void dsv41_attn::split_kernel<2, false, false, 4>(dsv41_attn::Args); +template __global__ void dsv41_attn::split_kernel<2, false, true, 4>(dsv41_attn::Args); +template __global__ void dsv41_attn::split_kernel<2, true, false, 4>(dsv41_attn::Args); +template __global__ void dsv41_attn::split_kernel<2, true, true, 4>(dsv41_attn::Args); +template __global__ void dsv41_attn::split_kernel<4, false, false, 4>(dsv41_attn::Args); +template __global__ void dsv41_attn::split_kernel<4, false, true, 4>(dsv41_attn::Args); +template __global__ void dsv41_attn::split_kernel<4, true, false, 4>(dsv41_attn::Args); +template __global__ void dsv41_attn::split_kernel<4, true, true, 4>(dsv41_attn::Args); diff --git a/zig/kernels/cuda/deepseek_v41/attn_split.cuh b/zig/kernels/cuda/deepseek_v41/attn_split.cuh new file mode 100644 index 0000000000..cd70b1594a --- /dev/null +++ b/zig/kernels/cuda/deepseek_v41/attn_split.cuh @@ -0,0 +1,538 @@ +// attn_cuda.cu's split_kernel (TF_DSV41_ATTN_SPLIT): its helpers and the cluster kernel, inside namespace dsv41_attn +#pragma once +// ---- split_kernel (TF_DSV41_ATTN_SPLIT=P, 1 / 4): a chunk over a CLUSTER of P CTAs, attn_kernel's bits at any P --- +constexpr int MAXP = 4; // the ticket's ranks a (row, head tile) + +__host__ __device__ constexpr int col2byte(int n) { return n <= NOPE ? n : NOPE + 2 * (n - NOPE); } + +template struct Split { + static constexpr int TPC = 4 / P; // S tiles a CTA + static constexpr int KS = 32 * TPC; // S keys a CTA + static constexpr int W = D / P; // PV / merge columns a CTA + static constexpr int J = 8 / P; // n8 blocks a warp in PV + static constexpr int NPW = TPC == 4 ? 2 : 1; // S chains a warp ((tile, n8) pairs: 4 TPC over 8 warps) + static constexpr int SLB = P == 1 ? VB : col2byte(D) - col2byte(D - W); // slice pitch (the widest rank's) + static constexpr int SM_Q = 0; // q A fragments 16384 + static constexpr int SM_KP = 16384; // PV slice: 128 keys x SLB (P = 1: the whole rows) + static constexpr int SM_KS = P == 1 ? SM_KP : SM_KP + CH * SLB; // S rows: KS keys x 576 (P = 1: the slice) + static constexpr int SM_SC = SM_KS + KS * VB; // 128 keys x 8 scale bytes 1024 + static constexpr int SM_KB = SM_SC + CH * SB; // P > 1: S's B operands as bf16 pairs [KS][256] words + static constexpr int SM_P = SM_KB + (P == 1 ? 0 : KS * D * 2); // P A fragments [4][2][32] uint4 4096 + static constexpr int SM_TMAX = SM_P + 4096; // [4 tiles][4 n8][16 rows] fp32 (merge: ca, cb, ...) + static constexpr int SM_WSUM = SM_TMAX + 1024; + static constexpr int SM_ROW = SM_WSUM + 1024; // [4 tiles][16 rows] nm, alpha, active + static constexpr int SM_PHYS = SM_ROW + 768; // 128 physical rows (-1: masked) + static constexpr int SM_MISC = SM_PHYS + 512; // final m, l [16] each + the ticket flag + static constexpr int SMEM = SM_MISC + 256; // P = 1: 98,816 B; P = 4: 100,864 B + static constexpr int SM_VB = SM_KB; // P > 1, after S: PV's B pairs [64 key pairs][W] words + static_assert(P == 1 || (CH / 2) * W * 4 <= KS * D * 2, "PV's pairs fit where S's were"); + static_assert(SMEM <= 101376, "shared memory"); + static_assert(SLB % 64 == 0 && W % 128 == 0, "16-byte swizzle groups; whole n8 blocks a warp"); +}; + +__device__ __forceinline__ int fsw(int key) { return ((key >> 1) ^ (key >> 3)) & 3; } +__device__ __forceinline__ int swz16(int key, int x, int pitch) { + const int c = x >> 4; + return key * pitch + (((c & ~3) | ((c & 3) ^ fsw(key))) << 4) + (x & 15); +} + +// S's B pair (k, k + 1) of the S row `key` (local to the S rows; `ck`: its chunk key, for the scale) +template +__device__ __forceinline__ uint32_t spair_k(const uint8_t* sm, int key, int ck, int k) { + using G = Split

; + if (k < NOPE) { + const uint32_t two = *reinterpret_cast(sm + G::SM_KS + swz16(key, k, VB)); + const float s = pow2b(sm[G::SM_SC + ck * SB + (k >> 6)]); + return deq2(two, s, s); + } + return *reinterpret_cast(sm + G::SM_KS + swz16(key, NOPE + 2 * (k - NOPE), VB)); +} + +// P > 1: S's B operands dequantized once into bf16 pairs, word k / 2 of S row `key` at w ^ bsw(key): conflict-free +template __device__ __forceinline__ int bsw(int key) { + return (KW == 4 ? (key & 1) : ((key & 1) << 2)) | (((key >> 1) & 3) << 3); +} +template +__device__ __forceinline__ uint32_t sbf_k(const uint8_t* sm, int key, int k) { + return reinterpret_cast(sm + Split

::SM_KB)[key * (D / 2) + ((k >> 1) ^ bsw(key))]; +} + +// PV's B pair: value (k0, n), (k1, n) of the slice (byte x0: the rank's first byte) +template +__device__ __forceinline__ uint32_t spair_n(const uint8_t* sm, int k0, int k1, int n, int x0) { + using G = Split

; + if (n < NOPE) { + const int x = n - x0; + const uint32_t two = sm[G::SM_KP + swz16(k0, x, G::SLB)] | + (static_cast(sm[G::SM_KP + swz16(k1, x, G::SLB)]) << 8); + return deq2(two, pow2b(sm[G::SM_SC + k0 * SB + (n >> 6)]), pow2b(sm[G::SM_SC + k1 * SB + (n >> 6)])); + } + const int x = col2byte(n) - x0; + return *reinterpret_cast(sm + G::SM_KP + swz16(k0, x, G::SLB)) | + (static_cast(*reinterpret_cast(sm + G::SM_KP + swz16(k1, x, G::SLB))) << 16); +} + +// P > 1: PV's B pairs dequantized once, word (key pair kp, column x) at x ^ vsw(kp): conflict-free +template __device__ __forceinline__ int vsw(int kp) { return (KW == 4 ? (kp >> 1) & 3 : kp & 3) << 3; } + +// TF_DSV41_ATTN_ROPE: q's rope pair j in place, (rn(e c) - rn(o s), rn(o c) + rn(e s)), no contraction, then bf16 +__device__ __forceinline__ uint32_t rope_word(uint32_t w, const float* csr, int j) { + const float e = __uint_as_float(w << 16), o = __uint_as_float(w & 0xffff0000u); + const float c = csr[j], sn = csr[32 + j]; + const __nv_bfloat162 b = __floats2bfloat162_rn(__fsub_rn(__fmul_rn(e, c), __fmul_rn(o, sn)), + __fadd_rn(__fmul_rn(o, c), __fmul_rn(e, sn))); + return *reinterpret_cast(&b); +} + +__device__ __forceinline__ void cp16(uint32_t dst, const void* src, bool ok) { + asm volatile("cp.async.cg.shared.global [%0], [%1], 16, %2;" :: "r"(dst), "l"(src), "r"(ok ? 16 : 0)); +} +__device__ __forceinline__ void cl_arrive() { asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); } +__device__ __forceinline__ void cl_wait() { asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); } +__device__ __forceinline__ uint32_t cl_rank() { + uint32_t r; + asm volatile("mov.u32 %0, %%cluster_ctarank;" : "=r"(r)); + return r; +} +// store word v at shared address `p` of cluster rank q (DSMEM) +__device__ __forceinline__ void st_peer(void* p, uint32_t q, uint32_t v) { + uint32_t a = static_cast(__cvta_generic_to_shared(p)), b; + asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" : "=r"(b) : "r"(a), "r"(q)); + asm volatile("st.shared::cluster.u32 [%0], %1;" :: "r"(b), "r"(v) : "memory"); +} +// word i of shared array `p` in cluster rank q (DSMEM) +__device__ __forceinline__ uint32_t ld_peer(const void* p, uint32_t q) { + uint32_t a = static_cast(__cvta_generic_to_shared(p)), b, v; + asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" : "=r"(b) : "r"(a), "r"(q)); + asm volatile("ld.shared::cluster.u32 %0, [%1];" : "=r"(v) : "r"(b) : "memory"); + return v; +} + +// stage row bytes [b0, b0 + 16 n) at key slot `key` (16-byte pieces, or 8-byte halves without v16); masked: zeros +__device__ __forceinline__ void stage_piece(uint32_t base, int region, int key, int pitch, int pc, const uint8_t* src, + bool ok, bool v16) { + const uint32_t dst = base + region + swz16(key, pc * 16, pitch); + if (v16) { + cp16(dst, src, ok); + } else { + cp8(dst, src, ok); + cp8(dst + 8, src + 8, ok); + } +} + +// the merge of columns [c0, c0 + NC) of (row r, head tile hb): attn_kernel's merge per element, loads first +template +__device__ void merge_cols(const Args& a, uint8_t* sm, int smt, int r, int hb, int c0) { + const int tid = threadIdx.x, H = a.H, R = a.R; + float* ca = reinterpret_cast(sm + smt); // [5][16] a, [5][16] b, [16] af, [16] den + float* cb = ca + NCH * BMQ; + float* fa = cb + NCH * BMQ; + float* fd = fa + BMQ; + bool cw[NCH]; +#pragma unroll + for (int c = 0; c < NCH; ++c) cw[c] = chunk_works(a, r, c); + constexpr int UB = BMQ * (NC / 4) / NT; // float4 units a thread (8 at P = 1, 2 at P = 4) + static_assert(UB * NT == BMQ * (NC / 4), "whole units"); + float4 pv[UB][NCH]; +#pragma unroll + for (int ub = 0; ub < UB; ++ub) { // the partials in flight first + const int u = tid + ub * NT, hl = u / (NC / 4), d = c0 + (u % (NC / 4)) * 4, h = hb * BMQ + hl; +#pragma unroll + for (int c = 0; c < NCH; ++c) + pv[ub][c] = cw[c] ? __ldcg(reinterpret_cast( + a.po + ((static_cast(c) * R + r) * H + h) * D + d)) + : make_float4(0.f, 0.f, 0.f, 0.f); + } + if (tid < BMQ) { // a head's scalar chain (Triton's _merge, unrolled) + const int h = hb * BMQ + tid; + float pmv[NCH], plv[NCH]; +#pragma unroll + for (int c = 0; c < NCH; ++c) { + const long long base = (static_cast(c) * R + r) * H + h; + pmv[c] = cw[c] ? __ldcg(a.pm + base) : -INFINITY; + plv[c] = cw[c] ? __ldcg(a.pl + base) : 0.f; + } + const float sink = a.sink[h]; + float m = -INFINITY, l = 0.f; +#pragma unroll + for (int c = 0; c < NCH; ++c) { + const float cm = pmv[c], cl = plv[c]; + const bool act = cl > 0.f; + float aa, bb; + if (c == 0) { + const float nm = act ? fmaxf(cm, -INFINITY) : -INFINITY; + bb = act ? texp(__fsub_rn(cm, nm)) : 0.f; + aa = 0.f; + l = __fmaf_rn(bb, cl, 0.f); + m = nm; + } else { + const float nm = act ? fmaxf(m, cm) : m; + aa = act ? (m == -INFINITY ? 0.f : texp(__fsub_rn(m, nm))) : 1.f; + bb = act ? texp(__fsub_rn(cm, nm)) : 0.f; + l = __fmaf_rn(l, aa, __fmul_rn(bb, cl)); + m = nm; + } + ca[c * BMQ + tid] = aa; + cb[c * BMQ + tid] = bb; + } + const float top = fmaxf(m, sink); + const float af = m == -INFINITY ? 0.f : texp(__fsub_rn(m, top)); + const float e = sink == -INFINITY ? 0.f : texp(__fsub_rn(sink, top)); + fa[tid] = af; + fd[tid] = __fmaf_rn(l, af, e); + } + __syncthreads(); + const float* csr = a.cs + static_cast(row_pos(a, r)) * a.csst; +#pragma unroll + for (int ub = 0; ub < UB; ++ub) { + const int u = tid + ub * NT, hl = u / (NC / 4), d = c0 + (u % (NC / 4)) * 4, h = hb * BMQ + hl; + float o[4]; +#pragma unroll + for (int c = 0; c < NCH; ++c) { + const float4 v = pv[ub][c]; + const float co[4] = {v.x, v.y, v.z, v.w}; + const float aa = ca[c * BMQ + hl], bb = cb[c * BMQ + hl]; +#pragma unroll + for (int i = 0; i < 4; ++i) + o[i] = c == 0 ? __fmaf_rn(bb, co[i], 0.f) : __fmaf_rn(o[i], aa, __fmul_rn(co[i], bb)); + } + float x[4]; +#pragma unroll + for (int i = 0; i < 4; ++i) x[i] = div_full(__fmul_rn(o[i], fa[hl]), fd[hl]); + uint16_t ob[4]; +#pragma unroll + for (int i = 0; i < 4; i += 2) { // pair (d + i, d + i + 1): GPT-J, inverse + const int p = (d + i) >> 1; + const float c = p >= 224 ? csr[p - 224] : 1.f, s = p >= 224 ? csr[32 + p - 224] : 0.f; + const float ne = __fmaf_rn(x[i], c, __fmul_rn(x[i + 1], s)); + const float no = __fmaf_rn(x[i + 1], c, -__fmul_rn(x[i], s)); + const __nv_bfloat16 b0 = __float2bfloat16_rn(ne), b1 = __float2bfloat16_rn(no); + ob[i] = *reinterpret_cast(&b0); + ob[i + 1] = *reinterpret_cast(&b1); + } + uint2 w; + w.x = ob[0] | (static_cast(ob[1]) << 16); + w.y = ob[2] | (static_cast(ob[3]) << 16); + *reinterpret_cast(a.out + (static_cast(r) * H + h) * D + d) = w; + } +} + +template +__global__ void __launch_bounds__(NT, 1) split_kernel(Args a) { + using G = Split

; + extern __shared__ __align__(16) uint8_t sm[]; + const int k = P == 1 ? 0 : static_cast(cl_rank()); + const int c = blockIdx.x / P, hb = blockIdx.y, r = blockIdx.z; + const int tid = threadIdx.x, lane = tid & 31, warp = tid >> 5, g = lane >> 2, t = lane & 3; + const int H = a.H, R = a.R; + if (!chunk_works(a, r, c)) return; // the whole cluster: no barrier is pending + uint4 qv[32 * 32 / NT]; // q's A fragments: every load in flight at once + { + const uint16_t* qb = a.q + (static_cast(r) * H + hb * BMQ) * D; +#pragma unroll + for (int u = 0; u < 32 * 32 / NT; ++u) { + const int i = tid + u * NT, s = i >> 5, ln = i & 31, gg = ln >> 2, tt = ln & 3; + const int k0 = klo(s, tt), k1 = khi(s, tt); + qv[u].x = *reinterpret_cast(qb + gg * D + k0); + qv[u].y = *reinterpret_cast(qb + (gg + 8) * D + k0); + qv[u].z = *reinterpret_cast(qb + gg * D + k1); + qv[u].w = *reinterpret_cast(qb + (gg + 8) * D + k1); + } + if (a.qrope) { // the rope dims (k >= 448): their pair's rotation + const float* csr = a.cs + static_cast(row_pos(a, r)) * a.csst; +#pragma unroll + for (int u = 0; u < 32 * 32 / NT; ++u) { + const int i = tid + u * NT, s = i >> 5, tt = (i & 31) & 3; + const int k0 = klo(s, tt), k1 = khi(s, tt); + if (k0 >= NOPE) { + qv[u].x = rope_word(qv[u].x, csr, (k0 - NOPE) >> 1); + qv[u].y = rope_word(qv[u].y, csr, (k0 - NOPE) >> 1); + } + if (k1 >= NOPE) { + qv[u].z = rope_word(qv[u].z, csr, (k1 - NOPE) >> 1); + qv[u].w = rope_word(qv[u].w, csr, (k1 - NOPE) >> 1); + } + } + } + } + int* phys = reinterpret_cast(sm + G::SM_PHYS); + const long long sl = ROWS ? a.sl[r] : 0; + if (tid < CH) { // attn_kernel's row map + int row = -1; + if (c < NCOMP) { + const int idx = c * CH + tid; + if (idx < a.cnt[r]) { + const int tk = a.tok[static_cast(r) * a.ts + idx]; + if (tk >= 0) { + row = tk; + if (a.psh > 0) { + const int* ptr = a.pt + (ROWS ? sl * a.pts : 0); + row = (ptr[tk >> a.psh] << a.psh) + (tk & ((1 << a.psh) - 1)); + } + } + } + } else { + const int hi = HAS_HI ? a.hi[r] : row_pos(a, r); + const int lo = max(a.lo[r], hi - (WINDOW - 1)); + const int p = hi - (WINDOW - 1) + tid; + if (p >= lo && p >= 0) row = (p % a.ring) + (ROWS ? static_cast(sl) * a.ring : 0); + } + phys[tid] = row; + } + __syncthreads(); + { // the rows in flight before q's fragments + const bool comp = c < NCOMP; + const uint8_t* V = comp ? a.cv : a.sv; + const uint8_t* S = comp ? a.csc : a.ssc; + const long long vs = comp ? a.cvs : a.svs, ss = comp ? a.css : a.sss; + const bool v16 = ((reinterpret_cast(V) | static_cast(vs)) & 15) == 0; + const uint32_t base = static_cast(__cvta_generic_to_shared(sm)); + const int x0 = col2byte(G::W * k), nsl = (col2byte(G::W * k + G::W) - x0) / 16; // the slice's pieces a key + for (int i = tid; i < CH * nsl; i += NT) { + const int key = i / nsl, pc = i % nsl, row = phys[key]; + const long long rr = row >= 0 ? row : 0; + stage_piece(base, G::SM_KP, key, G::SLB, pc, V + rr * vs + x0 + pc * 16, row >= 0, v16); + } + if (P > 1) { // the S tiles' whole rows + for (int i = tid; i < G::KS * (VB / 16); i += NT) { + const int key = i / (VB / 16), pc = i % (VB / 16), row = phys[G::KS * k + key]; + const long long rr = row >= 0 ? row : 0; + stage_piece(base, G::SM_KS, key, VB, pc, V + rr * vs + pc * 16, row >= 0, v16); + } + } + for (int key = tid; key < CH; key += NT) { + const int row = phys[key]; + cp8(base + G::SM_SC + key * SB, S + (row >= 0 ? row : 0) * ss, row >= 0); + } + asm volatile("cp.async.commit_group;" ::: "memory"); + } +#pragma unroll + for (int u = 0; u < 32 * 32 / NT; ++u) reinterpret_cast(sm + G::SM_Q)[tid + u * NT] = qv[u]; + asm volatile("cp.async.wait_group 0;" ::: "memory"); + __syncthreads(); + if (P > 1) { // S's operands -> bf16 pairs, 4 k a step (deq2) + uint32_t* kb = reinterpret_cast(sm + G::SM_KB); + for (int i = tid; i < G::KS * (D / 4); i += NT) { + const int key = i / (D / 4), k4 = (i % (D / 4)) * 4, w = k4 >> 1; + uint32_t w0, w1; + if (k4 < NOPE) { + const uint32_t b = *reinterpret_cast(sm + G::SM_KS + swz16(key, k4, VB)); + const float sc = pow2b(sm[G::SM_SC + (key + G::KS * k) * SB + (k4 >> 6)]); + w0 = deq2(b & 0xffffu, sc, sc); + w1 = deq2(b >> 16, sc, sc); + } else { + const uint2 b = *reinterpret_cast(sm + G::SM_KS + swz16(key, NOPE + 2 * (k4 - NOPE), VB)); + w0 = b.x; + w1 = b.y; + } + kb[key * (D / 2) + (w ^ bsw(key))] = w0; + kb[key * (D / 2) + ((w + 1) ^ bsw(key))] = w1; + } + __syncthreads(); + } + const float NINF = -INFINITY; + float* tmax = reinterpret_cast(sm + G::SM_TMAX); + float* wsum = reinterpret_cast(sm + G::SM_WSUM); + float* rowv = reinterpret_cast(sm + G::SM_ROW); + // S: chain pair w (+ 8): local tile w >> 2, block n8 = w & 3 (P = 1: attn_kernel's tiles ta and ta + 2) + const int n8 = warp & 3; + const bool sw = warp < 4 * G::TPC; // this warp has S chains + float sv[G::NPW][4]; + bool val[G::NPW][2]; + int tile_of[G::NPW]; +#pragma unroll + for (int j = 0; j < G::NPW; ++j) { + tile_of[j] = G::TPC * k + (warp >> 2) + 2 * j; + for (int e = 0; e < 4; ++e) sv[j][e] = 0.f; + } + if (sw) { + const uint4* qf = reinterpret_cast(sm + G::SM_Q); + int lk[G::NPW]; +#pragma unroll + for (int j = 0; j < G::NPW; ++j) lk[j] = 32 * (tile_of[j] - G::TPC * k) + 8 * n8 + g; +#pragma unroll 4 + for (int s = 0; s < 32; ++s) { + const uint4 af = qf[s * 32 + lane]; + const int k0 = klo(s, t), k1 = khi(s, t); +#pragma unroll + for (int j = 0; j < G::NPW; ++j) { + if (P == 1) { + mma(sv[j], af, spair_k

(sm, lk[j], lk[j], k0), spair_k

(sm, lk[j], lk[j], k1)); + } else { + mma(sv[j], af, sbf_k(sm, lk[j], k0), sbf_k(sm, lk[j], k1)); + } + } + } +#pragma unroll + for (int j = 0; j < G::NPW; ++j) { + const int tile = tile_of[j], kc = 32 * tile + 8 * n8 + 2 * t; + val[j][0] = phys[kc] >= 0; + val[j][1] = phys[kc + 1] >= 0; + for (int e = 0; e < 4; ++e) + sv[j][e] = val[j][e & 1] ? __fmul_rn(sv[j][e], __int_as_float(0x3D3504F3)) : NINF; + float m0 = fmaxf(sv[j][0], sv[j][1]), m1 = fmaxf(sv[j][2], sv[j][3]); + m0 = fmaxf(m0, __shfl_xor_sync(FULL, m0, 2)); + m0 = fmaxf(m0, __shfl_xor_sync(FULL, m0, 1)); + m1 = fmaxf(m1, __shfl_xor_sync(FULL, m1, 2)); + m1 = fmaxf(m1, __shfl_xor_sync(FULL, m1, 1)); + if (t == 0) { + tmax[(tile * 4 + n8) * 16 + g] = m0; + tmax[(tile * 4 + n8) * 16 + g + 8] = m1; + } + } + } + if (P > 1) { // this rank's tile maxima pushed to the others + __syncthreads(); + for (int i = tid; i < 64 * G::TPC * (P - 1); i += NT) { + const int e = 64 * G::TPC * k + i % (64 * G::TPC), q = (k + 1 + i / (64 * G::TPC)) % P; + st_peer(tmax + e, q, __float_as_uint(tmax[e])); + } + cl_arrive(); + const int x0 = col2byte(G::W * k); // meanwhile: PV's B pairs, the fp8 slice -> bf16, + uint32_t* vb = reinterpret_cast(sm + G::SM_VB); // 4 columns a step (deq2 per column) + for (int i = tid; i < (CH / 2) * (G::W / 4); i += NT) { + const int kp = i / (G::W / 4), x = (i % (G::W / 4)) * 4, n = G::W * k + x, k0 = 2 * kp, k1 = k0 + 1; + uint32_t o[4]; + if (n < NOPE) { + const uint32_t b0 = *reinterpret_cast(sm + G::SM_KP + swz16(k0, n - x0, G::SLB)); + const uint32_t b1 = *reinterpret_cast(sm + G::SM_KP + swz16(k1, n - x0, G::SLB)); + const float s0 = pow2b(sm[G::SM_SC + k0 * SB + (n >> 6)]), s1 = pow2b(sm[G::SM_SC + k1 * SB + (n >> 6)]); +#pragma unroll + for (int e = 0; e < 4; ++e) o[e] = deq2(((b0 >> (8 * e)) & 0xffu) | (((b1 >> (8 * e)) & 0xffu) << 8), s0, s1); + } else { + const uint2 a0 = *reinterpret_cast(sm + G::SM_KP + swz16(k0, col2byte(n) - x0, G::SLB)); + const uint2 a1 = *reinterpret_cast(sm + G::SM_KP + swz16(k1, col2byte(n) - x0, G::SLB)); + o[0] = (a0.x & 0xffffu) | (a1.x << 16); + o[1] = (a0.x >> 16) | (a1.x & 0xffff0000u); + o[2] = (a0.y & 0xffffu) | (a1.y << 16); + o[3] = (a0.y >> 16) | (a1.y & 0xffff0000u); + } +#pragma unroll + for (int e = 0; e < 4; ++e) vb[kp * G::W + ((x + e) ^ vsw(kp))] = o[e]; + } + cl_wait(); + } + __syncthreads(); + if (tid < BMQ) { // the running max over the tiles, a row a thread + float m = NINF; + for (int tile = 0; tile < 4; ++tile) { + const float* tm4 = tmax + tile * 64 + tid; + const float tm = fmaxf(fmaxf(tm4[0], tm4[16]), fmaxf(tm4[32], tm4[48])); + const bool act = tm != NINF; + const float nm = act ? fmaxf(m, tm) : m; + const float al = act ? (m == NINF ? 0.f : texp(__fsub_rn(m, nm))) : 1.f; + rowv[tile * 16 + tid] = nm; + rowv[64 + tile * 16 + tid] = al; + rowv[128 + tile * 16 + tid] = act ? 1.f : 0.f; + m = nm; + } + reinterpret_cast(sm + G::SM_MISC)[tid] = m; + } + __syncthreads(); + uint32_t* pf = reinterpret_cast(sm + G::SM_P); // [tile][step][lane][4] + if (sw) { +#pragma unroll + for (int j = 0; j < G::NPW; ++j) { // p, the quad sums, bf16(p) -> fragments + const int tile = tile_of[j]; + float p[4]; + for (int e = 0; e < 4; ++e) { + const int row = g + (e >> 1) * 8; + const bool act = rowv[128 + tile * 16 + row] != 0.f; + p[e] = (val[j][e & 1] && act) ? texp(__fsub_rn(sv[j][e], rowv[tile * 16 + row])) : 0.f; + } + float a0 = __fadd_rn(p[0], p[1]), a1 = __fadd_rn(p[2], p[3]); + a0 = __fadd_rn(a0, __shfl_xor_sync(FULL, a0, 2)); + a0 = __fadd_rn(a0, __shfl_xor_sync(FULL, a0, 1)); + a1 = __fadd_rn(a1, __shfl_xor_sync(FULL, a1, 2)); + a1 = __fadd_rn(a1, __shfl_xor_sync(FULL, a1, 1)); + if (t == 0) { + wsum[(tile * 4 + n8) * 16 + g] = a0; + wsum[(tile * 4 + n8) * 16 + g + 8] = a1; + } + int step, tq, hi; + frag_of(8 * n8 + 2 * t, step, tq, hi); + const __nv_bfloat162 lo2 = __floats2bfloat162_rn(p[0], p[1]), hi2 = __floats2bfloat162_rn(p[2], p[3]); + uint32_t* f = pf + ((tile * 2 + step) * 32 + g * 4 + tq) * 4; + f[2 * hi] = *reinterpret_cast(&lo2); + f[2 * hi + 1] = *reinterpret_cast(&hi2); + } + } + if (P > 1) { // this rank's row sums and fragments pushed + __syncthreads(); + constexpr int NW = G::TPC * (64 + 256); // words of this rank's tiles + for (int i = tid; i < NW * (P - 1); i += NT) { + const int e = i % NW, q = (k + 1 + i / NW) % P; + uint32_t* w = e < 64 * G::TPC ? reinterpret_cast(wsum) + 64 * G::TPC * k + e + : pf + 256 * G::TPC * k + (e - 64 * G::TPC); + st_peer(w, q, *w); + } + cl_arrive(); // the last cluster access: no barrier before exit + cl_wait(); + } + __syncthreads(); + if (tid < BMQ) { // l: fma(l, alpha, (W0 + W2) + (W1 + W3)) + float l = 0.f; + for (int tile = 0; tile < 4; ++tile) { + const float* w4 = wsum + tile * 64 + tid; + const float sum = __fadd_rn(__fadd_rn(w4[0], w4[32]), __fadd_rn(w4[16], w4[48])); + l = __fmaf_rn(l, rowv[64 + tile * 16 + tid], sum); + } + reinterpret_cast(sm + G::SM_MISC)[16 + tid] = l; + } + const int x0 = col2byte(G::W * k); + float acc[G::J][4]; +#pragma unroll + for (int j = 0; j < G::J; ++j) + for (int e = 0; e < 4; ++e) acc[j][e] = 0.f; + for (int tile = 0; tile < 4; ++tile) { // o = mma(P, kv, rn(o * alpha)), 2 k steps a tile + const float al0 = rowv[64 + tile * 16 + g], al1 = rowv[64 + tile * 16 + g + 8]; +#pragma unroll + for (int j = 0; j < G::J; ++j) { + acc[j][0] = __fmul_rn(acc[j][0], al0); + acc[j][1] = __fmul_rn(acc[j][1], al0); + acc[j][2] = __fmul_rn(acc[j][2], al1); + acc[j][3] = __fmul_rn(acc[j][3], al1); + } + for (int step = 0; step < 2; ++step) { + const uint4 af = reinterpret_cast(pf)[(tile * 2 + step) * 32 + lane]; + const int k0 = 32 * tile + klo(step, t), k1 = 32 * tile + khi(step, t); +#pragma unroll + for (int j = 0; j < G::J; ++j) { + const int n = 64 * (G::J * k + j) + 8 * warp + g; + if (P == 1) { + mma(acc[j], af, spair_n

(sm, k0, k0 + 1, n, x0), spair_n

(sm, k1, k1 + 1, n, x0)); + } else { + const uint32_t* vb = reinterpret_cast(sm + G::SM_VB); + const int x = n - G::W * k; + mma(acc[j], af, vb[(k0 >> 1) * G::W + (x ^ vsw(k0 >> 1))], + vb[(k1 >> 1) * G::W + (x ^ vsw(k1 >> 1))]); + } + } + } + } + const long long pbase = (static_cast(c) * R + r) * H + hb * BMQ; +#pragma unroll + for (int j = 0; j < G::J; ++j) { + const int n = 64 * (G::J * k + j) + 8 * warp + 2 * t; + *reinterpret_cast(a.po + (pbase + g) * D + n) = make_float2(acc[j][0], acc[j][1]); + *reinterpret_cast(a.po + (pbase + g + 8) * D + n) = make_float2(acc[j][2], acc[j][3]); + } + if (tid < BMQ) { // every rank: the same (m, l), so each merge sees it + a.pm[pbase + tid] = reinterpret_cast(sm + G::SM_MISC)[tid]; + a.pl[pbase + tid] = reinterpret_cast(sm + G::SM_MISC)[16 + tid]; + } + int* flag = reinterpret_cast(sm + G::SM_MISC) + 32; + int* tk = a.ticket + (r * (H / BMQ) + hb) * MAXP + k; + __threadfence(); + __syncthreads(); + if (tid == 0) { + int need = 0; + for (int cc = 0; cc < NCH; ++cc) need += chunk_works(a, r, cc) ? 1 : 0; + *flag = atomicAdd(tk, 1) == need - 1; + } + __syncthreads(); + if (!*flag) return; + __threadfence(); + if (tid == 0) atomicExch(tk, 0); // self-cleaning: the next launch / replay + merge_cols(a, sm, G::SM_TMAX, r, hb, G::W * k); +} diff --git a/zig/kernels/cuda/deepseek_v41/decode.cuh b/zig/kernels/cuda/deepseek_v41/decode.cuh new file mode 100644 index 0000000000..2a663ae1c9 --- /dev/null +++ b/zig/kernels/cuda/deepseek_v41/decode.cuh @@ -0,0 +1,166 @@ +// Copy of src/tensorfold/cuda/exl3/decode.cuh (blob 1acca8f0457c): sass_check.sh + +// After ExLlamaV3, MIT, Copyright (c) 2025 Turboderp. + +// EXL3 tiles: lane L decodes a tile's values 8L..8L+7 straight into its two mma.m16n8k16 B fragments, bit for bit. + +#pragma once + +#include +#include + +namespace tf_exl3 { + +enum Codebook : int { CB_3INST = 0, CB_MCG = 1, CB_MUL1 = 2 }; + +// 32-bit words of a tile: 256 values of K2 / 2 bits. +template +__host__ __device__ constexpr int tile_words() { + return 4 * K2; +} + +// E(p): one past the last stream bit of value p's 16-bit window (stream_ends). +template +__host__ __device__ constexpr int stream_end(int p) { + return (K2 & 1) ? ((p + 1) * K2 - ((p + 1) & 1)) / 2 : (p + 1) * (K2 / 2); +} + +// Words a lane reads for its 8 values: the most any lane's windows span (2, or 3 for 3.5 bits and 5 to 8 bits). +template +__host__ __device__ constexpr int lane_words() { + int most = 0; + for (int l = 0; l < 32; ++l) { + const int first = (stream_end(8 * l) - 16 + 128 * K2) % 32; + const int need = first + stream_end(8 * l + 7) - stream_end(8 * l) + 16; + most = need > most ? need : most; + } + return (most + 31) / 32; +} + +// The index of the first word lane `lane` reads, and the bit offset of its first window in that word. +template +__device__ __forceinline__ void lane_start(int lane, int& word, int& offset) { + const int first = 4 * lane * K2 + K2 / 2 - 16 + 128 * K2; // E(8 lane) - 16, made non-negative + word = (first >> 5) % tile_words(); + offset = first & 31; +} + +// This lane's words of a tile (tile: the tile's first 32-bit word; global or shared memory). +template +__device__ __forceinline__ void load_lane_words(const uint32_t* tile, int lane, uint32_t (&w)[lane_words()]) { + int word, offset; + lane_start(lane, word, offset); +#pragma unroll + for (int i = 0; i < lane_words(); ++i) w[i] = tile[(word + i) % tile_words()]; +} + +// The same from global memory through the read-only cache. +template +__device__ __forceinline__ void ldg_lane_words(const uint32_t* tile, int lane, uint32_t (&w)[lane_words()]) { + int word, offset; + lane_start(lane, word, offset); +#pragma unroll + for (int i = 0; i < lane_words(); ++i) w[i] = __ldg(tile + (word + i) % tile_words()); +} + +// Bits [d, d + 16) of the 96-bit big-endian stream hi:mid:lo (d a compile-time constant after unrolling). +__device__ __forceinline__ uint32_t window16(uint32_t hi, uint32_t mid, uint32_t lo, int d) { + if (d <= 16) return (hi >> (16 - d)) & 0xffffu; + if (d < 32) return __funnelshift_l(mid, hi, d) >> 16; + if (d <= 48) return (mid >> (48 - d)) & 0xffffu; + return __funnelshift_l(lo, mid, d - 32) >> 16; +} + +// The 16-bit states of this lane's 8 values (8 lane + j, j = 0..7) from its words. +template +__device__ __forceinline__ void lane_states(const uint32_t (&w)[lane_words()], int lane, uint32_t (&s)[8]) { + int word, offset; + lane_start(lane, word, offset); + const uint32_t c = lane_words() > 2 ? w[lane_words() > 2 ? 2 : 0] : 0u; + const uint32_t hi = __funnelshift_l(w[1], w[0], offset); // the stream from the lane's first window on + const uint32_t mid = __funnelshift_l(c, w[1], offset); + const uint32_t lo = c << offset; +#pragma unroll + for (int j = 0; j < 8; ++j) s[j] = window16(hi, mid, lo, stream_end(j) - stream_end(0)); +} + +// Two codebook values from two 16-bit states, as a half2 in a uint32 (the first state in the low half). +template +__device__ __forceinline__ uint32_t decode2(uint32_t s0, uint32_t s1) { + if constexpr (CB == CB_MUL1) { + const uint32_t x0 = s0 * 0x83DCD12Du, x1 = s1 * 0x83DCD12Du; + const uint32_t h0 = __dp4a(x0, 0x01010101u, 0x6400u); // fp16 bits of 1024 + the byte sum + const uint32_t h1 = __dp4a(x1, 0x01010101u, 0x6400u); + const uint32_t hh = (h0 & 0xffffu) | (h1 << 16); + const half2 r = __hfma2(*reinterpret_cast(&hh), __half2half2(__ushort_as_half(0x1eee)), + __half2half2(__ushort_as_half(0xc931))); + return *reinterpret_cast(&r); + } else { + uint32_t x0, x1; + if constexpr (CB == CB_MCG) { + x0 = s0 * 0xCBAC1FEDu; + x1 = s1 * 0xCBAC1FEDu; + } else { + x0 = s0 * 89226354u + 64248484u; + x1 = s1 * 89226354u + 64248484u; + } + x0 = (x0 & 0x8FFF8FFFu) ^ 0x3B603B60u; + x1 = (x1 & 0x8FFF8FFFu) ^ 0x3B603B60u; + const uint32_t lo = __byte_perm(x0, x1, 0x5410); + const uint32_t hi = __byte_perm(x0, x1, 0x7632); + const half2 r = __hadd2(*reinterpret_cast(&lo), *reinterpret_cast(&hi)); + return *reinterpret_cast(&r); + } +} + +// One codebook value. +template +__device__ __forceinline__ half decode1(uint32_t s) { + const uint32_t v = decode2(s, s); + return __ushort_as_half(static_cast(v & 0xffffu)); +} + +// This lane's 8 values of a tile as the B fragments of the tile's two n8 halves (columns 0-7 in b0, 8-15 in b1). +template +__device__ __forceinline__ void decode_lane(const uint32_t (&w)[lane_words()], int lane, uint32_t (&b0)[2], + uint32_t (&b1)[2]) { + uint32_t s[8]; + lane_states(w, lane, s); + b0[0] = decode2(s[0], s[1]); + b0[1] = decode2(s[2], s[3]); + b1[0] = decode2(s[4], s[5]); + b1[1] = decode2(s[6], s[7]); +} + +// Row and column in the 16x16 tile (row = k, column = n) of this lane's value j (0..7). +__device__ __forceinline__ int value_row(int lane, int j) { return 2 * (lane & 3) + (j & 1) + 8 * ((j >> 1) & 1); } +__device__ __forceinline__ int value_col(int lane, int j) { return (lane >> 2) + 8 * (j >> 2); } + +// mma.m16n8k16, fp16 inputs, fp32 accumulators: d += a @ b. +__device__ __forceinline__ void mma16816(float (&d)[4], const uint32_t (&a)[4], const uint32_t (&b)[2]) { + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 {%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, " + "{%0,%1,%2,%3};\n" + : "+f"(d[0]), "+f"(d[1]), "+f"(d[2]), "+f"(d[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1])); +} + +// Unscaled Walsh-Hadamard transform of 128 fp32 values, 4 a lane, in a fixed butterfly order (multiply by 1/sqrt(128)). +__device__ __forceinline__ void fwht128(float (&v)[4], int lane) { + const float a = v[0] + v[1], b = v[0] - v[1], c = v[2] + v[3], d = v[2] - v[3]; + v[0] = a + c; + v[1] = b + d; + v[2] = a - c; + v[3] = b - d; +#pragma unroll + for (int m = 1; m < 32; m <<= 1) { +#pragma unroll + for (int j = 0; j < 4; ++j) { + const float o = __shfl_xor_sync(0xffffffffu, v[j], m); + v[j] = (lane & m) ? o - v[j] : v[j] + o; + } + } +} + +constexpr float HAD_SCALE = 0.08838834764831845f; // 1 / sqrt(128) + +} // namespace tf_exl3 diff --git a/zig/kernels/cuda/deepseek_v41/exl3_experts.cu b/zig/kernels/cuda/deepseek_v41/exl3_experts.cu new file mode 100644 index 0000000000..a0959e216b --- /dev/null +++ b/zig/kernels/cuda/deepseek_v41/exl3_experts.cu @@ -0,0 +1,283 @@ +// Copy of src/tensorfold/cuda/exl3/experts.cu (blob 8ce66ca83bce): sass_check.sh + +// EXL3 routed experts, a width per expert: fixed-order splits, slots and butterflies, no atomics, so bit-exact. + +#include +#include +#include // what the cut torch includes provided: the fixed-width integers + +#include "experts_grouped.cuh" + +namespace tf_exl3_experts { + +constexpr float HAD_SCALE = 0.08838834764831845f; // 1 / sqrt(128) + +// Grouping in one block: distinct experts (< E) in id order, members row * 32 + slot in row order, -1 after the last. +constexpr int GROUP_THREADS = 1024; +constexpr int GROUP_PER_THREAD = 4; + +__global__ void __launch_bounds__(GROUP_THREADS) group_kernel(const int* __restrict__ pick, int* __restrict__ uids, + int* __restrict__ ucount, int* __restrict__ members, + int R, int slots, int E, int maxm) { + extern __shared__ int sh_pick[]; + __shared__ int warp_tot[GROUP_THREADS / 32]; + const int n = R * slots; + for (int i = threadIdx.x; i < n; i += GROUP_THREADS) sh_pick[i] = pick[i]; + __syncthreads(); + int cnt[GROUP_PER_THREAD]; + int used = 0; +#pragma unroll + for (int q = 0; q < GROUP_PER_THREAD; ++q) { + const int e = threadIdx.x * GROUP_PER_THREAD + q; + int c = 0; + if (e < E) + for (int i = 0; i < n; ++i) c += sh_pick[i] == e; + cnt[q] = c; + used += c > 0; + } + // exclusive scan of `used` over threads + const int lane = threadIdx.x & 31, warp = threadIdx.x >> 5; + int inc = used; +#pragma unroll + for (int o = 1; o < 32; o <<= 1) { + int v = __shfl_up_sync(0xffffffffu, inc, o); + if (lane >= o) inc += v; + } + if (lane == 31) warp_tot[warp] = inc; + __syncthreads(); + if (warp == 0) { + int v = warp_tot[lane]; + int s = v; +#pragma unroll + for (int o = 1; o < 32; o <<= 1) { + int x = __shfl_up_sync(0xffffffffu, s, o); + if (lane >= o) s += x; + } + warp_tot[lane] = s - v; // exclusive per warp + if (lane == 31) ucount[0] = s; + } + __syncthreads(); + int place = warp_tot[warp] + inc - used; +#pragma unroll + for (int q = 0; q < GROUP_PER_THREAD; ++q) { + if (cnt[q] == 0) continue; + const int e = threadIdx.x * GROUP_PER_THREAD + q; + uids[place] = e; + int j = 0; + for (int i = 0; i < n && j < maxm; ++i) + if (sh_pick[i] == e) members[place * maxm + j++] = (i / slots) * 32 + (i % slots); + for (; j < maxm; ++j) members[place * maxm + j] = -1; + ++place; + } +} + +// Walsh-Hadamard of 128 values, 4 a lane, fixed butterfly order (strides 1, 2 in registers, 4..64 across lanes). +__device__ __forceinline__ void fwht128(float (&v)[4], int lane) { + float a = v[0] + v[1], b = v[0] - v[1], c = v[2] + v[3], d = v[2] - v[3]; + v[0] = a + c; v[1] = b + d; v[2] = a - c; v[3] = b - d; +#pragma unroll + for (int m = 1; m < 32; m <<= 1) { +#pragma unroll + for (int j = 0; j < 4; ++j) { + float o = __shfl_xor_sync(0xffffffffu, v[j], m); + v[j] = (lane & m) ? o - v[j] : v[j] + o; + } + } +} + +template __device__ __forceinline__ float to_f(T v); +template <> __device__ __forceinline__ float to_f<__nv_bfloat16>(__nv_bfloat16 v) { return __bfloat162float(v); } +template <> __device__ __forceinline__ float to_f(half v) { return __half2float(v); } + +// Program (member row, 128-block of K, matrix): Xh = fp16((x * suh) @ H) for gate and up of every routed slot. +template +__global__ void rot_in_kernel(const TIN* __restrict__ x, int x_stride, const int* __restrict__ pick, + const half* __restrict__ suh0, const half* __restrict__ suh1, half* __restrict__ out0, + half* __restrict__ out1, int K, int slots, int E) { + const int p = blockIdx.x, blk = blockIdx.y, mat = blockIdx.z; + const int row = p / slots; + const int e = pick[p]; + if (e < 0 || e >= E) return; + const int lane = threadIdx.x; + const half* suh = (mat ? suh1 : suh0) + (size_t)e * K + blk * 128 + 4 * lane; + const TIN* xr = x + (size_t)row * x_stride + blk * 128 + 4 * lane; + float v[4]; +#pragma unroll + for (int j = 0; j < 4; ++j) v[j] = to_f(xr[j]) * __half2float(suh[j]); + fwht128(v, lane); + half* o = (mat ? out1 : out0) + (size_t)p * K + blk * 128 + 4 * lane; +#pragma unroll + for (int j = 0; j < 4; ++j) o[j] = __float2half_rn(v[j] * HAD_SCALE); +} + +__device__ __forceinline__ float bf16r(float x) { return __bfloat162float(__float2bfloat16_rn(x)); } + +// Program (row, 128-block of width): ordered splits, rotate, * svh, SwiGLU (0: bf16 roundings, 1: fp32), act -> Xd. +__global__ void gateup_epilogue_kernel(const float* __restrict__ Z, const int* __restrict__ pick, + const half* __restrict__ svh_g, const half* __restrict__ svh_u, + const half* __restrict__ suh_d, half* __restrict__ xd, int P, int N, int SK, + int E, float limit, int act_mode) { + const int p = blockIdx.x, blk = blockIdx.y; + const int e = pick[p]; + if (e < 0 || e >= E) return; + const int lane = threadIdx.x; + const int n = blk * 128 + 4 * lane; + float gv[4], uv[4]; +#pragma unroll + for (int j = 0; j < 4; ++j) { + float sg = 0.f, su = 0.f; + for (int s = 0; s < SK; ++s) { + sg += Z[((size_t)(0 * SK + s) * P + p) * N + n + j]; + su += Z[((size_t)(1 * SK + s) * P + p) * N + n + j]; + } + gv[j] = sg; + uv[j] = su; + } + fwht128(gv, lane); + fwht128(uv, lane); + float v[4]; +#pragma unroll + for (int j = 0; j < 4; ++j) { + float act; + if (act_mode == 0) { + float gg = fminf(bf16r(gv[j] * HAD_SCALE * __half2float(svh_g[(size_t)e * N + n + j])), limit); + float uu = fminf(fmaxf(bf16r(uv[j] * HAD_SCALE * __half2float(svh_u[(size_t)e * N + n + j])), -limit), + limit); + act = bf16r(bf16r(gg / (1.f + expf(-gg))) * uu); + } else { + float gg = fminf(gv[j] * HAD_SCALE * __half2float(svh_g[(size_t)e * N + n + j]), limit); + float uu = fminf(fmaxf(uv[j] * HAD_SCALE * __half2float(svh_u[(size_t)e * N + n + j]), -limit), limit); + act = gg / (1.f + expf(-gg)) * uu; + } + v[j] = act * __half2float(suh_d[(size_t)e * N + n + j]); + } + fwht128(v, lane); + half* o = xd + (size_t)p * N + n; +#pragma unroll + for (int j = 0; j < 4; ++j) o[j] = __float2half_rn(v[j] * HAD_SCALE); +} + +// Program (member row, 128-block of the model width): Y = (splits summed in order) @ H * svh_d, fp32. +__global__ void down_epilogue_kernel(const float* __restrict__ Z, const int* __restrict__ pick, + const half* __restrict__ svh_d, float* __restrict__ y, int P, int D, int SK, + int E) { + const int p = blockIdx.x, blk = blockIdx.y; + const int e = pick[p]; + if (e < 0 || e >= E) return; + const int lane = threadIdx.x; + const int n = blk * 128 + 4 * lane; + float v[4]; +#pragma unroll + for (int j = 0; j < 4; ++j) { + float s = 0.f; + for (int k = 0; k < SK; ++k) s += Z[((size_t)k * P + p) * D + n + j]; + v[j] = s; + } + fwht128(v, lane); + float* o = y + (size_t)p * D + n; +#pragma unroll + for (int j = 0; j < 4; ++j) o[j] = v[j] * HAD_SCALE * __half2float(svh_d[(size_t)e * D + n + j]); +} + +// out[r][d] = sum over slots in order of wts[r][k] * y[r * slots + k][d] (fp32, fma chain from 0). +__global__ void combine_kernel(const float* __restrict__ y, const float* __restrict__ wts, float* __restrict__ out, + int D, int slots) { + const int r = blockIdx.x; + const int d = blockIdx.y * blockDim.x + threadIdx.x; + if (d >= D) return; + float acc = 0.f; + for (int k = 0; k < slots; ++k) acc = fmaf(wts[r * slots + k], y[((size_t)r * slots + k) * D + d], acc); + out[(size_t)r * D + d] = acc; +} + +// down_epilogue_kernel then combine_kernel in one launch, the same arithmetic in the same order (the same bits). +__device__ __forceinline__ void store_out(float* o, float v) { *o = v; } +__device__ __forceinline__ void store_out(__nv_bfloat16* o, float v) { *o = __float2bfloat16_rn(v); } + +// OutT bf16: the fp32 sum rounded to nearest even as it is stored (same bits as fp32 out then to bfloat16) +template +__global__ void down_combine_kernel(const float* __restrict__ Z, const int* __restrict__ pick, + const half* __restrict__ svh_d, float* __restrict__ y, + const float* __restrict__ wts, OutT* __restrict__ out, int P, int D, int SK, + int E, int slots) { + __shared__ float4 part[32][32]; // [slot][lane]: the slot's 4 outputs of the lane + const int r = blockIdx.x, blk = blockIdx.y; + const int k = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int n = blk * 128 + 4 * lane; + const int p = r * slots + k; + const int e = pick[p]; + float o[4]; + if (e >= 0 && e < E) { + float v[4]; +#pragma unroll + for (int j = 0; j < 4; ++j) { + float s = 0.f; + for (int q = 0; q < SK; ++q) s += Z[((size_t)q * P + p) * D + n + j]; + v[j] = s; + } + fwht128(v, lane); +#pragma unroll + for (int j = 0; j < 4; ++j) { + o[j] = v[j] * HAD_SCALE * __half2float(svh_d[(size_t)e * D + n + j]); + y[(size_t)p * D + n + j] = o[j]; + } + } else { +#pragma unroll + for (int j = 0; j < 4; ++j) o[j] = y[(size_t)p * D + n + j]; + } + part[k][lane] = make_float4(o[0], o[1], o[2], o[3]); + __syncthreads(); + if (k != 0) return; + float acc[4] = {0.f, 0.f, 0.f, 0.f}; + for (int q = 0; q < slots; ++q) { + const float w = wts[r * slots + q]; + const float4 u = part[q][lane]; + acc[0] = fmaf(w, u.x, acc[0]); + acc[1] = fmaf(w, u.y, acc[1]); + acc[2] = fmaf(w, u.z, acc[2]); + acc[3] = fmaf(w, u.w, acc[3]); + } +#pragma unroll + for (int j = 0; j < 4; ++j) store_out(out + (size_t)r * D + n + j, acc[j]); +} + +} // namespace tf_exl3_experts + +// upstream grouped_kernel at mul1 for every (nt, warps, pf) x TF_RANGES, and dequant_kernel<2, K2> (tests) +#define TF_EXL3X_INST(NT, W, PF, LO, HI) template __global__ void tf_exl3x::grouped_kernel<2, NT, W, PF, LO, HI>( \ + const half*, const half*, const int64_t*, const int64_t*, const int*, const int*, const int*, const int*, \ + const int*, float*, int, int, int, int, int, int); +TF_EXL3X_INST(8, 4, 1, 8, 8) +TF_EXL3X_INST(8, 4, 1, 2, 10) +TF_EXL3X_INST(8, 4, 1, 2, 16) +TF_EXL3X_INST(8, 4, 2, 8, 8) +TF_EXL3X_INST(8, 4, 2, 2, 10) +TF_EXL3X_INST(8, 4, 2, 2, 16) +TF_EXL3X_INST(4, 4, 2, 8, 8) +TF_EXL3X_INST(4, 4, 2, 2, 10) +TF_EXL3X_INST(4, 4, 2, 2, 16) +template __global__ void tf_exl3x::dequant_kernel<2, 2>(const uint32_t*, half*, int, int); +template __global__ void tf_exl3x::dequant_kernel<2, 3>(const uint32_t*, half*, int, int); +template __global__ void tf_exl3x::dequant_kernel<2, 4>(const uint32_t*, half*, int, int); +template __global__ void tf_exl3x::dequant_kernel<2, 5>(const uint32_t*, half*, int, int); +template __global__ void tf_exl3x::dequant_kernel<2, 6>(const uint32_t*, half*, int, int); +template __global__ void tf_exl3x::dequant_kernel<2, 7>(const uint32_t*, half*, int, int); +template __global__ void tf_exl3x::dequant_kernel<2, 8>(const uint32_t*, half*, int, int); +template __global__ void tf_exl3x::dequant_kernel<2, 9>(const uint32_t*, half*, int, int); +template __global__ void tf_exl3x::dequant_kernel<2, 10>(const uint32_t*, half*, int, int); +template __global__ void tf_exl3x::dequant_kernel<2, 11>(const uint32_t*, half*, int, int); +template __global__ void tf_exl3x::dequant_kernel<2, 12>(const uint32_t*, half*, int, int); +template __global__ void tf_exl3x::dequant_kernel<2, 13>(const uint32_t*, half*, int, int); +template __global__ void tf_exl3x::dequant_kernel<2, 14>(const uint32_t*, half*, int, int); +template __global__ void tf_exl3x::dequant_kernel<2, 15>(const uint32_t*, half*, int, int); +template __global__ void tf_exl3x::dequant_kernel<2, 16>(const uint32_t*, half*, int, int); +// down_combine_kernel at both outputs the binding takes (fp32, bf16: TF_DSV41_MOE_BF16) +template __global__ void tf_exl3_experts::down_combine_kernel(const float*, const int*, const half*, float*, \ + const float*, float*, int, int, int, int, int); +template __global__ void tf_exl3_experts::down_combine_kernel<__nv_bfloat16>(const float*, const int*, const half*, float*, \ + const float*, __nv_bfloat16*, int, int, int, int, int); +// rot_in_kernel for both inputs the binding takes +template __global__ void tf_exl3_experts::rot_in_kernel<__nv_bfloat16>(const __nv_bfloat16*, int, const int*, const half*, const half*, half*, \ + half*, int, int, int); +template __global__ void tf_exl3_experts::rot_in_kernel(const half*, int, const int*, const half*, const half*, half*, \ + half*, int, int, int); diff --git a/zig/kernels/cuda/deepseek_v41/experts_grouped.cuh b/zig/kernels/cuda/deepseek_v41/experts_grouped.cuh new file mode 100644 index 0000000000..ab7591ac2e --- /dev/null +++ b/zig/kernels/cuda/deepseek_v41/experts_grouped.cuh @@ -0,0 +1,311 @@ +// Copy of src/tensorfold/cuda/exl3/experts_grouped.cuh (blob 57be4c027785): sass_check.sh + +// After ExLlamaV3, MIT, Copyright (c) 2025 Turboderp. + +// Grouped EXL3 expert GEMV: rows stay independent, K ranges fixed by shape, warps summed in order. + +#pragma once + +#include +#include +#include // what the cut torch includes provided: the fixed-width integers + +namespace tf_exl3x { + +// Two codebook values (CB 0 3inst, 1 mcg, 2 mul1) as a half2, bit-identical to ExLlamaV3's decode_3inst_2. +template +__device__ __forceinline__ uint32_t cb_pair(uint32_t s0, uint32_t s1) { + if constexpr (CB == 2) { + const uint32_t x0 = s0 * 0x83DCD12Du, x1 = s1 * 0x83DCD12Du; + const uint32_t sum0 = __dp4a(x0, 0x01010101u, 0x6400u); + const uint32_t sum1 = __dp4a(x1, 0x01010101u, 0x6400u); + const uint32_t hv = __byte_perm(sum0, sum1, 0x5410); + half2 h = *reinterpret_cast(&hv); + half2 r = __hfma2(h, __half2half2(__ushort_as_half(0x1eee)), __half2half2(__ushort_as_half(0xc931))); + return *reinterpret_cast(&r); + } else { + uint32_t x0, x1; + if constexpr (CB == 1) { + x0 = s0 * 0xCBAC1FEDu; + x1 = s1 * 0xCBAC1FEDu; + } else { + x0 = s0 * 89226354u + 64248484u; + x1 = s1 * 89226354u + 64248484u; + } + x0 = (x0 & 0x8FFF8FFFu) ^ 0x3B603B60u; + x1 = (x1 & 0x8FFF8FFFu) ^ 0x3B603B60u; + uint32_t lo = __byte_perm(x0, x1, 0x5410); + uint32_t hi = __byte_perm(x0, x1, 0x7632); + half2 r = __hadd2(*reinterpret_cast(&lo), *reinterpret_cast(&hi)); + return *reinterpret_cast(&r); + } +} + +// K2 half-bits a value: a tile is 4 * K2 words; a lane's eight windows fall in NG runs of GV within two words. +template +struct Fmt { + static constexpr int TW = 4 * K2; + static constexpr int LW = (TW + 31) / 32; + // windows sharing one 64-bit merge (tested for every K2) + static constexpr int GV = (K2 >= 13) ? 2 : ((K2 == 7 || (K2 >= 9 && K2 <= 12) || K2 == 16) ? 4 : 8); + static constexpr int NG = 8 / GV; + __host__ __device__ static constexpr int end(int p) { return (p >> 1) * K2 + ((p & 1) ? K2 : (K2 >> 1)); } + // right shift of window j of a run (run starts at an even position) relative to the run's last window + __host__ __device__ static constexpr int off(int j) { return end(GV - 1) - end(j); } +}; + +template +struct LaneMap { + int hi[Fmt::NG], lo[Fmt::NG], sh[Fmt::NG]; + __device__ __forceinline__ explicit LaneMap(int lane) { + constexpr int TW = Fmt::TW, GV = Fmt::GV; +#pragma unroll + for (int g = 0; g < Fmt::NG; ++g) { + const int last_end = Fmt::end(8 * lane + g * GV + GV - 1) + 128 * K2; + const int hr = (last_end - 1) >> 5; + hi[g] = hr % TW; + lo[g] = (hr + TW - 1) % TW; + sh[g] = (hr + 1) * 32 - last_end; + } + } +}; + +template +__device__ __forceinline__ uint32_t fetch(const uint32_t (&w)[LW], int idx) { + if constexpr (LW == 1) { + return __shfl_sync(0xffffffffu, w[0], idx); + } else { + const uint32_t a = __shfl_sync(0xffffffffu, w[0], idx & 31); + const uint32_t b = __shfl_sync(0xffffffffu, w[1], idx & 31); + return idx < 32 ? a : b; + } +} + +// This lane's eight values of a tile as the B fragments of its two n8 halves. +template +__device__ __forceinline__ void decode_tile(const uint32_t (&w)[Fmt::LW], const LaneMap& m, int lane, + uint32_t (&b0)[2], uint32_t (&b1)[2]) { + uint32_t st[8]; + if constexpr (K2 == 8) { + // 4 bits: lane L's windows are exactly words L-1 and L (as the 4-bit EXL3 decode) + const uint32_t p = __shfl_sync(0xffffffffu, w[0], (lane + 31) & 31); + const uint32_t s = __funnelshift_r(w[0], p, 20); + st[0] = (s >> 8) & 0xffffu; + st[1] = (s >> 4) & 0xffffu; + st[2] = s & 0xffffu; + st[3] = w[0] >> 16; + st[4] = (w[0] >> 12) & 0xffffu; + st[5] = (w[0] >> 8) & 0xffffu; + st[6] = (w[0] >> 4) & 0xffffu; + st[7] = w[0] & 0xffffu; + } else { + constexpr int GV = Fmt::GV, NG = Fmt::NG; +#pragma unroll + for (int g = 0; g < NG; ++g) { + const uint32_t whi = fetch::LW>(w, m.hi[g]); + const uint32_t wlo = fetch::LW>(w, m.lo[g]); + const uint64_t mm = ((((uint64_t)wlo) << 32) | whi) >> m.sh[g]; +#pragma unroll + for (int j = 0; j < GV; ++j) st[g * GV + j] = (uint32_t)(mm >> Fmt::off(j)) & 0xffffu; + } + } + b0[0] = cb_pair(st[0], st[1]); + b0[1] = cb_pair(st[2], st[3]); + b1[0] = cb_pair(st[4], st[5]); + b1[1] = cb_pair(st[6], st[7]); +} + +__device__ __forceinline__ void mma16816(float (&d)[4], const uint32_t (&a)[4], const uint32_t (&b)[2]) { + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 {%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, " + "{%0,%1,%2,%3};\n" + : "+f"(d[0]), "+f"(d[1]), "+f"(d[2]), "+f"(d[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1])); +} + +__device__ __forceinline__ uint32_t load_pair(const half* x, bool ok) { + return ok ? *reinterpret_cast(x) : 0u; +} + +template +__device__ __forceinline__ void load_words(uint32_t (&dst)[Fmt::LW], const uint32_t* p, int lane) { + constexpr int TW = Fmt::TW; +#pragma unroll + for (int l = 0; l < Fmt::LW; ++l) { + if constexpr ((TW % 32) == 0) + dst[l] = __ldg(p + l * 32); + else + dst[l] = (l * 32 + lane < TW) ? __ldg(p + l * 32) : 0u; + } +} + +// One warp's k tiles [kt0, kt0 + nkt) of an expert matrix into acc, PF tiles in flight. +template +__device__ __forceinline__ void warp_tiles(const uint32_t* __restrict__ T, int NTILES, int kt0, int nkt, int nt0, + const half* x0, const half* x1, bool ok0, bool ok1, int lane, + float (&acc)[NT][2][4]) { + constexpr int TW = Fmt::TW, LW = Fmt::LW; + const LaneMap map(lane); + const size_t kstride = (size_t)NTILES * TW; + const uint32_t* tp = T + ((size_t)kt0 * NTILES + nt0) * TW + lane; + + uint32_t pf[PF][NT][LW]; +#pragma unroll + for (int d = 0; d < PF; ++d) + if (d < nkt) +#pragma unroll + for (int i = 0; i < NT; ++i) load_words(pf[d][i], tp + d * kstride + i * TW, lane); + + for (int ib = 0; ib < nkt; ib += PF) { +#pragma unroll + for (int d = 0; d < PF; ++d) { + const int it = ib + d; + if (it < nkt) { + uint32_t w[NT][LW]; +#pragma unroll + for (int i = 0; i < NT; ++i) +#pragma unroll + for (int l = 0; l < LW; ++l) w[i][l] = pf[d][i][l]; + if (it + PF < nkt) +#pragma unroll + for (int i = 0; i < NT; ++i) + load_words(pf[d][i], tp + (size_t)(it + PF) * kstride + i * TW, lane); + const int k = (kt0 + it) * 16; + uint32_t a[4] = {load_pair(x0 + k, ok0), load_pair(x1 + k, ok1), load_pair(x0 + k + 8, ok0), + load_pair(x1 + k + 8, ok1)}; +#pragma unroll + for (int i = 0; i < NT; ++i) { + uint32_t b0[2], b1[2]; + decode_tile(w[i], map, lane, b0, b1); + mma16816(acc[i][0], a, b0); + mma16816(acc[i][1], a, b1); + } + } + } + } +} + +// The K2 values an instance covering [LO, HI] compiles (half-bits 2..16). +__host__ __device__ constexpr bool k2_supported(int k2) { + return k2 >= 2 && k2 <= 16; +} + +// Program (expert u, n block, split, member tile): up to 16 members x W_q over the split's K range; warps in order. +template +__global__ void __launch_bounds__(W * 32) grouped_kernel( + const half* __restrict__ X0, const half* __restrict__ X1, const int64_t* __restrict__ TP0, + const int64_t* __restrict__ TP1, const int* __restrict__ K2_0, const int* __restrict__ K2_1, + const int* __restrict__ uids, const int* __restrict__ ucount, const int* __restrict__ members, + float* __restrict__ Z, int K, int N, int P, int SK, int maxm, int slots) { + const int u = blockIdx.x; + if (u >= ucount[0]) return; + const int MT = (maxm + 15) / 16; + const int mtile = blockIdx.z % MT; + const int split = (blockIdx.z / MT) % SK; + const int mat = blockIdx.z / MT / SK; + const half* X = mat ? X1 : X0; + const int e = uids[u]; + const uint32_t* T = reinterpret_cast(mat ? TP1[e] : TP0[e]); + const int k2 = mat ? K2_1[e] : K2_0[e]; + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int g = lane >> 2, t = lane & 3; + const int KT = K >> 4, NTILES = N >> 4; + + __shared__ int rows_sh[16]; + if (threadIdx.x < 16) { + const int m = mtile * 16 + threadIdx.x; + const int code = m < maxm ? members[u * maxm + m] : -1; + rows_sh[threadIdx.x] = code >= 0 ? (code >> 5) * slots + (code & 31) : -1; + } + __syncthreads(); + if (rows_sh[0] < 0) return; // members come first, so this tile is empty + const int r0 = rows_sh[g], r1 = rows_sh[g + 8]; + const half* x0 = X + (size_t)(r0 < 0 ? 0 : r0) * K + 2 * t; + const half* x1 = X + (size_t)(r1 < 0 ? 0 : r1) * K + 2 * t; + + const int per_split = KT / SK, per_warp = per_split / W; + const int kt0 = split * per_split + warp * per_warp; + const int nt0 = blockIdx.y * NT; + + float acc[NT][2][4]; +#pragma unroll + for (int i = 0; i < NT; ++i) +#pragma unroll + for (int h = 0; h < 2; ++h) +#pragma unroll + for (int c = 0; c < 4; ++c) acc[i][h][c] = 0.f; + + switch (k2) { +#define TF_EXL3X_CASE(K2_) \ + case K2_: \ + if constexpr (K2_ >= LO && K2_ <= HI) \ + warp_tiles(T, NTILES, kt0, per_warp, nt0, x0, x1, r0 >= 0, r1 >= 0, lane, acc); \ + else \ + __trap(); \ + break; + TF_EXL3X_CASE(2) + TF_EXL3X_CASE(3) + TF_EXL3X_CASE(4) + TF_EXL3X_CASE(5) + TF_EXL3X_CASE(6) + TF_EXL3X_CASE(7) + TF_EXL3X_CASE(8) + TF_EXL3X_CASE(9) + TF_EXL3X_CASE(10) + TF_EXL3X_CASE(11) + TF_EXL3X_CASE(12) + TF_EXL3X_CASE(13) + TF_EXL3X_CASE(14) + TF_EXL3X_CASE(15) + TF_EXL3X_CASE(16) +#undef TF_EXL3X_CASE + default: + __trap(); + } + + // warps' partial sums through shared memory, added in warp order + __shared__ float red[W][16][NT * 16]; +#pragma unroll + for (int i = 0; i < NT; ++i) +#pragma unroll + for (int h = 0; h < 2; ++h) { + const int col = i * 16 + h * 8 + 2 * t; + red[warp][g][col] = acc[i][h][0]; + red[warp][g][col + 1] = acc[i][h][1]; + red[warp][g + 8][col] = acc[i][h][2]; + red[warp][g + 8][col + 1] = acc[i][h][3]; + } + __syncthreads(); + for (int idx = threadIdx.x; idx < 16 * NT * 16; idx += W * 32) { + const int row = idx / (NT * 16), col = idx % (NT * 16); + const int r = rows_sh[row]; + if (r < 0) continue; + float s = red[0][row][col]; +#pragma unroll + for (int w = 1; w < W; ++w) s += red[w][row][col]; + Z[(((size_t)mat * SK + split) * P + r) * N + nt0 * 16 + col] = s; + } +} + +// W_q [K, N] fp16 of one matrix through the same lane decode (tests; not on the forward path). +template +__global__ void dequant_kernel(const uint32_t* __restrict__ T, half* __restrict__ out, int K, int N) { + const int kt = blockIdx.x, nt = blockIdx.y, lane = threadIdx.x; + const int NTILES = N >> 4; + uint32_t w[Fmt::LW]; + load_words(w, T + ((size_t)kt * NTILES + nt) * Fmt::TW + lane, lane); + const LaneMap map(lane); + uint32_t b0[2], b1[2]; + decode_tile(w, map, lane, b0, b1); + const int g = lane >> 2, t = lane & 3; + uint32_t v[4] = {b0[0], b0[1], b1[0], b1[1]}; +#pragma unroll + for (int q = 0; q < 4; ++q) { + half2 h = *reinterpret_cast(&v[q]); + const int col = nt * 16 + g + 8 * (q >> 1); + const int row = kt * 16 + 2 * t + 8 * (q & 1); + out[(size_t)row * N + col] = __low2half(h); + out[(size_t)(row + 1) * N + col] = __high2half(h); + } +} + +} // namespace tf_exl3x diff --git a/zig/kernels/cuda/deepseek_v41/glue.cu b/zig/kernels/cuda/deepseek_v41/glue.cu new file mode 100644 index 0000000000..ad589080e5 --- /dev/null +++ b/zig/kernels/cuda/deepseek_v41/glue.cu @@ -0,0 +1,403 @@ +// DeepSeek-V4.1 window glue: torch's CUDA arithmetic per element (fp32 opmath, bf16 / f16 rn casts, --fmad=false). + +#include +#include +#include + +namespace dsv41_glue { + +__device__ __forceinline__ float bf(uint16_t b) { return __uint_as_float(static_cast(b) << 16); } +__device__ __forceinline__ uint16_t to_bf(float x) { return __bfloat16_as_ushort(__float2bfloat16(x)); } + +// rows [n, D] = bf16_rn(f32 embed[ids - lo]) from the rank's slice [V, D], +0 off it (NaN -> 0x7FFF) +__global__ void __launch_bounds__(256) embed_send_kernel(const long long* __restrict__ ids, int n, + const uint16_t* __restrict__ embed, long long vocab_lo, int V, + int D, uint16_t* __restrict__ rows) { + const long long total = (long long)n * D; + for (long long e = (long long)blockIdx.x * blockDim.x + threadIdx.x; e < total; e += (long long)gridDim.x * blockDim.x) { + const long long i = e / D, d = e - i * D; + const long long local = ids[i] - vocab_lo; + rows[e] = (local >= 0 && local < V) ? to_bf(bf(embed[local * D + d])) : (uint16_t)0; + } +} + +// recv [W, n, D] bf16 summed in rank order (fp32 adds) -> out [n, 4 D] bf16, the sum repeated x4 +__global__ void __launch_bounds__(256) embed_sum_kernel(const uint16_t* __restrict__ recv, int W, int n, int D, + uint16_t* __restrict__ out) { + const long long total = (long long)n * D, plane = total; + for (long long e = (long long)blockIdx.x * blockDim.x + threadIdx.x; e < total; e += (long long)gridDim.x * blockDim.x) { + float s = bf(recv[e]); + for (int w = 1; w < W; ++w) s = __fadd_rn(s, bf(recv[w * plane + e])); + const uint16_t b = to_bf(s); + const long long i = e / D, d = e - i * D; + uint16_t* o = out + i * 4 * D + d; +#pragma unroll + for (int r = 0; r < 4; ++r) o[r * D] = b; + } +} + +// y = bf16_rn(x) (the MoE partial) +__global__ void __launch_bounds__(256) cast_bf16_kernel(const float* __restrict__ x, uint16_t* __restrict__ y, + long long n) { + for (long long i = (long long)blockIdx.x * blockDim.x + threadIdx.x; i < n; i += (long long)gridDim.x * blockDim.x) + y[i] = to_bf(x[i]); +} + +// y = f32(f16_rn(x)) +__global__ void __launch_bounds__(256) kit_weights_kernel(const float* __restrict__ x, float* __restrict__ y, + long long n) { + for (long long i = (long long)blockIdx.x * blockDim.x + threadIdx.x; i < n; i += (long long)gridDim.x * blockDim.x) + y[i] = __half2float(__float2half(x[i])); +} + +// x = f32(bf16_rn(x)) in place +__global__ void __launch_bounds__(256) kit_logits_kernel(float* __restrict__ x, long long n) { + for (long long i = (long long)blockIdx.x * blockDim.x + threadIdx.x; i < n; i += (long long)gridDim.x * blockDim.x) + x[i] = bf(to_bf(x[i])); +} + +// pos[0] = start; pos64[i] = start + i; lo[i] = 0 +__global__ void __launch_bounds__(256) positions_kernel(int start, int n, int* __restrict__ pos, + long long* __restrict__ pos64, int* __restrict__ lo) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i == 0) pos[0] = start; + if (i < n) { + pos64[i] = (long long)start + i; + lo[i] = 0; + } +} + +// a graphed window's positions: `start` from the window's staged statics, so one captured graph serves every window +__global__ void __launch_bounds__(256) positions_dev_kernel(const int* __restrict__ start, int n, int* __restrict__ pos, + long long* __restrict__ pos64, int* __restrict__ lo) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + const int s = *start; + if (i == 0) pos[0] = s; + if (i < n) { + pos64[i] = (long long)s + i; + lo[i] = 0; + } +} + +// recv [W, n, c] -> out [n, W c] (bf16 words, any 2-byte type) +__global__ void __launch_bounds__(256) gather_cols_kernel(const uint16_t* __restrict__ recv, int W, int n, int c, + uint16_t* __restrict__ out) { + const long long total = (long long)W * n * c; + for (long long e = (long long)blockIdx.x * blockDim.x + threadIdx.x; e < total; e += (long long)gridDim.x * blockDim.x) { + const long long i = e / ((long long)W * c), r = e - i * W * c, w = r / c, j = r - w * c; + out[e] = recv[(w * n + i) * c + j]; + } +} + +// compressor carry [cols] fp32 <- bf16 row `row` of src (row stride ld); the row from the device when `accepted` is set +__global__ void __launch_bounds__(256) carry_kernel(const uint16_t* __restrict__ src, long long ld, + const int* __restrict__ accepted, int row, int cols, + float* __restrict__ carry) { + const int r = accepted ? *accepted : row; + for (int j = blockIdx.x * blockDim.x + threadIdx.x; j < cols; j += gridDim.x * blockDim.x) + carry[j] = bf(src[(long long)r * ld + j]); +} + +// bf16 -> fp32 (exact: the bits shifted into the high half), n elements +__global__ void __launch_bounds__(256) widen_f32_kernel(const uint16_t* __restrict__ x, float* __restrict__ y, + long long n) { + for (long long i = (long long)blockIdx.x * blockDim.x + threadIdx.x; i < n; i += (long long)gridDim.x * blockDim.x) + y[i] = bf(x[i]); +} + +// DSpark candidates: a CTA per fp32 logit row [rows, V] (stride ld); the best k by (value desc, column asc), exact. +constexpr int DS_NT = 1024, DS_KMAX = 128; + +__device__ __forceinline__ unsigned long long ds_key(float x, int col) { // unique in a row; -0.0 == +0.0 + const unsigned int b = __float_as_uint(x + 0.0f); // -0.0 + 0.0 = +0.0 + const unsigned int mono = (b & 0x80000000u) ? ~b : (b | 0x80000000u); + return (static_cast(mono) << 32) | (0xFFFFFFFFull - static_cast(col)); +} + +__global__ void __launch_bounds__(DS_NT) ds_cands_kernel(const float* __restrict__ lg, long long ld, int V, int k, + int id0, float* __restrict__ out) { + __shared__ int hist[256]; + __shared__ unsigned long long s_prefix, s_mask; + __shared__ int s_need, s_done, s_n; + __shared__ unsigned long long key_sh[DS_KMAX]; + const int tid = threadIdx.x; + const float* row = lg + (long long)blockIdx.x * ld; + float* o = out + (long long)blockIdx.x * 2 * k; // k values, then id0 + col as int bits; padding -inf, -1 + const int c = V < k ? V : k; + if (tid == 0) { s_prefix = 0; s_mask = 0; s_need = c; s_done = 0; } + __syncthreads(); + for (int shift = 56; shift >= 0 && c > 0; shift -= 8) { // radix select of the k-th key, 8 bits a pass + if (tid < 256) hist[tid] = 0; + __syncthreads(); + const unsigned long long prefix = s_prefix, mask = s_mask; + for (int i = tid; i < V; i += DS_NT) { + const unsigned long long u = ds_key(row[i], i); + if ((u & mask) == prefix) atomicAdd(&hist[(u >> shift) & 255], 1); + } + __syncthreads(); + if (tid == 0) { + int need = s_need, d = 255; + for (; d > 0 && hist[d] < need; --d) need -= hist[d]; + s_need = need; + s_prefix = prefix | ((unsigned long long)d << shift); + s_mask = mask | (255ull << shift); + s_done = hist[d] == need; // the digit's keys: all winners + } + __syncthreads(); + if (s_done) break; + } + const unsigned long long prefix = s_prefix, mask = s_mask; + if (tid == 0) s_n = 0; + __syncthreads(); + for (int i = tid; i < V && c > 0; i += DS_NT) { + const unsigned long long u = ds_key(row[i], i); + if ((u & mask) >= prefix) { + const int at = atomicAdd(&s_n, 1); + if (at < DS_KMAX) key_sh[at] = u; + } + } + __syncthreads(); + int m = 1; + while (m < c) m <<= 1; + for (int i = c + tid; i < m; i += DS_NT) key_sh[i] = 0; // sorts last + __syncthreads(); + for (int size = 2; size <= m; size <<= 1) { // bitonic sort, descending + for (int stride = size >> 1; stride > 0; stride >>= 1) { + for (int i = tid; i < m; i += DS_NT) { + const int j = i ^ stride; + if (j > i) { + const bool desc = (i & size) == 0; + const unsigned long long a = key_sh[i], b = key_sh[j]; + if (desc ? a < b : a > b) { key_sh[i] = b; key_sh[j] = a; } + } + } + __syncthreads(); + } + } + for (int i = tid; i < k; i += DS_NT) { + if (i < c) { + const int col = (int)(0xFFFFFFFFull - (key_sh[i] & 0xFFFFFFFFull)); + o[i] = row[col]; + o[k + i] = __int_as_float(id0 + col); + } else { + o[i] = -__int_as_float(0x7F800000); + o[k + i] = __int_as_float(-1); + } + } +} + +// pairs[r] = (f64 vals[r], f64 (cols[r] + id0)), a rank's best; one thread a row +__global__ void __launch_bounds__(256) pick_pack_kernel(const float* __restrict__ vals, const long long* __restrict__ cols, + int n, long long id0, double* __restrict__ pairs) { + const int r = blockIdx.x * blockDim.x + threadIdx.x; + if (r >= n) return; + pairs[2 * r] = static_cast(vals[r]); + pairs[2 * r + 1] = static_cast(cols[r] + id0); +} + +// out[r]: best id over ranks (tie / NaN: lower rank), then accepted, bonus, next position; one block, n <= 1024 +__global__ void __launch_bounds__(1024) pick_merge_kernel(const double* __restrict__ all, int W, int n, + const long long* __restrict__ ids, long long start, + long long* __restrict__ out) { + const int r = threadIdx.x; + if (r < n) { + int bk = 0; + for (int k = 1; k < W; k++) + if (all[2 * (k * n + r)] > all[2 * (bk * n + r)]) bk = k; + out[r] = static_cast(all[2 * (bk * n + r) + 1]); + } + __syncthreads(); + if (r == 0) { + int acc = 0; + if (ids != nullptr) + while (acc + 1 < n && ids[acc + 1] == out[acc]) acc++; + out[n] = acc; + out[n + 1] = out[acc]; + out[n + 2] = start + acc + 1; + } +} + +// one block: DSpark's one-slot statics from the device pick, (anchor, start) = (bonus, next position); rest as staged +__global__ void __launch_bounds__(256) ds_stage_kernel(const long long* __restrict__ pick, int nw, long long* __restrict__ s64, + int* __restrict__ s32, int n, int W, int ring, long long valid, + int noise, int o_pos, int o_tok, int o_cnt, int o_lo, int o_hi, + int o_anchor) { + const long long A = pick[nw + 1], P = pick[nw + 2]; + const long long lo_w = P > W ? P - W : 0; // p -| window + const long long first = lo_w > valid ? lo_w : valid; + const int count = P > first ? (int)(P - first) : 0; + const int t = threadIdx.x; + for (int i = t; i < n; i += blockDim.x) { + s64[i] = i == 0 ? A : (long long)noise; // ids (offset 0) + s64[o_pos + i] = P + i; + s32[o_cnt + i] = count; + s32[o_lo + i] = (int)P; + s32[o_hi + i] = (int)(P + n - 1); + } + for (int e = t; e < n * W; e += blockDim.x) { + const int j = e % W; + s32[o_tok + e] = j < count ? (int)((first + j) % ring) : -1; + } + if (t == 0) { + s32[0] = (int)P; // start (offset 0) + s32[o_anchor] = (int)A; + } +} + +// one thread a slot (segs [k][3]: first row, rows, start): pick_merge's chain walk -> acc [k][3] +__global__ void __launch_bounds__(32) ds_accept_rows_kernel(const long long* __restrict__ picks, const long long* __restrict__ ids, + const long long* __restrict__ segs, int k, long long* __restrict__ acc) { + const int s = threadIdx.x; + if (s >= k) return; + const long long a = segs[3 * s], n = segs[3 * s + 1], start = segs[3 * s + 2]; + long long c = 0; + while (c + 1 < n && ids[a + c + 1] == picks[a + c]) c++; + acc[3 * s] = c; + acc[3 * s + 1] = picks[a + c]; + acc[3 * s + 2] = start + c + 1; +} + +// one block: a pass's members (mem [g][3]: acc index, slot, valid), each one's n rows from its (bonus, next position) +__global__ void __launch_bounds__(256) ds_stage_rows_kernel(const long long* __restrict__ acc, const long long* __restrict__ mem, + int g, long long* __restrict__ s64, int* __restrict__ s32, int n, + int W, int ring, int noise, int o_ids, int o_pos, int o_tok, + int o_cnt, int o_lo, int o_hi) { + const int t = threadIdx.x; + for (int m = 0; m < g; m++) { + const long long j = mem[3 * m], slot = mem[3 * m + 1], valid = mem[3 * m + 2]; + const long long A = acc[3 * j + 1], P = acc[3 * j + 2]; + const long long lo_w = P > W ? P - W : 0; + const long long first = lo_w > valid ? lo_w : valid; + const int count = P > first ? (int)(P - first) : 0; + const int a = m * n; + for (int i = t; i < n; i += blockDim.x) { + s64[o_ids + a + i] = i == 0 ? A : (long long)noise; + s64[o_pos + a + i] = P + i; + s32[o_cnt + a + i] = count; + s32[o_lo + a + i] = (int)P; + s32[o_hi + a + i] = (int)(P + n - 1); + } + for (int e = t; e < n * W; e += blockDim.x) { + const int jj = e % W; + s32[o_tok + a * W + e] = jj < count ? (int)(slot * ring + (first + jj) % ring) : -1; + } + } +} + +__device__ __forceinline__ unsigned long long vs_mix(unsigned long long x) { + x ^= x >> 30; + x *= 0xBF58476D1CE4E5B9ull; + x ^= x >> 27; + x *= 0x94D049BB133111EBull; + return x ^ (x >> 31); +} + +__device__ __forceinline__ double vs_uniform(unsigned long long seed, unsigned long long pos, unsigned long long id) { + unsigned long long x = vs_mix(seed + 0x9E3779B97F4A7C15ull); + x = vs_mix(x ^ (pos * 0xD1B54A32D192ED03ull)); + x = vs_mix(x ^ id); + return (double)(x >> 11) * 0x1p-53 + 0x1p-54; +} + +__device__ __forceinline__ float vs_value(unsigned long long k) { + const unsigned int mono = (unsigned int)(k >> 32); + return __uint_as_float((mono & 0x80000000u) ? (mono & 0x7FFFFFFFu) : ~mono); +} + +// numpy's float64 pairwise sum (leaves of 8..128, halves cut to x8); depth a template, so no runtime recursion +__device__ __forceinline__ double vs_leaf(const double* a, int n) { + if (n < 8) { + double r = -0.0; + for (int i = 0; i < n; i++) r += a[i]; + return r; + } + double r[8]; + for (int j = 0; j < 8; j++) r[j] = a[j]; + int i = 8; + for (; i < n - (n % 8); i += 8) + for (int j = 0; j < 8; j++) r[j] += a[i + j]; + double res = ((r[0] + r[1]) + (r[2] + r[3])) + ((r[4] + r[5]) + (r[6] + r[7])); + for (; i < n; i++) res += a[i]; + return res; +} + +template __device__ __forceinline__ double vs_pairwise(const double* a, int n) { + if (n <= 128) return vs_leaf(a, n); + int n2 = n / 2; + n2 -= n2 % 8; + return vs_pairwise(a, n2) + vs_pairwise(a + n2, n - n2); +} + +template <> __device__ __forceinline__ double vs_pairwise<0>(const double* a, int n) { return vs_leaf(a, n); } + +constexpr int VS_NT = 256, VS_MAX = 2048; + +// one block a row: speculative keyed pick over g [W][n][2 k] (values, ids as int bits); the host's pick decides +__global__ void __launch_bounds__(VS_NT) vs_choose_kernel(const float* __restrict__ g, int W, int n, int k, int count, + unsigned long long seed, long long pos0, int top_k, double temp, + double top_p, double min_p, long long* __restrict__ out) { + __shared__ unsigned long long keys[VS_MAX]; + __shared__ double work[1024]; + const int r = blockIdx.x, t = threadIdx.x, C = W * k; + if (C > VS_MAX) { if (t == 0) out[r] = -1; return; } + for (int i = t; i < C; i += VS_NT) { + const int w = i / k, j = i - w * k; + const float* row = g + ((long long)w * n + r) * 2 * k; + const unsigned int bits = __float_as_uint(row[j] + 0.0f); + const unsigned int mono = (bits & 0x80000000u) ? ~bits : (bits | 0x80000000u); + const unsigned int id = __float_as_uint(row[k + j]); + keys[i] = ((unsigned long long)mono << 32) | (0xFFFFFFFFull - id); + } + int m = 1; + while (m < C) m <<= 1; + for (int i = C + t; i < m; i += VS_NT) keys[i] = 0; + __syncthreads(); + for (int size = 2; size <= m; size <<= 1) + for (int stride = size >> 1; stride > 0; stride >>= 1) { + for (int i = t; i < m; i += VS_NT) { + const int j = i ^ stride; + if (j > i) { + const bool desc = (i & size) == 0; + const unsigned long long a = keys[i], b = keys[j]; + if (desc ? a < b : a > b) { keys[i] = b; keys[j] = a; } + } + } + __syncthreads(); + } + if (t != 0) return; + const long long id0 = (long long)(0xFFFFFFFFull - (keys[0] & 0xFFFFFFFFull)); + if (!(temp > 0.0)) { out[r] = id0; return; } // greedy: the first merged candidate + if (top_k == 0 && top_p > 0.0 && top_p < 1.0) { out[r] = -1; return; } // a nucleus row: the host's statistics + const int width = (count > 0 && count < C) ? count : C; // count 0: all W k + int kk = top_k != 0 ? (top_k < width ? top_k : width) : width; + if (kk < 1) kk = 1; + if (kk > 1024) kk = 1024; + const double tt = temp > 1e-6 ? temp : 1e-6; + const double s0 = (double)vs_value(keys[0]) / tt; + int keep = kk; + if (top_p > 0.0 && top_p < 1.0) { + for (int j = 0; j < kk; j++) work[j] = exp((double)vs_value(keys[j]) / tt - s0); + const double total = vs_pairwise<3>(work, kk); // kk <= 1024 + double cum = 0.0; + int below = 0; + for (int j = 0; j < kk; j++) { + cum += work[j] / total; + if (cum < top_p) below++; + } + keep = below + 1; + } + const double floor_ = s0 + (min_p > 0.0 ? log(min_p) : -INFINITY); + const int lim = keep < kk ? keep : kk; + int best = 0; + double best_score = -INFINITY; + for (int j = 0; j < lim; j++) { + const double x = (double)vs_value(keys[j]) / tt; + if (min_p > 0.0 && x < floor_) continue; + const unsigned long long id = 0xFFFFFFFFull - (keys[j] & 0xFFFFFFFFull); + const double score = x - log(-log(vs_uniform(seed, (unsigned long long)(pos0 + r), id))); + if (j == 0 || score > best_score) { best = j; best_score = score; } + } + out[r] = (long long)(0xFFFFFFFFull - (keys[best] & 0xFFFFFFFFull)); +} + +} // namespace dsv41_glue diff --git a/zig/kernels/cuda/deepseek_v41/linear.cu b/zig/kernels/cuda/deepseek_v41/linear.cu new file mode 100644 index 0000000000..b30ad5e4fa --- /dev/null +++ b/zig/kernels/cuda/deepseek_v41/linear.cu @@ -0,0 +1,363 @@ +// Copy of src/tensorfold/cuda/exl3/linear.cu (blob 83a8dd081226): sass_check.sh + +// EXL3 linear, 1-128 rows: a row's bits depend only on it (mma keeps rows apart, fixed K ranges, fixed-order sums). + +#include +#include +#include // what the cut torch includes provided: the fixed-width integers +#include + +#include "decode.cuh" + +using namespace tf_exl3; + +namespace tf_exl3_linear { + +enum DType : int { F16 = 0, BF16 = 1, F32 = 2 }; + +__device__ __forceinline__ void load4(const void* p, int dtype, size_t i, float (&v)[4]) { + if (dtype == F32) { + const float4 u = *reinterpret_cast(static_cast(p) + i); + v[0] = u.x; v[1] = u.y; v[2] = u.z; v[3] = u.w; + } else if (dtype == BF16) { + const uint2 u = *reinterpret_cast(static_cast(p) + i); + const float2 a = __bfloat1622float2(*reinterpret_cast(&u.x)); + const float2 b = __bfloat1622float2(*reinterpret_cast(&u.y)); + v[0] = a.x; v[1] = a.y; v[2] = b.x; v[3] = b.y; + } else { + const uint2 u = *reinterpret_cast(static_cast(p) + i); + const float2 a = __half22float2(*reinterpret_cast(&u.x)); + const float2 b = __half22float2(*reinterpret_cast(&u.y)); + v[0] = a.x; v[1] = a.y; v[2] = b.x; v[3] = b.y; + } +} + +__device__ __forceinline__ void store4(void* p, int dtype, size_t i, const float (&v)[4]) { + if (dtype == F32) { + *reinterpret_cast(static_cast(p) + i) = make_float4(v[0], v[1], v[2], v[3]); + } else if (dtype == BF16) { + __nv_bfloat162 a = __floats2bfloat162_rn(v[0], v[1]), b = __floats2bfloat162_rn(v[2], v[3]); + uint2 u; + u.x = *reinterpret_cast(&a); + u.y = *reinterpret_cast(&b); + *reinterpret_cast(static_cast<__nv_bfloat16*>(p) + i) = u; + } else { + half2 a = __floats2half2_rn(v[0], v[1]), b = __floats2half2_rn(v[2], v[3]); + uint2 u; + u.x = *reinterpret_cast(&a); + u.y = *reinterpret_cast(&b); + *reinterpret_cast(static_cast(p) + i) = u; + } +} + +// The finished outputs of one row's 128 columns from their fp32 sums (4 a lane): H / sqrt(128), * svh, + bias. +__device__ __forceinline__ void finish(float (&v)[4], int lane, const half* svh, const half* bias, int col) { + fwht128(v, lane); +#pragma unroll + for (int j = 0; j < 4; ++j) { + v[j] = v[j] * HAD_SCALE * __half2float(__ldg(svh + col + j)); + if (bias) v[j] += __half2float(__ldg(bias + col + j)); + } +} + +// A lane's words of one k step: at 1 and 2 bits the warp loads the step and each lane takes its words by shuffle. +template +__host__ __device__ constexpr bool step_shuffled() { + return K2 == 2 || K2 == 4; +} + +template +__host__ __device__ constexpr int step_regs() { + return step_shuffled() ? tile_words() / 4 : 8 * lane_words(); +} + +template +__device__ __forceinline__ void load_step(const uint32_t* step, int lane, uint32_t (&raw)[step_regs()]) { + constexpr int TW = tile_words(), LW = lane_words(); + if constexpr (step_shuffled()) { +#pragma unroll + for (int c = 0; c < step_regs(); ++c) raw[c] = __ldg(step + c * 32 + lane); + } else { + int word, offset; + lane_start(lane, word, offset); +#pragma unroll + for (int j = 0; j < 8; ++j) +#pragma unroll + for (int q = 0; q < LW; ++q) raw[j * LW + q] = __ldg(step + j * TW + (word + q) % TW); + } +} + +// Tile j's lane words from a loaded step; prev: the lane's first window starts in the word before its own. +template +__device__ __forceinline__ void step_lane_words(const uint32_t (&raw)[step_regs()], int j, int lane, bool prev, + uint32_t (&w)[lane_words()]) { + constexpr int TW = tile_words(), LW = lane_words(); + if constexpr (step_shuffled()) { + static_assert(LW == 2, "a lane's windows span two words at 1 and 2 bits"); + constexpr int LPW = 8 / K2; // lanes whose windows end in the same word + const uint32_t r = raw[j * TW / 32]; + const int base = (j * TW) % 32; + const uint32_t own = __shfl_sync(0xffffffffu, r, base + lane / LPW); + const uint32_t before = __shfl_sync(0xffffffffu, r, base + (lane / LPW + TW - 1) % TW); + w[0] = prev ? before : own; + w[1] = prev ? own : 0u; // a window inside one word never reads w[1] + } else { +#pragma unroll + for (int q = 0; q < LW; ++q) w[q] = raw[j * LW + q]; + } +} + +__global__ void __launch_bounds__(128) rot_in_kernel(const void* __restrict__ x, int x_dtype, + const half* __restrict__ suh, half* __restrict__ xh, int K) { + const int blk = blockIdx.x * 4 + (threadIdx.x >> 5), row = blockIdx.y, lane = threadIdx.x & 31; + if (blk * 128 >= K) return; + const int k = blk * 128 + 4 * lane; + float v[4], s[4]; + load4(x, x_dtype, (size_t)row * K + k, v); + load4(suh, F16, k, s); +#pragma unroll + for (int j = 0; j < 4; ++j) v[j] *= s[j]; + fwht128(v, lane); +#pragma unroll + for (int j = 0; j < 4; ++j) v[j] *= HAD_SCALE; + store4(xh, F16, (size_t)row * K + k, v); +} + +template +__global__ void __launch_bounds__(WK * 32) linear_kernel( + const half* __restrict__ xh, const uint32_t* __restrict__ T, long long stride_k, long long stride_nb, + const half* __restrict__ svh, const half* __restrict__ bias, void* __restrict__ y, int y_dtype, + float* __restrict__ Z, int* __restrict__ counters, int M, int K, int N, int SK, const int* __restrict__ skip) { + if (skip != nullptr && *skip != 0) return; // a device-side skip (every program: no counter moves) + constexpr int TW = tile_words(); + constexpr int LW = lane_words(); + extern __shared__ __align__(16) float red[]; // WK * RH * 128 floats + __shared__ int last; + const int RH = min(M, 8); // rows of red a warp + + const int nb = blockIdx.x, split = blockIdx.y, NB = gridDim.x; + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int g = lane >> 2, t = lane & 3; + const int per_warp = (K >> 4) / SK / WK; + const int kt0 = split * (per_warp * WK) + warp * per_warp; + const uint32_t* tiles = T + nb * stride_nb; + const int col0 = nb * 128; + bool prev = false; + if constexpr (step_shuffled()) { + int word, offset; + lane_start(lane, word, offset); + prev = word != lane / (8 / K2); + } + + for (int m0 = 0, pass = 0; m0 < M; m0 += 16, ++pass) { + const int R = min(16, M - m0); + + float acc[8][2][4]; +#pragma unroll + for (int i = 0; i < 8; ++i) +#pragma unroll + for (int h = 0; h < 2; ++h) +#pragma unroll + for (int c = 0; c < 4; ++c) acc[i][h][c] = 0.f; + + // the walk's rows for the mma, clamped so rows past the pass read inside the buffer (their outputs are dropped) + const int r0 = m0 + (g < R ? g : R - 1), r1 = m0 + (g + 8 < R ? g + 8 : R - 1); + const half* x0 = xh + (size_t)r0 * K; + const half* x1 = xh + (size_t)r1 * K; + const uint32_t* tile = tiles + (size_t)kt0 * stride_k; + // up to 6 bits the next k step's words are loaded while this one is decoded; 7 and 8 bits load per tile + constexpr bool PF = K2 <= 12; + constexpr int SR = step_regs(); + uint32_t cur[PF ? SR : 1], nxt[PF ? SR : 1]; + if constexpr (PF) load_step(tile, lane, cur); +#pragma unroll 1 + for (int i = 0; i < per_warp; ++i) { + const int kt = kt0 + i; + uint32_t a[4]; + a[0] = __ldg(reinterpret_cast(x0 + kt * 16 + 2 * t)); + a[1] = __ldg(reinterpret_cast(x1 + kt * 16 + 2 * t)); + a[2] = __ldg(reinterpret_cast(x0 + kt * 16 + 2 * t + 8)); + a[3] = __ldg(reinterpret_cast(x1 + kt * 16 + 2 * t + 8)); + if constexpr (PF) { + if (i + 1 < per_warp) load_step(tile + (size_t)(i + 1) * stride_k, lane, nxt); + } +#pragma unroll + for (int j = 0; j < 8; ++j) { + uint32_t w[LW]; + if constexpr (PF) step_lane_words(cur, j, lane, prev, w); + else ldg_lane_words(tile + (size_t)i * stride_k + j * TW, lane, w); + uint32_t b0[2], b1[2]; + decode_lane(w, lane, b0, b1); + mma16816(acc[j][0], a, b0); + mma16816(acc[j][1], a, b1); + } + if constexpr (PF) { +#pragma unroll + for (int q = 0; q < SR; ++q) cur[q] = nxt[q]; + } + } + + // the warps' sums, added in warp order, rows 0-7 of the pass and then rows 8-15 + for (int rlo = 0; rlo < R; rlo += 8) { + const int rn = min(R - rlo, 8); + __syncthreads(); // red is reused by every half and pass + if (g < RH) { +#pragma unroll + for (int i = 0; i < 8; ++i) +#pragma unroll + for (int h = 0; h < 2; ++h) { + const int col = i * 16 + h * 8 + 2 * t; + *reinterpret_cast(red + (warp * RH + g) * 128 + col) = + rlo ? make_float2(acc[i][h][2], acc[i][h][3]) : make_float2(acc[i][h][0], acc[i][h][1]); + } + } + __syncthreads(); + + if (SK == 1) { + for (int r = warp; r < rn; r += WK) { + float v[4]; + const float4 u = *reinterpret_cast(red + r * 128 + 4 * lane); + v[0] = u.x; v[1] = u.y; v[2] = u.z; v[3] = u.w; +#pragma unroll + for (int w = 1; w < WK; ++w) { + const float4 q = *reinterpret_cast(red + (w * RH + r) * 128 + 4 * lane); + v[0] += q.x; v[1] += q.y; v[2] += q.z; v[3] += q.w; + } + finish(v, lane, svh, bias, col0 + 4 * lane); + store4(y, y_dtype, (size_t)(m0 + rlo + r) * N + col0 + 4 * lane, v); + } + } else { + for (int idx = threadIdx.x; idx < rn * 32; idx += WK * 32) { + const int r = idx >> 5, c = 4 * (idx & 31); + float4 s = *reinterpret_cast(red + r * 128 + c); +#pragma unroll + for (int w = 1; w < WK; ++w) { + const float4 q = *reinterpret_cast(red + (w * RH + r) * 128 + c); + s.x += q.x; s.y += q.y; s.z += q.z; s.w += q.w; + } + *reinterpret_cast(Z + ((size_t)split * M + m0 + rlo + r) * N + col0 + c) = s; + } + } + } + if (SK > 1) { + __threadfence(); + __syncthreads(); + if (threadIdx.x == 0) last = atomicAdd(counters + pass * NB + nb, 1) == SK - 1; + __syncthreads(); + if (last) { + __threadfence(); + for (int r = warp; r < R; r += WK) { + const size_t at = ((size_t)m0 + r) * N + col0 + 4 * lane; + float4 s = __ldcg(reinterpret_cast(Z + at)); + for (int q = 1; q < SK; ++q) { + const float4 u = __ldcg(reinterpret_cast(Z + (size_t)q * M * N + at)); + s.x += u.x; s.y += u.y; s.z += u.z; s.w += u.w; + } + float v[4] = {s.x, s.y, s.z, s.w}; + finish(v, lane, svh, bias, col0 + 4 * lane); + store4(y, y_dtype, (size_t)(m0 + r) * N + col0 + 4 * lane, v); + } + if (threadIdx.x == 0) counters[pass * NB + nb] = 0; // every program of the block has arrived + } + } + __syncthreads(); // red is reused in the next pass + } +} + +// W_q [K, N] fp16 from the trellis words (tile (kt, nt) at kt * stride_k + (nt / 8) * stride_nb): one warp per tile. +template +__global__ void __launch_bounds__(32) unpack_kernel(const uint32_t* __restrict__ T, half* __restrict__ W, int N, + int64_t stride_k, int64_t stride_nb) { + const int kt = blockIdx.y, nt = blockIdx.x, lane = threadIdx.x; + uint32_t w[lane_words()]; + ldg_lane_words(T + kt * stride_k + (nt >> 3) * stride_nb + (nt & 7) * tile_words(), lane, w); + uint32_t b[2][2]; + decode_lane(w, lane, b[0], b[1]); +#pragma unroll + for (int j = 0; j < 8; ++j) { + const uint32_t v = b[j >> 2][(j >> 1) & 1]; + const unsigned short h = (j & 1) ? (unsigned short)(v >> 16) : (unsigned short)(v & 0xffffu); + W[(size_t)(kt * 16 + value_row(lane, j)) * N + nt * 16 + value_col(lane, j)] = __ushort_as_half(h); + } +} + +} // namespace tf_exl3_linear + +// Launched at mul1 (TF_EXL3_WIDTHS(X, 2) + K2 3, 5, 7): linear_kernel at WK 2 / 4 / 8, unpack_kernel. +template __global__ void tf_exl3_linear::linear_kernel<2, 2, 2>(const half*, const uint32_t*, long long, long long, \ + const half*, const half*, void*, int, float*, int*, int, int, int, int, const int*); +template __global__ void tf_exl3_linear::linear_kernel<2, 2, 4>(const half*, const uint32_t*, long long, long long, \ + const half*, const half*, void*, int, float*, int*, int, int, int, int, const int*); +template __global__ void tf_exl3_linear::linear_kernel<2, 2, 8>(const half*, const uint32_t*, long long, long long, \ + const half*, const half*, void*, int, float*, int*, int, int, int, int, const int*); +template __global__ void tf_exl3_linear::unpack_kernel<2, 2>(const uint32_t*, half*, int, int64_t, int64_t); +template __global__ void tf_exl3_linear::linear_kernel<3, 2, 2>(const half*, const uint32_t*, long long, long long, \ + const half*, const half*, void*, int, float*, int*, int, int, int, int, const int*); +template __global__ void tf_exl3_linear::linear_kernel<3, 2, 4>(const half*, const uint32_t*, long long, long long, \ + const half*, const half*, void*, int, float*, int*, int, int, int, int, const int*); +template __global__ void tf_exl3_linear::linear_kernel<3, 2, 8>(const half*, const uint32_t*, long long, long long, \ + const half*, const half*, void*, int, float*, int*, int, int, int, int, const int*); +template __global__ void tf_exl3_linear::unpack_kernel<3, 2>(const uint32_t*, half*, int, int64_t, int64_t); +template __global__ void tf_exl3_linear::linear_kernel<4, 2, 2>(const half*, const uint32_t*, long long, long long, \ + const half*, const half*, void*, int, float*, int*, int, int, int, int, const int*); +template __global__ void tf_exl3_linear::linear_kernel<4, 2, 4>(const half*, const uint32_t*, long long, long long, \ + const half*, const half*, void*, int, float*, int*, int, int, int, int, const int*); +template __global__ void tf_exl3_linear::linear_kernel<4, 2, 8>(const half*, const uint32_t*, long long, long long, \ + const half*, const half*, void*, int, float*, int*, int, int, int, int, const int*); +template __global__ void tf_exl3_linear::unpack_kernel<4, 2>(const uint32_t*, half*, int, int64_t, int64_t); +template __global__ void tf_exl3_linear::linear_kernel<5, 2, 2>(const half*, const uint32_t*, long long, long long, \ + const half*, const half*, void*, int, float*, int*, int, int, int, int, const int*); +template __global__ void tf_exl3_linear::linear_kernel<5, 2, 4>(const half*, const uint32_t*, long long, long long, \ + const half*, const half*, void*, int, float*, int*, int, int, int, int, const int*); +template __global__ void tf_exl3_linear::linear_kernel<5, 2, 8>(const half*, const uint32_t*, long long, long long, \ + const half*, const half*, void*, int, float*, int*, int, int, int, int, const int*); +template __global__ void tf_exl3_linear::unpack_kernel<5, 2>(const uint32_t*, half*, int, int64_t, int64_t); +template __global__ void tf_exl3_linear::linear_kernel<6, 2, 2>(const half*, const uint32_t*, long long, long long, \ + const half*, const half*, void*, int, float*, int*, int, int, int, int, const int*); +template __global__ void tf_exl3_linear::linear_kernel<6, 2, 4>(const half*, const uint32_t*, long long, long long, \ + const half*, const half*, void*, int, float*, int*, int, int, int, int, const int*); +template __global__ void tf_exl3_linear::linear_kernel<6, 2, 8>(const half*, const uint32_t*, long long, long long, \ + const half*, const half*, void*, int, float*, int*, int, int, int, int, const int*); +template __global__ void tf_exl3_linear::unpack_kernel<6, 2>(const uint32_t*, half*, int, int64_t, int64_t); +template __global__ void tf_exl3_linear::linear_kernel<7, 2, 2>(const half*, const uint32_t*, long long, long long, \ + const half*, const half*, void*, int, float*, int*, int, int, int, int, const int*); +template __global__ void tf_exl3_linear::linear_kernel<7, 2, 4>(const half*, const uint32_t*, long long, long long, \ + const half*, const half*, void*, int, float*, int*, int, int, int, int, const int*); +template __global__ void tf_exl3_linear::linear_kernel<7, 2, 8>(const half*, const uint32_t*, long long, long long, \ + const half*, const half*, void*, int, float*, int*, int, int, int, int, const int*); +template __global__ void tf_exl3_linear::unpack_kernel<7, 2>(const uint32_t*, half*, int, int64_t, int64_t); +template __global__ void tf_exl3_linear::linear_kernel<8, 2, 2>(const half*, const uint32_t*, long long, long long, \ + const half*, const half*, void*, int, float*, int*, int, int, int, int, const int*); +template __global__ void tf_exl3_linear::linear_kernel<8, 2, 4>(const half*, const uint32_t*, long long, long long, \ + const half*, const half*, void*, int, float*, int*, int, int, int, int, const int*); +template __global__ void tf_exl3_linear::linear_kernel<8, 2, 8>(const half*, const uint32_t*, long long, long long, \ + const half*, const half*, void*, int, float*, int*, int, int, int, int, const int*); +template __global__ void tf_exl3_linear::unpack_kernel<8, 2>(const uint32_t*, half*, int, int64_t, int64_t); +template __global__ void tf_exl3_linear::linear_kernel<10, 2, 2>(const half*, const uint32_t*, long long, long long, \ + const half*, const half*, void*, int, float*, int*, int, int, int, int, const int*); +template __global__ void tf_exl3_linear::linear_kernel<10, 2, 4>(const half*, const uint32_t*, long long, long long, \ + const half*, const half*, void*, int, float*, int*, int, int, int, int, const int*); +template __global__ void tf_exl3_linear::linear_kernel<10, 2, 8>(const half*, const uint32_t*, long long, long long, \ + const half*, const half*, void*, int, float*, int*, int, int, int, int, const int*); +template __global__ void tf_exl3_linear::unpack_kernel<10, 2>(const uint32_t*, half*, int, int64_t, int64_t); +template __global__ void tf_exl3_linear::linear_kernel<12, 2, 2>(const half*, const uint32_t*, long long, long long, \ + const half*, const half*, void*, int, float*, int*, int, int, int, int, const int*); +template __global__ void tf_exl3_linear::linear_kernel<12, 2, 4>(const half*, const uint32_t*, long long, long long, \ + const half*, const half*, void*, int, float*, int*, int, int, int, int, const int*); +template __global__ void tf_exl3_linear::linear_kernel<12, 2, 8>(const half*, const uint32_t*, long long, long long, \ + const half*, const half*, void*, int, float*, int*, int, int, int, int, const int*); +template __global__ void tf_exl3_linear::unpack_kernel<12, 2>(const uint32_t*, half*, int, int64_t, int64_t); +template __global__ void tf_exl3_linear::linear_kernel<14, 2, 2>(const half*, const uint32_t*, long long, long long, \ + const half*, const half*, void*, int, float*, int*, int, int, int, int, const int*); +template __global__ void tf_exl3_linear::linear_kernel<14, 2, 4>(const half*, const uint32_t*, long long, long long, \ + const half*, const half*, void*, int, float*, int*, int, int, int, int, const int*); +template __global__ void tf_exl3_linear::linear_kernel<14, 2, 8>(const half*, const uint32_t*, long long, long long, \ + const half*, const half*, void*, int, float*, int*, int, int, int, int, const int*); +template __global__ void tf_exl3_linear::unpack_kernel<14, 2>(const uint32_t*, half*, int, int64_t, int64_t); +template __global__ void tf_exl3_linear::linear_kernel<16, 2, 2>(const half*, const uint32_t*, long long, long long, \ + const half*, const half*, void*, int, float*, int*, int, int, int, int, const int*); +template __global__ void tf_exl3_linear::linear_kernel<16, 2, 4>(const half*, const uint32_t*, long long, long long, \ + const half*, const half*, void*, int, float*, int*, int, int, int, int, const int*); +template __global__ void tf_exl3_linear::linear_kernel<16, 2, 8>(const half*, const uint32_t*, long long, long long, \ + const half*, const half*, void*, int, float*, int*, int, int, int, int, const int*); +template __global__ void tf_exl3_linear::unpack_kernel<16, 2>(const uint32_t*, half*, int, int64_t, int64_t); diff --git a/zig/kernels/cuda/deepseek_v41/mhc_cuda.cu b/zig/kernels/cuda/deepseek_v41/mhc_cuda.cu new file mode 100644 index 0000000000..92414ed276 --- /dev/null +++ b/zig/kernels/cuda/deepseek_v41/mhc_cuda.cu @@ -0,0 +1,539 @@ +// Copy of src/tensorfold/families/deepseek_v41/cuda/mhc_cuda.cu (blob caa24ba4bb72): sass_check.sh +#define DSV41_MHC_NO_TORCH + + + +// TF_DSV41_MHC_CUDA=1: one mHC boundary (post + site + finish + next RMSNorm) as ONE launch for decode rows (R <= 16). + +// Triton's arithmetic operation for operation (same bits); the last CTA's finish reads every partial, so order is moot. + +#include +#include + +namespace dsv41_mhc { + +constexpr int D = 5120; +constexpr int NB = 40; // column blocks: part of the arithmetic +constexpr int CB = D / NB; // 128 columns a block: one FMA chain +constexpr int G8 = CB / 8; // 16-byte groups a block row +constexpr int NMIX = 24; +constexpr int MAXR = 16; +constexpr int NT = 128; // threads a CTA +constexpr int PB = CB + 8; // bf16 pitch of a staged row (272 B: conflict-free 16-byte reads) +constexpr int PF = CB + 4; // fp32 pitch (fp32 fn) +constexpr int NE = 4 * NB; // partial entries a row, (stream, block) order +constexpr int PW = 32; // floats an entry +constexpr int NQ = 7; // float4 groups of an entry the finish reads: [0, 28) +constexpr int STAGE4 = 640; // float4 a finish stage (10 KiB; 2 stages) +constexpr int ROWB = 128; // dynamic shared memory a row (the finish's row sums; nsys reads R from it) +static_assert(D % (NB * 8) == 0 && CB == 128, "the chains are 128 columns"); + +struct Args { + const uint16_t* x; long long xs; // streams [R, 4 D] bf16, row stride xs + uint16_t* xout; // new streams [R, 4 D] (POST) + const void* g; long long gr; int world; int gbf16; // gathered partials [world, R, D], rank stride gr + const float* post; const float* comb; const float* pre; // [R, 4], [R, 16] (comb[i, j] at 4 i + j), [R, 4] + const void* fn; const float* base; const float* scale; // [24, 4 D] bf16 / fp32, [24], [3] + float* part; uint16_t* c; // [R, 4, NB, 32] fp32, the collapsed row [R, D] bf16 + uint16_t* tap; long long ts; // [R, D] row stride ts (or null) + const void* nw; int nwbf16; // norm weight [D] bf16 / fp32 + uint16_t* out; // the next sublayer's input [R, D] bf16 + float* opre; float* opost; float* ocomb; // the next site's coefficients (MIX) + int* cnt; // [0] the ticket, [1] the tail's items, [2] exits (self-cleaning) + int R; + float eps, hc_eps, post_alpha; int iters; + long long spin; // TAIL: SM cycles a CTA waits for the last arrival (0: it leaves) + int defer; // TAIL: the coefficient items run in coef_kernel (a side stream) +}; + +__device__ __forceinline__ uint4 ld_stream(const void* p) { + uint4 v; + asm volatile("ld.global.nc.L1::no_allocate.v4.u32 {%0,%1,%2,%3}, [%4];" + : "=r"(v.x), "=r"(v.y), "=r"(v.z), "=r"(v.w) : "l"(p)); + return v; +} +__device__ __forceinline__ uint4 ld_plain(const void* p) { return *reinterpret_cast(p); } +__device__ __forceinline__ void widen8(const uint4 v, float (&f)[8]) { // 8 bf16, column ascending -> fp32 + f[0] = __uint_as_float(v.x << 16); f[1] = __uint_as_float(v.x & 0xffff0000u); + f[2] = __uint_as_float(v.y << 16); f[3] = __uint_as_float(v.y & 0xffff0000u); + f[4] = __uint_as_float(v.z << 16); f[5] = __uint_as_float(v.z & 0xffff0000u); + f[6] = __uint_as_float(v.w << 16); f[7] = __uint_as_float(v.w & 0xffff0000u); +} +__device__ __forceinline__ uint32_t bfbits(float f) { // cvt.rn.bf16.f32: Triton's .to(tl.bfloat16) + unsigned short h; + asm("cvt.rn.bf16.f32 %0, %1;" : "=h"(h) : "f"(f)); + return h; +} +__device__ __forceinline__ float bfr(float f) { return __uint_as_float(bfbits(f) << 16); } +__device__ __forceinline__ uint4 pack8(const float (&f)[8]) { // values already bf16: exact + return make_uint4(bfbits(f[0]) | (bfbits(f[1]) << 16), bfbits(f[2]) | (bfbits(f[3]) << 16), + bfbits(f[4]) | (bfbits(f[5]) << 16), bfbits(f[6]) | (bfbits(f[7]) << 16)); +} +__device__ __forceinline__ void load8f(const void* p, bool bf, float (&f)[8]) { // 8 bf16 or 8 fp32, read-only + if (bf) { + widen8(ld_stream(p), f); + } else { + const float4 a = __ldg(reinterpret_cast(p)), b = __ldg(reinterpret_cast(p) + 1); + f[0] = a.x; f[1] = a.y; f[2] = a.z; f[3] = a.w; f[4] = b.x; f[5] = b.y; f[6] = b.z; f[7] = b.w; + } +} +__device__ __forceinline__ float ex2a(float x) { // tl.exp's ex2.approx.f32 (not .ftz) + float y; + asm("ex2.approx.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} +__device__ __forceinline__ float divf(float a, float b) { // Triton's fp32 '/': div.full.f32 + float y; + asm("div.full.f32 %0, %1, %2;" : "=f"(y) : "f"(a), "f"(b)); + return y; +} +__device__ __forceinline__ void cp16(void* s, const void* g) { + const uint32_t a = static_cast(__cvta_generic_to_shared(s)); + asm volatile("cp.async.cg.shared.global [%0], [%1], 16;" :: "r"(a), "l"(g) : "memory"); +} +__device__ __forceinline__ void cp_commit() { asm volatile("cp.async.commit_group;" ::: "memory"); } +template __device__ __forceinline__ void cp_wait() { + asm volatile("cp.async.wait_group %0;" :: "n"(N) : "memory"); +} + +// X'_i of 8 columns: the post expression, one rounding an operation +__device__ __forceinline__ void post1(const Args& a, int r, int i, const float (&br)[8], const float (&x)[4][8], + float (&y)[8]) { + const float p = __ldg(a.post + r * 4 + i); + const float c0 = __ldg(a.comb + r * 16 + i), c1 = __ldg(a.comb + r * 16 + 4 + i); + const float c2 = __ldg(a.comb + r * 16 + 8 + i), c3 = __ldg(a.comb + r * 16 + 12 + i); +#pragma unroll + for (int k = 0; k < 8; ++k) { + const float m = __fadd_rn(__fadd_rn(__fadd_rn(__fmul_rn(c0, x[0][k]), __fmul_rn(c1, x[1][k])), + __fmul_rn(c2, x[2][k])), __fmul_rn(c3, x[3][k])); + y[k] = bfr(__fadd_rn(__fmul_rn(p, br[k]), m)); + } +} + +template struct Smem { + static constexpr int MAIN = 2 * MAXR * PB * 2 + (FN32 ? NMIX * PF * 4 : NMIX * PB * 2); + static constexpr int BYTES = MAIN > 2 * STAGE4 * 16 ? MAIN : 2 * STAGE4 * 16; +}; + +// phase B: NS chains a thread at once (chain c = t + s NT -> row c / K, kind c % K; kind 25 = the collapsed row) +template +__device__ __forceinline__ void chains(const Args& a, const uint16_t* xs, const uint16_t* cs, const void* fns, int K, + int nch, int b, int jp, int t) { + const uint16_t* src[NS]; + const void* wr[NS]; + int row[NS], kind[NS]; + float acc[NS]; +#pragma unroll + for (int s = 0; s < NS; ++s) { + const int cc = min(t + s * NT, nch - 1); + row[s] = cc / K; + kind[s] = MIX ? cc % K : 25; + src[s] = (kind[s] == 25 ? cs : xs) + row[s] * PB; + wr[s] = kind[s] < NMIX ? (FN32 ? static_cast(static_cast(fns) + kind[s] * PF) + : static_cast(static_cast(fns) + kind[s] * PB)) + : static_cast(src[s]); + acc[s] = 0.0f; + } +#pragma unroll + for (int k8 = 0; k8 < G8; ++k8) { +#pragma unroll + for (int s = 0; s < NS; ++s) { + float v[8], w[8]; + widen8(*reinterpret_cast(src[s] + k8 * 8), v); + if (FN32 && kind[s] < NMIX) { + const float4 w0 = reinterpret_cast(wr[s])[2 * k8]; + const float4 w1 = reinterpret_cast(wr[s])[2 * k8 + 1]; + w[0] = w0.x; w[1] = w0.y; w[2] = w0.z; w[3] = w0.w; w[4] = w1.x; w[5] = w1.y; w[6] = w1.z; w[7] = w1.w; + } else { + widen8(reinterpret_cast(wr[s])[k8], w); + } +#pragma unroll + for (int i = 0; i < 8; ++i) acc[s] = __fmaf_rn(v[i], w[i], acc[s]); + } + } +#pragma unroll + for (int s = 0; s < NS; ++s) + if (t + s * NT < nch) a.part[((row[s] * 4 + jp) * NB + b) * PW + kind[s]] = acc[s]; +} + +// row r's coefficients on 16 lanes (comb (i, j) at lane bits 3..2, 1..0): _finish_k's ops and butterfly order +__device__ __forceinline__ void coefs_of(const Args& a, const float* mx, int r, bool store, int e) { + const float rinv = __fdiv_rn(1.0f, __fsqrt_rn(__fadd_rn(__fdiv_rn(mx[24], 4.0f * D), a.eps))); + const float LOG2E = __int_as_float(0x3FB8AA3B), NLOG2E = __int_as_float(0xBFB8AA3B); + if (e < 4) { + const float pl = __fadd_rn(__fmul_rn(__fmul_rn(mx[e], rinv), __ldg(a.scale + 0)), __ldg(a.base + e)); + const float ql = __fadd_rn(__fmul_rn(__fmul_rn(mx[4 + e], rinv), __ldg(a.scale + 1)), __ldg(a.base + 4 + e)); + const float pre = __fadd_rn(a.hc_eps, divf(1.0f, __fadd_rn(ex2a(__fmul_rn(pl, NLOG2E)), 1.0f))); + const float post = __fmul_rn(a.post_alpha, divf(1.0f, __fadd_rn(ex2a(__fmul_rn(ql, NLOG2E)), 1.0f))); + if (store) { a.opre[r * 4 + e] = pre; a.opost[r * 4 + e] = post; } + } + const float cl = __fadd_rn(__fmul_rn(__fmul_rn(mx[8 + e], rinv), __ldg(a.scale + 2)), __ldg(a.base + 8 + e)); + float mxv = fmaxf(cl, __shfl_xor_sync(0xffffffffu, cl, 2)); + mxv = fmaxf(mxv, __shfl_xor_sync(0xffffffffu, mxv, 1)); + const float ce = ex2a(__fmul_rn(__fsub_rn(cl, mxv), LOG2E)); + float s = __fadd_rn(ce, __shfl_xor_sync(0xffffffffu, ce, 2)); + s = __fadd_rn(s, __shfl_xor_sync(0xffffffffu, s, 1)); + float cm = __fadd_rn(a.hc_eps, divf(ce, s)); + s = __fadd_rn(cm, __shfl_xor_sync(0xffffffffu, cm, 8)); + s = __fadd_rn(s, __shfl_xor_sync(0xffffffffu, s, 4)); + cm = divf(cm, __fadd_rn(a.hc_eps, s)); + for (int it = 1; it < a.iters; ++it) { + s = __fadd_rn(cm, __shfl_xor_sync(0xffffffffu, cm, 2)); + s = __fadd_rn(s, __shfl_xor_sync(0xffffffffu, s, 1)); + cm = divf(cm, __fadd_rn(a.hc_eps, s)); + s = __fadd_rn(cm, __shfl_xor_sync(0xffffffffu, cm, 8)); + s = __fadd_rn(s, __shfl_xor_sync(0xffffffffu, s, 4)); + cm = divf(cm, __fadd_rn(a.hc_eps, s)); + } + if (store) a.ocomb[r * 16 + e] = cm; +} +__device__ __forceinline__ void coefs(const Args& a, const float* rowv, int r, bool store, int e) { + coefs_of(a, rowv + r * (ROWB / 4), r, store, e); +} + +__host__ __device__ constexpr int chunk_entries(int R) { // partial entries a finish stage holds a row + return R <= 1 ? 80 : R <= 2 ? 40 : R <= 4 ? 20 : R <= 9 ? 10 : 5; +} + +// the last CTA: row sums (F1), then the coefficients (F3) beside the normed input (F2) +template +__device__ void finish(const Args& a, float* rowv, float4* st, int t) { + const int R = a.R; + if (MIX) { // F1: 7 float4 sums a row, a thread each; partials staged through 2 cp.async stages + const int ch = chunk_entries(R), nck = NE / ch, per = R * ch * NQ; + const int ri = t / NQ, qi = t % NQ; + float4 acc = make_float4(0.f, 0.f, 0.f, 0.f); + float ssc = 0.0f; + auto issue = [&](int c) { + float4* dst = st + (c & 1) * STAGE4; + for (int i = t; i < per; i += NT) { + const int r = i / (ch * NQ), rem = i % (ch * NQ), e = rem / NQ, q = rem % NQ; + cp16(dst + i, a.part + (static_cast(r) * NE + c * ch + e) * PW + q * 4); + } + cp_commit(); + }; + issue(0); + if (nck > 1) issue(1); + for (int c = 0; c < nck; ++c) { + if (c + 1 < nck) cp_wait<1>(); else cp_wait<0>(); + __syncthreads(); + if (t < R * NQ) { + const float4* src = st + (c & 1) * STAGE4 + ri * ch * NQ + qi; +#pragma unroll 5 + for (int e = 0; e < ch; ++e) { + const float4 v = src[e * NQ]; + acc.x = __fadd_rn(acc.x, v.x); + if (qi < 6) { + acc.y = __fadd_rn(acc.y, v.y); acc.z = __fadd_rn(acc.z, v.z); acc.w = __fadd_rn(acc.w, v.w); + } else if (c * ch + e < NB) { + ssc = __fadd_rn(ssc, v.y); // [25]: stream 0's blocks only + } + } + } + __syncthreads(); + if (c + 2 < nck) issue(c + 2); + } + if (t < R * NQ) { + float* rv = rowv + ri * (ROWB / 4); + if (qi < 6) { + rv[4 * qi] = acc.x; rv[4 * qi + 1] = acc.y; rv[4 * qi + 2] = acc.z; rv[4 * qi + 3] = acc.w; + } else { + rv[24] = acc.x; + rv[25] = __fdiv_rn(1.0f, __fsqrt_rn(__fadd_rn(__fdiv_rn(ssc, 1.0f * D), a.eps))); + } + } + } else if (t < R) { // final: the collapsed row's square sum alone + const float* p = a.part + static_cast(t) * NE * PW + 25; + float v[NB]; +#pragma unroll + for (int e = 0; e < NB; ++e) v[e] = __ldcg(p + e * PW); + float ssc = 0.0f; +#pragma unroll + for (int e = 0; e < NB; ++e) ssc = __fadd_rn(ssc, v[e]); + rowv[t * (ROWB / 4) + 25] = __fdiv_rn(1.0f, __fsqrt_rn(__fadd_rn(__fdiv_rn(ssc, 1.0f * D), a.eps))); + } + __syncthreads(); + const int warp = t >> 5, lane = t & 31; + const int f3w = MIX ? min(4, (R + 1) / 2) : 0; // warps with coefficient rows (2 rows a warp a pass) + if (warp < f3w) { +#pragma unroll 1 + for (int p = 0; p < 2 && warp * 2 + 8 * p < R; ++p) { + const int r = warp * 2 + (lane >> 4) + 8 * p; + coefs(a, rowv, min(r, R - 1), r < R, lane & 15); + } + } + const bool split = f3w < 4; // F2 beside F3 on the other warps, else after it + if (split && warp < f3w) return; + const int wid = split ? t - f3w * 32 : t, nwk = split ? NT - f3w * 32 : NT; + const bool nwbf = a.nwbf16 != 0; +#pragma unroll 4 + for (int i = wid; i < R * (D / 8); i += nwk) { + const int r = i / (D / 8), d0 = (i % (D / 8)) * 8; + const float rn = rowv[r * (ROWB / 4) + 25]; + float c[8], w[8], o[8]; + widen8(__ldcg(reinterpret_cast(a.c + static_cast(r) * D + d0)), c); + load8f(static_cast(a.nw) + d0 * (nwbf ? 2 : 4), nwbf, w); +#pragma unroll + for (int k = 0; k < 8; ++k) o[k] = __fmul_rn(__fmul_rn(c[k], rn), w[k]); + *reinterpret_cast(a.out + static_cast(r) * D + d0) = pack8(o); + } +} + +#include "mhc_tail.cuh" + +// DEFER: a row's coefficient item on its own CTA, on a side stream that joins before the next boundary +__global__ void __launch_bounds__(NT) coef_kernel(const Args a) { + __shared__ __align__(16) float4 st[NE * NQ]; + __shared__ float mx[32]; + coef_item(a, blockIdx.x, st, mx, threadIdx.x); +} + +// POST: post from gathered partials; COLL 1 / 2: c = stream 0 / pre-mix; MIX: next site; ALL4: a CTA writes 4 streams +template +__global__ void __launch_bounds__(NT, 4) boundary_kernel(const Args a) { + extern __shared__ float4 dyn[]; // R x ROWB: the finish's row sums + __shared__ __align__(16) unsigned char sm[Smem::BYTES]; + __shared__ int last; + const int b = blockIdx.x, j = ALL4 ? 0 : blockIdx.y, t = threadIdx.x, R = a.R; + const bool lead = j == 0; + uint16_t* xs = reinterpret_cast(sm); // stream j's new values: the chains' input + uint16_t* cs = xs + MAXR * PB; // the collapsed row (lead) + void* fns = cs + MAXR * PB; // this CTA's fn columns + constexpr int FV = FN32 ? 6 : 3; // 16-byte fn vectors a thread + uint4 fr[FV]; + if (MIX) { // in flight through phase A +#pragma unroll + for (int i = 0; i < FV; ++i) { + const int v = t + i * NT, m = v / (FN32 ? 32 : 16), q = v % (FN32 ? 32 : 16); + fr[i] = ld_stream(static_cast(a.fn) + + (static_cast(m) * 4 * D + j * D + b * CB) * (FN32 ? 4 : 2) + q * 16); + } + } + // ---- phase A: the post, the collapsed row, the taps; 8 columns a thread a step, every row ---- + for (int w = t; w < R * G8; w += NT) { + const int r = w / G8, q = w % G8, col = b * CB + q * 8; + const uint16_t* xr = a.x + static_cast(r) * a.xs + col; + float x[4][8], vj[8]; + if (POST || lead) { +#pragma unroll + for (int s = 0; s < 4; ++s) widen8(ld_plain(xr + s * D), x[s]); + } + if (POST) { + float br[8], gv[8]; + const char* gp = static_cast(a.g); + const int gsz = a.gbf16 ? 2 : 4; + load8f(gp + (static_cast(r) * D + col) * gsz, a.gbf16, br); + for (int rk = 1; rk < a.world; ++rk) { + load8f(gp + (rk * a.gr + static_cast(r) * D + col) * gsz, a.gbf16, gv); +#pragma unroll + for (int k = 0; k < 8; ++k) br[k] = __fadd_rn(br[k], gv[k]); + } +#pragma unroll + for (int k = 0; k < 8; ++k) br[k] = bfr(br[k]); + uint16_t* xo = a.xout + static_cast(r) * 4 * D + col; + if (lead) { // every stream's X' (the collapse, the taps) + float y[4][8]; +#pragma unroll + for (int i = 0; i < 4; ++i) post1(a, r, i, br, x, y[i]); +#pragma unroll + for (int i = 0; i < 4; ++i) + if (ALL4 || i == 0) *reinterpret_cast(xo + i * D) = pack8(y[i]); +#pragma unroll + for (int s = 0; s < 4; ++s) +#pragma unroll + for (int k = 0; k < 8; ++k) x[s][k] = y[s][k]; +#pragma unroll + for (int k = 0; k < 8; ++k) vj[k] = y[0][k]; + } else { + post1(a, r, j, br, x, vj); + *reinterpret_cast(xo + j * D) = pack8(vj); + } + } else if (lead) { +#pragma unroll + for (int k = 0; k < 8; ++k) vj[k] = x[0][k]; + } else { + widen8(ld_plain(xr + j * D), vj); + } + if (MIX) *reinterpret_cast(xs + r * PB + q * 8) = pack8(vj); + if (lead) { + float cv[8]; + if (COLL == 2) { + const float q0 = __ldg(a.pre + r * 4), q1 = __ldg(a.pre + r * 4 + 1); + const float q2 = __ldg(a.pre + r * 4 + 2), q3 = __ldg(a.pre + r * 4 + 3); +#pragma unroll + for (int k = 0; k < 8; ++k) + cv[k] = bfr(__fadd_rn(__fadd_rn(__fadd_rn(__fmul_rn(q0, x[0][k]), __fmul_rn(q1, x[1][k])), + __fmul_rn(q2, x[2][k])), __fmul_rn(q3, x[3][k]))); + } else { +#pragma unroll + for (int k = 0; k < 8; ++k) cv[k] = x[0][k]; + } + const uint4 cp = pack8(cv); + *reinterpret_cast(a.c + static_cast(r) * D + col) = cp; + *reinterpret_cast(cs + r * PB + q * 8) = cp; + if (POST && a.tap != nullptr) { + float tp[8]; +#pragma unroll + for (int k = 0; k < 8; ++k) + tp[k] = bfr(__fmul_rn(__fadd_rn(__fadd_rn(__fadd_rn(x[0][k], x[1][k]), x[2][k]), x[3][k]), 0.25f)); + *reinterpret_cast(a.tap + static_cast(r) * a.ts + col) = pack8(tp); + } + } + } + if (MIX) { +#pragma unroll + for (int i = 0; i < FV; ++i) { + const int v = t + i * NT, m = v / (FN32 ? 32 : 16), q = v % (FN32 ? 32 : 16); + reinterpret_cast(fns)[m * (FN32 ? PF / 4 : PB / 8) + q] = fr[i]; + } + } + __syncthreads(); + // ---- phase B: the FMA chains ---- + const int K = (MIX ? 25 : 0) + (lead ? 1 : 0); + const int nch = R * K; + switch ((nch + NT - 1) / NT) { + case 1: chains<1, FN32, MIX>(a, xs, cs, fns, K, nch, b, j, t); break; + case 2: chains<2, FN32, MIX>(a, xs, cs, fns, K, nch, b, j, t); break; + case 3: chains<3, FN32, MIX>(a, xs, cs, fns, K, nch, b, j, t); break; + case 4: chains<4, FN32, MIX>(a, xs, cs, fns, K, nch, b, j, t); break; + default: break; + } + // ---- the ticket: the last CTA runs the finish ---- + __threadfence(); + __syncthreads(); + if (TAIL) { + tail(a, reinterpret_cast(dyn), reinterpret_cast(sm), t); + return; + } + if (t == 0) last = atomicAdd(a.cnt, 1) == static_cast(gridDim.x * gridDim.y) - 1; + __syncthreads(); + if (!last) return; + __threadfence(); + if (t == 0) atomicExch(a.cnt, 0); // self-cleaning: the next launch / replay starts at 0 + finish(a, reinterpret_cast(dyn), reinterpret_cast(sm), t); +} + +template +cudaError_t launch_one(const Args& a, cudaStream_t stream, bool tail = false) { + const dim3 grid(NB, ALL4 ? 1 : 4); + if (tail) boundary_kernel<<>>(a); + else boundary_kernel<<>>(a); + return cudaGetLastError(); +} + +// mode: 0 boundary (POST, pre-mix collapse, MIX), 1 site from stream 0, 2 site with a carried pre-mix, 3 final +inline cudaError_t coef_launch(const Args& a, cudaStream_t stream) { + if (a.R < 1 || a.R > MAXR) return cudaErrorInvalidValue; + coef_kernel<<>>(a); + return cudaGetLastError(); +} + +// tail: the parallel tail instead of the last CTA's finish (the same bits) +inline cudaError_t dispatch(const Args& a, int mode, bool fn32, cudaStream_t stream, bool tail = false) { + if (a.R < 1 || a.R > MAXR || a.world < 1 || a.world > 8) return cudaErrorInvalidValue; + switch (mode * 2 + (fn32 ? 1 : 0)) { + case 0: return launch_one(a, stream, tail); + case 1: return launch_one(a, stream, tail); + case 2: return launch_one(a, stream, tail); + case 3: return launch_one(a, stream, tail); + case 4: return launch_one(a, stream, tail); + case 5: return launch_one(a, stream, tail); + case 6: case 7: return launch_one(a, stream, tail); + default: return cudaErrorInvalidValue; + } +} + +} // namespace dsv41_mhc + +#ifndef DSV41_MHC_NO_TORCH // the compile test builds the kernels alone (nvcc, no torch headers) +#include +#include +#include + +namespace { +template T* ptr(const at::Tensor& t) { return t.numel() ? static_cast(t.data_ptr()) : nullptr; } +} + +void dsv41_mhc_run_cuda(const at::Tensor& x, at::Tensor& xout, const at::Tensor& g, const at::Tensor& post, + const at::Tensor& comb, const at::Tensor& pre, const at::Tensor& fn, const at::Tensor& base, + const at::Tensor& scale, at::Tensor& part, at::Tensor& c, at::Tensor& tap, + const at::Tensor& nw, at::Tensor& out, at::Tensor& opre, at::Tensor& opost, + at::Tensor& ocomb, at::Tensor& cnt, int64_t R, int64_t mode, double eps, double hc_eps, + double post_alpha, int64_t iters, int64_t spin, int64_t defer) { + using namespace dsv41_mhc; + Args a{}; + a.x = ptr(x); a.xs = x.stride(0); + a.xout = ptr(xout); + a.g = g.numel() ? g.data_ptr() : nullptr; a.gr = g.numel() ? g.stride(0) : 0; + a.world = g.numel() ? static_cast(g.size(0)) : 1; + a.gbf16 = g.scalar_type() == at::kBFloat16 ? 1 : 0; + a.post = ptr(post); a.comb = ptr(comb); a.pre = ptr(pre); + a.fn = fn.numel() ? fn.data_ptr() : nullptr; a.base = ptr(base); a.scale = ptr(scale); + a.part = ptr(part); a.c = ptr(c); + a.tap = ptr(tap); a.ts = tap.numel() ? tap.stride(0) : 0; + a.nw = nw.data_ptr(); a.nwbf16 = nw.scalar_type() == at::kBFloat16 ? 1 : 0; + a.out = ptr(out); + a.opre = ptr(opre); a.opost = ptr(opost); a.ocomb = ptr(ocomb); + a.cnt = ptr(cnt); + a.R = static_cast(R); + a.eps = static_cast(eps); a.hc_eps = static_cast(hc_eps); + a.post_alpha = static_cast(post_alpha); a.iters = static_cast(iters); + a.spin = spin; // < 0: the last CTA's finish; >= 0: the parallel tail + a.defer = spin >= 0 && defer ? 1 : 0; + const bool fn32 = fn.numel() && fn.scalar_type() == at::kFloat; + C10_CUDA_CHECK(dispatch(a, static_cast(mode), fn32, at::cuda::getCurrentCUDAStream(), spin >= 0)); +} + +void dsv41_mhc_coef_cuda(const at::Tensor& part, const at::Tensor& base, const at::Tensor& scale, at::Tensor& opre, + at::Tensor& opost, at::Tensor& ocomb, int64_t R, double eps, double hc_eps, double post_alpha, + int64_t iters) { + using namespace dsv41_mhc; + Args a{}; + a.part = ptr(part); a.base = ptr(base); a.scale = ptr(scale); + a.opre = ptr(opre); a.opost = ptr(opost); a.ocomb = ptr(ocomb); + a.R = static_cast(R); + a.eps = static_cast(eps); a.hc_eps = static_cast(hc_eps); + a.post_alpha = static_cast(post_alpha); a.iters = static_cast(iters); + C10_CUDA_CHECK(coef_launch(a, at::cuda::getCurrentCUDAStream())); +} + +std::vector dsv41_mhc_info_cuda(int64_t device) { // what the GPU offers for cross-CTA reductions + int cl = 0, sms = 0, optin = 0, maxcl = 0; + cudaDeviceGetAttribute(&cl, cudaDevAttrClusterLaunch, static_cast(device)); + cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, static_cast(device)); + cudaDeviceGetAttribute(&optin, cudaDevAttrMaxSharedMemoryPerBlockOptin, static_cast(device)); + if (cl) { + cudaLaunchConfig_t cfg = {}; + cfg.gridDim = dim3(dsv41_mhc::NB, 4); + cfg.blockDim = dim3(dsv41_mhc::NT); + cfg.dynamicSmemBytes = dsv41_mhc::MAXR * dsv41_mhc::ROWB; + cudaFuncSetAttribute(dsv41_mhc::boundary_kernel, + cudaFuncAttributeNonPortableClusterSizeAllowed, 1); + if (cudaOccupancyMaxPotentialClusterSize(&maxcl, dsv41_mhc::boundary_kernel, + &cfg) != cudaSuccess) { + maxcl = -1; + cudaGetLastError(); + } + } + int blocks = 0; + cudaOccupancyMaxActiveBlocksPerMultiprocessor(&blocks, dsv41_mhc::boundary_kernel, + dsv41_mhc::NT, dsv41_mhc::MAXR * dsv41_mhc::ROWB); + return {cl, maxcl, sms, optin, blocks}; +} +#endif + +// The instantiations dispatch() launches: modes 0-2 at bf16 / fp32 fn, mode 3 at bf16; tail off and on. +template __global__ void dsv41_mhc::boundary_kernel(const dsv41_mhc::Args); +template __global__ void dsv41_mhc::boundary_kernel(const dsv41_mhc::Args); +template __global__ void dsv41_mhc::boundary_kernel(const dsv41_mhc::Args); +template __global__ void dsv41_mhc::boundary_kernel(const dsv41_mhc::Args); +template __global__ void dsv41_mhc::boundary_kernel(const dsv41_mhc::Args); +template __global__ void dsv41_mhc::boundary_kernel(const dsv41_mhc::Args); +template __global__ void dsv41_mhc::boundary_kernel(const dsv41_mhc::Args); +template __global__ void dsv41_mhc::boundary_kernel(const dsv41_mhc::Args); +template __global__ void dsv41_mhc::boundary_kernel(const dsv41_mhc::Args); +template __global__ void dsv41_mhc::boundary_kernel(const dsv41_mhc::Args); +template __global__ void dsv41_mhc::boundary_kernel(const dsv41_mhc::Args); +template __global__ void dsv41_mhc::boundary_kernel(const dsv41_mhc::Args); +template __global__ void dsv41_mhc::boundary_kernel(const dsv41_mhc::Args); +template __global__ void dsv41_mhc::boundary_kernel(const dsv41_mhc::Args); diff --git a/zig/kernels/cuda/deepseek_v41/mhc_tail.cuh b/zig/kernels/cuda/deepseek_v41/mhc_tail.cuh new file mode 100644 index 0000000000..e3f236de73 --- /dev/null +++ b/zig/kernels/cuda/deepseek_v41/mhc_tail.cuh @@ -0,0 +1,94 @@ +// mhc_cuda.cu's parallel tail (TF_DSV41_MHC_TAIL=1): coef/norm items and tail(), inside namespace dsv41_mhc +#pragma once +// ---- parallel tail (TF_DSV41_MHC_TAIL=1): live CTAs claim finish()'s items from a counter, same chains, same bits +constexpr int F2C = NT * 8; // columns a normed-input item +constexpr int F2N = D / F2C; // items a row +static_assert(D % F2C == 0 && NE * NQ * 16 <= 2 * STAGE4 * 16, "tail items"); + +__device__ __forceinline__ int ld_acquire(const int* p) { + int v; + asm volatile("ld.acquire.gpu.b32 %0, [%1];" : "=r"(v) : "l"(p) : "memory"); + return v; +} +__device__ __forceinline__ float rn_of(float ssc, float eps) { + return __fdiv_rn(1.0f, __fsqrt_rn(__fadd_rn(__fdiv_rn(ssc, 1.0f * D), eps))); +} + +// row r's 25 sums over its 160 entries and the square sum over stream 0's 40 (finish()'s chains), then coefs() +__device__ void coef_item(const Args& a, int r, float4* st, float* mx, int t) { + for (int i = t; i < NE * NQ; i += NT) + cp16(st + i, a.part + (static_cast(r) * NE + i / NQ) * PW + (i % NQ) * 4); + cp_commit(); + cp_wait<0>(); + __syncthreads(); + if (t < 26) { + const float* sf = reinterpret_cast(st); + const int n = t == 25 ? NB : NE; + float acc = 0.0f; +#pragma unroll 8 + for (int e = 0; e < n; ++e) acc = __fadd_rn(acc, sf[e * (NQ * 4) + t]); + mx[t] = t == 25 ? rn_of(acc, a.eps) : acc; + } + __syncthreads(); + if (t < 32) coefs_of(a, mx, r, t < 16, t & 15); +} + +// normed input of row r, columns [F2C q, F2C q + F2C): bf16((c * rn) * w) +__device__ void norm_item(const Args& a, int r, int q, float* buf, int t) { + const int d0 = q * F2C + t * 8; + const bool nwbf = a.nwbf16 != 0; + float c[8], w[8], o[8]; + widen8(__ldcg(reinterpret_cast(a.c + static_cast(r) * D + d0)), c); + load8f(static_cast(a.nw) + d0 * (nwbf ? 2 : 4), nwbf, w); + if (t < NB) buf[t] = __ldcg(a.part + (static_cast(r) * NE + t) * PW + 25); + __syncthreads(); + if (t == 0) { + float ssc = 0.0f; +#pragma unroll 8 + for (int e = 0; e < NB; ++e) ssc = __fadd_rn(ssc, buf[e]); + buf[NB] = rn_of(ssc, a.eps); + } + __syncthreads(); + const float rn = buf[NB]; +#pragma unroll + for (int k = 0; k < 8; ++k) o[k] = __fmul_rn(__fmul_rn(c[k], rn), w[k]); + *reinterpret_cast(a.out + static_cast(r) * D + d0) = pack8(o); +} + +template +__device__ void tail(const Args& a, float* rowv, float4* st, int t) { + __shared__ int live, item; + __shared__ float nbuf[NB + 1]; // a normed-input item's 40 square sums + rn + const int total = static_cast(gridDim.x * gridDim.y), R = a.R; + if (t == 0) { + int n = atomicAdd(a.cnt, 1) + 1; + if (n < total && a.spin > 0) { + const long long t0 = clock64(); + while ((n = ld_acquire(a.cnt)) < total && clock64() - t0 < a.spin) __nanosleep(32); + } + live = n >= total; + } + __syncthreads(); + if (live) { + __threadfence(); + const int nc = MIX && !a.defer ? R : 0, nitems = nc + R * F2N; + for (;;) { + if (t == 0) item = atomicAdd(a.cnt + 1, 1); + __syncthreads(); + const int it = item; + if (it >= nitems) break; + if (it < nc) { + coef_item(a, it, st, rowv, t); + } else { + const int f = it - nc; + norm_item(a, f / F2N, f % F2N, nbuf, t); + } + __syncthreads(); + } + } + if (t == 0 && atomicAdd(a.cnt + 2, 1) == total - 1) { // the last to leave: every counter back at 0 + atomicExch(a.cnt, 0); + atomicExch(a.cnt + 1, 0); + atomicExch(a.cnt + 2, 0); + } +} diff --git a/zig/kernels/cuda/deepseek_v41/pointwise.cu b/zig/kernels/cuda/deepseek_v41/pointwise.cu new file mode 100644 index 0000000000..eefb387097 --- /dev/null +++ b/zig/kernels/cuda/deepseek_v41/pointwise.cu @@ -0,0 +1,40 @@ +// Torch pointwise ops of DeepSeek-V4.1's fast prefill MoE, bit-exact: fp32 opmath, NaN-propagating clamps. + +// swiglu: act = silu(min(g, limit)) * clamp(u, -limit, limit); add: out += b; add_bf16: y = bf16_rn(a + b). + +// Built with --fmad=false: torch runs each op as its own kernel with fp32 rounding, so nothing may contract. + +#include +#include + +namespace dsv41_pointwise { + +// torch's clamp_max / clamp: a NaN stays NaN +__device__ __forceinline__ float clamp_max(float v, float hi) { return isnan(v) ? v : fminf(v, hi); } +__device__ __forceinline__ float clamp(float v, float lo, float hi) { return isnan(v) ? v : fminf(fmaxf(v, lo), hi); } +// torch's silu (ActivationSiluKernel.cu): x / (1 + exp(-x)) in fp32 +__device__ __forceinline__ float silu(float x) { return x / (1.0f + expf(-x)); } + +// g [R, ldg], u [R, ldu], act [R, lda]; n columns a row +__global__ void __launch_bounds__(256) swiglu_kernel(const float* __restrict__ g, long long ldg, + const float* __restrict__ u, long long ldu, + float* __restrict__ act, long long lda, int n, float limit) { + const int r = blockIdx.y; + for (int c = blockIdx.x * blockDim.x + threadIdx.x; c < n; c += gridDim.x * blockDim.x) { + const float a = silu(clamp_max(g[r * ldg + c], limit)); + act[r * lda + c] = a * clamp(u[r * ldu + c], -limit, limit); + } +} + +__global__ void __launch_bounds__(256) add_kernel(float* __restrict__ out, const float* __restrict__ b, long long n) { + for (long long i = (long long)blockIdx.x * blockDim.x + threadIdx.x; i < n; i += (long long)gridDim.x * blockDim.x) + out[i] = out[i] + b[i]; +} + +__global__ void __launch_bounds__(256) add_bf16_kernel(const float* __restrict__ a, const float* __restrict__ b, + __nv_bfloat16* __restrict__ y, long long n) { + for (long long i = (long long)blockIdx.x * blockDim.x + threadIdx.x; i < n; i += (long long)gridDim.x * blockDim.x) + y[i] = __float2bfloat16_rn(a[i] + b[i]); +} + +} // namespace dsv41_pointwise diff --git a/zig/kernels/cuda/deepseek_v41/ptxas.sh b/zig/kernels/cuda/deepseek_v41/ptxas.sh new file mode 100755 index 0000000000..a90929d20d --- /dev/null +++ b/zig/kernels/cuda/deepseek_v41/ptxas.sh @@ -0,0 +1,27 @@ +#!/bin/bash +# Compiles DeepSeek-V4.1 kernel copies to fatbins in the NGC PyTorch image with -Xptxas -v; reports spills / stack. +set -euo pipefail +# usage: OUT=

SM=121,120 bash ptxas.sh name[:flags] ... (flags comma separated as in cuda.zig, e.g. "x3ld:-O3") +HERE="$(cd "$(dirname "$0")" && pwd)" +OUT="${OUT:?set OUT}" +SM="${SM:-121}" +IMAGE="${TF_ZIG_IMAGE:-nvcr.io/nvidia/pytorch:26.07-py3}" +mkdir -p "$OUT" +TORCH="-D__CUDA_NO_HALF_OPERATORS__ -D__CUDA_NO_HALF_CONVERSIONS__ -D__CUDA_NO_BFLOAT16_CONVERSIONS__ -D__CUDA_NO_HALF2_OPERATORS__ --expt-relaxed-constexpr -std=c++20" +cmds="" +for spec in "$@"; do + name="${spec%%:*}"; flags="${spec#*:}"; [ "$flags" = "$spec" ] && flags="-O3" + gen="" + for sm in ${SM//,/ }; do gen="$gen -gencode=arch=compute_${sm},code=sm_${sm}"; done + cmds="$cmds nvcc -fatbin $TORCH ${flags//,/ } $gen -Xptxas -v -o /out/$name.fatbin /k/$name.cu > /out/$name.ptxas 2>&1 && echo built $name || { echo FAILED $name; tail -20 /out/$name.ptxas; } ;" +done +docker run --rm --network none --memory "${TF_ZIG_DOCKER_MEMORY:-12g}" --memory-swap "${TF_ZIG_DOCKER_MEMORY:-12g}" -u "$(id -u):$(id -g)" -v "$HERE:/k:ro" -v "$OUT:/out" --entrypoint bash "$IMAGE" -c "$cmds" +for spec in "$@"; do + name="${spec%%:*}" + f="$OUT/$name.ptxas" + [ -f "$f" ] || continue + fns=$(grep -c "Compiling entry function" "$f" || true) + spill=$(grep -E "bytes spill (stores|loads)" "$f" | grep -vc " 0 bytes spill stores, 0 bytes spill loads" || true) + stack=$(grep -E "bytes stack frame" "$f" | grep -vc " 0 bytes stack frame" || true) + echo "$name: $fns entry functions; with spills: $spill; with a stack frame: $stack" +done diff --git a/zig/kernels/cuda/deepseek_v41/router_gemv.cu b/zig/kernels/cuda/deepseek_v41/router_gemv.cu new file mode 100644 index 0000000000..60db87747a --- /dev/null +++ b/zig/kernels/cuda/deepseek_v41/router_gemv.cu @@ -0,0 +1,521 @@ +// Copy of src/tensorfold/families/deepseek_v41/cuda/router_gemv.cu (blob 5841ced53bf5): sass_check.sh +#define DSV41_RG_NO_TORCH + + + +// DeepSeek-V4.1 MoE router, one GEMV + top-K launch: each logit is one fp32 FMA chain + xor butterfly, any tiling +#include +#include +#include + +namespace dsv41_rg { + +constexpr int NV = 20; // 16-byte vectors a lane a gate row: D = 32 x 8 x NV = 5,120 +constexpr int D = 32 * 8 * NV; +constexpr int JC = 2; // vectors a lane a chunk +constexpr int NC = NV / JC; // chunks of x a row tile +constexpr int CK = 32 * 8 * JC; // k a chunk (512) +constexpr int EPL = 16; // experts a lane in the selection: E <= 512 +constexpr int KMAX = 8; // top-k at most +static_assert(NV % JC == 0, "whole chunks"); + +__device__ __forceinline__ uint4 ldg_stream(const uint16_t* p) { + uint4 v; + asm volatile("ld.global.nc.L1::no_allocate.v4.u32 {%0,%1,%2,%3}, [%4];" + : "=r"(v.x), "=r"(v.y), "=r"(v.z), "=r"(v.w) + : "l"(p)); + return v; +} + +// 8 bf16 (k ascending: the low half of each word first) -> fp32, exactly +__device__ __forceinline__ void widen8(const uint4 v, float (&f)[8]) { + f[0] = __uint_as_float(v.x << 16); f[1] = __uint_as_float(v.x & 0xffff0000u); + f[2] = __uint_as_float(v.y << 16); f[3] = __uint_as_float(v.y & 0xffff0000u); + f[4] = __uint_as_float(v.z << 16); f[5] = __uint_as_float(v.z & 0xffff0000u); + f[6] = __uint_as_float(v.w << 16); f[7] = __uint_as_float(v.w & 0xffff0000u); +} + +// the reference router's safe softplus, operation for operation (no contraction: the intrinsics are never fused) +__device__ __forceinline__ float softplus(float z) { + const float e = expf(fminf(z, 20.0f)); + if (z > 20.0f) return z; + if (z < -10.0f) return __fsub_rn(e, __fmul_rn(0.5f, __fmul_rn(e, e))); + return logf(__fadd_rn(1.0f, e)); +} + +// TF_DSV41_RG_PRUNE: the top-p prune folded into the selection, its operations in order on the row's K weights +struct Prune { + float min_w, topp, scale; + int min_k, orig, drop, dup, on; +}; + +__device__ __forceinline__ void prune_row(const Prune& P, int K, float (&w)[KMAX], int (&p)[KMAX]) { + float total = 0.0f, f[KMAX], excl[KMAX]; + int rank[KMAX]; + bool keep[KMAX]; +#pragma unroll + for (int i = 0; i < KMAX; ++i) if (i < K) total = __fadd_rn(total, w[i]); +#pragma unroll + for (int i = 0; i < KMAX; ++i) if (i < K) f[i] = __fdiv_rn(w[i], total); +#pragma unroll + for (int i = 0; i < KMAX; ++i) { + rank[i] = 0; +#pragma unroll + for (int j = 0; j < KMAX; ++j) + if (i < K && j < K) rank[i] += (f[j] > f[i]) || (f[j] == f[i] && j < i); + } + float acc = 0.0f; +#pragma unroll + for (int r = 0; r < KMAX; ++r) +#pragma unroll + for (int i = 0; i < KMAX; ++i) + if (r < K && i < K && rank[i] == r) { excl[i] = acc; acc = __fadd_rn(acc, f[i]); } + int nk = 0, top = 0; + float ks = 0.0f; +#pragma unroll + for (int i = 0; i < KMAX; ++i) { + if (i >= K) break; + keep[i] = ((excl[i] < P.topp) && (f[i] >= P.min_w)) || (rank[i] < P.min_k); + nk += keep[i]; + if (rank[i] == 0) top = p[i]; + ks = __fadd_rn(ks, keep[i] ? w[i] : 0.0f); + } + if (nk >= K) return; +#pragma unroll + for (int i = 0; i < KMAX; ++i) { + if (i >= K) break; + w[i] = keep[i] ? (P.orig ? w[i] : __fmul_rn(P.scale, __fdiv_rn(w[i], ks))) : 0.0f; + if (!keep[i]) p[i] = P.dup ? top : P.drop; + } +} + +// Prune pass of one selected row (lane 0); noinline so its arrays never share a live range with select_row's +__device__ __noinline__ void prune_pass(const Prune& P, int K, int kit, int* pick, float* wts, int* psh) { + float wv[KMAX]; + int ip[KMAX]; +#pragma unroll + for (int t = 0; t < KMAX; ++t) { + wv[t] = t < K ? wts[t] : 0.0f; + ip[t] = t < K ? pick[t] : 0; + } + prune_row(P, K, wv, ip); +#pragma unroll + for (int t = 0; t < KMAX; ++t) { + if (t >= K) break; + pick[t] = ip[t]; + wts[t] = kit ? __half2float(__float2half_rn(wv[t])) : wv[t]; + psh[t] = ip[t]; + } +} + +// One row's selection by one warp (lane l holds experts l + 32 i). +__device__ __forceinline__ void select_row(const float* lg, const float* __restrict__ bias, int E, int K, int slots, + float scale, int kit, int* pick, float* wts, int* psh, int lane, + const Prune& P) { + float s[EPL], ch[EPL]; +#pragma unroll + for (int i = 0; i < EPL; ++i) { + const int e = lane + 32 * i; + if (e < E) { + s[i] = __fsqrt_rn(softplus(__ldcg(lg + e))); + ch[i] = __fadd_rn(s[i], bias[e]); + } else { + s[i] = 0.0f; + ch[i] = -INFINITY; + } + } + float sp[KMAX]; + int ip[KMAX]; + float total = 0.0f; +#pragma unroll + for (int t = 0; t < KMAX; ++t) { + if (t >= K) break; + float bv = -INFINITY; + int bi = 0x7fffffff; +#pragma unroll + for (int i = 0; i < EPL; ++i) // ids ascend with i: strict > keeps the lower id on a tie + if (ch[i] > bv) { bv = ch[i]; bi = lane + 32 * i; } +#pragma unroll + for (int o = 16; o > 0; o >>= 1) { // (score desc, id asc): a total order, every lane ends equal + const float ov = __shfl_xor_sync(0xffffffffu, bv, o); + const int oi = __shfl_xor_sync(0xffffffffu, bi, o); + if (ov > bv || (ov == bv && oi < bi)) { bv = ov; bi = oi; } + } + float mine = 0.0f; +#pragma unroll + for (int i = 0; i < EPL; ++i) + if (lane + 32 * i == bi) { mine = s[i]; ch[i] = -INFINITY; } + sp[t] = __shfl_sync(0xffffffffu, mine, bi & 31); + ip[t] = bi; + total = __fadd_rn(total, sp[t]); + } + if (lane == 0) { +#pragma unroll + for (int t = 0; t < KMAX; ++t) { + if (t >= K) break; + float w = __fmul_rn(scale, __fdiv_rn(sp[t], total)); + if (kit && !P.on) w = __half2float(__float2half_rn(w)); // pick.kit_weights (with prune: after it) + pick[t] = ip[t]; + wts[t] = w; + psh[t] = ip[t]; + } + if (slots > K) { pick[K] = E; wts[K] = 1.0f; psh[K] = E; } + } +} + +struct Args { + const uint16_t* x; // [R, x_stride] bf16, rows 16-byte aligned + long long xs; + const uint16_t* w; // [E, D] bf16 + const float* bias; // [E] + int* pick; // [R, slots] + float* wts; // [R, slots] + float* lg; // [R, E] fp32 (scratch, or the caller's logits) + int* cnt; // [>= grid.y] int32, 0 between launches + int R, E, K, slots, select; + float scale; + int kit; // the routed weights rounded to fp16 (pick.kit_weights, folded) + int* gid; // group != 0: upstream's group kernel folded (one row tile only): distinct ids (< GE) + int* gcnt; // in id order, their count, members [., maxm] row * 32 + slot in row order, -1 after + int* gmem; + int GE, maxm, group; + Prune prune; // the top-p prune rule folded into the selection (on = 0: off) +}; + +constexpr int GE_MAX = 512; // table entries the folded grouping takes +constexpr int PICKS_MAX = 16 * (KMAX + 1); +constexpr int TAIL_INTS = PICKS_MAX; // the tail's picks (both kernels) +constexpr int GROUP_INTS = TAIL_INTS + 2 * GE_MAX + PICKS_MAX + 8; // + the grouping's (the narrow kernel) + +// Upstream's group_kernel (exl3/experts.cu) on the R <= 16 rows of one row tile, by the selecting CTA; integers only +template +__device__ __forceinline__ void group_rows(const Args& a, const int* psh, int* sh, int tid) { + constexpr int T = NW * 32; + int* gc = sh; // the main loop's shared memory, reused + int* gpl = gc + GE_MAX; + int* cpl = gpl + GE_MAX; + int* wtot = cpl + PICKS_MAX; + const int n = a.R * a.slots, GE = a.GE, lane = tid & 31, warp = tid >> 5; + for (int e = tid; e < GE; e += T) gc[e] = 0; + __syncthreads(); + for (int i = tid; i < n; i += T) { + const int e = psh[i]; + if (e >= 0 && e < GE) atomicAdd(gc + e, 1); + } + __syncthreads(); + const int per = (GE + T - 1) / T, lo = min(GE, tid * per), hi = min(GE, lo + per); + int used = 0; + for (int e = lo; e < hi; ++e) used += gc[e] > 0; + int inc = used; // inclusive scan in the warp, then over warps +#pragma unroll + for (int o = 1; o < 32; o <<= 1) { + const int v = __shfl_up_sync(0xffffffffu, inc, o); + if (lane >= o) inc += v; + } + if (lane == 31) wtot[warp] = inc; + __syncthreads(); + int base = 0, total = 0; +#pragma unroll + for (int w = 0; w < NW; ++w) { + base += w < warp ? wtot[w] : 0; + total += wtot[w]; + } + int place = base + inc - used; + for (int e = lo; e < hi; ++e) + if (gc[e] > 0) { + a.gid[place] = e; + gpl[e] = place; + cpl[place] = gc[e]; + ++place; + } + if (tid == 0) a.gcnt[0] = total; + __syncthreads(); + for (int i = tid; i < n; i += T) { // rank of pick i among its expert's picks: row order + const int e = psh[i]; + if (e < 0 || e >= GE) continue; + int rank = 0; + for (int q = 0; q < i; ++q) rank += psh[q] == e; + if (rank < a.maxm) a.gmem[(size_t)gpl[e] * a.maxm + rank] = (i / a.slots) * 32 + i % a.slots; + } + for (int q = tid; q < total * a.maxm; q += T) + if (q % a.maxm >= cpl[q / a.maxm]) a.gmem[q] = -1; +} + +// Fused tail: the row tile's last CTA selects its rows, groups them (GROUP), resets the counter; sh: TAIL_INTS ints +template +__device__ __forceinline__ void tail(const Args& a, int r0, int nr, int tid, int* sh) { + __shared__ int last_sh; + int* psh = sh; + const int lane = tid & 31, warp = tid >> 5; + __threadfence(); // this CTA's logits visible device-wide before its arrival + __syncthreads(); // (also: every thread is past the main loop's shared memory) + if (tid == 0) last_sh = atomicAdd(a.cnt + blockIdx.y, 1) == (int)gridDim.x - 1; + __syncthreads(); + if (!last_sh) return; + __threadfence(); + for (int r = warp; r < nr; r += NW) + select_row(a.lg + (size_t)(r0 + r) * a.E, a.bias, a.E, a.K, a.slots, a.scale, a.kit, + a.pick + (size_t)(r0 + r) * a.slots, a.wts + (size_t)(r0 + r) * a.slots, psh + r * a.slots, lane, + a.prune); + if (a.prune.on && lane == 0) // the rule, then the kit rounding (each lane 0: its own rows) + for (int r = warp; r < nr; r += NW) + prune_pass(a.prune, a.K, a.kit, a.pick + (size_t)(r0 + r) * a.slots, a.wts + (size_t)(r0 + r) * a.slots, + psh + r * a.slots); + if (GROUP && a.group) { + __syncthreads(); + group_rows(a, psh, sh + TAIL_INTS, tid); + } + if (tid == 0) a.cnt[blockIdx.y] = 0; // self-cleaning +} + +// RT rows a tile (1 .. 16), NW warps a CTA, EW experts a warp. +template +__global__ void __launch_bounds__(NW * 32) gemv_kernel(const Args a) { + constexpr int PD = 8 / EW; // vectors in flight a lane, per expert + constexpr int XU = RT * CK / 8; // uint4 of x a chunk + constexpr int XT = (XU + NW * 32 - 1) / (NW * 32); + static_assert(PD < NV, "ring shorter than a row"); + __shared__ __align__(16) uint16_t xsh[2][RT][CK]; + + const int tid = threadIdx.x, lane = tid & 31, warp = tid >> 5; + const int r0 = blockIdx.y * RT; + const int nr = min(RT, a.R - r0); + const int e0 = (blockIdx.x * NW + warp) * EW; + + auto wload = [&](int j, uint4 (&dst)[EW]) { +#pragma unroll + for (int e = 0; e < EW; ++e) { + if (e0 + e < a.E) dst[e] = ldg_stream(a.w + (size_t)(e0 + e) * D + (size_t)(j * 32 + lane) * 8); + else dst[e] = make_uint4(0u, 0u, 0u, 0u); + } + }; + auto xload = [&](int c, uint4 (&xr)[XT]) { +#pragma unroll + for (int q = 0; q < XT; ++q) { + const int u = tid + NW * 32 * q, row = u / (CK / 8), col = (u % (CK / 8)) * 8; + xr[q] = make_uint4(0u, 0u, 0u, 0u); + if (u < XU && row < nr) + xr[q] = __ldg(reinterpret_cast(a.x + (size_t)(r0 + row) * a.xs + c * CK + col)); + } + }; + auto xstore = [&](int buf, const uint4 (&xr)[XT]) { +#pragma unroll + for (int q = 0; q < XT; ++q) { + const int u = tid + NW * 32 * q, row = u / (CK / 8), col = (u % (CK / 8)) * 8; + if (u < XU) *reinterpret_cast(&xsh[buf][row][col]) = xr[q]; + } + }; + + float acc[EW][RT]; +#pragma unroll + for (int e = 0; e < EW; ++e) +#pragma unroll + for (int r = 0; r < RT; ++r) acc[e][r] = 0.0f; + uint4 ring[PD][EW]; +#pragma unroll + for (int j = 0; j < PD; ++j) wload(j, ring[j]); + uint4 xr[XT]; + xload(0, xr); + xstore(0, xr); + __syncthreads(); + +#pragma unroll + for (int j = 0; j < NV; ++j) { + const int c = j / JC, jj = j % JC; + if (jj == 0 && c + 1 < NC) xload(c + 1, xr); // next chunk's x: L2, consumed at the chunk's end + float wf[EW][8]; +#pragma unroll + for (int e = 0; e < EW; ++e) widen8(ring[j % PD][e], wf[e]); + if (j + PD < NV) wload(j + PD, ring[j % PD]); // the slot refilled PD steps ahead +#pragma unroll + for (int r = 0; r < RT; ++r) { + float xf[8]; + widen8(*reinterpret_cast(&xsh[c & 1][r][(jj * 32 + lane) * 8]), xf); +#pragma unroll + for (int e = 0; e < EW; ++e) +#pragma unroll + for (int i = 0; i < 8; ++i) acc[e][r] = __fmaf_rn(xf[i], wf[e][i], acc[e][r]); + } + if (jj == JC - 1) { + if (c + 1 < NC) xstore((c + 1) & 1, xr); // the buffer chunk c - 1 read (all past its barrier) + __syncthreads(); + } + } + +#pragma unroll + for (int e = 0; e < EW; ++e) +#pragma unroll + for (int r = 0; r < RT; ++r) { + float v = acc[e][r]; +#pragma unroll + for (int o = 16; o > 0; o >>= 1) v = __fadd_rn(v, __shfl_xor_sync(0xffffffffu, v, o)); + if (lane == r && r < nr && e0 + e < a.E) a.lg[(size_t)(r0 + r) * a.E + e0 + e] = v; + } + static_assert(sizeof(xsh) >= TAIL_INTS * sizeof(int), "the tail's picks fit the staged x"); + if (a.select) tail(a, r0, nr, tid, reinterpret_cast(&xsh[0][0][0])); +} + +#include "router_narrow.cuh" + +template +cudaError_t launch_narrow(const Args& a, cudaStream_t stream) { + static const bool carved = [] { // the most shared memory a SM (8 / NW CTAs resident at once); a + cudaFuncSetAttribute(narrow_kernel, // preference only: an error here is cleared, never fatal + cudaFuncAttributePreferredSharedMemoryCarveout, (int)cudaSharedmemCarveoutMaxShared); + (void)cudaGetLastError(); + return true; + }(); + (void)carved; + const dim3 grid((unsigned)((a.E + NW - 1) / NW), (unsigned)((a.R + RT - 1) / RT)); + narrow_kernel<<>>(a); + return cudaGetLastError(); +} + +template +cudaError_t by_rt_narrow(const Args& a, int rt, cudaStream_t stream) { + switch (rt) { + case 1: return launch_narrow<1, NW>(a, stream); + case 2: return launch_narrow<2, NW>(a, stream); + case 4: return launch_narrow<4, NW>(a, stream); + case 8: return launch_narrow<8, NW>(a, stream); + case 16: return launch_narrow<16, NW>(a, stream); + default: return cudaErrorInvalidValue; + } +} + +template +cudaError_t launch(const Args& a, cudaStream_t stream) { + const dim3 grid((unsigned)((a.E + NW * EW - 1) / (NW * EW)), (unsigned)((a.R + RT - 1) / RT)); + gemv_kernel<<>>(a); + return cudaGetLastError(); +} + +template +cudaError_t by_rt(const Args& a, int rt, cudaStream_t stream) { + switch (rt) { + case 1: return launch<1, NW, EW>(a, stream); + case 2: return launch<2, NW, EW>(a, stream); + case 4: return launch<4, NW, EW>(a, stream); + case 8: return launch<8, NW, EW>(a, stream); + case 16: return launch<16, NW, EW>(a, stream); + default: return cudaErrorInvalidValue; + } +} + +// rt rows a tile (1, 2, 4, 8, 16); (nw, ew) in {4, 8} x {1, 2}, or ew == 0: narrow, nw in {1, 2, 4, 8} +inline cudaError_t dispatch(const Args& a, int rt, int nw, int ew, cudaStream_t stream) { + if (ew == 0) { + switch (nw) { + case 1: return by_rt_narrow<1>(a, rt, stream); + case 2: return by_rt_narrow<2>(a, rt, stream); + case 4: return by_rt_narrow<4>(a, rt, stream); + case 8: return by_rt_narrow<8>(a, rt, stream); + default: return cudaErrorInvalidValue; + } + } + if (nw == 8 && ew == 1) return by_rt<8, 1>(a, rt, stream); + if (nw == 4 && ew == 1) return by_rt<4, 1>(a, rt, stream); + if (nw == 8 && ew == 2) return by_rt<8, 2>(a, rt, stream); + if (nw == 4 && ew == 2) return by_rt<4, 2>(a, rt, stream); + return cudaErrorInvalidValue; +} + +} // namespace dsv41_rg + +#ifndef DSV41_RG_NO_TORCH // the compile test builds the kernels alone (nvcc, no torch headers) +#include +#include +#include + +void dsv41_rg_route_cuda(const at::Tensor& X, const at::Tensor& W, const at::Tensor& B, at::Tensor& pick, + at::Tensor& wts, at::Tensor& LG, at::Tensor& cnt, int64_t K, int64_t slots, double scale, + int64_t rt, int64_t nw, int64_t ew, bool select, bool kit, at::Tensor& gid, at::Tensor& gcnt, + at::Tensor& gmem, int64_t GE, bool group, const std::vector& prune) { + using namespace dsv41_rg; + TORCH_CHECK(prune.empty() || prune.size() == 7, "router_gemv: prune = [min_w, topp, scale, min_k, orig, drop, dup]"); + const int64_t R = X.size(0), E = W.size(0); + TORCH_CHECK(X.size(1) == D && W.size(1) == D, "router_gemv: D must be ", D); + TORCH_CHECK(X.stride(1) == 1 && X.stride(0) % 8 == 0 && reinterpret_cast(X.data_ptr()) % 16 == 0, + "router_gemv: x rows must be unit-stride and 16-byte aligned"); + TORCH_CHECK(E >= 1 && E <= 32 * EPL && K >= 1 && K <= KMAX && K <= E && slots >= K && slots <= K + 1, + "router_gemv: E <= ", 32 * EPL, ", 1 <= K <= ", KMAX, ", slots K or K + 1"); + TORCH_CHECK(LG.numel() >= R * E, "router_gemv: logits buffer too small"); + TORCH_CHECK(!select || cnt.numel() >= (R + rt - 1) / rt, "router_gemv: counters too few for ", R, " rows"); + const bool wide = (nw == 4 || nw == 8) && (ew == 1 || ew == 2); + const bool narrow = ew == 0 && (nw == 1 || nw == 2 || nw == 4 || nw == 8); + TORCH_CHECK((rt == 1 || rt == 2 || rt == 4 || rt == 8 || rt == 16) && (wide || narrow), + "router_gemv: no instance (rt, nw, ew) = (", rt, ", ", nw, ", ", ew, ")"); + TORCH_CHECK(!group || (select && narrow && R <= rt && GE >= 1 && GE <= GE_MAX && gid.numel() >= std::min(R * slots, GE) + && gcnt.numel() >= 1 && gmem.dim() == 2 && gmem.size(0) >= std::min(R * slots, GE) + && gmem.size(1) >= 1), + "router_gemv: the folded grouping takes one row tile, GE <= ", GE_MAX, " and ids / members sized"); + if (R == 0) return; + Args a; + a.x = reinterpret_cast(X.data_ptr()); + a.xs = X.stride(0); + a.w = reinterpret_cast(W.data_ptr()); + a.bias = B.data_ptr(); + a.pick = pick.data_ptr(); + a.wts = wts.data_ptr(); + a.lg = LG.data_ptr(); + a.cnt = cnt.data_ptr(); + a.R = (int)R; a.E = (int)E; a.K = (int)K; a.slots = (int)slots; a.select = select ? 1 : 0; + a.scale = (float)scale; + a.kit = kit ? 1 : 0; + a.group = group ? 1 : 0; + a.gid = group ? gid.data_ptr() : nullptr; + a.gcnt = group ? gcnt.data_ptr() : nullptr; + a.gmem = group ? gmem.data_ptr() : nullptr; + a.GE = (int)GE; + a.maxm = group ? (int)gmem.size(1) : 0; // upstream's Scratch.window: members [maxu, R] + a.prune = Prune{}; + if (!prune.empty()) { // fp32 thresholds: the rounding of Triton's float arguments + a.prune.min_w = (float)prune[0]; a.prune.topp = (float)prune[1]; a.prune.scale = (float)prune[2]; + a.prune.min_k = (int)prune[3]; a.prune.orig = (int)prune[4]; a.prune.drop = (int)prune[5]; + a.prune.dup = (int)prune[6]; a.prune.on = 1; + } + C10_CUDA_CHECK(dispatch(a, (int)rt, (int)nw, (int)ew, at::cuda::getCurrentCUDAStream())); +} +#endif + +// The instantiations dispatch() launches: gemv_kernel and narrow_kernel. +template __global__ void dsv41_rg::gemv_kernel<1, 8, 1>(const dsv41_rg::Args); +template __global__ void dsv41_rg::gemv_kernel<1, 4, 1>(const dsv41_rg::Args); +template __global__ void dsv41_rg::gemv_kernel<1, 8, 2>(const dsv41_rg::Args); +template __global__ void dsv41_rg::gemv_kernel<1, 4, 2>(const dsv41_rg::Args); +template __global__ void dsv41_rg::narrow_kernel<1, 1>(const dsv41_rg::Args); +template __global__ void dsv41_rg::narrow_kernel<1, 2>(const dsv41_rg::Args); +template __global__ void dsv41_rg::narrow_kernel<1, 4>(const dsv41_rg::Args); +template __global__ void dsv41_rg::narrow_kernel<1, 8>(const dsv41_rg::Args); +template __global__ void dsv41_rg::gemv_kernel<2, 8, 1>(const dsv41_rg::Args); +template __global__ void dsv41_rg::gemv_kernel<2, 4, 1>(const dsv41_rg::Args); +template __global__ void dsv41_rg::gemv_kernel<2, 8, 2>(const dsv41_rg::Args); +template __global__ void dsv41_rg::gemv_kernel<2, 4, 2>(const dsv41_rg::Args); +template __global__ void dsv41_rg::narrow_kernel<2, 1>(const dsv41_rg::Args); +template __global__ void dsv41_rg::narrow_kernel<2, 2>(const dsv41_rg::Args); +template __global__ void dsv41_rg::narrow_kernel<2, 4>(const dsv41_rg::Args); +template __global__ void dsv41_rg::narrow_kernel<2, 8>(const dsv41_rg::Args); +template __global__ void dsv41_rg::gemv_kernel<4, 8, 1>(const dsv41_rg::Args); +template __global__ void dsv41_rg::gemv_kernel<4, 4, 1>(const dsv41_rg::Args); +template __global__ void dsv41_rg::gemv_kernel<4, 8, 2>(const dsv41_rg::Args); +template __global__ void dsv41_rg::gemv_kernel<4, 4, 2>(const dsv41_rg::Args); +template __global__ void dsv41_rg::narrow_kernel<4, 1>(const dsv41_rg::Args); +template __global__ void dsv41_rg::narrow_kernel<4, 2>(const dsv41_rg::Args); +template __global__ void dsv41_rg::narrow_kernel<4, 4>(const dsv41_rg::Args); +template __global__ void dsv41_rg::narrow_kernel<4, 8>(const dsv41_rg::Args); +template __global__ void dsv41_rg::gemv_kernel<8, 8, 1>(const dsv41_rg::Args); +template __global__ void dsv41_rg::gemv_kernel<8, 4, 1>(const dsv41_rg::Args); +template __global__ void dsv41_rg::gemv_kernel<8, 8, 2>(const dsv41_rg::Args); +template __global__ void dsv41_rg::gemv_kernel<8, 4, 2>(const dsv41_rg::Args); +template __global__ void dsv41_rg::narrow_kernel<8, 1>(const dsv41_rg::Args); +template __global__ void dsv41_rg::narrow_kernel<8, 2>(const dsv41_rg::Args); +template __global__ void dsv41_rg::narrow_kernel<8, 4>(const dsv41_rg::Args); +template __global__ void dsv41_rg::narrow_kernel<8, 8>(const dsv41_rg::Args); +template __global__ void dsv41_rg::gemv_kernel<16, 8, 1>(const dsv41_rg::Args); +template __global__ void dsv41_rg::gemv_kernel<16, 4, 1>(const dsv41_rg::Args); +template __global__ void dsv41_rg::gemv_kernel<16, 8, 2>(const dsv41_rg::Args); +template __global__ void dsv41_rg::gemv_kernel<16, 4, 2>(const dsv41_rg::Args); +template __global__ void dsv41_rg::narrow_kernel<16, 1>(const dsv41_rg::Args); +template __global__ void dsv41_rg::narrow_kernel<16, 2>(const dsv41_rg::Args); +template __global__ void dsv41_rg::narrow_kernel<16, 4>(const dsv41_rg::Args); +template __global__ void dsv41_rg::narrow_kernel<16, 8>(const dsv41_rg::Args); diff --git a/zig/kernels/cuda/deepseek_v41/router_glue.cu b/zig/kernels/cuda/deepseek_v41/router_glue.cu new file mode 100644 index 0000000000..5e39b67d95 --- /dev/null +++ b/zig/kernels/cuda/deepseek_v41/router_glue.cu @@ -0,0 +1,115 @@ +// Decode router glue, arithmetic of exl3/experts.cu (blob 8ce66ca83bce); -O3, fmad on, ftz off; R <= 64, K == 5120. +#include +#include +#include +#include + +namespace dsv41_router_glue { +constexpr int THREADS = 1024; +constexpr int WARPS = THREADS / 32; +constexpr int MAX_E = 512; +constexpr int MAX_PICKS = 64 * 9; +constexpr float HAD_SCALE = 0.08838834764831845f; + +// Verbatim arithmetic from tf_exl3_experts::fwht128 at the source above. +__device__ __forceinline__ void fwht128(float (&v)[4], int lane) { + float a = v[0] + v[1], b = v[0] - v[1], c = v[2] + v[3], d = v[2] - v[3]; + v[0] = a + c; v[1] = b + d; v[2] = a - c; v[3] = b - d; +#pragma unroll + for (int m = 1; m < 32; m <<= 1) { +#pragma unroll + for (int j = 0; j < 4; ++j) { + float o = __shfl_xor_sync(0xffffffffu, v[j], m); + v[j] = (lane & m) ? o - v[j] : v[j] + o; + } + } +} + +// Integer-equal to upstream group_kernel (histogram, id scan, earlier equal picks; duplicates past maxm truncated). +__device__ __forceinline__ void group_rows(const int* pick, int* uids, int* ucount, + int* members, int R, int slots, int E, int maxm) { + __shared__ int picks[MAX_PICKS], count[MAX_E], place[MAX_E], warp_tot[WARPS]; + const int tid = threadIdx.x, lane = tid & 31, warp = tid >> 5; + const int n = R * slots; + if (tid < n) picks[tid] = pick[tid]; + if (tid < E) count[tid] = 0; + __syncthreads(); + if (tid < n) { + const int e = picks[tid]; + if (e >= 0 && e < E) atomicAdd(count + e, 1); + } + __syncthreads(); + const int used = tid < E && count[tid] > 0; + int inc = used; +#pragma unroll + for (int o = 1; o < 32; o <<= 1) { + const int v = __shfl_up_sync(0xffffffffu, inc, o); + if (lane >= o) inc += v; + } + if (lane == 31) warp_tot[warp] = inc; + __syncthreads(); + if (warp == 0) { + const int v = warp_tot[lane]; + int s = v; +#pragma unroll + for (int o = 1; o < 32; o <<= 1) { + const int x = __shfl_up_sync(0xffffffffu, s, o); + if (lane >= o) s += x; + } + warp_tot[lane] = s - v; + if (lane == 31) ucount[0] = s; + } + __syncthreads(); + if (used) { + const int p = warp_tot[warp] + inc - used; + uids[p] = tid; + place[tid] = p; + } + __syncthreads(); + if (tid < n) { + const int e = picks[tid]; + if (e >= 0 && e < E) { + int rank = 0; + for (int q = 0; q < tid; ++q) rank += picks[q] == e; + if (rank < maxm) members[(size_t)place[e] * maxm + rank] = (tid / slots) * 32 + tid % slots; + } + } + if (used) + for (int j = count[tid]; j < maxm; ++j) members[(size_t)place[tid] * maxm + j] = -1; +} +} // namespace dsv41_router_glue + +// One CTA, block 1024, no dynamic shared memory. Same argument order as group. +extern "C" __global__ void __launch_bounds__(1024) router_group( + const int* pick, int* uids, int* ucount, int* members, int R, int slots, int E, int maxm) { + dsv41_router_glue::group_rows(pick, uids, ucount, members, R, slots, E, maxm); +} + +// grid.x = 1 + ceil(R slots (K / 128) 2 / 32), block 1024: CTA 0 groups, other warps rotate one (pick, block, mat). +extern "C" __global__ void __launch_bounds__(1024) router_group_rot( + const uint16_t* x, int x_stride, const int* pick, const half* suh0, const half* suh1, + half* out0, half* out1, int* uids, int* ucount, int* members, + int R, int K, int slots, int E, int maxm) { + using namespace dsv41_router_glue; + if (blockIdx.x == 0) { + group_rows(pick, uids, ucount, members, R, slots, E, maxm); + return; + } + const int lane = threadIdx.x & 31, warp = threadIdx.x >> 5; + const int tile = (blockIdx.x - 1) * WARPS + warp; + const int P = R * slots, blocks = K / 128; + if (tile >= P * blocks * 2) return; // uniform within each warp + const int p = tile % P, blk = (tile / P) % blocks, mat = tile / (P * blocks); + const int row = p / slots, e = pick[p]; + if (e < 0 || e >= E) return; // invalid picks leave outputs untouched + const half* suh = (mat ? suh1 : suh0) + (size_t)e * K + blk * 128 + 4 * lane; + const __nv_bfloat16* xr = reinterpret_cast(x) + + (size_t)row * x_stride + blk * 128 + 4 * lane; + float v[4]; +#pragma unroll + for (int j = 0; j < 4; ++j) v[j] = __bfloat162float(xr[j]) * __half2float(suh[j]); + fwht128(v, lane); + half* o = (mat ? out1 : out0) + (size_t)p * K + blk * 128 + 4 * lane; +#pragma unroll + for (int j = 0; j < 4; ++j) o[j] = __float2half_rn(v[j] * HAD_SCALE); +} diff --git a/zig/kernels/cuda/deepseek_v41/router_narrow.cuh b/zig/kernels/cuda/deepseek_v41/router_narrow.cuh new file mode 100644 index 0000000000..2b38f41edb --- /dev/null +++ b/zig/kernels/cuda/deepseek_v41/router_narrow.cuh @@ -0,0 +1,96 @@ +// router_gemv.cu's narrow_kernel and its cp.async helpers, inside namespace dsv41_rg +#pragma once +// x through L1 (ld.global.nc), volatile so it stays in the step loop (a plain __ldg was hoisted, spilling registers) +__device__ __forceinline__ uint4 ldg_keep(const uint16_t* p) { + uint4 v; + asm volatile("ld.global.nc.v4.u32 {%0,%1,%2,%3}, [%4];" : "=r"(v.x), "=r"(v.y), "=r"(v.z), "=r"(v.w) : "l"(p)); + return v; +} + +__device__ __forceinline__ void cp_async16(void* dst, const void* src) { + const unsigned d = (unsigned)__cvta_generic_to_shared(dst); + asm volatile("cp.async.cg.shared.global [%0], [%1], 16;" ::"r"(d), "l"(src) : "memory"); +} +__device__ __forceinline__ void cp_commit() { asm volatile("cp.async.commit_group;" ::: "memory"); } +template +__device__ __forceinline__ void cp_wait() { asm volatile("cp.async.wait_group %0;" ::"n"(N) : "memory"); } +// n a constant once the step loop is unrolled: the switch folds to one wait_group +__device__ __forceinline__ void cp_wait_n(int n) { + switch (n) { +#define DSV41_RG_W(k) case k: cp_wait(); break; + DSV41_RG_W(0) DSV41_RG_W(1) DSV41_RG_W(2) DSV41_RG_W(3) DSV41_RG_W(4) DSV41_RG_W(5) DSV41_RG_W(6) DSV41_RG_W(7) + DSV41_RG_W(8) DSV41_RG_W(9) DSV41_RG_W(10) DSV41_RG_W(11) DSV41_RG_W(12) DSV41_RG_W(13) DSV41_RG_W(14) + DSV41_RG_W(15) DSV41_RG_W(16) DSV41_RG_W(17) DSV41_RG_W(18) DSV41_RG_W(19) +#undef DSV41_RG_W + default: cp_wait<0>(); break; + } +} + +// Narrow kernel (TF_DSV41_RG_NARROW = NW, rows <= 16): gemv_kernel's chains and bits; gate via cp.async, x via L1 +template +__global__ void __launch_bounds__(NW * 32) narrow_kernel(const Args a) { + constexpr int PD = NW <= 4 && RT <= 4 ? NV : NV / 2; // gate vectors in flight a lane + constexpr int XP = RT <= 4 ? 2 : 0; // steps of x in flight (0: loaded at use, L1) + constexpr int TS = (GROUP_INTS * 4 + NW * 512 - 1) / (NW * 512); // slots the tail needs + __shared__ __align__(16) uint4 wsh[NW][PD > TS ? PD : TS][32]; + const int tid = threadIdx.x, lane = tid & 31, warp = tid >> 5; + const int r0 = blockIdx.y * RT; + const int nr = min(RT, a.R - r0); + const int e = blockIdx.x * NW + warp; + const bool live = e < a.E; + const uint16_t* wrow = a.w + (size_t)(live ? e : 0) * D + (size_t)lane * 8; + + auto wload = [&](int j) { // a dead warp's slots are never read: committed empty + if (live) cp_async16(&wsh[warp][j % PD][lane], wrow + (size_t)j * 256); + cp_commit(); + }; + auto xload = [&](int j, uint4 (&dst)[RT]) { +#pragma unroll + for (int r = 0; r < RT; ++r) + dst[r] = (live && r < nr) + ? ldg_keep(a.x + (size_t)(r0 + r) * a.xs + (size_t)(j * 32 + lane) * 8) + : make_uint4(0u, 0u, 0u, 0u); + }; + +#pragma unroll + for (int j = 0; j < PD; ++j) wload(j); + uint4 xr[XP > 0 ? XP : 1][RT]; +#pragma unroll + for (int j = 0; j < XP; ++j) xload(j, xr[j]); + float acc[RT]; +#pragma unroll + for (int r = 0; r < RT; ++r) acc[r] = 0.0f; + +#pragma unroll + for (int j = 0; j < NV; ++j) { + // groups committed before step j: min(NV, PD + j); group j done = at most that - j - 1 pending + cp_wait_n(min(NV - j - 1, PD - 1)); + float wf[8]; + widen8(wsh[warp][j % PD][lane], wf); + uint4 xc[RT]; + if constexpr (XP == 0) { + xload(j, xc); + } else { +#pragma unroll + for (int r = 0; r < RT; ++r) xc[r] = xr[j % XP][r]; + if (j + XP < NV) xload(j + XP, xr[j % XP]); + } +#pragma unroll + for (int r = 0; r < RT; ++r) { + float xf[8]; + widen8(xc[r], xf); +#pragma unroll + for (int i = 0; i < 8; ++i) acc[r] = __fmaf_rn(xf[i], wf[i], acc[r]); + } + if (j + PD < NV) wload(j + PD); // the slot just read (its value is in wf: used above) + } + +#pragma unroll + for (int r = 0; r < RT; ++r) { + float v = acc[r]; +#pragma unroll + for (int o = 16; o > 0; o >>= 1) v = __fadd_rn(v, __shfl_xor_sync(0xffffffffu, v, o)); + if (lane == r && r < nr && live) a.lg[(size_t)(r0 + r) * a.E + e] = v; + } + if (a.select) tail(a, r0, nr, tid, reinterpret_cast(&wsh[0][0][0])); +} diff --git a/zig/kernels/cuda/deepseek_v41/topk_cuda.cu b/zig/kernels/cuda/deepseek_v41/topk_cuda.cu new file mode 100644 index 0000000000..f633f78e22 --- /dev/null +++ b/zig/kernels/cuda/deepseek_v41/topk_cuda.cu @@ -0,0 +1,296 @@ +// Copy of src/tensorfold/families/deepseek_v41/cuda/csa2/topk_cuda.cu (blob 839ee8a85f0e): sass_check.sh +#define DSV41_ATTN_NO_TORCH + + + +// Indexer decode top-k (TF_DSV41_ATTN_CUDA=1): top K of unique keys (score image, then lower position), exact. + +// Radix select of the K-th 32-bit image over a cluster of <= 8 CTAs a row (DSMEM histograms), then scan-order pick. + +// Image: s == 0 -> 0, NaN -> 0x7FC00000, negative -> bits ^ 0x7FFFFFFF, then ^ 0x80000000 (0 marks missing). + +#include +#include +#include + +namespace dsv41_topk { +namespace cg = cooperative_groups; + +constexpr int NT = 512, NW = NT / 32, MAXCL = 8, MAXEPT = 31; +constexpr unsigned FULL = 0xffffffffu; + +struct Job { + const float* s; long long ss; // scores [R, n] (unit stride) + const int* cand; long long cs; // mode 3: ascending candidate blocks [R, cols] (-1 last) + int* out; long long os; // [R, k] + int* cnt; // [R] visible counts or null + int nk, k, mode, ept; // entries, K, 0 dense / 2 blocks / 3 candidate positions, entries a thread +}; + +struct Args { + Job job[2]; + const void* pos; int pos64; // int32 [1] (POS + r) or int64 [R] + int ratio, bs; +}; + +__device__ __forceinline__ uint32_t image(float s) { + int b = __float_as_int(s); + if (s == 0.f) b = 0; + if (s != s) b = 0x7FC00000; + if (b < 0) b ^= 0x7FFFFFFF; + return static_cast(b) ^ 0x80000000u; +} + +// exclusive scan of one int a thread over the CTA (512 threads); returns the prefix, *total the sum +__device__ int block_scan(int v, int* ws, int* total) { + const int lane = threadIdx.x & 31, w = threadIdx.x >> 5; + int x = v; + for (int o = 1; o < 32; o <<= 1) { + const int y = __shfl_up_sync(FULL, x, o); + if (lane >= o) x += y; + } + if (lane == 31) ws[w] = x; + __syncthreads(); + if (w == 0) { + int z = lane < NW ? ws[lane] : 0; + for (int o = 1; o < 32; o <<= 1) { + const int y = __shfl_up_sync(FULL, z, o); + if (lane >= o) z += y; + } + if (lane < NW) ws[NW + lane] = z; + } + __syncthreads(); + const int pre = (w ? ws[NW + w - 1] : 0) + x - v; + *total = ws[2 * NW - 1]; + __syncthreads(); + return pre; +} + +__global__ void __launch_bounds__(NT, 1) topk_kernel(Args a) { + cg::cluster_group cl = cg::this_cluster(); + const Job J = blockIdx.z == 0 ? a.job[0] : a.job[1]; // (no dynamic indexing: no local memory) + const int CL = static_cast(cl.num_blocks()), rank = static_cast(cl.block_rank()); + const int r = blockIdx.y, tid = threadIdx.x, lane = tid & 31; + extern __shared__ uint32_t img[]; // SL = 512 x ept images (0: missing) + __shared__ int hist[2][256]; + __shared__ int tot[256]; + __shared__ int ws[2 * NW]; + __shared__ int ex[2]; // (above, equal) of this CTA, read by the cluster + __shared__ int dec[4]; + const int SL = NT * J.ept, lo = rank * SL, n = max(0, min(SL, J.nk - lo)); + const int q = a.pos64 ? static_cast(static_cast(a.pos)[r]) + : static_cast(a.pos)[0] + r; + const int nvis = (q + 1) / a.ratio; + if (rank == 0 && tid == 0 && J.cnt) J.cnt[r] = 0; // before the first cluster barrier (release) + const float* srow = J.s + static_cast(r) * J.ss; + const int* crow = J.cand ? J.cand + static_cast(r) * J.cs : nullptr; + for (int i = tid; i < SL; i += NT) { // images, coalesced + const int e = lo + i; + uint32_t u = 0; + if (i < n) { + if (J.mode == 2) { // block e: max over its visible members, newest +inf + float m = 0.f; // tl.max (maxnumf) over the members, masked -inf + for (int j = 0; j < a.bs; ++j) { + const int p = e * a.bs + j; + const float v = p < nvis ? srow[p] : -INFINITY; + m = j ? fmaxf(m, v) : v; + } + if (e == (nvis - 1) / a.bs) m = INFINITY; + if (e * a.bs < nvis) u = image(m); + } else if (J.mode == 3) { + if (crow[e / a.bs] >= 0) u = image(srow[e]); + } else { + u = image(srow[e]); + } + } + img[i] = u; + } + __syncthreads(); + uint32_t prefix = 0, pmask = 0; + int above = 0, quota_mode = 0; // 0: every match selected; 1: equal images capped + for (int p = 0; p < 4; ++p) { + int* h = hist[p & 1]; + if (tid < 256) h[tid] = 0; + __syncthreads(); + const int shift = 24 - 8 * p; + for (int i0 = 0; i0 < SL; i0 += NT) { + const uint32_t v = img[i0 + tid]; + const bool m = v != 0 && (v & pmask) == prefix; + const unsigned act = __ballot_sync(FULL, m); + if (m) { + const int d = (v >> shift) & 255; + const unsigned peers = __match_any_sync(act, d); + if (lane == __ffs(peers) - 1) atomicAdd(&h[d], __popc(peers)); + } + } + cl.sync(); + if (tid < 256) { + int v[MAXCL]; +#pragma unroll + for (int x = 0; x < MAXCL; ++x) v[x] = x < CL ? cl.map_shared_rank(h, x)[tid] : 0; + int sum = 0; +#pragma unroll + for (int x = 0; x < MAXCL; ++x) sum += v[x]; + tot[tid] = sum; + } + __syncthreads(); + if (tid < 32) { // suffix sums over 256 digits, 8 a lane; the digit + int loc[8], run = 0; + for (int x = 7; x >= 0; --x) { run += tot[lane * 8 + x]; loc[x] = run; } + int suf = run; // inclusive suffix over lanes >= this one + for (int o = 1; o < 32; o <<= 1) { + const int y = __shfl_down_sync(FULL, suf, o); + if (lane + o < 32) suf += y; + } + const int after = suf - run; // lanes above this one + int best = -1, bge = 0; + for (int x = 7; x >= 0; --x) + if (best < 0 && above + after + loc[x] >= J.k) { best = lane * 8 + x; bge = after + loc[x]; } + const unsigned has = __ballot_sync(FULL, best >= 0); + const int src = has ? 31 - __clz(has) : 0; + const int dg = __shfl_sync(FULL, best, src), ge = __shfl_sync(FULL, bge, src); + const int all = __shfl_sync(FULL, suf, 0); + if (lane == 0) { dec[0] = dg; dec[1] = ge; dec[2] = all; } + } + __syncthreads(); + const int dg = dec[0], ge = dec[1], all = dec[2]; + if (p == 0 && all <= J.k) break; // every valid entry (prefix / mask stay 0) + const int hd = tot[dg]; + above += ge - hd; + prefix |= static_cast(dg) << shift; + pmask |= 255u << shift; + if (above + hd == J.k) break; // the prefix's entries complete the top K + if (p == 3) quota_mode = 1; // ties at the K-th image: the first ones win + } + // classify: 2 above (selected), 1 equal (selected up to the quota, in scan order), 0 out + const int e0 = tid * J.ept; + int gt = 0, eq = 0; + for (int x = 0; x < J.ept; ++x) { + const uint32_t v = img[e0 + x]; + if (v == 0) continue; + const uint32_t mv = v & pmask; + gt += mv > prefix; + eq += mv == prefix; + } + int tgt, teq; + const int pre_eq = block_scan(eq, ws, &teq); + block_scan(gt, ws, &tgt); + if (tid == 0) { ex[0] = tgt; ex[1] = teq; } + cl.sync(); + int base = 0, eq_before = 0, total = 0, quota = teq; + { + const int need = J.k - above; // equal images to take (quota mode) + for (int x = 0; x < CL; ++x) { + const int* rx = cl.map_shared_rank(ex, x); + const int xg = rx[0], xe = rx[1]; + const int xq = quota_mode ? max(0, min(xe, need - eq_before)) : xe; + if (x < rank) base += xg + xq; + if (x == rank) quota = xq; + eq_before += xe; + total += xg + xq; + } + } + cl.sync(); // no DSMEM access after this: CTAs may exit + int sel = 0; + { + int er = pre_eq; + for (int x = 0; x < J.ept; ++x) { + const uint32_t v = img[e0 + x]; + if (v == 0) continue; + const uint32_t mv = v & pmask; + if (mv > prefix) ++sel; + else if (mv == prefix) { sel += er < quota; ++er; } + } + } + int tsel; + const int off = base + block_scan(sel, ws, &tsel); + int* orow = J.out + static_cast(r) * J.os; + int vis = 0; + { + int er = pre_eq, o = off; + for (int x = 0; x < J.ept; ++x) { + const uint32_t v = img[e0 + x]; + if (v == 0) continue; + const uint32_t mv = v & pmask; + bool s = mv > prefix; + if (!s && mv == prefix) { s = er < quota; ++er; } + if (!s) continue; + const int e = lo + e0 + x; + const int pp = J.mode == 3 ? crow[e / a.bs] * a.bs + e % a.bs : e; + orow[o++] = pp; + vis += pp < nvis; + } + } + if (J.cnt) { + int tv; + block_scan(vis, ws, &tv); + if (tid == 0 && tv) atomicAdd(J.cnt + r, tv); + } + if (rank == 0) + for (int i = total + tid; i < J.k; i += NT) orow[i] = -1; +} + +// Shared memory a CTA for ``ept`` entries a thread (the images; the rest is static). +inline size_t smem(int ept) { return static_cast(NT) * ept * 4; } + +inline cudaError_t launch(const Args& a, int jobs, int R, int CL, int ept, cudaStream_t stream) { + if (R < 1 || jobs < 1 || jobs > 2 || CL < 1 || CL > MAXCL || ept < 1 || ept > MAXEPT) + return cudaErrorInvalidValue; + static bool init = false; + if (!init) { + const cudaError_t e = cudaFuncSetAttribute(topk_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, + static_cast(smem(MAXEPT))); + if (e != cudaSuccess) return e; + init = true; + } + cudaLaunchConfig_t cfg = {}; + cfg.gridDim = dim3(CL, R, jobs); + cfg.blockDim = dim3(NT); + cfg.dynamicSmemBytes = smem(ept); + cfg.stream = stream; + cudaLaunchAttribute at[1]; + at[0].id = cudaLaunchAttributeClusterDimension; + at[0].val.clusterDim.x = CL; + at[0].val.clusterDim.y = 1; + at[0].val.clusterDim.z = 1; + cfg.attrs = at; + cfg.numAttrs = 1; + return cudaLaunchKernelEx(&cfg, topk_kernel, a); +} + +} // namespace dsv41_topk + +#ifndef DSV41_ATTN_NO_TORCH +#include +#include +#include + +namespace { +dsv41_topk::Job job_of(const at::Tensor& s, const at::Tensor& cand, at::Tensor& out, at::Tensor& cnt, int64_t nk, + int64_t k, int64_t mode, int64_t ept) { + dsv41_topk::Job j{}; + j.s = static_cast(s.data_ptr()); j.ss = s.stride(0); + j.cand = cand.numel() ? static_cast(cand.data_ptr()) : nullptr; + j.cs = cand.numel() ? cand.stride(0) : 0; + j.out = static_cast(out.data_ptr()); j.os = out.stride(0); + j.cnt = cnt.numel() ? static_cast(cnt.data_ptr()) : nullptr; + j.nk = static_cast(nk); j.k = static_cast(k); j.mode = static_cast(mode); + j.ept = static_cast(ept); + return j; +} +} + +void dsv41_topk_run_cuda(const at::Tensor& s0, const at::Tensor& c0, at::Tensor& o0, at::Tensor& n0, int64_t nk0, + int64_t k0, int64_t m0, int64_t e0, const at::Tensor& s1, const at::Tensor& c1, + at::Tensor& o1, at::Tensor& n1, int64_t nk1, int64_t k1, int64_t m1, int64_t e1, + int64_t jobs, const at::Tensor& pos, int64_t ratio, int64_t bs, int64_t R, int64_t CL) { + dsv41_topk::Args a{}; + a.job[0] = job_of(s0, c0, o0, n0, nk0, k0, m0, e0); + a.job[1] = jobs > 1 ? job_of(s1, c1, o1, n1, nk1, k1, m1, e1) : a.job[0]; + a.pos = pos.data_ptr(); a.pos64 = pos.scalar_type() == at::kLong ? 1 : 0; + a.ratio = static_cast(ratio); a.bs = static_cast(bs); + const int ept = static_cast(std::max(e0, jobs > 1 ? e1 : e0)); + C10_CUDA_CHECK(dsv41_topk::launch(a, static_cast(jobs), static_cast(R), static_cast(CL), ept, + at::cuda::getCurrentCUDAStream())); +} +#endif diff --git a/zig/kernels/cuda/deepseek_v41/topk_keys.cu b/zig/kernels/cuda/deepseek_v41/topk_keys.cu new file mode 100644 index 0000000000..1bf475ce08 --- /dev/null +++ b/zig/kernels/cuda/deepseek_v41/topk_keys.cu @@ -0,0 +1,94 @@ +// Row top-k of bf16-valued fp32 logits on the device: value descending, ties by lower column (torch.topk of keys). + +// keys = (mono(bits(x + 0.0)) & -65536) | (65535 - col): unique in a row (C <= 65,536), so any exact select matches. + +// One CTA a row: 4-pass 8-bit radix select of the k-th key, gather k winners, bitonic sort; cols int64, vals fp32. + +#include + +namespace dsv41_topk_keys { + +constexpr int NT = 1024, KMAX = 1024, COLS = 1 << 16; + +__device__ __forceinline__ uint32_t ukey(float x, int col) { + int b = __float_as_int(x + 0.0f); // -0.0 + 0.0 = +0.0 + const int mono = b >= 0 ? b : (b ^ 0x7FFFFFFF); + const int key = (mono & -COLS) | (COLS - 1 - col); + return static_cast(key) ^ 0x80000000u; // signed order as unsigned +} + +__global__ void __launch_bounds__(NT) topk_kernel(const float* __restrict__ lg, long long ld, int C, int k, + float* __restrict__ vals, long long* __restrict__ cols) { + __shared__ int hist[256]; + __shared__ uint32_t s_prefix, s_mask; + __shared__ int s_need, s_n; + __shared__ uint32_t key_sh[KMAX]; + __shared__ int col_sh[KMAX]; + const int tid = threadIdx.x; + const float* row = lg + (long long)blockIdx.x * ld; + + // radix select: the k-th largest key, 8 bits a pass from the top + if (tid == 0) { s_prefix = 0; s_mask = 0; s_need = k; } + __syncthreads(); + for (int shift = 24; shift >= 0; shift -= 8) { + if (tid < 256) hist[tid] = 0; + __syncthreads(); + const uint32_t prefix = s_prefix, mask = s_mask; + for (int c = tid; c < C; c += NT) { + const uint32_t u = ukey(row[c], c); + if ((u & mask) == prefix) atomicAdd(&hist[(u >> shift) & 255], 1); + } + __syncthreads(); + if (tid == 0) { // the digit holding the need-th largest, from the top + int need = s_need, d = 255; + for (; d > 0 && hist[d] < need; --d) need -= hist[d]; + s_need = need; + s_prefix = prefix | ((uint32_t)d << shift); + s_mask = mask | (255u << shift); + } + __syncthreads(); + } + const uint32_t kth = s_prefix; // the k-th largest key itself (keys are unique) + + // the winners (> kth, and kth): exactly k, in any order + if (tid == 0) s_n = 0; + __syncthreads(); + for (int c = tid; c < C; c += NT) { + const uint32_t u = ukey(row[c], c); + if (u >= kth) { + const int at = atomicAdd(&s_n, 1); + key_sh[at] = u; + col_sh[at] = c; + } + } + __syncthreads(); + int n = 1; + while (n < k) n <<= 1; + for (int i = k + tid; i < n; i += NT) { key_sh[i] = 0; col_sh[i] = -1; } // pads sort last (keys >= 1) + __syncthreads(); + + // bitonic sort, descending + for (int size = 2; size <= n; size <<= 1) { + for (int stride = size >> 1; stride > 0; stride >>= 1) { + for (int i = tid; i < n; i += NT) { + const int j = i ^ stride; + if (j > i) { + const bool desc = (i & size) == 0; + const uint32_t a = key_sh[i], b = key_sh[j]; + if (desc ? a < b : a > b) { + key_sh[i] = b; key_sh[j] = a; + const int t = col_sh[i]; col_sh[i] = col_sh[j]; col_sh[j] = t; + } + } + } + __syncthreads(); + } + } + for (int i = tid; i < k; i += NT) { + const int c = col_sh[i]; + cols[(long long)blockIdx.x * k + i] = c; + vals[(long long)blockIdx.x * k + i] = row[c]; + } +} + +} // namespace dsv41_topk_keys diff --git a/zig/kernels/cuda/deepseek_v41/x3gm.cu b/zig/kernels/cuda/deepseek_v41/x3gm.cu new file mode 100644 index 0000000000..3798f4e298 --- /dev/null +++ b/zig/kernels/cuda/deepseek_v41/x3gm.cu @@ -0,0 +1,477 @@ +// Copy of src/tensorfold/families/deepseek_v41/cuda/x3gm.cu (blob 7aed2b027afc): sass_check.sh + + + +// x3gm: routed EXL3 prefill experts (gate/up, down); each output is one mma chain over ascending k from +0, any tiling + +#include +#include + +#include "decode.cuh" +#include "experts_grouped.cuh" + +namespace dsv41_x3gm { + +using tf_exl3x::cb_pair; +using tf_exl3x::Fmt; +using tf_exl3x::LaneMap; +using tf_exl3x::mma16816; + +constexpr float HAD = 0.08838834764831845f; // 1 / sqrt(128) +constexpr int NB = 8; // column tiles an item: 128 columns, one Hadamard block +constexpr int MUL1 = 2; + +__device__ __forceinline__ void cp_async16(void* smem, const void* gmem, bool ok) { + const unsigned s = (unsigned)__cvta_generic_to_shared(smem); + asm volatile("cp.async.cg.shared.global [%0], [%1], 16, %2;\n" ::"r"(s), "l"(gmem), "r"(ok ? 16 : 0)); +} +__device__ __forceinline__ void cp_commit() { asm volatile("cp.async.commit_group;\n" ::); } +template +__device__ __forceinline__ void cp_wait() { asm volatile("cp.async.wait_group %0;\n" ::"n"(N)); } +__device__ __forceinline__ void ldsm4(uint32_t (&a)[4], const void* p) { + const unsigned s = (unsigned)__cvta_generic_to_shared(p); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0,%1,%2,%3}, [%4];\n" + : "=r"(a[0]), "=r"(a[1]), "=r"(a[2]), "=r"(a[3]) + : "r"(s)); +} + +// the 16-byte chunk c of member row r in a stage (CPR chunks a row): 8 rows of an ldmatrix phase in 8 bank groups +template +__device__ __forceinline__ int swz(int r, int c) { + static_assert(CPR == 8 || CPR == 4, "rows of 64 or 32 halves"); + return CPR == 8 ? (c ^ (r & 7)) : (c ^ ((r >> 1) & 3)); +} + +// A 16x16 W_q tile from shared words as the A fragment of W^T: upstream decode_tile, registers reordered +template +__device__ __forceinline__ void decode_a(const uint32_t* tile, const LaneMap& m, uint32_t (&a)[4]) { + constexpr int GV = Fmt::GV, NG = Fmt::NG; + uint32_t st[8]; +#pragma unroll + for (int g = 0; g < NG; ++g) { + const uint32_t whi = tile[m.hi[g]], wlo = tile[m.lo[g]]; + const uint64_t mm = ((((uint64_t)wlo) << 32) | whi) >> m.sh[g]; +#pragma unroll + for (int j = 0; j < GV; ++j) st[g * GV + j] = (uint32_t)(mm >> Fmt::off(j)) & 0xffffu; + } + a[0] = cb_pair(st[0], st[1]); // b0[0]: k 2t, 2t+1 of column g -> A row g, k 2t.. + a[1] = cb_pair(st[4], st[5]); // b1[0]: k 2t, 2t+1 of column g + 8 -> A row g + 8, k 2t.. + a[2] = cb_pair(st[2], st[3]); // b0[1]: k 2t+8, 2t+9 of column g -> A row g, k 2t+8.. + a[3] = cb_pair(st[6], st[7]); // b1[1] -> A row g + 8, k 2t+8.. +} + +template +struct Cfg { + static constexpr int BM = NG * 8, WPM = NB / MTL, W = MATS * WPM, THREADS = W * 32; + static constexpr int CPR = KS * 2, LDA = KS * 16, TW = 4 * K2; + static constexpr int LDE = 132, ER = NGR * 8, ROUNDS = NG / NGR; + static constexpr size_t A_BYTES = (size_t)XM * BM * LDA * sizeof(half); + static constexpr size_t W_BYTES = (size_t)MATS * KS * NB * TW * sizeof(uint32_t); + static constexpr size_t STAGE = A_BYTES + W_BYTES; + static constexpr size_t RING = (size_t)NSA * STAGE; + static constexpr size_t EPI = (size_t)MATS * ER * LDE * sizeof(float); + static constexpr size_t SMEM = RING > EPI ? RING : EPI; + // CTAs an SM the registers are capped for: 2 when two rings fit, <= 8 warps, <= 64 accumulators a lane; else 1 + static constexpr int MINB = 2 * (SMEM + 1536) <= 102400 && THREADS <= 256 && MTL * NG <= 16 ? 2 : 1; + static_assert(NG % 2 == 0 && NG % NGR == 0 && NGR % 2 == 0 && NB % MTL == 0, "tiling"); + static_assert(A_BYTES % 16 == 0 && (TW * 4) % 16 == 0, "16-byte stages"); + static_assert(XM == 1 || XM == MATS, "one shared input or one a matrix"); +}; + +struct Args { + const half* x0; // [P, K] rows of matrix 0 (gate / down input) + const half* x1; // [P, K] rows of matrix 1 (up) when XM == 2 + const uint32_t* t0; // [E, K/16, N/16, 4 K2] trellis words, stacked (when tp0 is null) + const uint32_t* t1; + const int64_t* tp0; // [E] expert e's words [K/16, N/16, 4 K2] (a ragged stack), or null + const int64_t* tp1; + const int* k2e; // [E] each expert's K2 (gm2_kernel; the launch's width table) + const int* order; // pair indices by expert (stable: rows ascending) + const int* pe; // pass -> expert + const int* poff; // pass -> first index into order + const int* pcnt; // pass -> members (<= BM) + const int* npass; // [1] passes in all + int* ticket; // [1] zeroed, or null (static stride) + const half* sv0; // gate/up: svh_g [E, N]; down: svh_d [E, N] + const half* sv1; // svh_u [E, N] + const half* sd; // suh_d [E, N] + half* xd; // gate/up out [P, N] + float* y; // down out [P, N] + int K, N; + float limit; +}; + +template +__global__ void __launch_bounds__(Cfg::THREADS, + Cfg::MINB) gm_kernel(const Args a) { + using C = Cfg; + constexpr int THREADS = C::THREADS, BM = C::BM, LDA = C::LDA, LDE = C::LDE, ER = C::ER, CPR = C::CPR, + TW = C::TW, W = C::W; + extern __shared__ __align__(16) unsigned char smem[]; + float* ep = reinterpret_cast(smem); // [MATS][ER][LDE] after the K loop (aliases the ring) + __shared__ int rows_sh[BM]; + __shared__ int item_sh; + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int mat = warp / C::WPM, slice = warp % C::WPM; // columns slice 16 MTL .. + 16 MTL of the item + const int xmat = XM == 1 ? 0 : mat; + const int g = lane >> 2, t = lane & 3; + const int K = a.K, N = a.N, KT = K >> 4, NTILES = N >> 4, S = KT / KS, NBLK = N >> 7; + const size_t se = (size_t)KT * NTILES * TW; // words of one expert's matrix + const int total = a.npass[0] * NBLK; + const int lr = (lane & 7) + ((lane >> 4) << 3), lc = (lane >> 3) & 1; + const LaneMap map(lane); + + int next = blockIdx.x; + for (;;) { + __syncthreads(); // the last item is done with rows_sh and shared memory + if (threadIdx.x == 0) item_sh = a.ticket ? atomicAdd(a.ticket, 1) : next; + next += gridDim.x; + __syncthreads(); + const int item = item_sh; + if (item >= total) break; + const int tt = item / NBLK, nb = item % NBLK; + const int e = a.pe[tt], off = a.poff[tt], cnt = a.pcnt[tt]; + const uint32_t* te0 = a.tp0 ? reinterpret_cast(a.tp0[e]) : a.t0 + (size_t)e * se; + const uint32_t* te1 = a.tp1 ? reinterpret_cast(a.tp1[e]) : a.t1 + (size_t)e * se; + for (int i = threadIdx.x; i < BM; i += THREADS) rows_sh[i] = i < cnt ? a.order[off + i] : -1; + __syncthreads(); + const int cnt16 = (cnt + 15) & ~15; // rows past it are never read (warp-uniform skip) + + // stage s: member rows' k [16 KS s, + 16 KS) (zero past the count), then the item's k tile words, block nb + auto load_stage = [&](int s) { + unsigned char* st = smem + (size_t)(s % NSA) * C::STAGE; + half* ra = reinterpret_cast(st); + const int k0 = s * KS * 16; + for (int i = threadIdx.x; i < XM * BM * CPR; i += THREADS) { + const int m = i / (BM * CPR), r = (i / CPR) % BM, c = i % CPR; + if (r < cnt16) { + const int row = rows_sh[r]; + const half* src = (m ? a.x1 : a.x0) + (size_t)(row < 0 ? 0 : row) * K + k0 + c * 8; + cp_async16(ra + ((size_t)m * BM + r) * LDA + swz(r, c) * 8, src, row >= 0); + } + } + uint32_t* rw = reinterpret_cast(st + C::A_BYTES); + constexpr int WCH = MATS * KS * NB * K2; // 16-byte chunks: K2 a tile + for (int i = threadIdx.x; i < WCH; i += THREADS) { + const int q = i % K2, n = (i / K2) % NB, kk = (i / (K2 * NB)) % KS, m = i / (K2 * NB * KS); + const uint32_t* src = (m ? te1 : te0) + ((size_t)(s * KS + kk) * NTILES + nb * NB + n) * TW + q * 4; + cp_async16(rw + ((m * KS + kk) * NB + n) * TW + q * 4, src, true); + } + }; + + float acc[MTL][NG][4]; +#pragma unroll + for (int l = 0; l < MTL; ++l) +#pragma unroll + for (int n = 0; n < NG; ++n) +#pragma unroll + for (int c = 0; c < 4; ++c) acc[l][n][c] = 0.f; + +#pragma unroll + for (int s = 0; s < NSA - 1; ++s) { + if (s < S) load_stage(s); + cp_commit(); + } + cp_wait(); + __syncthreads(); + for (int s = 0; s < S; ++s) { + if (s + NSA - 1 < S) load_stage(s + NSA - 1); + cp_commit(); + const unsigned char* st = smem + (size_t)(s % NSA) * C::STAGE; + const half* A = reinterpret_cast(st) + (size_t)xmat * BM * LDA; + const uint32_t* Wd = reinterpret_cast(st + C::A_BYTES) + (size_t)mat * KS * NB * TW + + slice * MTL * TW; +#pragma unroll + for (int kk = 0; kk < KS; ++kk) { + uint32_t af[MTL][4]; +#pragma unroll + for (int l = 0; l < MTL; ++l) decode_a(Wd + (kk * NB + l) * TW, map, af[l]); + const int chunk = swz(lr, kk * 2 + lc); +#pragma unroll + for (int np = 0; np < NG / 2; ++np) { + if (16 * np < cnt) { // warp-uniform + uint32_t b[4]; + ldsm4(b, A + (16 * np + lr) * LDA + chunk * 8); + const uint32_t b01[2] = {b[0], b[1]}, b23[2] = {b[2], b[3]}; +#pragma unroll + for (int l = 0; l < MTL; ++l) { + mma16816(acc[l][2 * np], af[l], b01); + mma16816(acc[l][2 * np + 1], af[l], b23); + } + } + } + } + cp_wait(); + __syncthreads(); + } + + // epilogue: ER members a round, accumulators -> ep[mat][member][column], then one warp a member row +#pragma unroll + for (int rd = 0; rd < C::ROUNDS; ++rd) { + if (rd) __syncthreads(); + if (ER * rd < cnt) { +#pragma unroll + for (int n = 0; n < NGR; ++n) { + const int ng = rd * NGR + n; + float* E = ep + ((size_t)mat * ER + 8 * n + 2 * t) * LDE + slice * 16 * MTL + g; +#pragma unroll + for (int l = 0; l < MTL; ++l) { + E[l * 16] = acc[l][ng][0]; + E[LDE + l * 16] = acc[l][ng][1]; + E[l * 16 + 8] = acc[l][ng][2]; + E[LDE + l * 16 + 8] = acc[l][ng][3]; + } + } + } + __syncthreads(); + for (int r = warp; r < ER && ER * rd + r < cnt; r += W) { + const int row = rows_sh[ER * rd + r]; + const int c0 = 4 * lane; + const size_t col = (size_t)e * N + nb * 128 + c0; + float v[4]; + *reinterpret_cast(v) = *reinterpret_cast(ep + (size_t)r * LDE + c0); + tf_exl3::fwht128(v, lane); + if constexpr (MATS == 2) { + float u[4]; + *reinterpret_cast(u) = *reinterpret_cast(ep + ((size_t)ER + r) * LDE + c0); + tf_exl3::fwht128(u, lane); +#pragma unroll + for (int j = 0; j < 4; ++j) { // upstream gateup_epilogue, ACT_F32 + const float gg = fminf(v[j] * HAD * __half2float(a.sv0[col + j]), a.limit); + const float uu = fminf(fmaxf(u[j] * HAD * __half2float(a.sv1[col + j]), -a.limit), a.limit); + const float act = gg / (1.f + expf(-gg)) * uu; + v[j] = act * __half2float(a.sd[col + j]); + } + tf_exl3::fwht128(v, lane); + half2* o = reinterpret_cast(a.xd + (size_t)row * N + nb * 128 + c0); + o[0] = __halves2half2(__float2half_rn(v[0] * HAD), __float2half_rn(v[1] * HAD)); + o[1] = __halves2half2(__float2half_rn(v[2] * HAD), __float2half_rn(v[3] * HAD)); + } else { + float4 o; // upstream down_epilogue + o.x = v[0] * HAD * __half2float(a.sv0[col + 0]); + o.y = v[1] * HAD * __half2float(a.sv0[col + 1]); + o.z = v[2] * HAD * __half2float(a.sv0[col + 2]); + o.w = v[3] * HAD * __half2float(a.sv0[col + 3]); + *reinterpret_cast(a.y + (size_t)row * N + nb * 128 + c0) = o; + } + } + } + } +} + +#include "x3gm2.cuh" + +// Rotated routed inputs: out_m[p] = fp16((x[p / slots] * suh_m[e]) H / sqrt 128), upstream rot_in's arithmetic +template __device__ __forceinline__ float to_f(TIN v); +template <> __device__ __forceinline__ float to_f<__nv_bfloat16>(__nv_bfloat16 v) { return __bfloat162float(v); } +template <> __device__ __forceinline__ float to_f(half v) { return __half2float(v); } + +template +__global__ void __launch_bounds__(256) rot_kernel(const TIN* __restrict__ x, int x_stride, const int* __restrict__ pick, + const half* __restrict__ suh0, const half* __restrict__ suh1, + half* __restrict__ out0, half* __restrict__ out1, int K, int slots, + int items) { + const int item = blockIdx.x * 8 + (threadIdx.x >> 5); + if (item >= items) return; // warp-uniform + const int KB = K >> 7, p = item / KB, blk = item % KB, lane = threadIdx.x & 31; + const int e = pick[p]; + const TIN* xr = x + (size_t)(p / slots) * x_stride + blk * 128 + 4 * lane; + float xv[4]; +#pragma unroll + for (int j = 0; j < 4; ++j) xv[j] = to_f(xr[j]); +#pragma unroll + for (int m = 0; m < MATS; ++m) { + const half* suh = (m ? suh1 : suh0) + (size_t)e * K + blk * 128 + 4 * lane; + float v[4]; +#pragma unroll + for (int j = 0; j < 4; ++j) v[j] = xv[j] * __half2float(suh[j]); + tf_exl3::fwht128(v, lane); + half* o = (m ? out1 : out0) + (size_t)p * K + blk * 128 + 4 * lane; +#pragma unroll + for (int j = 0; j < 4; ++j) o[j] = __float2half_rn(v[j] * HAD); + } +} + +// Tests: W_q [K, N] fp16 of one trellis through decode_a (words staged in shared memory as the gm kernels stage them). +template +__global__ void __launch_bounds__(32) dequant_kernel(const uint32_t* __restrict__ T, half* __restrict__ out, int N) { + constexpr int TW = 4 * K2; + __shared__ __align__(16) uint32_t tile[TW]; + const int kt = blockIdx.x, nt = blockIdx.y, lane = threadIdx.x, NTILES = N >> 4; + for (int i = lane; i < TW; i += 32) tile[i] = T[((size_t)kt * NTILES + nt) * TW + i]; + __syncwarp(); + const LaneMap map(lane); + uint32_t af[4]; + decode_a(tile, map, af); + const int g = lane >> 2, t = lane & 3; +#pragma unroll + for (int q = 0; q < 4; ++q) { // A row (q & 1 ? g + 8 : g) = W column; k 2t (+8) + const half2 h = *reinterpret_cast(&af[q]); + const int col = nt * 16 + g + 8 * (q & 1), row = kt * 16 + 2 * t + 8 * (q >> 1); + out[(size_t)row * N + col] = __low2half(h); + out[(size_t)(row + 1) * N + col] = __high2half(h); + } +} + + +// Tile table (MTL, NG, KS, NSA, NGR), mirrored by the host tile lists (a test parses these lines); 2.. = 4K tiles +#define GM_GU0 2, 8, 4, 3, 4 +#define GM_GU1 2, 8, 2, 4, 4 +#define GM_GU2 1, 16, 2, 4, 4 +#define GM_GU3 1, 16, 4, 3, 4 +#define GM_GU4 2, 16, 2, 3, 4 +#define GM_GU5 2, 8, 2, 6, 4 +#define GM_DN0 2, 8, 4, 4, 8 +#define GM_DN1 1, 8, 4, 4, 8 +#define GM_DN2 1, 16, 4, 4, 8 +#define GM_DN3 1, 16, 2, 6, 8 +#define GM_DN4 2, 16, 4, 3, 16 +#define GM_NGU 6 +#define GM_NDN 5 + + +} // namespace dsv41_x3gm + +#ifndef DSV41_X3GM_NO_INSTANCES +// Instances the bindings launch: gate/up (XM 1 shared input, 2 own) GM_GU0-5, down GM_DN0-4, K2 3-8 and 10 +#define DSV41_X3GM_GU(XM, K2, CFG) template __global__ void dsv41_x3gm::gm_kernel<2, XM, K2, CFG>(const dsv41_x3gm::Args); +#define DSV41_X3GM_DN(K2, CFG) template __global__ void dsv41_x3gm::gm_kernel<1, 1, K2, CFG>(const dsv41_x3gm::Args); +// v2 (gm2_kernel: every width in one launch): gate/up GM_GU0 / GM_GU1 at either input mode, down GM_DN0 +template __global__ void dsv41_x3gm::gm2_kernel<2, 1, GM_GU0>(const dsv41_x3gm::Args); +template __global__ void dsv41_x3gm::gm2_kernel<2, 1, GM_GU1>(const dsv41_x3gm::Args); +template __global__ void dsv41_x3gm::gm2_kernel<2, 2, GM_GU0>(const dsv41_x3gm::Args); +template __global__ void dsv41_x3gm::gm2_kernel<2, 2, GM_GU1>(const dsv41_x3gm::Args); +template __global__ void dsv41_x3gm::gm2_kernel<1, 1, GM_DN0>(const dsv41_x3gm::Args); +DSV41_X3GM_GU(1, 3, GM_GU0) +DSV41_X3GM_GU(1, 3, GM_GU1) +DSV41_X3GM_GU(1, 3, GM_GU2) +DSV41_X3GM_GU(1, 3, GM_GU3) +DSV41_X3GM_GU(1, 3, GM_GU4) +DSV41_X3GM_GU(1, 3, GM_GU5) +DSV41_X3GM_GU(2, 3, GM_GU0) +DSV41_X3GM_GU(2, 3, GM_GU1) +DSV41_X3GM_GU(2, 3, GM_GU2) +DSV41_X3GM_GU(2, 3, GM_GU3) +DSV41_X3GM_GU(2, 3, GM_GU4) +DSV41_X3GM_GU(2, 3, GM_GU5) +DSV41_X3GM_DN(3, GM_DN0) +DSV41_X3GM_DN(3, GM_DN1) +DSV41_X3GM_DN(3, GM_DN2) +DSV41_X3GM_DN(3, GM_DN3) +DSV41_X3GM_DN(3, GM_DN4) +DSV41_X3GM_GU(1, 4, GM_GU0) +DSV41_X3GM_GU(1, 4, GM_GU1) +DSV41_X3GM_GU(1, 4, GM_GU2) +DSV41_X3GM_GU(1, 4, GM_GU3) +DSV41_X3GM_GU(1, 4, GM_GU4) +DSV41_X3GM_GU(1, 4, GM_GU5) +DSV41_X3GM_GU(2, 4, GM_GU0) +DSV41_X3GM_GU(2, 4, GM_GU1) +DSV41_X3GM_GU(2, 4, GM_GU2) +DSV41_X3GM_GU(2, 4, GM_GU3) +DSV41_X3GM_GU(2, 4, GM_GU4) +DSV41_X3GM_GU(2, 4, GM_GU5) +DSV41_X3GM_DN(4, GM_DN0) +DSV41_X3GM_DN(4, GM_DN1) +DSV41_X3GM_DN(4, GM_DN2) +DSV41_X3GM_DN(4, GM_DN3) +DSV41_X3GM_DN(4, GM_DN4) +DSV41_X3GM_GU(1, 5, GM_GU0) +DSV41_X3GM_GU(1, 5, GM_GU1) +DSV41_X3GM_GU(1, 5, GM_GU2) +DSV41_X3GM_GU(1, 5, GM_GU3) +DSV41_X3GM_GU(1, 5, GM_GU4) +DSV41_X3GM_GU(1, 5, GM_GU5) +DSV41_X3GM_GU(2, 5, GM_GU0) +DSV41_X3GM_GU(2, 5, GM_GU1) +DSV41_X3GM_GU(2, 5, GM_GU2) +DSV41_X3GM_GU(2, 5, GM_GU3) +DSV41_X3GM_GU(2, 5, GM_GU4) +DSV41_X3GM_GU(2, 5, GM_GU5) +DSV41_X3GM_DN(5, GM_DN0) +DSV41_X3GM_DN(5, GM_DN1) +DSV41_X3GM_DN(5, GM_DN2) +DSV41_X3GM_DN(5, GM_DN3) +DSV41_X3GM_DN(5, GM_DN4) +DSV41_X3GM_GU(1, 6, GM_GU0) +DSV41_X3GM_GU(1, 6, GM_GU1) +DSV41_X3GM_GU(1, 6, GM_GU2) +DSV41_X3GM_GU(1, 6, GM_GU3) +DSV41_X3GM_GU(1, 6, GM_GU4) +DSV41_X3GM_GU(1, 6, GM_GU5) +DSV41_X3GM_GU(2, 6, GM_GU0) +DSV41_X3GM_GU(2, 6, GM_GU1) +DSV41_X3GM_GU(2, 6, GM_GU2) +DSV41_X3GM_GU(2, 6, GM_GU3) +DSV41_X3GM_GU(2, 6, GM_GU4) +DSV41_X3GM_GU(2, 6, GM_GU5) +DSV41_X3GM_DN(6, GM_DN0) +DSV41_X3GM_DN(6, GM_DN1) +DSV41_X3GM_DN(6, GM_DN2) +DSV41_X3GM_DN(6, GM_DN3) +DSV41_X3GM_DN(6, GM_DN4) +DSV41_X3GM_GU(1, 7, GM_GU0) +DSV41_X3GM_GU(1, 7, GM_GU1) +DSV41_X3GM_GU(1, 7, GM_GU2) +DSV41_X3GM_GU(1, 7, GM_GU3) +DSV41_X3GM_GU(1, 7, GM_GU4) +DSV41_X3GM_GU(1, 7, GM_GU5) +DSV41_X3GM_GU(2, 7, GM_GU0) +DSV41_X3GM_GU(2, 7, GM_GU1) +DSV41_X3GM_GU(2, 7, GM_GU2) +DSV41_X3GM_GU(2, 7, GM_GU3) +DSV41_X3GM_GU(2, 7, GM_GU4) +DSV41_X3GM_GU(2, 7, GM_GU5) +DSV41_X3GM_DN(7, GM_DN0) +DSV41_X3GM_DN(7, GM_DN1) +DSV41_X3GM_DN(7, GM_DN2) +DSV41_X3GM_DN(7, GM_DN3) +DSV41_X3GM_DN(7, GM_DN4) +DSV41_X3GM_GU(1, 8, GM_GU0) +DSV41_X3GM_GU(1, 8, GM_GU1) +DSV41_X3GM_GU(1, 8, GM_GU2) +DSV41_X3GM_GU(1, 8, GM_GU3) +DSV41_X3GM_GU(1, 8, GM_GU4) +DSV41_X3GM_GU(1, 8, GM_GU5) +DSV41_X3GM_GU(2, 8, GM_GU0) +DSV41_X3GM_GU(2, 8, GM_GU1) +DSV41_X3GM_GU(2, 8, GM_GU2) +DSV41_X3GM_GU(2, 8, GM_GU3) +DSV41_X3GM_GU(2, 8, GM_GU4) +DSV41_X3GM_GU(2, 8, GM_GU5) +DSV41_X3GM_DN(8, GM_DN0) +DSV41_X3GM_DN(8, GM_DN1) +DSV41_X3GM_DN(8, GM_DN2) +DSV41_X3GM_DN(8, GM_DN3) +DSV41_X3GM_DN(8, GM_DN4) +DSV41_X3GM_GU(1, 10, GM_GU0) +DSV41_X3GM_GU(1, 10, GM_GU1) +DSV41_X3GM_GU(1, 10, GM_GU2) +DSV41_X3GM_GU(1, 10, GM_GU3) +DSV41_X3GM_GU(1, 10, GM_GU4) +DSV41_X3GM_GU(1, 10, GM_GU5) +DSV41_X3GM_GU(2, 10, GM_GU0) +DSV41_X3GM_GU(2, 10, GM_GU1) +DSV41_X3GM_GU(2, 10, GM_GU2) +DSV41_X3GM_GU(2, 10, GM_GU3) +DSV41_X3GM_GU(2, 10, GM_GU4) +DSV41_X3GM_GU(2, 10, GM_GU5) +DSV41_X3GM_DN(10, GM_DN0) +DSV41_X3GM_DN(10, GM_DN1) +DSV41_X3GM_DN(10, GM_DN2) +DSV41_X3GM_DN(10, GM_DN3) +DSV41_X3GM_DN(10, GM_DN4) +#define DSV41_X3GM_ROT(T, M) template __global__ void dsv41_x3gm::rot_kernel(const T*, int, const int*, const half*, \ + const half*, half*, half*, int, int, int); +DSV41_X3GM_ROT(__nv_bfloat16, 1) +DSV41_X3GM_ROT(__nv_bfloat16, 2) +DSV41_X3GM_ROT(half, 1) +DSV41_X3GM_ROT(half, 2) +template __global__ void dsv41_x3gm::dequant_kernel<3>(const uint32_t*, half*, int); +template __global__ void dsv41_x3gm::dequant_kernel<4>(const uint32_t*, half*, int); +template __global__ void dsv41_x3gm::dequant_kernel<5>(const uint32_t*, half*, int); +template __global__ void dsv41_x3gm::dequant_kernel<6>(const uint32_t*, half*, int); +template __global__ void dsv41_x3gm::dequant_kernel<7>(const uint32_t*, half*, int); +template __global__ void dsv41_x3gm::dequant_kernel<8>(const uint32_t*, half*, int); +template __global__ void dsv41_x3gm::dequant_kernel<10>(const uint32_t*, half*, int); +#endif // DSV41_X3GM_NO_INSTANCES diff --git a/zig/kernels/cuda/deepseek_v41/x3gm2.cuh b/zig/kernels/cuda/deepseek_v41/x3gm2.cuh new file mode 100644 index 0000000000..526092524d --- /dev/null +++ b/zig/kernels/cuda/deepseek_v41/x3gm2.cuh @@ -0,0 +1,256 @@ +// x3gm.cu's v2 kernel (TF_DSV41_GM_V2): gm2_item, kloop2 and gm2_kernel, inside namespace dsv41_x3gm +#pragma once +// x3gm v2 (TF_DSV41_GM_V2): gm_kernel's items, chains and bits with less issue work (member groups, width per item) + +__device__ __forceinline__ void ldsm2(uint32_t (&a)[2], const void* p) { + const unsigned s = (unsigned)__cvta_generic_to_shared(p); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0,%1}, [%2];\n" : "=r"(a[0]), "=r"(a[1]) : "r"(s)); +} + +// The K loop of one item with N8 active n8 member groups (acc[l][0 .. N8 - 1]; the rest stay 0.0, never stored) +template +__device__ __forceinline__ void kloop2(float (&acc)[MTL][NG][4], const unsigned char* smem, int S, int xmat, int mat, + int slice, int lr, int lc, const LaneMap& map, LoadFn&& load_stage) { + using C = Cfg; + constexpr int BM = C::BM, LDA = C::LDA, CPR = C::CPR, TW = C::TW, PAIRS = N8 / 2; +#pragma unroll + for (int s = 0; s < NSA - 1; ++s) { + if (s < S) load_stage(s); + cp_commit(); + } + cp_wait(); + __syncthreads(); + for (int s = 0; s < S; ++s) { + if (s + NSA - 1 < S) load_stage(s + NSA - 1); + cp_commit(); + const unsigned char* st = smem + (size_t)(s % NSA) * C::STAGE; + const half* A = reinterpret_cast(st) + (size_t)xmat * BM * LDA; + const uint32_t* Wd = reinterpret_cast(st + C::A_BYTES) + (size_t)mat * KS * NB * TW + + slice * MTL * TW; +#pragma unroll + for (int kk = 0; kk < KS; ++kk) { + uint32_t af[MTL][4]; +#pragma unroll + for (int l = 0; l < MTL; ++l) decode_a(Wd + (kk * NB + l) * TW, map, af[l]); + const int chunk = swz(lr, kk * 2 + lc); +#pragma unroll + for (int np = 0; np < PAIRS; ++np) { + uint32_t b[4]; + ldsm4(b, A + (16 * np + lr) * LDA + chunk * 8); + const uint32_t b01[2] = {b[0], b[1]}, b23[2] = {b[2], b[3]}; +#pragma unroll + for (int l = 0; l < MTL; ++l) { + mma16816(acc[l][2 * np], af[l], b01); + mma16816(acc[l][2 * np + 1], af[l], b23); + } + } + if constexpr ((N8 & 1) != 0) { // members 16 PAIRS .. + 7 (x2: lanes 0-15's addresses) + uint32_t b[2]; + ldsm2(b, A + (16 * PAIRS + lr) * LDA + chunk * 8); +#pragma unroll + for (int l = 0; l < MTL; ++l) mma16816(acc[l][2 * PAIRS], af[l], b); + } + } + cp_wait(); + __syncthreads(); + } +} + +// One item (pass x 128-column block) at width K2, rows_sh filled by the caller; otherwise gm_kernel's body +template +__device__ __forceinline__ void gm2_item(const Args& a, unsigned char* smem, const int* rows_sh, int e, int cnt, + int nb) { + using C = Cfg; + constexpr int THREADS = C::THREADS, BM = C::BM, LDA = C::LDA, LDE = C::LDE, ER = C::ER, CPR = C::CPR, + TW = C::TW, W = C::W; + constexpr int ACH = XM * BM * CPR, WCH = MATS * KS * NB * K2; // 16-byte chunks a stage: member rows, words + constexpr int APT = (ACH + THREADS - 1) / THREADS, WPT = (WCH + THREADS - 1) / THREADS; + float* ep = reinterpret_cast(smem); // [MATS][ER][LDE] after the K loop (aliases the ring) + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int mat = warp / C::WPM, slice = warp % C::WPM; + const int xmat = XM == 1 ? 0 : mat; + const int g = lane >> 2, t = lane & 3; + const int K = a.K, N = a.N, KT = K >> 4, NTILES = N >> 4, S = KT / KS; + const size_t se = (size_t)KT * NTILES * TW; + const int lr = (lane & 7) + ((lane >> 4) << 3), lc = (lane >> 3) & 1; + const LaneMap map(lane); + const uint32_t* te0 = a.tp0 ? reinterpret_cast(a.tp0[e]) : a.t0 + (size_t)e * se; + const uint32_t* te1 = a.tp1 ? reinterpret_cast(a.tp1[e]) : a.t1 + (size_t)e * se; + const int cnt16 = (cnt + 15) & ~15; // rows past it are never read + + // this thread's chunks: row chunk element offset at k 0 (-1: none, -2: zero fill), word chunk offset at stage 0 + int a_src[APT], w_src[WPT]; +#pragma unroll + for (int j = 0; j < APT; ++j) { + const int i = threadIdx.x + j * THREADS, r = (i / CPR) % BM, c = i % CPR; + int v = -1; + if (i < ACH && r < cnt16) v = rows_sh[r] < 0 ? -2 : rows_sh[r] * K + c * 8; + a_src[j] = v; + } +#pragma unroll + for (int j = 0; j < WPT; ++j) { + const int i = threadIdx.x + j * THREADS, q = i % K2, n = (i / K2) % NB, kk = (i / (K2 * NB)) % KS; + w_src[j] = (kk * NTILES + nb * NB + n) * TW + q * 4; + } + auto load_stage = [&](int s) { + unsigned char* st = smem + (size_t)(s % NSA) * C::STAGE; + half* ra = reinterpret_cast(st); + uint32_t* rw = reinterpret_cast(st + C::A_BYTES); + const int k0 = s * KS * 16; +#pragma unroll + for (int j = 0; j < APT; ++j) { + const int i = threadIdx.x + j * THREADS, v = a_src[j]; + if (v != -1) { + const int m = i / (BM * CPR), r = (i / CPR) % BM, c = i % CPR; + cp_async16(ra + ((size_t)m * BM + r) * LDA + swz(r, c) * 8, (m ? a.x1 : a.x0) + (v < 0 ? 0 : v + k0), + v >= 0); + } + } + const size_t ws = (size_t)s * KS * NTILES * TW; +#pragma unroll + for (int j = 0; j < WPT; ++j) { + const int i = threadIdx.x + j * THREADS; + if (WCH % THREADS == 0 || i < WCH) { + const int q = i % K2, n = (i / K2) % NB, kk = (i / (K2 * NB)) % KS, m = i / (K2 * NB * KS); + cp_async16(rw + ((m * KS + kk) * NB + n) * TW + q * 4, (m ? te1 : te0) + ws + w_src[j], true); + } + } + }; + + float acc[MTL][NG][4]; +#pragma unroll + for (int l = 0; l < MTL; ++l) +#pragma unroll + for (int n = 0; n < NG; ++n) +#pragma unroll + for (int c = 0; c < 4; ++c) acc[l][n][c] = 0.f; + + static_assert(NG <= 16 && NG % 2 == 0, "n8 groups"); + const int n8 = (cnt + 7) >> 3; // >= 1 (a pass has members) +#define GM2_K(G) kloop2(acc, smem, S, xmat, mat, slice, lr, lc, map, load_stage) + switch (n8 <= 8 ? n8 : (n8 + 1) & ~1) { // 1..8, then even counts (NG 16) + case 1: GM2_K(1); break; + case 2: GM2_K(2); break; + case 3: GM2_K(3); break; + case 4: GM2_K(4); break; + case 5: if constexpr (NG >= 6) GM2_K(5); break; + case 6: if constexpr (NG >= 6) GM2_K(6); break; + case 7: if constexpr (NG >= 8) GM2_K(7); break; + case 8: if constexpr (NG >= 8) GM2_K(8); break; + case 10: if constexpr (NG >= 10) GM2_K(10); break; + case 12: if constexpr (NG >= 12) GM2_K(12); break; + case 14: if constexpr (NG >= 14) GM2_K(14); break; + default: if constexpr (NG >= 16) GM2_K(16); break; + } +#undef GM2_K + + // epilogue: gm_kernel's +#pragma unroll + for (int rd = 0; rd < C::ROUNDS; ++rd) { + if (rd) __syncthreads(); + if (ER * rd < cnt) { +#pragma unroll + for (int n = 0; n < NGR; ++n) { + const int ng = rd * NGR + n; + float* E = ep + ((size_t)mat * ER + 8 * n + 2 * t) * LDE + slice * 16 * MTL + g; +#pragma unroll + for (int l = 0; l < MTL; ++l) { + E[l * 16] = acc[l][ng][0]; + E[LDE + l * 16] = acc[l][ng][1]; + E[l * 16 + 8] = acc[l][ng][2]; + E[LDE + l * 16 + 8] = acc[l][ng][3]; + } + } + } + __syncthreads(); + for (int r = warp; r < ER && ER * rd + r < cnt; r += W) { + const int row = rows_sh[ER * rd + r]; + const int c0 = 4 * lane; + const size_t col = (size_t)e * N + nb * 128 + c0; + float v[4]; + *reinterpret_cast(v) = *reinterpret_cast(ep + (size_t)r * LDE + c0); + tf_exl3::fwht128(v, lane); + if constexpr (MATS == 2) { + float u[4]; + *reinterpret_cast(u) = *reinterpret_cast(ep + ((size_t)ER + r) * LDE + c0); + tf_exl3::fwht128(u, lane); +#pragma unroll + for (int j = 0; j < 4; ++j) { // upstream gateup_epilogue, ACT_F32 + const float gg = fminf(v[j] * HAD * __half2float(a.sv0[col + j]), a.limit); + const float uu = fminf(fmaxf(u[j] * HAD * __half2float(a.sv1[col + j]), -a.limit), a.limit); + const float act = gg / (1.f + expf(-gg)) * uu; + v[j] = act * __half2float(a.sd[col + j]); + } + tf_exl3::fwht128(v, lane); + half2* o = reinterpret_cast(a.xd + (size_t)row * N + nb * 128 + c0); + o[0] = __halves2half2(__float2half_rn(v[0] * HAD), __float2half_rn(v[1] * HAD)); + o[1] = __halves2half2(__float2half_rn(v[2] * HAD), __float2half_rn(v[3] * HAD)); + } else { + float4 o; // upstream down_epilogue + o.x = v[0] * HAD * __half2float(a.sv0[col + 0]); + o.y = v[1] * HAD * __half2float(a.sv0[col + 1]); + o.z = v[2] * HAD * __half2float(a.sv0[col + 2]); + o.w = v[3] * HAD * __half2float(a.sv0[col + 3]); + *reinterpret_cast(a.y + (size_t)row * N + nb * 128 + c0) = o; + } + } + } +} + +// The widths a v2 launch dispatches and, per width, the ring depth (cut where it would cost a second CTA an SM) +#define GM2_WIDTHS(X) X(3) X(4) X(5) X(6) X(7) X(8) X(10) +constexpr int GM2_NWIDTHS = 7; +template +struct Ring2 { + static constexpr bool fits2(int n) { + return 2 * (Cfg::STAGE * (size_t)n + 1536) <= 102400; + } + static constexpr int cut(int n) { return n <= 2 || fits2(n) ? n : cut(n - 1); } + static constexpr int value = Cfg::MINB == 2 ? cut(NSA) : NSA; +}; +template +struct Smem2 { +#define GM2_SM(K) Cfg::value, NGR>::SMEM, + static constexpr size_t list[GM2_NWIDTHS] = {GM2_WIDTHS(GM2_SM)}; +#undef GM2_SM + static constexpr size_t max_from(int i) { + return i == GM2_NWIDTHS ? 0 : (list[i] > max_from(i + 1) ? list[i] : max_from(i + 1)); + } + static constexpr size_t SMEM = max_from(0); + static constexpr int THREADS = Cfg::THREADS; + static constexpr int MINB = 2 * (SMEM + 1536) <= 102400 && THREADS <= 256 && MTL * NG <= 16 ? 2 : 1; +}; + +// One v2 launch over the passes of every width (uniform stack: one width; ragged: each expert its own, a.k2e) +template +__global__ void __launch_bounds__(Smem2::THREADS, + Smem2::MINB) gm2_kernel(const Args a) { + constexpr int BM = NG * 8, THREADS = Smem2::THREADS; + extern __shared__ __align__(16) unsigned char smem[]; + __shared__ int rows_sh[BM]; + __shared__ int item_sh; + const int NBLK = a.N >> 7, total = a.npass[0] * NBLK; + int next = blockIdx.x; + for (;;) { + __syncthreads(); // the last item is done with rows_sh and shared memory + if (threadIdx.x == 0) item_sh = a.ticket ? atomicAdd(a.ticket, 1) : next; + next += gridDim.x; + __syncthreads(); + const int item = item_sh; + if (item >= total) break; + const int tt = item / NBLK, nb = item % NBLK; + const int e = a.pe[tt], off = a.poff[tt], cnt = a.pcnt[tt]; + for (int i = threadIdx.x; i < BM; i += THREADS) rows_sh[i] = i < cnt ? a.order[off + i] : -1; + __syncthreads(); + switch (a.k2e[e]) { // block-uniform +#define GM2_CASE(K) \ + case K: \ + gm2_item::value, NGR>(a, smem, rows_sh, e, \ + cnt, nb); \ + break; + GM2_WIDTHS(GM2_CASE) +#undef GM2_CASE + default: break; // never: the host checks every width is in K2S + } + } +} diff --git a/zig/kernels/cuda/deepseek_v41/x3gm_plan.cu b/zig/kernels/cuda/deepseek_v41/x3gm_plan.cu new file mode 100644 index 0000000000..ee7355a061 --- /dev/null +++ b/zig/kernels/cuda/deepseek_v41/x3gm_plan.cu @@ -0,0 +1,104 @@ +// x3gm's pass tables on the device in one launch: same integers as the torch ops of x3gm.plan. + +// order = stable argsort(pick); cnt/off/npe/pend per expert; pass t -> expert e, poff, pcnt; npass = pend[E - 1]. + +// One CTA (E < 512): histogram, scan, stable chunked scatter (per-warp counts + match_any); integers only, so exact. + +#include + +namespace dsv41_x3gm_plan { + +constexpr int NT = 512, NW = NT / 32, EMAX = 512; // E <= 511: the routers take <= 512 experts (GE_MAX) + +__global__ void __launch_bounds__(NT) plan_kernel(const int* __restrict__ pick, int P, int E, int bm, + int* __restrict__ order, int* __restrict__ pe, + int* __restrict__ poff, int* __restrict__ pcnt, + int* __restrict__ npass) { + __shared__ int cnt[EMAX + 1]; // counts of values 0..E, then (scatter) the moving bases + __shared__ int pend[EMAX]; // inclusive scan of passes an expert + __shared__ int wc[NW][EMAX + 1]; // a chunk's per-warp counts of each value (scatter) + __shared__ int scan_tmp[NT]; + const int tid = threadIdx.x, lane = tid & 31, warp = tid >> 5, V = E + 1; + + for (int v = tid; v < V; v += NT) cnt[v] = 0; + __syncthreads(); + for (int i = tid; i < P; i += NT) atomicAdd(&cnt[pick[i]], 1); // counts: order-free + __syncthreads(); + + // exclusive scans: bases of values 0..E and pend over the first E; V <= NT * 2, two elements a thread + auto block_scan = [&](int x) { // inclusive scan of one value a thread + scan_tmp[tid] = x; + __syncthreads(); + for (int o = 1; o < NT; o <<= 1) { + const int y = tid >= o ? scan_tmp[tid - o] : 0; + __syncthreads(); + scan_tmp[tid] += y; + __syncthreads(); + } + const int r = scan_tmp[tid]; + __syncthreads(); + return r; + }; + int c0 = 2 * tid < V ? cnt[2 * tid] : 0, c1 = 2 * tid + 1 < V ? cnt[2 * tid + 1] : 0; + const int b_end = block_scan(c0 + c1); // inclusive over pairs of values + const int n0 = 2 * tid < E ? (c0 + bm - 1) / bm : 0, n1 = 2 * tid + 1 < E ? (c1 + bm - 1) / bm : 0; + const int p_end = block_scan(n0 + n1); + __syncthreads(); + if (2 * tid < V) cnt[2 * tid] = b_end - c0 - c1; // base of value 2 tid + if (2 * tid + 1 < V) cnt[2 * tid + 1] = b_end - c1; + if (2 * tid < E) pend[2 * tid] = p_end - n1; + if (2 * tid + 1 < E) pend[2 * tid + 1] = p_end; + __syncthreads(); + + // pass tables (counts recomputed from the bases: base[v + 1] - base[v], base[E + 1] = P) + const int tmax = P / bm + E + 1; + const int npe_last = E > 0 ? pend[E - 1] : 0; + for (int t = tid; t < tmax; t += NT) { + int lo = 0, hi = E; // searchsorted(pend, t, right): first pend > t + while (lo < hi) { + const int mid = (lo + hi) >> 1; + if (pend[mid] <= t) lo = mid + 1; else hi = mid; + } + const int e = lo < E - 1 ? lo : E - 1; + const int ce = (e + 1 < V ? cnt[e + 1] : P) - cnt[e]; + const int npe_e = (ce + bm - 1) / bm; + const long long j = (long long)t - (pend[e] - npe_e); + long long o = (long long)cnt[e] + j * bm; + const long long cap = P - 1 > 0 ? P - 1 : 0; + if (o > cap) o = cap; + long long c = (long long)ce - j * bm; + c = c < 0 ? 0 : (c > bm ? bm : c); + pe[t] = e; + poff[t] = (int)o; + pcnt[t] = (int)c; + } + if (tid == 0) npass[0] = npe_last; + __syncthreads(); + + // stable scatter, NT pairs a chunk in index order + for (int base = 0; base < P; base += NT) { + for (int i = tid; i < NW * V; i += NT) (&wc[0][0])[(i / V) * (EMAX + 1) + i % V] = 0; + __syncthreads(); + const int i = base + tid; + const bool live = i < P; + const int v = live ? pick[i] : -1; + const unsigned same = __match_any_sync(0xffffffffu, v); + const int before = __popc(same & ((1u << lane) - 1u)); // equal values in earlier lanes of the warp + if (live && before == 0) wc[warp][v] = __popc(same); // the warp's count of v, by its first lane + __syncthreads(); + if (live) { + int slot = cnt[v] + before; + for (int w = 0; w < warp; ++w) slot += wc[w][v]; + order[slot] = i; + } + __syncthreads(); + for (int u = tid; u < V; u += NT) { // the bases move past this chunk's pairs + int s = 0; + for (int w = 0; w < NW; ++w) s += wc[w][u]; + cnt[u] += s; + } + __syncthreads(); + } +} + +} // namespace dsv41_x3gm_plan diff --git a/zig/kernels/cuda/deepseek_v41/x3ld.cu b/zig/kernels/cuda/deepseek_v41/x3ld.cu new file mode 100644 index 0000000000..15bdfb4fcc --- /dev/null +++ b/zig/kernels/cuda/deepseek_v41/x3ld.cu @@ -0,0 +1,249 @@ +// Copy of src/tensorfold/families/deepseek_v41/cuda/x3ld.cu (blob 5cd25f2cdf3c): sass_check.sh + + + +// DeepSeek-V4.1-Flash routed experts (TF_DSV41_EXPERT_LOADS=1): grouped EXL3 expert GEMV with a faster load path. + +// Same bits as grouped_kernel at any row count; data movement changes (prologue, ring, PDL). + +// PROBE 3 (timing only, wrong results): no decode and no mma, the load path alone. + +#include "experts_grouped.cuh" + +namespace dsv41_x3ld { + +using tf_exl3x::decode_tile; +using tf_exl3x::Fmt; +using tf_exl3x::LaneMap; +using tf_exl3x::mma16816; + +constexpr int W = 4; // warps a CTA: upstream's default_config, which fixes the K ranges + +__device__ __forceinline__ uint4 ldg_v4(const uint32_t* p) { + uint4 v; + asm volatile("ld.global.nc.L1::no_allocate.v4.u32 {%0,%1,%2,%3}, [%4];" + : "=r"(v.x), "=r"(v.y), "=r"(v.z), "=r"(v.w) + : "l"(p)); + return v; +} + +// load_pair's value (ok ? *(const uint32_t*)x : 0), the load pinned where it stands +__device__ __forceinline__ uint32_t ldg_pair(const half* x, bool ok) { + uint32_t v = 0u; + if (ok) asm volatile("ld.global.nc.u32 %0, [%1];" : "=r"(v) : "l"(x)); + return v; +} + +__device__ __forceinline__ void pdl_wait() { +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + asm volatile("griddepcontrol.wait;" ::: "memory"); +#endif +} + +template +struct Step { + static constexpr int TW = Fmt::TW; // words a tile + static constexpr int LW = Fmt::LW; // words a lane a tile (load_words) + static constexpr int WORDS = NT * TW; // a warp's k step, contiguous + static constexpr int NV = (WORDS + 127) / 128; // 16-byte loads a lane a step +}; + +// One warp's k tiles [kt0, kt0 + nkt) into acc (warp_tiles' arithmetic, loads PD ahead); all CTA threads call it. +template +__device__ __forceinline__ void ld_tiles(const uint32_t* __restrict__ T, int NTILES, int kt0, int nkt, int nt0, + const half* __restrict__ X, int K, const int* rows_sh, uint32_t* stage, + int lane, float (&acc)[NT][2][4]) { + using S = Step; + const int g = lane >> 2, t = lane & 3; + const size_t kstride = (size_t)NTILES * S::TW; + const uint32_t* tp = T + ((size_t)kt0 * NTILES + nt0) * S::TW; + + uint4 ring[PD][S::NV]; + auto issue_w = [&](int d, int s) { + const uint32_t* src = tp + (size_t)s * kstride; +#pragma unroll + for (int v = 0; v < S::NV; ++v) { + const int off = 128 * v + 4 * lane; + if (off < S::WORDS) ring[d][v] = ldg_v4(src + off); + } + }; +#pragma unroll + for (int d = 0; d < PD; ++d) issue_w(d, d); // PD <= nkt (host) + + pdl_wait(); // X (and Z) belong to the kernels before + __syncthreads(); // rows_sh + const int r0 = rows_sh[g], r1 = rows_sh[g + 8]; + const half* x0 = X + (size_t)(r0 < 0 ? 0 : r0) * K + 2 * t; + const half* x1 = X + (size_t)(r1 < 0 ? 0 : r1) * K + 2 * t; + uint32_t a_ring[PD][4]; + auto issue_a = [&](int d, int s) { + const int k = (kt0 + s) * 16; + a_ring[d][0] = ldg_pair(x0 + k, r0 >= 0); + a_ring[d][1] = ldg_pair(x1 + k, r1 >= 0); + a_ring[d][2] = ldg_pair(x0 + k + 8, r0 >= 0); + a_ring[d][3] = ldg_pair(x1 + k + 8, r1 >= 0); + }; +#pragma unroll + for (int d = 0; d < PD; ++d) issue_a(d, d); + + const LaneMap map(lane); + for (int kb = 0; kb < nkt; kb += PD) { // nkt % PD == 0 (host) +#pragma unroll + for (int d = 0; d < PD; ++d) { + const int s = kb + d; + const bool more = s + PD < nkt; + __syncwarp(); // every lane has read step s - 1 back +#pragma unroll + for (int v = 0; v < S::NV; ++v) { + const int off = 128 * v + 4 * lane; + if (off < S::WORDS) *reinterpret_cast(stage + off) = ring[d][v]; + } + if (more) issue_w(d, s + PD); // the slot's registers are free again + __syncwarp(); + uint32_t w[NT][S::LW]; +#pragma unroll + for (int i = 0; i < NT; ++i) +#pragma unroll + for (int l = 0; l < S::LW; ++l) + w[i][l] = (S::TW % 32 == 0 || l * 32 + lane < S::TW) ? stage[i * S::TW + l * 32 + lane] : 0u; + uint32_t a[4]; +#pragma unroll + for (int j = 0; j < 4; ++j) a[j] = a_ring[d][j]; + if (more) issue_a(d, s + PD); +#pragma unroll + for (int i = 0; i < NT; ++i) { + if constexpr (PROBE == 3) { + uint32_t x = a[0] ^ a[1] ^ a[2] ^ a[3]; +#pragma unroll + for (int l = 0; l < S::LW; ++l) x ^= w[i][l]; + acc[i][0][0] = __uint_as_float(__float_as_uint(acc[i][0][0]) ^ x); + } else { + uint32_t b0[2], b1[2]; + decode_tile(w[i], map, lane, b0, b1); + mma16816(acc[i][0], a, b0); + mma16816(acc[i][1], a, b1); + } + } + } + } + __syncwarp(); // the staging area becomes red[warp] again +} + +// Program (u, n block, (mat * SK + split) * MT + member tile): grouped_kernel's work item. +template +__global__ void __launch_bounds__(W * 32, NT == 8 ? 3 : 4) ld_kernel( + const half* __restrict__ X0, const half* __restrict__ X1, const int64_t* __restrict__ TP0, + const int64_t* __restrict__ TP1, const int* __restrict__ K2_0, const int* __restrict__ K2_1, + const int* __restrict__ uids, const int* __restrict__ ucount, const int* __restrict__ members, + float* __restrict__ Z, int K, int N, int P, int SK, int maxm, int slots) { + static_assert(PD >= 1 && PD <= 4, "PD 1-4"); + const int u = blockIdx.x; + const int MT = (maxm + 15) / 16; + const int mtile = blockIdx.z % MT; + const int split = (blockIdx.z / MT) % SK; + const int mat = blockIdx.z / MT / SK; + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + + // round trip 1: none of these depends on another (uids / members are sized by the grid: always in bounds) + const int cnt = __ldcg(ucount); + const int e = __ldcg(uids + u); + const int m = mtile * 16 + threadIdx.x; + const int code = threadIdx.x < 16 && m < maxm ? __ldcg(members + u * maxm + m) : -1; + const int first = __ldcg(members + u * maxm + mtile * 16); + if (u >= cnt || first < 0) return; // grouped_kernel's exits (CTA-uniform) + + __shared__ int rows_sh[16]; + __shared__ __align__(16) float red[W][16][NT * 16]; + if (threadIdx.x < 16) rows_sh[threadIdx.x] = code >= 0 ? (code >> 5) * slots + (code & 31) : -1; + + // round trip 2: the expert's trellis and width + const uint32_t* T = reinterpret_cast(mat ? TP1[e] : TP0[e]); + const int k2 = mat ? K2_1[e] : K2_0[e]; + const half* X = mat ? X1 : X0; + const int KT = K >> 4, NTILES = N >> 4; + const int per_split = KT / SK, per_warp = per_split / W; + const int kt0 = split * per_split + warp * per_warp; + const int nt0 = blockIdx.y * NT; + uint32_t* stage = reinterpret_cast(&red[warp][0][0]); + + float acc[NT][2][4]; +#pragma unroll + for (int i = 0; i < NT; ++i) +#pragma unroll + for (int h = 0; h < 2; ++h) +#pragma unroll + for (int c = 0; c < 4; ++c) acc[i][h][c] = 0.f; + + switch (k2) { +#define DSV41_X3LD_CASE(K2_) \ + case K2_: \ + if constexpr (K2_ >= LO && K2_ <= HI) \ + ld_tiles(T, NTILES, kt0, per_warp, nt0, X, K, rows_sh, stage, lane, acc); \ + else \ + __trap(); \ + break; + DSV41_X3LD_CASE(2) + DSV41_X3LD_CASE(3) + DSV41_X3LD_CASE(4) + DSV41_X3LD_CASE(5) + DSV41_X3LD_CASE(6) + DSV41_X3LD_CASE(7) + DSV41_X3LD_CASE(8) + DSV41_X3LD_CASE(9) + DSV41_X3LD_CASE(10) + DSV41_X3LD_CASE(11) + DSV41_X3LD_CASE(12) + DSV41_X3LD_CASE(13) + DSV41_X3LD_CASE(14) + DSV41_X3LD_CASE(15) + DSV41_X3LD_CASE(16) +#undef DSV41_X3LD_CASE + default: + __trap(); + } + + // grouped_kernel's reduction, verbatim: the warps' partial sums through shared memory, added in warp order + const int g = lane >> 2, t = lane & 3; +#pragma unroll + for (int i = 0; i < NT; ++i) +#pragma unroll + for (int h = 0; h < 2; ++h) { + const int col = i * 16 + h * 8 + 2 * t; + red[warp][g][col] = acc[i][h][0]; + red[warp][g][col + 1] = acc[i][h][1]; + red[warp][g + 8][col] = acc[i][h][2]; + red[warp][g + 8][col + 1] = acc[i][h][3]; + } + __syncthreads(); + for (int idx = threadIdx.x; idx < 16 * NT * 16; idx += W * 32) { + const int row = idx / (NT * 16), col = idx % (NT * 16); + const int r = rows_sh[row]; + if (r < 0) continue; + float s = red[0][row][col]; +#pragma unroll + for (int w = 1; w < W; ++w) s += red[w][row][col]; + Z[(((size_t)mat * SK + split) * P + r) * N + nt0 * 16 + col] = s; + } +} + +} // namespace dsv41_x3ld + +// The instantiations dsv41_x3ld_grouped_cuda dispatches: mul1, (nt, pd) x probe x the (lo, hi) ranges. +#define DSV41_X3LD_INST(NT, PD, LO, HI, PROBE) template __global__ void dsv41_x3ld::ld_kernel<2, NT, PD, LO, HI, PROBE>( \ + const half*, const half*, const int64_t*, const int64_t*, const int*, const int*, const int*, const int*, \ + const int*, float*, int, int, int, int, int, int); +DSV41_X3LD_INST(8, 1, 8, 8, 0) +DSV41_X3LD_INST(8, 1, 2, 10, 0) +DSV41_X3LD_INST(8, 1, 2, 12, 0) +DSV41_X3LD_INST(8, 2, 8, 8, 0) +DSV41_X3LD_INST(8, 2, 2, 10, 0) +DSV41_X3LD_INST(8, 2, 2, 12, 0) +DSV41_X3LD_INST(4, 2, 8, 8, 0) +DSV41_X3LD_INST(4, 2, 2, 10, 0) +DSV41_X3LD_INST(4, 2, 2, 12, 0) +DSV41_X3LD_INST(8, 1, 8, 8, 3) +DSV41_X3LD_INST(8, 1, 2, 10, 3) +DSV41_X3LD_INST(8, 1, 2, 12, 3) +DSV41_X3LD_INST(8, 2, 8, 8, 3) +DSV41_X3LD_INST(8, 2, 2, 10, 3) +DSV41_X3LD_INST(8, 2, 2, 12, 3) From b17df971c61fd7f84eb0bfc38fb4600f0265c075 Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Sun, 11 Oct 2026 22:00:59 +0700 Subject: [PATCH 035/104] DeepSeek-V4.1 kernel tooling: nvcc in a container for any checkout layout. nvcc-docker.sh runs nvcc from the NGC PyTorch image for hosts without CUDA (zig build -Dnvcc=