Skip to content

Migrate Gemma 4 onto the step contract #956

Description

@xiaguan

Problem

Gemma 4 still launches on the legacy engine contract (LaunchedEngine::Handle + per-request TokenEvent), while Qwen3 and pegainfer-sim already speak the step contract (LaunchedEngine::Stepped).

The two generations currently coexist in pegainfer-frontend (bridge.rs vs bridge/stepped.rs). Every remaining Handle line keeps that dual stack alive. Gemma 4 is a good next cut: single GPU, one scheduler, no LoRA / TP / EP / P/D, and prefix cache is resolved inside the engine rather than via frontend KvPrefix.

Today:

  • pegainfer-gemma4/src/model_line.rs returns LaunchedEngine::Handle.
  • pegainfer-gemma4/src/engine.rs owns its own thread, blocking_recvs when idle, and emits TokenEvent from admit / mixed gather / async-prefill join / chunked walk / decode.
  • Cancellation is TokenSink send-failure / channel close, not RequestControl::abort.
  • There is no Scheduler::load() occupancy feed (running / waiting / KV pages).

Related: Gemma 4 roadmap #758. Contract source of truth: docs/subsystems/frontend/frontend-architecture.md. References: pegainfer-qwen3/src/frontend_adapter.rs (full adapter) and pegainfer-sim (smallest Scheduler). The architecture doc currently lists glm52 as the next pilot because it exercises P/D and multi-scheduler; this issue is the simpler topology and can land independently.

Direction

Stop sending TokenEvent from the Gemma 4 engine. Implement Scheduler and return LaunchedEngine::Stepped.

Keep GPU mechanics in engine.rs (or a contract-free sibling). Put contract I/O in one adapter file, Qwen3-style:

  • submit only parks the QueuedRequest.
  • step is the only emission site: admit / reject / tokens / finish / fail / retire via the ledger.
  • load reports running, waiting, and KV occupancy.

Do not rewrite prefix cache, mixed admission, the async prefill lane, or the chunked walk. Fold their existing first-token paths into the same step so each touched request still gets one RequestUpdate.

Idle policy: the contract driver busy-polls. Gemma 4 currently blocking_recvs when the roster is empty. Either accept the poll (matches Qwen3) or park briefly inside step the way sim does — pick one and document it.

Abort: retire silently on the next touch when is_aborted() is set. Do not emit a terminal for a frontend abort.

Acceptance criteria

  • Gemma4Line::launch returns LaunchedEngine::Stepped. The crate no longer constructs EngineHandle or sends TokenEvent.
  • HTTP serving through the existing vLLM stack still works (completions + streaming).
  • Prefix-cache hits still surface as cached_tokens on the scheduled update.
  • Mixed admission, async prefill join, and chunked-walk graduation still emit the first token once, on the step that produced it.
  • Client abort no longer depends on sink teardown; the request is retired without a scheduler terminal.
  • Scheduler::load() is implemented (at least running / waiting / page occupancy). lora stays None.
  • Existing ignored GPU gates that cover lifecycle, cancel, prefix cache, mixed admission, and async prefill are rewritten against StepOutputs / Terminal and still pass on a checkpoint machine.
  • docs/subsystems/frontend/frontend-architecture.md and docs/models/gemma4/serving.md name Gemma 4 as a stepped line.

Contributor guardrails

  • Do not migrate glm52 / qwen35 / kimi-k2 / deepseek-v2-lite in this PR, and do not delete the legacy handle modules.
  • Do not fold frontend KvPrefix into the step contract here. Gemma 4 already refuses a frontend-resolved prefix.
  • Do not change CUDA kernels, KV page layout, or prefix-cache match/capture semantics except where the adapter has to copy pages the engine already copies.
  • Follow the onboarding checklist in frontend-architecture.md. Illegal event order should be unrepresentable (typestate handles), not re-checked by hand.
  • GPU evidence is required for the serving gates above. CPU-only adapter tests (fake executor, Qwen3 frontend_adapter/tests.rs style) are welcome in addition, not as a substitute.

Metadata

Metadata

Assignees

Labels

enhancementNew feature or requesthelp wantedExtra attention is neededhw:1-gpuNeeds one consumer GPU to verify

Type

No type

Projects

No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions