Tối ưu hóa LLM: Cơ chế Hybrid Attention giúp tăng tốc độ suy luận lên 51 lần

Một kỹ sư gần đây đã chia sẻ quá trình xây dựng một mô hình ngôn ngữ nhỏ (Small Language Model) chuyên tập trung vào ngôn ngữ lập trình Rust ngay từ đầu (from scratch). Không phải là phiên bản tinh chỉnh (finetune) của các mô hình hiện có, dự án này sử dụng kiến trúc Transformer GPT nhưng áp dụng một cải cách lớn về cơ chế Attention nhằm giải quyết bài toán chi phí tính toán.

Kiến trúc HybridAttention: Thay đổi cuộc chơi về hiệu suất

Mô hình Transformer tiêu chuẩn sử dụng cơ chế "Full Attention" có độ phức tạp tính toán là $O(n^2)$, điều này khiến quá trình suy luận (inference) trở nên chậm chạp khi chiều dài ngữ cảnh tăng lên. Trong dự án này, tác giả đã thay thế bằng HybridAttention, một khối kiến trúc mới kết hợp giữa hai phương pháp:

• Cửa sổ cục bộ (Local Windowed Causal Attention): Chịu trách nhiệm xử lý các cú pháp ở phạm vi ngắn.
• Đường dẫn trạng thái đệ quy (Recurrent State Path): Tương tự như GRU, mang theo trạng thái đã nén cho các mối quan hệ phạm vi xa.

Hai thành phần này được hòa trộn thông qua một cổng học được (learned gate). Thiết kế này giúp chuyển đổi độ phức tạp tính toán hiệu quả từ dạng bậc hai ($O(n^2)$) sang gần tuyến tính ($O(n \cdot W + n \cdot D)$), trong đó $W$ là kích thước cửa sổ và $D$ là kích thước trạng thái.

Hiệu suất kỷ lục trên GPU tiêu dùng

Kết quả benchmark trên một card đồ họa RTX 4060 Ti 8GB cho thấy sự chênh lệch đáng kinh ngạc về hiệu suất:

• Full Attention tiêu chuẩn: 17,96 giây / 5,6 tokens mỗi giây.
• HybridAttention: 0,35 giây / 286,6 tokens mỗi giây.

Con số này tương ứng với việc tăng tốc độ suy luận lên 51 lần mà không gây ra sự suy giảm rõ rệt về chất lượng đầu ra. Bí mật nằm ở chiến lược KV Cache tùy chỉnh:

Hệ thống giữ lại một "cửa sổ nóng" (hot window) gồm 64 token mới nhất trong bộ nhớ VRAM với độ chính xác đầy đủ. Các token cũ hơn sẽ được nén xuống dạng 8-bit (lưu độ lớn và góc) và có thể được khôi phục lại độ chính xác đầy đủ một cách chọn lọc khi cần thiết.

Chi tiết huấn luyện và chất lượng mã

Mô hình này sở hữu 25,6 triệu tham số, sử dụng từ vựng cấp độ byte (256 vocab size) và chiều dài ngữ cảnh là 512. Quá trình huấn luyện diễn ra trong 30k bước trên một bộ dữ liệu Rust khổng lồ (173,5 triệu byte), được tổng hợp từ mã nguồn chính thức của Rust và 500 crate phổ biến nhất (như tokio, serde, clap...).

Về mặt chất lượng, cú pháp bề mặt của mã Rust được tạo ra trông khá ổn thỏa, các câu lệnh nhập khẩu (imports) và chữ ký hàm thường khả thi. Tuy nhiên, ngữ nghĩa (semantics) vẫn còn yếu và các mẫu lặp lại thường xuyên xuất hiện. Dù vậy, việc mở rộng quy mô bộ dữ liệu (corpus expansion) được tác giả đánh giá là mang lại hiệu quả lớn hơn cả việc thay đổi kiến trúc mô hình.

Dự án này là minh chứng cho thấy việc tối ưu hóa kiến trúc mô hình và chiến lược bộ nhớ đệm có thể mang lại lợi ích khổng lồ, cho phép các mô hình AI chạy mượt mà trên phần cứng tiêu dùng mà không cần đòi hỏi hệ thống đắt đỏ.