arXiv Science⌕ Search

arXiv subjects

Xuehai Zhou

Publications and source records attributed to Xuehai Zhou.

At least 19 recordsLinked to original sources
arXiv2026

A Measure-Theoretic Analysis of Reasoning in Transformers: Structural Generalization and Approximation Limits

Scaling laws for reasoning in large language models (LLMs) are well documented, but little theory explains their out-of-distribution (OOD) generalization. We formalize reasoning with optimal transport: each search trajectory is mapped to twelve structural coordinates, such as tree depth and backtracking rate, and dataset shift is measured by the Wasserstein-1 distance. Kantorovich duality bounds the OOD risk through this distance, the architecture's Lipschitz constant and approximation limits. First, under absolute positional encoding (APE), shifting part of a trajectory to later positions almost surely changes the attention scores within that part; if this misroutes exact-match search, the Lipschitz constant and risk are $Ω(1)$. Shift-invariant encodings such as rotary position embedding (RoPE) leave the scores unchanged and bound the error. Second, we map backtracking to a Dyck-$k$ language and, under an implicit-depth assumption, derive a risk floor for search trees deeper than a multiple of the number of layers. Enough layers remove this floor; more width only shrinks a separate Barron-space approximation term. Experiments with 54 trained Transformers on combinatorial search (Countdown) support these bounds. Appendix experiments measure structural shifts in public agent, embodied and robot trajectories, shift all position indices of 33 public language models, and split deeply nested tasks into separate model calls.

cs.LG↗
arXiv2026

CLEARMIND: Closed-Form Laplacian Embeddings of Arakelov-Green Resistance for Morphological Intrinsic Neuronal Distances

Tropical geometry turns a metric graph into a flat torus, its tropical Jacobian, and measures distances between the images of its points by a closest vector problem on the period lattice. We prove that for points of the graph this problem is solved by geodesics: the squared tropical polarization distance equals the path metric minus the Arakelov-Green (AG) distance, an effective-resistance kernel given in closed form by the Laplacian pseudoinverse. Equivalently, the AG distance and the path metric are the minimal energies of real and of integral unit flows, and the squared tropical distance is exactly their integrality gap. The full tropical distance matrix is thus computable in cubic time. On this identity we build CLEARMIND, a training-free descriptor of 3D neuronal morphology. A reconstruction is reduced to its branching skeleton, neurite tips near the soma are joined to it, short bridges are contracted, and the AG matrix of the result is summarized by its spectrum. Every step has an exact algebraic description; for instance, the period matrix records the shared soma-to-tip path lengths, and the spectrum has exactly one positive eigenvalue, so its absolute values determine it. The signature needs no lattice search, is invariant to vertex order, rigid motions and subdivision, and is Lipschitz in the edge lengths. On ACT-4, JML-4 and BIL-6, appending the 64-dimensional spectrum improves a point-cloud GNN and MorphVAE on every dataset, by up to 20.3 points; a Tree-LSTM with the spectrum outperforms every reimplemented baseline; and an MLP on the spectrum alone is a competitive classifier that surpasses every reimplemented deep baseline on ACT-4. On BREC, a GIN with AG eigenvector encodings distinguishes 70.0% of the pairs beyond the 1-WL limit, including 33% of the CFI pairs, more than PPGN (23%) and I$^2$-GNN (21%), and the spectrum alone separates 217 of 400 pairs without training.

cs.LG↗
arXiv2026

ActKV: Efficient LLM Agents through Action-Guided KV Cache Management

Agentic LLM inference accumulates long KV caches across iterative observation-reasoning-action loops, imposing substantial memory overhead and limiting serving throughput. Existing compression methods emphasize overall output quality, overlooking the asymmetric importance of actions in driving task progress. Our key idea is to establish a compression criterion that values KV entries by their contribution to action generation and prioritizes action quality. However, iterative execution, dynamic memory demands, and scattered action-critical entries pose challenges to eviction policies, budget allocation, and paged memory integration. To this end, we propose ActKV, the first KV cache compression framework tailored for agentic LLM inference. (i) Action-oriented KV cache eviction exploits stable action access patterns to retain entries critical to future actions, supporting reliable task progress under compression. (ii) Confidence-driven adaptive budget allocation uses LLM's intrinsic confidence to adapt the budget to evolving action-critical memory demands. (iii) Page-aware compression management standardizes compression into three primitives with customized kernels, realizing practical throughput gains. On long-trace tasks, ActKV retains an average of 98.53% of FullKV's accuracy with only 25.98% of its peak KV cache memory. It also achieves 3.97 times and 3.58 times FullKV's token and task throughput, delivering state-of-the-art performance.

cs.OS↗
arXiv2026

Window-Diffusion: Accelerating Diffusion Language Model Inference with Windowed Token Pruning and Caching

Diffusion language models (DLMs) generate text through iterative denoising, but inference requires full-sequence attention at every iteration, resulting in substantial redundant computation on masked tokens. Block-wise diffusion can reduce this cost, yet it typically relies on retraining and constrained update orders, limiting its direct applicability to pretrained DLMs. Our token-level analysis reveals pronounced structural locality in DLM inference. Decoding is driven by a small set of prefix-localized active tokens; the influence of distant undecoded context diminishes rapidly, and decoded tokens exhibit stage-wise temporal stability, enabling reuse of intermediate representations except for a brief post-decode transient. Motivated by these observations, we propose \textbf{\placeholder}\footnote{The source code is available at https://github.com/vhicrgit/Window-Diffusion.}, a window-based token pruning and caching method for inference. We maintain a local computation window that slides rightward as denoising progresses, and partition undecoded tokens into: (i) \textit{active tokens} that are computed online, (ii) \textit{buffer tokens} whose KV states are cached and periodically refreshed, and (iii) \textit{far-field tokens} that are pruned outside the window. Computation is restricted to active and buffer tokens within the window, while far-field tokens are omitted at each stage. Experiments on LLaDA and Dream show that, under matched compute budgets, our method achieves up to $99\times$ inference speedup while largely preserving generation performance.

cs.LG↗
arXiv2026

SeqMoE: Toward Full-Load Performance via Predictive and Graph-Compatible MoE Offloading

Mixture-of-Experts (MoE) creates a structural advantage for offloading: only a small fraction of activated experts need to reside in device memory, and if they can be loaded in time for computation, offloading can in principle approach full-load performance, where all model weights reside in device memory. Yet translating MoE's structural advantage into practical offloading gains remains challenging. We propose SeqMoE to bridge this gap. To maximize expert hits, we build predictive memory management: (i) Sequence-to-sequence prediction. We are the first to recast expert activation prediction as sequence modeling, enabling accurate multi-step, multi-layer forecasts that provide a long and reliable window for downstream decisions. (ii) Joint prefetch scheduling. We formulate prefetch scheduling as Job Sequencing with Deadlines to maximize expected expert hits and improve bandwidth efficiency. (iii) Forecast-driven caching. Leveraging the recursive nature of sequence modeling, we introduce a probabilistic Belady policy for future-aware eviction. To eliminate execution bottleneck, we develop (iv) Graph-compatible offloading runtime. We derive general runtime principles encompassing compute-transparent expert placement and synchronization-free orchestration disciplines for end-to-end graph capture. With 45% expert residency, SeqMoE averages a 96.97% hit rate and 80.22% of full-load performance, advancing the state of the art in MoE offloading.

cs.OS↗
arXiv2026

Beacon: LLM Multi-Agent Driven Hardware Design Space Exploration for Heterogeneous Multi-Chiplet Deep Learning Accelerators

Heterogeneous multi-chiplet accelerators allow chiplets to be configured independently to better match different operator characteristics and improve inference efficiency. However, heterogeneity makes simulator evaluation expensive, limiting the number of iterations affordable for hardware design space exploration (HW-DSE). Mainstream data-driven methods rely mainly on final metrics and a few predefined states, and require many search iterations to implicitly learn the relationships between input parameters and optimization objectives, making them less effective in this setting. In practice, evaluators also generate detailed reports on execution timelines, resource utilization, memory accesses, and communication behavior. Large language models (LLMs) can combine domain knowledge with these reports to explicitly identify bottleneck locations, degradation causes, and parameter adjustment directions, thereby improving each design decision under limited iteration budgets. Based on this observation, we propose Beacon, a report-driven LLM multi-agent framework for heterogeneous multi-chiplet HW-DSE. Beacon employs hierarchical agents for bottleneck localization, root-cause diagnosis, and hardware candidate generation, together with an Analysis Toolbox and RAG memory for closed-loop search. Under the same limited iteration budget, Beacon reduces the composite latency-energy-monetary-cost objective by 25.1\%--93.5\% compared with random search, Bayesian optimization, and reinforcement learning.

cs.AR↗
arXiv2026

UnionSparse: An Index-Efficient Sparsity Framework for Low-Bit Sparse LLM Inference on Edge

Edge LLM inference combines sparsity and low-bit quantization to meet device memory, latency, and power limits. Yet quantization shrinks weight payloads without proportionally reducing sparse metadata, so index traffic and nonzero extraction become critical SpMM bottlenecks. We introduce the Payload-to-Metadata Ratio (PMR) and show that improving PMR raises effective compute intensity in decoding. We present UnionSparse, an index-efficient framework that combines Index-Efficient Bitmap Encoding (IE-BME) with a SpMM kernel using Low-Bit Shared-Memory Parallel Decoding (LSPD). IE-BME amortizes metadata and aligns sparse traversal with fragment assembly, while LSPD improves small-batch execution. Under W4A4 quantization and 30%--70% sparsity, UnionSparse outperforms FlashLLM and SpInfer by 2.30x and 1.43x, and CUTLASS and cuBLAS Tensor Core by 1.56x and 3.46x, respectively. These results establish payload-extraction efficiency as a first-order concern for low-bit sparse inference on edge GPUs. Source code is available at: https://github.com/Victor-Alen/UnionSparse.

cs.DC↗
arXiv2026

Realizable N:M Sparse Transformer Inference via Search-Kernel Co-Design

Vision Transformers (ViTs) achieve strong accuracy but incur high inference latency. Semi-structured N:M sparsity can reduce arithmetic cost, yet its theoretical savings often fail to translate into proportional end-to-end speedups on modern GPUs. This mismatch arises because deployment latency depends not only on arithmetic reduction but also on execution regularity and hardware scheduling under sparsity. Achieving practical acceleration, therefore, requires coordinated design across sparse execution and sparsity configuration. To this end, we propose a hardware-software co-design framework for N:M sparse ViT inference. On the hardware side, we design MD-SpMM, an N:M sparse CUDA kernel that reorganizes sparse GEMM into micro-dense, Tensor-Core-aligned dataflow and uses inference-aware adaptive parallelism to sustain utilization. On the software side, we perform layer-wise sparsity search under explicit end-to-end latency budgets using a three-stage heuristic search with constraint relaxation to avoid premature convergence and enable deployment-aware sparsity allocation. Experiments on multiple ViT/Swin models and GPU platforms show that the framework achieves over 2.2x latency speedup while maintaining comparable accuracy and delivering superior accuracy under the same latency constraint. The source code is publicly available at https://github.com/liuganhuo/realizable-nm-sparse-transformer.

cs.AR↗
arXiv2026

Compass: Co-Exploration of Mapping and Hardware for Heterogeneous Multi-Chiplet Accelerators Targeting LLM Inference Service Workloads

Large language models (LLMs) bring huge computational demands, which makes multi-chiplet accelerators that can integrate large-scale computing resources a powerful solution. However, existing design space exploration (DSE) efforts for such accelerators primarily focus on traditional CNN/Transformer workloads and fall short in supporting the highly dynamic behavior of real-world LLM inference services. This dynamic nature manifests in two key aspects: 1) Mixed request types: the prefill and decode phases exhibit significantly different computational patterns and are frequently interleaved by modern system-level service schedulers; 2) Variable sequence lengths: the sequence length differences across requests can span several orders of magnitude, rendering padding-based assumptions inefficient. Moreover, many prior works assume homogeneous chiplets and overlook the potential beneficial interaction between LLM dynamics and heterogeneous chiplet architectures. To bridge this gap, we introduce Compass, a co-exploration framework designed to optimize mapping strategies and hardware design for multi-chiplet accelerators, specifically tailored for dynamic LLM workloads. First, we propose a computation execution graph-based mapping encoding scheme that decouples micro-batch and layer dimensions, enabling fine-grained execution control on heterogeneous chiplets and flexibly representing various parallelism strategies. Second, based on this scheme, we develop the Compass framework itself, which integrates an evaluation engine, a mapping generation engine based on genetic algorithm, and a hardware sampling engine based on Bayesian optimization, enabling fast and flexible cross-level co-design. Compared with the SOTA DSE works Gemini and MOHaM, Compass reduces latency by 63.92\% and energy by 40.32\% on average in various scenarios, with only a 3.11\% increase in monetary cost.

cs.AR↗
arXiv2026

Hermes: A Unified High-Performance NTT Architecture with Hybrid Dataflow

Fully Homomorphic Encryption (FHE) relies heavily on the Number Theoretic Transform (NTT), making NTT a major performance bottleneck due to its intensive polynomial computations. Hybrid Homomorphic Encryption (HHE), which integrates arithmetic and logic FHE, further requires support for multiple NTT lengths. However, existing accelerators mainly optimize NTT throughput and do not provide unified support for HHE. This paper presents Hermes, a unified high-performance NTT architecture based on hybrid dataflow. Hermes exploits parallelism along both temporal and spatial dimensions and incorporates a fully pipelined on-chip computing core. A conflict-free on-chip fragmentation algorithm is introduced to resolve bank conflicts and enable burst HBM access, while an efficient dataflow improves computational intensity through data reuse, reducing bandwidth demand. Experimental results show that Hermes supports multiple NTT lengths and achieves up to 13.6x and 1.3x higher throughput than state-of-the-art GPU and FPGA accelerators, respectively. Our source code is available at https://anonymous.4open.science/r/Hermes_conf-4E6F.

cs.AR↗
arXiv2026

Reducing End-to-End Latency of Cause-Effect Chains with Shared Cache Analysis

Cause-effect chains, as a widely used modeling method in real-time embedded systems, are extensively applied in various safety-critical domains. End-to-end latency, as a key real-time attribute of cause-effect chains, is crucial in many applications. But the analysis of end-to-end latency for cause-effect chains on multicore platforms with shared caches still presents an unresolved issue. Traditional methods typically assume that the worst-case execution time (WCET) of each task in the cause-effect chain is known. However, in the absence of scheduling information, these methods often assume that all shared cache accesses result in misses, leading to an overestimation of WCET and, consequently, affecting the accuracy of end-to-end latency. However, effectively integrating scheduling information into the WCET analysis process of the chains may introduce two challenges: first, how to leverage the structural characteristics of the chains to optimize shared cache analysis, and second, how to improve analysis accuracy while avoiding state space explosion. To address these issues, this paper proposes a novel end-to-end latency analysis framework designed for multi-chain systems on multicore platforms with shared caches. This framework extracts scheduling information and structural characteristics of cause-effect chains, constructing fine-grained and scalable inter-core memory access contexts at the basic block level for time-sensitive shared cache analysis. This results in more accurate WCET (TSC-WCET) estimates, which are then used to derive the end-to-end latency. Finally, we conduct experiments on dual-core and quad-core systems with various cache configurations, which show that under certain settings, the average maximum end-to-end latency of cause-effect chains is reduced by up to 34% and 26%.

eess.SY↗
arXiv2026

A Timing-Anomaly Free Dynamic Scheduling on Heterogeneous Systems

Heterogeneous systems commonly adopt dynamic scheduling algorithms to improve resource utilization and enhance scheduling flexibility. However, such flexibility may introduce timing anomalies, wherein locally reduced execution times can lead to an increase in the overall system execution time. This phenomenon significantly complicates the analysis of Worst-Case Response Time (WCRT), rendering conventional analysis either overly pessimistic or unsafe, and often necessitating exhaustive state-space exploration to ensure correctness. To address this challenge, this paper presents the first timing-anomaly-free dynamic scheduling algorithm for heterogeneous systems, referred to as Deterministic Dynamic Execution. It achieves a safe and tight WCRT estimate through a single offline simulation execution. The core idea is to apply deterministic execution constraints, which partially restrict the resource allocation and execution order of tasks at runtime. Based on a formally defined execution progress model for heterogeneous system scheduling, we prove the correctness of the proposed execution constraints and their ability to eliminate timing anomalies. Furthermore, we propose two methods to generate execution constraints. The first method derives execution constraints directly from the execution traces produced by existing scheduling algorithms. The second method is a heuristic-based approach that constructs execution constraints, enabling further reduction of the WCRT. Experimental results on synthetically generated DAG task sets under various system configurations demonstrate that, compared to traditional dynamic scheduling algorithms, our approach not only eliminates timing anomalies but also effectively reduces both the WCRT and response time jitter.

eess.SY↗
arXiv2026

Crystal-KV: Efficient KV Cache Management for Chain-of-Thought LLMs via Answer-First Principle

Chain-of-Thought (CoT) reasoning in large language models (LLMs) significantly improves accuracy on complex tasks, yet incurs excessive memory overhead due to the long think-stage sequences stored in the Key-Value (KV) cache. Unlike traditional generation tasks where all tokens are uniformly important, CoT emphasizes the final answer, rendering conventional KV compression strategies ineffective. In this paper, we present Crystal-KV, an efficient KV cache management framework tailored for CoT reasoning. Our key insight is the answer-first principle. By mapping answer preferences into think-stage attention map, we distinguish between SlipKV, which mainly maintains the reasoning flow but may occasionally introduce misleading context, and CrystalKV, which truly contributes to the correctness of the final answer. Next, we propose an attention-based Least Recently Frequently Used algorithm. It precisely identifies when a SlipKV entry's utility expires and evicts it, retaining CrystalKV without disrupting reasoning flow. Finally, we introduce an adaptive cache budget allocation algorithm. Based on the dynamic proportion of CrystalKV, it estimates the importance of each layer/head and adjusts the KV cache budget during inference, amplifying critical components to improve budget utilization. Results show that Crystal-KV achieves state-of-the-art KV cache compression, significantly improves throughput, and enables faster response time, while maintaining, or even improving, answer accuracy for CoT reasoning.

cs.CL↗
arXiv2025

ActionFlow: A Pipelined Action Acceleration for Vision Language Models on Edge

Vision-Language-Action (VLA) models have emerged as a unified paradigm for robotic perception and control, enabling emergent generalization and long-horizon task execution. However, their deployment in dynamic, real-world environments is severely hin dered by high inference latency. While smooth robotic interaction requires control frequencies of 20 to 30 Hz, current VLA models typi cally operate at only 3-5 Hz on edge devices due to the memory bound nature of autoregressive decoding. Existing optimizations often require extensive retraining or compromise model accuracy. To bridge this gap, we introduce ActionFlow, a system-level inference framework tailored for resource-constrained edge plat forms. At the core of ActionFlow is a Cross-Request Pipelin ing strategy, a novel scheduler that redefines VLA inference as a macro-pipeline of micro-requests. The strategy intelligently batches memory-bound Decode phases with compute-bound Prefill phases across continuous time steps to maximize hardware utilization. Furthermore, to support this scheduling, we propose a Cross Request State Packed Forward operator and a Unified KV Ring Buffer, which fuse fragmented memory operations into efficient dense computations. Experimental results demonstrate that ActionFlow achieves a 2.55x improvement in FPS on the OpenVLA-7B model without retraining, enabling real-time dy namic manipulation on edge hardware. Our work is available at https://anonymous.4open.science/r/ActionFlow-1D47.

cs.AI↗
arXiv2025

CoQMoE: Co-Designed Quantization and Computation Orchestration for Mixture-of-Experts Vision Transformer on FPGA

Vision Transformers (ViTs) exhibit superior performance in computer vision tasks but face deployment challenges on resource-constrained devices due to high computational/memory demands. While Mixture-of-Experts Vision Transformers (MoE-ViTs) mitigate this through a scalable architecture with sub-linear computational growth, their hardware implementation on FPGAs remains constrained by resource limitations. This paper proposes a novel accelerator for efficiently implementing quantized MoE models on FPGAs through two key innovations: (1) A dual-stage quantization scheme combining precision-preserving complex quantizers with hardware-friendly simplified quantizers via scale reparameterization, with only 0.28 $\%$ accuracy loss compared to full precision; (2) A resource-aware accelerator architecture featuring latency-optimized streaming attention kernels and reusable linear operators, effectively balancing performance and resource consumption. Experimental results demonstrate that our accelerator achieves nearly 155 frames per second, a 5.35$\times$ improvement in throughput, and over $80\%$ energy reduction compared to state-of-the-art (SOTA) FPGA MoE accelerators, while maintaining $<1\%$ accuracy loss across vision benchmarks. Our implementation is available at https://github.com/DJ000011/CoQMoE.

