fix(scheduler): PD decode admission cap and remote-KV backpressur#1647
Open
Jasen2201 wants to merge 9 commits into
Open
fix(scheduler): PD decode admission cap and remote-KV backpressur#1647Jasen2201 wants to merge 9 commits into
Jasen2201 wants to merge 9 commits into
Conversation
Production replay payloads contain tool-call arguments with invalid JSON
escape sequences (\k, \p, broken \uXXXX from data redaction). The
existing _normalize_tool_call_arguments would silently leave these as
raw strings, causing Jinja chat templates to crash on .items().
Add _fix_invalid_json_escapes() to double lone backslashes before
invalid escape chars, and fall back to wrapping unparseable arguments
in {"_raw": ...} so the template never receives a plain string.
- Gate waiting→running promotion on len(running) < max_num_seqs to prevent unbounded growth that caused KV thrash livelock on decode consumers - Add ATOM_PD_MAX_INFLIGHT_LOADS env var to cap concurrent remote-KV park slots, preventing HSA resource exhaustion under burst traffic - Propagate prefix_cache_hit_tokens through kv_transfer_params so PD decode consumers inherit the producer's genuine cache-hit count - Throttle PD backpressure log to every 100 schedule ticks (~2s)
…ode to max_num_seqs
Contributor
🏷️ CI GuideRuns automatically on every eligible PR before approval:
Heavy model tests:
|
Jasen2201
changed the title
Contributor
There was a problem hiding this comment.
Pull request overview
This PR tightens scheduler admission/backpressure behavior for remote-KV (prefill/decode disaggregation) so “ready-for-first-decode” consumers don’t overflow max_num_seqs, adds a prefix_cache_hit_tokens metadata field to KV-transfer handoff, and hardens OpenAI tool-call argument normalization against invalid JSON escapes.
Changes:
- Add a scheduler regression test ensuring remote-KV decode promotion does not exceed
max_num_seqs. - Track “parked for remote KV” sequences in the scheduler and apply admission backpressure; propagate
prefix_cache_hit_tokensthrough KV-transfer connectors. - Add a fallback path for parsing model-generated tool-call arguments that contain invalid JSON escape sequences.
Reviewed changes
Copilot reviewed 6 out of 6 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
| tests/test_scheduler.py | Adds a regression test for remote-KV decode promotion respecting max_num_seqs. |
| atom/model_engine/sequence.py | Initializes prefix_cache_hit_tokens from incoming kv_transfer_params. |
| atom/model_engine/scheduler.py | Adds remote-KV parking accounting/backpressure and extra logging; adjusts remote-KV state transitions. |
| atom/kv_transfer/disaggregation/moriio/moriio_connector.py | Includes prefix_cache_hit_tokens in KV-transfer completion metadata. |
| atom/kv_transfer/disaggregation/mooncake/mooncake_connector.py | Includes prefix_cache_hit_tokens in KV-transfer completion metadata. |
| atom/entrypoints/openai/protocol.py | Adds invalid-escape repair + safer fallback when deserializing tool-call arguments JSON. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
890
to
894
| if seq.status == SequenceStatus.ABORTED: | ||
| self._uncount_inflight_load(seq) | ||
| seq.status = SequenceStatus.FINISHED | ||
| seq.leave_reason = "aborted" | ||
| self._rejected.append(seq) |
Comment on lines
+31
to
+60
| def _fix_invalid_json_escapes(s: str) -> str: | ||
| """Fix invalid JSON escapes in model-generated tool-call arguments. | ||
|
|
||
| Models occasionally produce invalid escape sequences like ``\\k`` or | ||
| ``\\p`` in function.arguments JSON. ``json.loads`` rejects these. This | ||
| helper doubles any backslash not followed by a valid JSON escape char. | ||
| """ | ||
| _VALID = frozenset('"\\bfnrtu/') | ||
| out: list[str] = [] | ||
| i = 0 | ||
| while i < len(s): | ||
| if s[i] == "\\": | ||
| if i + 1 >= len(s): | ||
| out.append("\\\\") | ||
| i += 1 | ||
| elif s[i + 1] == "\\": | ||
| out.append("\\\\") | ||
| i += 2 | ||
| elif s[i + 1] in _VALID: | ||
| out.append("\\") | ||
| out.append(s[i + 1]) | ||
| i += 2 | ||
| else: | ||
| out.append("\\\\") | ||
| out.append(s[i + 1]) | ||
| i += 2 | ||
| else: | ||
| out.append(s[i]) | ||
| i += 1 | ||
| return "".join(out) |
valarLip
approved these changes
Jul 22, 2026
Comment on lines
1302
to
1307
| seq.status = SequenceStatus.WAITING | ||
| if self._connector_flag("is_offload"): | ||
| self._mark_offload_load_ready(seq) | ||
| return False | ||
| self._uncount_inflight_load(seq) | ||
| return True |
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
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.
Motivation
Technical Details
Test Plan
Test Result
Submission Checklist