Transformers documentation
Kiến trúc xếp lô liên tục
Kiến trúc xếp lô liên tục
Batching truyền thống xử lý một nhóm yêu cầu cố định cùng lúc và chờ đến khi yêu cầu chậm nhất hoàn thành trước khi bắt đầu nhóm tiếp theo. Điều này khiến GPU ở trạng thái rảnh giữa các batch.
Với continuous batching, ở mỗi bước sinh, trình lập lịch sẽ kiểm tra các yêu cầu đã hoàn thành và thay thế chúng ngay lập tức bằng các yêu cầu đang chờ. Các yêu cầu ngắn sẽ bị loại bỏ ngay khi hoàn tất, và GPU luôn được sử dụng liên tục. Kết quả là thông lượng tăng đáng kể và độ trễ trung bình giảm xuống.
Chu kỳ vòng đời của yêu cầu
Một yêu cầu trải qua bốn trạng thái từ khi được gửi đến khi hoàn tất.
WAIT IN QUEUE LOAD PROMPT INTO KV STREAM OUTPUT DONE
┌────────────────┐ ┌─────────────────────┐ ┌─────────────────┐ ┌────────┐
│ PENDING │───▶│ PREFILLING │───▶│ DECODING │───▶│FINISHED│
│ ○ ○ ○ · · · │ │ ████████░░░░░░░░░░ │ │ → → → → · · │ │ ✓ │
└────────────────┘ └─────────────────────┘ └─────────────────┘ └────────┘
others ahead prompt chunks +1 token/step- Đang chờ — yêu cầu đã được đưa vào hàng đợi và đang chờ được lên lịch.
- Prefilling — các token của prompt được xử lý trong một bước truyền thẳng.
- Giải mã — các token đầu ra được tạo ra lần lượt từng cái một.
- Hoàn tất — quá trình tạo đã kết thúc và kết quả đã sẵn sàng.
Một yêu cầu được chuyển từ trạng thái chờ sang trạng thái prefill khi bộ lập lịch tìm thấy đủ ngân sách token và không gian cache để chấp nhận yêu cầu đó. Nếu prompt quá dài để vừa trong một bước duy nhất, bộ lập lịch sẽ chia nó thành nhiều lần truyền tiến (chunk prefill).
Chunked prefill
Việc xử lý một prompt dài trong một bước duy nhất tốn kém vì nó chặn các yêu cầu khác khỏi việc tạo ra nội dung. Chunked prefill giải quyết vấn đề này bằng cách chia các prefill lớn thành nhiều bước.
Blocking (long prefill blocks the batch) Step: 1 2 3 4 Req A [▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓] ← one huge prefill Req B ················ wait ················ Req C ················ wait ················ Chunked (same prompt split and interleave with decode) Step: 1 2 3 4 Req A [▓▓ pre] [▓▓ pre] [▓▓ pre] [→ dec] Req B [→ dec] [→ dec] [▓▓ pre] [→ dec] Req C [→ dec] · idle · [→ dec] · idle ·
Khi prompt của một yêu cầu mới vượt quá ngân sách token khả dụng (được đặt bởi max_batch_tokens trong ContinuousBatchingConfig), trình lập lịch sẽ xử lý càng nhiều token càng tốt và giữ lại phần còn lại. Ở các bước tiếp theo, nó tiếp tục từ vị trí đã dừng, xen kẽ công việc prefill với các bước decode đang diễn ra. Điều này giúp batch duy trì hiệu suất và giảm thời gian đến token đầu tiên cho các yêu cầu khác.
Trình lập lịch
Trình lập lịch quyết định các yêu cầu nào tham gia mỗi bước truyền trực tiếp dựa trên hai ngân sách và một giới hạn yêu cầu.
- Ngân sách token — số lượng token truy vấn tối đa được xử lý trong một lần truyền tiến, được thiết lập bởi
max_batch_tokenstrong ContinuousBatchingConfig. - Ngân sách cache — tổng số trang KV có thể được đọc trong một lượt, bị giới hạn bởi kích thước cache tổng thể.
- Giới hạn yêu cầu — số lượng yêu cầu tối đa trong một lần truyền tiến, được đặt bởi
max_requests_per_batch. Tensor logits tỷ lệ với kích thước từ vựng, do đó giới hạn này ràng buộc bộ nhớ đỉnh cho các mô hình có từ vựng lớn.
FIFO
FIFOScheduler là trình lập lịch mặc định. Nó lấp đầy lô theo thứ tự ưu tiên. Các yêu cầu giải mã đang hoạt động được xử lý trước, sau đó là các yêu cầu prefill đang hoạt động, và cuối cùng là các yêu cầu chờ mới theo thứ tự đến. Điều này tối đa hóa thông lượng bằng cách đẩy các yêu cầu hiện có qua trước khi đưa các yêu cầu mới vào.
PrefillFirst
PrefillFirstScheduler hoàn thành các thao tác prefill dạng chunk trước khi tiếp tục decode. Điều này giúp giảm phân mảnh cho các khối công việc chủ yếu là các prompt dài.
Đặt scheduler trong ContinuousBatchingConfig.
cb_config = ContinuousBatchingConfig(scheduler_type="prefill_first")Quản lý bộ nhớ
Bộ đệm KV được phân trang để các yêu cầu có độ dài khác nhau có thể chia sẻ một vùng nhớ cố định mà không bị phân mảnh. Một vùng nhớ chung gồm các khối bộ nhớ, được triển khai dưới dạng CachePool, cung cấp các khối đó cho bộ cấp phát của từng loại attention, sau đó bộ cấp phát sẽ chia chúng thành các khối và trang.
Shared pool (first two chunks reserved for padding / never allocated)
chunk N [ ======================================== ]
| |
full attn: [ FULL BLOCK 0 | FULL BLOCK 1 | … ]
[ page L0 | page L0 | … ] (one page per full layer)
sliding: [ ----------- SLIDING BLOCK 0 ----------- ]
[ page L1 | page L2 | page L3 | … ] (one page per sliding layer)
page = page_size tokens of KV for one layer
block = one page per layer of that attention type
chunk = pool unit sized so full and sliding blocks tile the same bytesContinuous batching chỉ hỗ trợ hai loại attention: full_attention (FullAttentionCacheAllocator) và sliding_attention (SlidingAttentionCacheAllocator). Các mô hình hybrid sẽ nhận một allocator cho mỗi loại. Loại attention được lấy từ mô hình (layer_types / sliding_window), không phải từ ContinuousBatchingConfig.
Các lớp trượt (sliding layers) chỉ giữ một cửa sổ cố định các token gần nhất bằng cách tái sử dụng một tập hợp nhỏ các trang. Chúng không hỗ trợ bảng khối kernel (kernel block tables) hay chia sẻ khối (block sharing).
Bộ đệm dành trước hai khối phía trước các khối dữ liệu cho các thao tác đọc và ghi đệm mà các yêu cầu không bao giờ sử dụng. Những khối đó không bao giờ được cấp phát cho các yêu cầu.
Kích thước bộ nhớ đệm
Trình quản lý suy luận số lượng chunk có thể chứa khi khởi động dựa trên bộ nhớ GPU trống. Nó tính đến các tensor KV, attention mask, activations và các chỉ số ghi sổ, sau đó điều chỉnh kích thước pool để nằm trong max_memory_percent của bộ nhớ khả dụng. Khi không được đặt, giá trị mặc định sẽ được xác định theo phần cứng và các tham số continuous batching. MPS dành ít không gian hơn cho cache, nơi CPU và GPU chia sẻ bộ nhớ, để tránh swapping. num_blocks biểu diễn sức chứa dưới dạng các khối mô hình nguyên có page_size token khi bạn đặt nó một cách tường minh.
Bạn có thể cố định các giá trị trong ContinuousBatchingConfig.
cb_config = ContinuousBatchingConfig(
page_size=256,
num_blocks=4096,
max_batch_tokens=512,
max_memory_percent=0.8,
)Tiếp nhận
Trước khi một yêu cầu tham gia vào lô, trình lập lịch kiểm tra xem có đủ dung lượng trống cho mọi bộ cấp phát kiểu chú ý hay không. Nếu bất kỳ bộ cấp phát nào không đủ, yêu cầu sẽ bị từ chối và không có gì được cấp phát.
Để tránh việc lấp đầy bộ nhớ đệm đến mức offloading là lựa chọn duy nhất, trình lập lịch thực thi safety_margin. Khi dung lượng bộ nhớ đệm trống giảm xuống dưới safety_margin, các prefill mới sẽ bị giữ lại và chỉ các decode đang hoạt động tiếp tục. Trình lập lịch FIFO mặc định là 0.15 (giữ lại 15% dung lượng dự phòng), trong khi prefill_first mặc định là 0.0. Đặt safety_margin trong ContinuousBatchingConfig để điều chỉnh, trong đó 0 sẽ vô hiệu hóa biên an toàn.
Bộ nhớ đệm tiền tố
Khi hai yêu cầu chia sẻ một tiền tố prompt, chúng có thể chia sẻ các khối chứa KV cho tiền tố đó. Mỗi khối hoàn thành được băm theo nội dung. Một yêu cầu sau có tiền tố khớp sẽ tái sử dụng khối đó và bỏ qua bước prefill cho các token tương ứng. Các khối được chia sẻ được đếm tham chiếu và chỉ được trả về pool trống khi tất cả các yêu cầu đang sử dụng chúng đã hoàn tất.
allow_block_sharing mặc định là True. Chia sẻ tiền tố chỉ được kích hoạt khi mọi bộ cấp phát đều hỗ trợ chia sẻ, do đó chỉ các mô hình chỉ dùng full-attention mới sử dụng tính năng này. Các mô hình thuần sliding và hybrid sẽ tắt tính năng này. Đặt allow_block_sharing=False cho các tác vụ có prompt ngắn và sinh ra chuỗi dài, nơi chi phí quản lý vượt quá lợi ích tiết kiệm.
Gỡ bỏ
Các khối đã giải phóng không bị xóa và vẫn ở trạng thái “đã khởi tạo” với nội dung và hash nguyên vẹn để một lần khớp tiền tố sau đó có thể tái sử dụng chúng. Khi pool chưa khởi tạo sắp cạn, trình quản lý sẽ hạ cấp một cách lười biếng khối đã khởi tạo gần nhất trở về trạng thái chưa khởi tạo, đồng thời loại bỏ hash của chúng. Bộ nhớ đệm tiền tố là một lớp cố gắng hết sức mà trình cấp phát sẽ loại bỏ để tiếp tục phục vụ các yêu cầu mới.
Đặt lại mềm
Soft reset là đường dẫn dự phòng cho offloading. Khi trình quản lý không thể sao chép KV cache của một yêu cầu sang pool hoán đổi CPU, do offloading CPU bị vô hiệu hóa hoặc pool đã đầy, nó sẽ thực hiện soft reset cho yêu cầu đó. Nó nối các token đã tạo của yêu cầu vào prompt gốc, giải phóng các khối cache của nó và đưa nó trở lại hàng đợi. Phần Offloading nêu rõ những yêu cầu nào được chọn.
Khi bộ đệm có lại không gian, yêu cầu sẽ tiếp tục từ nơi nó dừng lại, với lịch sử sinh được mã hóa trong prompt.
Biểu đồ CUDA
Mỗi bước sinh đều phát sinh chi phí xử lý trên CPU, từ việc tập hợp batch đến việc điều phối các kernel GPU và đọc kết quả. CUDA graphs loại bỏ chi phí này bằng cách ghi lại toàn bộ trình tự thực thi trên GPU một lần và phát lại nó cho các batch có hình dạng khớp nhau.
Vì hình dạng của mỗi lô thay đổi ở mỗi bước, hệ thống liên tục xử lý điều này bằng cách đệm và lưu trữ bộ nhớ đệm.
- Chiều dài truy vấn được đệm lên bội số gần nhất của
q_padding_interval_size. - Chiều dài KV được đệm lên bội số gần nhất của
kv_padding_interval_size. - Các đồ thị được ghi lại được lưu trong bộ nhớ đệm.
Khi hình dạng đã đệm của một lô khớp với một đồ thị được lưu trong bộ nhớ đệm, đồ thị sẽ được phát lại mà không cần bất kỳ lệnh điều phối nào trên CPU. Các hình dạng mới sẽ kích hoạt việc bắt giữ một đồ thị mới.
CUDA graphs yêu cầu luồng điều khiển tĩnh và không tương thích với attention masks. Chúng được tự động phát hiện theo mặc định và bị vô hiệu hóa khi các điều kiện không được đáp ứng.
Đóng gói đồng bộ
Async batching sử dụng hai cặp bộ đệm I/O và hai luồng CUDA để chồng chéo công việc giữa CPU và GPU.
Sequential CPU ── [prep N] ·········idle········· [prep N+1] ·········idle········· GPU ·········idle········· [compute N] ·········idle········· [compute N+1] Async CPU ── [prep N] [prep N+1] [prep N+2] ── GPU ── ········ [compute N] [compute N+1] ──
Trong khi GPU tính toán batch N, CPU chuẩn bị batch N+1. Tuy nhiên, việc chồng chéo hai quá trình này đòi hỏi VRAM gấp đôi.
Async batching yêu cầu CUDA graphs phải được kích hoạt, vì việc phát lại đồ thị (graph replay) cung cấp các địa chỉ tensor ổn định cần thiết để việc chồng chéo luồng (stream overlap) diễn ra chính xác.
Offloading
Các yêu cầu tạo ra đầu ra dài có thể tăng dần cho đến khi bộ đệm KV không còn chỗ để cấp phát khối tiếp theo. Thay vì gặp lỗi, trình quản lý sẽ chuyển các yêu cầu đang hoạt động trở lại hàng chờ để giải phóng các khối cho các yêu cầu còn lại.
Offloading được điều khiển theo nhu cầu. Trình lập lịch báo cáo các yêu cầu stalled, tức các yêu cầu đang hoạt động không phân bổ được các khối cần thiết, kèm theo số khối mà mỗi yêu cầu đòi hỏi. Trình quản lý sẽ offload số lượng yêu cầu ít nhất cần thiết để giải phóng đủ khối nhằm đáp ứng nhu cầu đó. Việc offload được thực hiện theo thứ tự từ yêu cầu mới nhất trước, nhờ đó các yêu cầu đã được lập lịch vào batch hiện tại vẫn giữ được cache, và yêu cầu hoạt động còn lại cuối cùng sẽ không bao giờ bị offload.
Hai tình huống kích hoạt việc offload.
- Một batch được lên lịch, nhưng một số yêu cầu đang hoạt động bị tắc nghẽn. Trình quản lý sẽ offload đủ số lượng các yêu cầu bị tắc nghẽn để batch tiếp theo có thể cấp phát các block cho phần còn lại.
- Bộ đệm đã đầy và không có lô nào có thể được lên lịch. Trình quản lý sẽ chuyển các yêu cầu ra ngoài và thử lại việc lên lịch trong một vòng lặp trong cùng một bước, cho đến khi một lô vừa đủ hoặc không còn yêu cầu nào có thể chuyển ra ngoài. Nó không còn chờ đến lần truyền tiến tiếp theo để phục hồi nữa.
Một yêu cầu bị offload sẽ đi theo một trong hai đường dẫn. Khi cpu_offload_space lớn hơn 0.0, trình quản lý sẽ sao chép các khối KV cache của yêu cầu vào một bộ đệm CPU pinned đã được phân bổ trước, gom tất cả các yêu cầu vừa với swap pool thành một lần chuyển dữ liệu duy nhất. Sau khi không gian cache GPU trở nên khả dụng, trình quản lý sẽ sao chép các khối này trở lại GPU và tiếp tục yêu cầu mà không cần tính toán lại prompt và các token đã sinh ra. Các yêu cầu không vừa với pool, hoặc mọi yêu cầu bị offload khi CPU offloading bị tắt, sẽ chuyển sang soft reset.
Các bước tiếp theo
- Bài đăng trên blog Continuous batching đề cập đến KV caching, chunked prefill và dynamic scheduling kèm theo các con số benchmark hiệu năng.
- Để xem các ví dụ sử dụng, hãy xem tài liệu Continuous batching.