Skip to content

feat: stream RNGD KV layers into Torch and ATOM decode - #266

Closed
gspark-etri wants to merge 29 commits into
perf/47-chunk-copyfrom
perf/265-rngd-stream
Closed

gspark-etri wants to merge 29 commits into
perf/47-chunk-copyfrom
perf/265-rngd-stream

Conversation

@gspark-etri

@gspark-etri gspark-etri commented Oct 23, 2026 •

Copy link
Copy Markdown
Collaborator

Background

The subprocess PD path previously exposed KV only after whole-prefill completion. This draft connects incremental source export and Torch/ATOM decoder import while keeping incomplete cache state private. Whole-result transfer remains the default and the same-engine control.

Changes

  • Validate disjoint, complete layer/token coverage. Import dense spans into grouped/hybrid paged caches without exposing incomplete lengths or overwriting an admitted prefix.
  • Negotiate source capability, copy BF16 spans to SHM before notification, and wake the decoder owner for incremental import. Fence CUDA copies before cancellation can reuse pages. Reject malformed response IDs and fail invalid capacity through the request Future.
  • Add explicit trusted checkpoint factories and streaming adapters for Qwen3-0.6B, Qwen3-4B, Gemma3-1B and partitioned KONI-8B. Preserve their original whole-result methods, padding, attention inputs and output heads. Normalize server input to a token list; real integration exposed the prior tuple/list mismatch.
  • Expose --pd-prefill-stream-layers for subprocess producers with Torch or ATOM decoders. Unsupported decoder owners are rejected.

Includes the dense-hybrid prerequisite from #259 and B1 prefetch from #255 on #254's branch. Experimental model factories depend on the existing serving_matrix/qwen3_rngd modules at commit04dab7c and the measured rngd_pd snapshot; paths and source hashes are documented in STREAMING.md/evidence265.

