Nhà nghiên cứu hệ thống AI được đánh giá cao Tri Dao – đồng sáng tạo của FlashAttention và Mamba – vừa hé lộ một thủ thuật tối ưu hệ thống mang tính đột phá, giúp giảm đáng kể nút thắt về suy luận (inference bottleneck) cho các mô hình trạng thái không gian (SSMs) như Mamba và các kiến trúc attention tuyến tính như Gated-DeltaNet. Khám phá này đã được ông chia sẻ trên X (trước đây là Twitter).
Cải tiến đơn giản nhưng sâu sắc này giúp các quy trình làm việc theo tác nhân (agentic workflows) chạy trên kiến trúc lai (hybrid architectures) – như Qwen 3.5 và Nemotron Ultra – nhanh hơn tới 2 lần, cuối cùng khai phá được sức mạnh thực sự của Giải mã suy luận (Speculative Decoding) cho các kiến trúc không phải Transformer.
Bối Cảnh: Vì Sao Mamba và Gated-DeltaNet Gặp Khó Khăn với Giải Mã Suy Luận? 📉
Giải mã suy luận là một kỹ thuật tăng tốc suy luận LLM phổ biến, trong đó một mô hình “nháp” nhỏ đề xuất một số token ứng cử viên (thường là từ 3 đến 12), và một mô hình “mục tiêu” lớn hơn sẽ xác minh chúng trong một lần truyền tới song song duy nhất. Đối với các Transformer tiêu chuẩn, điều này làm giảm đáng kể thời gian giải mã tự hồi quy (autoregressive decoding) bị giới hạn bởi bộ nhớ.
Tuy nhiên, khi áp dụng cho các Mô hình Trạng thái Không gian (SSMs) như Mamba hoặc các lớp attention tuyến tính như Gated-DeltaNet (đang ngày càng được sử dụng trong các kiến trúc lai như Qwen 3.5 và Nemotron Ultra), giải mã suy luận gặp phải một nút thắt hệ thống lớn:
1. Dấu Chân Trạng Thái Khổng Lồ: Không giống như Transformer lưu trữ bộ nhớ đệm KV, SSMs và Gated-DeltaNet duy trì một trạng thái ẩn đang chạy được cập nhật tuần tự. 2. Chi Phí Bộ Nhớ Lớn (Kế Hoạch I - Lưu Trữ Tất Cả Các Trạng Thái): Trong quá trình xác minh song song các token suy đoán của mô hình nháp, các triển khai tiêu chuẩn sẽ hiện thực hóa và lưu trữ các trạng thái SSM ẩn trung gian cho mọi token ứng cử viên được đề xuất vào bộ nhớ băng thông cao ngoài chip (DRAM). 3. Nút Thắt Cổ Chai: Vì kích thước trạng thái của SSM và attention tuyến tính rất lớn, việc di chuyển các trạng thái trung gian này qua lại giữa SRAM GPU nhanh và bộ nhớ ngoài chip tạo ra độ trễ truyền bộ nhớ cực lớn. Chi phí giao tiếp này hoàn toàn xóa sạch mọi lợi ích tốc độ có được từ giải mã suy luận.
Giải Pháp Của Tri Dao: Kỹ Thuật Tái Tính Toán “Tải, Tính Toán, Nhưng Không Lưu Trữ” (Load, Compute, but Don't Store) 💡
Giải pháp thanh lịch của Tri Dao tận dụng thiết kế nhân (kernel) nhận biết phần cứng và chiến lược tái tính toán trạng thái có chọn lọc:
* Loại Bỏ Trạng Thái Trung Gian: Trong quá trình xác minh song song các token ứng cử viên, hệ thống tải trạng thái ban đầu vào SRAM GPU nhanh, tính toán các phép chiếu đầu ra để xác minh các token nháp, nhưng không ghi các trạng thái ẩn trung gian trở lại bộ nhớ ngoài chip. * Vì Sao Hiệu Quả: Vì giải mã suy luận thường xuyên từ chối các token được đề xuất, việc ghi các trạng thái của các token bị từ chối ra ngoài chip là hoàn toàn lãng phí công sức. Loại bỏ chúng ngay lập tức giúp giảm đáng kể lưu lượng truy cập bộ nhớ ngoài chip của GPU. * Tái Phát Hoạt (Activation Replay) (Tái Tính Toán Tức Thì): Khi mô hình mục tiêu xác thực các token và quyết định tiền tố được chấp nhận dài nhất, hệ thống sử dụng kỹ thuật gọi là tái phát hoạt (activation replay) để tái tính toán trạng thái chính xác chỉ cho chuỗi được chấp nhận vào đầu lần lặp suy luận tiếp theo.
Bằng cách tránh lưu trữ ngoài chip và chỉ tái tính toán các trạng thái được chấp nhận ngay lập tức, lưu lượng bộ nhớ được giảm đi một yếu tố bằng kích thước trạng thái. Điều này loại bỏ nút thắt băng thông bộ nhớ, tăng tốc pha xác minh lên 2 lần và hỗ trợ đầy đủ giải mã cây suy luận trên Mamba-3, Mamba-2 và Gated-DeltaNet.
Ý Nghĩa Quan Trọng Cho Kỷ Nguyên AI Agent 🤖
Khi các AI agent thế hệ mới thực hiện các tác vụ phức tạp, dài hạn, các mô hình đòi hỏi các cửa sổ ngữ cảnh khổng lồ. Các kiến trúc lai hiện đại – như Gated-DeltaNet kết hợp với các lớp attention cục bộ – đã nổi lên như những lựa chọn thay thế hiệu quả cao cho các Transformer thuần túy.
Các tối ưu hóa hệ thống nhận biết phần cứng của Tri Dao đảm bảo rằng các mô hình không phải Transformer giờ đây có thể hưởng lợi từ giải mã suy luận mà không bị tê liệt bởi các nút thắt bộ nhớ. Đây là một bước tiến lớn trong việc biến các AI agent cục bộ, độ trễ thấp và có khả năng cao thành hiện thực trong sản xuất.