deepseek_v41: four-rank tensor parallel (whole-block shares, padded prefill heads)

zig/build/dsv41_cuda.zig

@@ -37,6 +37,9 @@ pub fn module(b: *std.Build, target: std.Build.ResolvedTarget, file: []const u8,
 pub fn steps(b: *std.Build) void {
     _ = @import("dsv41_images.zig").images(b, b.option([]const u8, "dsv41-nvcc", "Optional family nvcc compiler"), b.option([]const u8, "dsv41-fatbins", "Optional complete family prebuilt image directory"));
     _ = @import("dsv41_images.zig").module(b, b.graph.host, .Debug, runtime(b, b.graph.host, false), null);
+    const check_target = b.resolveTargetQuery(.{ .cpu_arch = b.graph.host.result.cpu.arch, .os_tag = .linux, .abi = .gnu });
+    const heads_check = b.addExecutable(.{ .name = "tf-dsv41-prefill-heads", .root_module = module(b, check_target, "prefill_heads_check.zig", false) });
+    b.step("dsv41-prefill-heads", "Build the AOT-only padded attention comparison for Linux, never run").dependOn(&b.addInstallArtifact(heads_check, .{}).step);
     const planner = module(b, b.graph.host, "cuda_test.zig", false);
     b.step("test-dsv41-bounds-host", "DeepSeek resource bounds and real wrapper controls without devices").dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = planner, .filters = &.{"DeepSeek bounds:"} })).step);
     const binding_tests = b.addRunArtifact(b.addTest(.{ .root_module = planner.import_table.get("dsv41_kernels").? }));

zig/src/families/deepseek_v41/aot_needs.zig

