nemotron cuda: a host test counts the RADIO tower's passes in a lane prefill (follow-up to #1057)

CHANGELOG.md

@@ -2,6 +2,10 @@
 
 Each release's page on GitHub has its notes and measurements. See [the 1.0.0 release notes](RELEASE-NOTES-1.0.0.md) for the native binary's supported models and migration details.
 
+## Unreleased
+
+- A host test runs Nemotron's CUDA lane prefill with two images through a test-only stand-in for the RADIO tower (`Vision.test_encode`) and checks that each image goes through the tower once per prompt pass; before #1057 removed the second `Vision.begin`, each went through twice. Release builds compile no hook.
+
 ## 1.1.0
 
 - Structured output and forced tool calls work with Nemotron 3.5 Super VL (`nemotron_h_omni`), whose text vocabulary is in `llm_config`: the server read `vocab_size` only from `text_config` and the top level, and refused those requests (reported in #951). A `vocab_size` that is not a positive integer within the C ABI's range is skipped rather than cast (#953). Thanks to @swbiggart and @BobClawblaw.

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;

zig/src/families/nemotron/cuda_vision.zig

@@ -10,6 +10,8 @@ pub const video = radio.video;
 pub const Limits = struct { image_tokens: u32, max_images: u32, video_tokens: u32 = 0 };
 
 pub const Vision = struct {
+    // Host fixtures replace only the tower pass; planning and the rows budget stay real.
+    pub var test_encode: if (@import("builtin").is_test) ?*const fn ([]const radio.Span) anyerror!void else void = if (@import("builtin").is_test) null else {};
     gpa: std.mem.Allocator,
     io: std.Io,
     cfg: radio.Config,
@@ -64,6 +66,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 (comptime @import("builtin").is_test) if (test_encode) |run| return run(v.spans);
         try v.tower.encode(s, images, v.spans, &v.rows);
     }
 

zig/src/families/nemotron/cuda_vision_test.zig

@@ -1,8 +1,10 @@
-//! Actual Vision.begin reaches a no-CUDA tower shape sentinel after checking the production reservation.
+//! Actual Vision.begin on the host: the reservation up to the tower's shape check, and a lane prefill's tower passes.
 const std = @import("std");
 const vision = @import("cuda_vision.zig");
 const radio = @import("radio");
 const lanes = @import("lanes");
+const Engine = @import("cuda_engine.zig").Engine;
+const Cuda = @import("cuda_lanes.zig").Cuda;
 test "low independent image and video budgets reach the actual tower boundary" {
     for ([_][3]u32{ .{ 16, 512, 256 }, .{ 256, 16384, 2048 }, .{ 256, 512, 768 } }) |c| {
         var v: vision.Vision = undefined;
@@ -23,3 +25,40 @@ test "reservation rejects overflow and keeps image-only budget" {
     try std.testing.expectError(error.BadVisionLimits, vision.reservation(.{ .image_tokens = 1, .video_tokens = std.math.maxInt(u32), .max_images = 1 }, true));
     try std.testing.expectEqual(@as(u32, 256), try vision.reservation(.{ .image_tokens = 256, .video_tokens = 512, .max_images = 1 }, false));
 }
+
+const Tower = struct {
+    var encoded: [2]u32 = @splat(0);
+    fn encode(spans: []const radio.Span) !void {
+        for (spans) |sp| encoded[sp.image] += 1;
+    }
+};
+
+test "a lane prefill runs each image of its prompt through the tower once" {
+    const a = std.testing.allocator;
+    vision.Vision.test_encode = Tower.encode;
+    defer vision.Vision.test_encode = null;
+    var v: vision.Vision = undefined;
+    v.gpa = a;
+    v.spans = &.{};
+    v.rows.limit = 8;
+    var e: Engine = undefined;
+    e.max_len = 4096;
+    e.serial = .{ .d = undefined, .handle = null };
+    e.own = .{ .arena = null, .ptr = @splat(0), .cap = 1 }; // too few rows: the pass stops before its first chunk
+    e.bound = &e.own;
+    e.head = null;
+    e.vision = &v;
+    var b: Cuda = undefined;
+    b.gpa = a;
+    b.e = &e;
+    b.lanes = .empty;
+    defer b.lanes.deinit(a);
+    b.own_free = true;
+    b.paged_only = false;
+    b.tapped = null;
+    const images = [_]lanes.stream.Image{ .{ .pixels = &.{}, .gh = 2, .gw = 6, .at = 2, .tokens = 3 }, .{ .pixels = &.{}, .gh = 2, .gw = 4, .at = 6, .tokens = 2 } };
+    var s = try lanes.Stream.init(a, .{ .id = "images", .prompt = &.{ 1, 2, 3, 3, 3, 4, 3, 3, 5, 6 }, .max_new = 4, .images = &images });
+    defer s.deinit(a);
+    try std.testing.expectError(error.ContextFull, b.backend().prefill(&s));
+    try std.testing.expectEqual([2]u32{ 1, 1 }, Tower.encoded);
+}