07/10/2026 · VietTechies · Biên tập
PyTorch đã nâng cấp toàn bộ tập lệnh Table Batched Embedding (TBE) bằng cách thay thế các mẫu CUDA C++ truyền thống bằng kernel FBTriton được viết trực tiếp bằng Python. Các toán tử TBE thực hiện tác vụ truy xuất và gộp (pooling) nhúng thưa trên nhiều bảng đặc trưng trong cùng một lần khởi chạy GPU. Đây là khối lượng công việc cốt lõi trong các hệ thống gợi ý quy mô lớn đang vận hành trên hàng nghìn GPU được phân mảnh.
Cải thiện hiệu năng trong lượt truyền tiến và truyền lùi
Trong các bài đo hiệu năng trên 307 cấu hình phân mảnh chạy trên trình tăng tốc Nvidia GB200 (sử dụng trọng số FP16 và thuật toán Adagrad theo dòng), FBTriton đạt tốc độ xử lý lượt truyền tiến (forward pass) nhanh hơn 1,28 lần (trung vị) so với các kernel CUDA tiêu chuẩn. Mức tăng hiệu năng này chủ yếu đến từ việc tối ưu hóa truy xuất bộ nhớ trong lượt truyền lùi (backward pass). Kernel CUDA TBE thường chuyển sang chế độ đồng bộ hóa mảng luồng hợp tác (CTA) khi độ dài phân đoạn đạt 32 lần truy xuất mỗi dòng, trong khi FBTriton tiếp tục duy trì xử lý luồng liên tục với độ dài lên tới 256. Sự thay đổi về cấu trúc này giúp FBTriton đạt băng thông bộ nhớ 3.948 GB/s, vượt trội so với mức 678 GB/s của CUDA trên cùng một mức dung lượng bộ nhớ sử dụng.
Thay đổi kiến trúc và phân loại dữ liệu trên GPU
Để loại bỏ chi phí đồng bộ hóa giữa thiết bị và máy chủ do thao tác đọc ngược dữ liệu về CPU, FBTriton lưu giữ toàn bộ dữ liệu phân loại khối lượng công việc ngay trên GPU. Tuyến trình xử lý áp dụng các bộ đếm nguyên tử (atomic counter) và phép cộng dồn (prefix sum) bên trong các khối tính toán, chia độ dài dòng phụ thuộc vào dữ liệu thành các đoạn cố định gồm 256 lần truy xuất mà không làm kích hoạt quá trình đồng bộ hóa luồng ở cấp CPU.
Bản cập nhật này bổ sung tùy chọn kiểm tra ranh giới hợp nhất (fused bounds-checking), giúp tăng tốc tới 1,24 lần trong bước xác thực trên Nvidia B200. Ngoài ra, một đường dẫn tùy chọn trong mô-đun cốt lõi đã chuyển các thao tác sắp xếp chỉ số và mã hóa độ dài chạy (run-length encoding) sang lượt truyền tiến. Trên các hệ thống B200 quy mô lớn, sự dịch chuyển này giảm tổng độ trễ thực thi của cả hai lượt truyền tiến và lùi từ 79,537 ms xuống 66,183 ms, tương đương mức giảm 16,8%. FBTriton còn tích hợp các tính năng phần cứng chuyên dụng cho kiến trúc Nvidia Blackwell như cơ chế tranh thủ công việc (hardware work stealing), rào chắn đồng bộ hóa phạm vi thiết bị và giảm thiểu nguyên tử số lượng lớn TMA (TMA bulk atomic reductions).
Hiệu suất phát triển và các trường hợp biên
Việc chuyển đổi các thao tác thưa trong lượt truyền tiến và truyền lùi sang Python giúp PyTorch giảm đáng kể chi phí bảo trì so với việc dùng các mẫu Jinja CUDA phức tạp, đồng thời chuẩn hóa logic kernel trên nhiều kiến trúc GPU khác nhau như AMD, Hopper và Blackwell.
Dù vậy, FBTriton không áp đảo hoàn toàn CUDA cũ trong mọi khối lượng công việc. Trong số 307 phân mảnh thử nghiệm, CUDA vẫn duy trì tốc độ xử lý nhỉnh hơn ở 11 cấu hình đoạn ngắn với độ dài phân đoạn dưới 4 lần truy xuất mỗi dòng. Lý do là cơ chế phân bổ warp theo từng dòng của CUDA giảm độ trễ nạp dữ liệu phụ thuộc tốt hơn khi xử lý các đợt dữ liệu nhỏ.
Ảnh minh họa do AI tạo, không phải ảnh tư liệu.