[MrlixiangWE]

Purpose

Part of #7679: add the generic Scheduler/Worker path for request-scoped KV contexts in SenseNova-U1.5 B1. The model consumer and its two-request parity and throughput validation remain the following B1 slice.

Propose

1. Deduplicate contexts shared by execution sequences and allocate primary and context rows as one transaction. Roll back every attempted row if native allocation raises after changing block ownership; retain prefix pins until rollback finishes.
2. Enforce the per-request row limit at Scheduler admission; a request larger than that limit or than the whole Worker row pool fails before allocating blocks. Worker rows hold only scheduled requests, and the factory sizes the pool as max_num_seqs * diffusion_kv_max_rows_per_request, so prefetched and draining reservations, which hold blocks only, do not count against it.
3. Keep context-only native shims outside the primary sequence collection, so sequence-only prefix publication cannot consume context allocations.
4. Carry cache roles through Worker installation and native attention dispatch. Reject unregistered roles, mixed-role rows in one forward and layers consuming another role's cache.
5. Release primary and context rows together, preserving transfer-drain and allocation-generation ownership. Run both row types in the existing CUDA ready lane.

Test Plan

vLLM Version: 0.32.0.

vLLM-Omni Commit: 2227443 (410f0ab merged with main 3266a6f).

python -m pytest -q -o addopts= --strict-markers --strict-config \
  tests/diffusion/diffusion_kv tests/diffusion/test_diffusion_scheduler.py \
  tests/diffusion/models/hunyuan_image3/test_diffusion_kv_request.py -m 'not cuda'
CUDA_VISIBLE_DEVICES=0 python -m pytest -q -o addopts= \
  --strict-markers --strict-config tests/diffusion/diffusion_kv/ \
  -m 'core_model and cuda and L4 and cards_1' --run-level core_model
SKIP=actionlint pre-commit run --files <changed files>

Test Result

Executed on NVIDIA A800 80GB with each worktree's import path verified.

Check	Result
KV, Scheduler and HunyuanImage3 KV-request suites, not cuda selector	548 passed, 4 deselected
diffusion_kv CUDA L4 selection, including primary and ar_decode context rows against SDPA	4 passed
Prefetch, manager and Scheduler suites with the aggregate Worker-row count restored in manager.py	8 failed, 196 passed
Changed-file hooks, including mypy	Passed

The native allocation regression injects an exception after real allocation at each of two primary and two context positions. All four fail before the rollback fix and pass after it, including empty ownership tables, restored free capacity and successful retry with the same request id. The not cuda run is performed on a GPU host; it includes an existing graph test without the CUDA marker.

[vllm-omni-review-bot]

This PR appears to belong to: docs/design/module/diffusion/continuous_batching.md, docs/design/module/diffusion/index.md.

Module owners: @Isotr0py @princepride @david6666666

Routing: @Isotr0py via module of the changed files, module named in the PR description, semantic router; @princepride 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, CODEOWNERS

@MrlixiangWE, 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.

[MrlixiangWE]

Self-reviewed on 6a04cb7. Focused regressions cover shared-context deduplication, atomic rollback and release, rejection of over-limit requests before allocate_slots, and a request whose contexts can never fit the empty pool finishing with an error while the next request is scheduled.

#7719's prefix publication still sees only primary sequences, while the A800 CUDA test compares primary and ar_decode context identities with SDPA.

With vLLM 0.30.0, the Diffusion KV suite reports 335 passed and 2 deselected, the Scheduler suite 135 passed, and native CUDA 2 passed; model integration, two-request parity, and aggregate throughput stay in the following B1 slice.

[Levius-Fubuki]

I exercised request-scoped context ownership on head 732abf5 with vLLM 0.29.0 and the native KV manager. Five additional Scheduler cases cover waiting cancellation, abort after reservation but before output update, preempt/resume, capacity backpressure and retry, and request-ID reuse. They all pass; the Scheduler + diffusion KV suite reports 299 passed. I found no production defect and did not touch the SenseNova consumer work. The tests-only commit is 7055b3d. @MrlixiangWE, would you prefer these regression tests folded into #7882, or submitted as a small follow-up PR after it lands? Happy to adapt.

[MrlixiangWE]

I exercised request-scoped context ownership on head 732abf5 with vLLM 0.29.0 and the native KV manager. Five additional Scheduler cases cover waiting cancellation, abort after reservation but before output update, preempt/resume, capacity backpressure and retry, and request-ID reuse. They all pass; the Scheduler + diffusion KV suite reports 299 passed. I found no production defect and did not touch the SenseNova consumer work. The tests-only commit is 7055b3d. @MrlixiangWE, would you prefer these regression tests folded into #7882, or submitted as a small follow-up PR after it lands? Happy to adapt.

