GEMM optimizasyonu, GPU’larda modüler bir sorun alanıdır. Performanslı uygulamalar, tile şekilleri, matematiksel ve kopyalama komutları ile warp uzmanlaşma şemaları gibi hiperparametreler belirlemeyi gerektirir. Bu hiperparametreler bir ölçüde bağımsızdır ve en iyi seçimler donanım, problem şekli veya diğer kullanıcı ihtiyaçlarına bağlı olarak önemli ölçüde değişiklik gösterebilir.
CUTLASS 3.x tasarımında, GEMM uygulamalarının kapsamının genişletilmesi hedeflenmiştir. Bu, kompozit ve ortogonal yapı taşları sistemine dayalı hiyerarşik bir yapı ile gerçekleştirilmektedir. Bu design felsefesi, GPU’nun hiyerarşik donanım tasarımıyla bağlantılı olup, diğer GPU uygulamaları için de iyi bir seçenek olabilir. Örneğin, FlashAttention-3, tasarımında tanıdık CUTLASS soyutlamalarını kullanmaktadır.
Bu CUTLASS blog serisinin ikinci yazısında, kutlASS 3.x’deki GEMM’in hiyerarşik sisteminin tasarım ilkelerini keşfedeceğiz. Ayrıca, CUTLASS’ın, ilk bölümde tanıtılan düşük seviyeli CuTe soyutlamalarından GEMM çekirdeklerini nasıl oluşturduğunu açıklayacağız.
CUTLASS 3.x’te Yeni Bir Kavramsal GEMM Hiyerarşisi
CUTLASS 3.x, belirli donanım özelliklerinden bağımsız bir kavramsal GEMM hiyerarşisi geliştirir. Bu hiyerarşi, beş katmandan oluşmaktadır:
- Atom katmanı: Mimariye özgü komutlar ve ilişkili meta bilgilere sahiptir.
cute::Mma_Atom<>vecute::Copy_Atom<>
- Tiled MMA/Kopya: Belirli bir mimariye özgü atomların keyfi karıştırılması ve tile biçimlendirme işlemlerini yapabilen mekansal mikro çekirdekler.
cute::TiledMma<>vecute::TiledCopy<>
- Toplu katman: Bir veya daha fazla mekansal mikro çekirdeğin yürütülmesini senkronize etmek için mimariye özgü zamanlama kullanarak tek bir çıktı tile’ı hesaplayan geçici mikro çekirdekler.
cutlass::gemm::collective::CollectiveMma<>,cutlass::epilogue::collective::CollectiveEpilogue<>
- Çekirdek katmanı: Bir dizi threadblock/küme üzerinden bir çekirdek çalıştırmak için cihaz kodu.
cutlass::gemm::kernel::GemmUniversal<>
- Cihaz katmanı: Ana bilgisayar tarafında kurulum ve arayüz.
cutlass::gemm::device::GemmUniversalAdapter<>
Her katman, özelleştirilebilir soyutlamalar için bir bileşim noktası işlevi görmektedir. Kullanıcılar, CUTLASS’ın derleme zamanı mantığına güvenerek en üst katmanları kullanabilir veya hiyerarşinin alt seviyelerinden sunulan gelişmiş değişikliklere yönelebilirler. Atom ve Tiled MMA/Kopya katmanlarındaki mekansal mikro çekirdekler, CuTe domainidir ve ilk bölümde tartışılmıştır. Bu yazının geri kalanı, daha yüksek katmanlarda sunulan geçici ve çekirdek seviyesi GEMM organizasyonunu kapsayacaktır.
Temel Çekirdek Tanımı
CUTLASS 3.x’te bir GEMM çekirdeği tanımlamak, aşağıdaki gibi basit bir örnekle yapılabilir:
// Adım 1: Gerekli toplu katman ana döngüsü özel düzenlemesini oluştur
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
ArchTag, OperatorClass,
ElementA, LayoutA, AlignmentA,
ElementB, LayoutB, AlignmentB,
ElementAccumulator,
TilesShape, ClusterShape,
cutlass::gemm::collective::StageCountAuto,
cutlass::gemm::collective::KernelScheduleAuto
>::CollectiveOp;
// Adım 2: Toplu katman epilog türünü belirt
using CollectiveEpilogue = cutlass::epilogue::collective::DefaultEpilogue<
cutlass::gemm::TagToStrideC_t<LayoutC>,
cutlass::gemm::TagToStrideC_t<LayoutC>,
cutlass::epilogue::thread::LinearCombination<ElementC, 1, ElementAccumulator, ElementAccumulator>>;
// Adım 3: Ana döngü ve epilogu çekirdek katmanında bir araya getir
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
cute::Shape<int,int,int,int>, // Problem Şekli [M,N,K,L]
CollectiveMainloop,
CollectiveEpilogue
>;
// Adım 4: Kernel sınıfını sarmalayın ve ana bilgisayar tarafında öğeye erişim elde
using GemmHandle = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
Toplu Katman: Ana Döngü
Bir topluluk, birlikte çalışan thread’ler grubudur ve bu işlem, ana çekirdeği oluşturmak için tekrarlanabilir. Genellikle, bu bir threadblock veya kümeyi ifade eder. TiledMMA ve TiledCopy nesneleri, paralel işçilerin hesaplama ve kopyalama işlemlerine atanmasını sağlarken, Toplu katman bu işlemleri zamansal olarak organize etme sorumluluğuna sahiptir. Bu, boru hatları ve warp uzmanlaşma şemaları oluşturmayı ve boru hatlarını yönetmek için donanım hızlandırmalı senkronizasyon temelini kullanmayı içerir.
CUTLASS 3.x GEMM çekirdekleri, bir toplu ana döngü içerir. Bu, bir topluluğun gerçekleştirdiği tek bir ana döngü iterasyonunun temel bileşenlerini tanımlayan GEMM sınıf şablonudur; bu bileşenler arasında yükleme ve MMA prosedürleri yer alır. Bir topluluk ana döngüsü şu şekilde tanımlanabilir:
using CollectiveMainloop = cutlass::gemm::collective::CollectiveMma<
DispatchPolicy,
TileShape,
ElementA, // veri türü, örneğin float
StrideA, // örneğin M-major için Stride<_1, int>
ElementB, StrideB,
TiledMma,
GmemTiledCopyA, SmemLayoutAtomA, SmemCopyAtomA, TransformA,
GmemTiledCopyB, SmemLayoutAtomB, SmemCopyAtomB, TransformB
>;
Toplu ana döngü, daha düşük katmanlardan gelen soyutlamalar için bir bileşim noktasıdır; bir TiledMMA, her operand için GMEM’den SMEM yükleme için bir TiledCopy ve kaydedici kaynaklı MMA’lar için kullanılacak SMEM’e yükleme atomları gibi. Bu soyutlamalar büyük ölçüde ortogonal olup, farklı MMA işlemlerinin farklı kopya işlemleri ile birleştirilmesine ve kod yeniden kullanımının en üst düzeye çıkarılmasına olanak tanır.
Toplu Yapıcı
CollectiveMma, bir GEMM ana döngüsünü TiledCopy ve TiledMma nesneleri açısından belirlemek için çeşitli ayar düğmeleri sunar. Ancak bu esneklik bazı karmaşıklıkları da beraberinde getirir. Genellikle kullanıcılar, boru hatları, donanım yetenekleri ve kaynak kullanımı hakkında daha üst düzey düşünceler üzerinden bu nesneleri türetmek ister. CUTLASS, CollectiveBuilder arabirimi ile bu türetmeyi gerçekleştirebilir. Toplu ana döngü bildirimi CollectiveBuilder ile şu şekilde görünür:
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
ArchTag, // örneğin cute::arch::Sm90 Hopper için
OpClass, // örneğin cute::arch::OpClassTensorOp Tensor Çekirdekleri için
ElementA, LayoutA, AlignmentA,
ElementB, LayoutB, AlignmentB,
ElementAccumulator,
TileShape, ClusterShape,
StageCount, // örneğin cutlass::gemm::collective::StageCountAuto
KernelSchedule // örneğin cutlass::gemm::collective::KernelScheduleAuto
>::CollectiveOp;
Şablon argümanları, kullanıcı dostu ölçütlerden seçilir ve bunları toplu Mma şablonuna daha düşük seviyeli parametreleri türetmek için kullanır:
- Mimari özelleştirme: GPU mimarisi ve MMA operatörünün türü (örneğin, SIMT veya Tensor Çekirdekleri).
- Operand ve toplayıcı bilgileri: Operandlar ve toplayıcı için veri türleri, ve global bellek içerisindeki operandlar için hizalama ve derleme zamanı düzenleme bilgileri (örneğin, satır veya sütun-mayör).
- Tile şekilleri: TiledMma ve TiledCopy nesneleri ve SMEM düzenlemeleri deduktesinde kullanılır.
- Planlama bilgileri: Küme şekli, boru hattı aşama sayısı ve çekirdek planlaması, planlama algoritmasında kullanılır. Aşama sayısı ve çekirdek planlaması için varsayılan otomatik seçenekler vardır, bu da CUTLASS’a belirli bir mimari ve parametreler için en iyi seçeneği otomatik olarak seçmesini söyler.
Toplu Katman: Epilog
Toplu epilog, Toplu API’nin diğer yarısını temsil eder. Her ana döngü iterasyonundan sonra iş parçacıklarının çıktısını ve saklanmasını işlemekle yükümlüdür. Ana döngü ile benzer şekilde, toplu epilog, (genellikle eleman bazında işlemlerin yanı sıra, bazen de azaltma işlemlerini içerebilen) matematik işlemleri için bir kopyalama işlemi için bir bileşim noktasını temsil eder. CUTLASS’ın toplu epilogları, bu aktivasyon fonksiyonunu netha işlemine entegre eder ve gereksiz veri hareketlerini ortadan kaldırır.
CUTLASS’ta birçok epilog bulunmaktadır; bu epiloglar GitHub’da buradatanımlanmıştır. Şablon argümanları uygulamalara göre önemli ölçüde değişiklik gösterse de, genellikle aşağıdaki bilgileri içerir:
- İşlemler A ve D’yle ilgili veri türü ve derleme zamanı düzenleme bilgileri.
- Herhangi bir ek işleme işlemi belirten bir birleşim işlemi.
- GMEM mağazası ve herhangi bir SMEM sahnelenmesi için TiledCopy işlemleri.
- Toplu ana döngü ile aynı şekilde, küme boyutu, TMA kullanımı, warp uzmanlaşması vb. hakkında bilgi içeren dağıtım politikaları.
CollectiveBuilder epilog için, daha tutarlı ve yüksek bir seviyeli arayüz sunmaktadır:
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
ArchTag,
OpClass,
TileShape,
ClusterShape,
EpilogueTileType,
ElementAccumulator,
ElementCompute,
ElementC, GmemLayoutTagC, AlignmentC,
ElementD, GmemLayoutTagD, AlignmentD,
EpilogueScheduleType,
FusionOpOrCallbacks
>::CollectiveOp;
Bu argümanlardan bazıları ana döngü yapıcısında daha önce tanıdık olsa da, bazı yeni argümanlar bulunur:
- Epilog, matematik-kopya çakışmasını iyileştirmek için bir CTA tile’ını daha küçük parçalara ayırabilir.
- Toplayıcı, ana döngünün çıktısı artık epiloga bir girdi haline gelir. Epilog hesaplamaları, farklı bir ara veri türünde (verilen
ElementCompute) yapılabilir. - CUTLASS, yaygın birleşim işlemleri sunmaktadır. Kullanıcı, ayrıca Epilogue Visitor Trees kullanarak özel bir birleşim işlemi geliştirebilir. Epilogue Visitor Trees hakkında daha fazla bilgi için, bu Colfax öğretici kaynağını inceleyebilirsiniz.
- Epilog planlama türleri, TMA ve warp uzmanlaşmasının kullanımını tanımlar. Varsayılan
EpilogueScheduleAuto, CUTLASS’a en iyi seçeneği türetmesini söyler.
Her iki Toplu Yapıcının kullanımını görmek için, CUTLASS’ın örnek 49’una ve örnek 71’ye başvurabiliriz.
Çekirdek Katmanı
Toplu katman, bir topluluğun gerçekleştirdiği hesaplamayı tamamen tanımlar. Çekirdek katmanının görevi, toplu ana döngüsünü ve toplu epilogu bir cihaz çekirdeği haline getirmektir. Çekirdek katmanının giriş noktası cutlass::gemm::kernel::GemmUniversal sınıfıdır. Bu, bir durumsuz evrensel cihaz çekirdek olup, GEMM’i bir toplu ana döngü ve bir toplu epilogun bileşimi olarak uygular. Durumsuz demek, çağıranın çekirdeğin durumunu bunun için parametreler geçerek yönettiği anlamına gelir. Evrensel demekse, GemmUniversal‘ın hem 2.x hem de 3.x GEMM çekirdeklerine giriş noktası olduğu anlamına gelir. 3.x API’sinin temel kullanımı ise şöyle görünür:
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
ProblemShape, // örneğin tamamen genel bir GEMM için Shape<int, int, int>
CollectiveMainloop,
CollectiveEpilogue
>;
TiledMma ve TiledCopy gibi, CollectiveMainloop ve CollectiveEpilogue da GemmUniversal aracılığıyla birleştirilmiş ortogonal soyutlamalardır. İlk şablon argümanı, problem şekli, genellikle sıradan GEMM (bir sıralama ile 3. düzlem şekli) ile toplu GEMM (bir sıralama ile 4. düzlem şekli) arasında seçim yapılması için kullanılır ama ayrıca bazı problem boyutlarını statik olarak kısıtlama amacıyla da kullanılabilir.
GemmUniversal instansiyonları, cutlass/gemm/kernel/sm*_gemm_*.hpp biçimindeki dosyalarda bulunur. GemmUniversal, büyük ölçüde toplu ana döngünün planlama KernelSchedule parametreye dayanarak dağıtım yapar. Tüm instansiyonlar tutarlı bir arayüze sahiptir:
- Çekirdeğe geçecek argümanları (problem şekli, donanım hakkında bilgiler, tensörlerin işaretçileri ve düzenleri, epilog parametreleri dahil) geçmek için bir arayüz sağlamaktadır.
- Çekirdek boyutlarını almak, verilen donanımda çalışabilirliğini kontrol etmek ve epilog veya tile planlayıcı tarafından gereken global bellek çerçevesi ayarlamak için statik başlatma fonksiyonları sunmaktadır.
- En önemlisi,
operator()olarak tanımlanan çekirdek mantığını uygulamaktadır. Bu bir cihaz işlevidir; çekirdek katmanı tüm çekirdek yürütme mantığını içeriyor olsa da, henüz ev sahipliğinden yürütmek için bir yol sağlamaz.
Örneğin, Blackwell için TMA warp-uzmanlaşmış çekirdek burada tanımlanmıştır.
Tile Planlama
Çekirdek katmanı, bir tile planlayıcının belirtilmesi için de bir bileşim noktasıdır. Çekirdek planlaması, bir topluluğun içindeki işin zamanlamasını tanımlarken, tile planlayıcı, topluluklar arasındaki işin sırasını ve dağıtımını tanımlar. En temel tile planlayıcısı, bir çıktı tile için bir CTA atanır. CUTLASS 3.x, Hopper için ek olarak, her SM’ye bir CTA atayan ve her CTA’nın (potansiyel olarak) ömürleri boyunca birden fazla çıktı tile’ı hesapladığı kalıcı bir planlayıcıyı ve aynı zamanda bir çıktı tile işini K modu boyunca daha iyi yük dengelemesi için bölen Stream-K planlayıcısını uygular. Blackwell mimarisinde ise, bunun yerine, Küme Başlatma Kontrolü ile aynı türden planlayıcılar kullanılır. Daha detaylı bilgi için bu Colfax öğretici kaynağına göz atabilirsiniz.
Aşağıdaki çekirdeği, bir Stream-K tile planlayıcısı kullanacak şekilde genişletebiliriz:
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
cute::Shape<int,int,int,int>,
CollectiveMainloop,
CollectiveEpilogue,
cutlass::gemm::StreamKScheduler
>;
CUTLASS örnek 74, Stream-K planlayıcısı kullanarak daha ayrıntılı bir örnektir.
Cihaz Katmanı
Çekirdek başlatma dahil, küme desteği ile veya farklı cihazlar veya CUDA akışları üzerinde başlatma mantığı Cihaz katmanı ile gerçekleştirilir. Cihaz katmanına giriş noktası, cutlass::gemm::device::GemmUniversalAdapter sınıfıdır. Bu sınıf, GemmUniversal çekirdeğini durumlu ve yeniden kullanılabilir bir kolla sarar. Durumlu demek, kolla ifade edilen bir örneğin çalıştığı için gereksinim duyduğu durumu yönetmelidir. Yeniden kullanılabilir demekse, aynı kolla birden fazla kez farklı argümanlarla çekirdeği çağırabilmelidir.
GemmUniversalAdapter‘ın nasıl kullanılabileceğine dair örnek:
using GemmHandle = cutlass::gemm::kernel::GemmUniversalAdapter<GemmKernel>;
using Arguments = typename GemmHandle::Arguments; // GemmKernel'dan çıktısı
Arguments args {
cutlass::Gemm::kBatched, // mod (burada toplu GEMM)
cute::make_shape(M, N, K, L), // problem şekli
{A, stride_A, B, stride_B}, // ana döngü argümanları
{{alpha, beta}, C, stride_C, D, stride_D}, // epilog argümanları
make_kernel_hardware_info(device_id), // donanım bilgisi
{} // planlayıcı argümanları (burada varsayılan)
};
GemmHandle gemm;
// Verilen şekil ve donanım ile çalışılabilirliği kontrol et
cutlass::Status status;
status = GemmHandle::can_implement(args);
if (status != cutlass::Status::kSuccess) {
std::cerr << "Problem desteklenmiyorn";
exit(EXIT_FAILURE);
}
// Global bellek çalışma alanını ayarlayın
size_t workspace_size = GemmHandle::get_workspace_size(args);
cutlass::device_memory::allocation<uint8_t> workspace(workspace_size);
// GemmHandle durumunu argümanlardan başlatın
status = gemm.initialize(args, workspace.get());
if (status != cutlass::Status::kSuccess) {
std::cerr << "GEMM çekirdeğini başlatmada hatan";
exit(EXIT_FAILURE);
}
// Çekirdeği başlat
status = gemm.run(); // burada CUDA akışını ve CUDA ana bilgisayarını belirtebilirsiniz
if (status != cutlass::Status::kSuccess) {
std::cerr << "GEMM çekirdeğini başlatmada hatan";
exit(EXIT_FAILURE);
}
Özet
Bu yazıda, CUTLASS kütüphanesinin kavramsal olarak hiyerarşide düzenlendiğini ve her katmanda yer alan nesnelerin aşağıdaki katmanlardan ortogonal nesneler ile bileştirildiğini tartıştık. Bu tasarım, yüksek düzeyde kod yeniden kullanımına olanak tanıyan son derece özelleştirilebilir GEMM uygulamalarının geniş bir yelpazesini mümkün kılmaktadır. Serinin bir sonraki ve son yazısında, CUTLASS 4.0’daki değişiklikleri, özellikle de CuTe Python DSL‘yi inceleyeceğiz.
Daha fazla bilgi için yazılımı GitHub üzerinden indirerek, belgelerimiz‘i okuyabilir veya daha derin tartışmalar için Geliştirici Forumumuza katılabilirsiniz.
Teşekkürler
Bu yazıya katkılarından dolayı Cris Cecka, Jack Kosaian, Mark Hoemmen, Haicheng Wu ve Matt Nicely’e teşekkürler.
