# Affinity-Aware Sharding for Delayed Tensor Parallelism arXiv: https://arxiv.org/abs/2609.13846 · HTML: https://mlsanity.com/paper/2609.13846/ ## Key facts - No citing papers in the map yet (data through 14 Oct 2026). - Builds on 28 of its 38 references that are in the map. - First posted on arXiv 12 Sep 2026. ## Builds on (most cited first, 28 of 28) - [Distilling the Knowledge in a Neural Network](https://mlsanity.com/paper/1503.02531/) (2015) - [Pointer Sentinel Mixture Models](https://mlsanity.com/paper/1609.07843/) (2016) - [Efficient Streaming Language Models with Attention Sinks](https://mlsanity.com/paper/2309.17453/) (2023) - [Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism](https://mlsanity.com/paper/1909.08053/) (2019) - [GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints](https://mlsanity.com/paper/2305.13245/) (2023) - [Transformer Feed-Forward Layers Are Key-Value Memories](https://mlsanity.com/paper/2012.14913/) (2020) - [The FineWeb Datasets: Decanting the Web for the Finest Text Data at Scale](https://mlsanity.com/paper/2406.17557/) (2024) - [Git Re-Basin: Merging Models modulo Permutation Symmetries](https://mlsanity.com/paper/2209.04836/) (2022) - [Massive Activations in Large Language Models](https://mlsanity.com/paper/2402.17762/) (2024) - [The Role of Permutation Invariance in Linear Mode Connectivity of Neural Networks](https://mlsanity.com/paper/2110.06296/) (2021) - [MoEfication: Transformer Feed-forward Layers are Mixtures of Experts](https://mlsanity.com/paper/2110.01786/) (2021) - [Compact Language Models via Pruning and Knowledge Distillation](https://mlsanity.com/paper/2407.14679/) (2024) - [FLUX: Fast Software-based Communication Overlap On GPUs Through Kernel Fusion](https://mlsanity.com/paper/2406.06858/) (2024) - [The Super Weight in Large Language Models](https://mlsanity.com/paper/2411.07191/) (2024) - [Emergent Modularity in Pre-trained Transformers](https://mlsanity.com/paper/2305.18390/) (2023) - [The Spike, the Sparse and the Sink: Anatomy of Massive Activations and Attention Sinks](https://mlsanity.com/paper/2603.05498/) (2026) - [Interpreting Context Look-ups in Transformers: Investigating Attention-MLP Interactions](https://mlsanity.com/paper/2402.15055/) (2024) - [Flash Communication: Reducing Tensor Parallelization Bottleneck for Fast Large Language Model Inference](https://mlsanity.com/paper/2412.04964/) (2024) - [Domino: Eliminating Communication in LLM Training via Generic Tensor Slicing and Overlapping](https://mlsanity.com/paper/2409.15241/) (2024) - [Ladder-residual: parallelism-aware architecture for accelerating large model inference with communication overlapping](https://mlsanity.com/paper/2501.06589/) (2025) - [Communication Compression for Tensor Parallel LLM Inference](https://mlsanity.com/paper/2411.09510/) (2024) - [SPD: Sync-Point Drop for Efficient Tensor Parallelism of Large Language Models](https://mlsanity.com/paper/2502.20727/) (2025) - [Kraken: Inherently Parallel Transformers For Efficient Multi-Device Inference](https://mlsanity.com/paper/2408.07802/) (2024) - [Tensor-Parallelism with Partially Synchronized Activations](https://mlsanity.com/paper/2506.19645/) (2025) - [Functional Equivalence in Attention: A Comprehensive Study with Applications to Linear Mode Connectivity](https://mlsanity.com/paper/2606.17830/) (2026) - [Parallel Track Transformers: Enabling Fast GPU Inference with Reduced Synchronization](https://mlsanity.com/paper/2602.07306/) (2026) - [SiFAR: Synchronization-Free All-Reduce for Low-Latency LLM Inference](https://mlsanity.com/paper/2607.08973/) (2026) - [Sparse Inter-Layer Dependencies of Transformer FFN Neurons](https://mlsanity.com/paper/2607.11990/) (2026) ## Abstract (arXiv) Delayed Tensor Parallelism (DTP) removes the blocking all-reduce of tensor-parallel Transformer inference. Every device adds its own partial output to its residual stream (and broadcasts it) immediately, but only gathers (receives) the other devices' partials δ modules later. A TP to DTP change therefore amounts to a real architecture change, and dense Transformer models need to be retrained or distilled after adaptation. We show that DTP breaks the permutation symmetry of neurons inside FFNs and of KV heads inside attention modules, and that this symmetry breakage makes the sharding itself a modelling decision. We show that maximising the affinity between the KV heads and the FFN neurons co-located on a device, by permuting the dense model before sharding, speeds up the distillation or retraining process. The affinity is measured with a first-order approximation of the damage that losing a head's contribution does to each neuron's output, and the co-located affinity is maximised with a coordinate-ascent optimiser that alternates an exact balanced assignment of neurons with an exhaustive search over the KV head partitions. The whole procedure takes under two minutes on one GPU for Qwen3-0.6B and Danube3-500M. On these models, at δ=1, the affinity-optimised layouts reach any distillation target in about half to two thirds of the steps needed by the naive contiguous layouts, over the whole 10k-step range we tested, and every optimised seed beats every contiguous seed and all but one of the sixteen random layouts. We also show that the co-located affinity score at initialisation predicts the KL to the base model after training, across seventeen layouts ranging from anti-optimised to optimised (Pearson -0.81 and -0.89).