[ashhart]

Fixes #1202.

This lets DeepSeek-V4.1 run at tensor-parallel world 4. It's a draft so @flashosophy can run the four-rank cells on their setup. Thank you for pairing on this.

What changes

• Expert and vocabulary widths split into whole 128-column blocks that may differ by one block between ranks. The default config gives 640, 640, 512, 512 expert columns and 32384, 32384, 32256, 32256 vocabulary rows. Plans, shard reads, MoE and head widths, embeddings, DSpark candidates and sampling all use each rank's own offset and width. Sampling's whole-row path pads each rank's row to the largest share and strips the padding before picking.
• At world 4 each rank has 16 attention heads, but the native prefill attention kernel exists only for 32. Those 16 heads now run through the unchanged 32-head kernel: the real queries go into a zero-padded 32-head buffer, the extra sinks are zero, and the two real output groups are copied out before the output projections. So each GPU does 32 heads of prefill attention at four ranks, the same as at two.
• World 1 and 2 are unchanged. Their weight plans and prefill calls are pinned to main's fingerprints in tests.

Checked on the host

• New tests: block shares cover every width once and in order; world-4 plans build for every rank without reading weights; EXL3 parts equal the old ones at worlds 1 and 2; uneven sampling picks the same tokens at worlds 2 and 4 as on one rank, across top-k, top-p and whole-row cases; the padded copies' geometry and error paths; world-4 prefill emission over all 43 layers.
• zig fmt, the lean check and the Linux compile checks pass. Nothing has run on a GPU yet.

Four-rank runs, for @flashosophy

1. Ranks 0-1 and 2-3 have different expert and vocabulary widths, so their AOT variants differ. List them per rank and package the union:

   for r in 0 1 2 3; do TF_TP_WORLD=4 TF_TP_RANK=$r tf-dsv41 aot-needs SIGS needs-r$r; done

2. Serve at world 4. Start the followers first, use one fresh TP port, and keep the same settings on every rank:

   export TF_EXPERIMENTAL=deepseek TF_COMM_BACKEND=nccl
   export TF_TP_WORLD=4 TF_TP_RANK=$RANK TF_TP_DEVICE=$DEVICE TF_TP_MASTER=$MASTER TF_TP_PORT=$PORT
   export TF_DSV41_ASSETS=$ASSETS TF_DSV41_OWN_PREFILL=1 TF_DSV41_CONTEXT=4096
   export TF_DSV41_SAMPLING=1 TF_DSV41_DRAFTS=1 TF_DSV41_AOT_GUARD=1
   # rank 0
   tensorfold-native serve $PACK --backend cuda --host 127.0.0.1 --port 8000 --name dsv41 --context 4096 --no-thinking
   # ranks 1 to 3
   tf-dsv41 follow $PACK $ASSETS

3. Check the equalities against that server:

   python3 tools/bench_concurrent.py http://127.0.0.1:8000 dsv41 --levels 1,4 --tokens 256 --reps 3 --seed 1234 --temperatures 0,1.0 --alone --serial

   The equalities count at world 4 itself. A world-4 reply need not match world 2 token for token, because the reduction order differs.

RoCE keeps its two-rank limit, so four ranks use NCCL.

On our side we'll run the one-GPU comparator for the padded heads (zig build dsv41-prefill-heads, then tf-dsv41-prefill-heads AOT_DIR OUT_DIR) and check that world 2 stays byte-identical to main.

• Synthema, Ash's AI Assistant.

[flashosophy]

Thank you for this. We ran it on four GB10s (NCCL, world 4, a fresh port, followers first) with the env from the
description. Two small local additions so our 2.9 bpw pack loads: its router bias and DSpark confidence head are
stored fp16, so we widen them exactly to fp32 (a separate commit on top of 9e2f85b; nothing else changed). We also
set TF_DSV41_ENGRAM_DIR, without which the Engram rows are synthetic.

Step 1. With the description's env, aot-needs fails Unsupported on all four ranks. The default prefill
segment is 2,048 rows. With 1,024-row segments it passes, and ranks 0-1 and 2-3 differ as you expected: two
_gemm variants each (K/N 640 vs 512), 97 variants a rank, a union of 99. With a 32K context and 8 slots, 118 a rank
and a union of 120. One union set on every rank, no misses.

Step 2. With the default 2,048 rows, rank 3 refuses at boot: prefill: a 2048-row segment is not on the prefill path at any position (block_prefill: Unsupported). PF_4K=1 (with GM_V2=1 and PREFILL=replay) fails the same
way, because the base 2,048-row segment is still planned. 1,536 rows also fail; 1,024 rows plan and serve. We think
the limit is the MoE scratch bound in block_prefill_moe.zig, not the attention padding. Without PF_4K a segment
needs n·topk·D <= 1024·(topk+1)·max(8·I, D). With D = 5120, topk = 6 and I = 640 or 512 at world 4, that allows
about 1,170 rows. At world 2 (I = 1152) 2,048 fits. This predicts exactly the sizes we saw pass and fail.

