cs.CLOct 22, 2026

Compiling Autoregressive LLM Control into Policy-Exact Parallel Inference for Hierarchical RAG

Authors: Zehao Lin, Qianfeng Shen, Xixuan Hao, Shaobo Cui, Chunyu Li, Kai Chen, Bo Tang, Zhenjun Wang, 

Organizations: MemTensor (Shanghai) Technology Co., Ltd. · China Haisum Engineering Co., Ltd. · The Hong Kong University of Science and Technology (Guangzhou) · Shanghai Jiao Tong University

Abstract

Large language models increasingly act as controllers in structured decision pipelines, yet their decisions are still executed by autoregressive text decoding. Hierarchical RAG is one such setting: at each frontier of a document tree, an LLM commits to bb of kk sections under an input-dependent action grammar, executed as either failure-prone free-form JSON or reliable but strictly serial grammar-constrained generation. We show that this serialization is unnecessary for bounded-horizon control: the reachable future forms a shallow action-prefix tree, which we compile --- one root prefill plus one tree-causal probe evaluate every prefix-conditioned action distribution with serial decoding's tokens, positions, and causal ancestry. A prefix-state lemma makes compiled control policy-exact in exact arithmetic: the same policy with a two-forward critical path and zero output tokens. One further probe pass certifies the exact order-marginalized MAP set. Across four document-disjoint splits and two scales on the primary MetaX device, compiled control matches the serial policy on 8484--89%89\% of questions --- every instrumented divergence a floating-point near-tie --- and cuts control latency by 3535--47%47\% versus token-wise serial decoding and 66--7%7\% versus an idealized chunked engine; an H800 repetition preserves the ordering.

Figures & tables

Explore similar work

CardsList
  1. Latent Memory Palace: Reasoning for Control as Autoregressive Variational Inference

    Jul 9, 2026Chuning Zhu, Eva Xu, Jose Barreiros +3RL ControlLatent Action Learning

  2. Multi-Head Latent Control: A Unified Interface for LLM Agent Decision Making

    Jul 15, 2026Amirhosein Ghasemabadi, Ruichen Chen, Bahador Rashidi +1Multi-Agent LLM SystemsLanguage Model-Based Control