Инженеры удвоили скорость prefill Gemma 4 31B на TPU v6e

🤖 Инженеры удвоили скорость prefill Gemma 4 31B на TPU v6e

Sail Research опубликовала разбор оптимизации prefill на узле из 4 чипов TPU v6e со стеком SGLang-JAX: MFU вырос с ~32% до ~63%, throughput при промпте 8192 токена — с 18 228 до 36 669 tok/s, а TTFT упал с 449 до 223 мс. Приёмы: query-блок 512 вместо 32 в Ragged Paged Attention, collective matmuls вместо монолитного AllReduce и собственные Pallas-кернелы — XLA не справился с автослиянием коллективов.

🌍 HBM даёт более 63% стоимости ускорителя, поэтому TPU v6e с FLOPs как у H100, но в 2,5 раза меньшей памятью рентабелен для prefill-нагруженного инференса малых и средних моделей.

👤 Главный overhead инференса — не FLOPs, а коммуникации и мелкий «хвост» операций: в кейсе профили Perfetto/XProf и разбор, когда XLA бессилен и нужен свой кернел.

Источник 1: https://www.sailresearch.com/blog/tpu-v6e-gemma Источник 2: https://www.sailresearch.com/blog/htdym