At 1,024 rows, world 4 boots in 335-345 s (weights from the pack, no prepared cache) with no panics on any rank.
Ranks 0-1 hold 55.1 GiB of weights and ranks 2-3 hold 44.6 GiB.

Step 3, as written, on one slot:

prompt	temperature	streams	tok/s	alone equal	serial equal
code	0	1	97.7	3/3	1/1
code	0	4	91.7	12/12	1/1
code	1.0	1	47.8	3/3	4/4
code	1.0	4	45.7	12/12	4/4

The chat prompt can't run on our pack. The server asks for a chat template, and the V4 checkpoints ship an
encoding/ folder rather than a Jinja template, so every chat request was refused.

Other numbers, from completions:

• A temp-0 reply of 88 tokens decodes at 34.6-39.5 tok/s. It is token-identical across two boots with different
  context and slot settings.
• A 21,358-token prompt prefills in 17.4 s (1,226 tok/s); 2.5K-6.6K-token prompts run at 1,116-1,212 tok/s.

Not done yet: a concurrency run with several live slots (it needs TF_DSV41_SLOT_DRAFTS=1 with drafts on; we'll
send that next), and no world-1 or world-2 run, because this pack doesn't fit one or two GB10s.

[flashosophy]

Follow-up with the concurrency run we promised. Same four GB10s, same build (9e2f85b plus the fp16→fp32 widen
commit), 1,024-row prefill segments, TF_DSV41_DRAFTS=1, TF_DSV41_SLOT_DRAFTS=1, 8 slots, 32K context,
TF_DSV41_ENGRAM_DIR set. Boot was 335 s again, with no panics.
Eight streams, all busy: /v1/completions, a 5,262-token prompt, max_tokens 1024, temp 0, ignore_eos. Two
identical 120 s passes:

pass	aggregate tok/s	per-stream median tok/s	max/min	TTFT median / max	ok / attempted
cold	90.5	11.9-18.6	1.56	19.8 s / 35.3 s	16 / 16
steady	91.3	12.3-18.1	1.47	19.8 s / 35.2 s	16 / 16
Every request finished with all 1,024 tokens, and there were no errors. Before each sweep we checked that one
request returns more than 32 tokens and that 8 parallel requests all return 200. Both checks passed.
An earlier run without SLOT_DRAFTS=1 had only one of the eight streams producing tokens (about 6 tok/s in
total), so that flag is what makes multi-slot decoding work here.
For scale: one stream alone was 97.7 tok/s in step 3. Eight streams together come to about 91 tok/s, so
concurrency splits the decode rather than adding to it on this setup. Streams are within about 1.5× of each other.
Not done: the 4-busy cell (we ran out of time in the window), and still no world-1 or world-2 run.

[flashosophy]

The 2,048-row refusal at world 4 is the MoE scratch bound in block_prefill_moe.zig, not the attention side. Below PF_4K, x3gm's row block has to fit the decode-sized expert scratch. At world 4 the experts are 640 / 512 wide (8·I ≤ D), so that bound admits only about 1,170 rows, where world 2 admits 2,048. That's why 1,024 boots and 1,536 / 2,048 don't.

#1205's branch fixes it: neko-legends:pr/dsv41-prefill-tp4 @ 11a24e45, rebased on main c9fab624. Past TP=2 the block is the whole segment, as it already is under PF_4K. TP=2 is unchanged. The MoE part is three lines:

const whole_block = e.o.pf4k or e.o.world > 2;
if (!whole_block and (P * D > zb or P * I > rows * (topk + 1) * D)) return error.Unsupported;

I tried it on top of this PR's head:

• this PR plus DeepSeek-V4.1 at four ranks: no prefill segment plans (BMQ 32 against 16 heads a rank; x3gm's block against the decode scratch) #1205's test only (prefill_tp4_test.zig, which plans 1,024 / 2,048-row segments on ranks 0..3 at 640 / 512 widths): test-dsv41-forward 122/126, with 3 failing at block_prefill_moe.zig:103
• the same plus the MoE hunk: 125/126 (1 skipped, the same skip as this PR alone, 121/122)

The branch also changes block_prefill_attn.zig (it clamps BMQ to the rank's heads). That overlaps your padded-heads approach, and the two conflict textually. If you'd rather keep padding, the MoE hunk and the test stack cleanly without it. Either way it's offered as a companion, so this PR can run at its own 2,048-row default. Happy to rebase it onto this PR, or keep it separate, whichever is easier for you.

[ashhart]

Landed on main (e8faf00) for the experimental DeepSeek-V4.1 engine: four-rank prefill through the existing 32-head core, with whole-block expert and vocabulary shares and 16-to-32-head padding, together with #1205's view-driven expert-scratch bound. World 1 and 2 host goldens match main byte for byte, and the DeepSeek host suites pass. Four-GB10 device runs are still owed: we have three boxes. The release CUDA engine still refuses DeepSeek-V4.1 (#827). Closing as landed.

- Synthema, Ash's AI Assistant.