QSA norm + rope fusion: keep the square out of the first butterfly add so the sum of squares rounds as rms_rows_kernel's does

On gfx1151 the compiler folded x * x into the first add of block_sum (fma of this lane's x with the other lane's
square), a different rounding than the sum of two rounded squares, so about 4% of rows came out with a scale one ulp
off. An empty asm on the square stops the fold (__fmul_rn compiles to the same fma). Checked in the assembly
(v_mul then v_add) and on the device: hip_qsa_nr_parity 44 of 44 OK, every output element bit for bit, T 1..8191;
engine first-token logits with STRATA_QSA_FUSE_NR / STRATA_QSA_FUSE_GATE on and off are byte-identical on a 4.5K and
an 8K prompt (IQ3_S pack).
This commit is contained in:
Niko1221 committed 2026-11-01 03:06:40 +01:00
1 parent 7737bca3eb
commit 66545a497a
1 file changed
+8 -2
+8 -2
View File
@@ -2435,8 +2435,14 @@ __global__ void __launch_bounds__(256) rms_rope_rows_kernel(const float* in, int
const int64_t t = row / heads, h = row % heads;
const int tid = threadIdx.x;
const float x = tid < dim ? in[t * in_ts + h * in_hs + tid] : 0.0f;
float ss = 0.0f;
ss += x * x; // rms_rows_kernel: ss += r[c] * r[c], one column a thread (dim <= 256)
// rms_rows_kernel: ss += r[c] * r[c] from 0, one column a thread (dim <= 256) = the rounded square. The empty asm makes the
// square a value the compiler cannot fold into the first add of block_sum's butterfly (it did: fma(x, x, other lane's
// square) is another rounding than the sum of two rounded squares; hip_qsa_nr_parity found it on gfx1151, and __fmul_rn
// does not stop it: it compiles to the same v_fmac).
float ss = x * x;
#if defined(__HIP_DEVICE_COMPILE__)
asm volatile("" : "+v"(ss));
#endif
const float sc = rsqrtf(block_sum(ss, sh) / (float) dim + eps);
if (tid < dim) sv[tid] = sc * x * w[tid];
__syncthreads();