tiny-vLLM: CUDA로 직접 쓴 LLM 추론 엔진 읽기 on Jaehun's Blog

1. CUDA 1,500줄로 만든 추론 엔진, 전체 지도

Thu, 17 Sep 2026 00:00:00 +0900

jmaczan/tiny-vllm 은 Llama 3.2 1B Instruct 를 GPU 에서 돌리는 추론 엔진을 C++ 와 CUDA 로 처음부터 짠 저장소다. PyTorch 도, Hugging Face 도, 심지어 토크나이저도 쓰지 않는다. safetensors 파일을 직접 열어 가중치를 GPU 로 올리고, 어텐션과 RMSNorm 과 softmax 를 전부 자기 커널로 계산한다.

2. 가중치 적재와 GPU 버퍼 설계

Thu, 17 Sep 2026 00:00:00 +0900

이 엔진에서 main() 은 cuBLAS 핸들을 만든 뒤 곧바로 두 가지를 한다(src/main.cpp:557-569 ). model.safetensors 파일에서 가중치를 꺼내 GPU 메모리에 올리는 일, 그리고 앞으로 모든 계산이 쓸 버퍼를 미리 잡아 두는 일. 이번 편은 그 둘을 읽는다.

3. 프롬프트 전체를 한 번에 흘리는 prefill 경로

Thu, 17 Sep 2026 00:00:00 +0900

1편에서 prefill() 은 프롬프트 전체를 한 번에 처리하는 함수라고 했다. 17개 토큰이 한 호출 안에서 임베딩부터 다음 토큰 선택까지 통과한다. 이번 편은 그 함수의 안을 본다. 토큰들이 레이어 하나를 지나는 동안 어떤 커널을, 어떤 순서로 거치는지, 그리고 그 순서가 왜 그 모양인지를 각 커널 코드로 따라간다.

4. cuBLAS 전치 트릭 — 왜 행렬을 뒤집어서 넘기는가

Thu, 17 Sep 2026 00:00:00 +0900

3편에서 Q/K/V 투영을 볼 때 이런 문장을 적어 두고 넘어갔다. “활성과 가중치는 행 우선인데 cuBLAS 는 열 우선을 가정하므로, CUBLAS_OP_T 전치 플래그와 lda/ldb/ldc 만으로 호출한다. 왜 이렇게 부르는지는 4편이 다룬다.” 이 편이 그 “왜"다. 3편이 커널이 어떤 순서로 흐르는지를 보았다면, 이 편은 같은 cuBLAS 호출을 다시 꺼내 인자를 하나씩 읽는다.

5. 토큰 하나를 만드는 decode 경로 — 왜 커널을 따로 만들었나

Thu, 17 Sep 2026 00:00:00 +0900

prefill 이 프롬프트 전체를 한 번에 흘려 보내면, 그다음부터는 지금까지 만든 토큰을 바탕으로 토큰을 하나씩 만들어야 한다. 이 단계가 decode 다. 1편에서 본 대로 prefill 과 decode 는 연산의 모양이 달라서 — 전자는 행렬 × 행렬, 후자는 벡터 × 행렬 — 이 저장소는 커널을 따로 만든다. embeddingGatherKernelDecode, ropeKernelDecode, softmaxKernelDecode 처럼 이름에 Decode 가 붙은 변형들이고, 여기에 prefill 에 대응이 없는 pagedAttentionKernel 이 하나 더 붙는다. 이 편은 그 커널들을 prefill 과 나란히 놓고 무엇이 다르고 왜 다른 모양인지, 그리고 새 토큰의 K/V 가 어떻게 캐시에 붙는지를 읽는다.

6. KV cache 를 16토큰 블록으로 쪼개기 — block table 로 비연속 블록을 읽는 법

Thu, 17 Sep 2026 00:00:00 +0900

decode 가 토큰을 하나 만들 때마다 어텐션은 그 시퀀스가 지금까지 쌓은 모든 K/V 를 다시 읽는다. 이 K/V 를 시퀀스마다 연속된 큰 버퍼 하나로 잡으면, 시퀀스들이 저마다 다른 시점에 끝나면서 안 쓰는 구간이 생겨 메모리가 낭비된다. vLLM 을 유명하게 만든 PagedAttention 의 답은 KV cache 를 운영체제의 페이지처럼 작은 블록으로 쪼개 필요할 때만 할당하고, block table 이라는 배열이 “어느 시퀀스의 어느 묶음이 어느 블록인지"를 추적하게 하는 것이다(Kwon et al., SOSP 2023 ). 이 저장소는 BLOCK_SIZE = 16 토큰짜리 블록, block table, 그리고 pagedAttentionKernel 커널 하나로 그 아이디어를 구현했다.

7. 슬롯과 큐로 여러 요청을 겹치기 — continuous batching 의 구현

Thu, 17 Sep 2026 00:00:00 +0900

GPU 는 한 번에 한 시퀀스만 처리하면 그 시퀀스가 prefill 을 마치고 decode 로 넘어가는 사이에서 쉬게 된다. continuous batching 은 처리 중이던 요청이 끝나는 대로 그 자리에 큐에서 기다리던 다음 요청을 넣어 GPU 를 놀리지 않는 방식이다. 이 저장소의 README 도 그 원리를 한 문단으로 정리한다 — “슬롯에 프롬프트를 채우고, 어느 슬롯의 생성이 끝나면 그 결과를 돌려준 뒤 큐에서 기다리던 프롬프트를 갓 비워진 슬롯에 채운다.”