Tối ưu Mamba-2 với các kernel forward pass hiệu năng cao trong NVIDIA CUDA Tile

NVIDIA CUDA Tile (cuTile) là một ngôn ngữ DSL hướng Python và mô hình lập trình song song cho GPU của NVIDIA, tự động khai thác các khả năng phần cứng nâng cao để mang lại hiệu suất tối đa cho lập trình ở cấp độ tile. Bằng cách trừu tượng hóa lịch trình ở cấp độ thread và warp porver các nguyên tắc tile-centric cấp cao, cuTile cho phép bạn viết kernel GPU bằng Python trong khi các tối ưu hóa trình biên dịch tự động nhắm vào các tính năng phần cứng nâng cao như Tensor Cores. Là một nhà sáng lập sớm của cuTile, các kỹ sư học máy tại Jane Street đã cài đặt Mamba-2, một mô hình chuỗi trạng thái tiên tiến, bằng cuTile và hợp tác với NVIDIA để tối ưu kernel. Jane Street là một công ty giao dịch toàn cầu rất kỹ thuật, sử dụng deep learning để thúc đẩy đổi mới trong nghiên cứu định lượng nâng cao. Bài viết này sẽ hướng dẫn bạn qua ba kernel tùy chỉnh đang cầm quyền cho forward pass của Mamba-2 và các tối ưu cuTile đằng sau chúng. Để sử dụng cuTile, hãy đảm bảo môi trường của bạn đáp ứng các yêu cầu được nêu trong hướng dẫn nhanh Python cuTile. Nhập khẩu gói như cuda.tile và hiển thị nó dưới dạng ct trong toàn bộ cài đặt kernel được đề cập trong bài viết này. Các kiến trúc mô hình sequence modeling phải đối mặt với một sự thỏa hiệp tính toán: attention softmax truyền thống huấn luyện nhanh bằng cách nhân ma trận được tăng tốc bởi Tensor Cores, nhưng trong quá trình suy diễn, yêu cầu tính toán và bộ nhớ tăng tuyến tính theo độ dài chuỗi. Ngược lại, các mạng nơron lặp lại (RNNs) gặp khó khăn trong việc huấn luyện chậm và không thể song song hóa nhưng yêu cầu tính toán không đổi và dung lượng bộ nhớ cố định bất kể độ dài chuỗi.