Folded all five lifecycle cases into 07f839e, with the request-ID reuse case exercised through the public Scheduler path. The commit credits you through its Co-authored-by trailer.

[MrlixiangWE]

@shihongzhi could you also take a look at this one when you have time? It is the scheduler side of B1: request-scoped KV contexts in the diffusion scheduler. Thanks!

[shihongzhi]

LGTM

The core design is solid — I verified the transactional allocation/rollback path, the prefix-cache isolation (contexts never go through get_computed_blocks, and delay_cache_blocks=True short-circuits before cache_blocks() publishes hashes), the admission-bound math in _get_empty_pool_required_blocks, and the end-to-end cache_role validation chain. Lifecycle coverage (abort/preempt/backpressure/ID reuse) looks thorough.

Two non-blocking nits, happy to see them in a follow-up:

1. free_request vs _rollback inconsistency for contexts (manager.py): free_request frees context shims via native_manager.free() directly, while _rollback routes them through free_kv_blocks_in_physical_order. Harmless today (contexts never join KV transfers), but worth unifying or adding a comment on why they differ.

2. _ContextKVRequest.skip_reading_prefix_cache = not self.block_hashes: contexts never call get_computed_blocks, so this field is effectively dead — but if someone later passes block_hashes on a context, False would imply prefix-cache reads, contradicting the "contexts never read the prefix cache" contract. Consider hardcoding True or asserting not context.block_hashes.

[MrlixiangWE]

@shihongzhi Both done in 4c43ab0: free_request now releases context shims through free_kv_blocks_in_physical_order like _rollback, and a context shim always skips prefix-cache reads; test_context_with_block_hashes_never_reads_the_prefix_cache covers the second.

[shihongzhi]

LGTM ～

[vllm-omni-review-bot]

Omni ReviewBot: no human activity for 7 days

@MrlixiangWE this pull request has had no human commit, comment or review since 2026-09-26. Please confirm the current plan and next step. The author or a maintainer decides whether to change the PR state.

To keep it moving, any one of these is enough: push an update, reply to the open blocker, or post the current plan and timeline.

[RuixiangMa]

Could admission account for aggregate Worker row capacity, not just the per-request limit? With max_num_seqs=4, two requests needing three rows each pass max_rows_per_request=4, but together exceed the Worker’s four-row pool and cause installation to fail.

Please enforce a global row budget or size the pool for concurrent demand？

[MrlixiangWE]

The factory already sizes the native Worker pool as max_num_seqs * diffusion_kv_max_rows_per_request (16 rows for 4 × 4), so the two three-row requests in your example take 6 of 16 rows. Worker rows are installed only for scheduled requests, at most max_num_seqs of them and each within diffusion_kv_max_rows_per_request, so installation cannot exceed the pool; a request larger than either limit fails admission before allocating blocks.

I first added an aggregate count in 77d733c, but it also counted prefetched reservations, which hold blocks and no rows, so at max_num_seqs=1 it refused every native prefetch while a request was running. 410f0ab removes it. At 2227443 with vLLM 0.32.0 the KV, Scheduler and HunyuanImage3 KV-request suites pass (548 passed) and the diffusion_kv CUDA L4 selection passes (4 passed); restoring the aggregate count turns 8 prefetch and manager tests red.

[MrlixiangWE]

@hsliuustc0106 Ready for review.

[vllm-omni-review-bot]

Omni ReviewBot routing record

Assigned Strict on cursor (cursor-grok-4.6-high) under experiment fleet-strict-cursor-grok46-zcode-glm53flash-5050-c5-z10-20261002.

[vllm-omni-review-bot]

Omni ReviewBot attempt record

Review attempt ended as failed (failed; retrying strict/cursor/cursor-grok-4.6-high in 120s (try 2 of 3)).

[vllm-omni-review-bot]

Omni ReviewBot attempt record

Review attempt ended as failed (failed; retrying strict/cursor/cursor-grok-4.6-high in 600s (try 3 of 3)).

[vllm-omni-review-bot]

Omni ReviewBot attempt record

Review attempt ended as failed (failed; falling back to direct/cursor/auto).

[vllm-omni-review-bot]

Omni ReviewBot review

Changes since the previous review

• 0 new inline finding(s); 0 finding(s) below.

