Google tăng tốc AI Python bằng cách loại bỏ Python

Google vừa mã nguồn mở XSE (XLA Serving Engine), một thành phần trong hệ sinh thái OpenXLA dùng để chạy các mô hình học máy đã huấn luyện. Điểm đáng chú ý: công cụ này loại bỏ hoàn toàn trình thông dịch Python khỏi đường ống phục vụ (serving pipeline).

Đây là bước đi quan trọng bởi tại Google, XSE đã xử lý hơn 6 triệu truy vấn mỗi giây trên các dịch vụ như YouTube Shorts, Google Ads và Search. Việc loại bỏ Python khỏi khâu suy luận giúp giảm độ trễ và chi phí vận hành cho khối lượng truy vấn khổng lồ này.

Triết lý "huấn luyện bằng Python, phục vụ bằng C++"

Python dễ dùng với các nhà khoa học dữ liệu, nhưng giữ trình thông dịch của nó trong đường ống phục vụ lại tạo ra gánh nặng về hiệu năng và triển khai. Google vốn đã quen với mô hình "train in Python, serve in C++" (huấn luyện bằng Python, phục vụ bằng C++) trong nhiều dự án ML.

Theo cách làm mới, các kỹ sư ML vẫn viết mã mô hình bằng Python, sử dụng JAX của Google để xử lý các phép tính. Sau đó, một công cụ xuất (export) sẽ đóng gói các phép tính của mô hình thành StableHLO — một dạng biểu diễn trung gian (intermediate representation) có thể chuyển đổi giữa các nền tảng — cùng với trọng số đã huấn luyện, tạo thành một "XLA SavedModel".

Ở phía chạy thực tế, runtime C++ dùng XLA để biên dịch các phép tính đó thành mã máy (native code) cho phần cứng đích. Không còn trình thông dịch Python nào trong quy trình.

Không dành cho LLM cỡ lớn

XSE không nhắm tới suy luận LLM quy mô lớn — lĩnh vực mà các công cụ như vLLM hay SGLang phù hợp hơn. Nó tập trung vào các mô hình nhỏ, chuyên biệt cho từng tác vụ, vốn rất phổ biến tại Google.

Nathan Beach, giám đốc quản lý sản phẩm của Google, viết trên LinkedIn: các mô hình xếp hạng, gợi ý, embedding và thị giác vẫn chiếm phần lớn suy luận trong môi trường production, dù ít được chú ý hơn so với LLM. Đặc điểm của chúng là yêu cầu ngắn, số truy vấn mỗi giây rất cao và ngân sách độ trễ chặt chẽ.

Trong thử nghiệm của Google với các yêu cầu ngắn, XSE đạt thông lượng gấp 2 đến 5 lần so với máy chủ Python được tối ưu tốt nhất, trong khi chỉ dùng một phần tư bộ nhớ.

Chạy trên GPU, CPU và TPU

XLA SavedModel có thể chạy trên GPU và CPU — Google đã hợp tác với cả Arm và Intel để phát triển hệ sinh thái XLA. Tuy nhiên, không khó đoán rằng Google muốn các mô hình chạy trên dòng chip tăng tốc TPU của chính mình.

Google cũng có kế hoạch tích hợp với TorchTPU để người dùng PyTorch có thể xuất mô hình phục vụ bằng C++ trên TPU.

Cuộc đua thoát khỏi Python

Tốc độ chậm của Python vẫn là vấn đề với cộng đồng ML, và nhiều hướng tiếp cận đã ra đời nhằm loại bỏ mã Python càng sớm càng tốt. PyTorch cung cấp runtime C++ riêng có tên LibTorch. Microsoft cùng các đối tác cung cấp ONNX Runtime, hỗ trợ định dạng đồ thị mở (open graph format) với binding Python.

Ngay cả bản tham chiếu chính thức của Python cũng dựa nhiều vào C, và giờ Rust đang bắt đầu chen chân vào.

Góc nhìn

Với các kỹ sư và doanh nghiệp Việt Nam đang triển khai mô hình AI trong production, XSE là lời nhắc rằng giai đoạn huấn luyện và giai đoạn phục vụ nên được tách bạch về mặt kỹ thuật. Nhiều đội nhóm hiện giữ nguyên pipeline Python từ lúc thử nghiệm đến khi chạy thật, dẫn tới chi phí máy chủ cao và độ trễ khó kiểm soát — trong khi phần lớn mô hình nghiệp vụ như xếp hạng, gợi ý hay embedding lại đúng là loại có lưu lượng lớn và yêu cầu phản hồi nhanh.

Việc XSE được mã nguồn mở tạo thêm lựa chọn cho những ai muốn biên dịch mô hình ra mã máy thay vì thuê thêm tài nguyên. Tuy vậy, công cụ này gắn chặt với hệ sinh thái XLA và JAX, nên các đội dùng PyTorch sẽ cần cân nhắc kỹ trước khi chuyển đổi, hoặc chờ tích hợp TorchTPU hoàn thiện.