CuTeDSL tại Perplexity

Cách chúng tôi sử dụng CuTeDSL để xây dựng các kernel GPU giúp các mô hình của chúng tôi đạt hiệu năng đỉnh cao trên GPU NVIDIA.

Tác giảPerplexity Research

Tại Perplexity, chúng tôi sử dụng Công cụ Phục vụ Tối ưu hóa Thời gian Chạy (ROSE) tự phát triển để phục vụ các mô hình từ embedding cho đến các LLM nghìn tỷ tham số. Nằm đằng sau tất cả các API, bao gồm Sonar, Tìm kiếm và Embedding, hệ thống thời gian chạy của chúng tôi được xây dựng nhằm trao quyền cho các nhà nghiên cứu khám phá các mô hình và kiến trúc phù hợp nhất, đồng thời nhanh chóng triển khai các ứng viên tốt nhất vào môi trường sản xuất. Để đạt được điều này, chúng tôi cần có khả năng triển khai, vận hành và tối ưu hóa quá trình suy luận cho các kiến trúc tối tân với tốc độ nhanh chóng. Trong bài viết blog này, chúng tôi khám phá cách các kernel CuTeDSL được tích hợp vào công cụ suy luận của mình cho phép chúng tôi xây dựng các kernel GPU cần thiết để đưa các mô hình đạt hiệu năng đỉnh cao trên GPU NVIDIA.

Suy luận tại Perplexity AI

Các mô hình tùy chỉnh, được xây dựng bởi các nhóm khác nhau tại Perplexity, được lưu trữ nội bộ trên GPU NVIDIA Hopper và Blackwell bằng cách sử dụng công cụ suy luận ROSE do nhóm Suy luận AI xây dựng. ROSE ban đầu được xây dựng để phục vụ tốt hơn các mô hình Llama tùy chỉnh cho cả việc giải mã ngôn ngữ và phân loại, với giao diện tương thích với NVIDIA Triton Inference Server. Theo thời gian, nó phát triển thành một công cụ suy luận toàn diện phục vụ các LLM đáng kể, cũng như các mô hình xếp hạng, phân loại, chấm điểm và các mô hình dựa trên transformer khác cho một loạt các nhóm nhằm cung cấp năng lượng cho các sản phẩm của Perplexity. Mục tiêu thiết kế của ROSE là sự đơn giản: theo thời gian, nỗ lực được đầu tư để đạt được hiệu năng đỉnh cao, nhưng ưu tiên chính là thích ứng với các kiến trúc và trường hợp sử dụng không ngừng phát triển cho các mô hình của chúng tôi.

Không gian cấu hình kernel trên các kiểu dữ liệu, số chiều đầu, hằng số thời gian biên dịch và cài đặt khởi chạy.

Ở mức độ cao, ROSE bao gồm một công cụ thích ứng mô hình với giao diện mà các client có thể tương tác. Các công cụ này có khả năng tái sử dụng cao, xử lý việc lên lịch và gộp nhóm yêu cầu, đồng thời truy xuất các token đã giải mã, embedding hoặc điểm số từ các mô hình cơ sở. Việc khởi tạo các thiết bị, các kênh giao tiếp giữa các tiến trình và giữa các nút, cùng với việc tải trọng số đều do chúng xử lý. Đối với các LLM, công cụ này chịu trách nhiệm gộp nhóm, chia nhỏ, lấy mẫu. Nó cũng chịu trách nhiệm phân bổ bộ nhớ KV, cho cả cơ chế chú ý toàn phần và tuyến tính, cũng như khớp tiền tố. Các công cụ embedding đơn giản hơn, vì chúng chỉ thực hiện gộp nhóm trực tuyến. Quá trình thực thi được tối ưu hóa sao cho công cụ chuẩn bị vòng công việc tiếp theo trong khi mô hình đang chạy song song trên bộ tăng tốc.

ROSE cung cấp một bộ sưu tập phong phú các lớp tùy chỉnh để tạo mô hình, bao bọc các kernel GPU được xây dựng bằng CuTeDSL, Triton, CUDA, CUTLASS và cuBLAS. Đối với các phép nhân ma trận và triển khai cơ chế chú ý, chúng tôi chủ yếu dựa vào các kernel NVIDIA từ CUTLASS và cuBLAS. Phần còn lại của pipeline, bao gồm embedding, chuẩn hóa, MoE và các kernel kích hoạt, là một phần của thư viện kernel của ROSE. Để hỗ trợ tốt nhất cho phần cứng mà chúng tôi chạy suy luận, ROSE được trang bị các nguyên thủy giao tiếp giữa các thiết bị của riêng mình cho việc định tuyến, phân phối và kết hợp MoE. Mặc dù các kernel khác tương đối đơn giản, chúng phải được chuyên biệt hóa một cách triệt để để đạt hiệu năng đỉnh cao trên tất cả các chiều ẩn và thiết bị khác nhau mà chúng chạy trên đó. Ngoài ra, chúng tôi dựa vào việc chuyên biệt hóa tại thời điểm biên dịch để đáp ứng tất cả các núm điều chỉnh và tinh chỉnh của các mô hình khác nhau (có trọng số bias hay không cho chuẩn hóa RMS, các hàm kích hoạt, v.v.). Do đó, chúng tôi quyết định sử dụng các DSL để triển khai chúng: ban đầu chúng tôi xây dựng ROSE bằng Triton. Tuy nhiên, kể từ đó chúng tôi đã áp dụng CuTeDSL làm môi trường chính cho lập trình GPU.

