cuda gemma4: admission counts vision and the drafter, init cleanup, DFlash config bounds, contract build fix

zig/kernels/cuda/contract_controls.cu

@@ -76,6 +76,8 @@ extern "C" __global__ void __launch_bounds__(tf_qk_round::WIDE) tf_ctl_qk_round(
 }
 
 // Gate 3: act_rows.cu's tf_softcap_rows keeping only the rounding points in `keep`; the exact check must go red.
+#include "act_round.cuh"
+
 extern "C" __global__ void __launch_bounds__(256) tf_ctl_softcap(const __nv_bfloat16* x, __nv_bfloat16* out, int n,
                                                                   float cap, int keep) {
     tf_act_round::softcap_rows(x, out, n, cap, keep);
@@ -93,8 +95,6 @@ extern "C" __global__ void __launch_bounds__(256) tf_ctl_embed_affine(const int*
 
 
 // Gate 3: act_rows.cu's MLX activations keeping only the rounding points in `keep`; the exact check must go red.
-#include "act_round.cuh"
-
 extern "C" __global__ void __launch_bounds__(256) tf_ctl_act_round(const __nv_bfloat16* gate, long long g_ld,
                                                                     const __nv_bfloat16* up, long long u_ld,
                                                                     __nv_bfloat16* out, float* xs, int n, int geglu,

zig/src/families/gemma4/cuda.zig

@@ -45,6 +45,7 @@ test {
     _ = @import("cuda_assistant_weights.zig");
     _ = @import("cuda_vision.zig");
     _ = @import("cuda_dflash.zig");
+    _ = @import("cuda_dflash_test.zig");
     _ = @import("cuda_tower.zig");
     _ = @import("cuda_image.zig");
     _ = @import("cuda_native.zig");

zig/src/families/gemma4/cuda_dflash.zig

@@ -50,7 +50,7 @@ pub const Config = struct {
         const archs = o.get("architectures") orelse return error.BadDraftConfig;
         if (archs != .array or archs.array.items.len != 1 or archs.array.items[0] != .string or !std.mem.eql(u8, archs.array.items[0].string, "DFlashDraftModel")) return error.NotDFlashV1;
         const d = if (o.get("dflash_config")) |v| (if (v == .object) v.object else return error.BadDraftConfig) else return error.BadDraftConfig;
-        var c: Config = .{ .layers = @intCast(try integer(o, "num_hidden_layers")), .hidden = @intCast(try integer(o, "hidden_size")), .intermediate = @intCast(try integer(o, "intermediate_size")), .heads = @intCast(try integer(o, "num_attention_heads")), .kv_heads = @intCast(try integer(o, "num_key_value_heads")), .head_dim = @intCast(try integer(o, "head_dim")), .block = @intCast(if (d.get("block_size") != null) try integer(d, "block_size") else try integer(o, "block_size")), .mask = @intCast(try integer(d, "mask_token_id")), .window = @intCast(try integer(o, "sliding_window")), .eps = @floatCast(try number(o, "rms_norm_eps")), .theta = try number(o, "rope_theta") };
+        var c: Config = .{ .layers = try size(o, "num_hidden_layers"), .hidden = try size(o, "hidden_size"), .intermediate = try size(o, "intermediate_size"), .heads = try size(o, "num_attention_heads"), .kv_heads = try size(o, "num_key_value_heads"), .head_dim = try size(o, "head_dim"), .block = try size(if (d.get("block_size") != null) d else o, "block_size"), .mask = try integer(u32, d, "mask_token_id"), .window = try size(o, "sliding_window"), .eps = @floatCast(try number(o, "rms_norm_eps")), .theta = try number(o, "rope_theta") };
         if (c.hidden != geo.hidden) return error.DraftForAnotherTarget;
         const taps = d.get("target_layer_ids") orelse return error.BadDraftConfig;
         if (taps != .array or taps.array.items.len == 0 or taps.array.items.len > 8) return error.BadDraftConfig;
@@ -66,16 +66,22 @@ pub const Config = struct {
             for (kinds.array.items, 0..) |k, i| c.sliding[i] = k == .string and std.mem.eql(u8, k.string, "sliding_attention");
         }
         if (o.get("is_causal") != null) return error.UnsupportedDraftMask; // the vendored default decides the block's mask
-        if (c.head_dim != dim or c.heads % c.kv_heads != 0 or c.heads / c.kv_heads > 8 or c.block != block_rows or c.window < 32) return error.UnsupportedDraftShape;
+        if (c.head_dim != dim or c.kv_heads == 0 or c.heads % c.kv_heads != 0 or c.heads / c.kv_heads > 8 or c.block != block_rows or c.window < 32) return error.UnsupportedDraftShape;
         if (c.hidden % 64 != 0 or c.intermediate % 64 != 0 or (c.heads * c.head_dim) % 64 != 0 or (c.kv_heads * c.head_dim) % 64 != 0) return error.UnsupportedDraftShape;
         if (c.mask >= cfg.vocab) return error.BadDraftConfig;
         return c;
     }
 };
 
-fn integer(o: std.json.ObjectMap, key: []const u8) !i64 {
+fn integer(comptime T: type, o: std.json.ObjectMap, key: []const u8) !T {
     const v = o.get(key) orelse return error.BadDraftConfig;
-    return if (v == .integer) v.integer else error.BadDraftConfig;
+    return if (v == .integer) std.math.cast(T, v.integer) orelse error.BadDraftConfig else error.BadDraftConfig;
+}
+
+/// A count: positive and within 32 bits, so no product of two of them overflows.
+fn size(o: std.json.ObjectMap, key: []const u8) !usize {
+    const n = try integer(u32, o, key);
+    return if (n == 0) error.BadDraftConfig else n;
 }
 
 fn number(o: std.json.ObjectMap, key: []const u8) !f64 {
@@ -176,10 +182,12 @@ pub const Drafter = struct {
         const x = try gpa.create(Drafter);
         errdefer gpa.destroy(x);
         x.* = .{ .gpa = gpa, .io = io, .d = gb.d, .m = m, .k = gb.k, .gb = gb, .s = gb.s, .c = c, .w = .{ .fc = undefined, .hidden_norm = 0, .norm = 0, .layers = undefined }, .ba = undefined };
+        errdefer x.w.buffers.deinit(gpa);
         errdefer for (x.w.buffers.items) |*b| b.free();
         try load(gpa, io, x.d, x.s, x.k, dir, &x.c, &x.w);
         x.ba = try cuda.block_attention.BlockAttention.load(x.d);
         errdefer x.ba.unload();
+        errdefer x.buffers.deinit(gpa);
         errdefer for (x.buffers.items) |*b| b.free();
         const h = c.hidden;
         const qw = c.heads * dim;

zig/src/families/gemma4/cuda_dflash_test.zig

@@ -0,0 +1,25 @@
+//! The DFlash drafter's config parse on hostile values: errors, never a trap.
+const std = @import("std");
+const cfg = @import("cuda_config.zig");
+const dfl = @import("cuda_dflash.zig");
+
+const stock =
+    \\{"architectures":["DFlashDraftModel"],"dflash_config":{"mask_token_id":4,"target_layer_ids":[1,12,23,35,46,57]},"block_size":16,
+    \\"hidden_size":5376,"intermediate_size":10752,"num_hidden_layers":5,"num_attention_heads":64,"num_key_value_heads":8,"head_dim":128,
+    \\"sliding_window":2048,"rms_norm_eps":1e-6,"rope_theta":1000000}
+;
+
+fn refused(from: []const u8, to: []const u8) !void {
+    const text = try std.mem.replaceOwned(u8, std.testing.allocator, stock, from, to);
+    defer std.testing.allocator.free(text);
+    try std.testing.expectError(error.BadDraftConfig, dfl.Config.parse(std.testing.allocator, text, cfg.dense31b));
+}
+
+test "a zero, negative or oversized count is a bad config, not a trap" {
+    try refused("\"num_key_value_heads\":8", "\"num_key_value_heads\":0");
+    try refused("\"hidden_size\":5376", "\"hidden_size\":-1");
+    try refused("\"block_size\":16", "\"block_size\":1099511627776");
+    try refused("\"num_attention_heads\":64", "\"num_attention_heads\":9223372036854775807");
+    try refused("\"sliding_window\":2048", "\"sliding_window\":0");
+    try refused("\"mask_token_id\":4", "\"mask_token_id\":-4");
+}

zig/src/families/gemma4/cuda_native.zig

@@ -44,7 +44,7 @@ const native = @This();
 /// The server's reading of the room left after the weights; null if unreadable.
 pub const Room = struct { ctx: *anyopaque, read: *const fn (*anyopaque) ?u64 };
 /// No `precision` field: the server refuses --precision for this family by name. --vision: the 12B's image embedder.
-pub const Options = struct { context: usize, drafts: bool, segments: usize = 1, slide: bool = false, lanes: u32 = 8, lanes_fixed: bool = false, room: ?Room = null, vision: ?cv.Limits = null };
+pub const Options = struct { context: usize, drafts: bool, segments: usize = 1, slide: bool = false, lanes: u32 = 8, lanes_fixed: bool = false, room: ?Room = null, vision: ?cv.Limits = null, drafter: ?[]const u8 = null };
 
 pub const Loaded = struct {
     backend: lanes.backend.Backend,
@@ -59,7 +59,7 @@ pub const Loaded = struct {
     vision: ?*cv.Vision = null,
 };
 
-const Owned = struct { gpa: std.mem.Allocator, w: wts.Weights, k: kern.Kernels, s: cuda.Stream, m: fwd.Model, gb: be.Gemma, v: ?*cv.Vision = null };
+const Owned = struct { gpa: std.mem.Allocator, w: wts.Weights, k: kern.Kernels, s: cuda.Stream, m: fwd.Model, gb: be.Gemma, v: ?*cv.Vision = null, room: ?Room = null };
 
 /// Streams to allocate in `room`: as asked, else (auto) the most that fit; a fixed count past them or a round refuses.
 fn fit(room: u64, asked: u32, fixed: bool, plan: anytype) !usize {
@@ -76,13 +76,33 @@ const Plan = struct {
     own: *const Owned,
     d: *const cuda.Driver,
     context: usize,
+    /// What open allocates beside the streams: the vision workspace and the drafter's weights.
+    beside: u64,
 
     fn bytes(p: Plan, n: usize) !u64 {
         const o = p.own;
-        return try fwd.Model.openBytes(o.gpa, p.d, o.s, &o.w, p.context, n) + be.Gemma.openBytes(n);
+        return try fwd.Model.openBytes(o.gpa, p.d, o.s, &o.w, p.context, n) + be.Gemma.openBytes(n) + p.beside;
     }
 };
 
+/// The device bytes of the image path's workspace (the 31B's tower or the 12B's embedder), 0 without --vision.
+fn visionBytes(vision: bool, tower: bool) u64 {
+    if (!vision) return 0;
+    return if (tower) tw.Tower.workspaceBytes() else cv.Embedder.workspaceBytes();
+}
+
+/// The drafter's weights on the device: the checkpoint's files, a DFlash drafter's bf16 linears made 4-bit (5/16).
+fn drafterBytes(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) !u64 {
+    var d = try std.Io.Dir.cwd().openDir(io, dir, .{ .iterate = true });
+    defer d.close(io);
+    var total: u64 = 0;
+    var it = d.iterate();
+    while (try it.next(io)) |e| {
+        if (std.mem.endsWith(u8, e.name, ".safetensors")) total += (try d.statFile(io, e.name, .{})).size;
+    }
+    return if (try dfl.isDFlash(gpa, io, dir)) total / 16 * 5 else total;
+}
+
 /// Positions a stream's caches hold: whole 1,024s.
 fn positions(context: usize) usize {
     return (context + 1023) / 1024 * 1024;
@@ -101,7 +121,9 @@ pub fn open(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: [
     own.w = try wts.load(gpa, io, ctx.d, own.s, &own.k, dir, null, o.vision != null);
     errdefer own.w.deinit();
     const context = @min(positions(o.context), own.w.config.context);
-    const plan: Plan = .{ .own = own, .d = ctx.d, .context = context };
+    const beside = visionBytes(o.vision != null, own.w.tower != null) + if (o.drafter) |dd| (if (o.drafts) try drafterBytes(gpa, io, dd) else 0) else 0;
+    const plan: Plan = .{ .own = own, .d = ctx.d, .context = context, .beside = beside };
+    own.room = o.room;
     const room: u64 = if (o.room) |r| r.read(r.ctx) orelse 0 else std.math.maxInt(u64);
     const n = fit(room, o.lanes, o.lanes_fixed, plan) catch |e| {
         if (e == error.ParallelPastMemory) {
@@ -184,6 +206,7 @@ pub fn explain(_: ?*anyopaque, err: anyerror) ?[]const u8 {
         error.CapturedKernels => "Gemma 4 runs on the engine's own kernels: unset TENSORFOLD_CUDA_KERNELS",
         error.ParallelPastRound => std.fmt.comptimePrint("--parallel: Gemma 4 on CUDA holds at most {d} streams (a round's rows): serve --parallel {d} or fewer, or auto", .{ be.round_rows, be.round_rows }),
         error.ParallelPastMemory => "--parallel: the CUDA memory budget fits fewer Gemma 4 streams than asked (each stream's caches are allocated at load): lower --parallel or --context, serve --parallel auto, or adjust TENSORFOLD_MEMORY_RESERVE_GIB / TENSORFOLD_CUDA_MEMORY_LIMIT_GB",
+        error.DrafterPastMemory => "--drafter: the CUDA memory budget leaves less room than the drafter's weights need after the streams: lower --parallel or --context, or adjust TENSORFOLD_MEMORY_RESERVE_GIB / TENSORFOLD_CUDA_MEMORY_LIMIT_GB",
         error.NoStreamFits => "the CUDA memory budget fits no Gemma 4 stream after the weights: lower --context, free device memory, or adjust TENSORFOLD_MEMORY_RESERVE_GIB / TENSORFOLD_CUDA_MEMORY_LIMIT_GB",
         else => null,
     };
@@ -194,6 +217,7 @@ pub fn openDrafter(gpa: std.mem.Allocator, io: std.Io, p: *anyopaque, dir: []con
     if (bits != 4) return error.DrafterBitsUnsupported;
     const own: *Owned = @ptrCast(@alignCast(p));
     if (own.gb.lanes.count() != 0) return error.TapsAfterStreams;
+    if (own.room) |r| if (r.read(r.ctx)) |left| if (left < try drafterBytes(gpa, io, dir)) return error.DrafterPastMemory;
     if (try dfl.isDFlash(gpa, io, dir)) {
         const d = try dfl.Drafter.init(gpa, io, &own.gb, dir);
         errdefer d.deinit();
@@ -243,7 +267,7 @@ fn release(p: *anyopaque) void {
 }
 
 test "every refusal of its own is worded, the rest left to the server" {
-    for ([_]anyerror{ error.PastContext, error.NoFreeSlot, error.NoDraftHead, error.StructuresUnsupported, error.NonFiniteLogits, error.CapturedKernels, error.ParallelPastRound, error.ParallelPastMemory, error.NoStreamFits, error.GemmaDrafterNotServed, error.AssistantForAnotherTarget, error.DrafterBitsUnsupported, error.NoGemmaVision, error.GemmaVisionDenseQuantized }) |e| {
+    for ([_]anyerror{ error.PastContext, error.NoFreeSlot, error.NoDraftHead, error.StructuresUnsupported, error.NonFiniteLogits, error.CapturedKernels, error.ParallelPastRound, error.ParallelPastMemory, error.NoStreamFits, error.DrafterPastMemory, error.GemmaDrafterNotServed, error.AssistantForAnotherTarget, error.DrafterBitsUnsupported, error.NoGemmaVision, error.GemmaVisionDenseQuantized }) |e| {
         try std.testing.expect(explain(null, e) != null);
     }
     try std.testing.expect(explain(null, error.OutOfMemory) == null);
@@ -266,3 +290,15 @@ test "the streams allocated at open fit the room the weights leave: auto takes w
     try std.testing.expectError(error.NoStreamFits, fit(g - 1, 8, false, one));
     try std.testing.expectError(error.ParallelPastRound, fit(1024 * g, be.round_rows + 1, true, one));
 }
+
+test "admission counts the vision workspace: more bytes per open, fewer streams under a tight room" {
+    const g: u64 = 1 << 30;
+    try std.testing.expectEqual(@as(u64, 0), visionBytes(false, true));
+    try std.testing.expect(visionBytes(true, true) > 0 and visionBytes(true, false) > 0);
+    const plain: Fake = .{ .stream = g, .fixed = visionBytes(false, true) };
+    const with: Fake = .{ .stream = g, .fixed = visionBytes(true, true) };
+    try std.testing.expect(try with.bytes(2) > try plain.bytes(2));
+    const room = 4 * g + visionBytes(true, true) / 2;
+    try std.testing.expectEqual(@as(usize, 4), try fit(room, 8, false, plain));
+    try std.testing.expectEqual(@as(usize, 3), try fit(room, 8, false, with));
+}

zig/src/families/gemma4/cuda_tower.zig

@@ -33,6 +33,7 @@ pub const Tower = struct {
         errdefer t.attn.unload();
         t.gemms = try vit.Gemms.init(gpa, d);
         errdefer t.gemms.deinit();
+        errdefer t.buffers.deinit(gpa);
         errdefer for (t.buffers.items) |*b| b.free();
         const p = max_patches;
         const n = image.max_budget;
@@ -58,6 +59,13 @@ pub const Tower = struct {
         return t;
     }
 
+    /// The device bytes init allocates for its scratch.
+    pub fn workspaceBytes() u64 {
+        const p: u64 = max_patches;
+        const n: u64 = image.max_budget;
+        return p * wts.patch_width * 2 + 11 * p * W * 2 + 3 * p * wts.tower_inter * 2 + 2 * p * 4 + 2 * p * D * 2 + 2 * n * W * 2 + p * W * 4;
+    }
+
     fn alloc(t: *Tower, d: *const cuda.Driver, bytes: usize) !u64 {
         const b = try cuda.DeviceBuffer.alloc(d, bytes);
         errdefer {

zig/src/families/gemma4/cuda_vision.zig

@@ -66,6 +66,7 @@ pub const Embedder = struct {
         errdefer e.vit.unload();
         e.gemms = try vit.Gemms.init(gpa, d);
         errdefer e.gemms.deinit();
+        errdefer e.buffers.deinit(gpa);
         errdefer for (e.buffers.items) |*b| b.free();
         const t = image.max_budget;
         for (&e.x) |*p| p.* = try e.alloc(d, t * wts.merged_width * 2);
@@ -75,6 +76,12 @@ pub const Embedder = struct {
         return e;
     }
 
+    /// The device bytes init allocates for its scratch.
+    pub fn workspaceBytes() u64 {
+        const t: u64 = image.max_budget;
+        return 2 * t * wts.merged_width * 2 + 6 * t * wts.vision_width * 2 + t * wts.vision_width * 4 + 2 * t * 4;
+    }
+
     fn alloc(e: *Embedder, d: *const cuda.Driver, bytes: usize) !u64 {
         const b = try cuda.DeviceBuffer.alloc(d, bytes);
         errdefer {