fix(qwen3): preserve explicit stop-token causes

[RicardoMin]

Related to #865

Qwen3: preserve explicit stop-token causes in the stepped contract

Why this is a separate PR

This is the scoped follow-up requested by the maintainers during review of
#865. They asked
that the broad stop-contract change be split so that the shared boundary,
frontend bridges, and one model can be reviewed and validated independently.
This PR therefore extracts the Qwen3 migration from that work; Qwen3.5 and the
other model schedulers remain on their existing contract for now.

Summary

The previous vLLM dependency update fixed several frontend compatibility issues,
but the stepped Qwen3 path still collapsed two independent controls into the
single legacy ignore_eos flag. That made an explicit stop_token_ids request
indistinguishable from a model-EOS request and forced the bridge to guess a
synthetic stop token after the scheduler had already discarded the real one.

This PR gives Qwen3 a typed stop contract. It preserves the sampled trigger
token and its logprob, carries the concrete StopCause through the scheduler
and stepped bridge, and keeps EOS handling independent from request-provided
stop IDs.

What was wrong

The old contract exposed only FinishReason::Stop or FinishReason::Length.
When a request stopped, the bridge could not tell whether the model emitted EOS
or an explicit request stop token. It therefore reconstructed a sentinel (EOS
first, otherwise the first configured stop ID). That reconstruction can report
the wrong token, loses the token's logprob, and is incorrect for a speculative
span where the first terminal token is followed by additional accepted tokens.

The old boolean also could not express the valid combination "ignore model EOS,
but still stop on these explicit request token IDs".

Contract change

Event	Legacy behavior	Qwen3 stepped behavior
Model EOS	FinishReason::Stop; trigger may be suppressed or reconstructed	FinishReason::Stop + StopCause::Eos(id); token is retained internally; wire stop_reason is absent
Explicit stop_token_ids match	Often treated as ordinary output or replaced by a guessed sentinel	FinishReason::Stop + StopCause::Token(id); actual ID is reported as wire stop_reason
ignore_eos=true	Could inadvertently disable explicit stops	Disables only model EOS; explicit request stops remain active
Length limit	FinishReason::Length	FinishReason::Length + no stop cause; final sampled token is retained
Speculative span	Trigger/suffix ownership was ambiguous	Commit the prefix through the first trigger and discard only the suffix after it

EOS has precedence when the same ID is both the active EOS token and an explicit
request stop. Completion-token accounting is incremented once, including the
trigger token.

Scope and compatibility

Following the maintainer's scope request on #865, this PR intentionally
migrates Qwen3 only. The shared request/step types
accept an optional typed cause, while the existing legacy event path remains
available for models that have not been audited. The legacy bridge keeps its
synthetic-sentinel fallback only when an old producer supplies no typed cause.
Therefore Qwen3.5 and other model schedulers are not changed in this PR and do
not need to adopt the new resolver contract yet. If the maintainers agree with
the semantics, the remaining model lines can be migrated one at a time with
their own lifecycle tests.

Implementation

• Added StopPolicy, EosPolicy, and StopCause at the frontend engine
  boundary.
• Converted wire EOS and explicit stop fields without collapsing them into one
  boolean for the stepped path.
• Propagated the policy through Qwen3 request state, ledger updates, prefill,
  ordinary decode, and speculative verification.
• Emitted the trigger token before terminal metadata and retained its logprob.
• Mapped StopCause::Token(id) to the vLLM-compatible wire stop_reason.
• Left the legacy bridge fallback and un-migrated model implementations intact.
• Updated two existing test fixtures (K3 and simulator) only to supply the new
  shared default field; no legacy model runtime behavior is changed.

Automated verification

Check	Result
cargo test --release -p pegainfer-frontend --lib	73 passed, 0 failed
cargo test --release -p pegainfer-qwen3 --lib	93 passed, 0 failed
Qwen3 request-stop focused tests	4 passed, 0 failed
Qwen3 speculative-stop focused tests	2 passed, 0 failed
cargo test --release -p pegainfer-sim --tests -- --test-threads=1	6 passed, 0 failed
cargo check --release -p pegainfer-qwen35 --features qwen35	Passed (control build only)
cargo build --release -p pegainfer-server --bin pegainfer	Passed
cargo fmt --all -- --check	Passed
git diff --check	Passed

HTTP A/B verification

The comparison used two already-running OpenAI-compatible endpoints on the
same validation host. The explicit stop set covered the complete vocabulary,
so the first generated token was guaranteed to exercise the request-stop path.
This is a deterministic contract probe, not a generation-quality benchmark.

Results are shown as finish_reason / stop_reason / completion_tokens:

Target	Control (ignore_eos=true)	Explicit stop + EOS ignored	Explicit stop + EOS enabled
Qwen3-0.6B (adapted)	length / null / 8	stop / 12095 / 1	stop / 12095 / 1
Qwen3.5-0.8B (legacy control)	length / null / 8	length / null / 8	length / null / 8

Contract check	Qwen3 adapted	Qwen3.5 legacy control
Baseline control	pass	pass
Explicit stop with EOS ignored	pass	fail
Explicit stop with EOS enabled	pass	fail
Stop-set order invariant	pass	not satisfied (no typed stop)
Trigger logprob present	pass	fail
Streaming typed stop	pass	fail
Mixed controls (3/3)	3/3	3/3
Mixed explicit stops (3/3)	3/3	0/3
Overall new-contract checks	8/8	2/8

The Qwen3.5 rows are an intentional legacy comparison: ordinary generation
still works, but its un-migrated scheduler does not yet satisfy the new typed
explicit-stop contract. They are not a claim that every legacy model fails in
all workloads.

Reproduction

Build and start each server independently. Qwen3.5 requires its feature-gated
Triton build environment; it does not support or require a
--gpu-memory-utilization CLI argument.

# Qwen3 stepped path
cargo run --release -p pegainfer-server -- \
  --model-path "$QWEN3_MODEL" \
  --served-model-name qwen3-adapted \
  --port 18081

# Qwen3.5 legacy control path (set PEGAINFER_TRITON_PYTHON if needed)
cargo run --release -p pegainfer-server --features qwen35 -- \
  --model-path "$QWEN35_MODEL" \
  --served-model-name qwen35-legacy \
  --port 18082

Then run the attached script (Python standard library only):

python3 pr865_qwen3_stop_contract_ab.py \
  --qwen3-url http://127.0.0.1:18081 \
  --qwen3-model qwen3-adapted \
  --qwen35-url http://127.0.0.1:18082 \
  --qwen35-model qwen35-legacy \
  --out stop-contract-ab.json

The script prints a compact comparison table and writes machine-readable JSON.
Use --stop-token-id ID to replace the full-vocabulary deterministic set with
a single known token when reproducing on a different prompt/model pair.

Follow-up

As requested during review of #865, this PR deliberately stops at the Qwen3
migration boundary. After the
maintainers confirm that the independent EOS/request-stop semantics are wanted,
the same policy propagation and resolver audit can be applied to Qwen3.5 and the
other legacy model lines in separate, model-scoped changes.

pr865_qwen3_stop_contract_ab.py

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 4a4d324ec8

ℹ️ 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".

[FeathBow]

Thanks. Please first rebase this PR onto the current main and handle CI error :)

[RicardoMin]

Hi @FeathBow, I have rebased this PR onto the latest main and addressed the CI error you mentioned. Could you please take another look when you have a chance? Thank you!

[FeathBow]

Thanks for separating the Qwen3 stop contract and for preserving the trigger token, logprob, completion count, EOS priority, and legacy fallback.

First, a successful hedged DFlash verify returns before terminal truncation. It selects and copies a winner, advances DFlash hidden context, updates acceptance accounting, and ticks the hedge controller from the untruncated span. The later executor truncation protects the final returned/KV length, but cannot undo those worker-side decisions. Apply each request's stop policy to every A/B candidate before winner comparison and add a hedge + mid-span terminal gate.

Second, nonzero min_tokens is silently accepted even though the new policy starts EOS/explicit-stop classification at token one. Restore the fail-early rejection until sampler-side masking exists; one assertion in the shared validator test is sufficient because both bridges invoke that validator before submission.

Please also revert the frontend architecture text that describes nonexistent ActiveRequest/StepEmitter code. Finally, normalize and share large stop sets: the current Vec::contains is linear in a per-token path and the full vector is deep-copied three times per speculative step. A shared sorted slice with binary search removes the request-size-linear scan and bulk copies; verify both the common zero/one-ID and full-vocabulary shapes before choosing anything more elaborate.

Keep the patch narrow: remove the unconstructed EosPolicy::Token branch, collapse the two identical resolver wrappers, and replace the one-off fake speculative executor/test with the missing production hedge gate. The existing bridge mapping and prefill/decode logprob tests cover distinct local contracts and should remain.

[github-actions]

All commit attributions previously flagged on this pull request are resolved.

[RicardoMin]

Summary

This update narrows the stop-contract change to Qwen3 while retaining the legacy frontend path for models that have not migrated yet.

The previous implementation could classify a terminal token after speculative candidate selection and KV-state preparation. When a stop token appeared in the middle of a verify span, the accepted suffix could therefore remain visible in the candidate state or be copied back before truncation.

This update applies terminal truncation before hedge winner selection and before speculative KV/state commit.

Stop Contract

Condition	finish_reason	Internal cause	Wire stop_reason
Model EOS	stop	Eos(token_id)	null
Explicit stop_token_ids match	stop	Token(token_id)	Actual token ID
Output limit reached	length	None	null

The contract preserves the triggering token, its logprob, and completion-token accounting. ignore_eos=true disables only model EOS termination; explicit request stop tokens remain active.

Review Fixes

• Classify and truncate every speculative candidate at its first terminal token.
• Select the hedge winner using the retained prefix length.
• Use the retained prefix for KV-page copy-back, hidden-state compaction, DFlash context, and counters.
• Keep an idempotent truncation check immediately before the final speculative commit.
• Fix the accepted-draft-token boundary so a terminal token inside the accepted draft prefix is counted correctly.
• Restore early rejection for non-zero min_tokens, since the current scheduler does not yet implement EOS/stop-token masking.
• Remove fake speculative scaffolding and duplicate helper tests.
• Keep legacy bridge behavior for un-migrated model lines.

Only Qwen3 consumes the typed stop policy in this change. The legacy frontend conversion path remains compatible with other models.

Verification

Test	Result
cargo fmt --all -- --check	Passed
git diff --check	Passed
cargo test --release -p pegainfer-frontend --lib -- --nocapture --test-threads=1	74 passed, 0 failed
cargo check --release -p pegainfer-qwen3 --lib	Passed
cargo check --release -p pegainfer-qwen3 --tests	Passed
Strict DSpark hedge gate with Qwen3-4B and dspark_qwen3_4b_block7	Passed
Qwen3 HTTP request with ignore_eos=true, max_tokens=8	finish_reason=length, completion_tokens=8
Qwen3 HTTP request with explicit stop token coverage	finish_reason=stop, numeric stop_reason, trigger token preserved

The change does not modify model mathematics, attention kernels, sampling kernels, or CUDA Graph shapes.

[RicardoMin]

Hi @FeathBow, your requested changes have been implemented in the latest commits.

Please take another look when you get a chance. Thank you for the detailed review!

[xiaguan]

Thank you for the careful revision. The production ordering is much closer to the required contract, and the min_tokens, stop-set lookup, and documentation changes address several parts of the previous review. I still do not think the new hedge test protects the worker-side bug, and there is one invariant leak in the new public type.

1. Please make the hedge regression fail against the previous implementation.

   dflash_hedged_midspan_stop_retains_trigger only asserts the final emitted tokens and Terminal. Before the latest worker-side fix, execute_speculative_verify_impl already called truncate_after_terminal after run_step and before RequestKv::apply_speculative, so those assertions still pass even if the worker ranks A/B candidates, copies KV/hidden state, and records DFlash context from the untruncated result.

   The test also derives its stop ID from a baseline run on the same hedged engine. Because the stop policy does not affect sampling, the old worker reproduces the same candidates and winner; the later executor truncation then produces exactly the expected external prefix. The parent gate requires each child to execute a hedge span, but total_wins > 0 is aggregated across all three children. The stop child is not required to have a B win, and it never demonstrates a case where terminal truncation changes the A/B ordering.

   Please add a real hedge case that observes the request-local raw and retained candidate lengths and the selected winner, with a fixture where truncation changes the winner decision. It should also verify that the retained winner is what feeds the KV/hidden context and committed/controller accounting. As a practical mutation check, this gate must fail when the new pre-selection truncation in try_execute_hedged_verify is reverted while the executor-side safety truncation remains.

2. Please keep the sorted stop-set invariant inside StopPolicy.

   StopPolicy::classify relies on binary_search, while both eos and token_ids are public. Any caller can therefore construct StopPolicy { token_ids: Arc::from([7, 3]), ... } and get a silently incorrect classification, bypassing the normalization promised by the type. The production callers only need new, default, and classify; please make the fields private and add a narrow read-only accessor only if a real caller needs one.

3. Please reduce the low-value tests and refresh the comments to match the final flow.

   normalizes_stop_sets_across_common_sizes repeats the one-ID behavior already covered above it and allocates an entire 151,936-ID set while primarily retesting sort_unstable, dedup, and binary_search. It does not prove that per-step clones share the allocation or that the production wire conversion preserves the contract. A focused dedup/membership test plus the meaningful wire/hedge gates is enough.

   Several comments now overstate or contradict the implementation: the hedge test claims it proves pre-winner selection and commit ordering although it only checks the final stream; the parent gate still says it runs “two” losslessness tests and “one child per lossless suite” after adding a third, non-losslessness child; and execute_speculative_verify_impl says the worker returns the mathematically accepted span even though the worker now applies the stop policy and the executor only rechecks the invariant. Please remove or update these while narrowing the tests.

The main implementation direction looks reasonable, but the current regression can stay green with the original worker-side bug restored, so I cannot approve this head yet. Thank you.