Open source de fusión de kernels Triton que acelera hasta 5 veces la inferencia de Qwen3-TTS
(github.com/newgrit1004)Hola, comparto una librería de fusión de kernels Triton que resuelve los cuellos de botella de inferencia del modelo Qwen3-TTS 1.7B y logra una mejora de velocidad de alrededor de 5 veces.
1. ¿Por qué lo hice? (contexto)
Quienes han usado audio TTS en entornos reales ya lo sabrán: los modelos estocásticos como Qwen3-TTS producen resultados distintos en cada generación (entonación, tono, etc.).
Al final, en producción se vuelve indispensable una estrategia de generación múltiple que permita sacar rápidamente varias opciones de audio y elegir la más natural entre ellas, pero con la velocidad anterior el pipeline de trabajo se volvía demasiado lento, así que decidí optimizarlo por mi cuenta. Con esta optimización, ahora es posible generar 5 candidatos en el tiempo en que antes se generaba 1.
2. ¿Cómo lo hice? (Claude Code + pruebas intensivas)
Para ser sincero, yo conocía la potencia de OpenAI Triton como librería de optimización de kernels, pero nunca había escrito código de kernels directamente. Por eso, la mayor parte del código de kernels de este proyecto fue escrita con ayuda de Claude Code.
En cambio, para compensar mi falta de experiencia implementando en Triton y garantizar al 100% la confiabilidad del modelo, en vez de concentrarme en escribir código invertí toda mi energía en hacer pruebas realmente exhaustivas.
- Escribí 90 pruebas unitarias para garantizar resultados matemáticos perfectamente idénticos al modelo original.
- Se logró una Cosine Similarity > 0.997 en todas las capas principales de checkpoint y en la salida final.
3. Puntos de ingeniería y resultados
Inspirado en Liger Kernel de LinkedIn, fusioné en kernels Triton 4 operaciones que eran cuellos de botella durante la inferencia (RMSNorm, M-RoPE, Norm+Residual, SwiGLU).
[Benchmark de rendimiento - con RTX 5090]
- Base (PyTorch): 3,902 ms
- Hybrid (Faster+Triton): 919 ms (~4.7x más rápido)
(※ El modo Hybrid es el resultado de montar esta fusión de kernels Triton sobre faster-qwen3-tts, basado en CUDA Graph.)
4. Cierre
Por ahora solo terminé las pruebas en mi equipo personal con RTX 5090. Si alguien lo prueba en otro hardware, como servidores (A100, H100) o una RTX 4090, y me deja feedback en GitHub o en los comentarios, de verdad sería de gran ayuda.
¡Gracias por leer hasta aquí!
Aún no hay comentarios.