Tối ưu hóa Prefill ngữ cảnh dài trên nhiều nút GPU thế hệ cũ

Tổng quan

Khi các quy trình làm việc dựa trên tác tử ngày càng phổ biến trong các ứng dụng AI, mô hình cần xử lý độ dài ngữ cảnh dài hơn đáng kể. Sự thay đổi này đặt ra một thách thức quan trọng: duy trì các mục tiêu cấp độ dịch vụ (SLO) trong khi tối đa hóa việc sử dụng hệ thống. Đặc biệt, tối ưu hóa hiệu suất prefill rất quan trọng để giảm Time-To-First-Token (TTFT) cho các yêu cầu ngữ cảnh dài. Để giải quyết vấn đề này, chúng tôi đã phát triển SLO-Driven Prefill Engine (SLOPE).

SLOPE là một prefill engine chuyên dụng áp dụng các kỹ thuật context parallelism (Ulysses + Ring Attention) trên các cụm GPU đa nút để tối ưu hóa hướng SLO cho các đầu vào ngữ cảnh dài. Các LLM engine hiện tại bị hạn chế cơ bản trong việc giảm TTFT xuống dưới một ngưỡng nhất định khi xử lý các prompt đầu vào dài, vì tính toán prefill bị giới hạn trong một nút duy nhất. SLOPE vượt qua giới hạn này bằng cách song song hóa tính toán prefill trên nhiều nút, đưa TTFT xuống thấp hơn nhiều so với ngưỡng SLO mục tiêu trong khi tối đa hóa số lượng yêu cầu người dùng đồng thời.

SLOPE có thể được tích hợp với các LLM engine hiện có ở chế độ prefill/decode disaggregation, hoạt động như một prefill worker chuyên dụng.

Giới thiệu Context Parallelism

Tensor parallelism (TP) là kỹ thuật song song hóa được sử dụng rộng rãi nhất trong suy luận LLM. TP phân phối các tham số lớp tuyến tính trên các thiết bị, chia attention heads trên các thiết bị và yêu cầu giao tiếp AllReduce. Tuy nhiên, việc đơn giản tăng kích thước TP để giảm TTFT có những hạn chế:

• Các phép toán GEMM không hiệu quả: Việc phân chia tham số quá mức dẫn đến các phép toán GEMM không hiệu quả.
• Tính toán dư thừa: Khi số lượng attention heads nhỏ hơn kích thước TP, tính toán dư thừa xảy ra.
• Chi phí giao tiếp: Khi kích thước TP tăng, AllReduce liên nút trở nên cần thiết, làm tăng đáng kể chi phí giao tiếp.

Context parallelism phân chia prompt đầu vào theo chiều chuỗi. Đặc biệt đối với các kịch bản ngữ cảnh dài, context parallelism hiệu quả hơn tensor parallelism nhờ tính song song quy mô lớn có sẵn trong chiều chuỗi. Hầu hết các phép toán, ngoại trừ lớp attention, không có sự phụ thuộc giữa các token và có thể được xử lý hoàn toàn độc lập. Sự phụ thuộc giữa các token trong lớp attention có thể được giải quyết thông qua hai phương pháp:

• Ulysses: thực hiện giao tiếp all-to-all trên query, key và value trước khi tính toán attention, phân chia chúng theo chiều head để cho phép tính toán độc lập.
• Ring Attention: thực hiện tính toán attention với query, key và value một phần trên mỗi thiết bị trong khi đồng thời gửi/nhận key và value đến/từ các thiết bị lân cận.

Hai phương pháp này có những đánh đổi bổ sung cho nhau. Việc áp dụng Ulysses bị hạn chế khi sử dụng Grouped Query Attention (GQA) hoặc Multi-Head Latent Attention (MLA) do số lượng key và value heads giảm. Mặt khác, Ring Attention tạo ra khối lượng giao tiếp lớn hơn Ulysses. Tuy nhiên, với độ dài chuỗi lớn hơn, tính toán attention trở nên chiếm ưu thế theo tỷ lệ, cho phép giao tiếp được ẩn bởi tính toán attention. Do đó, chúng nên được kết hợp phù hợp dựa trên độ dài ngữ cảnh và yêu cầu SLO.

