PagedAttention

이 문서는 vLLM의 PagedAttention 커널을 높은 수준에서 설명하는 내용이에요. 아래에서 중요하게 짚는 블록(block) 과 GPU의 스레드 블록(thread block) 은 서로 다른 개념이라는 걸 먼저 기억해 두세요. 이 문서에서는 vLLM의 페이지 방식 캐시 블록을 "블록", GPU 쪽을 "스레드 블록"이라고 부를게요.

주의: 이 문서는 vLLM의 원래 논문을 바탕으로 한 역사적 문서입니다. 지금 vLLM이 쓰는 코드를 그대로 설명하지는 않아요.

현재 vLLM은 자체적으로 구현한 멀티헤드 쿼리 어텐션 커널(csrc/attention/attention_kernels.cu)을 사용합니다. 이 커널은 vLLM의 페이지 방식 KV 캐시와 호환되도록 설계됐어요. 여기서 키(key)와 값(value) 캐시는 각각 별도의 블록에 저장됩니다.

이 커널이 높은 성능을 내는 핵심은, 특히 스레드가 글로벌 메모리에서 공유 메모리로 데이터를 읽을 때 쓰는 특별한 메모리 레이아웃과 접근 방식에 있어요. 이 문서의 목적은 커널 구현을 단계별로 높은 수준에서 설명해서, vLLM의 멀티헤드 쿼리 어텐션 커널을 배우고 싶은 분들이 실제 코드를 더 쉽게 따라갈 수 있도록 돕는 거예요. 이 문서를 읽고 나면 실제 구현을 따라가기가 훨씬 수월해질 거예요.

이 문서가 모든 세부 사항을 다루지는 않아요. 예를 들어 데이터의 올바른 인덱스를 계산하는 방법이나 내적(dot product) 구현은 생략돼요. 하지만 높은 수준의 로직 흐름을 익히고 나면 실제 코드를 읽으면서 세부 내용을 이해하기 훨씬 쉬워진답니다.

Inputs

커널 함수는 현재 스레드가 할당된 작업을 수행하는 데 필요한 여러 인자를 받아요. 그중 가장 중요한 세 인자는 q, k_cache, v_cache 입력 포인터입니다. 각각 글로벌 메모리에 있는 쿼리, 키, 값 데이터를 가리키죠. 출력 포인터 out은 결과를 써야 하는 글로벌 메모리를 가리켜요. 이 네 포인터는 사실 다차원 배열을 가리키는데, 각 스레드는 자기에게 할당된 일부 데이터에만 접근합니다. 여기서는 이해를 돕기 위해 나머지 런타임 인자들은 모두 생략할게요.

template<typename scalar_t, int HEAD_SIZE, int BLOCK_SIZE, int NUM_THREADS, int PARTITION_SIZE = 0>

__device__ void paged_attention_kernel(

    ... // Other side args.

    const scalar_t* __restrict__ out,       // [num_seqs, num_heads, max_num_partitions, head_size]

    const scalar_t* __restrict__ q,         // [num_seqs, num_heads, head_size]

    const scalar_t* __restrict__ k_cache,   // [num_blocks, num_kv_heads, head_size/x, block_size, x]

    const scalar_t* __restrict__ v_cache,   // [num_blocks, num_kv_heads, head_size, block_size]

    ... // Other side args.

)

함수 시그니처 위쪽에는 컴파일 타임에 정해지는 템플릿 인자 목록도 있어요. scalar_t는 쿼리·키·값 데이터 요소들의 데이터 타입을 뜻합니다(예: FP16). HEAD_SIZE는 각 헤드에 들어 있는 요소 수를, BLOCK_SIZE는 각 블록의 토큰 수를 나타내요. NUM_THREADS는 각 스레드 블록 안의 스레드 수이고, PARTITION_SIZE는 텐서 병렬 GPU 수를 의미해요(간단히 하기 위해 여기서는 0으로 두고 텐서 병렬을 비활성화했다고 가정할게요).

이 인자들을 받으면 여러 준비 작업을 수행해야 해요. 여기에는 현재 헤드 인덱스, 블록 인덱스 등 필요한 변수들을 계산하는 일이 포함됩니다. 다만 지금은 이 준비 과정은 잠시 무시하고 실제 계산으로 바로 들어가 볼게요. 전체 흐름을 이해하면 이 부분들도 훨씬 쉽게 이해할 수 있어요.

Concepts

계산 흐름에 들어가기 전에, 뒤에서 필요할 몇 가지 개념을 먼저 짚고 갈게요. 혹시 뒤에서 용어가 헷갈리면 이 섹션으로 돌아와도 됩니다.

Sequence (시퀀스): 시퀀스는 하나의 클라이언트 요청을 나타내요. 예를 들어 q가 가리키는 데이터는 [num_seqs, num_heads, head_size] 모양을 갖는데, 이는 q가 총 num_seqs 개의 쿼리 시퀀스 데이터를 가리킨다는 뜻입니다. 이 커널은 싱글 쿼리 어텐션 커널이라 각 시퀀스에는 쿼리 토큰이 딱 하나만 있어요. 따라서 num_seqs는 배치에서 처리되는 총 토큰 수와 같아요.

Context (컨텍스트): 컨텍스트는 시퀀스에서 생성된 토큰들로 구성됩니다. 예를 들어 ["What", "is", "your"]가 컨텍스트 토큰이고, 입력 쿼리 토큰은 "name"이에요. 모델은 아마 "?"라는 토큰을 생성할 거예요.

Vec (벡터): vec은 함께 가져와서 계산되는 요소들의 묶음입니다. 쿼리와 키 데이터의 경우 스레드 그룹이 한 번에 16바이트 데이터를 가져와 계산하도록 vec 크기(VEC_SIZE)가 정해져요. 값 데이터의 경우 각 스레드가 한 번에 16바이트를 처리하도록 vec 크기(V_VEC_SIZE)가 정해지죠. 예를 들어 scalar_t가 FP16(2바이트)이고 THREAD_GROUP_SIZE가 2라면 VEC_SIZE는 4, V_VEC_SIZE는 8이 됩니다.

Thread group (스레드 그룹): 스레드 그룹은 한 번에 하나의 쿼리 토큰과 하나의 키 토큰을 가져와 계산하는 작은 스레드 묶음(THREAD_GROUP_SIZE)입니다. 각 스레드는 토큰 데이터의 일부만 처리해요. 하나의 스레드 그룹이 처리하는 요소의 총 수를 x라고 부릅니다. 예를 들어 스레드 그룹에 2개의 스레드가 있고 헤드 크기가 8이라면, 스레드 0은 인덱스 0, 2, 4, 6에 있는 쿼리·키 요소를, 스레드 1은 인덱스 1, 3, 5, 7에 있는 요소를 처리해요.

Block (블록): vLLM의 키·값 캐시 데이터는 블록으로 나뉩니다. 각 블록은 한 헤드에서 고정된 수(BLOCK_SIZE)의 토큰 데이터를 저장해요. 각 블록은 전체 컨텍스트 토큰의 일부만 담을 수도 있습니다. 예를 들어 블록 크기가 16이고 헤드 크기가 128이라면, 한 헤드에 대해 블록 하나는 16 × 128 = 2048개의 요소를 저장할 수 있어요.

Warp (워프): 워프는 스트리밍 멀티프로세서(SM)에서 동시에 실행되는 32개 스레드(WARP_SIZE)의 묶음입니다. 이 커널에서 각 워프는 한 번에 하나의 쿼리 토큰과 한 블록 전체의 키 토큰 사이 계산을 처리해요(여러 블록은 여러 반복에서 처리할 수 있습니다). 예를 들어 한 컨텍스트에 워프 4개와 블록 6개가 있다면, 워프 0은 0번과 4번 블록, 워프 1은 1번과 5번 블록, 워프 2는 2번 블록, 워프 3은 3번 블록을 담당하게 됩니다.

Thread block (스레드 블록): 스레드 블록은 같은 공유 메모리에 접근할 수 있는 스레드(NUM_THREADS)들의 그룹입니다. 각 스레드 블록은 여러 워프(NUM_WARPS)를 담고 있어요. 이 커널에서 각 스레드 블록은 하나의 쿼리 토큰과 전체 컨텍스트의 키 토큰 사이 계산을 처리합니다.

Grid (그리드): 그리드는 스레드 블록들의 모음으로, 그 모음의 모양을 정의해요. 이 커널에서 그리드의 모양은 (num_heads, num_seqs, max_num_partitions)입니다. 따라서 각 스레드 블록은 하나의 헤드, 하나의 시퀀스, 하나의 파티션에 대한 계산만 담당해요.

Query

이번엔 쿼리 데이터가 메모리에 어떻게 저장되고 각 스레드가 어떻게 가져오는지 볼게요. 앞서 말했듯 각 스레드 그룹은 하나의 쿼리 토큰 데이터를 가져오고, 각 스레드는 쿼리 토큰 데이터의 일부만 처리합니다. 워프 안의 모든 스레드 그룹은 같은 쿼리 토큰 데이터를 가져오지만, 그걸 서로 다른 키 토큰 데이터와 곱해요.

const scalar_t* q_ptr = q + seq_idx * q_stride + head_idx * HEAD_SIZE;

각 스레드는 글로벌 메모리에서 자기에게 할당된 쿼리 토큰 데이터를 가리키는 q_ptr을 정의합니다. 예를 들어 VEC_SIZE가 4이고 HEAD_SIZE가 128이라면, q_ptr은 총 128개 요소를 가리키고, 이 요소들은 128 / 4 = 32개의 vec으로 나뉘어요.

__shared__ Q_vec q_vecs[THREAD_GROUP_SIZE][NUM_VECS_PER_THREAD];

이제 q_ptr이 가리키는 글로벌 메모리 데이터를 공유 메모리에 q_vecs로 읽어 들여야 해요. 여기서 중요한 점은 각 vec이 서로 다른 행에 할당된다는 것입니다. 예를 들어 THREAD_GROUP_SIZE가 2라면 스레드 0은 0번 행의 vec, 스레드 1은 1번 행의 vec을 처리해요. 이렇게 쿼리 데이터를 읽으면 스레드 0과 스레드 1처럼 이웃한 스레드들이 서로 인접한 메모리를 읽을 수 있어서, 메모리 코얼레싱(memory coalescing)을 달성해 성능을 높일 수 있습니다.

Key

"Query" 섹션과 비슷하게, 이번엔 키의 메모리 레이아웃과 할당 방식을 살펴볼게요. 각 스레드 그룹은 커널 실행 한 번에 쿼리 토큰 하나만 처리하지만, 여러 번의 반복을 통해 여러 키 토큰을 처리할 수 있어요. 한편 각 워프는 여러 번의 반복에서 여러 블록의 키 토큰을 처리하면서, 커널 실행이 끝났을 때 전체 컨텍스트 토큰을 스레드 그룹 전체가 처리하도록 보장합니다. 여기서 "처리"란 쿼리 데이터와 키 데이터 사이의 내적을 수행하는 것을 뜻해요.

const scalar_t* k_ptr = k_cache + physical_block_number * kv_block_stride

                    + kv_head_idx * kv_head_stride

                    + physical_block_offset * x;

q_ptr과는 달리, 각 스레드의 k_ptr은 반복이 달라질 때마다 다른 키 토큰을 가리켜요. 위 코드에서 보듯 k_ptr은 k_cache에서 할당된 블록, 할당된 헤드, 할당된 토큰을 기준으로 키 토큰 데이터를 가리킵니다.

위 다이어그램은 키 데이터의 메모리 레이아웃을 보여줘요. BLOCK_SIZE가 16, HEAD_SIZE가 128, x가 8, THREAD_GROUP_SIZE가 2, 총 4개의 워프가 있다고 가정합니다. 각 사각형은 한 헤드에서 하나의 키 토큰에 대한 모든 요소를 나타내며, 이것은 하나의 스레드 그룹이 처리해요. 왼쪽 절반은 워프 0이 처리하는 키 토큰 데이터 16블록 전체를, 오른쪽 절반은 나머지 워프나 반복이 처리할 나머지 키 토큰 데이터를 보여줍니다. 각 사각형 안에는 2개의 스레드(하나의 스레드 그룹)가 나누어 처리할 총 32개의 vec(한 토큰의 128개 요소)이 있어요.

K_vec k_vecs[NUM_VECS_PER_THREAD]

