[gspark-etri]

Background

PD with a subprocess prefill producer (RNGD through prefill_server) transfers the whole prompt KV after prefill finishes. #265 / #266 add an opt-in path: the producer exports KV one layer or layer group at a time into the shared-memory slab, and the Torch or ATOM decoder imports each finished span. The decoder publishes cache length and prefix state only after full coverage.

This PR lands the library and test part of #266 on current optimus-rm, plus three review fixes.

Changes

Commit 1 (port). From #266 (branch perf/265-rngd-stream), path-limited and resolved against optimus-rm:

#266 commit	taken
42f59c3	new pd/layer_progress.py (LayerProgress); PagedKV.write_layer_span in paged.py; tests/test_pd_layer_progress.py
8e0d271	pd_batching.py: _KVProgress, SubprocessPrefillWorker(stream_layers=) (capability kv_progress: 1, kv_span messages), PDBatcher partial import and wake-up event; layer_progress.py pending/finished; tests/test_pd_stream_admission.py, tests/test_subprocess_layer_progress.py
952f491	prefill_server.py: _RngdFull.prefill_stream, _last_logits, stream_result, ready capability, stream_layers request handling; serve_batched.py --pd-prefill-stream-layers; tests/test_prefill_source_progress.py; layer-stream tests in tests/test_serving_prefill_cache.py
171e8c4	only prefill_server.py (factory_engine, --engine factory, --engine-factory, prefill(list(tokens))) and tests/test_prefill_factory.py
33d6cb0	atom/batching.py: streamed import through the ATOM owner (reserved lane, B1 prefetch waits for the lane, hybrid/multi-lane deferral); serve_batched.py allows --backend atom; tests/test_atom_stream_progress.py; test_serving_prefill_cache.py parametrization
69180aa	tests/test_atom_stream_progress.py (partial-prefetch faults)

Not taken:

• Already on optimus-rm: 8903b73 (= fix: map dense subprocess KV into hybrid GPU cache groups #258 / pd: import dense subprocess KV into hybrid GPU cache groups (#258) #283) and 2e69570 (= Prefill one queued request during single-lane ATOM decode #255 / atom: prefill one queued request during B1 remote-only decode (#253) #285).
• The Overlap GPU prefill token chunks with bounded KV copies #254 chunk-copy code that feat: stream RNGD KV layers into Torch and ATOM decode #266 sits on: _ChunkCopy, TorchPrefillEngine(copy_overlap=), --pd-prefill-copy-overlap, test_copy_overlap_rejects_unsupported_route_before_model_load, and the ThreadPoolExecutor import. These were stripped while resolving conflicts.
• examples/pd/streaming_factories.py and tests/test_streaming_factories.py. They import the serving_matrix / qwen3_rngd example modules, which are not on optimus-rm.
• All http_matrix/ evidence, experiment records, evidence scripts and the docs-only commits.

Commit 2 (review: recurrent models fail early). Layer streaming carries attention KV only. Before this commit, a recurrent (linear-attention) model failed at the first span import. It is now rejected at setup:

• PDBatcher.__init__ rejects a stream_layers worker over a cache with rstate.
• SubprocessPrefillWorker(stream_layers=True, recurrent_state=True) raises before it allocates the slab or spawns.
• serve_batched.py prefill_worker rejects --pd-prefill-stream-layers when describe_model(config).recurrent is set, before the producer starts.
• The import-time check stays as a guard for producers that set kv_progress without stream_layers.
• New test: tests/test_pd_stream_recurrent.py.

Commit 3 (review: docs).

• --engine-factory help text and the factory_engine docstring now say it is a trusted research flag (it imports and calls any module:callable from PYTHONPATH), and that it loads the whole checkpoint into host memory as fp32 first (4 bytes per parameter, about 32 GB for 8B).
• examples/pd/STREAMING.md is cut down to a short design note: protocol, decode-side invariants, and the measurement attribution below.

Verification

Targeted files only. CUDA cases run on one GPU; full suite: CI.

cd opt_prime && TORCH_DEVICE_BACKEND_AUTOLOAD=0 CUDA_VISIBLE_DEVICES=1 PYTHONPATH=$PWD python -m pytest -q \
  tests/test_pd_layer_progress.py tests/test_pd_stream_admission.py tests/test_pd_stream_recurrent.py \
  tests/test_subprocess_layer_progress.py tests/test_prefill_source_progress.py tests/test_prefill_factory.py \
  tests/test_atom_stream_progress.py tests/test_serving_prefill_cache.py tests/test_atom_batching.py \
  tests/test_subprocess_prefill_failure.py tests/test_pd_dense_hybrid.py tests/test_pd_paged_empty_span.py

• torch 2.13: 128 passed.
• torch 2.10: 123 passed, 2 skipped (that transformers has no qwen3_5), 3 failed. The failures are the CUDA cases of test_pd_stream_admission.py, with an Inductor NoValidChoicesError in flex-attention. A plain whole-KV CUDA Gemma3 PDBatcher on unmodified code fails the same way under that torch, so the cause is the environment, not this change.
• Recurrent fix, test first: before commit 2, all 3 tests in test_pd_stream_recurrent.py fail (no raise, unknown recurrent_state keyword, producer started). After it, they pass.
• Against the feat: stream RNGD KV layers into Torch and ATOM decode #266 head, the ported files differ only by the stripped chunk-copy code, the B1-prefetch prefix-cache note from atom: prefill one queued request during B1 remote-only decode (#253) #285, and newer optimus-rm code (pd: import dense subprocess KV into hybrid GPU cache groups (#258) #283 dense-hybrid import, ProcessPrefillWorker.ready, subprocess close order).
• No hardware rerun for this PR.

Compatibility and limits

• Off by default. Without --pd-prefill-stream-layers / stream_layers=True, the whole-KV path is unchanged.
• Requires a subprocess producer that advertises kv_progress: 1 (in-tree: --engine rngd-full, or a trusted --engine factory callable) and a Torch or ATOM decoder. Same-host shared memory only; no cross-host streaming.
• Recurrent (linear-attention) models are rejected.
• The research factories used for the feat: stream RNGD KV layers into Torch and ATOM decode #266 measurements are not included, so the measured setup cannot be reproduced from this branch alone.
• No overlap-benefit claim. feat: stream RNGD KV layers into Torch and ATOM decode #266 measured completion-time reductions against whole-KV transfer of -20% to -35% for Qwen3-0.6B, Qwen3-4B and KONI-8B, and -2% to -10% for Gemma3-1B. Its notification-delay control kept per-layer export but held every progress message until the source finished. That control still reduced completion time by 22.7 / 19.6 / 2.8 / 21.3% (Qwen3-0.6B / Qwen3-4B / Gemma3-1B / KONI-8B). Sending early added only 4.6 / 3.0 / 0.4 / 1.7%. So the gain mostly comes from the research factories' slower whole-result export, which the per-layer path avoids, not from overlapping transfer with prefill. These numbers come from the source branch and were not re-measured here.

Related Issue

Refs #265
Refs #47
Supersedes #266 (library part)

🤖 Generated with Claude Code