Skip to content

perf(pd): P/D first-token latency and API-server GC fixes (from Kimi-K3 AgentX 1M) - #2786

Open
wanzhenchn wants to merge 7 commits into
mainfrom
zwan/k3-pd-api
Open

wanzhenchn wants to merge 7 commits into
mainfrom
zwan/k3-pd-api

Conversation

@wanzhenchn

@wanzhenchn wanzhenchn commented Oct 28, 2026 •

Copy link
Copy Markdown
Contributor

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:

  1. A decode schedules a request one engine step after its KV lands.
  2. Prompt ids reach the decode as a JSON int list of up to 670K entries, parsed and validated on the API server's event loop.
  3. The decode streams the prefill's first token (T0) two forwards late.
  4. The API server never frees finished requests, and it kept a metric series for every request it served. Gen-2 GC passes over the growing heap blocked the event loop 45% (prefill) and 35% (decode) of the time after 80 minutes of load. main has 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 main without #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.

Run tok/s/GPU TTFT P50 / P90 / mean (s) Probe
baseline 18,052 4.62 / 14.67 / 7.32 KV landed → schedulable: 121 ms
+ fresh KV poll (3) 18,073 4.60 / 13.35 / 7.07 KV landed → schedulable: 44 ms
+ base64 prompt ids, early T0 (4, 5) 17,952 4.23 / 13.67 / 6.85 last chunk → first token: 3.56 → 3.29 s

Changes

1. Finished sequences the OpenAI server never reads: dropped, now on main (#2912)

CoreManager put every finished batch on outputs_queue for get_output(), which the OpenAI and atomesh servers never read, so every request's Sequence stayed 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 on main the other way round: Config.keep_finished_outputs is off by default, and the queue's readers (LLMEngine.add_request, asyncio mode) turn it on. This PR now leaves it as main has 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_tokens kept 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 on main (--enable-request-context-metrics), so by default there are no per-request series. This PR now leaves them as main has 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)

  • Problem. The completion poll an engine step reads is taken right after the previous forward, so a receive that lands during this forward is seen one step late. At C88 that is 121 ms on average (P90 170 ms) from every rank holding the KV to the request being schedulable, about one decode step.
  • Change. While any request waits on a connector load (a P/D receive, an offload or LMCache state load), the engine blocks for up to N ms at each completion poll, after a forward or on an idle tick, waiting on the poll it re-armed. The workers are idle at that point and answer at once. Each such step can take up to N ms longer. Fresh-poll counts and blocking time are logged every 30 s. PP engines ignore the flag. engine_core.py, async_proc.py, envs.py: +87/−10.
  • Benefit (N = 5). KV landed → schedulable: 121 → 44 ms. TTFT P90 −9% (14.67 → 13.35 s). Throughput is flat.

4. Pass a prefill's prompt ids to decode as base64 int32 (ATOM_PD_PROMPT_IDS_B64, default 0)

  • Problem. The prefill returns prompt ids as a top-level JSON int list (up to 670K ints). The router forwards it, and the decode's FastAPI parses and validates it on the event loop: ~130 ms per 670K-token request. The decode's loop lag tracked the prompt tokens arriving (r = 0.46, ~0.56 s per 1M tokens).
  • Change. When the prefill answers a router's P/D prefill request (kv_transfer_params.do_remote_decode), it packs the ids into kv_transfer_params.prompt_token_ids_b64, which ATOM's router forwards to decode as is. Every other request that asks for return_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.
  • Benefit. The decode's parse drops from ~130 ms to ~15 ms per 670K-token request. At equal elapsed time, the decode event loop's p99 lag is 35–55% lower and its blocked fraction drops by about 40% (at 40 min: 184 → 105 ms, 31.8% → 18.9%).

5. Stream the injected T0 when the decode's KV lands (ATOM_PD_EARLY_T0, default 0)

  • Problem. The decode receives the prefill's sampled first token with the KV, but streams it with its first local decode output. Deferred output surfaces that output two forwards later.
  • Change. When the KV lands, the scheduler queues T0 as a stream output and records TTFT, the first-token trace event and the generation-token count for it there. The engine steps whose scheduler can take a remote prefill's KV, the base step (shared by DP) and the PP head's, put T0 on the output queue ahead of the forward; postprocess does not send it again. RapidServe's steps never queue one. Stop checks never looked at T0 alone, so they are unchanged. max_tokens=1 and the other edge cases are unit-tested. scheduler.py, engine_core.py, pp_engine_core.py, envs.py: +63/−4.
  • Benefit. T0 leaves two decode steps earlier: 27 ms measured at low load, ~150 ms at full C88 load (about 78 ms per step). Together with 4 at C88: last chunk → first token 3.56 → 3.29 s, TTFT P50 −8% (4.60 → 4.23 s).

6. Serve Kimi-K3 mono decode behind a Mooncake P/D decode

7. Gzip torch traces at level 3

  • Problem. 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.
  • Change. _on_trace_ready passes compresslevel=3 to its existing inline gzip.open; nothing else in the export changes, and it only runs while profiling. model_runner.py: +4/−2.
  • Benefit. Re-gzipping one 265 MB PP8 stage trace offline, level 3 takes 1.3 s against 8.4 s at level 9, for a file about 25% larger.

Tests

  • Unit tests: tests/test_kv_fresh_poll.py (3), tests/entrypoints/test_pd_prompt_token_ids.py and tests/entrypoints/test_protocol.py (4), tests/test_pd_early_t0.py (5), tests/models/test_kimi_k3_mono_config_refusal.py (6).
  • Pre Checkin on 5329537 (rebased onto main d2ab58f): 11,439 passed, 1,170 skipped and 3 xfailed (41991127058). Black and ruff pass there.

Related

@github-actions

Copy link
Copy Markdown
Contributor

🏷️ CI Guide

Runs automatically on every eligible PR before approval:

  • ✅ Pre Checkin: Black, Ruff, catalog schema validation, non-GPU unit tests

Heavy model tests:

  • ✅ Run after the PR is approved and Pre Checkin passes
  • ✅ Run immediately when an approval review is submitted
  • ✅ Can be requested before approval with labels
Label Tests
ci:full Run all heavy PR model tests: native ATOM, vLLM, and SGLang
ci:atom Run native ATOM model accuracy tests
ci:vllm Run ATOM vLLM OOT model accuracy tests
ci:sglang Run ATOM SGLang model accuracy tests

Heavy jobs are skipped when the PR is not approved and no matching ci:* label is present.
Add labels via the sidebar or gh pr edit 2786 --add-label <label>

@valarLip valarLip left a comment •

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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_outputs grows 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_ids can now return array('i'), but generate_async, generate_async_stream, generate_async_fanout and setup_streaming_request still annotate prompt_or_tokens as str | list[int] (api_server.py:775). reasoning.py already needed a list() because of this. The next + [...] on that value would raise TypeError.
  • The call_func_with_aggregation docstring still says "subsequent calls non-blockingly drain its per-rank replies" (model_engine/async_proc.py:924), but the new wait_s parameter makes the call blocking.
  • _record_fresh_poll creates its stats lazily with self.__dict__.setdefault(...) (engine_core.py:641) and rebuilds it as a literal after each log line (:661), and _poll_kv_transfer_progress reads _fresh_kv_poll_s through getattr(..., 0.0) (:668) even though __init__ assigns it. A small dataclass initialized in __init__ would hold the counters in one place.

@wanzhenchn

Copy link
Copy Markdown
Contributor Author

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. EngineCore._put_early_outputs() now drains take_early_outputs(). Both _process_engine_step_inner and _pp_head_step call it right after take_rejected(), so T0 also goes out ahead of the forward on a PP decode. test_pp_head_loop_streams_early_t0 fails without the call.

