어텐션: MHA·MQA·GQA (Attention)

어텐션: MHA·MQA·GQA (Multi-Head, Multi-Query, and Group-Query Attention)

자기회귀(autoregressive) 모델을 만들다 보면 어텐션이 성능의 핵심 자리를 차지해요. 이 문서에서는 TensorRT-LLM의 PyTorch 백엔드가 멀티 헤드 어텐션(MHA), 멀티 쿼리 어텐션(MQA), 그룹 쿼리 어텐션(GQA)을 어떻게 구현하는지 설명합니다. 어떤 백엔드를 골라야 하고, 새 백엔드는 어떻게 만들면 되는지도 함께 살펴봐요.

출처: 공식 문서 - Multi-Head, Multi-Query, and Group-Query Attention

어텐션 백엔드 (Attention Backends)

현재 사용 가능한 어텐션 백엔드는 세 가지예요. vanilla 백엔드, TRT-LLM 백엔드, Flashinfer 백엔드입니다. 원하는 백엔드는 PyTorchConfig.attn_backend로 지정해요. 예를 들어 Flashinfer를 쓰려면 LLM(attn_backend="flashinfer")처럼 LLM 생성자에 넘기면 됩니다.

VanillaAttention(vanilla 백엔드)은 in-flight 배칭과 선형 KV 캐시 지원을 위해 설계된 참조 구현이에요. 유용한 기준선 역할을 하지만 최적화가 제한적이라 운영 환경에는 권장되지 않아요. 반면 FlashInferAttention(Flashinfer 백엔드)은 성능 최적화가 되어 있고 in-flight 배칭과 paged KV 캐시를 모두 지원해요. 추가 기능은 다음과 같습니다.

  1. FP8 양자화: 입력과 KV 캐시를 FP8 형식으로 양자화해 메모리 사용을 크게 줄이고 연산 처리량을 높여요.
  2. RoPE 퓨전: rotary position embedding(RoPE)을 어텐션 계산에 직접 통합해 효율을 높이고 오버헤드를 줄여요.

TrtllmAttention(TRT-LLM 백엔드)은 기본 백엔드로, Flashinfer의 모든 기능을 지원하면서 성능이 더 최적화되어 운영 환경에서 권장돼요. 추가로 이런 고급 기능을 제공합니다.

  1. Fused QKV 입력: Q, K, V를 각각 받는 대신 하나의 QKV 텐서를 받아들여 더 효율적이에요.
  2. FP8 출력: 어텐션 결과를 FP8 형식으로 출력하며, 어텐션 계산 과정에 양자화를 퓨전해요.

새 어텐션 백엔드 구현하기 (Implement a New Attention Backend)

다른 어텐션 라이브러리를 통합하려면 새 어텐션 백엔드를 구현할 수 있어요. 어텐션 백엔드는 AttentionBackend 클래스와 AttentionMetadata 클래스로 이뤄집니다. PyTorch 백엔드에서 어텐션 백엔드가 관여하는 단계는 세 가지예요.

  1. 모델 구성: 모델의 __init__에서 레이어마다 AttentionBackend.__init__을 호출해 어텐션 백엔드를 만듭니다.
  2. 메타데이터 준비: 모델의 각 forward 단계 전에,
    1. 메타데이터가 초기화되지 않았다면 AttentionMetadata.__init__으로 어텐션 메타데이터를 만듭니다.
    2. CUDA 그래프를 쓴다면 AttentionMetadata.create_cuda_graph_metadata를 호출해 메타데이터를 CUDA 그래프 메타데이터로 변환해요. 이는 모든 텐서를 미리 할당하고 CUDA 그래프 캡처에 사용할 수 있게 합니다. CUDA 그래프를 쓸 때는 최초 워밍업 실행 후에 AttentionMetadata 안의 텐서를 다시 할당하면 안 됩니다.
    3. 입력과 KV 캐시의 파라미터를 준비하려면 AttentionMetadata.prepare를 호출해 기존 메타데이터와 KV 캐시 매니저에서 변환해요.
  3. 단일 스텝 forward: 각 어텐션 레이어의 forward 패스에서 AttentionBackend.forward를 호출해 어텐션 연산을 수행합니다. AttentionMetadata는 forward 인자로 제공돼요.

AttentionMetadata 구현

AttentionMetadata 클래스는 배치된 입력과 KV 캐시의 메타데이터를 어텐션 백엔드에 저장해요. 미리 정의된 필드를 포함하며, AttentionMetadata.__init__에서 새 어텐션 메타데이터를 위한 추가 필드를 초기화할 수 있어요. 예를 들어 Flashinfer 메타데이터는 여기서 decode_wrapper를 초기화합니다. AttentionMetadata.prepare에서 런타임이 미리 정의된 모든 필드를 채우고, 여러분은 이 필드들을 바탕으로 커스텀 필드를 채우면 돼요. Flashinfer 메타데이터는 여기서 context_lens와 num_generations를 조합해 qo_indptr을 채웁니다.

AttentionBackend 구현

AttentionBackend는 어텐션 연산을 백엔드 구현에 위임해요. __init__과 forward가 받는 인자들은 문서에서 확인할 수 있고, 예를 들어 Flashinfer 백엔드는 append_paged_kv_cache를 호출한 뒤 wrapper.run을 호출해 어텐션 연산을 수행합니다.

TrtllmAttention 백엔드의 기능 (The Features of the TrtllmAttention Backend)

기본 백엔드인 TrtllmAttention의 기능 몇 가지를 소개할게요.

패킹된 텐서 (Packed Tensors)

TrtllmAttention 백엔드에서 어텐션 연산자는 패킹된(즉 패딩되지 않은) QKV 입력을 지원해요. 순진한 QKV 레이아웃은 max_sequence_length보다 짧은 시퀀스를 최대 길이로 패딩하는데, 이는 과도한 메모리 사용과 패딩 토큰에 대한 불필요한 연산을 일으킬 수 있어요. 이 문제를 해결하기 위해 TensorRT-LLM은 패딩이 없는 모드를 지원하며, 서로 다른 토큰을 패킹하고 사용자가 서로 다른 시퀀스 길이를 담은 1D 텐서를 제공합니다.

컨텍스트와 생성 단계 (Context and Generation Phases)

TrtllmAttention 백엔드는 컨텍스트와 생성 두 단계에 대한 서로 다른 구현을 하나의 커스텀 torch op에 담아요.

컨텍스트 단계 (Context Phase)

최적화 없는 컨텍스트 단계 구현은 중간 Q*K^T 텐서를 메모리에 저장한 뒤 softmax 연산자를 호출하는 GPU 커널 시퀀스에 해당해요. 가장 느린 방법이고 메모리 사용량도 시퀀스 길이의 제곱에 비례해 커집니다. TrtllmAttention 백엔드는 대신 단일 커널로 MHA/MQA 블록을 수행해요. 짧은 시퀀스에는 vanilla 구현을, 긴 시퀀스에는 FlashAttention 계열 알고리즘을 사용합니다. 현재는 RoPE 같은 전처리와 KV 캐시 채우기를 위한 추가 커널이 트리거되는데, 향후 릴리스에서 이 커널 수를 줄여 성능을 개선할 계획이에요.

FP8 컨텍스트 FMHA

FP8 양자화가 활성화되면 FP8 Context FMHA를 켜 어텐션을 더 가속할 수 있어요. FP8 Paged Context FMHA도 fp8 양자화 워크플로에서 지원되며, 어텐션 연산자에 use_paged_context_fmha=True를 지정해야 합니다. 이 기능은 Ada, Hopper 이상에서만 지원된다는 점을 알아 두세요.

생성 단계 (Generation Phase)

생성 단계는 TensorRT-LLM에서 masked multi-head attention이라 부르는 단일 커널로 구현돼요. 이 커널은 Q, K, V 요소에 on-the-fly로 전처리를 적용합니다. QKV 바이어스를 더하고, RoPE를 적용하며, 역양자화와 양자화를 수행하죠. 미래 릴리스에서는 IA3 지원 같은 기능이 추가될 예정입니다.

masked MHA 커널에는 GPU 점유율이 낮은 경우 작업을 여러 CUDA 스레드 블록에 분산하는 특별한 버전이 있어요. multi-block이라 부르는 이 모드는 항상 활성화됩니다. NVIDIA는 배치 크기와 모델의 헤드 수가 상대적으로 작은 시나리오에서 이 모드를 테스트해 보길 권장해요. 여기서 '작다'의 정의는 GPU 모델에 따라 달라 정량화하기 어렵지만, 현재는 batch_size*num_heads가 GPU의 멀티프로세서 수보다 작을 때 테스트를 권장합니다.

multi-block 모드가 켜져 있어도 어텐션 연산자가 즉시 multi-block 버전의 GPU 커널을 트리거하는 건 아니에요. multi-block 버전이 헤드당 단일 CUDA 스레드 블록을 쓰는 "vanilla" 구현보다 효율적이 되려면 최소한의 토큰 수(입력+생성)가 필요하며, 내부 휴리스틱으로 이를 제어합니다. 또 masked MHA 커널은 시퀀스 길이에 비례하는 공유 메모리를 쓰므로, multi-block 모드가 꺼졌을 때 GPU 공유 메모리가 부족한 경우가 있을 수 있어요. 그럴 때는 multi-block 모드가 강제로 켜지고 로그에 경고가 출력됩니다.

XQA 최적화

XQA 최적화는 생성 단계에서 MQA/GQA를 위한 또 다른 최적화예요. 현재 LLAMA2 70B 모델 같은 제한된 모델 구성만 지원합니다. XQA 최적화의 지원 범위는 다음과 같아요.

  • FP16 / BF16 계산 데이터 타입
  • FP16 / BF16 / FP8 / INT8 KV 캐시 데이터 타입
  • Paged KV 캐시 (블록당 8 / 16 / 32 / 64 / 128 토큰)

기본적으로 이 기능은 켜져 있어요. XQA 커널과 masked MHA 커널 중 뭘 쓸지도 휴리스틱 알고리즘으로 정합니다. 모델 구성이 지원될 때 XQA 커널을 항상 쓰고 싶다면 TRTLLM_FORCE_XQA=1로 강제할 수 있고, 지원되는 구성은 cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/decoderXQARunner.h의 DecoderXQARunner 클래스 shouldUse 함수에서 확인할 수 있어요.

In-flight Batching

TensorRT-LLM은 더 높은 서빙 처리량을 위해 요청의 in-flight 배칭(연속 배칭, iteration-level 배칭이라고도 함)을 지원해요. 컨텍스트 단계의 시퀀스가 생성 단계의 시퀀스와 함께 처리될 수 있습니다. 이 기법의 목적은 요청을 더 잘 인터리빙해 지연을 줄이고 GPU를 더 잘 활용하는 거예요. 효율성 때문에 in-flight 배칭 지원은 입력 텐서가 패킹(패딩 없음)되어 있어야 합니다.

현재 구현에서는 컨텍스트 단계를 거치는 시퀀스가 입력 텐서에서 생성 단계 시퀀스보다 앞에 있어야 해요. 예를 들어 S0, S1, S2 시퀀스에서 S0와 S2가 컨텍스트 단계(그리고 S1이 생성 단계)라면, S0와 S2의 토큰이 S1의 토큰보다 입력 텐서에서 앞에 나와야 합니다.

청크 컨텍스트 (Chunked Context)

