Databricks chia sẻ cách tối ưu hóa hiệu suất suy luận AI nhờ tạo kernel GPU chuyên biệt

Theo truyền thống, các hệ thống suy luận (inference) trong môi trường sản xuất thường dựa vào các kernel chung để xử lý đa dạng các mô hình và khối lượng công việc. Cách tiếp cận này không tối ưu vì hình dạng (shape) của các thao tác trên GPU được quyết định bởi sự kết hợp giữa các tham số mô hình tĩnh và các yếu tố động tại thời điểm yêu cầu; ví dụ, trong khi một mô hình xác định một trong các chiều của phép nhân ma trận, thì chiều còn lại sẽ biến động dựa trên số lượng token cụ thể của mỗi yêu cầu. Hiện nay, sự quan tâm đến việc tạo kernel GPU bằng tác nhân (agentic GPU kernel generation) đang ngày càng tăng và các nỗ lực gần đây đã cho thấy nhiều hứa hẹn. Chúng tôi đã khám phá một câu hỏi cốt lõi: nếu việc tạo kernel có thể được tự động hóa, tại sao các mô hình với kích thước khác biệt lớn (từ 1 tỷ đến 1 nghìn tỷ tham số) lại phải dựa vào cùng một kernel? Bằng cách chuyên biệt hóa các kernel cho các hình dạng cụ thể gặp phải trong quá trình thực thi (runtime), chúng ta có thể đạt được hiệu suất cực cao.

Trong bài blog này, chúng tôi chia sẻ những thành công và hiểu biết sâu sắc từ việc sử dụng các tác nhân để tạo kernel GPU. Chúng tôi đã xây dựng Proteus, một hệ thống được thiết kế để đạt được sự chuyên biệt hóa cực độ, đòi hỏi một bộ khung (harness) được tùy chỉnh để tối ưu hóa, xác thực và quản lý ngữ cảnh một cách nghiêm ngặt.

Các bộ khung lập trình thông thường thường thất bại ở điểm này vì các tác nhân có xu hướng "hack phần thưởng" (reward-hack): tuân thủ câu chữ của quy tắc thay vì tinh thần của nó. Nếu bạn đưa cho một tác nhân một bộ benchmark, nó có thể tối ưu hóa bộ benchmark đó thay vì thao tác dự định ban đầu.

Để giải quyết vấn đề này, Proteus đề xuất các kernel, xác minh chúng dựa trên một bản thực thi tham chiếu có kiểm soát, đo thời gian các kernel thành công và cải thiện dần dần dựa trên những kết quả tốt nhất. Mặc dù quy trình này khá đơn giản, nhưng sự thành công của nó hoàn toàn phụ thuộc vào việc giải quyết hai thách thức nền tảng. Hình 1 cho thấy kiến trúc đơn giản hóa trong thiết kế của chúng tôi. Sử dụng bộ khung Proteus, chúng tôi đã tạo ra các kernel cho Qwen 3.5 122B với tốc độ nhanh hơn từ 1,8–5,2 lần so với các kernel tốt nhất hiện có trong vLLM.

Xác thực

Ban đầu, chúng tôi coi việc tìm kiếm kernel là phần khó nhất: làm thế nào để khám phá một không gian chương trình rộng lớn mà không bị mắc kẹt ở một trạng thái bão hòa không thể cải thiện? Trên thực tế, câu hỏi đầu tiên lại cơ bản hơn: Liệu chúng ta có đang đo lường đúng những gì chúng ta nghĩ là mình đang đo hay không?

Một mô hình sẽ tối ưu hóa điểm số mà bạn cung cấp cho nó. Nó không cần một thủ thuật tinh vi nào: vấn đề đơn giản có thể nằm ở chỗ đánh giá đang dựa trên một giả định sai lệch. Một ví dụ là các kernel cho rotary position embeddings (RoPE), một bước phổ biến trong các lớp attention. Một ứng viên có thể tái sử dụng mã đã biên dịch còn sót lại từ lần thử trước đó và trông có vẻ "rẻ" hơn so với việc xây dựng lại từ đầu một cách công bằng. Một ứng viên khác có thể ghi lại một loạt các lệnh khởi chạy GPU vào một đồ thị (ví dụ: CUDA graph) và phát lại chúng như một đơn vị duy nhất, trong khi baseline mà chúng tôi so sánh vẫn khởi chạy từng phần riêng biệt, dẫn đến việc hai bên không thực hiện cùng một khối lượng công việc. Một trường hợp khác là mô hình hoạt động tốt trên các kích thước đầu vào có trong tập kiểm tra hiển thị nhưng lại yếu trên các kích thước chưa từng được thấy.

Vì vậy, chúng tôi dành thời gian thiết kế ban đầu cho bộ kiểm tra (checker), thay vì cho prompt. Chúng tôi đo thời gian cả hai bên theo cùng một cách, bao gồm việc sử dụng nhiều hơn một bộ đếm thời gian (ví dụ: bộ đếm sự kiện CUDA, thời gian thực tế và bộ đếm CUPTI) khi cần đối chiếu. Chúng tôi xóa sạch trạng thái biên dịch còn sót lại không nên tồn tại, đồng thời giữ cho thứ tự thiết lập và dọn dẹp nhất quán để một bên không thể bỏ qua công việc mà bên kia vẫn phải thực hiện. Chúng tôi đo lại thời gian của những người chiến thắng trước khi sử dụng chúng làm điểm khởi đầu cho vòng tiếp theo. Chúng tôi giữ lại một số bài kiểm tra mà ứng viên không thể nhìn thấy để nó không thể chỉ "học vẹt" cho bài thi. Để ngăn chặn việc "gian lận" đánh giá bằng hiệu suất được thổi phồng nhân tạo, chúng tôi triển khai các kiểm tra tính nhất quán tự động để gắn cờ các mức tăng tốc không thể xảy ra về mặt lý thuyết (ví dụ: >100x) vượt quá giới hạn băng thông và tính toán vật lý của GPU. Điều này giúp bảo vệ khỏi những cạm bẫy "hack phần thưởng" tương tự như các trường hợp trong quá khứ, nơi các tác nhân tối ưu hóa cho các chỉ số của bộ khung thay vì đạt được hiệu suất thực sự. Nếu không có những ràng buộc này, việc tạo thêm kernel chủ yếu chỉ tạo ra thêm nhiễu.

Việc chú trọng vào bộ kiểm tra cũng làm thay đổi nút thắt cổ chai của quá trình tạo kernel bằng tác nhân. Trong công việc tìm kiếm chương trình đơn thuần (tức là tối ưu hóa lặp đi lặp lại nơi hệ thống tìm kiếm qua các chương trình bằng cách tạo ra các biến thể), các ứng viên tốt rất hiếm, vì vậy việc viết chúng chiếm phần lớn chi phí. Chúng tôi có thể tạo ra nhiều bản nháp song song, nhưng chúng tôi không thể bỏ qua bước xác thực. Chúng tôi phải xây dựng quy trình xác thực một cách cẩn thận, và việc kiểm tra phải chạy trên GPU thực, trong môi trường cô lập và thực hiện nhiều hơn một lần. Hệ thống chỉ có thể tiến nhanh như mức độ tin cậy của nó vào một kernel, chứ không phải nhanh như tốc độ nó viết ra một kernel.

Quản lý ngữ cảnh

Một thách thức khác là xác định những gì mô hình tạo kernel được phép nhìn thấy. Đây là một sự đánh đổi. Cung cấp cho mô hình một prompt lớn hơn đồng nghĩa với việc nó có nhiều thông tin hơn: kernel tốt nhất hiện tại, các thất bại gần đây, gợi ý từ profiler, ghi chú từ các lần chạy trước. Điều đó có thể hữu ích. Nhưng nó cũng tốn kém hơn vì chúng ta phải trả phí cho mỗi token mà mô hình đọc. Và khi prompt ngày càng dài, các lần thử tiếp theo dễ bị chệch hướng hơn. Các tín hiệu hữu ích bị trộn lẫn với những lời khuyên cũ kỹ, các mẹo mâu thuẫn và các chi tiết chỉ áp dụng cho một kích thước đầu vào hoặc một thao tác khác. Mô hình không phải lúc nào cũng biết nên tin vào câu nào, vì vậy nó thường làm theo những câu "to tiếng" nhất, hoặc làm theo tất cả một chút.

Nếu cung cấp quá ít thông tin, điều ngược lại sẽ xảy ra. Mọi nỗ lực đều bắt đầu từ con số không. Những ngõ cụt cũ lại xuất hiện. Không có gì được kế thừa từ lần chạy trước hoặc từ một thao tác liên quan, và vòng lặp không tiến triển.

Chúng tôi muốn có một lớp tri thức để hỗ trợ việc đó: ghi nhớ những gì hiệu quả, tái sử dụng sau này và thực hiện mà không cần sự can thiệp của con người. Lớp đó có một sự đánh đổi thứ hai, giữa mức độ chi tiết của bài học được lưu trữ và phạm vi áp dụng của nó.