Tại sao lại là CuTeDSL?

Từ góc độ toán học, các kernel suy luận thường không phức tạp. Do đó, khi phát triển một kernel, độ phức tạp thường không nằm ở việc diễn đạt phép tính nào cần thực hiện, mà nằm ở cách chính xác phần cứng nên thực hiện nó như thế nào. Để xây dựng các kernel tốt nhất, chúng tôi yêu cầu một ngôn ngữ cung cấp các sự trừu tượng hóa phù hợp đối với phần cứng cơ sở. Ngôn ngữ phù hợp cho phép chúng tôi chỉ định các thao tác cấp cao để thực hiện phép tính, chuyên biệt hóa và tối ưu hóa chúng một cách mạnh mẽ trên tất cả các điểm trong một không gian cấu hình rộng.

CuTeDSL là một DSL dựa trên Python dành cho GPU NVIDIA được xây dựng trên nền tảng đại số bố cục CuTe và MLIR. Nó biên dịch tức thời (just-in-time) thành PTX được tối ưu hóa cao, đồng thời duy trì quyền kiểm soát đối với các nguyên thủy phần cứng cấp thấp. Nó nhằm mục đích cung cấp các tính năng và hiệu năng giống như CUTLASS C++, mà không có thời gian biên dịch tốn kém làm cản trở tốc độ phát triển. Trong số các ngôn ngữ mà chúng tôi đã khảo sát hoặc làm việc cùng (bao gồm CUDA, cuTile, TVM IR và Triton), chúng tôi coi đây là sự phù hợp nhất cho việc phát triển các kernel của mình.

Nhìn chung, mặc dù không ở cấp cao và ngắn gọn như Triton hay cuTile, chúng tôi thấy rằng Đại số Bố cục CuTe và các thành phần của nó (được cung cấp thông qua cute.Tensor) là những sự thay thế hợp lý. Ngoài ra, CuTeDSL cung cấp quyền truy cập không giới hạn vào các nguyên thủy phần cứng, mang lại cho chúng tôi mức độ kiểm soát cần thiết để đạt hiệu năng đỉnh cao trong mọi trường hợp ngoại lệ. Với CUDA, các cơ hội chuyên biệt hóa tại thời điểm biên dịch bị hạn chế, vì mỗi tham số bổ sung lại thêm một chi phí bất hợp lý vào thời gian biên dịch hoặc làm tăng độ phức tạp cho hệ thống xây dựng. Ngược lại, với một DSL được biên dịch tức thời, chúng tôi có thể chuyên biệt hóa mạnh mẽ hơn mà không hy sinh tốc độ lặp lại. Không giống như TVM IR, sử dụng CUDA làm bước trung gian, chúng tôi thấy pipeline dựa trên MLIR hạ cấp xuống PTX mạnh mẽ hơn và dễ gỡ lỗi hơn so với pipeline phát ra CUDA làm bước trung gian. Ngoài ra, với việc sử dụng các toán tử tùy chỉnh để bao bọc các lệnh PTX, chúng tôi nhận thấy rằng mình có thể tận dụng triệt để tất cả các tính năng của phần cứng cơ sở.

Trải nghiệm nhà phát triển

Phần lớn chi phí biên dịch của các kernel CUDA thường được tiêu tốn trong giai đoạn mở rộng template của các kernel, vì việc tối ưu hóa và biên dịch từ giai đoạn PTX trở đi có thể hoàn thành trong vài giây. Với nhiều tham số tại thời điểm biên dịch, giai đoạn biên dịch này có thể bùng nổ lên tới vài phút đối với một kernel không tầm thường. Nút thắt cổ chai này càng trở nên trầm trọng hơn bởi thực tế là các cơ hội biên dịch song song chỉ xuất hiện sau đó, ở cấp độ IR và PTX, trong khi việc xử lý template C++ diễn ra tuần tự. Mặc dù có các thủ thuật để giải quyết vấn đề này, bằng cách chuyển việc khởi tạo template sang các đơn vị biên dịch riêng biệt, chúng lại đòi hỏi những tinh chỉnh đáng kể đối với hệ thống xây dựng và việc sử dụng các macro cồng kềnh. Ngược lại, các kernel CuTeDSL được biên dịch Tức thời (JIT) không đòi hỏi các quy tắc xây dựng phức tạp và không cản trở thời gian biên dịch, vì chu kỳ phát triển có thể tập trung vào một điểm duy nhất trong không gian cấu hình, trong khi các bài kiểm tra đơn vị cẩn thận sau đó có thể bao phủ toàn bộ không gian thiết kế của kernel.

Mặc dù việc biên dịch JIT chuyển chi phí từ thời gian phát triển sang thời gian chạy, chi phí chung này hoàn toàn có thể chấp nhận được. Khi một mô hình được triển khai, nó chỉ kích hoạt một vài điểm từ toàn bộ không gian cấu hình, gây ra chi phí chung ở mức vài giây trong quá trình khởi động/làm nóng của công cụ suy luận. Ngược lại, việc biên dịch trước tất cả các cấu hình của một kernel có thể kéo dài tới vài phút. Ngoài ra, chi phí chung có thể được giảm thiểu hơn nữa bằng cách lưu bộ nhớ đệm (cache) các tạo phẩm được biên dịch JIT.

