docs: define vime integration design

[inaniloquentee]

Summary

• Add the WS5 RL-Kernel to Vime integration design doc for [WS5] RL-Kernel to vime integration design doc #119.
• Define the non-intrusive activation model, ownership boundaries, lifecycle, tensor/API contracts, and rollout/training-side responsibilities.
• Capture the change-surface estimate across RL-Kernel, Vime/slime adapter work, and upstream vLLM/vllm-router requirements.
• Document risk/fallback behavior and the recommended [WS5] Minimal vime PoC with one RL-Kernel operator #120 PoC path.

Scope

This is a documentation-only PR.

It does not add a production integration, change Vime/slime/vLLM/Megatron source, inject RL-Kernel operators, or claim real-engine train-inference consistency.

Depends on #118
Closes #119
Part of #83

[coderabbitai]

Important

Review skipped

Auto reviews are disabled on base/target branches other than the default branch.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: 0b230eb8-467f-44e4-ba90-273b83405a6b

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch docs/ws5-vime-integration-design

---

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[z1ying]

I think this could be replaced with:

• Temperature scaling: the adapter divides logits by args.rollout_temperature before invoking the op. The RL-Kernel logp op is temperature-unaware and must not apply scaling itself; otherwise the parity comparison against Vime's native compute_log_probs would be made on differently-scaled logits.

[z1ying]

Who implements this? get_op raises RuntimeError (registry.py:315) and shape errors raise at call time (pytorch/loss/logp.py:28) — nothing in RL-Kernel catches them, so shouldn't the doc say the adapter owns this try/except, not "RL-Kernel owns local fallback"?

[z1ying]

Doesn't this import already violate the constraint below? registry.py:337 instantiates KernelRegistry() at import time, which calls torch.cuda.get_device_capability() (registry.py:255) → _lazy_init() → CUDA context on the driver process.

[z1ying]

This snippet no longer matches main: both NativeGRPOLossOp and TritonGRPOLossOp now take (policy_logits, ref_logits, action_ids, old_logps, rewards, completion_mask, ...) — logits-based via the fused ratio/KL op, not current_logps/ref_logps.

Two contract points worth capturing here:

1. __call__/forward expect raw rewards and group-normalize internally, while L238 shows Vime supplies already-normalized batch["advantages"] — feeding those as rewards double-normalizes. Both ops expose apply(..., sample_advantages, ...) for precomputed advantages ("apply takes the per-sequence advantage vector directly", ops/triton/loss/grpo_loss.py:65-66); the doc should direct adapters there.
2. The logits-based signature also changes the tensor contract below ([B, T] → logits are [B, T, V]), which affects the packing-adapter estimate.

[z1ying]

"fp32 preferred" is worth pinning down: the default apply/__call__ entries return the input dtype (e.g. NativeLogpOp.apply uses output_dtype=logits.dtype); fp32 output requires calling apply_fp32 explicitly. Should the contract state that parity comparisons use apply_fp32? Relatedly, since this section motivates parity/drift testing, the separate logp_deterministic op type (batch-invariant, fp32 by design, registry.py:158-160) seems like the natural reference implementation to mention here.

[z1ying]

Thanks for the PR! I've left a few suggestions for clarification.