NVIDIA dùng AI Agent chuyển đổi mã cuTile Python và Triton-TileIR sang Rust

cuTile Rust (cutile-rs) là một hệ thống dựa trên tile (tile-based) giúp viết các GPU kernel an toàn và chuẩn hóa (idiomatic) bằng ngôn ngữ lập trình Rust. Bằng cách mở rộng mô hình sở hữu (ownership model) của Rust cho các GPU kernel dạng tile, nó chia các đầu ra có thể thay đổi (mutable outputs) thành các phần riêng biệt và bảo toàn hợp đồng sở hữu phía host trong suốt quá trình khởi chạy kernel. Nó cũng cho phép lập trình viên chủ động bỏ qua các cơ chế này khi cần quyền kiểm soát cấp thấp hơn, cho phép thực thi trực tiếp các thao tác Tile IR.

Thư viện TileGym CUDA tile kernel đã tích lũy được một kho lớn các kernel sản xuất được viết bằng CUDA Tile Python (cuTile Python) và Triton-TileIR (nvtriton). Để tất cả các kernel này cũng có thể sử dụng được trong Rust, đội ngũ của chúng tôi đã xây dựng một kỹ năng AI (AI agent skill) có khả năng chuyển đổi các kernel từ cuTile Python và Triton-TileIR sang cuTile Rust.

Sử dụng kỹ năng này, chúng tôi đã chuyển đổi toàn bộ 24 toán tử công khai của TileGym sang cuTile Rust và đạt hiệu suất trung bình bằng 99,5% so với cuTile Python. Chúng bao gồm khoảng 40 GPU kernel, từ các thao tác phần tử (element-wise) đến flash-attention decode, Multi-head Latent Attention (MLA) và các mô hình mixture-of-experts (MoE). Lưu ý rằng một số toán tử cần nhiều biến thể kernel khác nhau.

Mỗi quá trình chuyển đổi bắt đầu từ bất kỳ bản triển khai tham chiếu nào mà toán tử đó có (cuTile Python hoặc Triton-TileIR) và chạy qua một quy trình đa tác nhân (multi-agent pipeline) có giới hạn, bao gồm phân tích, device kernel, mã host và FFI, cùng với việc đo kiểm hiệu năng (benchmarking). Mỗi giai đoạn kết thúc bằng một kết luận có thể kiểm chứng bằng máy, với các tập lệnh xác thực và các tệp diff Tile IR quyết định xem quá trình chuyển đổi có được tiếp tục hay không. Thách thức chính là việc biên dịch JIT của cuTile Python chuyên biệt hóa từng kernel một cách ngầm định tại thời điểm gọi, trong khi Rust yêu cầu bạn phải khai báo mọi sự chuyên biệt hóa trong chữ ký (signature) của kernel.

Bài viết này giải thích cách chúng tôi phát triển quy trình làm việc đa tác nhân để chuyển đổi các kernel cuTile Python và Triton-TileIR sang cuTile Rust, với các bước kiểm tra tính đúng đắn và hiệu năng ở từng giai đoạn. Bài viết bao gồm việc khoảng cách giữa các phiên bản trông như thế nào trong một kernel thực tế, cách cấu trúc kỹ năng để không giai đoạn nào phải dựa trên sự tin tưởng mù quáng, và cách các kernel thu được hoạt động so với các bản tham chiếu. Kỹ năng này được tích hợp trong repo TileGym, vì vậy bạn có thể áp dụng nó cho các kernel của riêng mình.

Chuyển đổi kernel giữa các front end của Tile IR

cuTile Python, Triton-TileIR và cuTile Rust là ba front end cho cùng một IR: CUDA Tile IR, hay dialect cuda_tile. Cả ba đều cung cấp dữ liệu cho cùng một trình biên dịch tileiras, thực hiện các tối ưu hóa ở cấp độ tile và xuất ra tệp nhị phân GPU. Nền tảng chung này giúp việc chuyển đổi giữa các thành viên trong gia đình CUDA Tile trở nên thực tế và quan trọng không kém là có thể kiểm chứng được.

Các tile kernel sản xuất của TileGym được viết dựa trên hai front end đầu tiên. Vì cả ba đều gặp nhau tại cùng một IR, việc chuyển đổi một kernel sang cuTile Rust không phải là vấn đề tối ưu hóa lại. Đó là việc diễn đạt lại cùng một chương trình tile trong một ngôn ngữ host an toàn hơn, với cùng một trình biên dịch và cùng một mô hình hiệu năng bên dưới. IR chung giúp cho việc chuyển đổi có thể kiểm tra được.

Một bản chuyển đổi trung thực phải tái tạo được cấu trúc IR của kernel tham chiếu: cùng các họ thao tác bộ nhớ (memory-op families), cùng hình dạng tile và cùng các phép rút gọn (reductions). Vì cả ba front end đều xuất ra cùng một dialect, điều này có thể được xác minh trực tiếp bằng cách xuất (dump) Tile IR của kernel tham chiếu và Tile IR của kernel đã chuyển đổi, sau đó thực hiện “diff” chúng trước khi bất kỳ bài kiểm tra nào được thực thi.

Điều này cho phép kiểm tra đầu ra của tác nhân về mặt cấu trúc, không chỉ về mặt chức năng. Một bản chuyển đổi sai nhưng có vẻ hợp lý (ví dụ: một lệnh tải TMA với gợi ý chi phí sai hoặc thiếu thuộc tính chia hết) có thể vượt qua các bài kiểm tra nhưng vẫn không chính xác ngoài phạm vi kiểm tra và có thể gây ra sự suy giảm hiệu năng. Những vấn đề này có thể dễ dàng được kiểm tra và sửa chữa bằng cách so sánh với IR tham chiếu. Giai đoạn diff IR là trọng tâm của quy trình được mô tả trong bài viết này.

Hai khía cạnh bổ sung của front end Rust cần lưu ý cho cuộc thảo luận này. Thứ nhất, mã nguồn Rust được biên dịch trước (ahead of time). Hình dạng tile và kiểu phần tử được kiểm tra bởi rustc. Crate này nhúng AST của kernel, và tại lần khởi chạy đầu tiên, runtime sẽ chuyên biệt hóa nó với các giá trị const-generic cụ thể và biên dịch thành một tệp cubin (sau đó được lưu vào bộ nhớ đệm). Bản thân tệp nhị phân GPU vẫn được biên dịch JIT, nhưng tính ngầm định đã biến mất: không có gì được chuyên biệt hóa trừ khi chữ ký của kernel khai báo nó. Thứ hai, trong TileGym, cuTile Rust đơn giản là một backend khác. tilegym.set_backend("cutile-rs") định tuyến cùng một API toán tử đến các kernel Rust.

Làm cho sự chuyên biệt hóa trở nên tường minh

Hai front end khác nhau ở nơi diễn ra sự chuyên biệt hóa. cuTile Python JIT chuyên biệt hóa dựa trên bất cứ thứ gì nó thấy tại thời điểm gọi. cuTile Rust chỉ chuyên biệt hóa dựa trên những gì chữ ký của kernel khai báo. Hầu hết công việc chuyển đổi đến từ việc làm rõ những gì mã nguồn Python để ở dạng ngầm định. Các trường hợp chính được tóm tắt trong bảng dưới đây.

Phần sau minh họa những khác biệt này bằng một ví dụ kernel thực tế.

Ví dụ chuyển đổi Softmax

Kernel ví dụ này cố tình được làm đơn giản để bạn có thể so sánh hai phiên bản theo từng dòng. Đầu tiên, trong cuTile Python:

Và cùng một kernel đó trong cuTile Rust:

Bạn có thể đọc các sự tương ứng trực tiếp. Chúng sạch sẽ như vậy vì cả hai front end đều là các lớp bề mặt mỏng trên cùng các thao tác Tile IR:

Sau đó, IR diff xác nhận rằng Rust biên dịch ra cùng một tập hợp thao tác như bản gốc Python: một view load, reduce_max/reduce_sum trên trục phải, một view store và TMA ở cả hai đầu. Lưu ý rằng không phải tất cả các kernel được cung cấp trong TileGym đều sử dụng kiểu an toàn hoàn toàn này. Mỗi bản chuyển đổi phải tái tạo chính xác Tile IR của kernel tham chiếu, vì vậy ở những nơi chỉ có API không được kiểm tra mới tái tạo được nó, bản chuyển đổi sẽ sử dụng API đó. Chúng tôi vẫn đang chuyển các kernel đó sang bề mặt an toàn được hiển thị trong bài viết này.

Vượt qua C ABI

Tệp kernel.rs ví dụ đã là một kernel cuTile Rust hoàn chỉnh, hạng nhất. Một ứng dụng Rust có thể phụ thuộc vào crate cutile, bao gồm module kernel và khởi chạy entry point của nó trực tiếp thông qua API định kiểu của crate (bao gồm kiểm tra sở hữu, kiểu tile, v.v.) mà không cần FFI.

Lớp C-ABI phục vụ một mục đích hẹp hơn: cắm các kernel đó vào framework điều phối và kiểm tra Python của TileGym (và thông qua cùng cơ chế đó, bất kỳ host nào không phải Rust).

Mỗi toán tử xuất ra một C symbol từ cdylib tổng hợp (một tệp libcutile_kernels.so cho toàn bộ thư viện). Các tensor đi qua dưới dạng một struct mô tả đơn giản (ptr, ndim, shape[], strides[]) được phản chiếu giữa Rust và Python:

Ở phía Python, cffi liên kết symbol đó từ một chuỗi cdef, vốn là nguồn sự thật duy nhất cho chữ ký. Wrapper là một lớp mỏng với các bước kiểm tra xác thực:

Lưu ý rằng trình khởi chạy không bao giờ sao chép, không bao giờ cấp phát và không bao giờ chiếm quyền sở hữu. borrow_f32 bao bọc con trỏ thiết bị PyTorch trong một ManuallyDrop<Tensor>, vì vậy Rust có thể đưa các tensor cho kernel mà không bao giờ giải phóng bộ nhớ mà nó không sở hữu, và kernel khởi chạy không đồng bộ trên luồng CUDA của người gọi. Từ góc độ PyTorch, điều này trông giống như bất kỳ toán tử mở rộng nào khác.

Điều này cũng không gây trở ngại trong TileGym, vì cuTile Rust biên dịch một cách lười biếng (lazily). Backend theo dõi độ mới của mã nguồn, vì vậy việc chỉnh sửa bất kỳ tệp kernel.rs nào (hoặc tệp manifest của crate) sẽ khiến lệnh gọi tiếp theo tự động xây dựng lại thư viện chia sẻ trước khi điều phối, mà không cần lệnh cargo build rõ ràng trong chu kỳ phát triển-kiểm tra. Việc lặp lại trên một tile kernel trong Rust cũng dễ dàng như trong Python: thay đổi kernel, chạy kiểm tra và tệp nhị phân mới đã sẵn sàng.

Kỹ năng tác nhân hoạt động như thế nào?

Kỹ năng tác nhân tilegym-converting-python-to-rust, được cung cấp trong repo GitHub NVIDIA/TileGym, được xây dựng dựa trên một quyết định thiết kế: tác nhân tải nó không thực hiện bất kỳ công việc kỹ thuật nào cả. Việc đọc SKILL.md biến tác nhân cấp cao nhất thành một bộ điều phối thuần túy, có thẩm quyền duy nhất là định tuyến; công việc thực tế diễn ra trong các tác nhân phụ chuyên biệt mà nó tạo ra, mỗi tác nhân chỉ tải các tài liệu tham chiếu mà giai đoạn của nó cần. Chúng ta sẽ đi qua từng loại tác nhân phụ và vai trò của nó trong quá trình chuyển đổi.

Bộ phân tích giải quyết vấn đề “JIT che giấu đặc tả”. Trong một kernel tham chiếu, các hằng số được cố định khi DSL hạ cấp xuống dialect cuda_tile, các nhánh không được thực thi sẽ biến mất và các tham số khởi chạy nằm trong mã host. Bộ phân tích cũng chọn baseline: một toán tử thường có cả bản triển khai cuTile Python và Triton-TileIR, vì vậy bộ phân tích đo kiểm từng bản, so sánh chúng và chọn bản nhanh hơn cho mỗi biến thể cấu trúc làm tham chiếu mà bản chuyển đổi phải khớp.

Trước khi bất kỳ mã Rust nào tồn tại, nó xuất ra Tile IR của tham chiếu đó cho mỗi biến thể (sự thật cơ bản cho người viết kernel) và viết tệp analysis.json, một đặc tả có thể đọc bằng máy về các biến thể, hằng số, dtypes, dung sai, lưới khởi chạy, không gian autotune và baseline đã chọn. Mọi thứ phía sau đều định tuyến từ tệp này.

Người viết kernel chỉ tạo ra kernel.rs và không gì khác. Bị cấm truy cập mã host, các lỗi của nó vẫn có thể quy trách nhiệm được. Vấn đề khó khăn của nó chính là khoảng cách chuyển đổi, được chắt lọc thành 49 quy tắc viết mã của kỹ năng này. Nó chứng minh công việc của mình hai lần. Đầu tiên là về mặt chức năng, với một bài kiểm tra quy trình trong Rust chạy kernel mà không cần FFI và không cần Python, vì vậy một lỗi số học không thể ẩn sau các đường ống host. Thứ hai là về mặt cấu trúc, bằng cách vượt qua bước tự kiểm tra IR so với bản xuất tham chiếu của bộ phân tích.

Trình xây dựng host/FFI làm cho các kernel đã được xác thực có thể gọi được từ TileGym (trình khởi chạy C-ABI cộng với wrapper Python) và chịu trách nhiệm kiểm tra tính đúng đắn, chạy bộ kiểm tra TileGym thực tế của toán tử trên tất cả các dtype và hình dạng, và chỉ kết luận ALL_PASS của nó mới mở khóa việc đo kiểm hiệu năng. Đây là điểm đầu tiên mà toàn bộ stack (kernel, trình khởi chạy và wrapper) chạy từ đầu đến cuối.

Trình xác thực hiệu năng chạy giao thức đo kiểm CUPTI (đo thời gian thiết bị, ghép cặp theo cấu hình so với tham chiếu trên cùng GPU) và yêu cầu trung bình nhân phải nằm trong phạm vi 5% so với tham chiếu. Công việc của nó không phải là tối ưu hóa mà là đo lường một cách trung thực.

Hai chuyên gia chỉ tham gia khi có lỗi. Cả hai đều không chỉnh sửa mã; cả hai đều chẩn đoán bằng cách đọc IR. Nhà phân tích IR-diff được tạo ra khi một bài kiểm tra tính đúng đắn thất bại hoặc một kết quả đo kiểm có vẻ không ổn. Nó so sánh Tile IR tham chiếu với biến thể IR được tạo ra theo từng biến thể và phân loại từng sự khác biệt. Quan trọng là, điều này tách biệt một lỗi chuyển đổi (định tuyến lại cho người viết kernel với một bản sửa lỗi cụ thể) khỏi một lỗi trình biên dịch upstream mà không thay đổi kernel nào có thể sửa được.

Điều tra viên hiệu năng dư thừa lấy một kernel đúng nhưng chậm trên một số hình dạng đầu vào và tìm nguyên nhân gốc rễ của khoảng cách ở cả hai phía của ranh giới: phía thiết bị (họ memory-op, codegen) và phía host (cấu hình khởi chạy, autotune và logic wrapper). Nó đưa ra một báo cáo để người viết kernel hành động.

Hai lý do chính thúc đẩy thiết kế này. Thứ nhất, một quá trình chuyển đổi đầy đủ chạy với quy mô hàng triệu token. Thứ hai, sự phân tách giúp cô lập trách nhiệm. Vì kernel được chứng minh một cách độc lập trước khi bất kỳ mã host nào tồn tại, một lỗi xảy ra sau đó sẽ có chủ sở hữu rõ ràng.

Ba lựa chọn làm cho sự phân tách này hoạt động. Các tác nhân phụ chỉ giao tiếp thông qua các artifact với lược đồ cố định, không bao giờ thông qua hội thoại. Mỗi giai đoạn kết thúc bằng một kết luận có thể kiểm chứng bằng máy mà bộ điều phối định tuyến dựa trên đó mà không cần đọc văn bản. Và dialect cuda_tile chung làm cho IR diff trở thành xương sống của việc xác minh—được sử dụng cả như bước tự kiểm tra của người viết kernel trước khi chạy kiểm tra và như sự so sánh sâu của nhà phân tích IR-diff khi có lỗi—loại bỏ các bản chuyển đổi sai về cấu trúc (phép rút gọn trên sai trục, mất mask) mà nếu không sẽ vượt qua như thể hợp lý.

Vòng lặp điều phối

Một lần chạy chuyển đổi là một máy trạng thái nhỏ, và các hướng dẫn của chính bộ điều phối nằm gọn trong tệp SKILL.md tinh gọn. Các bước được trình bày chi tiết trong Hình 1 và các phần tiếp theo.