Chi phí biên dịch nặng nề của CUDA càng trở nên trầm trọng hơn bởi khó khăn trong việc truy tìm nguồn gốc của lỗi khi việc điều phối và chuyên biệt hóa dựa trên template được sử dụng, như trường hợp của CuTe và CUTLASS. Với một DSL, các lỗi có thể được chỉ ra chính xác dòng mã đã tạo ra nó, giúp việc gỡ lỗi trở nên dễ dàng hơn đáng kể. Ngoài ra, việc ánh xạ mã được tạo với câu lệnh DSL đã tạo ra nó thông qua các giai đoạn MLIR trung gian hoặc dựa vào thông tin gỡ lỗi đính kèm với mã PTX được tạo sẽ đơn giản hơn nhiều. Nhờ những khía cạnh này của DSL, chúng tôi thấy mình làm việc hiệu quả hơn rất nhiều khi phát triển các kernel.

Không gian cấu hình lớn

Mặc dù các kernel có nhiều điểm tương đồng giữa các mô hình, chúng vẫn phải được trang bị một loạt các núm cấu hình có thể ảnh hưởng đến hiệu năng. Các tham số có liên quan bao gồm:

Chiều ẩn / Chiều đầu / Số lượng đầu chú ý

Kiểu dữ liệu đầu vào/đầu ra

Lược đồ lượng tử hóa

Hằng số tại thời điểm biên dịch (bias, hàm kích hoạt, v.v.)

Kiến trúc ROSE, bao gồm engine, mô hình, thư viện kernel và dịch vụ.

Chuyên biệt hóa tại thời điểm biên dịch rất hữu ích cho cả hiệu năng và việc gỡ lỗi. Đặc biệt trong một bộ xử lý theo dây chuyền (pipelined), như trường hợp của GPU, một nhánh không tồn tại sẽ tốt hơn là một nhánh luôn được thực hiện hoặc không bao giờ được thực hiện. Ngoài ra, việc biên dịch một kernel cho một tập hợp các núm cấu hình cụ thể sẽ tạo ra ít lệnh PTX/SASS hơn, dễ phân tích và gỡ lỗi hơn khi sử dụng cuda-gdb trong quá trình điều tra.

Đối với một kích thước lô (batch size) và cấu hình cho trước, lưu lượng dự kiến đến DRAM và số lượng các phép tính số học có thể được ước tính. Sau đó, các quyết định tối ưu có thể được đưa ra cho từng trường hợp cụ thể để phân bổ lưới, đồng bộ hóa và sử dụng các tính năng phần cứng nhất định một cách tối ưu. CuTeDSL cho phép chúng tôi định nghĩa và tái sử dụng các thành phần để chuyên biệt hóa sâu mà không phát sinh chi phí biên dịch quá mức không thể chấp nhận được.

Chuyên môn hóa điền trước và giải mã

Nói chung, hiệu năng của các kernel suy luận phải được tối ưu hóa cho hai trường hợp sử dụng riêng biệt:

Chuẩn bị (Prefill): kích thước lô nhỏ (ít khối, chuỗi), nhiều token cho mỗi chuỗi. Thông lượng cần được tối ưu hóa cho cả embedding và các token đầu vào của các LLM. Độ trễ trong khoảng hàng chục đến hàng trăm mili giây.

Giải mã (Decode): kích thước lô từ nhỏ đến lớn, từ một (không giải mã dự đoán) đến một vài token cho mỗi chuỗi (dự đoán mục tiêu nháp, dự đoán đa token). Độ trễ gần bằng độ trễ khởi chạy kernel.

Đối với GEMM và attention, giai đoạn chuẩn bị chạy trên một số lượng lớn token và hiệu năng bị giới hạn bởi thông lượng tính toán của thiết bị. Quá trình giải mã liên quan đến cùng một thao tác trên kích thước lô nhỏ, với thời gian thực thi bị chi phối bởi chi phí tải các kích hoạt và trọng số từ DRAM vào các bộ đa xử lý dòng (streaming multiprocessors) và các lõi tensor, thay vào đó bị giới hạn bởi băng thông bộ nhớ của thiết bị. Mặc dù kích thước lô hoặc tổng số token trong một lô là một chỉ báo tương đối rõ ràng, ngưỡng chuyển đổi giữa chuẩn bị và giải mã tinh tế hơn vì băng thông bộ nhớ cũng được quyết định bởi chiều ẩn của mô hình và kiểu dữ liệu của kích hoạt. Do đó, các kernel được tách thành các bản triển khai chuẩn bị và giải mã nếu sự khác biệt về thuật toán hoặc việc gán lưới/khối mang lại lợi thế cho một trong hai.

Tham số khác biệt phổ biến giữa các kernel chuẩn bị và giải mã là số lượng warp trong một khối (số lượng chuỗi, chia cho kích thước warp là 32 chuỗi cho mỗi warp). Việc biến tham số này thành một hằng số tại thời điểm biên dịch là điều đáng mong đợi, vì nó có thể hỗ trợ phân bổ thanh ghi. Trong CUDA, điều này thường được thực hiện bằng cách chỉ định số lượng chuỗi thông qua __launch_bounds__ và cung cấp số lượng warp như một tham số template:

cpp
template <size_t NUM_WARPS> __global__ void __launch_bounds__(NUM_WARPS * WARP_SIZE, 1) some_kernel(...) { }