이제 k_ptr에서 키 토큰 데이터를 읽어 레지스터 메모리에 k_vecs로 저장할 차례입니다. k_vecs에는 레지스터 메모리를 사용하는데, 그 이유는 이 값이 여러 스레드가 여러 번 접근하는 q_vecs와 달리 한 스레드가 딱 한 번만 접근하기 때문이에요. 각 k_vecs는 이후 계산을 위한 여러 벡터를 담고 있고, 각 vec은 안쪽 반복(inner iteration)마다 설정됩니다. vec의 할당 방식 덕분에 워프 안의 이웃한 스레드들이 함께 이웃한 메모리를 읽을 수 있어서, 이 역시 메모리 코얼레싱을 촉진해요. 예를 들어 스레드 0은 vec 0, 스레드 1은 vec 1을 읽습니다. 다음 안쪽 반복에서는 스레드 0이 vec 2, 스레드 1이 vec 3을 읽는 식이죠.

아직 전체 흐름이 조금 헷갈릴 수 있을 거예요. 걱정하지 마세요. 다음 "QK" 섹션을 계속 읽어보면 쿼리와 키 계산 흐름을 더 명확하고 높은 수준으로 보여드릴게요.

QK

아래 의사코드처럼, 전체 for 루프 블록이 시작되기 전에 한 토큰의 쿼리 데이터를 가져와 q_vecs에 저장해요. 그다음 바깥 for 루프에서 서로 다른 토큰을 가리키는 여러 k_ptr을 순회하며 안쪽 for 루프에서 k_vecs를 준비합니다. 마지막으로 q_vecs와 각 k_vecs 사이의 내적을 수행해요.

q_vecs = ...

for ... {

    k_ptr = ...

    for ... {

        k_vecs[i] = ...

    }

    ...

    float qk = scale * Qk_dot<scalar_t, THREAD_GROUP_SIZE>::dot(q_vecs[thread_group_offset], k_vecs);

}

앞서 말했듯 각 스레드는 쿼리와 키 토큰 데이터의 일부만 가져와요. 하지만 Qk_dot<>::dot 안에서는 스레드 그룹 간 리덕션(cross thread group reduction)이 일어납니다. 그래서 여기서 반환되는 qk는 쿼리·키 토큰 데이터 일부 사이의 내적 결과가 아니라, 실제로는 전체 쿼리·키 토큰 데이터 사이의 내적 결과예요.

예를 들어 HEAD_SIZE가 128이고 THREAD_GROUP_SIZE가 2라면, 각 스레드의 k_vecs는 총 64개 요소를 담고 있어요. 하지만 반환되는 qk는 실제로 128개의 쿼리 요소와 128개의 키 요소 사이의 내적 결과입니다. 내적과 리덕션의 세부 내용이 궁금하다면 Qk_dot<>::dot 구현을 참고하세요. 다만 이 문서에서는 간단히 하기 위해 다루지 않을게요.

Softmax

다음으로 모든 qk에 대한 정규화된 소프트맥스를 계산해야 해요. 각 𝑥는 하나의 qk를 나타냅니다. 이를 위해 모든 qk의 리덕션 값인 qk_max(𝑚(𝑥))와 exp_sum(ℓ(𝑥))을 구해야 해요. 이 리덕션은 전체 스레드 블록에 걸쳐 수행되어야 하며, 쿼리 토큰과 모든 컨텍스트 키 토큰 사이의 결과를 포함합니다.

𝑚(𝑥) := max𝑖 𝑥𝑖
𝑓(𝑥) := [𝑒𝑥1−𝑚(𝑥) … 𝑒𝑥𝐵−𝑚(𝑥)]
ℓ(𝑥) := ∑𝑖 𝑓(𝑥)𝑖
softmax(𝑥) := 𝑓(𝑥) / ℓ(𝑥)

qk_max and logits

qk 결과를 얻은 직후, 임시 logits 결과를 qk로 설정할 수 있어요(결국 logits에는 정규화된 소프트맥스 결과가 저장되어야 합니다). 또한 현재 스레드 그룹이 계산한 모든 qk에 대해 qk_max를 비교하고 수집할 수 있습니다.

if (thread_group_offset == 0) {

    const bool mask = token_idx >= context_len;

    logits[token_idx - start_token_idx] = mask ? 0.f : qk;

    qk_max = mask ? qk_max : fmaxf(qk_max, qk);

}

