Tăng tốc suy luận AI cho việc tạo nội dung 3D trên Roblox
Tạo đối tượng 3D nhanh hơn 7,8 lần và phản hồi tốt hơn
- Roblox đã triển khai CUDA Graphs và bộ nhớ đệm KV để tăng tốc độ tạo lưới 3D, giúp quá trình lặp lại trở nên mượt mà hơn.
- Khi ra mắt, mô hình Cube 3D có thể tạo ra các token trong 7,8 mili giây (giảm từ 60,5 mili giây) và các đối tượng hoàn chỉnh trong 4 giây (giảm từ 31 giây).
Đầu năm nay, Roblox đã chia sẻ khả năng đầu tiên của mô hình nền tảng Cube 3D của chúng tôi. Với Cube 3D, các nhà sáng tạo có thể tạo ra các mô hình và môi trường 3D trực tiếp từ các lời nhắc văn bản. Ngay từ đầu, chúng tôi đã ưu tiên tối ưu hóa độ trễ, nhận thức rằng thời gian tạo chậm sẽ làm gián đoạn quá trình lặp vốn dĩ vốn có. Trước khi Cube 3D ra mắt vào tháng 3, chúng tôi đã làm cho bước suy luận nhanh hơn 7,8 lần và phản hồi tốt hơn cho cả nhà phát triển lẫn người dùng.
Kể từ khi ra mắt, hơn 578.000 đối tượng đã được tạo ra trong nhiều trải nghiệm nổi bật. Các nhà phát triển cũng bày tỏ sự quan tâm đến việc cho phép người dùng tạo đối tượng 3D trong các trải nghiệm thông qua các lời nhắc văn bản như “mèo”, “hamburger”, v.v. Đáng chú ý nhất là Mic Up, một trò chơi giao lưu phổ biến sử dụng chat giọng nói, đã tận dụng Cube 3D để mang đến cho người chơi cách thức vui nhộn và tương tác để tạo đối tượng. Trong quá trình triển khai, người chơi có thể mở menu bên trái để truy cập các tính năng bổ sung, bao gồm biểu tượng AI. Sau khi nhấp vào biểu tượng đó, người chơi có thể nhập lời nhắc văn bản để tạo đối tượng 3D. Đối với người dùng, thời gian tạo đối tượng kéo dài gây ra sự cản trở, khiến họ mất đi sự kỳ diệu khi thấy ý tưởng của mình được biến thành 3D theo thời gian thực.
Chúng tôi muốn biến trải nghiệm tạo 3D từ một tương tác phải dừng lại và chờ đợi thành một trải nghiệm phản hồi nhanh và tự nhiên, cho phép thử nghiệm nhanh chóng. Khả năng thêm nhanh các đối tượng vào một cảnh là rất quan trọng đối với các nhà phát triển. Để tăng tốc Cube 3D, trước tiên chúng tôi đã phân tích đường ống suy luận để xác định các điểm nghẽn hiệu suất. Mặc dù sử dụng GPU mạnh mẽ, chúng tôi vẫn phát hiện ra thời gian nhàn rỗi đáng kể giữa các hoạt động.
Giải quyết nút thắt cổ chai trong việc lập lịch CPU-GPU
Các khung công tác học sâu hiện đại dựa vào CPU để lên lịch và khởi chạy các hoạt động (hoặc hạt nhân) trên GPU. CPU chuẩn bị từng hoạt động, gửi nó đến GPU và chờ xác nhận trước khi chuẩn bị hoạt động tiếp theo. Việc chờ đợi này tạo ra một điểm nghẽn trong việc lên lịch, khiến GPU có thể ở trạng thái nhàn rỗi trong khi CPU chuẩn bị đợt công việc tiếp theo. Lý tưởng nhất là chúng ta muốn CPU chạy trước GPU, chuẩn bị và xếp hàng các hoạt động để GPU luôn có việc để làm.
Điều này đặc biệt gây khó khăn cho các bộ giải mã tự hồi quy trong các mô hình kiểu transformer như Cube 3D, vốn cần xử lý đầu vào và tạo ra các token theo thứ tự. Các mô hình này yêu cầu hàng nghìn thao tác riêng lẻ cho mỗi lần tạo ra dữ liệu, và gánh nặng tính toán tích lũy theo từng bước trong chuỗi.
“Chúng tôi muốn xây dựng một thứ gì đó cho phép tương tác bốn chiều,” Phó Chủ tịch Kỹ thuật Anupam Singh cho biết, giải thích lý do tại sao Roblox chọn phương pháp tự hồi quy. “Chúng tôi không chỉ muốn chế tạo chiếc xe; chúng tôi còn muốn có thể mở cửa xe và bước vào trong.”
Mỗi thao tác phát sinh:
- Thời gian CPU để chuẩn bị mỗi kernel
- Chi phí phát sinh từ việc khởi chạy kernel
- Thời gian thực thi GPU (tính toán thực tế)
- Chi phí đồng bộ hóa khi kiểm tra hoàn thành
Trong trường hợp các thao tác nhỏ được thực thi nhanh trên GPU, chi phí này có thể chiếm phần lớn thời gian suy luận. GPU có thể chỉ thực hiện tính toán tích cực trong một phần nhỏ của tổng thời gian suy luận.
Triển khai đồ thị CUDA: Loại bỏ trung gian
Để giải quyết vấn đề này, chúng tôi đã tận dụng CUDA Graphs — một tính năng cho phép ghi lại và phát lại các chuỗi hoạt động của GPU mà không cần sự can thiệp của CPU. Thành phần bộ giải mã tự hồi quy trong kiến trúc của Cube 3D xử lý các lời nhắc văn bản và tạo ra các mã hình dạng thông qua một vectơ có độ dài cố định.
Mặc dù về mặt chức năng tương tự như mô hình ngôn ngữ lớn (LLM) truyền thống, kiến trúc bộ giải mã hai luồng của chúng tôi có một điểm khác biệt quan trọng — nó sử dụng hai luồng chú ý song song. Một luồng dành riêng cho các token điều kiện và luồng còn lại dành cho các token hình dạng. Các công cụ suy luận LLM có sẵn trên thị trường không phù hợp với nhu cầu của chúng tôi, và chúng tôi cần một triển khai tùy chỉnh phù hợp với kiến trúc cụ thể của mình.
Hãy tưởng tượng CUDA Graphs như việc ghi lại một macro cho GPU. Thay vì CPU phát hành từng lệnh riêng lẻ, nó ghi lại toàn bộ chuỗi các thao tác GPU (đồ thị) và khởi chạy toàn bộ đồ thị bằng một lệnh CPU duy nhất. Cách tiếp cận này giảm đáng kể chi phí khởi chạy kernel bằng cách loại bỏ nhu cầu CPU phải lên lịch từng thao tác riêng lẻ trong quá trình suy luận. Một khi đồ thị được khởi chạy, GPU sẽ tự động thực thi toàn bộ chuỗi mà không cần chờ đợi các lệnh tiếp theo.
CUDA Graphs cũng có một số hạn chế. Vì cấu trúc đồ thị cần được xác định trước, chúng yêu cầu kích thước lô và kích thước đầu vào cố định. Điều này có nghĩa là cần tạo các đồ thị riêng biệt cho mỗi kích thước lô hoặc hình dạng đầu vào. Đối với trường hợp sử dụng của chúng tôi với mô hình Cube 3D, hạn chế này là có thể chấp nhận được, vì chúng tôi có thể tiêu chuẩn hóa quá trình suy luận xung quanh các kích thước đầu vào chung.
Chúng tôi phải điều chỉnh phương pháp của mình để triển khai CUDA Graphs cho mô hình Cube 3D. Trong các mô hình ngôn ngữ lớn (LLMs) truyền thống, các thao tác chú ý luôn được thực hiện với độ dài chuỗi cố định, cung cấp một hình dạng tĩnh để làm việc. Tuy nhiên, trong kiến trúc hai luồng tùy chỉnh của chúng tôi, một số lớp chú ý hoạt động dựa trên độ dài chuỗi, trong khi những lớp khác hoạt động dựa trên sự kết hợp giữa độ dài chuỗi và độ dài điều kiện.
Bất chấp những thách thức này, chúng tôi đã thu được kết quả đáng chú ý sau khi triển khai CUDA Graphs. Chúng tôi sử dụng thời gian trên mỗi token đầu ra (TPOT) để đo thời gian tạo ra từng token trong quá trình suy luận. Sau khi triển khai CUDA Graphs, TPOT của chúng tôi đã cải thiện từ 60,5 mili giây xuống còn 20,5 mili giây, tương đương mức cải thiện 2,9 lần. Tổng thời gian tạo ra nội dung giảm 66%, từ 31 giây xuống còn 10,5 giây.
KV Caching: Phát huy thành công của chúng tôi
Để cải thiện độ trễ hơn nữa, chúng tôi đã triển khai bộ nhớ đệm KV, một phương pháp tiêu chuẩn trong suy luận LLM đã được chứng minh là rất hiệu quả trong toàn ngành.
Trong các mô hình dựa trên biến đổi như Cube 3D, mỗi lần tạo mã thông báo đều yêu cầu tính toán ma trận khóa (K) và giá trị (V) dựa trên tất cả các mã thông báo được tạo trước đó. Khi chuỗi ngày càng dài, việc tính toán lại các ma trận này cho mỗi mã thông báo trở nên ngày càng kém hiệu quả.
Bộ nhớ đệm KV giải quyết vấn đề này bằng cách:
- Lưu trữ các ma trận K và V cho tất cả các token đã được tạo trước đó
- Chỉ tính toán ma trận K và V cho các token mới
- Ghép nối các ma trận mới này vào các giá trị đã lưu trữ
Cách tiếp cận này loại bỏ các tính toán dư thừa, giảm bớt công việc cần thiết cho mỗi token mới. Điều này trở nên đặc biệt có tác động khi chuỗi được tạo ra ngày càng dài.
Cách tiếp cận của chúng tôi để tích hợp bộ nhớ đệm KV với triển khai CUDA Graph tương tự như suy luận LLM truyền thống. Việc bổ sung bộ nhớ đệm KV đã giảm TPOT của chúng tôi xuống chỉ còn 7,8 mili giây. Thời gian tạo tổng thể giảm 87%, từ 31 giây ban đầu xuống chỉ còn 4 giây. Việc giảm thời gian đáng kể này giúp công cụ này hiệu quả hơn nhiều cho những người sáng tạo sử dụng công cụ này.
Đánh giá tác động thực tế đối với các nhà phát triển và người dùng
Chúng tôi đang nghiên cứu các kỹ thuật nhằm giảm thiểu độ trễ và nâng cao trải nghiệm người dùng, bao gồm các kernel được tối ưu hóa, lượng tử hóa mô hình để suy luận nhanh hơn, tối ưu hóa dành riêng cho phần cứng và tạo token song song.
Công việc này trở nên phức tạp hơn khi chúng tôi mở rộng sang việc tạo và hiểu toàn bộ cảnh, nơi nhiều yếu tố 3D cần phải phối hợp với nhau trong bối cảnh của bố cục. Chúng tôi cũng muốn các đối tượng và thế giới 3D mà chúng tôi tạo ra phải hoạt động đầy đủ chức năng, để cửa có thể mở và đóng, bánh xe có thể quay, v.v. Để đạt được điều đó, chúng tôi cần tạo và lặp lại nhanh chóng để mở rộng quy mô lên toàn bộ cảnh, các đối tượng hoạt động đầy đủ chức năng và các avatar. Chúng tôi rất hào hứng chia sẻ những cải tiến tiếp theo và tính năng mới khi mở rộng mô hình cơ sở Cube 3D của mình — và mong chờ được chiêm ngưỡng những thế giới đắm chìm mà cộng đồng sáng tạo của chúng tôi xây dựng từ đó.
