Enable GLM recurrent boundary restore with DCP and DFlash - #682
Merged
lukealonso merged 2 commits intoSep 7, 2026
Conversation
Snapshot each attention group using its own context-parallel page geometry. Preserve reachable draft sliding-window pages and reject incomplete windows. Restore DFlash context KV without rewriting shared prefix pages; its fixed-size proposal query uses the ordinary CUDA graph path. GLM static DFlash and DCP configurations can use request-boundary retention. External KV connectors, dynamic speculation, and unsupported parallel layouts retain aligned retention. The existing collective worker wait remains the publication barrier. Validation: 42 DCP snapshot and 18 scalar-restore GPU cases; 260 scheduler, prefix-cache, and DFlash CPU cases; all pre-commit hooks. DFlash2 TP4/DCP1 and TP4/DCP4 passed exact prompt, shared instruction, and response continuation probes. Eight concurrent DCP4 conversations restored all prompt tokens with correct answers. The graph-replayed restore and MTP/no-spec DCP4 serving gates remain in progress. Co-authored-by: OpenAI Codex <noreply@openai.com> Signed-off-by: Martin Vit <martin@voipmonitor.org> (cherry picked from commit fa4720e)
Cover automatic boundary selection for no-speculation, MTP and DFlash under DCP1, DCP2 and DCP4. External cache connectors must retain aligned retention until a checkpoint persistence contract is implemented. All 46 prefix primitive cases and pre-commit checks pass. Co-authored-by: OpenAI Codex <noreply@openai.com> Signed-off-by: Martin Vit <martin@voipmonitor.org> (cherry picked from commit 826121d)
|
Important Review skippedNo new commits to review since the last review. ⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Team Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Team Run ID: 📒 Files selected for processing (9)
Included review availability: Your plan provides up to 1 included review per hour; 0 remain after this review. 📝 WalkthroughWalkthroughThe PR adds DCP-aware boundary checkpoint tests, supports DFlash checkpoints for GLM5, handles evicted sliding-window pages, updates attention-tail page mapping, and changes DFlash replay to reuse restored context K/V. ChangesDFlash boundary checkpoints
Estimated code review effort: 4 (Complex) | ~45 minutes Merge Risk: ⚪ Minimal · up to No actionable merge risk remains in the reviewed checkpoint restoration changes. Sequence Diagram(s)sequenceDiagram
participant replay_draft
participant DFlashSpeculator
participant ContextKV
replay_draft->>DFlashSpeculator: propose(context_kv_is_restored=True)
DFlashSpeculator->>ContextKV: reuse restored context K/V
DFlashSpeculator-->>replay_draft: generate draft without precompute
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
Author
|
@coderabbitai review Please check the rank-local versus replicated attention-page mapping, immutable DFlash context handling, sliding-window reachability, and external-cache fail-closed gate. Packaged TP4/DCP4 no-spec FP8 and DFlash2 NVFP4 serving passed the exact/instruction/assistant/tool/concurrency matrix. Performance and LMCache release qualification are separate and still running. |
🧩 Analysis chain🏁 Script executed: #!/bin/bash
set -euo pipefail
printf '%s\n' '=== PR diff summary ==='
git diff --stat HEAD^ HEAD
git diff -- \
vllm/config/vllm.py \
vllm/v1/core/kv_cache_manager.py \
vllm/v1/worker/gpu/boundary_checkpoint.py \
vllm/v1/worker/gpu/model_runner.py \
vllm/v1/worker/gpu/spec_decode/dflash/speculator.py \
tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py \
tests/v1/core/prefix_cache/test_partial_prefix_cache_primitives.py \
tests/v1/spec_decode/test_dflash_dcp.py \
tests/v1/worker/test_boundary_checkpoint_attention.py
printf '%s\n' '=== File outlines ==='
ast-grep outline vllm/config/vllm.py --items all --type class,function --match 'VllmConfig|use_request_boundary_checkpoints' || true
ast-grep outline vllm/v1/core/kv_cache_manager.py --items all --type class,function --match 'KVCacheManager|publish_boundary_checkpoint' || true
ast-grep outline vllm/v1/worker/gpu/boundary_checkpoint.py --items all --type class,function --match 'Boundary|replay_draft|copy_attention' || true
ast-grep outline vllm/v1/worker/gpu/model_runner.py --items all --type class,function --match 'GPUModelRunner|execute_model|sample_tokens' || true
ast-grep outline vllm/v1/worker/gpu/spec_decode/dflash/speculator.py --items all --type class,function --match 'DFlashSpeculator|propose|precompute_and_store_context_kv' || true
printf '%s\n' '=== Definitions and call sites ==='
rg -n -C 5 -P '\b(use_request_boundary_checkpoints|publish_boundary_checkpoint|_copy_attention_tails_kernel|replay_draft|context_kv_is_restored|precompute_and_store_context_kv|kv_transfer_config|get_num_skipped_tokens)\b' \
vllm tests
printf '%s\n' '=== Related page geometry and snapshot/restore code ==='
rg -n -C 5 -P '\b(group_cp_sizes_ptr|cp_size|context_parallel|page_table|physical_block|virtual_block|attention.*pool|copy_attention)\b' \
vllm/v1/worker/gpu/boundary_checkpoint.py \
vllm/v1/worker/gpu \
vllm/v1/core \
tests/v1/worker/test_boundary_checkpoint_attention.py \
tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.pyLength of output: 202577 🏁 Script executed: #!/bin/bash
set -euo pipefail
printf '%s\n' '=== Changed hunks with line numbers ==='
git diff --unified=80 HEAD^ HEAD -- \
vllm/config/vllm.py \
vllm/v1/core/kv_cache_manager.py \
vllm/v1/worker/gpu/boundary_checkpoint.py \
vllm/v1/worker/gpu/model_runner.py \
vllm/v1/worker/gpu/spec_decode/dflash/speculator.py \
tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py \
tests/v1/core/prefix_cache/test_partial_prefix_cache_primitives.py \
tests/v1/spec_decode/test_dflash_dcp.py \
tests/v1/worker/test_boundary_checkpoint_attention.py
printf '%s\n' '=== Boundary checkpoint implementation ==='
sed -n '220,370p' vllm/v1/worker/gpu/boundary_checkpoint.py
sed -n '500,730p' vllm/v1/worker/gpu/boundary_checkpoint.py
printf '%s\n' '=== Block table DCP mapping ==='
sed -n '1,260p' vllm/v1/worker/gpu/block_table.py
printf '%s\n' '=== Boundary checkpoint publish ==='
sed -n '650,785p' vllm/v1/core/kv_cache_manager.py
printf '%s\n' '=== Feature gate and initialization callers ==='
sed -n '575,640p' vllm/config/vllm.py
sed -n '580,710p' vllm/v1/worker/gpu/model_runner.py
sed -n '1560,1620p' vllm/v1/worker/gpu/model_runner.py
sed -n '2035,2090p' vllm/v1/worker/gpu/model_runner.py
sed -n '340,375p' vllm/v1/core/sched/scheduler.py
printf '%s\n' '=== DFlash propose and metadata construction ==='
sed -n '300,500p' vllm/v1/worker/gpu/spec_decode/dflash/speculator.py
printf '%s\n' '=== New and modified tests ==='
sed -n '1,220p' tests/v1/worker/test_boundary_checkpoint_attention.py
sed -n '1,205p' tests/v1/spec_decode/test_dflash_dcp.py
sed -n '120,190p' tests/v1/core/prefix_cache/test_partial_prefix_cache_primitives.py
rg -n -C 8 -P '\b(test_request_boundaries_reuse_exact_prompt_and_response_with_private_state|test_request_boundary_retains_complete_dflash_windows)\b' tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.pyLength of output: 81658 🏁 Script executed: #!/bin/bash
set -euo pipefail
printf '%s\n' '=== PR changed lines only ==='
git diff --unified=12 HEAD^ HEAD -- \
vllm/config/vllm.py \
vllm/v1/core/kv_cache_manager.py \
vllm/v1/worker/gpu/boundary_checkpoint.py \
vllm/v1/worker/gpu/model_runner.py \
vllm/v1/worker/gpu/spec_decode/dflash/speculator.py
printf '%s\n' '=== Sliding-window block removal and request block representation ==='
rg -n -C 12 -P '\b(def _remove_blocks_in_range|def remove_skipped_blocks|req_to_blocks\[|req_to_blocks\.get|def get_num_skipped_tokens)\b' \
vllm/v1/core/single_type_kv_cache_manager.py \
vllm/v1/core/kv_cache_manager.py
printf '%s\n' '=== Supported DCP layout derivation ==='
rg -n -C 10 -P '\b(group_cp_sizes|dcp_replicated|decode_context_parallel_size|dcp_size|dcp_rank|cp_interleave)\b' \
vllm/v1/worker/gpu/model_runner.py \
vllm/v1/core/kv_cache_manager.py \
vllm/v1/core/single_type_kv_cache_manager.py \
vllm/v1/kv_cache_interface.py \
vllm
printf '%s\n' '=== DFlash propose completion ==='
sed -n '480,590p' vllm/v1/worker/gpu/spec_decode/dflash/speculator.py
sed -n '720,825p' vllm/v1/worker/gpu/spec_decode/dflash/speculator.py
printf '%s\n' '=== Boundary snapshot allocation and restore callers ==='
rg -n -C 10 -P '\b(boundary_checkpoint_blocks|capture_attention|wait_for_copies|boundary_checkpoint_tokens|take_kv_cache_block_copies|copy.*boundary|boundary.*copy)\b' \
vllm/v1 \
tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py
printf '%s\n' '=== Complete DFlash window test ==='
sed -n '200,330p' tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py
printf '%s\n' '=== Test changes as compact diffs ==='
git diff --unified=20 HEAD^ HEAD -- \
tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py \
tests/v1/core/prefix_cache/test_partial_prefix_cache_primitives.py \
tests/v1/spec_decode/test_dflash_dcp.py \
tests/v1/worker/test_boundary_checkpoint_attention.pyLength of output: 199443 🏁 Script executed: #!/bin/bash
set -euo pipefail
printf '%s\n' '=== Block removal implementation ==='
rg -n -A 65 -B 8 -P '^\s*def _remove_blocks_in_range\s*\(' \
vllm/v1/core/single_type_kv_cache_manager.py
rg -n -A 45 -B 8 -P '^\s*def remove_skipped_blocks\s*\(' \
vllm/v1/core/single_type_kv_cache_manager.py
printf '%s\n' '=== Request block-list storage and allocation updates ==='
rg -n -A 45 -B 8 -P '^\s*(self\.)?req_to_blocks\[.*\]\s*=|req_to_blocks.*append|blocks\[.*\]\s*=\s*self\.block_pool\.null_block' \
vllm/v1/core/single_type_kv_cache_manager.pyLength of output: 12122
The added tests cover DCP sizes, replicated DFlash groups, eager and full-graph DFlash paths, reachable-window holes, copy-on-write, and the external-cache gate. I could not run GPU tests in this environment.
|
Author
|
Additional no-spec/DCP1 qualification found a limitation of the greedy-output TP4, four RTX PRO 6000 Blackwell Workstation GPUs at stock clocks, FP8 target
The reproducible input is token pattern No numerical kernel was changed to make this probe pass. Strict failures are |
lukealonso
merged commit Sep 7, 2026
99c9436
into
local-inference-lab:dev/jovian-judgement
2 of 3 checks passed
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
2 participants
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.
Purpose
Enable GLM-5.3-Flash semantic recurrent checkpoints with static DFlash2 and
decode context parallelism. Shared leading instructions, completed prompts,
and committed response endpoints can be reused without retaining a recurrent
state every 256 tokens.
Restore contract
Target pages are sharded under DCP; DFlash draft pages remain replicated.
through its ordinary query graph. It must not rewrite the last context row
into an immutable shared page.
the checkpoint. A hole inside that window invalidates the checkpoint.
the scheduler publishes the checkpoint. No additional handshake is introduced.
External KV connectors, dynamic speculation, and unsupported parallel layouts
continue to select aligned retention. This PR does not implement semantic
checkpoint persistence in LMCache or change physical serving-page defaults.
Dependencies and scope
The serving qualification stack includes #674 for scalar-safe recurrent
restoration and #664 for isolated logits-only scheduler steps. It also includes
#676, which fixes partial physical attention-page lookup independently.
This does not duplicate those changes: it supplies the missing DCP address
mapping and DFlash context/window restore contract. Open-PR searches in this
fork and upstream found no equivalent request-boundary implementation.
Validation
Status: implemented; GPU-resident cache correctness is qualified for the
configurations below. Matched release-performance qualification is in progress.
Four RTX PRO 6000 Blackwell Workstation Edition GPUs, TP4, stock clocks,
4,096-token scheduler budget, 2,048-token target/recurrent pages, and full plus
piecewise CUDA graphs:
logical page edges, exact copied bytes, and untouched neighboring pages.
12 partial-attention-rewind cases pass.
8,192, 8,193, and 32,768 prompt tokens and preserve all 256 greedy output
token IDs. Response continuations match independent cold controls.
matrix. Five seeded 4,096-token coding requests show no acceptance collapse
after restoration: median accepted draft tokens/step is 2.095 cold and
2.157 restored; verifier throughput is 92.32 and 92.45 steps/s.
including actual assistant/user continuation, tool history, and divergent
instructions. The assistant continuation restores 11,385 tokens; the tool
history replays all 11,520 prompt tokens.
matrix, including the first restored proposal through the query CUDA graph.
selection under DCP1/2/4 and fail-closed external-cache selection. The complete
prefix-primitive test file passes all 46 cases.
llama-benchy0.4.0 at depth 8,192 passes coherence and prefix reuse withDFlash2/DCP1, DFlash2/DCP4, and MTP3/DCP4.
Representative reproduction commands, in a built vLLM environment:
Serving probes use
/v1/completionswith token-ID prompts,temperature=0,max_tokens=256,ignore_eos=true,return_token_ids=true, and an unchangedcache_salt. Each exact prompt is sent three times. The two replays must reportthe entire prompt as cached and return the same output token IDs as the first
request. Appended-response probes are compared with separately salted cold
controls. Chat probes check literal record values through instruction branches,
assistant continuations, and tool histories.
Matched release-performance and external-cache qualification remain required
before recommending the image for deployment. The stochastic acceptance medians
above are not a speedup claim.
AI assistance was used for implementation, review, and validation. Human
maintainer review is required before merge.
Summary by CodeRabbit
New Features
Bug Fixes