[LingBot World] Eleven bit-identical removals of repeated work in the realtime path

[tzhouam]

Purpose

Eleven changes to the LingBot-World realtime path that leave the output bit-identical. Each removes work that reproduces a result an earlier forward already had, recomputes something that cannot change within an AR block or a session, or converts data on the wrong side of a device boundary.

change	what it removes
Decode a repeated source image once	re-decoding the same source image every tick
Build the camera modulation once per AR block	the camera injector in all 40 blocks on 4 of 5 forwards
Drop both cross-attention all-to-alls under Ulysses	2 collectives per layer per forward
Restage only the K/V that moved within an AR block	re-gathering the visible window every forward (opt-in, reuse_history_staging)
Copy the block-table metadata to the device asynchronously	three synchronous pageable H2D copies per AR block
Deliver realtime chunks as uint8 frames straight from the device	a float32 host copy, a numpy multiply/round/cast and twelve PIL images per chunk, then the endpoint's coercion back to uint8
Stop re-encoding the condition once the encoder is at a fixed point	three causal-encoder passes per block that, after the opening block, return the same condition and the same history every time
Open the camera reuse window across a stepwise block too	the camera injector on all five forwards of a served block: the window used to be opened only by generate_block, which the realtime server never calls
Run a realtime block's steps to the chunk boundary in one runner call (pipeline capability + AR-runner policy)	three scheduler/executor round trips per chunk, each of which also let the GPU queue drain
Stage the timestep projection in one static buffer per shape	200 x 72 MB device-to-device copies per chunk: every regional block graph re-copied the fresh per-forward projection into its placeholder
Assemble the sharded Wan causal conv input in one pass	two of the three full writes of the activation per causal conv call (cat, pad, cat)

All are verified bit-identical against the unoptimized path: the first four over 40 saved chunks (max_abs 0.0) teacher-forced against a bitwise-repeatable eager reference; the uint8 conversion against VideoProcessor itself and byte-for-byte on the fragmented-MP4 output of the real server; the fixed point by construction (a block whose returned encoder history is bit-identical to the history it was given proves the encoder is at a fixed point, and every block is still checked against the whole-clip encode in the tests).

Measured: served path, cumulative

arm (cumulative)	steady median ms	std	p99	delta	RTF @12 fps	RTF @16 fps	stall
stock main (e515cad)	1326	35	1396		1.32	1.77	43%
stock again, end of chain (run-to-run)	1340	40	1394	+14	1.34	1.79	43%
+ source-image cache, camera modulation once per block, cross-attention without all-to-alls, K/V restage (reuse_history_staging: true)	1338	39	1417	+12	1.35	1.80	44%
+ async block-table metadata	1326	37	1354	-13	1.32	1.76	42%
+ uint8 frames from the device	1209	43	1305	-117	1.21	1.61	37%
+ condition fixed point (lossless total)	1026	60	1211	-183	1.05	1.40	29%
+ camera window on the stepwise block (measured last, on top of #7651's width shard at 849 ms)	787	60	968	-62	0.81	1.08	9%
same, three measured sessions (n=99), as the base for the next row	777	54	967		0.80	1.06	8%
+ steps grouped per chunk in the AR runner, three measured sessions (n=99)	736	54	925	-41	0.76	1.01	4%
base again after it, three sessions (run-to-run control)	779	54	980	+2	0.80	1.07	8%
+ static timestep buffer, three sessions, other card set (control on that set: 715)	726	55	918		0.78	1.03	5%
+ single-pass conv input, three sessions, same card set	715	54	906		0.76	1.01	4%
both PRs at their current heads (this PR's review rounds included, #7651's shard on top), three sessions; run again as control: 718	721	54	906		0.74	0.99	4%
This PR alone (every change above, no #7651 options), three sessions	916	56	1115		0.94	1.25	21%
this branch at bf3be8b run ungrouped (one scheduler cycle per step), three sessions	969	55	1156	+53	0.99	1.32	25%
review round 2: grouping loop moved into the AR runner, batch rebuilt per step, three sessions	966	53	1153	+50	0.99	1.31	25%
+ sync-free step-batch index mapping (lossless, every stepwise pipeline), three sessions (run again as control: 916)	920	57	1131	-46	0.94	1.26	22%

Run-to-run on this protocol: the same code measured 736 and 715 ms on two different four-card sets (three sessions each), so the resolution of any single row is about ±10 ms. The last two commits sit inside that floor served: the static timestep buffer is justified by the campaign's same-run ledger (200 -> 5 device copies per chunk, -7 ms of DiT wall in its harness) rather than by a served delta, and the conv input assembly removes memory traffic the sharded decode turns out not to be bound by at 480x832. Both are bit-exact and tested; neither is claimed as a served gain.

This PR alone, with the stock VAE decode, takes the served path from 1326 to 920 ms per chunk at the current head (1.44x; 916 before the review rounds, inside the run-to-run floor), which is real time at the 12 fps basis (RTF 0.94). With #7651 on top, both PRs at their current heads, it is 721 ms (718 on the control run); the same combination measured 715-736 ms before this PR's review rounds, so those rounds cost nothing on the all-in path either.
Protocol: 4xH200 (GPUs on one NUMA node), 480x832, benchmarks/lingbot_world/configs/usp4_compiled.yaml (Ulysses SP-4, regional torch.compile + CUDA graphs, dtype: bfloat16), the WebSocket realtime endpoint, benchmark_lingbot_world_realtime.py with one warm-up session and one measured 40-chunk session per arm, steady state = chunks after the sixth. Every arm is a fresh server on the exact commit, in one held reservation, on an idle host (arms whose std exceeded 60 ms were re-measured; the stock arm was measured at both ends of the chain and moved 14 ms). Single-chunk end-to-end latency (session tick to muxed bytes) tracks cadence plus ~45 ms of overlapped delivery on every arm.

Reading the first rows honestly: the four original changes are worth nothing on the served path as they stood. The camera cache's window was only ever opened by generate_block, which the stepwise server path does not call, so the injector still ran five times per block; that is what the last commit fixes, and it is the -62 ms row. The source-image cache cannot fire either, because a realtime tick carries its image inline rather than as a path. The cross-attention change and the K/V restage are DiT-side and within the run-to-run spread here. The lossless set is stock 1326 -> 1026 ms (-300, 1.29x); with the camera window it is -362 on the DiT/host side. In-process numbers on an older main that this PR used to quote are gone; these replace them.

Driver-rank stage medians per chunk (ms), from a wall-clock overlay; the pipeline profiler (--enable-diffusion-pipeline-profiler, device-synchronised) puts the width-sharded streaming decode at 104 ms of GPU time where the unsynchronised span shows 54, so the difference sits in the "GPU tail" row:

stage	stock	lossless (E)	all (I)
denoise_step x4 (launch wall)	314	309	293
commit_block_kv (5th forward, derived)	128	95	68
decode_to_pixels (streaming VAE + uint8, includes the sync drain)	477	409	225
prepare_next_chunk (condition encode until it freezes)	182	3	3
GPU tail drained at the chunk barrier (derived)	201	192	177
outside the model runner	25	20	17
chunk interval	1327	1027	783

Decisions worth reviewing

Camera reuse is scoped to one AR block. The caller constructs a CameraModulationCache per block and passes it explicitly as camera_modulation_cache=. generate_block shares it across its forwards; the stepwise pipeline keeps it in request state across the four probes and commit, then drops it. Interleaved sessions therefore hold separate caches. The caller must keep the camera fixed for that block. The retained scale and shift tensors are included in ARDiffusionKVCacheSpec.model_owned_state_bytes_per_session before KV admission: at 480x832, 40 layers and bf16 they cost 958,464,000 bytes per session per rank under Ulysses SP-4, or 3,833,856,000 bytes without sequence sharding.

Cross-attention replicates text K/V instead of sharding it. Every query token attends the whole text sequence independently, so the sequence split Ulysses already provides is the only split needed. The cost is text K/V per rank rather than head-sharded (about 314 MB/rank at this size). ARDiffusionCrossAttentionKVSpec grows an optional num_kv_heads; absent, it still means "same as self-attention".

K/V restaging is off by default (ARDiffusionKVConfig.reuse_history_staging): two max-sequence buffers per layer per rank, allocated by the KV manager alongside the pools outside any capture, mark_static_address, declared in the custom op's mutates_args. The measurement row above is the first for this implementation (the earlier number came from a Triton runtime patch).

Realtime ticks always deliver (F, H, W, 3) uint8. The streaming endpoint muxes uint8 frames, so pil and np only ever differed in how much work stood between the decoder and those bytes; only latent keeps its meaning. The conversion is VideoProcessor's arithmetic in its order, on the device.

The condition fixed point is detected, never assumed. It is per session, cleared with the session's encoder history, and a session restarted from a new source image encodes again until it settles. On the served path it settles by the eleventh block.

Step grouping is a pipeline capability plus an AR-runner policy, not a public switch, and the shared runner stays single-step. Step execution hands the runner one denoise step per scheduler cycle; for a lone streaming request that is four hops per chunk through the scheduler, the executor and the workers, with the GPU queue draining at each. Grouping the steps to the chunk boundary removes the hops but also means the scheduler gets no control between them, so it is not folded into step_execution and not exposed as a user-facing Boolean either (an earlier revision had stepwise_chunk_per_call; review argued, rightly, that a single-request batch is not a scheduling argument). Instead: a step-execution pipeline may declare the optional supports_chunk_step_grouping = True (LingBot World does: a block is four probes and a commit over one prepared chunk state), and everything else lives in experimental/ar_diffusion/runner.py. DiffusionModelRunner._execute_stepwise_core() keeps its original one-step-per-call semantics with no policy hook; ARDiffusionModelRunner overrides it to run the parent's core once with the scheduler's request and then, for a lone streaming request of a pipeline that declares the capability, to keep calling that same core with a continuation of the scheduler output (the request as a cached request, nothing new to admit, nothing to retire, no prefetch job or connector metadata) until a call emits a result, finishes the request or leaves it without state. Grouped execution is capped at the declared chunk_num_steps; failure to reach a boundary returns a request error and drops its state. Every grouped step is thus exactly what the next scheduler cycle would have asked for, batch rebuild included, and it all runs inside the one session KV binding and the fail-closed scope the AR runner already wraps around a stepwise call; the chunk barrier is unchanged and no per-step device sync is added. The AR path regains control at every chunk boundary, where its interactions are applied, and the pipeline's interrupt is still checked between grouped steps. Other pipelines keep step granularity until they declare the capability and are covered by equivalent tests. Measured served with three sessions per arm: 777 -> 736 ms with the loop in the shared runner; re-measured after moving it into the AR runner: 966 ms, because the per-step batch rebuild that review asked for re-entered InputBatch.make_batch, whose index mapping was built on the device and read straight back (tolist() and cpu()), a device synchronisation between every two steps that gave the grouping gain back. Building that mapping on the host (next paragraph) brings the same code to 920 ms.

Nothing model-specific in shared modules. The shared changes are generic AR-diffusion ones (paged attention staging, cross-attention K/V head count, one pinned async copy helper already used by the slot mappings). A benchmarking env override that an earlier revision of this branch had put into the shared model runner is gone.

Review fixes (2026-09-17)

Three findings on the opt-in paths, all reproduced on CPU before fixing:

The step batch's index mapping is built without a device read-back. InputBatch.make_batch created its identity index mapping with torch.arange on the latents' device and then called tolist() and cpu() on it, two host waits for everything queued on the stream, on every batch build, which is every step of every stepwise request. That is the ~170 ms per chunk of the host waiting inside the runner that the ungrouped served path always showed, and it is what re-appeared when the batch was rebuilt between grouped steps. The identity mapping is known on the host: the device tensor is now a plain launch and the host array comes from NumPy; only a caller-supplied mapping is still read back. Same ids, order, device tensor and host array as before; the test forbids Tensor.tolist and Tensor.cpu during the default build and the cached repack. Measured served, three sessions per arm: grouped 966 -> 920 ms; the same fix applies to the ungrouped path of every step-execution pipeline.

• Step grouping design. Replaced the public stepwise_chunk_per_call Boolean with the capability/policy split described above; no user-facing setting is added. A second round moved the grouping loop itself out of the shared runner into ARDiffusionModelRunner._execute_stepwise_core(), so diffusion_model_runner.py carries no model-path code for it.
• Whole-chunk call reused a stale batch. The loop ran every step of a chunk over the InputBatch built before it; a denoiser that reads latents or timesteps from the batch (Helios stage 1) saw the first step's values throughout. The batch is now rebuilt between steps exactly as the next scheduler cycle would; the test drives a denoiser that depends on both fields through both modes and checks the per-step inputs and the final latent are identical.
• K/V restage refreshed the wrong blocks. The reuse path took the current chunk to be the last current_blocks entries of the padded table, which are padding (the action-capacity block, plus unused window capacity while the window grows), so the current rows from the previous probe stayed in place and were read. The refresh now starts at the visible history's block count, carried from the host-side metadata build_block_table already computes, no device sync. Test: two probes of one block through the real table layout with an empty, partial, full and slid window, staged buffers equal to a full gather byte for byte.
• Staging buffers were outside the KV budget and allocated without a consumer. They are now reserved before capacity and pool size are chosen (tests show capacity dropping and the one-session check rejecting at the boundary), and both the reservation and the allocation are gated on the contiguous-gather path that is their only consumer.

Prerequisites

The branch carries [LingBot-World] Cut launch/collective count under Ulysses + CUDA graphs and [AR-Diffusion] Allocate the flat K/V slot pools as the base tensors underneath, which the K/V restaging extends; they can be split into a prerequisite PR if preferred.

Test Plan

• pytest tests/diffusion/ar_diffusion tests/diffusion/models/lingbot_world tests/diffusion/distributed/test_comm.py on CPU.
• pytest tests/diffusion/models/lingbot_world/test_lingbot_world_transformer.py -k "sp2 or sp4 or tp2" on 4xH200 under a GPU reservation.
• pre-commit run ruff-check --all-files and ruff-format --all-files.

New tests pin the properties each change relies on: that a decode cache never serves a file rewritten under it; that the camera modulation depends only on the camera tokens and the injector runs once per window; that cross-attention K/V is replicated rather than sharded; that staged reuse equals a full gather byte-for-byte; that every pixel output type streams uint8 frames and the bytes equal the PIL path's; that the encoder stops once its history stops changing, restarts with the session, and every block still equals the whole-clip encode; that the five forwards of a stepwise block share one camera cache and the next block gets its own; that a chunk run in one call emits the same outputs at the same step indices as one run step by step; that two forwards' timestep projections land in one buffer per shape; that the single-pass conv input equals the cat/pad/halo sequence byte for byte on one rank and on both ranks of a pair, and equals the stock conv for a halo-free kernel.

Test Result

• CPU: tests/diffusion/ar_diffusion tests/diffusion/models/lingbot_world tests/diffusion/distributed/test_comm.py -- 318 passed, the 3 multi-GPU parity tests deselected; the lingbot_world suites alone -- 263 passed; tests/diffusion/test_diffusion_step_pipeline.py tests/diffusion/test_diffusion_config_fields.py tests/config -- 619 passed, 1 skipped.
• CI on the pushed head: DCO, pre-commit (ruff, typos and the rest), both builds and the docs build pass. The repo's mypy hook is CI-skipped by its own config; its reports on the touched files are pre-existing signatures on main.
• 4xH200: test_lingbot_world_transformer.py -k "sp2 or sp4 or tp2" -- 3 passed (164 s) on the pushed head under a GPU reservation.
• pre-commit run ruff-check --all-files / ruff-format --all-files -- Passed.
• Every commit carries a Signed-off-by (the two prerequisite commits did not, which is what DCO was failing on).

vLLM Version: 0.29.0
vLLM-Omni Commit: branched from 507cb1d83

🤖 Generated with Claude Code

CPU property tests provide the standing regression coverage; the 4xH200 parity, CUDA-graph and served timings above are author-run measurements, not results from ordinary CI. They have not been remeasured for this merge and review cleanup.

Validation after merging main: 295 passed, 5 skipped across the LingBot pipeline, stepwise runner, Wan spatial sharding, AR KV cache and capability-runner tests. Other pre-commit checks passed; mypy still reports existing typing errors in surrounding code.

[chatgpt-codex-connector]

Codex Review Summary

This comment shows the latest Codex review activity on this pull request.

Review	Status	Commit	Review trigger
📝 Code Review	✅ Completed 2026-09-16T11:48:31.647465Z	1197ef8	PR opened

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review" or "@codex security review".

Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings.

[vllm-omni-review-bot]

This PR appears to belong to: docs/design/module/cache_management.md, docs/design/module/diffusion/offloader.md, docs/design/module/diffusion/index.md, docs/design/module/diffusion/diffusion_model_integration.md.

Module owners: @RuixiangMa @wtomin @david6666666

Routing: @RuixiangMa via module of the changed files, module named in the PR description, semantic router; @wtomin via module of the changed files, module named in the PR description, semantic router; @david6666666 via module of the changed files, module named in the PR description

@tzhouam, please review your own changes and leave a short self-review comment describing what you checked. PRs without author self-review may not be assigned a reviewer.

Please take a look when you have a chance. If you would like an automated review, mention @vllm-omni-review-bot in a comment.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 1197ef8a79

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[chatgpt-codex-connector]

Refresh current blocks at their actual table offset

When reuse_history_staging=True and LINGBOT_KV_GATHER=1, the second probe of an AR block takes this reuse path, but n_blocks includes the action-capacity padding block that build_block_table() reserves even when action_len == 0 (and may include additional unused window padding). Consequently, the current video blocks are not the last current_blocks entries: this copies null padding while leaving the staged current K/V from the previous probe unchanged, and those stale rows are still read because seq_lens includes the current video tokens. Compute the refresh start from the current blocks' position in block_ids, not from the padded table width.

Useful? React with 👍 / 👎.

[tzhouam]

Fixed in c5b20b0. The refresh now starts at the visible history's block count (kv_len // block_size - num_current_video_blocks, carried from the host-side metadata build_block_table already computes, as stage_first_block on the layer inputs) instead of n_blocks - current_blocks, so the padding blocks are never mistaken for the current chunk and nothing is synchronised inside attention. Test: two probes of one block through the real build_block_table layout with an empty, partial, full and slid window; the staged buffers equal a full gather byte for byte, and the old end-of-table guess is asserted to differ from the live offset.

[chatgpt-codex-connector]

Reserve staging buffers in the KV memory budget

With reuse_history_staging=True, this lazily allocates a full-length K and V buffer for every layer, but ARDiffusionKVCache._required_bytes() accounts only for scratch, cross-attention, model-owned state, and managed blocks. On memory-constrained configurations—especially when the one-session calculation already raises the budget close to available_bytes—capacity admission succeeds and the first prepared forward can then OOM while allocating these unbudgeted buffers. Include this worker-wide staging allocation in the preallocation budget before selecting session capacity and pool size.

Useful? React with 👍 / 👎.

[tzhouam]

Fixed in c5b20b0. The manager now reserves the staging pair before choosing session capacity and pool size: 2 x layers x (sink + window + the action-capacity block build_block_table always reserves) tokens x heads x head size x element size, added to _required_bytes and subtracted from the self-attention pool budget. Both the reservation and the lazy allocation are gated on the contiguous-gather path that is the only consumer (LINGBOT_KV_GATHER=1, via one helper in config.py), so the config flag alone no longer keeps the buffers for nothing. Tests: with the same geometry as the capacity-two case, staging drops the capacity from two to one, and the one-session check rejects the boundary budget that fit without staging.

[vllm-omni-review-bot]

Omni ReviewBot triage note

Automated triage of commit 4caf8bb33b9f produced:

• Priority: high. Prompt maintainer attention is suggested.

These are automated triage suggestions only — the final decision belongs to the maintainers.

[princepride]

Reviewed the full diff against 507cb1d8, focusing on dead code, over-defensive code and over-abstraction. Inline comments below; this is the summary.

Scope. By the PR's own measurements, several of the eleven changes produce no served gain or never execute on the served path:

• the source-image cache "cannot fire" (the tick image is inline);
• cross-attention without all-to-all and the K/V restage are within run-to-run noise;
• the static timestep buffer and the single-pass conv input are explicitly "not claimed as a served gain".

The K/V restage is the largest of these: ~250 lines in paged_attention.py plus budget logic in the manager, off by default, gated by both a config flag and an env var. The 1326 → 920 ms comes from uint8 output, the condition fixed point, the camera window, step grouping and the sync removals.

I'd suggest splitting the unmeasured items out into follow-ups that come with served numbers. That roughly halves this PR and makes the rest much easier to review.

Summary of inline comments:

• Dead code: source-image cache; write-only session_state.prompt; leftover Ulysses state, num_sp_heads alias and shard_kv_heads in cross-attention; unused device parameter; stale reuse_history type/comment; a redundant slice branch.
• Over-defensive code: getattr defaults on attributes that always exist; a silent fallback in the restage path; data_ptr-keyed observe(); get/pop/nullcontext for an always-present camera_cache; a requires_grad guard in inference; the env var read in two places.
• Over-abstraction: two mechanisms for camera reuse; staging state monkey-patched onto the KV cache as a dict; the grouping capability placed in the shared protocol; _assemble_input.
• Behavior change: stepwise vs. request mode now disagree on output_type.

[princepride]

Dead code on the served path. Per the PR description, this cache "cannot fire" on the realtime server because a tick carries its image inline rather than as a path. Yet it brings a module-level dict + threading.Lock, a 5-field stat fingerprint, a second stat after decode, a copy on both get and put, and a broad except fallthrough.

Suggest dropping this commit from the PR. If the offline/request path really needs it, functools.lru_cache(maxsize=1) keyed on the fingerprint would cover it; a one-entry cache implemented as dict + clear() is also more machinery than needed.

[tzhouam]

Kept, but reduced to what it needs: the module-level dict, lock, second stat and broad except are gone; it is now functools.lru_cache(maxsize=1) on the stat fingerprint, with a copy on return (5127705). On the served path it indeed never fires (the tick carries the image inline); it is there for request mode, where the same source path is replayed per request: in-process A/B on 4xH200 480x832, 732.7/732.3 ms original vs 714.6/716.2 ms cached (n=20 each, two independent launches per arm, 40/40 bit-exact eager gate), i.e. -17 ms. The PR description states the served result (no gain) for it.

[princepride]

_LingBotARSessionState.prompt is now write-only: its only reader was prompt_changed, which this PR replaced with _session_text_caches / text_cache_key, but session_state.prompt = inputs.prompt (L1585) is still assigned. Please remove the field and the assignment.

[tzhouam]

Removed the field and the assignment (5127705). The one remaining prompt read on that path is on StepRequestState, a different object.

[princepride]

Behavior change worth making explicit: on the stepwise path every pixel output_type (pil, np, pt) now returns a (F, H, W, 3) uint8 ndarray, while request mode (L543) still goes through VideoProcessor and honors output_type. The same pipeline now gives different types for the same parameter depending on execution mode. If intended, please validate/document it at the request level rather than only in this docstring.

[tzhouam]

Documented at the request boundary (5127705): the _LingBotRequestInputs.output_type field now states that step execution returns every pixel type as (F, H, W, 3) uint8 frames (what the streaming endpoint muxes) and only latent differs, while request mode goes through VideoProcessor. The endpoint consumes uint8 frames whatever the client named, so I kept the behaviour rather than making a pixel output_type an error on sessions; happy to reject non-np pixel types there instead if you prefer a hard boundary.

[princepride]

Leftover dead state after dropping the cross-attention all-to-alls:

• self.ulysses_world_size / ulysses_rank / ulysses_group have no readers in this class anymore.
• self.num_sp_heads = self.num_local_heads (L440) is just an alias.
• shard_kv_heads (L490) no longer shards; its docstring says it is "kept as a seam". It now only does clone(), which is an extra copy of the text K/V per layer: linear(...).unflatten(2, ...) is already contiguous.

Suggest removing the method and the alias and using the tensors directly (or .contiguous() at most).

[tzhouam]

Removed (5127705): the Ulysses state, the num_sp_heads alias and shard_kv_heads are gone; the cross-attention cache takes the projection outputs directly (they are already contiguous), the pipeline's pre-population does the same, and the test that pinned the clone went with the method.

[tzhouam]

Follow-up: 5127705 missed one reader of the removed alias, the cross-attention K/V spec in the pipeline (hidden from the CPU pipeline tests by their stand-in transformer; the served workers failed at model load). Fixed in 6fa891b, which reads num_local_heads there and in the tests' stand-ins.

[princepride]

observe() is both redundant and not a sound guard:

• Redundant: both paths already scope the cache to exactly one block (stepwise creates one per block; generate_block uses the with window), and the camera cannot change within it.
• Not sound: keying on data_ptr() can alias under the CUDA caching allocator — a different camera tensor of the same shape can land at the same address — so the docstring's claim that a misplaced window "costs the speedup but never correctness" doesn't hold.

Either trust the caller-owned lifetime and delete observe, or compare by value if a guard is really wanted.

[tzhouam]

Deleted observe() (5127705) and with it the data_ptr key: the cache now only holds entries, its lifetime is the caller's, and the docstring says so instead of claiming a guard.

[princepride]

block_ids if first_block == 0 else block_ids[first_block:] — both branches are equivalent (block_ids[0:]); just use block_ids[first_block:].

[tzhouam]

Done (5127705).

[princepride]

LINGBOT_KV_GATHER is read two different ways: frozen at import here as _KV_GATHER_ENABLED, but called at runtime via contiguous_kv_gather_enabled() in manager.py L302. If the env var is set after import (common in tests), the budget reservation and the actual allocation disagree. Please read it in one place.

Separately, kv_cache/config.py is a shared AR-diffusion module but now carries a LINGBOT_* env var reader, which contradicts "nothing model-specific in shared modules". Having both an env var and a config flag (reuse_history_staging) gate one feature is also one knob too many.

[tzhouam]

One reader now (5127705): contiguous_kv_gather_enabled() in kv_cache/config.py is called where the decision is made, by the manager when it budgets and allocates staging and by the attention dispatch, so a value set after import is seen by both; the frozen module constant is gone. The switch is renamed VLLM_OMNI_AR_DIFFUSION_KV_GATHER: it selects an experimental attention path of the shared AR-diffusion module, not anything LingBot-specific. On the two knobs: staging is a sub-option of that path (it only has a consumer there), so the manager now warns when reuse_history_staging is set while the path is off, instead of silently doing nothing. I kept them separate because the gather path is useful on its own and the staging buffers cost two max-sequence buffers per layer per rank.

[princepride]

Silent fallback: if any of these four conditions fails (current_tokens > 0, divisibility, 0 <= stage_first_block, upper bound), we quietly do a full restage. None of them should fail for a correct caller; if one does, the metadata is wrong, and this hides it as "the optimization didn't kick in" — the offset bug Codex found is exactly this class. Please make these asserts (or drop them).

[tzhouam]

Asserts now (5127705), only on the reuse branch: a block-aligned current chunk and stage_first_block + current_blocks <= n_blocks. A wrong offset fails instead of restaging; the non-reuse branch still restages the full window, which is its job.

[princepride]

supports_chunk_step_grouping extends the shared step-execution protocol, but its only consumer is ARDiffusionModelRunner and its only declarer is LingBot. Suggest moving the capability into the AR protocol in experimental/ar_diffusion/capability.py. Also, the supports_step_execution(pipeline) check in this helper is redundant: anything reaching _execute_stepwise_core already supports step execution.

[tzhouam]

Moved to experimental/ar_diffusion/capability.py (5127705) and the redundant supports_step_execution check dropped; the runner imports it from there.

[princepride]

Per the PR description this is not claimed as a served gain ("the sharded decode turns out not to be bound by" this traffic at 480x832), yet it replaces cat / F.pad / halo_exchange with ~75 lines of hand-written indexing. Two details also don't match the description: the halo region is zeroed and then overwritten, and halos are received into separate recv buffers and then copied in, rather than "received into their columns" as the commit message says.

A definite readability cost with no measured benefit. Suggest dropping this commit from this PR.

[tzhouam]

Kept, with the two inaccuracies fixed (5127705): the halo slots are no longer zeroed before being written, and the docstring now says what the code does: the neighbours' rows are received into contiguous buffers because a P2P op needs a contiguous tensor and the halo slot of a 5D tensor is a strided view, then copied into their slots. On the measurement: it removes two of the three full writes per causal conv call (cat, pad, cat), which is memory traffic that the ledger shows and that matters below 480x832 sharding as well; on the served 4xH200 480x832 arm it sits inside the ±10 ms floor, as the description says. I would rather keep it here than carry it as a follow-up, because the width-sharded decode of #7651 and the decoder fast path in #7749 both build on this one-pass input.

[princepride]

LGTM

[princepride]

Reviewed the full diff at 917715e1b, with an eye on whether the two earlier rounds (Codex's P1/P2, @princepride's set) actually landed rather than just got answered.

What I checked and believe holds:

• The restage offset fix. stage_first_block = kv_len // block_size - num_current_video_blocks really is the visible-history block count, so the refresh no longer lands on the trailing padding blocks; the new empty/partial/full/slid-window test drives the real table layout.
• The staging budget. history_staging_tokens = sliding_window + sink_chunks*chunk_size + block_size is exactly max_seq_len when action_len == 0, so the capacity check in _prepare_history_staging cannot spuriously fire, and the reservation is taken before capacity selection.
• Staging under CFG. The buffers are per-layer and shared across KV branches, but HistoryStagingState.adapter is a weakref to the per-(session, branch) adapter, so an alternating branch or session drops reuse_history back to a full gather. Correct, and the weakref also rules out an adapter-address ABA. (Cost is that the optimisation never fires with more than one branch, which seems fine for an opt-in path.)
• The condition fixed point. For start_frame > 0 the encoder's only inputs are the zero frames and pending_cache, and _condition_from_latents(first_frame=False) does not read start_frame, so "same input + same state implies same output and same state" is sound rather than empirical.
• _assemble_input's halo layout. buf[halo:2*halo] / buf[extent:extent+halo] are the same rows F.pad + halo_exchange sent, the send and receive slices do not overlap, and local_extent matches the old x.shape[dim] after padding.
• _to_device_async. pin_memory().to(non_blocking=True) is lifetime-safe here; the caching host allocator records a stream event on the copy.

Three comments inline. Only the first is one I'd want an answer on before merge.

Small extra: kv_cache/manager.py now defines both logger = init_logger(__name__) (new) and the pre-existing _log = init_logger(__name__). Two names for the same logger; worth collapsing to one.

[Dong1017]

Nice work — no blockers from my read. Each of the eleven changes removes work that looks provably redundant within one AR block or one session, and each carries a targeted property test. I traced the load-bearing invariants: the single-pass conv input matches the old cat → F.pad → halo_exchange layout exactly, the staging-reuse signature and stage_first_block offset are consistent with the padded block-table layout (and action tokens disable staging), the staging buffers are reserved inside the KV budget before admission, the camera modulation cache is request-scoped and rebuilt when the camera changes, and the condition fixed point is detected by bit-identical history comparison rather than assumed. The new text-cache key also replaces the removed session_state.prompt tracking with equivalent-or-stricter invalidation, as far as I can tell.

Two small notes:

• The PR description mentions transformer.reuse_camera_modulation() and new_camera_modulation_cache; I couldn't find either name at head. The implemented contract seems to be the caller-held CameraModulationCache passed via camera_modulation_cache= (transformer.py:771, per-block construction in generate_block / _prepare_next_chunk). Behavior matches the description as far as I can tell, but the names might be worth syncing.
• My guess is the 4xH200 parity and CUDA-graph paths are author-run only (CI here doesn't exercise them), with the CPU property tests as the standing guard — worth keeping in mind for future regressions.

[Dong1017]

Small nit: this adds a second module logger right above the existing _log = init_logger(__name__) two lines down, and the new staging warning uses logger while the rest of the module uses _log. Might be simplest to drop this binding and call _log.warning at line 339.

[Dong1017]

Non-blocking thought: _uint8_frames is now the third near-copy of the denorm→uint8 conversion (postprocess/device_reduction.py, dreamzero/video_export_worker.py:_latents_to_uint8_frames). This one is deliberately pinned byte-for-byte to VideoProcessor's arithmetic, so no change requested — just noting that a shared "PIL-identical" helper could keep the three from drifting apart someday.