Speculative Decoding in vLLM on AMD GPUs!

Accelerating Inference with Speculative Decoding on AMD Hardware: A Technical Deep Dive

The computational cost of autoregressive transformer inference remains the primary bottleneck for large-scale language model deployment. While throughput optimization techniques like Continuous Batching and PagedAttention have become standard, the fundamental latency constraint of sequential token generation persists. Speculative Decoding emerges as a transformative approach to this constraint, and its implementation within the vLLM framework on AMD Instinct accelerators presents unique architectural considerations regarding memory bandwidth, compute scheduling, and kernel optimization.

The Mechanism of Speculative Decoding

Speculative Decoding operates on the principle that many language modeling tasks possess low entropy during the initial stages of sequence generation. Instead of relying solely on a large, parameter-heavy "target" model to generate each token sequentially, a smaller, low-latency "draft" model is used to propose a sequence of candidate tokens. The target model then verifies these tokens in parallel.

The formal process can be defined as follows:

1. Draft Generation: The draft model, $M_{draft}$, generates a sequence of $k$ tokens ${t_1, t_2, \dots, t_k}$ autoregressively.
2. Parallel Verification: The target model, $M_{target}$, processes the input prompt concatenated with the proposed tokens ${t_1, t_2, \dots, t_k}$ in a single forward pass.
3. Acceptance: Based on the target model’s logits, we apply a rejection sampling scheme (typically based on the rejection criterion defined by Leviathan et al. or Chen et al.). If the draft token matches the target model's distribution within a specified probability threshold, it is accepted.
4. Correction: The first token rejected by the target model is replaced by the target model’s output, and the sequence continues.

AMD ROCm and Hardware Constraints

Deploying this architecture on AMD GPUs requires a deep understanding of the AMD ROCm (Radeon Open Compute) stack, specifically the interaction between the HipBLAS library and the underlying Compute Units (CUs) of CDNA architectures (e.g., MI250X, MI300X).

Unlike NVIDIA’s CUDA ecosystem where kernel fusion for speculative decoding is heavily optimized via tools like Triton or specialized CUTLASS kernels, AMD environments require explicit management of the memory hierarchy. The primary challenge on AMD hardware is the management of the KV cache when running two distinct models concurrently within the same memory space.

Memory Partitioning and KV Cache Allocation

In vLLM, PagedAttention provides a sophisticated mechanism for managing KV cache memory. When implementing Speculative Decoding, we must partition the GPU VRAM between the draft and target model caches.

# Conceptual memory allocation for Speculative Decoding in vLLM/AMD
def allocate_kv_cache(model_config, device_memory_pool):
    target_memory = model_config.target_model_size * target_kv_overhead
    draft_memory = model_config.draft_model_size * draft_kv_overhead

    # AMD specific: ensure block alignment for MI300X memory controllers
    alignment = 256 * 1024  # 256KB alignment for optimal coalescing

    target_pool = device_memory_pool.allocate(target_memory, alignment)
    draft_pool = device_memory_pool.allocate(draft_memory, alignment)

    return target_pool, draft_pool

On AMD GPUs, the L2 cache utilization is critical. Because Speculative Decoding involves two models, the cache contention increases. If the draft model is too large, it risks evicting the weights of the target model from the L2 cache, resulting in a significant latency penalty that negates the speed-up gained from speculation.

Kernel Fusion and Performance Optimization

The efficiency of Speculative Decoding on AMD hardware is dictated by the transition between the draft generation and the target verification phases. Traditional implementations often suffer from latency overhead due to GPU-CPU synchronization when deciding which tokens to accept.

To mitigate this, the vLLM implementation on AMD utilizes specialized fused kernels that perform the rejection sampling directly on the GPU. By avoiding a copy-back to the host, we minimize synchronization primitives (like hipStreamSynchronize).

Optimizing the Verification Step

The verification step involves a batch of tokens being processed as a single prompt. This is a "batch-level" attention operation. On CDNA 3 architectures, we can leverage the matrix core units (WMMA - Wave Matrix Multiply-Accumulate) to accelerate this verification.

