How Pie thinks

Most inference engines take a prompt and return text. Pie gives you the generation loop instead.

You do not need these internals to serve an API, but four concepts explain how Pie works: inferlets, forward passes, working sets, and epilogues.

1. The loop belongs to your programCopy Anchor Link

Every LLM engine runs the same basic loop: run a forward pass, pick a token from the output logits, append it, and repeat. Most engines seal that loop inside the server and only let you tweak it through request flags.

In Pie, the loop is a user program. The engine only executes the forward passes your program submits.

As the figure above shows, speculative decoding, beam search, constrained decoding, and agent loops are simply different loops over forward passes. In Pie, each one is an ordinary program rather than a hardcoded engine feature.

2. InferletsCopy Anchor Link

An inferlet is the program that handles one request from start to finish. It reads the request input, runs the model as many times as needed, and returns or streams the output.

• You write an inferlet in Python, JavaScript, or Rust. Pie compiles it to WebAssembly and runs it in a sandbox next to the model.
• Its entry point is async def main(input: dict) -> dict, which takes JSON and returns JSON.
• While running, it can stream partial text, call external tools over HTTP, and stay alive across conversation turns.

The OpenAI, Anthropic, and Gemini endpoints exposed by pie serve are built-in inferlets that ship with Pie.

3. Forward passesCopy Anchor Link

A ForwardPass runs the model once. Your inferlet builds a pass by specifying which tokens go in, which sequence state they update, and how to pick the output token, then submits the pass to a Pipeline:

fwd = ForwardPass()
fwd.embed(new_tokens)                       # input tokens for this pass
fwd.state([ws], StateGeometry.folding())    # working set to extend and freeze
token = fwd.epilogue(lambda: reduce_argmax(intrinsics.output(0, "logits")))
pipe.submit(fwd)

pipe.submit(fwd) does not block. Passes queued on a pipeline execute in order, so an inferlet can submit future passes before earlier results come back.

4. Working setsCopy Anchor Link

A WorkingSet holds one sequence's state inside the model: the KV cache for attention layers and the recurrent state for linear-attention or SSM layers. Your code uses the same handle regardless of the model architecture.

As shown below, a working set has two parts:

• Frozen prefix: Committed tokens that never change again.
• Writable buffer: Uncommitted tokens, such as speculative drafts that you might roll back.

Calling ws.fork(pipe) creates a new working set that shares the frozen prefix without copying it. Branching one long prompt into ten completions computes the prompt once and pays only for the ten short suffixes. The runtime also deduplicates frozen state across requests: if another request has already computed your prefix, your working set reuses those pages instead of recomputing them.

5. EpiloguesCopy Anchor Link

A forward pass produces a logits vector with one score per vocabulary token. Copying all logits from the GPU back to the WebAssembly host on every token would be slow. Instead, Pie sends your sampling code to the logits.

An epilogue is a small tensor function attached to a pass, such as argmax, temperature sampling, or a grammar mask. The first time you submit an epilogue, Pie traces its operations, compiles them into a native GPU kernel (CUDA, ROCm, Metal, Vulkan, WebGPU, or XLA), and runs that kernel right after the forward pass. Only its result (usually a single token ID) comes back to the host.

# Scale logits by temperature 0.7 and sample; rng is a Channel holding the random seed
fwd.epilogue(lambda: gumbel_max(intrinsics.output(0, "logits") / 0.7, rng.take()))

An epilogue can also write its chosen token directly into the next pass's input on the GPU, so a decode loop can run for many steps without waiting on the host.

Batching is still the engine's jobCopy Anchor Link

You write an inferlet for a single request. You do not manage batching in your code.

As the figure above shows, the runtime collects pending forward passes from all active inferlets and packs them into shared GPU kernel launches. You write single-request code and still get continuous batching across users.

Complete exampleCopy Anchor Link

Here is a complete inferlet that encodes a prompt and greedily generates 8 tokens, one forward pass per token:

from inferlet import chat, model
from inferlet.eta import ForwardPass, Pipeline, StateGeometry, WorkingSet, intrinsics, reduce_argmax


async def main(input: dict) -> dict:
    tokens = chat.prefix() + model.encode(input.get("prompt", "The capital of France is"))
    ws = WorkingSet()
    pipe = Pipeline()

    done = 0
    for _ in range(8):
        fwd = ForwardPass()
        fwd.embed(tokens[done:])                     # feed only the new tokens
        fwd.state([ws], StateGeometry.folding())     # append to state and freeze
        out = fwd.epilogue(lambda: reduce_argmax(intrinsics.output(0, "logits")))
        pipe.submit(fwd)
        done = len(tokens)
        tokens.append(await out)
    pipe.close()

    return {"text": model.decode(tokens[-8:])}

Every other page in the docs builds on these four pieces: