fix(qwen3): preserve explicit stop-token causes

docs/subsystems/frontend/frontend-architecture.md

@@ -1,6 +1,6 @@
 # Frontend architecture: pegainfer-frontend and the engine boundary
 
-**TL;DR:** `pegainfer-frontend` owns everything north of the model schedulers: the engine contract, the vLLM protocol stack, and the `ModelLine` dispatch trait. The contract now has two generations living side by side: the **step contract** (`StepOutputs` wire + `RequestLedger` lifecycle + a contract-owned polling driver — Qwen3, Gemma 4 and `pegainfer-sim` are migrated) and the **legacy handle contract** (`EngineHandle` + `TokenEvent` per-request events — glm52/qwen35/kimi-k2/deepseek-v2-lite still launch through it). **Next step: migrate glm52, then delete the legacy contract.**
+**TL;DR:** `pegainfer-frontend` owns everything north of the model schedulers: the engine contract, the vLLM protocol stack, and the `ModelLine` dispatch trait. The contract now has two generations living side by side: the **step contract** (`StepOutputs` wire + `RequestLedger` lifecycle + a contract-owned polling driver — Qwen3, Gemma 4, K3 and `pegainfer-sim` are migrated) and the **legacy handle contract** (`EngineHandle` + `TokenEvent` per-request events — glm52/qwen35/kimi-k2/deepseek-v2-lite still launch through it). **Next step: migrate glm52, then delete the legacy contract.**
 
 Last touched: 2026-09
 