cs.AR↗
arXiv2025

QiMeng-Xpiler: Transcompiling Tensor Programs for Deep Learning Systems with a Neural-Symbolic Approach

Heterogeneous deep learning systems (DLS) such as GPUs and ASICs have been widely deployed in industrial data centers, which requires to develop multiple low-level tensor programs for different platforms. An attractive solution to relieve the programming burden is to transcompile the legacy code of one platform to others. However, current transcompilation techniques struggle with either tremendous manual efforts or functional incorrectness, rendering "Write Once, Run Anywhere" of tensor programs an open question. We propose a novel transcompiler, i.e., QiMeng-Xpiler, for automatically translating tensor programs across DLS via both large language models (LLMs) and symbolic program synthesis, i.e., neural-symbolic synthesis. The key insight is leveraging the powerful code generation ability of LLM to make costly search-based symbolic synthesis computationally tractable. Concretely, we propose multiple LLM-assisted compilation passes via pre-defined meta-prompts for program transformation. During each program transformation, efficient symbolic program synthesis is employed to repair incorrect code snippets with a limited scale. To attain high performance, we propose a hierarchical auto-tuning approach to systematically explore both the parameters and sequences of transformation passes. Experiments on 4 DLS with distinct programming interfaces, i.e., Intel DL Boost with VNNI, NVIDIA GPU with CUDA, AMD MI with HIP, and Cambricon MLU with BANG, demonstrate that QiMeng-Xpiler correctly translates different tensor programs at the accuracy of 95% on average, and the performance of translated programs achieves up to 2.0x over vendor-provided manually-optimized libraries. As a result, the programming productivity of DLS is improved by up to 96.0x via transcompiling legacy tensor programs.

cs.CL↗
arXiv2025

UbiMoE: A Ubiquitous Mixture-of-Experts Vision Transformer Accelerator With Hybrid Computation Pattern on FPGA

Compared to traditional Vision Transformers (ViT), Mixture-of-Experts Vision Transformers (MoE-ViT) are introduced to scale model size without a proportional increase in computational complexity, making them a new research focus. Given the high performance and reconfigurability, FPGA-based accelerators for MoE-ViT emerge, delivering substantial gains over general-purpose processors. However, existing accelerators often fall short of fully exploring the design space, leading to suboptimal trade-offs between resource utilization and performance. To overcome this problem, we introduce UbiMoE, a novel end-to-end FPGA accelerator tailored for MoE-ViT. Leveraging the unique computational and memory access patterns of MoE-ViTs, we develop a latency-optimized streaming attention kernel and a resource-efficient reusable linear kernel, effectively balancing performance and resource consumption. To further enhance design efficiency, we propose a two-stage heuristic search algorithm that optimally tunes hardware parameters for various FPGA resource constraints. Compared to state-of-the-art (SOTA) FPGA designs, UbiMoE achieves 1.34x and 3.35x throughput improvements for MoE-ViT on Xilinx ZCU102 and Alveo U280 platforms, respectively, while enhancing energy efficiency by 1.75x and 1.54x. Our implementation is available at https://github.com/DJ000011/UbiMoE.

cs.AR↗
arXiv2023

Emergent Communication for Rules Reasoning

Research on emergent communication between deep-learning-based agents has received extensive attention due to its inspiration for linguistics and artificial intelligence. However, previous attempts have hovered around emerging communication under perception-oriented environmental settings, that forces agents to describe low-level perceptual features intra image or symbol contexts. In this work, inspired by the classic human reasoning test (namely Raven's Progressive Matrix), we propose the Reasoning Game, a cognition-oriented environment that encourages agents to reason and communicate high-level rules, rather than perceived low-level contexts. Moreover, we propose 1) an unbiased dataset (namely rule-RAVEN) as a benchmark to avoid overfitting, 2) and a two-stage curriculum agent training method as a baseline for more stable convergence in the Reasoning Game, where contexts and semantics are bilaterally drifting. Experimental results show that, in the Reasoning Game, a semantically stable and compositional language emerges to solve reasoning problems. The emerged language helps agents apply the extracted rules to the generalization of unseen context attributes, and to the transfer between different context attributes or even tasks.

cs.AI↗