arXiv subjects

Yarong Mu

Publications and source records attributed to Yarong Mu.

8 recordsLinked to original sources

arXiv2026

AoiZora: Topology-Aware Auto-Parallel Optimization for Inference of Diffusion Transformers

Video diffusion has quickly grown into a key generative serving workload, yet producing each clip demands many denoising iterations over large spatio-temporal latents, which puts low-latency inference out of reach on a single device. A denoising step is therefore typically distributed across multiple accelerators, and TPU sub-slices have become an attractive and practical fabric for doing so. Current auto-parallel systems, however, search almost exclusively over logical device meshes and disregard how a chosen sharding is actually laid out on the physical TPU interconnect -- an oversight that leaves large, topology-dependent performance on the table. We address this gap with AoiZora, a compiler-mediated topology planner built for low-latency video diffusion inference on TPU sub-slices. Its guiding principle is to reconnect logical sharding with physical placement by drawing on different points in the compilation flow: AoiZora first eliminates weak sharding candidates from inexpensive pre-compilation IRs, then compiles only the ones that survive and orders their physical placements using compiled HLO together with a topology-aware communication model. The winning plan is realized along the ordinary compiler path, leaving model code, compiler lowering, collective kernels, and network routing entirely intact. On TPU v5e sub-slices, AoiZora reduces Wan 2.1 one-step denoising latency by as much as 1.42x relative to existing solutions.

cs.DC↗

arXiv2026

Ragged Paged Attention: A High-Performance and Flexible LLM Inference Kernel for TPU

Large Language Model (LLM) deployment is increasingly shifting to cost-efficient accelerators like Google's Tensor Processing Units (TPUs), prioritizing both performance and total cost of ownership (TCO). However, existing LLM inference kernels and serving systems remain largely GPU-centric, and there is no well-established approach for efficiently mapping LLM workloads onto TPU architectures--particularly under the dynamic and ragged execution patterns common in modern serving. In this paper, we present Ragged Paged Attention (RPA), a high-performance and flexible attention kernel for TPUs, implemented using Pallas and Mosaic. RPA addresses these challenges through three key techniques: (1) fine-grained tiling to enable efficient dynamic slicing over ragged memory, (2) a custom software pipeline that fuses KV cache updates with attention computation, and (3) a distribution-aware compilation strategy that generates specialized kernels for decode, prefill, and mixed workloads. Evaluated on Llama 3 8B on TPU7x, RPA achieves up to 86% memory bandwidth utilization (MBU) in decode and 73% model FLOPs utilization (MFU) in prefill. Integrated as the primary TPU backend in vLLM and SGLang, RPA provides a production-grade foundation for efficient TPU inference and offers practical insights into kernel design.

cs.PF↗

arXiv2022

Recycling Scraps: Improving Private Learning by Leveraging Intermediate Checkpoints

In this work, we focus on improving the accuracy-variance trade-off for state-of-the-art differentially private machine learning (DP ML) methods. First, we design a general framework that uses aggregates of intermediate checkpoints \emph{during training} to increase the accuracy of DP ML techniques. Specifically, we demonstrate that training over aggregates can provide significant gains in prediction accuracy over the existing state-of-the-art for StackOverflow, CIFAR10 and CIFAR100 datasets. For instance, we improve the state-of-the-art DP StackOverflow accuracies to 22.74\% (+2.06\% relative) for $\epsilon=8.2$, and 23.90\% (+2.09\%) for $\epsilon=18.9$. Furthermore, these gains magnify in settings with periodically varying training data distributions. We also demonstrate that our methods achieve relative improvements of 0.54\% and 62.6\% in terms of utility and variance, on a proprietary, production-grade pCVR task. Lastly, we initiate an exploration into estimating the uncertainty (variance) that DP noise adds in the predictions of DP ML models. We prove that, under standard assumptions on the loss function, the sample variance from last few checkpoints provides a good approximation of the variance of the final model of a DP run. Empirically, we show that the last few checkpoints can provide a reasonable lower bound for the variance of a converged DP model. Crucially, all the methods proposed in this paper operate on \emph{a single training run} of the DP ML technique, thus incurring no additional privacy cost.

cs.LG↗

arXiv2022

Strings And Colorings Of Topological Coding Towards Asymmetric Topology Cryptography

We, for anti-quantum computing, will discuss various number-based strings, such as number-based super-strings, parameterized strings, set-based strings, graph-based strings, integer-partitioned and integer-decomposed strings, Hanzi-based strings, as well as algebraic operations based on number-based strings. Moreover, we introduce number-based string-colorings, magic-constraint colorings, and vector-colorings and set-colorings related with strings. For the technique of encrypting the entire network at once, we propose graphic lattices related with number-based strings, Hanzi-graphic lattices, string groups, all-tree-graphic lattices. We study some topics of asymmetric topology cryptography, such as topological signatures, Key-pair graphs, Key-pair strings, one-encryption one-time and self-certification algorithms. Part of topological techniques and algorithms introduced here are closely related with NP-complete problems or NP-hard problems.

