Mã nguồn mở hợp nhất kernel Triton giúp tăng tốc suy luận Qwen3-TTS lên tới 5 lần

Xin chào, tôi xin chia sẻ một thư viện hợp nhất kernel Triton đã giải quyết nút thắt cổ chai suy luận của mô hình Qwen3-TTS 1.7B và đạt được mức tăng tốc khoảng 5 lần.

1. Vì sao tôi làm nó? (Bối cảnh)
Những ai từng dùng audio TTS trong công việc thực tế hẳn đều biết, các mô hình mang tính xác suất (stochastic) như Qwen3-TTS sẽ cho kết quả tạo sinh khác nhau mỗi lần chạy (ngữ điệu, tông giọng, v.v.).

Vì vậy trong thực tế, chiến lược sinh nhiều phương án rồi chọn ra phương án tự nhiên nhất là điều gần như bắt buộc. Nhưng với tốc độ trước đây, pipeline làm việc quá chậm chạp nên tôi đã trực tiếp bắt tay vào tối ưu hóa. Nhờ lần tối ưu này, trong cùng khoảng thời gian trước đây chỉ tạo được 1 mẫu, giờ có thể tạo ra 5 ứng viên.

2. Tôi đã làm như thế nào? (Claude Code + kiểm thử cực gắt)
Nói thật thì, tôi biết sức mạnh của OpenAI Triton, một thư viện tối ưu hóa kernel, nhưng chưa từng tự tay viết mã kernel trước đây. Vì vậy, phần lớn mã kernel trong dự án này được viết với sự hỗ trợ của Claude Code.

Thay vào đó, để bù đắp cho kinh nghiệm triển khai Triton còn hạn chế và bảo đảm độ tin cậy của mô hình ở mức 100%, tôi đã dồn toàn bộ năng lượng vào việc kiểm thử cực kỳ gắt gao thay vì chỉ tập trung viết mã.

• Tôi đã viết 90 bài kiểm thử đơn vị để bảo đảm kết quả toán học hoàn toàn giống với mô hình gốc.
• Đạt Cosine Similarity > 0.997 ở toàn bộ các lớp checkpoint chính và cả đầu ra cuối cùng.

3. Điểm nhấn kỹ thuật và thành quả
Lấy cảm hứng từ Liger Kernel của LinkedIn, tôi đã hợp nhất 4 phép toán gây nghẽn khi suy luận (RMSNorm, M-RoPE, Norm+Residual, SwiGLU) thành các kernel Triton.

[Benchmark hiệu năng - trên RTX 5090]

• Base (PyTorch): 3,902 ms
• Hybrid (Faster+Triton): 919 ms (~tăng tốc 4.7x)
  (※ Chế độ Hybrid là kết quả của việc chồng phần hợp nhất kernel Triton lần này lên faster-qwen3-tts dựa trên CUDA Graph.)

4. Kết lời
Hiện tại tôi mới chỉ hoàn tất kiểm thử trên thiết bị cá nhân là RTX 5090. Nếu ai đang dùng phần cứng khác như máy chủ cấp doanh nghiệp (A100, H100) hoặc RTX 4090 và có thể chạy thử rồi phản hồi qua GitHub hoặc phần bình luận thì sẽ là sự giúp đỡ rất lớn.

Cảm ơn mọi người đã đọc bài viết dài này!

• GitHub: https://github.com/newgrit1004/qwen3-tts-triton
• PyPI: https://pypi.org/project/qwen3-tts-triton/