Verification

  • Source/factory/adapter and real-pipe tests:27 passed after the tuple/list regression was first reproduced. Cancellation, slot reuse, incomplete coverage and active-decode isolation tests are also retained.

  • Actual RNGD source:24/24 conditions have exact K/V and logits across four models. Short/full-width inputs and full/half/empty suffixes; width128 for Qwen/Gemma, width512 for KONI.

  • Actual RNGD→GPU:8/8 generated-token pairs exact, two concurrent requests/model, fixed8 output tokens, page64, caches off. Stream imports56/72/52/64 spans respectively versus0 for whole mode. All batches drain and owned processes exit.

  • Actual long RNGD→GPU, runtime171e8c4: compiled width8192, mixed2K/8K inputs, concurrency2, fixed32 output, three alternating batches/mode after warmups. Median completion reductions: Qwen0.6 28.2%, Qwen4 22.1%, Gemma4.3%, KONI22.9%. All64 sequences including warmups exact;732 span imports completed before source completion. Published auditor reproduces raw-record medians and checks output/timing/trace invariants.

  • Failed legacy KONI initialization (missing native head dependency) and first GPU tuple/list failure are retained. Validated partitioned/host-head factory replaces that legacy path for the actual KONI test.

  • Long source:24/24 additional conditions at2K/8K have exact K/V and logits.

  • Actual RNGD→ATOM: four models × four whole/stream and prefetch off/on modes,16batches/32sequences exact against the same-route whole reference, short/512 inputs, fixed8 output. All drain and imports preserve B1 ownership. Owned Pod deleted and absence independently verified.

  • ATOM per-layer importer validates selected destinations, defers multi-lane hybrid import until final admission, and prevents failed/partial prefills publishing prefixes. Cancellation/reuse and SDK failure are covered by owner tests.

  • Actual streaming GPU faults, runtime33d6cb0: allfour models pass partial-import cancellation/reuse and source death with an active512-token survivor. Kill points5/6/4/10; complete survivor outputs equal same-route no-fault references. All child exits0; local devices idle. Auditor rechecks recorded assertions/triggers and process exits. Full512-token outputs are not retained; equality assertions are in the executed driver. Initial harness-only AttributeError/interrupted run retained.

  • Actual long RNGD→ATOM: allfour64batches/128sequences exact. Compiledwidth8192, B1/two queued2K/8K prompts, fixed32outputs, prefixcacheoff,3measured batches/mode after warmups. Whole-vs-stream completion reductions35.66/21.26/9.79/19.19% without prefetch and39.25/22.90/11.21/16.84% with prefetch (Qwen0.6/Qwen4/Gemma/KONI). Auditor verifies ordering, outputs, timings and B1 import ownership; timing Pod deletion confirmed. Actual ATOM streaming fault matrix now passes allfour models/12cases; see the fault evidence below.

  • Four-model GPU notification control:48batches/96sequences exact,3 measured batches/mode after warmups. Early-vs-deferred reductions4.64/2.99/0.40/1.72%; deferred-vs-whole22.74/19.62/2.84/21.25%. Most Qwen/KONI gain remains when notification is delayed until source logits are ready. This does not isolate source assembly from destination per-layer import, permits cross-request overlap, and is not proof of device DMA overlap. Gemma sample ranges overlap. Exact test-only control and independent auditor included; production unchanged.

  • Higher-load GPU c4/fixed128 extension: allfour32batches/128sequences exact,2K/8K/2K/8K inputs,3measured batches/mode. Whole-vs-stream completion reductions20.72/16.44/2.26/19.34% (Qwen0.6/Qwen4/Gemma/KONI). Independent auditor validates raw outputs, ordering, arrival times/span traces and reproduces summaries. All exits0; local devices idle. STREAMING.md includes a review map and explicitly scopes this same-host research-factory SHM path.

  • Added partial-prefetch error/cancel regressions after hardware fault work:39 related tests passed in5.00s; production code unchanged.

  • Final full regression (before these two added tests):626 passed,15 skipped,253 subtests passed,27 warnings,168.94s. Exact drivers, raw JSON/logs and source hashes are in evidence265. Scoped sensitive scan and whitespace checks run before push.

  • Actual ATOM streaming faults (evidenceb59e7fc, runtime33d6cb0): allfour models pass partial-import cancellation/reuse, partial-import producer death, and prefetched producer death preserving active512-token decode. Full retained16/512-token arrays independently audited; survivor kill points13/6/10/5 (Qwen0.6/Qwen4/Gemma/KONI). All child/copy exits0; queue exit0 and owned Pod absence independently verified. Reproduction and raw evidence: opt_prime/examples/pd/http_matrix/experiments/2026-10-24-stream265-atom-faults.md.

  • Higher queue pressure (evidence515a404): allfour models complete32batches/512sequences exact atc16 queued behind4execution/source slots, mixed2K/8K and fixed128output. Whole-to-stream completion reductions25.20/20.02/2.33/20.91%,3measured batches/mode. Full auditor and malformed-evidence self-check pass on published copies; queue0/devices idle verified. See opt_prime/examples/pd/http_matrix/experiments/2026-10-24-stream265-c16.md. HTTPc16 runs are separate and not included in these results.

  • Completed GPU HTTPc16 matrix (evidence897373c): allfour models x native/whole/stream,96conditions,96warmups,1536retained-token checks and1536timed SSE text checks; all12child exits0 and queue41030 exit0. Whole/stream references exact8/8 permodel. Native/whole exact Qwen0.6=8/8,Qwen4=4/8,Gemma=4/8(throughEOS8/8),KONI=5/8; preEOS Qwen4/KONI differences and earlier quality failures remain. Observed single-batch throughput gains: Qwen4+22.1–27.2%,Gemma+1.51–3.17%,KONI+20.94–33.93%; these are not statistical or equal-resource speedup claims. See opt_prime/examples/pd/http_matrix/experiments/2026-10-24-stream265-gpu-http-complete.md. Subsequent ATOM coverage is recorded below.

  • ATOM HTTP follow-up (evidence 7646a31): Qwen3-0.6B, Qwen3-4B and Gemma3-1B whole/stream controls pass 48 conditions, 768 retained-token checks and 768 timed SSE text checks; all eight reference arrays per model match. At 8K/128, whole→stream throughput is 33.07→43.35, 12.67→14.53 and 39.49→43.57 tok/s respectively. Gemma remains slower than native ATOM (57.11 tok/s). KONI native hits the original 180-second response wait at 8K/128/c16; server alive, no OOM/listener drops, zero timed batches. A separate 600-second native diagnostic completed (evidence f1d9117): 264 responses across eight conditions pass, all original reference arrays match, OOM/listen-drop deltas remain zero. At 8K/128, max first SSE is 253.93s and max completion 263.08s, with six timed responses exceeding 180s. This retains the original deadline failure; KONI whole also completed (evidence cb25ff6): 264 response checks pass, 8K/128 throughput is 9.31 tok/s versus native 7.78, but shorter inputs remain slower. Native/whole references match only 6/8, including EOS-aware comparison; timing uses unequal device counts and one batch per condition. Streaming subsequently completed (evidence eaf7daa): 264 responses audited, whole/stream references exact8/8. At8K/128 stream10.83tok/s versus whole9.31 and native7.78; median completion100.70s, maximum189.11s. The original180s failure remains unresolved. All children/copies/outer queue exit0 and owned Pod absence verified. Stream used another physical ATOM card; single-batch timing and unequal native/PD resources remain limitations. All four models now have completed B1/c16 ATOM whole/stream controls. Raw three-route KONI JSON and independent summary are published. The HTTP listener candidate and width-1024 numerical experiments are not integrated into this PR. See opt_prime/examples/pd/http_matrix/experiments/2026-10-24-stream265-atom-http-progress.md for resource differences, native parity limits, source hashes and remaining work.

  • Native GPU coverage follow-up (03054a1, runtime eaf7daa): four models × input512/2048 × output32/128 × concurrency1/4,32conditions/176responses. All80retained-token and80timed SSE comparisons pass; all16reference arrays match the earlier c16 native references. Raw arrays/counts/timing summaries independently audited. Initial coordinator PYTHONPATH omission failed before inference and is retained; corrected environment retry passed allfour children. This uses the same historical process-local HTTP module as c16 (MSG_NOSIGNAL writer and older readiness predicate), so it does not certify default HTTP readiness or current-main behavior. Cache-off, one timed batch per condition; no speedup claim. GPU compute-process query empty after completion. See opt_prime/examples/pd/http_matrix/experiments/2026-10-25-native-gpu-mid201.md.

  • Default HTTP qualification (a06a36f, runtime03054a1): repeated the four-model native GPU middle-length32conditions/176responses using direct serve_batched.py and the branch default HTTP module, with no process-local module override.80/80token and80/80SSE checks pass;16/16references equal the prior overridden-module and c16 references. Same GPU identity and decoder flags verified, saved arrays/timing summaries independently audited, all children/coordinator0 and GPUs idle. This removes the override-only limitation for these native successful-request cells; it does not certify default-module cancellation/readiness faults, heterogeneous routes or main. See opt_prime/examples/pd/http_matrix/experiments/2026-10-25-default-http201.md.

  • Evidence summary correction (162c2f6): preserve per-request output length separately from aggregate output tokens. Corrected the native GPU intermediate table's c4 column, which previously selected32tokens/request under a128tokens/request heading; raw measurements and correctness results are unchanged. GPU PD and ATOM timing selections were checked against raw fields and remain unchanged. Three published GPU summaries now include both fields. Portable ATOM auditor and10selfchecks cover the length collision, duplicate arms and corrupt/missing raw results. The current final-HTTP ATOM intermediate matrix is still in progress and is not certified by this correction.

  • Final-HTTP ATOM intermediate comparison (ef8cc09): four selected models × native/whole/stream × input512/2048 × output32/128 × c1/c4,96conditions/528responses. Actual saved token/SSE checks pass; whole/stream references match4/4 per model. Native/PD matches2/4 for each Qwen and4/4 for Gemma/KONI; Qwen pre-EOS differences and prior quality limitations remain. This uses one ATOM decode lane and historical research engines/artifacts with the final HTTP module, not current-main or B4 qualification. All children/coordinator exit0, UID cleanup0 and fresh owned-Pod absence verified. See opt_prime/examples/pd/http_matrix/experiments/2026-10-30-atom-mid201-final.md.

  • Native GPU short/long controls (ef8cc09, runtime162c2f6): four models × input128/8192 × output32/128 × c1/c4,32conditions/176responses. All saved token/SSE checks and16 historical c16 references match; all exits0 and fresh GPU compute-process query empty. Portable raw auditor and8 selfchecks accompany the summary. ATOM auditor10 selfchecks also pass; staged sensitive/whitespace checks clean. No production engine change. See opt_prime/examples/pd/http_matrix/experiments/2026-10-30-native-gpu-edge201.md. ATOM short/long low-concurrency, heterogeneous GPU short/long, cache-on and supported-export work remain outside these completed cells.

  • RNGD-to-GPU short/long controls (8f9f2a0, runtime ef8cc09): four models × whole/stream × input128/8192 × output32/128 × c1/c4,64conditions/352responses. All saved token/SSE checks pass, whole/stream and historical same-route references match4/4 per model. Fresh-native equality is Qwen0.6B4/4,Qwen4B2/4,Gemma4/4,KONI2/4; Qwen4/KONI pre-EOS differences remain. At8K/128/c4, stream throughput is75.41/36.29/56.54/39.71tok/s versus native60.84/40.92/58.83/40.46. Single batches and additional RNGD resources prevent stable/equal-resource claims. All8 children/coordinator exit0; fresh local GPU/RNGD idle checks pass. Portable auditor and10 selfchecks included; raw-copy, sensitive and whitespace checks pass. Remote ATOM short/long sweep is now complete as detailed below. See opt_prime/examples/pd/http_matrix/experiments/2026-10-30-pd-gpu-edge201.md.

