arXiv subjects

Haojun Sun

Publications and source records attributed to Haojun Sun.

5 recordsLinked to original sources

arXiv2026

ForgeKernel: Self-Evolving Kernels Catching and Surpassing Experts and Agents

Purpose-built kernel agents can write high-performance GPU kernels, but existing systems generally lack cross-run knowledge accumulation and architectural evolution. Each run starts from the same initial knowledge, and later iterations largely optimize within an architecture fixed early in the search. We present ForgeKernel, a kernel-generation system that accumulates and reuses knowledge across runs while enabling continued evolution of kernel architectures. It analyzes the hardware behavior of a high-performance reference and distills validated architectural decisions into reusable knowledge of hardware mechanisms rather than imitating source code. It then seals a kernel that passes correctness and performance validation and uses it as the reference for the next optimization run, enabling architectural decisions to be revisited. Across 27 operators and 29 input configurations on NVIDIA Hopper and Blackwell GPUs, ForgeKernel matches or beats its expert reference on 25 configurations, with a geometric-mean performance gain of 10%. On Ascend 910C, it outperforms the per-operator public best on the public case set of all 8 level-4 operators in the CANN-Bench Season 1 final leaderboard, with a geometric-mean gain of 25.1%, while its full-suite submission ranks first on the A3 leaderboard. As of 2026-10-30, public submissions generated by ForgeKernel rank first on all three CANN-Bench 2.0 solution leaderboards (Ascend 910C, 950DT, and 950PR) and on the SOL-ExecBench FlashInfer-Bench track.

cs.DC↗

arXiv2026

JustGoal: With Only a Goal, the AI Baseline Is Above the Expert Kernels

A small number of GPU kernels often determines the efficiency of large language model training and inference. High-performance kernel development typically relies on expert knowledge through manual implementation, collaboration with general-purpose models, or specialized agents equipped with domain knowledge, evaluation methods, and search workflows. We investigate how much value this knowledge provides as foundation models become stronger. We introduce JustGoal, a goal-only baseline in which a general-purpose agent receives only the optimization objective and constraints, without kernel-specific infrastructure, knowledge bases, or predefined workflows, and builds its own evaluation loop. On NVIDIA B200 GPUs, JustGoal outperforms open-source libraries on 8 of 10 operators, achieving per-operator geometric mean speedups of 1.14x to 2.34x. It also surpasses the fastest publicly released specialized-agent kernel on 5 of 9 operators with available comparisons. In a controlled same-model comparison using Fable 5.1 on single shapes, the KDA-0.5 harness reaches 0.95x of JustGoal's performance, while single-mechanism ablations show no consistent benefit from domain knowledge or predefined workflows. These results indicate that specialized machinery should be justified by clear gains over a same-model, same-scope goal-only baseline and reevaluated whenever the foundation model changes.

cs.DC↗

arXiv2026

ForgeTrain: Forging Production-Grade Training Frameworks via Harness-Driven AI Development

Training large models still relies on general-purpose frameworks such as Megatron-LM, whose generality tax constrains scenario-specific optimization and adds runtime overhead through accumulated abstraction. AI code generation reduces the cost of building a framework, and makes it affordable to forge one per scenario. We propose Forge Engineering: building a dedicated implementation from scratch for each scenario and iteratively optimizing it toward peak performance under correctness and usability constraints. Dedicated implementations inherit no abstraction boundaries, so they can integrate optimizations across the stack and reach a higher performance ceiling. We instantiate this paradigm for training frameworks as ForgeTrain, which holds a trusted framework as a golden reference and relaxes equivalence monotonically from Bit-for-Bit to Surpass. Experiments across multiple model--hardware configurations show that ForgeTrain consistently produces correct training engines and improves MFU over established training frameworks by 4.7--33.2%. To our knowledge this is the first production-grade training framework forged end-to-end by AI to match or surpass its human reference.

cs.SE↗

arXiv2026

MiniCPM-SALA: Hybridizing Sparse and Linear Attention for Efficient Long-Context Modeling

The evolution of large language models (LLMs) towards applications with ultra-long contexts faces challenges posed by the high computational and memory costs of the Transformer architecture. While existing sparse and linear attention mechanisms attempt to mitigate these issues, they typically involve a trade-off between memory efficiency and model performance. This paper introduces MiniCPM-SALA, a 9B-parameter hybrid architecture that integrates the high-fidelity long-context modeling of sparse attention (InfLLM-V2) with the global efficiency of linear attention (Lightning Attention). By employing a layer selection algorithm to integrate these mechanisms in a 1:3 ratio and utilizing a hybrid positional encoding (HyPE), the model maintains efficiency and performance for long-context tasks. Furthermore, we introduce a cost-effective continual training framework that transforms pre-trained Transformer-based models into hybrid models, which reduces training costs by approximately 75% compared to training from scratch. Extensive experiments show that MiniCPM-SALA maintains general capabilities comparable to full-attention models while offering improved efficiency. On a single NVIDIA A6000D GPU, the model achieves up to 3.5x the inference speed of the full-attention model at the sequence length of 256K tokens and supports context lengths of up to 1M tokens, a scale where traditional full-attention 8B models fail because of memory constraints.

cs.CL↗

arXiv2024

TMPQ-DM: Joint Timestep Reduction and Quantization Precision Selection for Efficient Diffusion Models

Diffusion models have emerged as preeminent contenders in the realm of generative models. Distinguished by their distinctive sequential generative processes, characterized by hundreds or even thousands of timesteps, diffusion models progressively reconstruct images from pure Gaussian noise, with each timestep necessitating full inference of the entire model. However, the substantial computational demands inherent to these models present challenges for deployment, quantization is thus widely used to lower the bit-width for reducing the storage and computing overheads. Current quantization methodologies primarily focus on model-side optimization, disregarding the temporal dimension, such as the length of the timestep sequence, thereby allowing redundant timesteps to continue consuming computational resources, leaving substantial scope for accelerating the generative process. In this paper, we introduce TMPQ-DM, which jointly optimizes timestep reduction and quantization to achieve a superior performance-efficiency trade-off, addressing both temporal and model optimization aspects. For timestep reduction, we devise a non-uniform grouping scheme tailored to the non-uniform nature of the denoising process, thereby mitigating the explosive combinations of timesteps. In terms of quantization, we adopt a fine-grained layer-wise approach to allocate varying bit-widths to different layers based on their respective contributions to the final generative performance, thus rectifying performance degradation observed in prior studies. To expedite the evaluation of fine-grained quantization, we further devise a super-network to serve as a precision solver by leveraging shared quantization results. These two design components are seamlessly integrated within our framework, enabling rapid joint exploration of the exponentially large decision space via a gradient-free evolutionary search algorithm.

cs.CV↗