Starrick Liu
421db17d53
This release introduces significant performance optimizations, memory efficiency improvements, and enhanced system robustness: 🚀 Performance Optimizations: - Add three new fused CUDA kernels (rope_index, rot_pos_emb, get_window_index) for accelerated multimodal preprocessing - Implement FSDP2 support for distributed training with improved memory efficiency - Add Torch.compile integration for additional performance gains - Optimize memory usage: reduce peak allocation from 48GB to 24GB on 8-GPU setup 🔧 System Robustness: - Fix missing token position inputs in prediction pipeline - Add type-robust negation operations in RoPE CUDA kernels (half/bfloat16 support) - Fix dataset root parameter initialization in LeRobot data loader - Enhanced error handling and input validation across fusion operators 📚 Documentation & Usability: - Add comprehensive memory usage benchmarks and hardware recommendations - Update citation format with proper arXiv reference - Improve training configuration documentation with quick start guide - Add detailed API documentation for new fusion operators 🛠️ Technical Details: - Version bump to 1.0.1 - New CUDA kernels: rope_index.cu, rot_pos.cu, window_index.cu - FSDP2 state dict loading with distribute_tensor support - Enhanced multimodal RoPE with 3D position encoding - Window attention optimization for Vision Transformers Breaking Changes: None - all changes are backward compatible
22 lines
963 B
Plaintext
22 lines
963 B
Plaintext
#include "dual_asym_grouped_gemm.h"
|
|
#include "permute.h"
|
|
#include "rope.h"
|
|
#include "rope_index.h"
|
|
#include "rot_pos.h"
|
|
#include "window_index.h"
|
|
|
|
#include <torch/extension.h>
|
|
|
|
|
|
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
|
|
m.def("asym_dual_gmm", &AsymmetricDualExpertGemm, "Asymmetric Dual Expert Grouped GEMM.");
|
|
m.def("permute", &moe_permute_topK_op, "Token permutation kernel");
|
|
m.def("unpermute", &moe_recover_topK_op, "Token un-permutation kernel");
|
|
m.def("unpermute_bwd", &moe_recover_topK_bwd_op, "Token un-permutation backward kernel");
|
|
m.def("rope", &launch_multimodal_rope_forward, "Multimodal RoPE forward kernel");
|
|
m.def("rope_bwd", &launch_multimodal_rope_backward, "Multimodal RoPE backward kernel");
|
|
m.def("rope_index", &get_rope_index, "Get RoPE index kernel");
|
|
m.def("rot_pos_emb", &fused_rot_pos_emb_cuda, "Fused Rotary Position Embedding kernel");
|
|
m.def("get_window_index", &get_window_index_cuda, "Get window index kernel");
|
|
}
|