Batch Inference at Scale: How to Maximize GPU Utilization

This post walks through a sample architecture for running 100M+ LLM inference calls with near-maximal GPU utilization on Union using Flyte’s orchestration primitives.

GPUs are brutally expensive to idle. And yet, that’s what most batch inference pipelines do: load a model, process some data, hit a straggler, wait, eventually finish. If you’re running evals at scale, generating synthetic datasets, scoring large corpora or doing multi-round LLM judging, you’ve probably already felt the cost of this the hard way.

The core problem isn’t that inference is slow. It’s that the orchestration around inference is naive. Data gets materialized into memory all at once. Workers finish at different rates and sit idle waiting for each other. Stragglers block entire rounds. Models load and unload repeatedly on every task invocation. All of this is infrastructure overhead eating into what should be GPU compute time.

The architecture in this post solves each of these problems directly. It’s built on Flyte, runs on Union, and gives you two concrete dispatch strategies depending on the shape of your workload. By the end you’ll understand not just how it works, but why the design choices were made, and how to tune it for your own pipeline.

Why naive batch inference falls apart

Before we get into the architecture, let’s be precise about what goes wrong with the obvious approach. The obvious approach looks like this: take your dataset, split it into chunks, spin up workers, run inference, collect results. Four things break when you take this to scale.

Memory pressure from materializing the dataset. If you load 100M samples upfront, you’re allocating gigabytes of data into memory before a single inference call has happened. Even with distributed storage, the overhead is real and entirely avoidable.

GPU waste from uneven workloads. LLM inference isn’t uniform. Some prompts are longer. Some outputs trigger more decoding steps. Some requests hit the model at an awkward point in its batching window. When one worker is stuck on a difficult batch, every other worker either waits or drains its work queue and idles.

The straggler problem. In any multi-round judging or approval pipeline, some samples get rejected and need another pass. If your architecture doesn’t handle stragglers gracefully, a handful of slow samples can hold the entire pipeline hostage, blocking the round from closing while GPUs sit mostly unused.

Cold start overhead. Loading a 70B parameter model onto a GPU takes time. On fast NVMe, you’re looking at 30–60 seconds per load. In a naive architecture, every new pod pays this cost upfront. If you’re spinning up hundreds of pods, that’s a lot of model loading that never touches a single real sample.

The architecture we’re looking at fixes all four. Let’s go through the primitives it’s built on, then the design itself.

The platform: Flyte on Union

Flyte and Union together give you four primitives that do most of the heavy lifting in this architecture. Worth understanding each one clearly, because they solve different problems at different layers.

`JsonlDir` as a streaming primitive. Flyte’s `JsonlDir` type lets you work with JSONL data stored in object storage (S3, GCS) without loading it all into memory. The dataset is split into shards and iterated lazily, one shard at a time. Memory footprint stays constant regardless of whether you’re processing 1M samples or 100M. The I/O and the compute overlap, so your GPU is never waiting on data. More on how this works under the hood in the next section.

Reusable containers. By default, every task execution in Flyte spins up a fresh container and discards it when done. For batch inference that’s a problem. Reusable containers fix this by maintaining a persistent pool via `ReusePolicy`. The model loads once and stays hot. One thing worth getting right upfront: `scaledown_ttl` needs to match your actual round duration. Set it too short and you’re cold-starting containers mid-pipeline; too long and you’re paying for idle GPU time at the tail end of a run.

Dynamic batchers. This operates inside each container, at inference time, and is a separate concern from container lifecycle entirely. When multiple concurrent tasks are submitting samples asynchronously, the dynamic batcher accumulates those individual submissions and groups them into batched model calls before they hit the GPU. A single forward pass over 256 samples is far cheaper than 256 individual forward passes. The container handles scheduling; the batcher handles GPU utilization within that container.

Replica-level parallelism. Tasks run with multiple replicas, which is effectively a pool of identical GPU workers. Combined with the `concurrency` setting on reusable containers, this is what determines how many chunks the system processes in parallel at any given moment. Total throughput is `replicas × concurrency`.

Deep dive: JsonlDir and backpressure

`JsonlDir` is worth spending time on because it’s doing more work than it looks like at first glance.

At the surface level, it’s a directory of JSONL shard files stored in object storage: `part-00000.jsonl`, `part-00001.jsonl` and so on. Tasks read from it by iterating shards in sorted order. But the implementation has two features that make it genuinely fast at scale: one-shard-ahead prefetching and bounded queue backpressure.

One-shard-ahead prefetching

When your task is processing shard `N`, `JsonlDir` is already fetching shard `N+1` in the background. This overlaps I/O with computation while the GPU is crunching through 2,000 samples, the next 2,000 are being streamed from object storage. By the time the current shard finishes, the next one is already waiting.

The implementation uses an async background task that streams records from the incoming shard into a bounded `asyncio.Queue`. The consumer task reads from that queue. The key detail: the queue is bounded.

_DEFAULT_PREFETCH_QUEUE_SIZE = 8192  # ~8 MB at ~1 KB per record

Backpressure via a bounded queue

A bounded queue is how the system avoids overwhelming downstream consumers. If your inference task is slower than the I/O layer (which is often the case because model forward passes are slow), the queue fills up. Once full, the prefetch task blocks on `queue.put()`. It doesn’t fetch more. It waits.

This is classical backpressure: the consumer signals to the producer “slow down, I’m not ready” without any explicit coordination code. The queue size is the knob that controls how far ahead the system reads.

async def _prefetch_shard(shard, on_error, queue):
    # Streams records into the queue. Blocks when the queue is full,
    # which naturally throttles how far ahead we read.
    async for record in shard:
        await queue.put(record)
    await queue.put(_PREFETCH_SENTINEL)

Why does this matter for GPU utilization? Because without backpressure, a fast I/O layer racing ahead of a slow GPU layer results in memory bloat. You’re buffering hundreds of thousands of records you can’t process yet. With backpressure, the system self-regulates. The GPU is always the bottleneck, which is exactly what you want.

Automatic shard rotation on writes

`JsonlDir` also handles the write side. When your inference task writes results, `JsonlDir` manages shard rotation automatically, splitting output into new shard files when a shard hits a configurable size limit (default: 256 MB uncompressed). You don’t manage file handles or shard indices. You just write records.

async with JsonlDir.new_remote("results").writer() as w:
    for result in results:
        await w.write(result)
# Shards rotate automatically at 256 MB; output lands in object storage

Compression is supported too with per-shard zstd via `.jsonl.zst` extension. On a 100M-sample dataset, this is the difference between terabytes and hundreds of gigabytes in storage costs.

Deep dive: DynamicBatcher and TokenBatcher

The batching layer is where most of the GPU efficiency actually lives, and the way it works is less obvious than it looks.

How DynamicBatcher works

The batcher acts as an accumulator between your async inference coroutines and the model. Coroutines submit individual samples via `batcher.submit(record)` which is non-blocking and returns a `Future` immediately, and the batcher groups them into optimally-sized batches before sending to the model.

The batcher accumulates across all concurrent tasks in the container, not just one. With 10-way concurrency, you have 10 tasks all submitting samples simultaneously. The batcher sees submissions from all of them, groups them together, and dispatches in cost-budgeted batches. That’s how you get high GPU utilization even when individual tasks are processing unevenly-sized chunks.

Three signals control when a batch gets dispatched:

1. Cost budget. When the accumulated cost of pending records reaches `target_batch_cost`, the batcher dispatches immediately.
2. Hard cap. `max_batch_size` sets an absolute ceiling on records per batch regardless of cost.
3. Timeout. When a record has been waiting longer than `batch_timeout_s`, the batcher dispatches whatever it has accumulated, even a partial batch. This prevents stalls when arrival rates aren't perfectly steady.

from flyte.extras import DynamicBatcher


async def process(batch: list[dict]) -> list[str]:
    """Your batch processing function. 
    Must return results in the same order as input."""
    return [heavy_computation(item) for item in batch]

async with DynamicBatcher(
    process_fn=process,
    target_batch_cost=1000,   # dispatch when accumulated cost reaches this
    max_batch_size=64,        # hard cap on records per batch
    batch_timeout_s=0.05,     # dispatch partial batches after 50ms
    max_queue_size=5_000,     # queue size for backpressure
) as batcher:
    futures = []
    for record in my_records:
        future = await batcher.submit(record, estimated_cost=10)
        futures.append(future)
    results = await asyncio.gather(*futures)

When the queue is full, `submit()` blocks until space is available, providing natural backpressure that prevents producers from overwhelming the GPU.

Cost estimation

The batcher uses cost estimates to decide how many records to group into each batch. You can provide them in several ways, checked in order:

1. Explicit: pass `estimated_cost` to `submit()`
2. Estimator function: pass `cost_estimator` to the constructor
3. Protocol: implement `estimate_cost()` on your record type
4. Default: falls back to `default_cost` (default: `1`)

TokenBatcher for LLM inference

For LLM workloads, `TokenBatcher` is a convenience subclass that uses token-aware parameter names and checks the `TokenEstimator` protocol.

This matters because LLM inference cost scales with total token count in the batch, not sample count. A batch of 256 short prompts (50 tokens each) and a batch of 64 long prompts (200 tokens each) both fit in the same 12,800-token budget, but naive count-based batching would treat these very differently, either wasting memory padding short sequences to match the longest one in the batch, or getting artificially constrained by that longest sequence. The token batcher handles this by construction: it fills the GPU's token budget regardless of how lengths vary across the batch.

from dataclasses import dataclass
from flyte.extras import TokenBatcher


@dataclass
class Prompt:
    text: str
    def estimate_tokens(self) -> int:
        """Rough token estimate (~4 chars per token)."""
        return len(self.text) // 4 + 1
        
async def inference(batch: list[Prompt]) -> list[str]:
    texts = [p.text for p in batch]
    outputs = model.generate(texts, sampling_params)
    return [o.outputs[0].text for o in outputs]
    
async with TokenBatcher(
    inference_fn=inference,
    target_batch_tokens=32_000,  # token budget per batch
    max_batch_size=256,
    batch_timeout_s=0.05,
    max_queue_size=5_000,
) as batcher:
    future = await batcher.submit(Prompt(text="What is 2+2?"))
    result = await future

Monitoring utilization

The batcher exposes a `stats` property with real-time metrics:

stats = batcher.stats

print(f"Utilization:    {stats.utilization:.1%}")       # fraction of time spent processing
print(f"Records done:   {stats.total_completed}")
print(f"Batches sent:   {stats.total_batches}")
print(f"Avg batch size: {stats.avg_batch_size:.1f}")
print(f"Busy time:      {stats.busy_time_s:.1f}s")
print(f"Idle time:      {stats.idle_time_s:.1f}s")

Target utilization is above 0.9. If it’s lower, the typical levers are: more concurrent producers, a shorter `batch_timeout_s` to dispatch partial batches faster or a larger `max_queue_size` to let more records buffer ahead of the GPU.

The data pipeline: streaming at scale

With those primitives in mind, here’s how the data flows through the system.

The example dataset is 100M samples stored as a `JsonlFile` in object storage. The first task in the workflow chunks it into batches of 2,000 samples each, producing a `JsonlDir`. At 100M samples, that’s 50,000 shards.

jsonl_dir: JsonlDir = chunk_dataset(jsonl_file, chunk_size=2000)
# Result: 50,000 shards of 2,000 samples each, stored in object storage

Flyte’s scheduler dispatches the shards across 32 replicas with 10-way concurrency per replica: 320 concurrent shard processors. At 50,000 tasks / 320 concurrent workers, you get roughly 156 scheduling rounds to process the full dataset.

The critical property: at any point in time, only the shards actively being processed by the 320 workers are in memory. Everything else sits in object storage. Memory footprint is bounded by `active_workers × shard_size`, not by dataset size.

The two dispatch strategies

This is the core of the architecture. The system offers two fundamentally different approaches to dispatching work, and choosing correctly between them is the difference between efficient utilization and expensive stragglers.

Strategy 1: re-dispatch remaining samples each round

Use this when approval rates are high and most samples complete within 1–3 rounds.

In this strategy, each round dispatches all pending samples together. After the round finishes, the workflow runs the judge, identifies which samples didn’t pass, and re-dispatches the remaining set as input to the next round.

Round 1: Infer 100M samples -> Judge -> Approved: 90M, Pending: 10M
Round 2: Re-dispatch 10M  -> Judge -> Approved: 9M,  Pending: 1M
Round 3: Re-dispatch 1M   -> Judge -> Done

The re-dispatch loop in the workflow:

@env.task
async def batch_inference_workflow(data: JsonlDir) -> JsonlDir:
    output = JsonlDir.new_remote("results")
    sem = asyncio.Semaphore(320)  # 32 replicas × 10 concurrency
    pending = data

    for round_num in range(MAX_ROUNDS):
        # Phase 1: inference
        inferred = JsonlDir.new_remote(f"round_{round_num}/inferred")

        async def dispatch_infer(batch):
            async with sem:
                return await infer_batch(batch)

        tasks = set()
        async with inferred.writer() as w:
            async for batch in pending.iter_batches(2000):
                task = asyncio.create_task(dispatch_infer(batch))
                tasks.add(task)
                task.add_done_callback(tasks.discard)
                
                for t in [t for t in tasks if t.done()]:
                    await w.write_many(t.result())
                    
            for coro in asyncio.as_completed(tasks):
                await w.write_many(await coro)

        # Phase 2: judge — split into approved and rejected
        approved = JsonlDir.new_remote(f"round_{round_num}/approved")
        rejected = JsonlDir.new_remote(f"round_{round_num}/rejected")

        async def dispatch_judge(batch):
            async with sem:
                return await judge_batch(batch)

        tasks = set()
        async with approved.writer() as wa, rejected.writer() as wr:
            async for batch in inferred.iter_batches(2000):
                task = asyncio.create_task(dispatch_judge(batch))
                tasks.add(task)
                task.add_done_callback(tasks.discard)
                
                for t in [t for t in tasks if t.done()]:
                    for item in t.result():
                        await (wa if item["approved"] else wr).write(item)
                        
            for coro in asyncio.as_completed(tasks):
                for item in await coro:
                    await (wa if item["approved"] else wr).write(item)

        # Accumulate approved, feed rejected into next round
        async with output.writer() as w:
            async for record in approved.iter_records():
                await w.write(record)

        pending = rejected
        has_rejected = False
        
        async for _ in rejected.iter_records():
            has_rejected = True
            break
            
        if not has_rejected:
            break

    return output

The `batch_inference_workflow` driver fans out across 32 replicas with 10-way concurrency, resulting in 320 chunks in flight simultaneously. Inside each `infer_batch` call, sample submissions flow into the batcher, which accumulates across all concurrent tasks in the container and dispatches to the model. The `judge_batch` phase works identically: the same fan-out, the same batcher, splitting results into `approved` and `rejected` `JsonlDir` objects. The rejected set becomes `pending` for the next round.

Why it works: On high-approval workloads, the pending set shrinks fast. Round 1 clears 90% of samples. Round 2 clears 90% of what’s left. By round 3 you’re dealing with a negligible tail. GPU utilization stays high because every round is a full fan-out, and the stragglers never accumulate enough to stall progress.

The tradeoff: If approval rates are low, this strategy amplifies costs. A pipeline with 50% approval per round leaves 50M samples pending after round 1, 25M after round 2. You’re doing significantly more total inference work, and each round carries the full overhead of scheduling, dispatching, and writing intermediate results. That’s when you reach for Strategy 2.

Strategy 2: in-place retry with per-chunk blocking

Use this when approval rates are low, samples frequently need multiple passes, and you want GPU utilization to stay high regardless of how many samples are churning.

Both strategies use the same reusable container setup with 32 replicas, 10-way concurrency, where the model is loaded once and kept hot. What changes is where retry logic lives. In Strategy 1, a rejected sample exits the current round and gets re-queued into the next. The pod is done with it. In Strategy 2, the pod owns the sample until it’s approved or exhausted. Retry happens inside `process_one`, against the same warm model, without ever touching the scheduler.

@gpu_env.task
async def infer_batch(samples: list[dict]) -> list[dict]:
    batcher = await get_batcher()

    async def process_one(sample):
        try:
            state = await load_checkpoint(sample["id"])
            result = state.get("last_result")

            for attempt in range(state.get("attempt", 0), MAX_ATTEMPTS):
                if result is None:
                    future = await batcher.submit(build_infer_prompt(sample))
                    result = await future
                    await save_checkpoint(
                        sample["id"], attempt=attempt, last_result=result
                    )

                future = await batcher.submit(build_judge_prompt(sample, result))
                judgment = await future

                if judgment["approved"]:
                    await mark_complete(sample["id"])
                    return {**sample, "result": result,
                            "status": "approved", "attempts": attempt + 1}

                result = None
                await save_checkpoint(sample["id"], attempt=attempt + 1)

            return {**sample, "result": result,
                    "status": "max_attempts", "attempts": MAX_ATTEMPTS}
        except Exception as e:
            return {**sample, "result": None, "status": "error", "error": str(e)}

    return list(await asyncio.gather(*(process_one(s) for s in samples)))

A few things worth paying attention to in this code.

`asyncio.gather` over all samples means inference and judging for every sample in the chunk run concurrently. The batcher accumulates submissions from all of them and dispatches to the GPU in groups. No sample waits for another to finish before it gets scheduled.

The checkpoint pattern is doing more work than it looks like. Each attempt writes its state before moving on. If the container goes down mid-attempt due to spot preemption or OOM, the next container picks up from the last saved state rather than starting over.

Compare that to Strategy 1, where you don't need explicit checkpointing at all. The round structure gives you fault tolerance for free. Each round writes `approved` and `rejected` to object storage as it goes. If a container dies, those samples just weren't written to `approved` yet, so they’re still in `pending` and get picked up next round. The `JsonlDir` write pattern is your checkpoint, so nothing extra is needed.

Strategy 2 has no such round boundary to fall back on. Retry lives inside the task, which means a container failure mid-attempt loses all progress on that sample’s retry loop. That’s why `save_checkpoint` is the mechanism that makes in-place retry safe to run on interruptible infrastructure. An eviction at attempt 4 of 5 picks up at attempt 4, not attempt 0.

Then there’s the straggler behavior. Say 1,999 of 2,000 samples in a chunk approve on the first attempt. The one difficult sample keeps retrying inside `process_one`. The other 9 concurrency slots in the same container are processing entirely different chunks. The GPU never idles waiting for one sample, hence the retry becomes background noise rather than a bottleneck.

Why this works for low-approval workloads: In Strategy 1, a 50% approval rate means round 2 is scheduling 50M samples with full fan-out, scheduling overhead and intermediate writes. Strategy 2 handles those retries where they happen, inside the container, against a model that’s already warm. The scheduler never sees them. What you give up is simplicity. The code is more involved, the checkpoint infrastructure needs to exist and you need confidence that your retry loop eventually converges.

One more thing: always put a hard cap on `max_rounds`. If you leave retries open-ended and your judge has a flaw, the pipeline can keep running indefinitely without it being immediately obvious. You usually only notice once the costs start adding up. And when the pipeline does hit the retry cap, that should be treated as something worth investigating, not as a normal completion path.

The capacity math stays the same across both strategies:

Replicas:             32
Concurrency:          10     (tasks per container)
Chunk size:           2,000  (samples per task)
Max batch size:       256    (samples per model call)

Chunks in flight:     32 × 10       = 320
Samples in flight:    320 × 2,000   = 640,000
Model calls:          640,000 / 256 ≈ 2,500 batched forward passes
Scheduling rounds:    50,000 / 320  ≈ 156

The numbers are identical across both strategies. What changes is entirely where retry work happens: at the scheduler level or inside the container. That single decision shapes your fault tolerance model, your infrastructure requirements, and how gracefully the pipeline degrades when approval rates are low.

Choosing between the two strategies

The decision comes down to one number: your approval rate per round.

If most samples pass on the first or second attempt, above roughly 70% approval per round, Strategy 1 is the right call. The pending set shrinks fast, the round structure keeps things simple, and you get fault tolerance for free from the `JsonlDir` write pattern. The code is straightforward and the failure modes are easy to reason about.

If approval rates are low, or your judging criteria are strict enough that samples regularly need three, four, or five attempts, Strategy 1 starts to break down. Every retry goes back through the scheduler, so the overhead keeps piling up and the tail latency only gets worse. That’s where Strategy 2 starts making a lot more sense. Retries happen inside the container itself against a model that’s already warm, without constantly bouncing back through the scheduler.

That fault tolerance story is also what makes spot instances viable here. Strategy 1 recovers via the round structure; Strategy 2 recovers via checkpointing. Either way, a preempted container picks up where it left off rather than starting over. With `interruptible=True`, Flyte handles the retry automatically and at 60–70% cost reduction on GPU instances, it’s hard to justify not using it.

The judge in both strategies doesn’t have to be an LLM either. Union's human-in-the-loop support lets you insert a human review gate at the round boundary. The workflow pauses after each inference phase, surfaces the pending samples for a human reviewer, and resumes once approvals come back. The `JsonlDir` round structure maps cleanly onto this: each round’s `rejected` set is exactly the queue of samples waiting for human judgment. You get the same fan-out and streaming behavior, with a human in the approval loop instead of, or alongside, an automated judge. This is useful when judging criteria are genuinely subjective or when you want human oversight on low-confidence outputs before they get re-queued.

There’s also a hybrid worth considering. Run Strategy 1 for the first round or two to clear the easy approvals cheaply, then switch to Strategy 2 for the stubborn tail. Both strategies read from and write to `JsonlDir`, so the output of one is a valid input to the other:

@env.task
async def hybrid_batch_inference(data: JsonlDir) -> JsonlDir:
    output = JsonlDir.new_remote("results")
    sem = asyncio.Semaphore(320)
    pending = data

    # Strategy 1: clear easy approvals across a few rounds
    for round_num in range(STRATEGY_1_ROUNDS):
        inferred = JsonlDir.new_remote(f"round_{round_num}/inferred")
        ...  # same fan-out pattern as Strategy 1

        pending = rejected  # feed rejected into next round
        if not has_rejected:
            return output

    # Strategy 2: in-place retry for the remaining tail
    async def dispatch(batch):
        async with sem:
            return await infer_batch(batch)  # the version with process_one + checkpointing

    async with output.writer() as w:
        async for batch in pending.iter_batches(2000):
            task = asyncio.create_task(dispatch(batch))
            ...

    return output

The crossover point is mostly a judgment call. A reasonable rule of thumb is to switch to Strategy 2 once the pending set stops shrinking meaningfully between rounds. At that point, the remaining samples are unlikely to clear cheaply no matter how many times you keep re-dispatching them.

Closing thoughts

Not every batch inference workload needs this level of complexity. If you’re processing a few million samples with a strong judge and relatively low retry rates, a simple fan-out approach is usually enough.

But once datasets get large, approval rates become unpredictable, and GPU costs start mattering, these decisions begin to compound quickly. Streaming instead of materializing keeps memory usage under control. Backpressure ensures the GPU stays the bottleneck instead of the scheduler. Reusable containers help avoid repeated cold starts. And the two dispatch strategies let you adapt to actual retry behavior instead of forcing one approach onto every workload.

The nice part is that most of these decisions are encapsulated in the primitives themselves. The pipeline code ends up expressing intent rather than infrastructure concerns, which is usually what good orchestration should feel like. You can