feat(rl-kernel): add CUDA linear logp provider integration

[inaniloquentee]

Summary

This PR adds the Vime-side provider boundary required to run RL-Kernel train/rollout consistency on CUDA while keeping Vime responsible for token layout, masking, loss composition, and framework orchestration.

The integration is fail-closed in strict mode, preserves Vime's native implementation when the provider is not configured, and supports the Qwen3-8B TP4/CP2 workload used by the CUDA reference experiment.

The end-to-end implementation and 200-step reference results are documented in RL-Kernel PR #377.

Motivation

Vime currently computes selected-token log probabilities inside its Megatron loss path. RL-Kernel needs the same logical token layout and model projection metadata to provide a deterministic CUDA implementation without taking ownership of Vime's training loop.

This PR introduces a narrow provider ABI instead of adding RL-Kernel-specific kernel code to Vime:

• Vime builds a structured request from its existing logits, target IDs, tensor-parallel group, token layout, masks, and optional projection context.
• The configured provider returns log probabilities and optional entropy.
• auto falls back only when the provider explicitly reports that it is unavailable.
• strict rejects unavailable providers, malformed outputs, detached gradients, and route mismatches.
• An unconfigured run stays on the existing native Vime path.

Changes

Linear-logp provider contract

• Adds LinearLogpRequest, LinearLogpContext, LinearProjection, VocabPartition, and TokenLayout.
• Adds --linear-logp-provider and --linear-logp-provider-mode.
• Supports auto, strict, audit, and off execution modes through the CUDA launcher.
• Validates tensor shape, dtype, device, vocabulary partition, identity, and autograd requirements before accepting provider output.

Megatron integration

• Captures output-layer hidden states, local projection weight and bias, tensor-parallel vocabulary range, and real/padded vocabulary sizes.
• Reuses strict LM-head local logits when available, avoiding a second projection for the provider path.
• Passes structural context through the existing loss and log-probability paths without changing Vime's policy-loss semantics.
• Keeps native fallback behavior unchanged when no provider is configured.

CUDA compatibility and launch support

• Adds aligned Qwen3-8B TP2/CP2 and TP4/CP2 launcher scripts for one 8-GPU H100 node.
• Supports colocated actor/rollout execution, deterministic CUDA inference, CUDA Graph capture, and FlashAttention routing.
• Keeps context parallelism as row ownership; tensor parallelism remains the vocabulary reduction group.
• Records native/provider route evidence and per-step train/rollout log-probability diagnostics.

Quick Start

Run the strict CUDA TP4/CP2 path on one node with 8 NVIDIA H100 GPUs. The launcher submits train.py directly through Ray; it does not run an ablation-matrix wrapper. Set NUM_ROLLOUT=200 for the full reference horizon.

export VIME_ROOT="$PWD"
export RL_KERNEL_ROOT=/path/to/RL-Kernel
export MEGATRON_ROOT=/root/Megatron-LM
export MODEL_ROOT=/path/to/Qwen3-8B
export TORCH_DIST_ROOT=/path/to/Qwen3-8B_torch_dist
export PROMPT_DATA=/path/to/dapo-math-17k.jsonl

export RL_KERNEL_MODE=strict
export RL_KERNEL_ALIGNED=1
export TP_SIZE=4
export CP_SIZE=2
export ACTOR_GPUS=8
export ROLLOUT_GPUS=8
export ROLLOUT_GPUS_PER_ENGINE=4
export COLOCATE=1
export NUM_ROLLOUT="${NUM_ROLLOUT:-1}"
export ROLLOUT_BATCH_SIZE=8
export N_SAMPLES_PER_PROMPT=16
export GLOBAL_BATCH_SIZE=128
export MAX_RESPONSE_LEN=7168
export MAX_TOKENS_PER_GPU=4096
export VLLM_GPU_MEMORY_UTILIZATION=0.4

scripts/run-qwen3-8B-rlkernel-tp4-cp2.sh

To run the matched native control, use RL_KERNEL_MODE=off with the same workload settings.

200-step reference result

The matched CUDA experiment used:

Item	Value
Hardware	1 node, 8x NVIDIA H100 80GB
Training topology	TP4 / CP2 / PP1
Rollout topology	2 engines, TP4 each
Horizon	200 rollout/training steps
Sampling	8 prompts x 16 samples = global batch 128
Response limit	7,168 tokens
Dynamic batching	maximum 4,096 tokens/GPU
vLLM memory utilization	0.4
KL loss	enabled, coefficient 0.001

The optimized strict R/R arm completed all 200 steps with mismatch_count == 0 and max_abs_diff == 0 at every step. The historical paired end-to-end result was 8.2% faster for G11, with a 95% paired-bootstrap interval of [7.2%, 9.2%]. These figures are descriptive results from the matched artifact bundle in RL-Kernel PR #377, not a new benchmark claim for this Vime patch.

TP4/CP2 module ablation

The short CUDA TP4/CP2 module matrix used Vime eb1f5a0, RL-Kernel d78ef96, 8 rollouts, batch size 1 x 8 samples, maximum response length 7,168, seed 1234, and KL coefficient 0.001 on one node with 8 NVIDIA H100 80GB GPUs.

Group	Attention	FFN	Logp	Status	Mean abs dlogp	Max abs dlogp	Mismatch rate	KL	tok/GPU/s	Step time (s)
M000	0	0	0	PASS	0.013502	0.499342	0.651841	0.00069839	105.71	87.693
M100	1	0	0	PASS	0.0134265	0.715138	0.647915	0.000734545	100.88	91.17
M010	0	1	0	PASS	0.0134155	0.596962	0.645922	0.000793217	112.03	84.557
M001	0	0	1	PASS	0.0138536	0.747852	0.545247	0.000723473	99.166	92.329
M110	1	1	0	PASS	9.42655e-08	2.08616e-06	0.409312	7.85817e-08	116.73	83.515
M101	1	0	1	PASS	0.0136217	0.648769	0.541777	0.000644915	104.1	89.991
M011	0	1	1	PASS	0.0136526	0.936146	0.540853	0.000798233	108.1	87.406
M111	1	1	1	PASS	0	0	0	0	105.45	87.878

M111 was bitwise equal over all 390,806 compared elements, with zero mismatches and zero KL drift. M110 removed nearly all average drift (mean abs dlogp 9.43e-8 and max abs dlogp 2.09e-6), while the other partial combinations retained measurable log-probability differences. All eight module arms completed successfully with exact snapshot validation.

Validation

• Strict R/R completed 200/200 steps with zero train/rollout log-probability mismatches.
• Exact snapshot validation required eight rank files for every completed step.
• CUDA Graph, provider route, fallback, traceback, and step-count checks are recorded by the reference validation scripts.
• Existing native execution remains unchanged when no provider is configured.

Compatibility and scope

• The provider is optional and disabled by default.
• No CUDA kernel implementation is added to Vime; kernels and strict runtime routing remain in RL-Kernel.
• This PR does not change reward computation, policy-loss composition, optimizer behavior, or rollout sampling semantics.
• The 200-step figures are reference evidence from RL-Kernel PR #377; they are included to document the validated contract, not to claim an independently rerun benchmark here.

[read-the-docs-community]

Documentation build overview

📚 vime | 🛠️ Build #34526198 | 📁 Comparing 6c03169 against latest (ce92eff)

  🔍 Preview build  

45 files changed · ± 44 modified · - 1 deleted

± Modified

• genindex.html
• index.html
• search.html
• advanced/arch-support-beyond-megatron.html
• advanced/delta-weight-sync.html
• advanced/external-rollout-engines.html
• advanced/fault-tolerance.html
• advanced/low-precision.html
• advanced/megatron-config.html
• advanced/observability.html
• and 34 more...

- Deleted

• _examples_synced/dspark/README.html

[gemini-code-assist]

Code Review

This pull request introduces an optional provider boundary for Megatron linear_logp computation, allowing custom numerical implementations, tensor-parallel reductions, and runtime validation. It integrates this provider into the loss computation and model forward passes, adds corresponding CLI arguments, and updates vLLM compatibility layers and weight version tracking. The review feedback highlights two important improvements: first, separating the provider loading and execution phases in compute_linear_logp to prevent internal provider errors from triggering silent fallbacks in auto mode; second, adding device and data type validation for the returned entropy tensor to avoid potential runtime crashes during loss computation.

[Flink-ddd]

LGTM. Tested the CUDA integration end-to-end on 8x H100 GPUs for three rounds. R/R strict achieved bitwise-identical log probabilities with 0/126,184 mismatches and no fallback, while P/P remained on the native vime path. Let’s merge this PR first.