Describe the GPTBigCode and Yuan shared-QK layouts (#8575)
Three of the four layouts in `AUTOTP_UNSUPPORTED_PARAMETER_PATTERNS` now
describe themselves, so universal checkpoint conversion stops refusing
them. Follows #8519, which made AutoTP layers emit affine maps for the
layouts a generic rule already covered.
`codegentype` is not included — its reshape-and-interleave needs more
than a segment list or a block selection, and it is the next piece of
work.
### What was actually missing
Both refusals came from the same limit: a parameter carries one
partition dimension, which cannot say that one block of it is sharded
while the next is held whole, nor that a rank takes blocks which are not
adjacent.
**GPTBigCode** splits the query rows across ranks and hands every rank
the whole key/value block. `segmented_map` describes a parameter as an
ordered list of segments, each split or replicated, so the two end up as
pieces that differ in which ranks hold them rather than in kind. That is
the case the per-piece location set was introduced for in #8385.
**Yuan shared-QK** gives a rank the value heads that pair with its query
heads, plus the matching second-half heads. Those land in two runs
rather than one span, so `block_gather_map` takes a selection of
equal-sized blocks and merges consecutive ids into single pieces.
### Why the tests are worth reading
A layer builds its map analytically, from head counts and the tp degree.
The test harness derives one independently by running the **real**
partition function on a marker tensor and reading back where each slice
came from. The tests require the two to agree, so what ships is checked
against what the partition code actually does rather than against a
reading of it.
I confirmed they fail when the map is wrong: corrupting `segmented_map`
so the replicated segment claims a single owner produces two failures
rather than passing quietly.
`shared_qk_value_head_ids` is extracted from `shard_value_with_share_qk`
and used by both the partition and the map, so the description cannot
drift from the layout it describes.
### What stays unsupported, deliberately
- **A layout with no head counts.** Describing Yuan needs
`num_kv_heads`; without it the layer publishes the refusal rather than
guessing, and there is a test for that.
- **The Yuan o_proj bias.** Its partition pre-divides by the world size
and replicates, which needs a scaled piece — and that code path does not
currently produce a usable bias, so there is no working layout to
describe.
### Validation
Validated on CPU/gloo (DS_ACCELERATOR=cpu LOCAL_SIZE=4): 90 passed
across the affine suite, the producer and coverage tests, the resume
matrix and tests/unit/runtime/tensor_parallel/.
Related: #8252, #8230.
cc @delock
---------
Signed-off-by: Achyuthan Sivasankar <achyuthan.sivasankar@gmail.com>
Co-authored-by: Ma, Guokai <guokai.ma@gmail.com>