Số lượng thanh ghi khả dụng cho tất cả các chuỗi trong một khối là một giới hạn phần cứng, 64K trên Hopper và Blackwell. Nếu số lượng chuỗi được biết trước, trình biên dịch biết chính xác số lượng thanh ghi nó có thể gán cho mỗi chuỗi và có thể biên dịch thành công các kernel được tối ưu hóa mạnh mẽ hơn. Tuy nhiên, điều này đi kèm với chi phí tại thời điểm biên dịch, vì bây giờ kernel phải được biên dịch lại cho các kích thước warp khác nhau. Với CuTeDSL, đây không còn là vấn đề nữa, vì chúng tôi không cần phải trả chi phí biên dịch cho toàn bộ không gian cấu hình trước. DSL cũng mong đợi kích thước khối là một biểu thức hằng số, tự động thực hiện việc tối ưu hóa được ngụ ý bởi __latch_bounds__ trong CUDA.

Bên cạnh việc làm cho số lượng warp trở thành một hằng số, việc thay đổi nó giữa quá trình chuẩn bị và giải mã là điều nên làm. Để tận dụng triệt để băng thông bộ nhớ của GPU, việc luôn gán số lượng warp cao cho mỗi khối là chưa đủ, vì băng thông tối đa chỉ có thể đạt được bằng cách sử dụng càng nhiều Bộ đa xử lý dòng (SM) khả dụng càng tốt. Do đó, đối với một số lượng token nhất định cùng với các kiểu dữ liệu và chiều ẩn đã biết, trước tiên chúng tôi chọn một chiều vectơ trên mỗi chuỗi lý tưởng (8 cho bfloat16, 4 cho float32) để đảm bảo rằng mã được tạo sử dụng các lệnh tải/lưu 128-bit (256-bit trên Blackwell). Biết được kích thước vectơ, chúng tôi có thể tính toán số lượng warp mà kernel yêu cầu. Trước tiên, chúng tôi cố gắng tối đa hóa số lượng SM: nếu số lượng warp ít hơn số lượng SM, một warp duy nhất cho mỗi khối sẽ được chọn. Đây thường là biến thể giải mã của kernel. Nếu có đủ warp, chúng tôi dần dần tăng số lượng warp lên 8, 16 hoặc 32, thực hiện benchmark giới hạn trên sau đó chúng tôi thường không thấy hiệu năng cải thiện thêm nữa.

Mặc dù Triton đã thực hiện một số tối ưu hóa này bên dưới lớp vỏ, chúng tôi nhận thấy rằng mình có thể đạt được kết quả tương tự bằng cách sử dụng cẩn thận đại số bố cục CuTe, mà không bị ràng buộc bởi những thiếu sót của một trình biên dịch cấp cao. Đặc biệt khi có sự tham gia của các phép thu gọn ngang dưới mức warp (sub-warp horizontal reductions), chúng tôi thấy rằng CuTeDSL cung cấp cho chúng tôi khả năng kiểm soát chi tiết cần thiết để tạo ra mã tối ưu.

Đồng bộ hóa lưới

Trong một số tình huống, để phân phối công việc trên càng nhiều khối càng tốt, cooperative_groups::this_grid().sync() có thể được sử dụng để đồng bộ hóa tất cả các khối tham gia. Các hạn chế được áp dụng: vì một khối không thể bị hủy lập lịch khỏi một SM một khi nó bắt đầu thực thi, để đảm bảo tiến độ chuyển tiếp, số lượng khối không thể vượt quá số lượng SM khả dụng trên thiết bị, nhằm đảm bảo rằng tại bất kỳ thời điểm nào tất cả các khối đều có thể cập nhật và thăm dò các vị trí bộ nhớ toàn cục tham gia vào quá trình đồng bộ hóa. Hạn chế này được thực thi bởi cudaLaunchCooperativeKernel. Chi phí của một rào cản lưới là khoảng 3us trên Hopper, cộng với chi phí xả tất cả các thao tác bộ nhớ trước đó thông qua một rào cản bộ nhớ ngầm định trong một số tình huống. Nhìn chung, điều này có thể so sánh với chi phí chung của việc tách một kernel thành hai lần khởi chạy, dẫn đến một quyết định thiết kế nhạy cảm về hiệu năng. Mặc dù thực tế là Khởi chạy phụ thuộc theo chương trình (Programmatic Dependent Launch - PDL) có thể làm giảm chi phí khởi chạy từ dòng Hopper trở đi, đối với các kernel giải mã, việc tránh chi phí chung này là rất đáng giá.

Nếu một trong các khối lượng công việc yêu cầu nhiều khối hơn số lượng SM, kernel nên được tách ra, dựa vào rào cản ngầm định của việc khởi chạy kernel. Đây chủ yếu là trường hợp dành cho quá trình chuẩn bị. Nếu số lượng khối lý tưởng cho tất cả các khối lượng công việc nằm vừa trong lưới, rào cản lưới sẽ được ưu tiên, như trường hợp của các kernel giải mã hoạt động trên số lượng token nhỏ hơn.

Trong CuTeDSL, những quyết định này có thể được đưa ra trong mã khởi chạy phía máy chủ (host-side) và các kernel có thể được soạn thảo và đan xen với các rào cản lưới theo nhu cầu, như minh họa trong ví dụ dưới đây. Với chiến lược lai như vậy, các kernel giải mã có thể được tối ưu hóa để đạt độ trễ tối thiểu và các kernel chuẩn bị có thể được tối ưu hóa cho thông lượng tối đa. Không giống như các giải pháp dựa trên CUDA, logic này có thể được định nghĩa ngắn gọn hơn và ít chi phí chung tại thời điểm biên dịch hơn.

python
@cute.kernel def dispatch_kernel_decode(...): route_tokens(...) grid_barrier(grid_barrier_ptr) copy_tokens_to_experts(...) @cute.kernel def dispatch_kernel_prefill_route(...): route_tokens(...) @cute.kernel def dispatch_kernel_prefill_copy(...): copy_tokens_to_experts(...) @cute.jit def dispatch(...): num_route_blocks = ... num_copy_blocks = ... if num_copy_blocks < NUM_SMS: dispatch_kernel_decode(...).launch( grid=(min(max(num_route_blocks, num_copy_blocks), NUM_SMS), 1, 1) ) else: dispach_kernel_prefill_route(...).launch( grid=(num_route_blocks, 1, 1) ) dispatch_kernel_prefill_copy(...).launch( grid=(num_copy_blocks, 1, 1) )

Trong trường hợp các kernel MoE all-to-all của chúng tôi, như được chi tiết bên dưới, việc có sự linh hoạt để xây dựng các kernel có hoặc không có rào cản lưới và tối ưu hóa chúng tại thời điểm biên dịch đã giúp chúng tôi đạt được những cải thiện hiệu năng đáng kể thông qua CuTeDSL.

Kernel CuTeDSL

Chuẩn hóa QK (QK Norm)

Kernel đầu tiên chúng tôi xây dựng bằng CuTeDSL là một kernel chuẩn hóa RMS chuyên biệt, được tối ưu hóa cho các chiều nhỏ. Chuẩn hóa QK, được sử dụng trong lớp tự chú ý của các mô hình Qwen gần đây, được áp dụng trên cơ sở từng đầu (per-head), trên 64, 128 hoặc 256 phần tử. Không giống như chuẩn hóa RMS thông thường, hoạt động trên các chiều ẩn vượt quá hàng nghìn, đòi hỏi một phép thu gọn ngang (horizontal reduction) qua cả bộ nhớ chia sẻ và thanh ghi, chuẩn hóa QK chỉ cần thu gọn trên các warp, với khả năng dồn nhiều đầu vào một warp duy nhất. Điều này đã trở thành một trường hợp sử dụng nơi chúng tôi có thể tối ưu hóa mạnh mẽ tại thời điểm biên dịch để tìm ra các phân bổ token cho warp tốt nhất cho mọi chiều đầu và kiểu dữ liệu.

Chuẩn hóa RMS là một thao tác khá đơn giản, mặc dù nó đòi hỏi một phép cộng ngang để tính bình phương chuẩn của một vectơ:

cpp
r = w * (x * ((x * x).sum(dim=-1) / x.shape[-1] + EPS))

Vì các đầu chú ý thường nhỏ và các thao tác trên mỗi chuỗi được vector hóa mạnh mẽ, một warp duy nhất có thể tính toán bình phương chuẩn. Để tối đa hóa băng thông bộ nhớ, mỗi chuỗi nên xử lý 128 bit dữ liệu đầu vào, tương ứng với 8 giá trị bfloat16. Do đó, một nửa warp là đủ để xử lý một đầu và bình phương chuẩn có thể được xác định bằng cách sử dụng một phép thu gọn ngang được thực hiện bởi shuffle_sync_bfly. Bắt đầu từ chiều đầu và kiểu dữ liệu, chúng tôi có thể xác định số lượng đầu trên mỗi warp và chỉ định các đối số cũng như số vòng thu gọn tại thời điểm biên dịch.

Sau đó, chúng tôi có thể xây dựng một kernel chuyên biệt với số dòng mã tương đối ít, dựa vào các lệnh tải/lưu được vector hóa tự động và sử dụng các thao tác đặc thù cho kiến trúc cấp thấp hơn cho việc thu gọn warp. Kernel xác định đầu nào (Q hoặc K) để làm việc và xác định các trọng số tương ứng để áp dụng. Nó tải toàn bộ một vectơ từ đầu và trước tiên cộng dồn cục bộ. Sau đó, dựa trên số lượng chuỗi cho mỗi đầu, bình phương chuẩn được tính toán. Cuối cùng, chuẩn hóa RMS, được vector hóa tự động, được xác định và ghi ngược lại vào bộ nhớ. Các tham số cụ thể của mô hình, chẳng hạn như WEIGHT_BIAS và EPS, được gấp hằng số (constant-folded).

python
@cute.kernel def _qk_norm_kernel( q: cute.Tensor, k: cute.Tensor, q_weight: cute.Tensor, k_weight: cute.Tensor, WEIGHT_BIAS: cutlass.Constexpr[float], EPS: cutlass.Constexpr[float], ) -> None: # Find the block and thread indices. bidx, _, _ = cute.arch.block_idx() bdim, _, _ = cute.arch.block_dim() tidx, _, _ = cute.arch.thread_idx() _, num_q_heads, (threads_per_head, vecs_per_thread, vec_size) = q.shape _, num_k_heads, _ = k.shape head_dim = threads_per_head * vecs_per_thread * vec_size # Determine the thread and the vector to work on. vec_id = (bidx * bdim + tidx) // threads_per_head thread_id = tidx % threads_per_head # Determine the token and head indices. token = vec_id // (num_q_heads + num_k_heads) head_offset = vec_id % (num_q_heads + num_k_heads) is_q = head_offset < num_q_heads # Determine which tensor to work with. thread_coord = (thread_id, None, None) head_coord = (token, head_offset if is_q else head_offset - num_q_heads, thread_coord) # Load the head into registers. x_ref = q[head_coord] if is_q else k[head_coord] w = (q_weight if is_q else k_weight)[thread_coord].load() # Compute the squared sum locally. x = x_ref.load() x_f32 = x.to(cute.Float32) x_square = x_f32 * x_f32 x_square_sum = x_square.reduce(cute.ReductionOp.ADD, init_val=0.0) # Sum up across the warp. mask = activemask() start = cute.arch.WARP_SIZE // (cute.arch.WARP_SIZE // threads_per_head * 2) for offset in [start >> i for i in range(5) if (start >> i) > 0]: x_square_sum += cute.arch.shuffle_sync_bfly( x_square_sum, offset=offset, mask=mask, ) # RMS norm. rms = cute.math.rsqrt(x_square_sum / head_dim + EPS) x_norm = ((w + WEIGHT_BIAS) * (x * rms)).to(x.element_type) x_ref.store(x_norm)

So với một kernel chuẩn hóa RMS thông thường, biến thể chuyên biệt này mang lại thông lượng và độ trễ tốt hơn gấp 2-3 lần cho cả quá trình chuẩn bị và giải mã.

Phân phối / Kết hợp Hỗn hợp Chuyên gia (Mixture-of-Experts Dispatch/Combine)

Việc phân phối và kết hợp MoE xáo trộn các token đầu vào đến các chuyên gia mà chúng được định tuyến tới, thực hiện việc tích lũy có trọng số sau khi các kernel MoE GEMM thực hiện các phép nhân ma trận có liên quan. Hiệu năng là mối quan tâm đặc biệt, vì các sản phẩm khác nhau phụ thuộc vào các mô hình MoE có kích thước khác nhau. Dựa trên các loại triển khai, chúng tôi phân biệt các kernel khác nhau để tối ưu hóa:

TP=1: Kernel cục bộ cho một GPU duy nhất. Việc phân phối xáo trộn các token đến các chuyên gia có liên quan và việc kết hợp tích lũy chúng. Được sử dụng chủ yếu cho việc triển khai Qwen3 30B trên H100.

DP=1, TP ≠ 1: Kernel song song tensor, đơn nút (singlenode). Đầu vào được nhân bản, việc phân phối chọn các token cho cấp bậc hiện tại và việc kết hợp tích lũy chúng trên tất cả các cấp bậc. Tích lũy qua NVLink hoặc NVLS. Được sử dụng cho các bộ chuẩn bị lớn.

DP ≠ 1, NVLink: Kernel song song dữ liệu với tùy chọn song song tensor. Mỗi cấp bậc gửi token đến từng cấp bậc khác lưu trữ các chuyên gia cụ thể. Thông tin định tuyến được tổng hợp, các token được truyền qua NVIDIA NVLink. Được sử dụng cho việc triển khai H200 hoặc GB200 song song dữ liệu.

DP ≠ 1, InfiniBand: Kernel song song dữ liệu và song song tensor Điểm tới Điểm (Peer-to-Peer) hoạt động trên cả NVLink và NVIDIA Quantum InfiniBand. Được sử dụng cho việc triển khai Hopper đa nút.

Cục bộ (Local)

Kernel cục bộ đếm số lượng token được gán cho mỗi chuyên gia dựa trên bảng định tuyến và gán một chỉ mục duy nhất cho mỗi token trong danh sách các token thuộc về mỗi chuyên gia. Sau bước định tuyến, các token được sao chép vào các bộ đệm có thể được đệm tùy chọn để chuyển tiếp đến kernel MoE GEMM. Sử dụng các vị trí token, kernel kết hợp sau đó thực hiện việc tích lũy trong các thanh ghi, ghi các token cuối cùng vào một bộ đệm đầu ra.

Giai đoạn phân phối thay đổi giữa quá trình chuẩn bị và giải mã: đối với số lượng token ít, có ít thông tin định tuyến để đọc ra, vì vậy chúng tôi xử lý nó trong một khối duy nhất và báo hiệu phần còn lại của lưới, với các khối riêng lẻ sao chép các khối token. Phân bổ token được xác định bằng cách tăng nguyên tử các bộ đếm trong bộ nhớ chia sẻ. Vì quá trình chuẩn bị phải xử lý nhiều thông tin định tuyến hơn, việc định tuyến được chia nhỏ trên nhiều khối trong một kernel riêng biệt. Các chỉ mục trước tiên được gán cục bộ bằng cách sử dụng các bộ đếm trong bộ nhớ chia sẻ và sau đó được tổng hợp trên tất cả các khối trong bộ nhớ toàn cục để xác định vị trí của từng token. Lưới của kernel này được kẹp (clamp) ở số lượng SM để có thể sử dụng các rào cản theo cấp độ lưới (grid-wise barriers). Một kernel tiếp theo thực hiện các bản sao, được khởi chạy trên số lượng lớn hơn nhiều các khối và warp để tận dụng triệt để băng thông bộ nhớ khả dụng.

Kernel kết hợp tương đối đơn giản, vì nó thực hiện tích lũy được vector hóa, dựa trên thông tin định tuyến và các khoảng bù (offset) được tính toán bởi việc phân phối. CuTeDSL tỏ ra có giá trị ở đây, vì phép tính có thể được biểu diễn thông qua các tensor CuTe đơn giản biên dịch thành mã hiệu quả cho tất cả các kiểu dữ liệu và chiều ẩn.

Đơn nút (Singlenode)

Trong một triển khai song song tensor, thông tin định tuyến và các token đầu vào được nhân bản. Kernel phân phối cục bộ có thể được tái sử dụng hoàn toàn với một thay đổi nhỏ để định tuyến và chỉ chọn các token được gán cho các chuyên gia đặt tại cấp bậc hiện tại. Với một số tham số và chuyên biệt hóa tại thời điểm biên dịch, các kernel phân phối cục bộ được tái sử dụng.

Vì số lượng cấp bậc thường thấp hơn số lượng chuyên gia mà một token được định tuyến tới, trong trường hợp trung bình, có thể giả định rằng mỗi token có ít nhất một trong các bản sao của nó được định tuyến đến mỗi cấp bậc. Trong tình huống này, kernel kết hợp cục bộ có thể được tái sử dụng để tổng hợp các bản sao cục bộ của một token, tiếp theo là một phép all-reduce để cộng dồn trên các cấp bậc tham gia. Trong trường hợp của kernel chuẩn bị, chúng tôi khởi chạy các kernel riêng biệt, tuy nhiên quá trình giải mã kết hợp cả hai thao tác thành một lần khởi chạy duy nhất.

Đối với all-reduce, chúng tôi sử dụng NVLink thông thường nếu chỉ có hai cấp bậc tham gia, chuyển sang NVLS khi có từ 4 cấp bậc trở lên tham gia để tăng thông lượng thêm 10-20%. Nhờ có DSL, chúng tôi có thể chia sẻ phần lớn logic thu gọn và các trao đổi bộ đệm ngang hàng giữa MoE và các kernel all-reduce thông thường. Đây là một tình huống mà việc kiểm soát chặt chẽ phần cứng mang lại lợi ích, vì các thao tác DSL trên multimem.ld_reduce và multimem.st là bắt buộc đối với NVLS. Việc triển khai bản thân quá trình thu gọn phần lớn tuân theo ví dụ all_reduce từ các bài kiểm tra NCCL.

Song song dữ liệu (Data Parallel)

Các kernel Song song dữ liệu dựa vào NVLink để tổng hợp token cho nhiều thiết bị. Trong giai đoạn phân phối, trước khi trao đổi token, các bên ngang hàng trước tiên định tuyến cục bộ để xác định khoảng bù mục tiêu trên mỗi cấp bậc cho mỗi token mà họ đang gửi. Sau đó, họ bắt đầu sao chép các token vào các bộ đệm nhận được phân bổ cho từng bên gửi để tránh xung đột. Các lệnh ghi release-store trên mỗi token được sử dụng để báo hiệu hoàn thành. Ở phía người nhận, thông tin định tuyến được tổng hợp để có thể cô đọng tất cả các token thành một bộ đệm duy nhất cho MoE GEMM. Người nhận thăm dò cờ do bên ngang hàng ghi bằng lệnh acquire-load để đợi các lệnh ghi hoàn tất.

Ở phía kết hợp, các token được sao chép qua NVLink về các cấp bậc nguồn, với một cờ trên mỗi token biểu thị việc đã nhận. Kernel kết hợp tương tự như các biến thể cục bộ và đơn nút, với những khác biệt nhỏ về cách giải quyết địa chỉ và một số logic đồng bộ hóa bổ sung.

Chúng tôi đã xây dựng một thư viện các wrapper PTX để cung cấp các lệnh tải và lưu release-acquire có phạm vi ở cả cấp độ hệ thống và cấp độ thiết bị để triển khai kernel này. Bằng cách cải thiện độ chi tiết của các cờ đồng bộ hóa và phân bổ nhiều khối hơn để gửi và nhận token qua NVLink, chúng tôi đã đạt được mức cải thiện gấp 2 lần so với các kernel CUDA trước đây của mình.

Điểm tới Điểm (Peer-to-Peer)

Các kernel Điểm tới Điểm của chúng tôi trước đây đã đạt được hiệu năng tối tân trên GPU Hopper được ghép nối với các bộ điều hợp ConnectX-7 InfiniBand. Kể từ đó, chúng tôi đã chuyển các kernel kết hợp sang CuTeDSL để đạt được hiệu năng cao hơn nữa. Vì các phiên bản CUDA đã phụ thuộc rất nhiều vào chuyên biệt hóa tại thời điểm biên dịch, chúng đã chứng minh là các nút thắt cổ chai về thời gian biên dịch trong kho lưu trữ của chúng tôi, buộc chúng tôi phải bỏ qua một số tham số. Bằng cách chuyển sang DSL được biên dịch JIT, chúng tôi đã chuyên biệt hóa cho nhiều tham số hơn và tinh chỉnh phân bổ lưới, giới thiệu thêm những khác biệt giữa chuẩn bị và giải mã để đạt được mức tăng độ trễ khoảng 3%.

Hỗ trợ Blackwell

Từ góc độ của các kernel ROSE, việc hỗ trợ Blackwell được chia thành hai loại: các kernel bất khả tri về kiến trúc (architecture-agnostic) và đặc thù cho kiến trúc. Các kernel đặc thù cho kiến trúc, phụ thuộc vào các lõi Tensor, đòi hỏi công việc rõ ràng để thích ứng. Tuy nhiên, nhờ việc sử dụng DSL, các kernel bất khả tri về kiến trúc thích ứng một cách dễ dàng.

