Skip to content

Stream RNGD layer-group KV into reserved decoder cache before prefill completes #265

Open

Description

@gspark-etri

Problem

The verified GPU chunk-copy optimization (#254) and ATOM next-request prefetch (#255) do not stream KV from the currently executing RNGD prefill into its decoder. SubprocessPrefillWorker currently exposes only a whole-request Future. This leaves the original within-request overlap requirement in #47 incomplete.

Evidence

The public research FullPrefillEngine.prefill_stream(tokens, sink) emits each layer group on the host, but the generic worker253/factory adapters return all layers before copying into the shared-memory slab and sending the completion message. GPU/ATOM owners admit KV after that completion. Existing full-source slabs provide bounded ownership, which should be preserved. RNGD runtime calls are not assumed reentrant: overlap must be demonstrated between producer compute and decoder-side import, not inferred from Python threads.

Hypothesis and scope

An opt-in validated layer-progress message can expose completed shared-memory spans before final logits arrive. Import only into cache storage reserved for that request; publish a reusable cache entry and start decoding only after complete coverage, final token readiness, and transfer completion. Keep legacy all-at-once behavior. Support dense model-layer ordering and grouped hybrid caches explicitly. Validate model layer/position geometry, coverage, IDs, and slot lifetime before mutation. Cancellation or producer death must drain transfers before freeing/reusing storage. ATOM's next-request prefetch must not overwrite the active single-lane decode cache; defer import when storage is occupied.

First establish a portable protocol and GPU owner path, then add source adapters and ATOM admission under the same ownership rules. Benchmark all four requested model families where supported. Do not change active benchmark environments or shared vendor SDKs.

Done when

  • Malformed/overlapping/incomplete progress cannot publish a partial cache or corrupt other requests.
  • Cancellation, producer failure, and slot reuse pass regression and real-device checks.
  • Source/decoder timelines show actual within-request progress/import overlap; no kernel-overlap claim from host intervals alone.
  • Same-route whole-transfer versus streamed output is exact, including hybrid Gemma.
  • Long-context TTFT/transfer-tail A/B and the 0.1–4GB transfer-size curve from [PD serving] Layer-wise KV transfer overlapped with prefill, and large-KV handling #47 are recorded, including regressions; keep opt-in unless benefit and compatibility are established.

Owner / related

Codex; maintainer reviews/merges. Refs #47, #201, #253, #254, #255, #258, #259. This issue does not replace pending four-model held-out tests in #263.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions