[Issue]: rmsnorm2d_fwd_with_add truncates fp32->bf16 on gfx942, biasing residual_out and compounding across layers

Problem Description

On gfx942, rmsnorm2d_fwd_with_add converts fp32→bf16 by truncation instead of round-to-nearest-even. Truncation is one-sided, and it is applied to residual_out — the value that becomes the next layer's residual_in — so the error never cancels. After 72 calls (a 36-layer model × 2 norms) the residual stream has shrunk 8.98%.

Per call the error is ≤ 1 ULP, so every tolerance-based test passes. The problem is that the sign never changes.

Source

For hidden_size <= 8192, rmsnorm2d_fwd_with_add routes to add_rmsnorm in module_rmsnorm_quant, compiled with -DOPUS_FP32_to_BF16_DEFAULT=2. In csrc/include/opus/opus.hpp mode 2 is a plain shift:

// 0:standard, 1:truncate_with_nan, 2:truncate, 3:standard asm, 4:rta_asm
else if constexpr (rm == 2) { u32_t z = bit_cast<u32_t>(x);
                              return bit_cast<bf16_t>((unsigned short)(z >> 16)); }

residual_out is bit-identical to a truncation reference (0 of 2,097,152 elements differ) and differs from torch.Tensor.to(torch.bfloat16) in 27–50% of elements, with a mean signed magnitude error of −1.53e-03 … −2.80e-03.

Impact: RL train/inference consistency

Inference alone is fine. The failure mode is that RL frameworks compare rollout logprobs (SGLang → aiter) against actor logprobs recomputed by the training engine (torch/HF). A biased kernel makes the two disagree systematically, and the policy then optimizes into that gap (arXiv:2506.13585).

Qwen3-8B GRPO under verl, 8× gfx942, SGLANG_USE_AITER=1 in both arms, identical config, 20 steps — only the rounding mode differs:

	truncate (current)	round-to-nearest-even
rollout_actor_probs_pearson_corr mean	0.9715	0.9972
drift over 20 steps	−0.0060	−0.0004
inside the healthy 0.997 ± 0.001 band	never	all 20 steps
rollout_probs_diff_mean	0.0140	0.0041

verl flags rollout_probs_diff_mean > 0.01 as a precision anomaly. The truncating arm was run twice independently and the two agree to 8.0e-05 in mean corr, so the decline is not noise. Wall-clock is not comparable between the two arms — the truncating one shared the node with a second job — so only the logprob arithmetic is being compared here, not throughput.

Cost of fixing

The kernel writes two values and only one needs unbiased rounding:

• residual_out — fed forward, so its bias compounds
• out — consumed by the next GEMM, never fed forward; in the quantized configs it does not even reach fp32_to_bf16 (scaled_cast handles int8/fp8/fp4, where a 0.5 ULP bf16 bias is far below the quantization error)

Note add_rmsnorm (FUSE_QUANT=false) and add_rmsnorm_quant (FUSE_QUANT=true) share add_rmsnorm_quant_kernel, so the two paths have to be priced separately. gfx942, .so prebuilt per configuration, measured round-robin so clock drift cancels, best-of-5 × 3 interleaved rounds, mean:

n	m	bf16 no-quant	i8 quant
2048	4096	+4.7%	+30.8%
4096	8192	+3.4%	+13.4%
8192	2048	+2.9%	+19.3%
8192	4096	+5.1%	+11.2%

So RNE is cheap on the bf16 no-quant path — the one vLLM and SGLang call twice per layer — and expensive on the quantized tiles, which is presumably the +181 VALU on the i8 256×32 add tile cited in #4080. Gating on FUSE_QUANT keeps the quantized path at baseline (−0.8% … +1.2%, i.e. noise) while still fixing the no-quant path.

For scale on the no-quant path: the two fused add+norm calls are 1.06–1.74% of the four GEMMs in a Qwen3-8B decoder layer at 256–4096 tokens per forward. That denominator excludes attention, RoPE, SiLU and the q/k norms, so the real share is smaller.