Final-HTTP ATOM short/long validation

All four models complete native/whole/stream at input128/8192, output32/128, concurrency1/4:12arms96conditions528responses. Whole/stream16/16 reference pairs match. Native/PD references match4/4 for each Qwen,2/4 for Gemma (other2equal throughEOS),and2/4 for KONI (two pre-EOS differences retained). One ATOM execution lane; c4 means queued clients and PD adds RNGD.

At8K/128/c4, native/whole/stream tok/s:Qwen0.6 50.65/33.11/44.24,Qwen4 12.30/12.72/14.47,Gemma57.68/40.19/43.87,KONI7.78/9.32/10.44. Single ordered batches;no stable/equal-resource claim,and KONI native comparison has different outputs. Single-client Qwen4 firstSSE4.970/4.601/3.219s with estimatedTPOT42.76/42.74/42.85ms localizes the observed benefit before first SSE without isolating DMA/compute/import.

Historical stream265 B1 runtime plus finalHTTP SHA256e3449979b88d03edb6d5bf0d402a708e5d98054fdf203e22b480f3d7e6fd759a; exact versions/source/raw hashes and geometry in the evidence. Gemma uses a research corrected hybrid artifact. Eleven raw-auditor selfchecks pass;all children/coordinator/UIDcleanup exit0 and fresh owned-Pod absence verified. Portable auditor,summary and scoped research event included;staged sensitive and whitespace checks pass. Cache-on and broader quality work continue separately. See opt_prime/examples/pd/http_matrix/experiments/2026-10-30-atom-edge201-final.md.

