Repository navigation
[Rollout] Add experimental deferred exact-KV Ouro decoding - #475
Camellia86 wants to merge 1 commit into
Conversation
Documentation build overview
64 files changed ·
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces an experimental deferred exact-KV decoding feature for Ouro inference, adding a runner implementation, a CLI example, and unit tests. A high-severity memory retention issue was identified in the deferred decoding logic where slicing the 'hidden' tensor without cloning it keeps the entire tensor's storage alive in memory across steps, potentially causing significant GPU memory overhead. Using '.clone()' on the slice is recommended to resolve this.
| depth = 4 if request.response_inputs % self.interval == 0 else 2 | ||
| readouts.append(depth) | ||
| if self.defer: | ||
| request.pending_hidden.append(hidden[row : row + 1]) |
There was a problem hiding this comment.
Appending a slice of hidden (i.e., hidden[row : row + 1]) to request.pending_hidden keeps the entire hidden tensor (which contains hidden states for all active requests at that step) alive in memory because PyTorch slices share the underlying storage. Since request.pending_hidden persists across multiple steps (up to interval steps), this prevents the large hidden tensors of those steps from being garbage collected, leading to a significant GPU memory footprint overhead that scales with the batch size and interval. Using .clone() on the slice breaks the reference to the larger tensor's storage, allowing it to be reclaimed immediately.
| request.pending_hidden.append(hidden[row : row + 1]) | |
| request.pending_hidden.append(hidden[row : row + 1].clone()) |
Signed-off-by: Camellia86 <127925265+Camellia86@users.noreply.github.com> Assisted-by: OpenAI Codex
661ece6 to
758d0b9
Compare
Ouro's deep loops can be grouped across response positions instead of running
them immediately after every shallow readout. This adds an inference-only runner
that retains loop-two states, batches loops three and four at a fixed deep-readout
boundary, and preserves each request's original positions and true per-loop KV.
The example exposes intervals 1/2/4/8 and an eager same-readout reference. It keeps
full-depth prompt prefill, per-answer RNG, selected processed logprobs, raw tokens,
readout depths and stop/length reasons. Unread deep tails are discarded on stop.
Validation:
engine documented by the Ouro recipe: readouts, all committed KV planes, GQA,
staggered blocks, batch shrink, termination, processed scores and reproducible
sampling.
The fixed mixed-readout policy needs quality comparison with native K4. BF16 batch
shapes can change scores and sampled trajectories. This PR does not wire the
experimental policy into GRPO; full-model end-to-end performance and training
probability replay remain separate validation work.
Follow-up to #459.
AI assistance: OpenAI Codex assisted with the implementation, tests, and documentation.