arXiv subjects

Lingchao Zheng

Publications and source records attributed to Lingchao Zheng.

11 recordsLinked to original sources

arXiv2026

CANN Bench: Benchmarking Agent Generated Kernels against Real NPU and Algorithmic Limits

AI agents are now capable of writing, compiling, and iteratively optimizing low-level operator kernels on different hardware platforms. Existing benchmarks, however, focus almost exclusively on CUDA and Triton, leaving hardware ecosystems with less-exposed programming models without a common evaluation baseline. We present CANN Bench, an open benchmark for AI-generated operator code on Huawei's Ascend NPU. The current release covers 53 operators and 1060 test cases organized into four difficulty tiers -- from simple elementwise primitives to MoE dispatch and FlashAttention kernels -- spanning FP16, BF16, FP32, and INT8 precision formats. Evaluation adopts a \textbf{three-dimensional weighted composite score} that treats compilation, functional correctness, and performance as independent axes, providing a principled reward signal for kernel-generation agents. Performance is graded against an out-of-the-box PyTorch-on-Ascend baseline and an analytical per-case Hardware-Anchored Performance (HAP) limit on real NPU hardware, ensuring scores reflect genuine optimization headroom rather than measurement artifacts. The evaluation harness is designed to resist reward hacking from the ground up. CANN Bench is versioned within the official CANN repository and is designed for long-term community co-construction, providing the Ascend ecosystem with a quantitative, reproducible, and sustainably maintained yardstick for AI operator-authoring capability.

cs.AI↗

arXiv2026

AIS: Adaptive Importance Sampling for Quantized RL

Reinforcement learning (RL) for large language models (LLMs) is dominated by the cost of rollout generation, which has motivated the use of low-precision rollouts (e.g., FP8) paired with a BF16 trainer to improve throughput and reduce memory pressure. This introduces a rollout-training mismatch that biases the policy gradient and can cause training to collapse outright on reasoning benchmarks. We show that the mismatch is non-stationary and acts as a double-edged sword: early in training it provides a stochastic exploration bonus, exposing the gradient to trajectories the trainer would otherwise under-sample, but the same perturbation transitions into a destabilizing source of bias as the policy concentrates. To solve this, we propose Adaptive Importance Sampling (AIS), a correction framework that adjusts the strength of its intervention on a per-batch basis. AIS combines three real-time diagnostics, namely weight reliability, divergence severity, and variance amplification, into a single mixing coefficient that interpolates between the uncorrected and fully importance-weighted gradients, suppressing the destabilizing component of the mismatch while preserving its exploratory benefit. We integrate AIS into GRPO and evaluate it on the diffusion-based LLaDA-8B-Instruct and the autoregressive Qwen3-8B and Qwen3.5-9B across mathematical reasoning and planning benchmarks. AIS matches the BF16 baseline on most tasks while retaining the 1.5 to 2.76x rollout speedup of FP8.

stat.ML↗

arXiv2026

Multi-Scale Dequant: Eliminating Dequantization Bottleneck via Activation Decomposition for Efficient LLM Inference

Quantization is essential for efficient large language model (LLM) inference, yet the dequantization step-converting low-bit weights back to high-precision for matrix multiplication has become a critical bottleneck on modern AI accelerators. On architectures with decoupled compute units (e.g., Ascend NPUs), dequantization operations can consume more cycles than the matrix multiplication itself, leaving the high-throughput tensor cores underutilized. This paper presents Multi-Scale Dequant (MSD), a quantization framework that removes weight/KV dequantization from the GEMM critical path. Instead of lifting low-bit weights to BF16 precision, MSD decomposes high-precision BF16 activations into multiple low-precision components, each of which can be multiplied directly with quantized weights via native hardware-accelerated GEMM. This approach shifts the computational paradigm from precision conversion to multi-scale approximation, avoiding INT8-to-BF16 weight conversion before GEMM. We instantiate MSD for two weight formats and derive tight error bounds for each. For INT8 weights (W4A16), two-pass INT8 decomposition achieves near 16 effective bits. For MXFP4 weights (W4A16), two-pass MXFP4 decomposition yields near 6.6 effective bits with error bound 1/64 per block surpassing single-pass MXFP8(5.24 bits) while maintaining the same effective GEMM compute time. We further derive closed-form latency and HBM traffic models showing that MSD avoids the Vector-Cube pipeline stall caused by dequantization and reduces KV cache HBM traffic by up to 2.5 times in attention. Numerical simulations on matrix multiplication and Flash Attention kernels confirm that MSD does not degrade accuracy compared to dequantization baselines, and in many settings achieves lower L2 error.

stat.ML↗

arXiv2025

LoPT: Lossless Parallel Tokenization Acceleration for Long Context Inference of Large Language Model

Long context inference scenarios have become increasingly important for large language models, yet they introduce significant computational latency. While prior research has optimized long-sequence inference through operators, model architectures, and system frameworks, tokenization remains an overlooked bottleneck. Existing parallel tokenization methods accelerate processing through text segmentation and multi-process tokenization, but they suffer from inconsistent results due to boundary artifacts that occur after merging. To address this, we propose LoPT, a novel Lossless Parallel Tokenization framework that ensures output identical to standard sequential tokenization. Our approach employs character-position-based matching and dynamic chunk length adjustment to align and merge tokenized segments accurately. Extensive experiments across diverse long-text datasets demonstrate that LoPT achieves significant speedup while guaranteeing lossless tokenization. We also provide theoretical proof of consistency and comprehensive analytical studies to validate the robustness of our method.

cs.CL↗

arXiv2025

AMLA: MUL by ADD in FlashAttention Rescaling

Multi-head Latent Attention (MLA) significantly reduces KVCache memory usage in Large Language Models while introducing substantial computational overhead and intermediate variable expansion. This poses challenges for efficient hardware implementation -- especially during the decode phase. This paper introduces Ascend MLA (AMLA), a high-performance kernel specifically optimized for Huawei's Ascend NPUs. AMLA is built on two core innovations: (1) A novel FlashAttention-based algorithm that replaces floating-point multiplications with integer additions for output block rescaling, leveraging binary correspondence between FP32 and INT32 representations; (2) A Preload Pipeline strategy with hierarchical tiling that maximizes FLOPS utilization: the Preload Pipeline achieves Cube-bound performance, while hierarchical tiling overlaps data movement and computation within the Cube core. Experiments show that on Ascend 910 NPUs (integrated in CloudMatrix384), AMLA achieves up to 614 TFLOPS, reaching 86.8% of the theoretical maximum FLOPS, outperforming the state-of-the-art open-source FlashMLA implementation, whose FLOPS utilization is up to 66.7% on NVIDIA H800 SXM5. The AMLA kernel has been integrated into Huawei's CANN and will be released soon.

cs.LG↗

arXiv2021

Direct Flux Gradient Approximation to Close Moment Model for Kinetic Equations

To close the moment model deduced from kinetic equations, the canonical approach is to provide an approximation to the flux function not able to be depicted by the moments in the reduced model. In this paper, we propose a brand new closure approach with remarkable advantages than the canonical approach. Instead of approximating the flux function, the new approach close the moment model by approximating the flux gradient. Precisely, we approximate the space derivative of the distribution function by an ansatz which is a weighted polynomial, and the derivative of the closing flux is computed by taking the moments of the ansatz. Consequently, the method provides us an improved framework to derive globally hyperbolic moment models, which preserve all those conservative variables in the low order moments. It is shown that the linearized system at the weight function, which is often the local equilibrium, of the moment model deduced by our new approach is automatically coincided with the system deduced from the classical perturbation theory, which can not be satisfied by previous hyperbolic regularization framework. Taking the Boltzmann equation as example, the linearlization of the moment model gives the correct Navier-Stokes-Fourier law same as that the Chapman-Enskog expansion gives. Most existing globally hyperbolic moment models are re-produced by our new approach, and several new models are proposed based on this framework.

physics.comp-ph↗

arXiv2020

A Nonlinear Moment Model for Radiative Transfer Equation

We derive a nonlinear moment model for radiative transfer equation in 3D space, using the method to derive the nonlinear moment model for the radiative transfer equation in slab geometry. The resulted 3D HMPN model enjoys a list of mathematical advantages, including global hyperbolicity, rotational invariance, physical wave speeds, spectral accuracy, and correct higher-order Eddington approximation. Simulation examples are presented to validate the new model numerically.

physics.comp-ph↗

arXiv2020

A Nonlinear Hyperbolic Model for Radiative Transfer Equation in Slab Geometry

Linear models for the radiative transfer equation have been well developed, while nonlinear models are seldom investigated even for slab geometry due to some essential difficulties. We have proposed a moment model in MPN for slab geometry which combines the ideas of the classical PN and MN model. Though the model is far from perfect, it was demonstrated to be quite efficient in numerically approximating the solution of the radiative transfer equation, that we are motivated to further improve this model. Consequently we propose in this paper a new model following the chartmap in MPN with some significant theoretic progresses. The new model is derived with global hyperbolicity, and meanwhile some necessary physical properties are preserved. We give a complete analysis to the characteristic structure and propose a numerical scheme for the new model. Numerical examples are presented to demonstrate the numerical performance of the new model.

math.NA↗

arXiv2019

A Nonlinear Three-Moment Model for Radiative Transfer in Spherical Symmetry

We study the approximation of the radiative transfer equation with a relatively few moments in the spherically symmetric case. We propose a three-moment model based on choosing the beta distribution as the ansatz for the specific intensity. This ansatz enables our model to capture the anisotropy in the distribution function. The characteristic structure of the Riemann problem of the model is studied in detail. Numerical simulations demonstrate its validity in approximating the radiative transfer equation in the spherically symmetric case and its advantage in approximating highly anisotropic distribution functions in comparison to the $P_n$ method.

physics.comp-ph↗

arXiv2018

A Nonlinear Moment Model for Radiative Transfer Equation in Slab Geometry

This paper is concerned with the approximation of the radiative transfer equation for a grey medium in the slab geometry by the moment method. We develop a novel moment model inspired by the classical $P_N$ model and $M_N$ model. The new model takes the ansatz of the $M_1$ model as the weight function and follows the primary idea of the $P_N$ model to approximate the specific intensity by expanding it around the weight function in terms of orthogonal polynomials. The weight function uses the information of the first two moments, which brings the new model the capability to approximate an anisotropic distribution. Mathematical properties of the moment model are investigated, and particularly the hyperbolicity and the characteristic structure of the Riemann problem of the model with three moments are studied in detail. Some numerical simulations demonstrate its numerical efficiency and show its superior in comparison to the $P_N$ model.

physics.class-ph↗

arXiv2017

Linear Stability of Hyperbolic Moment Models for Boltzmann Equation

Grad's moment models for Boltzmann equation were recently regularized to globally hyperbolic systems, and thus the regularized models attain local well-posedness for Cauchy data. The hyperbolic regularization is only related to the convection term in Boltzmann equation. We in this paper studied the regularized models with the presentation of collision terms. It is proved that the regularized models are linear stability at the local equilibrium and satisfy Yong's first stability condition with commonly used approximate collision terms, and particularly with Boltzmann's binary collision model.

math-ph↗