Gemma-4 (2B/4B/12B) AWS Inferentia2'ye Nasıl Taşındı
Gemma-4'ü AWS Inferentia2'ye taşıma raporu: karışık dikkat başlıkları, KV paylaşımı ve vLLM/NxD sınırlamalarının nasıl aşıldığı.
Bir mühendislik ekibi, Google'ın Gemma-4 ailesini (E2B, E4B, 12B) AWS Inferentia2 üzerinde çalıştırmak için resmi vLLM/optimum-neuron/neuronx-distributed yığınını devre dışı bırakıp Hugging Face forward geçişini doğrudan izleyen özel bir çözüm geliştirdi. Sorun kökeninde Gemma-4'ün TPU/XLA için tasarlanmış katman-arası KV paylaşımı, karışık dikkat başlığı sayıları (GQA) ve kayan pencere/global dikkat karışımı gibi özelliklerinin AWS'nin statik grafik oluşturucusunda ifade edilememesiydi; sonuç olarak vendor vLLM ucu akıcı görünen ama anlamsız çıktı üretiyordu.
Ekip 'karışık başlıklar' sorununu üç ayrı biçimde çözdü: E2B/E4B'de KV sahibi olmayan katmanları tespit edip yalnızca gerçek KV üreten katmanlara bellek ayırarak; E4B'de tensor paralelliği altında GQA'yı KV başlıklarını TP'ye bölünebildiğinde parçalayıp num_key_value_groups'u sabit tutarak; 12B'de ise bölünemeyen tek KV başlıklı global katmanları çoğaltılmış bırakıp grup sayısını rank başına düşen sorgu başlığı sayısına göre küçülterek. Sıralama hatası ve hiç test edilmemiş bir kod yolu, saatler süren hata ayıklamaya yol açan iki kritik tuzak olarak öne çıktı.
Sonuçta greedy decode çıktıları CPU referansıyla token-token birebir eşleşti; E2B tek çekirdekte ~44 tok/s, E4B TP=2 ile ~33-39 tok/s, 12B TP=2 ile ~15 tok/s hız verdi. Bu, üretici destekli çıkarım yığınının kapsamadığı yeni nesil mimarileri özel model izleme yoluyla üretime taşımak isteyen donanım mühendisleri için somut bir yol haritası sunuyor.