feat: restore linear-state slots in base step_for_schedule_overlap. - #2295
Draft
yingxudeng wants to merge 3 commits into
Draft
feat: restore linear-state slots in base step_for_schedule_overlap.#2295yingxudeng wants to merge 3 commits into
yingxudeng wants to merge 3 commits into
Conversation
VLMWorkerImpl does not override step_for_schedule_overlap, so the base version was used when a VLM backend ran with enable_schedule_overlap and enable_prefix_cache both on. The base version only invoked step() and skipped the linear-state slot restore, while prepare_work_before_execute already gated its own restore on !enable_schedule_overlap(). The net effect was that every chunked-prefill boundary read a stale recurrent slot on VLM frontends (e.g. GLM-5.3-Flash), which produced non- deterministic accuracy regressions with symptoms matching KV / linear state cross-request contamination. Move the restore into the base step_for_schedule_overlap and gate it by owns_recurrent_cache so speculative/MTP outer workers -- which own no kv_caches_ -- fall through unchanged. LLMWorkerImpl keeps its override because it also drops the default-stream sync via execute_no_sync_on_stream, which is an LLM-only optimization. Verified on GLM-5.3-Flash 45L with a 5x5 rigorous GSM8K run (serial batch=1, temperature=0, MAX_CHUNK_TOKENS=128): Group run1 run2 run3 run4 run5 mean --------------- ---- ---- ---- ---- ---- ---- prefix OFF 0.9 0.9 0.9 1.0 0.9 0.92 prefix ON (fix) 0.9 1.0 0.9 1.0 0.8 0.92 prefix ON (pre) 0.4 0.7 0.9 0.8 0.6 0.68 Cache hits confirmed unchanged (input_tokens_cache_read=512 per request).
… ordering. Extract has_linear_attention_layers + owns_recurrent_cache + restore_linear_state_slots into WorkerImpl::try_restore_linear_state_slots and replace both callsites (base step_for_schedule_overlap and prepare_work_before_execute_on_stream) so the invariants stay in sync. Widen the compute_stream_ StreamGuard in the base overlap path to cover both the restore and the following step() so a future pooled compute_stream_ cannot desync chunk-boundary restores from the forward. Change step_for_schedule_overlap signature from const ForwardInput& to ForwardInput& (base + LLMWorkerImpl override) so the write-through of linear_state_validity_mask is honest instead of const_cast. Turn the compute_stream_ nullptr check into a CHECK so a future regression cannot silently revert to the stale-slot bug.
…event. Base step_for_schedule_overlap only enters compute_stream_ when the model has linear-attention layers, so dense frontends inherit the historical stream selection unchanged. Inside the guard, wait on input.metadata_ready_event before restore or step() so the prepare_stream_ H2D copies are observable on compute_stream_ regardless of which stream published them. Mirror the same has_linear_attention_layers gate in LLMWorkerImpl::step_for_schedule_overlap. Update the MLU overlap test helper to take ForwardInput& so it matches the tightened base signature.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Related Issues
Change Type
Pull Request Checklist
Thank you for contributing to xLLM. Before requesting review, please make sure the following items are complete.
PR Title and Commit Messages
<type>: <subject>.Pre-commit Checks
pre-commitby runningpip install pre-commitor an equivalent command.pre-commit install.pre-commit run --all-filesand fixed any reported issues.Self Review
.agents/skills/code-review/references/custom-code-style.md, especially code written or assisted by AI.mainbranch.Build and Test Coverage
python setup.py build testhas passed on a CUDA machine.python setup.py build testhas passed on an NPU machine.python setup.py build testhas passed on an MLU machine.Reviewer Notes