2. ATOM_PD_PROMPT_IDS_B64 gate: fixed, gated per request. The b64 form is used only when the request carries kv_transfer_params.do_remote_decode. ATOM's router sets that on every P/D prefill request (AtomAdapter::inject_prefill_fields), and mooncake uses the same flag to decide whether to send. A decode request never carries it, so a decode's response keeps the top-level list and gets no blob, whichever connector filled kv_transfer_params_output. A client that calls the prefill directly also keeps the list. The test covers four cases: router prefill, prefill without metadata, decode, and direct client.

I did not keep the top-level field for every request that asks for return_token_ids. The router sets return_token_ids=true on each prefill request itself (same function), so a request with do_remote_decode is the router's. If both forms came back, carry_prompt_token_ids would copy the list into decode's kv_transfer_params, and decode would parse the 670K ints again, which is the cost this flag removes. The env doc now says to enable it only behind a router that forwards kv_transfer_params as is. A router that reads only the top-level list gets none, and its decode tokenizes the prompt again: correct, just slower.

3. Engine TTFT with early T0: fixed. When T0 is queued, the scheduler now records first_token_time, tracing.first_token and one generation token, and postprocess skips them, as it already did for the stream. The Prometheus atom:time_to_first_token_seconds is timed on the API server from the stream, so it already showed the gain. The engine-side TTFT, the trace span and the status line were what lagged.

4. Conflict message: fixed. It now names the field it compared, prompt_token_ids or prompt_token_ids_b64.

5. keep_finished_outputs: fixed, it is now a constructor argument. It goes create_engine() → LLMEngine → CoreManager / DisaggCoreManager → _init_shared_state, defaulting to True for offline generate() and rollout, which read the queue. Your concern was already real: the atomesh standalone server (atom/entrypoints/atomesh/server.py) also answers through callbacks and never reads get_output(), so it had the same leak. It now passes False as well. A rule test, written like the GC-policy one in test_gc_utils.py, fails any engine host that neither reads the queue nor turns it off.

6. request_contexts every interval: left as is, with numbers.

  • Worst case, the 1024-row cap: the push is 92 KiB. It takes 0.26 ms to pickle in the engine and 0.23 ms to unpickle on the API server's output thread, once per second per DP rank.
  • It does not grow with PP: non-head stages call push_metrics(scheduler_metrics=False).
  • The cap is reached only above about 8.5 requests/s per scheduler sustained (1024 rows / 120 s).
  • The rows are tuples of str and int, which the collector does not track (on 3.12, 0 of 1024 are tracked after unpickling). So they add none of the gen-2 work this commit is about.
  • Pulling on scrape is what push_metrics replaced (see its docstring). Pushing deltas would move the retention window, and its clock, to the API-server side, plus state shared across threads, to save at most about 0.5 ms/s. If you would still rather have deltas, I can do it as a follow-up.

Smaller

  • Annotations: prompt_or_tokens is now str | list[int] | array.array in generate_async, generate_async_fanout, setup_streaming_request and setup_streaming_request_fanout. There is no generate_async_stream. At runtime preprocess_fanout only checks isinstance(..., str), so the array was already handled.
  • The call_func_with_aggregation docstring now describes wait_s.
  • Fresh-poll counters: a FreshKVPollStats dataclass created in __init__, and _fresh_kv_poll_s is read directly. The two tests that build a bare EngineCore now set it, as __init__ does.

Tests: each of the six commits passes the PR's tests on its own, and black is clean. Ruff reports nothing new.

