[RFC] Companion vLLM-rollout workstream to #3743: zero train-generation mismatch via upstream batch-invariant kernels

## Relationship to #3743

#3743 is the tracking issue for zero train-generation mismatch and already covers three complementary routes. This issue is a **companion workstream for the vLLM rollout path**, not a second umbrella: it proposes how NeMo RL adopts the upstream batch-invariant vLLM path and the matching actor-side recompute. Requests: link this issue from #3743 and back. #3743's current success criterion is `gen_kl_error → 0`; this issue **proposes** that #3743 adopt the stricter oracle below for the shared pieces (oracle, WS3), and defers to #3743 on the decision.

## Motivation

On the vLLM rollout path NeMo RL today *corrects* for train-generation mismatch — optional token-level importance sampling and `gen_kl_error` monitoring — rather than removing it.

Public evidence that removal matters at scale is now available, still as an open preview: the DeepSeek-V4 true-on-policy results (verl recipe [verl-project/verl#7688](https://github.com/verl-project/verl/pull/7688), built on [vllm-project/vllm#54955](https://github.com/vllm-project/vllm/pull/54955) and [NVIDIA/Megatron-LM#7050](https://github.com/NVIDIA/Megatron-LM/pull/7050)). The numbers below come from the authors' public write-up ("Consistency practice", https://my.feishu.cn/wiki/HMacwKvaki8BqokUmYFcqoAsn3g), not from the PR text itself; in the reported 30-step DAPO run (32×GB200 / 64×H100):

- with mismatch (`VLLM_BATCH_INVARIANT=0`): `critic/score/mean` showed no score improvement in this reported 30-step run (flat at about −0.85); `rollout_corr/kl` grew to 1.32 and `log_ppl_diff_max` to 13.8;
- with exact 0-diff (`VLLM_BATCH_INVARIANT=1`): score rose from −0.95 to +0.08 and all mismatch metrics were exactly 0 (not an epsilon);
- reported end-to-end overhead was about 10%, still being optimized. The same report notes that the separate logprob-recompute pass becomes unnecessary and that outputs were 10–20% shorter in that run; the latter is an observation, not claimed here as a benefit, and it also complicates throughput comparisons.

This is a large-MoE + FP8 (W8A8) result. Earlier small-model / BF16 reports of "no reward gain from zero mismatch" do not settle this regime, which is where our GB200 recipes are heading.

## Dependencies (upstream, in flight)

- vllm-project/vllm#54955 is a preview (marked do-not-merge) to be split into focused PRs: batch-invariant kernels (fused SiLU + FP8 group quant), deterministic sparse top-k, MoE/DeepEP (low-latency fixed-capacity staging, FP32 router gate, fixed expert capacity), layerwise reload. Default vLLM behavior is unchanged; the aligned path is gated by `VLLM_BATCH_INVARIANT=1` and fails closed on missing kernel support.
- NVIDIA/Megatron-LM#7050 (`experimental/lite`, "mLite"): a composable actor whose `impl=vllm` protocol recomputes logprobs with the same batch-invariant primitives as rollout; also to be split.
- vllm-project/vllm#55627 (draft, replaces vllm-project/vllm#55525): under `VLLM_BATCH_INVARIANT=1` the Mamba2 backend keeps the SSM state in fp32 and runs every chunked-scan call, including each decode step, from the sequence's last chunk boundary, so prefill, chunked prefill and decode produce the same bits. The mechanism applies to every model built on `MambaMixer2`, but has been exercised on two models so far (WS5); design in vllm-project/vllm#55524 (see its 2026-09-06 design update). Not merged; current limitations: TP=PP=1, no prefix caching for the Mamba2 layers, no speculative decoding, no micro-batching (`--enable-dbo` or `--ubatch-size > 1`), no KV connectors, Triton mamba backend only, not together with `--use-replayssm`.
- Both DeepSeek-V4 items are previews today. None of the three is merged. This RFC proposes the NeMo RL design now so that the plumbing and the exactness oracle are ready, and so that the rollout/oracle smoke test (WS2) can run on the batch-invariant path that upstream vLLM already ships for dense models.

## Proposed workstreams

**WS1 — Rollout side (MVP).** A first-class typed option, e.g. `policy.generation.vllm_cfg.batch_invariant: bool` (a validated config field, not a raw environment passthrough) that (a) sets `VLLM_BATCH_INVARIANT=1` for the vLLM workers, (b) forces `logprobs_mode: raw_logprobs` so the rollout side reports the unprocessed log-softmax the oracle compares, and (c) applies a **composite fail-closed gate at construction and per request**: model architecture (hybrid Mamba excluded until WS5 is adopted and certified), attention and MoE backends, quantization scheme, parallelism (TP/EP/DCP), speculative decoding, prefix caching, CUDA-graph mode, and sampling configuration. Raw logprobs are taken before the logits processors, so a true-on-policy claim requires the certified sampling configuration to keep the realized distribution equal to the raw softmax: temperature 1.0, top-p 1.0, top-k off, min-p off, no repetition/presence/frequency penalties, no logit bias, no bad-words or allowed-token constraints, no guided decoding. Greedy decoding (temperature 0 / top-k 1) is a deterministic smoke configuration, not a true-on-policy one. Because the OpenAI-compatible HTTP and NeMo-Gym paths accept sampling parameters per request, every request is validated against the certified configuration and rejected otherwise; construction-time checks alone are not sufficient. The certification matrix pins the vLLM commit, the Model Runner (V1/V2), `max_num_seqs` and the concurrency cap: upstream batch invariance is still beta, and vllm-project/vllm#51187 reports repeat-to-repeat logprob differences at about 44 co-resident sequences that disappear at `--max-num-seqs 8`.

**WS2 — Actor side: logprob recompute on the same numerics (core path).** First reference MVP: mLite `impl=vllm` inside the Megatron policy worker (reuses Megatron-LM's composable primitives; aligns with the DS4 recipe). Later candidate: vLLM-kernel-backed recompute in the DTensor worker (TorchTitan/IsoExec pattern). The contract is broader than "same kernels, same reduction order"; both sides must match on:

- kernel selection and reduction structure (the upstream batch-invariant set);
- sequence packing / padding and varlen layouts;
- context-parallel boundaries (or CP disabled in the certified configuration);
- MoE routing metadata: top-k tie-break rule, expert capacity, dispatch/combine order;
- FP8 weight and activation scales and their cache lifecycle (when scales are recomputed versus reused across refits);
- the log-softmax implementation used to turn logits into logprobs;
- generation-time logits processors (temperature, top-p/k, penalties, logit bias, custom processors): disabled or identity only in aligned mode. The certification validates the final effective `SamplingParams` of every request against a positive allowlist (temperature 1, top-p 1, top-k off, no penalties, no logit bias, no custom processors) and fails closed otherwise; the oracle compares raw logprobs, so a distribution-changing processor would not break the comparison but would make it certify an off-policy sample.

**First rollout/oracle smoke test:** a dense BF16 model already in NeMo RL recipes (proposal: Qwen3-8B, five `grpo-qwen3-8b-*` recipes exist), on the batch-invariant path upstream vLLM ships today. This exercises the WS1 gate and the oracle plumbing; exact zero is **not** expected there until an actor path exists for that model family, so it is a smoke test, not a certification. **First actor path and first certification:** mLite `impl=vllm` inside the Megatron policy worker on a DeepSeek-V4-class FP8 MoE, once the vllm-project/vllm#54955 / NVIDIA/Megatron-LM#7050 splits merge — those kernels are model-family specific, so "adopt the day the splits land" holds for that family only.

**WS3 — Weight-transfer integrity (necessary, not sufficient).** PR #3961 adds opt-in digest verification for the colocated CUDA-IPC refit (`off | log | enforce`). It is a **necessary transport-integrity gate, not a sufficient proof of post-load model equality**: it verifies that the receiver's staging tensor — the input to the vLLM weight loader — has the same bytes, dtype and shape as the tensor the policy workers sent (docs/guides/refit.md, "Verify a Transfer"); it does not cover the loader itself, weight fusion, TP resharding, online quantization, post-load caches, or the loaded model state. Follow-ups: extend the gate to the collective transports; a separate post-load equality check (digest of the materialized vLLM parameters against a reference) is its own deliverable.

**WS4 — Parallel-topology freedom (hypothesis / follow-up).** Measured on GB200 with real DeepEP kernels (4096 tokens × hidden 7168, top-k 8, vLLM finalize semantics): the low-latency combine is bit-identical across EP2↔EP4 (4096/4096 token rows), the high-throughput two-stage combine is not (98.4% of rows change); in that single-node microbenchmark the LL dispatch+combine communication kernels took 1.1–1.5× less time than HT despite moving 2.2–4× more bytes (kernel time of the two communication ops only; not a layer-level or end-to-end number). This covers **one layer's combine only** — not EP1 or EP8, not the router or expert GEMMs, not FP8, not end-to-end logprobs. The evidence supports the hypothesis that trainer EP ≠ rollout EP can be made exact with LL-style combine; the deliverable is a full-stack EP8↔EP1 certification under the oracle below before any topology claim is documented.

**WS5 — Hybrid Mamba models (Nemotron): summary only.** Not covered by the DS4 work; the blocker is kernel-level (Mamba2 chunked-prefill and recurrent-decode kernels differ bitwise, vllm-project/vllm#54993 thread; vllm-project/vllm#38561 closed on the same wall). Our direct-kernel measurements on GB200 quantify the gap (mean |Δ| 0.07–0.22% of mean |output| over 22–45% of elements) and show that replaying the current partial chunk through the prefill kernels from an fp32 boundary state reproduces single-shot prefill bitwise — data posted in [vllm-project/vllm#54993](https://github.com/vllm-project/vllm/pull/54993#issuecomment-5551095754). The upstream fix is now proposed at kernel granularity: draft PR vllm-project/vllm#55627 (replaces vllm-project/vllm#55525; design in vllm-project/vllm#55524, 2026-09-06 design update) makes this the behaviour of the Mamba2 backend under `VLLM_BATCH_INVARIANT=1` for every model built on `MambaMixer2` — no separate option, no per-model opt-in, so WS1's single `batch_invariant` switch covers it once merged. Upstream-side evidence so far is a rollout-side analogue of the oracle below (vLLM decode against vLLM single-shot prefill, not the three-way rollout/recompute/train check), collected with an out-of-tree harness on the PR branch: `AntonV/mamba2-130m-hf`, 8 concurrent requests with mixed chunk offsets, with and without interleaved chunked prefill, 0 of 3666 prefill and 0 of 1190 decode positions differ, in eager mode and with CUDA graphs plus async scheduling; `ibm-granite/granite-4.0-h-350m` (Mamba2 + attention), 4 concurrent requests under a real scheduler preemption, 0 of 1859 and 0 of 5750 positions differ in three combinations (TRITON_ATTN eager, TRITON_ATTN eager with async scheduling, FLASH_ATTN with CUDA graphs). The mechanism applies to every `MambaMixer2` model, but only these two have been exercised. Two findings matter for the actor side: (a) the Triton tile configuration of the SSD kernels can change the bits: forcing every autotune candidate gave identical results for three model shapes (chunk size 256, 48 or more heads), but with chunk size 64 and 8 heads the dt-cumsum kernel's fp32 prefix sums differ between head tiles and the difference reaches the outputs and the final states (GB200, Triton 3.7.1). The trainer and the engine therefore have to run the same tile configurations; vllm-project/vllm#55895 pins them on the vLLM side under the batch-invariant flag, and the actor-side kernels must be pinned to the same values (correction of 2026-09-08; the earlier text here said pinning was unnecessary); (b) the mode costs an fp32 SSM state plus three `chunk_size`-token input buffers (x, raw dt, B) per layer per sequence (about 38 MiB per sequence for a 32-head, head_dim 64, dstate 128, 27-layer model at chunk 256) and runs decode through the chunked-scan kernels (roughly 2–3× the single-token update per layer). Current upstream limitations: pipeline-parallel size 1 (tensor parallelism is supported since 2026-09-08), prefix caching off for the Mamba2 layers, no speculative decoding, no micro-batching (DBO or `--ubatch-size > 1`), no KV connectors, Triton mamba backend only, not together with `--use-replayssm`. WS1 keeps hybrid Mamba models excluded until vllm-project/vllm#55627 lands, NeMo RL adopts it, the actor side runs the same chunk size and kernels, and the end-to-end certification below passes.

## Acceptance criteria: realized-trajectory sampled-token raw-logprob bitwise consistency oracle

Scope of the claim: the oracle compares, for every generated token, the raw logprob of *that* token — one vocabulary coordinate per position — across the three computations. It proves consistency along the realized trajectory; it does **not** prove that the full-vocabulary distributions are identical. A distribution-level check (full logits rows on a sample of positions) is a separate, stronger deliverable.

For the same policy version and over all valid response tokens (generated tokens only; prompt and padding masked out; aligned by absolute position in the response):

- `rollout_raw_logprob == prev_recompute_raw_logprob == first_train_forward_raw_logprob`, where all three are fp32 and compared as bit patterns; `mismatch_count == 0`; all three finite.
- `prev_recompute_raw_logprob` must actually be computed: the certification forbids `loss_fn.force_on_policy_ratio=true`, which skips the recompute and would make the check vacuous. `first_train_forward_raw_logprob` is taken in the training forward in which each token first contributes to the loss, before any `optimizer.step()` that changes that policy version.
- "Same policy version": the first certification disables async rollout, so every trajectory is generated and scored under one weight version. If async rollout is certified later, every trajectory must satisfy `rollout_start_version == rollout_end_version == trainer_version` (NeMo RL already tracks `trainer_version` and `weight_version`), and refit needs a drain/cancel barrier for in-flight generations.
- Aligned mode forces `raw_logprobs` on the rollout side and the certified sampling configuration keeps the realized distribution equal to the raw softmax (see WS1); the recompute and the first training forward compare the same raw log-softmax outputs.
- `gen_kl_error` and the existing mismatch metrics remain runtime observations, not the gate.
- Sustained across refits: proposal N = 10 consecutive refits in the certification run (to be confirmed with #3743).
- Overhead ≤ 15% end-to-end step time on GB200 for the certification recipe, measured against the same recipe with the aligned path off (batch-invariant kernels off, no oracle), with the oracle and refit verification enabled on the aligned side since they are part of the certified configuration.
- Fail-closed behavior verified: every excluded configuration raises at construction, and every non-conforming request is rejected at request time.

## Non-goals

Not a replacement for #3743's Megatron-inference route; the two share WS3 and the oracle and can share WS2's mLite reference.

## Open questions

1. Where the oracle lives: a correctness-suite test versus an always-on assertion in aligned mode.
2. Default EP topology under WS4; whether to expose LL versus HT combine as an explicit knob once certified.
3. FP8: the DS4 path is W8A8 (BF16 master, E4M3 block-FP8 deploy, dynamic FP8 activations) — relationship to the MXFP8 workstream in #3743.
4. Blackwell-specific kernel coverage as the splits land (GB200 is our primary target).
5. Mamba2 under TP > 1 is not yet supported by vllm-project/vllm#55627; Nemotron-H sizes need it before WS5 can be certified, and the partial-chunk buffers are sharded per head so the extension is expected to be mechanical but untested.

Start with #3743 and the proposed WS1/WS2 boundaries, then inspect the five grpo-qwen3-8b-* recipes and docs/guides/refit.md. Review PR #3961 for the existing transfer-integrity gate and trace the policy.generation.vllm_cfg entry point. Done requires an agreed design, certified upstream dependencies, and an exactness oracle, but this RFC does not identify a bounded implementation task.

Written by the indexing model from the issue text.