$ show microsoft/Tutel
Tutel is a high-performance library for Mixture-of-Experts (MoE) in PyTorch, focused on both training and inference of large-scale sparsely-gated models. The codebase is organized into distinct layers: the core `MOELayer` configuration (gating, experts, dispatch), custom NCCL-based distributed communication primitives (`net.py`), JIT-compiled CUDA/ROCm kernels for routing and sparse operations, and utility modules for checkpoint resharding across different distributed topologies. Heavily exploited features include low-precision quantized inference (FP4, FP8, block quantization), fused Expert comunication via all-to-all collectives, and a experimental auto-paritioning solver for adaptive parallelism. The repository is dominated by software engineering artifacts (Python library design, C++/CUDA/ROCm extensions and binaries, structured examples) rather than research notebooks or analytic pipeline, but its origins include published papers (MLSys 2023) and active development for contemporary models such as DeepSeek, Qwen3, and Kimi.
The LLM identified 3 files worth highlighting for reuse:
Provides a reusable abstraction for distributed MoE communication, including all-to-all, scatter, all-to-all, gather, all-reduce, and 2D all-to-all, and 2D all-to-all operations customized for variable-length token batches. It uses NCCL and custom stream management to overlap communication with computation. This module can be extracted into a project needing efficient, variable-length collective communication beyond standard all-to-all, and abstracts away complex multi-stream synchronization and buffer allocation.
The checkpoint package (gather.py and scatter.py) provides functions to redistributing MoE experts weights between different numbers of GPUs or world sizes. It supports both state-sharing formats and SwinTransformer-style checkpoints. This is a self-contained tool that can be extracted and adapted to other distributed models and frameworks that require conversion of sharded backend format as model is deployed on different compute topologies.
Implements a runtime kernel compilation infrastructure using nvcc/nvrtc. It provides classes to :class:`JitCompiler`, :class:`torch::RegisterJIT`, and :class:`two` that dynamically generate and load CUDA and ROCm kernels. This pattern allows architecture and shape-specific code without compile-time overhead. It could be repackaged as a standalone JIT library for PyTorch custom operations, supporting both `nvcc` and `nvrtc` backends.