Skip to content

[Feature] Add experimental deferred Ouro deep batching - #125

Draft
Camellia86 wants to merge 1 commit into
InferMatrixLab:mainfrom
Camellia86:feat/ouro-deferred-decode
Draft

Camellia86 wants to merge 1 commit into
InferMatrixLab:mainfrom
Camellia86:feat/ouro-deferred-decode

Conversation

@Camellia86

Copy link
Copy Markdown

Required human contributor answers

The contributor's approved scope and PR decision are recorded in the
issue #70 proposal. This is an opt-in experimental delivery.

1. What outcome are you trying to achieve?

Improve Ouro rollout throughput under a measured quality budget, by making the
deferred deep-readout experiment available without a training-framework dependency.
The linked proposal identifies this outcome in its second paragraph.

2. What is the highest-leverage variable for that outcome?

The effective batch size of loops three and four. This execution variable is
explicitly identified in the second paragraph of the linked proposal.

3. How will this work move that variable rather than create busywork?

Retain loop-two states and group their genuine deep history at a fixed boundary.
For an aligned group of N requests at interval B, each deep call contains N*B
rows rather than N. Completed boundaries execute the same deep rows. The linked
proposal specifies this mechanism and the eager same-readout control.

4. Who is the customer or user, and what do they actually need?

Ouro rollout developers working with verl/vime need a reusable execution path
with real per-loop KV, aligned token probabilities and explicit comparison
controls. The linked proposal identifies these users and the engine-side delivery.

PR necessity and bugfix evidence

The contributor explicitly chose a PR so the execution path lives in vllm-rlt
without a dependency on either training framework; see the final paragraph of
the linked proposal. This is a feature PR, not a bugfix claim. Related to #70;
it does not close the RFC or deliver its training/weight-publication contract.

Purpose

Add a standalone fixed-interval Ouro runner, CLI and recipe. Prompt prefill
executes all four loops. Mixed readout uses loop two except at a deep boundary;
the boundary batches loops three and four using original positions and true KV.
Unread deep tails are discarded on termination. Interval one and the eager
same-readout path provide explicit controls.

Reuse the public native recurrent API, Sampler and RngRegistry. Per-request
retained rows own their storage; recurrent failures close the cohort. Reservation
excludes the final sampled token's unused forward position. The normal LLM and
serving paths are unchanged.

The alternative is to keep the prototype in vime, where the execution path is
coupled to a training repository. A native opt-in runner makes the experiment
usable from either framework. Integrating it into the main scheduler is deferred
until BF16 quality/performance and feature-combination evidence justify it.

Test Plan

vllm-rlt Version / Commit: pinned current main plus this PR's seven files;
the exact base and resulting code are identified by the PR's parent and head.
Validation used a separate snapshot whose in-scope files are checked against
the staged patch before publication.

  • CUDA_VISIBLE_DEVICES="" OMP_NUM_THREADS=1 python -m pytest -q tests/test_deferred_ouro.py
  • pre-commit run --all-files --verbose with CUDA hidden and one CPU thread.
  • python -m vllm_rlt.entrypoints.ouro_deferred --help
  • git diff --cached --check

Environment: Linux x86_64, Python 3.12, PyTorch 2.7.0a0 NVIDIA build, Torch
attention backend for causal tests, Ruff 0.14.7 and pre-commit 4.7.0. Test tools
were installed in a task-local directory; the existing runtime was not modified.
The causal tests use tiny random FP32 Ouro weights (vocabulary 31, hidden size
32, two physical layers, four loops); they do not measure real-model task accuracy.

Test Result

  • New targeted CPU suite: 36 passed.
  • CPU-mode full regression: 631 passed, 185 skipped. Skips do not qualify device or optional-feature paths.
  • Full pre-commit run --all-files --verbose: Ruff lint, Ruff formatting and CPU regression hooks passed.
  • CLI help/import and staged whitespace checks passed.
  • All seven tested file contents match the staged submission.

The new tests cover intervals 1/2/4/8, three GQA configurations, eager/deferred
hidden readouts and every populated KV plane against the independent functional
dense oracle. They verify deep call/row counts, staggered blocks, stopping and
batch shrink, actual processed logprobs, request-local RNG, retained storage,
failure cleanup and the final executable position.

Qualification and self-review

This is a draft. Real-weight BF16 task quality, GPU memory, end-to-end throughput,
and separate ops-only/full GPU profiles remain unrun. No acceleration factor is
claimed for this implementation. CPU causal correctness does not qualify GPU
kernels or differentiable current-policy replay.

The batching tests show fewer, larger deep calls at completed boundaries without
reducing the number of deep rows. GPU profiling must establish whether this
covers the per-loop metadata preparation, retained-row copies and synchronous
Python sampling/output costs. The eager/deferred pair holds the readout policy
fixed; the mixed/K4 pair requires a paired task-quality comparison.

Code self-review applied the repository's ownership, public-operation, cleanup,
loop/sampling and bounded-storage rules (R03/R11/R16/R19/R21/R25). Existing
interfaces and the functional oracle are reused. Adaptive gating, speculation,
CUDA Graphs, PD, prefix reuse, preemption, weight updates and RL learner replay
are not integrated into this standalone experiment. The recipe lists these gaps.

The existing prototype's 28 CPU compatibility checks passed against current main
before this port; the new checks above validate the submitted implementation.
Follow-up to vllm-project/vime#475.

AI assistance: OpenAI Codex assisted with implementation, tests and documentation.

Signed-off-by: Camellia86 <127925265+Camellia86@users.noreply.github.com>
Assisted-by: OpenAI Codex
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant