Device epilogues

An epilogue is a function in Python, JavaScript, or Rust that operates on tensors. The first time you submit it, Pie traces the function, compiles it into a native kernel for your GPU, and fuses it onto the end of the forward pass. The full logits vector stays in GPU memory; only what the epilogue outputs leaves the device.

TracingCopy Anchor Link

The function you pass to fwd.epilogue does not run on the CPU at each token. As the figure above shows, Pie calls it once with placeholder tensors, records every operation into an eta program, and compiles that program into device code (CUDA, HIP, Metal, SPIR-V, WGSL, or StableHLO). The compiled kernel is cached by content hash, so subsequent passes with the same epilogue reuse the kernel directly.

Because tracing records operations rather than executing them on live values, standard Python if statements run at trace time and choose which program graph gets traced (for example, if temperature == 1.0:). To branch on a tensor value inside the kernel, use select(cond, if_true, if_false).

Writing a samplerCopy Anchor Link

Here is top-a sampling, which keeps tokens whose probability is at least a · p_max² and samples from them:

from inferlet.eta import broadcast, ge, gumbel_max, reduce_max, select, softmax

def top_a(logits, vocab, a, rng):
    probs = softmax(logits)
    p_max = reduce_max(probs)
    keep = ge(probs, broadcast(a * p_max * p_max, [vocab]))
    masked = select(keep, logits, broadcast(float("-inf"), [vocab]))
    return gumbel_max(masked, rng)

The compiler fuses the entire top_a function into a single kernel that runs right after the model's final layer. The complete inferlet is top-a-sampling-py.

OperationsCopy Anchor Link

The inferlet.eta module provides operations for samplers, guidance, speculative verification, and watermarks:

See the Inferlet API reference for the full list.

ChannelsCopy Anchor Link

A Channel is a typed queue in GPU memory. Both the host and device epilogues can put values into a channel and take values from it:

• Host to device: The host puts a grammar mask or random seed into a channel, and the epilogue takes it.
• Device to host: The epilogue puts a sampled token ID into a channel, and the host awaits take_scalar().
• Device to device: The epilogue puts its chosen token directly into the next pass's input channel.

token_in = Channel.from_([first], dtype.i32)               # next pass's input on the GPU
token_out = Channel([1], dtype.i32).capacity(channel_capacity())

fwd = ForwardPass()
fwd.embed(token_in, Channel.from_([0, 1], dtype.u32))
fwd.state([ws], StateGeometry.folding())

@fwd.epilogue
def _():
    token = reshape(reduce_argmax(intrinsics.output(0, "logits")), [1])
    token_in.put(token)    # feeds the next pass on the GPU
    token_out.put(token)   # streams the token ID back to the host

A loop that stays on the deviceCopy Anchor Link

When token_in feeds each chosen token back on the GPU, the runtime can submit the same pass repeatedly without waiting on the host between steps. pipe.run_ahead keeps a window of passes in flight while your callback reads the output tokens:

async def on_token():
    token = await token_out.take_scalar()
    return token not in stop_tokens        # False stops decoding

await pipe.run_ahead(fwd, max_tokens, on_token)

As the figure above shows, run_ahead eliminates host round trips during decoding. Pie's built-in API endpoints use this pattern. When a loop must wait on the host after every token (for example, when the host updates a grammar mask), it cannot use run_ahead and submits one pass at a time instead.