CI at 65258877f040 (2026-10-10T13:31:41.851073+00:00): required check(s) blocking: buildkite/vllm-omni (missing).

Note: The assigned review arm strict/cursor/cursor-grok-4.6-high could not complete this review, so it was produced by the fallback arm direct/cursor/auto. It is excluded from the routing experiment.

Full review analysis

PR description

This change turns request-scoped diffusion KV contexts into real Scheduler and Worker rows for the paged path. Shared contexts are deduplicated, then primary sequences and context rows are reserved in one transaction: a native allocation failure rolls back every attempted row, and prefix-cache pins stay until that rollback finishes. Admission rejects a request that can never fit the per-request row limit, the Worker row pool, or an empty block pool, and it waits when live or prefetched reservations already consume the row pool. Context shims stay out of prefix publication and KV transfer, and Worker installation plus attention dispatch reject an unregistered role, mixed roles in one forward, and a layer that reads another role's rows. The SenseNova model consumer is not in this diff.

Change flow

flowchart LR
  Req["[EXISTING] Public request with kv_contexts"]:::existing
  Sched["[CHANGED] Scheduler admission and row budget"]:::changed
  Mgr["[CHANGED] KV manager transactional reserve"]:::changed
  Shim["[NEW] Context-only native shim"]:::new
  Worker["[CHANGED] Worker install and role dispatch"]:::changed
  Req --> Sched --> Mgr --> Shim
  Mgr --> Worker
  classDef existing fill:#e5e7eb,stroke:#6b7280,color:#111827
  classDef changed fill:#fef3c7,stroke:#d97706,color:#451a03,stroke-width:2px
  classDef new fill:#dcfce7,stroke:#16a34a,color:#052e16,stroke-width:2px
  classDef removed fill:#fee2e2,stroke:#dc2626,color:#450a0a,stroke-width:2px

Loading

---

_{🤖 This review was generated by InferMatrix Copilot, an open-source repo-maintenance agent for PR review, CI debugging and issue triage. Try it on your own repo, and ⭐ star it if it helped!}

[vllm-omni-review-bot]

Omni ReviewBot: no human activity for 13 days

@MrlixiangWE this pull request has had no human commit, comment or review since 2026-10-04. Please confirm the current plan and next step. The author or a maintainer decides whether to change the PR state.

To keep it moving, any one of these is enough: push an update, reply to the open blocker, or post the current plan and timeline.

[vllm-omni-review-bot]

Omni ReviewBot triage note

Automated triage of commit 222744379a23 produced:

• Priority: high. Prompt maintainer attention is suggested.

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

[MrlixiangWE]

@hsliuustc0106 this pr have resove all conflect. could you help in review and merge?

[0z5a]

Reviewed head: 2227443 (code reading plus surrounding Scheduler/Worker/engine paths; nothing run).

Thanks @MrlixiangWE, this is a careful slice. I read it against the cache-ownership work in #8264. I checked these and found no issues:

• Attempted rows are appended before allocate_slots, so rollback covers a native mutation that raises.
• Contexts never reach get_computed_blocks, and delay_cache_blocks=True keeps them out of prefix publication.
• Context blocks are counted in the empty-pool bound.
• The request-ID reuse and drain guards hold: _add_request_with_request_id rejects IDs still in _kv_draining_requests.

My main point is an ownership-contract question for the AR→DiT side (item 2). The rest are smaller. None of them repeat the RuixiangMa or shihongzhi threads.

1. [Low] The Worker-row invariant that replaced the aggregate gate is held only by FIFO ordering of the waiting queue. manager.py:307-311, base_scheduler.py:538

The 410f0ab comment says Worker rows are held only by scheduled requests. preempt_request keeps a request's reservation and its installed Worker rows, frees its _running slot, and puts it at the head of _waiting with appendleft. The request then resumes through scheduled_cached_reqs with the same generation.

Today that is safe only because schedule() admits the head of the queue first and breaks when the head cannot be scheduled. So a new request can never take a preempted request's slot while that request's rows are still installed. If anyone later reorders the queue (priority, skip-ahead, or a deferral that also covers already-reserved requests), the row pool overflows. For example, with max_num_seqs=1 and 2 rows per request, the Worker would raise "only 0 rows are free".

Could the comment at L307 name this dependency? Please also extend test_request_scoped_context_preemption_reuses_allocation so it adds a second request while the first is preempted and asserts that the preempted one is re-admitted first. An assertion in schedule() would also work: installed rows (running plus preempted) never exceed the pool.

2. [Medium, question] AR→DiT ownership: contexts can never be transfer destinations, but nothing rejects a request that expects one. base_scheduler.py:708, kv_connector.py:344, base_scheduler.py:394/442/636

All of the transfer paths cover sequences only:

• prepare_kv_requests stamps kv_transfer_params on sequences only.
• commit_kv_load registers only state.diffusion_kv_requests.
• _loading_sequence_ids, update_kv_connector_output and _release_request_kv all track sequence IDs only.

So a remote-prefill request that also carries kv_contexts gets context rows over freshly allocated blocks, which may hold stale data from earlier requests. Nothing loads those rows, and nothing marks them unwritten.

If the B1 consumer expects the ar_decode context to come from the AR stage, it will read garbage silently. Please choose one of these and pin it with a test:

• Document that contexts are always produced in-stage by the DiT Worker, and reject kv_transfer_params + kv_contexts at admission with DiffusionKVAdmissionError for now.
• Extend those four paths to context shims before the consumer slice lands.

Related question, also for the consumer slice. A shared context is deduplicated to one row. prepare_batch rejects the same row twice in one forward, and every forward writes K/V through slot_mapping. Which forward is the single writer of a shared context, and are CFG branches expected to read it in separate forwards? A sentence in the DiffusionKVContext docstring would make that contract explicit.

3. [Low-Medium] Error taxonomy on reservation failure. base_scheduler.py:215, base_scheduler.py:368, manager.py:216/263/266

except Exception maps every reservation failure to FINISHED_ERROR and returns str(exc) to the client. That covers request-caused errors, but also invariant violations:

• "internal request IDs are already allocated" means bookkeeping leaked.
• A native allocator exception after rollback is also an invariant failure.

Two related problems:

• The conflicting-context ValueError at L216 is request-caused but is not a DiffusionKVAdmissionError.
• The prefetch path at L368 drops every exception at debug level.

So a leak shows up as a stream of per-request user errors with no engine-level signal. I'd suggest three outcomes:

• Request-caused failures raise DiffusionKVAdmissionError, and the request finishes with an error.
• Capacity returns None, and the request retries.
• Anything else, after rollback, is logged with logger.exception and escalated through _fail_engine. Prefetch should log non-admission failures at warning or above.

4. [Test gap] Native prefetch with contexts. tests/diffusion/diffusion_kv/test_native_prefetch.py

_request() never carries kv_contexts, so the prefetch→admit and abort-during-load→drain→release_kv_drains paths are untested with context shims. Two cases would close this:

• A prefetched B with one context is admitted and reuses the same generation, including metadata.contexts.
• B is aborted while loading. Its context blocks must stay held until release_kv_drains, then return the pool to its baseline.

5. [Nit] manager.py:245

The num_rows > self.max_num_rows check can't fire. max_num_rows comes from kv_vllm_config.scheduler_config.max_num_seqs, which vllm_config.py:272 sets to max_num_seqs * max_rows_per_request, so it is always at least max_rows_per_request. Please either assert max_num_rows % max_rows_per_request == 0 at construction, mirroring model_runner_backend.py:255, or drop the check. Its "Worker capacity" message suggests an aggregate gate that no longer exists.

Happy to re-check once item 2 is settled. I'd like the admission guard in this PR; the rest can be follow-ups.

[MrlixiangWE]

Thanks @0z5a. I'll keep this PR as it is; per item:

1. The diffusion scheduler's preempt_request has no production caller yet, only tests. When it is called, the preempted request goes back to the head of _waiting with its reservation, and deferred only takes heads that are not reserved yet, so it is re-admitted before any new request. A later production preemption or queue reorder would need to account for its rows.
2. Nothing builds kv_contexts yet: the only production DiffusionKVRequest (HunyuanImage-3) passes kv_contexts=(), so a request carrying both kv_transfer_params and contexts cannot reach admission today. Which forward writes a context, and whether a context can be a transfer destination, belongs to the consumer slice, which will add the guard or the transfer paths together with its test.
3. The except Exception around reserve_request is on main and this PR does not change it. The conflicting-context ValueError takes the same branch as DiffusionKVAdmissionError, so the request finishes with the same error either way. Separating request errors from invariant failures is a change to that handler on its own.
4. Prefetch reserves through the same reserve_request, and the drain release goes through free_request, which also frees the context rows; test_manager.py covers context reservation, release, sharing and rollback. With no context producer, prefetch with contexts is not reachable yet.
5. Yes, with the factory sizing max_num_rows is a multiple of max_rows_per_request, so the second check never fires in serving. I'll drop it or turn it into an assert the next time this file changes.