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

[inaniloquentee]

Summary

This PR adds the Vime-side provider boundary required to run RL-Kernel train/rollout consistency on ROCm 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 current vLLM CLI and server layouts used by the ROCm TP4/CP2 workload.

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

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 an operator-level 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 mode falls back only when the provider explicitly reports that it is unavailable.
• strict mode 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 and strict execution modes.
• Validates tensor shape, dtype, device, vocabulary partition, identity, and autograd requirements before accepting provider output.

Megatron integration

• Captures the output-layer hidden states, local projection weight and bias, tensor-parallel vocabulary range, and real/padded vocabulary sizes.
• Passes the structural context through the existing loss and log-probability paths without changing Vime's policy-loss semantics.
• Keeps temperature scaling on the native fallback path exactly once.

Current vLLM compatibility

• Accepts both current and legacy locations of FlexibleArgumentParser.
• Accepts both current launcher CLI and OpenAI CLI server argument layouts.
• Derives the accepted server field set from the active vLLM parser when the dataclass-based launcher API is unavailable.

Quick Start

Run the strict ROCm TP4/CP2 path on one node with 8 AMD Instinct MI300X GPUs. Set NUM_ROLLOUT=200 for the full reference horizon.

export VIME_ROOT=/workspace/vime
export RL_KERNEL_ROOT=/tmp/pr396-rollout
export MEGATRON_ROOT=/workspace/Megatron-LM-vime
export MODEL_ROOT=/app/model/Qwen3-8B
export TORCH_DIST_ROOT=/app/model/Qwen3-8B_torch_dist
export PROMPT_DATA=/app/model/dapo-math-17k/dapo-math-17k.jsonl
export NUM_ROLLOUT="${NUM_ROLLOUT:-1}"
export PYTHONPATH="${RL_KERNEL_ROOT}:${PYTHONPATH:-}"

python -m examples.vime_rocm_attention_ablation.run_pr377_workload \
  --case R/R \
  --num-rollout "${NUM_ROLLOUT}" \
  --run-dir "/app/model/vime-runs/pr423-rocm-rr-${NUM_ROLLOUT}" \
  --rl-kernel-root "${RL_KERNEL_ROOT}" \
  --vime-root "${VIME_ROOT}" \
  --megatron-root "${MEGATRON_ROOT}" \
  --model-root "${MODEL_ROOT}" \
  --reference-checkpoint "${TORCH_DIST_ROOT}" \
  --prompt-data "${PROMPT_DATA}" \
  --samples-per-prompt 8 \
  --global-batch-size 8 \
  --max-response-length 7168 \
  --max-tokens-per-gpu 4096 \
  --seed 1234 \
  --rollout-seed 1234 \
  --fixed-paged-tile 128 \
  --paged-kv-max-tokens 8192 \
  --vllm-gpu-memory-utilization 0.38

To run the matched native control, change --case R/R to --case P/P.

200-step reference result

The matched ROCm experiment used:

Item	Value
Hardware	1 node, 8x AMD Instinct MI300X 192GB
Training topology	TP4 / CP2 / PP1
Rollout topology	2 engines, TP4 each
Horizon	200 rollout/training steps
Sampling	1 prompt x 8 samples
Response limit	7,168 tokens
KL loss	enabled, coefficient 0.001

The strict R/R arm achieved bitwise equality for all 200 steps over 9,400,614 compared elements from 1,600 samples.

Training stability

The following figure compares the reported gradient norm and training loss from the two authoritative W&B runs:

• G10 production P/P
• G11 strict R/R

The bold curves are trailing 9-step medians and the lighter curves are the raw per-step values. Both runs remain finite. The mean gradient norm is 0.1552 for G10 and 0.2450 for G11. The reported mean training loss is 6.013e-4 for G10 and 2.433e-6 for G11; this scalar difference is descriptive and is not treated as a reward or optimization-quality comparison.

Raw W&B system metrics

The raw 200-step W&B means are 93.79 versus 72.60 tokens/GPU/s for G10 and G11, and 84.52 versus 99.21 seconds for end-to-end step time. These raw trajectories show a 22.6% G11 rollout-throughput reduction and a 17.4% step-time increase. They are included for transparency, not as a normalized performance claim: response lengths differ across the free-running arms, and the result bundle separately reports warmup-excluded, token-normalized, and equal-length views.

TP4/CP2 module ablation

The short integration matrix used Vime 9404528, RL-Kernel a588f11, 8 rollouts, batch size 1, 8 samples per prompt, maximum response length 7,168, seed 1234, and KL coefficient 0.001.

Group	Attention	FFN	Logp	Status	Mean abs dlogp	Max abs dlogp	tok/GPU/s
M000	P/P	P/P	P/P	NOT_COMPARABLE	0.033510	20.3915	56.63
M100	R/R	P/P	P/P	NOT_COMPARABLE	0.032906	4.03697	53.85
M010	P/P	R/R	P/P	PASS	0.033806	32.6212	99.91
M001	P/P	P/P	R/R	NOT_COMPARABLE	0.033032	4.52478	38.58
M110	R/R	R/R	P/P	PASS	2.41e-6	0.133394	63.61
M101	R/R	P/P	R/R	NOT_COMPARABLE	0.033176	43.8511	43.04
M011	P/P	R/R	R/R	PASS	0.032853	5.17654	74.88
M111	R/R	R/R	R/R	PASS	0	0	56.80

M111 was bitwise equal over 354,543 active tokens. Its rollout throughput was 0.31% above M000 in this short run. M110 removed nearly all average drift, but its 0.133394 maximum remained above the 0.06 numerical threshold; strict Logp removed the remaining tail in M111.

M000, M100, M001, and M101 are marked NOT_COMPARABLE only because production vLLM FFN execution inside the HIP graph did not emit a Python execution record. These arms had no launcher, fallback, sidecar, or numerical errors, but they do not provide the same route-proof strength as the four PASS arms.

Validation

• Target rl-kernel base port: 98 targeted unit tests passed on Linux.
• Latest-main preview patch: 102 targeted unit tests passed on Linux.
• The latest-main ROCm Vime revision completed all eight TP4/CP2 module arms.
• M111 passed exact sidecar validation with torch.equal == true.
• The 200-step strict reference run completed 200/200 steps.
• Existing native execution remains unchanged when no provider is configured.

Compatibility and scope

• The provider is optional and disabled by default.
• No ROCm 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.
• The 200-step figures are reference evidence from RL-Kernel PR #400. The latest-main port is additionally covered by the unit and 8-rollout validation listed above.

[read-the-docs-community]

Documentation build overview

📚 vime | 🛠️ Build #34515813 | 📁 Comparing eeda543 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 provider dispatching, validation, and fallback mechanisms, along with compatibility updates for vLLM. The review feedback focuses on improving robustness and type safety across the new boundary. Key recommendations include safely handling optional fields in _normalize_result, separating provider loading from execution to avoid masking internal bugs, adding explicit torch.Tensor type checks for request and result attributes, handling None values gracefully during vocabulary size parsing, and filtering out non-string action.dest values when collecting vLLM server field names.

[Flink-ddd]

LGTM. I tested the ROCm integration end-to-end on 8× MI300X GPUs. The R/R strict path achieved bitwise-identical training and rollout log probabilities across three rounds, with 0/192 mismatches, while P/P remained on the native vime path. Let’s merge this PR first.