Cached-prefix prefill runs 10-40x slower than a cold prefill on a hybrid MoE model (scheduler spends 90% in _invalidate_prefill_buffer + stream sync) #501

Labels sugeridas: bug, performance

Environment

• FreeToken built from source, commit b123202 (branch work, includes feat(kvcache): store the KV cache as fp8 e4m3 codes (--kv-cache-dtype… #354 and two local fixes described below)
• Model: Qwen3.8-Flash-Next-mixed-FP8 (hybrid GDN + QSA sparse attention, 512 experts/layer, 48 layers)
• Hardware: EPYC 7K62 (48c/96t), RTX 4070 Ti Super 16 GB, 251 GB DDR4-3200, PCIe Gen4
• Server flags: --memory-ratio 0.92 --max-prefill-length 4096 --moe-cache-size 1024 --ple-backend pinned --max-running-requests 2 --enable-special-token-ckpt --enable-cache-report --kv-cache-dtype fp8
• No KV host offload on this instance (all KV in VRAM, 189,376 tokens)

Symptom

We run WorkBuddy (an agentic coding/chat client) against a single FreeToken instance for
multi-hour sessions. The prefix cache works: every turn matches 82k-170k cached tokens.
But the wall time per turn grows with the cached context, and the client waits 20-150 s
before the first token. Real numbers from the production log (one conversation, same server):

For comparison, a cold 4096-token prefill on the same server runs at ~1,238 tok/s. So the
per-chunk cost is 10-40x higher when the chunk attends over a large cached prefix.

Profile evidence (py-spy, engine scheduler process)

Captured during a slow cached-prefix turn (4,096 new over ~101k cached), 9,194 samples:

nvidia-smi --query-gpu=utilization.gpu sampled at 1 Hz during the turn reads 0-3%.
The GPU is idle while the scheduler blocks in _invalidate_prefill_buffer and in the
copy_done.synchronize() wait.

What _invalidate_prefill_buffer does

def _invalidate_prefill_buffer(self, buffer_id: int) -> None:
    slot_start = buffer_id * self.num_experts
    slot_end = slot_start + self.num_experts
    old_ids = self.id_of_slot[slot_start:slot_end]
    self.slot_for_id.view(-1)[old_ids[old_ids >= 0].long()] = -1
    old_ids.fill_(-1)
    self.usage[slot_start:slot_end].zero_()

This runs twice per chunk (once per double-buffer slot). The boolean-mask index
old_ids[old_ids >= 0] produces a data-dependent shape, so every call hides a
device-to-host synchronization. Each of those syncs drains ALL work already enqueued on
the device, which includes the attention over the full cached context. The prefill-overlap
pipeline (copy stream + compute stream) therefore serializes: per layer, the host waits
for the enqueued kernels before prefetching the next one, and the wait grows with the
cached-context length. That matches the profile: ~90% of scheduler time in the
invalidation plus its sync, ~3% doing actual work.

We reproduced the same anti-pattern in a second place:
kvcache/kv_host_offload.py::ensure_write_pages (also a boolean-mask index on a
data-dependent shape, 48 calls per chunk). We patched both locally:

1. ensure_write_pages: replaced torch.unique + boolean indexing with a fixed-shape
   Triton kernel (stride samples + one boundary sample per request run). After the fix,
   the offloaded-KV path boots with --max-running-requests 2 (it previously died at
   CUDA graph capture with AssertionError: ensure query of 1340 pages exceeds the 742 physical slots).
2. _invalidate_prefill_buffer: single fixed-shape Triton launch (scatter with a sink
   entry instead of the boolean index). The function disappears from the profile.

A flamegraph of the engine scheduler during a slow turn is attached
(freetoken-flamegraph-prefix-cache.svg).

Second finding: the tokenizer re-tokenizes the whole prompt every turn

In the same workload, a separate py-spy capture of the cached-prefix turn shows the
scheduler spending 67.7% of its samples in tokenize <- encode <- _encode_plus and
~29% in as_tensor/flatten (tensor conversion), with the GPU at 3.4% utilization.
The client resends the full conversation text every turn, and the tokenizer worker
re-tokenizes all 70k-170k tokens from scratch. An incremental tokenizer prefix cache
(keep the token ids per conversation prefix, tokenize only the delta, like SGLang/vLLM
do) would remove most of that cost for agentic sessions.

Third note: single-request context is capped by the rope table

The KV pool can be sized beyond the model's native context (we measured 424k logical
tokens with --kv-host-pages 5000 --num-pages 1635), but the engine refuses requests
past the rope table:

ValueError: --max-seq-len-override 520192 exceeds the model's rope table (262144
positions). Serving past it would read out of bounds; extend the checkpoint's
rope_scaling

That is reasonable behavior; it just means the extra pool capacity is only usable by
concurrent requests and by the prefix cache, not by a single longer request. Worth
documenting, since "KV capacity" reads like "single-request context".

Related

• hybrid MoE strategy 20-40x slower than offload on an offload-heavy box (24 GB card, PCIe gen 3) #436 (hybrid 20-40x slower than offload on PCIe gen 3; same model family)
• Long-context decode: client wall-clock is ~half the scheduler gen throughput #306 (long-context wall-clock gap)

Questions

1. Is the per-layer buffer invalidation required for correctness, or can the two
   double-buffer slot ranges be invalidated once per chunk at begin_prefill?
2. Would a sink entry in slot_for_id (one extra id, so the scatter is fixed-shape)
   be acceptable upstream? Happy to send the patch.
3. Is an incremental tokenizer cache on the roadmap?

Engine scheduler flamegraph during a cached-prefix turn (46 s window, 9,194 samples):

SVG source - raw py-spy stacks