Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
15 commits
Select commit Hold shift + click to select a range
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 3 additions & 2 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

6 changes: 3 additions & 3 deletions docs/models/gemma4/serving.md
Original file line number Diff line number Diff line change
Expand Up @@ -2,11 +2,11 @@

**TL;DR:** The engine schedules per iteration: up to the configured decode slots (16 by default) hold requests, each prompt prefills whole at a step boundary by default, and every active request advances one token per batched step. Prompt plus output past the ceiling — 8192 by default, raised up to the checkpoint's 262144 by `PEGAINFER_MAX_CONTEXT` — is refused at admission, while a request that only has to wait for a decode slot queues instead. The two KV families are budgeted separately (7.27 GiB sliding + 2.00 GiB global at 12B, at the defaults). **The default configuration needs a 48 GiB card**: it sits at 32.2 GiB before it serves anything, so a 32 GiB device cannot start it — smaller envelopes are a matter of the slots and ceiling knobs below. A row's output moves with the bucket widths it decodes at, but not with what its companions contain. An opt-in conversation prefix cache (`PEGAINFER_PREFIX_CACHE=K`) resumes multi-turn prompts at the cost of a pre-allocated page budget, and an opt-in overlap lane (`PEGAINFER_ASYNC_PREFILL=green:NN`) trades prefill latency for decode-tail protection under long-prompt admissions. An opt-in chunked walk (`PEGAINFER_MIX_CHUNK_TOKENS=N`) bounds how many prompt rows a mixed admission computes per step, so live streams advance per segment instead of waiting out whole prompts, and a raised ceiling (`PEGAINFER_MAX_CONTEXT`, with `PEGAINFER_DECODE_SLOTS` trading concurrency for context) serves long-context workloads on the same card. Dense and routed checkpoints use the same startup-precaptured decode graphs.

Last touched: 2026-08
Last touched: 2026-09

## What a step is

The engine thread runs one loop. Each turn it admits whatever the pools can hold, up to the slot ceiling. With streams in flight, admissions share one mixed step with them — the prompts' rows sit in the step's row prefix, each as its own segment, while every active request advances its token in the suffix; a prompt that arrives with nothing active prefills alone as its own step. Up to four coincident prompts gather into the same step, bounded at 512 unseen prompt rows (a warm resume is priced at its suffix), with every popped candidate — gathered, rejected or cancelled — consuming the turn's shared admission budget: gathering amortizes only the ~27 ms step floor while every live stream's inter-token gap pays the whole gathered step (~0.2 ms per row), so short bursts fold their admission staircase and a longer prompt keeps its own step. With the opt-in overlap lane enabled (below), an admission into a live batch prefills asynchronously on its own stream instead of sharing the mixed step. Between admissions, every active request advances exactly one token in a single batched decode step that shares the weight pass. A request that arrives while all slots are taken waits at the head of the queue. It is refused only when nothing is active — when there is no other request whose pages could free up, the pools genuinely cannot hold it and saying so is the honest answer.
The engine thread runs one loop. Each turn it admits whatever the pools can hold, up to the slot ceiling. With streams in flight, admissions share one mixed step with them — the prompts' rows sit in the step's row prefix, each as its own segment, while every active request advances its token in the suffix; a prompt that arrives with nothing active prefills alone as its own step. Up to four coincident prompts gather into the same step; a follower joins only while the gathered unseen rows stay under 512 (a warm resume is priced at its suffix), the leader itself is never bounded — a long prompt still shares the mixed step with the live batch, just without followers — and the chunked walk bypasses the follower budget, pricing rows per chunked step instead. Every popped candidate — gathered, rejected or cancelled — consumes the turn's shared admission budget: gathering amortizes only the step floor while every live stream's inter-token gap pays the whole gathered step, so short bursts fold their admission staircase (floor and per-row figures are 12B measurements recorded in the benchmark notes). With the opt-in overlap lane enabled (below), an admission into a live batch prefills asynchronously on its own stream instead of sharing the mixed step. Between admissions, every active request advances exactly one token in a single batched decode step that shares the weight pass. A request that arrives while all slots are taken waits at the head of the queue. It is refused only when nothing is active — when there is no other request whose pages could free up, the pools genuinely cannot hold it and saying so is the honest answer.

Rows retire independently. Requests in one batch have their own frontiers, their own page tables and, for the sliding family, their own released window front, so a short request finishing does not disturb the rows that continue.

Expand All @@ -25,7 +25,7 @@ Alignment uses 16-row blocks below a 1024-row step (8192 routed slots) and 64-ro

The register router accepts exactly 128 experts and from 1 through 32 picks. A non-finite row fails closed: every emitted index stays valid and row-unique, while every emitted weight is NaN so downstream computation remains loud.

The checkpoint-backed `the_routed_block_matches_the_reference_formulas` gate owns the scratch-capacity, companion-route and coarse-block evidence, including a narrow block replay after a coarse block on one scratch. On shared rows, it proves that the 16-row and 64-row block pick the same experts with the same weight bits and produce the same gate, weighted-down and block bits. `router_topk_matches_the_exact_128_expert_contract` owns the register-router boundary and non-finite rows. The kernels-owned `kimi_marlin_align_boundary_matches_vllm_contract` oracle owns stable counts, offsets, padding and expert-local order on both sides of the alignment dispatch boundary. `scripts/gemma4_gates.sh` owns the Gemma crate's ignored gates and the kernels crate's Gemma router contract (the device-only test under the `gemma4` feature), holding both crates' ignored sets against its manifest; the Kimi alignment oracle needs the `kimi-k2` feature and an `sm_90` device and is run by hand.
The checkpoint-backed `the_routed_block_matches_the_reference_formulas` gate owns the scratch-capacity, companion-route and coarse-block evidence, including a narrow block replay after a coarse block on one scratch. On shared rows, it proves that the 16-row and 64-row block pick the same experts with the same weight bits and produce the same gate, weighted-down and block bits. `router_topk_matches_the_exact_128_expert_contract` owns the register-router boundary and non-finite rows. The kernels-owned `kimi_marlin_align_boundary_matches_vllm_contract` oracle owns stable counts, offsets, padding and expert-local order on both sides of the alignment dispatch boundary. `scripts/gemma4_gates.sh` owns three scopes and holds each ignored set against its manifest by exact name: the Gemma crate's library and integration gates, the kernels crate's Gemma contracts (router, suppression, the norm parities and the hd256 fp8 pool suite, all device-only under the `gemma4` feature), and the frontend crate's chat-render parity gate; the Kimi alignment oracle needs the `kimi-k2` feature and an `sm_90` device and is run by hand.

## The decode pipeline

Expand Down
27 changes: 13 additions & 14 deletions docs/models/gemma4/tokenizer.md
Original file line number Diff line number Diff line change
@@ -1,8 +1,8 @@
# Gemma 4 tokenizer and chat template

**TL;DR:** All five chat renders reproduce the Hugging Face reference when content is flattened to strings, and token ids agree across the layers that can differ — gated by `pegainfer-gemma4/tests/tokenizer_parity.rs` against the pinned 12B checkpoint, with the other two sizes covered by inspection rather than by running. One divergence is open: under the frontend's default content format the system turn gains a trailing space. Contracts the engine must honour: BOS comes only from the chat template, EOS is declared in three places with three different values, the published generation defaults are sampled rather than greedy, and text-only serving rejects modality tokens before embedding and suppresses them before sampling.
**TL;DR:** All five chat renders reproduce the Hugging Face reference when content is flattened to strings — gated by `pegainfer-frontend/tests/gemma4_tokenizer_parity.rs` (runner-owned) against the pinned 12B checkpoint, with the other two sizes covered by inspection rather than by running. The token-id probes are retired: both sides run the same `tokenizers` crate, so they gated the Python wrapper's version skew rather than behavior. One divergence is open: under the frontend's default content format the system turn gains a trailing space. Contracts the engine must honour: BOS comes only from the chat template, EOS is declared in three places with three different values, the published generation defaults are sampled rather than greedy, and text-only serving rejects modality tokens before embedding and suppresses them before sampling.

Last touched: 2026-08
Last touched: 2026-09

## The gate runs on 12B; the result carries to the other sizes by inspection

Expand Down Expand Up @@ -47,19 +47,17 @@ except those two and the transformers version, which are provenance — they rec
reference came from and cannot be checked from here. No local path is recorded.

Both sides tokenize with the same `tokenizers` crate — Python's fast tokenizer wraps it and
`vllm-tokenizer` calls it directly — so broad script coverage would re-run one implementation
twice. The probes instead target the layers that can genuinely disagree: the Python wrapper's
added-token and `add_special_tokens` handling, and version skew between the crate this workspace
pins and the one the transformers wheel bundles. One case per algorithm class (whitespace, digits,
multibyte, multi-codepoint graphemes, combining marks, byte fallback), every special token
standalone, and one embedded in a sentence to cover adjacency.
`vllm-tokenizer` calls it directly — which is why the token-id probes retired: they compared one
implementation with itself and could only gate the Python wrapper's version skew. What remains,
and what the gate asserts, is the chat render: minijinja against the reference's Jinja2, the one
comparison with two genuinely independent sides.

The parity tests carry `#[ignore]` because they need the checkpoint; run them explicitly against
the pinned 12B one:
The parity test carries `#[ignore]` because it needs the checkpoint; the maintainer runner
executes it, or run it directly against the pinned 12B one:

```bash
OPENINFER_TEST_MODEL_PATH=<pinned-12B-checkpoint-dir> \
cargo test --release -p pegainfer-gemma4 --test tokenizer_parity -- --ignored
PEGAINFER_TEST_MODEL_PATH=<pinned-12B-checkpoint-dir> \
cargo test --release -p pegainfer-frontend --test gemma4_tokenizer_parity -- --ignored
```

## The chat template lives in its own file
Expand Down Expand Up @@ -96,8 +94,9 @@ frontend today, so changing it touches every model line and needs its own review

## BOS comes from the template, not the tokenizer

`add_special_tokens` is a no-op for this tokenizer: every probe encodes identically with it on
and off. The leading `<bos>` in a chat request comes from the template. An engine that adds BOS
`add_special_tokens` was established as a no-op for this tokenizer by the retired token-id
probes (every probe encoded identically with it on and off); the fixture no longer carries them.
The leading `<bos>` in a chat request comes from the template. An engine that adds BOS
itself would double it.

## EOS is declared three times
Expand Down
3 changes: 3 additions & 0 deletions pegainfer-frontend/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,10 @@ vllm-server = { workspace = true }
zeromq = { workspace = true }

[dev-dependencies]
sha2 = { workspace = true }
tokio = { workspace = true, features = ["macros", "rt"] }
vllm-chat = { workspace = true }
vllm-text = { workspace = true }

[lints]
workspace = true
Original file line number Diff line number Diff line change
@@ -1,8 +1,12 @@
//! Chat-render parity against a Hugging Face reference dumped by
//! `tools/accuracy/dump_gemma4_tokenizer_golden.py`. Point
//! `PEGAINFER_TEST_MODEL_PATH` at the pinned 12B checkpoint the reference was
//! dumped from and run with `--ignored`; the file-hash guard binds it to
//! exactly that checkpoint.
//! `tools/accuracy/dump_gemma4_tokenizer_golden.py`. Owned by the frontend
//! crate because the contract under test is the chat-render path (the
//! vendored vllm renderer against the reference's Jinja2); the fixture is
//! shared test data and the checkpoint comes from the environment. The
//! maintainer runner executes it; directly, point `PEGAINFER_TEST_MODEL_PATH`
//! at the pinned 12B checkpoint the reference was dumped from and run
//! `cargo test -p pegainfer-frontend --test gemma4_tokenizer_parity -- --ignored`;
//! the file-hash guard binds it to exactly that checkpoint.
//!
//! Only the render comparison lives here, because only it compares two
//! implementations: minijinja against the reference's Jinja2. The token-id
Expand Down
2 changes: 0 additions & 2 deletions pegainfer-gemma4/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -35,8 +35,6 @@ serde_json = { workspace = true }
sha2 = { workspace = true }
tempfile = { workspace = true }
tokio = { workspace = true, features = ["macros", "rt"] }
vllm-chat = { workspace = true }
vllm-text = { workspace = true }

[lints]
workspace = true
Loading
Loading