Decode Token Parallelism: Memory-Efficient Distributed Attention Parallelism for Scalable LLM Decoding
Abstract
Serving large language models at high concurrency with long contexts requires KV-cache capacity and attention bandwidth that grow with the decode pool. Existing parallelization strategies trade one against the other: data parallelism replicates model weights, limiting KV capacity; tensor parallelism stops scaling attention at the model's KV-head count and pushes dense collectives across nodes; and decode context parallelism (DCP) recovers capacity by sharding the sequence but adds cross-node aggregation of partial attention results. We observe that during decode, dense layers and attention scale differently: dense layers amortize weight reads across the batch and run efficiently within a single NVLink domain, whereas attention cost grows with every active request's KV cache. We present Decode Token Parallelism (TKNP), a decode-pool architecture that decouples weight placement from attention ownership. TKNP distributes a single model copy across nodes in balanced layer blocks. The active node holding a block executes its projections and MLP for the full batch using intra-node tensor parallelism, while every node computes attention for its own requests against locally resident KV cache. This extends request-local attention across nodes without replicating weights: only current-token activations cross the network, dense all-reduces remain within NVLink, and attention requires no cross-node merging of partial results. Implemented in vLLM and evaluated on Llama 3.1 405B and Nemotron-4 340B on up to 64 H100 GPUs, TKNP improves completion throughput for long-context, high-concurrency workloads by up to 1.7$\times$ over DCP and by up to 2.06$\times$ over data parallelism at matched GPU count.
The following analysis is generated by AI and may contain errors. It is separate from the authors’ abstract and is not peer review.
Research summary
The paper introduces Decode Token Parallelism (TKNP), a memory-efficient distributed attention parallelism strategy that decouples weight placement from attention ownership, achieving up to 2.06x throughput improvement over data parallelism for large-scale LLM decoding on multi-node H100 clusters.
Methodology
The paper proposes Decode Token Parallelism (TKNP), a novel parallelization strategy for LLM decoding that decouples weight placement from attention ownership. Unlike standard Tensor Parallelism (TP) which shards attention heads and requires cross-node communication for attention, or Data Parallelism (DP) which replicates weights, TKNP distributes model layers across nodes. Each node holds a subset of layers and the KV cache for its assigned tokens. During decoding, the node holding the current layer block executes the dense computations (MLP/Projections) for the entire batch using intra-node TP, while all nodes compute attention locally for their respective tokens. This avoids cross-node aggregation of partial attention results, limiting network traffic to current-token activations. The approach is theoretically sound, leveraging the observation that dense layers amortize well over batches within a node, while attention cost scales with KV cache size.
Results & evaluation
The authors evaluate TKNP on Llama 3.1 405B and Nemotron-4 340B using up to 64 H100 GPUs. The results show significant throughput improvements for long-context, high-concurrency workloads: up to 1.7x over Decode Context Parallelism (DCP) and 2.06x over Data Parallelism. The choice of models (400B+ scale) and hardware (H100) is highly relevant to current industrial needs. The comparison against DCP is particularly strong as DCP is a recent, sophisticated baseline. The metrics focus on completion throughput, which is a key serving metric.
Reproducibility
The paper states the implementation is in vLLM, a widely used open-source inference engine. This significantly boosts reproducibility compared to proprietary stacks. However, specific hyperparameters for the "balanced layer blocks" and the exact configuration of intra-node TP within TKNP would need to be detailed in the code or appendix for full replication. The use of standard models and hardware makes it feasible for other groups to reproduce.
Limitations
The primary limitation is the assumption that the network latency for passing current-token activations between nodes is negligible compared to the compute time of the dense layers. This holds for high-bandwidth interconnects (like NVLink or InfiniBand) but may degrade on slower networks. Additionally, the load balancing of "balanced layer blocks" might be complex if layer compute costs vary significantly (e.g., MoE layers vs. dense layers, though the paper focuses on dense LLMs). The paper does not extensively discuss the impact of prefill phase latency, focusing primarily on decode.
Broader impact
This work has high practical impact for LLM serving infrastructure. As models grow larger and contexts longer, efficient memory and bandwidth management is critical. TKNP offers a scalable path for serving 100B+ parameter models on multi-node clusters without excessive weight replication or complex attention sharding. It is likely to be adopted by inference providers looking to optimize cost-per-token for long-context applications. The paper introduces Decode Token Parallelism (TKNP), a memory-efficient distributed attention parallelism strategy that decouples weight placement from attention ownership, achieving up to 2.06x throughput improvement over data parallelism for large-scale LLM decoding on multi-node H100 clusters.
Author of this paper? Send a correction.