[WS2][Logprob] Deterministic SM90 linear_logp: fused lm_head projection + frozen LSE reduction contract

[inaniloquentee]

Summary

Make the SM90 fused linear_logp CUDA operator the single deterministic owner of the selected-token log-probability computation, so that training-side scoring and rollout-side sampling record bitwise-identical FP32 selected-token logprobs at the same weight version.

Today the logp stack is aligned after the language-model head has already produced logits through two different projections:

• training: Megatron materializes FP32 logits over its padded vocab shard ([rows, 76032], padded vocab 152064, real vocab 151936),
• rollout: vLLM produces BF16 logits over the real vocab ([1, 151936]) and processed_logprobs are derived from them.

No amount of alignment downstream of those two projections can make the final logp bytes agree. The operator boundary must move before the projection: both engines should call the same linear_logp(hidden, lm_head_weight, target_ids, ...) CUDA op, which fuses the projection GEMM and the streaming vocabulary LSE and never materializes the [N, V] logits tensor on the training path.

This is the CUDA-side closure of the logp contract tracked by #241 and the drift reported in #329.

Background

#329 records the current failure mode for Qwen3-8B TP=2 CP=2 BF16 on 8xH100 (validation artifacts live on that machine; numbers below are from the linked report):

train/train_rollout_logprob_abs_diff              = 0.00813760794699192
train/train_current_rollout_logprob_max_abs_diff  = 0.23483610153198242
train/train_current_rollout_logprob_mismatch_count = 96 / 128

Readbacks show the RL-Kernel integration is active with no fallback, but both sides run backend = pytorch-vocab-parallel-logp-ws2 on top of different logits:

training : logits [256, 76032] float32  (padded vocab 152064, TP=2 local shard)
rollout  : source_logits [1, 151936] bfloat16 (vLLM native)

and the rollout-side native-vs-RL-Kernel selected-logprob deltas are already nonzero before any training-side comparison:

native_vs_rlkernel_selected_diff.max = 0.2781919240951538    # prompt / special-token path
native_vs_rlkernel_selected_diff.max = 7.71600753068924e-07  # generated-token path

The conclusion in #329 is correct: the selected-logp math is aligned, but the lm_head projection feeding it is not. This issue specifies what the deterministic projection+LSE operator must look like, and what has to change in the existing SM90 kernel to earn a strict bitwise contract.

Why the current operators do not qualify

The repo already has three linear_logp implementations and a batch-invariant logp over materialized logits:

Path	Files	Status
PyTorch reference	rl_engine/kernels/ops/pytorch/loss/linear_logp.py	F.linear (cuBLAS) + torch.log_softmax. cuBLAS is shape-keyed: a different M can select a different kernel, split-K count, and K-accumulation order. Cannot be made cross-shape deterministic without replacing the GEMM. Must be barred from strict mode.
Triton	rl_engine/kernels/ops/triton/loss/linear_logp.py	Streaming online-softmax over fixed vocab tiles with input_precision="ieee"; row-local, so batch-invariant in structure. Uses tl.sum in bit-relevant positions (compiler chooses the tree) and an implicit serial rescale chain across tiles. Not the performance path for SM90.
CUDA SM90 fused	csrc/cuda/fused_linear_logp_sm90.cu, rl_engine/kernels/ops/cuda/loss/linear_logp.py	WGMMA/TMA fused projection + streaming softmax, FP32 accumulators, split-V partials + combine kernel. Closest to the goal, but the reduction-shape parameters are not pinned (details below), so its merge tree changes with batch size and it cannot yet carry a strict contract.
batch-invariant logp	rl_engine/kernels/ops/pytorch/loss/batch_invariant_logp.py, batch_invariant_logp_sm90	Operates on already-materialized logits, so it inherits whatever projection produced them. Solves a different boundary.

Specific defects in the SM90 kernel

1. n_split is batch- and occupancy-dependent. csrc/cuda/fused_linear_logp_sm90.cu (~line 1825):

   int n_split = std::max(1, std::min(target_ctas / std::max(row_blocks, 1), total_vtiles));

   The combine kernel then folds n_split partial (max, sum) states. When M changes, row_blocks changes, n_split changes, and the number of rescaled additions in the merge changes with it. This is a reduction whose shape depends on the batch — exactly the property a bitwise contract cannot tolerate.

2. In-tile sum reduction is an unpinned serial chain. The online-softmax epilogue accumulates tsum += __expf(val - tmax) in a plain loop and folds tiles with a running rescale (sSum * __expf(old_max - new_max) + tsum * __expf(tmax - new_max)). Serial-in-index-order is a legitimate choice, but it is currently an accident of the code, not a documented contract constant, and the transcendental choice (__expf vs expf) is likewise unpinned.

3. TP merge defers to torch.logsumexp / .sum(dim=0). _merge_tp_local_logp (pytorch/loss/linear_logp.py:266-293) gathers rank-local stats in rank order (good) and then reduces with torch built-ins whose internal tree is an implementation detail. The merge tree must be owned by the contract, not by the framework.

