Technology · Discipline

Token-aware GPU batch scheduling

The token is the unit of GPU scheduling in my broker: shape-pure batches packed to each card's declared limits, across two different GPUs.

Rust broker, four cratesInference Fabric and FAFO CoreTwo NVIDIA Blackwell GPUs
One shape signature · its lanes, sorted by length
the lane it joins
  • 0.421%padding waste across both cards
  • 99.9%key-value cache hit rate on both cards, zero shape switches
  • 2,280lanes in use at 22 percent lane pressure
  • ≈17,800lane reuses a second, one new lane every two seconds

The token as the unit of scheduling

Inference Fabric started from a backlog of billions of tokens and two GPUs of different sizes. A queue that hands each card whatever arrived next mixes short and long prompts, fills the gaps with padding, thrashes the engine's key-value cache and sends the smaller card work it cannot hold. I made the token the broker's unit of scheduling and split the problem into two separate axes: a signature for cacheable structure and lanes for length.

Whatever arrived next
every row padded to the longest
One shape, one lane
rows of nearly one length
tokenspaddingillustrative

Shape: a signature, not a parse

Shape is declared by the producer, which knows its own prompt template. Each cacheable component of a prompt is hashed and appended in order, and the list is reduced to one 64-bit signature. The broker never parses a prompt and never sees the model.

Four prompts from one template
system prompt, examples, tool schemathe request

Length: rolling lanes

Length is handled inside each shape by what I call rolling lanes. A shape owns up to 2,048 lanes kept sorted by a representative length; I raised the budget from 128 after busy shapes ran out of lanes and padding climbed steadily. A new item binary-searches for its two neighbours and takes the one that pads it less.

It opens a new lane instead when the padding is above a threshold that rises as the shape's lane budget fills, with a floor so short prompts do not spawn lanes over trivial waste and a hard cap so a long prompt never absorbs large padding because its percentage looks small. Lanes only grow their length, are never merged, and go quiet after a few seconds without traffic so their ids can be reused.

When a new lane is worth opening
10%1%
lane budget emptymeasured: threshold near 3%full
no new lane under 16 tokens of paddinga new lane above 128 tokens, whatever the percentage

The rule I wrote at the top of that code: spend structure before spending waste.

A four-number contract with each card

Each GPU describes itself to the broker in four numbers. A batch pads to its longest item rounded up to 8 tokens and fills in padded tokens. Anything that does not fit, or belongs to another model, goes back to the front of the queue in order for the next consumer. I took input-length filtering out of the broker: it stranded work no consumer would take, and the runtime handles an oversized item more cheaply.

What each card tells the broker
Larger card
Smaller card
each batch filled to the card's own maximum tokens per step; model, maximum batch and input length complete the contract
Padding, share of every batch
Power-of-two buckets≈ a third Lanes, aligned to 8

On disk every segment holds a single model, shape and lane, so purity is a property of the file. Dispatch stays on one shape until it has nothing left, then moves to the shape with the most waiting tokens, and the broker stops itself before it would emit a batch that mixes shapes. The dispatch hot path is one owner on dedicated threads with no async machinery.

The broker is about 24,000 lines of Rust in four crates, and its figures come from serving a 7-billion-parameter dense model in FP8 from one backlog to two NVIDIA Blackwell GPUs, each on its own engine build. The inference engines are NVIDIA's; the scheduling layer, the consumer contract and the broker are mine. The design is in A message broker that understands tokens; its effect on throughput, up to 163,000 tokens per second across those two cards, is in From 44K to 163K tokens per second; a failure in its accounting is in Millions of queued items that were not there.

The producer side

FAFO Core and Function Forge had the same problem from the producer side. Its job conductor has a batching mode keyed to a GPU token budget next to the usual row-count mode, and a decoupled mode in which a producer submits and exits while a separate completion path writes the results back.

Row count
GPU token budget
Submit and exit
fixed rows per batchbatches sized by tokensthe producer leaves; a separate completion path writes back