Hạt nhân là gì và tại sao mọi người lại viết hạt nhân mới

Hạt nhân là một chương trình nhỏ chạy trên GPU. FlashAttention, Triton và FlashInfer đều là tên trong lớp này và tất cả chúng đều tồn tại vì cùng một lý do: số học chưa bao giờ là phần đắt tiền. Ngăn xếp phân phát có lớp hạt nhân.Phần ba trong số bảy. Một lớp phía trên silicon và là nguồn gốc của hầu hết những cái tên mà bạn đã nghe. Phần 1: Bắt đầu tại đây: Những từ mọi người sử dụng về suy luận LLM Phần 2: Bộ đệm KV của bạn lớn hơn mô hình của bạn Phần một khẳng định rằng GPU có thể tính toán nhanh hơn nhiều so với khả năng tìm nạp và phần hai cho thấy số lượng cần tìm nạp. Phần này nói về lớp mà ai đó thực sự làm gì đó với nó. Đó cũng là nơi gần như tất cả các từ vựng đều đến từ đó. FlashAttention, FlashInfer, Triton, CUTLASS, CUDA: mỗi tên trong số đó là một tên trong lớp duy nhất này và chúng sẽ có ý nghĩa hơn nhiều khi bạn biết chức năng của lớp đó. Hạt nhân là gì Hạt nhân là một chức năng duy nhất chạy trên GPU. Bạn khởi chạy nó, hàng ngàn luồng thực thi nó song song và nó kết thúc. Nhân hai ma trận là một hạt nhân. Áp dụng softmax cũng vậy, bước biến một hàng điểm thô thành các phân số có tổng bằng một. Chạy toàn bộ mô hình có nghĩa là khởi chạy hàng trăm mô hình theo trình tự. CUDA là nền tảng của NVIDIA để viết chúng và mọi thứ khác có tên ở trên đều là thư viện hoặc ngôn ngữ nằm trên đó. Từ này nghe có vẻ khó hơn nó. Điều khiến hạt nhân trở nên thú vị là hệ thống phân cấp bộ nhớ chứa chúng và độ dốc của nó: Bốn bậc được vẽ theo tỷ lệ trên trục nhật ký: thanh ghi 33 MiB, bộ nhớ dùng chung 29 MiB, L2 50 MB và HBM3 80 GB, với khoảng cách 1.600 lần giữa hai cấp cuối cùng được đánh dấu. Vẽ theo tỷ lệ vì hầu hết các mô tả khác về điều này đều sai. Nó không phải là một cái thang trơn tru. Hình dạng không phải là những gì mọi người mô tả. Ba trong số các cấp đó có dung lượng gần như nhau: 33 MiB, 29 MiB và 50 MB. Tệp đăng ký thực sự lớn hơn bộ nhớ dùng chung bên dưới nó, do đó, các tầng thậm chí không lớn hơn khi bạn giảm xuống. Có một bước thực sự và nó nằm giữa bộ nhớ đệm L2 và HBM: 1.600 lần. Bộ nhớ dùng chung là cấp quan trọng ở đây: một bàn di chuột vật lý trên chip, có tốc độ truy cập nhanh hơn bộ nhớ chính và được chia sẻ bởi các luồng của một khối. Điều đáng chú ý là hầu như không có gì trong số đó. SM (bộ xử lý đa luồng phát trực tuyến) là một trong những đơn vị tính toán độc lập của chip và mỗi đơn vị nhận được 228 KiB. Một khối duy nhất chỉ nhận được 48 KiB trong số đó trừ khi hạt nhân chọn tham gia thêm một cách rõ ràng. Vì vậy, một hạt nhân tốt không phải là một hạt nhân thực hiện các phép tính thông minh, mà là hạt nhân đưa dữ liệu vào bảng ghi nhớ đó và xử lý càng nhiều công việc càng tốt trước khi quay lại để tìm thêm. Tại sao lại có vấn đề? Cách rõ ràng để tính toán sự chú ý là làm theo phương trình. Đối với một chuỗi có độ dài S, bạn tính điểm giữa mọi mã thông báo và mọi mã thông báo khác, đó là ma trận S x S. Viết nó vào bộ nhớ. Đọc lại để áp dụng softmax. Viết cái đó. Đọc lại để nhân với các giá trị. Bốn chuyến đi qua ranh giới chậm nhất trong máy, để có một ma trận tăng theo bình phương độ dài chuỗi của bạn. Với 64K mã thông báo, ma trận đó là 8,6 GB cho một đầu truy vấn trong một lớp. Dấu chân bộ nhớ chú ý ở bốn độ dài chuỗi, ma trận điểm cụ thể hóa so với các ô được truyền trực tuyến. Với 64K mã thông báo, ma trận điểm là 8,6 GB, cho một đầu truy vấn trong một lớp. gpt-oss-120b có 64 đầu truy vấn và 18 lớp tham gia vào toàn bộ chuỗi. Việc che dấu nhân quả làm giảm một nửa số học, vì mã thông báo không bao giờ quan tâm đến tương lai của nó e, nhưng nó không giảm một nửa 8,6 GB đó: một cách triển khai đơn giản vẫn phân bổ toàn bộ hình vuông và lấp đầy một nửa hình vuông đó bằng âm vô cực. FlashAttention không có ma trận để phân bổ, do đó, nó bỏ qua hoàn toàn các ô được che kín và đó là nguồn gốc của phần lớn khả năng tăng tốc trên các mô hình bộ giải mã. Đó là con số đã biến lớp này thành một lĩnh vực nghiên cứu. FlashAttention thực sự làm được những gì Nó không bao giờ viết ra ma trận điểm số. Không phải một lần. Thay vào đó, nó xử lý trình tự theo từng ô: tải một khối truy vấn và một khối khóa vào bộ nhớ dùng chung, tính điểm của chúng ở đó, sử dụng chúng, loại bỏ chúng, rồi tiếp tục. Ma trận đầy đủ không bao giờ tồn tại. Chỉ có ô đang được xử lý mới xuất hiện trên chip trong thời gian ngắn. Thủ thuật giúp thực hiện được điều này là softmax trực tuyến và nó xuất hiện trước FlashAttention. Milkov và Gimelshein đã mô tả nó vào năm 2018, và vào năm 2021, Rabe và Staats cho thấy sự chú ý có thể được tính toán mà không cần hiện thực hóa ma trận đầy đủ. Softmax thường cần xem mọi điểm số trước khi có thể bình thường hóa, đó chính xác là lý do tại sao phiên bản đơn giản phải lưu trữ tất cả. Phiên bản trực tuyến giữ mức tối đa đang chạy và tổng hiện có, đồng thời điều chỉnh lại kết quả một phần của nó mỗi khi có khối mới. Khi khối cuối cùng đã được xử lý, câu trả lời là đúng — chính xác, không phải gần như vậy. Cả hai bên đều tính toán các ô 4x4 giống nhau theo cùng một thứ tự, bên trái tích lũy mọi ô trong bộ nhớ trong khi bên phải chỉ giữ ô hiện tại. Các ô giống nhau, cùng thứ tự, cùng một câu trả lời. Đường viền vàng là ô đang được tính toán. Bên trái giữ lại mọi việc đã làm xong; bên phải giữ bất cứ thứ gì nó đang giữ và bỏ phần còn lại, đó là lý do tại sao một bộ đếm tăng lên còn bộ đếm kia thì không. Điều này được mô tả một cách lỏng lẻo và hai xác nhận quyền sở hữu khác nhau được hợp nhất thành một, vì vậy cần phải chính xác. Nó tính toán sự chú ý chính xác. Nó không gần đúng. Các giấy tờ “chú ý hiệu quả” mà nó được nộp cùng với tất cả đều chạy nhanh hơn bằng cách bỏ một phần phép tính hoặc nén nó. Cái này không có gì cả. Nó không giống một chút nào với phiên bản đơn giản, vì việc tính tổng theo một thứ tự khác sẽ thay đổi một số bit cuối cùng của kết quả dấu phẩy động. Trong thực tế, nó tiến gần đến tham chiếu có độ chính xác kép hơn so với phiên bản đơn giản vì nó tích lũy trong 32 bit. Và trên các ô nó tính toán, nó tính toán nhiều hơn chứ không ít hơn. Truyền phát một ô tại một thời điểm có nghĩa là số softmax phải chia cho không được biết cho đến khi ô cuối cùng xuất hiện, do đó tổng số đang chạy phải được điều chỉnh lại khi mỗi ô mới xuất hiện