cs.IT↗

arXiv2019

Topological Coding and Topological Matrices Toward Network Overall Security

A mathematical topology with matrix is a natural representation of a coding relational structure that is found in many fields of the world. Matrices are very important in computation of real applications, s ce matrices are easy saved in computer and run quickly, as well as matrices are convenient to deal with communities of current networks, such as Laplacian matrices, adjacent matrices in graph theory. Motivated from convenient, useful and powerful matrices used in computation and investigation of today's networks, we have introduced Topcode-matrices, which are matrices of order $3\times q$ and differ from popular matrices applied in linear algebra and computer science. Topcode-matrices can use numbers, letters, Chinese characters, sets, graphs, algebraic groups \emph{etc.} as their elements. One important thing is that Topcode-matrices of numbers can derive easily number strings, since number strings are text-based passwords used in information security. Topcode-matrices can be used to describe topological graphic passwords (Topsnut-gpws) used in information security and graph connected properties for solving some problems coming in the investigation of Graph Networks and Graph Neural Networks proposed by GoogleBrain and DeepMind. Our topics, in this article, are: Topsnut-matrices, Topcode-matrices, Hanzi-matrices, adjacency ve-value matrices and pan-Topcode-matrices, and some connections between these Topcode-matrices will be proven. We will discuss algebraic groups obtained from the above matrices, graph groups, graph networking groups and number string groups for encrypting different communities of dynamic networks. The operations and results on our matrices help us to set up our overall security mechanism to protect networks.

cs.IT↗

arXiv2019

Using Chinese Characters To Generate Text-Based Passwords For Information Security

Graphical passwords (GPWs) are in many areas of the current world. Topological graphic passwords (Topsnut-gpws) are a new type of cryptography, and they differ from the existing GPWs. A Topsnut-gpw consists of two parts: one is a topological structure (graph), and one is a set of discrete elements (a graph labelling, or coloring), the topological structure connects these discrete elements together to form an interesting "story". Our idea is to transform Chinese characters into computer and electronic equipments with touch screen by speaking, writing and keyboard for forming Hanzi-graphs and Hanzi-gpws. We will use Hanzigpws to produce text-based passwords (TB-paws). We will introduce flawed graph labellings on disconnected Hanzi-graphs.

cs.IT↗

arXiv2018

Text-based Passwords Generated From Topological Graphic Passwords

Topological graphic passwords (Topsnut-gpws) are one of graph-type passwords, but differ from the existing graphical passwords, since Topsnut-gpws are saved in computer by algebraic matrices. We focus on the transformation between text-based passwords (TB-paws) and Topsnut-gpws in this article. Several methods for generating TB-paws from Topsnut-gpws are introduced; these methods are based on topological structures and graph coloring/labellings, such that authentications must have two steps: one is topological structure authentication, and another is text-based authentication. Four basic topological structure authentications are introduced and many text-based authentications follow Topsnut-gpws. Our methods are based on algebraic, number theory and graph theory, many of them can be transformed into polynomial algorithms. A new type of matrices for describing Topsnut-gpws is created here, and such matrices can produce TB-paws in complex forms and longer bytes. Estimating the space of TB-paws made by Topsnut-gpws is very important for application. We propose to encrypt dynamic networks and try to face: (1) thousands of nodes and links of dynamic networks; (2) large numbers of Topsnut-gpws generated by machines rather than human's hands. As a try, we apply spanning trees of dynamic networks and graphic groups (Topsnut-groups) to approximate the solutions of these two problems. We present some unknown problems in the end of the article for further research.

cs.IT↗

arXiv2018

Topological Graphic Passwords And Their Matchings Towards Cryptography

Graphical passwords (GPWs) are convenient for mobile equipments with touch screen. Topological graphic passwords (Topsnut-gpws) can be saved in computer by classical matrices and run quickly than the existing GPWs. We research Topsnut-gpws by the matching of view, since they have many advantages. We discuss: configuration matching partition, coloring/labelling matching partition, set matching partition, matching chain, etc. And, we introduce new graph labellings for enriching Topsnut-matchings and show that these labellings can be realized for trees or spanning trees of networks. In theoretical works we explore Graph Labelling Analysis, and show that every graph admits our extremal labellings and set-type labellings in graph theory. Many of the graph labellings mentioned are related with problems of set matching partitions to number theory, and yield new objects and new problems to graph theory.

cs.CR↗