Skip to content

Price a batch from Go

Point the kernel at its inputs, open a kernel for a deployment, and price a batch. The later sections build variants of a deployment in memory, read its memory and its evidence, and say what to do when a deployment is refused or a step must be priced often.

go get github.com/inference-sim/blis-latency-kernel@latest

Point it at its inputs

A scenario names its model, its chip, its fabric and its coefficient sets, but not where they live. Repos says where: a directory of scenario files, a blis-catalog checkout, and a blis-registry checkout.

// The roots a scenario's names resolve against: its model, chip, fabric and coefficient sets.
var repos = latencykernel.Repos{
    Scenarios: "testdata",
    Catalog:   "testdata/catalog",
    Registry:  "testdata/registry",
}

These are the paths in this repository after scripts/fetch-testdata.sh has fetched the catalog and registry at the commits the kernel is tested against. Elsewhere, point Catalog and Registry at checkouts of the two repositories, at the versions on Sources of truth. Scenarios and their inputs describes the scenario files.

Open a kernel and price a step

Open reads a scenario file and everything it names, and builds a kernel for the deployment's first pool. OpenPool does the same for any pool of a disaggregated deployment, which states a prefill pool and a decode pool, each with its own kernel.

A batch lists, for each request, the tokens it schedules this step (Scheduled), the tokens it has already computed (Computed) and its prompt length (PromptLen). A steady-state decode schedules one token against its whole context:

// DeepSeek-V3, served in fp8 on eight H200s at tensor-parallel width 8.
const scenario = "aisimulate/deepseek-v3-h200-fp8-sglang-tp8.yaml"
k, err := latencykernel.Open(scenario, repos)
if err != nil {
    panic(err)
}

// 32 requests, each decoding one token against a context of 4,096 tokens.
batch := kernel.Batch{DecodeThreshold: 8}
for range 32 {
    batch.Reqs = append(batch.Reqs, kernel.ReqShape{
        Scheduled: 1, Computed: 4095, PromptLen: 4096,
    })
}

e := k.StepTime(batch)
fmt.Printf("between %.1f and %.1f ms, bottleneck %s\n",
    ms(e.Overlap), ms(e.NoOverlap), e.Bottleneck)
for _, r := range []kernel.Resource{kernel.ResourceSM, kernel.ResourceHBM,
    kernel.ResourceNVLink, kernel.ResourceHost} {
    fmt.Printf("%-7s %5.1f ms\n", r, ms(e.PerResource[r]))
}

between 24.1 and 25.7 ms, bottleneck hbm
sm        4.5 ms
hbm      16.5 ms
nvlink    1.2 ms
host      3.5 ms

DecodeThreshold is the engine's: a request scheduling at most that many tokens is priced as a decode. In vLLM v0.31.0 the attention backend sets it: 1 by default, 128 on FlashMLA and 512 on FlashAttention-MLA. Under speculative decoding, a backend that runs speculative tokens as decodes raises it to one more than their number, and with decode-context parallelism most backends lower it to 1. A simulator should pass its engine's value. The examples here use 8, which every one of those values classifies the same way for a one-token decode.

Overlap assumes the resources within a layer work at once and NoOverlap that they take turns; when one number is wanted, take NoOverlap. Pricing a step explains both, and the five resources.

Build a variant without writing a file

OpenInputs loads the documents without building a kernel. Change them, and call New. A configuration search does this to try many deployments without writing any of them to disk.

in, err := latencykernel.OpenInputs("aisimulate/deepseek-v3-h200-fp8-sglang-tp8.yaml",
    repos, 0)
if err != nil {
    panic(err)
}
base, err := latencykernel.New(in)
if err != nil {
    panic(err)
}

// The same deployment with a 16-bit KV cache instead of fp8. New copied what it needed,
// so changing the documents now leaves base as it was.
in.Deployment.Pools[0].Engine.CacheDType = "bfloat16"
wide, err := latencykernel.New(in)
if err != nil {
    panic(err)
}
fmt.Printf("fp8:  %d bytes of KV per token per rank\n", base.SequenceVariableBytes(64)/64)
fmt.Printf("bf16: %d bytes of KV per token per rank\n", wide.SequenceVariableBytes(64)/64)

fp8:  35136 bytes of KV per token per rank
bf16: 70272 bytes of KV per token per rank

New copies what it needs, so changing the documents after it returns does not change the kernel it returned.

Read the memory

const gib = 1 << 30
m := k.FixedBytes()
fmt.Printf("weights      %5.1f GiB\n", float64(m.Weights)/gib)
fmt.Printf("activations  %5.1f GiB\n", float64(m.ActivationPeak)/gib)
fmt.Printf("CUDA graphs  %5.1f GiB\n", float64(m.CUDAGraph)/gib)
fmt.Printf("comm buffers %5.1f GiB\n", float64(m.CommBuffers)/gib)
fmt.Printf("EPLB         %5.1f GiB\n", float64(m.EPLBRedundant)/gib)
fmt.Printf("total        %5.1f GiB per rank\n", float64(m.Total())/gib)
// The communicator and engine workspace, held before the budget is measured.
fmt.Printf("before the budget %4.1f GiB\n", float64(k.PreSnapshotBytes())/gib)

// KV cache for one request, rounded up to whole pages.
fmt.Printf("a 4,096-token request: %.0f MiB of KV cache per rank\n",
    float64(k.SequenceVariableBytes(4096))/(1<<20))

weights       78.2 GiB
activations    1.1 GiB
CUDA graphs    0.8 GiB
comm buffers   0.0 GiB
EPLB           0.0 GiB
total         80.1 GiB per rank
before the budget  3.9 GiB
a 4,096-token request: 137 MiB of KV cache per rank

Memory explains each figure.

Read the evidence

measured, total, assumed := k.Evidence()
fmt.Printf("%d of %d values rest on evidence; %d were assumed\n",
    measured, total, len(assumed))

// The same entries, counted by the method that produced each.
byMethod := map[string]int{}
for _, o := range k.Provenance() {
    byMethod[o.Method]++
}
for _, m := range slices.Sorted(maps.Keys(byMethod)) {
    fmt.Printf("  %-12s %3d\n", m, byMethod[m])
}

// The kernel's own assumptions, as distinct from registry values whose method is
// "assumed", carry the set name KernelAssumptionSet.
for _, o := range k.Provenance() {
    if o.Set == latencykernel.KernelAssumptionSet {
        fmt.Println("the kernel assumed", o.Name)
    }
}

113 of 155 values rest on evidence; 41 were assumed
  assumed       41
  measured      99
  not_charged    1
  vendor_spec   14
the kernel assumed max_cudagraph_capture_size
the kernel assumed engine_version

Refuse, assume, override explains what these entries mean.

When New refuses

Open, OpenPool and New return an error rather than a kernel for a deployment the engine would not start, one the kernel cannot price truthfully, or one missing a coefficient it needs. The message names the setting at fault. What New refuses lists every case.

Calling it often

A kernel holds no mutable state: it can be shared between goroutines, and the same batch always prices the same. A step is cheap because everything that does not depend on the batch is computed once, in New. A step then walks the model's distinct kinds of layer rather than its layers, and multiplies by how many of each there are: three kinds instead of 108 layers for the deepest model in the catalog. On a laptop, bench_test.go measures a step at a few hundred nanoseconds for one request and a few microseconds for a 256-request batch.

A caller pricing millions of steps can use StepTimeInto, which writes the per-resource breakdown into a map the caller reuses instead of allocating one; the step then allocates only a few small scratch slices.