Optimized Speculative Sampling for GPU Hardware Accelerators (2024)

Abstract: In this work, we optimize speculative sampling for parallel hardware accelerators to improve sampling speed. We notice that substantial portions of the intermediate matrices necessary for speculative sampling can be computed concurrently. This allows us to distribute the workload across multiple GPU threads, enabling simultaneous operations on matrix segments within thread blocks. This results in profiling time improvements ranging from 6% to 13% relative to the baseline implementation, without compromising accuracy. To further accelerate speculative sampling, probability distributions parameterized by softmax are approximated by sigmoid. This approximation approach results in significantly greater relative improvements in profiling time, ranging from 37% to 94%, with a minor decline in accuracy. We conduct extensive experiments on both automatic speech recognition and summarization tasks to validate the effectiveness of our optimization methods.

## introduction large foundational speech and language models based on autoregressive transformer @xcite architectures have demonstrated remarkable proficiency across a variety of downstream tasks @xcite @xcite . these models frequently increase in size, consequently requiring more memory and computational resources. however, downstream applications, such as dialogue systems, have strict wallclock constraints and are often required to generate long sequences @xcite @xcite . due to the sequential token generation in autoregressive decoding, latency increases with both the length of the sequence and the size of the model, resulting in a significant barrier to widespread deployment. on many generalpurpose gpu hardware accelerator architectures, the increasing size of models leads to more read and write operations between high-bandwidth memory (hbm) and on-chip shared memory (sram) at each decoding step, necessitating more memory bandwidth to meet latency constraints @xcite @xcite . consequently, the speed of autoregressive decoding becomes primarily limited by the available memory bandwidth and not by the number of computations that need to be executed on the dedicated hardware @xcite . in many cases, however, tokens may be accurately generated by much smaller models that require fewer resources. motivated by this hypothesis, speculative sampling has been developed to accelerate autoregressive sampling @xcite @xcite . speculative sampling employs a small draft model to generate tokens, which are potential future outputs of a larger target model. these drafted tokens are then verified in parallel by the target model, and only tokens that meet the validation criteria are retained as final outputs to ensure generation accuracy. this approach has been shown to significantly reduce the frequency of time-consuming operations, thereby improving inference latency @xcite . in this paper, we focus on optimizing the validation part of speculative sampling to further increase the inference speed. inspired by recent advances in accelerating computations in the attention mechanism @xcite , we explore two faster methods for assessing drafted tokens by leveraging the parallelism capabilities of modern gpus. we identify that a significant portion of the intermediate matrices required for the sampling process can be computed independently. exploiting this observation, we distribute the workload across multiple gpu threads and simultaneously compute portions of the intermediate output matri-ces within thread blocks. this optimization method is faster than the non-optimized baseline implementation and exact with regard to the decoding outputs, i.e., it generates the same outputs as the non-optimized method. to further accelerate speculative sampling, we propose using sigmoid as an element-wise approximation to softmax @xcite , which is used to parameterize distributions of target and draft models. since sigmoid is applied to logits in elementwise fashion, it can be computed in parallel and fused with other sampling-related computations. this enables significant acceleration of the overall process, but results in a small accuracy decline due to the non-exact nature of the method. we evaluate our two optimized algorithms on automatic speech recognition (asr) and summarization tasks, covering draft model sizes between 166m and 2b parameters and target model sizes between 244m and 13b parameters. the exact optimization method reduces profiling time between 6% and 13% relative to the baseline implementation without compromising accuracy. moreover, the non-exact optimization method improves profiling time by 37% to 94%, albeit with a small reduction in accuracy. we summarize our main contributions as follows: foot_0 • implementation of an exact and consistently faster variant of speculative sampling optimized for gpu hardware accelerators. ## related work techniques such as quantization @xcite @xcite , pruning @xcite @xcite and knowledge distillation @xcite @xcite have proven effective in reducing inference latency with minimal performance impact. however, these approaches often require architectural changes or custom training procedures. efforts specifically targeting the reduction of memory bandwidth bottlenecks during decoding include methods like multi-query attention, which aims to optimize memory usage per attention layer @xcite , or flashattention @xcite , which aims to reduce the number of read/write operations between hbm and sram on gpus. @xcite , achieve improvements in large-scale inference latency by partitioning models and workload across multiple accelerators combined with various lowlevel optimizations to improve communication efficiency between devices. speculative sampling approaches can be broadly categorized based on how drafting and verification are conducted @xcite . drafting refers to the efficient prediction of multiple future tokens with a draft model and verification refers to the methods used to verify the token sequence with the target model. some works use specialized draft models @xcite . others employ an existing smaller model from the same series @xcite @xcite or leverage the target model directly for drafting, e.g., by skipping intermediate layers @xcite , using special look-ahead tokens @xcite , or additional modeling heads @xcite @xcite . verification approaches first supported greedy decoding @xcite @xcite and were subsequently extended to support other methods such as nucleus sampling @xcite . recently, methods to verify multiple draft sequences in parallel have also been explored @xcite @xcite . several studies have experimented with relu and sigmoid as alternatives to softmax in the attention mechanism @xcite @xcite @xcite @xcite . they either require training new models from scratch or maintain the computational overhead of gathering information along the full sequence length. other softmax-related optimizations are tailored to reduce the memory requirement during training by computing only smaller fractions of the full softmax output in the backward pass @xcite or leverage word frequencies in the training data to speed up computation @xcite . @xcite approximate softmax by computing only a fraction of the full input with singular value decomposition. other approximation methods are specifically designed for custom hardware such as field-programmable gate arrays @xcite and application specific integrated circuits @xcite . first, given the prefix (x 1 , . . . , x i+c-1 ), γ ∈ n + draft tokens are sequentially sampled with m draft q : we denote the numerator of eq. 3 by a(x) and the denominator by b, as these terms are treated separately in subsequent sections. the underlying concept of speculative sampling is similar to rejection sampling @xcite . intuitively, a new token for the target model m target p is generated by first sampling from a smaller draft model m draft q , which shares the same support (v) as m target p . the token sampled from m draft q is then evaluated in parallel with m target p , and its acceptance is determined based on the probability ratio defined in eq. 1. if the token is rejected, a new one is drawn using the modified distribution in eq. 2. gpu memory and execution model. we briefly describe the memory components and parts of the execution model of gpu hardware accelerators relevant to this work. gpu memory has a hierarchical layout, consisting of various types of memory that differ in size and read/write bandwidth @xcite . recent gpus (e.g., nvidia's a100 series) typically feature several gigabytes of highbandwidth memory (hbm) and only a few hundred kilobytes of on-chip shared memory (sram) per streaming multiprocessor (sm) (nvidia corporation, 2020). while hbm provides substantial capacity, its memory bandwidth is lower compared to sram. the execution model of gpu hardware accelerators involves a large number of threads executing operations known as kernels @xcite . these threads are organized into thread blocks and assigned to sms. each sm partitions its assigned thread blocks into warps of 32 threads, which are then queued for execution on available hardware resources. each kernel typically follows a pattern: loading inputs from hbm into registers and sram, performing computations, and writing the outputs back to hbm. ## conclusions we introduced two optimization methods to accelerate speculative sampling for autoregressive models on hardware accelerators. by computing significant portions of intermediate matrices across multiple gpu threads within thread blocks, our exact optimization method led to improved sampling speed without compromising accuracy. additionally, we employed an approximation technique using element-wise sigmoid instead of softmax, to enable parallel computation of probabilities. this approximation further accelerated the decoding process but resulted in a small degradation of sampling quality.