4. No temperature inside the op. Callers scale logits after materialization (ViME divides by rollout_temperature on the training side), while the rollout applies temperature inside its sampler. Two different rounding points.

5. No dual entry for serving. The rollout needs the FP32 logits row for sampling; the trainer needs only (logp, lse). Today these are different code paths with different materialization boundaries.

Design: the numeric contract

The operator computes, for each row n with final hidden state h ∈ R^D, vocab-parallel weight shard W ∈ R^{V_local × D}, optional bias b, and target id t:

z_v    = <h, W_v> + b_v                  (FP32 accumulator, fixed K order)
m      = max over vocabulary of z_v       (order-free)
S      = Σ_v exp(z_v − m)                 (explicit frozen tree)
lse    = m + log(S)
z_t    = selected logit (single owner)
logp   = min(z_t − lse, 0)                (boundary clamp)

Contract constants (bit-relevant; changes require a version bump)

Constant	Value	Rationale
VOCAB_TILE (vocab tile width)	64 (current BN)	Fixes which vocab columns share a tile summary and the in-tile tree width.
K_SLAB (K-direction slab)	32 (current BK)	Fixes the GEMM K accumulation order: K is walked left-to-right, one FP32 accumulator, no split-K, no atomics.
n_split policy	function of V only	Decouples the merge-tree shape from M, batch, padding, and occupancy.
in-tile sum tree	adjacent-pairwise balanced binary tree over the tile, then a sequential scalar chain across tiles in ascending index order	Every bit-relevant addition has exactly one association; nothing is left to the compiler or launch geometry.
cross-split merge	fixed-order fold over splits in ascending index order; global max first (order-free), then rescaled S and single-owner z_t	Deterministic combine independent of batch.
transcendental	one pinned choice (expf/logf) used by every entry point	__expf vs expf round differently; the contract picks one and both sides use it.
padded-vocab lanes	contribute exactly 0.0 to S and −inf to m; targets are masked to the real vocab	Padded 152064 vs real 151936 boundary stays bit-neutral.
temperature	applied in the epilogue to m/S/z_t (and thus logp), not to stored logits	Same rounding point on both sides; serving logits stay unscaled for the sampler.
output clamp	logp = min(z_t − lse, 0)	Keeps the p ≈ 1 boundary stable.

Contract identity: cuda-fused-linear-logp-sm90-contract-v1 (see provenance below).

Bit-relevant vs bit-neutral launch axes

Not every launch parameter must be frozen — only those that change the shape of a reduction:

• Frozen (contract): vocab tile width, K slab width, n_split policy, reduction trees, transcendental choice, bias add point, temperature application point, padding semantics.
• Free (performance): row-block size / BLOCK_M, CTA count along M, warp count, pipeline stages, shared-memory staging, L2 group swizzle — provided every candidate walks K identically and produces identical per-row bytes.

An M-bucketed launch table is expected (decode rows vs prefill batches), but it may only vary the free axes, and every bucket must pass the same golden gates. OOM fallbacks follow the same rule.

Dual entry points, one arithmetic

// trainer:  never materializes [N, V]
fused_linear_logp_sm90(hidden, weight, bias, targets, /*return_logits=*/false, ...)
    -> {logp[N] fp32, lse[N] fp32}

// serving:  same GEMM + same stats, additionally stores the FP32 logits row
fused_linear_logp_sm90(hidden, weight, bias, targets, /*return_logits=*/true, ...)
    -> {logp[N] fp32, lse[N] fp32, logits[N, V] fp32}

The sampler consumes the stored (unscaled) FP32 logits; the trainer consumes logp/lse directly. Memory and bandwidth for [N, V] are paid only where the row is actually needed.

Tensor parallelism: rank order + owned merge

• Each rank computes rank-local tile summaries over its vocab shard with the same contract constants (the existing fused_linear_logp_sm90_global_target entry is the right place to grow this).
• Partials are exchanged without in-transit reduction (all-gather), placed in rank order.
• The global merge runs the same frozen tree over rank-ordered partials — owned by the CUDA op or a pinned helper, not torch.logsumexp/.sum.
• CP remains row ownership only. A vocabulary LSE must never be folded into a CP reduction (the provider boundary in the ViME integration already states this invariant; keep it).

Implementation plan

Kernel (csrc/cuda/fused_linear_logp_sm90.cu)

• Replace the occupancy-derived n_split with a V-only policy; document it as a contract constant and assert it at launch.
• Rewrite the in-tile sum as an explicit adjacent-pairwise tree (fixed width = VOCAB_TILE), keep the cross-tile fold as an ascending-index scalar chain; remove accidental reliance on thread-stride order.
• Pin the transcendental (expf/logf) at every bit-relevant site (currently __expf in the streaming fold and combine).
• Add the epilogue temperature path (scale stats + selected logit; leave stored logits unscaled).
• Add return_logits dual entry; guarantee identical logp/lse bytes with the flag either way.
• Extend the global-target/TP path with the rank-ordered owned merge kernel.
• Embed the contract version in the extension API (argument or returned provenance struct).
• Keep backward correctness (existing chunked/fused backward variants) — bitwise backward is explicitly out of scope.

Wrapper and registry (rl_engine/kernels/ops/cuda/loss/linear_logp.py, rl_engine/kernels/registry.py)

• Backend id cuda-fused-linear-logp-sm90-contract-v1, reported with (arch, cc, VOCAB_TILE, K_SLAB, n_split_policy, tree_id, transcendental).
• Strict mode selects only contract-bearing backends; the PyTorch cuBLAS path must fail loudly in strict mode instead of silently serving non-deterministic bits (this is how the current pytorch-vocab-parallel-logp-ws2 pairing in [Bugfix] Close the Logp contract for Qwen3-8B TP=2 CP=2 strict bitwise consistency #329 slips through).
• Replace _merge_tp_local_logp's torch.logsumexp/.sum with the owned merge.
• Env knobs only for bit-neutral axes; any attempt to override a contract constant is rejected.

Integration (tracked here, landed with the ViME-side work)

• Training: feed final hidden states + lm_head shard + targets to the op; stop materializing [T, V] logits on the logp path.
• Rollout: compute selected logp through the same op at the same weight version (weight sync must guarantee byte-identical lm_head shards and record a weight-version fingerprint).
• Relay sampling transforms (temperature → top-k → top-p → min-p) as a replayed program with pinned order on both sides — replay the transform program, not a mask.

Validation ladder

Each level must pass before the next; all comparisons are byte-exact on FP32 outputs.

Level	Check	Existing tests to extend
L1	Same program, same inputs, repeated calls → identical bytes	tests/test_deterministic_logp.py
L2	Vary M, batch position, padding, mixed batches, launch bucket → same rows byte-identical	tests/test_batch_invariant_logp.py, tests/test_linear_logp.py
L3	Both engines, same inputs + byte-identical weights → identical logp/lse	tests/test_logprob_comparison.py, tests/test_alignment_model_wrappers.py
L4	Live rollout at a fixed weight version: recorded logprob vs trainer recomputation	ViME diagnostic harness from #329
L5	Full training run: K3 = ρ − 1 − log ρ remains exactly zero for every token	validation report exact-zero metrics

Golden-value gates: per (arch, cc, contract version), freeze reference fingerprints produced under both engines' environments. Agreement is established by behavior (the golden replay), not by source identity, so the two sides may keep separate builds of the same arithmetic.

Acceptance criteria

• Strict mode reports cuda-fused-linear-logp-sm90-contract-v1 active on both training and rollout, with no fallback observed.
• L1–L3 byte-exact in CI; L4 reproduced on the 8xH100 validation topology (TP=2 CP=2, Qwen3-8B, BF16 hidden/weights, FP32 logp).
• The [Bugfix] Close the Logp contract for Qwen3-8B TP=2 CP=2 strict bitwise consistency #329 deltas collapse to exact zeros (mismatch_count = 0, max_abs_diff = 0) at matched weight versions.
• n_split, tile trees, and transcendental choice are assert-pinned at launch; overriding them is impossible without a version bump.
• Claim boundary in the validation report upgrades the logp row from unclaimed to a readback-backed exact-zero claim, without weakening the attention/FFN boundaries ([Bugfix] Define Attention/FFN strict-contract acceptance for Qwen3-8B TP=2 CP=2 #330, [Perf] Add an operator drift/performance trace report for Attention, FFN, Logp, AG, and RS #331).

Performance expectations

• Pinning n_split removes an occupancy-adaptive knob; small-M decode may lose some parallelism. Mitigate with bit-neutral M-bucketed launches (more CTAs along M, unchanged V/K/trees).
• The training path stops paying [N, V] materialization (for Qwen3-8B: ~152k floats/row) — expected net win on the trainer side.
• Report decode tok/s and trainer scoring time against the current non-contract path; keep the non-strict lane available for throughput-critical, non-bitwise use.

Non-goals

• Bitwise backward pass (forward logp bytes are what ρ and K3 consume; backward stays correct, not bit-frozen).
• ROCm/HIP backends — tracked in feat(logprob): add deterministic ROCm vocab-parallel path #328.
• FP8/quantized lm_head, MoE router/expert paths, PP consistency.
• Replacing the Triton implementation; it remains the reference/fallback lane and may adopt the same tree rules separately.

Related

• Closes the logp half of [Bugfix] Close the Logp contract for Qwen3-8B TP=2 CP=2 strict bitwise consistency #329 when L4/L5 pass (parent: [WS2] TP-aware deterministic logprob for cross-config alignment (Qwen3-8B TP=2 CP=2 BF16) #241; ROCm sibling: feat(logprob): add deterministic ROCm vocab-parallel path #328; acceptance framing: [Bugfix] Define Attention/FFN strict-contract acceptance for Qwen3-8B TP=2 CP=2 #330, [Perf] Add an operator drift/performance trace report for Attention, FFN, Logp, AG, and RS #331).
• Existing batch-invariance groundwork: issue [WS1][kernels] Batch-invariant logprob (selected, locked reduction) #148.
• Related reading on train/infer mismatch sources and reduction-order discipline: https://kiddyboots216.github.io/mismatch/