wanzhenchn added a commit that referenced this pull request Oct 30, 2026
wanzhenchn added a commit that referenced this pull request Oct 30, 2026
wanzhenchn added a commit that referenced this pull request Oct 30, 2026
wanzhenchn added a commit that referenced this pull request Oct 30, 2026
wanzhenchn added a commit that referenced this pull request Oct 30, 2026
wanzhenchn added a commit that referenced this pull request Oct 30, 2026
wanzhenchn added a commit that referenced this pull request Oct 30, 2026
wanzhenchn added a commit that referenced this pull request Oct 30, 2026
wanzhenchn added a commit that referenced this pull request Oct 31, 2026
wanzhenchn added a commit that referenced this pull request Oct 31, 2026
wanzhenchn added a commit that referenced this pull request Oct 31, 2026
wanzhenchn added a commit that referenced this pull request Oct 31, 2026
wanzhenchn added a commit that referenced this pull request Oct 31, 2026
wanzhenchn added a commit that referenced this pull request Oct 31, 2026
@wanzhenchn

Copy link
Copy Markdown
Contributor Author

Rebased onto main (99ebf63), which had come to conflict with item 6. The new head is 1236e95.

#2871 moved K3 mono's CU-count and KV-connector refusals into the shared resident_grid_refusals, and K3 mono now serves decode CP. Item 6 used to override the connector check in K3's own list. It now passes _nic_only_pd_consumer to resident_grid_refusals, so the other caller, the DSA indexer, still refuses every connector.

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 main at 99ebf63 (41355578710). The 37 more are this PR's tests.

@valar-cop valar-cop left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

a few small things..

accuracy for the kimi-k3 mono decode behind mooncake P/D? only throughput reported

ci green.

Comment thread atom/model_engine/engine_core.py Resolved
@wanzhenchn

Copy link
Copy Markdown
Contributor Author

accuracy for the kimi-k3 mono decode behind mooncake P/D? only throughput reported

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 Kimi-K3 mono decode on, and it fails if the decode ever logs Kimi-K3 mono decode off or declines the 8-row build, the width a concurrency-1 decode serves. So the eval's decode steps ran on mono.

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.

Yuechguo
Yuechguo previously approved these changes Nov 2, 2026
@valar-cop

Copy link
Copy Markdown

@wanzhenchn

Copy link
Copy Markdown
Contributor Author

Rebased onto main 825b63b to clear the conflicts with #2862. Nothing in this PR changes.

git range-diff shows only those context changes. This PR's test files and #2862's pass (3644 passed). The full suite's remaining failures (66 failed, 21 collection errors, CPU-only environment) are the same tests that fail on main.

@Yuechguo the force-push dismissed your approval. Could you take another look?

@wanzhenchn
wanzhenchn requested a review from Yuechguo November 3, 2026 09:07
@wanzhenchn
wanzhenchn force-pushed the zwan/k3-pd-api branch 2 times, most recently from f40ee40 to 3bf08d8 Compare November 4, 2026 07:26
@wanzhenchn

Copy link
Copy Markdown
Contributor Author

Rebased twice onto main, now at 3bf08d8 on 3ff3d68, with two follow-up commits from a re-review.

#2610 (atom/mono/runtime/deployment.py). #2610 gave resident_grid_refusals a connectors admit list for Qwen's offload; this PR had added a served_connector predicate to the same check. Two mechanisms for one check is one too many, so I dropped the predicate and kept #2610's list: deployment.py is now #2610's version unchanged, and the K3 config passes ("mooncake",) only for a Mooncake kv_consumer without decode CP (_served_connectors, which no longer needs the function-local factory import). The refusal now reads KV connector <name> as on main, so the K3 refusal tests assert that, and their multi case carries kv_role: kv_consumer so it fails if a multi is ever admitted. One narrowing: a multi whose only leaf is a Mooncake consumer is now refused; the predicate used to admit it.

#2910 (atom/entrypoints/openai/api_server.py). Import-list conflict only: both sides added one name from .protocol, and both are kept.