Compatibility and limits

This is experimental opt-in work. Real device-overlap traces remain outstanding. Actual ATOM streaming cancellation and source-death isolation are now verified for the recorded four-model B1 cases. Direct KV/logit parity now includes width8192. Long GPU timing applies only to the recorded c2/fixed32, c4/fixed128 and c16/fixed128 workloads; c16 queues behind4 execution slots. ATOM checks now include B1 and2K/8K inputs with prefetch off/on. Fixed8-token parity ignores EOS and is not answer-quality evaluation or nativeGPU equivalence. Direct-engine timing uses three measured samples/mode; HTTPc16 timing uses one batch per condition and needs repeats for reliable speedup claims; the default remains whole transfer. Recurrent-state streaming is rejected; serial producers must preserve immutable published spans until completion.

Related Issue

Refs #265, #47, #201, #234, #254, #255, #259. Does not close them.

Default HTTP RNGD-to-GPU intermediate conditions

Runtime a06a36f; evidence a3497e5. Four selected models, input512/2048, output32/128, c1/c4, whole/stream:64conditions/352responses (32warmups,160token checks,160SSE checks). All arms complete; within-route checks and whole/stream reference arrays pass. Native reference equality remains Qwen0.6B4/4, Qwen4B2/4, Gemma0/4 fixed-length but4/4 throughEOS, KONI3/4. Qwen4/KONI pre-EOS differences remain unresolved. GPU and selected RNGD have no remaining process holders. Single-batch throughput is descriptive, not a stable speedup or fresh overlap proof. See opt_prime/examples/pd/http_matrix/experiments/2026-10-25-pd-gpu-mid201.md and its sanitized evidence summary.

@gspark-etri gspark-etri changed the title feat: validate partial KV imports for RNGD layer streaming feat: stream RNGD KV layers into Torch decode Oct 23, 2026
@gspark-etri gspark-etri changed the title feat: stream RNGD KV layers into Torch decode feat: stream RNGD KV layers into Torch and ATOM decode Oct 23, 2026
gspark-etri added a commit that referenced this pull request Oct 30, 2026
…#265) (#287)

* pd: stream RNGD prefill KV layers into Torch and ATOM decode (opt-in) (#265)

Port the library and test parts of #266 onto optimus-rm:
- layer_progress.LayerProgress: validated, non-overlapping per-layer token
  coverage of a reserved slab; finish() requires complete coverage.
- PagedKV.write_layer_span: import a model-layer subset into reserved pages
  without publishing length/prefix/state.
- SubprocessPrefillWorker(stream_layers=True): requires producer capability
  kv_progress=1, accepts {id, kv_span} messages, exposes future.kv_progress.
- PDBatcher and AtomBatcher import published spans while the producer runs
  and admit only after completion and full coverage.
- prefill_server: stream_result() writes validated spans into the slab
  before notifying; _RngdFull.prefill_stream; --engine factory loader.
- serve_batched.py --pd-prefill-stream-layers (off by default).

From #266 commits 42f59c3, 8e0d271, 952f491, 171e8c4 (prefill_server.py and
test_prefill_factory.py only), 33d6cb0, 69180aa. The #254 chunk-copy code
(_ChunkCopy, --pd-prefill-copy-overlap and its test) is stripped;
streaming_factories.py, STREAMING.md evidence and scripts are not taken.

Refs #265

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* pd: reject layer streaming for recurrent models at setup (#265)

Layer streaming carries attention KV only, not linear-attention state, so
PDBatcher imported nothing until the first span and then raised. Reject the
combination when it is configured instead: PDBatcher with a stream_layers
worker over a recurrent cache, SubprocessPrefillWorker(stream_layers=True,
recurrent_state=True) before it allocates the slab or spawns, and
serve_batched.py --pd-prefill-stream-layers for a recurrent model before the
producer starts. The import-time check stays as a guard for other producers.

Refs #265

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* pd: document layer streaming scope and factory-engine trust/memory cost (#265)

Add a short STREAMING.md design note: protocol, decode-side invariants, and
the measured attribution. The notification-delay control kept most of the
completion-time gain (22.7/19.6/2.8/21.3%) while early sending added only
4.6/3.0/0.4/1.7%, so the gain comes from the research factories' slower
whole-result export, not overlap. State in --engine-factory help and the
factory_engine docstring that it is a trusted research flag that imports
arbitrary module:callable code and loads the checkpoint as fp32 on the host.

Refs #265

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

---------

Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
@gspark-etri

Copy link
Copy Markdown
Collaborator Author

Landed on optimus-rm via #287. Closing; the branch is archived locally (experiment records and evidence included).

@gspark-etri gspark-etri closed this Nov 3, 2026
@gspark-etri
gspark-etri deleted the perf/265-rngd-stream branch November 3, 2026 17:06
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