AgentStage: Exploiting LLM Thinking for Data Staging in Scientific Agents

Abstract

LLM agents on scientific data tasks spend most of their session wall time inside tool calls reading data from storage. Yet before each tool call, the agent reasons in a thinking phase, during which the storage layer sits idle and the next tool's files begin to surface in the reasoning. Exploiting this is hard: the agent rarely names complete paths, so those files must be inferred and then staged before the tool fires. We present AgentStage, which reads the streaming reasoning and the agent's filesystem probes to identify those files, then stages them into fast local storage with in-flight decompression so the tool's reads land hot. We evaluate AgentStage on three curated I/O-bearing scientific tasks and one data-intensive task from each of three community benchmarks (MLE-bench, KramaBench, DSBench), across four reasoning models (three closed-weight, one open-weight). Under a paired trajectory-replay protocol that fixes each agent's tool sequence and varies only cold versus staged execution, AgentStage delivers an average 1.74× end-to-end speedup on the curated tasks, up to 2.34× per run. On the three community benchmarks, AgentStage delivers up to 1.36× speedup. The results show that thinking-phase staging helps when I/O share is high and prefetch completes within the thinking window, while compute-bound or many-file workloads see smaller gains. We show both conditions can be estimated before any staging run, so AgentStage's benefit is predictable in advance. AgentStage adds negligible LLM-side overhead and needs no per-model or per-task tuning.

Tags