feat(gemma4): serve requests through the engine contract - #890
Conversation
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: eaf3e5bf24
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
|
@codex review please |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: eaf3e5bf24
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
Signed-off-by: Feathbow <feathbow@gmail.com>
eaf3e5b to
a0a55cc
Compare
|
@codex review again please |
|
Codex Review: Didn't find any major issues. 🚀 Reviewed commit: ℹ️ About Codex in GitHubCodex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback". |
Description
Closes #889.
One owned thread serves Gemma 4 requests serially through the KV-backed forward. Admission runs per step, and a request whose prompt plus requested output exceeds the serving ceiling is refused at this boundary; next-token selection and per-token logprobs go through
pegainfer-sampleon device logits; EOS and length finish reasons are real from the first request. Cancellation is the sink refusing an event: the request retires mid-stream, its KV returns by drop, and the engine serves the next one.The stop token retires a request without being emitted, which is the protocol the frontend documents: it appends its own sentinel for the terminal
Stopand unconditionally drops the last id, so an engine that emits EOS costs the client its final visible token. The ids the checkpoint lists undersuppress_tokensare dropped out of contention before anything reads the logits row, so sampling and the logprob normalization see the same distribution and a placeholder id can neither reach the client nor return as the next step's input. Both sets are read fromgeneration_config.jsonrather than assumed for the family: 12B declares the end-of-image and end-of-audio ids, 26B and 31B declare none.At the Gemma engine boundary, everything this stage does not implement is refused loudly rather than downgraded. At load: CUDA graph capture, multi-device ordinals, parallel topologies. Per request: a resolved KV prefix,
kv_transfer_params, a nonzerodata_parallel_rank, a LoRA adapter, prompt echo.The model line stops being registration only.
--cuda-graphdefaults to true for every line, so an explicit--cuda-graph=trueis refused inModelLine::validatebefore the weight load, and everything else launches eagerly. The handle publishes an 8192-token serving ceiling, which is what lets the frontend reject an over-length prompt with its own message; without it the frontend advertises the checkpoint's 262144max_position_embeddingsand forwards a request the engine can only fail mid-stream.That ceiling is what bounds the rope tables and the pool budget: each KV pool is sized for one full-context request plus its padding page, since a prompt is prefilled in one step and both families hold every position of it before the window releases anything. The 262k the checkpoint declares needs a table and KV budget design of its own.
Test Env
Single GPU (sm_89, x86_64), CUDA 12.9, pinned 12B checkpoint.
Verification
-D warningsover all gemma4 targets with and without the feature; the CI package set with--locked; thepegainfer-kernelshd256 smoke; gemma4 lib tests;pegainfer-serverbuilds with--features gemma4.--ignored --test-threads=1): 8 passed, including all three certified generate fixture prompts matching token for token./v1/modelsreportsmax_model_len8192. All three certified fixture prompts, sent as/v1/completionstext attemperature 0, return text byte-identical to the fixture continuations (3/3), with prompt token counts matching the fixture (18/13/16).'111.1.'over 8 chunks, andParisrespectively); the chat turn ends withfinish_reason=stopand no stop token in the text; per-token logprobs come back with the picked token's logprob equal to its top entry; dropping a stream mid-decode aborts the request and the next one is served normally.this model's maximum context length is 8192 tokens, but the prompt contains 9001 input tokens. A one-token prompt asking for 9000 output tokens is not refused: the frontend clamps it to the remaining budget and the request completes withfinish_reason=lengthat 8191 completion tokens.--cuda-graph=trueexits withGemma 4 serves eagerly; --cuda-graph=true is not supportedand no weight-load line in its output;--cuda-graph=falseserves.Type of Change