[RFC] Batch-Invariant RL Kernel Suite for Train-Inference Consistency

[inaniloquentee]

This RFC defines the batch-invariance contract for RL-Kernel's train-inference consistency suite.

The target identity is:

same model + same token sequence + same policy state
=> aligned logprobs / hidden states regardless of surrounding batch shape

This is the P0.3 roadmap item from #83. It is broader than the CUDA deterministic logprob kernel in #96: #96 is one concrete operator path, while this RFC should define the cross-kernel and cross-engine contract that rollout and training backends must satisfy.

Problem

RL rollout and training rarely execute the same sequence in the same shape:

• rollout engines use dynamic batching, chunked prefill, paged KV, and prefix cache;
• reference/reward scoring may use different packing and padding layouts;
• training engines use micro-batches, gradient accumulation, FSDP/TP partitioning, and different row order;
• kernel dispatch may choose different CUDA/ROCm/Triton implementations by shape or hardware.

Small logprob drift can be amplified by GRPO/PPO ratio and KL terms. The suite should make drift visible and locate whether it comes from logprob reduction, attention/prefix reuse, packing layout, TP reduction, dtype rounding, or engine integration.

Scope

Define the expected invariance contract for the same logical sequence across:

• batch size changes: B=1 / small batch / large batch;
• batch position changes and unrelated neighboring samples;
• chunked prefill on/off;
• prefix cache on/off;
• padding layout and varlen packing changes;
• dense vs indexed selected-logprob paths;
• CUDA / ROCm / Triton backend parity, with backend-specific tolerance policy where bitwise identity is not realistic;
• rollout engines such as vLLM/sglang vs training engines such as Megatron/DeepSpeed/FSDP.

Questions To Resolve

• Which outputs require bitwise identity, and which require explicit tolerance?
• What is the canonical fixture format for prompt, completion, token ids, position ids, masks, cache metadata, and policy state?
• Which kernels are included in the first suite: selected logprob, attention with exported LSE, prefix-shared attention, sampling, loss reductions?
• How should failures report the first divergent layer/operator?
• How should TP>1 reductions compare against FSDP(TP=1)?
• What command should contributors run locally before opening PRs?
• Which checks belong in CI, and which require GPU/nightly/benchmark lanes?

Proposed Deliverables

• RFC document describing the invariance contract and tolerance policy.
• A reusable fixture generator for logically identical sequences under different batch/cache/layout configurations.
• A batch-invariance test matrix for selected logprob and attention-adjacent paths.
• A cross-engine benchmark plan comparing rollout logprobs with training logprobs.
• A drift report format that includes max error, sequence id, token position, layer/operator, backend, dtype, and launch configuration.
• Links to implementation issues for concrete kernels and tooling.

Acceptance Criteria

• The RFC clearly separates bitwise-invariant requirements from tolerance-based parity requirements.
• The RFC names the minimal P0 suite and the deferred P1/P2 extensions.
• The RFC references Implement Batch-Invariant Deterministic LogProb CUDA Kernel to Eliminate Batch-Size Drift #96 as the CUDA deterministic selected-logprob implementation path.
• The RFC defines at least one local command shape for contributors and one CI/nightly command shape for maintainers.
• Follow-up issues can be opened directly from the RFC without re-litigating the contract.

[Billy1900]

Drafted an RFC for this: #264 — docs/design/ws1_batch_invariance_contract.md.

Filed as WS1 rather than a new workstream, since WS1 in #83 already is "full batch-invariant forward chain." The RFC is the missing contract document over WS1's existing implementation issues, plus the localisation half of WS4.

How it answers the questions in this issue

"Which outputs require bitwise identity, and which require explicit tolerance?" — three relations, so the threshold follows from the kind of check:

Relation	Varies	Threshold
I (invariance)	batch shape only	bitwise (0, 0)
A (accuracy)	impl vs fp32 gold	#108 accuracy table
P (parity)	backend vs backend	#108 accuracy table

Bitwise on I because kernel, dtype, data and device are held fixed and only the surrounding shape moves — no legitimate source of difference exists, so a tolerance would hide the reduction-schedule defect the suite is meant to catch. Cross-backend parity stays tolerance-based, as this issue anticipated.

Worth noting: this is closer to ratification than proposal. tolerance_contract.json already declares "batch_invariance": {"atol": 0.0, "rtol": 0.0} — and nothing reads it; op_checks.py::_resolve_tolerance only consults contract["accuracy"]. Five test files carry a copy-pasted Axis A -- batch invariance, bitwise comment. The policy exists in practice but is unwired and undocumented.

"Which kernels are included in the first suite?" — P0 covers embedding, rope, rms_norm, attention (+LSE), det_gemm, silu/swiglu, lm_head, linear_logp, logp/batch_invariant_logp, grpo_loss, ratio_kl. Prefix-shared attention and sampling are P1: the former has no operator yet, the latter first needs a decision on what invariance means for a stochastic op (most likely distribution- and greedy-invariance, not sampled tokens).

"How should failures report the first divergent layer/operator?" — a static chain_index over the forward chain; on multi-operator failure the report names the lowest index only, since anything upstream propagates. Operator granularity here; model-level layer probing stays WS4.

"What command should contributors run locally, and which checks belong in CI?" — an I-relation needs no reference implementation, since the kernel is its own oracle. So the P0 suite is single-device, engine-free, and runs in seconds, which makes it viable as a blocking PR gate. Nightly carries large shapes, bucket boundaries, backend parity, and the engine-dependent prefix-cache axis.

"How should TP>1 compare against FSDP(TP=1)?" — left to WS2/#111. WS1 is single-device by construction.

Relationship to #96 and #111

#96 is referenced as the CUDA deterministic selected-logprob implementation path, per the acceptance criteria.

On #111: WS1 is a precondition, not a competitor. The WS2 ablation ladder starts at A0: fully aligned reference and says a failing A0 is not yet a cross-config problem. That is only true if the kernels underneath are batch-invariant — otherwise A0 is nondeterministic under batching and every measurement above it is noise. The six transform classes are named as source_class values compatible with the WS2 KnobDefinition model, so a WS1 axis lifts into a WS2 ablation without redefinition.

Current coverage

Measured, not assumed. Coverage is genuinely deep — T1/T2/T4 hold across 6-9 operators — but it lives in ~30 hand-written tests across a dozen files with no shared axis vocabulary, so there is no way to answer "which (operator x transform x backend) cells are green?" without grepping. grpo_loss and ratio_kl have no invariance tests at all, despite sitting at the end of the chain where drift amplifies into the policy ratio. The work is mostly consolidation rather than new verification.

Deliberately not filed

The RFC's roadmap lists 12 sized follow-up items but I have not opened them as issues — the table is reviewable as-is, and how to slice the work is a maintainer call. Happy to file them once the contract settles.

Most worth challenging

The RFC asserts bitwise as the P0 requirement rather than leaving it open. If invariance should be tolerance-based instead, the suite collapses into a WS2 knob and the separate PR-blocking lane loses its justification.

🤖 Generated with Claude Code