Notes

• Which arches are affected. fp32_to_bf16 short-circuits to static_cast<bf16_t> under (__gfx950__ || __gfx1200__ || __gfx1201__ || __gfx1250__) && __clang_major__ >= 20, ignoring rm entirely, so those targets are exempt. Everything else takes the mode dispatch — gfx942/CDNA3 and also RDNA3 (gfx1100/gfx115x), which is not "older" than gfx942. Whether the exempt path is genuinely RNE is an inference from __bf16 narrowing semantics, not something the source asserts.
• A ULP-tolerance test cannot catch this. Truncation and RNE are both within 1 ULP, so checkAllclose passes either way — which is why test_rmsnorm2d.py is green today. A regression test has to assert on the mean signed error.
• Of the five documented modes, only 0 works. Built and measured one at a time on gfx942: mode 1 emits NaN (never shifts z, so static_cast<unsigned short> keeps the low half of the word); mode 3 rounds correctly in the asm but then converts instead of bit-casting on the way out (return bf16_t(u.i); vs mode 0's __builtin_bit_cast), so 0x3F80 comes back as 16256 rather than 1.0 — max rel err 4.9e+07; mode 4 has no if constexpr branch and does not compile.
• hidden > 8192 is a separate case. It leaves add_rmsnorm for the opus norm kernel, which pins OPUS_FP32_to_BF16_DEFAULT=2 in csrc/include/opus/rmsnorm_opus_kernel.hpp and truncates residual_out there too. Same for the bf16 residual stores in fused_qk_rmsnorm_group_quant.cu and mhc_kernels.cu, which were not audited.
• Related: #4059 states "bf16 stores truncate to match CK", and #4080 commit 9587f1e moves the opus residual_out store from RNE back to truncate citing +181 VALU on the i8 256×32 add tile. That measurement is on the quantized path; the bf16 no-quant path that vLLM/SGLang call twice per layer is the +3.2% row above. "Match the reference" is also circular here, since the reference is the biased implementation.

Operating System

Ubuntu 24.04.4 LTS

CPU

x86_64

GPU

gfx942 (CDNA3)

ROCm Version

ROCm 7.14 (HIP 7.14.60850)

ROCm Component

No response

Steps to Reproduce

torch 2.12.0+rocm7.14, aiter 0.1.12.post2.dev214+gb5e03ed19. Self-contained, no framework dependency:

import torch
import aiter

DEV, EPS, HIDDEN, DEPTH = "cuda", 1e-6, 4096, 72  # hidden<=8192 -> add_rmsnorm; 36 layers x 2

def bf16_truncate(x):   # OPUS_FP32_to_BF16_DEFAULT=2
    return (x.view(torch.int32) >> 16).to(torch.int16).view(torch.bfloat16)

def reference(x, residual, weight, cast):
    acc = x.float() + residual.float()
    var = acc.pow(2).mean(-1, keepdim=True)
    return cast((acc * torch.rsqrt(var + EPS)) * weight.float()), cast(acc)

def call_aiter(x, residual, weight):
    out, residual_out = torch.empty_like(x), torch.empty_like(x)
    aiter.rmsnorm2d_fwd_with_add(out, x, residual, residual_out, weight, EPS)
    torch.cuda.synchronize()
    return out, residual_out

def magnitude_bias(got, exp):  # negative => shrunk toward zero
    g, e = got.float(), exp.float()
    nz = e != 0
    return ((g.abs() - e.abs()) / e.abs().clamp_min(1e-30))[nz].mean().item()

print("1) which rounding mode does residual_out use?")
torch.manual_seed(0)
for m, xs, rs in [(512, 1.0, 1.0), (512, 1.0, 30.0), (4096, 3.0, 3.0)]:
    x = (torch.randn(m, HIDDEN, device=DEV) * xs).to(torch.bfloat16)
    res = (torch.randn(m, HIDDEN, device=DEV) * rs).to(torch.bfloat16)
    w = (torch.randn(HIDDEN, device=DEV) * 0.05 + 1.0).to(torch.bfloat16)
    _, ro = call_aiter(x, res, w)
    _, r_trunc = reference(x, res, w, bf16_truncate)
    _, r_rne = reference(x, res, w, lambda t: t.to(torch.bfloat16))
    print(f"   m={m:<5d} trunc_bitexact={torch.equal(ro, r_trunc)} "
          f"rne_bitexact={torch.equal(ro, r_rne)} bias={magnitude_bias(ro, r_rne):+.3e}")

print(f"\n2) feed residual_out back {DEPTH} times, the way a 36-layer model does")
torch.manual_seed(1)
m = 256
h = torch.randn(m, HIDDEN, device=DEV).to(torch.bfloat16)
w = (torch.randn(HIDDEN, device=DEV) * 0.05 + 1.0).to(torch.bfloat16)
res_k, out_k = h.clone(), h.clone()
res_r, out_r = h.clone(), h.clone()
for step in range(DEPTH):
    g = torch.Generator(device=DEV).manual_seed(1000 + step)  # same delta for both paths
    delta = (torch.randn(m, HIDDEN, device=DEV, generator=g) * 0.5).to(torch.bfloat16)
    out_k, res_k = call_aiter(delta, res_k, w)
    out_r, res_r = reference(delta, res_r, w, lambda t: t.to(torch.bfloat16))
    if (step + 1) % 24 == 0:
        rk = res_k.float().pow(2).mean().sqrt().item()
        rr = res_r.float().pow(2).mean().sqrt().item()
        cos = torch.nn.functional.cosine_similarity(
            out_k.float().flatten(), out_r.float().flatten(), dim=0).item()
        print(f"   after {step+1:3d}: ratio={rk/rr:.6f} "
              f"(shrink {100*(1-rk/rr):+.3f}%)  out cos_sim={cos:.8f}")

(Optional for Linux users) Output of /opt/rocm/bin/rocminfo --support

1) which rounding mode does residual_out use?
   m=512   trunc_bitexact=True rne_bitexact=False bias=-1.532e-03
   m=512   trunc_bitexact=True rne_bitexact=False bias=-2.564e-03
   m=4096  trunc_bitexact=True rne_bitexact=False bias=-1.531e-03

2) feed residual_out back 72 times, the way a 36-layer model does
   after  24: ratio=0.967077 (shrink +3.292%)  out cos_sim=0.99978364
   after  48: ratio=0.938011 (shrink +6.199%)  out cos_sim=0.99924225
   after  72: ratio=0.910199 (shrink +8.980%)  out cos_sim=0.99836612

Rebuilding with OPUS_FP32_to_BF16_DEFAULT=0 turns this clean: rne_bitexact=True, bias=+0.000e+00, shrink +0.000%, cos_sim=1.00000000. Note aiter's JIT cache is keyed only on whether <module>.so exists, never on the compile flags, so the env var alone silently reuses the truncating binary:

rm -f  $AITER_JIT/module_rmsnorm_quant.so
rm -rf $AITER_JIT/build/module_rmsnorm_quant
OPUS_FP32_to_BF16_DEFAULT=0 python3 repro.py

Additional Information

Suggested fix, in order of preference:

1. Force RNE on the residual_out store, gated on FUSE_QUANT so only the no-quant path pays for it. That costs +2.9% … +5.1% there and leaves the quantized tiles at baseline.
2. If even that is unwanted as a default, gate it behind an env var or template flag so RL users can opt in without changing the inference default.
3. Separately, modes 1/3/4 of fp32_to_bf16 look broken and are silent traps for anyone picking from the mode list.

Happy to open a PR for whichever shape you prefer. Since #4080 replaces module_rmsnorm_quant with add_rmsnorm_quant_opus and reverts aiter_opus_plus.h to untouched, the change probably belongs in csrc/include/opus/rmsnorm_opus_io.hpp on top of that PR rather than in the current kernel — happy to rebase and re-measure on the opus kernel.