Attention head partitions: keep each head’s dependencies intact

For ordinary multi-head attention, each head compares its queries with the permitted keys, normalizes scores over those keys, and combines the corresponding values. Different heads can be evaluated on different tensor-parallel ranks while retaining their full head-dimension calculations. The attention scaling uses the key/query head dimension, not the number of devices. Position encodings and masks must use the same global token positions as the reference. A rank owning only some heads does not justify ignoring keys from other token positions.

Projecting hidden states into Q, K, and V can be organized with output-feature partitions aligned to head boundaries. Each rank then holds its head outputs as a feature shard. A compatible output projection partitions its input rows along those concatenated head features and sums partial contributions to the final hidden output. A gather of all head outputs is one possible implementation boundary, but it is not universally necessary before that row-parallel projection.

With H_q query heads and H_kv key/value heads, a regular GQA layout groups H_q/H_kv query heads per KV head when that ratio is integral. A simple even partition over P ranks can assign H_q/P query heads and H_kv/P KV heads per rank only when the counts and grouping support it. If P exceeds the KV-head count, an implementation may replicate KV heads or use another supported arrangement. Dividing to obtain half a head is not a valid execution plan.

Replicated KV changes memory and gradient bookkeeping. Its storage is not reduced by the same factor as evenly sharded query features, and contributions to shared parameters must be combined across the actual replica groups. Multi-query attention, with one KV head, makes this difference especially visible. Frameworks impose configuration constraints that can be stricter than the mathematical possibility. Check the selected version’s supported head mapping rather than claiming that every GQA configuration supports an arbitrary tensor-parallel size.

Test each head with distinguishable inputs so an accidental permutation cannot hide behind identical values. Preserve the mapping from global query heads to their KV group, then reconstruct or project the output in the documented order. Compare forward attention, output projection, and gradients to shared KV parameters. Include a causal mask and a nontrivial sequence length; a single-token test cannot reveal a missing-key or future-token visibility bug. Disable dropout initially before testing the intended random-stream mapping.

The matrix lab illustrates an inspectable map without implementing attention. Reset, set a=1 and b=0, predict the identity transform and determinant 1, then Step. This exposes what an unchanged coordinate map means; head reordering in the worked layout must likewise preserve its declared correspondence. The triangular 2D family cannot represent arbitrary head permutations, attention scores, GQA sharing, or distributed projection. Use the explicit head-count and global-ID table to establish the actual partition.