@@ -12,14 +12,14 @@ An engine is a set of schedulers, each a `Scheduler` implementation driven by th
 
 ```
 pegainfer-frontend/src/engine/
-├── step.rs              # the wire: RequestId, Request, QueuedRequest,
-│                        #   StepOutputs { Vec<RequestUpdate> },
-│                        #   RequestUpdate { scheduled, tokens, logprobs, cached_tokens,
-│                        #   prompt_echo, kv_transfer, terminal }, Terminal
+├── step.rs              # the wire: RequestId, Request, StepOutputs { Vec<RequestUpdate> },
+│                        #   Request { ..., stop_policy }, RequestUpdate { scheduled, tokens,
+│                        #   logprobs, cached_tokens, prompt_echo, kv_transfer, terminal },
+│                        #   Terminal { ..., stop_cause }
 ├── request_lifecycle.rs # submission envelope, abort control and step sender plumbing;
 │                        #   DeferredFinish remains available for P/D handoff
-├── ledger.rs            # RequestLedger: admit/reject/push/finish/fail/retire,
-│                        #   prompt/completion tallies, one merged update per touched id
+├── ledger.rs             # RequestLedger: admit/reject/push/finish/fail/retire,
+│                         #   prompt/completion tallies, one merged update per touched id
 ├── wiring.rs            # scheduler_pair, SchedulerHandle (submit/take_steps/load),
 │                        #   Engine { schedulers, info, lora }, LiveScheduler,
 │                        #   EngineInfo, LaunchedEngine { Handle | Stepped }
@@ -32,10 +32,11 @@ pegainfer-frontend/src/engine/
 Design decisions worth knowing before touching it:
 
 - **Step-batched wire.** One message per scheduler step, not one channel per request: the scheduler's natural output unit is the step batch, and per-request channels were tried and rejected (the scheduler for-loop over N channels was the bottleneck). The protocol stack demuxes.
-- **Flat `RequestUpdate`.** All facts a step produced for one request travel in one struct, so intra-request ordering is structure, not convention. The ledger merges every write for an id into that one record before the driver commits the step.
-- **Ledger lifecycle.** `RequestLedger` owns one account for every unanswered submission. A scheduler receives only `QueuedRequest { id, request }` and writes every lifecycle transition by id. The ledger rejects touches after closure, derives completion counts from `push_tokens`, and writes off every open account when an engine-fatal `step` error ends the driver.
-- **Ledger as single writer.** Schedulers never touch the step channel; they call ledger methods. `admit` stamps `ScheduledInfo` from the registered prompt length, `push_tokens` tallies completions, and `commit_step` publishes the merged statement once per driver iteration.
-- **Polling driver, scheduler-owned park.** `spawn_scheduler` owns the serve loop: drain submissions, `Scheduler::step`, publish metrics, commit. An idle iteration ends in a `spin_loop` hint. Gemma 4 has one deliberate park inside this policy: while async prefill is the only remaining work, its scheduler drains the lane and joins it rather than hot-polling the completion (a drain failure is engine-fatal); when decode or queued work exists it keeps polling the lane without blocking.
+- **Flat `RequestUpdate`.** All facts a step produced for one request travel in one struct, so intra-request ordering is structure, not convention. This is what makes `defer_finish` safe: a P/D prefill executor can withhold a request's `Finished` until its KV saves are peer-visible and send it later from any thread — the deferred message carries the request's entire buffered update, so late delivery cannot reorder.
+- **Independent stop policy and cause.** `Request.stop_policy` keeps model EOS handling separate from explicit request stop IDs. A token-driven finish retains the triggering token in `RequestUpdate` and carries `Terminal::Finished.stop_cause`; decode paths include its real logprob when available, so the stepped bridge can report the actual explicit stop ID without reconstructing it. `ignore_eos` affects only model EOS. Legacy producers may leave the cause empty while they are migrated individually.
+- **Ledger lifecycle.** Schedulers carry plain `RequestId`s and mutate `RequestLedger`; `RequestEnvelope` answers submissions dropped before registration, while `DeferredFinish` preserves a complete buffered update when a terminal is delivered after a P/D handoff. `RequestControl` is the frontend abort flag. The ledger remains the single writer for admission, token accounting, terminal transitions, and step publication.
+- **Pure polling driver.** `spawn_scheduler` owns the serve loop: drain submissions, `Scheduler::step`, publish load, commit. No idle/park distinction — the scheduler owns the GPU and spinning on it costs nothing anyone else could use; async KV I/O (prefetch, decode-overlap prefill) is naturally absorbed by polling. An idle iteration ends in a `spin_loop` hint (relaxes the core's issue slots, no latency cost — busy iterations never pause). The loop exits when the frontend drops the handle and the queue drains.
+- **Gemma 4 async prefill exception.** While asynchronous prefill is the only remaining work, Gemma 4 drains and joins that lane rather than hot-polling its completion; decode or queued work keeps the normal polling path.
 - **Abort is a flag, not channel teardown.** `SchedulerHandle::submit` returns a `RequestControl`; the frontend flips its boolean abort flag and the scheduler retires the request silently on its next touch (no terminal — the frontend already dropped its state for that id).
 - **Channels:** the submit channel is crossbeam (sync consumer on the scheduler thread), steps are tokio mpsc (async consumer in the bridge); load is a shared cell read via `SchedulerHandle::load()` — pull-only by design, "notify me on load change" is deliberately unrepresentable (the driver busy-polls, so a subscription edge would fire per spin). All channels unbounded on purpose — admission control is the scheduler's job, expressed as `Rejected`, never as backpressure on submit.
 - **Control plane lives outside the contract.** `Scheduler` has no control method and the contract carries no control channel. A capability like LoRA is a private channel the model crate mints *before* `spawn_scheduler` — the scheduler closes over the receiver, the `LoraClient` sender surfaces as `Engine.lora: Option<LoraClient>`, and the `Option` *is* the capability (no `bool` flag, no registry until a second capability exists). The vocabulary (`LoraControl`, `LoraClient`) is still defined in the frontend crate because the frontend must speak it without holding model structs; only the wiring is the model's business.
@@ -91,7 +92,7 @@ All six lines are onboarded. Adding a model line = write `model_line.rs` in the
 
 ## Protocol stacks
 
-**`vllm` (current default, fleet-proven).** Impersonates a vLLM EngineCore process over in-process ZMQ/msgpack because upstream `vllm-server` assumes the engine is a separate process. HTTP routes, OpenAI types, tokenizer, chat templates, Prometheus live in the external `vllm-server`/`vllm-metrics`/`vllm-text` crates. `SteppedEngineBridge` translates each `RequestUpdate` 1:1 into an EngineCore output (wall-clock timestamps are reconstructed from the contract's `Instant`s via a per-bridge unix anchor; a `Finished{Stop}` appends the stop sentinel token, which is how usage keeps counting the suppressed EOS).
+**`vllm` (current default, fleet-proven).** Impersonates a vLLM EngineCore process over in-process ZMQ/msgpack because upstream `vllm-server` assumes the engine is a separate process. HTTP routes, OpenAI types, tokenizer, chat templates, Prometheus live in the external `vllm-server`/`vllm-metrics`/`vllm-text` crates. `SteppedEngineBridge` translates each `RequestUpdate` 1:1 into an EngineCore output (wall-clock timestamps are reconstructed from the contract's `Instant`s via a per-bridge unix anchor; a typed `StopCause` reports the actual request stop token, while the synthetic sentinel remains only as a compatibility fallback for producers that provide no cause).
 
 **`dynamo` (planned second stack).** dynamo's `lib/llm` in-process path removes the wire protocol entirely (`EngineConfig::InProcessTokens` + `run_input`). The step contract was shaped so this stack can consume `StepOutputs` directly without impersonation overhead. Decision gate: prototype, A/B against the vllm stack, let TTFT/step-overhead numbers pick the default.
 

pegainfer-frontend/src/engine/driver.rs

@@ -108,6 +108,7 @@ mod tests {
     use super::super::step::Request;
     use super::super::step::RequestId;
     use super::super::step::Terminal;
+    use super::super::stop::StopPolicy;
     use super::*;
     use crate::engine::FinishReason;
 
@@ -163,6 +164,7 @@ mod tests {
         Request {
             prompt_tokens: vec![1, 2],
             params: crate::sampler::SamplingParams::default(),
+            stop_policy: StopPolicy::default(),
             max_tokens,
             lora_adapter: None,
             kv_transfer_params: None,
@@ -197,6 +199,7 @@ mod tests {
             terminal,
             Some(Terminal::Finished {
                 reason: FinishReason::Length,
+                stop_cause: None,
                 prompt_tokens: 2,
                 completion_tokens: 3,
             })

pegainfer-frontend/src/engine/ledger.rs

@@ -38,6 +38,7 @@ use super::step::RequestUpdate;
 use super::step::ScheduledInfo;
 use super::step::StepOutputs;
 use super::step::Terminal;
+use super::stop::StopCause;
 
 /// One open account: the request's admission facts and running tally. The
 /// payload is not here — it went to the scheduler at `submit`; the account is
@@ -239,12 +240,23 @@ impl RequestLedger {
 
     /// Finish the request. Token counts come from the ledger's tally.
     pub fn finish(&mut self, id: RequestId, reason: FinishReason) {
+        self.finish_with_cause(id, reason, None);
+    }
+
+    /// Finish a request while preserving a typed token-level stop cause.
+    pub fn finish_with_cause(
+        &mut self,
+        id: RequestId,
+        reason: FinishReason,
+        stop_cause: Option<StopCause>,
+    ) {
         let account = self.close(id);
         let AccountState::Active { completion_tokens } = account.state else {
             panic!("finish on {id} before admission");
         };
         self.statement.entry(id).terminal = Some(Terminal::Finished {
             reason,
+            stop_cause,
             prompt_tokens: account.prompt_len,
             completion_tokens,
         });
@@ -279,6 +291,16 @@ impl RequestLedger {
     /// this step — tokens included — folds into the returned message, so late
     /// delivery cannot reorder against the step stream.
     pub fn defer_finish(&mut self, id: RequestId, reason: FinishReason) -> DeferredFinish {
+        self.defer_finish_with_cause(id, reason, None)
+    }
+
+    /// Defer a finish while preserving a typed token-level stop cause.
+    pub fn defer_finish_with_cause(
+        &mut self,
+        id: RequestId,
+        reason: FinishReason,
+        stop_cause: Option<StopCause>,
+    ) -> DeferredFinish {
         let account = self.close(id);
         let AccountState::Active { completion_tokens } = account.state else {
             panic!("defer_finish on {id} before admission");
@@ -289,6 +311,7 @@ impl RequestLedger {
             .unwrap_or_else(|| RequestUpdate::empty(id));
         update.terminal = Some(Terminal::Finished {
             reason,
+            stop_cause,
             prompt_tokens: account.prompt_len,
             completion_tokens,
         });
@@ -374,6 +397,7 @@ mod tests {
     use super::super::request_lifecycle::StepReceiver;
     use super::super::step::Request;
     use super::super::step::Terminal;
+    use super::super::stop::StopPolicy;
     use super::super::wiring::SchedulerHandle;
     use super::super::wiring::scheduler_pair;
     use super::*;
@@ -382,6 +406,7 @@ mod tests {
         Request {
             prompt_tokens: prompt,
             params: crate::sampler::SamplingParams::default(),
+            stop_policy: StopPolicy::default(),
             max_tokens: 8,
             lora_adapter: None,
             kv_transfer_params: None,
@@ -402,7 +427,9 @@ mod tests {
         backend.ledger.admit(id);
         backend.ledger.push_tokens(id, &[10, 11], &[]);
         backend.ledger.set_cached_tokens(id, 2);
-        backend.ledger.finish(id, FinishReason::Stop);
+        backend
+            .ledger
+            .finish_with_cause(id, FinishReason::Stop, Some(StopCause::Token(11)));
         backend.ledger.commit_step();
 
         let mut steps = handle_steps(handle);
@@ -419,6 +446,7 @@ mod tests {
             update.terminal,
             Some(Terminal::Finished {
                 reason: FinishReason::Stop,
+                stop_cause: Some(StopCause::Token(11)),
                 prompt_tokens: 3,
                 completion_tokens: 2,
             })
@@ -479,7 +507,9 @@ mod tests {
         let id = backend.ledger.register(envelope).id;
         backend.ledger.admit(id);
         backend.ledger.push_tokens(id, &[7], &[]);
-        let deferred = backend.ledger.defer_finish(id, FinishReason::Length);
+        let deferred = backend
+            .ledger
+            .defer_finish_with_cause(id, FinishReason::Length, None);
         backend.ledger.commit_step();
 
         let mut steps = handle_steps(handle);
@@ -498,6 +528,7 @@ mod tests {
             update.terminal,
             Some(Terminal::Finished {
                 reason: FinishReason::Length,
+                stop_cause: None,
                 prompt_tokens: 2,
                 completion_tokens: 1,
             })

pegainfer-frontend/src/engine/mod.rs

@@ -39,6 +39,7 @@ mod request;
 mod request_lifecycle;
 mod sink;
 mod step;
+mod stop;
 mod wiring;
 
 pub use control::*;
@@ -52,4 +53,5 @@ pub use request::*;
 pub use request_lifecycle::*;
 pub use sink::*;
 pub use step::*;
+pub use stop::*;
 pub use wiring::*;

pegainfer-frontend/src/engine/step.rs

@@ -14,6 +14,8 @@ use std::time::Instant;
 
 use super::event::FinishReason;
 use super::event::TokenLogprob;
+use super::stop::StopCause;
+use super::stop::StopPolicy;
 
 /// In-process routing id for one generate request, minted by
 /// [`super::SchedulerHandle::submit`] from a per-scheduler counter. `Copy` and
@@ -49,6 +51,7 @@ impl std::fmt::Display for RequestId {
 pub struct Request {
     pub prompt_tokens: Vec<u32>,
     pub params: crate::sampler::SamplingParams,
+    pub stop_policy: StopPolicy,
     pub max_tokens: usize,
     pub lora_adapter: Option<String>,
     /// Opaque router/P-D metadata from the request's
@@ -233,6 +236,10 @@ impl fmt::Display for RejectReason {
 pub enum Terminal {
     Finished {
         reason: FinishReason,
+        /// Present for token-driven stop finishes. The triggering token remains
+        /// in `RequestUpdate.tokens`, with its real logprob in the matching
+        /// `RequestUpdate.logprobs` entry.
+        stop_cause: Option<StopCause>,
         prompt_tokens: usize,
         completion_tokens: usize,
     },

pegainfer-frontend/src/engine/stop.rs

@@ -0,0 +1,121 @@
+use std::sync::Arc;
+
+/// How a request treats end-of-sequence tokens.
+#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
+pub enum EosPolicy {
+    /// Do not stop on model EOS tokens.
+    Ignore,
+    /// Use the model executor's configured EOS set.
+    #[default]
+    ModelDefault,
+}
+
+/// Request-scoped token stopping policy.
+///
+/// EOS is kept separate from caller stop tokens because the vLLM protocol
+/// reports them differently: EOS has no 'stop_reason', while a request stop
+/// reports the actual matching token ID.
+#[derive(Clone, Debug, Default, Eq, PartialEq)]
+pub struct StopPolicy {
+    eos: EosPolicy,
+    /// Sorted and deduplicated explicit stop IDs.
+    ///
+    /// Requests clone this policy while building a plan and sending it to
+    /// worker ranks. Keeping the normalized set behind an `Arc` makes those
+    /// clones cheap and lets classification use binary search for large stop
+    /// sets without regressing the common one-ID case.
+    token_ids: Arc<[u32]>,
+}
+
+impl StopPolicy {
+    /// Build a policy from wire-provided stop IDs.
+    ///
+    /// Normalization happens once at the request boundary. Internal copies can
+    /// then share the immutable slice instead of repeatedly sorting, deduping,
+    /// or cloning the caller's vector.
+    #[must_use]
+    pub fn new(eos: EosPolicy, mut token_ids: Vec<u32>) -> Self {
+        token_ids.sort_unstable();
+        token_ids.dedup();
+        Self {
+            eos,
+            token_ids: token_ids.into(),
+        }
+    }
+
+    /// Classify a token using vLLM's priority: EOS first, then the request's
+    /// explicit stop-token set.
+    #[must_use]
+    pub fn classify(
+        &self,
+        token_id: u32,
+        is_model_eos: impl FnOnce(u32) -> bool,
+    ) -> Option<StopCause> {
+        let is_eos = match self.eos {
+            EosPolicy::Ignore => false,
+            EosPolicy::ModelDefault => is_model_eos(token_id),
+        };
+
+        if is_eos {
+            Some(StopCause::Eos(token_id))
+        } else if self.token_ids.binary_search(&token_id).is_ok() {
+            Some(StopCause::Token(token_id))
+        } else {
+            None
+        }
+    }
+}
+
+/// The token-level cause of a normal stop finish.
+#[derive(Clone, Copy, Debug, Eq, PartialEq)]
+pub enum StopCause {
+    /// A primary or model-default EOS token.
+    Eos(u32),
+    /// A token from the request's explicit stop-token set.
+    Token(u32),
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn model_default_classifies_model_eos() {
+        let policy = StopPolicy::default();
+
+        assert_eq!(
+            policy.classify(99, |token_id| token_id == 99),
+            Some(StopCause::Eos(99))
+        );
+    }
+
+    #[test]
+    fn ignored_eos_does_not_disable_an_explicit_stop() {
+        let policy = StopPolicy::new(EosPolicy::Ignore, vec![99]);
+
+        assert_eq!(
+            policy.classify(99, |token_id| token_id == 99),
+            Some(StopCause::Token(99))
+        );
+    }
+
+    #[test]
+    fn normalizes_unsorted_duplicate_stop_ids() {
+        let policy = StopPolicy::new(EosPolicy::Ignore, vec![7, 3, 7, 1]);
+
+        assert_eq!(policy.classify(1, |_| false), Some(StopCause::Token(1)));
+        assert_eq!(policy.classify(3, |_| false), Some(StopCause::Token(3)));
+        assert_eq!(policy.classify(7, |_| false), Some(StopCause::Token(7)));
+        assert!(policy.classify(8, |_| false).is_none());
+    }
+
+    #[test]
+    fn model_eos_has_priority_over_explicit_stop() {
+        let policy = StopPolicy::new(EosPolicy::ModelDefault, vec![99]);
+
+        assert_eq!(
+            policy.classify(99, |token_id| token_id == 99),
+            Some(StopCause::Eos(99))
+        );
+    }
+}