[RFC]: Training support and verl/vime integration for looped-transformer inference optimization

@MinhaoLi0318

Progress update — 2026-10-22

Baseline: main 4a350344027960a48384626aa646f1765fda6d4f.

[FEAT] Rollout outputs: selected-token logprobs, effective seeds and stop token IDs (part of #70) #83 remains open for selected-token logprobs, effective seeds and custom stop IDs. Its latest description includes the speculative-runner scoring contribution; the GPU suite has not been rerun on that combined head.
Expose rollout logprobs and full policy publication on shared engines #78 remains open for rollout scores and transactional full-weight publication. Its broader real-model/PD/training evidence is historical and must not be transferred automatically to a rebased implementation.
Interface/order decision remains open: the coordination question asks whether to land [FEAT] Rollout outputs: selected-token logprobs, effective seeds and stop token IDs (part of #70) #83 then publication, or use Expose rollout logprobs and full policy publication on shared engines #78 as the integration base. This refresh does not select either route or assign duplicate implementation work.
[Architecture] Unify stage execution with loop-level scheduling and mixed batching #118/[Refactor] Move runner result application to Scheduler (M2/R1) #107 changed the engine/stage and result-application boundaries. Integrate the selected training contract through those current interfaces, then revalidate output/depth/score alignment, publication versions, prefix invalidation and failure recovery.
[Feature] Add experimental deferred Ouro deep batching #125 is a separate draft mixed-readout deferred-deep experiment. Its CPU tests do not complete rollout/weight publication, learner replay or real-weight BF16 performance qualification.
Token-facing serving/control work and end-to-end framework replay remain separate acceptance items; no complete training integration is marked merged here.

Existing contributor scopes and external adapter work are retained below.

Motivation.

Enable training-driven improvements to vllm-rlt's looped-transformer inference efficiency. Optimize throughput, latency, and memory use under an explicit quality constraint.

Looped transformers provide specific opportunities: predict exits ahead of execution, align shallow draft predictions with deeper target computation, and make better use of recurrent state and shared weights. Exploring these ideas requires reliable rollout and weight-update interfaces before integrating a training framework.

The proposed order is:

Add framework-independent training support to vllm-rlt through repository PRs.
Integrate and run training loops on contributor forks of verl/vime.
Use those forks to train and evaluate loop-specific inference optimizations; upstream useful engine capabilities with evidence.

Proposed Change.

Scope, ownership, and progress

Initial integration target: Ouro-1.4B, BF16, synchronous training iterations, and full weight updates. Training interfaces and framework adapters must preserve existing inference features, including early exit, PD, speculative decoding, and CUDA Graphs, from the start. Do not impose a fixed-depth, single-device, non-speculative restriction on the integration. Raise concrete compatibility issues for discussion; this requirement does not imply that currently unsupported feature combinations are already implemented. Asynchronous training, incremental weight transfer, and colocated training/inference memory switching remain separate extensions.

Start with verl + FSDP as the proposed first integration; add vime with its training stack through a second adapter. Pin framework versions and validate model compatibility in the integration itself. Train/inference consistency and gradient checks are required integration tests.

✅ Completed for the stated scope.
🚧 Claimed or in progress, with a linked claim, PR, or fork branch; not yet accepted as complete.
⬜ Not yet delivered in this tracker; existing components should be reused where applicable.

Contributors may split tasks or collaborate. Record Assignee(s), engine PRs, and fork branch/results links. Phase 1 is a joint development scope for @MinhaoLi0318 and @IDEA-V, coordinated with @0z5a and @Levius-Fubuki on the shared interfaces. @Levius-Fubuki has confirmed ownership of N+2 exit optimization, including predictor training and runtime integration, planned after establishing the verl/FSDP loop; implementation has not yet been reported as started. @IDEA-V has confirmed participation in Phase 1. @CXJorz is onboarding and has no implementation task assigned yet.

Phase	Deliverable	Code destination	Status	Assignee(s)	PR / branch / evidence
1	Training-facing engine capabilities	vllm-rlt PRs	🚧	@MinhaoLi0318, @IDEA-V	Plan, co-author coordination
2A	Ouro training integration with verl/FSDP	Contributor verl fork	🚧	@Levius-Fubuki	Claim
2B	Equivalent integration with vime	Contributor vime fork	🚧	@0z5a	vime prototype, engine prototype
3	Loop-specific exit and speculative training experiments	Framework forks; experimental engine branches as needed	🚧	@IDEA-V (speculation); @Levius-Fubuki (exit, confirmed; planned after 2A)	Claim, Ownership confirmation

Implementation references

Reference vLLM's rollout and weight-synchronization implementation, and adapt verl/vime's existing vLLM integrations. Preserve recurrent shared-weight and depth-aware KV semantics, documenting necessary differences and the referenced versions.

vLLM weight transfer and source.
verl vLLM rollout and extension interfaces.
vime, customization, and model adaptation.

Implementation details, experiment procedures, training methods and PR boundaries are left to contributors. The requirements below define the deliverables and acceptance criteria.

Phase 1: training-facing capabilities in vllm-rlt

Provide a shared engine contract based on the relevant vLLM interfaces, usable by both frameworks without adding verl/vime dependencies to normal inference installation. Prefer compatible field and operation names; document required deviations in the implementation PRs.

Task	Scope and acceptance	Status	Assignee(s)	PR / plan
Training rollout interface	Token-ID input/output, selected-token logprobs, effective sampling configuration, concurrent requests and request lifecycle. Preserve token alignment and termination reasons; support cancellation/error cleanup; document model-versus-sampling probability conventions.	🚧	@MinhaoLi0318, @IDEA-V	Plan, co-author coordination
Weight update and synchronization	Support reliable full-weight updates with model-version tracking and correct cache/state handling. Generation must use a consistent weight version, and update failures must be reported safely.	🚧	@MinhaoLi0318, @IDEA-V	Plan, co-author coordination

Acceptance requires correct token/probability alignment, concurrent request handling, and consistent generation after weight updates. Document probability conventions, preserve recurrent weight sharing, and prevent reuse of incompatible cached state. Keep the interfaces framework-independent and include relevant correctness and overhead measurements.

The agreed Phase 1 delivery is split into three PRs: rollout outputs; full-weight update and synchronization; and concurrent token-in/token-out API with abort/pause. Build on vLLM and @0z5a’s existing prototype, with shared development and appropriate co-author credit. Existing inference-feature compatibility is part of this delivery, not a later phase.

Phase 2: integration on verl/vime forks

After the engine contract is usable, adapt each framework’s existing vLLM integration to vllm-rlt and run end-to-end experiments on contributor forks. Upstream PRs to verl or vime are not required to complete this phase. Missing generic engine capabilities should be proposed back to vllm-rlt rather than maintained as two incompatible private interfaces.

Task	Scope and acceptance	Status	Assignee(s)	Branch / evidence
verl/FSDP integration	Adapt Ouro training and the vllm-rlt rollout backend, map/export shared recurrent weights, coordinate updates, and provide a reproducible multi-iteration training recipe with consistency and save/resume checks.	🚧	@Levius-Fubuki	Claim
vime integration	Adapt Ouro to the training stack and reuse the engine rollout/update contract; provide parameter conversion and a reproducible multi-iteration recipe with consistency and save/resume checks.	🚧	@0z5a	vime prototype, engine prototype

Provide a reproducible end-to-end training run that verifies correct gradients, weight updates, synchronization, and generation with the updated model. Include train/inference probability consistency and checkpoint recovery evidence. Link the runnable fork branch, configuration and results.

Phase 3: loop-specific training experiments for inference efficiency

Use the integrated training environment to optimize recurrent execution.

Task	Scope and acceptance	Status	Assignee(s)	Branch / results
Lookahead exit optimization	Train and deploy the N+2 exit predictor, including data/labels, checkpoint loading and runtime decisions. Evaluate task quality, selected versus executed depth, and low/high-concurrency throughput and latency.	⬜	@Levius-Fubuki (confirmed; planned after 2A)	Ownership confirmation
Loop-aware speculative optimization	Train shallow draft heads/adapters or depth-specific fine-tuning, integrate the resulting draft path, and evaluate quality, acceptance lengths, computation reuse, memory and end-to-end speedup.	🚧	@IDEA-V	Claim

Experiment goals and acceptance

N+2 exit: use information available at round N to predict whether the token can exit at round N+2; execute through N+2 and use that round's hidden state for output. Preserve valid exit/KV semantics and evaluate quality, selected versus actual executed depth, and throughput/latency against relevant existing baselines.
Loop-aware speculation: improve shallow drafting while preserving valid target verification and identifying reusable recurrent computation. Evaluate output quality, acceptance lengths, extra computation/memory and end-to-end performance against target-only and current speculative inference.

Training architecture, labels, losses, datasets and experimental procedures are chosen by contributors and documented with results. Add reusable trace collection, trained-component loading or other engine support as required by the experiment. Adaptive training/rollout must preserve consistent trajectory and KV semantics.

Code location and PR scope

Code / artifact	Destination
Generic rollout, probability, versioning and update interfaces	vllm-rlt PRs
Framework adapters, differentiable model adaptation, weight conversion and training configuration	Contributor verl/vime forks initially
Labels, losses, training scripts and experimental results	Framework forks / linked experiment artifacts
Experimental exit/draft execution changes	Linked engine experiment branches, paired with framework forks
Useful trace collection or trained-component loading	vllm-rlt PRs after validating the shared requirement
New runtime exit/speculation policy	vllm-rlt PR with correctness, quality and performance evidence

Negative or inconclusive experiments should retain results without requiring a production runtime merge. Keep framework-specific dependencies outside the normal inference installation. Update README/relevant docs for new public engine capabilities and supply reproducible integration/experiment instructions.

Alternatives and impact

Use a shared engine contract for both framework adapters to avoid duplicating rollout and weight-update protocols. Keep framework dependencies in the fork integrations.

Implementation PRs must document API changes, probability-collection overhead, update downtime, memory use and state lifecycle behavior. Experimental heads and policies are opt-in; performance evaluation must include matched quality measurements.

Feedback Period.

At least one week after publication of this draft. Keep the issue open as a living implementation and experiment tracker.

CC List.

@bjf-frz @hsliuustc0106

Anything else.

Related work: #69 for model/backend/runtime support; #56 for adaptive-exit evaluation and execution semantics; #43 for speculative inference.

Training frameworks: verl and vime.

Before submitting a new issue...

Make sure you already searched for relevant issues in the issue tracker, and looked for the answer in the documentation.

I've tried combination of vllm-rlt and Vime on GRPO rollout with few steps reward(0z5a/vime#4 and 0z5a#1). I would be glad to take further combination work of Vime integration.

@bjf-frz

Hi @bjf-frz @hsliuustc0106, I'd like to take ownership of Phase 2A: Ouro training integration with verl/FSDP, and develop it into a reusable training–rollout path.

My related contributions include merged Triton attention optimization in vLLM, MoE-LoRA computation and routing optimizations in VeOmni, and shared diffusion runtime integration. I also have an open DSpark training-loss optimization PR in Vime.

For the first delivery, I'll follow the RFC's Ouro-1.4B/BF16, fixed four-loop, synchronous-iteration scope, with a standalone non-speculative single-device rollout instance and full weight updates. The implementation would cover:

Differentiable Ouro adaptation in the verl fork, preserving recurrent parameter sharing and correct gradient accumulation under FSDP.
A vllm-rlt rollout backend with aligned token IDs/logprobs, explicit sampling conventions, and request lifecycle handling through the shared engine contract.
Shared-weight mapping/export and update coordination, including model-version consistency and checkpoint save/resume.
A runnable multi-iteration training recipe with gradient and train/inference probability checks, followed by optimization of measured rollout/update bottlenecks.

@0z5a, I saw your Vime integration and engine prototype. I'd like to coordinate on and reuse the common rollout/weight-publication contract, contributing missing framework-independent capabilities back to vllm-rlt rather than creating a separate verl-specific engine interface. My primary ownership request is Phase 2A.

Once that loop works, I'd also like to pursue Phase 3's trained N+2 lookahead exit predictor as the follow-on direction: training and component loading through runtime exit decisions, coordinated with the existing exit-policy work. This would use round-N information to predict an exit at N+2, actually execute through N+2, and use that round's hidden state for output; any runtime proposal would include quality and end-to-end performance evidence.

I'll start by pinning the framework versions, mapping the shared-parameter and engine interfaces, and implementing the first fixed-depth adapter increment. Please record me for Phase 2A if this scope fits; happy to coordinate the shared Phase 1 boundaries and subsequent N+2 work.

I'm new to infra, but I also want to make some contributions to the community

I’d like to take on the loop-aware speculative optimization task in Phase 3. It aligns closely with my ongoing work on intermediate predictions across loops and speculative execution.

I’m happy to own the experiments and runtime integration for this task, and coordinate with the Phase 1/2 contributors on the required interfaces. Please put me down for this part!

@bjf-frz

Hi @bjf-frz @hsliuustc0106, I'd like to take Phase 1: training-facing engine capabilities, covering both the training rollout interface and weight update/synchronization, so that 2A (@Levius-Fubuki), 2B (@0z5a) and the Phase 3 items can build on one shared contract.

Background: I own the adaptive-depth rows in #69 (#56, with the harness merged in #58). On the vLLM side I've worked on RL rollout integration, for example Tencent-Hunyuan/UniRL rollout through vLLM-Omni and vllm-project/vllm-omni#8196.

@0z5a, your RLEngine prototype (0z5a#1) already has the minimal shape. I'd build on its semantics in the core LLMEngine/LLM path and credit it, so your vime adapter can use the upstream contract. If you'd rather upstream part of it yourself, I'm happy to split the work.

Proposed PRs. They use vLLM-compatible names, and each reports overhead, update downtime and memory, as the RFC asks.

Rollout outputs:
- selected-token logprobs aligned 1:1 with token_ids;
- logprobs_mode raw by default, as in vLLM, with processed also supported because verl and vime request it;
- seed=None and stop_token_ids;
- the weight version and effective sampling settings in RequestOutput;
- a short doc on probability conventions for looped models.
Full-weight update:
- start_weight_update / update_weights / finish_weight_update / get_weight_version, plus reset_prefix_cache;
- validate names and shapes before mutating anything;
- copy in place, so captured CUDA graphs keep their parameter addresses;
- invalidate prefix-cache entries;
- stamp the weight version on each output;
- refuse new generation after a failed update.
Concurrent token-in/token-out API with abort/pause. I'd first agree with the 2A/2B owners on where the transport sits: in-process only, or HTTP/IPC inside vllm-rlt.

Two things on main matter for RL. Both reproduce on the tiny CPU model at 9e3d13d, and both would get tests:

With the default seed=0, identical prompts at T=1 sample identical outputs, so a GRPO group would get zero advantage.
The prefix cache is keyed on tokens only, so after an in-place weight change it serves stale KV. The output differs from a fresh engine with the same weights.

Questions:

For the first delivery: fixed four loops only, PD and speculative decoding rejected when updates are used, and updates only while the engine is idle, with pause/abort coming in PR 3?
seed: add seed=None as an opt-in, or change the default? [REFACT] [M7 1/3] Extract the sampling module and pin request-local RNG behavior #45 pinned the current seeded behavior.
Logprobs under early exit depend on the executed exit depth and the KV layout. Document that and leave adaptive-exit rollouts out of the first delivery?

I can start once the scope fits. CPU tests will use the tiny model; GPU checks will run on an RTX 5090 with the Triton backend.

@bjf-frz

Hi @bjf-frz @hsliuustc0106, I'd like to take Phase 1: training-facing engine capabilities, covering both the training rollout interface and weight update/synchronization, so that 2A (@Levius-Fubuki), 2B (@0z5a) and the Phase 3 items can build on one shared contract.

Background: I own the adaptive-depth rows in #69 (#56, with the harness merged in #58). On the vLLM side I've worked on RL rollout integration, for example Tencent-Hunyuan/UniRL rollout through vLLM-Omni and vllm-project/vllm-omni#8196.

@0z5a, your RLEngine prototype (0z5a#1) already has the minimal shape. I'd build on its semantics in the core LLMEngine/LLM path and credit it, so your vime adapter can use the upstream contract. If you'd rather upstream part of it yourself, I'm happy to split the work.

Proposed PRs. They use vLLM-compatible names, and each reports overhead, update downtime and memory, as the RFC asks.

Rollout outputs:

selected-token logprobs aligned 1:1 with token_ids;

logprobs_mode raw by default, as in vLLM, with processed also supported because verl and vime request it;

seed=None and stop_token_ids;

the weight version and effective sampling settings in RequestOutput;

a short doc on probability conventions for looped models.

Full-weight update:

start_weight_update / update_weights / finish_weight_update / get_weight_version, plus reset_prefix_cache;

validate names and shapes before mutating anything;

copy in place, so captured CUDA graphs keep their parameter addresses;

invalidate prefix-cache entries;

stamp the weight version on each output;

refuse new generation after a failed update.

Concurrent token-in/token-out API with abort/pause. I'd first agree with the 2A/2B owners on where the transport sits: in-process only, or HTTP/IPC inside vllm-rlt.

Two things on main matter for RL. Both reproduce on the tiny CPU model at 9e3d13d, and both would get tests:

With the default seed=0, identical prompts at T=1 sample identical outputs, so a GRPO group would get zero advantage.

The prefix cache is keyed on tokens only, so after an in-place weight change it serves stale KV. The output differs from a fresh engine with the same weights.

Questions:

For the first delivery: fixed four loops only, PD and speculative decoding rejected when updates are used, and updates only while the engine is idle, with pause/abort coming in PR 3?

seed: add seed=None as an opt-in, or change the default? [REFACT] [M7 1/3] Extract the sampling module and pin request-local RNG behavior #45 pinned the current seeded behavior.

Logprobs under early exit depend on the executed exit depth and the KV layout. Document that and leave adaptive-exit rollouts out of the first delivery?

I can start once the scope fits. CPU tests will use the tiny model; GPU checks will run on an RTX 5090 with the Triton backend.

Sounds good.

@MinhaoLi0318

I've tried combination of vllm-rlt and Vime on GRPO rollout with few steps reward(0z5a/vime#4 and 0z5a#1). I would be glad to take further combination work of Vime integration.

Hi @bjf-frz @hsliuustc0106, I'd like to take Phase 1: training-facing engine capabilities, covering both the training rollout interface and weight update/synchronization, so that 2A (@Levius-Fubuki), 2B (@0z5a) and the Phase 3 items can build on one shared contract.

Background: I own the adaptive-depth rows in #69 (#56, with the harness merged in #58). On the vLLM side I've worked on RL rollout integration, for example Tencent-Hunyuan/UniRL rollout through vLLM-Omni and vllm-project/vllm-omni#8196.

@0z5a, your RLEngine prototype (0z5a#1) already has the minimal shape. I'd build on its semantics in the core LLMEngine/LLM path and credit it, so your vime adapter can use the upstream contract. If you'd rather upstream part of it yourself, I'm happy to split the work.

Proposed PRs. They use vLLM-compatible names, and each reports overhead, update downtime and memory, as the RFC asks.

Rollout outputs:

selected-token logprobs aligned 1:1 with token_ids;

logprobs_mode raw by default, as in vLLM, with processed also supported because verl and vime request it;

seed=None and stop_token_ids;

the weight version and effective sampling settings in RequestOutput;

a short doc on probability conventions for looped models.

Full-weight update:

start_weight_update / update_weights / finish_weight_update / get_weight_version, plus reset_prefix_cache;

validate names and shapes before mutating anything;

copy in place, so captured CUDA graphs keep their parameter addresses;

invalidate prefix-cache entries;

stamp the weight version on each output;

refuse new generation after a failed update.

Concurrent token-in/token-out API with abort/pause. I'd first agree with the 2A/2B owners on where the transport sits: in-process only, or HTTP/IPC inside vllm-rlt.

Two things on main matter for RL. Both reproduce on the tiny CPU model at 9e3d13d, and both would get tests:

With the default seed=0, identical prompts at T=1 sample identical outputs, so a GRPO group would get zero advantage.

The prefix cache is keyed on tokens only, so after an in-place weight change it serves stale KV. The output differs from a fresh engine with the same weights.

Questions:

For the first delivery: fixed four loops only, PD and speculative decoding rejected when updates are used, and updates only while the engine is idle, with pause/abort coming in PR 3?

seed: add seed=None as an opt-in, or change the default? [REFACT] [M7 1/3] Extract the sampling module and pin request-local RNG behavior #45 pinned the current seeded behavior.

Logprobs under early exit depend on the executed exit depth and the KV layout. Document that and leave adaptive-exit rollouts out of the first delivery?

I can start once the scope fits. CPU tests will use the tiny model; GPU checks will run on an RTX 5090 with the Triton backend.

Hi @MinhaoLi0318. I agree with splitting the work into the three PRs you proposed. The training interfaces do not inherently conflict with existing features such as early exit, PD, speculative decoding, and CUDA Graphs, so I suggest preserving support for them from the start rather than deferring compatibility to a later phase. Please refer to vLLM’s implementation and @0z5a’s existing work, and coordinate the shared interfaces with @0z5a and @Levius-Fubuki.

@Levius-Fubuki

Hi @bjf-frz @hsliuustc0106, I'd like to take ownership of Phase 2A: Ouro training integration with verl/FSDP, and develop it into a reusable training–rollout path.

My related contributions include merged Triton attention optimization in vLLM, MoE-LoRA computation and routing optimizations in VeOmni, and shared diffusion runtime integration. I also have an open DSpark training-loss optimization PR in Vime.

For the first delivery, I'll follow the RFC's Ouro-1.4B/BF16, fixed four-loop, synchronous-iteration scope, with a standalone non-speculative single-device rollout instance and full weight updates. The implementation would cover:

Differentiable Ouro adaptation in the verl fork, preserving recurrent parameter sharing and correct gradient accumulation under FSDP.

A vllm-rlt rollout backend with aligned token IDs/logprobs, explicit sampling conventions, and request lifecycle handling through the shared engine contract.

Shared-weight mapping/export and update coordination, including model-version consistency and checkpoint save/resume.

A runnable multi-iteration training recipe with gradient and train/inference probability checks, followed by optimization of measured rollout/update bottlenecks.

@0z5a, I saw your Vime integration and engine prototype. I'd like to coordinate on and reuse the common rollout/weight-publication contract, contributing missing framework-independent capabilities back to vllm-rlt rather than creating a separate verl-specific engine interface. My primary ownership request is Phase 2A.

Once that loop works, I'd also like to pursue Phase 3's trained N+2 lookahead exit predictor as the follow-on direction: training and component loading through runtime exit decisions, coordinated with the existing exit-policy work. This would use round-N information to predict an exit at N+2, actually execute through N+2, and use that round's hidden state for output; any runtime proposal would include quality and end-to-end performance evidence.

I'll start by pinning the framework versions, mapping the shared-parameter and engine interfaces, and implementing the first fixed-depth adapter increment. Please record me for Phase 2A if this scope fits; happy to coordinate the shared Phase 1 boundaries and subsequent N+2 work.

Hi @Levius-Fubuki, the plan looks good overall. As with Phase 1, I suggest making the verl integration work with existing inference features from the start, including early exit, PD, speculative decoding, and CUDA Graphs. Please refer to verl’s existing vLLM integration and coordinate the shared interfaces with @0z5a and @MinhaoLi0318.

@IDEA-V

I’d like to take on the loop-aware speculative optimization task in Phase 3. It aligns closely with my ongoing work on intermediate predictions across loops and speculative execution.

I’m happy to own the experiments and runtime integration for this task, and coordinate with the Phase 1/2 contributors on the required interfaces. Please put me down for this part!

Hi @IDEA-V, you could also collaborate on Phase 1 as a co-author with @MinhaoLi0318, coordinating the shared interfaces with @0z5a and @Levius-Fubuki. This should also help support your Phase 3 work on loop-aware speculative decoding.

@0z5a

Thanks! I’ll coordinate with @0z5a and @MinhaoLi0318 and preserve existing inference-feature support from the start. I also confirm my intention to own N+2 exit optimization, including predictor training and runtime integration, after establishing the verl/FSDP loop.

I’d like to take on the loop-aware speculative optimization task in Phase 3. It aligns closely with my ongoing work on intermediate predictions across loops and speculative execution.
I’m happy to own the experiments and runtime integration for this task, and coordinate with the Phase 1/2 contributors on the required interfaces. Please put me down for this part!

Hi @IDEA-V, you could also collaborate on Phase 1 as a co-author with @MinhaoLi0318, coordinating the shared interfaces with @0z5a and @Levius-Fubuki. This should also help support your Phase 3 work on loop-aware speculative decoding.

Sure. I'm happy to participate.

@IDEA-V

Following the developer guide (#75), here is the design for PR 1 (rollout outputs), written against main ecb1f8b, before I start coding. @IDEA-V @0z5a @Levius-Fubuki, please point out anything your adapters need that is missing.

Contract (vLLM names; deviations marked)

SamplingParams.logprobs: int | None = None. 0 returns the selected token only, which is what verl and vime request. Values above 0 (top-k alternatives) are rejected for now.
SamplingParams.seed: int | None = 0. With None, add_request draws a per-request seed and reports it, so a rollout can be replayed exactly. The default stays 0, as pinned by [REFACT] [M7 1/3] Extract the sampling module and pin request-local RNG behavior #45. vLLM defaults to None; if you want that here, flipping the default can be a separate one-line PR.
SamplingParams.stop_token_ids. EOS is currently checked in three places: sync coda, async coda delivery and speculative commit. One shared helper will check both EOS and stop_token_ids.
Engine-level logprobs_mode: raw_logprobs (the default, as in vLLM) or processed_logprobs. The logits modes are rejected.
New RequestOutput fields:
- logprobs: list[float] | None, aligned 1:1 with token_ids (and with exit_depths) for stop, length and abort. Deviation: flat floats rather than vLLM's list[dict[int, Logprob]], since only the selected token is returned.
- sampling_params: the resolved parameters, including the effective seed.
weight_version moves to PR 2, where versions actually exist.
Mapping from @0z5a's prototype (Add synchronous fixed-budget Ouro rollouts for RL 0z5a/vllm-rlt#1): its RLCompletion fields map onto RequestOutput as follows. token_ids, exit_depths and finish_reason keep the same names. log_probs becomes logprobs. loop_budget becomes sampling_params.min_loops/max_loops. policy_version becomes weight_version (PR 2). Its logprob, log_softmax(logits / T)[token], is exactly processed_logprobs with no top-k/top-p. Differences from the prototype: the support is in the core engine rather than an RLEngine subclass, and it doesn't add the subclass's restrictions (T>0 only, unfiltered sampling only, fixed depth, no prefix cache/async/spec).

Probability convention (this goes into the docs)

raw is log_softmax(coda(h_d))[token], where d = exit_depths[i] is the depth the token actually exited at.
processed is the log-probability under the distribution actually sampled from: after temperature and top-k/top-p, renormalized. For greedy requests no temperature is applied. With T=1 and no filters, raw and processed are equal.
The first output token always comes from full-depth prefill. Per-request min_loops/max_loops affect decode tokens only.
A trainer that recomputes log π at full depth matches the engine only for fixed full-depth rollouts. With early exit, the trainer must replay exit_depths. Under last_exited KV, it must also reproduce the copied deeper-loop KV. I'll report the measured mismatch on Ouro-1.4B rather than block PR 1 on it.

Compatibility, all in this PR

Early exit (ouro, ouro_delayed, random_lookahead, trace): logprobs are computed at coda, and the exit logic is unchanged. Test: they equal a trace-mode replay with the same depths.
Async scheduling: add one float32 readback on the coda submission's existing event, only when some request in the batch asks for logprobs. No extra host sync.
Speculative decoding: each committed token (accepted, correction or bonus) gets log p_target[token] from its verification row. Greedy spec must match non-spec greedy, with exact tokens and logprobs within tolerance; this extends test_greedy_matches_native_*. Sampled spec is exact in distribution, so log p_target is the right policy logprob.
PD: only D samples. Logprobs are stored on Request like exit_depths and travel in the existing output message. This gets a 2-GPU test.
CUDA Graphs: logprobs are computed outside the recurrent graphs. The static logits of the spec CodaGraphs are read before the next replay. Test: graphs on vs off.
Preemption: snapshots keep the request, so logprobs survive. I'll add an alignment assertion to the existing preemption tests.

Overlap with #62: I'll put the "processed logits" step in one function shared by Sampler and the speculative probabilities(). Whichever PR lands second adds its processors there, and processed logprobs then include repetition penalty and depth temperature automatically. cc @cs-ai-li

Validation:

CPU tiny-model tests for everything above.
RTX 5090 with Triton for async, CUDA Graphs and spec.
A rented 2-GPU box for PD.
Overhead per the developer guide: unprofiled throughput and latency for logprobs=None vs 0 (raw and processed) in sync, async and spec, plus ops-only and full profiles.

@IDEA-V, would you take the speculative path: target-row logprobs plus the greedy-equality and sampled-replay tests, as co-authored commits on my branch? You know that runner from #51, and your Phase 3 work needs those target logprobs anyway. I'll do params, sampler, output, sync and async coda, PD, CUDA Graphs and the docs.

@0z5a, your set_loop_budget lowers total_ut_steps, so prefill also runs at the reduced depth. Does the vime pilot depend on reduced-depth prefill? If it does, I'd add that as an explicit engine option rather than a sampling parameter.

@MinhaoLi0318

@MinhaoLi0318 The current VIME pilot needs four-loop prompt prefill, so no reduced-prefill option is needed. VIME draft #5 supersedes the older RLEngine.set_loop_budget prototype with public LLM.generate / full-weight publication. Checkpoint total_ut_steps stays four; decode K uses SamplingParams.max_loops. The trainer replays [4] * prompt_length + exit_depths[1:] + [1] (last input loss-masked), including last_exited deeper KV.

I support your RequestOutput.logprobs and engine-level raw_logprobs / processed_logprobs names, default seed=0, and opt-in seed=None. VIME uses processed scores. My reference draft #4 currently differs in output spelling (log_probs) and request-level mode (raw / processed); I will adapt the VIME consumer to your PR 1 when its branch is available, without introducing a duplicate compatibility API. The pilot's first consumer is in-process LLM; existing PD worker IPC can remain inside the engine.

You and @IDEA-V can lead the proposed upstream Phase 1 output/stop/processor work. The companion draft provides tested reference code and full-publication checks reusable for PR 2; I am not claiming ownership of your Phase 1 PRs. These Thor checks test the proposal's semantics on the reference implementation, not your unpublished PR 1 code.

Your effective-seed requirement exposed a real omission: entropy sampling returned seed=None. The minimal fix resolves/stores a 63-bit seed when Request is created, including abort before sampling. Seven CPU contract/PD-control tests, six CUDA contract tests and two VIME replay/gradient/conversion CPU tests pass. Exact token/score/depth replay covers ordinary, async and speculative paths, and preemption preserves scores/RNG; the unfixed seed regression fails as expected. Earlier full training/save/resume evidence retains exact 272 checkpoint tensors/Adam metadata and 269 export tensors.

Full Ouro-1.4B now passes exact effective-seed token/score replay for K=2/3/4, early exit and delayed async, all with four-loop first-token prefill. Actual-depth recomputation reduces early-exit score MAE from 3.622689 (uniform four loops) to 0.011142. Residual differences remain: fixed K=4 also has MAE 0.026607, versus 0.025122 for uniform-four recomputation. Five-case error table and raw evidence.

Full-model selected-score time overhead vs logprobs=None, raw/processed respectively: sync +1.46%/+1.47%, async +5.15%/+5.42%, spec +0.67%/+0.83%. Each mode has five rotated-order unprofiled repeats, plus separate ops-only/full-detail profiling (18 operator-summary windows). Markdown latency/throughput/allocation table and raw data. These observations include Qwen-Image training contention (initial GPU utilization 98%). Two UVM allocation waits resolved naturally; driver memory-allocation errors and stacks are retained. Initialization/capture and profiled windows are excluded from throughput medians. No processes were killed or environments/drivers changed.

Please share the PR 1 branch/commit when ready so I can validate VIME against that exact revision. Effective seed, selected scores, physical exit trace and committed weight version remain the consumer requirements, with version publication landing separately in PR 2.

@MinhaoLi0318

@MinhaoLi0318 Sounds good! I’m happy to take the speculative path. Please share your branch or draft PR when ready, and I’ll build on it.

@0z5a

@0z5a @IDEA-V Thanks both. PR 1 (rollout outputs) is up as draft #83 from MinhaoLi0318/vllm-rlt:rl/rollout-outputs at 57d2308, on main ecb1f8b. Everything except the speculative path is implemented. CPU tests pass; GPU validation and overhead measurements are next.

@0z5a, that settles the open points: four-loop prefill stays, the names stay as proposed, and VIME reads processed scores. When you validate VIME against 57d2308:

Seed: seed=None is resolved in add_request before the Request exists, and also in the PD coordinator. An abort before the first sample therefore reports the int seed; there's a test for it.
Processed scores are log_softmax(processed_logits)[token] rather than log(softmax(...)[token]). The values agree up to fp32 rounding, and very low-probability tokens stay finite instead of underflowing to -inf.
Async: the logprob goes to a pinned host buffer on the same CUDA event as the sampled ID and is delivered with it. No device tensor is kept per token and stacked when outputs are built. I expect less than your +5.15% async overhead, and I'll measure it with your method: five rotated-order unprofiled repeats, plus ops-only/full profiles.
Replay depths: the docs give the per-position depths as [depth] * len(prompt) + exit_depths[1:], which matches your trainer.

@IDEA-V, great, the speculative path is yours. To build on the branch, open PRs against MinhaoLi0318/vllm-rlt:rl/rollout-outputs; I'll merge them so your commits land in #83 under your name. If you'd rather push to the branch directly, I can add you as a collaborator on the fork. What's already in place:

SpeculativeResult.logprobs exists. The engine slices it together with the emitted tokens on EOS, stop IDs and length (LLMEngine._append_output).
add_request currently rejects logprobs with speculative_config. Your change lifts that rejection, and test_admission_rejects_out_of_vocabulary_stop_ids_and_speculative_logprobs needs the matching update.
The selected logprob should come from vllm_rlt.worker.sampling.token_logprob(row, params, token, mode) on the target verification row. Draws already go through the shared processed_logits, so processed scores match what was sampled. 0z5a's reference (Expose rollout logprobs and full policy publication on shared engines 0z5a/vllm-rlt#4) already scores verification rows, which could save you time (with his credit).
Heads-up: on main, the GPU test from [FEAT] [PERF] Remove per-candidate and metadata host syncs from speculative decoding #51 at tests/test_speculative.py:398 calls OuroConfig.tiny(...). OuroConfig isn't imported there (ruff F821), and no tiny constructor exists; the test helpers provide tiny_ouro_config. So the test would fail on a GPU run.

@bjf-frz, could you link #83 in the Phase 1 rows of the tracker?

@0z5a

Following the developer guide (#75), here is the design for PR 2 (full-weight update and synchronization), before I start coding. It is stacked on #83 (GPU-validated at c5a7019, results in its description), since both add RequestOutput fields and touch pd/engine.py.

It upstreams the publication contract of @0z5a's #78: the same methods and call order, the idle precondition, the prefix reset at start, and the PD broadcast. The changes are listed below. @0z5a @Levius-Fubuki @IDEA-V, please point out anything your adapters need. The full write-up will go into docs/weight_update.md.

Contract (vLLM names, identical on LLM, LLMEngine and PDEngine; deviations in italics)

start_weight_update(weight_version: int | None = None).
- Unlike vLLM, it accepts a version. vime #459/#463/#464 pass the version here.
- A new start discards an open update. Their retry path relies on this.
update_weights(weights: Iterable[tuple[str, Tensor]]), which may be called any number of times, so callers can send buckets.
- Unlike vLLM, it takes named in-process tensors rather than transport update_info.
finish_weight_update(weight_version: int | None = None) and get_weight_version() -> int.
- The version is a monotonic int, starting at 0 for the constructor weights, instead of vLLM's opaque str. One version then names exactly one weight state.
reset_prefix_cache(reset_running_requests=False, reset_connector=False) -> bool, with vLLM's signature.
- It drops index entries only, never the KV pools.
get_weight_manifest() -> list[(name, shape, dtype)]: the deduplicated physical parameters (for Ouro and Nanbeige, exactly the HF checkpoint names; tied aliases and buffers are rejected).
RequestOutput.weight_version: the version committed when add_request ran. It appears on every streamed, finished and aborted output.

Lifecycle and failures

The states are READY, UPDATING and FAILED. add_request raises unless the state is READY.
start closes admission and requires the engine to be fully idle (no waiting, running, preempted or PD-in-flight request, no live KV allocation). It then syncs the runner streams, clears async tickets and exit signals, and resets the prefix cache.
update_weights validates the whole call before the first copy.
- It checks exact names, no duplicates, exact shapes, and a plain fp32/fp16/bf16 tensor on CPU or CUDA.
- A rejected call changes nothing.
- Accepted tensors are written with param.copy_(src) in place and followed by a device sync. The caller may then reuse its buffers.
finish requires every manifest name exactly once. It then rechecks idleness and parameter addresses, and commits.
Failures are fenced, not rolled back: after any error past the first copy, admission stays closed until a complete new update commits. Rollback would need a second model copy in memory; the trainer can simply republish. No staging copy, so no extra GPU memory.
Invariant: every token, logprob, exit depth and KV entry behind an output stamped v comes from the weights committed as v. This includes reused prefix KV and P's prompt KV.

State handling per feature

Prefix cache: this fixes a bug on main. After an in-place update, a repeated prompt reuses stale prefix KV and differs from a fresh engine under ouro, ouro_delayed and random_lookahead. The reset at start fixes it.
CUDA Graphs: in-place copies keep every captured address, so nothing is recaptured. The address check also covers buffers such as inv_freq.
Early exit: the ouro gate is a published parameter. The engine-owned random_lookahead probe is not updated.
Speculative decoding: the draft loops reuse the target's parameters, so one update covers both.
Preemption: preempted requests must be aborted first. Restoring their KV under new weights would mix policies.

PD (one update across all workers, driven by the coordinator)

The coordinator validates each chunk once. It then sends the chunk to every peer through a flat CPU shared-memory buffer, using a (name, offset, shape) index.
It commits only after every peer reports the same version.
A peer error in start or update moves the engine to FAILED; a new start recovers.
A finish-phase error, a control timeout or a dead worker is fatal, as PD request failures are today.

Out of scope

PR 3: pause/resume, token-in/token-out with abort, and HTTP routes.
Later: init_weight_transfer_engine with NCCL/IPC transports, start_draft_weight_update, sleep/wake for colocated trainers, incremental or LoRA updates, Huginn, and Nanbeige on PD.

Validation

CPU tests with tiny Ouro and Nanbeige models: the stale-prefix regression; bitwise parameters with unchanged data_ptr; rejections leave the state unchanged; an injected copy failure followed by a retry; versions on every output; every exit mode and self-speculation vs a fresh engine with the new weights; PD control via @0z5a's harness.
GPU: addresses and graph captures stay unchanged on the RTX PRO 6000 across the graph, async, speculative and prefix-cache arms. 2-GPU NIXL PD must match a local engine.
Measured unprofiled, then with ops-only and full profiles: downtime per phase, memory deltas, and overhead against a bare copy_ loop.

Open questions

@bjf-frz @hsliuustc0106:
- an int version instead of vLLM's str;
- CPU shared-memory chunks for PD;
- a fatal close on finish-phase errors rather than a prepare/commit round.
@0z5a:
- Can #463/#464 skip pause_generation/resume_generation until PR 3?
- May I add you as Co-authored-by on the commits that port Expose rollout logprobs and full policy publication on shared engines #78's code and tests, with Expose rollout logprobs and full policy publication on shared engines #78 then narrowed to what [FEAT] Rollout outputs: selected-token logprobs, effective seeds and stop token IDs (part of #70) #83 and PR 2 don't cover?
@Levius-Fubuki: verl publishes global_steps=0 first, but here 0 means the constructor weights. Could the 2A adapter publish global_steps + 1 and bucket its per-tensor generator?
@IDEA-V:
- Would you take the self-speculation test arms, as in [FEAT] Rollout outputs: selected-token logprobs, effective seeds and stop token IDs (part of #70) #83?
- If Phase 3 adds trainable draft parameters, should they join the same manifest and update?

feat: expose selected-token logprobs, effective seeds and stop token IDs

Phase 2A update: draft PR in my verl fork, with reproduction instructions and experiment results. Two-GPU Ouro-1.4B BF16 runs completed three full-depth and three adaptive-exit RL iterations, exact next-step checkpoint recovery, and weight-publication checks across existing serving modes, including NIXL P/D and CUDA Graphs.

Current limits: this uses the pinned engine prototype and an in-process torchrun recipe; migration to the final shared API and production Ray integration remain. BF16 train/rollout selected logprobs differ by up to 0.341, so strict numerical parity is not established. These short runs make no convergence or speedup claim; the draft is ready for review, not yet accepted as complete.

I remain responsible for Phase 3 N+2 lookahead exit optimization, including predictor training and runtime integration. That is the next development stage and is not implemented in this PR.

@Levius-Fubuki

@Levius-Fubuki, these are the changes verl#1 (e901c2c) runs into when it moves from the #78 prototype (bdce287) to #83 (head df66711) and PR 2:

verl#1 today	#83 / PR 2
`RequestOutput.log_probs`; `SamplingParams(logprobs_mode="raw")`, read back in `rollout_to_tensordict`	`RequestOutput.logprobs`; engine-level `LLM(..., logprobs_mode="raw_logprobs")`, the default (#83, as your README notes)
`logprobs=0` with `speculative_config` (matrix modes `speculative`, `graphs_speculative`)	rejected at `add_request` in #83 until the speculative-runner logprobs land, in #83 or as a follow-up PR
`RequestOutput.weight_version`, `start_weight_update(version=...)`, `update_weights`, `finish_weight_update()`, `get_weight_version()`	PR 2, same names and arguments
`LLM(PDEngine(...))` in the P/D publication matrix (`validate_adapter.py --pd`), which `bdce287` accepts	not in main `a32c89a` or #83, whose `LLM` takes only a model or a path; my 10-06 design adds the methods to `PDEngine` itself but does not cover `LLM(PDEngine(...))`
`pause_generation()` / `resume_generation()` around publication	PR 3 in the design above; in PR 2, `start_weight_update` closes admission until `finish_weight_update` commits

An update to my 10-06 design: PR 2 will name the keyword version, not weight_version, in both start_weight_update and finish_weight_update, as #78 does for start_weight_update. It stays optional, so OuroWeightPublisher can keep its calls as written. The effective seed stays at RequestOutput.sampling_params.seed, where your rollout reads it. verl#1 also already publishes global_steps + 1 in buckets, which answers question 3 of the design.

Could you confirm this table covers everything verl#1 needs from #83 and PR 2, before #83 leaves draft?

Thanks!

AI assistance: Claude Code drafted this and checked each row against verl#1 e901c2c, #78 bdce287, #83 df66711 and main a32c89a; I reviewed it before posting.

@MinhaoLi0318

@MinhaoLi0318 Confirmed: the table covers our current Phase 2A interface needs. I'll migrate to RequestOutput.logprobs and engine-level raw_logprobs, including the worker's probability checks. version=, effective seeds, and bucketed global_steps + 1 publication work for us. For synchronous iterations, I'll use PR 2's admission fencing plus adapter-local lifecycle gating until PR 3.

For P/D, I'll adapt the consumer around PDEngine's public request/step/abort APIs; LLM(PDEngine(...)) is not required. Speculative-logprob migration tests will remain pending until that support lands in #83 or its follow-up, then be rerun against the new revision. Existing speculative results apply only to the pinned prototype.

I'd like to contribute an opt-in, experimental deferred Ouro decode path related to this RFC.

The intended users are Ouro rollout developers working with verl/vime. The goal is to improve rollout throughput under a measured quality budget. The main execution variable is the effective batch size of loops three and four: retain loop-two states across response positions, then execute a complete deep block at each fixed deep-readout boundary. Original token positions and separate, genuinely computed per-loop KV histories are preserved. This is a mixed shallow/deep readout policy; comparison with native K4 is a quality/performance tradeoff, while an eager implementation with the same readout schedule isolates the effect of deferring deep work.

I checked the existing prototype from vime PR #475 against current vllm-rlt main: all 28 CPU compatibility tests pass, covering readouts, every populated KV plane, grouped-query attention, staggered blocks, batch shrink, stopping, processed selected-token logprobs and independent seeded sampling. This establishes API and small-model causal correctness, not BF16 GPU performance or task quality.

The proposed PR will provide a standalone engine-side runner, a reproducible CLI example, correctness/lifecycle tests and documentation. A PR is useful because the reusable execution path should live in vllm-rlt without a dependency on a training framework. It will leave the normal inference entry points and their feature selection intact. I will open it as a draft with real-weight BF16 quality/performance/profiling and feature-combination qualification listed as unfinished. It will not claim completion of this RFC's training, weight-update or differentiable probability-replay deliverables.

feat: expose selected-token logprobs, effective seeds and stop token IDs

@bjf-frz

@bjf-frz one scoping question before I start PR 2. #78 (0z5a, now ready) exposes rollout logprobs and full policy publication on shared engines, and its description says weight_version and publication belong to PR 2 in the split from this RFC. #83 (PR 1) and #78 touch several of the same files, so I'd like to know which way you want to go:

A: land [FEAT] Rollout outputs: selected-token logprobs, effective seeds and stop token IDs (part of #70) #83 first, then PR 2 adds publication / weight_version on top, as Expose rollout logprobs and full policy publication on shared engines #78's description already suggests, or
B: take Expose rollout logprobs and full policy publication on shared engines #78 as the base for publication, and I rebase [FEAT] Rollout outputs: selected-token logprobs, effective seeds and stop token IDs (part of #70) #83 onto it and drop PR 2.

Either works for me; I just don't want two PRs implementing the same interface. cc @0z5a

AI assistance: drafted with Claude Code; I reviewed it before posting.

feat: add rollout sampling fields and logprob scoring primitives

Navigation Menu

Repository navigation

[RFC]: Training support and verl/vime integration for looped-transformer inference optimization #70

Progress update — 2026-10-22

Motivation.

Proposed Change.

Feedback Period.

CC List.

Anything else.

Before submitting a new issue...

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

[RFC]: Training support and verl/vime integration for looped-transformer inference optimization #70

Description

Progress update — 2026-10-22

Motivation.

Proposed Change.

Feedback Period.

CC List.

Anything else.

Before submitting a new issue...

Activity

0z5a commented on Sep 29, 2026

Levius-Fubuki commented on Sep 29, 2026

CXJorz commented on Sep 29, 2026

IDEA-V commented on Sep 29, 2026

MinhaoLi0318 commented on Sep 29, 2026

0z5a commented on Sep 29, 2026

bjf-frz commented on Sep 29, 2026

bjf-frz commented on Sep 29, 2026

bjf-frz commented on Sep 29, 2026

Levius-Fubuki commented on Sep 29, 2026

IDEA-V commented on Sep 29, 2026

MinhaoLi0318 commented on Oct 1, 2026

0z5a commented on Oct 1, 2026

IDEA-V commented on Oct 3, 2026

MinhaoLi0318 commented on Oct 3, 2026

MinhaoLi0318 commented on Oct 6, 2026

Levius-Fubuki commented on Oct 13, 2026

MinhaoLi0318 commented on Oct 14, 2026

Levius-Fubuki commented on Oct 14, 2026

Camellia86 commented on Oct 14, 2026

MinhaoLi0318 commented on Oct 15, 2026

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions