NVIDIA và Jane Street tối ưu các kernel forward pass của Mamba-2 bằng CUDA Tile
NVIDIA và Jane Street đã triển khai mô hình chuỗi Mamba-2 bằng NVIDIA CUDA Tile (cuTile), một mô hình lập trình dựa trên Python để viết kernel GPU. Hai bên tối ưu ba kernel dùng trong lượt tính toán xuôi của Mamba-2: chunkcumsum, fwdpassstate và computeout. Các kernel này lần lượt tính tổng tích lũy liên quan đến độ suy giảm, cập nhật trạng thái ẩn của Mamba-2 qua các đoạn chuỗi bằng phép nhân ma trận trên Tensor Core, đồng thời gộp trạng thái, tương tác có che causal và kết nối bỏ qua vào một phép tính. Công việc sử dụng nhiều kỹ thuật như giữ kết quả trung gian trong thanh ghi, chạy Tensor Core ở định dạng TF32, sắp xếp dữ liệu dạng chuyển vị để truy cập bộ nhớ nhanh hơn, cung cấp gợi ý về độ trễ khi truy cập DRAM, đưa toán hạng vào Tensor Memory và tự động dò kích thước block cùng mức độ chiếm dụng GPU. Mamba-2 hướng tới việc kết hợp ưu điểm suy luận với bộ nhớ cố định của mô hình hồi quy với khả năng huấn luyện và xử lý phần đầu chuỗi song song hơn nhờ tính toán theo từng đoạn, qua đó tránh một phần chi phí tăng theo độ dài chuỗi của attention truyền thống. Về thực tế, bài viết cho thấy lập trình viên có thể dùng giao diện Python cấp cao hơn trong khi trình biên dịch cuTile vẫn nhắm tới các tính năng GPU như Tensor Core, giúp việc phát triển và tối ưu kernel Mamba-2 tùy biến trở nên dễ tiếp cận hơn.
Các kỹ sư AI và nhóm xây dựng hạ tầng GPU có thể xem mã mẫu cuTile để thử viết kernel Mamba-2 bằng Python và tận dụng Tensor Cores mà không phải tự quản lý toàn bộ lịch chạy thread và warp. Đây chủ yếu là tài liệu kỹ thuật cho việc tối ưu hiệu năng, chưa phải một sản phẩm người dùng cuối.
Đọc bài gốc