Sampled, Not Yet Cached: A GPU Trace of Mini-SGLang's KV Cache

In my previous Mini-SGLang post, I traced a request from the HTTP endpoint through the tokenizer, scheduler, GPU forward pass, and streaming response.

There was one detail I wanted to check more carefully:

If prefill has already generated token D, why doesn't D have its own KV cache yet?

I could follow the code and reason about the answer, but I wanted to see the actual request state and GPU KV slots—not just infer the behavior from function names.

So I added a small amount of instrumentation to Mini-SGLang and traced one request end to end.

Where prefill stops and decode starts

Let's say a prompt becomes three tokens after tokenization:

A B C

Prefill runs A, B, and C through the Transformer. At each layer, the model computes Q, K, and V for these input tokens, runs attention, and stores their K/V in the cache. The logits are then used to sample the next token, D.

Input: A B C
      │
      ▼
Transformer
      ├── K/V for A, B, C → KV Cache
      └── Logits → Sampler → D

The distinction is easy to miss: D is the output of this forward pass, not one of its inputs. The model computed K/V for A, B, and C, but not for D.

In the next decode step, D becomes the input. The model computes Q/K/V for this one new token and uses the cached K/V of A, B, and C during attention. It then samples E, and the process continues.

Mini-SGLang tracks this with three request-state fields:

Without a prefix-cache hit:

Before prefill:   cached/device/extend = 0/3/3
After prefill:    cached/device/extend = 3/4/1   ← D sampled, D's KV still missing
After decoding D: cached/device/extend = 4/5/1   ← D's KV stored, E just sampled

I already knew that extend_len=1 usually means we're decoding one token for this request. What I wanted to understand was how that 1 turns into an actual input token and a physical KV slot.

Who picks the slot, and who writes to it?

It helps to separate three operations:

1. Allocate: reserve physical KV slots and update the mapping.
2. Write: store the newly computed K/V tensors in those slots.
3. Read: use the cached K/V during attention through the page-table mapping and backend metadata.

These are not interchangeable.

In Mini-SGLang, TableManager allocates a page-table row for a request, while CacheManager manages the available physical KV slots. They are different kinds of indices.

page_table[table_idx, logical_token_position] → physical_kv_slot

In Scheduler._prepare_batch():

self.cache_manager.allocate_paged(batch.reqs)
# prepare positions, input_mapping and write_mapping ...
batch.out_loc = self.engine.page_table[input_mapping]

allocate_paged() allocates slots for the tokens in [cached_len, device_len). The resulting batch.out_loc tells the attention backend where the new K/V from this forward pass should go.

At this point, the slots have been assigned, but the new K/V hasn't been computed or written. The KV tensor pool itself was created during engine initialization; this is not a new CUDA allocation for every generated token.

The actual write happens in the FlashInfer attention backend:

self.kvcache.store_kv(k, v, batch.out_loc, layer_id)
# ...
return metadata.wrapper.run(q=q, paged_kv_cache=kv_cache)

store_kv() places the new K/V in the appropriate layer's KV pool. The FlashInfer wrapper then runs attention using the cache.

One more detail: each Transformer layer has its own K/V tensors. out_loc=[13] is a token-slot index used across layers, not a claim that the whole model stores only one copy of K/V.

How I ran the trace

I wasn't trying to benchmark performance here. I wanted one request with as few moving parts as possible.

I used a naive cache to avoid prefix reuse and page size 1 to make the slot mapping easier to inspect. Disabling overlap and CUDA Graphs made the Python path easier to trace. None of these choices were intended to maximize throughput.

The request was:

{
  "prompt": "The capital of France is Paris, and the capital of Japan is",
  "max_tokens": 4,
  "ignore_eos": true
}

The tokenizer returned 13 prompt tokens (L=13).

I added an opt-in flag, MINISGL_TRACE_REQUEST=1, to emit structured JSONL events. This flag was added for the experiment; it is not an upstream Mini-SGLang option.

The logs include request UID, process ID, phase, state lengths, out_loc, model input IDs, KV writes, and sampled token IDs.

I also read back K-cache values at layers 0 and 27 to compare them with the new K tensors. These GPU-to-CPU reads introduce synchronization, so the trace is useful for behavior and state, not performance measurements.

What the trace actually showed

Here is the useful part, with timestamps and some fields removed:

tokenized             L=13