긴 시퀀스를 다룰 때 KV 캐시 메모리 사용량을 줄이는 데 도움이 되는 기능이에요. 사이클릭(cyclic) KV 캐시 기능은 현재 빔 탐색(beam searching)과 함께 동작하지 않는다는 점을 알아 두세요. 컨텍스트 KV 캐시가 빔 간에 공유되기 때문이에요.

StreamingLLM

StreamingLLM 기능은 윈도우 어텐션을 사용해 긴 텍스트에서 효율적이고 안정적인 LLM을 수행해요. 즉 KV 캐시에 N개 토큰만 저장하면 됩니다. TensorRT-LLM의 사이클릭 KV 캐시 기능과 유사하게 attention_window_size 파라미터가 N을 정해요. 차이점은 싱크 토큰(sink token)이라 부르는 처음 S개 토큰이 항상 어텐션 윈도우에 유지된다는 것이고, S는 sink_token_length 파라미터로 정합니다. 공식 구현에서 컨텍스트 단계의 self-attention은 dense하게, 모든 토큰을 사용해 계산하고 N개 토큰만 KV 캐시에 저장해요. 또 상대 위치 임베딩도 달라지는데, 상대 거리를 정하고 토큰에 위치 정보를 더할 때 원본 텍스트가 아니라 캐시 안의 위치를 사용합니다. sink_token_length는 이 기능을 활성화하는 데도 쓰여요.

어텐션 연산자는 빔 탐색을 지원해요. 컨텍스트 단계에서는 입력 시퀀스당 단일 빔이 계산됩니다. 생성 단계에서는 MHA/MQA/GQA 커널이 추가 텐서를 사용해 각 빔의 올바른 경로를 재구성하는데, 이 텐서를 cache_indirection이라 부르고 형태는 [batch_size,beam_width,max_seqlen]이에요. 시퀀스 si, 빔 bi, 토큰 ti에 대해 cache_indirection[si][bi][ti]는 0에서 beam_width-1 사이의 정수로, KV 캐시에서 K와 V 요소를 읽을 빔의 어느 경로를 쓸지 가리킵니다. 이 텐서는 샘플링 단계에서 채워집니다.

입력 QKV 텐서 (Input QKV tensor)

입력 QKV 텐서는 히든 스테이트 투영 후에 Q, K, V 텐서(마지막 차원을 따라 연결)를 패킹해요. 3D 텐서이며, RoPE와 필요 시 INT8 또는 FP8 양자화는 GPT 어텐션 연산자가 수행합니다. 패킹 모드에서 형태는 [num_tokens,3*hidden_dim]이고 num_tokens은 배치의 총 토큰 수예요. 컨텍스트 단계 시퀀스의 토큰 수는 입력 길이에 해당하고, 생성 단계 시퀀스는 시퀀스당 beam_width개의 토큰이 있습니다. 빔 폭은 시퀀스마다 다를 수 있어요. 토큰 수를 계산하는 의사 코드는 이렇게 됩니다.

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

ALiBi

어텐션 연산자는 Q*K^T 결과에 ALiBi를 적용할 수 있어요. 바이어스는 최적화된 커널에서 ALiBi slopes로부터 on-the-fly로 계산됩니다.

스케일링 인자 (Scaling factor)

MHA에서 Q*K^T 출력은 상수 값으로 스케일링되는데, 다음과 같이 계산해요.

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

크로스 어텐션 (Cross Attention)

GPT 스타일 디코더 전용 모델이 필요로 하는 self attention으로서의 MHA 외에도, 어텐션 연산자는 크로스 어텐션을 지원해요. 이 덕분에 어텐션 연산자를 일반적인 디코더 컴포넌트로 더 폭넓게 쓸 수 있습니다. 예를 들어 Encoder-Decoder 모델은 디코더 안에서 self attention과 cross attention 모듈을 모두 발행할 때 이를 사용해요.

더 알아보기 (Learn more)