ImplementAttentionMetadata#

TheAttentionMetadata class stores metadata from the batched input and KV cache for the attention backend.It contains the following predefined fields:

DuringAttentionMetadata.__init__, you can initialize additional fields for the new attention metadata.For example, the Flashinfer metadata initializesdecode_wrapper here.DuringAttentionMetadata.prepare, the runtime will fill all predefined fields, and you can fill your customized fields according to these predefined fields.For example, the Flashinfer metadata fillsqo_indptr by combiningcontext_lens andnum_generations here.

ImplementAttentionBackend#

TheAttentionBackend delegates the attention operation to the backend implementation.

Its__init__ accepts the following arguments:

Itsforward accepts the following arguments:

For example, the Flashinfer backend callsappend_paged_kv_cache and thenwrapper.run to perform the attention operation here.

Packed Tensors#

In theTrtllmAttention backend, the attention operator supports the packed (i.e. non padded) QKV inputs.A naive layout for the QKV inputs is padding the sequencesthat are shorter than themax_sequence_length to the maximumlength. It may result in excessive memory consumption as well as unneededcomputations on padding tokens (in the various matrix multiplications thatsurround the MHA block).To overcome that problem, TensorRT LLM supports a mode without padding wherethe different tokens are packed together and the user provides the operatorwith a 1D tensor containing the lengths of the different sequences.

Context and Generation Phases#

TheTrtllmAttention backend encapsulates different implementations for bothcontext and generation phases into a single custom torch op.

In-flight Batching#

TensorRT LLM supports in-flight batching of requests (also known as continuousbatching or iteration-level batching) for higher serving throughput. With this feature,sequences in the context phase can be processed together with sequences inthe generation phase. The purpose of that technique is to better interleaverequests to reduce latency as well as make better use of the GPUs.For efficiency reasons (1), the support for inflight batchingrequires theinput tensors to be packed (no padding).

In the current implementation, the sequences that are going through thecontext phase must be before the sequences in the generation phase in the inputtensor. For example, for sequencesS0,S1 andS2, ifS0 andS2 are incontext phase (andS1 in generation), tokens fromS0 andS2 must appearbefore the tokens ofS1 in the input tensor.

(1) Padding sequences in the generation phase, that contain a single token, tothe length of the maximum input sequence is inefficient use of resources.

Chunked Context#

In the original state, the common behavior was to process all context tokens atonce. This feature splits the context into several chunks. In this way, thecontext chunks can be batched with more tokens during the generation phase,which is expected to increase the total throughput. Chunking contexts also removesconstraints on input length. Except for the last one, the size of the context chunk needsto be an integer multiple of the kv-cache block size.

To enable this feature, the FMHA paged kv-cache also needs to be enabled.

KV Cache#

In the generation phase, a common optimization is to provide the MHA kernelwith a cache containing the values of the past K and V elements that havealready been computed. That cache is known as the KV cache. TensorRT LLM usesthat technique to accelerate its generation phase. In TensorRT LLM, there isone KV cache per Transformer layer, which means that there are as many KVcaches as layers in a model. The current version of TensorRT LLM supports twodifferent types of KV caches:contiguous andpaged KV caches.

[max_batch_size*max_beam_width,2,num_heads,max_seqlen,hidden_dim_per_head].

Sliding Window Attention, Cyclic (Rolling Buffer) KV Cache#

TensorRT LLM has a feature calledCyclicKVCache, which treats the kv cacheas a circular buffer. This means that it only stores the kv cache for the last Ntokens, where N is determined by theattention_window_size parameter inTrtllmAttention.forward. When the cache is full, new tokens’ kv cache willoverwrite the “least recently used” caches.

In the context phase, if the input length surpasses theattention_window_size,SlidingWindowAttention will be activated. This serves the same function asthe sliding window size.

This feature helps to reduce the memory footprint of the kv cache whendealing with very long sequences.

_Note that the cyclic kv cache feature doesn’t work with beam searching currently asthe context kv cache are shared across beams.

StreamingLLM#

The StreamingLLM feature uses a window attention to perform efficient and stable LLMon long texts, which means that onlyN tokens need to be stored in the KV cache.Similar to the cyclic KV cache feature in TensorRT LLM,attention_window_sizeparameter is used to determineN. Different from the cyclic KV cache feature,the firstS tokens, called sink tokens, are always kept in the attention window,whereS is determined bysink_token_length parameter.But in context phase, the self-attentions are dense in the official implementation ofStreamingLLM. It uses all of the tokens for computation and only savesN tokensto the KV cache.

In addition, the relative position embedding is also changed in StreamingLLM.When determining the relative distance and adding positional information to tokens,StreamingLLM use the positions within the cache rather than those in the original text.

sink_token_length is also used to enable this feature.

Beam-Search#

The attention operator supports beam-search. In the context phase, a singlebeam is computed per input sequence. In the generation phase, the MHA/MQA/GQAkernel uses an additional tensor to reconstruct the correct path for each beam.That tensor is called thecache_indirection. Its shape is[batch_size,beam_width,max_seqlen].

For a sequencesi, a beambi and a tokenti, the elementcache_indirection[si][bi][ti] is an integer between0 andbeam_width-1that indicates which path in the beam to read the K and V elements from in theKV cache. This tensor is populated in the sampling stage.

Input QKV tensor#

The input QKV tensor packs the Q, K and V tensors (concatenated along the lastdimension) after the projection of the hidden states. It is a 3D tensor. RoPEand quantization to INT8 or FP8 (when needed) are performed by the GPTattention operator.

In packed mode, its shape is[num_tokens,3*hidden_dim] wherenum_tokens is the total number of tokens in the batch. For the sequences incontext phase, the number of tokens of a sequence corresponds to its inputlength (even if the beam width is greater than1 for beam search). For thesequences in generation phase, there arebeam_width tokens per sequence. Thebeam width can be different for each sequence.

The following pseudo code explains how the number of tokens is computed:

num_tokens=0# Add the length of each sequence in context phase.forseqincontext_phase:num_tokens+=seq.length# Add the width of the beam for each sequence in generation phase.forseqingeneration_phase:num_tokens+=seq.beam_width

Rotary Positional Embedding (RoPE)#

The attention operator can perform the computation of the RotaryPositional Embedding (RoPE). When that operation is enabled,rotary_embedding_dim is set to a value greater than 0, it is fused with otheroperations. The GPT operator supports GPT-NeoX and GPT-J forms of RoPE bysettingposition_embedding_type toPositionEmbeddingType.rope_gpt_neoxorPositionEmbeddingType.rope_gptj.

ALiBi#

The attention operator can apply ALiBi to the result of theQ*K^Tproduct. The bias is computed on-the-fly from the ALiBi slopes in the optimizedkernel.

Scaling factor(s)#

In MHA, the output of theQ*K^T product is scaled by a constant value thatis computed as:

norm_factor=1.f/(q_scaling*sqrt(head_size)).

Cross Attention#

On top of the MHA as self attention needed by GPT-style decoder-only models, the attention operator also supports cross attention.

This enables the attention operator to be more broadly used as a generic decoder component. For example, the Encoder-Decoder model uses it to issue both the self attention and cross attention modules in its Decoder.