Một GPU Blackwell có tới 160 SM, trái ngược với 132 của Hopper. Đối với hầu hết các kernel, đây là một hằng số tại thời điểm biên dịch và logic trong mã trình khởi chạy phía máy chủ tính đến yếu tố này để xác định kernel tối ưu (chuẩn bị so với giải mã) để chạy. Ngoài ra, Blackwell cũng giới thiệu các lệnh tải 256-bit, có thể dễ dàng được phát ra bằng cách tăng chiều vectơ (chiều cuối cùng của các tensor CuTe) để cải thiện việc sử dụng băng thông bộ nhớ.

Blackwell giới thiệu hỗ trợ phần cứng cho một số lượng đáng kể các lược đồ lượng tử hóa (MXFP8, MXFP4, NVFP4). Hỗ trợ cho những tính năng này chủ yếu được triển khai bằng cách tinh chỉnh các kernel hiện có cho việc chia tỷ lệ theo khối FP8. Hỗ trợ JIT là rất quan trọng ở đây, vì các chế độ mới có thể được hỗ trợ mà không làm bùng nổ không gian cấu hình tại thời điểm biên dịch. Trong quá khứ, chúng tôi thấy Triton lúng túng trong việc biểu diễn các phép thu gọn ngang trên một vài phần tử, trên một nửa hoặc một phần tư của một warp. Việc kiểm soát chi tiết hơn đối với các phép thu gọn là rất quan trọng để đạt được hiệu năng tối ưu khi tính toán giá trị cực đại tuyệt đối trên 16 hoặc 32 phần tử cho NVFP4 và MXFP8/MXFP4 tương ứng.

Tạo Kernel được hỗ trợ bởi AI

Là một ngôn ngữ tương đối mới, các LLM ít được tiếp xúc với mã CuTeDSL trong quá trình huấn luyện trước (pre-training) của chúng. Ngoài ra, phong cách mã hiện tại của chúng tôi, với các kernel được tối ưu hóa cho các kích hoạt và chuẩn hóa, hơi khác so với các ví dụ CUTLASS hoặc Quack hiện có, vốn được tối ưu hóa mạnh mẽ cho TMA và các lệnh nhân ma trận đặc thù cho kiến trúc. Đạt đến thời điểm mà các LLM có thể hỗ trợ viết kernel là một quá trình dần dần.

Chúng tôi đã thủ công xây dựng một tập hợp ban đầu các kernel và hàm tiện ích, theo các mẫu nhất quán, trước khi có thể dựa vào các tác nhân (agents) để chuyển đổi kernel. Khi tập hợp này tiến triển, chúng tôi có thể thực hiện các bản dịch thô từ CUDA sang CuTeDSL, với việc làm lại thủ công đáng kể theo sau. Đặc biệt, chúng tôi thấy rằng các mô hình gặp khó khăn với các chuyển đổi con trỏ sang tensor và việc vector hóa, bị mắc kẹt khi cố gắng sửa các lỗi kiểu dữ liệu khi hướng tới việc biên dịch một bài kiểm tra. Chúng thể hiện sự hiểu biết rất ít về bản chất dựa trên lưới của các kernel, nhầm lẫn giữa đồng bộ hóa chuỗi và lưới. Đáng thất vọng nhất là có sự cám dỗ muốn tạo ảo giác các bản dịch không tồn tại của các thư viện đặc thù cho CUDA thay vì xây dựng dựa trên các bản thay thế CuTeDSL của riêng chúng tôi. Theo thời gian, khi tập hợp các kernel CuTeDSL phát triển, chúng tôi thấy rằng các LLM có thể xây dựng phần lớn mã khung (skeleton code) và điền vào các chỗ trống, cho phép chúng tôi tập trung vào thiết kế cấp cao.

Kết luận và Công việc trong tương lai

So với CUDA, chúng tôi thấy rằng CuTeDSL mang lại hiệu năng tại thời điểm biên dịch tốt hơn nhiều, cho phép chúng tôi chuyên biệt hóa cho nhiều tham số hơn và đạt được những cải thiện về hiệu năng. Mặc dù các kernel có phần dài dòng hơn do thiếu một số sự trừu tượng hóa so với Triton, chúng tôi thấy rằng độ phức tạp là có lý vì khả năng kiểm soát chi tiết hơn đối với phần cứng cho phép chúng tôi xây dựng các kernel nhanh hơn.

Tại thời điểm viết bài này, khoảng cách hiệu năng giữa CUTLASS C++ và các kernel CuTeDSL là tối thiểu, mở ra những cơ hội thú vị để soạn thảo và kết hợp các kernel GEMM hiệu năng cao với giao tiếp giữa các thiết bị, chuẩn hóa hoặc kích hoạt. Bằng cách cam kết với DSL này, chúng tôi sẽ có thể trích xuất thêm hiệu năng từ các thiết bị của mình để cải thiện thông lượng trên tất cả các khối lượng công việc, từ embedding đến các LLM.

Tài liệu tham khảo

https://github.com/Dao-AILab/quack

https://github.com/NVIDIA/cutlass

https://github.com/NVIDIA/nccl-tests

https://docs.nvidia.com/cuda/parallel-thread-execution/

Trích dẫn

BibTeX
@online{perplexity2026CutedslAtPerplexity,
  author = {{Perplexity} and Nandor Licker and Bex Brooker and Jesse Dwyer and Phi Hoang},
  title = {{CuTeDSL at Perplexity}},
  date = {2026-05-06},
  url = {https://www.perplexity.ai/hub/blog/cutedsl-at-perplexity}
}