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

The limitation

At tensor-parallel world 4, DeepSeek-V4.1 plans no prefill segment of any size. Every rank logs:

prefill: a 1024-row segment is not on the prefill path at any position (block_prefill: Unsupported)
prefill: a 2048-row segment is not on the prefill path at any position (block_prefill: Unsupported)
prefill: a 4096-row segment is not on the prefill path at any position (block_prefill: Unsupported)

(forward_prefill.zig:242). Today world 4 is refused earlier, at admission (#1202). This is what comes next once admission passes: boot stops before serving.

Current main (958efa3b) and v1.1.2 (4ff78441) have the same two checks at the same lines. git rev-list --count v1.1.2..958efa3b on both files is 0, and nothing under zig/src/families/deepseek_v41 changed between 1.1.1 (d1018b31) and v1.1.2.

Why world 4 hits it

Two checks in the prefill emitter assume a TP=2 shape.

1. _fused's head tile must divide a rank's query heads. block_prefill.zig:33 sets bmq = 32, and block_prefill_attn.zig:221 refuses otherwise:

if (@mod(H, bmq) != 0 or @mod(H, groups) != 0) return error.Unsupported;

A rank holds 64 / 4 = 16 query heads at world 4, so this refuses every segment (1,024 / 2,048 / 4,096 rows) before MoE is reached. The grid and the BMQ constexpr at :313 and :334 use the same bmq.

2. Below PF_4K, x3gm's row block must fit the decode-sized expert scratch. block_prefill_moe.zig:101-103:

const zb = rows * (topk + 1) * @max(4 * 2 * I, D); // the scratch z's floats (block_moe)
// past the 1,024-row scratch rows halve; under TF_DSV41_PF_4K the block is the whole segment
if (!e.o.pf4k and (P * D > zb or P * I > rows * (topk + 1) * D)) return error.Unsupported;

With rows = expert_rows = 1,024, topk = 6, D = 5,120 and P = n × 6:

Here 640 / 512 is the whole-block split from #1202, and 576 is the even split. With the head tile fixed alone, a 1,024-row segment fits (31,457,280 ≤ 36,700,160), but a 2,048-row segment, the forward's default prefill_rows, is still refused. PF_4K doesn't help either: it demands 4,096-row chunks and TF_DSV41_PF_4K_MEMORY_OK=1, and the forward plans its own segments with pf4k off.

A fix

A fix is on neko-legends/TensorFold:pr/dsv41-prefill-tp4 (ce84b44e), one commit on 958efa3b. It changes two files (+11 / −5), and adds a test file:

• block_prefill_attn.zig: attnBmq(H) = @min(bmq, H) at the check, the grid and BMQ. The kernel is unchanged and a 16-head rank launches it at BMQ = 16, a separate Triton variant. Two-rank launches (H = 32) are unchanged.
• block_prefill_moe.zig: whole_block = e.o.pf4k or e.o.world > 2. Past world 2 the block is the whole segment, as under PF_4K, and the buffer plan sizes s.ex.z / s.ex.xg from the views that follow. That is about 100 MB more a rank at 2,048 rows. World 2 keeps the bound and its launches exactly.
• prefill_tp4_test.zig, imported from state_test.zig so test-dsv41-forward runs it, with four cases:
  • on ranks 0..3, at 640 and 512 expert columns, 1,024 / 2,048 / 17 / 16-row segments are planned at positions 0 and 300,000. Each 2,048-row segment has one x3gm block (P = 12,288), the plan's s.ex.z / s.ex.xg hold it, and CED replay and a two-span emitMulti plan too;
  • staged replay segments get ring 4,096 with pf4k off and 8,192 with pf4k on;
  • world 2 still refuses a 4,096-row segment without PF_4K and still plans 2,048 rows;
  • image prompts at world 4: the plan's image options cover a segment's roles at any split, and runs are planned at the text options.

These tests call the planner directly with world-4-shaped configs that today's checks accept: expert_width 2,560 / 2,048 (640 / 512 a rank) and vocab 129,536 (253 blocks a rank). The shipped 2,304 / 129,280 are refused by block_admission.zig:5 and plan.zig:270 (#1202). Without a vocab that splits evenly, the head emit in block_prefill_mhc.zig:199 (vocab / world) panics on @divExact. That's the same limit, which #1202's split covers.

Fail-before (Zig 0.17.0, zig build test-dsv41-forward):

The world-2 case passes on all three trees.

Gates, 958efa3b → ce84b44e:

These are host tests only. No four-rank token, fidelity or speed run has been made on this branch, because world 4 can't boot on main until #1202 lands. We'll run the four-rank cells on top of #1202's fix and post them here.

Question

Is this split acceptable for the gate, or would you prefer another shape? Two alternatives we considered:

• keep the decode-sized scratch and halve the block past it, as Python's x3gm does: two 1,024-row passes fit at world 4 (31,457,280 ≤ 36,700,160). That needs no extra memory, but it means more x3gm launches a segment, and the emitter would need the halving it currently refuses instead;
• let a PF_4K workspace serve 1,024 / 2,048-row segments, instead of requiring an explicit 4,096 chunk.

The BMQ = 16 variant of _fused is the same kernel at a different constexpr. If the kernel contract needs its own records for it, we'll add them in the pull request with the backend's SUMMARY line.

Related: #1202 (prefill sits behind admission at world 4).