Speculative Programmatic Tool Calling

Speculative programmatic tool calling is a class of techniques for overlapping tool call computation with the code being generated by a harness.

It is no surprise from my previous writing and work on Recursive Language Models (RLMs) that I believe (1) code in a REPL is the only “tool” a system needs; (2) all other tools should be functions in that code tool. When code becomes the primary action space of your system, you need to start thinking about overlapping these specialized tool calls with the code being generated and executed.

Inspired by speculative execution in CPUs and speculative decoding in LLMs, a relatively simple and useful trick I’m proposing in harness design is speculative programmatic tool calling (sPTC), in which we speculate and pre-launch tool calls from partially generated REPL calls as the harness is still generating tokens, rather than waiting for it to finish the entire generation. In particular, this is mostly useful when the tool of interest is a sub-LLM or sub-agent call, such as in the RLM. If the fully generated REPL actually ends up calling these tools, they immediately return with their cached outputs from the speculated call.

A key observation is that LLM tools, such as sub-agents or search APIs, are often high-latency and are usually the bottleneck in harnesses that rely on code execution as actions. Furthermore, the actual generation of the main context is often also a significant latency bottleneck that blocks intermediate calls from happening.

For the rest of this blog, I’ll mostly be discussing inference time savings with respect to an RLM, but note that this generally applies to harnesses that generate code (i.e. CodeAct-like or “code mode” harnesses). There are two obvious areas where you save time:

1. Overlapping already-generated tool calls during token streaming. Most harness designs will wait for the entire model generation to complete its generation before executing the tools. This design is likely a consequence of JSON-style tool-calling, where this was not really a bottleneck. Because the generation of the main context per turn is often slow, this represents a significant portion of time that can be cut, especially for models that think for a long time.
2. Acting as a JIT compiler over REPL calls. Even without streaming enabled, an obvious optimization is that many REPL programs contain blocking tool calls that aren’t actually blocking — e.g. two independent sub-agent calls that the code does not write as asynchronous, can still be run in parallel. sPTC acts as a really naive JIT compiler to prevent such cases, and likely can be improved further across languages and REPL designs.

On locally running LLMs running one or a few chat instances at a time, your inference engine is often highly memory-bound from decoding just the main context, and speculation can help increase the arithmetic intensity. For high-volume serving systems (e.g. if you’re using a frontier lab or model router API), batched requests are abstracted away to various, likely disjoint serving engines, so the gain is purely from either overlapping computation with the main context, or overlapping execution time with a slow REPL call.

Some Comparison Numbers on RLMs

I’ll start by highlighting some basic runtime numbers, although you can probably reason through the runtime benefit. The actual implementation details are quite simple and easy to follow in the codebase, which I will talk about in the latter half of the blog.

You should keep in mind that it’s very difficult to estimate the exact speed-ups because it’s highly dependent on the latency of the tools, the number of tokens generated, the load of your serving engine, and the actual choices the harness makes.

The speed-ups for the RLM are generally on the order of 1-1.2x. We also tested a more “deterministic” suite of LM programs where we’d observe hypothetical or realistic runtime speed-ups in the codebase, but we omit them here because they are too specific.

Designing Speculative PTC Methods

The code is quite short and readable, but I’ll talk a bit about the design philosophy below.

The high-level design is to have imported tool calls in our REPL have a hook that can be called earlier and replaced with a cached output when needed. We want a library and frontend contract that allows us to define (1) what tools should be speculated and what shouldn’t (e.g. maybe we want sub-LLM calls to be speculated, but a sub-RLM call is too costly); and (2) a mechanism for speculating on tool calls whose inputs rely on variables in memory that were computed previously, rather than just literals.

@spec.tool(speculatable=True, pure=True)
def tool(...) -> OutputType:
    ...

# Speculated version
def tool_spec(...):
    promise = launch(tool(...))
    register_speculation(promise)

# Hooked version run in real REPL
def tool_real(...):
    if exists(promise, ID(...)):
        return promise
    else:
        return tool(...)

The contract above lets us define a “shadowed” namespace which invokes these modified tool calls while the LLM output is being parsed, and save them as futures in some store that can be invoked by the real tool. This gives us the relatively simple logic below:

real_ns = {**locals, real_tools}
shadow_ns = replace_tools(real_ns)

# speculate while LLM is streaming
while not LLM.done:
    code += LLM.next_tokens()
    parse_and_peek(code, shadow_ns)  # queue without code execution
    parse_and_speculate(code, shadow_ns)  # re-run shadow REPL

# real tools now route to promised tools
exec(code, real_ns)

Shadowed execution for speculation. The easy cases for speculation are when you can parse a tool call and infer the inputs from the tokens directly, i.e. when the inputs are literals. The nastier cases come when tool calls are embedded inside conditional or looping logic, or the inputs are dependent on some variable in memory that was computed earlier in the REPL call.

What you can speculate and “run” ahead-of-time.

There’s a bit of consideration into what exactly can be speculated, how aggressive you want to be, and the overall safety of running partial code.

Case 1: Literals. String, integer, or other literals can immediately be parsed and converted into a tool call even without shadowed execution of each line.

title = llm_query("Give a title for: The Odyssey")  # parses
blurb = llm_query("One-line blurb for: The Odyssey")  # parses
print(title, blurb)

Case 2: Input dependencies. When input dependencies are involved, as long as all inputs are safe (i.e. pure functions, no side effects), they can be speculated. Furthermore, tools that rely on dependencies that are speculated will wait on the dependencies to be computed first, and are then executed even if the LLM is still streaming.

a = llm_query("Triage: " + doc)  # parses then executes
c = llm_query("Summarize: " + str(a))  # parses, waits on a, then executes
print(c)

if len(doc) > 10_000:  # will evaluate and speculate if safe
    extra = llm_query("Also outline it: " + doc)

Case 3: Peekable and non-peekable dependencies. During streaming at any step, we have a working namespace of the shadow REPL. When a complete tool is parsed, there are cases when we can speculate the inputs, even if the dependencies are variables in memory and not literals.

def gist(t):
    return llm_query("One-line gist: " + t)  # non-peekable

parts = [gist(c) for c in chunks]

side = llm_query("Give me a random title for:", chunks[0])  # peekable
print(side, parts)

Case 4: Blocked speculation calls. While trying to speculate, there is an allowlist of keywords and function calls that can be used to compute dependencies for inputs. We also can specify new tools that we do want to specify are not pure functions. Any speculatable tools that have dependencies that are blocked will not be speculated.

a = llm_query("Triage: " + doc)  # speculated
notes = open("/tmp/scratch.txt").read()  # blocked
b = llm_query("Annotate with notes: " + notes)  # blocked
c = llm_query("Summarize: " + a)  # speculated after a

Related Ideas

Speculative algorithms at the level of LLM decoding (i.e. the streamed outputs) have been explored a little bit for the older tool-calling designs, although not extensively. The likely case is that for standard tool-calling, these techniques were not that useful or the overhead was not worth the marginal latency improvements.

Ending Note

Speculative programmatic tool calling is a natural trick that arises from programmatic tool calling itself having more potential for overlap than traditional tool calling itself. I want to clarify that beyond overlapping with streamed token generations of the root LLM in a harness, the real value in the long run will come from more clever JIT compilation tricks to overlap tool calls in the PTC setting with the actual REPL execution itself, which may become more expensive as harnesses generate more complex programs.