batch_scheduled       prefill
kv_slots_allocated    0/13/13  out_loc=[0..12]  table_idx=255
kv_store_invoked      prefill  layer=0   out_loc=[0..12]  K matches
kv_store_invoked      prefill  layer=27  out_loc=[0..12]  K matches
request_state_updated prefill  0/13/13 -> 13/14/1
token_sampled         prefill  next_token=26194

batch_scheduled       decode
kv_slots_allocated    13/14/1 out_loc=[13]
forward_started       decode  input_ids=[26194]
kv_store_invoked      decode  layer=0   out_loc=[13]  K matches
kv_store_invoked      decode  layer=27  out_loc=[13]  K matches
request_state_updated decode  13/14/1 -> 14/15/1
token_sampled         decode  next_token=13

table_idx=255 is the request's page-table row, not physical KV slot 255. The physical slots used in this run started at 0.

Prefill only allocated slots for the prompt

Prefill started at 0/13/13, with out_loc=[0..12]: exactly 13 slots for 13 prompt tokens.

Both traced layers wrote to these slots. There was no slot allocated for D during prefill.

The sampler then produced:

D = token ID 26194

The state advanced to 13/14/1. The prompt's K/V was available, while D was waiting to be processed in the next pass.

D got its slot in decode

The next forward pass had input_ids=[26194]. That is the point where D changed from a sampled output into a model input.

The scheduler then allocated:

out_loc = [13]
page_table[255, 13] = 13

At layer 0, the K-cache check reported:

k_norm_before     = 0.0
k_norm_after      = 768.829
slot_equals_new_k = True

The important check is slot_equals_new_k=True: after the store, the K value in slot 13 matched the K computed for D in this forward pass.

The experiment did not independently compare V. The source code passes K and V through the same store_kv(k, v, ...) write path, but the direct tensor-content check covered K.

Also, the observed k_norm_before=0.0 does not mean freshly allocated cache memory is guaranteed to be zero-initialized. That number alone would not prove the slot had never been used.

The next two iterations followed the same pattern:

There were 38 trace events, and the predefined validation script reported 12/12 PASS. These were 12 checks on one execution—not 12 independent runs.

A couple of details I almost missed

Req.complete_one() runs before the sampler call in Python.

My simplified mental model had been “model forward → sample → update request state.” But in this pinned version of Engine.forward_batch(), with the CUDA Graph branch omitted, the order is:

logits = self.model.forward()
for req in batch.reqs:
    req.complete_one()
next_tokens_gpu = self.sampler.sample(...)

complete_one() doesn't need the next token ID. It marks the currently processed token prefix as cached and advances the logical length for the next token.

The trace's token_sampled event was logged after copy_done.synchronize()—when the CPU had the result—not at the exact moment the GPU sampler started.

The final generated token G never got its own KV in this run.

The request used max_tokens=4. Once G=6722 was sampled, the request finished. There was no next decode pass to process G.

The final state was:

cached_len = 16
device_len = 17
extend_len = 1

So device_len=17 does not mean 17 tokens already have computed K/V. Here, K/V had been computed for 13 prompt tokens plus D, E, and F: 16 tokens.

That made the difference between cached_len and device_len much more concrete for me.

The process IDs also matched the architecture from my first post: PID 885 handled the API, PID 939 tokenization, and PID 938 scheduling, engine execution, and attention. Different responsibilities do not necessarily mean different OS processes.

What I wouldn't infer from these logs

The initial forward passes triggered FlashInfer JIT compilation. On top of that, the instrumentation read GPU tensors back to the CPU using operations such as .tolist() and K-cache comparisons, forcing host synchronization.

So the roughly 76-second prefill interval and the long first-decode wait in this trace are not representative TTFT or TPOT numbers. I did not measure clean serving latency, throughput, exact GPU kernel completion times, or HBM read volume.

This was one request on one NVIDIA L4, in one execution. I inspected K writes at layers 0 and 27, not every layer. The result also comes from a deliberately simplified configuration—naive cache, page size 1, overlap disabled, and CUDA Graph disabled.

Within that scope, though, the trace answers the original question: D was sampled during prefill without its own KV slot. The next decode pass took D as input, allocated slot 13, and wrote the newly computed K/V there.

That was the distinction I wanted to verify beyond source-code reading.

The instrumentation patch, scripts, and raw trace are in ChihTsungLu/mini-sglang-experiments.