Skip to content

fix(server): keep chat prefill alive with protocol chunks - #865

Closed
Avicennasis wants to merge 2 commits into
FlashML-org:mainfrom
Avicennasis:rm52587-chat-stream-heartbeat
Closed

Avicennasis wants to merge 2 commits into
FlashML-org:mainfrom
Avicennasis:rm52587-chat-stream-heartbeat

Conversation

@Avicennasis

@Avicennasis Avicennasis commented Nov 5, 2026 •

Copy link
Copy Markdown

A queued or prefilling Chat Completions request currently sends its role chunk and then stays silent until generation produces output. This reuses with_keepalive with the same 15-second interval policy as Anthropic and Responses, sending chat.completion.chunk data frames with choices: [{"index": 0, "delta": {}, "finish_reason": null}]. Event-based idle timers can count these frames. Heartbeats stop after the first nonempty content/reasoning delta or tool output. A heartbeat is a chunk like any other after the role chunk: with continuous_usage_stats it carries the usage so far, and with timings_per_token the timings of the ack behind it, so none while queued or prefilling. All chat chunks, including the optional usage chunk, retain the same id, model and creation time; finish reasons, timings and [DONE] retain their existing behavior. Non-streaming paths are untouched.

Idea from Wallawalla47/Infernix commit 94953a26 (UE-05), reviewed against donor head 048d484da86144ad10f27fcabe218d528530c9de. Both projects use Apache-2.0. No donor code was copied, so no third-party notice was added.

Verified on Linux x86_64, Python 3.12, torch 2.14.1+cpu, OpenAI SDK 2.54.0, from base f69d50686d4f9bba578f7ac91135c609bb9f9e90. No GPU, driver or model checkpoint was used. The two new event-gated regressions sit immediately beside the existing shared keepalive test. Removing prefill heartbeats, sending heartbeats after content, changing their wire shape, drifting metadata, changing terminal/usage frames, or dropping the running usage or timings from a heartbeat that should carry them can turn them red.

On the unchanged implementation, the focused command failed because no case received a heartbeat before its one-second guard: 4 failed, 145 deselected (exit 1). After implementation: 4 passed, 145 deselected (exit 0). From the repository root:

export PYTHONPATH=python
python -m pytest tests/server/test_streaming_model_matrix.py -k test_chat_keepalive -q
python -m pytest tests/server/ -q

Final server-subsystem gate: 692 passed, 0 failed, 0 skipped (exit 0); main alone gives 688. The full suite gives 20 failed, 1721 passed, 977 skipped against 20 failed, 1717 passed, 977 skipped on main in the same environment: the same twenty test ids, which need a compiled extension or flashinfer, or are parametrised over the local model cache. A local probe (not part of this change) also passed one real localhost HTTP request through the registered route: two SDK-decoded prefill heartbeats, none during the later decode pause, unchanged length finish reason, usage total 4, and termination on [DONE]. A second probe passed six adapter checks for content/reasoning/tool suppression, closing streams during prefill and after output, and errors after a heartbeat.

NOT VERIFIED: CUDA, a real model or checkpoint, real client idle-timeout behavior, production-length 15-second/multi-minute prefill timing, or a combined checkout with other pending PRs. Some declared runtime dependencies were missing from the test environment and were installed locally; no package/dependency changes are in this diff.

Based directly on origin/main f69d50686d4f9bba578f7ac91135c609bb9f9e90. File-list checks show no overlap with the open #334, #459, #565, #592, #593 or draft #656; none is a stack prerequisite. Open upstream #572 adds SSE comments and disconnect polling in api_server.py; this change supplies protocol data events in openai_api.py, so the paths are separate. Their combined runtime behavior has not been tested. The companion connection-level keep-alive window change is separate and not included here.

Question for the maintainer: is prefill-only heartbeat suppression also the preferred policy for reasoning and tool output? This implementation treats either as generation having started, consistent with stopping once tokens are flowing.

A keep-alive is a chat.completion.chunk after the role chunk, so it goes
through the same emit() as every other one: with continuous_usage_stats it
carries the usage so far, and with timings_per_token the snapshot of the ack
behind it. Before the first ack nothing is measured, so no timings block.
@Avicennasis
Avicennasis force-pushed the rm52587-chat-stream-heartbeat branch from 5ccef3e to 8d6b772 Compare November 13, 2026 06:41
@Avicennasis

Copy link
Copy Markdown
Author

Restacked onto main at f69d506, 16 commits past the fork point (628b88f). main has moved five commits since (to 58b3ad7); none of them touches this PR's two files and git merge-tree of this head against 58b3ad7 is clean, but the numbers below were measured on f69d506. In the restack itself git merge-tree named one conflicting file, openai_api.py, and in it one region: the event loop of stream_chat_completion_chunks, which this PR re-indents under try/finally and which #868 (continuous_usage_stats) and #887 (timings_per_token) edit line by line. Four conflict hunks, all resolved the same way, as this PR's shape of the loop with upstream's edits re-applied onto it: generate_events(..., usage=usage) now sits inside with_keepalive(...); the six delta sites and the finish chunk go through #868's emit(...) instead of _sse(...), each still passing created; GenDone no longer copies the prompt/completion/cached counters, since GenUsage carries them; the error chunk stays on _sse, as on main. The rest merged on its own: the imports, created on the role chunk and on the final usage chunk beside _reported_usage(state, usage), _chat_chunk's signature, and the test file (#871 added three lines elsewhere in it). git range-diff of the old commit against the first new one shows only those edits.

