Repository navigation
Overlap GPU prefill token chunks with bounded KV copies - #254
Draft
gspark-etri wants to merge 4 commits into
Draft
gspark-etri wants to merge 4 commits into
gspark-etri wants to merge 4 commits into
Conversation
This was referenced 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>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Background
Torch GPU PD currently gathers and transfers the complete KV span only after producer prefill finishes. This draft adds opt-in token-chunk transfer so the next prefill chunk can run while earlier KV copies are pending.
Changes
TorchPrefillEngine(copy_overlap=True)andserve_batched.py --pd-prefill-copy-overlap: one copy worker, dedicated CUDA streams, readiness events, pinned-host staging, at most two pending spans.transfer_smeans the unhidden transfer tail, not total DMA time.Verification
Runtime: torch 2.13.0+cu130 on the GPU host. Code: 0870dd3, based on 273952a. Commands, setup, failures and logs are in
opt_prime/examples/pd/http_matrix/experiments/2026-10-24-chunk-copy47.mdandevidence47/.copy_overlapconstructor option; first implementation then exposed the existing empty-span index error.PYTHONPATH=opt_prime /home/gspark/venvs/pt213/bin/python -m pytest opt_prime/tests/test_pd_chunk_copy.py opt_prime/tests/test_serving_prefill_cache.py -q: 17 passed, 4.79 s, using two CUDA devices. Exact KV/token parity across dense/hybrid models, prefix reuse and suffix offsets, empty spans, delayed-copy producer progress, transfer failure/recovery, consumer source-reuse and unsupported route rejection.opt_prime/:CUDA_VISIBLE_DEVICES='' PYTHONPATH=. /home/gspark/venvs/pt213/bin/python -m pytest -q: 504 passed, 43 skipped, 253 subtests, 151.19 s; one existing unfused flex-attention warning. This full run preceded adding the final GPU-only source-reuse test, which passed separately and in the focused rerun.git diff --cached --check: clean. Ruff unavailable in this environment.Compatibility and limits
Distinct Torch CUDA producer/decoder devices only. The decoder handoff contract stays complete-KV. RNGD/ATOM transport and layer-wise export are not implemented by this change. Recurrent-model coverage is still pending.
No serving speedup claim yet. On two A100 80GB PCIe GPUs (driver580.178.04), Qwen3-0.6B input8192/chunk256/cache-off paired prefill+handoff median fell997.98→955.64ms (-4.24%, three alternating repeats/mode); all six runs had exact KV/first-token parity. Both modes transferred939,524,096 bytes. Separate CUDA traces show27.91ms DtoH/source-kernel overlap, mainly27.47ms KV index gather; GEMM intersection was zero. Unhidden transfer-tail median fell55.61→2.62ms. Profiling runs are excluded from latency comparisons. Results/reproducers are in evidence47; large raw traces remain local with SHA256s in the experiment record.
Four-model HTTP blockA completed96timed batches with all token/text parity checks passing. At8192input/128output/c8, whole→chunk tok/s: Qwen0.6 78.85→76.39(-3.1%), Qwen4 54.45→55.81(+2.5%), KONI52.67→53.93(+2.4%), Gemma72.47→71.36(-1.5%). This is not a universal serving improvement; reverse model/mode order blockB also completed. Across both blocks, all 192 timed batches passed parity. Long/c8 changes in blockB were Qwen0.6 +5.80%, Qwen4 +17.39%, KONI -1.53%, Gemma -0.28%. Qwen4 baseline accumulated prefill increased 69.90→82.21s while transfer stayed 3.35s, so the larger gain is confounded. No general repeatable throughput benefit is established.
The separate0.1/0.5/1/2/4GB export-only curve completed30constant-payload integrity checks. With no model compute to overlap, chunked export is slower: at4GB253.60→341.48ms. Incremental source GPU allocation drops4001.6→67.1MB; this is not pinned-host RSS or total device memory. Both paths include gather and host-staged DMA; exact script/config/results are in evidence47.
Keep opt-in and in draft. Controlled resolution of serving variance, recurrent coverage and original layer-wise RNGD/ATOM transport requirements remain open.
Related Issue
Partial implementation toward #47; does not close it. Stacked on #250 for producer timing counters.