Running dense, open-weight large language models on commercially accessible workstations or single-GPU cloud setups is increasingly desirable due to cost and privacy constraints, particularly on tasks like coding and reasoning. Starting at around 24B parameters, however, they exceed the available GPU VRAM on these setups, forcing reliance on offloading methods that limit throughput, while leaving substantial CPU compute underutilised. Heterogeneous CPU–GPU speculative decoding is the dominant lossless-acceleration framework in this regime: a draft model proposes tokens for a verifier to validate in parallel.
In work accepted at NeurIPS 2026, researchers Viren Luke Radhakrishnan from AMD, alongside IISc project associates Dhruva Kashyap and Pranav K Nayak, under the guidance of Chiranjib Bhattacharyya at the Machine Learning Lab, Department of Computer Science and Automation, IISc, and Prakash Raghavendra at AMD, identified a limitation in existing methods, which they refer to as the heterogeneity gap.
The ‘heterogeneity gap’: In a heterogeneous (CPU + GPU), speculative decoding-based inference system, at any instant during a generation step, only one device performs useful work, and any draft weights held separately on the GPU sit idle during verification. The researchers asked: can we design a system that closes the heterogeneity gap, thereby realising improvements in throughput on performance-sensitive workloads like coding and reasoning?
Challenges in designing such a system
To close the heterogeneity gap in the speculative decoding context, they identified three conditions that a method must satisfy:
(i) Co-execution: Execute productive CPU and GPU compute simultaneously within each generation step, which necessarily requires a split along the intermediate dimension of each layer rather than by layer (ruling out cross-layer pipelining).
(ii) Subnetwork draft: Derive the draft as a subnetwork of the verifier itself rather than a separate set of weights; otherwise, the draft weights occupy GPU memory that sits idle at verify time. One might attempt this by placing some transformer layers on the GPU and the rest on the CPU, but layer-level partitioning forces the two devices to execute sequentially; channel-level splitting within a layer is harder, since the GPU shard must double as a high-quality draft of the full model.
(iii) Maximal VRAM utilisation: At any GPU memory budget, use a near-full fraction (around 90%) of the available VRAM, reserving the remainder for the KV cache and activations.
No prior method satisfies all conditions, leaving the corresponding throughput improvements substantially unrealised.
Speculative tensor parallelism for closing the heterogeneity gap
(a) The solution observes that a saliency-based pruner such as ModHiFi (Kashyap et al., NeurIPS 2025) can be used to re-order the channels of each feed-forward network (FFN) layer in an LLM with no change in output, so that the top-k channels form a prefix closely approximating the full layer.
ModHiFi/Wanda reorder channels by saliency so the top-k form a contiguous prefix (dashed arrow: saliency-ordered axis; orange triangle: truncation point).
(b) This re-ordering allows for a heterogeneous, self-speculative partitioning strategy that closes the heterogeneity gap. Building on it, the solution introduces Speculative Tensor Parallelism (STeP), a training-free method that extracts a subnetwork of the verifier as a GPU-resident draft, places the remainder on the CPU, and verifies via concurrent CPU and GPU computation, provably preserving the verifier’s sampling distribution.
Overview of STeP for a single FFN layer. The top-k saliency-ordered channels (active ratio r) reside on the GPU and serve as both the draft and the GPU portion of the verifier; the remaining channels reside on the CPU. Their element-wise sum recovers the exact verifier output.
Results
Across dense models from 24B to 123B, GPU memory budgets from 32 to 192 GB, and five benchmarks, STeP outperforms the heterogeneous speculative methods SpecExec and SubSpec by up to 1.6x under greedy decoding and 2.2x at sampling temperature 0.6, surpasses either tensor parallelism or speculative decoding alone, maximally utilises the available VRAM, and closes the heterogeneity gap. The official implementation of STeP is available as part of AMD PACE.
Acknowledgements
The authors gratefully acknowledge the financial and technical support of AMD Research and Advanced Development (RAD), as well as Ramaswamy Govindarajan (CSA, IISc) and Arun Ramachandran (AMD) for their insight and assistance in this work.
Paper: https://neurips.cc/virtual/2026/poster/155402
Project page: https://github.com/amd/AMD-PACE
I joined the BTech programme in Mathematics and Computing at the Indian Institute of Science in 2022. My faculty advisor, Professor Viraj Kumar, has guided my academic journey, and I had the opportunity to work under his mentorship during the summer of 2024.
I joined BTech at IISc in 2024.
I joined the BTech programme in 2024. My faculty advisor is Professor Viraj Kumar. I am interested in computer science and would like to do research in computer science in the future.
Anmol Gill pursues the MTech (Artificial Intelligence) programme at IISc, in the Robert Bosch Centre for Cyber-Physical Systems (
I recently completed my BTech in August 2024 from the Vellore Institute of Technology, Chennai, with a major in Electronics and Communication Engineering. From May 2024, I am a research intern, supported by the Kotak-IISc AI/ML Centre, in the Human-Interactive Robotics (HiRo) Lab at the Robert Bosch Centre for Cyber-Physical Systems, IISc. I am working under the guidance of Professor Ravi Prakash. My essential area of interest lies in robot learning, which lies under the general area of AI and robotics.
