Phân tách Prefill-Decode: Tại sao suy luận LLM đang bắt đầu giống một hệ thống phân tán

Chào bạn, tôi là Shrijith Venkatramana. Tôi đang phát triển git-lrc, một công cụ đánh giá mã nguồn AI chạy trên mỗi commit. Hãy gắn dấu sao (Star) cho chúng tôi để giúp các nhà phát triển biết đến dự án này. Xin hãy dùng thử và chia sẻ phản hồi của bạn để cải thiện sản phẩm. **Phân tách Prefill-Decode: Vì sao suy luận LLM đang dần giống một hệ thống phân tán** Có một điều kỳ lạ khi vận hành một mô hình ngôn ngữ lớn (LLM). Một người dùng gửi 4.000 token và yêu cầu thêm 200 token. Bộ xử lý đồ họa (GPU) dành một loại tính toán để xử lý 4.000 token đó – và sau đó dành 200 bước tiếp theo để thực hiện một việc hoàn toàn khác. Tuy nhiên, trong một hệ thống vận hành thông thường, chúng ta thường đặt cả hai khối lượng công việc này trên cùng một GPU, dưới cùng một bộ lập lịch, tranh giành cùng một tài nguyên. Điều này đang dần trở nên ngày càng không hợp lý. Ý tưởng phân tách prefill-decode rất đơn giản: Chạy phần suy luận đọc lời nhắc (prompt) tách biệt với phần tạo ra câu trả lời. Điều này nghe có vẻ là một tối ưu hóa hiển nhiên khi bạn nhìn thấy nó. Nhưng để thực hiện đúng, cần liên quan đến bộ nhớ đệm KV (KV caches), bộ nhớ GPU, băng thông mạng, phân lô (batching), các mục tiêu mức độ dịch vụ về độ trễ (latency SLOs), lập lịch và một số yếu tố kinh tế học thú vị đáng ngạc nhiên. Ý tưởng này đã được khám phá một cách có hệ thống trong công trình Splitwise của Microsoft, được công bố tại ISCA 2024, và sau đó được phát triển thêm cho việc vận hành nhạy cảm về độ trễ bởi DistServe, được công bố tại OSDI 2024. Hãy xây dựng ý tưởng này từ những nguyên tắc cơ bản. **1. Một yêu cầu LLM thực chất chứa hai công việc rất khác nhau** Giả sử bạn gửi điều này đến một mô hình: *Giải thích cách thức hoạt động của kiểm soát tắc nghẽn TCP.* *Giả sử tôi đã hiểu định tuyến IP.* *Hãy cho tôi một lời giải thích chi tiết kèm theo ví dụ.* Hãy tưởng tượng rằng điều đó trở thành 25 token đầu vào và mô hình tạo ra 300 token đầu ra. Suy luận trông đại khái như sau: ``` ĐẦU VÀO | v +-----------------+ | PREFILL | | | | Xử lý tất cả | | các token đầu vào | +--------+--------+ | | Bộ nhớ đệm KV v +-----------------+ | DECODE | | | | token -> token | | -> token -> ... | +--------+--------+ | v ĐẦU RA ``` **Prefill** Mô hình xử lý toàn bộ lời nhắc. Nếu lời nhắc chứa 4.000 token, bộ biến đổi (transformer) có thể xử lý các token đó phần lớn song song. Đây là một phép tính lớn, dày đặc. GPU cực kỳ giỏi trong việc này. **Decode** Bây giờ mô hình tạo ra câu trả lời từng token một: ``` token 1 | v token 2 | v token 3 | v token 4 | v ... ``` Có một sự phụ thuộc tự hồi quy không thể tránh khỏi: P(token_n | token_1 ... token_(n-1)) Bạn thường không thể tạo token 400 trước khi biết token 399. Vì vậy, decode bao gồm nhiều phép toán GPU tương đối nhỏ. Điều này tạo ra một sự bất đối xứng quan trọng: Prefill cần thông lượng tính toán. Decode cần băng thông bộ nhớ và độ trễ có thể dự đoán được. Các nhà nghiên cứu của Microsoft là Pratyush Patel, Esha Choukse, Chaojie Zhang và các đồng nghiệp đã mô tả chính xác sự khác biệt này trong công trình Splitwise của họ. Họ nhận thấy việc xử lý lời nhắc là tốn nhiều tính toán, trong khi việc tạo token lại tốn nhiều bộ nhớ hơn và sử dụng kém hiệu quả các tài nguyên bộ tăng tốc đắt tiền. Đây là quan sát cơ bản đằng sau sự phân tách. **2. Phần đáng ngạc nhiên: cùng một GPU là một sự thỏa hiệp tồi** Hãy tưởng tượng một nhà hàng có hai loại khách hàng. Một khách hàng đến với 100 người và muốn đặt một đơn hàng lớn. Một khách hàng khác đến một mình và muốn một chiếc bánh sandwich cứ sau 30 giây trong mười phút tiếp theo. Nếu bạn buộc cả hai nhóm phải trải qua cùng một quy trình làm việc trong bếp, việc tối ưu hóa cho nhóm này sẽ làm hại nhóm kia. Việc vận hành LLM cũng có một vấn đề tương tự. Hãy xem xét một lô decode: yêu cầu A -> tạo mã thông báo tiếp theo yêu cầu B -> tạo mã thông báo tiếp theo yêu cầu C -> tạo mã thông báo tiếp theo yêu cầu D -> tạo mã thông báo tiếp theo ... Xử lý theo lô (batching) cực kỳ hữu ích trong trường hợp này. Bộ xử lý đồ họa (GPU) có thể tái sử dụng trọng số mô hình trên nhiều yêu cầu, do đó chi phí đọc trọng số được phân bổ. Bây giờ, một yêu cầu mới đến với một lời nhắc (prompt) dài 20.000 mã thông báo. Quá trình tiền xử lý (prefill) của nó là một hoạt động tính toán nặng. Nếu bộ lập lịch chèn quá trình tiền xử lý đó vào cùng một khối lượng công việc của GPU, nó có thể tạm thời tiêu thụ một lượng lớn tài nguyên tính toán và băng thông bộ nhớ. Những người dùng đang tạo mã thông báo không quan tâm đến việc yêu cầu mới có một lời nhắc khổng lồ. Họ chỉ trải nghiệm: mã thông báo mã thông báo mã thông báo <- tiền xử lý khổng lồ ... mã thông báo Độ trễ giữa các mã thông báo của họ tăng vọt. Đây là lý do tại sao TTFT và TPOT/ITL lại quan trọng. TTFT — Thời gian đến mã thông báo đầu tiên (Time To First Token): thời gian người dùng chờ đợi trước khi thấy bất kỳ điều gì. ITL — Độ trễ giữa các mã thông báo (Inter-Token Latency): thời gian người dùng chờ đợi giữa các mã thông báo được tạo. TPOT — Thời gian trên mỗi mã thông báo đầu ra (Time Per Output Token): liên quan chặt chẽ đến ITL trong các phép đo phục vụ. Một hệ thống có thể có thông lượng trung bình xuất sắc nhưng lại mang lại trải nghiệm tương tác tệ hại. Luận điểm chính của DistServe là quá trình tiền xử lý và giải mã can thiệp vào nhau đủ mạnh đến mức việc coi chúng là một khối lượng công việc đồng nhất khiến việc đồng thời đáp ứng các yêu cầu TTFT và TPOT trở nên khó khăn. 3. Điều gì thực sự đang được "phân tách"? Đây là nơi triển khai trở nên thú vị. Bạn không chia mạng thần kinh thành: GPU A: các lớp 1-40 GPU B: các lớp 41-80 Đó là song song hóa mô hình thông thường. Thay vào đó, bạn chạy về cơ bản cùng một mô hình trong hai nhóm khác nhau: Yêu cầu | v +----------------+ | Nhóm tiền xử lý | | | | GPU GPU GPU | +-------+--------+ | | Bộ nhớ đệm KV | v +----------------+ | Nhóm giải mã | | | | GPU GPU GPU | +-------+--------+ | v Mã thông báo Các nhân viên tiền xử lý xử lý lời nhắc. Họ tạo ra trạng thái trung gian cần thiết cho bộ giải mã: chủ yếu là bộ nhớ đệm khóa/giá trị (KV cache). Bộ nhớ đệm đó được chuyển đến một nhân viên giải mã. Nhân viên giải mã sau đó tiếp tục tạo. Vì vậy, kiến trúc trở thành: mạng / kết nối -----------------------------> GPU TIỀN XỬ LÝ GPU GIẢI MÃ ---------- ---------- lời nhắc | v bộ biến đổi | v Bộ nhớ đệm KV ------------------------> Bộ nhớ đệm KV | v tạo mã thông báo 1 | mã thông báo 2 | mã thông báo 3 Đây là lý do tại sao kỹ thuật này thú vị hơn nhiều so với việc chỉ chạy hai hàng đợi. Bộ nhớ đệm KV trở thành một đối tượng hệ thống phân tán. Và điều đó thay đổi mọi thứ. 4. Bộ nhớ đệm KV là thứ gắn kết toàn bộ hệ thống lại với nhau Durin