Add RL-Kernel vime integration roadmap - #5
Conversation
📝 WalkthroughWalkthroughAdds a 768-line roadmap for integrating RL-Kernel with vime through optional fast and consistency modes, explicit logprob contracts, operator hooks, provenance and rollout metadata, distributed requirements, phased implementation, and runtime observability. ChangesRL-Kernel–vime integration
Estimated code review effort: 2 (Simple) | ~10 minutes Possibly related issues
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
There was a problem hiding this comment.
Actionable comments posted: 6
🧹 Nitpick comments (1)
RL-Kernel-vime-integration-roadmap.md (1)
295-297: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winSeparate attention output dtype from LSE dtype.
The contract currently requires both
(out, lse)buffers to use fp32. Please verify that fp32 is intended for the attention output as well as LSE; if only LSE requires fp32, document the output dtype and downcast point separately to avoid unnecessary memory traffic and an unintended parity requirement.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@RL-Kernel-vime-integration-roadmap.md` around lines 295 - 297, Update the deterministic attention merge contract to distinguish the attention output buffer dtype from the LSE buffer dtype. Confirm and document whether attention output must remain fp32; if only LSE requires fp32, specify the output dtype and the exact downcast point separately, and remove the unintended output-parity requirement.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@RL-Kernel-vime-integration-roadmap.md`:
- Around line 84-87: The deterministic operator claim in the roadmap should
limit bitwise stability to identical input bytes on a fixed backend, device
architecture, and build fingerprint. Remove GPU architecture and build changes
from the bitwise variation axes, and explicitly keep cross-architecture or
cross-build comparisons tolerance-based unless a separate cross-platform bitwise
contract is declared.
- Around line 82-87: Revise the batch-invariance definition to cover
operator-local recomputed logp and per-token loss fragments only after
algorithm-level tensors and normalization inputs are fixed. Remove the
unconditional claim about unchanged training gradient contributions, and
separately identify any final-gradient guarantees while preserving vime-owned
global normalization and GRPO/group semantics.
- Around line 46-55: Clarify the precedence between strict acceleration and
contract-preserving reference fallback in the behavior table and corresponding
strict-mode guidance. Define whether --rlk-fast strict must always fail when
optimized operators are unavailable, or may use reference execution when
--rlk-consistency strict permits it, and ensure both sections state the same
rule.
- Around line 519-525: Update the strict/audit reporting and pass/fail
aggregation requirements to explicitly handle empty active-token sets: emit a
not_applicable/zero_active_tokens status, avoid calculating or emitting NaN
max_abs_dlogp and percentile metrics, and exclude these cases from pass/fail
aggregates so strict acceptance cannot pass vacuously. Anchor the changes to the
acceptance-criteria reporting and aggregation definitions covering active-token
counts and worst-token locations.
- Around line 199-203: Clarify the custom rollout contract to fail closed when
strict-mode consistency metadata is missing: quarantine the affected sample,
downgrade the entire batch/run with an explicit status, or fail strict
execution. Replace the current “disable strict consistency for that sample”
behavior in the metadata requirements so unknown-provenance samples cannot
continue in a supposedly strict batch.
- Around line 508-517: Update the strict eligibility and RL-Kernel per-dtype
tolerance lookup requirements to fail closed when the dtype, contract ID, or
RL-Kernel version lacks a tolerance entry. Require strict mode to reject or
explicitly downgrade with a structured reason, and prohibit continuing with any
implicit local default.
---
Nitpick comments:
In `@RL-Kernel-vime-integration-roadmap.md`:
- Around line 295-297: Update the deterministic attention merge contract to
distinguish the attention output buffer dtype from the LSE buffer dtype. Confirm
and document whether attention output must remain fp32; if only LSE requires
fp32, specify the output dtype and the exact downcast point separately, and
remove the unintended output-parity requirement.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 8a797a81-3731-44fc-8ea6-93afa83e01d7
⛔ Files ignored due to path filters (2)
assets/rlk-vime-a0-a5-ablation-matrix.svgis excluded by!**/*.svgassets/rlk-vime-layered-architecture.svgis excluded by!**/*.svg
📒 Files selected for processing (1)
RL-Kernel-vime-integration-roadmap.md
| | `--rlk-fast` | `--rlk-consistency` | Behavior | | ||
| |---|---|---| | ||
| | `off` | `off` | Native vime behavior. | | ||
| | `auto` | `off` | Use validated RL-Kernel optimized operators when eligible; warn and fall back otherwise. | | ||
| | `strict` | `off` | Require enabled RL-Kernel optimized operators; error when unavailable. | | ||
| | `off` | `audit` | Keep native vime execution, but compute and log consistency diagnostics. | | ||
| | `off` | `strict` | Require complete metadata and contract checks; fail on mismatch. | | ||
| | `auto` | `audit` | Use optimized operators when eligible and report consistency metrics. | | ||
| | `auto` | `strict` | Use only optimized operators that satisfy the consistency contract; otherwise use contract-preserving reference execution or fail according to policy. | | ||
| | `strict` | `strict` | Require both acceleration and strict consistency support for enabled operators. | |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Resolve the meaning of strict fallback behavior.
The mode table says --rlk-fast strict must error when an optimized backend is unavailable, while the desired end state says strict mode may select contract-preserving reference execution. Define the precedence explicitly; otherwise implementations cannot determine whether reference fallback is compliant or a failure.
Also applies to: 758-768
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@RL-Kernel-vime-integration-roadmap.md` around lines 46 - 55, Clarify the
precedence between strict acceleration and contract-preserving reference
fallback in the behavior table and corresponding strict-mode guidance. Define
whether --rlk-fast strict must always fail when optimized operators are
unavailable, or may use reference execution when --rlk-consistency strict
permits it, and ensure both sections state the same rule.
| Batch invariance is part of strict consistency. For a given sample/token, the recomputed logp and the training gradient contribution should not change when the same sample is scored alone, placed in a different batch, padded differently, packed in a different order, assigned to a different microbatch, or grouped by dynamic sampling differently, as long as the sample-level inputs and the declared contract are the same. If an optimized backend changes math based on batch composition, it can still run in `off` or `audit`, but it should not be advertised as `strict` until the batch-invariance contract is validated. | ||
|
|
||
| There are two levels of batch-invariance claims: | ||
|
|
||
| - **End-to-end strict consistency**: same-sample `dlogp` must stay within the per-dtype tolerance across the covered batch-layout matrix. This is the default production contract. | ||
| - **Deterministic operator claim**: an RL-Kernel backend that advertises deterministic/batch-invariant `logp` must prove same-backend bitwise stability for identical input bytes across batch size, row position, dense vs indexed rows, sparse-row density, stream/repeat runs, GPU architecture, and build fingerprint. Cross-backend or cross-platform parity can remain tolerance-based unless a separate bitwise contract is declared. |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | 🏗️ Heavy lift
Scope batch invariance to operator-local quantities.
The requirement that a sample’s “training gradient contribution” remain unchanged across batch layouts is too broad: global loss normalization and GRPO/group semantics can legitimately depend on batch composition, and the loss section later identifies those semantics as vime-owned. Define invariance for recomputed logp and per-token loss fragments after algorithm-level tensors and normalization inputs are fixed; separately specify which final gradients are covered.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@RL-Kernel-vime-integration-roadmap.md` around lines 82 - 87, Revise the
batch-invariance definition to cover operator-local recomputed logp and
per-token loss fragments only after algorithm-level tensors and normalization
inputs are fixed. Remove the unconditional claim about unchanged training
gradient contributions, and separately identify any final-gradient guarantees
while preserving vime-owned global normalization and GRPO/group semantics.
| There are two levels of batch-invariance claims: | ||
|
|
||
| - **End-to-end strict consistency**: same-sample `dlogp` must stay within the per-dtype tolerance across the covered batch-layout matrix. This is the default production contract. | ||
| - **Deterministic operator claim**: an RL-Kernel backend that advertises deterministic/batch-invariant `logp` must prove same-backend bitwise stability for identical input bytes across batch size, row position, dense vs indexed rows, sparse-row density, stream/repeat runs, GPU architecture, and build fingerprint. Cross-backend or cross-platform parity can remain tolerance-based unless a separate bitwise contract is declared. |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Do not require bitwise stability across architecture and build changes.
The deterministic claim lists GPU architecture and build fingerprint as variation axes, then requires bitwise stability, while the next sentence permits tolerance-based cross-platform parity. Restrict bitwise guarantees to a fixed device/build/backend, and keep architecture or build changes under a tolerance-based contract unless a separate cross-platform bitwise guarantee is intentionally required.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@RL-Kernel-vime-integration-roadmap.md` around lines 84 - 87, The
deterministic operator claim in the roadmap should limit bitwise stability to
identical input bytes on a fixed backend, device architecture, and build
fingerprint. Remove GPU architecture and build changes from the bitwise
variation axes, and explicitly keep cross-architecture or cross-build
comparisons tolerance-based unless a separate cross-platform bitwise contract is
declared.
| Required additions: | ||
|
|
||
| - **Metadata lives on vime samples and the Data Buffer**: consistency fingerprints should travel with `Sample` objects and survive Data Buffer storage, filtering, dynamic sampling, and partial-rollout continuation. | ||
| - **Custom rollout contract**: custom rollout/generate functions should either provide the fields needed for audit, or vime should report a structured "missing consistency metadata" reason and disable strict consistency for that sample. | ||
| - **Sample fields that matter**: `tokens`, `response_length`, `loss_mask`, `status`, `reward`, `rollout_id`, `session_id`, sampling params, finish/truncation reason, old logp source, model name, router/engine identity when available, and weight version. |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
Fail closed when strict-mode metadata is missing.
“Disable strict consistency for that sample” permits a sample with unknown provenance to continue through a supposedly strict batch. This conflicts with the strict-mode requirement for complete metadata and can produce a mixed batch whose overall status is misleading. Specify whether the sample is quarantined, the whole batch/run is downgraded with an explicit status, or strict execution fails.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@RL-Kernel-vime-integration-roadmap.md` around lines 199 - 203, Clarify the
custom rollout contract to fail closed when strict-mode consistency metadata is
missing: quarantine the affected sample, downgrade the entire batch/run with an
explicit status, or fail strict execution. Replace the current “disable strict
consistency for that sample” behavior in the metadata requirements so
unknown-provenance samples cannot continue in a supposedly strict batch.
| - Add `--rlk-fast` and `--rlk-consistency` parsing. | ||
| - Add structured execution-decision records. | ||
| - Add RL-Kernel capability query. | ||
| - Add RL-Kernel per-dtype tolerance lookup without copying threshold values into vime. | ||
| - Add compact rollout metadata fingerprints. | ||
| - Add a Data Buffer-compatible metadata schema for default rollout, custom rollout, dynamic sampling, and partial rollout. | ||
| - Add batch-layout fingerprints for strict/audit replay: sequence lengths, padding/packing, active-token density, microbatch/global-batch shape, and dynamic-sampling group IDs. | ||
| - Add requested-vs-actual provenance records for backend, router, engine topology, dtype, quantization, parallel placement, and lifecycle-sensitive knobs. | ||
| - Add contract IDs to `old_logp` and recomputed logp. | ||
| - Add audit-only `dlogp` diagnostics without enabling fast operators. |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
Define fail-closed behavior for unavailable tolerance contracts.
Strict eligibility depends on an RL-Kernel per-dtype tolerance lookup, but the roadmap does not specify what happens when the dtype, contract ID, or RL-Kernel version has no tolerance entry. Require strict mode to reject or downgrade with a structured reason; it must not proceed with an implicit local default.
🧰 Tools
🪛 LanguageTool
[style] ~510-~510: Three successive sentences begin with the same word. Consider rewording the sentence or use a thesaurus to find a synonym.
Context: ...tructured execution-decision records. - Add RL-Kernel capability query. - Add RL-Ke...
(ENGLISH_WORD_REPEAT_BEGINNING_RULE)
[style] ~511-~511: Three successive sentences begin with the same word. Consider rewording the sentence or use a thesaurus to find a synonym.
Context: ...ds. - Add RL-Kernel capability query. - Add RL-Kernel per-dtype tolerance lookup wi...
(ENGLISH_WORD_REPEAT_BEGINNING_RULE)
[style] ~512-~512: Three successive sentences begin with the same word. Consider rewording the sentence or use a thesaurus to find a synonym.
Context: ...t copying threshold values into vime. - Add compact rollout metadata fingerprints. ...
(ENGLISH_WORD_REPEAT_BEGINNING_RULE)
[style] ~513-~513: Three successive sentences begin with the same word. Consider rewording the sentence or use a thesaurus to find a synonym.
Context: ...ompact rollout metadata fingerprints. - Add a Data Buffer-compatible metadata schem...
(ENGLISH_WORD_REPEAT_BEGINNING_RULE)
[style] ~514-~514: Three successive sentences begin with the same word. Consider rewording the sentence or use a thesaurus to find a synonym.
Context: ...ynamic sampling, and partial rollout. - Add batch-layout fingerprints for strict/au...
(ENGLISH_WORD_REPEAT_BEGINNING_RULE)
[style] ~515-~515: Three successive sentences begin with the same word. Consider rewording the sentence or use a thesaurus to find a synonym.
Context: ...hape, and dynamic-sampling group IDs. - Add requested-vs-actual provenance records ...
(ENGLISH_WORD_REPEAT_BEGINNING_RULE)
[style] ~516-~516: Three successive sentences begin with the same word. Consider rewording the sentence or use a thesaurus to find a synonym.
Context: ...ement, and lifecycle-sensitive knobs. - Add contract IDs to old_logp and recomput...
(ENGLISH_WORD_REPEAT_BEGINNING_RULE)
[style] ~517-~517: Three successive sentences begin with the same word. Consider rewording the sentence or use a thesaurus to find a synonym.
Context: ...Ds to old_logp and recomputed logp. - Add audit-only dlogp diagnostics without ...
(ENGLISH_WORD_REPEAT_BEGINNING_RULE)
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@RL-Kernel-vime-integration-roadmap.md` around lines 508 - 517, Update the
strict eligibility and RL-Kernel per-dtype tolerance lookup requirements to fail
closed when the dtype, contract ID, or RL-Kernel version lacks a tolerance
entry. Require strict mode to reject or explicitly downgrade with a structured
reason, and prohibit continuing with any implicit local default.
| Acceptance criteria: | ||
|
|
||
| - Native vime behavior is unchanged when both switches are `off`. | ||
| - Audit mode can run on native vime and report `dlogp` metrics. | ||
| - Custom rollout samples without required metadata produce structured audit warnings rather than silent strict-mode claims. | ||
| - Strict/audit reports identify active-token count, zero-active-token cases, and the worst-token location for `max_abs_dlogp`. | ||
| - Every operator decision is explainable from logs. |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Define the zero-active-token result explicitly.
max_abs_dlogp and percentile metrics are undefined for an empty active-token set. Add an explicit not_applicable/zero_active_tokens status and exclude the case from pass/fail aggregates so strict acceptance cannot pass vacuously or emit NaN-based results.
Also applies to: 555-560
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@RL-Kernel-vime-integration-roadmap.md` around lines 519 - 525, Update the
strict/audit reporting and pass/fail aggregation requirements to explicitly
handle empty active-token sets: emit a not_applicable/zero_active_tokens status,
avoid calculating or emitting NaN max_abs_dlogp and percentile metrics, and
exclude these cases from pass/fail aggregates so strict acceptance cannot pass
vacuously. Anchor the changes to the acceptance-criteria reporting and
aggregation definitions covering active-token counts and worst-token locations.
|
Closing this because the RL-Kernel/vime integration roadmap is better tracked as an issue/RFC rather than a repository document PR. |
Summary
Notes
Based on RL-Align/RL-Kernel#222.
Thanks @CyberSecurityErial for the cross-config alignment roadmap work in RL-Align/RL-Kernel#222, which informed the consistency-contract and A0-A5 ablation sections.
Validation
Summary by CodeRabbit