CP LSE merge switches compiled variants on ROCm, breaking forward reproducibility

[JessicaJiang-123]

Summary

The two CP LSE merge functions use @jit_fuser, which calls torch.compile:

• flash_attn_fwd_softmax_lse_correction
• flash_attn_fwd_second_half_softmax_lse_correction

Dynamo may keep multiple compiled variants of the same function in one process. These variants can select different Triton launch configurations. On gfx950, libdevice.log1p is not bit-exact across those configurations. As a result, the reference and actor forwards can disagree even when their inputs and weights are bit-identical.

Symptom

At RL step 0, the reference model and actor are initialized from the same checkpoint, and no optimizer update has happened yet. When evaluated on the same tokens with the same inputs, their forward passes should therefore produce bit-identical logits and log-probabilities.

On 4× MI355X with CP=2, they differ by roughly 1e-5 to 2e-4. The same test is bit-exact on H200 and MI355X with CP=1. With CP=1, the LSE merge is never called.

Root cause

The first mismatch appears in:

layers.0.self_attention.core_attention.flash_attention

The attention inputs and both LSE merge operands are bit-identical. Only the merge result differs. The reference and actor forwards use different compiled variants of the same merge function:

• one takes a scalar log1p path;
• the other takes a vectorized log1p path.

On identical inputs, the two paths differ by about 1 ULP on gfx950.

Testing 48 combinations of XBLOCK and num_warps produces exactly two result groups. The split occurs when LLVM vectorizes the kernel body. Among 18 tested elementwise operations, only these depend on the launch configuration:

log1p    14.0% of elements differ
cosh      3.0% of elements differ

Operations such as exp, log, sqrt, tanh, erf, pow, and fma are invariant. Assembly confirms that vectorization changes the multiply-add contraction pattern inside the OCML implementation.

In the real training run:

• the reference forward exactly matches the scalar-path result;
• the actor forward exactly matches the vectorized-path result.

This roughly 1-ULP difference is amplified across the model's attention layers into the observed 1e-5 to 2e-4 log-probability mismatch.

Why this appears on ROCm

Both ROCm results are within OCML's documented 2 ULP accuracy bound. The problem is that torch.compile allows the process to switch between two valid but non-bit-identical implementations.

Minimal reproducer

import torch
import triton
import triton.language as tl
from triton.language.extra import libdevice

N = 225280


@triton.jit
def kern(in_ptr, out_ptr, xnumel, XBLOCK: tl.constexpr):
    i = tl.program_id(0) * XBLOCK + tl.arange(0, XBLOCK)
    mask = i < xnumel
    x = tl.load(in_ptr + i, mask)
    tl.store(out_ptr + i, libdevice.log1p(x), mask)


def run(x, XBLOCK, num_warps):
    out = torch.empty_like(x)
    kern[(triton.cdiv(N, XBLOCK),)](x, out, N, XBLOCK=XBLOCK, num_warps=num_warps)
    torch.cuda.synchronize()
    return out


torch.manual_seed(0)
x = torch.rand(N, device="cuda", dtype=torch.float32)

a = run(x, XBLOCK=256, num_warps=4)      # 256 / (4 * 64) = 1 element per thread
b = run(x, XBLOCK=1024, num_warps=4)     # 1024 / (4 * 64) = 4 elements per thread

print(f"{int((a != b).sum())} / {N} elements differ, max abs {(a - b).abs().max():.3e}")

ref = torch.log1p(x.double())
ulp = torch.finfo(torch.float32).eps * ref.abs()
for name, o in (("1 element per thread ", a), ("4 elements per thread", b)):
    err = (o.double() - ref).abs() / ulp
    print(f"  {name}: max {err.max():.3f} ULP, mean {err.mean():.4f} ULP")

Output on MI355X (gfx950), ROCm 7.2.0, torch 2.9.1+rocm7.2.0, Triton 3.6.0:

37933 / 225280 elements differ, max abs 5.960e-08
  1 element per thread : max 0.966 ULP, mean 0.2211 ULP
  4 elements per thread: max 1.394 ULP, mean 0.2594 ULP

The same split occurs for CP=2/4/8/16 and THD shapes.

Proposed fix

Both options prevent these functions from switching between static and dynamic compiled variants.

Option A: remove @jit_fuser

Run these two functions eagerly.

Option B: compile these two functions with dynamic=True

Compile these two functions as shape-dynamic from the start, avoiding the initial static specialization.