From 6777b0a80060ed7550fe0f553ff75e88fc148c52 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Sat, 24 Oct 2026 18:01:30 +0200 Subject: [PATCH 1/4] cuda gemma4: admission counts the vision workspace and the drafter --- zig/src/families/gemma4/cuda_native.zig | 46 ++++++++++++++++++++++--- zig/src/families/gemma4/cuda_tower.zig | 7 ++++ zig/src/families/gemma4/cuda_vision.zig | 6 ++++ 3 files changed, 54 insertions(+), 5 deletions(-) diff --git a/zig/src/families/gemma4/cuda_native.zig b/zig/src/families/gemma4/cuda_native.zig index f2ed6d787f..8c459dea81 100644 --- a/zig/src/families/gemma4/cuda_native.zig +++ b/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)); +} diff --git a/zig/src/families/gemma4/cuda_tower.zig b/zig/src/families/gemma4/cuda_tower.zig index a986064ce0..5417da9cac 100644 --- a/zig/src/families/gemma4/cuda_tower.zig +++ b/zig/src/families/gemma4/cuda_tower.zig @@ -58,6 +58,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 { diff --git a/zig/src/families/gemma4/cuda_vision.zig b/zig/src/families/gemma4/cuda_vision.zig index 9535be930a..edeaf673f2 100644 --- a/zig/src/families/gemma4/cuda_vision.zig +++ b/zig/src/families/gemma4/cuda_vision.zig @@ -75,6 +75,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 { From 08212dddade8a087be52523e8e026d5e071f88e7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Sat, 24 Oct 2026 18:01:56 +0200 Subject: [PATCH 2/4] cuda contract: the softcap control includes the act_round helpers it calls --- zig/kernels/cuda/contract_controls.cu | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/zig/kernels/cuda/contract_controls.cu b/zig/kernels/cuda/contract_controls.cu index 91b3171bbf..fd66e8873c 100644 --- a/zig/kernels/cuda/contract_controls.cu +++ b/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, From 332bf8755711a1d080f41d20bad27af01e0549ac Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Sat, 24 Oct 2026 18:40:35 +0200 Subject: [PATCH 3/4] cuda gemma4: a failed init frees the buffer lists as well as the device buffers --- zig/src/families/gemma4/cuda_dflash.zig | 2 ++ zig/src/families/gemma4/cuda_tower.zig | 1 + zig/src/families/gemma4/cuda_vision.zig | 1 + 3 files changed, 4 insertions(+) diff --git a/zig/src/families/gemma4/cuda_dflash.zig b/zig/src/families/gemma4/cuda_dflash.zig index de1449a36c..45b83688d1 100644 --- a/zig/src/families/gemma4/cuda_dflash.zig +++ b/zig/src/families/gemma4/cuda_dflash.zig @@ -176,10 +176,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; diff --git a/zig/src/families/gemma4/cuda_tower.zig b/zig/src/families/gemma4/cuda_tower.zig index 5417da9cac..297892507a 100644 --- a/zig/src/families/gemma4/cuda_tower.zig +++ b/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; diff --git a/zig/src/families/gemma4/cuda_vision.zig b/zig/src/families/gemma4/cuda_vision.zig index edeaf673f2..5da7886fe7 100644 --- a/zig/src/families/gemma4/cuda_vision.zig +++ b/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); From d2f382b475171224a8eeb6c27c351256f1641cd6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Sat, 24 Oct 2026 18:41:01 +0200 Subject: [PATCH 4/4] cuda gemma4: the DFlash config is range-checked before it is cast --- zig/src/families/gemma4/cuda.zig | 1 + zig/src/families/gemma4/cuda_dflash.zig | 14 +++++++---- zig/src/families/gemma4/cuda_dflash_test.zig | 25 ++++++++++++++++++++ 3 files changed, 36 insertions(+), 4 deletions(-) create mode 100644 zig/src/families/gemma4/cuda_dflash_test.zig diff --git a/zig/src/families/gemma4/cuda.zig b/zig/src/families/gemma4/cuda.zig index 3680b78d7e..20a40b20d7 100644 --- a/zig/src/families/gemma4/cuda.zig +++ b/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"); diff --git a/zig/src/families/gemma4/cuda_dflash.zig b/zig/src/families/gemma4/cuda_dflash.zig index 45b83688d1..57476ea50c 100644 --- a/zig/src/families/gemma4/cuda_dflash.zig +++ b/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 { diff --git a/zig/src/families/gemma4/cuda_dflash_test.zig b/zig/src/families/gemma4/cuda_dflash_test.zig new file mode 100644 index 0000000000..648e78da86 --- /dev/null +++ b/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"); +}