// Optimized verification kernel fragment for AMD CDNA
__global__ void verify_tokens_kernel(
    const float* draft_logits,
    const float* target_logits,
    int* accepted_count,
    float threshold) {

    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    // Calculate rejection sampling criteria using ROCm vector intrinsics
    // Avoid branching to keep warps synchronized
    float p_draft = draft_logits[idx];
    float p_target = target_logits[idx];

    bool accept = (p_target >= p_draft) || (random_uniform() < (p_target / p_draft));
    // Atomic operations to update the accepted token chain
}

Challenges with AMD-Specific Tooling

While the ROCm port of vLLM has reached functional parity with CUDA, performance tuning for Speculative Decoding presents unique obstacles:

1. Compiler Optimization: The hipcc compiler optimization flags for different CDNA generations are less forgiving than nvcc. Aggressive inlining in custom kernels can lead to register pressure, causing spills to local memory, which drastically reduces throughput during the verification phase.
2. Collective Communication: In multi-GPU setups (e.g., using Infinity Fabric to bridge MI300X accelerators), the latency of the all-reduce operations during the verification phase can be prohibitive if the draft model is running on one GPU and the target on another. It is generally recommended to keep both models on the same socket, if memory capacity permits.
3. KV Cache Fragmentation: With two models managing distinct KV caches via PagedAttention, the memory manager must ensure that pages are not fragmented in a way that interferes with the hardware prefetchers.

Benchmarking and Latency Trade-offs

The effectiveness of Speculative Decoding is measured by the "Acceptance Rate." If the draft model is too weak (e.g., a 125M parameter model drafting for a 70B parameter model), the acceptance rate often drops below 30%, resulting in a net latency increase due to the compute overhead of the verification pass.

On AMD Instinct MI300X, we have observed that the sweet spot for Speculative Decoding involves draft models that are approximately 1/10th to 1/20th the size of the target model. This ratio allows the verification pass to complete within the time window required to generate a single token on the target model, maximizing the effective tokens per second (TPS).

Architectural Considerations for Future Scaling

As model sizes increase and we move toward mixture-of-experts (MoE) architectures, Speculative Decoding becomes more complex. The verification pass on an MoE model requires fetching the relevant expert weights into the L2 cache. On AMD hardware, the high bandwidth of HBM3 memory becomes the primary enabler for this technique.

Future optimizations in the vLLM/ROCm stack should focus on:

• Quantized Verification: Executing the verification pass in FP8 or INT8 format to utilize the high-throughput matrix units of the MI300 series.
• Speculative Prefetching: Using the draft model not just for token generation, but to prefetch the relevant MoE expert weights for the target model’s next expected forward pass.
• Asynchronous Rejection Sampling: Decoupling the verification kernel from the generation loop through stream-based asynchronous execution to hide the latency of the rejection logic.

Summary of Implementation Strategy

To successfully implement Speculative Decoding on AMD infrastructure, engineering teams must prioritize:

1. Memory Alignment: Ensure KV cache allocations are tuned to the 256KB-512KB alignment requirements of the MI-series memory controllers.
2. Kernel Fusing: Consolidate the rejection sampling logic into the attention kernels to avoid host-side synchronization.
3. Draft Model Selection: Rigorously validate the draft-to-target size ratio to ensure the acceptance rate compensates for the increased FLOPs in the verification pass.

The integration of Speculative Decoding within vLLM on AMD GPUs signifies a maturing of the ecosystem, moving beyond simple model compatibility toward performance-optimized production deployments. By leveraging the underlying hardware capabilities of CDNA architectures and optimizing memory-intensive kernels, inference throughput can be significantly scaled without compromising model precision or availability.

For organizations seeking to implement high-performance large language model inference architectures on AMD hardware, strategic guidance on infrastructure optimization and custom kernel development is essential. Visit https://www.mgatc.com for consulting services.

---

Originally published in Spanish at www.mgatc.com/blog/speculative-decoding-vllm-amd-gpus/