arXiv subjects

Jiang Shao

Publications and source records attributed to Jiang Shao.

4 recordsLinked to original sources

arXiv2026

FalconGEMM: Surpassing Hardware Peaks with Lower-Complexity Matrix Multiplication

Peak breaking Matrix Multiplication is a promising technique to improve the performance of DL, especially in LLM training and inference. We present FalconGEMM, a cross-platform framework that automates the deployment, optimization, and selection of Lower-Complexity Matrix Multiplication Algorithms (LCMAs) across diverse hardware. There are three key innovations: (1) a Deployment Module that enables portable execution across various hardware and input configurations through code generation; (2) an Execution Module with Group-Parallel Optimizations that maximizes on-chip data reuse, utilizes parallel resources, and reduces bandwidth overhead; and (3) a Decision Module featuring a lightweight analytical performance model to select the optimal strategy based on matrix shapes and hardware profiles. Extensive evaluation is conducted on LLM workloads across GPU (H20, A100) and CPU (ARM, x86) architectures with multiple data types. FalconGEMM succeeds in delivering peak breaking performance and outperforms GEMM libraries (e.g., cuBLAS, CUTLASS, Intel MKL, etc) by 7.59%-17.85% and LCMA competitors like AlphaTensor by 12.41%-55.61%. Our framework makes the theoretical promise of LCMAs practical for production deployment across the heterogeneous landscape of modern hardware.

cs.DC↗

arXiv2026

Scalable Training of Mixture-of-Experts Models with Megatron Core

Scaling Mixture-of-Experts (MoE) training introduces systems challenges absent in dense models. Because each token activates only a subset of experts, this sparsity allows total parameters to grow much faster than per-token computation, creating coupled constraints across memory, communication, and computation. Optimizing one dimension often shifts pressure to another, demanding co-design across the full system stack. We address these challenges for MoE training through integrated optimizations spanning memory (fine-grained recomputation, offloading, etc.), communication (optimized dispatchers, overlapping, etc.), and computation (Grouped GEMM, fusions, CUDA Graphs, etc.). The framework also provides Parallel Folding for flexible multi-dimensional parallelism, low-precision training support for FP8 and NVFP4, and efficient long-context training. On NVIDIA GB300 and GB200, it achieves 1,233/1,048 TFLOPS/GPU for DeepSeek-V3-685B and 974/919 TFLOPS/GPU for Qwen3-235B. As a performant, scalable, and production-ready open-source solution, it has been used across academia and industry for training MoE models ranging from billions to trillions of parameters on clusters scaling up to thousands of GPUs. This report explains how these techniques work, their trade-offs, and their interactions at the systems level, providing practical guidance for scaling MoE models with Megatron Core.

cs.DC↗

arXiv2025

ByteQC: GPU-Accelerated Quantum Chemistry Package for Large-Scale Systems

Applying quantum chemistry algorithms to large-scale systems requires substantial computational resources scaled with the system size and the desired accuracy. To address this, ByteQC, a fully-functional and efficient package for large-scale quantum chemistry simulations, has been open-sourced at https://github.com/bytedance/byteqc, leveraging recent advances in computational power and many-body algorithms. Regarding computational power, several standard algorithms are efficiently implemented on modern GPUs, ranging from mean-field calculations (Hartree-Fock and density functional theory) to post-Hartree-Fock methods such as M{\o}ller-Plesset perturbation theory, random phase approximation, coupled cluster methods, and quantum Monte Carlo methods. For the algorithmic approach, we also employ a quantum embedding method, which significantly expands the tractable system size while preserving high accuracy at the gold-standard level. All these features have been systematically benchmarked. For standalone algorithms, the benchmark results demonstrate up to a 60$\times$ speedup when compared to 100-core CPUs. Additionally, the tractable system sizes have been significantly expanded: 1,610 orbitals for coupled cluster with single and double excitations (1,380 orbitals with perturbative triple excitations), 11,040 orbitals for M{\o}ller-Plesset perturbation theory of second order, 37,120 orbitals for mean-field calculations under open boundary conditions, and over 100,000 orbitals for periodic boundary conditions. For the advanced quantum embedding feature, two representative examples are demonstrated: the water cluster problem (2,752 orbitals) and a water monomer adsorbed on a boron nitride surface (3,929 orbitals), achieving the gold-standard accuracy.

physics.chem-ph↗

arXiv2023

SDT: A Low-cost and Topology-reconfigurable Testbed for Network Research

Network experiments are essential to network-related scientific research (e.g., congestion control, QoS, network topology design, and traffic engineering). However, (re)configuring various topologies on a real testbed is expensive, time-consuming, and error-prone. In this paper, we propose \emph{Software Defined Topology Testbed (SDT)}, a method for constructing a user-defined network topology using a few commodity switches. SDT is low-cost, deployment-friendly, and reconfigurable, which can run multiple sets of experiments under different topologies by simply using different topology configuration files at the controller we designed. We implement a prototype of SDT and conduct numerous experiments. Evaluations show that SDT only introduces at most 2\% extra overhead than full testbeds on multi-hop latency and is far more efficient than software simulators (reducing the evaluation time by up to 2899x). SDT is more cost-effective and scalable than existing Topology Projection (TP) solutions. Further experiments show that SDT can support various network research experiments at a low cost on topics including but not limited to topology design, congestion control, and traffic engineering.

cs.NI↗