diff --git a/zig/src/families/nemotron/cuda.zig b/zig/src/families/nemotron/cuda.zig index 4c8667fe1c..221d80757e 100644 --- a/zig/src/families/nemotron/cuda.zig +++ b/zig/src/families/nemotron/cuda.zig @@ -62,6 +62,7 @@ test { test { _ = @import("cuda_vision_test.zig"); + _ = @import("cuda_vision_lanes_test.zig"); } test { diff --git a/zig/src/families/nemotron/cuda_lanes.zig b/zig/src/families/nemotron/cuda_lanes.zig index 1479ab00c4..492846b31d 100644 --- a/zig/src/families/nemotron/cuda_lanes.zig +++ b/zig/src/families/nemotron/cuda_lanes.zig @@ -309,8 +309,6 @@ pub const Cuda = struct { e.prompt_rows = .{ .k = k, .rows = rows, .logits = if (self.prompt_logits) |b| b.ptr else e.b.logits, .picks = words[picks_at..][0..rows], .picks_dev = self.prompt_dev + picks_at * 4, .out = words[0..picks_at], .out_dev = self.prompt_dev, .stride = row_words, .ctx = &sink, .call = PromptSink.take }; } defer e.prompt_rows = null; - if (s.images.len > 0) try (e.vision orelse return error.NoVision).begin(e.stream, s.images, from); - defer if (e.vision) |v| v.end(); var tap: TapSink = .{ .self = self, .taps = &gop.value_ptr.taps }; tap.taps.prompt.clearRetainingCapacity(); tap.taps.from = from; diff --git a/zig/src/families/nemotron/cuda_vision.zig b/zig/src/families/nemotron/cuda_vision.zig index 495ee05de8..a7b00df8ca 100644 --- a/zig/src/families/nemotron/cuda_vision.zig +++ b/zig/src/families/nemotron/cuda_vision.zig @@ -9,6 +9,12 @@ pub const video = radio.video; /// What a request may hold: image tokens in all (the rows buffer), and images. pub const Limits = struct { image_tokens: u32, max_images: u32, video_tokens: u32 = 0 }; +/// Takes a prompt pass's images in the tower's place, so a host test can count them without a GPU. +pub const Encoder = struct { + ctx: *anyopaque, + encode: *const fn (ctx: *anyopaque, s: cuda.Stream, images: []const lanes.stream.Image, spans: []const radio.Span, rows: *const radio.Rows) anyerror!void, +}; + pub const Vision = struct { gpa: std.mem.Allocator, io: std.Io, @@ -19,6 +25,7 @@ pub const Vision = struct { limits: Limits, mark_ids: [3]u32, // the placeholder and its wrapper: text that spells one stays text spans: []radio.Span = &.{}, // the prompt pass in progress + encoder: ?Encoder = null, // null: the tower /// The checkpoint's tower and a rows buffer for `limits`; null with `problem` when it cannot take images. pub fn init(gpa: std.mem.Allocator, io: std.Io, d: *const cuda.Driver, dir: []const u8, limits: Limits, problem: *[]const u8) !?*Vision { @@ -64,6 +71,7 @@ pub const Vision = struct { if (images.len == 0) return; v.spans = try radio.plan(v.gpa, images, @intCast(from)); if (radio.rows.total(v.spans) > v.rows.limit) return error.ImageRows; + if (v.encoder) |x| return x.encode(x.ctx, s, images, v.spans, &v.rows); try v.tower.encode(s, images, v.spans, &v.rows); } diff --git a/zig/src/families/nemotron/cuda_vision_lanes_test.zig b/zig/src/families/nemotron/cuda_vision_lanes_test.zig new file mode 100644 index 0000000000..dd273242f6 --- /dev/null +++ b/zig/src/families/nemotron/cuda_vision_lanes_test.zig @@ -0,0 +1,45 @@ +//! The actual lane prompt pass with a counting encoder: each request sends each of its images through the tower once. +const std = @import("std"); +const cuda = @import("cuda"); +const lanes = @import("lanes"); +const radio = @import("radio"); +const vision = @import("cuda_vision.zig"); +const Engine = @import("cuda_engine.zig").Engine; +const Cuda = @import("cuda_lanes.zig").Cuda; +const a = std.testing.allocator; + +const Count = struct { + passes: usize = 0, + images: usize = 0, + + fn encode(ctx: *anyopaque, _: cuda.Stream, _: []const lanes.stream.Image, spans: []const radio.Span, _: *const radio.Rows) anyerror!void { + const c: *Count = @ptrCast(@alignCast(ctx)); + c.passes += 1; + c.images += spans.len; + } +}; + +test "a lane prompt pass encodes each image of a request once" { + for ([_]usize{ 1, 3 }) |n| { + var count: Count = .{}; + var v: vision.Vision = .{ .gpa = a, .io = undefined, .cfg = undefined, .tower = undefined, .rows = .{ .d = undefined, .buf = undefined, .hidden = 1, .limit = 64 }, .marks = undefined, .limits = undefined, .mark_ids = undefined, .encoder = .{ .ctx = &count, .encode = Count.encode } }; + var e: Engine = .{ .gpa = a, .io = undefined, .ctx = undefined, .stream = undefined, .k = undefined, .w = undefined, .b = undefined, .c = undefined, .max_len = 4096, .pinned = undefined, .history = undefined }; + // the own sequence (no device memory), too short for the prompt: the pass stops before its first chunk + e.serial = .{ .d = undefined, .handle = null }; + e.own = .{ .arena = null, .ptr = @splat(0), .cap = 1 }; + e.bound = &e.own; + e.vision = &v; + var b: Cuda = .{ .gpa = a, .e = &e, .head = null, .pinned = undefined, .rows_out = undefined, .rows_dev = 0, .prompt_out = undefined, .prompt_dev = 0 }; + defer b.lanes.deinit(a); + var images: [3]lanes.stream.Image = undefined; + for (images[0..n], 0..) |*img, i| img.* = .{ .pixels = &.{}, .gh = 4, .gw = 4, .tokens = 4, .at = @intCast(1 + i * 6) }; + const prompt: [20]u32 = @splat(1); + var s = try lanes.Stream.init(a, .{ .id = "images", .prompt = prompt[0 .. 6 * n + 2], .max_new = 4, .images = images[0..n] }); + defer s.deinit(a); + for (1..3) |request| { + try std.testing.expectError(error.ContextFull, b.backend().prefill(&s)); + try std.testing.expectEqual(request, count.passes); + try std.testing.expectEqual(request * n, count.images); + } + } +} diff --git a/zig/src/families/nemotron/cuda_vision_test.zig b/zig/src/families/nemotron/cuda_vision_test.zig index ffc7e5dbc7..da23f0652d 100644 --- a/zig/src/families/nemotron/cuda_vision_test.zig +++ b/zig/src/families/nemotron/cuda_vision_test.zig @@ -8,6 +8,7 @@ test "low independent image and video budgets reach the actual tower boundary" { var v: vision.Vision = undefined; v.gpa = std.testing.allocator; v.spans = &.{}; + v.encoder = null; defer v.end(); v.rows = .{ .d = undefined, .buf = undefined, .hidden = 4096, .limit = try vision.reservation(.{ .image_tokens = c[0], .video_tokens = c[1], .max_images = 4 }, true) }; v.tower.c.text_hidden = 1; diff --git a/zig/src/native/cuda_vision_offer.zig b/zig/src/native/cuda_vision_offer.zig index ceabf820f5..3d2b54c10a 100644 --- a/zig/src/native/cuda_vision_offer.zig +++ b/zig/src/native/cuda_vision_offer.zig @@ -158,6 +158,7 @@ test "actual offer prepares bounded video and plans mixed rows through the share var v: nemotron.native.vision.Vision = undefined; v.gpa = a; v.spans = &.{}; + v.encoder = null; defer v.end(); v.cfg = .{ .hidden = 1280, .layers = 32, .heads = 16, .ffn = 5120, .eps = 1e-6, .patch = 16, .prefix = 10, .side = 128, .projector_hidden = 20480, .text_hidden = 4096, .final_norm = false, .mean = .{ 0.5, 0.5, 0.5 }, .std = .{ 0.25, 0.25, 0.25 }, .min_patches = 1024, .max_patches = 13312, .max_model_len = 16384, .image_token = 18, .video_temporal = 2 }; v.marks = .{ .image = 18, .open = 19, .close = 20, .video = 30 };