Embedding nhanh trên GPU
Tìm kiếm nhanh chóng và chính xác đóng vai trò quan trọng đối với toàn bộ Perplexity, từ Tìm kiếm (Search) và Máy tính (Computer) đến Nền tảng API (API Platform) của chúng tôi. Đằng sau hậu trường, công việc nặng nhọc được thực hiện bởi các mô hình embedding và xếp hạng, giúp hệ thống của chúng tôi xác định các kết quả phù hợp nhất cho một truy vấn cụ thể. Chúng tôi đạt được trạng thái tiên tiến
Tìm kiếm nhanh chóng và chính xác đóng vai trò quan trọng đối với toàn bộ Perplexity, từ Tìm kiếm (Search) và Máy tính (Computer) đến Nền tảng API (API Platform) của chúng tôi. Đằng sau hậu trường, công việc nặng nhọc được thực hiện bởi các mô hình embedding và xếp hạng (ranking), giúp hệ thống của chúng tôi xác định các kết quả phù hợp nhất cho một truy vấn cụ thể. Chúng tôi đạt được chất lượng và độ trễ tiên tiến nhất bằng cách tự huấn luyện và phục vụ các mô hình của riêng mình, chẳng hạn như pplx-embed.
Bài viết này trình bày cái nhìn chi tiết bên trong cơ sở hạ tầng phục vụ của Perplexity dành cho loại mô hình đặc biệt này. Chúng tôi thảo luận về các kỹ thuật của mình để giải quyết hiệu quả các nhu cầu suy luận của tính năng tìm kiếm gốc bằng AI (AI-native search), cho phép tạo mẫu và đánh giá mô hình nhanh chóng đồng thời vận hành chỉ mục tìm kiếm quy mô exabyte. Các kỹ thuật này cùng nhau mở rộng ranh giới Pareto về chất lượng và hiệu quả tìm kiếm, giúp chúng tôi phục vụ các agent và người dùng với kết quả tốt nhất có thể ở mức chi phí và độ trễ thấp nhất.
Embedding cho Tìm kiếm
Trong một thiết lập tìm kiếm điển hình, các tài liệu đã được lập chỉ mục được ánh xạ tới một không gian véc-tơ số chiều cao bằng cách sử dụng một mô hình embedding và được lưu trữ trong cơ sở dữ liệu véc-tơ. Bằng cách nhúng một truy vấn sử dụng cùng một mô hình, các tài liệu tương tự có thể được định vị bằng cách tìm các véc-tơ gần nhất với véc-tơ của truy vấn. Điều này tạo ra hai mẫu lưu lượng truy cập khác nhau để một engine suy luận phục vụ:
- Batch Embedding: khi xây dựng, mở rộng hoặc đánh chỉ mục lại cơ sở dữ liệu, các tài liệu lớn phải được nhúng vào không gian véc-tơ, giúp tối đa hóa thông lượng để giảm thiểu chi phí.
Sau bước tìm kiếm véc-tơ, các lô tài liệu lớn phải được chấm điểm, tạo sự cân bằng giữa thông lượng và độ trễ.
- Online Embedding: khi truy vấn cơ sở dữ liệu, một truy vấn ngắn phải được nhúng để tra cứu, giúp giảm thiểu độ trễ.
Chúng tôi đã xây dựng cơ sở hạ tầng suy luận của mình để tận dụng tối đa các thành phần chung giữa các trường hợp sử dụng. Vì chúng tôi thường sử dụng các mô hình Transformer nhỏ để tạo ra các embedding, chúng tôi chia sẻ phần lớn việc triển khai với mã suy luận LLM của mình: các lô embedding tương tự như quá trình prefill bị ràng buộc bởi tính toán (compute-bound), trong khi các online embedding, vốn thường chạy trên một vài token, lại tương tự về mặt tính toán với quá trình decode bị ràng buộc bởi bộ nhớ (memory-bound). Do đó, chúng tôi tái sử dụng các kernel prefill và decode đã được tối ưu hóa của mình để phục vụ các mô hình embedding. Nhờ đó, chúng tôi có thể đạt được thông lượng suy luận theo lô lớn với công sức kỹ thuật bổ sung tối thiểu, đồng thời vẫn duy trì độ trễ thấp cho khối lượng công việc online embedding.
Tulip, Rose và một chút Ivy
Chúng tôi cung cấp khả năng suy luận thông qua các API tiêu chuẩn hóa, cả nội bộ và bên ngoài thông qua Nền tảng API (API Platform) của mình. Đằng sau hậu trường, nhiều dịch vụ tham gia vào quá trình xử lý một yêu cầu embedding:
- Ivy là một cổng HTTP Rust mà các dịch vụ của Perplexity gọi tới.
Nó xử lý công việc phía CPU cho các yêu cầu như phân tích cú pháp JSON, token hóa, tạo mẫu đầu vào (input templating) và chia tách lô, dịch các yêu cầu sang giao thức gRPC tùy chỉnh cho các máy chủ hạ nguồn. Sự tách biệt này cho phép chúng tôi cấu hình các tham số nhất định xung quanh việc token hóa và định dạng đầu vào mà không cần phải chạm vào các thực thể suy luận nặng hơn.
- Tulip là giao diện máy chủ suy luận.
Đây là một máy chủ gRPC được triển khai bằng Rust, tokio và `tonic`. Tulip nhận các yêu cầu suy luận gRPC, xử lý việc lập lịch và gom lô (batching). Sau đó, nó gửi các lô đến engine ROSE và trả về các phản hồi đã hoàn thành cho client.
- **ROSE** (Runtime-Optimized Serving Engine) thực hiện suy luận mô hình.
Nó chủ yếu được định nghĩa bằng Python, cung cấp các kernel, lớp và định nghĩa cho rất nhiều loại mô hình. ROSE thực hiện các lượt truyền xuôi (forward pass) qua các mô hình, đồng thời cung cấp khả năng quản lý đồ thị CUDA chuyên biệt cho các embedding. Nó được kết nối với Tulip thông qua một hàm step(), nhận vào một lô (batch) và trả về một tham chiếu đến phép tính mà nó thực hiện trên thiết bị tăng tốc.
Chú ý vượt ra ngoài Kernel
Cả các mô hình dựa trên Transformer lẫn các kiến trúc Hopper/Blackwell nền tảng đều là những công nghệ trưởng thành, do đó việc suy luận nhúng phía GPU đã hội tụ về một triển khai tối ưu phần lớn trên các công cụ suy luận khác nhau. Dù vậy, chúng tôi vẫn khám phá thêm các cơ hội cải thiện thời gian chạy và các công cụ giúp đưa mô hình đến client một cách toàn diện. Cụ thể, chúng tôi nhận thấy có thể cải thiện độ trễ bằng cách quản lý cẩn thận đồ thị CUDA và xây dựng một cỗ trừu tượng LazyTensor để theo dõi bất đồng bộ kết quả ở phía GPU trong engine Rust gốc. Chúng tôi đã triển khai các tính năng này trong Tulip để nó có thể tương tác hiệu quả với các triển khai mô hình của ROSE.
Tulip
Chúng tôi thiết kế Tulip sao cho trở thành một giao diện nhẹ nhất có thể đối với việc phục vụ mô hình của mình. Nó xử lý các yêu cầu đến trong các tác vụ bất đồng bộ của Tokio, duy trì một nhóm các yêu cầu mà nó theo dõi và lập lịch các lô từ đó để phân phối đến thiết bị tăng tốc. Cơ chế lập lịch trong Tulip rất đơn giản: các yêu cầu tích lũy trong khi Tulip đang phân phối công việc hoặc đang chờ kết quả. Từ các yêu cầu tích lũy này, các chuỗi được chọn theo nguyên tắc đến trước, phục vụ trước để chạy qua mô hình.
Cơ chế lập lịch đơn giản được thúc đẩy bởi một quan sát về hiệu năng của mô hình. Đối với các mô hình embedding nhỏ, ở các độ dài chuỗi mà chúng tôi phục vụ, chúng tôi nhận thấy rằng 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 attention. Do đó, độ trễ hầu như tỷ lệ thuận với số lượng token, chứ không phải số lượng chuỗi. Do đó, một khi một lô đủ lớn để bão hòa GPU (khoảng 512 token trên một mô hình dưới một tỷ tham số), việc nhồi nhét thêm nhiều chuỗi vào đó sẽ không làm tăng hiệu quả.
Để tương tác hiệu quả với mô hình, Tulip dựa vào các đồ thị CUDA và tính năng theo dõi kết quả trễ để chồng chéo công việc giữa GPU và CPU nhằm tận dụng tối đa các tài nguyên có sẵn.
Quản lý Đồ thị CUDA
Việc chạy lượt truyền xuôi của một mô hình bao gồm cả công việc ở phía CPU và phía GPU. CPU chịu trách nhiệm lên lịch cho các lô và khởi chạy các kernel với các tham số thích hợp, trong khi GPU thực thi các phép nhân ma trận, attention, chuẩn hóa (norm) hoặc các kernel kích hoạt tương ứng. Đối với khối lượng công việc thông lượng cao như huấn luyện và đánh chỉ mục lại (reindexing), chi phí chung ở phía CPU là không đáng kể vì kích thước lô và độ trễ ở phía GPU đều lớn. Tuy nhiên, đối với các kích thước lô nhỏ hơn, công việc ở phía CPU có thể lớn hơn công việc ở phía GPU.
Để giảm thiểu chi phí chung, thay vì khởi chạy các kernel độc lập, một đồ thị CUDA có thể được xây dựng để chụp lại siêu dữ liệu cần thiết nhằm khởi chạy tất cả các kernel của một lượt truyền xuôi chỉ với một lệnh gọi duy nhất đến trình điều khiển CUDA. Điều này loại bỏ sự cần thiết phải chạy lại các đoạn mã Python và PyTorch đắt đỏ đối với các cấu hình mà đồ thị CUDA có thể chụp được.
Đối với mỗi mô hình, chúng tôi theo dõi điểm uốn, xác định số lượng token tối thiểu mà tại đó việc thực thi trên GPU tốn kém hơn so với việc khởi chạy kernel phía CPU. Vì các mô hình embedding rất nhỏ, chúng tôi nhận thấy điểm uốn này xuất hiện ở các lô hàng nghìn token và hàng chục chuỗi. Một số triển khai attention dựa vào các đầu vào động ở phía host để cấu hình việc khởi chạy kernel, ngăn cản việc sử dụng đầy đủ đồ thị CUDA cho prefill/dense của toàn bộ mô hình. Chúng tôi đã đưa ngược (upstream) các thay đổi vào các kernel liên quan để bật chúng trong engine suy luận của mình.
Để giải quyết các chi phí chung, chúng tôi xây dựng đồ thị CUDA cho toàn bộ mô hình cho tất cả các mô hình embedding và chồng chéo công việc của CPU với công việc của GPU. Vì đồ thị CUDA giảm thiểu các chi phí chung ở phía CPU, một khi đồ thị đã được khởi chạy, chúng tôi có thời gian rảnh để bắt đầu và xếp hàng thực thi lô tiếp theo bất cứ khi nào nó sẵn sàng. Kết quả của lô đang chờ được theo dõi bằng một LazyTensor, cho phép một tác vụ bất đồng bộ trong Rust chặn lại cho đến khi lô trước đó thực thi xong. Đồ thị CUDA hỗ trợ phục vụ độ trễ thấp bằng cách đảm bảo chúng tôi không bị kìm hãm bởi chi phí khởi chạy kernel và tạo điều kiện lập lịch cải thiện trong trường hợp thông lượng cao vì chúng giải phóng CPU để thực hiện công việc trên lô tiếp theo sớm hơn.
Đồ thị CUDA phải được chụp lại cho từng cấu hình riêng biệt, đối với các embedding thì có nghĩa là mỗi đồ thị tương ứng với một tổ hợp số lượng chuỗi và số lượng token. Do lưới này rất rộng, chúng tôi đệm số lượng token thành các khoảng (bucket) là bội số của 64 hoặc 256. Điều này vẫn dẫn đến hàng nghìn đồ thị có thể mất vài phút để chụp đối với một mô hình điển hình. Chi phí chụp xuất phát từ hai nguồn: một lượt truyền xuôi háo hức (eager forward pass) phải được thực thi để biên dịch các kernel và thiết lập bộ đệm cho các kernel cần chúng, theo sau là lượt chạy chụp (capture run) thực thi lại mã Python.
Chúng tôi giảm thiểu chi phí khởi động bằng cách chụp các đồ thị CUDA một cách trễ (lazily) khi engine phục vụ. Chúng tôi theo dõi từng cấu hình và đảm bảo rằng nó trải qua một lượt chạy khởi động háo hức (eager warmup run) trước khi kích hoạt việc chụp và phát lại đồ thị ở lần gọi thứ hai. Tất cả các lần thực thi tiếp theo của cùng một cấu hình đồ thị sau đó sẽ đi qua quá trình phát lại đồ thị CUDA. Việc chụp đồ thị trễ có tác động đến độ trễ p99 trong quá trình khởi động; tuy nhiên, nó rất có giá trị trong việc phân tán vài phút công việc háo hức thành nhiều giờ. Thời gian khởi động nhanh hơn cho phép chúng tôi mở rộng quy mô và quản lý các lượt triển khai embedding tốt hơn.
Lazy Tensor
Thông qua CUDA, công việc trên GPU là bất đồng bộ. Vì việc khởi chạy một kernel một cách bất đồng bộ sẽ đưa nó vào hàng đợi trên một luồng, mã host phải đồng bộ hóa một cách tường minh để đọc ra các véc-tơ kết quả. Để tạo điều kiện cho mức độ song song cao hơn và có thể khởi động các lô tương lai trong khi chờ lô trước hoàn tất trên thiết bị, chúng tôi dựa vào cỗ trừu tượng LazyTensor để theo dõi các giá trị.
LazyTensor theo dõi một bộ đệm host trong bộ nhớ được khóa trang (page-locked memory) và một thao tác cudaMemcpyAsync thông qua một sự kiện (event) sao chép dữ liệu từ thiết bị. Nó được khởi động sau khi lượt truyền xuôi bắt đầu trên cùng một luồng (stream). Do thao tác sao chép phải đợi tất cả các kernel trước đó trên luồng thực thi xong, sự kiện liên quan sẽ theo dõi cả việc hoàn tất lượt truyền xuôi và tính sẵn sàng của kết quả trên CPU.
Chúng tôi tận dụng LazyTensor trong engine mã hóa ROSE của mình để chồng chéo công việc giữa GPU và CPU. Thay vì mỗi lần gọi step() phải chạy đồ thị CUDA và chờ nó hoàn tất, step() trả về một LazyTensor để theo dõi kết quả một cách bất đồng bộ. Kết hợp với đồ thị CUDA, điều này giúp chúng tôi đạt được độ trễ thấp hơn và thông lượng tốt hơn.
ROSE
Chúng tôi đã điều chỉnh engine ROSE của mình, vốn ban đầu được xây dựng để phục vụ LLM, để xử lý cả việc thực thi các mô hình embedding. Để giảm thiểu công sức cần thiết nhằm hỗ trợ các mô hình embedding, ROSE tích cực tái sử dụng mã giữa LLM và embedding. Ví dụ: việc phục vụ pplx-embed và giải mã LLM Qwen3.5 đều đi qua các kernel giống nhau. Việc chia sẻ này cho phép chúng tôi dễ dàng phục vụ một mô hình embedding vốn ban đầu được tinh chỉnh từ một LLM để tạo mẫu, đánh giá và suy luận trong môi trường sản xuất.
Đối với các lớp dày đặc (dense layer), việc suy luận embedding và LLM là giống nhau vì các véc-tơ token được xử lý độc lập. Trong các lớp attention, sự khác biệt được xử lý bằng cách thêm hỗ trợ cho các đầu vào dạng rách (ragged input), bên cạnh các thiết lập prefill và decode được phân trang (paged) theo yêu cầu của LLM. Khi phục vụ một mô hình embedding, chúng tôi không khởi tạo bộ nhớ đệm KV (KV cache) và chuyển hướng đến các biến thể của kernel attention hỗ trợ định dạng ragged để tránh việc đệm (padding). Các quy trình chuyển đổi và hiệu chuẩn hỗ trợ cũng được chia sẻ với các LLM.
Ivy
Ivy, lớp proxy HTTP suy luận của chúng tôi, cũng đóng một vai trò quan trọng đối với hiệu năng. Do tải trọng yêu cầu thay đổi trong môi trường sản xuất, việc định tuyến các yêu cầu riêng lẻ đến các bản sao (replica) riêng biệt có thể gây ra hiện tượng mất cân bằng tải. Ivy chia các yêu cầu lô lớn thành các khối nhỏ và cân bằng tải chúng giữa các bản sao, giúp cải thiện hiệu suất sử dụng và làm mượt độ trễ. Công trình gần đây của chúng tôi về token hóa unigram tự phát triển, được triển khai đầy đủ trong Ivy, cải thiện đáng kể độ trễ so với các bộ token hóa có sẵn.
...nhưng các Kernel vẫn rất quan trọng
ROSE hỗ trợ nhiều loại phần phụ trợ attention (attention backend). Các kernel khác nhau có thể phù hợp với các kích thước bài toán cụ thể. Theo thời gian, chúng tôi đã tích hợp các kernel FlashInfer 2, FlashInfer 3 và FlashAttention 4 để triển khai ragged attention.
Nhìn chung, chúng tôi nhận thấy FlashAttention 4 hoạt động nhanh hơn. Tuy nhiên, FlashInfer 3 lại vượt trội hơn trên các mô hình dựa trên Qwen ở các độ dài chuỗi rất lớn. Do hiệu năng và việc tinh chỉnh có thể thay đổi tùy theo số lượng và chiều của các attention head, chúng tôi tiếp tục duy trì hỗ trợ cho nhiều cấu hình và đưa ra quyết định cho từng trường hợp cụ thể khi phục vụ.
Điểm chuẩn (Benchmarks)
Chúng tôi tiến hành so sánh chuẩn (benchmark) với vLLM v0.22.0, chạy suy luận ở độ chính xác BF16 trên trọng số mô hình thực tế và các đầu vào được lấy từ bộ dữ liệu đánh giá. Tất cả các lượt chạy đo thời gian đều được đi trước bởi các lượt chạy khởi động (warmup) nhằm xác minh rằng độ lệch về độ tương đồng cô-sin nằm trong giới hạn 0.1%.
Embedding độ trễ thấp (p50 / p90 / p99 / tối đa ms)
Chúng tôi báo cáo thời gian chạy cho kích thước lô yêu cầu đã được token hóa trước là 1, các yêu cầu hoàn toàn tuần tự, với các độ dài chuỗi là 128, 512 và 4096 token.
Chấm điểm độ trễ thấp (p50 / p90 / p99 / tối đa ms)
Kích thước lô yêu cầu đã được token hóa trước là 5, 25 và 50, với độ dài chuỗi là 512 token.
Embedding thông lượng cao (emb/s)
Kích thước lô yêu cầu là 100, bốn tiến trình đồng thời gửi yêu cầu, độ dài chuỗi lần lượt là 512, 1024 và 4096 token.
Embedding độ trễ thấp / đồng thời cao (p50 / p90 / p99 / tối đa ms)
Độ dài chuỗi 512, kích thước lô là 1, nhưng chúng tôi gửi 1, 2, 4, 8 và 16 yêu cầu đồng thời. Điểm chuẩn này cũng bao gồm chi phí token hóa thông qua Ivy, cùng với chi phí mạng giữa Ivy và Tulip.
Kết luận và Hướng phát triển trong tương lai
Cơ sở hạ tầng phục vụ bao gồm Ivy, Tulip và ROSE cho phép chúng tôi phục vụ các embedding cho Perplexity với độ trễ thấp hơn và thông lượng tốt hơn, mang lại kết quả tìm kiếm chính xác hơn với chi phí giảm đi so với các giải pháp có sẵn.
Bằng cách tập trung vào các mô hình cụ thể và nắm trọn toàn bộ ngăn xếp, chúng tôi có được sự tự do cần thiết để tạo ra sự cân bằng hiệu quả giữa hiệu năng và tính linh hoạt, kết hợp các thành phần nguyên thủy bằng Rust có khả năng tái sử dụng cao và hiệu năng cao cùng với mã mô hình Python chung hơn. Nhiều công cụ suy luận mã nguồn mở như vLLM, SGLang và TokenSpeed đang tích hợp các ngôn ngữ như Rust và C++ vào ngăn xếp của họ. Chúng tôi đã đầu tư vào Rust trong hai năm qua và gặt hái được những thành quả lớn về cả hiệu năng lẫn khả năng bảo trì. Bằng cách chia sẻ phần lớn việc triển khai embedding với ngăn xếp phục vụ LLM của mình, chúng tôi cũng đạt được những cải thiện về thông lượng mà không cần tốn nhiều công sức kỹ thuật cho việc duy trì các mô hình embedding.
Khi các mô hình tiến hóa, chúng tôi sẽ tiếp tục cải thiện từng lớp trong ngăn xếp của mình để giảm thiểu cả độ trễ bị ràng buộc bởi CPU lẫn GPU. Các giao thức tùy chỉnh dựa trên gRPC của chúng tôi trong Ivy và Tulip cho phép chúng tôi tinh chỉnh việc giao tiếp để giảm độ trễ mạng, trong khi ROSE cung cấp nền tảng để cải thiện thông lượng tính toán. Ngoài ra, khi sự hỗ trợ cho Python tự do luồng (free-threaded Python) phát triển trong toàn bộ hệ sinh thái, chúng tôi sẽ có thể cải thiện hơn nữa khả năng tương tác giữa Python và Rust để giảm thiểu chi phí chung.