(Accelerated Model Architectures)

XMA is a repository comprising of fast kernels for model training.
We are planning on adding lots of experimental and fun model architectures with support for multiple accelerators like NVIDIA, AMD GPUs and Google TPUs.

layers

Module	Triton	CUDA
GRU	✅	❌
MoE	✅	✅
RNN	✅	❌

functional

Module	Triton	CUDA
bmm	✅	❌
continuous_count	❌	✅
cross_entropy	✅	❌
fused_linear_cross_entropy	✅	❌
fused_residual_add_rmsnorm	✅	❌
grouped_gemm	❌	✅
rmsnorm	✅	❌
pack_sequence	✅	✅
softmax	✅	❌
swiglu	✅	✅
swiglu_packed	✅	❌
unpack_sequence	✅	✅

Discord Server

Join the discord server if you are interested in LLM architecture or distributed training/inference research.