Ingegneria dei kernel Mixture-of-Experts su Amazon SageMaker AI
aws ·
Questo articolo descrive le ottimizzazioni di ingegneria dei kernel per i modelli Mixture-of-Experts (MoE) su Amazon SageMaker AI, con un focus sui miglioramenti delle prestazioni per l'inferenza dei grandi modelli linguistici. Gli autori illustrano tre tecniche chiave: la trasposizione Swap-AB per allineare le dimensioni delle matrici ai requisiti hardware WGMMA di NVIDIA Hopper, regolazioni del routing backend per utilizzare Triton invece di CUTLASS per la quantizzazione FP8 e l'eliminazione delle ricompilazioni JIT non necessarie di Triton per ridurre la latenza. Queste ottimizzazioni affrontano colli di bottiglia specifici nelle architetture MoE, come i batch size sottili durante la decodifica e il numero variabile di token. L'implementazione di queste modifiche ha portato a un aumento della throughput fino al 10% e a un miglioramento del 9% nel Time to First Token (TTFT) sulle GPU H200. Diversi di questi potenziamenti sono stati integrati nella comunità open-source vLLM. Il lavoro dimostra come una messa a punto profonda specifica per l'hardware possa migliorare significativamente l'efficienza nei modelli ad attivazione sparsa.
aws