여기서 logits는 공유 메모리에 있으므로, 각 스레드 그룹은 자기에게 할당된 컨텍스트 토큰에 대한 필드를 설정해요. 전체적으로 logits의 크기는 컨텍스트 토큰의 개수여야 합니다.

for (int mask = WARP_SIZE / 2; mask >= THREAD_GROUP_SIZE; mask /= 2) {

    qk_max = fmaxf(qk_max, VLLM_SHFL_XOR_SYNC(qk_max, mask));

}

if (lane == 0) {

    red_smem[warp_idx] = qk_max;

}

그다음 각 워프 안에서 qk_max 리덕션을 수행해야 해요. 핵심 아이디어는 워프 안의 스레드들이 서로 통신해서 최종 max qk를 얻는 것입니다.

for (int mask = NUM_WARPS / 2; mask >= 1; mask /= 2) {

    qk_max = fmaxf(qk_max, VLLM_SHFL_XOR_SYNC(qk_max, mask));

}

qk_max = VLLM_SHFL_SYNC(qk_max, 0);

마지막으로, 이 스레드 블록의 모든 워프의 qk_max를 비교해서 전체 스레드 블록의 리덕션된 qk_max를 얻을 수 있어요. 그리고 이 최종 결과를 각 스레드로 브로드캐스트합니다.

exp_sum

qk_max와 비슷하게, 전체 스레드 블록의 리덕션된 합 값을 구해야 합니다.

for (int i = thread_idx; i < num_tokens; i += NUM_THREADS) {

    float val = __expf(logits[i] - qk_max);

    logits[i] = val;

    exp_sum += val;

}

...

exp_sum = block_sum<NUM_WARPS>(&red_smem[NUM_WARPS], exp_sum);

먼저 각 스레드 그룹의 모든 exp 값을 더하고, 동시에 logits의 각 항목을 qk에서 exp(qk - qk_max)로 바꿔요. 여기서 사용하는 qk_max는 이미 전체 스레드 블록에 걸친 최대 qk입니다. 그리고 qk_max에서 했던 것처럼 exp_sum을 전체 스레드 블록에 걸쳐 리덕션합니다.

const float inv_sum = __fdividef(1.f, exp_sum + 1e-6f);

for (int i = thread_idx; i < num_tokens; i += NUM_THREADS) {

    logits[i] *= inv_sum;

}

마지막으로, 리덕션된 qk_max와 exp_sum을 이용해 정규화된 최종 소프트맥스 결과를 logits로 얻을 수 있어요. 이 logits 변수는 이후 단계에서 값 데이터와의 내적에 사용됩니다. 이제 logits에는 할당된 모든 컨텍스트 토큰에 대한 qk의 정규화된 소프트맥스 결과가 저장되어 있어야 해요.

Value

이제 값 데이터를 가져와 logits와 내적을 수행할 차례입니다. 쿼리와 키와 달리 값 데이터에는 스레드 그룹 개념이 없어요. 다이어그램에서 보듯 키 토큰 메모리 레이아웃과는 달리, 값 데이터에서는 같은 열의 요소들이 같은 값 토큰에 해당합니다. 값 데이터 한 블록에는 HEAD_SIZE개의 행과 BLOCK_SIZE개의 열이 있고, 이것이 여러 개의 v_vec로 나뉘어요.

각 스레드는 항상 같은 V_VEC_SIZE개의 토큰에서 V_VEC_SIZE개의 요소를 가져옵니다. 그 결과 하나의 스레드는 여러 안쪽 반복을 통해 서로 다른 행의, 그러나 같은 열에 있는 여러 v_vec을 가져와요. 각 v_vec은 logits에서 같은 V_VEC_SIZE개 요소로 이루어진 해당 logits_vec과 내적되어야 합니다. 전체적으로 여러 안쪽 반복을 거치면 각 워프가 값 토큰 한 블록을 처리하고, 여러 바깥 반복을 거치면 전체 컨텍스트 값 토큰을 처리하게 됩니다.

float accs[NUM_ROWS_PER_THREAD];

for ... { // Iteration over different blocks.

    logits_vec = ...

    for ... { // Iteration over different rows.

        v_vec = ...

        ...

        accs[i] += dot(logits_vec, v_vec);

    }

}

