6 papers
OmniMoE: An Efficient MoE by Orchestrating Atomic Experts at Scale
Jingze Shi, Zhangyang Peng, Yizhang Zhu +3
Mixture-of-Experts (MoE) architectures are evolving towards finer granularity to improve parameter efficiency. However, existing MoE designs face an inherent trade-off between the…
Towards Automated Kernel Generation in the Era of LLMs
Yang Yu, Peiyu Zang, Chi Hsu Tsai +11
The performance of modern AI systems is fundamentally constrained by the quality of their underlying GPU kernels, which translate high-level algorithmic semantics into low-level ha…
Trainable Dynamic Mask Sparse Attention
Jingze Shi, Yifan Wu, Yiran Peng +4
The increasing demand for long-context modeling in large language models (LLMs) is bottlenecked by the quadratic complexity of the standard self-attention mechanism. The community…
TransXSSM: A Hybrid Transformer State Space Model with Unified Rotary Position Embedding
Bingheng Wu, Jingze Shi, Yifan Wu +2
Transformers exhibit proficiency in capturing long-range dependencies, whereas State Space Models (SSMs) facilitate linear-time sequence modeling. Notwithstanding their synergistic…
Concise Reasoning, Big Gains: Pruning Long Reasoning Trace with Difficulty-Aware Prompting
Yifan Wu, Jingze Shi, Bingheng Wu +4
Existing chain-of-thought (CoT) distillation methods can effectively transfer reasoning abilities to base models but suffer from two major limitations: excessive verbosity of reaso…
Wonderful Matrices: Combining for a More Efficient and Effective Foundation Model Architecture
Jingze Shi, Bingheng Wu
In order to make the foundation model more efficient and effective, our idea is combining sequence transformation and state transformation. First, we prove the availability of rota…