The second commit is new and is the part worth a look. A keep-alive is a chat.completion.chunk after the role chunk, and main now says two things about those: with continuous_usage_stats, "the chunks after the role chunk also carry the usage so far" (api_models.py), and with timings_per_token the block rides "every chunk that follows an ack" (docs/cli.md), while the role chunk "precedes the first ack: nothing measured, never a block" (test_stream_timings_per_token_rides_every_chat_content_chunk). So a keep-alive goes through the same emit() as every other chunk instead of being exempt. With continuous_usage_stats it carries the usage so far, which is 0/0/0 while the request is queued or prefilling. With timings_per_token it carries the snapshot of the ack behind it: none before the first ack, and the latest one if the stream goes quiet after an ack that produced no delta (text a parser is still holding back). Without those options it is the bare chunk it was. test_chat_keepalive_carries_running_usage_and_timings_when_requested[True|False] holds the three cases; it fails on the mechanical resolution (KeyError: 'usage' on the first keep-alive) and passes with the one-line change. ShellClient, the in-tree reader of continuous usage, yields the same deltas, running totals and TurnDone through two such keep-alives (a local probe over real HTTP, not part of the diff). If you would rather keep keep-alives bare, that is the one line back to _sse and the [True] case inverted. docs/cli.md's list of timings carriers does not name keep-alives; say if you want that sentence extended.

Re-ran on a CPU-only Ubuntu 24.04 x86_64 box (AMD Ryzen 7 5800X, 16 threads, 121 GiB, no GPU; Python 3.12.3, torch 2.14.1+cpu, triton 3.8.0, transformers 5.18.0, flashlib 0.4.0, numpy 2.5.4, safetensors 0.8.0, pytest 9.2.0, OpenAI SDK 2.54.0; PYTHONPATH=python, no C++ extension, no flashinfer; huggingface_hub 2.4.0 and pillow 13.0.0 in this venv sit above the declared ranges, the same venv for both trees):

  • This PR's tests (-k test_chat_keepalive in tests/server/test_streaming_model_matrix.py): 4 passed, 145 deselected; the same four fail against main's openai_api.py (4 failed, each waiting for a keep-alive that never comes). The whole file: 149 passed, main 145.
  • tests/server: 692 passed against main's 688, no failures or skips on either.
  • Full pytest -q tests from the repo root: this branch 20 failed / 1721 passed / 977 skipped (2718 collected, 1:47), against main at f69d506 in the same venv 20 failed / 1717 passed / 977 skipped (2714 collected, 1:53). Identical failure sets by test id and identical failure messages (+4 passed is this PR's tests), and the same twenty ids as on 9cacd08, all environment-bound: the eleven tests/moe/test_cpu_moe_fp8_block.py cases and tests/kernels/test_swiglu_clamp.py::test_cpu_extension_supports_swiglu_clamp need the _cpu_moe C++ extension; tests/engine/test_cache_budget.py x2 and tests/moe/test_offload.py x1 need flashinfer; tests/models/deepseek_v41/test_engram.py::test_ftw_side_files_carry_the_tables_and_reopen needs the _row_store extension; and the four tests/models/test_quant_config.py::test_scheme_for_agrees_with_the_stored_tensors cases are parametrised over this box's HF cache (two BERTs: assert 0 > 0; two MLX quantisations: quantization method None is not supported).

Not exercised here: CUDA, a real model, a real client's idle timer, and the 15-second interval at production length (the tests and probes run a 20 ms interval against a stub engine). #572's SSE comments were not combined with this branch.

One thing changed upstream of the idea since this was filed: the donor project moved its Chat Completions keep-alive back to an SSE comment (Wallawalla47/Infernix 05f89c4a). There the role chunk used to leave at admission, and it wanted the first JSON chunk to mean the first real event, so tools reading time to first token from the first chunk stop reading zero. Here the role chunk already leaves before prefill on main and this PR does not move it; it adds empty-delta chunks only when a prefill outlasts the 15-second interval, and none once output has started. #572 is the comment-based route through the same silence, for every stream at once. If you prefer that route for chat, this PR can close; if event-counting idle timers matter on the chat route as they did for response.in_progress, this is the chat counterpart.

@Avicennasis

Copy link
Copy Markdown
Author

Closing this draft. #572 covers the same silence with SSE comments on every stream at once, which is also where the donor project ended up. The branch stays on my fork in case the chat-side keep-alive chunks are wanted later.

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

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant