From 4e7efa552d5156fbd35e233469ad1bd52480a38a Mon Sep 17 00:00:00 2001
From: ashhart <71890694+ashhart@users.noreply.github.com>
Date: Thu, 22 Oct 2026 22:04:24 +0100
Subject: [PATCH 1/2] deepseek_v41: prepare whole-block shares for four ranks

Use contiguous 128-column shares for expert and vocabulary planning,
reads, local dimensions and sampling exchanges. Preserve the complete
world-1/2 plan fingerprint and test default world-4 plans without weights.

Native four-rank prefill remains held at its 32-head tile contract.
Refuse that unsupported mode before opening pack payloads; this host-only
preparation does not qualify four-rank generation or close #1202.
---
 zig/src/families/deepseek_v41/aot_needs.zig   |   2 +-
 zig/src/families/deepseek_v41/batch_graph.zig |   5 +-
 zig/src/families/deepseek_v41/block.zig       |   3 +-
 .../families/deepseek_v41/block_admission.zig |  12 +-
 zig/src/families/deepseek_v41/block_moe.zig   |   4 +-
 .../families/deepseek_v41/block_prefetch.zig  |   2 +-
 .../deepseek_v41/block_prefill_mhc.zig        |   2 +-
 .../deepseek_v41/block_prefill_moe.zig        |   2 +-
 zig/src/families/deepseek_v41/calls_check.zig |   6 +
 zig/src/families/deepseek_v41/cuda_model.zig  |   3 +-
 zig/src/families/deepseek_v41/dspark_dev.zig  |   8 +-
 zig/src/families/deepseek_v41/dspark_emit.zig |   2 +-
 zig/src/families/deepseek_v41/dspark_gpu.zig  |   2 +-
 .../deepseek_v41/dspark_gpu_device.zig        |   7 +-
 .../families/deepseek_v41/dspark_slots.zig    |   4 +-
 .../deepseek_v41/dspark_slots_device.zig      |   7 +-
 .../families/deepseek_v41/forward_glue.zig    |  10 +-
 .../families/deepseek_v41/forward_sample.zig  |   9 +-
 .../families/deepseek_v41/gate_decode_tp.zig  |  11 +-
 zig/src/families/deepseek_v41/gate_kv.zig     |   2 +-
 zig/src/families/deepseek_v41/host_test.zig   |   2 +-
 zig/src/families/deepseek_v41/named.zig       |   8 +-
 zig/src/families/deepseek_v41/partition.zig   |  29 +++++
 zig/src/families/deepseek_v41/plan.zig        |  12 +-
 zig/src/families/deepseek_v41/prepared.zig    |   2 +-
 .../deepseek_v41/review_constructors_test.zig |   4 +-
 .../families/deepseek_v41/sampling_dev.zig    |   8 +-
 .../families/deepseek_v41/sampling_gpu.zig    |   7 +-
 .../families/deepseek_v41/sampling_pipe.zig   |  80 ++++++++++----
 .../families/deepseek_v41/sampling_test.zig   |  75 ++++++++++++-
 zig/src/families/deepseek_v41/state_test.zig  |   5 +
 zig/src/families/deepseek_v41/tool_main.zig   |   2 +-
 .../deepseek_v41/world4_plan_test.zig         |  50 +++++++++
 .../deepseek_v41/world4_stage_test.zig        |  48 ++++++++
 zig/src/families/deepseek_v41/world4_test.zig | 104 ++++++++++++++++++
 35 files changed, 458 insertions(+), 81 deletions(-)
 create mode 100644 zig/src/families/deepseek_v41/partition.zig
 create mode 100644 zig/src/families/deepseek_v41/world4_plan_test.zig
 create mode 100644 zig/src/families/deepseek_v41/world4_stage_test.zig
 create mode 100644 zig/src/families/deepseek_v41/world4_test.zig

diff --git a/zig/src/families/deepseek_v41/aot_needs.zig b/zig/src/families/deepseek_v41/aot_needs.zig
index b9e7a8f45c..e34111958e 100644
--- a/zig/src/families/deepseek_v41/aot_needs.zig
+++ b/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;
diff --git a/zig/src/families/deepseek_v41/batch_graph.zig b/zig/src/families/deepseek_v41/batch_graph.zig
index 2b496863e2..be20d918ce 100644
--- a/zig/src/families/deepseek_v41/batch_graph.zig
+++ b/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));
diff --git a/zig/src/families/deepseek_v41/block.zig b/zig/src/families/deepseek_v41/block.zig
index 6fad115d2e..5857b83210 100644
--- a/zig/src/families/deepseek_v41/block.zig
+++ b/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)
diff --git a/zig/src/families/deepseek_v41/block_admission.zig b/zig/src/families/deepseek_v41/block_admission.zig
index c3efe966e8..cd27e01e39 100644
--- a/zig/src/families/deepseek_v41/block_admission.zig
+++ b/zig/src/families/deepseek_v41/block_admission.zig
@@ -1,8 +1,16 @@
 //! 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) % 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;
@@ -10,7 +18,7 @@ pub fn options(c: *const Config, o: anytype, layers: []const u32, n: i64) !void
     }
     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;
     }
diff --git a/zig/src/families/deepseek_v41/block_moe.zig b/zig/src/families/deepseek_v41/block_moe.zig
index 34168fa012..a9ff68162f 100644
--- a/zig/src/families/deepseek_v41/block_moe.zig
+++ b/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});
diff --git a/zig/src/families/deepseek_v41/block_prefetch.zig b/zig/src/families/deepseek_v41/block_prefetch.zig
index 47a9d4bcd2..69613b648a 100644
--- a/zig/src/families/deepseek_v41/block_prefetch.zig
+++ b/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);
diff --git a/zig/src/families/deepseek_v41/block_prefill_mhc.zig b/zig/src/families/deepseek_v41/block_prefill_mhc.zig
index 4f0cfe427b..8dfebad6d8 100644
--- a/zig/src/families/deepseek_v41/block_prefill_mhc.zig
+++ b/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);
diff --git a/zig/src/families/deepseek_v41/block_prefill_moe.zig b/zig/src/families/deepseek_v41/block_prefill_moe.zig
index ca57efb0af..4583dc49ce 100644
--- a/zig/src/families/deepseek_v41/block_prefill_moe.zig
+++ b/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 });
diff --git a/zig/src/families/deepseek_v41/calls_check.zig b/zig/src/families/deepseek_v41/calls_check.zig
index 1c3c9cf1e9..eefd444caf 100644
--- a/zig/src/families/deepseek_v41/calls_check.zig
+++ b/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;
diff --git a/zig/src/families/deepseek_v41/cuda_model.zig b/zig/src/families/deepseek_v41/cuda_model.zig
index 8877244be3..1d7043d321 100644
--- a/zig/src/families/deepseek_v41/cuda_model.zig
+++ b/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)
diff --git a/zig/src/families/deepseek_v41/dspark_dev.zig b/zig/src/families/deepseek_v41/dspark_dev.zig
index aad74c3a96..1f8da959ed 100644
--- a/zig/src/families/deepseek_v41/dspark_dev.zig
+++ b/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]);
diff --git a/zig/src/families/deepseek_v41/dspark_emit.zig b/zig/src/families/deepseek_v41/dspark_emit.zig
index 3fe14a6a07..b9bb348c17 100644
--- a/zig/src/families/deepseek_v41/dspark_emit.zig
+++ b/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 }) });
diff --git a/zig/src/families/deepseek_v41/dspark_gpu.zig b/zig/src/families/deepseek_v41/dspark_gpu.zig
index 55ab590e29..077364dafb 100644
--- a/zig/src/families/deepseek_v41/dspark_gpu.zig
+++ b/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;
diff --git a/zig/src/families/deepseek_v41/dspark_gpu_device.zig b/zig/src/families/deepseek_v41/dspark_gpu_device.zig
index 43dc052ef5..fc0134a8d2 100644
--- a/zig/src/families/deepseek_v41/dspark_gpu_device.zig
+++ b/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]);
diff --git a/zig/src/families/deepseek_v41/dspark_slots.zig b/zig/src/families/deepseek_v41/dspark_slots.zig
index 29910d7132..cb44ad70b8 100644
--- a/zig/src/families/deepseek_v41/dspark_slots.zig
+++ b/zig/src/families/deepseek_v41/dspark_slots.zig
@@ -133,7 +133,7 @@ pub const SlotPass = struct {
         const Rr = n * @max(group, nslots);
         const c: usize = f.comm.world() * one.k;
         const lay = emit.Layout.of(group, @intCast(n), f.cfg.window);
-        const V: usize = f.cfg.vocab / f.comm.world();
+        const V: usize = (try @import("partition.zig").share(f.cfg.vocab, f.opts.rank, f.comm.world())).len;
         const valid = try gpa.alloc(u64, nslots);
         errdefer if (!assembled) gpa.free(valid);
         const h64 = try gpa.alloc(i64, @intCast(lay.len64));
@@ -199,7 +199,7 @@ pub const SlotPass = struct {
                 if (!quiet) std.log.scoped(.dsv41).warn("draft graphs: TF_DSV41_DRAFT_GRAPHS set, but the glue fatbin has no ds_cands kernel: the host step", .{});
             } else {
                 _ = try x.device(&x.send, 2 * R * f.cfg.hidden); // the embedding's and exchanges' send, at its largest
-                x.dev = try dsd.Dev.init(gpa, r.d, f.comm, r.stream, mode, one.k, f.cfg.hidden, n * nslots, one.conf != null, @intCast(8 * lay.len64 + 4 * lay.len32), nslots);
+                x.dev = try dsd.Dev.init(gpa, r.d, f.comm, r.stream, mode, one.k, f.cfg.hidden, n * nslots, one.conf != null, @intCast(8 * lay.len64 + 4 * lay.len32), nslots, (try @import("partition.zig").share(f.cfg.vocab, f.opts.rank, f.comm.world())).lo);
                 // the speculative round from the row window's device pick: this pass takes the forward's hook over
                 const ko = r.kernels.others(r.stream).f;
                 if (ko.ds_accept_rows != null and ko.ds_stage_rows != null and nslots <= 16) {
diff --git a/zig/src/families/deepseek_v41/dspark_slots_device.zig b/zig/src/families/deepseek_v41/dspark_slots_device.zig
index 3db0e7863b..062c1fb5f2 100644
--- a/zig/src/families/deepseek_v41/dspark_slots_device.zig
+++ b/zig/src/families/deepseek_v41/dspark_slots_device.zig
@@ -445,10 +445,11 @@ pub fn glueFn(ctx: *anyopaque, r: *run.Runner, c: *const calls.Call) anyerror!vo
     const Wd: usize = comm.world();
     if (std.mem.eql(u8, step, "ds_embed")) {
         const n: usize = numel(c.args[2].arg.t);
-        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 rws = 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, rws);
+        try k.embedSend(try r.tensorAddr(c.args[2].arg.t), n, embed.ptr, @intCast(vocab_part.lo), V, D, rws);
         const recv = try r.tensorAddr(c.args[1].arg.t);
         try comm.allGather(rws, recv, n * D, .bf16, s);
         return k.embedSum(recv, Wd, n, D, try r.tensorAddr(c.args[0].arg.t));
@@ -490,7 +491,7 @@ pub fn hostStep(x: *SlotPass, r: *run.Runner, logits_t: calls.Tensor, hid_t: cal
     try cuda.DeviceBuffer.download(.{ .d = r.d, .ptr = try r.tensorAddr(hid_t), .len = 2 * hid16.len }, 0, std.mem.sliceAsBytes(hid16));
     for (hid16, x.hid[0 .. R * D]) |b, *h| h.* = dspark.bf16(b);
     const packed_ = x.packed_[0 .. R * 2 * k];
-    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..R) |i| {
         const row = packed_[i * 2 * k ..][0 .. 2 * k];
         try dspark.candidates(x.gpa, logits[i * V ..][0..V], id0, x.ids, row[0..k]);
diff --git a/zig/src/families/deepseek_v41/forward_glue.zig b/zig/src/families/deepseek_v41/forward_glue.zig
index b75a0f1823..5634bf667f 100644
--- a/zig/src/families/deepseek_v41/forward_glue.zig
+++ b/zig/src/families/deepseek_v41/forward_glue.zig
@@ -80,9 +80,10 @@ pub fn glueFn(ctx: *anyopaque, r: *run.Runner, c: *const calls.Call) anyerror!vo
         // graphed: the ids staged before the launch; no host step inside the window
         const n: usize = p.n;
         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 * n * D);
-        try k.embedSend(g.?.statics.idsDev(), n, embed.ptr, @intCast(f.comm.rank() * V), V, D, rows);
+        try k.embedSend(g.?.statics.idsDev(), n, 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, n * D, .bf16, s);
         return k.embedSum(recv, W, n, D, try r.tensorAddr(c.args[0].arg.t));
@@ -101,9 +102,10 @@ pub fn glueFn(ctx: *anyopaque, r: *run.Runner, c: *const calls.Call) anyerror!vo
         for (f.ids, 0..) |id, i| std.mem.writeInt(i64, h[8 * i ..][0..8], id, .little);
         try f.ids_stage.copyTo(f.dev_ids.?, 0, r.stream);
         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 * n * D);
-        try k.embedSend(f.dev_ids.?.ptr, n, embed.ptr, @intCast(f.comm.rank() * V), V, D, rows);
+        try k.embedSend(f.dev_ids.?.ptr, n, embed.ptr, @intCast(vocab_part.lo), V, D, rows);
         if (f.vision) |vr| try hooks.set.vision.?.fill(vr, r.stream, f.ids, rows, D); // rank 0: the image positions' rows (vision.fill)
         const recv = try r.tensorAddr(c.args[1].arg.t);
         try f.comm.allGather(rows, recv, n * D, .bf16, s);
diff --git a/zig/src/families/deepseek_v41/forward_sample.zig b/zig/src/families/deepseek_v41/forward_sample.zig
index 6782c2c255..2de801e419 100644
--- a/zig/src/families/deepseek_v41/forward_sample.zig
+++ b/zig/src/families/deepseek_v41/forward_sample.zig
@@ -19,7 +19,8 @@ pub fn greedy(f: *Forward, out: []u32) !void {
     const r = f.runner;
     const n = out.len;
     const W: usize = f.comm.world();
-    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;
     // pinned: the logits [n, V], then the pairs [2n] and every rank's [W][2n] (reused, no allocation a window)
     const lb = std.mem.alignForward(usize, 4 * n * V, 8);
     const need = lb + 16 * n * (W + 1);
@@ -49,7 +50,7 @@ pub fn greedy(f: *Forward, out: []u32) !void {
             // the pairs, their exchange and the merge on the device; one copy back (glue.cu pick_*)
             const ids_dev = cols_dev + 8 * n;
             const out_dev = ids_dev + 8 * n;
-            try k.pickPack(vals_dev, cols_dev, n, f.comm.rank() * V, dev.ptr);
+            try k.pickPack(vals_dev, cols_dev, n, vocab_part.lo, dev.ptr);
             try f.comm.allGather(dev.ptr, dev.ptr + 16 * n, 2 * n, .f64, r.stream.handle);
             const chain = if (f.slot.pending) |p| p.n == n and f.ids.len >= n else false;
             const start: i64 = if (f.slot.pending) |p| @intCast(p.start) else 0;
@@ -79,7 +80,7 @@ pub fn greedy(f: *Forward, out: []u32) !void {
         try r.stream.synchronize();
         for (0..n) |row| {
             pairs[2 * row] = vh[row];
-            pairs[2 * row + 1] = @floatFromInt(@as(usize, @intCast(ch[row])) + f.comm.rank() * V);
+            pairs[2 * row + 1] = @floatFromInt(@as(usize, @intCast(ch[row])) + vocab_part.lo);
         }
     } else {
         try r.d.check(r.d.api.cuMemcpyDtoHAsync_v2(host.ptr, logits, 4 * n * V, r.stream.handle), "cuMemcpyDtoHAsync");
@@ -89,7 +90,7 @@ pub fn greedy(f: *Forward, out: []u32) !void {
             const vals = host[row * V ..][0..V];
             const best = argmax(vals);
             pairs[2 * row] = vals[best];
-            pairs[2 * row + 1] = @floatFromInt(best + f.comm.rank() * V);
+            pairs[2 * row + 1] = @floatFromInt(best + vocab_part.lo);
         }
     }
     try dev.uploadAsync(0, std.mem.sliceAsBytes(pairs[0 .. 2 * n]), r.stream.handle);
diff --git a/zig/src/families/deepseek_v41/gate_decode_tp.zig b/zig/src/families/deepseek_v41/gate_decode_tp.zig
index 30952a7528..de9a9c6a3b 100644
--- a/zig/src/families/deepseek_v41/gate_decode_tp.zig
+++ b/zig/src/families/deepseek_v41/gate_decode_tp.zig
@@ -1,4 +1,4 @@
-//! `tf-dsv41 decode-tp PACK REF OUT` (a process a rank): the decode TP check, our TP=2 greedy decode vs REF's tokens.
+//! `tf-dsv41 decode-tp PACK REF OUT`: one process per rank, against that world's reference tokens and digests.
 
 const std = @import("std");
 const pf = @import("forward_prefill.zig");
@@ -40,6 +40,7 @@ pub fn main(gpa: std.mem.Allocator, io: std.Io, pack_dir: []const u8, ref: []con
     const cwd = std.Io.Dir.cwd();
     const tcfg = try tp.Config.fromEnv();
     const rank = tcfg.rank;
+    const world = tcfg.world;
     try cwd.createDirPath(io, out_dir);
     const meta = try std.json.parseFromSliceLeaky(Value, a, try cwd.readFileAlloc(io, try std.fs.path.join(a, &.{ ref, "ref.json" }), a, .limited(1 << 28)), .{});
     var lo: u32 = 0;
@@ -86,14 +87,14 @@ pub fn main(gpa: std.mem.Allocator, io: std.Io, pack_dir: []const u8, ref: []con
         var ds: [8]u32 = undefined;
         break :blk try std.mem.concat(a, u32, &.{ layers, dspark_emit.blocks(&cfg, &ds) });
     } else layers;
-    var pl = try plan.build(gpa, &cfg, &pack, .{ .rank = rank, .world = 2, .blocks = blocks, .top = true, .dspark = lanes_mode == 3 });
+    var pl = try plan.build(gpa, &cfg, &pack, .{ .rank = rank, .world = world, .blocks = blocks, .top = true, .dspark = lanes_mode == 3 });
     defer pl.deinit();
     var w: load.Weights = .{ .gpa = gpa };
     defer @import("close_or_exit.zig").finish(&w);
     for (pl.layers) |*l| {
         var b: named.Builder = .{ .gpa = gpa, .io = io, .pack = &pack };
         defer b.deinit();
-        try b.layer(l, cfg.o_groups, rank, 2);
+        try b.layer(l, cfg.o_groups, rank, world);
         try w.upload(&driver, &b);
     }
     {
@@ -119,7 +120,7 @@ pub fn main(gpa: std.mem.Allocator, io: std.Io, pack_dir: []const u8, ref: []con
         const v = env.object.get("TF_DSV41_MHC_DEFER") orelse break :blk false;
         break :blk v == .string and std.mem.eql(u8, v.string, "1");
     };
-    var f = fwd.Forward.init(gpa, &cfg, &wd, .{ .expert_topp = topp, .r1 = r1 }, &runner, comm);
+    var f = fwd.Forward.init(gpa, &cfg, &wd, .{ .world = world, .rank = rank, .expert_topp = topp, .r1 = r1 }, &runner, comm);
     defer @import("close_or_exit.zig").finish(&f);
     try @import("engine_hooks.zig").open(&f);
     f.layers = layers;
@@ -255,7 +256,7 @@ pub fn main(gpa: std.mem.Allocator, io: std.Io, pack_dir: []const u8, ref: []con
     // greedy decoding against the reference's tokens
     const mine = meta.object.get("ranks").?.array.items[rank].object;
     const want_logits = mine.get("logits_sha256").?.array.items;
-    const V: usize = cfg.vocab / 2;
+    const V: usize = (try @import("partition.zig").share(cfg.vocab, rank, world)).len;
     const host_logits = try a.alloc(u8, 4 * V);
     var trace: std.AutoArrayHashMapUnmanaged(u32, [64]u8) = .empty;
     defer trace.deinit(gpa);
diff --git a/zig/src/families/deepseek_v41/gate_kv.zig b/zig/src/families/deepseek_v41/gate_kv.zig
index bdafd45176..9d115418f8 100644
--- a/zig/src/families/deepseek_v41/gate_kv.zig
+++ b/zig/src/families/deepseek_v41/gate_kv.zig
@@ -136,7 +136,7 @@ fn onePrompt(gpa: std.mem.Allocator, io: std.Io, a: std.mem.Allocator, m: *model
     };
     // the decode steps (park: a RAM round trip a third in, an NVMe one two thirds in)
     const steps = want_tokens.len - 1;
-    const V: usize = m.cfg.vocab / f.comm.world();
+    const V: usize = (try @import("partition.zig").share(m.cfg.vocab, f.comm.rank(), f.comm.world())).len;
     const host_logits = try gpa.alloc(u8, 4 * V);
     defer gpa.free(host_logits);
     var tok: u32 = @intCast(want_tokens[0].integer);
diff --git a/zig/src/families/deepseek_v41/host_test.zig b/zig/src/families/deepseek_v41/host_test.zig
index 39a843b278..b2b4d5c1db 100644
--- a/zig/src/families/deepseek_v41/host_test.zig
+++ b/zig/src/families/deepseek_v41/host_test.zig
@@ -41,7 +41,7 @@ test "the release config.json: topology, routing and rope" {
 }
 
 /// Header JSON of a pack for `c`, named as the release packs are (their index, pattern by pattern), header only.
-fn tinyPack(a: std.mem.Allocator, c: *const Config, extra: []const u8, skip: []const u8) ![]u8 {
+pub fn tinyPack(a: std.mem.Allocator, c: *const Config, extra: []const u8, skip: []const u8) ![]u8 {
     var n: u64 = 0;
     return tinyPackLen(a, c, extra, skip, &n);
 }
diff --git a/zig/src/families/deepseek_v41/named.zig b/zig/src/families/deepseek_v41/named.zig
index 88f78d3559..1c1d735253 100644
--- a/zig/src/families/deepseek_v41/named.zig
+++ b/zig/src/families/deepseek_v41/named.zig
@@ -153,10 +153,12 @@ pub const Builder = struct {
     /// The rank's slice of a native tensor, as stored.
     fn raw(b: *Builder, v: plan.NativePlan, rank: u32, world: u32) ![]u8 {
         const whole = v.slice == .whole;
-        const len = if (whole) v.info.nbytes else v.info.nbytes / world;
+        const s = if (v.slice == .vocab_rows) try @import("partition.zig").share(@intCast(v.info.shape[0]), rank, world) else null;
+        const len = if (s) |part| v.info.nbytes / v.info.shape[0] * part.len else if (whole) v.info.nbytes else v.info.nbytes / world;
+        const offset = if (s) |part| v.info.nbytes / v.info.shape[0] * part.lo else if (whole) 0 else rank * len;
         const bytes = try b.gpa.alloc(u8, @intCast(len));
         errdefer b.gpa.free(bytes);
-        try b.pack.read(b.io, v.info, .{ .offset = if (whole) 0 else rank * len, .len = len }, bytes);
+        try b.pack.read(b.io, v.info, .{ .offset = offset, .len = len }, bytes);
         return bytes;
     }
 
@@ -165,7 +167,7 @@ pub const Builder = struct {
         var shape: [4]usize = @splat(1);
         const dims = v.info.dims();
         @memcpy(shape[0..dims.len], dims);
-        if (v.slice != .whole) shape[0] /= world;
+        if (v.slice == .vocab_rows) shape[0] = (try @import("partition.zig").share(@intCast(shape[0]), rank, world)).len else if (v.slice != .whole) shape[0] /= world;
         if (v.info.dtype != .f32 and v.info.dtype != .bf16 and v.info.dtype != .f16) return error.UnexpectedTensor;
         if (b.metadata_only) {
             if (rule.out == .qk and std.mem.endsWith(u8, v.name, "q_weight")) return;
diff --git a/zig/src/families/deepseek_v41/partition.zig b/zig/src/families/deepseek_v41/partition.zig
new file mode 100644
index 0000000000..ccb3d54ecd
--- /dev/null
+++ b/zig/src/families/deepseek_v41/partition.zig
@@ -0,0 +1,29 @@
+const exl3 = @import("core").exl3_layout;
+
+pub const Share = struct { lo: u32, len: u32 };
+
+pub fn share(width: u32, rank: usize, world: usize) !Share {
+    if (world == 0 or rank >= world or width % 128 != 0 or width / 128 < world) return error.SplitNotOnBlocks;
+    const blocks: usize = width / 128;
+    const base = blocks / world;
+    const extra = blocks % world;
+    return .{ .lo = @intCast((rank * base + @min(rank, extra)) * 128), .len = @intCast((base + @intFromBool(rank < extra)) * 128) };
+}
+
+pub fn group(g: exl3.Group, split: exl3.Split, rank: u32, world: u32) !exl3.Part {
+    if (split == .whole) return g.part(split, rank, world);
+    const s = try share(if (split == .row) g.k() else g.n(), rank, world);
+    var p = try g.part(.whole, rank, world);
+    const tiles = s.len / 16;
+    if (split == .row) {
+        p.group.k_tiles = tiles;
+        const row = @as(u64, g.n_tiles) * g.tileBytes();
+        p.trellis = .{ .offset = @as(u64, s.lo / 16) * row, .len = @as(u64, tiles) * row };
+        p.suh = .{ .offset = 2 * @as(u64, s.lo), .len = 2 * @as(u64, s.len) };
+    } else {
+        p.group.n_tiles = tiles;
+        p.trellis = .{ .offset = @as(u64, s.lo / 16) * g.tileBytes(), .len = @as(u64, tiles) * g.tileBytes(), .stride = @as(u64, g.n_tiles) * g.tileBytes(), .count = g.k_tiles };
+        p.svh = .{ .offset = 2 * @as(u64, s.lo), .len = 2 * @as(u64, s.len) };
+    }
+    return p;
+}
diff --git a/zig/src/families/deepseek_v41/plan.zig b/zig/src/families/deepseek_v41/plan.zig
index 659fc3734d..db78fbf525 100644
--- a/zig/src/families/deepseek_v41/plan.zig
+++ b/zig/src/families/deepseek_v41/plan.zig
@@ -8,6 +8,7 @@ const exl3 = @import("core").exl3_layout;
 const names = @import("names.zig");
 const Pack = @import("pack.zig").Pack;
 const Info = @import("pack.zig").Info;
+const partition = @import("partition.zig");
 
 /// One EXL3 group of the rank: its stored name and the rank's part.
 pub const GroupPlan = struct {
@@ -147,7 +148,7 @@ const Builder = struct {
         inline for (.{ "trellis", "suh", "svh" }) |part| try b.mark(try std.fmt.bufPrint(&b.buf, "{s}." ++ part, .{prefix}));
         const marker = try b.pack.marker(prefix, &b.buf);
         if (marker) |m| try b.mark(m);
-        return .{ .proj = proj, .prefix = try b.dupe(prefix), .split = split, .full = g, .part = try g.part(split, b.o.rank, b.o.world), .marker = if (marker) |m| try b.dupe(m) else null };
+        return .{ .proj = proj, .prefix = try b.dupe(prefix), .split = split, .full = g, .part = try partition.group(g, split, b.o.rank, b.o.world), .marker = if (marker) |m| try b.dupe(m) else null };
     }
 
     fn native(b: *Builder, name: []const u8, spec: names.NativeSpec) !NativePlan {
@@ -164,7 +165,8 @@ const Builder = struct {
         try b.mark(name);
         const bytes = switch (spec.slice) {
             .whole => info.nbytes,
-            .heads, .vocab_rows => blk: {
+            .vocab_rows => info.nbytes / info.shape[0] * (try partition.share(@intCast(info.shape[0]), b.o.rank, b.o.world)).len,
+            .heads => blk: {
                 if (info.shape[0] % b.o.world != 0) return error.SplitNotEven;
                 break :blk info.nbytes / b.o.world;
             },
@@ -265,9 +267,11 @@ const Builder = struct {
     }
 };
 
-/// The rank's plan; heads, groups and the vocabulary must split over `world` (blocks of 128).
+/// Heads and groups split evenly; expert and vocabulary slices retain whole 128-column blocks.
 pub fn build(gpa: std.mem.Allocator, c: *const Config, pack: *const Pack, o: Options) !Plan {
-    if (o.world == 0 or o.rank >= o.world or c.heads % o.world != 0 or c.o_groups % o.world != 0 or c.vocab % (128 * o.world) != 0) return error.BadWorld;
+    if (o.world == 0 or o.rank >= o.world or c.heads % o.world != 0 or c.o_groups % o.world != 0) return error.BadWorld;
+    _ = partition.share(c.vocab, o.rank, o.world) catch return error.BadWorld;
+    _ = partition.share(c.expert_width, o.rank, o.world) catch return error.BadWorld;
     var arena = std.heap.ArenaAllocator.init(gpa);
     errdefer arena.deinit();
     var b: Builder = .{ .a = arena.allocator(), .c = c, .pack = pack, .o = o };
diff --git a/zig/src/families/deepseek_v41/prepared.zig b/zig/src/families/deepseek_v41/prepared.zig
index 1b28c8ae6e..6aa95a4044 100644
--- a/zig/src/families/deepseek_v41/prepared.zig
+++ b/zig/src/families/deepseek_v41/prepared.zig
@@ -20,7 +20,7 @@ pub const Key = struct {
 };
 
 /// The builders' source, so a change to how images are made never reads an old folder.
-const sources = [_][]const u8{ @embedFile("named.zig"), @embedFile("plan.zig"), @embedFile("stage.zig"), @embedFile("pack.zig"), @embedFile("projection.zig"), @embedFile("names.zig"), @import("core").exl3_layout.source_bytes, @import("prepare.zig").source_bytes };
+const sources = [_][]const u8{ @embedFile("named.zig"), @embedFile("plan.zig"), @embedFile("partition.zig"), @embedFile("stage.zig"), @embedFile("pack.zig"), @embedFile("projection.zig"), @embedFile("names.zig"), @import("core").exl3_layout.source_bytes, @import("prepare.zig").source_bytes };
 
 pub const Stats = struct { hits: u32 = 0, misses: u32 = 0, written: u32 = 0, read_bytes: u64 = 0, read_ns: u64 = 0, write_ns: u64 = 0 };
 
diff --git a/zig/src/families/deepseek_v41/review_constructors_test.zig b/zig/src/families/deepseek_v41/review_constructors_test.zig
index 67868c9ac7..915c0e7399 100644
--- a/zig/src/families/deepseek_v41/review_constructors_test.zig
+++ b/zig/src/families/deepseek_v41/review_constructors_test.zig
@@ -54,7 +54,7 @@ test "Completion ownership: failed DSpark event construction unwinds its precedi
     vt.world = world;
     const comm: @import("tp").collective.Collective = .{ .ptr = &d, .vtable = &vt };
     mock.fail_event = true;
-    try t.expectError(error.NotReady, @import("dspark_dev.zig").Dev.init(t.allocator, &d, comm, .{ .d = &d, .handle = null }, .device, 1, 8, 1, false, 16, 1));
+    try t.expectError(error.NotReady, @import("dspark_dev.zig").Dev.init(t.allocator, &d, comm, .{ .d = &d, .handle = null }, .device, 1, 8, 1, false, 16, 1, 0));
     try t.expectEqual(mock.allocated, mock.freed);
     try t.expect(mock.pinned == null);
 }
@@ -90,7 +90,7 @@ test "Completion ownership: SlotPass host-array allocation sweep leaves callback
         const a = fail.allocator();
         var vt: @import("tp").collective.Collective.VTable = undefined;
         vt.world = world;
-        var cfg: @import("config.zig").Config = .{ .vocab = 8, .hidden = 16, .window = 2, .dspark_block = 1 };
+        var cfg: @import("config.zig").Config = .{ .vocab = 128, .hidden = 16, .window = 2, .dspark_block = 1 };
         var widths: @import("block.zig").Widths = .{};
         var runner: run.Runner = .{ .gpa = a, .d = &d, .stream = .{ .d = &d, .handle = null }, .kernels = undefined, .triton = undefined, .weights = undefined };
         var f = fwd.Forward.init(a, &cfg, &widths, .{}, &runner, .{ .ptr = &d, .vtable = &vt });
diff --git a/zig/src/families/deepseek_v41/sampling_dev.zig b/zig/src/families/deepseek_v41/sampling_dev.zig
index a1b0fa498b..4ab6d95397 100644
--- a/zig/src/families/deepseek_v41/sampling_dev.zig
+++ b/zig/src/families/deepseek_v41/sampling_dev.zig
@@ -42,7 +42,7 @@ pub const Gpu = struct {
     }
 
     pub fn device(g: *Gpu) pipe.Device {
-        return .{ .ptr = g, .vtable = &.{ .top = top, .stats = stats, .scratch = scratch, .download = download, .stream = streamOf } };
+        return .{ .ptr = g, .vtable = &.{ .top = top, .stats = stats, .scratch = scratch, .download = download, .stream = streamOf, .pad = pad } };
     }
 
     fn self_(p: *anyopaque) *Gpu {
@@ -84,4 +84,10 @@ pub const Gpu = struct {
     fn streamOf(p: *anyopaque) cuda.abi.Stream {
         return self_(p).stream.handle;
     }
+
+    fn pad(p: *anyopaque, src: u64, cols: usize, stride: usize, out: u64) anyerror!void {
+        const g = self_(p);
+        try g.d.check(g.d.api.cuMemsetD32Async(out, @bitCast(-std.math.inf(f32)), stride, g.stream.handle), "cuMemsetD32Async");
+        try g.d.check(g.d.api.cuMemcpyDtoDAsync_v2(out, src, 4 * cols, g.stream.handle), "cuMemcpyDtoDAsync");
+    }
 };
diff --git a/zig/src/families/deepseek_v41/sampling_gpu.zig b/zig/src/families/deepseek_v41/sampling_gpu.zig
index 494ca003ef..58e67ff4cc 100644
--- a/zig/src/families/deepseek_v41/sampling_gpu.zig
+++ b/zig/src/families/deepseek_v41/sampling_gpu.zig
@@ -40,10 +40,9 @@ pub const GpuSampler = struct {
     pub fn init(gpa: std.mem.Allocator, f: *fwd.Forward) !*GpuSampler {
         const x = try gpa.create(GpuSampler);
         errdefer gpa.destroy(x);
-        const W = f.comm.world();
         x.* = .{ .gpa = gpa, .f = f, .gpu = try Gpu.init(f.runner.d, f.runner.stream, f.runner.kernels), .p = undefined };
         errdefer x.gpu.deinit();
-        x.p = try pipe.Pipe.init(gpa, f.comm, f.cfg.vocab / W, try pipe.nucleusEnv());
+        x.p = try pipe.Pipe.initVocabAt(gpa, f.comm, f.opts.rank, f.cfg.vocab, try pipe.nucleusEnv());
         f.sampler = .{ .ptr = x, .follow = follow };
         x.p.after_gather = .{ .ctx = x, .run = afterGather };
         x.batch = try batchEnv();
@@ -168,7 +167,7 @@ pub const GpuSampler = struct {
     fn note(x: *GpuSampler, s: Sampling, batched: bool) void {
         if (x.f.comm.rank() != 0) return;
         const V: usize = x.p.V;
-        const count = pipe.countOf(s, @intCast(V * x.f.comm.world()), x.p.nucleus);
+        const count = pipe.countOf(s, x.p.vocab, x.p.nucleus);
         const kind: u3 = if (@min(count, V) > pipe.max_k or V > pipe.max_cols) 0 else if (isNucleusRow(s)) 1 else 2;
         const bit = @as(u8, 1) << (kind + @as(u3, if (batched) 3 else 0));
         if (x.seen & bit != 0) return;
@@ -208,7 +207,7 @@ pub const GpuSampler = struct {
 
     fn run(x: *GpuSampler, row0: u32, n: u32, s: Sampling, draws: ?[]const u64, out: []u32, pos0: ?u64) !void {
         const f = x.f;
-        const V: u64 = f.cfg.vocab / f.comm.world();
+        const V: u64 = x.p.V;
         const lg = (f.runner.addressOf("w.logits") orelse return error.Unbound) + 4 * V * row0;
         if (f.grammar) |g| try g.apply(g.ptr, row0, n);
         x.cur = .{ .row0 = row0, .s = s, .pos0 = pos0 };
diff --git a/zig/src/families/deepseek_v41/sampling_pipe.zig b/zig/src/families/deepseek_v41/sampling_pipe.zig
index b2c889cb40..70e3e8457a 100644
--- a/zig/src/families/deepseek_v41/sampling_pipe.zig
+++ b/zig/src/families/deepseek_v41/sampling_pipe.zig
@@ -26,6 +26,7 @@ pub const Device = struct {
         download: *const fn (ptr: *anyopaque, src: u64, dst: []u8) anyerror!void,
         /// The stream the collectives run on.
         stream: *const fn (ptr: *anyopaque) tp.collective.Stream,
+        pad: ?*const fn (ptr: *anyopaque, src: u64, cols: usize, stride: usize, out: u64) anyerror!void = null,
     };
 };
 
@@ -127,8 +128,12 @@ pub fn nucleusEnv() !u32 {
 pub const Pipe = struct {
     gpa: std.mem.Allocator,
     comm: tp.collective.Collective,
-    /// this rank's vocabulary slice (the head's columns): global ids rank x V + column
+    /// This rank's vocabulary slice; global ids are lo + column.
     V: u32,
+    lo: ?u32,
+    vocab: u32,
+    stride: u32,
+    min_cols: u32,
     nucleus: u32,
     picker: smp.Picker,
     values: []f32,
@@ -142,12 +147,27 @@ pub const Pipe = struct {
     whole_rows: usize = 0,
 
     pub fn init(gpa: std.mem.Allocator, comm: tp.collective.Collective, V: u32, nucleus: u32) !Pipe {
-        const vocab: usize = @as(usize, V) * comm.world();
+        return initPart(gpa, comm, V, null, @intCast(@as(usize, V) * comm.world()), V, V, nucleus);
+    }
+
+    pub fn initVocab(gpa: std.mem.Allocator, comm: tp.collective.Collective, vocab: u32, nucleus: u32) !Pipe {
+        return initVocabAt(gpa, comm, @intCast(comm.rank()), vocab, nucleus);
+    }
+
+    pub fn initVocabAt(gpa: std.mem.Allocator, comm: tp.collective.Collective, rank: u32, vocab: u32, nucleus: u32) !Pipe {
+        const partition = @import("partition.zig");
+        const part = try partition.share(vocab, rank, comm.world());
+        const first = try partition.share(vocab, 0, comm.world());
+        const last = try partition.share(vocab, comm.world() - 1, comm.world());
+        return initPart(gpa, comm, part.len, part.lo, vocab, first.len, last.len, nucleus);
+    }
+
+    fn initPart(gpa: std.mem.Allocator, comm: tp.collective.Collective, V: u32, lo: ?u32, vocab: u32, stride: u32, min_cols: u32, nucleus: u32) !Pipe {
         var picker = try smp.Picker.init(gpa, vocab);
         errdefer picker.deinit();
         const values = try gpa.alloc(f32, vocab);
         errdefer gpa.free(values);
-        return .{ .gpa = gpa, .comm = comm, .V = V, .nucleus = nucleus, .picker = picker, .values = values, .ids = try gpa.alloc(u32, vocab) };
+        return .{ .gpa = gpa, .comm = comm, .V = V, .lo = lo, .vocab = vocab, .stride = stride, .min_cols = min_cols, .nucleus = nucleus, .picker = picker, .values = values, .ids = try gpa.alloc(u32, vocab) };
     }
 
     pub fn deinit(p: *Pipe) void {
@@ -169,11 +189,11 @@ pub const Pipe = struct {
         std.debug.assert(out.len >= n and n > 0);
         const W: usize = p.comm.world();
         const V: usize = p.V;
-        const count: usize = countOf(s, @intCast(V * W), p.nucleus);
-        const k = @min(count, V);
+        const count: usize = countOf(s, p.vocab, p.nucleus);
+        const k = @min(count, p.min_cols);
         p.whole_rows = 0;
         try p.picks.resize(p.gpa, n);
-        if (k > max_k or V > max_cols) {
+        if (k > max_k or p.stride > max_cols or (p.stride != p.min_cols and count > p.min_cols)) {
             // every row whole (top_k 0 without the nucleus path, a top_k past the device bound, or one rank's 129,280)
             p.full.clearRetainingCapacity();
             for (0..n) |r| try p.full.append(p.gpa, @intCast(r));
@@ -185,7 +205,7 @@ pub const Pipe = struct {
         const row_f32 = 2 * k;
         const packed_ = try vt.scratch(dev.ptr, 0, 4 * n * row_f32);
         const gathered = try vt.scratch(dev.ptr, 1, 4 * W * n * row_f32);
-        try vt.top(dev.ptr, lg, n, V, k, @intCast(p.comm.rank() * V), packed_);
+        try vt.top(dev.ptr, lg, n, V, k, p.lo orelse @intCast(p.comm.rank() * V), packed_);
         try p.comm.allGather(packed_, gathered, n * row_f32, .f32, stream);
         var stats_dev: u64 = 0;
         if (nuc) {
@@ -229,14 +249,29 @@ pub const Pipe = struct {
         const W: usize = p.comm.world();
         const V: usize = p.V;
         const vt = dev.vtable;
-        const recv = try vt.scratch(dev.ptr, 4, 4 * W * V);
-        const host = try p.hostBytes(4 * W * V);
+        const stride: usize = p.stride;
+        const recv = try vt.scratch(dev.ptr, 4, 4 * W * stride);
+        const host = try p.hostBytes(4 * W * stride);
         const vals: []const f32 = @alignCast(std.mem.bytesAsSlice(f32, host));
-        for (p.ids[0 .. W * V], 0..) |*id, i| id.* = @intCast(i); // rank order: global id = the index
-        try p.comm.allGather(row, recv, V, .f32, vt.stream(dev.ptr));
+        for (p.ids, 0..) |*id, i| id.* = @intCast(i);
+        var send = row;
+        if (stride != p.min_cols) {
+            const padded = try vt.scratch(dev.ptr, 5, 4 * stride);
+            const pad = vt.pad orelse return error.UnevenDeviceUnsupported;
+            try pad(dev.ptr, row, V, stride, padded);
+            send = padded;
+        }
+        try p.comm.allGather(send, recv, stride, .f32, vt.stream(dev.ptr));
         try vt.download(dev.ptr, recv, host);
-        @memcpy(p.values[0 .. W * V], vals);
-        const r: smp.Row = .{ .values = p.values[0 .. W * V], .ids = p.ids[0 .. W * V], .count = count };
+        if (stride == p.min_cols) {
+            @memcpy(p.values, vals);
+        } else {
+            for (0..W) |rank| {
+                const part = try @import("partition.zig").share(p.vocab, rank, W);
+                @memcpy(p.values[part.lo..][0..part.len], vals[rank * stride ..][0..part.len]);
+            }
+        }
+        const r: smp.Row = .{ .values = p.values, .ids = p.ids, .count = count };
         p.whole_rows += 1;
         return p.picker.pick(r, position, s).token; // no statistics: never .full_row
     }
@@ -246,8 +281,9 @@ pub const Pipe = struct {
 
     /// Whether `chooseSegments` takes a segment under `s`: a keyed candidates-path row within the device bounds.
     pub fn batchable(p: *const Pipe, s: Sampling) bool {
-        if (!sampled(s) or p.V > max_cols) return false;
-        return @min(countOf(s, @intCast(@as(usize, p.V) * p.comm.world()), p.nucleus), p.V) <= max_k;
+        if (!sampled(s) or p.stride > max_cols) return false;
+        const count = countOf(s, p.vocab, p.nucleus);
+        return @min(count, p.min_cols) <= max_k and (p.stride == p.min_cols or count <= p.min_cols);
     }
 
     /// `chooseSegments`' hook over the gathered candidates before their download; true when it downloaded `host`.
@@ -258,7 +294,7 @@ pub const Pipe = struct {
         std.debug.assert(segs.len > 0);
         const W: usize = p.comm.world();
         const V: usize = p.V;
-        const vocab: u32 = @intCast(V * W);
+        const vocab = p.vocab;
         var count: usize = 0;
         var nuc = false;
         const r0: usize = segs[0].row0;
@@ -271,7 +307,7 @@ pub const Pipe = struct {
             nuc = nuc or smp.isNucleus(g.s);
         }
         const n = end - r0;
-        const k = @min(count, V);
+        const k = @min(count, p.min_cols);
         p.whole_rows = 0;
         const vt = dev.vtable;
         const stream = vt.stream(dev.ptr);
@@ -279,7 +315,7 @@ pub const Pipe = struct {
         const span = lg + 4 * @as(u64, r0) * V;
         const packed_ = try vt.scratch(dev.ptr, 0, 4 * n * row_f32);
         const gathered = try vt.scratch(dev.ptr, 1, 4 * W * n * row_f32);
-        try vt.top(dev.ptr, span, n, V, k, @intCast(p.comm.rank() * V), packed_);
+        try vt.top(dev.ptr, span, n, V, k, p.lo orelse @intCast(p.comm.rank() * V), packed_);
         try p.comm.allGather(packed_, gathered, n * row_f32, .f32, stream);
         var stats_dev: u64 = 0;
         if (nuc) {
@@ -337,7 +373,7 @@ pub const HostDevice = struct {
     }
 
     pub fn device(h: *HostDevice) Device {
-        return .{ .ptr = h, .vtable = &.{ .top = top, .stats = stats, .scratch = scratch, .download = download, .stream = stream } };
+        return .{ .ptr = h, .vtable = &.{ .top = top, .stats = stats, .scratch = scratch, .download = download, .stream = stream, .pad = pad } };
     }
 
     fn self_(p: *anyopaque) *HostDevice {
@@ -398,6 +434,12 @@ pub const HostDevice = struct {
         @memcpy(dst, @as([*]const u8, @ptrFromInt(src))[0..dst.len]);
     }
 
+    fn pad(_: *anyopaque, src: u64, cols: usize, stride: usize, out: u64) anyerror!void {
+        const dst = f32s(out, stride);
+        @memset(dst, -std.math.inf(f32));
+        @memcpy(dst[0..cols], f32s(src, cols));
+    }
+
     fn stream(_: *anyopaque) tp.collective.Stream {
         return null;
     }
diff --git a/zig/src/families/deepseek_v41/sampling_test.zig b/zig/src/families/deepseek_v41/sampling_test.zig
index 0bfd39dbdc..4a24ac50a7 100644
--- a/zig/src/families/deepseek_v41/sampling_test.zig
+++ b/zig/src/families/deepseek_v41/sampling_test.zig
@@ -7,9 +7,44 @@ const json = @import("json");
 const serve = @import("dsv41_serve");
 const pipe = @import("sampling_pipe.zig");
 
+test "world four: uneven candidate and whole-row sampling match one rank" {
+    const vocab = 1280;
+    const n = 4;
+    const a = testing.allocator;
+    const logits = try a.alloc(f32, n * vocab);
+    defer a.free(logits);
+    var random = std.Random.DefaultPrng.init(1202);
+    for (logits) |*v| v.* = random.random().float(f32) * 4 - 2;
+    const peaks = [_]u32{ 383, 384, 1023, 1279 };
+    for (peaks, 0..) |id, row| logits[row * vocab + id] = 8;
+    const positions = [_]u64{ 7, 8, 9, 10 };
+    for ([_]pipe.Sampling{
+        .{ .seed = 1202, .temperature = 0.9, .top_k = 1, .top_p = 1 },
+        .{ .seed = 1202, .temperature = 0.9, .top_k = 23, .top_p = 0.9 },
+        .{ .seed = 1202, .temperature = 0.9, .top_k = 400, .top_p = 1 },
+        .{ .seed = 1202, .temperature = 0.9, .top_k = 0, .top_p = 1 },
+        .{ .seed = 1202, .temperature = 0.9, .top_k = 0, .top_p = 0.9 },
+    }) |sampling| {
+        var expected: [n]u32 = undefined;
+        var expected_whole: [1]usize = undefined;
+        var one: Case = .{ .logits = logits, .world = 1, .uneven = true, .rows = n, .sampling = sampling, .nucleus = 64, .positions = &positions, .got = &expected, .whole = &expected_whole };
+        try run(&one);
+        for ([_]u32{ 2, 4 }) |world| {
+            var got: [n]u32 = undefined;
+            var whole: [1]usize = undefined;
+            var split: Case = .{ .logits = logits, .world = world, .uneven = true, .rows = n, .sampling = sampling, .nucleus = 64, .positions = &positions, .got = &got, .whole = &whole };
+            try run(&split);
+            try testing.expectEqualSlices(u32, &expected, &got);
+            if (sampling.top_k == 1) try testing.expectEqualSlices(u32, &peaks, &got);
+            if (world == 4 and sampling.top_k == 400) try testing.expectEqual(@as(usize, n), whole[0]);
+        }
+    }
+}
+
 const Case = struct {
     logits: []const f32,
     world: u32,
+    uneven: bool = false,
     rows: usize,
     sampling: pipe.Sampling,
     positions: []const u64,
@@ -42,13 +77,14 @@ const Mirror = struct {
 fn rankBody(comm: tp.collective.Collective, me: u32, c: *Case) anyerror!void {
     const a = testing.allocator;
     const vocab = c.logits.len / c.rows;
-    const V = vocab / c.world;
+    const part = if (c.uneven) try @import("partition.zig").share(@intCast(vocab), me, c.world) else @import("partition.zig").Share{ .lo = @intCast(me * (vocab / c.world)), .len = @intCast(vocab / c.world) };
+    const V = part.len;
     const mine = try a.alloc(f32, c.rows * V);
     defer a.free(mine);
-    for (0..c.rows) |r| @memcpy(mine[r * V ..][0..V], c.logits[r * vocab + me * V ..][0..V]);
+    for (0..c.rows) |r| @memcpy(mine[r * V ..][0..V], c.logits[r * vocab + part.lo ..][0..V]);
     var dev: pipe.HostDevice = .{ .gpa = a };
     defer dev.deinit();
-    var p = try pipe.Pipe.init(a, comm, @intCast(V), c.nucleus);
+    var p = if (c.uneven) try pipe.Pipe.initVocab(a, comm, @intCast(vocab), c.nucleus) else try pipe.Pipe.init(a, comm, @intCast(V), c.nucleus);
     defer p.deinit();
     var mh: Mirror = .{ .c = c };
     if (me == 0 and c.mirror != null) p.after_gather = .{ .ctx = &mh, .run = Mirror.run };
@@ -230,10 +266,36 @@ test "vs_choose's mirror == the pipeline's picks over top_k x top_p x min_p x te
 
 const SegSpec = struct { row0: u32, n: u32, s: pipe.Sampling, positions: []const u64 };
 
+test "world four: batched uneven sampling keeps segment offsets and global ids" {
+    const vocab = 1280;
+    const rows = 4;
+    const a = testing.allocator;
+    const logits = try a.alloc(f32, rows * vocab);
+    defer a.free(logits);
+    var random = std.Random.DefaultPrng.init(12024);
+    for (logits) |*v| v.* = random.random().float(f32) * 4;
+    const segs = [_]SegSpec{
+        .{ .row0 = 0, .n = 2, .s = .{ .seed = 7, .temperature = 0.7, .top_k = 7, .top_p = 0.9 }, .positions = &.{ 40, 41 } },
+        .{ .row0 = 2, .n = 2, .s = .{ .seed = 9, .temperature = 1.1, .top_k = 11, .top_p = 0.8 }, .positions = &.{ 900, 901 } },
+    };
+    var expected: [rows]u32 = undefined;
+    var one: Multi = .{ .logits = logits, .world = 1, .uneven = true, .rows = rows, .segs = &segs, .batched = false, .got = &expected };
+    try tp.host.run(1, &one, multiBody);
+    for ([_]u32{ 2, 4 }) |world| {
+        var got: [rows]u32 = undefined;
+        var mirror: [rows]i64 = @splat(-2);
+        var split: Multi = .{ .logits = logits, .world = world, .uneven = true, .rows = rows, .segs = &segs, .batched = true, .got = &got, .mirror = &mirror };
+        try tp.host.run(world, &split, multiBody);
+        try testing.expectEqualSlices(u32, &expected, &got);
+        for (got, mirror) |id, mirrored| try testing.expectEqual(@as(i64, id), mirrored);
+    }
+}
+
 /// Several segments of one forward's rows (rows not in a segment are another window's, never picked).
 const Multi = struct {
     logits: []const f32,
     world: u32,
+    uneven: bool = false,
     rows: usize,
     segs: []const SegSpec,
     batched: bool,
@@ -273,13 +335,14 @@ const SegMirror = struct {
 fn multiBody(comm: tp.collective.Collective, me: u32, c: *Multi) anyerror!void {
     const a = testing.allocator;
     const vocab = c.logits.len / c.rows;
-    const V = vocab / c.world;
+    const part = if (c.uneven) try @import("partition.zig").share(@intCast(vocab), me, c.world) else @import("partition.zig").Share{ .lo = @intCast(me * (vocab / c.world)), .len = @intCast(vocab / c.world) };
+    const V = part.len;
     const mine = try a.alloc(f32, c.rows * V);
     defer a.free(mine);
-    for (0..c.rows) |r| @memcpy(mine[r * V ..][0..V], c.logits[r * vocab + me * V ..][0..V]);
+    for (0..c.rows) |r| @memcpy(mine[r * V ..][0..V], c.logits[r * vocab + part.lo ..][0..V]);
     var dev: pipe.HostDevice = .{ .gpa = a };
     defer dev.deinit();
-    var p = try pipe.Pipe.init(a, comm, @intCast(V), serve.sampling.nucleus_count);
+    var p = if (c.uneven) try pipe.Pipe.initVocab(a, comm, @intCast(vocab), serve.sampling.nucleus_count) else try pipe.Pipe.init(a, comm, @intCast(V), serve.sampling.nucleus_count);
     defer p.deinit();
     const lg: u64 = @intFromPtr(mine.ptr);
     const out = try a.alloc(u32, c.rows);
diff --git a/zig/src/families/deepseek_v41/state_test.zig b/zig/src/families/deepseek_v41/state_test.zig
index cb89009adf..adcc7f29cf 100644
--- a/zig/src/families/deepseek_v41/state_test.zig
+++ b/zig/src/families/deepseek_v41/state_test.zig
@@ -1,4 +1,9 @@
 const std = @import("std");
+test {
+    _ = @import("world4_test.zig");
+    _ = @import("world4_plan_test.zig");
+    _ = @import("world4_stage_test.zig");
+}
 test {
     _ = @import("kv_io_test.zig");
     inline for (.{ @import("forward.zig").Forward, @import("run.zig").Runner, @import("graphs.zig").Statics, @import("graphs.zig").EngramStage, @import("cuda_kv.zig").Kv, @import("forward_prefill.zig"), @import("longpf.zig"), @import("ced.zig"), @import("sessions_gpu.zig").Sessions, @import("slots.zig").SlotSet, @import("engine_hooks.zig") }) |m| std.testing.refAllDecls(m);
diff --git a/zig/src/families/deepseek_v41/tool_main.zig b/zig/src/families/deepseek_v41/tool_main.zig
index cceca1700c..fd8d1f9535 100644
--- a/zig/src/families/deepseek_v41/tool_main.zig
+++ b/zig/src/families/deepseek_v41/tool_main.zig
@@ -37,7 +37,7 @@ const usage =
     \\                   recipe/triton_fill.py compile; SIGS from its `sigs`)
     \\       tf-dsv41 kv paged|split|compact|park PACK REF OUT (a rank of TP=2 on the paged pool vs the pool reference)
     \\       tf-dsv41 forward-digests PACK REF OUT (our forward vs recipe/dsv41_forward_digests.py's run, one GPU)
-    \\       tf-dsv41 decode-tp PACK REF OUT (a rank of TP=2 over NCCL vs recipe/dsv41_decode_ref.py; TF_TP_RANK ...)
+    \\       tf-dsv41 decode-tp PACK REF OUT (one NCCL rank, against matching-world reference tokens and digests)
     \\       tf-dsv41 follow PACK ASSETS (a served model's rank > 0: rank 0's forward operations until it stops)
     \\       tf-dsv41 generate PACK ASSETS TRACE (rank 0: TRACE's prompts through the served engine, ids == recorded)
     \\       tf-dsv41 slot-drafts PACK ASSETS TRACE... (rank 0: drafts over several live slots, batched == alone == recorded)
diff --git a/zig/src/families/deepseek_v41/world4_plan_test.zig b/zig/src/families/deepseek_v41/world4_plan_test.zig
new file mode 100644
index 0000000000..aa34c745f1
--- /dev/null
+++ b/zig/src/families/deepseek_v41/world4_plan_test.zig
@@ -0,0 +1,50 @@
+const std = @import("std");
+const Config = @import("config.zig").Config;
+const Pack = @import("pack.zig").Pack;
+const plan = @import("plan.zig");
+
+fn defaultPack(a: std.mem.Allocator) !Pack {
+    const cfg: Config = .{};
+    const header = try @import("host_test.zig").tinyPack(a, &cfg, "", "");
+    defer a.free(header);
+    var p = Pack.init(a);
+    errdefer p.deinit();
+    try p.addHeader("model.safetensors", header, std.math.maxInt(u64));
+    return p;
+}
+
+fn fingerprint(p: *const plan.Plan, hash: *std.crypto.hash.sha2.Sha256) !void {
+    var out: std.Io.Writer.Allocating = .init(std.testing.allocator);
+    defer out.deinit();
+    try std.json.Stringify.value(.{ .rank = p.rank, .world = p.world, .layers = p.layers, .top = p.top, .head = p.head, .dspark = p.dspark, .leftovers = p.leftovers }, .{}, &out.writer);
+    hash.update(out.written());
+}
+
+test "world four: default configuration plans every rank without payload reads" {
+    const cfg: Config = .{};
+    var pack = try defaultPack(std.testing.allocator);
+    defer pack.deinit();
+    for (0..4) |rank| {
+        var p = try plan.build(std.testing.allocator, &cfg, &pack, .{ .rank = @intCast(rank), .world = 4 });
+        defer p.deinit();
+        try std.testing.expectEqual(@as(usize, cfg.allLayers()), p.layers.len);
+        try std.testing.expectEqual(@as(usize, 0), p.leftovers.len);
+    }
+}
+
+test "world four: world one and two plans retain their main fingerprint" {
+    const cfg: Config = .{};
+    var pack = try defaultPack(std.testing.allocator);
+    defer pack.deinit();
+    var hash = std.crypto.hash.sha2.Sha256.init(.{});
+    for ([_]u32{ 1, 2 }) |world| for (0..world) |rank| {
+        var p = try plan.build(std.testing.allocator, &cfg, &pack, .{ .rank = @intCast(rank), .world = world });
+        defer p.deinit();
+        try fingerprint(&p, &hash);
+    };
+    var digest: [32]u8 = undefined;
+    hash.final(&digest);
+    const hex = std.fmt.bytesToHex(digest, .lower);
+    std.debug.print("world one/two plan fingerprint: {s}\n", .{hex});
+    try std.testing.expectEqualStrings("0b2e082e4d244eec2d631a0da6f4495a5ccf73259f9384faea2167507169b5f5", &hex);
+}
diff --git a/zig/src/families/deepseek_v41/world4_stage_test.zig b/zig/src/families/deepseek_v41/world4_stage_test.zig
new file mode 100644
index 0000000000..5f609a82fc
--- /dev/null
+++ b/zig/src/families/deepseek_v41/world4_stage_test.zig
@@ -0,0 +1,48 @@
+const std = @import("std");
+const pk = @import("pack.zig");
+const partition = @import("partition.zig");
+const named = @import("named.zig");
+
+test "world four: captured programs retain their rank and world" {
+    const a = std.testing.allocator;
+    var tmp = std.testing.tmpDir(.{});
+    defer tmp.cleanup();
+    try tmp.dir.writeFile(std.testing.io, .{ .sub_path = "meta.json", .data = "{\"env\":{\"TF_TP_WORLD\":\"4\",\"TF_TP_RANK\":\"3\"}}" });
+    const dir = try tmp.dir.realPathFileAlloc(std.testing.io, ".", a);
+    defer a.free(dir);
+    var arena = std.heap.ArenaAllocator.init(a);
+    defer arena.deinit();
+    const o = try @import("calls_check.zig").optionsOf(arena.allocator(), std.testing.io, dir);
+    try std.testing.expectEqual(@as(u32, 4), o.world);
+    try std.testing.expectEqual(@as(u32, 3), o.rank);
+}
+
+test "world four: embedding shard reads use each rank's global offset" {
+    const a = std.testing.allocator;
+    const io = std.testing.io;
+    const width = 768;
+    const hidden = 2;
+    const file = try pk.image(a, &.{.{ .name = "embed.weight", .dtype = "BF16", .shape = &.{ width, hidden } }});
+    defer a.free(file);
+    const start = 8 + std.mem.readInt(u64, file[0..8], .little);
+    for (0..width * hidden) |i| std.mem.writeInt(u16, file[start + i * 2 ..][0..2], @intCast(i), .little);
+    var tmp = std.testing.tmpDir(.{});
+    defer tmp.cleanup();
+    try tmp.dir.writeFile(io, .{ .sub_path = "model.safetensors", .data = file });
+    const dir = try tmp.dir.realPathFileAlloc(io, ".", a);
+    defer a.free(dir);
+    var pack = try pk.Pack.open(a, io, dir);
+    defer pack.deinit();
+    for (0..4) |rank| {
+        const part = try partition.share(width, rank, 4);
+        var b: named.Builder = .{ .gpa = a, .io = io, .pack = &pack };
+        defer b.deinit();
+        var top = [_]@import("plan.zig").NativePlan{.{ .name = "embed.weight", .info = try pack.need("embed.weight"), .slice = .vocab_rows, .bytes = part.len * hidden * 2 }};
+        const plan: @import("plan.zig").Plan = .{ .arena = undefined, .rank = @intCast(rank), .world = 4, .layers = &.{}, .top = &top, .head = null, .dspark = null, .leftovers = &.{} };
+        try b.top(&plan);
+        const embed = b.find("embed").?;
+        try std.testing.expectEqual(@as(usize, part.len), embed.shape[0]);
+        try std.testing.expectEqual(@as(usize, part.len * hidden * 2), embed.bytes.len);
+        for (0..part.len * hidden) |i| try std.testing.expectEqual(@as(u16, @intCast(part.lo * hidden + i)), std.mem.readInt(u16, embed.bytes[i * 2 ..][0..2], .little));
+    }
+}
diff --git a/zig/src/families/deepseek_v41/world4_test.zig b/zig/src/families/deepseek_v41/world4_test.zig
new file mode 100644
index 0000000000..715ad6b158
--- /dev/null
+++ b/zig/src/families/deepseek_v41/world4_test.zig
@@ -0,0 +1,104 @@
+const std = @import("std");
+const admission = @import("block_admission.zig");
+const Config = @import("config.zig").Config;
+const Options = @import("block.zig").Options;
+const partition = @import("partition.zig");
+
+test "world four: the default expert width is admitted" {
+    const c: Config = .{};
+    try admission.options(&c, Options{ .world = 4 }, &.{0}, 1);
+    try std.testing.expectEqual(@as(u32, 2304), c.expert_width);
+}
+
+test "world four: block shares cover each width once in order" {
+    for ([_]u32{ 128, 256, 512, 2304, 7168, 129280 }) |width| {
+        for ([_]u32{ 1, 2, 4 }) |world| {
+            if (width / 128 < world) continue;
+            var cursor: u32 = 0;
+            var previous: u32 = std.math.maxInt(u32);
+            for (0..world) |rank| {
+                const p = try partition.share(width, rank, world);
+                try std.testing.expectEqual(cursor, p.lo);
+                try std.testing.expect(p.len % 128 == 0 and p.len > 0);
+                try std.testing.expect(p.len <= previous and previous - p.len <= 128 or rank == 0);
+                previous = p.len;
+                cursor += p.len;
+            }
+            try std.testing.expectEqual(width, cursor);
+        }
+    }
+    try std.testing.expectError(error.SplitNotOnBlocks, partition.share(2304, 0, 0));
+    try std.testing.expectError(error.SplitNotOnBlocks, partition.share(2305, 0, 4));
+    try std.testing.expectError(error.SplitNotOnBlocks, partition.share(128, 0, 4));
+    try std.testing.expectError(error.SplitNotOnBlocks, partition.share(2304, 4, 4));
+}
+
+test "world four: default expert and vocabulary shares keep whole blocks" {
+    const cfg: Config = .{};
+    for ([_]u32{ 640, 640, 512, 512 }, [_]u32{ 32384, 32384, 32256, 32256 }, 0..) |expert, vocab, rank| {
+        try std.testing.expectEqual(expert, (try partition.share(cfg.expert_width, rank, 4)).len);
+        try std.testing.expectEqual(vocab, (try partition.share(cfg.vocab, rank, 4)).len);
+        try admission.options(&cfg, Options{ .world = 4, .rank = @intCast(rank) }, &.{0}, 1);
+    }
+}
+
+test "world four: trellis and scale spans retain all blocks and legacy parts" {
+    const Group = @import("core").exl3_layout.Group;
+    for ([_]Group{ .{ .k_tiles = 448, .n_tiles = 144, .words = 32 }, .{ .k_tiles = 144, .n_tiles = 448, .words = 48 }, .{ .k_tiles = 448, .n_tiles = 8080, .words = 96 } }) |g| {
+        for ([_]@import("core").exl3_layout.Split{ .row, .col }) |split| {
+            for ([_]u32{ 1, 2 }) |world| for (0..world) |rank| {
+                try std.testing.expectEqualDeep(try g.part(split, @intCast(rank), world), try partition.group(g, split, @intCast(rank), world));
+            };
+            var bytes: u64 = 0;
+            var scale: u64 = 0;
+            for (0..4) |rank| {
+                const p = try partition.group(g, split, @intCast(rank), 4);
+                try std.testing.expectEqual(bytes, p.trellis.offset);
+                try std.testing.expectEqual(scale, if (split == .row) p.suh.offset else p.svh.offset);
+                bytes += p.trellis.len;
+                scale += if (split == .row) p.suh.len else p.svh.len;
+            }
+            try std.testing.expectEqual(if (split == .row) g.trellisBytes() else @as(u64, g.n_tiles) * g.tileBytes(), bytes);
+            try std.testing.expectEqual(2 * @as(u64, if (split == .row) g.k() else g.n()), scale);
+        }
+    }
+}
+
+test "world four: decode widths are local and native prefill stays explicitly held" {
+    const block = @import("block.zig");
+    const cfg: Config = .{};
+    for (0..4) |rank| for ([_]bool{ false, true }) |prefill| {
+        var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
+        defer arena.deinit();
+        const a = arena.allocator();
+        var widths: block.Widths = .{};
+        try @import("widths_listing.zig").parseDense(a, &widths, @embedFile("fixtures/dense-k2.txt"));
+        try @import("widths_listing.zig").parseExperts(a, &widths, @embedFile("fixtures/expert-k2.txt"));
+        const o: Options = .{ .world = 4, .rank = @intCast(rank) };
+        if (prefill) {
+            try std.testing.expectError(error.Unsupported, @import("block_prefill.zig").emitPrefill(a, &cfg, &widths, o, &.{ 0, 40 }, 17, 0, true));
+            try std.testing.expectError(error.UnsupportedPrefillHeads, admission.nativePrefill(&cfg, 4));
+            continue;
+        }
+        const calls = try block.emit(a, &cfg, &widths, o, &.{ 0, 40 }, 1, 0, true);
+        const expert = try partition.share(cfg.expert_width, rank, 4);
+        const vocab = try partition.share(cfg.vocab, rank, 4);
+        var seen_expert = false;
+        var seen_logits = false;
+        for (calls) |call| for (call.args) |arg| {
+            if (arg.arg != .t or arg.arg.t.role != .buf) continue;
+            const t = arg.arg.t;
+            if (std.mem.endsWith(u8, t.role.buf, "svh_g")) {
+                try std.testing.expectEqual(@as(i64, expert.len), t.shape[1]);
+                seen_expert = true;
+            }
+            if (std.mem.eql(u8, t.role.buf, "w.logits")) {
+                try std.testing.expectEqual(@as(i64, vocab.len), t.shape[1]);
+                seen_logits = true;
+            }
+        };
+        try std.testing.expect(seen_expert and seen_logits);
+    };
+    try admission.nativePrefill(&cfg, 1);
+    try admission.nativePrefill(&cfg, 2);
+}

From 9e2f85bcb15e7ea2d737db2ae294a008451bcdbe Mon Sep 17 00:00:00 2001
From: ashhart <71890694+ashhart@users.noreply.github.com>
Date: Thu, 22 Oct 2026 22:35:06 +0100
Subject: [PATCH 2/2] deepseek_v41: pad 16-head prefill through the 32-head
 core

Copy real queries after RoPE and pad queries and sinks for the existing
32-head attention specialization. Extract the real groups with the full
run stride and segment offset before the output projections.

Keep world-1/2 calls and weight plans unchanged. Cover copy geometry,
default four-rank programs and legacy fingerprints on the host, and add
an AOT-only comparison runner for device qualification.
---
 zig/build/dsv41_cuda.zig                      |   3 +
 .../families/deepseek_v41/block_admission.zig |   2 +-
 .../deepseek_v41/block_prefill_attn.zig       |  22 ++-
 .../deepseek_v41/forward_prefill_glue.zig     |   1 +
 .../families/deepseek_v41/prefill_heads.zig   |  69 +++++++
 .../deepseek_v41/prefill_heads_check.zig      | 183 ++++++++++++++++++
 .../deepseek_v41/prefill_heads_copy_test.zig  |  79 ++++++++
 .../deepseek_v41/prefill_heads_test.zig       | 130 +++++++++++++
 zig/src/families/deepseek_v41/state_test.zig  |   1 +
 zig/src/families/deepseek_v41/world4_test.zig |  14 +-
 10 files changed, 488 insertions(+), 16 deletions(-)
 create mode 100644 zig/src/families/deepseek_v41/prefill_heads.zig
 create mode 100644 zig/src/families/deepseek_v41/prefill_heads_check.zig
 create mode 100644 zig/src/families/deepseek_v41/prefill_heads_copy_test.zig
 create mode 100644 zig/src/families/deepseek_v41/prefill_heads_test.zig

diff --git a/zig/build/dsv41_cuda.zig b/zig/build/dsv41_cuda.zig
index 7967ba0ef1..6fafa3319e 100644
--- a/zig/build/dsv41_cuda.zig
+++ b/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").? }));
diff --git a/zig/src/families/deepseek_v41/block_admission.zig b/zig/src/families/deepseek_v41/block_admission.zig
index cd27e01e39..c65c337439 100644
--- a/zig/src/families/deepseek_v41/block_admission.zig
+++ b/zig/src/families/deepseek_v41/block_admission.zig
@@ -5,7 +5,7 @@ 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) % tile != 0) return error.UnsupportedPrefillHeads;
+    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 {
diff --git a/zig/src/families/deepseek_v41/block_prefill_attn.zig b/zig/src/families/deepseek_v41/block_prefill_attn.zig
index e13c038762..3f30369058 100644
--- a/zig/src/families/deepseek_v41/block_prefill_attn.zig
+++ b/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() } });
 }
diff --git a/zig/src/families/deepseek_v41/forward_prefill_glue.zig b/zig/src/families/deepseek_v41/forward_prefill_glue.zig
index 51935ca81d..41aeaa6e09 100644
--- a/zig/src/families/deepseek_v41/forward_prefill_glue.zig
+++ b/zig/src/families/deepseek_v41/forward_prefill_glue.zig
@@ -174,6 +174,7 @@ pub fn glueFn(ctx: *anyopaque, r: *run.Runner, c: *const calls.Call) anyerror!vo
 }
 
 pub fn glueOne(ctx: *anyopaque, r: *run.Runner, c: *const calls.Call) anyerror!void {
+    if (try @import("prefill_heads.zig").glue(r, c)) return;
     const f: *Forward = @ptrCast(@alignCast(ctx));
     const st = pf.state(f);
     const step = c.name["glue.".len..];
diff --git a/zig/src/families/deepseek_v41/prefill_heads.zig b/zig/src/families/deepseek_v41/prefill_heads.zig
new file mode 100644
index 0000000000..90f1608025
--- /dev/null
+++ b/zig/src/families/deepseek_v41/prefill_heads.zig
@@ -0,0 +1,69 @@
+//! Sixteen real heads through the unchanged 32-head prefill core; copies preserve the run's group-major layout.
+const std = @import("std");
+const cuda = @import("cuda");
+const calls = @import("block_calls.zig");
+const E = @import("block.zig").Emitter;
+const Arg = calls.Arg;
+
+pub fn query(e: *E, q: Arg, hd: i64) !Arg {
+    const padded = try e.buf("L.pf.q32", .{}, .bf16, &.{ e.n, 32 * hd });
+    try e.glue("pf_heads_q", &.{ q, padded });
+    return padded;
+}
+
+pub fn sink(e: *E, real: Arg) !Arg {
+    const padded = try e.buf("L.pf.sink32", .{}, .f32, &.{32});
+    try e.glue("pf_heads_sink", &.{ real, padded });
+    return padded;
+}
+
+pub fn extract(e: *E, source: Arg, dest: Arg) !void {
+    try e.glue("pf_heads_out", &.{ source, dest });
+}
+
+/// Only these copy calls use this dispatcher; all storage belongs to the existing buffer plan and stream.
+pub fn glue(r: anytype, c: *const calls.Call) !bool {
+    if (!std.mem.startsWith(u8, c.name, "glue.pf_heads_")) return false;
+    if (c.args.len != 2 or c.args[0].arg != .t or c.args[1].arg != .t) return error.BadGlue;
+    const src = c.args[0].arg.t;
+    const dst = c.args[1].arg.t;
+    if (src.dt != dst.dt or src.shape.len != src.stride.len or dst.shape.len != dst.stride.len) return error.BadGlue;
+    var width: usize = undefined;
+    var height: usize = undefined;
+    var sp: usize = undefined;
+    var dp: usize = undefined;
+    var zero: usize = 0;
+    if (std.mem.eql(u8, c.name, "glue.pf_heads_q")) {
+        if (src.dt != .bf16 or src.shape.len != 2 or dst.shape.len != 2 or src.shape[0] <= 0 or src.shape[1] <= 0 or
+            dst.shape[0] != src.shape[0] or dst.shape[1] != 2 * src.shape[1] or src.stride[1] != 1 or dst.stride[1] != 1 or
+            src.stride[0] != src.shape[1] or dst.stride[0] != dst.shape[1]) return error.BadGlue;
+        width = @intCast(2 * src.shape[1]);
+        height = @intCast(src.shape[0]);
+        sp = width;
+        dp = 2 * width;
+        zero = height * dp;
+    } else if (std.mem.eql(u8, c.name, "glue.pf_heads_sink")) {
+        if (src.dt != .f32 or src.shape.len != 1 or dst.shape.len != 1 or src.shape[0] != 16 or dst.shape[0] != 32 or
+            src.stride[0] != 1 or dst.stride[0] != 1) return error.BadGlue;
+        width = 64;
+        height = 1;
+        sp = width;
+        dp = 128;
+        zero = dp;
+    } else if (std.mem.eql(u8, c.name, "glue.pf_heads_out")) {
+        if (src.dt != .bf16 or src.shape.len != 3 or dst.shape.len != 3 or dst.shape[0] <= 0 or dst.shape[1] <= 0 or dst.shape[2] <= 0 or
+            src.shape[0] != 2 * dst.shape[0] or src.shape[1] != dst.shape[1] or src.shape[2] != dst.shape[2] or
+            src.stride[2] != 1 or dst.stride[2] != 1 or src.stride[1] != src.shape[2] or dst.stride[1] != dst.shape[2] or
+            src.stride[0] < src.shape[1] * src.shape[2] or dst.stride[0] < dst.shape[1] * dst.shape[2]) return error.BadGlue;
+        width = @intCast(2 * dst.shape[1] * dst.shape[2]);
+        height = @intCast(dst.shape[0]);
+        sp = @intCast(2 * src.stride[0]);
+        dp = @intCast(2 * dst.stride[0]);
+    } else return error.BadGlue;
+    const source = try r.tensorAddr(src);
+    const dest = try r.tensorAddr(dst);
+    if (zero != 0) try r.d.check(r.d.api.cuMemsetD8Async(dest, 0, zero, r.stream.handle), "cuMemsetD8Async");
+    const copy: cuda.abi.Memcpy2D = .{ .src = source, .src_pitch = sp, .dst = dest, .dst_pitch = dp, .width = width, .height = height };
+    try r.d.check(r.d.api.cuMemcpy2DAsync_v2(&copy, r.stream.handle), "cuMemcpy2DAsync");
+    return true;
+}
diff --git a/zig/src/families/deepseek_v41/prefill_heads_check.zig b/zig/src/families/deepseek_v41/prefill_heads_check.zig
new file mode 100644
index 0000000000..94f4b855e3
--- /dev/null
+++ b/zig/src/families/deepseek_v41/prefill_heads_check.zig
@@ -0,0 +1,183 @@
+//! One GPU: production padding/copy calls and the unchanged fused core against a direct 32-head launch.
+const std = @import("std");
+const cuda = @import("cuda");
+const calls = @import("block_calls.zig");
+const block = @import("block.zig");
+const bp = @import("block_prefill.zig");
+const heads = @import("prefill_heads.zig");
+const tc = @import("triton_call.zig");
+const Buffer = struct { dev: cuda.DeviceBuffer, host: []u8 };
+const Case = struct { total: i64, rows: i64, row0: i64 = 0 };
+const Receipt = struct { total: i64, rows: i64, row0: i64, sha256: [64]u8 };
+
+const Memory = struct {
+    a: std.mem.Allocator,
+    d: *const cuda.Driver,
+    stream: cuda.Stream,
+    buffers: std.StringArrayHashMapUnmanaged(Buffer) = .empty,
+
+    fn close(m: *Memory) void {
+        m.stream.synchronize() catch std.process.exit(1);
+        for (m.buffers.values()) |*b| b.dev.free();
+    }
+    fn add(m: *Memory, name: []const u8, len: usize, poison: u8) !void {
+        if (m.buffers.contains(name)) return;
+        const host = try m.a.alloc(u8, len + 32);
+        @memset(host, 0xA5);
+        @memset(host[16..][0..len], poison);
+        var dev = try cuda.DeviceBuffer.alloc(m.d, host.len);
+        errdefer dev.free();
+        try m.buffers.put(m.a, name, .{ .dev = dev, .host = host });
+    }
+    fn data(m: *Memory, name: []const u8) []u8 {
+        const host = m.buffers.get(name).?.host;
+        return host[16 .. host.len - 16];
+    }
+    pub fn tensorAddr(m: *Memory, t: calls.Tensor) !u64 {
+        const name = switch (t.role) {
+            .buf => |s| s,
+            .weight => |s| s,
+            .empty => return 0,
+        };
+        return (m.buffers.get(name) orelse return error.Unbound).dev.ptr + 16 + @as(u64, @intCast(t.offset));
+    }
+    fn address(ctx: *anyopaque, t: calls.Tensor) anyerror!u64 {
+        const m: *Memory = @ptrCast(@alignCast(ctx));
+        return m.tensorAddr(t);
+    }
+    fn launch(m: *Memory, set: *const cuda.aot.Set, c: *const calls.Call) !void {
+        var specs: std.ArrayList(cuda.aot.Spec) = .empty;
+        for (set.variants) |v| if (std.mem.eql(u8, v.spec.@"fn", c.name)) try specs.append(m.a, v.spec);
+        const split = try tc.split(m.a, c, specs.items, .{ .ctx = m, .of = address });
+        try set.run(m.stream, c.name, .{ @intCast(c.grid[0]), @intCast(c.grid[1]), @intCast(c.grid[2]) }, split.args, split.consts);
+    }
+};
+
+fn word(bytes: []u8, i: usize, value: f32) void {
+    std.mem.writeInt(u32, bytes[i * 4 ..][0..4], @bitCast(value), .little);
+}
+fn bf16(bytes: []u8, i: usize, value: f32) void {
+    std.mem.writeInt(u16, bytes[i * 2 ..][0..2], @import("core").contract.spec.toBf16(value), .little);
+}
+fn nameOf(t: calls.Tensor) []const u8 {
+    return switch (t.role) {
+        .buf => |s| s,
+        .weight => |s| s,
+        .empty => "",
+    };
+}
+
+fn one(gpa: std.mem.Allocator, d: *const cuda.Driver, stream: cuda.Stream, set: *const cuda.aot.Set, shape: Case) !Receipt {
+    var arena = std.heap.ArenaAllocator.init(gpa);
+    defer arena.deinit();
+    const a = arena.allocator();
+    var m: Memory = .{ .a = a, .d = d, .stream = stream };
+    defer m.close();
+    const cfg: @import("config.zig").Config = .{};
+    var widths: block.Widths = .{};
+    var e: block.Emitter = .{ .a = a, .cfg = &cfg, .w = &widths, .o = .{ .world = 4 }, .n = shape.total, .start = 128 };
+    const q = try e.buf("L.q", .{}, .bf16, &.{ shape.total, 8192 });
+    const padded = try heads.query(&e, q, 512);
+    const out = try e.buf("L.o", .{}, .bf16, &.{ 2, shape.total, 4096 });
+    const kv = try e.buf("L.kv", .{}, .bf16, &.{ shape.total, 512 });
+    e.row0 = shape.row0;
+    e.n = shape.rows;
+    var seg: bp.Seg = .{ .e = &e, .staged = false, .ring = 256 };
+    try seg.core(0, padded, kv, .none, false, false, out, shape.total);
+    var selected: std.ArrayList(calls.Call) = .empty;
+    var fused: ?calls.Call = null;
+    var sizes: std.StringArrayHashMapUnmanaged(usize) = .empty;
+    for (e.out.items) |c| {
+        if (!std.mem.startsWith(u8, c.name, "glue.pf_heads_") and !std.mem.eql(u8, c.name, "_fused")) continue;
+        try selected.append(a, c);
+        if (c.triton) fused = c;
+        for (c.args) |x| if (x.arg == .t) {
+            const t = x.arg.t;
+            const size: usize = @intCast(@import("buffers.zig").extent(t));
+            const entry = try sizes.getOrPut(a, nameOf(t));
+            if (!entry.found_existing) entry.value_ptr.* = 0;
+            entry.value_ptr.* = @max(entry.value_ptr.*, size);
+        };
+    }
+    for (sizes.keys(), sizes.values()) |name, size| try m.add(name, size, 0);
+    // The real output is poisoned differently from the direct output, so two unwritten outputs cannot agree.
+    try m.add("ref.q", @intCast(shape.total * 32768), 0);
+    try m.add("ref.sink", 128, 0);
+    try m.add("ref.out", @intCast(4 * shape.total * 8192), 0x5A);
+    @memset(m.data("L.o"), 0xA5);
+    @memset(m.data("L.pf.o32"), 0xA5);
+    const real_q = m.data("L.q");
+    for (0..real_q.len / 2) |i| bf16(real_q, i, @as(f32, @floatFromInt(i % 29)) / 32 - 0.4);
+    const total: usize = @intCast(shape.total);
+    for (0..total) |r| @memcpy(m.data("ref.q")[r * 32768 ..][0..16384], real_q[r * 16384 ..][0..16384]);
+    for (0..16) |h| word(m.data("L0.attn.sink"), h, @as(f32, @floatFromInt(h)) / 16 - 2);
+    @memcpy(m.data("ref.sink")[0..64], m.data("L0.attn.sink"));
+    const v = m.data("s.L0.swa.v");
+    for (0..v.len / 576) |r| {
+        @memset(v[r * 576 ..][0..448], 0x30);
+        for (0..64) |j| bf16(v[r * 576 + 448 ..], j, @as(f32, @floatFromInt((j + r) % 13)) / 32);
+    }
+    @memset(m.data("s.L0.swa.s"), 127);
+    const cs = m.data("s.rope.main");
+    for (0..cs.len / 256) |r| for (0..32) |j| word(cs, r * 64 + j, 1);
+    std.mem.writeInt(i32, m.data("w.pos")[0..4], 128, .little);
+    for (m.buffers.values()) |b| try b.dev.upload(0, b.host);
+    for (selected.items) |*c| {
+        if (c.triton) try m.launch(set, c) else if (!try heads.glue(&m, c)) return error.NoBinding;
+    }
+    var direct = fused orelse return error.NoFused;
+    const args = try a.dupe(calls.Named, direct.args);
+    for (args) |*x| {
+        if (std.mem.eql(u8, x.name, "Q")) x.arg.t.role = .{ .buf = "ref.q" };
+        if (std.mem.eql(u8, x.name, "SINK")) x.arg.t.role = .{ .buf = "ref.sink" };
+        if (std.mem.eql(u8, x.name, "OUT")) x.arg.t.role = .{ .buf = "ref.out" };
+    }
+    direct.args = args;
+    try m.launch(set, &direct);
+    try stream.synchronize();
+    for (m.buffers.values()) |b| {
+        try b.dev.download(0, b.host);
+        if (!std.mem.allEqual(u8, b.host[0..16], 0xA5) or !std.mem.allEqual(u8, b.host[b.host.len - 16 ..], 0xA5)) return error.GuardChanged;
+    }
+    var hash = std.crypto.hash.sha2.Sha256.init(.{});
+    const actual = m.data("L.o");
+    const reference = m.data("ref.out");
+    const row0: usize = @intCast(shape.row0);
+    const rows: usize = @intCast(shape.rows);
+    for (0..2) |g| for (0..total) |r| {
+        const at = (g * total + r) * 8192;
+        const bytes = actual[at..][0..8192];
+        if (r >= row0 and r < row0 + rows) {
+            if (!std.mem.eql(u8, bytes, reference[at..][0..8192])) return error.DifferentHeads;
+            hash.update(bytes);
+        } else if (!std.mem.allEqual(u8, bytes, 0xA5)) return error.WrongRows;
+    };
+    var digest: [32]u8 = undefined;
+    hash.final(&digest);
+    return .{ .total = shape.total, .rows = shape.rows, .row0 = shape.row0, .sha256 = std.fmt.bytesToHex(digest, .lower) };
+}
+
+pub fn main(init: std.process.Init) !u8 {
+    const args = try init.minimal.args.toSlice(init.arena.allocator());
+    if (args.len != 3) return error.ExpectedAotAndOutputDirectory;
+    const cwd = std.Io.Dir.cwd();
+    try cwd.createDirPath(init.io, args[2]);
+    const output = try std.fs.path.join(init.arena.allocator(), &.{ args[2], "summary.json" });
+    try cwd.writeFile(init.io, .{ .sub_path = output, .data = "{\"equal\":false,\"state\":\"running\"}\n" });
+    errdefer cwd.writeFile(init.io, .{ .sub_path = output, .data = "{\"equal\":false,\"state\":\"failed\"}\n" }) catch {};
+    var d = try cuda.Driver.open();
+    defer d.close();
+    var ctx = try cuda.Context.init(&d, 0);
+    defer ctx.deinit();
+    var stream = try cuda.Stream.init(&d, false);
+    defer stream.deinit();
+    var set = try cuda.aot.Set.load(init.gpa, init.io, &d, ctx.device, args[1]);
+    defer set.deinit();
+    const shapes = [_]Case{ .{ .total = 1, .rows = 1 }, .{ .total = 16, .rows = 16 }, .{ .total = 17, .rows = 17 }, .{ .total = 33, .rows = 33 }, .{ .total = 129, .rows = 129 }, .{ .total = 50, .rows = 17, .row0 = 33 } };
+    var receipts: [shapes.len]Receipt = undefined;
+    for (shapes, &receipts) |shape, *receipt| receipt.* = try one(init.gpa, &d, stream, &set, shape);
+    const text = try std.json.Stringify.valueAlloc(init.arena.allocator(), .{ .sm = try ctx.capability(), .equal = true, .cases = receipts }, .{});
+    try cwd.writeFile(init.io, .{ .sub_path = output, .data = text });
+    std.debug.print("PASS padded prefill heads: {d} cases, real bytes and row guards equal\n", .{shapes.len});
+    return 0;
+}
diff --git a/zig/src/families/deepseek_v41/prefill_heads_copy_test.zig b/zig/src/families/deepseek_v41/prefill_heads_copy_test.zig
new file mode 100644
index 0000000000..c1c7b7c6c6
--- /dev/null
+++ b/zig/src/families/deepseek_v41/prefill_heads_copy_test.zig
@@ -0,0 +1,79 @@
+const std = @import("std");
+const cuda = @import("cuda");
+const calls = @import("block_calls.zig");
+const t = std.testing;
+const Mock = struct {
+    var fail = false;
+    var copies: usize = 0;
+    var zeros: usize = 0;
+    fn zero(dst: u64, byte: u8, len: usize, _: cuda.abi.Stream) callconv(.c) cuda.abi.Result {
+        zeros += 1;
+        @memset(@as([*]u8, @ptrFromInt(dst))[0..len], byte);
+        return 0;
+    }
+    fn copy(c: *const cuda.abi.Memcpy2D, _: cuda.abi.Stream) callconv(.c) cuda.abi.Result {
+        copies += 1;
+        if (fail) return cuda.abi.error_not_ready;
+        for (0..c.height) |row| @memcpy(@as([*]u8, @ptrFromInt(c.dst + row * c.dst_pitch))[0..c.width], @as([*]const u8, @ptrFromInt(c.src + row * c.src_pitch))[0..c.width]);
+        return 0;
+    }
+};
+
+fn call(name: []const u8, args: *[2]calls.Named, src: calls.Tensor, dst: calls.Tensor) calls.Call {
+    args.* = .{ .{ .arg = .{ .t = src } }, .{ .arg = .{ .t = dst } } };
+    return .{ .triton = false, .glue = true, .name = name, .args = args };
+}
+
+test "prefill heads: production copy dispatch pads finite heads and extracts pitched groups without touching guards" {
+    var d: cuda.Driver = undefined;
+    d.api.cuMemsetD8Async = Mock.zero;
+    d.api.cuMemcpy2DAsync_v2 = Mock.copy;
+    Mock.fail = false;
+    Mock.copies = 0;
+    Mock.zeros = 0;
+    var r: @import("run.zig").Runner = .{ .gpa = t.allocator, .d = &d, .stream = .{ .d = &d, .handle = null }, .kernels = undefined, .triton = undefined, .weights = undefined };
+    defer r.deinit();
+    var source: [3 * 32]u16 align(16) = undefined;
+    for (&source, 0..) |*x, i| x.* = @intCast(i + 1);
+    var padded: [3 * 64 + 16]u16 align(16) = @splat(0xA5A5);
+    try r.external("source", @intFromPtr(&source));
+    try r.external("padded", @intFromPtr(&padded));
+    var args: [2]calls.Named = undefined;
+    var c = call("glue.pf_heads_q", &args, .{ .role = .{ .buf = "source" }, .dt = .bf16, .shape = &.{ 3, 32 }, .stride = &.{ 32, 1 } }, .{ .role = .{ .buf = "padded" }, .dt = .bf16, .shape = &.{ 3, 64 }, .stride = &.{ 64, 1 }, .offset = 16 });
+    try @import("forward_prefill_glue.zig").glueOne(undefined, &r, &c);
+    for (0..3) |row| {
+        try t.expectEqualSlices(u16, source[row * 32 ..][0..32], padded[8 + row * 64 ..][0..32]);
+        try t.expect(std.mem.allEqual(u16, padded[8 + row * 64 + 32 ..][0..32], 0));
+    }
+    try t.expect(std.mem.allEqual(u16, padded[0..8], 0xA5A5));
+    try t.expect(std.mem.allEqual(u16, padded[padded.len - 8 ..], 0xA5A5));
+    var sink: [16]f32 align(16) = undefined;
+    for (&sink, 0..) |*x, i| x.* = @as(f32, @floatFromInt(i)) - 8;
+    var sink32: [36]f32 align(16) = @splat(99);
+    try r.external("sink", @intFromPtr(&sink));
+    try r.external("sink32", @intFromPtr(&sink32));
+    c = call("glue.pf_heads_sink", &args, .{ .role = .{ .buf = "sink" }, .dt = .f32, .shape = &.{16}, .stride = &.{1} }, .{ .role = .{ .buf = "sink32" }, .dt = .f32, .shape = &.{32}, .stride = &.{1}, .offset = 8 });
+    try @import("forward_prefill_glue.zig").glueOne(undefined, &r, &c);
+    try t.expectEqualSlices(f32, &sink, sink32[2..18]);
+    try t.expect(std.mem.allEqual(f32, sink32[18..34], 0));
+    try t.expect(sink32[0] == 99 and sink32[35] == 99);
+    var full: [4 * 11 * 16]u16 align(16) = undefined;
+    for (&full, 0..) |*x, i| x.* = @intCast(i);
+    var real: [2 * 11 * 16]u16 align(16) = @splat(0xA5A5);
+    try r.external("full", @intFromPtr(&full));
+    try r.external("real", @intFromPtr(&real));
+    c = call("glue.pf_heads_out", &args, .{ .role = .{ .buf = "full" }, .dt = .bf16, .shape = &.{ 4, 5, 16 }, .stride = &.{ 11 * 16, 16, 1 }, .offset = 3 * 16 * 2 }, .{ .role = .{ .buf = "real" }, .dt = .bf16, .shape = &.{ 2, 5, 16 }, .stride = &.{ 11 * 16, 16, 1 }, .offset = 3 * 16 * 2 });
+    try @import("forward_prefill_glue.zig").glueOne(undefined, &r, &c);
+    for (&real, 0..) |value, i| {
+        const row = (i / 16) % 11;
+        try t.expectEqual(if (row >= 3 and row < 8) full[i] else @as(u16, 0xA5A5), value);
+    }
+    try t.expectEqual(@as(usize, 3), Mock.copies);
+    try t.expectEqual(@as(usize, 2), Mock.zeros);
+    Mock.fail = true;
+    defer Mock.fail = false;
+    try t.expectError(error.NotReady, @import("forward_prefill_glue.zig").glueOne(undefined, &r, &c));
+    args[1].arg.t.stride = &.{ 1, 16, 1 };
+    try t.expectError(error.BadGlue, @import("forward_prefill_glue.zig").glueOne(undefined, &r, &c));
+    try t.expectEqual(@as(usize, 4), Mock.copies);
+}
diff --git a/zig/src/families/deepseek_v41/prefill_heads_test.zig b/zig/src/families/deepseek_v41/prefill_heads_test.zig
new file mode 100644
index 0000000000..15993d48b8
--- /dev/null
+++ b/zig/src/families/deepseek_v41/prefill_heads_test.zig
@@ -0,0 +1,130 @@
+const std = @import("std");
+const t = std.testing;
+const block = @import("block.zig");
+const bp = @import("block_prefill.zig");
+const calls = @import("block_calls.zig");
+const buffers = @import("buffers.zig");
+const Config = @import("config.zig").Config;
+
+fn widths(a: std.mem.Allocator) !block.Widths {
+    var w: block.Widths = .{};
+    try @import("widths_listing.zig").parseDense(a, &w, @embedFile("fixtures/dense-k2.txt"));
+    try @import("widths_listing.zig").parseExperts(a, &w, @embedFile("fixtures/expert-k2.txt"));
+    for (0..43) |layer| try w.gm.put(a, @intCast(layer), .{ (1 << 4) | (1 << 6) | (1 << 8), (1 << 4) | (1 << 10) });
+    for (0..43) |layer| for (4..8) |g| {
+        const first = try std.fmt.allocPrint(a, "L{d}.attn.wo_a.0", .{layer});
+        try w.dense.put(a, try std.fmt.allocPrint(a, "L{d}.attn.wo_a.{d}", .{ layer, g }), try w.k2(first));
+    };
+    return w;
+}
+
+fn arg(c: calls.Call, name: []const u8) calls.Arg {
+    for (c.args) |x| if (std.mem.eql(u8, x.name, name)) return x.arg;
+    unreachable;
+}
+
+test "prefill heads: legacy world one and two calls and buffer plans keep their fingerprint" {
+    var hash = std.crypto.hash.sha2.Sha256.init(.{});
+    const cfg: Config = .{};
+    for ([_]u32{ 1, 2 }) |world| for (0..world) |rank| for ([_]i64{ 16, 17, 129 }) |rows| {
+        var arena = std.heap.ArenaAllocator.init(t.allocator);
+        defer arena.deinit();
+        const a = arena.allocator();
+        var w = try widths(a);
+        const opts: block.Options = .{ .world = world, .rank = @intCast(rank) };
+        const cs = try bp.emitPrefill(a, &cfg, &w, opts, &.{ 0, 2, 20, 24, 40 }, rows, 2048, true);
+        var plan: buffers.Plan = .{ .a = a };
+        try plan.add(cs);
+        var text: std.Io.Writer.Allocating = .init(a);
+        try std.json.Stringify.value(.{ .calls = cs, .roles = plan.sizes.keys(), .sizes = plan.sizes.values(), .weights = plan.weights.keys() }, .{}, &text.writer);
+        hash.update(text.written());
+    };
+    var digest: [32]u8 = undefined;
+    hash.final(&digest);
+    const hex = std.fmt.bytesToHex(digest, .lower);
+    std.debug.print("legacy prefill fingerprint: {s}\n", .{hex});
+    try t.expectEqualStrings("66015c45a28104778de630ab14e49aa4c871cefb95fad2e3a06c142c05449b58", &hex);
+}
+
+test "prefill heads: world four emits and plans the padded 32-head operation" {
+    const cfg: Config = .{};
+    try @import("block_admission.zig").nativePrefill(&cfg, 4);
+    for (0..4) |rank| {
+        var arena = std.heap.ArenaAllocator.init(t.allocator);
+        defer arena.deinit();
+        const a = arena.allocator();
+        var w = try widths(a);
+        const opts: block.Options = .{ .world = 4, .rank = @intCast(rank) };
+        var layers: [43]u32 = undefined;
+        for (&layers, 0..) |*layer, i| layer.* = @intCast(i);
+        const cs = try bp.emitPrefill(a, &cfg, &w, opts, &layers, 17, 0, true);
+        var plan: buffers.Plan = .{ .a = a };
+        try plan.add(cs);
+        try t.expectEqual(@as(u64, 557056), plan.sizes.get("L.pf.q32").?);
+        try t.expectEqual(@as(u64, 128), plan.sizes.get("L.pf.sink32").?);
+        try t.expectEqual(@as(u64, 557056), plan.sizes.get("L.pf.o32").?);
+        var seen: usize = 0;
+        for (cs) |c| if (std.mem.eql(u8, c.name, "_fused")) {
+            seen += 1;
+            try t.expectEqual(@as(i64, 32), arg(c, "H").i);
+            try t.expectEqual(@as(i64, 32), arg(c, "BMQ").i);
+            try t.expectEqual(@as(i64, 8), arg(c, "HG").i);
+            try t.expectEqualSlices(i64, &.{ 17, 32, 512 }, arg(c, "Q").t.shape);
+            try t.expectEqualSlices(i64, &.{32}, arg(c, "SINK").t.shape);
+            try t.expectEqualSlices(i64, &.{ 4, 17, 4096 }, arg(c, "OUT").t.shape);
+            try t.expectEqual([3]i64{ 17, 1, 1 }, c.grid);
+        };
+        try t.expectEqual(@as(usize, 43), seen);
+    }
+}
+
+test {
+    _ = @import("prefill_heads_copy_test.zig");
+}
+
+test "prefill heads: several segments preserve row0 and group strides before the output projection" {
+    var arena = std.heap.ArenaAllocator.init(t.allocator);
+    defer arena.deinit();
+    const a = arena.allocator();
+    const cfg: Config = .{};
+    var w = try widths(a);
+    const cs = try bp.emitMulti(a, &cfg, &w, .{ .world = 4 }, &.{0}, &.{ .{ .slot = 0, .start = 0, .n = 33 }, .{ .slot = 1, .start = 128, .n = 17 } }, .whole);
+    var seen: usize = 0;
+    var copied: usize = 0;
+    var query_at: usize = 0;
+    for (cs, 0..) |c, i| {
+        if (std.mem.eql(u8, c.name, "glue.pf_heads_q")) {
+            query_at = i;
+            try t.expectEqualSlices(i64, &.{ 50, 8192 }, c.args[0].arg.t.shape);
+            try t.expectEqualSlices(i64, &.{ 50, 16384 }, c.args[1].arg.t.shape);
+            var rope_at: ?usize = null;
+            for (cs[0..i], 0..) |old, j| if (std.mem.eql(u8, old.name, "_rope")) {
+                rope_at = j;
+            };
+            try t.expect(rope_at != null and rope_at.? < i);
+        }
+        if (std.mem.eql(u8, c.name, "_fused")) {
+            try t.expect(query_at < i);
+            const offset: i64 = if (seen == 0) 0 else 33 * 4096 * 2;
+            const rows: i64 = if (seen == 0) 33 else 17;
+            const out = arg(c, "OUT").t;
+            try t.expectEqualSlices(i64, &.{ 4, rows, 4096 }, out.shape);
+            try t.expectEqualSlices(i64, &.{ 50 * 4096, 4096, 1 }, out.stride);
+            try t.expectEqual(offset, out.offset);
+            try t.expectEqual(offset * 4, arg(c, "Q").t.offset);
+            seen += 1;
+        }
+        if (std.mem.eql(u8, c.name, "glue.pf_heads_out")) {
+            try t.expectEqual(seen, copied + 1);
+            const src = c.args[0].arg.t;
+            const dst = c.args[1].arg.t;
+            try t.expectEqualSlices(i64, &.{ 2, src.shape[1], 4096 }, dst.shape);
+            try t.expectEqualSlices(i64, src.stride, dst.stride);
+            try t.expectEqual(src.offset, dst.offset);
+            copied += 1;
+        }
+        if (std.mem.eql(u8, c.name, "_gemm") and arg(c, "OUT") == .t and std.mem.eql(u8, arg(c, "OUT").t.role.buf, "L.z")) try t.expectEqual(@as(usize, 2), copied);
+    }
+    try t.expectEqual(@as(usize, 2), seen);
+    try t.expectEqual(seen, copied);
+}
diff --git a/zig/src/families/deepseek_v41/state_test.zig b/zig/src/families/deepseek_v41/state_test.zig
index adcc7f29cf..85bbbbcfb7 100644
--- a/zig/src/families/deepseek_v41/state_test.zig
+++ b/zig/src/families/deepseek_v41/state_test.zig
@@ -1,6 +1,7 @@
 const std = @import("std");
 test {
     _ = @import("world4_test.zig");
+    _ = @import("prefill_heads_test.zig");
     _ = @import("world4_plan_test.zig");
     _ = @import("world4_stage_test.zig");
 }
diff --git a/zig/src/families/deepseek_v41/world4_test.zig b/zig/src/families/deepseek_v41/world4_test.zig
index 715ad6b158..a1c3632570 100644
--- a/zig/src/families/deepseek_v41/world4_test.zig
+++ b/zig/src/families/deepseek_v41/world4_test.zig
@@ -64,7 +64,7 @@ test "world four: trellis and scale spans retain all blocks and legacy parts" {
     }
 }
 
-test "world four: decode widths are local and native prefill stays explicitly held" {
+test "world four: decode and padded prefill widths are local" {
     const block = @import("block.zig");
     const cfg: Config = .{};
     for (0..4) |rank| for ([_]bool{ false, true }) |prefill| {
@@ -74,13 +74,9 @@ test "world four: decode widths are local and native prefill stays explicitly he
         var widths: block.Widths = .{};
         try @import("widths_listing.zig").parseDense(a, &widths, @embedFile("fixtures/dense-k2.txt"));
         try @import("widths_listing.zig").parseExperts(a, &widths, @embedFile("fixtures/expert-k2.txt"));
+        for ([_]u32{ 0, 40 }) |L| try widths.gm.put(a, L, .{ (1 << 4) | (1 << 6) | (1 << 8), (1 << 4) | (1 << 10) });
         const o: Options = .{ .world = 4, .rank = @intCast(rank) };
-        if (prefill) {
-            try std.testing.expectError(error.Unsupported, @import("block_prefill.zig").emitPrefill(a, &cfg, &widths, o, &.{ 0, 40 }, 17, 0, true));
-            try std.testing.expectError(error.UnsupportedPrefillHeads, admission.nativePrefill(&cfg, 4));
-            continue;
-        }
-        const calls = try block.emit(a, &cfg, &widths, o, &.{ 0, 40 }, 1, 0, true);
+        const calls = if (prefill) try @import("block_prefill.zig").emitPrefill(a, &cfg, &widths, o, &.{ 0, 40 }, 17, 0, true) else try block.emit(a, &cfg, &widths, o, &.{ 0, 40 }, 1, 0, true);
         const expert = try partition.share(cfg.expert_width, rank, 4);
         const vocab = try partition.share(cfg.vocab, rank, 4);
         var seen_expert = false;
@@ -97,8 +93,10 @@ test "world four: decode widths are local and native prefill stays explicitly he
                 seen_logits = true;
             }
         };
-        try std.testing.expect(seen_expert and seen_logits);
+        try std.testing.expect(seen_logits and (prefill or seen_expert));
     };
     try admission.nativePrefill(&cfg, 1);
     try admission.nativePrefill(&cfg, 2);
+    try admission.nativePrefill(&cfg, 4);
+    try std.testing.expectError(error.UnsupportedPrefillHeads, admission.nativePrefill(&cfg, 8));
 }