Một ghi chú rất cụ thể (“trên kernel này, với kích thước đầu vào này, hãy unroll vòng lặp này”) có thể chính xác là những gì lần thử tiếp theo cần. Nhưng nó cũng dễ bị sử dụng sai cho thao tác tiếp theo, GPU tiếp theo hoặc kích thước đầu vào khác. Một ghi chú rất chung chung (“tận dụng bộ nhớ trên chip tốt hơn”) áp dụng cho hầu hết mọi nơi nhưng lại không chỉ cho mô hình biết phải làm gì. Chúng tôi đã thấy cả hai kiểu thất bại này. Khi các bài học quá chung chung, chúng chỉ nêu lại một thất bại mà không có hành động cụ thể. Khi chúng tôi lưu trữ chi tiết hơn, chúng thường quá gắn liền với một lần chạy cụ thể nên không giúp ích được cho lần sau. Trong một lần chạy dài, phần lớn những gì mô hình đọc và viết chỉ dành cho việc tìm nạp và định tuyến bộ nhớ đó thay vì viết kernel. Lớp bộ nhớ đang làm việc quá nhiều nhưng không làm cho ứng viên tiếp theo tốt hơn. Hình 2 cho thấy sự phân bổ chi phí token của một hệ thống như vậy. Chi phí bị chi phối bởi các lớp tri thức.

Phiên bản tri thức đáng giữ lại là phiên bản nhỏ hơn và tạo ra sự cân bằng giữa tính tổng quát và tính cụ thể. Khi mô hình chuẩn bị viết một kernel, prompt của nó chỉ nên bao gồm ngữ cảnh có độ tin cậy cao: những bài học có thể hành động được, ghép nối các tình huống cụ thể với các hành động (được chắt lọc từ việc ánh xạ sửa đổi-đến-tác động trong quá khứ) và các ghi chú thất bại ngắn gọn từ các lần chạy liên quan chặt chẽ. Được truy xuất thông qua lọc thẻ phân cấp kết hợp với tìm kiếm lai (từ khóa + ngữ nghĩa), các bài học phải đủ cụ thể để thực hiện và đủ phạm vi để làm rõ nơi nào chúng không áp dụng được. Các thao tác sâu hơn như tổ chức lại và chắt lọc thêm kho bài học nên thuộc về các công việc chạy nền, không phải là duyệt qua các lần chạy trước một cách đồng bộ trên mỗi lần thử. Nếu một bài học không thể nêu tên tình huống và hành động, nó không đáng để đưa vào prompt. Hình 3 cho thấy sự phân bổ chi phí token sau khi sửa lớp tri thức, và hầu hết các token được dành cho việc tạo ứng viên sau khi đã khắc phục.

Nghiên cứu điển hình: Gated DeltaNet packed decode

Một ví dụ cụ thể là kernel packed decode trên đường dẫn Gated DeltaNet trong Qwen 3.5 122B. Thao tác này cập nhật một trạng thái đệ quy và ghi đầu ra decode từ các đầu vào QKV đã đóng gói (packed), các tham số cổng (gate) và các chỉ số trạng thái. Chúng tôi đã sử dụng tác vụ này để thực hiện toàn bộ vòng lặp Proteus trên GPU NVIDIA B200 với backend Triton: xác thực hợp đồng tác vụ, đo lường bản thực thi tham chiếu, yêu cầu các tác nhân tạo kernel ứng viên, chạy các kiểm tra tĩnh và build, xác minh tính đúng đắn so với tham chiếu có kiểm soát, benchmark các ứng viên đã được xác minh, và sau đó đo lại các ứng viên tốt nhất.

Hình 4 đọc từ trái sang phải. Nút baseline neo benchmark ở mức 0,025 ms. Ứng viên 0000 là hạt giống an toàn: nó tái tạo cấu trúc packed-decode và vượt qua xác thực, nhưng nó chậm hơn tham chiếu, vì vậy Proteus giữ nó như một "cha mẹ" đã được đo lường thay vì coi đó là một chiến thắng. Từ đó, Proteus ngừng tối ưu hóa một kernel chung cho mọi hình dạng và chia việc tìm kiếm thành các đường dẫn cụ thể theo hình dạng.

