Perplexity hé lộ hạ tầng GPU: Cách Ivy, Tulip và ROSE vận hành mô hình pplx-embed

Chất lượng truy xuất trong một sản phẩm tìm kiếm AI bị giới hạn bởi hai yếu tố: chất lượng của mô hình embedding và chi phí vận hành mô hình đó trên một chỉ mục (index). Tuần này, đội ngũ kỹ thuật của Perplexity đã công bố bài viết "Fast Embeddings on GPUs", một bản tường thuật chuyên sâu về yếu tố thứ hai — cơ sở hạ tầng phục vụ (serving infrastructure) đằng sau pplx-embed và các mô hình xếp hạng được sử dụng trên Perplexity Search, Computer và nền tảng API.

Đội ngũ Perplexity cho biết việc suy luận (inference) embedding trên GPU đã đạt đến sự đồng nhất giữa các engine trên phần cứng Hopper và Blackwell thế hệ mới. Những cải tiến nằm ở runtime và các công cụ hỗ trợ xung quanh mô hình: quản lý CUDA graph, một cơ chế trừu tượng theo dõi kết quả bất đồng bộ (async) và đường dẫn yêu cầu (request path) được viết bằng Rust.

Hai mô hình lưu lượng, một engine

Perplexity định nghĩa việc phục vụ embedding thành hai khối lượng công việc. Batch embedding diễn ra khi xây dựng hoặc lập chỉ mục lại cơ sở dữ liệu vector, nơi thông lượng (throughput) giúp tối thiểu hóa chi phí. Online embedding diễn ra tại thời điểm truy vấn, nơi một truy vấn ngắn cần được chuyển đổi thành vector (embed) nhanh chóng. Scoring (chấm điểm) nằm ở giữa: sau khi tìm kiếm vector, các lô tài liệu lớn sẽ được xếp hạng, cân bằng cả hai yếu tố trên.

Quyết định then chốt là Perplexity không xây dựng một engine embedding riêng biệt. Vì các mô hình embedding là các Transformer nhỏ, batch embedding giống với quá trình prefill bị giới hạn bởi tính toán (compute-bound), còn online embedding (thường chỉ vài token) lại giống với quá trình decode bị giới hạn bởi bộ nhớ (memory-bound). Do đó, đội ngũ nghiên cứu đã tái sử dụng các kernel prefill và decode từ stack LLM của họ.

Ba dịch vụ xử lý một yêu cầu:

Tại sao bộ lập lịch (scheduler) lại được thiết kế đơn giản một cách có chủ đích

Tulip chọn các chuỗi theo thứ tự đến trước, phục vụ trước (first-come, first-served) trong khi các yêu cầu tích lũy. Sự đơn giản đó được chứng minh bằng một phép đo: đối với các mô hình embedding nhỏ ở độ dài chuỗi mà Perplexity phục vụ, chi phí tuyến tính của các lớp dense chiếm ưu thế hơn chi phí bậc hai của cơ chế attention. Do đó, độ trễ tỷ lệ thuận với số lượng token chứ không phải số lượng chuỗi. Khi một batch đã làm đầy GPU, khoảng 512 token trên một mô hình dưới một tỷ tham số, việc nhồi thêm các chuỗi không giúp cải thiện hiệu suất.

CUDA graphs và LazyTensors

Trên các batch nhỏ, việc khởi chạy kernel từ phía CPU có thể tốn kém hơn cả quá trình thực thi trên GPU. Perplexity xây dựng các CUDA graph toàn mô hình cho tất cả các mô hình embedding, gom mọi lệnh khởi chạy vào một lệnh gọi driver duy nhất. Vì các mô hình embedding nhỏ, điểm bùng phát nơi công việc trên GPU vượt quá chi phí khởi chạy đạt được ở các batch hàng nghìn token và hàng chục chuỗi. Một số triển khai attention chặn các graph toàn mô hình do phụ thuộc vào các đầu vào động từ phía host; Perplexity đã đóng góp các thay đổi ngược dòng (upstream) cho FlashInfer để cho phép thực hiện việc ghi lại (capture).

Các graph phải được ghi lại theo từng cấu hình, vì vậy số lượng token được đệm (pad) thành các bucket là bội số của 64 hoặc 256. Điều đó vẫn tạo ra hàng nghìn graph và mất vài phút để ghi lại cho mỗi mô hình. Giải pháp là lazy capture: mỗi cấu hình sẽ có một lần chạy khởi động (warmup) eager, sau đó kích hoạt ghi lại và phát lại ở lần truy cập thứ hai. Điều này ảnh hưởng đến độ trễ p99 khi khởi động nhưng giúp phân bổ thời gian thực hiện eager từ vài phút ra hàng giờ.

Phần thứ hai là LazyTensor, giúp theo dõi một bộ đệm host được khóa trang (page-locked) cộng với cudaMemcpyAsync và một sự kiện CUDA. Thay vì step chặn trên thiết bị, nó trả về một LazyTensor, cho phép một tác vụ async của Rust chờ đợi ở batch N trong khi CPU xếp hàng cho batch N+1.

Các kernel vẫn đóng vai trò quan trọng

ROSE hỗ trợ nhiều backend attention cho các đầu vào không đồng nhất (ragged inputs): FlashInfer 2, FlashInfer 3 và FlashAttention 4. Đội ngũ Perplexity báo cáo rằng FlashAttention 4 nhìn chung nhanh hơn, nhưng FlashInfer 3 lại vượt trội hơn trên các mô hình dựa trên Qwen ở độ dài chuỗi rất lớn, vì vậy việc lựa chọn backend được thực hiện tùy theo từng trường hợp. Đáng chú ý, khi phục vụ một mô hình embedding, ROSE không khởi tạo KV cache và điều phối đến các biến thể ragged attention để tránh việc đệm (padding).

Điểm chuẩn (Benchmarks)

Perplexity thực hiện benchmark so với vLLM v0.22.0 ở định dạng BF16 trên các trọng số thực và đầu vào được trích xuất từ đánh giá, với các lần chạy warmup để xác nhận độ lệch tương đồng cosine trong phạm vi 0.1%. Bốn bộ thử nghiệm được lập biểu đồ: embedding độ trễ thấp (batch 1; 128/512/4096 token), scoring độ trễ thấp (batch 5/25/50 tại 512 token), embedding thông lượng cao (batch 100, bốn tiến trình đồng thời) và embedding độ đồng thời cao (1 đến 16 yêu cầu đồng thời, bao gồm cả quá trình token hóa Ivy và chi phí mạng).

Những điểm chính cần lưu ý

Hãy xem các chi tiết kỹ thuật. Ngoài ra, đừng ngần ngại theo dõi chúng tôi trên Twitter và đừng quên tham gia cộng đồng 150k+ ML SubReddit của chúng tôi và đăng ký nhận bản tin (Newsletter). Khoan đã! Bạn có dùng Telegram không? Bây giờ bạn cũng có thể tham gia cùng chúng tôi trên Telegram.

Bạn cần hợp tác với chúng tôi để quảng bá GitHub Repo, trang Hugging Face, bản phát hành sản phẩm hoặc hội thảo trực tuyến (Webinar), v.v.? Hãy kết nối với chúng tôi.

Asif Razzaq là CEO của Marktechpost Media Inc. Là một doanh nhân và kỹ sư có tầm nhìn, Asif cam kết khai thác tiềm năng của Trí tuệ nhân tạo vì lợi ích xã hội. Nỗ lực gần đây nhất của ông là ra mắt nền tảng truyền thông về Trí tuệ nhân tạo, Marktechpost, nổi bật với việc đưa tin chuyên sâu về tin tức học máy (machine learning) và học sâu (deep learning), vừa đảm bảo tính kỹ thuật vừa dễ hiểu đối với đông đảo khán giả. Nền tảng này tự hào với hơn 2 triệu lượt xem mỗi tháng, minh chứng cho sự phổ biến của nó đối với độc giả.