New commits.

  • b82fdfd: with keep_finished_outputs=False nothing is ever queued, so get_output() and get_output_async() now raise, naming the flag, instead of blocking forever. No host does this today. The test that the queue is filled only when kept used to scan the AST for a guarded put_nowait; it now runs the real output thread with the flag on and off, and fails with the guard removed. The host rule test counts .generate as a reader, so atom/examples is checked instead of skipped. The InputOutputProcessor prompt annotations admit the array("i") the API server passes since the base64 prompt ids.
  • 3bf08d8: the ATOM_KV_POLL_FRESH_WAIT_MS comment and doc now say the poll runs while any connector load is outstanding (a P/D receive, an offload or LMCache state load) and on idle ticks as well as after a forward. They also say why it and ATOM_PD_EARLY_T0 are off by default: each such step can take up to N ms longer, and both have been measured only on a Kimi-K3 P/D decode.

Everything else is unchanged (git range-diff shows context only). On a CPU host, tests/entrypoints and this PR's tests pass: 3,473 passed, 95 skipped, 3 xfailed. Pre Checkin passed on 3bf08d8: 11,297 passed, 1,158 skipped (41792340963).

@Yuechguo this still needs your re-approval (the earlier rebase dismissed it). Could you take another look?

@wanzhenchn

Copy link
Copy Markdown
Contributor Author

@valarLip would you review again at the latest head?

@wanzhenchn

Copy link
Copy Markdown
Contributor Author

Rebased onto main df61c5d, now at 5497095.

#2912 and item 1. #2912 fixes the finished-sequence leak on main with Config.keep_finished_outputs, off by default and turned on by the queue's readers (LLMEngine.add_request, asyncio mode). That is the leak this PR's 69b55fc fixed the other way round, by having the OpenAI and atomesh servers construct the engine with the queue off. Two mechanisms for one flag would be one too many, so I dropped 69b55fc and b82fdfd: the get_output() guard and its tests assume the flag is fixed at construction, which it no longer is. The description now lists item 1 as dropped, as item 2 is.

Overlap with #2912 in api_server.py. #2912 widened the prompt annotations to PrefixEncoding, and this PR widened them to array.array. Both are kept, in api_server.py and in InputOutputProcessor; the latter is folded into the base64 commit (844ee62), where the array prompts come from. generate_async's fallback prompt count, used only when there is no sequence, now counts a PrefixEncoding by its token_ids instead of taking len() of the tuple.

The other five commits are unchanged (git range-diff shows them identical). On a CPU host, tests/entrypoints, this PR's tests and the engine tests #2912 touched pass: 3,359 passed, 76 skipped, 3 xfailed. Ruff and black are clean. Pre Checkin passed on 5497095: 11,307 passed, 1,160 skipped, 3 xfailed (41968432794).

@Yuechguo this still needs your re-approval, and the rebase drops item 1 in favour of #2912. Could you take another look?

wanzhenchn and others added 7 commits November 5, 2026 07:56
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>
@wanzhenchn

Copy link
Copy Markdown
Contributor Author

Rebased onto main d2ab58f, now at 5329537, after #2856 (Chat token logprobs) made the PR conflict.

The one conflict is in generate_async in api_server.py, where #2856 changed the text line next to this PR's fallback prompt count. Both are kept: main's text (empty when logprobs are requested) and this PR's count, which handles str, PrefixEncoding, list and array.array prompts. Nothing else this PR changes needed adjusting:

  • feat(api): return Chat token logprobs with DSpark support #2856's single streaming setup still passes prompt_or_tokens and kv_transfer_params through, so the base64 prompt ids reach setup_streaming_request as before.
  • validate_logprob_options refuses logprobs with P/D transfer, and the base64 ids and early T0 (items 4 and 5) run only on P/D requests, so the two never meet. The early T0's RequestOutput takes the new logprobs field's None default.

git range-diff shows the other six commits unchanged. On a CPU host, tests/entrypoints, this PR's tests and the scheduler, postprocess and logprob tests #2856 touched pass: 3,691 passed, 72 skipped, 3 xfailed. Pre Checkin passed on 5329537: 11,439 passed, 1,170 skipped, 3 xfailed (41991127058).

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants