lanes: a family's round planner drives the lane host with prompt pieces, and a halt ends the rounds at a drain deadline (DeepSeek-V4.1 stack 6/14)

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

build.zig

@@ -5,6 +5,8 @@ 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");
+const sessions_build = @import("zig/build/sessions.zig");
 
 comptime {
     const required = std.mem.trim(u8, @embedFile(".zig-version"), "\r\n");
@@ -175,6 +177,8 @@ 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);
+    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).

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);
+}

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);
+}

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 };
 
@@ -381,6 +387,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 +410,8 @@ test {
     _ = @import("serial_host_test.zig");
     _ = @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");

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);
+}