[BDE] Phase 1: engine-level KV cache management for DreamZero (RFC #4366) by tzhouam · Pull Request #4534 · vllm-project/vllm-omni

tzhouam · 2026-06-18T08:40:44Z

What this does

Phase-1 of the AR-Diffusion Engine KV-cache RFC (#4366): DreamZero's persistent attention KV runs on vLLM's paged KV stack, managed at the engine level — no new attention kernel. DreamZero now runs exclusively on the engine (its model-local cat+slice KV path is removed); the engine was validated bit-exact against that path before removing it.

Design

flowchart TB
  REQ["DreamZero request<br/>engine_backend = ar_diffusion (default)"] --> ENG["ARDiffusionEngine : DiffusionEngine"]
  ENG --> RUN["ARDiffusionModelRunner — worker / GPU<br/>KV keyed by session_id"]
  RUN -->|"sets _ar_diffusion_kv_state before forward()"| PIPE["DreamZero pipeline.forward()"]
  PIPE -->|"contiguous window tensor — no new kernel"| ATTN["model self_attn / cross_attn"]
  PIPE <-->|"_kv_get (memoized gather)<br/>_kv_update / _kv_reset"| STATE["ARDiffusionKVState — per-session bridge"]
  STATE --> CACHE["ARDiffusionKVCache — engine-owned pools"]
  CACHE --> SELF["Self-attn paged pool<br/>1 block = 1 frame · sliding-window evict"]
  CACHE --> CROSS["Cross-attn static pool<br/>write-once · read every step"]
  SELF --> VLLM["vLLM KVCacheManager / BlockPool / SlidingWindowManager<br/>via ARDiffusionRequestAdapter + ChunkWindowSpec"]

ARDiffusionEngine(DiffusionEngine), selected by the engine_backend command/deploy arg. DreamZero defaults to it (engine_backend = ar_diffusion) and is engine-only — its model-local KV fallback was removed. Other diffusion models are byte-for-byte unchanged.
Reuses vLLM's KVCacheManager / BlockPool / SlidingWindowManager as a library (via a duck-typed ARDiffusionRequestAdapter), with a chunk-granular sliding window (ChunkWindowSpec / ChunkWindowManager). No diffusion-local allocator.
Self-attention — frame-granular paging (1 pool block = 1 frame). Each forward writes only the new tokens into the pool and gathers the resident window into the contiguous tensor the existing self_attn already consumes; ChunkWindowManager evicts out-of-window frames.
Cross-attention — also moved into the engine pool (static: write-once from the text encoder, read every denoise step).
Session-scoped state keyed by session_id, window-boundary reset, independent CFG pos/neg adapters. KV parameters come from the ar_diffusion_kv_config deploy/command arg (defaults work out-of-box).

Memoized gather. The resident window is frozen across a forward's denoise steps (KV is written back only after the last step), so re-gathering it every step is redundant. The engine gathers once per forward and reuses it (write-back invalidates it), building a vectorized slot index once and sharing it across all 40 layers. This keeps latency at parity and holds down memory — gathers drop 136 → 34 per run (~4×).

read #1 (miss):  gather resident pool blocks → one window-set → memoize
reads #2..N:     return the memoized set — NO pool access, NO allocation
write-back:      scatter new tokens into pool blocks in place → invalidate memo
next forward:    one fresh gather

Test plan

pytest tests/ar_diffusion/ — 59 CPU tests (slot mapping, pool write/gather, chunk-window eviction, orchestration lifecycle, engine-owned-KV guards, routing).
Bit-exact validation (before removing the model-local path): one 13-forward DreamZero session (1 prefill + 12 chunks), engine vs the model-local path, identical inputs/seed, on the production dreamzero.yaml (H200, 1 GPU, enforce_eager, DiT step-caching on, CFG-parallel off). The engine matched bit-for-bit — which is what justified making DreamZero engine-only. Absolute latency is HW/config-specific; the delta is the meaningful figure. Perf/rollout harnesses are local dev tooling, out of this PR.

Results (H200, 1 GPU, DreamZero, bf16)

Measured against the former model-local path (removed in this PR) — the validation that justified the removal.

Precision — bit-exact:

Artifact	engine vs model-local
Output latents `(1,16,30,44,80)`	`max_abs_diff = 0`, PSNR = ∞
Decoded video (117 frames)	all frames identical
Final `.mp4`	byte-identical (same SHA-256)

Generated video — former model-local path (left) | AR-Diffusion engine (right), visually identical:

Performance — at parity (mean ± σ over 5 runs/condition):

Config	E2E / chunk (steady)	vs model-local	gathers / run
model-local (now removed)	1644 ± 5 ms	1.000×	—
engine, no memoization	1688 ± 3 ms	1.027× (+44 ms)	136
AR-Diffusion (current)	1667 ± 8 ms	1.014× — at parity (+23 ms)	34

Memory — +4.6 GiB peak vs the model-local path. The pools are preallocated once and written in place; the only KV memory allocated during inference is the transient gathered window — gathered once per forward instead of re-gathered every denoise step, which keeps the overhead modest. Without memoization the per-step re-gather inflates the allocator reserve to +13.3 GiB.

Config	peak VRAM	vs model-local
model-local (now removed)	78.0 GiB	+0.0
engine, no memoization	91.3 GiB	+13.3 GiB
AR-Diffusion (current)	82.6 GiB	+4.6 GiB

Status

Phase-1 milestone for RFC #4366. Engine/KV selection is via the engine_backend / ar_diffusion_kv_config command/deploy args (the former env-var stopgap is gone; DreamZero defaults to the engine). Remaining: CFG-parallel / 2-GPU validation; the fused paged kernel that removes the gather/copy entirely is Phase 2.

Update (2026-07-04): performance & memory optimization stack

This update lands #4863's paged self-attention (@Acerak01-fy's commits, coordinated offline) plus the optimizations built on top of it:

Startup warm-up: torch.compile/CUDAGraph capture moved off the serving hot path — full speed from the first request (previously ~100 s cold-inductor + per-shape warm-in on the serving path).
Session caches: per-session text-embed cache; window-persistent VAE stream and text cross-attn KV across window restarts (only the observation-dependent image half re-projects).
torch.compile-native paged attention: as cherry-picked, the paged path graph-breaks every DiT block on compiled stacks (+6% latency, GPU idle 24%→34%). A fused write+attend custom op + once-per-forward fixed-shape block-table metadata (layer identity carried by absolute slot/block values into a stacked all-layer pool — no per-layer graph specialization) keeps blocks fullgraph=True: CFG-parallel ×2 steady 613→611 ms/chunk, TP1 1104→1095 ms, paged memory win intact.
Branch-aware KV pool sizing: under CFG-parallel each rank executes exactly one branch; sizing the preallocated pool per local branch cuts 14 GiB/GPU (76.3 → 62.2 GiB reserved) at unchanged latency — leaner than both NVIDIA upstream (73.7 GiB) and the released v0.22.0 pair (67.8 GiB) on the same model.

Measured on H200 (num_chunks=8, seed 0, per-chunk steady mean over 8 distinct DROID scenes): 624 ms/chunk — 2.69× vs NVIDIA upstream FA2, 4.24× vs the released v0.22.0 pair at the same 2-GPU CFG-parallel topology. tests/ar_diffusion: 72 passed (+8 GPU paged-vs-dense parity combos).

Reproducing the numbers

The benchmark/profiling harness is kept off this PR to stay slim (per earlier review) and lives on the companion branch bde-kv-phase1-bench — this PR's head plus benchmarks/ar_diffusion/ (PSNR/SSIM/LPIPS parity gates, KV-memory plateau profiler), examples/offline_inference/dreamzero/ar_diffusion_perf_compare.py (the per-chunk latency driver behind every number here), the long-rollout driver, the multi-GPU deploy configs used for the measurements (dreamzero_ardiffusion_{cfg2,tp2}.yaml — the PR itself ships only the single dreamzero.yaml; adjust devices/parallel_config for multi-GPU), an env-gated memory profiler (AR_DIFFUSION_MEM_PROFILE=<dir>: per-rank allocator checkpoints, per-component weight breakdown, allocation-history snapshots) backing the memory numbers, and a torch-profiler analysis skill. Example:

git checkout bde-kv-phase1-bench
CUDA_VISIBLE_DEVICES=0,1 HF_HOME=<hf_cache> \
  python examples/offline_inference/dreamzero/ar_diffusion_perf_compare.py \
    --num-chunks 8 --warmup 2 --seed 0 --cuda-graph \
    --deploy-config vllm_omni/deploy/dreamzero_ardiffusion_cfg2.yaml \
    --video-dir <scene_dir> --prompt "<scene prompt>" --timing out.json

Credits

Joint work of the AR-Diffusion / BDE group: @tzhouam @amy-why-3459 @TKONIY @Acerak01-fy @AbelSara @matchyc @loveysuby @yiwen101 @wsl2000 @asukaqaq-s @zwhzzz0821 @wtz2333 — co-author trailers are carried on the branch so a squash merge credits everyone.

hsliuustc0106

DCO and pre-commit checks failing. Fix those before requesting full review.

chatgpt-codex-connector · 2026-06-18T10:45:03Z

Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits.
Credits must be used to enable repository wide code reviews.

hsliuustc0106 · 2026-06-18T19:01:52Z

Review findings:

[P1] BDE session KV is never evicted or freed for old session IDs. BDEModelRunner keeps _bde_states forever (vllm_omni/bde/runner.py, around line 186), while the pipeline LRU evicts DreamZero states at vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py, around line 366. Old BDE adapters keep owning pool blocks, so a long-running server or session switch can exhaust the finite BDE KV pool even after model-local state was evicted. Please add matching LRU/free behavior, or an explicit session teardown path that calls end_request for both branches.
[P2] The DreamZero config helper still uses causal-block granularity, but the runner uses frame granularity. benchmarks/bde/dreamzero_config.py derives chunk_size = num_frame_per_block * frame_seqlen and window_chunks = local_attn_size // num_frame_per_block, while vllm_omni/bde/runner.py overwrites chunk_size to frame_seqlen and preserves any configured window_chunks. If a helper-produced config sets window_chunks=7, the runner turns that into a 7-frame window instead of the intended 21-frame window. Please update the helper/tests to chunk_size=frame_seqlen, window_chunks=local_attn_size.
[P2] BDE infers frame_seqlen from deploy image_resolution rather than the loaded transformer/runtime shape. If deploy config and model geometry drift, BDE allocates/writes with the wrong chunk size while the pipeline computes seq_len from actual latents. Safer to use t.frame_seqlen and assert any deploy-derived value matches.

I did not run pytest locally: this active environment does not have vllm/torch importable, and the repo instructions ask for a fresh uv env before local testing. Static diff check passed.

hsliuustc0106 · 2026-06-18T19:03:32Z

+        _log.info("BDE RESET [%s/%s] session KV cleared (window boundary)", pos_id, neg_id)
+
+
+class BDEPipelineMixin:


This mixin is exported (kv_cache/__init__.py) but never actually used: DreamZeroPipeline defines its own _kv_get / _kv_create / _kv_update / _kv_commit / _kv_get_cross / _kv_reset proxies directly and doesn't inherit BDEPipelineMixin. The two have already diverged — _bde_kv_update here is missing the seq_len arg the real _kv_update in pipeline_dreamzero.py passes. Suggest deleting the mixin (avoids a misleading duplicate surface that can silently drift further), or wiring the pipeline to inherit it.

hsliuustc0106 · 2026-06-18T19:03:34Z

        x = self.o(x)
        return x

+    def project_kv(self, context: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:


Both project_kv methods (here and on WanI2VCrossAttention) are never called — _kv_populate_cross inlines the projection (ca.norm_k(ca.k(projected)).unflatten(2, (n, d)) …) rather than calling block.cross_attn.project_kv(...). Only a docstring references it. Either drop these, or call them from _kv_populate_cross so the K/V projection lives in one place.

hsliuustc0106 · 2026-06-18T19:03:35Z

+                        # BDE engine (engine-level KV cache management) unless the
+                        # caller explicitly selected an engine backend.
+                        bde_default = default_engine_backend_for_model(self.model_class_name)
+                        if bde_default is not None and self.engine_backend == "default":


Design question on blast radius: both enrich_config (here) and from_kwargs rewrite engine_backend "default" -> "bde" for DreamZeroPipeline, so every DreamZero load now builds BDEModelRunner — even though the KV feature is off by default and the production topology (CFG-parallel + CUDA graph + torch.compile) is still unvalidated per the remaining-work doc.

Have you considered gating the routing behind the same enable switch (BDE_KV_ENABLE / the config field) so DreamZero stays on base DiffusionEngine until BDE KV is explicitly opted in? That would shrink the blast radius of a half-validated runner to zero for anyone not testing BDE. Side note: because "default" gets overridden here, there's currently no way to force base DiffusionEngine for DreamZero.

hsliuustc0106 · 2026-06-18T19:03:37Z

+        """Resident (non-null) blocks the read path gathers the window from."""
+        return resident_block_ids(self.block_table(adapter), self.null_block_id)
+
+    def commit_chunk(self, adapter: BDERequestAdapter) -> None:


Minor, but easy to misread: commit_chunk is only reached from tests in practice. In production, completed_chunks advances via adapter.on_chunk_committed() directly inside BDEKVState.update_kv_cache (gather.py), and the pipeline's _kv_commit -> BDEKVState.commit_chunk is a no-op log. A one-line docstring noting the real commit path (or collapsing the two) would help the next reader.

hsliuustc0106 · 2026-06-18T19:06:47Z

Re-review with #4425 (VllmOmniConfig) as background:

Context: #4425 introduces the additive structured config layer and a DiffusionConfig mirror of the current OmniDiffusionConfig surface. That changes the compatibility bar for this PR: new diffusion config fields added here should also be represented in the structured config path, or future consumers of VllmOmniConfig can silently diverge from the legacy config behavior.

Findings:

[P1] BDE session KV still has no eviction/free path for old sessions. BDEModelRunner keeps _bde_states forever (vllm_omni/bde/runner.py, around line 186), while the DreamZero pipeline caps model-local DreamZeroState with an LRU (vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py, around line 366). Old BDE adapters can keep owning pool blocks after the corresponding model-local state was evicted, so a long-running server or session churn can exhaust the finite BDE KV pool. Please add matching LRU/free behavior, or an explicit session teardown path that calls end_request for both branches.
[P2] Considering [Config] Add structured VllmOmniConfig classes #4425, the new diffusion fields from this PR are missing from the structured config surface. This PR adds engine_backend and diffusion_model_runner_cls to OmniDiffusionConfig (vllm_omni/diffusion/data.py, around line 555), but [Config] Add structured VllmOmniConfig classes #4425's DiffusionConfig does not include them (vllm_omni/config/omni_config.py, around line 392). Its enrich_config() bridge only copies fields present on DiffusionConfig, so a future structured-config consumer would silently drop DreamZero's BDE routing and fall back toward the default engine path. Please add these fields/classification/tests when rebasing with [Config] Add structured VllmOmniConfig classes #4425.
[P2] The DreamZero helper still derives chunk/window at causal-block granularity while the runner uses frame granularity. benchmarks/bde/dreamzero_config.py derives chunk_size = num_frame_per_block * frame_seqlen and window_chunks = local_attn_size // num_frame_per_block, while vllm_omni/bde/runner.py overwrites chunk_size to frame_seqlen and preserves any configured window_chunks. If a helper-produced config sets window_chunks=7, the runner turns that into a 7-frame window instead of the intended 21-frame window. Please update the helper/tests to chunk_size=frame_seqlen, window_chunks=local_attn_size.
[P2] BDE infers frame_seqlen from deploy image_resolution rather than the loaded transformer/runtime shape (vllm_omni/bde/runner.py, around line 120). If deploy config and model geometry drift, BDE allocates/writes with the wrong chunk size while the pipeline computes seq_len from actual latents. Safer to use t.frame_seqlen and assert any deploy-derived value matches.

I did not run pytest locally: this active environment does not have vllm/torch importable, and the repo instructions ask for a fresh uv env before local testing. Static diff check passed in the PR worktree.

hsliuustc0106 · 2026-06-18T19:06:58Z

        filtered_kwargs = {k: v for k, v in kwargs.items() if k in valid_fields}

-        return cls(**filtered_kwargs)
+        instance = cls(**filtered_kwargs)


Heads-up on coordination with #4425 (RFC #4021 Phase 2 structured config), which is also open and touches this exact region:

Merge conflict is certain. [Config] Add structured VllmOmniConfig classes #4425 refactors OmniDiffusionConfig.from_kwargs here — it extracts the legacy kwarg-normalization into a module-level normalize_omni_diffusion_kwargs() and keeps the tail as return cls(**filtered_kwargs). This block (which changes that tail to instance = cls(...) + the per-model routing) will conflict textually, and [Config] Add structured VllmOmniConfig classes #4425 also adds fields to OmniDiffusionConfig. Whichever lands second rebases; worth syncing order with @Acerak01-fy.

Mirror-config debt. [Config] Add structured VllmOmniConfig classes #4425's new structured DiffusionConfig is a hand-maintained explicit mirror of OmniDiffusionConfig's diffusion knobs (with a parity test), not a dynamic fields() enumeration — so the engine_backend / diffusion_model_runner_cls fields added here won't auto-propagate into the structured config. That's exactly the fragmentation RFC [RFC]: Unified VllmOmniConfig #4021 is trying to retire. Either mirror them into [Config] Add structured VllmOmniConfig classes #4425's DiffusionConfig (probably under the engine/runtime group, where these belong) or note the deferral deliberately.

Env-var direction. The BDE_KV_ENABLE stopgap + standalone BDEKVConfig go the opposite direction from [Config] Add structured VllmOmniConfig classes #4425's structured-config push. The stated next-step config-field enablement would land most naturally inside the new DiffusionConfig rather than as another env var / parallel dataclass.

hsliuustc0106 · 2026-06-18T19:09:19Z

Review summary — BDE Phase 1 (#4534)

High-quality PR: clean engine/runner separation, honest eager-mode perf methodology, bit-exact parity evidence, and a solid 72-test CPU suite. No hard blockers from me — leaving as COMMENT. Five inline notes:

BDEPipelineMixin (gather.py) — exported but never inherited. DreamZeroPipeline defines its own _kv_* proxies and the two have already diverged (mixin's _bde_kv_update is missing seq_len). Delete it, or wire the pipeline to inherit it.
project_kv (causal_wan_model.py) — both cross-attn project_kv methods are never called (_kv_populate_cross inlines the projection). Remove or actually call them.
Routing default-on (data.py) — consider gating the DreamZero→BDE engine/runner routing behind the same enable switch, so the swap isn't default-on for every DreamZero deployment while KV is off and the production topology (CFG-parallel + CUDA graph + compile) is still unvalidated.
commit_chunk (manager.py) — test-only in production; the real commit advances completed_chunks directly in BDEKVState.update_kv_cache. Two commit paths is easy to misread.
Config coordination with [Config] Add structured VllmOmniConfig classes #4425 (RFC [RFC]: Unified VllmOmniConfig #4021) (data.py) — both PRs rewrite OmniDiffusionConfig.from_kwargs in data.py (certain merge conflict). Also, [Config] Add structured VllmOmniConfig classes #4425's structured DiffusionConfig is a hand-maintained mirror, so the new engine_backend / diffusion_model_runner_cls fields won't auto-propagate into it. The stated next-step config-field enablement should land in [Config] Add structured VllmOmniConfig classes #4425's structured DiffusionConfig rather than as another env var / standalone BDEKVConfig.

Verified: the gather memoization is correct against DreamZero's denoise cadence (base diffuse runs predict_noise(..., update_kv_cache=False), so the memo holds across the 16 steps; write-back fires ≤1×/forward — consistent with the 434→26 gather count). Gates green on ef3b67e2; note CI ran no pytest, so parity is PR-description evidence only.

Address the PR vllm-project#4534 review findings, relocate BDE under a new experimental package, and add memory/memoization profiling support. Fixes (from the bde_phase1 review): - P1 session-KV leak: bound BDEModelRunner._bde_states with an LRU at the same cap as DreamZeroPipeline._states and free pool blocks on eviction via a new BDEKVState.close(), so session-id churn can no longer leak the KV pool until exhaustion. - P2 frame_seqlen source of truth: read it from the loaded transformer; the deploy-config inference is now only a drift cross-check (warns on mismatch). This surfaced a real 880-vs-220 mismatch on DreamZero-DROID. - KV pool floor: guarantee 2*(window_chunks+1)+2 blocks so the memory fraction can't under-size the pool (fixes exhaustion exposed by the corrected 880-token frame granularity). - P2 benchmark config helper: frame-granular (chunk_size=frame_seqlen, window_chunks=local_attn_size), matching the runtime. - P3: remove the drifted, unused BDEPipelineMixin (its update path was missing the required seq_len argument). Refactor: - Move vllm_omni/bde -> vllm_omni/experimental/bde; update all imports, the engine resolver, and the dynamic runner-class path string. Profiling: - BDE_KV_NO_MEMO env toggle to A/B the per-forward gather memoization. - gpu_mem_stats() worker RPC + peak-VRAM capture in bde_perf_compare. Validated on H200 (DreamZero-DROID): tests/bde 73 passing; bit-exact parity (max_abs_diff=0, PSNR=inf) for KV-on and for both memo/no-memo vs the KV-off baseline. Three-way profile confirms the memoization win (-9.25 GiB peak VRAM, -101 ms/chunk, sawtooth flattened). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Taichang Zhou <tzhouam@connect.ust.hk>

…commit_chunk docs) Close inline review findings 2-5 from hsliuustc0106: - Inline 2: delete the two unused project_kv methods in causal_wan_model.py (_kv_populate_cross inlines the identical projection); fix the stale comment. - Inline 3: gate DreamZero->BDE engine/runner routing behind BDE_KV_ENABLE via resolve_default_engine_backend(), so non-BDE deploys stay on base DiffusionEngine and base is selectable again. +4 unit tests. - Inline 4: docstrings clarifying the two commit_chunk paths (test-only vs the real advance inside BDEKVState.update_kv_cache). - Inline 5: vllm-project#4425 structured-config coordination is blocked (not in this tree); deferral marker left at the field defs + a note in the remaining-work doc. Validated on H200 (DreamZero-DROID): tests/bde 77 passing; routing gate confirmed end-to-end (KV-off -> base engine, BDE_KV_ENABLE=1 -> BDE); bit-exact parity (max_abs_diff=0, PSNR=inf) preserved. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Taichang Zhou <tzhouam@connect.ust.hk>

Bring PR vllm-project#4534 up to date with upstream main (DreamZero CUDA graph / torch.compile / DiT step-cache vllm-project#4213, tp & cross-attn cache fix vllm-project#4154, GR00T, stepcache, etc.). Conflict resolution: - pipeline_dreamzero.py: * diffuse loop: adopt upstream's run_dit (DiT step-cache) wrapper, keep BDE's _kv_get / _kv_get_cross proxies for all KV reads. * reset: adopt upstream's reset_reason (session/inference) model; route both session and inference (window-boundary) resets through _kv_reset so the BDE pool resets too (parity). _kv_reset gains clear_video_latents. * take upstream's state.append_video_latents(video_out). * cross-attn: BDE now uses the model-local cross-attn cache (upstream vllm-project#4154 added k_img/v_img image-token caching that BDE's text-only pool lacked); _kv_get_cross returns model-local, _kv_populate_cross is a no-op. BDE's self-attn paged pool is unchanged. - causal_wan_model.py: silent auto-merge duplicated the crossattn_cache arg in _forward_blocks (upstream vllm-project#4154 and our wiring-gap fix are the same change) — de-duplicated to one param. - video_export_worker.py: keep both gpu_mem_stats (ours) and _latents_to_uint8_frames (upstream); switch gpu_mem_stats off the now-banned torch.cuda.* to torch.accelerator.* (TID251). - bde_perf_compare.py: update to upstream's new _build_observations signature. Validated on H200 (DreamZero-DROID): tests/bde 77 passing; whole package compiles; KV-off (base engine) and BDE_KV_ENABLE=1 both run on the merged code and stay bit-exact (max_abs_diff=0, PSNR=inf). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Taichang Zhou <tzhouam@connect.ust.hk>

…name) The experimental/bde imports sorted before vllm_omni.diffusion.* once the package moved under experimental/; ruff isort wants diffusion first. Fixes the pre-commit (ruff-check) CI check on PR vllm-project#4534. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Taichang Zhou <tzhouam@connect.ust.hk>

tzhouam · 2026-06-20T12:05:39Z

Inline findings:

BDEPipelineMixin: deleted (unused, and already drifted from the pipeline's own kv* proxies).
project_kv: deleted both unused methods; _kv_populate_cross was the only (inlined) projection path. Fixed the stale comment.
Routing default-on: DreamZero to BDE routing is now gated behind BDE_KV_ENABLE via a resolve_default_engine_backend() helper. With it unset, DreamZero stays on the base DiffusionEngine/runner (so the half-validated runner is off the default path, and base is selectable again); an explicit engine_backend still bypasses the gate. Added unit tests for both states.
commit_chunk: docstrings on both BDEKVCache.commit_chunk and BDEKVState.commit_chunk noting they're the test-only/no-op paths and that the real advance happens in update_kv_cache via on_chunk_committed().
[Config] Add structured VllmOmniConfig classes #4425 structured config: deferred deliberately (it isn't merged yet). Left a marker at the engine_backend/diffusion_model_runner_cls field defs and a remaining-work note to mirror the fields into the structured DiffusionConfig + parity test, and to sync the from_kwargs rebase order with whichever lands second.

Earlier review-comment items:

P1 session-KV leak: _bde_states is now an LRU bounded at the same cap as the pipeline, with a BDEKVState.close() that frees both CFG branches' pool blocks on eviction.
P2 config-helper granularity: frame-granular (chunk_size=frame_seqlen, window_chunks=local_attn_size).
P2 frame_seqlen: read from the loaded transformer; the deploy-config value is only a drift check now.

From the main rebase:

[Feat] DreamZero fix tp & cross attn cache #4154 added image-token k_img/v_img caching to the I2V cross-attn. I extended BDE's cross-attn pool to cache those too (alongside text k/v), so BDE again manages both self- and cross-attn KV (the section 5.2 design). Self-attn paged KV is unchanged.
Resolved conflicts with [Perf][DreamZero]dreamzero supports CUDA graph, torch.compile, and DiT Caching. #4213 (DiT step-cache / CUDA graph / compile) and #415op keeps BDE's KV proxies inside the new run_dit wrapper, and the new reset_reason (session/inference) model routes both reset types through BDE's pool reset.

This PR is the joint work of the AR-Diffusion / BDE group; crediting all contributors so the squash merge carries every co-author. Co-authored-by: amy-why-3459 <221522807+amy-why-3459@users.noreply.github.com> Co-authored-by: Yangshen Deng <47519522+TKONIY@users.noreply.github.com> Co-authored-by: Wang Fuyin <wfy2003324@163.com> Co-authored-by: Honghan Zhu <nenbaying@163.com> Co-authored-by: Match_yc <57976772+matchyc@users.noreply.github.com> Co-authored-by: Hyoseop Song <lvsy@proton.me> Co-authored-by: Wang Yiwen <121547057+yiwen101@users.noreply.github.com> Co-authored-by: Suli Wang <50828270+wsl2000@users.noreply.github.com> Co-authored-by: Wu Jiazhen <83007646+asukaqaq-s@users.noreply.github.com> Co-authored-by: Yukim1 <121286183+zwhzzz0821@users.noreply.github.com> Co-authored-by: wtz2333 <wtz2333@mail.ustc.edu.cn> Signed-off-by: Taichang Zhou <tzhouam@connect.ust.hk>

test_dreamzero_state_cache_access_requires_initialization exercised the model-local get_kv_caches/update_kv_cache/get_crossattn_caches accessors that the paged self-attention backend removed (all KV is engine-owned now). Replace it with the inverse contract: those accessors must be GONE from DreamZeroState so nothing can silently bypass the engine pool. Fixes the Simple unit-test CI failure; the remaining two failures in that bucket (serve_cli mock, voxcpm2 engine_extras) reproduce on origin/main and are unrelated to this PR. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Taichang Zhou <tzhouam@connect.ust.hk>

Gaohan123

LGTM. Thanks

main independently landed DreamZero prompt/neg-embed caching as the default (state.prompt_embeds + self._negative_prompt_embeds_cache, vllm-project#4534 BDE KV-cache + v0.24 rebase), superseding W7's flag-gated embed caches. Adopt main's caching as canonical and drop the now-redundant W7 paths: - remove _encode_positive_prompt / _encode_negative_prompt helpers and their VLLM_OMNI_CACHE_{,NEG_}PROMPT_EMBED flags + _POS_CACHE_MAX - drop tests/dreamzero/test_{pos,neg}_embed_cache_equiv.py (equivalence is moot now that caching is the baseline) - keep W7's unique contributions: host-side profiling harness (_hostprof_*) wrapping the tokenize/encode/neg/prefill/diffuse stages, and the scheduler deepcopy skip (VLLM_OMNI_SCHEDULER_FRESH_BUILD) - coexist with main's DZ_PHASE_TIMING (_pt) markers Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Since PR vllm-project#4534 DreamZero's attention KV (self- and cross-) lives in the AR-Diffusion engine's paged pool, so the adapter now mirrors only what DreamZeroState still owns: frame stitching, accumulated video latents, prompt-embed cache, the incremental VAE encoder stream, counters, and reset heuristics. The dense PagedKV placeholder is removed; RFC vllm-project#4480 Phase 1 reintroduces PagedKV as a handle over the engine's session KV state instead of a buffer. Equivalence tests updated for the new surface, including the window-reset preservation semantics. Signed-off-by: Wijnand Suijlen <wijnand.suijlen@huawei.com>

…lm-project#4366) (vllm-project#4534)

…, so the contract shows where the AR self-attention KV fits: a handle wrapping the engine-owned paged state from vllm-project#4534, not a plain buffer here. Signed-off-by: Wijnand Suijlen <wijnand.suijlen@huawei.com>

…lm-project#4366) (vllm-project#4534) Signed-off-by: Wojciech Kutak <wkutak@nvidia.com>

Since PR vllm-project#4534 DreamZero's attention KV (self- and cross-) lives in the AR-Diffusion engine's paged pool, so the adapter now mirrors only what DreamZeroState still owns: frame stitching, accumulated video latents, prompt-embed cache, the incremental VAE encoder stream, counters, and reset heuristics. The dense PagedKV placeholder is removed; RFC vllm-project#4480 Phase 1 reintroduces PagedKV as a handle over the engine's session KV state instead of a buffer. Equivalence tests updated for the new surface, including the window-reset preservation semantics. Signed-off-by: Wijnand Suijlen <wijnand.suijlen@huawei.com>

…, so the contract shows where the AR self-attention KV fits: a handle wrapping the engine-owned paged state from vllm-project#4534, not a plain buffer here. Signed-off-by: Wijnand Suijlen <wijnand.suijlen@huawei.com>

Since PR vllm-project#4534 DreamZero's attention KV (self- and cross-) lives in the AR-Diffusion engine's paged pool, so the adapter now mirrors only what DreamZeroState still owns: frame stitching, accumulated video latents, prompt-embed cache, the incremental VAE encoder stream, counters, and reset heuristics. The dense PagedKV placeholder is removed; RFC vllm-project#4480 Phase 1 reintroduces PagedKV as a handle over the engine's session KV state instead of a buffer. Equivalence tests updated for the new surface, including the window-reset preservation semantics. Signed-off-by: Wijnand Suijlen <wijnand.suijlen@huawei.com>

…, so the contract shows where the AR self-attention KV fits: a handle wrapping the engine-owned paged state from vllm-project#4534, not a plain buffer here. Signed-off-by: Wijnand Suijlen <wijnand.suijlen@huawei.com>

… deploy config DreamZeroPipeline requires the AR-Diffusion engine and raises at construction when engine_backend resolves to anything else. vllm-project#4534 added engine_backend to vllm_omni/deploy/dreamzero.yaml but not to dreamzero_tp1_cfg2.yaml, so every launch through the CFG-parallel config died during worker startup with: ValueError: DreamZeroPipeline requires the AR-Diffusion engine; set engine_backend: ...ARDiffusionEngine in the deploy config (got engine_backend='default'). This broke the weekly DreamZero e2e serving test, which passes that config explicitly, and examples/online_serving/dreamzero/run_server.sh, which defaults to it. Add the missing engine_backend, plus a CPU test asserting that every shipped dreamzero*.yaml stage selects an AR-Diffusion engine. The test resolves the import path through DiffusionEngine.resolve_engine_class rather than matching a string, so a stale or typo'd path fails too. Previously only an H100x2 slow-marked e2e test covered this. Fixes vllm-project#5046 Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

…lm-project#4366) (vllm-project#4534)

hsliuustc0106 reviewed Jun 18, 2026

View reviewed changes

tzhouam force-pushed the bde-kv-phase1 branch from 7bb9a64 to 500d7f8 Compare June 18, 2026 10:43

tzhouam marked this pull request as ready for review June 18, 2026 10:44

tzhouam requested review from Gaohan123, Isotr0py, RuixiangMa, SamitHuang, ZJY0516, david6666666, princepride, wtomin, yenuo26 and ywang96 as code owners June 18, 2026 10:44

tzhouam force-pushed the bde-kv-phase1 branch from 500d7f8 to ef3b67e Compare June 18, 2026 10:50

hsliuustc0106 added the core related to core module: cache, scheduler, engine, worker, modelrunner label Jun 18, 2026

hsliuustc0106 reviewed Jun 18, 2026

View reviewed changes

hsliuustc0106 mentioned this pull request Jun 19, 2026

[New Model]: DreamX-World #4570

Open

tzhouam added ready label to trigger buildkite CI and removed ready label to trigger buildkite CI labels Jul 4, 2026

Gaohan123 approved these changes Jul 4, 2026

View reviewed changes

Gaohan123 added ready label to trigger buildkite CI and removed ready label to trigger buildkite CI labels Jul 4, 2026

Merge branch 'main' into bde-kv-phase1

bd50533

Gaohan123 added ready label to trigger buildkite CI and removed ready label to trigger buildkite CI labels Jul 4, 2026

Gaohan123 merged commit 86bdcaf into vllm-project:main Jul 4, 2026
9 checks passed

This was referenced Jul 7, 2026

[Bug]: DreamZero crashes at startup on current main: ARDiffusionModelRunner.execute_model() got an unexpected keyword argument 'kv_prefetch_jobs'` #4940

Closed

[Bugfix] Accept kv_prefetch_jobs in ARDiffusionModelRunner.execute_model #4941

Merged

SamitHuang pushed a commit to SamitHuang/vllm-omni that referenced this pull request Jul 7, 2026

[BDE] Phase 1: engine-level KV cache management for DreamZero (RFC vl…

57af0bf

…lm-project#4366) (vllm-project#4534)

This was referenced Jul 8, 2026

vLLM-Omni Q3 Roadmap #4901

Closed

[Core] Session state manager for AR-diffusion world models (RFC #4480 Phase 0) #4487

Merged

wkutak pushed a commit to wkutak/vllm-omni-cosmos-genai-nim that referenced this pull request Jul 13, 2026

[BDE] Phase 1: engine-level KV cache management for DreamZero (RFC vl…

a423f9a

…lm-project#4366) (vllm-project#4534) Signed-off-by: Wojciech Kutak <wkutak@nvidia.com>

tzhouam mentioned this pull request Jul 27, 2026

[Bugfix] Select the AR-Diffusion engine in the DreamZero CFG-parallel deploy config #5439

Merged

ZJLi2013 mentioned this pull request Aug 6, 2026

[Perf][DreamZero] Fix vae encoder's torch.compile #4750

Closed

khairulkabir1661 pushed a commit to khairulkabir1661/vllm-omni that referenced this pull request Sep 25, 2026

[BDE] Phase 1: engine-level KV cache management for DreamZero (RFC vl…

93b8add

…lm-project#4366) (vllm-project#4534)

Navigation Menu

Repository navigation

[BDE] Phase 1: engine-level KV cache management for DreamZero (RFC #4366) - #4534

[BDE] Phase 1: engine-level KV cache management for DreamZero (RFC #4366)#4534
Gaohan123 merged 53 commits into
vllm-project:mainfrom
tzhouam:bde-kv-phase1

tzhouam commented Jun 18, 2026 •

edited

Loading

Uh oh!

hsliuustc0106 left a comment

Uh oh!

chatgpt-codex-connector Bot commented Jun 18, 2026

Uh oh!

hsliuustc0106 commented Jun 18, 2026

Uh oh!

hsliuustc0106 Jun 18, 2026

Uh oh!

hsliuustc0106 Jun 18, 2026

Uh oh!

hsliuustc0106 Jun 18, 2026

Uh oh!

hsliuustc0106 Jun 18, 2026

Uh oh!

hsliuustc0106 commented Jun 18, 2026

Uh oh!

hsliuustc0106 Jun 18, 2026

Uh oh!

hsliuustc0106 commented Jun 18, 2026

Uh oh!

tzhouam commented Jun 20, 2026

Uh oh!

Gaohan123 left a comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

		_log.info("BDE RESET [%s/%s] session KV cleared (window boundary)", pos_id, neg_id)


		class BDEPipelineMixin:

Conversation

tzhouam commented Jun 18, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What this does

Design

Test plan

Results (H200, 1 GPU, DreamZero, bf16)

Status

Update (2026-07-04): performance & memory optimization stack

Reproducing the numbers

Credits

Uh oh!

hsliuustc0106 left a comment

Choose a reason for hiding this comment

Uh oh!

chatgpt-codex-connector Bot commented Jun 18, 2026

Uh oh!

hsliuustc0106 commented Jun 18, 2026

Uh oh!

hsliuustc0106 Jun 18, 2026

Choose a reason for hiding this comment

Uh oh!

hsliuustc0106 Jun 18, 2026

Choose a reason for hiding this comment

Uh oh!

hsliuustc0106 Jun 18, 2026

Choose a reason for hiding this comment

Uh oh!

hsliuustc0106 Jun 18, 2026

Choose a reason for hiding this comment

Uh oh!

hsliuustc0106 commented Jun 18, 2026

Uh oh!

hsliuustc0106 Jun 18, 2026

Choose a reason for hiding this comment

Uh oh!

hsliuustc0106 commented Jun 18, 2026

Uh oh!

tzhouam commented Jun 20, 2026

Uh oh!

Gaohan123 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

tzhouam commented Jun 18, 2026 •

edited

Loading