Một nhóm nghiên cứu từ CMU , UIUC , Berkeley , NVIDIA và Đại học Tsinghua đã phát triển một hệ thống trình biên dịch đột phá giúp cải thiện đáng kể tốc độ suy luận của các mô hình ngôn ngữ lớn. Hệ thống Mirage Persistent Kernel ( MPK ) tự động chuyển đổi các phép toán LLM thành một megakernel duy nhất chạy hoàn toàn trên phần cứng GPU , mang lại hiệu suất nhanh hơn từ 1.2 đến 1.7 lần so với các hệ thống hiện tại.
Sự đổi mới này giải quyết một vấn đề không hiệu quả cơ bản trong các hệ thống LLM hiện tại, nơi nhiều phép toán GPU riêng biệt tạo ra các nút thắt cổ chai và lãng phí sức mạnh xử lý. Thay vì khởi chạy hàng chục kernel riêng lẻ cho các tác vụ khác nhau, MPE kết hợp tất cả thành một phép toán liên tục duy nhất không bao giờ ngừng chạy trên GPU .
So sánh hiệu suất trên NVIDIA A100 (40GB)
- Các hệ thống truyền thống ( vLLM , SGLang ): 51.5ms mỗi token
- Hệ thống MPE : 33.1ms mỗi token
- Mức tối thiểu lý thuyết: 30ms mỗi token
- Cải thiện hiệu suất: nhanh hơn 1.2-1.7 lần
Kernel Fusion cách mạng loại bỏ các nút thắt cổ chai hiệu suất
Các hệ thống suy luận LLM truyền thống gặp phải một lỗ hổng nghiêm trọng: chúng xử lý mỗi phép toán như một tác vụ riêng biệt, hoàn chỉnh phải kết thúc trước khi tác vụ tiếp theo bắt đầu. Điều này tạo ra những khoảng thời gian chờ đợi không cần thiết khi phần cứng GPU mạnh mẽ phải ngồi không. Cộng đồng đã nhận ra đây là một vấn đề tồn tại lâu dài, với các nhà phát triển lưu ý rằng các giải pháp hiện có như CUDA Graphs vẫn duy trì các ranh giới cứng nhắc giữa các phép toán.
Cách tiếp cận của MPE phá vỡ những rào cản nhân tạo này bằng cách tạo ra các tác vụ chi tiết có thể chồng chéo và chạy đồng thời. Ví dụ, trong khi một phần của GPU xử lý các phép tính ma trận, một phần khác có thể ngay lập tức bắt đầu các tác vụ truyền thông với kết quả một phần, thay vì chờ đợi toàn bộ phép tính hoàn thành.
Hệ thống hoạt động bằng cách tự động phân tích các yêu cầu tính toán của LLM và xây dựng lại chúng như một mạng lưới các tác vụ nhỏ, kết nối với nhau. Mỗi tác vụ chạy trên các bộ xử lý đa luồng streaming riêng lẻ của GPU , với các bộ lập lịch tích hợp điều phối quy trình làm việc mà không cần bất kỳ can thiệp bên ngoài nào.
Các thành phần kiến trúc hệ thống MPE
- Workers: Thực thi các tác vụ trên từng bộ xử lý đa luồng streaming ( SMs ) của GPU
- Schedulers: Chạy trên các warp đơn lẻ, tối đa 4 trên mỗi SM , quản lý hàng đợi tác vụ
- Task Graph: Các đơn vị chi tiết thay thế mô hình kernel-per-operator truyền thống
- Events: Các điểm đồng bộ hóa giữa các tác vụ cho phép chồng lấp pipeline
Cải thiện hiệu suất ấn tượng trên cả thiết lập GPU đơn và đa GPU
Thử nghiệm trên phần cứng NVIDIA A100 cho thấy những cải thiện đáng kể trong các tình huống thực tế. Đối với cấu hình GPU đơn, MPE đã giảm thời gian tạo mỗi token từ 51.5 mili giây xuống 33.1 mili giây, tiến gần đến tốc độ tối đa lý thuyết có thể đạt được với các giới hạn băng thông bộ nhớ.
Lợi thế hiệu suất trở nên rõ rệt hơn với nhiều GPU . Trong khi các hệ thống truyền thống gặp khó khăn với chi phí điều phối giữa các card đồ họa riêng biệt, cách tiếp cận thống nhất của MPE cho phép chồng chéo liền mạch giữa tính toán và truyền thông giữa các GPU . Điều này làm cho công nghệ đặc biệt có giá trị cho các triển khai quy mô lớn nơi nhiều GPU làm việc cùng nhau.
Phản hồi từ cộng đồng đã vô cùng tích cực, với các nhà phát triển làm việc chặt chẽ với các hệ thống hiện có như vLLM và SGLang bày tỏ sự hứng thú về cách tiếp cận phân tích biểu đồ phụ thuộc.
Thách thức kỹ thuật và phát triển tương lai
Mặc dù có kết quả ấn tượng, MPE hiện tại vẫn đối mặt với những hạn chế mà nhóm nghiên cứu đang tích cực giải quyết. Hệ thống hoạt động tốt nhất với các biểu đồ tính toán tĩnh, khiến nó không tương thích với các mô hình động như kiến trúc Mixture-of-Experts thay đổi hành vi dựa trên nội dung đầu vào.
Nhóm cũng đang làm việc để hỗ trợ các kiến trúc GPU thế hệ tiếp theo như NVIDIA Blackwell , mang đến những cơ hội và thách thức tối ưu hóa mới. Các chính sách lập lịch nâng cao có thể cải thiện thêm hiệu suất cho các trường hợp sử dụng cụ thể như ứng dụng nhạy cảm với độ trễ hoặc các tình huống batching lai.
Dự án này đến từ CMU . Hazy Research tại Stanford cũng đã nói về megakernel... Thật tốt khi thấy sự cạnh tranh trong lĩnh vực này.
Bối cảnh cạnh tranh bao gồm các nỗ lực song song từ các tổ chức nghiên cứu khác, đặc biệt là nhóm Hazy Research của Stanford , mặc dù trọng tâm của MPE vào biên dịch tự động khiến nó khác biệt với các cách tiếp cận tối ưu hóa thủ công.
Trình biên dịch MPE đại diện cho một bước tiến đáng kể hướng tới việc làm cho suy luận LLM hiệu suất cao trở nên dễ tiếp cận với các nhà phát triển mà không cần kiến thức lập trình GPU sâu. Chỉ với vài chục dòng code Python , các nhà phát triển giờ đây có thể biên dịch mô hình của họ thành các megakernel được tối ưu hóa, dân chủ hóa việc tiếp cận các tối ưu hóa hiệu suất tiên tiến trước đây đòi hỏi kiến thức chuyên môn.
Tham khảo: Compiling LLMs into a MegaKernel: A Path to Low-Latency Inference