@@ -179,7 +179,7 @@ pub const Served = struct {
     /// The engine's knobs as it reads them (PF_DENSE is CUDA, so the `_gemm` variants it replaces are listed too).
     pub fn fromEnv() !Served {
         const knobs = @import("serve_knobs.zig");
-        var o: block.Options = .{ .world = @intCast(try envInt("TF_TP_WORLD", 2)), .r1 = knobs.r1(&knobs.env) };
+        var o: block.Options = .{ .world = @intCast(try envInt("TF_TP_WORLD", 2)), .rank = @intCast(try envInt("TF_TP_RANK", 0)), .r1 = knobs.r1(&knobs.env) };
         if (std.c.getenv("TF_DSV41_EXPERT_TOPP")) |v| o.expert_topp = try std.fmt.parseFloat(f64, std.mem.span(v));
         o.limit = try envInt("TF_DSV41_CONTEXT", o.limit);
         o.rope_rows = o.limit + 2048;

zig/src/families/deepseek_v41/batch_graph.zig

@@ -55,9 +55,10 @@ pub fn glueFn(ctx: *anyopaque, r: *run.Runner, c: *const calls.Call) anyerror!vo
     }
     if (std.mem.eql(u8, step, "embed")) {
         const embed = r.weights.get("embed") orelse return error.MissingWeight;
-        const V: usize = f.cfg.vocab / W;
+        const vocab_part = try @import("partition.zig").share(f.cfg.vocab, f.comm.rank(), W);
+        const V: usize = vocab_part.len;
         const rows = try f.sendBuffer(r, 2 * R * D);
-        try k.embedSend(b.tab.dev.ptr + rowtab.colOffset(.ids, b.cap), R, embed.ptr, @intCast(f.comm.rank() * V), V, D, rows);
+        try k.embedSend(b.tab.dev.ptr + rowtab.colOffset(.ids, b.cap), R, embed.ptr, @intCast(vocab_part.lo), V, D, rows);
         const recv = try r.tensorAddr(c.args[1].arg.t);
         try f.comm.allGather(rows, recv, R * D, .bf16, s);
         return k.embedSum(recv, W, R, D, try r.tensorAddr(c.args[0].arg.t));

zig/src/families/deepseek_v41/block.zig

@@ -33,8 +33,9 @@ pub const Options = struct {
     image_keep: bool = false,
     /// TF_DSV41_BIAS_VL: the segment's image rows (MoE call apart, gate.bias_vl); 0: none; < 0: the buffer plan's
     image_rows: i64 = 0,
-    /// TP ranks (a boot constant): every per-rank width below is the model's / world; the plan checked divisibility
+    /// TP world and local rank; expert and vocabulary slices keep whole 128-column blocks.
     world: u32 = 2,
+    rank: u32 = 0,
     /// KV positions a slot holds (pool rows = limit / ratio + 1)
     limit: i64 = 4096,
     /// rope table rows (cos / sin pairs of 32)

zig/src/families/deepseek_v41/block_admission.zig

@@ -1,16 +1,24 @@
 //! Decode admission bounds every scratch dimension before emitting a view or launch.
 const std = @import("std");
 const Config = @import("config.zig").Config;
+
+pub fn nativePrefill(c: *const Config, world: u32) !void {
+    if (world == 0 or c.heads % world != 0) return error.BadWorld;
+    const tile: u32 = @intCast(@import("block_prefill.zig").bmq);
+    if (c.heads / world != 16 and (c.heads / world) % tile != 0) return error.UnsupportedPrefillHeads;
+}
+
 pub fn options(c: *const Config, o: anytype, layers: []const u32, n: i64) !void {
-    if (c.expert_width % o.world != 0 or (c.expert_width / o.world) % 128 != 0) return error.InvalidBlock;
+    const rank = if (@hasField(@TypeOf(o), "rank")) o.rank else 0;
+    const part = @import("partition.zig").share(c.expert_width, rank, o.world) catch return error.InvalidBlock;
     if (o.expert_rows < 0 or o.dspark_rows < 0 or o.branch_rows < 0 or o.prefill_rows < 0 or o.image_rows < -1) return error.InvalidBlock;
     for (layers) |l| {
         if (l >= c.allLayers()) return error.InvalidBlock;
         if ((if (c.isBackbone(l)) o.expert_rows else o.dspark_rows) < n) return error.InvalidBlock;
     }
     inline for (.{ .{ o.expert_rows, @as(i64, c.topk) + 1 }, .{ o.dspark_rows, @as(i64, c.dspark_topk) + 1 } }) |cap| {
         const slots = std.math.mul(i64, cap[0], cap[1]) catch return error.InvalidBlock;
-        const width = @max(8 * @as(i64, c.expert_width / o.world), @as(i64, c.hidden));
+        const width = @max(8 * @as(i64, part.len), @as(i64, c.hidden));
         const items = std.math.mul(i64, slots, width) catch return error.InvalidBlock;
         _ = std.math.mul(i64, items, 4) catch return error.InvalidBlock;
     }

zig/src/families/deepseek_v41/block_moe.zig

@@ -142,7 +142,7 @@ pub fn moe(e: *E, L: u32) !void {
     const ex = e.cfg.expertsOf(L);
     const slots: i64 = ex.topk + 1;
     const Et: i64 = ex.count + 1;
-    const I: i64 = e.cfg.expert_width / e.o.world;
+    const I: i64 = (try @import("partition.zig").share(e.cfg.expert_width, e.o.rank, e.o.world)).len;
     const P = n * slots;
     const rows: i64 = if (backbone) e.o.expert_rows else e.o.dspark_rows;
     const sc: []const u8 = if (backbone) "s.ex" else "s.dx";
@@ -312,7 +312,7 @@ pub fn engram(e: *E, L: u32) !void {
 pub fn finish(e: *E) !void {
     try mhc(e, null, "", .final);
     const D: i64 = e.cfg.hidden;
-    const V: i64 = e.cfg.vocab / e.o.world;
+    const V: i64 = (try @import("partition.zig").share(e.cfg.vocab, e.o.rank, e.o.world)).len;
     const logits = try e.buf("w.logits", .{}, .f32, &.{ e.n, V });
     try e.upstream("head", try e.buf("w.out", .{}, .bf16, &.{ e.n, D }), D, V, logits);
     try e.glue("kit_logits", &.{logits});

zig/src/families/deepseek_v41/block_prefetch.zig

@@ -36,7 +36,7 @@ pub fn l2pfSources(e: *Emitter, L: u32, site: L2Site) ![]const L2Src {
         },
         .f => {
             const D: i64 = e.cfg.hidden;
-            const I: i64 = e.cfg.expert_width / e.o.world;
+            const I: i64 = (try @import("partition.zig").share(e.cfg.expert_width, e.o.rank, e.o.world)).len;
             const ex = e.cfg.expertsOf(L);
             const sh = try e.fmt("L{d}.moe.shared.0.w1", .{L});
             const k2: u64 = try e.w.k2(sh);

zig/src/families/deepseek_v41/block_prefill_attn.zig

@@ -12,6 +12,7 @@ const pieceAt = prefill.Seg.pieceAt;
 const selection = prefill.Seg.selection;
 const rowsOf = prefill.Seg.rowsOf;
 const stashArgs = prefill.stashArgs;
+const pf_heads = @import("prefill_heads.zig");
 
 // -- attention -----------------------------------------------------------------------------------------------
 
@@ -218,16 +219,17 @@ pub fn attention(s: *Seg, L: u32) !void {
     }
     // the CSA2 steps a segment at a time, into the group-major core output
     const groups: i64 = e.cfg.o_groups / e.o.world;
-    if (@mod(H, bmq) != 0 or @mod(H, groups) != 0) return error.Unsupported;
+    if ((H != 16 and @mod(H, bmq) != 0) or @mod(H, groups) != 0) return error.Unsupported;
+    const core_q = if (H == 16) try pf_heads.query(e, q, hd) else q;
     const hg = @divExact(H, groups);
     const o = try e.buf("L.o", .{}, .bf16, &.{ groups, n, hg * hd });
     if (s.multi) |segs| {
         for (segs, 0..) |sp, j| {
             try s.enter(sp, j);
-            try s.core(L, q, kv, cproj, own_kv, indexer, o, n);
+            try s.core(L, core_q, kv, cproj, own_kv, indexer, o, n);
         }
         try s.leave(n);
-    } else try s.core(L, q, kv, cproj, own_kv, indexer, o, n);
+    } else try s.core(L, core_q, kv, cproj, own_kv, indexer, o, n);
     // _out_pf: each wo_a group's GEMM over its contiguous slice into its columns of z, then wo_b (bf16 partial)
     const ol: i64 = e.cfg.o_lora;
     const gk = hg * hd;
@@ -270,10 +272,11 @@ pub fn core(s: *Seg, L: u32, q: Arg, kv: Arg, cproj: Arg, own_kv: bool, indexer:
     const e = s.e;
     const n = e.n;
     const ratio: i64 = e.cfg.compressRatio(L);
-    const H: i64 = e.cfg.heads / e.o.world;
+    const real_heads: i64 = e.cfg.heads / e.o.world;
+    const H: i64 = if (real_heads == 16) 32 else real_heads;
     const hd: i64 = e.cfg.head_dim;
     const groups: i64 = e.cfg.o_groups / e.o.world;
-    const hg = @divExact(H, groups);
+    const hg = @divExact(real_heads, groups);
     // blocks._segment: a segment the slot's ring cannot hold runs on the staging ring, the window before it in
     const rows = try s.swa(L);
     const slot = try slotRing(e, L);
@@ -310,6 +313,10 @@ pub fn core(s: *Seg, L: u32, q: Arg, kv: Arg, cproj: Arg, own_kv: bool, indexer:
             fpg = try e.paging(e.cfg.compressRatio(src), false);
         }
     }
+    const real_sink = try e.weight("L{d}.attn.sink", .{L}, .f32, &.{real_heads});
+    const sink = if (real_heads == 16) try pf_heads.sink(e, real_sink) else real_sink;
+    const output = if (real_heads == 16) try e.buf("L.pf.o32", .{}, .bf16, &.{ 2 * groups, total, hg * hd }) else o;
+    const out_view = try e.view(output.t.role, .bf16, &.{ @divExact(H, hg), n, hg * hd }, &.{ total * hg * hd, hg * hd, 1 }, e.row0 * hg * hd * 2);
     try e.triton("_fused", .{ n, @divExact(H, bmq), 1 }, &.{
         .{ .name = "Q", .arg = try e.view(q.t.role, .bf16, &.{ n, H, hd }, &.{ H * hd, hd, 1 }, e.row0 * H * hd * 2) },
         .{ .name = "CV", .arg = cv },
@@ -325,10 +332,10 @@ pub fn core(s: *Seg, L: u32, q: Arg, kv: Arg, cproj: Arg, own_kv: bool, indexer:
         .{ .name = "sss", .arg = .{ .i = 8 } },
         .{ .name = "LO", .arg = lo },
         .{ .name = "POS", .arg = try pos(e) },
-        .{ .name = "SINK", .arg = try e.weight("L{d}.attn.sink", .{L}, .f32, &.{H}) },
+        .{ .name = "SINK", .arg = sink },
         .{ .name = "CS", .arg = try e.ropeTable() },
         .{ .name = "cs_stride", .arg = .{ .i = 64 } },
-        .{ .name = "OUT", .arg = try e.view(o.t.role, .bf16, &.{ groups, n, hg * hd }, &.{ total * hg * hd, hg * hd, 1 }, e.row0 * hg * hd * 2) },
+        .{ .name = "OUT", .arg = out_view },
         .{ .name = "H", .arg = .{ .i = H } },
         .{ .name = "KT", .arg = .{ .i = kt } },
         .{ .name = "BMQ", .arg = .{ .i = bmq } },
@@ -342,5 +349,6 @@ pub fn core(s: *Seg, L: u32, q: Arg, kv: Arg, cproj: Arg, own_kv: bool, indexer:
         .{ .name = "HG", .arg = .{ .i = hg } },
         .{ .name = "OR", .arg = .{ .i = total } },
     });
+    if (real_heads == 16) try pf_heads.extract(e, out_view, try e.view(o.t.role, .bf16, &.{ groups, n, hg * hd }, &.{ total * hg * hd, hg * hd, 1 }, e.row0 * hg * hd * 2));
     if (s.staged) try e.glue("stage_out", &.{ slot[0], slot[1], .{ .i = ring0 }, rows[0], rows[1], .{ .i = s.ring }, .{ .i = @max(0, e.end() - W) }, .{ .i = e.end() } });
 }

zig/src/families/deepseek_v41/block_prefill_mhc.zig

@@ -196,7 +196,7 @@ pub fn final(s: *Seg) !void {
     try s.siteTriton(x, x, try gathered(e, ""), prev, prev[0], null, part, c, null, 2, false);
     try s.finish(null, prev, try e.weight("norm", .{}, .f32, &.{D}), part, c, 0);
     // the head: the decode path's linear (prefill_mm.fast off), upstream's 128 rows a launch
-    const V: i64 = e.cfg.vocab / e.o.world;
+    const V: i64 = (try @import("partition.zig").share(e.cfg.vocab, e.o.rank, e.o.world)).len;
     var a: i64 = 0;
     while (a < n) : (a += 128) {
         const m = @min(n - a, 128);

zig/src/families/deepseek_v41/block_prefill_moe.zig

@@ -55,7 +55,7 @@ pub fn moeRows(s: *Seg, L: u32, n: i64, x: Arg, comptime bias: []const u8, part:
     const topk: i64 = ex.topk;
     const slots: i64 = topk + 1;
     const Et: i64 = ex.count;
-    const I: i64 = e.cfg.expert_width / e.o.world;
+    const I: i64 = (try @import("partition.zig").share(e.cfg.expert_width, e.o.rank, e.o.world)).len;
     const P = n * topk;
     const pick = try e.buf("L.pick", .{}, .i32, &.{ n, slots });
     const wts = try e.buf("L.wts", .{}, .f32, &.{ n, slots });

zig/src/families/deepseek_v41/calls_check.zig

@@ -164,6 +164,12 @@ pub fn optionsOf(a: std.mem.Allocator, io: std.Io, dir: []const u8) !block.Optio
     const meta = try std.json.parseFromSliceLeaky(Value, a, text, .{});
     var o: block.Options = .{};
     const env = meta.object.get("env") orelse return o;
+    if (env.object.get("TF_TP_WORLD")) |v| if (v == .string) {
+        o.world = try std.fmt.parseInt(u32, v.string, 10);
+    };
+    if (env.object.get("TF_TP_RANK")) |v| if (v == .string) {
+        o.rank = try std.fmt.parseInt(u32, v.string, 10);
+    };
     const on = struct {
         fn f(e: Value, k: []const u8) bool {
             const v = e.object.get(k) orelse return false;

zig/src/families/deepseek_v41/cuda_model.zig

@@ -188,6 +188,7 @@ pub const Model = struct {
         const comm = m.session.comm();
         const world: u32 = @intCast(comm.world());
         m.cfg = try Config.read(gpa, io, pack_dir);
+        if (o.own_prefill) try @import("block_admission.zig").nativePrefill(&m.cfg, world);
         m.pack = try Pack.open(gpa, io, pack_dir);
         const backbone = if (o.layers) |ls| try a.dupe(u32, ls) else blk: {
             const all = try a.alloc(u32, m.cfg.layers);
@@ -264,7 +265,7 @@ pub const Model = struct {
         const topp: ?f64 = if (std.c.getenv("TF_DSV41_EXPERT_TOPP")) |v| try std.fmt.parseFloat(f64, std.mem.span(v)) else null;
         // TF_DSV41_R1=1, or the eight separate R1 knobs (the knob reader's r1)
         const r1 = @import("serve_knobs.zig").r1(&@import("serve_knobs.zig").env);
-        m.f = fwd.Forward.init(gpa, &m.cfg, &m.widths, .{ .world = world, .expert_topp = topp, .r1 = r1, .taps = o.drafts }, &m.runner, comm);
+        m.f = fwd.Forward.init(gpa, &m.cfg, &m.widths, .{ .world = world, .rank = m.tcfg.rank, .expert_topp = topp, .r1 = r1, .taps = o.drafts }, &m.runner, comm);
         try @import("engine_hooks.zig").open(&m.f); // the forward's later-layer hooks and prefill state
         m.f.layers = backbone;
         // the rope tables cover the limit + a prefill segment (the assets' rope.bin)

zig/src/families/deepseek_v41/dspark_dev.zig

@@ -35,6 +35,7 @@ pub const Dev = struct {
     k: usize,
     world: usize,
     hidden: usize,
+    vocab_lo: u32,
     /// rows of every pass of a round together (the regions' size)
     rows_max: usize,
     conf: bool,
@@ -56,7 +57,7 @@ pub const Dev = struct {
     /// TF_DSV41_DRAFT_CHAIN=1: the chain on the device (Chain); null: the host's
     chain: ?*Chain = null,
 
-    pub fn init(gpa: std.mem.Allocator, d: *const cuda.Driver, comm: tp.collective.Collective, stream: cuda.Stream, want: Mode, k: usize, hidden: usize, rows_max: usize, conf: bool, region_bytes: usize, regions: usize) !*Dev {
+    pub fn init(gpa: std.mem.Allocator, d: *const cuda.Driver, comm: tp.collective.Collective, stream: cuda.Stream, want: Mode, k: usize, hidden: usize, rows_max: usize, conf: bool, region_bytes: usize, regions: usize, vocab_lo: u32) !*Dev {
         const x = try gpa.create(Dev);
         errdefer gpa.destroy(x);
         const W: usize = comm.world();
@@ -85,6 +86,7 @@ pub const Dev = struct {
             .k = k,
             .world = W,
             .hidden = hidden,
+            .vocab_lo = vocab_lo,
             .rows_max = rows_max,
             .conf = conf,
             .send = send,
@@ -144,7 +146,7 @@ pub const Dev = struct {
         const ops = r.kernels.others(r.stream);
         const send = x.send.ptr + at * 2 * k * 4;
         const dst = x.gathered.ptr + at * 2 * k * 4 * x.world;
-        try ops.dsCandidates(try r.tensorAddr(logits_t), V, rows, V, k, @intCast(x.comm.rank() * V), send);
+        try ops.dsCandidates(try r.tensorAddr(logits_t), V, rows, V, k, x.vocab_lo, send);
         try x.comm.allGather(send, dst, rows * 2 * k, .f32, r.stream.handle);
         if (x.chain) |ch| try ch.layout(r, x.gathered.ptr, x.world, rows, k, at, try r.tensorAddr(hid_t));
         try r.d.check(r.d.api.cuMemcpyDtoHAsync_v2(x.host.bytes.ptr + at * 2 * k * 4 * x.world, dst, pk * x.world, r.stream.handle), "cuMemcpyDtoHAsync");
@@ -194,7 +196,7 @@ pub const Dev = struct {
         const rank = x.comm.rank();
         const all: []const f32 = @alignCast(std.mem.bytesAsSlice(f32, x.host.bytes[at * 2 * k * 4 * x.world ..][0 .. rows * 2 * k * 4 * x.world]));
         for (0..rows) |i| {
-            try dspark.candidates(x.gpa, host[i * V ..][0..V], @intCast(rank * V), ids, vals);
+            try dspark.candidates(x.gpa, host[i * V ..][0..V], x.vocab_lo, ids, vals);
             const got = all[(rank * rows + i) * 2 * k ..][0 .. 2 * k];
             for (0..k) |j| {
                 const gid: i32 = @bitCast(got[k + j]);

zig/src/families/deepseek_v41/dspark_emit.zig

@@ -337,7 +337,7 @@ const Ds = struct {
             try e.glue("exchange_f32", &.{ try e.buf("L.moe", .{}, .f32, &.{ n, D }), try e.buf("w.recv.moe", .{}, .bf16, &.{ e.o.world, n, D }) });
         }
         try d.final();
-        const V: i64 = d.cfg().vocab / e.o.world;
+        const V: i64 = (try @import("partition.zig").share(d.cfg().vocab, e.o.rank, e.o.world)).len;
         const logits = try e.buf("w.ds.logits", .{}, .f32, &.{ n, V });
         try e.upstream("head", try e.buf("w.out", .{}, .bf16, &.{ n, D }), D, V, logits);
         try e.glue("ds_out", &.{ logits, try e.buf("L.mhc.c", .{}, .bf16, &.{ n, D }) });

zig/src/families/deepseek_v41/dspark_gpu.zig

@@ -150,7 +150,7 @@ pub const GpuPass = struct {
             } else {
                 // every buffer a graphed pass bakes in, at its largest: the embedding's and exchanges' send
                 _ = try x.device(&x.send, 2 * @as(usize, @intCast(n)) * cfg.hidden);
-                x.dev = try dsd.Dev.init(gpa, r.d, f.comm, r.stream, mode, k, cfg.hidden, @intCast(n), conf != null, @intCast(8 * st.len64() + 4 * st.len32()), 1);
+                x.dev = try dsd.Dev.init(gpa, r.d, f.comm, r.stream, mode, k, cfg.hidden, @intCast(n), conf != null, @intCast(8 * st.len64() + 4 * st.len32()), 1, (try @import("partition.zig").share(cfg.vocab, f.opts.rank, world)).lo);
                 // the speculative pass from the device pick (needs TF_DSV41_GREEDY_GPU's pick and ds_stage)
                 if (r.kernels.others(r.stream).f.ds_stage != null) {
                     publish_pick = true;

zig/src/families/deepseek_v41/dspark_gpu_device.zig

@@ -161,10 +161,11 @@ pub fn glueFn(ctx: *anyopaque, r: *run.Runner, c: *const calls.Call) anyerror!vo
     if (std.mem.eql(u8, step, "ds_embed")) {
         // the block's ids (staged) through this rank's vocabulary rows, their all-sum, x hc_mult into the streams
         const n: usize = @intCast(x.st.n);
-        const V: usize = x.f.cfg.vocab / Wd;
+        const vocab_part = try @import("partition.zig").share(x.f.cfg.vocab, comm.rank(), Wd);
+        const V: usize = vocab_part.len;
         const embed = r.weights.get("embed") orelse return error.MissingWeight;
         const rows = try x.device(&x.send, 2 * n * D);
-        try k.embedSend(try r.tensorAddr(c.args[2].arg.t), n, embed.ptr, @intCast(comm.rank() * V), V, D, rows);
+        try k.embedSend(try r.tensorAddr(c.args[2].arg.t), n, embed.ptr, @intCast(vocab_part.lo), V, D, rows);
         const recv = try r.tensorAddr(c.args[1].arg.t);
         try comm.allGather(rows, recv, n * D, .bf16, s);
         return k.embedSum(recv, Wd, n, D, try r.tensorAddr(c.args[0].arg.t));
@@ -213,7 +214,7 @@ pub fn hostStep(x: *GpuPass, r: *run.Runner, logits_t: calls.Tensor, hid_t: call
     defer gpa.free(packed_);
     const ids = try gpa.alloc(i32, k);
     defer gpa.free(ids);
-    const id0: u32 = @intCast(comm.rank() * V);
+    const id0: u32 = (try @import("partition.zig").share(x.f.cfg.vocab, comm.rank(), Wd)).lo;
     for (0..n) |i| {
         const row = packed_[i * 2 * k ..][0 .. 2 * k];
         try dspark.candidates(gpa, logits[i * V ..][0..V], id0, ids, row[0..k]);