Repository navigation
perf(pd): P/D first-token latency and API-server GC fixes (from Kimi-K3 AgentX 1M) - #2786
wanzhenchn wants to merge 7 commits into
Conversation
🏷️ CI GuideRuns automatically on every eligible PR before approval:
Heavy model tests:
|
d6b810b to
5c807a8
Compare
5c807a8 to
f086c7e
Compare
There was a problem hiding this comment.
Thanks. I reviewed f086c7e by reading the code; I did not run anything. I found no correctness problems in the two leak fixes or in ATOM_KV_POLL_FRESH_WAIT_MS. The K3 mono Mooncake-consumer change also holds: DCP is already refused by common_refusals, and mla_landing_configured requires DCP > 1. Two of the three opt-in switches have problems, described below.
Must change
1. ATOM_PD_EARLY_T0 drops the first token under PP (model_engine/scheduler.py:3244). When the KV lands, _schedule_first_decode_after_remote_kv appends T0 to self._early_outputs and sets seq._injected_t0 = None, so postprocess no longer emits T0. The only reader of _early_outputs is EngineCore (engine_core.py:491, take_early_outputs). The PP loop's _pp_head_step (pp_engine_core.py:110) also calls schedule() but never collects the list. With PP > 1 P/D decode (kv_consumer, which config.py accepts) and ATOM_PD_EARLY_T0=1:
- every response is missing its first token;
_early_outputsgrows by one entry per request until the process exits.
The env doc says PP ignores the fresh poll, but it says nothing about early T0. Either drain the list in the PP loop or refuse / ignore the switch under PP, the same way the fresh poll is handled.
2. ATOM_PD_PROMPT_IDS_B64 checks the wrong condition (entrypoints/openai/api_server.py:905). The env doc says "Prefill only", but the branch is chosen by isinstance(kv_transfer_output_meta_info, dict), not by whether this instance is the producer. Consumer-side connectors set that output too: moriio does it regardless of role (moriio_connector.py:1821), and mooncake's main request_finished path does not check for a producer either. A deployment that uses the same env for P and D therefore also changes decode's responses: with return_token_ids=true, the top-level prompt_token_ids disappears and kv_transfer_params carries a base64 blob, about 3.6 MB at 670K tokens. Even on prefill, a client that explicitly asks for return_token_ids loses the top-level field, and so does any router or proxy that reads only that field. Please gate the switch on the producer role. I would also keep the top-level field whenever the client asked for it.
Worth addressing
3. With early T0, the engine's TTFT no longer matches what the client sees (scheduler.py:4432). T0 now reaches the client from schedule(), but seq.first_token_time and the tracing first_token are still recorded in postprocess of the first local decode, about two decode steps later. num_new_generation_tokens (:4444) also stops counting T0, so the status line undercounts by one per request. As a result, the engine's own metrics do not show the TTFT gain this switch is for.
4. The wrong field in the conflict error (entrypoints/openai/protocol.py:126-129). When prompt_token_ids_b64 disagrees with the top-level prompt_token_ids, the 400 still says "prompt_token_ids and kv_transfer_params['prompt_token_ids'] disagree", which names a field that request does not have.
5. keep_finished_outputs is a post-construction attribute flip (api_server.py:3088). The description says "CoreManager takes keep_finished_outputs", but CoreManager still defaults to queueing every finished Sequence, and the API server turns that off afterwards. The next host that delivers through callbacks and never reads get_output() will bring the leak back if it misses this line. Passing it to create_engine / CoreManager as a constructor argument would make the default explicit and match the description.
6. request_contexts is pushed every metrics interval (model_engine/engine_utility.py:848). collect_metrics copies up to 1024 context tuples every ATOM_METRICS_UPDATE_INTERVAL_S (1 s), pickles them, and sends them over ZMQ, whether anyone scrapes or not. Prometheus usually scrapes every 15–30 s, and the cost grows with the DP × PP scheduler count. Pushing only the entries added since the last push, or pulling them on scrape, would avoid that.
Smaller
resolve_prompt_token_idscan now returnarray('i'), butgenerate_async,generate_async_stream,generate_async_fanoutandsetup_streaming_requeststill annotateprompt_or_tokensasstr | list[int](api_server.py:775).reasoning.pyalready needed alist()because of this. The next+ [...]on that value would raiseTypeError.- The
call_func_with_aggregationdocstring still says "subsequent calls non-blockingly drain its per-rank replies" (model_engine/async_proc.py:924), but the newwait_sparameter makes the call blocking. _record_fresh_pollcreates its stats lazily withself.__dict__.setdefault(...)(engine_core.py:641) and rebuilds it as a literal after each log line (:661), and_poll_kv_transfer_progressreads_fresh_kv_poll_sthroughgetattr(..., 0.0)(:668) even though__init__assigns it. A small dataclass initialized in__init__would hold the counters in one place.
f086c7e to
a7b864b
Compare
|
Thanks for the careful read. All points are addressed in a7b864b except 6. Each fix is folded into the commit it belongs to, and the description is updated. 1. Early T0 under PP: fixed. 2. I did not keep the top-level field for every request that asks for 3. Engine TTFT with early T0: fixed. When T0 is queued, the scheduler now records 4. Conflict message: fixed. It now names the field it compared, 5. 6.
Smaller
Tests: each of the six commits passes the PR's tests on its own, and black is clean. Ruff reports nothing new. |
1671386 to
a5f0fc4
Compare
a5f0fc4 to
1236e95
Compare
|
Rebased onto #2871 moved K3 mono's CU-count and KV-connector refusals into the shared A Mooncake decode running decode CP stays refused, as it was before #2871 (as "decode CP" then, as "KV connector" now). Mono has never served one behind a connector. Two new tests cover that and the DSA indexer's refusal. The other five commits applied unchanged. Pre Checkin passed with 10,693 passed and 1,127 skipped (41370603841), against 10,656 and 1,127 on |
valar-cop
left a comment
There was a problem hiding this comment.
a few small things..
accuracy for the kimi-k3 mono decode behind mooncake P/D? only throughput reported
ci green.
GSM8K was measured on that path in the C1 cell of the #2787 verify run 41196318546: 0.9629 (±0.0052), strict 0.9629, at concurrency 1, against a gate of 0.94. That cell is a Mooncake P/D with a TP8EP8 prefill and a TP8 DSpark7 decode. It fails unless the decode logs That run carried this PR at a5f0fc4, before the rebase onto #2871, so it measured the mono main had then. #2871 has since taken K3 mono to 64 rows and added plain decode and ReplaySSM. |
|
ci triage for 4aef0f7:
|
4aef0f7 to
6608ac1
Compare
|
Rebased onto main 825b63b to clear the conflicts with #2862. Nothing in this PR changes.
@Yuechguo the force-push dismissed your approval. Could you take another look? |
f40ee40 to
3bf08d8
Compare
|
Rebased twice onto #2610 ( #2910 ( New commits.
Everything else is unchanged ( @Yuechguo this still needs your re-approval (the earlier rebase dismissed it). Could you take another look? |
|
@valarLip would you review again at the latest head? |
3bf08d8 to
5497095
Compare
|
Rebased onto #2912 and item 1. #2912 fixes the finished-sequence leak on Overlap with #2912 in The other five commits are unchanged ( @Yuechguo this still needs your re-approval, and the rebase drops item 1 in favour of #2912. Could you take another look? |
Mono refused every KV connector, since a connector's copy kernels can hold CUs the all-resident mono kernels wait on. A Mooncake kv_consumer queues no GPU work of its own for K3: its receives are NIC writes into the paged KV and state pools, and the only consumer kernels (the slot scatter, the DCP landing pool) are for DeepSeek-V4 and DCP decode. Serve that one transport; every other, or an offload tier beside it, is still refused. K3 mono now serves decode CP on its own, but a Mooncake decode running it stays refused: mono has never served one behind a connector. The K3 config passes Mooncake in resident_grid_refusals' `connectors` (the admit list #2610 added for Qwen's offload) only for such a decode, so the DSA indexer and every other mono model still refuse it. Co-authored-by: Cursor <cursoragent@cursor.com>
…sts wait for KV The completion poll an engine step reads was taken right after the previous forward, so a receive that landed during this forward is seen one step late. While requests are parked for remote KV, ATOM_KV_POLL_FRESH_WAIT_MS blocks up to that many ms after the forward for the poll it re-armed; the workers are idle then and answer at once. Co-authored-by: Cursor <cursoragent@cursor.com>
…_PD_PROMPT_IDS_B64)
A P/D prefill returns the prompt ids to the router, which forwards them
to decode. As a JSON int list a 670K-token prompt costs the decode's
event loop ~130 ms to parse and validate. With ATOM_PD_PROMPT_IDS_B64
the prefill answers a router's P/D prefill request (do_remote_decode)
with them in kv_transfer_params as base64 little-endian int32 (~15 ms);
the decode resolves them to an array("i") without boxing an id. Any
other request that asks for return_token_ids, a decode's included,
keeps the top-level list. Every decode accepts both forms.
Co-authored-by: Cursor <cursoragent@cursor.com>
…RLY_T0) A P/D decode receives the prefill's sampled first token with the KV but streams it with its first local decode output, which deferred output makes visible two forwards later. With ATOM_PD_EARLY_T0 the scheduler queues T0 as a stream output when the KV lands, and the engine loop, the PP head's included, puts it on the output queue ahead of the forward. TTFT, the first-token trace event and the generation-token count are taken at that point too. Co-authored-by: Cursor <cursoragent@cursor.com>
stop_profile returns only once the trace is exported and gzipped, and gzip's default level 9 dominates that. Under PP every stage waits for the next one's export, so on the Kimi-K3 PP8 prefill level 9 was about 4.5 s of each 6.5 s export. On a 265 MB PP8 stage trace, level 3 takes 1.3 s against 8.4 s, for a file about 25% larger. The level is passed to the existing inline gzip.open in _on_trace_ready; nothing else in the export changes. Co-authored-by: Cursor <cursoragent@cursor.com>
The comment said every loop that calls schedule() must drain the early outputs. RapidServe's prefill and decode steps call schedule() without draining, and rightly: their schedulers never admit a remote prefill's KV, the only path that queues an early T0. Say which steps must. Co-authored-by: Cursor <cursoragent@cursor.com>
The fresh poll runs while any request waits on a connector load -- a P/D receive, an offload or LMCache state load -- and on idle ticks as well as after a forward; the env comment and doc said only "parked for remote KV" and "after each forward". State why it and early T0 default to off: the poll can lengthen each such step by up to N ms, and both have been measured only on a Kimi-K3 P/D decode. Co-authored-by: Cursor <cursoragent@cursor.com>
5497095 to
5329537
Compare
|
Rebased onto The one conflict is in
|
Motivation
Serving Kimi-K3 1P1D on the AgentX 1M trace (88 concurrent agentic sessions, prompts up to ~670K tokens), per-request P/D timing and event-loop lag probes showed four costs on paths every P/D deployment uses:
mainhas since fixed both: fix(metrics): make per-request context gauges opt-in via CLI #2836 made those series opt-in, and perf(api-server): stop keeping finished sequences, keep prompt rendering and tokenizing off the event loop #2912 stopped queueing finished requests no one reads. This PR no longer carries either fix.The changes are independent of each other and of PP (they apply to
mainwithout #2365). The three latency changes alter hot-path behavior and have been measured only on a Kimi-K3 P/D decode, so each is an opt-in flag that defaults to today's behavior. The CI cells that enable them are in #2787.How it was measured. Kimi-K3 MXFP4, PP8 prefill → TP8×DCP8 DSpark4 decode, MI355X, C88, 3600 s window; each row adds one change to the row above. Runs land on different node pairs, which moves throughput by ±1–2%; the per-request probe timings are node-independent.
Changes
1. Finished sequences the OpenAI server never reads: dropped, now on
main(#2912)CoreManagerput every finished batch onoutputs_queueforget_output(), which the OpenAI and atomesh servers never read, so every request'sSequencestayed alive for the life of the process. This PR had the two servers construct the engine with the queue off. #2912 has since fixed the same leak onmainthe other way round:Config.keep_finished_outputsis off by default, and the queue's readers (LLMEngine.add_request, asyncio mode) turn it on. This PR now leaves it asmainhas it, and the item numbers below are unchanged. Measured here at C88, together with item 2's fix, it gave +1.4% throughput and 4–5% lower TTFT, and the API-server event loop was blocked for under 0.2% of the run.2. Per-request context gauges: dropped, now on
main(#2836)atom:{prefill,decode}_request_context_tokenskept one series per request ever served, and every scrape re-read them all (0.43 s of API-server Python after ~9.5K requests). This PR exported them for a bounded window. #2836 has since made both gauges opt-in onmain(--enable-request-context-metrics), so by default there are no per-request series. This PR now leaves them asmainhas them, and the item numbers below are unchanged.3. Poll KV completions again after the forward (
ATOM_KV_POLL_FRESH_WAIT_MS, default 0 = off)engine_core.py,async_proc.py,envs.py: +87/−10.4. Pass a prefill's prompt ids to decode as base64 int32 (
ATOM_PD_PROMPT_IDS_B64, default 0)kv_transfer_params.do_remote_decode), it packs the ids intokv_transfer_params.prompt_token_ids_b64, which ATOM's router forwards to decode as is. Every other request that asks forreturn_token_ids, a decode's included, keeps the top-level list. Every decode accepts both forms, so prefill and decode can be upgraded independently.protocol.py,api_server.py,reasoning.py,llm_engine.py,envs.py: +111/−30.5. Stream the injected T0 when the decode's KV lands (
ATOM_PD_EARLY_T0, default 0)max_tokens=1and the other edge cases are unit-tested.scheduler.py,engine_core.py,pp_engine_core.py,envs.py: +63/−4.6. Serve Kimi-K3 mono decode behind a Mooncake P/D decode
kv_consumer, for K3. Its receives are NIC writes into the paged KV and state pools, and K3 queues no consumer kernels: the slot scatter and the DCP landing pool exist only for DeepSeek-V4 and DCP decode. Every other connector, or an offload tier beside it, is still refused. So is a Mooncake decode running decode CP: K3 mono serves decode CP on its own since feat(mono): shared DSA indexer mono launch; Kimi-K3 mono to 64 rows, plain decode and ReplaySSM #2871, but has never served one behind a connector. The K3 config passes Mooncake in theconnectorsadmit list ofresident_grid_refusals(added by perf(qwen): integrate Mono decode and optimized TP1 MTP3 serving #2610) only for such a decode, so the DSA indexer and every other mono model still refuse it.models/kimi_k3/mono/config.py: +24/−1.7. Gzip torch traces at level 3
stop_profilereturns only once the trace is exported and gzipped, and gzip's default level 9 dominates that. Under PP every stage waits for the next one's export, so on the Kimi-K3 PP8 prefill level 9 was about 4.5 s of each 6.5 s export._on_trace_readypassescompresslevel=3to its existing inlinegzip.open; nothing else in the export changes, and it only runs while profiling.model_runner.py: +4/−2.Tests
tests/test_kv_fresh_poll.py(3),tests/entrypoints/test_pd_prompt_token_ids.pyandtests/entrypoints/test_protocol.py(4),tests/test_pd_early_t0.py(5),tests/models/test_kimi_k3_mono_config_refusal.py(6).maind2ab58f): 11,439 passed, 1,170 skipped and 3 xfailed (41991127058). Black and ruff pass there.Related
mainon its own; it does not depend on [feature] support pp>1 for Kimi-K3 #2365 or feat(kimi-k3): PP8 prefill for P/D serving — KDA state under PP, PP→DCP transfer, MLA staging #2709.