Tăng tốc mô hình lai với replayssm cho State space model

Các mô hình ngôn ngữ kết hợp giữa cấu trúc trạng thái không gian và cơ chế chú ý đang gặp một điểm nghẽn lớn về tốc độ giải mã. Đội ngũ nghiên cứu từ NVIDIA phối hợp cùng Dao AI Lab tại Đại học Princeton vừa công bố giải pháp ReplaySSM để giải quyết triệt để bài toán này. Phương pháp mới thay đổi cách lưu trữ trạng thái, giúp các mô hình lai đạt tốc độ vượt trội ở kích thước batch lớn và mở đường cho việc ứng dụng giải mã dự đoán.

Các mô hình như NVIDIA Nemotron 3, dòng Qwen từ bản 3.5 đến 3.8, Kimi-K3 hay GLM-5.3-Flash đều sử dụng kiến trúc lai, đan xen các lớp trạng thái không gian như Mamba-2, Gated DeltaNet hoặc Kimi Delta Attention vào giữa các tầng chú ý. Về lý thuyết, các lớp trạng thái không gian giúp giải mã nhanh bằng cách dùng một trạng thái kích thước cố định thay cho bộ nhớ đệm khóa giá trị truyền thống vốn phình ra theo ngữ cảnh. Tuy nhiên, ở các quy mô phục vụ thực tế, chính những lớp này lại trở thành nút thắt cổ chai về độ trễ với ngữ cảnh lên tới 100.000 token, đồng thời làm phức tạp hóa các kỹ thuật tối ưu như giải mã dự đoán và lượng tử hóa.

Nguyên nhân gốc rễ nằm ở cơ chế vận hành truyền thống khi mỗi bước giải mã đều phải đọc toàn bộ trạng thái từ bộ nhớ băng thông cao, cập nhật rồi ghi lại ngay lập tức. Thao tác này tiêu tốn băng thông bộ nhớ trong khi thực hiện rất ít phép tính số học. ReplaySSM khắc phục điểm yếu này bằng cách tách bạch việc tính toán đầu ra khỏi việc cập nhật trạng thái. Thay vì lưu giữ một ma trận trạng thái cồng kềnh, hệ thống giữ một điểm kiểm tra kết hợp cùng một bộ đệm nhỏ chứa các đầu vào gần nhất.

Trong quá trình giải mã thông thường, bộ đệm ghi nhận các đầu vào mới và tính toán kết quả trực tiếp mà không cần ghi đè trạng thái liên tục. Khi bộ đệm này đầy, thuật toán mới thực hiện một bước làm sạch để dồn dữ liệu vào điểm kiểm tra chính. Kết quả thử nghiệm trên bốn mô hình lai từ hai họ Nemotron 3 và Qwen 3.5, với quy mô từ các mô hình bốn tỷ tham số chạy trên một GPU NVIDIA H100 cho tới các mô hình chuyên gia hỗn hợp 120 tỷ tham số chạy trên một GPU NVIDIA Blackwell Ultra B300, cho thấy hiệu suất rõ rệt. Ở kích thước batch 256, nhân xử lý chạy nhanh hơn từ 1,46 đến 1,80 lần, giúp toàn bộ bước giải mã tăng tốc từ 1,20 đến 1,47 lần.

Giải pháp này cũng giải quyết bài toán khó trong việc áp dụng giải mã dự đoán cho các mô hình trạng thái không gian. Trước đây, việc khôi phục trạng thái sau khi kiểm tra các token dự đoán rất tốn kém vì các cập nhật không có nghịch đảo đơn giản, buộc các hệ thống phải lưu bản chụp trạng thái liên tục và làm giảm dung lượng bộ nhớ GPU. Với ReplaySSM, các đầu vào dự đoán chỉ đơn giản được ghi thêm vào bộ đệm và con trỏ sẽ dịch chuyển khi token được chấp nhận mà không cần tạo bản chụp hay khôi phục trạng thái phức tạp. Trên các mô hình 120 tỷ tham số với kích thước batch 512, ReplaySSM mang lại thông lượng cao gấp 1,85 đến 2,10 lần so với cấu hình chuẩn của vLLM.

Việc tách biệt tính toán và cập nhật còn đem lại lợi ích lớn trong bài toán lượng tử hóa trạng thái nhằm giảm dung lượng bộ nhớ. Phương pháp mới duy trì một bản sao độ chính xác cao cho vòng lặp trạng thái và một bản sao độ chính xác thấp chỉ để đọc đầu ra, qua đó ngăn lỗi làm tròn lan truyền qua các bước lặp. Các thử nghiệm trên bộ dữ liệu kiểm định dài cho thấy việc giữ vòng lặp ở định dạng dấu phẩy động 32 bit trong khi đọc trạng thái 4 bit giúp giảm mạnh sai số phân kỳ. Trên điểm chuẩn GPQA-Diamond, cách tiếp cận này giúp khôi phục phần lớn độ chính xác bị mất, giữ cho các mô hình bám sát kết quả gốc.

Giải pháp hiện đã được đưa vào thực tế sản xuất. Các khung suy luận mã nguồn mở phổ biến bao gồm vLLM, TensorRT-LLM, SGLang và FlashInfer đều đã tích hợp ReplaySSM, trong khi các mô hình như GLM-5.3-Flash và Kimi-K3 cũng đang ứng dụng kỹ thuật này để tối ưu hóa hiệu năng tầng trạng thái không gian.