贡献 / #1887

#1887 Remote stage: conversation checkpoints across the link (--prompt-cache), stacked on #1748

open · @lask3802 · 0 评论 · 去 GitHub 看

Server & APIAMD / HIPNVIDIA / CUDAModels & quantsDocumentationLinux

说明

**Stacked on #1748.** The first four commits are #1748's, and this PR adds the last two: the feature, then its docs.
Both PRs are rebased onto `main` 61b3fb5 (engine 0.1.42).

## What it changes

With `--remote-stage`, every prompt was read from token 0. A conversation checkpoint holds only the running state of
the process's own layers, so the main process could not restore the worker's. A chat client sends the whole
conversation every turn, so every turn of a long chat waited for the whole prompt again: about 15 s at 16K tokens on
the test pair below.

With `--prompt-cache N` given **explicitly** on the main process, the worker keeps its layers' part of each checkpoint
under an id the main process hands it. The chain then works as on one GPU, at the prompt's synchronous cut points
(`turn_at`, `root_at`, `message_at`, `pin_at`):

- **Protocol 4 -> 5.** New messages `CkptSave` / `CkptRestore` carry an id, a position and the list of ids to keep,
  and `Reset` carries that list too.
- **Retention.** Only the main process decides what is kept; the worker prunes when an operation succeeds.
- **Limits at the hello.** The worker refuses more than 64 parts, or more parts than fit in half of its available RAM,
  or any part when it cannot read its RAM (as the engine's own RAM guard refuses parking then).
- **Disconnect.** The worker clears its store.
- **Errors.** A checkpoint the worker does not hold (it restarted or reconnected) makes the request read from token 0;
  any other error ends the engine as before.
- **A live continuation settles the worker first.** A request that continues the live session now waits for the
  worker's last Commit and for the adaptive tier's swaps in flight before it lends cache slots to the prompt
  (`hs.prefill`), as a Reset and a restore do. Without that, a loan could hand the prompt path slots that a swap was
  still writing, and the reply degenerated into `!`. Before this PR every request started with a Reset, so this never
  showed.

Design, scope and where it lives: `docs/remote-stage/CHECKPOINTS.md`.

## What it leaves alone

- **The default path.** Without `--prompt-cache`, a remote main keeps `prompt_cache` at 0: no new message is sent,
  the prompt is not split, and a Reset is the same header-only message. Checked on GPUs at exact flags (below): this
  build's answers agree with #1748's about as closely as #1748's own runs agree with one another, 8 of 8 in 2 of 6
  pairs. One-GPU and in-process splits see two mechanical refactors only (the
  eviction loop and the zero-and-reset at request start moved into lambdas, the same calls in the same order). On one
  GPU with the conversation cache, upstream `main`'s, #1748's and this branch's builds gave byte-identical answers.
- Mid-prompt checkpoints (`--prompt-cache-every`, `--prompt-cache-tail`) and parking (`--conversation-cache-mib`) stay
  off with a remote stage; session files are refused.
- The hello says protocol 5, so a main and a worker must both run this version, with or without checkpoints.
- A relay (three PCs) forwards the new messages, and it links its next worker again after each main process's hello,
  so the next worker checks the real ckpt_max. This was read from the code and unit-tested, not run on GPUs.

## Measured

An RTX 3080 main (layers 0-25, the head, the drafter) and an RTX 2080 Ti worker (layers 26-47) over 10 GbE, RVN
IQ3_S, Linux, engine 0.1.42 (after the rebase); #1748's rebased build is the base. The replay is two chats of four
turns on one shared system prompt of about 16,000 tokens, greedy. Exact flags: `--pcie-frac 0 --adapt-every 0
--adapt-swaps 0` on both PCs, `--mtp-max-t 1 --suffix-draft 0` on the main, and `STRATA_IQ_MT_MIN=1
STRATA_PREFILL_CPU_SHARE=0` on both.

| Check | Result |
|---|---|
| Default path (no `--prompt-cache`, exact flags) vs #1748 | Of the 6 pairs between two runs of this build and three of #1748's, 2 were 8 of 8 byte-identical and the rest agreed on 4-7 of 8 turns, each the same until one answer diverged mid-way. #1748's own three runs agreed pairwise on 5, 5 and 8 of 8. 0.1.42 does not repeat from run to run at these flags, even on one GPU: `main`'s own two runs there agreed on 1 of 8 turns |
| Every request sent twice, the next turn resuming through a checkpoint (exact flags) | 8 of 8 repeats byte-identical; with thinking on, 8 of 8; with each request first cancelled mid-decode, 8 of 8, and equal to the run without cancels. 53 restores on both PCs, 0 errors |
| Worker store emptied before each restore | the request reads from token 0 and finishes |
| Live continuations with the worker's adaptive tier on, 15 min, 100-2,500 new tokens a turn | 161 requests, 154 live continuations (98 reading 1,000+ tokens): 0 replies of `!`. Without the settle, on an earlier build with the same load: about 4 in 10 |

Time to the first token (served flags, thinking on):

| | #1748 (every prompt from 0) | This PR, `--prompt-cache 6` |
|---|---|---|
| First turn, 16,079 tokens | 14.68 s | 16.00 s |
| The 7 later turns | 14.57-15.06 s | **0.28-0.54 s** |

A turn read from 0 was 1.3 s slower with checkpoints on here (0.7-1.3 s on 0.1.41): the prompt is read in parts that
end at its cut points, and a checkpoint is saved at each. With `--prompt-cache`, answers need not be byte-identical to a run without it,
because the read order differs; the repeated-turn check is the exactness test.

The same checks on 0.1.41, before the rebase, gave the same results: 8 of 8 on each line, 53 restores, 163 requests
without a `!`, and later turns in 0.33-0.49 s instead of 14.2-14.9 s. That build also carried a guard for
non-finite logits (#879), which is proposed in its own PR, #1892; the 0.1.42 build did not.

## Builds

- **CUDA 13.0 on Linux** (GCC, sm_75 + sm_86): the engine, the CPU tests (`stage_ckpt_store_test`,
  `stage_link_test`, `conv_cache_test`, `conversation_cache_test`, `conversation_validation_test`,
  `conversation_transfer_test`, `draft_kv_plan_test`) and the serve tests (690, OK).
- **MSVC 2022 + CUDA 13.0:** the touched files compile, and the link, store, conversation-cache and draft_kv_plan
  tests pass.
- **HIP** (ROCm 7.0, gfx1100) and the **SYCL port** (oneAPI 2026.1.1): compile. SYCL has no remote stage.

本站相关内容

相关页面的快捷入口。