Khi áp dụng context parallelism, tất cả các thiết bị giữ các tham số mô hình trùng lặp, do đó mức sử dụng bộ nhớ tăng so với TP. Chúng tôi giải quyết vấn đề này bằng cách áp dụng pipeline parallelism cùng với context parallelism.

Đánh giá

Chúng tôi đánh giá hiệu suất prefill của SLOPE trên các cấu hình song song hóa khác nhau kết hợp Ulysses, Ring Attention và Pipeline Parallel trong các tổ hợp khác nhau, và so sánh với SGLang sử dụng TP=8 trên cụm 4 nút AMD MI250 (8 thiết bị mỗi nút). Trong tất cả các thí nghiệm, chúng tôi sử dụng mô hình openai/gpt-oss-120b để đánh giá hiệu suất prefill, và kiểm tra với ba độ dài ngữ cảnh: 32K, 64K và 100K.

Như bảng dưới đây cho thấy, SLOPE đáp ứng các SLO TTFT rất thấp trên tất cả các độ dài ngữ cảnh, trong khi khả năng đáp ứng các yêu cầu này của SGLang thì khác nhau tùy trường hợp.

Hình 1 cho thấy các phép đo TTFT và thông lượng trên các cấu hình song song hóa và mức độ đồng thời khác nhau. Kết quả xác nhận rằng SLOPE đạt được SLO độ trễ rất thấp dưới 1 giây TTFT cho cả ngữ cảnh 32K và 64K. So với SGLang sử dụng TP=8, SLOPE đạt thông lượng cao hơn ở cùng TTFT, nghĩa là nó có thể xử lý nhiều yêu cầu người dùng đồng thời hơn trong khi vẫn đáp ứng các yêu cầu SLO. Ngay cả ở độ dài ngữ cảnh 100K, SLOPE đạt TTFT dưới 2 giây, trong khi SGLang sử dụng TP=8 bị giới hạn ở khoảng 9 giây. Điều này cho thấy SLOPE cung cấp một giải pháp thực tế cho suy luận LLM ngữ cảnh dài.

SLOPE Engine trong các cụm không đồng nhất

Việc sử dụng SLOPE engine không chỉ đơn thuần là giảm độ trễ phản hồi và tăng thông lượng token cho suy luận ngữ cảnh dài. Nó còn cung cấp cách để sử dụng hiệu quả các máy chủ GPU thế hệ cũ đã tồn tại trong các trung tâm dữ liệu.

Nói chung, suy luận LLM khó song song hóa trên nhiều nút, đây là một trong những lý do tại sao các GPU ngày càng nhanh hơn, dung lượng lớn hơn và đắt hơn tiếp tục được giới thiệu để xử lý các LLM quy mô lớn. Tuy nhiên, với SLOPE engine, ít nhất trong giai đoạn prefill, nhiều máy chủ GPU thế hệ cũ có thể được tổng hợp để đạt được độ trễ đủ thấp và thông lượng cao.

Để chứng minh rằng kiểu sử dụng này là khả thi, chúng tôi đã đánh giá SLOPE trên cụm GPU AMD MI250. Trong một hệ thống mà GPU AMD MI250 (tương đương với NVIDIA A100) và GPU MI300-series mới hơn cùng tồn tại, GPU MI250 có thể được phân bổ cho SLOPE, trong khi GPU MI300-series thực hiện decoding bằng inference engine truyền thống, cho phép hiệu quả hệ thống tổng thể cao hơn. Các khả năng này sẽ được hỗ trợ thông qua MoAI Inference Framework từ Q1 2026.