Cơ chế chuyển KV theo mô hình push của vLLM giúp giảm độ trễ trong suy luận phân tách

Amazon Web Services đã giới thiệu cơ chế chuyển key-value (KV) theo mô hình push cho engine suy luận mã nguồn mở vLLM. Bản cập nhật này nhắm vào các kiến trúc prefill-decode phân tách, với mục tiêu giảm đáng kể thời gian phản hồi token đầu tiên (TTFT) và cải thiện thông lượng trong quá trình tạo token. Tính năng này đã được kiểm chứng trên các instance Amazon EC2 P5en sử dụng mô hình Qwen3-32B.

Chuyện gì đã xảy ra

Quá trình suy luận của mô hình ngôn ngữ lớn bao gồm hai giai đoạn riêng biệt: prefill, xử lý prompt đầu vào để xây dựng cache KV, và decode, tạo ra các token đầu ra bằng cách sử dụng cache đó. Trong mô hình phục vụ phân tách (disaggregated serving), hai giai đoạn này chạy trên các instance engine riêng biệt để cho phép mở rộng và tối ưu hóa độc lập. Tuy nhiên, sự phân tách này đòi hỏi phải chuyển cache KV từ instance prefill sang instance decode trước khi bắt đầu quá trình tạo sinh, khiến hiệu quả của việc chuyển dữ liệu trở nên quan trọng đối với tổng thể độ trễ.

Trước đây, vLLM sử dụng phương pháp pull thông qua connector NVIDIA Inference Xfer Library (NIXL). Ở chế độ mặc định này, instance decode chờ prefill hoàn tất, nhận metadata từ router, cấp phát bộ nhớ, sau đó mới pull dữ liệu. Các bước này diễn ra tuần tự sau khi prefill kết thúc, gây thêm độ trễ đáng kể vào đường dẫn tới hạn (critical path). Connector mới dựa trên cơ chế push thay đổi luồng hoạt động này bằng cách chồng lấn các tác vụ phối hợp với chính quá trình tính toán prefill.

Các kỹ sư của AWS đã phát triển và thử nghiệm tính năng này trên các cụm được kết nối bằng AWS Elastic Fabric Adapter (EFA), hỗ trợ Remote Direct Memory Access (RDMA). Mặc dù ban đầu được xây dựng cho các instance AWS Trainium, đóng góp này không phụ thuộc vào loại accelerator cụ thể và cũng hoạt động tốt với GPU. Các bài kiểm tra benchmark cho thấy dưới tải cao, chế độ push có thể giảm TTFT lên đến ba lần so với chế độ pull, đồng thời cải thiện thời gian cho mỗi token đầu ra (TPOT) khoảng 30 phần trăm.

Cách thức hoạt động

Kiến trúc push loại bỏ các nút thắt cổ chai do tuần tự hóa bằng cách cho phép instance decode đăng ký trước các block bộ nhớ đích. Thay vì chờ prefill hoàn tất mới chuẩn bị nhận dữ liệu, phía decode gửi thông báo đăng ký tới instance prefill ngay khi nó cấp phát các block cục bộ. Điều này có nghĩa là đích đến đã được xác định trước khi dữ liệu nguồn sẵn sàng.

Ở phía prefill, một thread nền chuyên dụng xử lý việc chuẩn bị chuyển giao. Khi quá trình tính toán prefill hoàn tất và nhận được thông báo đăng ký từ decoder, vòng lặp engine chính sẽ bàn giao các block KV cho thread nền này. Thread thiết lập kết nối NIXL, ánh xạ các ID block và thực hiện thao tác NIXL WRITE để sao chép dữ liệu trực tiếp vào bộ nhớ đã đăng ký của decoder. Cách làm này giữ cho các công việc nặng nề như chuẩn bị descriptor và gửi yêu cầu chuyển giao nằm ngoài các vòng lặp chính nhạy cảm về độ trễ của cả hai instance.

Hệ thống sử dụng giao thức rendezvous để xử lý các chênh lệch về thời điểm. Nếu thông báo đăng ký đến trước khi prefill hoàn tất, producer sẽ giữ lại nó. Nếu prefill hoàn tất trước, nó sẽ giữ các block cho đến khi nhận được thông báo đăng ký. Sự kiện nào xảy ra thứ hai sẽ kích hoạt việc chuyển giao ngay lập tức. Decoder đếm các thông báo hoàn tất từ mỗi rank producer và chỉ tiến hành tạo token khi toàn bộ dữ liệu dự kiến đã đến nơi.

Chi tiết chính

• Hiệu suất tăng: Giảm TTFT lên đến 3 lần và cải thiện TPOT 30% cho mô hình Qwen3-32B trong các triển khai 1P1D ở tốc độ yêu cầu cao.
• Phần cứng đã kiểm tra: Instance Amazon EC2 P5en với EFA, sử dụng GPUDirect RDMA để chuyển bộ nhớ trực tiếp từ GPU sang GPU.
• Cơ chế: Sử dụng các thao tác NIXL WRITE do instance prefill khởi tạo, thay vì NIXL READ do instance decode khởi tạo.
• Threading: Một thread nền trên worker prefill xử lý việc thiết lập chuyển giao, ngăn chặn tình trạng nghẽn trong vòng lặp thực thi mô hình chính.
• Đánh đổi: Chế độ push yêu cầu chọn instance decode sớm hơn trong quy trình, điều này có thể làm giảm tính linh hoạt cho việc định tuyến dựa trên tải so với chế độ pull.
• Cấu hình: Người dùng bật tính năng này bằng cách đặt cả server prefill và decode sử dụng NixlPushConnector thay vì NixlConnector mặc định.

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

Đối với các kỹ sư xây dựng dịch vụ LLM có thông lượng cao, phục vụ phân tách là một chiến lược phổ biến để quản lý việc sử dụng tài nguyên hiệu quả. Tuy nhiên, chi phí di chuyển các cache KV lớn giữa các node có thể triệt tiêu lợi ích nếu không được xử lý cẩn thận. Mô hình pull truyền thống giới thiệu một khoảng thời gian chờ đồng bộ, trong đó instance decode nhàn rỗi, chờ metadata rồi mới pull dữ liệu. Độ trễ này trở nên rõ rệt khi độ dài ngữ cảnh tăng và tốc độ yêu cầu cao hơn, ảnh hưởng trực tiếp đến độ trễ mà người dùng cảm nhận.

Bằng cách chuyển công việc phối hợp sang một thread nền và chồng lấn nó với quá trình tính toán prefill, mô hình push đảm bảo rằng instance decode sẵn sàng nhận dữ liệu ngay khi nó có sẵn. Điều này giảm thời gian nhàn rỗi trên đường dẫn tới hạn, dẫn đến thời gian phản hồi nhanh hơn cho người dùng. Nó cũng ngăn worker decode bị sa lầy bởi các tác vụ thiết lập chuyển giao, cho phép duy trì tốc độ tạo token ổn định ngay cả dưới mức độ đồng thời cao.

Sự phát triển này nhấn mạnh tầm quan trọng của các tối ưu hóa hạ tầng ở cấp thấp trong việc phục vụ AI. Khi các mô hình trở nên lớn hơn và ngữ cảnh dài hơn, hiệu quả của việc di chuyển dữ liệu giữa các node tính toán trở nên quan trọng không kém tốc độ của bản thân các phép nhân ma trận. Các kỹ sư triển khai vLLM trên AWS hoặc hạ tầng tương tự hỗ trợ RDMA nên cân nhắc tùy chọn này để tối đa hóa hiệu suất.

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

• Đánh giá triển khai vLLM hiện tại của bạn để xem liệu phục vụ prefill-decode phân tách có phù hợp với khối lượng công việc của bạn hay không.
• Thử nghiệm NixlPushConnector trong môi trường staging với các mẫu lưu lượng thực tế để đo lường mức độ cải thiện TTFT.
• Đảm bảo hạ tầng của bạn hỗ trợ RDMA, chẳng hạn như AWS EFA, để tận dụng các tính năng truy cập bộ nhớ trực tiếp.
• Giám sát sự đánh đổi giữa việc chọn instance decode sớm và nhu cầu cân bằng tải động trong cấu hình router của bạn.
• Tắt prefix caching trong quá trình benchmarking để cô lập hiệu suất chuyển giao, vì caching có thể che giấu chi phí chuyển giao.
• Xem xét tài liệu vLLM để biết các cờ cấu hình cụ thể cần thiết để chuyển từ chế độ pull sang push.