From db05b45831a5a534b74510797832dcf9b3c7e7ab Mon Sep 17 00:00:00 2001 From: "Wang, Jian4" <61138589+hzjane@users.noreply.github.com> Date: Sat, 9 May 2026 03:16:46 +0000 Subject: [PATCH] esimd: fix resadd_norm_gemv_int4 race on large N The fused ResAddNormGEMV INT4 kernel has every work-group read residual in pass 1 (for sum_sq), but only n==0 writes the updated residual back to global memory. When N is large enough that work-groups cannot all run concurrently, the later WGs observe the already-updated residual and compute added = h + (h + r_old) instead of h + r_old, which garbles the GEMV output for all downstream rows. Observed on Qwen3.6-27B dense gate_up (N=8704, K=5120, TP=4): kernel rel-error vs IPEX is ~0.10 (quant noise) at N=128 but jumps to 0.28 once N crosses ~512, even though the per-row math is otherwise identical. Qwen3.6-35B MoE router (N=256) stays below the threshold, which is why the bug did not surface there. Fix: for N > 512, split into a prepass kernel that owns the residual / normed_out update (single WG, no cross-WG dep), followed by a pure GEMV-from-normed kernel that just reads normed_out. For N <= 512 keep the original fused path (faster and race-free because all WGs run concurrently on the existing hardware). Verified: * Standalone correctness sweep: rel-error at N=5120/8704, K up to 5120 drops from 0.28-0.33 back to ~0.10. * Qwen3.6-27B sym_int4 TP=4 generates coherent text. * Qwen3.6-35B-A3B sym_int4 TP=4 unchanged (still uses the fused path for N=256). --- .../xpu/esimd_kernels/resadd_norm_gemv_int4.h | 191 +++++++++++++++++- 1 file changed, 190 insertions(+), 1 deletion(-) diff --git a/vllm/custom-esimd-kernels-vllm/csrc/xpu/esimd_kernels/resadd_norm_gemv_int4.h b/vllm/custom-esimd-kernels-vllm/csrc/xpu/esimd_kernels/resadd_norm_gemv_int4.h index d7db26e5..7754455b 100644 --- a/vllm/custom-esimd-kernels-vllm/csrc/xpu/esimd_kernels/resadd_norm_gemv_int4.h +++ b/vllm/custom-esimd-kernels-vllm/csrc/xpu/esimd_kernels/resadd_norm_gemv_int4.h @@ -485,6 +485,174 @@ struct ResAddNormGEMV_int4_ksplit_kernel { } }; +/* ================================================================ + * Prepass kernel: resadd + Gemma-style RMSNorm, K split across WGs. + * + * Motivation: the fused ResAddNormGEMV kernel above performs the + * residual += hidden update only in the n==0 WG, while every WG reads + * residual independently in pass-1 sum_sq. When N is large (more WGs + * than concurrently resident on the device), later WGs read the + * already-updated residual and compute added = h + (h+r_old) instead + * of added = h + r_old — producing incorrect GEMV output. The actual + * bug manifests at N≳512 even for small K. + * + * To avoid a cross-WG sync (not expressible in SYCL), split the work: + * Pass A (this kernel): each WG owns a K-chunk, computes added = h+r, + * writes back residual and normed_out. No cross-WG dependency. + * Pass B (GEMV kernel): reads normed_out, writes output. + * + * Grid: K_WG work-groups × 1 thread, each owns VL=512 elements. + * sum_sq is reduced across WGs via a workgroup-local atomic on an + * SLM-staged float accumulator, then Gemma-style rsqrt is applied. + * + * For simplicity and correctness, this prepass uses a two-phase design + * mirroring fused_add_rms_norm.h: phase 1 computes partial sums with + * atomic add into a device-global scratch; phase 2 reads the final + * inv_rms and writes normed_out. But to avoid adding a scratch buffer + * to the public ABI, we use the caller-provided normed_out as scratch + * for the partial-sum accumulator during phase 1 (word[0] = sum_sq), + * which is then overwritten in phase 2. + * + * Simplest correct implementation: single WG (K<=8192 covers all + * callers), one thread, full K pass. For K=5120 fp16 this runs in a + * few µs and is not a bottleneck relative to the subsequent N×K GEMV. + * ================================================================ */ +struct ResAddNormPrepass_int4_kernel { + fp16* hidden_ptr; + fp16* residual_ptr; // [K] — updated in-place + const fp16* norm_w_ptr; // [K] + fp16* normed_out; // [K] — written + int K; + float eps; + + static constexpr int VL = 512; + + void operator()(sycl::nd_item<1> /*item*/) const SYCL_ESIMD_KERNEL { + const int n_chunks = K / VL; + const int tail = K - n_chunks * VL; + + // Pass 1: residual += hidden, accumulate sum_sq. + float sum_sq = 0.0f; + for (int c = 0; c < n_chunks; c++) { + int off = c * VL; + simd h = block_load(hidden_ptr + off); + simd r = block_load(residual_ptr + off); + simd added = h + r; + block_store(residual_ptr + off, simd(added)); + simd sq = added * added; + sq.select<256,1>(0) += sq.select<256,1>(256); + sq.select<128,1>(0) += sq.select<128,1>(128); + sq.select<64,1>(0) += sq.select<64,1>(64); + sq.select<32,1>(0) += sq.select<32,1>(32); + sq.select<16,1>(0) += sq.select<16,1>(16); + sq.select<8,1>(0) += sq.select<8,1>(8); + sq.select<4,1>(0) += sq.select<4,1>(4); + sq.select<2,1>(0) += sq.select<2,1>(2); + sum_sq += (float)sq[0] + (float)sq[1]; + } + // Handle remainder (callers currently always have K % VL == 0; + // guard for safety). + if (tail > 0) { + for (int i = n_chunks * VL; i < K; i++) { + float h = (float)hidden_ptr[i]; + float r = (float)residual_ptr[i]; + float a = h + r; + residual_ptr[i] = fp16(a); + sum_sq += a * a; + } + } + + float inv_rms = sycl::ext::intel::esimd::rsqrt( + simd(sum_sq / (float)K + eps))[0]; + + // Pass 2: normed = residual * inv_rms * norm_w → normed_out. + for (int c = 0; c < n_chunks; c++) { + int off = c * VL; + simd r = block_load(residual_ptr + off); + simd w = block_load(norm_w_ptr + off); + simd n = r * inv_rms * w; + block_store(normed_out + off, simd(n)); + } + if (tail > 0) { + for (int i = n_chunks * VL; i < K; i++) { + normed_out[i] = fp16((float)residual_ptr[i] * inv_rms * (float)norm_w_ptr[i]); + } + } + } +}; + +/* ================================================================ + * GEMV-from-normed kernel: reads the pre-computed normed_out and does + * INT4 GEMV. Identical math to ResAddNormGEMV_int4_pert_kernel's + * pass-2 without the resadd/norm prelude or the n==0 write-back. + * Grid: N work-groups × 1 thread. + * ================================================================ */ +struct GEMV_int4_from_normed_kernel { + const fp16* normed_ptr; // [K] — pre-computed normed hidden + const int32_t* gemv_weight; // [N, K/8] + const fp16* gemv_scale; // [N, K/128] + fp16* output; // [N] + int N, K; + + static constexpr int BLOCK_SIZE = 128; + static constexpr int PACK = 8; + + void operator()(sycl::nd_item<1> item) const SYCL_ESIMD_KERNEL { + int n = item.get_group(0); + if (n >= N) return; + + constexpr int VL = 512; + constexpr int BLOCKS_PER_VL = VL / BLOCK_SIZE; + constexpr int PACKED_PER_VL = VL / PACK; + const int packed_K = K / PACK; + const int num_blocks_per_row = K / BLOCK_SIZE; + const int n_chunks = K / VL; + + simd acc = 0.0f; + + if (n_chunks > 0) { + xesimd::lsc_prefetch( + gemv_weight + (size_t)n * packed_K); + } + + for (int c = 0; c < n_chunks; c++) { + int offset = c * VL; + if (c + 1 < n_chunks) { + xesimd::lsc_prefetch( + gemv_weight + (size_t)n * packed_K + (c + 1) * PACKED_PER_VL); + } + + simd normed = block_load(normed_ptr + offset); + + simd all_packed = block_load( + gemv_weight + (size_t)n * packed_K + c * PACKED_PER_VL); + + #pragma unroll + for (int blk = 0; blk < BLOCKS_PER_VL; blk++) { + int blk_off = blk * BLOCK_SIZE; + float s = (float)gemv_scale[(size_t)n * num_blocks_per_row + c * BLOCKS_PER_VL + blk]; + float neg_8s = -8.0f * s; + simd blk_packed = all_packed.select<16, 1>(blk * 16); + simd u32 = convert( + blk_packed.template bit_cast_view().read()); + simd w_lo = convert(u32 & 0xFu) * s + neg_8s; + simd w_hi = convert((u32 >> 4) & 0xFu) * s + neg_8s; + acc += normed.select<64, 2>(blk_off) * w_lo + + normed.select<64, 2>(blk_off + 1) * w_hi; + } + } + + acc.select<32,1>(0) += acc.select<32,1>(32); + acc.select<16,1>(0) += acc.select<16,1>(16); + acc.select<8,1>(0) += acc.select<8,1>(8); + acc.select<4,1>(0) += acc.select<4,1>(4); + acc.select<2,1>(0) += acc.select<2,1>(2); + output[n] = fp16((float)acc[0] + (float)acc[1]); + } +}; + /* Host dispatcher — auto-selects K_SPLIT based on N and K */ inline void resadd_norm_gemv_int4_pert_host( fp16* hidden_ptr, fp16* residual_ptr, const fp16* norm_w_ptr, @@ -503,7 +671,28 @@ inline void resadd_norm_gemv_int4_pert_host( while (ks > 1 && (K % (ks * 128) != 0)) ks /= 2; if (ks <= 1) { - // Original optimized kernel (VL=512 register-cached path) + // For large N (> 512) the in-kernel "n==0 writes residual" path + // races with readers in other WGs — split into a prepass + a + // GEMV-only kernel that reads normed_out. For N ≤ 512 the fused + // path is safe (all WGs execute concurrently) and faster. + if (N > 512) { + q.submit([&](sycl::handler& cgh) { + cgh.parallel_for( + sycl::nd_range<1>(1, 1), + ResAddNormPrepass_int4_kernel{ + hidden_ptr, residual_ptr, norm_w_ptr, + normed_out, K, eps}); + }); + q.submit([&](sycl::handler& cgh) { + cgh.parallel_for( + sycl::nd_range<1>(N, 1), + GEMV_int4_from_normed_kernel{ + normed_out, gemv_weight, gemv_scale, + output, N, K}); + }); + return; + } + // Original fused path (safe for small N) q.submit([&](sycl::handler& cgh) { cgh.parallel_for( sycl::nd_range<1>(N, 1),