Biên dịch LLM thành MegaKernel để hiện thực suy luận độ trễ thấp
(zhihaojia.medium.com)- Đã phát triển một trình biên dịch tự động chuyển suy luận LLM thành một megakernel duy nhất
- Phương pháp MegaKernel (kernel persistent) cho phép độ trễ rất thấp bằng cách hợp nhất hoàn toàn tính toán và giao tiếp trong suy luận LLM vào một kernel GPU duy nhất
- Có một vấn đề là do cấu trúc phân tán của các framework ML hoặc thư viện kernel hiện có, việc kernel hóa toàn bộ pipeline thành một kernel duy nhất là rất khó
- Mirage Persistent Kernel (MPK) tự động chuyển suy luận LLM đa GPU thành megakernel hiệu năng cao thông qua trình biên dịch và hệ thống runtime
- MPK chuyển đồ thị tính toán thành đồ thị tác vụ chi tiết, từ đó tối đa hóa software pipelining và chồng lấp giữa tính toán - giao tiếp
- Khi áp dụng MPK, độ trễ sinh token giảm xuống so với các hệ thống hiện có, và mức cải thiện hiệu năng càng lớn khi số lượng GPU tăng lên
Tổng quan và ưu điểm của phương pháp MegaKernel
- Trong suy luận mô hình ngôn ngữ lớn (LLM), một trong những cách hiệu quả để giảm độ trễ là hợp nhất toàn bộ quá trình tính toán và giao tiếp vào một megakernel (kernel nhất quán) duy nhất
- Cách tiếp cận này cho phép một kernel GPU duy nhất thực hiện liền mạch toàn bộ xử lý, từ các phép toán theo từng lớp của mô hình đến giao tiếp giữa các GPU
- Các lợi ích chính gồm có
- Bỏ qua các lần gọi kernel lặp lại để loại bỏ overhead khởi chạy kernel
- Có thể hiện thực software pipelining trên toàn bộ các lớp
- Thực hiện đồng thời tính toán và giao tiếp để ẩn độ trễ
Những giới hạn trước đây và sự xuất hiện của MPK
- Các framework ML hiện có như PyTorch, Triton, TVM về bản chất không hỗ trợ tự động tạo end-to-end megakernel
- Các hệ thống LLM thực tế được tạo thành từ sự kết hợp của nhiều thư viện kernel như NCCL/NVSHMEM (giao tiếp), FlashInfer/FlashAttention (attention), CUDA/Triton (phép toán tùy biến), nên rất khó hợp nhất thành một kernel duy nhất
- Trong bối cảnh đó, các nhà nghiên cứu từ CMU, UW, Berkeley, NVIDIA và Tsinghua đã phát triển Mirage Persistent Kernel (MPK)
- MPK kết hợp trình biên dịch và runtime để tự động chuyển toàn bộ pipeline suy luận LLM thành megakernel hiệu năng cao
Giá trị cốt lõi của MPK
- MPK loại bỏ hoàn toàn overhead khởi chạy kernel, đồng thời tối đa hóa chồng lấp giữa tính toán/tải dữ liệu/giao tiếp giữa các tầng để hiện thực môi trường suy luận LLM độ trễ cực thấp
- Trong thử nghiệm thực tế (prompt 39 token, sinh 512 token, không dùng speculative decoding),
- Trong môi trường một NVIDIA A100 40GB GPU, so với độ trễ giải mã mỗi token của các hệ thống tối ưu hóa hiện có như vLLM/SGLang (14,5ms), MPK giảm xuống còn 12,5ms
- Con số này tiệm cận giới hạn lý thuyết (10ms) dựa trên băng thông bộ nhớ 1,6TB/s và tải trọng số 16GB
- Trong môi trường đa GPU, nhờ tích hợp hoàn toàn tính toán và giao tiếp, ưu thế hiệu năng của MPK càng nổi bật khi số lượng GPU tăng lên
Cấu trúc vận hành chi tiết của MPK
Phần 1. Trình biên dịch – chuyển đồ thị tính toán LLM → đồ thị tác vụ
- Thông thường, phép tính LLM được biểu diễn dưới dạng đồ thị tính toán, trong đó mỗi phép toán (ví dụ: nhân ma trận, attention) hoặc phép giao tiếp (ví dụ: all-reduce) là một nút, còn phụ thuộc dữ liệu là các cạnh
- Trong thiết kế truyền thống, cách chạy kernel riêng cho từng toán tử là phổ biến, nhưng điều này chỉ phản ánh phụ thuộc ở cấp kernel chứ không phản ánh đơn vị phụ thuộc dữ liệu thực tế, nên cơ hội pipelining bị hạn chế
- Ví dụ: nếu có all-reduce sau một phép nhân ma trận, thì all-reduce chỉ bắt đầu khi toàn bộ phép nhân ma trận kết thúc. Trong thực tế, có thể chia nhỏ dữ liệu và tận dụng thực thi/phụ thuộc từng phần
- Trình biên dịch MPK tinh chỉnh đồ thị tính toán và tự động chuyển thành fine-grained task graph phù hợp với đơn vị dữ liệu thực tế
- Mỗi tác vụ (hình chữ nhật) là một đơn vị tính toán/giao tiếp được gán cho từng GPU SM riêng lẻ
- Mỗi sự kiện (hình tròn) là một điểm đồng bộ giữa các tác vụ
- Các cạnh giữa tác vụ và sự kiện biểu diễn hiệu quả phụ thuộc dữ liệu/điều khiển
- Nhờ đồ thị tác vụ này, MPK cho phép tính toán và giao tiếp chồng lấp từng phần hoặc song song nhiều hơn
- Với Mirage kernel superoptimizer, hệ thống cũng tự động sinh ra triển khai CUDA hiệu năng cao phù hợp cho từng tác vụ
Phần 2. Runtime – thực thi đồ thị tác vụ bên trong megakernel
- Runtime của MPK thực thi hoàn toàn đồ thị tác vụ chỉ bên trong một kernel GPU (megakernel)
- Toàn bộ SM (Streaming Multiprocessors) của GPU được phân chia tĩnh thành vai trò worker và scheduler
Worker
- Mỗi worker hoạt động ở cấp SM và quản lý hàng đợi tác vụ chuyên dụng
- Theo cơ chế vòng lặp
- Lấy tác vụ tiếp theo từ hàng đợi
- Thực thi (ví dụ: matmul, attention, truyền dữ liệu)
- Khi hoàn thành thì thông báo cho sự kiện
- Lặp lại xử lý
- Nhờ đó có thể tối ưu hóa việc sử dụng tài nguyên của từng worker và cho phép tính toán phân tầng bất đồng bộ
Scheduler
- Scheduler phân tán hoạt động ở cấp một warp trong mỗi SM, và có thể chạy đồng thời tối đa 4 scheduler
- Mỗi scheduler quản lý hàng đợi các sự kiện đã được kích hoạt và gán các tác vụ thỏa điều kiện cho worker
- Nhờ vậy có thể phân phối tác vụ quy mô lớn mà không cần overhead đồng bộ tập trung
Cơ chế thực thi dựa trên sự kiện
- Khi một tác vụ hoàn thành, nó sẽ tăng bộ đếm sự kiện tương ứng. Khi bộ đếm đạt ngưỡng, sự kiện được kích hoạt và được chèn vào hàng đợi của scheduler
- Scheduler sẽ thực thi các tác vụ kế tiếp có quan hệ phụ thuộc với sự kiện đó
- Nhờ vậy, software pipelining fine-grained và chồng lấp tính toán - giao tiếp diễn ra một cách tự nhiên
- Ví dụ: matmul của một lớp và attention của lớp khác có thể chạy đồng thời
- Ngay khi có kết quả matmul hoàn thành một phần, có thể bắt đầu giao tiếp all-reduce
- Vì toàn bộ lập lịch và chuyển đổi tác vụ đều diễn ra trong ngữ cảnh một kernel duy nhất, overhead giữa các tác vụ rất thấp, chỉ ở mức 1–2 micro giây (μs)
Hướng đi tương lai
-
Mục tiêu của MPK: hỗ trợ để nhà phát triển chỉ cần viết một lượng nhỏ mã Python (khoảng vài chục dòng) cũng có thể dễ dàng biên dịch LLM thành megakernel và khai thác hiệu năng tối đa
-
Các hướng phát triển chính
- Hỗ trợ kiến trúc GPU mới nhất: ví dụ như NVIDIA Blackwell, các phương thức tối ưu chuyên biệt ở cấp warp
- Xử lý workload động: nghiên cứu chiến lược biên dịch cho các mô hình cần luồng điều khiển động như mixture-of-experts (MoE)
- Lập lịch tác vụ nâng cao: theo đuổi khả năng nghiên cứu và áp dụng các chính sách hiện đại như dựa trên mức ưu tiên, tối ưu thông lượng
-
MPK đưa ra một bước ngoặt mang tính nền tảng trong cách biên dịch và thực thi suy luận LLM trên GPU, đồng thời kỳ vọng mở rộng hợp tác với cộng đồng
Tài liệu thêm
- Có thể xem mã nguồn, tài liệu và các kết quả nghiên cứu mới nhất của MPK (Mirage Persistent Kernel) trên GitHub (https://github.com/mirage-project/mirage)
1 bình luận
Ý kiến Hacker News
Gửi tác giả, thật thú vị khi cách tiếp cận trình thông dịch chạy trực tiếp trên GPU có vẻ là một hướng đi rất hứa hẹn cho tương lai. Cũng có những nghiên cứu khác với cách tiếp cận gần như tương tự, nên khuyến nghị tham khảo bài viết liên quan. Mô hình lập trình nền tảng của CUDA (ví dụ: kernel launch) đang bị lách qua để phục vụ kiểu song song hóa dựa trên tác vụ rất nhỏ, và tôi đã trực tiếp chứng kiến cách này giúp tận dụng phần cứng tốt hơn. Tự hỏi liệu CUDA có phải đã kìm hãm chúng ta ở nhiều khía cạnh hay không. Cũng rất mong chờ khả năng nghiên cứu của tác giả có thể đi vào backend thử nghiệm của PyTorch. Ngoài ra, hai đoạn ở phần đầu gần như giống hệt nhau nên xin chỉ ra một lỗi nhỏ.
Tôi đã làm việc khá sát với vLLM và SGLang trong một thời gian, và tin rằng dự án này chính là hình mẫu lý tưởng cho bước tiếp theo. Phần phân tích đồ thị phụ thuộc phép toán, rồi fusion các phép toán hoặc lên lịch tác vụ thông minh hơn thực sự rất ấn tượng. Chúc mừng cả nhóm.
Tôi đã đọc lướt bài viết và README trên github, và thấy đây là một dự án thực sự tuyệt vời. Tôi tò mò liệu các hướng tối ưu hóa như thế này có thể áp dụng không chỉ cho suy luận mà còn cho cả giai đoạn huấn luyện hay không. Đặc biệt, tôi hiểu rằng việc fusion giữa phép toán backward và giao tiếp gradient sẽ là một thách thức. Theo tôi biết thì hiện tại chưa hỗ trợ dynamic workload (ví dụ: MoE), và xin nhắc đến bài báo gần đây xử lý MoE trong một kernel duy nhất: FlashDMoE: Fast Distributed MoE in a Single Kernel.
Cảm ơn bạn đã đọc cả bài viết lẫn README. Hỗ trợ giai đoạn huấn luyện là khả thi, nhưng nhìn chung kernel cho huấn luyện lớn hơn nên overhead của kernel launch không phải vấn đề quá lớn; vì thế suy luận (đặc biệt là độ trễ thấp) sẽ là bên hưởng lợi nhiều hơn. Chúng tôi cũng thấy bài báo FlashDMoE bạn chia sẻ rất thú vị, và hỗ trợ mô hình MoE cũng là mục tiêu tiếp theo của chúng tôi.
Cá nhân tôi hơi hoài nghi về việc đầu tư nhiều thời gian vào tối ưu hóa huấn luyện dựa trên gradient. Trên thực tế, nhiều tác vụ huấn luyện mang đặc tính giá trị rời rạc, và tôi cho rằng học dựa trên gradient không xử lý chúng tốt.
Bước tiếp theo thì phải là biên dịch thẳng sang Verilog rồi lên aliexpress mua luôn phần cứng LLM về dùng mới đúng mơ ước.
Xin chia sẻ một bài viết giới thiệu công nghệ phần cứng như Chisel. Trước khi AI và GPU xuất hiện, ý tưởng chuyển trực tiếp từ phần mềm sang phần cứng như vậy từng là một hướng tiếp cận đầy hứa hẹn. Sự phát triển của CPU hiện đang chững lại, và khát vọng tối ưu thêm lớp trung gian giữa phần mềm và phần cứng vẫn luôn tồn tại, nhưng nhiều khả năng kiểu tính toán song song theo phong cách GPU sẽ tiếp tục là cách tăng tốc chủ đạo. CPU phổ thông rốt cuộc có lẽ sẽ chỉ còn đóng vai trò như bộ não nhỏ điều phối GPU. Dù vậy, tôi dự đoán cách đi trực tiếp từ phần mềm sang phần cứng sẽ khó trở thành xu hướng chính.
Có thể trong 5–10 năm nữa, khi cấu trúc của LLM ổn định hơn, việc ánh xạ trực tiếp lên phần cứng sẽ trở nên thực tế. Với công nghệ hiện nay, thậm chí một mô hình hàng chục tỷ tham số cũng có thể vừa trên một wafer đơn nếu chỉ dùng logic gate siêu thấp chính xác quanh mức 1.5 bit. Khi độ chính xác tăng lên, số lượng gate tăng theo cấp số nhân, nên hiện tại cách giữ bộ nhớ trọng số và chia sẻ đơn vị tính toán vẫn hiệu quả hơn. Trong tương lai, phát triển LLM siêu thấp chính xác sẽ là một bài toán bắt buộc.
Chi phí huấn luyện vốn đã cao, mà cộng thêm chi phí mask nữa thì chỉ càng khó khăn hơn — một câu đùa, nhưng đồng thời cũng là đánh giá khá tỉnh táo rằng các startup phần cứng AI thực chất đã thử những hướng như vậy từ lâu rồi.
Nếu thực sự có kiểu LLM-in-a-box thì sẽ rất hấp dẫn. Tôi sắp có cơ hội làm việc trong một môi trường offline (air-gap), và một giải pháp như thế có vẻ sẽ cực kỳ hữu ích.
Tôi đã tự chạy mã trong môi trường Modal GPU, và các con số cải thiện hiệu năng mà nghiên cứu nêu ra thực sự tái hiện được. Xin chia sẻ mã kết quả của dự án mirage. Với tổ hợp Triton + FlashInfer, độ trễ mỗi token ở mức khoảng 19.2ms, còn với MPK thì trong cùng điều kiện chỉ còn 7.7ms, cải thiện rất lớn.
Trước đây tôi từng tham gia một cuộc thi CUDA nhỏ. Đó là một thuật toán song song trong lĩnh vực ảnh hoặc thị giác, và khi đó tôi nghĩ mình thông minh nên đã cache kết quả trung gian vào bộ nhớ. Sau khi xem kết quả cuộc thi, tôi ngạc nhiên vì người khác nộp mã nhanh hơn tôi rất nhiều. Lý do là họ không cache mấy kết quả trung gian đó mà cứ tính lại liên tục. Chi phí tính toán hóa ra nhỏ hơn rất nhiều so với việc đi-về bộ nhớ. Tôi đoán dự án này có lẽ cũng tương tự. Khi biên dịch thành megakernel, ranh giới giữa các layer biến mất, nên việc chia sẻ kết quả trung gian giảm đi còn lượng tính toán lại tăng lên, nhưng xét tổng thể thì giảm được rất nhiều lượt truy cập bộ nhớ nên lợi ích lớn. Đặc biệt với mạng tích chập thì chắc sẽ có một sweet spot nào đó, nhưng tôi không rõ megakernel xử lý phần này như thế nào.
Ngay cả bây giờ vẫn liên tục xuất hiện những phép so sánh mới về LLM. Tôi chợt nghĩ, liệu có thể xem LLM như transistor không? Hiện tại có cảm giác giống thời kỳ máy tính cỡ căn phòng chỉ biết làm phép nhân bằng thẻ đục lỗ. Cứ thử tưởng tượng nếu có thể chạy đồng thời 1 triệu truy vấn o3-pro thì mọi chuyện sẽ ra sao, cũng khá thú vị.
Dự án này xuất thân từ CMU (Carnegie Mellon). Ở Stanford, Hazy Research cũng có một blog về megakernel là No Bubbles. Rất ấn tượng khi thấy sự cạnh tranh diễn ra sôi động trong lĩnh vực này. (Bổ sung) Cũng có một bài báo nói về bức tranh lớn hơn của dự án "mirage", nhưng không đề cập tới cách tiếp cận megakernel: liên kết bài báo
Chính tác giả bài đăng đã trực tiếp trả lời. Đồng ý rằng nghiên cứu với Stanford đang diễn ra song song. Khác biệt chính là chúng tôi tập trung vào compiler sinh megakernel một cách tự động.
Cũng xin nhắc rằng ThunderKittens của Hazy Research là một thư viện cực kỳ ngầu. Gần đây có rất nhiều nỗ lực dồn vào việc formal hóa, pipelining, chia để trị, tối đa hóa hiệu suất, và phát triển compiler/DSL chuyên dụng để tận dụng tối đa các dòng GPU mới của NVIDIA.
Nếu các con số hiệu năng của Qwen 8B được xác minh thì thực sự rất ấn tượng. So với các cách làm megakernel trước đây, hướng này có vẻ thực tiễn hơn. Kiểu kernel được duy trì một cái trên mỗi SM này gợi nhớ đến Larrabee ngày xưa. Tò mò không biết thế giới bây giờ sẽ ra sao nếu trước đây người ta không đi theo CUDA mà chọn con đường process-thread-SIMD truyền thống.
Một ý tưởng là thay vì suy luận dựa trên phần mềm, hãy làm một LLM cố định hoàn toàn bằng ASIC thuần túy. Lợi thế chi phí ra sao? Liệu có thể cung cấp thêm một lớp để phần mềm xử lý bổ sung hoặc tinh chỉnh vi mô không? Vì chúng ta dường như đã gần đạt đến mức “đủ tốt” trên thực tế, nên cũng có khả năng trong 2–4 năm tới sẽ có người quyết định đóng cứng lên chip chuyên dụng rồi dùng luôn. Tôi rất tò mò không biết từ thời điểm nào thì lợi ích của phần cứng siêu chuyên biệt mới thực sự phát huy.