ReplaySSM đẩy nhanh quá trình giải mã mô hình hybrid bằng cách lưu đệm đầu vào thay vì trạng thái

NVIDIA và Dao AI Lab tại Đại học Princeton đã giới thiệu ReplaySSM, một phương pháp mới nhằm tối ưu hóa giai đoạn giải mã của các mô hình ngôn ngữ lớn hybrid. Được công bố vào tháng 10 năm 2026, cách tiếp cận này xem xét lại cách thức mà các Mô hình Không gian Trạng thái (SSMs) lưu trữ và cập nhật trạng thái nội bộ trong quá trình suy luận. Bằng cách lưu đệm các token đầu vào gần đây thay vì liên tục ghi đè lên toàn ma trận trạng thái, kỹ thuật này giảm đáng kể mức sử dụng băng thông bộ nhớ và cho phép quá trình giải mã suy đoán hiệu quả hơn.

Chuyện gì đã xảy ra

Các mô hình hybrid, xen kẽ giữa các lớp SSM với các lớp cơ chế chú ý (attention) truyền thống, đang trở nên phổ biến trong các kiến trúc hiện đại như NVIDIA Nemotron 3, Qwen3.5 và Kimi-K3. Những mô hình này sử dụng SSM để xử lý ngữ cảnh dài một cách hiệu quả, thay thế cache khóa-giá trị (key-value cache) ngày càng phình to của cơ chế attention bằng một trạng thái có kích thước cố định. Tuy nhiên, ở quy mô batch phục vụ lớn, các lớp SSM thường trở thành nút thắt cổ chai chính về độ trễ do yêu cầu đọc và ghi thường xuyên các ma trận trạng thái lớn từ bộ nhớ băng thông cao.

Nhóm nghiên cứu nhận thấy rằng quá trình giải mã thông thường gộp chung việc tính toán đầu ra và cập nhật trạng thái vào một bước duy nhất, gây ra lưu lượng bộ nhớ quá mức. ReplaySSM tách rời các thao tác này bằng cách biểu diễn trạng thái dưới dạng một điểm kiểm tra (checkpoint) cộng với một bộ đệm chứa các đầu vào gần đây. Thay đổi này cho phép hệ thống tính toán đầu ra mà không cần liên tục ghi toàn bộ trạng thái trở lại bộ nhớ, chỉ cập nhật điểm kiểm tra khi bộ đệm đầu vào đầy. Phương pháp này hiện đã được tích hợp vào các framework suy luận chính bao gồm TensorRT-LLM, vLLM và FlashInfer.

Cách hoạt động

Trong quá trình giải mã SSM tiêu chuẩn, mỗi bước tạo token sẽ tải trạng thái trước đó, áp dụng sự suy giảm và cập nhật, tính toán đầu ra và lưu trữ trạng thái mới. Quá trình này bị giới hạn bởi bộ nhớ vì cường độ tính toán thấp so với lượng dữ liệu di chuyển. ReplaySSM thay đổi điều này bằng cách mở rộng vòng lặp hồi quy trên một cửa sổ nhỏ các token. Thay vì cập nhật trạng thái ở mỗi bước, nó thêm các vector đầu vào của token hiện tại vào bộ đệm. Đầu ra được tính toán bằng cách sử dụng điểm kiểm tra đã lưu và các đầu vào trong bộ đệm.

Khi bộ đệm đạt đến kích thước xác định trước, một thao tác flush sẽ gộp tất cả các đầu vào trong bộ đệm vào điểm kiểm tra và ghi nó vào bộ nhớ chỉ một lần. Điều này giảm số lần ghi trạng thái xuống khoảng một nửa. Đối với giải mã suy đoán, nơi nhiều token nháp được tạo ra và xác minh, ReplaySSM đơn giản hóa việc hoàn tác. Thay vì lưu ảnh chụp nhanh (snapshot) của toàn bộ trạng thái cho mỗi token nháp, hệ thống chỉ cần thêm các đầu vào nháp vào bộ đệm. Nếu một bản nháp bị từ chối, con trỏ bộ đệm sẽ được đặt lại, loại bỏ các đầu vào chưa cam kết mà không cần bất kỳ thao tác khôi phục trạng thái tốn kém nào.

Chi tiết chính

• ReplaySSM giảm lưu lượng bộ nhớ trạng thái khoảng 50% bằng cách chỉ ghi các điểm kiểm tra sau khi bộ đệm chứa các đầu vào gần đây đầy.
• Trên GPU NVIDIA Blackwell Ultra B300 với kích thước batch 256, các kernel giải mã SSM chạy nhanh hơn 1,46x đến 1,80x so với các phương pháp thông thường.
• Thông lượng giải mã suy đoán tăng 1,85x đến 2,10x đối với các mô hình cỡ 120B vì việc hoàn tác không còn đòi hỏi phải lưu và khôi phục ảnh chụp nhanh trạng thái.
• Kỹ thuật này hỗ trợ lưu trữ trạng thái độ chính xác thấp bằng cách giữ một điểm kiểm tra độ chính xác cao cho vòng lặp hồi quy trong khi sử dụng các bản sao độ chính xác thấp cho việc đọc, hạn chế sự tích lũy lỗi.
• Các mô hình được đánh giá bao gồm các biến thể Nemotron-3 và Qwen3.5, cho thấy lợi ích nhất quán trên các kiến trúc SSM khác nhau như Mamba-2 và Gated DeltaNet.
• Phương pháp này đã được triển khai trong các hệ thống sản xuất như GLM-5.3-Flash và Kimi-K3.

Tại sao điều này quan trọng

Đối với các kỹ sư xây dựng và phục vụ các mô hình hybrid lớn, băng thông bộ nhớ thường là yếu tố giới hạn thông lượng. Bằng cách cắt giảm lưu lượng bộ nhớ chiếm ưu thế trong các lớp SSM, ReplaySSM cho phép tăng kích thước batch mà không làm tăng độ trễ trên mỗi token. Điều này rất quan trọng cho việc phục vụ chi phí hiệu quả, vì nó cho phép tận dụng phần cứng tốt hơn trên các GPU đắt tiền như H100 và B300. Lợi ích hiệu năng tăng theo kích thước batch, khiến nó đặc biệt có giá trị trong các môi trường sản xuất có độ đồng thời cao.

Ngoài ra, ReplaySSM làm cho giải mã suy đoán khả thi đối với các mô hình dựa trên SSM. Trước đây, chi phí lưu ảnh chụp nhanh trạng thái để hoàn tác đã triệt tiêu lợi ích của việc suy đoán trong các kiến trúc này. Bằng cách coi bộ đệm đầu vào giống như một con trỏ cache khóa-giá trị, phương pháp này khôi phục hiệu quả của dự đoán đa token. Điều này cho phép các nhà phát triển sử dụng các chiến lược suy đoán mạnh mẽ để giảm thêm độ trễ, đưa hiệu năng SSM tiến gần hơn đến các giới hạn lý thuyết của phần cứng.

Bạn có thể làm gì

• Đọc bài báo ReplaySSM đầy đủ trên arXiv để hiểu rõ nguồn gốc toán học và các chỉ số đánh giá chi tiết.
• Khám phá triển khai mã nguồn mở trên GitHub để thử nghiệm cơ chế lưu đệm trong các dự án của riêng bạn.
• Nâng cấp lên các phiên bản mới nhất của TensorRT-LLM hoặc vLLM để bật hỗ trợ ReplaySSM cho các mô hình hybrid tương thích.
• Chạy benchmark các khối lượng công việc SSM hiện tại ở kích thước batch cao hơn để định lượng tiềm năng cải thiện thông lượng nhờ giảm lưu lượng bộ nhớ.
• Nghiên cứu các cấu hình giải mã suy đoán với độ dài bản nháp lớn hơn, vì ReplaySSM loại bỏ các hình phạt bộ nhớ trước đây liên quan đến việc hoàn tác.
• Giám sát lỗi lượng tử hóa nếu sử dụng trạng thái độ chính xác thấp, lưu ý rằng sự phân tách giữa điểm kiểm tra và bộ đệm giúp duy trì độ chính xác gần với mức FP32.