Describe the CodeGen fused-QKV layout as a block selection

[Achyuthan-S]

CodeGen was the last layout in AUTOTP_UNSUPPORTED_PARAMETER_PATTERNS whose geometry had no description. Its weight now publishes an affine map, so conversion stops refusing it.

Follows #8519 (layers emit the map), #8575 (GPTBigCode and Yuan describe themselves) and #8622 (restore reads it).

What the layout actually is

_codegen_type_transpose views the weight as codegen_mp_num blocks, cuts each into q/k/v, subdivides every third across the ranks, then reassembles the pieces q, k, v at a time. The result is one fixed permutation of the whole weight — rank-independent — from which each rank takes a contiguous slice.

So the outer split is a plain even row split, and the interleaving lives entirely inside each rank's region, where the rows are a permutation of equal-sized blocks. That is why one partition_dim cannot say it and a block selection can.

block_gather_map already expresses this, so there is no new IR primitive here. I said when #8575 landed that codegen "needs more than a segment list or a block selection"; that was wrong, and this is the correction. The work is deriving the block order, not extending the language.

How the description is kept honest

codegen_block_layout derives the order independently, because the partition reaches it through reshapes and never names a block. That departs from #8575, where shared_qk_value_head_ids is shared between the partition and its description. Sharing here would mean refactoring _codegen_type_transpose itself, which is riskier than the description it would protect, so the drift guard is instead a test requiring the derived map to equal one probed from the real partition function. That checks behaviour rather than structure. @delock — if you would rather have the refactor so there is one definition, say so and I will do it that way.

The helper refuses every case the partition does not produce: no head count, a head count it cannot divide (the same condition _codegen_type_transpose asserts), fewer than two ranks, and rows that do not divide into equal blocks. An earlier version of this change published a map for AutoTPMeta() with no head count — a config whose own assert cannot be evaluated — and test_collect_records_fused_layouts_without_a_sub_param_split caught it.

The tests discriminate

corruption	result
q/k/v loop order swapped with the block loop	6 fail
block size missing its world_size factor	6 fail
hook renamed so it never fires	1 fail — the wiring test, which is its only job

The piece-count test is the one that matters beyond correctness. A compressor that missed the permutation would emit one piece per row, still cover the tensor and still round-trip, so only the count reveals it. Across rows 24 → 192 at TP 2 the count holds at 11.

Validation

The derived map equals the probed map for 24x8 TP2, 48x16 TP2, 96x32 TP2, 96x8 TP4 and 192x8 TP8. In every case extract is bit-exact against the real partition, rebuild round-trips, the map covers F, and both validate() and validate_coverage() pass. 11 pieces per rank at TP 2, 12 at TP 4 and TP 8 — matching the count in §5 of affine_ir_spec.md.

126 passed on this rebase across the affine suite, the resume matrix, the coverage test, tests/unit/runtime/tensor_parallel/ and the AutoEP affine suite, on CPU/gloo (DS_ACCELERATOR=cpu LOCAL_SIZE=4) — now on top of #8624's stricter validate(), which every codegen piece passes with its destination checked.

With #8623 merged, codegen maps also compose through plan_transfer exactly: TP 2↔4, 2↔8 and 4↔8, with source shards taken from the real partition function and every target shard compared element for element. #8623's own suite passes here too (119).

Earlier, at base 581ac41bd, I ran tests/unit/checkpoint/ and tests/unit/runtime/tensor_parallel/ on this branch and on a pristine worktree of that commit: 244 passed here against 236 on master — the 8 this PR adds — with 74 failed and 106 errors on both, and the failing set identical by name rather than merely in count, 180 names either way with none on one side only. All of them are environmental: 126 FusedAdam JIT (no ninja), 54 fp16-on-CPU, 4 cpu_adam JIT. Master has moved since, so that pairing is against 581ac41bd rather than today's tip; I can redo it if the numbers matter for review.

What this does not close

Yuan's o_proj bias is still refused, so the unsupported set is not empty yet. That one is not a description gap: its partition pre-divides by the world size and that path returns an empty bias rather than a usable one, so there is no layout to describe until the partition is fixed. Separate change.

Related: #8252, #8230.

cc @delock

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

[Copilot]

Copilot review overview

🟢 Approval recommended

The implementation matches the partition ordering and includes comprehensive validation across layouts and TP sizes.

Review effort: Balanced
Findings: None

[Achyuthan-S]

@delock gentle ping when you have a moment. This is still clean against master — including against #8680, which touches layers.py and applies alongside this without conflict — and the CI workflows are waiting on approval to run.