Đường dẫn sửa lỗi Batch-1 đã tạo ra một kernel cụ thể theo hình dạng ở Ứng viên 012, đạt tốc độ 1,5x trên hình dạng decode đơn batch. Kết quả mạnh mẽ nhất đến từ đường dẫn serving-decode: Ứng viên 030 tìm thấy độ trễ kernel thấp nhất ở mức 0,018 ms, và Ứng viên 036 tạo ra mức tăng tốc hình dạng tốt nhất là 1,6x. Đoạn mã serving chiến thắng đó được chuyên biệt hóa cho layout Batch=4, Key=128, Value=128 và xử lý chiều giá trị theo các khối rộng 64, vì vậy nó là một kernel an toàn cho hình dạng cụ thể đó thay vì là một sự thay thế phổ quát.

Sự chuyển hướng cuối cùng trong dòng thời gian cho thấy lý do tại sao dấu vết (trace) lại quan trọng. Các lần tạo C++ sau đó (thay vì Triton) gặp phải lỗi build và tạo, và lần chạy dài kết thúc sau khi cạn kiệt ngân sách thử nghiệm của nhánh đó. Do đó, sản phẩm hữu ích không chỉ là ứng viên nhanh nhất. Đó là toàn bộ con đường được hiển thị trong hình: các lỗi ngữ nghĩa bị loại bỏ, các kernel đúng nhưng chậm hơn được đo lường, và các chiến thắng hiệu suất thực sự được giữ lại gắn liền với hình dạng giúp chúng an toàn để kết hợp thành một kernel sản xuất.

Chúng tôi đang làm gì tiếp theo

Vòng lặp tiến hóa mà chúng tôi xây dựng rất nghiêm ngặt với tư cách là người viết. Nó thường gọi mô hình theo một khuôn mẫu cố định: lấy kernel tốt nhất hiện tại, thử một chỉnh sửa nhỏ, kiểm tra, lặp lại. Điều đó lấy đi sự tự chủ mà tác nhân cần. Nó không thể dễ dàng thay đổi cấu trúc, chuyển đổi ngôn ngữ hoặc từ bỏ một thiết kế đã chết.

Vòng lặp vẫn là cần thiết. Không phải để viết kernel, mà để cung cấp cho tác nhân một gợi ý tiếp theo đáng tin cậy. Gợi ý đó phải đến từ hai nơi.

Thứ nhất, giao tiếp với lớp tri thức: một vài bài học đủ cụ thể để hành động và có phạm vi rõ ràng để chúng ta biết nơi nào chúng không áp dụng được. Nếu không có điều đó, mọi nỗ lực đều bắt đầu từ con số không.

Thứ hai, kết quả từ một bộ kiểm tra đáng tin cậy: tính đúng đắn và thời gian mà chính tác nhân không tự đo lường được. Những con số đó là gợi ý cho lần thử tiếp theo. Chúng cũng là những điểm số duy nhất chúng ta nên tin tưởng. Nếu tác nhân tự đo lường công việc của chính mình, chúng ta sẽ quay lại với các bộ nhớ đệm cũ, các so sánh không tương xứng và các bài kiểm tra mà nó có thể nhìn thấy.

Vì vậy, sự phân chia mà chúng tôi muốn hẹp hơn là "tác nhân so với vòng lặp". Hãy trao cho tác nhân quyền tự chủ về cách viết một kernel. Giữ vòng lặp làm kênh cho bộ nhớ và đánh giá. Tác nhân đề xuất. Vòng lặp trả về những gì nó được phép nhìn thấy và liệu đề xuất cuối cùng có thực sự chiến thắng hay không.

Proteus đã xây dựng các kernel chuyên biệt cho các phần của Qwen 3.5 122B trên đường dẫn Gated DeltaNet (một khối kiểu linear-attention) chạy trên GPU NVIDIA B200. Mức tăng tốc trên các kernel riêng lẻ nằm trong khoảng từ 1,8x đến 5,2x.

Bài học rút ra là việc tạo ra (generation) là bước rẻ tiền. Xác thực và quản lý ngữ cảnh mới là phần khó. Đó là nơi cần thời gian thiết kế cẩn thận và những đổi mới.

Việc tạo kernel GPU bằng tác nhân đã mở ra tiềm năng đáng kinh ngạc của sự chuyên biệt hóa cực độ, nhưng việc xây dựng các bộ khung đáng tin cậy, sẵn sàng cho sản xuất vẫn là một biên giới đầy thách thức. Chúng tôi đang giải quyết những thách thức khó khăn nhất tại điểm giao thoa giữa AI và hệ thống, và chúng tôi đang tìm kiếm những kỹ sư táo bạo để cùng chúng tôi định hình tương lai của suy luận hiệu quả. Nếu bạn đam mê vượt qua các giới hạn của những gì có thể, chúng tôi đang tuyển dụng!