위 의사코드처럼, 바깥 루프에서는 k_ptr과 유사하게 logits_vec이 서로 다른 블록을 순회하며 logits에서 V_VEC_SIZE개의 요소를 읽어요. 안쪽 루프에서는 각 스레드가 같은 토큰들에서 V_VEC_SIZE개의 요소를 v_vec으로 읽고 내적을 수행합니다. 중요하게 기억할 점은 각 안쪽 반복에서 스레드가 같은 토큰들에 대해 서로 다른 헤드 위치의 요소를 가져온다는 것입니다. 내적 결과는 accs에 누적되어요. 따라서 accs의 각 항목은 현재 스레드에 할당된 헤드 위치 하나에 대응합니다.

예를 들어 BLOCK_SIZE가 16이고 V_VEC_SIZE가 8이라면, 각 스레드는 한 번에 8개 토큰에 대한 8개의 값 요소를 가져와요. 각 요소는 같은 헤드 위치에 있는 서로 다른 토큰의 것입니다. HEAD_SIZE가 128이고 WARP_SIZE가 32라면, 안쪽 루프 한 번마다 워프는 WARP_SIZE * V_VEC_SIZE = 256개의 요소를 가져와야 해요. 즉 값 토큰 한 블록 전체를 처리하려면 워프가 128 × 16 / 256 = 8번의 안쪽 반복을 수행한다는 뜻입니다. 그리고 각 스레드의 각 accs에는 서로 다른 8개 헤드 위치에서 누적된 8개 요소가 들어 있어요. 스레드 0의 경우 accs 변수는 할당된 8개 토큰 전체에서 누적된 값 헤드의 0번째, 32번째 … 224번째 요소 8개를 갖게 됩니다.

LV

이제 각 워프 안에서 accs에 대한 리덕션을 수행해야 해요. 이 과정을 통해 각 스레드가 한 블록의 모든 토큰에 대해 할당된 헤드 위치의 accs를 누적할 수 있게 됩니다.

for (int i = 0; i < NUM_ROWS_PER_THREAD; i++) {

    float acc = accs[i];

    for (int mask = NUM_V_VECS_PER_ROW / 2; mask >= 1; mask /= 2) {

        acc += VLLM_SHFL_XOR_SYNC(acc, mask);

    }

    accs[i] = acc;

}

다음으로 모든 워프에 걸쳐 accs 리덕션을 수행해서, 각 스레드가 모든 컨텍스트 토큰에 대해 할당된 헤드 위치의 accs 누적값을 갖도록 해요. 주의할 점은 각 스레드의 각 accs는 모든 컨텍스트 토큰에 대해 전체 헤드 요소의 일부만 누적해서 저장한다는 것입니다. 하지만 전체적으로 보면 출력을 위한 모든 결과가 계산되어 있고, 다만 서로 다른 스레드의 레지스터 메모리에 나누어 저장되어 있을 뿐이에요.

Output

이제 계산된 모든 결과를 로컬 레지스터 메모리에서 최종 출력 글로벌 메모리로 써야 합니다.

scalar_t* out_ptr = out + seq_idx * num_heads * max_num_partitions * HEAD_SIZE

                + head_idx * max_num_partitions * HEAD_SIZE

                + partition_idx * HEAD_SIZE;

먼저 할당된 시퀀스와 할당된 헤드의 시작 주소를 가리키는 out_ptr 변수를 정의합니다.

for (int i = 0; i < NUM_ROWS_PER_THREAD; i++) {

    const int row_idx = lane / NUM_V_VECS_PER_ROW + i * NUM_ROWS_PER_ITER;

    if (row_idx < HEAD_SIZE && lane % NUM_V_VECS_PER_ROW == 0) {

        from_float(*(out_ptr + row_idx), accs[i]);

    }

}

마지막으로 서로 다른 할당된 헤드 위치를 순회하면서, out_ptr을 기준으로 해당하는 누적 결과를 써 내려갑니다.

Citation

@inproceedings{kwon2023efficient,

  title={Efficient Memory Management for Large Language Model Serving with PagedAttention},

  author={Woosuk Kwon and Zhuohan Li and Siyuan Zhuang and Ying Sheng and Lianmin Zheng and Cody Hao Yu and Joseph E. Gonzalez and Hao Zhang and Ion Stoica},

  booktitle={Proceedings of the ACM SIGOPS 29th Symposium on Operating Systems Principles},

  year={2023}

}