Tự xây dựng engine suy luận LLM từ đầu bằng C++ và CUDA

• Cách xây dựng một engine suy luận LLM bằng C++ và CUDA mà không dùng thư viện
• Qua đó có thể hiểu toàn bộ stack suy luận LLM và cảm nhận trực tiếp tác động của nhiều kiểu tối ưu hóa lên tốc độ suy luận
• Mục tiêu: triển khai mô hình để suy luận nhanh với single batch trên máy chủ đơn CPU + GPU và đạt tốc độ xử lý token nhanh hơn llama.cpp

1. Tổng quan về kiến trúc và suy luận LLM

• Hầu hết các LLM chủ đạo đều theo cùng một kiến trúc sử dụng các khối transformer nối tiếp nhau.
• Việc nạp mô hình là định nghĩa một lớp khối transformer có thể tùy biến, ghép chúng thành một chuỗi và khởi tạo bằng trọng số safetensors.
• Suy luận chủ yếu diễn ra dưới dạng single batch, và "giai đoạn decode" chiếm phần lớn thời gian thực thi.

1.1 Tổng quan suy luận

• Suy luận được chia thành giai đoạn prefill, nơi các token của prompt được đưa vào mô hình để điền KV cache, và giai đoạn decode, nơi mô hình được chạy lặp lại để sinh token
  • Giai đoạn Prefill: xử lý các token của prompt và khởi tạo KV cache
  • Giai đoạn Decode: sinh từng token một
• KV cache: lưu các cặp key/value trước đó để tính attention với ngữ cảnh quá khứ nhanh hơn
• Forward pass của mô hình dùng bảng embedding để ánh xạ token ID sang vector embedding, rồi biến đổi trạng thái qua chuỗi khối transformer

1.2 Nút thắt cổ chai và benchmark

• Nút thắt cổ chai: trên phần cứng hiện đại, băng thông bộ nhớ là yếu tố giới hạn
  • Khi suy luận mô hình, mỗi lần sinh một token đều phải đọc toàn bộ mô hình, nên băng thông bộ nhớ là ràng buộc lớn hơn cả tính toán
• Lượng tử hóa mô hình rất hiệu quả trong việc cải thiện tốc độ suy luận
• Thông lượng token tối đa về lý thuyết khác nhau tùy phần cứng, và hiệu năng thực tế có thể kiểm tra qua nhiều engine suy luận
• Giới hạn tốc độ lý thuyết:
  • AMD EPYC 7702P: tối đa 13.6 tok/s (FP16)
  • RTX 4090: tối đa 67.1 tok/s (FP16)
• Benchmark:
  • llama.cpp: CPU 8.7 tok/s, GPU 61 tok/s
  • calm: GPU 66 tok/s

2. Suy luận trên CPU

• Bản triển khai ban đầu trên CPU chạy đơn luồng và chỉ hỗ trợ trọng số FP32
• Có thể bắt đầu song song hóa mã bằng multithreading và cải thiện hiệu năng bằng SIMD

2.1 Multithreading

• Dùng OpenMP để song song hóa phép nhân ma trận-vectơ (matmul) và multi-head attention nhằm cải thiện hiệu năng
• Kết quả tối ưu hóa: tốc độ tăng từ 0.6 tok/s → 4.4 tok/s

2.2 Lượng tử hóa trọng số và tối ưu SIMD

• Lượng tử hóa: lượng tử hóa trọng số FP32 xuống FP16 để giảm một nửa mức dùng bộ nhớ và cải thiện hiệu năng
• SIMD: dùng AVX2 để tối ưu xử lý đồng thời 8 giá trị FP32
• Kết quả: đạt 8.4 tok/s

3. Suy luận trên GPU

• Có thể lượng tử hóa mô hình sang FP16, nạp lên RTX 4090 và bắt đầu triển khai suy luận trên GPU
• Với CUDA, có thể chạy song song các hàm C++ (kernel) trên GPU

3.1 Port đơn giản sang CUDA

• Có thể triển khai backend GPU bằng cách chuyển các phép tính trên CPU sang kernel CUDA theo kiểu 1-1
• Kernel CUDA chạy bất đồng bộ, nhưng vẫn thực thi tuần tự trong cùng một stream
• Vấn đề: hiệu quả luồng thấp khiến tài nguyên GPU không được tận dụng đầy đủ → chậm ở mức 2.9 tok/s

3.2 Phép nhân ma trận (matmul) tốt hơn

• Phép nhân ma trận chiếm phần lớn thời gian chạy trên CPU và có thể được tối ưu bằng OpenMP
• Trên GPU, có thể tăng mức sử dụng luồng bằng cách để mỗi block xử lý một hàng
• Cách tối ưu:
  1. Mỗi block xử lý một hàng, các luồng trong block phối hợp để tính toán
  2. Áp dụng warp reduction
• Kết quả: tốc độ tăng lên 51.7 tok/s

3.3 Hợp nhất kernel và tối ưu bổ sung

• Có thể cải thiện hiệu năng bằng cách hợp nhất kernel
  • Kernel fusion: gộp các phép tính liên tiếp vào một kernel để giảm truy cập bộ nhớ và thời gian tính toán
• Tối ưu mẫu truy cập bộ nhớ và tái sử dụng không gian giúp đạt 56.1 tok/s

3.4 Tối ưu attention và xử lý ngữ cảnh dài

• Vấn đề: với ngữ cảnh dài, kernel attention trở thành nút thắt hiệu năng
• Giải pháp:
  1. Tối ưu truy cập bộ nhớ: thiết kế lại để đọc các khối bộ nhớ liên tiếp
  2. Dùng bộ nhớ chia sẻ thay cho atomicAdd để giải quyết vấn đề thất thoát giá trị số thực
• Kết quả tối ưu hóa:
  • Ngữ cảnh ngắn: 63.8 tok/s (nhanh hơn 61.0 tok/s của llama.cpp)
  • Ngữ cảnh dài: đạt 58.8 tok/s

3.5 Lượng tử hóa KV cache và vấn đề tối ưu của compiler

• Lượng tử hóa KV cache sang FP16 làm suy giảm hiệu năng (do compiler tối ưu chưa tốt)
• Giải pháp: tự unroll loop và áp dụng memory prefetching
• Kết quả: nhanh hơn khoảng 2 lần so với FP32 và vẫn giữ hiệu năng ngữ cảnh dài ở mức 58.8 tok/s

4. Hướng cải thiện tiếp theo

• Tối ưu prefill của prompt: xử lý nhiều token cùng lúc để rút ngắn thời gian sinh token đầu tiên
• Hợp nhất kernel attention: áp dụng các kỹ thuật tối ưu như FlashAttention
• Lượng tử hóa cao hơn: áp dụng FP8, INT8, INT4 cùng lượng tử hóa activation/cache
• Tối ưu kernel: đưa vào các kỹ thuật nâng cao để tối đa hóa băng thông bộ nhớ và hiệu quả tính toán
• Sử dụng thư viện: tận dụng các thư viện như cuDNN, cuBLAS để rút ngắn thời gian tối ưu

Tóm tắt kết quả:

• Đạt tốc độ 63.8 tok/s nhờ nhiều tối ưu hóa trên cả CPU và GPU
• Ghi nhận hiệu năng tiệm cận hoặc tốt hơn llama.cpp và calm
• Triển khai thành công engine suy luận LLM hiệu năng cao chỉ với C++ và CUDA, không dùng thư viện

[GN⁺]

Ý kiến trên Hacker News

• Tác giả vui vì bài viết trên blog của mình đã thu hút được sự chú ý và muốn nghe phản hồi
• Một độc giả khen bài viết rất hay và tò mò về thời gian tác giả đã dành để viết
  • Là người làm việc trong lĩnh vực GPGPU, họ cũng muốn viết một bài tương tự nhưng còn do dự vì không chắc sẽ tốn bao nhiêu thời gian
• Một độc giả khác cho rằng đoạn mã không tận dụng tensor cores hay lệnh wgmma
  • Họ giải thích rằng kiểu lập trình này rất khó vì phải xử lý đồng thời nhiều tác vụ
  • Họ cũng nhắc rằng do giới hạn băng thông nên có thể không cần thêm phép tính nào nữa
  • Họ đánh giá mã trong bài blog có khả năng hoạt động tốt khi chuyển sang các bộ tăng tốc khác
  • Họ lo ngại rằng việc dùng wgmma có thể làm giảm tính di động giữa các thế hệ Nvidia
• Một độc giả khác đang tìm tài liệu Python tương tự để chia sẻ với nhóm của mình
  • Họ muốn tài liệu ngắn gọn theo phong cách tutorial và đầy đủ về mặt khái niệm hơn là tối ưu hiệu năng
• Một người dùng muốn so sánh phiên bản Mistral của mình và hiệu năng token/giây
  • Có người khuyên nên tham khảo phần lượng tử hóa trong README
• Có ý kiến cho rằng __shfl_down hiện nay không còn được khuyến nghị vì các vấn đề đồng bộ warp