diff --git a/python/sglang/srt/models/qwen3_5.py b/python/sglang/srt/models/qwen3_5.py index 217cece66..fed78cf88 100644 --- a/python/sglang/srt/models/qwen3_5.py +++ b/python/sglang/srt/models/qwen3_5.py @@ -143,6 +143,10 @@ if _is_cuda: if _is_cpu: fused_sigmoid_mul = torch.ops.sgl_kernel.fused_sigmoid_mul_cpu + fused_qk_gemma_rmsnorm = torch.ops.sgl_kernel.fused_qk_gemma_rmsnorm_cpu + fused_qk_gemma_rmsnorm_with_gate = ( + torch.ops.sgl_kernel.fused_qk_gemma_rmsnorm_with_gate_cpu + ) if _is_npu: from sgl_kernel_npu.norm.split_qkv_rmsnorm_rope import ( @@ -876,7 +880,7 @@ class Qwen3_5AttentionDecoderLayer(nn.Module): k_by_head = k.reshape(-1, self.head_dim) k_by_head = self.k_norm(k_by_head) current_stream.wait_stream(self.alt_stream) - elif _is_hip or _is_xpu: + elif _is_hip or _is_xpu or _is_cpu: q_by_head, k_by_head = fused_qk_gemma_rmsnorm( q, k, @@ -1001,7 +1005,7 @@ class Qwen3_5AttentionDecoderLayer(nn.Module): positions=positions, hidden_states=hidden_states, ) - elif (_is_hip or _is_xpu) and self.attn_output_gate: + elif (_is_hip or _is_xpu or _is_cpu) and self.attn_output_gate: q, k, v, gate = self.forward_prepare_fused_gate( positions=positions, hidden_states=hidden_states, diff --git a/sgl-kernel/csrc/cpu/norm.cpp b/sgl-kernel/csrc/cpu/norm.cpp index c009d1d1d..9c54c0c46 100644 --- a/sgl-kernel/csrc/cpu/norm.cpp +++ b/sgl-kernel/csrc/cpu/norm.cpp @@ -3,559 +3,490 @@ namespace { -// NB: avoid using `at::vec::map<>` on bfloat16 or half -// Llama4TextL2Norm -template -void l2norm_kernel_impl( - scalar_t* __restrict__ output, - const scalar_t* __restrict__ input, - int64_t batch_size, - int64_t seq_len, - int64_t hidden_size, - int64_t input_strideB, - int64_t input_strideS, - int64_t output_strideB, - int64_t output_strideS, - float eps = 1e-5) { - using bVec = at::vec::Vectorized; - using fVec = at::vec::Vectorized; +struct NormParams { + // Treat all tensors as [B, H, T, D]: + // 2D -> [B, 1, 1, D] + // 3D -> [B, 1, T, D] + // 4D -> [B, H, T, D] + // + // Input: last dimension contiguous. + // Output: contiguous. - constexpr int kVecSize = bVec::size(); - at::parallel_for(0, batch_size * seq_len, 0, [&](int64_t begin, int64_t end) { - int64_t bi{0}, si{0}; - data_index_init(begin, bi, batch_size, si, seq_len); - for (int64_t i = begin; i < end; ++i) { - // local ptrs - scalar_t* __restrict__ out_ptr = output + bi * output_strideB + si * output_strideS; - const scalar_t* __restrict__ input_ptr = input + bi * input_strideB + si * input_strideS; + int ndim{0}; + int64_t B{1}, H{1}, T{1}, D{1}; + int64_t i_strideB{0}, i_strideH{0}, i_strideT{0}; + float eps{1e-5f}; + float shift{0.f}; - fVec sum_fvec = fVec(float(0)); - float sum_val = float(0); + const void* weight{nullptr}; + const void* bias{nullptr}; - int64_t d; -#pragma GCC unroll 4 - for (d = 0; d <= hidden_size - kVecSize; d += kVecSize) { - bVec x_bvec = bVec::loadu(input_ptr + d); - fVec x_fvec0, x_fvec1; - std::tie(x_fvec0, x_fvec1) = at::vec::convert_to_float(x_bvec); + explicit NormParams(const at::Tensor& input, float eps_) : ndim(input.dim()), eps(eps_) { + TORCH_CHECK(ndim >= 2 && ndim <= 4, "Expected a 2D/3D/4D tensor, got ", ndim, "D."); - sum_fvec += x_fvec0 * x_fvec0; - sum_fvec += x_fvec1 * x_fvec1; - } -#pragma GCC unroll 4 - for (; d < hidden_size; ++d) { - float x_val = static_cast(input_ptr[d]); - sum_val += x_val * x_val; - } - - sum_val += vec_reduce_sum(sum_fvec); - float rsqrt_var = float(1) / std::sqrt(sum_val / hidden_size + eps); - const fVec scale_fvec = fVec(rsqrt_var); - -#pragma GCC unroll 4 - for (d = 0; d <= hidden_size - kVecSize; d += kVecSize) { - bVec x_bvec = bVec::loadu(input_ptr + d); - fVec x_fvec0, x_fvec1; - std::tie(x_fvec0, x_fvec1) = at::vec::convert_to_float(x_bvec); - - x_fvec0 = x_fvec0 * scale_fvec; - x_fvec1 = x_fvec1 * scale_fvec; - - bVec out_bvec = convert_from_float_ext(x_fvec0, x_fvec1); - out_bvec.store(out_ptr + d); - } -#pragma GCC unroll 4 - for (; d < hidden_size; ++d) { - float x_val = static_cast(input_ptr[d]); - out_ptr[d] = static_cast(x_val * rsqrt_var); - } - // move to the next index - data_index_step(bi, batch_size, si, seq_len); + B = input.size(0); + D = input.size(ndim - 1); + i_strideB = input.stride(0); + switch (ndim) { + case 2: + break; + case 3: + T = input.size(1); + i_strideT = input.stride(1); + break; + case 4: + H = input.size(1); + T = input.size(2); + i_strideH = input.stride(1); + i_strideT = input.stride(2); + break; + default: + TORCH_INTERNAL_ASSERT(false); } - }); -} + } -template -void rmsnorm_kernel_impl( - scalar_t* __restrict__ output, - const scalar_t* __restrict__ input, - const scalar_t* __restrict__ weight, - int64_t batch_size, - int64_t seq_len, - int64_t hidden_size, - int64_t input_strideB, - int64_t input_strideS, - int64_t output_strideB, - int64_t output_strideS, - const func_t& f, - const vec_func_t& vf, - float eps = 1e-5) { - using bVec = at::vec::Vectorized; - using fVec = at::vec::Vectorized; + inline int64_t rows() const { + return B * H * T; + } + inline int64_t input_offset(int64_t b, int64_t h, int64_t t) const { + return b * i_strideB + h * i_strideH + t * i_strideT; + } + inline int64_t output_offset(int64_t b, int64_t h, int64_t t) const { + return ((b * H + h) * T + t) * D; + } +}; - constexpr int kVecSize = bVec::size(); - at::parallel_for(0, batch_size * seq_len, 0, [&](int64_t begin, int64_t end) { - int64_t bi{0}, si{0}; - data_index_init(begin, bi, batch_size, si, seq_len); - for (int64_t i = begin; i < end; ++i) { - // local ptrs - scalar_t* __restrict__ out_ptr = output + bi * output_strideB + si * output_strideS; - const scalar_t* __restrict__ input_ptr = input + bi * input_strideB + si * input_strideS; +enum class NormMode { + L2Norm, // y = x / sqrt(mean(x^2) + eps) + RMSNorm, // y = x * weight / sqrt(mean(x^2) + eps) + GemmaNorm, // y = x * (weight + scale_shift) / sqrt(mean(x^2) + eps) + LayerNorm, // y = (x - mean(x)) * weight / sqrt(var(x) + eps) + bias + RMSNormGated, // y = x * weight / sqrt(mean(x^2) + eps) * SiLU(gate) +}; - fVec sum_fvec = fVec(float(0)); - float sum_val = float(0); +struct NormTraitsBase { + static constexpr bool has_weight = false; + static constexpr bool has_bias = false; + static constexpr bool has_shift = false; + static constexpr bool has_mean = false; + static constexpr bool has_gate = false; - int64_t d; -#pragma GCC unroll 4 - for (d = 0; d <= hidden_size - kVecSize; d += kVecSize) { - bVec x_bvec = bVec::loadu(input_ptr + d); - fVec x_fvec0, x_fvec1; - std::tie(x_fvec0, x_fvec1) = at::vec::convert_to_float(x_bvec); + template + static inline VT apply_weight(VT x, VT w) { + return x * w; + } +#if defined(CPU_CAPABILITY_AVX512) + static inline __m512 apply_weight(__m512 x, __m512 w) { + return _mm512_mul_ps(x, w); + } +#endif +}; - sum_fvec += x_fvec0 * x_fvec0; - sum_fvec += x_fvec1 * x_fvec1; - } -#pragma GCC unroll 4 - for (; d < hidden_size; ++d) { - float x_val = static_cast(input_ptr[d]); - sum_val += x_val * x_val; +template +struct NormTraits : NormTraitsBase {}; + +template <> +struct NormTraits : NormTraitsBase { + static constexpr bool has_weight = true; +}; + +template <> +struct NormTraits : NormTraitsBase { + static constexpr bool has_weight = true; + static constexpr bool has_shift = true; + + template + static inline VT apply_shift(VT w, VT shift) { + return w + shift; + } +#if defined(CPU_CAPABILITY_AVX512) + static inline __m512 apply_shift(__m512 w, __m512 shift) { + return _mm512_add_ps(w, shift); + } +#endif +}; + +// LayerNorm: Var(X) = E(X^2) - (E(X))^2, refer to FlashInfer impl: +// https://github.com/flashinfer-ai/flashinfer/blob/main/include/flashinfer/norm.cuh#L552 +template <> +struct NormTraits : NormTraitsBase { + static constexpr bool has_weight = true; + static constexpr bool has_bias = true; + static constexpr bool has_mean = true; + + template + static inline VT apply_bias(VT x, VT bias) { + return x + bias; + } +#if defined(CPU_CAPABILITY_AVX512) + static inline __m512 apply_bias(__m512 x, __m512 bias) { + return _mm512_add_ps(x, bias); + } +#endif +}; + +template <> +struct NormTraits : NormTraitsBase { + static constexpr bool has_weight = true; + static constexpr bool has_gate = true; + + static inline float apply_gate(float x, float gate) { + return x * (gate / (1.f + std::exp(-gate))); + } + static inline at::vec::Vectorized apply_gate(at::vec::Vectorized x, at::vec::Vectorized gate) { + const auto one = at::vec::Vectorized(1.f); + return x * (gate / (one + gate.neg().exp_u20())); + } +#if defined(CPU_CAPABILITY_AVX512) + static inline __m512 apply_gate(__m512 x, __m512 gate) { + __m512 minus_gate = _mm512_xor_ps(_mm512_set1_ps(-0.f), gate); + __m512 denom = _mm512_add_ps(_mm512_exp_u20_ps(minus_gate), _mm512_set1_ps(1.0f)); + // NOTE: avoid vdivps -> use reciprocal + __m512 sigmoid = _mm512_mul_ps(gate, _mm512_rcp14_ps(denom)); + return _mm512_mul_ps(x, sigmoid); + } +#endif +}; + +template +struct NormReduce; + +#if defined(CPU_CAPABILITY_AVX512) +template +struct NormReduce { + static inline void apply( + at::BFloat16* __restrict__ out, + const at::BFloat16* __restrict__ input, + const at::BFloat16* __restrict__ gate, + const NormParams& params) { + static_assert(D % 32 == 0); + constexpr int COLS = D / 32; + + const bool use_bias = params.bias != nullptr; + + __m512bh va[COLS]; + __m512 vmean, vrscale; + const __m512 vshift = _mm512_set1_ps(params.shift); + + // step 1: load input and do reduce with avx512-bf16 + __m512 vsum = _mm512_set1_ps(0.f); + __m512 vsum2 = _mm512_set1_ps(0.f); + Unroll{}([&](auto col) { + va[col] = (__m512bh)(_mm512_loadu_si512(input + col * 32)); + if constexpr (NormTraits::has_mean) { + vsum = _mm512_add_ps(vsum, CVT_BF16_TO_FP32(_mm512_extracti32x8_epi32((__m512i)va[col], 0))); + vsum = _mm512_add_ps(vsum, CVT_BF16_TO_FP32(_mm512_extracti32x8_epi32((__m512i)va[col], 1))); } + vsum2 = _mm512_dpbf16_ps(vsum2, va[col], va[col]); + }); - sum_val += vec_reduce_sum(sum_fvec); - float rsqrt_var = float(1) / std::sqrt(sum_val / hidden_size + eps); - const fVec scale_fvec = fVec(rsqrt_var); - -#pragma GCC unroll 4 - for (d = 0; d <= hidden_size - kVecSize; d += kVecSize) { - bVec x_bvec = bVec::loadu(input_ptr + d); - fVec x_fvec0, x_fvec1; - std::tie(x_fvec0, x_fvec1) = at::vec::convert_to_float(x_bvec); - - bVec w_bvec = bVec::loadu(weight + d); - fVec w_fvec0, w_fvec1; - std::tie(w_fvec0, w_fvec1) = at::vec::convert_to_float(w_bvec); - - x_fvec0 = x_fvec0 * scale_fvec * vf(w_fvec0); - x_fvec1 = x_fvec1 * scale_fvec * vf(w_fvec1); - - bVec out_bvec = convert_from_float_ext(x_fvec0, x_fvec1); - out_bvec.store(out_ptr + d); - } -#pragma GCC unroll 4 - for (; d < hidden_size; ++d) { - float x_val = static_cast(input_ptr[d]); - float w_val = static_cast(weight[d]); - out_ptr[d] = static_cast(x_val * rsqrt_var * f(w_val)); - } - // move to the next index - data_index_step(bi, batch_size, si, seq_len); + // compute mean (if has_mean) and rscale + float sum2 = _mm512_reduce_add_ps(vsum2); + float variance = sum2 / D; + if constexpr (NormTraits::has_mean) { + float sum = _mm512_reduce_add_ps(vsum); + float mean = sum / D; + variance -= mean * mean; + vmean = _mm512_set1_ps(mean); } - }); -} + float rscale = 1.f / std::sqrt(variance + params.eps); + vrscale = _mm512_set1_ps(rscale); -template -void gemma3_rmsnorm_kernel_4d_impl( - scalar_t* __restrict__ output, - const scalar_t* __restrict__ input, - const scalar_t* __restrict__ weight, - int64_t batch_size, - int64_t num_head, - int64_t seq_len, - int64_t hidden_size, - int64_t input_strideB, - int64_t input_strideH, - int64_t input_strideS, - int64_t output_strideB, - int64_t output_strideH, - int64_t output_strideS, - float eps = 1e-5) { - using bVec = at::vec::Vectorized; - using fVec = at::vec::Vectorized; - - constexpr int kVecSize = bVec::size(); - at::parallel_for(0, batch_size * num_head * seq_len, 0, [&](int64_t begin, int64_t end) { - int64_t bi{0}, hi{0}, si{0}; - data_index_init(begin, bi, batch_size, hi, num_head, si, seq_len); - for (int64_t i = begin; i < end; ++i) { - // local ptrs - scalar_t* __restrict__ out_ptr = output + bi * output_strideB + hi * output_strideH + si * output_strideS; - const scalar_t* __restrict__ input_ptr = input + bi * input_strideB + hi * input_strideH + si * input_strideS; - - fVec sum_fvec = fVec(float(0)); - float sum_val = float(0); - fVec one_fvec = fVec(float(1)); - - int64_t d; -#pragma GCC unroll 4 - for (d = 0; d <= hidden_size - kVecSize; d += kVecSize) { - bVec x_bvec = bVec::loadu(input_ptr + d); - fVec x_fvec0, x_fvec1; - std::tie(x_fvec0, x_fvec1) = at::vec::convert_to_float(x_bvec); - - sum_fvec += x_fvec0 * x_fvec0; - sum_fvec += x_fvec1 * x_fvec1; + // step 2: apply scale to output + Unroll{}([&](auto col) { + __m512i a16 = (__m512i)va[col]; + __m512 va0 = CVT_BF16_TO_FP32(_mm512_extracti32x8_epi32(a16, 0)); + __m512 va1 = CVT_BF16_TO_FP32(_mm512_extracti32x8_epi32(a16, 1)); + if constexpr (NormTraits::has_mean) { + va0 = _mm512_sub_ps(va0, vmean); + va1 = _mm512_sub_ps(va1, vmean); } -#pragma GCC unroll 4 - for (; d < hidden_size; ++d) { - float x_val = static_cast(input_ptr[d]); - sum_val += x_val * x_val; + va0 = _mm512_mul_ps(va0, vrscale); + va1 = _mm512_mul_ps(va1, vrscale); + if constexpr (NormTraits::has_weight) { + // TODO: need to block B to hide weight reload + const at::BFloat16* weight = static_cast(params.weight); + __m512i w16 = (__m512i)(_mm512_loadu_si512(weight + col * 32)); + __m512 w0 = CVT_BF16_TO_FP32(_mm512_extracti32x8_epi32(w16, 0)); + __m512 w1 = CVT_BF16_TO_FP32(_mm512_extracti32x8_epi32(w16, 1)); + if constexpr (NormTraits::has_shift) { + w0 = NormTraits::apply_shift(w0, vshift); + w1 = NormTraits::apply_shift(w1, vshift); + } + va0 = NormTraits::apply_weight(va0, w0); + va1 = NormTraits::apply_weight(va1, w1); } - - sum_val += vec_reduce_sum(sum_fvec); - float rsqrt_var = float(1) / std::sqrt(sum_val / hidden_size + eps); - const fVec scale_fvec = fVec(rsqrt_var); - -#pragma GCC unroll 4 - for (d = 0; d <= hidden_size - kVecSize; d += kVecSize) { - bVec x_bvec = bVec::loadu(input_ptr + d); - fVec x_fvec0, x_fvec1; - std::tie(x_fvec0, x_fvec1) = at::vec::convert_to_float(x_bvec); - - bVec w_bvec = bVec::loadu(weight + d); - fVec w_fvec0, w_fvec1; - std::tie(w_fvec0, w_fvec1) = at::vec::convert_to_float(w_bvec); - - x_fvec0 = x_fvec0 * scale_fvec * (w_fvec0 + one_fvec); - x_fvec1 = x_fvec1 * scale_fvec * (w_fvec1 + one_fvec); - - bVec out_bvec = convert_from_float_ext(x_fvec0, x_fvec1); - out_bvec.store(out_ptr + d); + if constexpr (NormTraits::has_bias) { + if (use_bias) { + const at::BFloat16* bias = static_cast(params.bias); + __m512i b16 = (__m512i)(_mm512_loadu_si512(bias + col * 32)); + __m512 vbias0 = CVT_BF16_TO_FP32(_mm512_extracti32x8_epi32(b16, 0)); + __m512 vbias1 = CVT_BF16_TO_FP32(_mm512_extracti32x8_epi32(b16, 1)); + va0 = NormTraits::apply_bias(va0, vbias0); + va1 = NormTraits::apply_bias(va1, vbias1); + } } -#pragma GCC unroll 4 - for (; d < hidden_size; ++d) { - float x_val = static_cast(input_ptr[d]); - float w_val = static_cast(weight[d]); - out_ptr[d] = static_cast(x_val * rsqrt_var * (w_val + 1)); + if constexpr (NormTraits::has_gate) { + __m512i g16 = (__m512i)(_mm512_loadu_si512(gate + col * 32)); + __m512 vgate0 = CVT_BF16_TO_FP32(_mm512_extracti32x8_epi32(g16, 0)); + __m512 vgate1 = CVT_BF16_TO_FP32(_mm512_extracti32x8_epi32(g16, 1)); + va0 = NormTraits::apply_gate(va0, vgate0); + va1 = NormTraits::apply_gate(va1, vgate1); } - // move to the next index - data_index_step(bi, batch_size, hi, num_head, si, seq_len); - } - }); -} + _mm512_storeu_si512(out + col * 32, (__m512i)(_mm512_cvtne2ps_pbh(va1, va0))); + }); + } +}; +#endif -template -void fused_add_rmsnorm_kernel_impl( - scalar_t* __restrict__ input, - scalar_t* __restrict__ residual, - const scalar_t* __restrict__ weight, - float* __restrict__ buffer, - int64_t batch_size, - int64_t seq_len, - int64_t hidden_size, - int64_t input_strideB, - int64_t input_strideS, - const func_t& f, - const vec_func_t& vf, - float eps = 1e-5) { - using bVec = at::vec::Vectorized; - using fVec = at::vec::Vectorized; +template +struct NormReduceGeneric { + static inline void apply( + scalar_t* __restrict__ out, + const scalar_t* __restrict__ input, + const scalar_t* __restrict__ gate, + scalar_t* __restrict__ residual, + const NormParams& params, + int D) { + using bVec = at::vec::Vectorized; + using fVec = at::vec::Vectorized; + constexpr int kVecSize = bVec::size(); - constexpr int kVecSize = bVec::size(); - at::parallel_for(0, batch_size * seq_len, 0, [&](int64_t begin, int64_t end) { - int64_t bi{0}, si{0}; - data_index_init(begin, bi, batch_size, si, seq_len); - int tid = at::get_thread_num(); - float* __restrict__ buffer_ptr = buffer + tid * hidden_size; + const bool use_bias = params.bias != nullptr; + fVec sum_fvec{0.f}, sum2_fvec{0.f}; + float sum_val{0.f}, sum2_val{0.f}; - for (int64_t i = begin; i < end; ++i) { - // local ptrs - scalar_t* __restrict__ input_ptr = input + bi * input_strideB + si * input_strideS; - scalar_t* __restrict__ residual_ptr = residual + i * hidden_size; - - fVec sum_fvec = fVec(float(0)); - float sum_val = float(0); - - int64_t d; + int d; #pragma GCC unroll 4 - for (d = 0; d <= hidden_size - kVecSize; d += kVecSize) { - bVec x_bvec = bVec::loadu(input_ptr + d); - fVec x_fvec0, x_fvec1; - std::tie(x_fvec0, x_fvec1) = at::vec::convert_to_float(x_bvec); - - bVec r_bvec = bVec::loadu(residual_ptr + d); - fVec r_fvec0, r_fvec1; - std::tie(r_fvec0, r_fvec1) = at::vec::convert_to_float(r_bvec); - + for (d = 0; d <= D - kVecSize; d += kVecSize) { + auto [x_fvec0, x_fvec1] = load_float_vec2(input + d); + if constexpr (has_residual) { + auto [r_fvec0, r_fvec1] = load_float_vec2(residual + d); x_fvec0 += r_fvec0; x_fvec1 += r_fvec1; - - bVec out_bvec = convert_from_float_ext(x_fvec0, x_fvec1); - out_bvec.store(residual_ptr + d); - - sum_fvec += x_fvec0 * x_fvec0; - sum_fvec += x_fvec1 * x_fvec1; - - x_fvec0.store(buffer_ptr + d); - x_fvec1.store(buffer_ptr + d + fVec::size()); } -#pragma GCC unroll 4 - for (; d < hidden_size; ++d) { - float x_val = static_cast(input_ptr[d]); - float r_val = static_cast(residual_ptr[d]); - - x_val += r_val; - residual_ptr[d] = static_cast(x_val); - - sum_val += x_val * x_val; - buffer_ptr[d] = x_val; - } - - sum_val += vec_reduce_sum(sum_fvec); - float rsqrt_var = float(1) / std::sqrt(sum_val / hidden_size + eps); - const fVec scale_fvec = fVec(rsqrt_var); - -#pragma GCC unroll 4 - for (d = 0; d <= hidden_size - kVecSize; d += kVecSize) { - fVec x_fvec0 = fVec::loadu(buffer_ptr + d); - fVec x_fvec1 = fVec::loadu(buffer_ptr + d + fVec::size()); - - bVec w_bvec = bVec::loadu(weight + d); - fVec w_fvec0, w_fvec1; - std::tie(w_fvec0, w_fvec1) = at::vec::convert_to_float(w_bvec); - - x_fvec0 = x_fvec0 * scale_fvec * vf(w_fvec0); - x_fvec1 = x_fvec1 * scale_fvec * vf(w_fvec1); - bVec x_bvec = convert_from_float_ext(x_fvec0, x_fvec1); - x_bvec.store(input_ptr + d); - } -#pragma GCC unroll 4 - for (; d < hidden_size; ++d) { - float x_val = buffer_ptr[d] * rsqrt_var * static_cast(f(weight[d])); - input_ptr[d] = x_val; - } - // move to the next index - data_index_step(bi, batch_size, si, seq_len); - } - }); -} - -template -void fused_rmsnorm_gated_kernel_impl( - scalar_t* __restrict__ output, - const scalar_t* __restrict__ input, - const scalar_t* __restrict__ weight, - const scalar_t* __restrict__ gate, - int64_t batch_size, - int64_t hidden_size, - int64_t input_strideN, - float eps = 1e-5) { - using bVec = at::vec::Vectorized; - using fVec = at::vec::Vectorized; - const fVec one = fVec(1.f); - - constexpr int kVecSize = bVec::size(); - at::parallel_for(0, batch_size, 0, [&](int64_t begin, int64_t end) { - for (int64_t i = begin; i < end; ++i) { - // local ptrs - scalar_t* __restrict__ out_ptr = output + i * hidden_size; - const scalar_t* __restrict__ input_ptr = input + i * input_strideN; - const scalar_t* __restrict__ gate_ptr = gate + i * hidden_size; - - fVec sum_fvec = fVec(float(0)); - float sum_val = float(0); - - int64_t d; -#pragma GCC unroll 4 - for (d = 0; d <= hidden_size - kVecSize; d += kVecSize) { - bVec x_bvec = bVec::loadu(input_ptr + d); - fVec x_fvec0, x_fvec1; - std::tie(x_fvec0, x_fvec1) = at::vec::convert_to_float(x_bvec); - - sum_fvec += x_fvec0 * x_fvec0; - sum_fvec += x_fvec1 * x_fvec1; - } -#pragma GCC unroll 4 - for (; d < hidden_size; ++d) { - float x_val = static_cast(input_ptr[d]); - sum_val += x_val * x_val; - } - - sum_val += vec_reduce_sum(sum_fvec); - float rsqrt_var = float(1) / std::sqrt(sum_val / hidden_size + eps); - const fVec scale_fvec = fVec(rsqrt_var); - -#pragma GCC unroll 4 - for (d = 0; d <= hidden_size - kVecSize; d += kVecSize) { - bVec x_bvec = bVec::loadu(input_ptr + d); - fVec x_fvec0, x_fvec1; - std::tie(x_fvec0, x_fvec1) = at::vec::convert_to_float(x_bvec); - - bVec w_bvec = bVec::loadu(weight + d); - fVec w_fvec0, w_fvec1; - std::tie(w_fvec0, w_fvec1) = at::vec::convert_to_float(w_bvec); - - bVec g_bvec = bVec::loadu(gate_ptr + d); - fVec g_fvec0, g_fvec1; - std::tie(g_fvec0, g_fvec1) = at::vec::convert_to_float(g_bvec); - g_fvec0 = g_fvec0 / (one + g_fvec0.neg().exp_u20()); - g_fvec1 = g_fvec1 / (one + g_fvec1.neg().exp_u20()); - - x_fvec0 = x_fvec0 * scale_fvec * w_fvec0 * g_fvec0; - x_fvec1 = x_fvec1 * scale_fvec * w_fvec1 * g_fvec1; - - bVec out_bvec = convert_from_float_ext(x_fvec0, x_fvec1); - out_bvec.store(out_ptr + d); - } -#pragma GCC unroll 4 - for (; d < hidden_size; ++d) { - float x_val = static_cast(input_ptr[d]); - float w_val = static_cast(weight[d]); - float g_val = static_cast(gate_ptr[d]); - - out_ptr[d] = static_cast(x_val * rsqrt_var * w_val * g_val / (1.f + std::exp(-g_val))); - } - } - }); -} - -} // anonymous namespace - -template -void fused_add_layernorm_kernel_impl( - scalar_t* __restrict__ output, - const scalar_t* __restrict__ input, - scalar_t* __restrict__ residual, - const scalar_t* __restrict__ weight, - const scalar_t* __restrict__ bias, - float* __restrict__ buffer, - int64_t batch_size, - int64_t seq_len, - int64_t hidden_size, - int64_t input_strideN, - float eps = 1e-5) { - using bVec = at::vec::Vectorized; - using fVec = at::vec::Vectorized; - constexpr int kVecSize = bVec::size(); - - const bool has_residual{residual != nullptr}; - const bool has_bias{bias != nullptr}; - const int64_t parallel_size{batch_size * seq_len}; - at::parallel_for(0, parallel_size, 0, [&](int64_t begin, int64_t end) { - float* __restrict__ buffer_ptr = buffer + at::get_thread_num() * hidden_size; - - for (int64_t i = begin; i < end; ++i) { - scalar_t* __restrict__ out_ptr = output + i * hidden_size; - const scalar_t* __restrict__ input_ptr = input + i * input_strideN; - scalar_t* __restrict__ residual_ptr{(scalar_t*)nullptr}; - if (has_residual) { - residual_ptr = residual + i * hidden_size; - } - - // First pass: compute mean and var - fVec sum_fvec{fVec(0.0)}, sum_sq_fvec{fVec(0.0)}; - float sum_val{0.0}, sum_sq_val{0.0}; - int64_t d{0}; - -#pragma GCC unroll 4 - for (; d <= hidden_size - kVecSize; d += kVecSize) { - bVec x_bvec = bVec::loadu(input_ptr + d); - fVec x_fvec0, x_fvec1; - std::tie(x_fvec0, x_fvec1) = at::vec::convert_to_float(x_bvec); - - if (has_residual) { - bVec r_bvec = bVec::loadu(residual_ptr + d); - fVec r_fvec0, r_fvec1; - std::tie(r_fvec0, r_fvec1) = at::vec::convert_to_float(r_bvec); - - x_fvec0 += r_fvec0; - x_fvec1 += r_fvec1; - - bVec out_bvec = convert_from_float_ext(x_fvec0, x_fvec1); - out_bvec.store(residual_ptr + d); - } - + sum2_fvec += x_fvec0 * x_fvec0; + sum2_fvec += x_fvec1 * x_fvec1; + if constexpr (NormTraits::has_mean) { sum_fvec += x_fvec0; sum_fvec += x_fvec1; - sum_sq_fvec += x_fvec0 * x_fvec0; - sum_sq_fvec += x_fvec1 * x_fvec1; - - x_fvec0.store(buffer_ptr + d); - x_fvec1.store(buffer_ptr + d + fVec::size()); } + } #pragma GCC unroll 4 - for (; d < hidden_size; ++d) { - float x_val = static_cast(input_ptr[d]); - if (has_residual) { - float r_val = static_cast(residual_ptr[d]); - x_val += r_val; - residual_ptr[d] = static_cast(x_val); - } - + for (; d < D; ++d) { + float x_val = static_cast(input[d]); + if constexpr (has_residual) { + x_val += static_cast(residual[d]); + } + sum2_val += x_val * x_val; + if constexpr (NormTraits::has_mean) { sum_val += x_val; - sum_sq_val += x_val * x_val; - buffer_ptr[d] = x_val; } + } - // Var(X) = E(X^2) - (E(X))^2 - // Refer to FlashInfer impl: - // https://github.com/flashinfer-ai/flashinfer/blob/6bb01d19c2d9ab3b6a3a5e9e97448891a5ed2844/include/flashinfer/norm.cuh#L554 + float mean = 0.f; + float variance = sum2_val + vec_reduce_sum(sum2_fvec); + variance /= D; + if constexpr (NormTraits::has_mean) { sum_val += vec_reduce_sum(sum_fvec); - sum_sq_val += vec_reduce_sum(sum_sq_fvec); + mean = sum_val / D; + variance -= mean * mean; + } - float mean{sum_val / hidden_size}; - float mean_sq{sum_sq_val / hidden_size}; - float variance{mean_sq - (mean * mean)}; - float rsqrt_var{float(1) / std::sqrt(variance + eps)}; + float rsqrt_var = float(1) / std::sqrt(variance + params.eps); + const fVec mean_fvec = fVec(mean); + const fVec scale_fvec = fVec(rsqrt_var); + const fVec shift_fvec = fVec(params.shift); - const fVec mean_fvec = fVec(mean); - const fVec scale_fvec = fVec(rsqrt_var); - - // Second pass: apply normalization #pragma GCC unroll 4 - for (d = 0; d <= hidden_size - kVecSize; d += kVecSize) { - fVec x_fvec0 = fVec::loadu(buffer_ptr + d); - fVec x_fvec1 = fVec::loadu(buffer_ptr + d + fVec::size()); - bVec w_bvec = bVec::loadu(weight + d); - fVec w_fvec0, w_fvec1; - std::tie(w_fvec0, w_fvec1) = at::vec::convert_to_float(w_bvec); - - x_fvec0 = (x_fvec0 - mean_fvec) * scale_fvec * w_fvec0; - x_fvec1 = (x_fvec1 - mean_fvec) * scale_fvec * w_fvec1; - - if (has_bias) { - bVec b_bvec = bVec::loadu(bias + d); - fVec b_fvec0, b_fvec1; - std::tie(b_fvec0, b_fvec1) = at::vec::convert_to_float(b_bvec); - x_fvec0 += b_fvec0; - x_fvec1 += b_fvec1; - } - - bVec o_bvec = convert_from_float_ext(x_fvec0, x_fvec1); - o_bvec.store(out_ptr + d); + for (d = 0; d <= D - kVecSize; d += kVecSize) { + auto [x_fvec0, x_fvec1] = load_float_vec2(input + d); + if constexpr (has_residual) { + auto [r_fvec0, r_fvec1] = load_float_vec2(residual + d); + x_fvec0 += r_fvec0; + x_fvec1 += r_fvec1; + convert_from_float_ext(x_fvec0, x_fvec1).store(residual + d); } -#pragma GCC unroll 4 - for (; d < hidden_size; ++d) { - float normalized = (buffer_ptr[d] - mean) * rsqrt_var; - float x_val = normalized * static_cast(weight[d]); - if (has_bias) { - x_val += static_cast(bias[d]); + if constexpr (NormTraits::has_mean) { + x_fvec0 = x_fvec0 - mean_fvec; + x_fvec1 = x_fvec1 - mean_fvec; + } + x_fvec0 = x_fvec0 * scale_fvec; + x_fvec1 = x_fvec1 * scale_fvec; + if constexpr (NormTraits::has_weight) { + auto [w_fvec0, w_fvec1] = load_float_vec2(static_cast(params.weight) + d); + if constexpr (NormTraits::has_shift) { + w_fvec0 = NormTraits::apply_shift(w_fvec0, shift_fvec); + w_fvec1 = NormTraits::apply_shift(w_fvec1, shift_fvec); } - out_ptr[d] = static_cast(x_val); + x_fvec0 = NormTraits::apply_weight(x_fvec0, w_fvec0); + x_fvec1 = NormTraits::apply_weight(x_fvec1, w_fvec1); + } + if constexpr (NormTraits::has_bias) { + if (use_bias) { + auto [b_fvec0, b_fvec1] = load_float_vec2(static_cast(params.bias) + d); + x_fvec0 = NormTraits::apply_bias(x_fvec0, b_fvec0); + x_fvec1 = NormTraits::apply_bias(x_fvec1, b_fvec1); + } + } + if constexpr (NormTraits::has_gate) { + auto [g_fvec0, g_fvec1] = load_float_vec2(static_cast(gate) + d); + x_fvec0 = NormTraits::apply_gate(x_fvec0, g_fvec0); + x_fvec1 = NormTraits::apply_gate(x_fvec1, g_fvec1); + } + bVec out_bvec = convert_from_float_ext(x_fvec0, x_fvec1); + out_bvec.store(out + d); + } +#pragma GCC unroll 4 + for (; d < D; ++d) { + float x_val = static_cast(input[d]); + if constexpr (has_residual) { + x_val += static_cast(residual[d]); + residual[d] = static_cast(x_val); + } + if constexpr (NormTraits::has_mean) { + x_val -= mean; + } + x_val *= rsqrt_var; + if constexpr (NormTraits::has_weight) { + float w_val = static_cast(static_cast(params.weight)[d]); + if constexpr (NormTraits::has_shift) { + w_val = NormTraits::apply_shift(w_val, params.shift); + } + x_val = NormTraits::apply_weight(x_val, w_val); + } + if constexpr (NormTraits::has_bias) { + if (use_bias) { + float b_val = static_cast(static_cast(params.bias)[d]); + x_val = NormTraits::apply_bias(x_val, b_val); + } + } + if constexpr (NormTraits::has_gate) { + float g_val = static_cast(static_cast(gate)[d]); + x_val = NormTraits::apply_gate(x_val, g_val); + } + out[d] = static_cast(x_val); + } + } +}; + +// TODO: add generic avx512-bf16 path here + +#define LAUNCH_PARALLEL_LOOP(...) \ + at::parallel_for(0, p.rows(), 0, [&](int64_t begin, int64_t end) { \ + int64_t b{0}, h{0}, t{0}; \ + data_index_init(begin, b, p.B, h, p.H, t, p.T); \ + for (int64_t i = begin; i < end; ++i) { \ + __VA_ARGS__; \ + data_index_step(b, p.B, h, p.H, t, p.T); \ + } \ + }) + +#define LAUNCH_PARALLEL_LOOP_HD(DIM) \ + case DIM: \ + LAUNCH_PARALLEL_LOOP( \ + const scalar_t* __restrict__ gate_ptr{nullptr}; if constexpr (NormTraits::has_gate) { \ + gate_ptr = gate + p.output_offset(b, h, t); \ + } NormReduce:: \ + apply(out + p.output_offset(b, h, t), input + p.input_offset(b, h, t), gate_ptr, p)); \ + return + +template +void norm4d_kernel_impl( + scalar_t* __restrict__ out, + const scalar_t* __restrict__ input, + const NormParams& p, + const scalar_t* __restrict__ gate = nullptr) { + // fast path only applies to bfloat16 when D in {32, 64, 128, 256, 512} + if constexpr (std::is_same_v) { + switch (p.D) { + LAUNCH_PARALLEL_LOOP_HD(32); + LAUNCH_PARALLEL_LOOP_HD(64); + LAUNCH_PARALLEL_LOOP_HD(128); + LAUNCH_PARALLEL_LOOP_HD(256); + LAUNCH_PARALLEL_LOOP_HD(512); + default: + break; + } + } + + // generic path + LAUNCH_PARALLEL_LOOP( + const scalar_t* __restrict__ gate_ptr{nullptr}; if constexpr (NormTraits::has_gate) { + gate_ptr = gate + p.output_offset(b, h, t); + } NormReduceGeneric:: + apply(out + p.output_offset(b, h, t), input + p.input_offset(b, h, t), gate_ptr, nullptr, p, p.D)); +} + +template +void fused_add_norm4d_kernel_impl( + scalar_t* __restrict__ out, + const scalar_t* __restrict__ input, + scalar_t* __restrict__ residual, + const NormParams& p, + bool output_uses_input_stride = false) { + LAUNCH_PARALLEL_LOOP( + const int64_t out_offset = output_uses_input_stride ? p.input_offset(b, h, t) : p.output_offset(b, h, t); + scalar_t* __restrict__ residual_ptr = residual + p.output_offset(b, h, t); + NormReduceGeneric::apply( + out + out_offset, input + p.input_offset(b, h, t), nullptr, residual_ptr, p, p.D)); +} + +template +void fused_qk_norm4d_kernel_impl( + scalar_t* __restrict__ q_out, + scalar_t* __restrict__ k_out, + scalar_t* __restrict__ gate_out, + const scalar_t* __restrict__ q, + const scalar_t* __restrict__ k, + const NormParams& params_q, + const NormParams& params_k) { + at::parallel_for(0, params_q.B, 0, [&](int64_t begin, int64_t end) { + for (int64_t b = begin; b < end; ++b) { + for (int64_t h = 0; h < params_q.H /*num_head*/; ++h) { + const int64_t q_offset = params_q.input_offset(b, h, /*t*/ 0); + const int64_t out_offset = params_q.output_offset(b, h, /*t*/ 0); + NormReduceGeneric::apply( + q_out + out_offset, q + q_offset, nullptr, nullptr, params_q, params_q.D); + if constexpr (copy_gate) { + std::memcpy(gate_out + out_offset, q + q_offset + params_q.D, params_q.D * sizeof(scalar_t)); + } + } + for (int64_t h = 0; h < params_k.H /*num_head_kv*/; ++h) { + NormReduceGeneric::apply( + k_out + params_k.output_offset(b, h, /*t*/ 0), + k + params_k.input_offset(b, h, /*t*/ 0), + nullptr, + nullptr, + params_k, + params_k.D); } } }); +} + +#undef LAUNCH_PARALLEL_LOOP +#undef LAUNCH_PARALLEL_LOOP_HD } // anonymous namespace +template +inline void CHECK_INPUT_ND(const at::Tensor& tensor) { + static_assert(sizeof...(Dims) > 0); + CHECK_LAST_DIM_CONTIGUOUS_INPUT(tensor); + const int64_t dim = tensor.dim(); + const bool dim_ok = ((dim == Dims) || ...); + TORCH_CHECK(dim_ok, "Expected input dim to match template constraints, got ", dim); +} + // input : {batch_size, hidden_size} at::Tensor l2norm_cpu(at::Tensor& input, double eps) { - CHECK_INPUT(input); - CHECK_DIM(2, input); - int64_t batch_size = input.size(0); - int64_t hidden_size = input.size(1); - at::Tensor output = at::empty_like(input); + const auto st = input.scalar_type(); + CHECK_INPUT_ND<2>(input); + NormParams p{input, static_cast(eps)}; - AT_DISPATCH_REDUCED_FLOATING_TYPES(input.scalar_type(), "l2norm_kernel", [&] { - l2norm_kernel_impl( - output.data_ptr(), - input.data_ptr(), - batch_size, - 1, - hidden_size, - hidden_size, - 0, - hidden_size, - 0, - eps); + at::Tensor output = at::empty_like(input); + AT_DISPATCH_REDUCED_FLOATING_TYPES(st, "l2norm_kernel", [&] { + norm4d_kernel_impl(output.data_ptr(), input.data_ptr(), p); }); return output; } @@ -563,43 +494,75 @@ at::Tensor l2norm_cpu(at::Tensor& input, double eps) { // input : {batch_size, hidden_size} or {batch_size, seq_len, hidden_size} // weight: {hidden_size} at::Tensor rmsnorm_cpu(at::Tensor& input, at::Tensor& weight, double eps) { - CHECK_LAST_DIM_CONTIGUOUS_INPUT(input); - CHECK_INPUT(weight); - int64_t inp_dim{input.dim()}; - TORCH_CHECK(inp_dim == 2 || inp_dim == 3, "Expected input dim to be 2 or 3, but got ", inp_dim); - CHECK_DIM(1, weight); - CHECK_EQ(input.size(-1), weight.size(0)); + const auto st = input.scalar_type(); + CHECK_INPUT_ND<2, 3>(input); + CHECK_INPUT_SHAPE_DTYPE(weight, {input.size(-1)}, st); + + NormParams p{input, static_cast(eps)}; + p.weight = weight.data_ptr(); - int64_t batch_size = input.size(0); - int64_t seq_len = 1; - int64_t hidden_size = input.size(-1); - int64_t input_strideB = input.stride(0); - int64_t input_strideS = 0; at::Tensor output = at::empty_like(input); - int64_t output_strideB = output.stride(0); - int64_t output_strideS = 0; - if (inp_dim == 3) { - seq_len = input.size(1); - input_strideS = input.stride(1); - output_strideS = output.stride(1); - } + AT_DISPATCH_REDUCED_FLOATING_TYPES(st, "rmsnorm_kernel", [&] { + norm4d_kernel_impl(output.data_ptr(), input.data_ptr(), p); + }); + return output; +} - AT_DISPATCH_REDUCED_FLOATING_TYPES(input.scalar_type(), "rmsnorm_kernel", [&] { - using Vec = at::vec::Vectorized; - rmsnorm_kernel_impl( - output.data_ptr(), - input.data_ptr(), - weight.data_ptr(), - batch_size, - seq_len, - hidden_size, - input_strideB, - input_strideS, - output_strideB, - output_strideS, - [](float x) { return x; }, - [](Vec x) { return x; }, - eps); +// input : {batch_size, hidden_size} +// weight: {hidden_size} +at::Tensor gemma_rmsnorm_cpu(at::Tensor& input, at::Tensor& weight, double eps) { + CHECK_INPUT_ND<2>(input); + const auto st = input.scalar_type(); + CHECK_INPUT_SHAPE_DTYPE(weight, {input.size(-1)}, st); + + NormParams p{input, static_cast(eps)}; + p.weight = weight.data_ptr(); + p.shift = 1.f; + + at::Tensor output = at::empty_like(input); + AT_DISPATCH_REDUCED_FLOATING_TYPES(st, "gemma_rmsnorm_kernel", [&] { + norm4d_kernel_impl(output.data_ptr(), input.data_ptr(), p); + }); + return output; +} + +// input : {batch_size, hidden_size} or {batch_size, num_head, seq_len, head_dim} +// weight: {hidden_size} +at::Tensor gemma3_rmsnorm_cpu(at::Tensor& input, at::Tensor& weight, double eps) { + const auto st = input.scalar_type(); + CHECK_INPUT_ND<2, 4>(input); + CHECK_INPUT_SHAPE_DTYPE(weight, {input.size(-1)}, st); + + NormParams p{input, static_cast(eps)}; + p.weight = weight.data_ptr(); + p.shift = 1.f; + + at::Tensor output = at::empty_like(input); + AT_DISPATCH_REDUCED_FLOATING_TYPES(st, "gemma3_rmsnorm_kernel", [&] { + norm4d_kernel_impl(output.data_ptr(), input.data_ptr(), p); + }); + return output; +} + +// Gemma4RMSNorm: with_scale ? norm(x) * (weight + scale_shift) : norm(x) +// input : {batch_size, hidden_size} or {batch_size, seq_len, hidden_size} +// weight: {hidden_size} +at::Tensor gemma4_rmsnorm_cpu(at::Tensor& input, at::Tensor& weight, double eps, double scale_shift, bool with_scale) { + const auto st = input.scalar_type(); + CHECK_INPUT_ND<2, 3>(input); + CHECK_INPUT_SHAPE_DTYPE(weight, {input.size(-1)}, st); + + NormParams p{input, static_cast(eps)}; + p.weight = weight.data_ptr(); + p.shift = static_cast(scale_shift); + + at::Tensor output = at::empty_like(input); + AT_DISPATCH_REDUCED_FLOATING_TYPES(st, "gemma4_rmsnorm_kernel", [&] { + if (with_scale) { + norm4d_kernel_impl(output.data_ptr(), input.data_ptr(), p); + } else { + norm4d_kernel_impl(output.data_ptr(), input.data_ptr(), p); + } }); return output; } @@ -609,235 +572,43 @@ at::Tensor rmsnorm_cpu(at::Tensor& input, at::Tensor& weight, double eps) { // bias : {hidden_size} at::Tensor layernorm_cpu(const at::Tensor& input, const at::Tensor& weight, const std::optional& bias, double eps) { - CHECK_LAST_DIM_CONTIGUOUS_INPUT(input); - CHECK_INPUT(weight); - int64_t inp_dim{input.dim()}; - TORCH_CHECK(inp_dim == 2 || inp_dim == 3, "Expected input dim to be 2 or 3, but got ", inp_dim); - CHECK_DIM(1, weight); + const auto st = input.scalar_type(); + const int64_t hidden_size = input.size(-1); + CHECK_INPUT_ND<2, 3>(input); + CHECK_INPUT_SHAPE_DTYPE(weight, {hidden_size}, st); if (bias.has_value()) { - CHECK_DIM(1, bias.value()); - CHECK_EQ(bias.value().size(0), weight.size(0)); + CHECK_INPUT_SHAPE_DTYPE(bias.value(), {hidden_size}, st); } - int64_t batch_size{input.size(0)}, seq_len{1}, hidden_size{input.size(1)}, input_strideN{input.stride(0)}; - if (inp_dim == 3) { - CHECK_EQ(input.size(2), weight.size(0)); - seq_len = input.size(1); - hidden_size = input.size(2); - input_strideN = input.stride(1); - } else { - CHECK_EQ(input.size(1), weight.size(0)); - } + NormParams p{input, static_cast(eps)}; + p.weight = weight.data_ptr(); + p.bias = bias.has_value() ? bias.value().data_ptr() : nullptr; at::Tensor output = at::empty_like(input); - int64_t num_threads = at::get_num_threads(); - at::Tensor buffer = at::empty({num_threads, hidden_size}, input.options().dtype(at::kFloat)); - - AT_DISPATCH_REDUCED_FLOATING_TYPES(input.scalar_type(), "layernorm_kernel", [&] { - fused_add_layernorm_kernel_impl( - output.data_ptr(), - input.data_ptr(), - nullptr, - weight.data_ptr(), - conditional_data_ptr(bias), - buffer.data_ptr(), - batch_size, - seq_len, - hidden_size, - input_strideN, - eps); + AT_DISPATCH_REDUCED_FLOATING_TYPES(st, "layernorm_kernel", [&] { + norm4d_kernel_impl(output.data_ptr(), input.data_ptr(), p); }); return output; } -at::Tensor gemma_rmsnorm_cpu(at::Tensor& input, at::Tensor& weight, double eps) { - CHECK_LAST_DIM_CONTIGUOUS_INPUT(input); - CHECK_INPUT(weight); - CHECK_DIM(2, input); - CHECK_DIM(1, weight); - CHECK_EQ(input.size(1), weight.size(0)); - int64_t batch_size = input.size(0); - int64_t hidden_size = input.size(1); - at::Tensor output = at::empty_like(input); - int64_t input_strideN = input.stride(0); - int64_t output_strideN = output.stride(0); - - AT_DISPATCH_REDUCED_FLOATING_TYPES(input.scalar_type(), "gemma_rmsnorm_kernel", [&] { - using Vec = at::vec::Vectorized; - Vec one_vec = Vec(float(1)); - rmsnorm_kernel_impl( - output.data_ptr(), - input.data_ptr(), - weight.data_ptr(), - batch_size, - 1, - hidden_size, - input_strideN, - 0, - output_strideN, - 0, - [](float x) { return x + 1; }, - [one_vec](Vec x) { return x + one_vec; }, - eps); - }); - return output; -} - -// input : {batch_size, hidden_size} or {batch_size, num_head, seq_len, head_dim} -// weight: {hidden_size} -at::Tensor gemma3_rmsnorm_cpu(at::Tensor& input, at::Tensor& weight, double eps) { - CHECK_LAST_DIM_CONTIGUOUS_INPUT(input); - CHECK_INPUT(weight); - TORCH_CHECK( - input.dim() == 2 || input.dim() == 4, "gemma3_rmsnorm_cpu: input must be 2D or 4D, got ", input.dim(), "D"); - CHECK_DIM(1, weight); - CHECK_EQ(input.size(-1), weight.size(0)); - int64_t batch_size = input.size(0); - int64_t hidden_size = weight.size(0); - at::Tensor output = at::empty_like(input); - if (input.dim() == 2) { - int64_t input_strideN = input.stride(0); - int64_t output_strideN = output.stride(0); - - AT_DISPATCH_REDUCED_FLOATING_TYPES(input.scalar_type(), "gemma3_rmsnorm_kernel", [&] { - using Vec = at::vec::Vectorized; - Vec one_vec = Vec(float(1)); - rmsnorm_kernel_impl( - output.data_ptr(), - input.data_ptr(), - weight.data_ptr(), - batch_size, - 1, - hidden_size, - input_strideN, - 0, - output_strideN, - 0, - [](float x) { return x + 1; }, - [one_vec](Vec x) { return x + one_vec; }, - eps); - }); - } else { - int64_t input_strideB = input.stride(0); - int64_t input_strideH = input.stride(1); - int64_t input_strideS = input.stride(2); - int64_t output_strideB = output.stride(0); - int64_t output_strideH = output.stride(1); - int64_t output_strideS = output.stride(2); - int64_t num_head = input.size(1); - int64_t seq_len = input.size(2); - AT_DISPATCH_REDUCED_FLOATING_TYPES(input.scalar_type(), "gemma3_rmsnorm_kernel", [&] { - gemma3_rmsnorm_kernel_4d_impl( - output.data_ptr(), - input.data_ptr(), - weight.data_ptr(), - batch_size, - num_head, - seq_len, - hidden_size, - input_strideB, - input_strideH, - input_strideS, - output_strideB, - output_strideH, - output_strideS, - eps); - }); - } - return output; -} - -// Gemma4RMSNorm: with_scale ? norm(x) * (weight + scale_shift) : norm(x) -// input : {batch_size, hidden_size} or {batch_size, seq_len, hidden_size} -// weight: {hidden_size} -at::Tensor gemma4_rmsnorm_cpu(at::Tensor& input, at::Tensor& weight, double eps, double scale_shift, bool with_scale) { - CHECK_LAST_DIM_CONTIGUOUS_INPUT(input); - CHECK_INPUT(weight); - int64_t inp_dim{input.dim()}; - TORCH_CHECK(inp_dim == 2 || inp_dim == 3, "gemma4_rmsnorm_cpu: expected input dim 2 or 3, got ", inp_dim); - CHECK_DIM(1, weight); - CHECK_EQ(input.size(-1), weight.size(0)); - - int64_t hidden_size = input.size(-1); - at::Tensor output = at::empty_like(input); - int64_t batch_size = input.size(0); - int64_t seq_len = 1; - int64_t input_strideB = input.stride(0); - int64_t input_strideS = 0; - int64_t output_strideB = output.stride(0); - int64_t output_strideS = 0; - if (inp_dim == 3) { - seq_len = input.size(1); - input_strideS = input.stride(1); - output_strideS = output.stride(1); - } - - if (with_scale) { - float shift = static_cast(scale_shift); - AT_DISPATCH_REDUCED_FLOATING_TYPES(input.scalar_type(), "gemma4_rmsnorm_kernel", [&] { - using Vec = at::vec::Vectorized; - Vec shift_vec = Vec(shift); - rmsnorm_kernel_impl( - output.data_ptr(), - input.data_ptr(), - weight.data_ptr(), - batch_size, - seq_len, - hidden_size, - input_strideB, - input_strideS, - output_strideB, - output_strideS, - [shift](float x) { return x + shift; }, - [shift_vec](Vec x) { return x + shift_vec; }, - eps); - }); - } else { - AT_DISPATCH_REDUCED_FLOATING_TYPES(input.scalar_type(), "gemma4_rmsnorm_kernel", [&] { - l2norm_kernel_impl( - output.data_ptr(), - input.data_ptr(), - batch_size, - seq_len, - hidden_size, - input_strideB, - input_strideS, - output_strideB, - output_strideS, - eps); - }); - } - return output; -} - // input : {batch_size, hidden_size} // weight: {hidden_size} // gate: {batch_size, hidden_size} at::Tensor fused_rmsnorm_gated_cpu(at::Tensor& input, at::Tensor& weight, at::Tensor& gate, double eps) { - CHECK_LAST_DIM_CONTIGUOUS_INPUT(input); - CHECK_INPUT(weight); - CHECK_INPUT(gate); - CHECK_DIM(2, input); - CHECK_DIM(1, weight); - CHECK_DIM(2, gate); - CHECK_EQ(input.size(1), weight.size(0)); - int64_t batch_size = input.size(0); - int64_t hidden_size = input.size(1); - CHECK_EQ(input.size(0), gate.size(0)); - CHECK_EQ(input.size(1), gate.size(1)); - at::Tensor output = at::empty_like(input); - int64_t input_strideN = input.stride(0); + const auto st = input.scalar_type(); + const int64_t batch_size = input.size(0); + const int64_t hidden_size = input.size(-1); + CHECK_INPUT_ND<2>(input); + CHECK_INPUT_SHAPE_DTYPE(weight, {hidden_size}, st); + CHECK_INPUT_SHAPE_DTYPE(gate, {batch_size, hidden_size}, st); - AT_DISPATCH_REDUCED_FLOATING_TYPES(input.scalar_type(), "fused_rmsnorm_gated_kernel", [&] { - fused_rmsnorm_gated_kernel_impl( - output.data_ptr(), - input.data_ptr(), - weight.data_ptr(), - gate.data_ptr(), - batch_size, - hidden_size, - input_strideN, - eps); + NormParams p{input, static_cast(eps)}; + p.weight = weight.data_ptr(); + + at::Tensor output = at::empty_like(input); + AT_DISPATCH_REDUCED_FLOATING_TYPES(st, "fused_rmsnorm_gated_kernel", [&] { + norm4d_kernel_impl( + output.data_ptr(), input.data_ptr(), p, gate.data_ptr()); }); return output; } @@ -846,47 +617,22 @@ at::Tensor fused_rmsnorm_gated_cpu(at::Tensor& input, at::Tensor& weight, at::Te // residual: {batch_size, hidden_size} or {batch_size, seq_len, hidden_size} // weight : {hidden_size} void fused_add_rmsnorm_cpu(at::Tensor& input, at::Tensor& residual, at::Tensor& weight, double eps) { - CHECK_LAST_DIM_CONTIGUOUS_INPUT(input); - CHECK_INPUT(residual); - CHECK_INPUT(weight); - int64_t inp_dim{input.dim()}, res_dim{residual.dim()}; - CHECK_EQ(inp_dim, res_dim); - TORCH_CHECK(inp_dim == 2 || inp_dim == 3, "Expected input dim to be 2 or 3, but got ", inp_dim); - CHECK_DIM(1, weight); - CHECK_EQ(input.size(0), residual.size(0)); - CHECK_EQ(input.size(-1), residual.size(-1)); - CHECK_EQ(input.size(-1), weight.size(0)); + const auto st = input.scalar_type(); + CHECK_INPUT_ND<2, 3>(input); + CHECK_EQ(input.sizes(), residual.sizes()); + CHECK_EQ(st, residual.scalar_type()); + CHECK_INPUT_SHAPE_DTYPE(weight, {input.size(-1)}, st); - int64_t batch_size = input.size(0); - int64_t seq_len = 1; - int64_t hidden_size = input.size(-1); - int64_t input_strideB = input.stride(0); - int64_t input_strideS = 0; - if (inp_dim == 3) { - seq_len = input.size(1); - input_strideS = input.stride(1); - } + NormParams p{input, static_cast(eps)}; + p.weight = weight.data_ptr(); - // allocate temp buffer to store x in float32 per thread - // TODO: implement a singleton for context - int64_t num_threads = at::get_num_threads(); - at::Tensor buffer = at::empty({num_threads, hidden_size}, input.options().dtype(at::kFloat)); - - AT_DISPATCH_REDUCED_FLOATING_TYPES(input.scalar_type(), "fused_add_rmsnorm_kernel", [&] { - using Vec = at::vec::Vectorized; - fused_add_rmsnorm_kernel_impl( + AT_DISPATCH_REDUCED_FLOATING_TYPES(st, "fused_add_rmsnorm_kernel", [&] { + fused_add_norm4d_kernel_impl( + input.data_ptr(), input.data_ptr(), residual.data_ptr(), - weight.data_ptr(), - buffer.data_ptr(), - batch_size, - seq_len, - hidden_size, - input_strideB, - input_strideS, - [](float x) { return x; }, - [](Vec x) { return x; }, - eps); + p, + /*output_uses_input_stride=*/true); }); } @@ -894,40 +640,23 @@ void fused_add_rmsnorm_cpu(at::Tensor& input, at::Tensor& residual, at::Tensor& // residual: {batch_size, hidden_size} // weight : {hidden_size} void gemma_fused_add_rmsnorm_cpu(at::Tensor& input, at::Tensor& residual, at::Tensor& weight, double eps) { - CHECK_LAST_DIM_CONTIGUOUS_INPUT(input); - CHECK_INPUT(residual); - CHECK_INPUT(weight); - CHECK_DIM(2, input); - CHECK_DIM(2, residual); - CHECK_DIM(1, weight); - CHECK_EQ(input.size(0), residual.size(0)); - CHECK_EQ(input.size(1), residual.size(1)); - CHECK_EQ(input.size(1), weight.size(0)); - int64_t batch_size = input.size(0); - int64_t hidden_size = input.size(1); - int64_t input_strideN = input.stride(0); + const auto st = input.scalar_type(); + CHECK_INPUT_ND<2>(input); + CHECK_EQ(input.sizes(), residual.sizes()); + CHECK_EQ(st, residual.scalar_type()); + CHECK_INPUT_SHAPE_DTYPE(weight, {input.size(-1)}, st); - // allocate temp buffer to store x in float32 per thread - // TODO: implement a singleton for context - int64_t num_threads = at::get_num_threads(); - at::Tensor buffer = at::empty({num_threads, hidden_size}, input.options().dtype(at::kFloat)); + NormParams p{input, static_cast(eps)}; + p.weight = weight.data_ptr(); + p.shift = 1.f; - AT_DISPATCH_REDUCED_FLOATING_TYPES(input.scalar_type(), "gemma_fused_add_rmsnorm_kernel", [&] { - using Vec = at::vec::Vectorized; - Vec one_vec = Vec(float(1)); - fused_add_rmsnorm_kernel_impl( + AT_DISPATCH_REDUCED_FLOATING_TYPES(st, "gemma_fused_add_rmsnorm_kernel", [&] { + fused_add_norm4d_kernel_impl( + input.data_ptr(), input.data_ptr(), residual.data_ptr(), - weight.data_ptr(), - buffer.data_ptr(), - batch_size, - 1, - hidden_size, - input_strideN, - 0, - [](float x) { return x + 1; }, - [one_vec](Vec x) { return x + one_vec; }, - eps); + p, + /*output_uses_input_stride=*/true); }); } @@ -941,54 +670,126 @@ at::Tensor fused_add_layernorm_cpu( const at::Tensor& weight, const std::optional& bias, double eps) { - CHECK_LAST_DIM_CONTIGUOUS_INPUT(input); - CHECK_INPUT(residual); - CHECK_INPUT(weight); - int64_t inp_dim{input.dim()}, res_dim{residual.dim()}; - CHECK_EQ(inp_dim, res_dim); - TORCH_CHECK(inp_dim == 2 || inp_dim == 3, "Expected input dim to be 2 or 3, but got ", inp_dim); - TORCH_CHECK(res_dim == 2 || res_dim == 3, "Expected residual dim to be 2 or 3, but got ", res_dim); - - CHECK_DIM(1, weight); + const auto st = input.scalar_type(); + const int64_t hidden_size = input.size(-1); + CHECK_INPUT_ND<2, 3>(input); + CHECK_EQ(input.sizes(), residual.sizes()); + CHECK_EQ(st, residual.scalar_type()); + CHECK_INPUT_SHAPE_DTYPE(weight, {hidden_size}, st); if (bias.has_value()) { - CHECK_DIM(1, bias.value()); - CHECK_EQ(bias.value().size(0), weight.size(0)); - } - CHECK_EQ(input.size(0), residual.size(0)); - CHECK_EQ(input.size(1), residual.size(1)); - if (inp_dim == 3) { - CHECK_EQ(input.size(2), residual.size(2)); - CHECK_EQ(input.size(2), weight.size(0)); - } else { - CHECK_EQ(input.size(1), weight.size(0)); + CHECK_INPUT_SHAPE_DTYPE(bias.value(), {hidden_size}, st); } - int64_t batch_size{input.size(0)}, seq_len{1}, hidden_size{input.size(1)}, input_strideN{input.stride(0)}; - if (inp_dim == 3) { - seq_len = input.size(1); - hidden_size = input.size(2); - input_strideN = input.stride(1); - } + NormParams p{input, static_cast(eps)}; + p.weight = weight.data_ptr(); + p.bias = bias.has_value() ? bias.value().data_ptr() : nullptr; + at::Tensor output = at::empty_like(input); - - // Allocate temp buffer to store x in float32 per thread - // It is necessary to store FP32 precision of residual-add results to pass UT acc test - int64_t num_threads = at::get_num_threads(); - at::Tensor buffer = at::empty({num_threads, hidden_size}, input.options().dtype(at::kFloat)); - - AT_DISPATCH_REDUCED_FLOATING_TYPES(input.scalar_type(), "fused_add_layernorm_kernel", [&] { - fused_add_layernorm_kernel_impl( - output.data_ptr(), - input.data_ptr(), - residual.data_ptr(), - weight.data_ptr(), - conditional_data_ptr(bias), - buffer.data_ptr(), - batch_size, - seq_len, - hidden_size, - input_strideN, - eps); + AT_DISPATCH_REDUCED_FLOATING_TYPES(st, "fused_add_layernorm_kernel", [&] { + fused_add_norm4d_kernel_impl( + output.data_ptr(), input.data_ptr(), residual.data_ptr(), p); }); return output; } + +// q : {batch_size, num_head * head_dim} 2D +// k : {batch_size, num_head_kv * head_dim} 2D +std::tuple fused_qk_gemma_rmsnorm_cpu( + const at::Tensor& q, + const at::Tensor& k, + const at::Tensor& q_weight, + const at::Tensor& k_weight, + double eps, + int64_t head_dim) { + const auto st = q.scalar_type(); + CHECK_INPUT_ND<2>(q); + CHECK_INPUT_ND<2>(k); + + int64_t batch_size = q.size(0); + int64_t num_head = q.size(1) / head_dim; + int64_t num_head_kv = k.size(1) / head_dim; + CHECK_EQ(k.size(0), batch_size); + CHECK_EQ(k.scalar_type(), st); + CHECK_INPUT_SHAPE_DTYPE(q_weight, {head_dim}, st); + CHECK_INPUT_SHAPE_DTYPE(k_weight, {head_dim}, st); + + NormParams q_params{q, static_cast(eps)}; + q_params.H = num_head; + q_params.D = head_dim; + q_params.i_strideH = head_dim; + q_params.weight = q_weight.data_ptr(); + q_params.shift = 1.f; + + NormParams k_params{k, static_cast(eps)}; + k_params.H = num_head_kv; + k_params.D = head_dim; + k_params.i_strideH = head_dim; + k_params.weight = k_weight.data_ptr(); + k_params.shift = 1.f; + + at::Tensor q_out = at::empty_like(q); + at::Tensor k_out = at::empty_like(k); + AT_DISPATCH_REDUCED_FLOATING_TYPES(st, "fused_qk_gemma_rmsnorm_kernel", [&] { + fused_qk_norm4d_kernel_impl( + q_out.data_ptr(), + k_out.data_ptr(), + nullptr, + q.data_ptr(), + k.data_ptr(), + q_params, + k_params); + }); + return std::make_tuple(q_out, k_out); +} + +// q_gate : {batch_size, num_head * head_dim * 2} 2D, interleaved per head as [q_h, gate_h] +// k : {batch_size, num_head_kv * head_dim} 2D +std::tuple fused_qk_gemma_rmsnorm_with_gate_cpu( + const at::Tensor& q_gate, + const at::Tensor& k, + const at::Tensor& q_weight, + const at::Tensor& k_weight, + double eps, + int64_t head_dim, + int64_t num_head) { + const auto st = q_gate.scalar_type(); + CHECK_INPUT_ND<2>(q_gate); + CHECK_INPUT_ND<2>(k); + + int64_t batch_size = q_gate.size(0); + int64_t num_head_kv = k.size(1) / head_dim; + CHECK_EQ(q_gate.size(1), num_head * head_dim * 2); + CHECK_EQ(k.size(0), batch_size); + CHECK_EQ(k.scalar_type(), st); + CHECK_INPUT_SHAPE_DTYPE(q_weight, {head_dim}, st); + CHECK_INPUT_SHAPE_DTYPE(k_weight, {head_dim}, st); + + NormParams q_params{q_gate, static_cast(eps)}; + q_params.H = num_head; + q_params.D = head_dim; + q_params.i_strideH = head_dim * 2; + q_params.weight = q_weight.data_ptr(); + q_params.shift = 1.f; + + NormParams k_params{k, static_cast(eps)}; + k_params.H = num_head_kv; + k_params.D = head_dim; + k_params.i_strideH = head_dim; + k_params.weight = k_weight.data_ptr(); + k_params.shift = 1.f; + + at::Tensor q_out = at::empty({batch_size * num_head, head_dim}, q_gate.options()); + at::Tensor k_out = at::empty({batch_size * num_head_kv, head_dim}, k.options()); + at::Tensor gate_out = at::empty_like(q_out); + AT_DISPATCH_REDUCED_FLOATING_TYPES(st, "fused_qk_gemma_rmsnorm_with_gate_kernel", [&] { + fused_qk_norm4d_kernel_impl( + q_out.data_ptr(), + k_out.data_ptr(), + gate_out.data_ptr(), + q_gate.data_ptr(), + k.data_ptr(), + q_params, + k_params); + }); + return std::make_tuple(q_out, k_out, gate_out); +} diff --git a/sgl-kernel/csrc/cpu/torch_extension_cpu.cpp b/sgl-kernel/csrc/cpu/torch_extension_cpu.cpp index 1a889b225..6aedaf536 100644 --- a/sgl-kernel/csrc/cpu/torch_extension_cpu.cpp +++ b/sgl-kernel/csrc/cpu/torch_extension_cpu.cpp @@ -58,6 +58,23 @@ at::Tensor fused_add_layernorm_cpu( const std::optional& bias, double eps); +// fused_qk_gemma_rmsnorm +std::tuple fused_qk_gemma_rmsnorm_cpu( + const at::Tensor& q, + const at::Tensor& k, + const at::Tensor& q_weight, + const at::Tensor& k_weight, + double eps, + int64_t head_dim); +std::tuple fused_qk_gemma_rmsnorm_with_gate_cpu( + const at::Tensor& q_gate, + const at::Tensor& k, + const at::Tensor& q_weight, + const at::Tensor& k_weight, + double eps, + int64_t head_dim, + int64_t num_head); + // topk std::tuple topk_sigmoid_cpu(at::Tensor& hidden_states, at::Tensor& gating_output, int64_t topk, bool renormalize); @@ -468,6 +485,15 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) { "fused_add_layernorm_cpu(Tensor input, Tensor residual, Tensor weight, Tensor? bias, float eps) -> " "Tensor"); m.impl("fused_add_layernorm_cpu", torch::kCPU, &fused_add_layernorm_cpu); + m.def( + "fused_qk_gemma_rmsnorm_cpu(Tensor q, Tensor k, Tensor q_weight, Tensor k_weight, float eps, int head_dim) -> " + "(Tensor, Tensor)"); + m.impl("fused_qk_gemma_rmsnorm_cpu", torch::kCPU, &fused_qk_gemma_rmsnorm_cpu); + m.def( + "fused_qk_gemma_rmsnorm_with_gate_cpu(Tensor q_gate, Tensor k, Tensor q_weight, Tensor k_weight, float eps, int " + "head_dim, int num_head) -> " + "(Tensor, Tensor, Tensor)"); + m.impl("fused_qk_gemma_rmsnorm_with_gate_cpu", torch::kCPU, &fused_qk_gemma_rmsnorm_with_gate_cpu); // topk m.def("topk_sigmoid_cpu(Tensor hidden_states, Tensor gating_output, int topk, bool renormalize) -> (Tensor, Tensor)"); diff --git a/sgl-kernel/csrc/cpu/vec.h b/sgl-kernel/csrc/cpu/vec.h index a235cb4a7..22c02a9ae 100644 --- a/sgl-kernel/csrc/cpu/vec.h +++ b/sgl-kernel/csrc/cpu/vec.h @@ -451,6 +451,62 @@ inline std::tuple<__m512i, __m512i> transpose_2x32_16bit(__m512i r0, __m512i r1) } #pragma GCC diagnostic pop +// Note: mapped from aten exp_u20 +inline __attribute__((always_inline)) __m512 _mm512_exp_u20_ps(const __m512 values) { + const __m512 vec_factorial_1 = _mm512_set1_ps(0.999999701f); + const __m512 vec_factorial_2 = _mm512_set1_ps(0.499991506f); + const __m512 vec_factorial_3 = _mm512_set1_ps(0.166676521f); + const __m512 vec_factorial_4 = _mm512_set1_ps(0.0418978221f); + const __m512 vec_factorial_5 = _mm512_set1_ps(0.00828929059f); + const __m512 vec_exp_log2ef = _mm512_castsi512_ps(_mm512_set1_epi32(0x3fb8aa3b)); // log2(e) + const __m512 vec_half = _mm512_set1_ps(0.5f); + const __m512 vec_one = _mm512_set1_ps(1.f); + const __m512 vec_zero = _mm512_set1_ps(0.f); + const __m512 vec_two = _mm512_set1_ps(2.f); + const __m512 vec_ln2f = _mm512_castsi512_ps(_mm512_set1_epi32(0x3f317218)); + const __m512 vec_ln_flt_min = _mm512_castsi512_ps(_mm512_set1_epi32(0xc2aeac50)); + const __m512 vec_ln_flt_max = _mm512_castsi512_ps(_mm512_set1_epi32(0x42b17218)); + const __m512i vec_127 = _mm512_set1_epi32(0x0000007f); + const int n_mantissa_bits = 23; + + // exp(x) = + // = exp(n * ln(2) + r) // divide x by ln(2) and get quot and rem + // = 2^n * exp(r) // simplify the exp(n*ln(2)) expression + + auto less_ln_flt_min_mask = _mm512_cmp_ps_mask(values, vec_ln_flt_min, 1 /*_CMP_LT_OS*/); + auto vec_src = _mm512_min_ps(values, vec_ln_flt_max); + vec_src = _mm512_max_ps(vec_src, vec_ln_flt_min); + + // fx = floorf(x * log2ef + 0.5) + auto vec_fx = _mm512_fmadd_ps(vec_src, vec_exp_log2ef, vec_half); + auto vec_fx_i = _mm512_cvt_roundps_epi32(vec_fx, _MM_FROUND_TO_NEG_INF | _MM_FROUND_NO_EXC); + vec_fx = _mm512_cvtepi32_ps(vec_fx_i); + + // x = x - fx * ln2 + auto vec_exp_poly = _mm512_fnmadd_ps(vec_fx, vec_ln2f, vec_src); + + // compute polynomial + auto vec_res = _mm512_fmadd_ps(vec_exp_poly, vec_factorial_5, vec_factorial_4); + vec_res = _mm512_fmadd_ps(vec_exp_poly, vec_res, vec_factorial_3); + vec_res = _mm512_fmadd_ps(vec_exp_poly, vec_res, vec_factorial_2); + vec_res = _mm512_fmadd_ps(vec_exp_poly, vec_res, vec_factorial_1); + vec_res = _mm512_fmadd_ps(vec_exp_poly, vec_res, vec_one); + + // compute 2^(n-1) + auto vec_exp_number = _mm512_sub_ps(vec_fx, vec_one); + auto vec_exp_number_i = _mm512_cvtps_epi32(vec_exp_number); + auto vec_two_pow_n_i = _mm512_add_epi32(vec_exp_number_i, vec_127); + vec_two_pow_n_i = _mm512_slli_epi32(vec_two_pow_n_i, n_mantissa_bits); + auto vec_two_pow_n = _mm512_castsi512_ps(vec_two_pow_n_i); + vec_two_pow_n = _mm512_mask_blend_ps(less_ln_flt_min_mask, vec_two_pow_n, vec_zero); + + // y = y * 2^n + vec_res = _mm512_mul_ps(vec_res, vec_two_pow_n); + vec_res = _mm512_mul_ps(vec_res, vec_two); + return vec_res; +} + +// Note: mapped from aten fexp_u20 inline __attribute__((always_inline)) __m512 _mm512_fexp_u20_ps(const __m512 values) { const __m512 vec_c0 = _mm512_set1_ps(0.00010703434948458272f); const __m512 vec_c1 = _mm512_set1_ps(0.30354260500649682f); diff --git a/test/registered/cpu/test_norm.py b/test/registered/cpu/test_norm.py index cdcc2b369..0061908a6 100644 --- a/test/registered/cpu/test_norm.py +++ b/test/registered/cpu/test_norm.py @@ -1,26 +1,29 @@ -import itertools -import unittest +import sys from typing import Optional, Tuple, Union +import pytest import torch -from utils import make_non_contiguous, parametrize, precision +from utils import make_non_contiguous, precision from sglang.test.ci.ci_register import register_cpu_ci -from sglang.test.test_utils import CustomTestCase register_cpu_ci(est_time=10, suite="base-b-test-cpu") register_cpu_ci(est_time=10, suite="base-b-test-cpu-arm64") torch.manual_seed(1234) +DTYPES = [torch.float16, torch.bfloat16] +DTYPE_IDS = ["float16", "bfloat16"] +eps = 1e-6 -class TestNorm(CustomTestCase): + +class TestNorm: def _forward_native( self, x: torch.Tensor, weight: torch.Tensor, - variance_epsilon: float = 1e-6, + variance_epsilon: float = eps, residual: Optional[torch.Tensor] = None, ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: orig_dtype = x.dtype @@ -41,7 +44,7 @@ class TestNorm(CustomTestCase): return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + eps) def _gemma3_rmsnorm_native( - self, x: torch.Tensor, weight: torch.Tensor, variance_epsilon: float = 1e-6 + self, x: torch.Tensor, weight: torch.Tensor, variance_epsilon: float = eps ): output = self._norm(x.float(), variance_epsilon) output = output * (1.0 + weight.float()) @@ -51,7 +54,7 @@ class TestNorm(CustomTestCase): self, x: torch.Tensor, weight: torch.Tensor, - variance_epsilon: float = 1e-6, + variance_epsilon: float = eps, residual: Optional[torch.Tensor] = None, ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: orig_dtype = x.dtype @@ -66,144 +69,91 @@ class TestNorm(CustomTestCase): x = x.to(orig_dtype) return x if residual is None else (x, residual) - @parametrize( - m=[4096, 1024], - n=[4096, 4109], - dtype=[torch.float16, torch.bfloat16], - ) - def test_norm(self, m, n, dtype): + @pytest.mark.parametrize("dtype", DTYPES, ids=DTYPE_IDS) + @pytest.mark.parametrize("hidden_size", [2048, 512]) + @pytest.mark.parametrize("batch_size", [32, 121]) + def test_l2norm(self, batch_size, hidden_size, dtype): - x = torch.randn([m, n], dtype=dtype) - x = make_non_contiguous(x) - hidden_size = x.size(-1) - weight = torch.randn(hidden_size, dtype=dtype) - variance_epsilon = 1e-6 - - out = torch.ops.sgl_kernel.rmsnorm_cpu(x, weight, variance_epsilon) - ref_out = self._forward_native(x, weight, variance_epsilon) - - atol = rtol = precision[ref_out.dtype] - torch.testing.assert_close(ref_out, out, atol=atol, rtol=rtol) - - ref_x = x.clone() - residual = torch.randn([m, hidden_size], dtype=dtype) - ref_residual = residual.clone() - - torch.ops.sgl_kernel.fused_add_rmsnorm_cpu( - x, residual, weight, variance_epsilon - ) - - ref_x, ref_residual = self._forward_native( - ref_x, weight, variance_epsilon, ref_residual - ) - - torch.testing.assert_close(x, ref_x, atol=atol, rtol=rtol) - torch.testing.assert_close(residual, ref_residual, atol=atol, rtol=rtol) - - @parametrize( - l=[1, 2], - m=[4096, 1024], - n=[4096, 4109], - dtype=[torch.float16, torch.bfloat16], - ) - def test_norm_3d(self, l, m, n, dtype): - - x = torch.randn([l, m, n], dtype=dtype) - x = make_non_contiguous(x) - hidden_size = x.size(-1) - weight = torch.randn(hidden_size, dtype=dtype) - variance_epsilon = 1e-6 - - out = torch.ops.sgl_kernel.rmsnorm_cpu(x, weight, variance_epsilon) - ref_out = self._forward_native(x, weight, variance_epsilon) - - atol = rtol = precision[ref_out.dtype] - torch.testing.assert_close(ref_out, out, atol=atol, rtol=rtol) - - ref_x = x.clone() - residual = torch.randn([l, m, hidden_size], dtype=dtype) - ref_residual = residual.clone() - - torch.ops.sgl_kernel.fused_add_rmsnorm_cpu( - x, residual, weight, variance_epsilon - ) - - ref_x, ref_residual = self._forward_native( - ref_x, weight, variance_epsilon, ref_residual - ) - - torch.testing.assert_close(x, ref_x, atol=atol, rtol=rtol) - torch.testing.assert_close(residual, ref_residual, atol=atol, rtol=rtol) - - @parametrize( - m=[4096, 1024], - n=[4096, 4109], - dtype=[torch.float16, torch.bfloat16], - ) - def test_l2norm(self, m, n, dtype): - - x = torch.randn([m, n], dtype=dtype) - hidden_size = x.size(-1) + x = torch.randn([batch_size, hidden_size], dtype=dtype) fake_ones_weight = torch.ones(hidden_size, dtype=dtype) - variance_epsilon = 1e-6 - out = torch.ops.sgl_kernel.l2norm_cpu(x, variance_epsilon) - ref_out = self._forward_native(x, fake_ones_weight, variance_epsilon) + out = torch.ops.sgl_kernel.l2norm_cpu(x, eps) + ref_out = self._forward_native(x, fake_ones_weight, eps) atol = rtol = precision[ref_out.dtype] torch.testing.assert_close(ref_out, out, atol=atol, rtol=rtol) - @parametrize( - m=[4096, 1024], - n=[4096, 4109], - dtype=[torch.float16, torch.bfloat16], - ) - def test_gemma_rmsnorm(self, m, n, dtype): + @pytest.mark.parametrize("dtype", DTYPES, ids=DTYPE_IDS) + @pytest.mark.parametrize("hidden_size", [2048, 512]) + @pytest.mark.parametrize("batch_size", [32, 121]) + @pytest.mark.parametrize("seq_len", [None, 2], ids=["2d", "3d"]) + def test_rmsnorm(self, seq_len, batch_size, hidden_size, dtype): - x = torch.randn([m, n], dtype=dtype) + if seq_len is None: + x = torch.randn([batch_size, hidden_size], dtype=dtype) + else: + x = torch.randn([batch_size, seq_len, hidden_size], dtype=dtype) x = make_non_contiguous(x) - hidden_size = x.size(-1) + residual = torch.randn(x.shape, dtype=dtype) weight = torch.randn(hidden_size, dtype=dtype) - variance_epsilon = 1e-6 - out = torch.ops.sgl_kernel.gemma_rmsnorm_cpu(x, weight, variance_epsilon) - ref_out = self._gemma_rmsnorm_native(x, weight, variance_epsilon) + out = torch.ops.sgl_kernel.rmsnorm_cpu(x, weight, eps) + ref_out = self._forward_native(x, weight, eps) atol = rtol = precision[ref_out.dtype] torch.testing.assert_close(ref_out, out, atol=atol, rtol=rtol) ref_x = x.clone() - residual = torch.randn([m, hidden_size], dtype=dtype) ref_residual = residual.clone() - torch.ops.sgl_kernel.gemma_fused_add_rmsnorm_cpu( - x, residual, weight, variance_epsilon - ) + torch.ops.sgl_kernel.fused_add_rmsnorm_cpu(x, residual, weight, eps) + + ref_x, ref_residual = self._forward_native(ref_x, weight, eps, ref_residual) + + torch.testing.assert_close(x, ref_x, atol=atol, rtol=rtol) + torch.testing.assert_close(residual, ref_residual, atol=atol, rtol=rtol) + + @pytest.mark.parametrize("dtype", [torch.bfloat16], ids=["bfloat16"]) + @pytest.mark.parametrize("hidden_size", [2048, 256, 33]) + @pytest.mark.parametrize("batch_size", [32, 121]) + def test_gemma_rmsnorm(self, batch_size, hidden_size, dtype): + + x = torch.randn([batch_size, hidden_size], dtype=dtype) + x = make_non_contiguous(x) + weight = torch.randn(hidden_size, dtype=dtype) + + out = torch.ops.sgl_kernel.gemma_rmsnorm_cpu(x, weight, eps) + ref_out = self._gemma_rmsnorm_native(x, weight, eps) + + atol = rtol = precision[ref_out.dtype] + torch.testing.assert_close(ref_out, out, atol=atol, rtol=rtol) + + ref_x = x.clone() + residual = torch.randn([batch_size, hidden_size], dtype=dtype) + ref_residual = residual.clone() + + torch.ops.sgl_kernel.gemma_fused_add_rmsnorm_cpu(x, residual, weight, eps) ref_x, ref_residual = self._gemma_rmsnorm_native( - ref_x, weight, variance_epsilon, ref_residual + ref_x, weight, eps, ref_residual ) torch.testing.assert_close(x, ref_x, atol=atol, rtol=rtol) torch.testing.assert_close(residual, ref_residual, atol=atol, rtol=rtol) - @parametrize( - m=[4096, 1024], - n=[4096, 4109], - dtype=[torch.float16, torch.bfloat16], - ) - def test_gemma3_rmsnorm(self, m, n, dtype): + @pytest.mark.parametrize("dtype", [torch.bfloat16], ids=["bfloat16"]) + @pytest.mark.parametrize("hidden_size", [128, 256]) + @pytest.mark.parametrize("batch_size", [32, 121]) + def test_gemma3_rmsnorm(self, batch_size, hidden_size, dtype): x_list = [ - torch.randn([m, n], dtype=dtype), - torch.randn([1, m, 2, n], dtype=dtype), + torch.randn([batch_size, hidden_size], dtype=dtype), + torch.randn([batch_size, 16, 2, hidden_size], dtype=dtype), ] for x in x_list: x = make_non_contiguous(x) - hidden_size = x.size(-1) weight = torch.randn(hidden_size, dtype=dtype) - variance_epsilon = 1e-6 - out = torch.ops.sgl_kernel.gemma3_rmsnorm_cpu(x, weight, variance_epsilon) - ref_out = self._gemma3_rmsnorm_native(x, weight, variance_epsilon) + out = torch.ops.sgl_kernel.gemma3_rmsnorm_cpu(x, weight, eps) + ref_out = self._gemma3_rmsnorm_native(x, weight, eps) atol = rtol = precision[ref_out.dtype] torch.testing.assert_close(ref_out, out, atol=atol, rtol=rtol) @@ -212,7 +162,7 @@ class TestNorm(CustomTestCase): self, x: torch.Tensor, weight: torch.Tensor, - variance_epsilon: float = 1e-6, + variance_epsilon: float = eps, scale_shift: float = 0.0, with_scale: bool = True, ): @@ -221,53 +171,41 @@ class TestNorm(CustomTestCase): output = output * (weight.float() + scale_shift) return output.type_as(x) - @parametrize( - m=[4096, 1024], - n=[4096, 4109], - dtype=[torch.float16, torch.bfloat16], - ) - def test_gemma4_rmsnorm(self, m, n, dtype): - for scale_shift, with_scale in [ - (0.0, True), - (1.0, True), - (0.0, False), - (1.0, False), - ]: - x_list = [ - torch.randn([m, n], dtype=dtype), - torch.randn([4, m, n], dtype=dtype), - ] - # Add non-block-contiguous 3D input - base = torch.randn([4, 2 * m, n], dtype=dtype) - x_list.append(base[:, :m, :]) + @pytest.mark.parametrize("dtype", [torch.bfloat16], ids=["bfloat16"]) + @pytest.mark.parametrize("hidden_size", [128, 2048]) + @pytest.mark.parametrize("batch_size", [32, 121]) + @pytest.mark.parametrize("scale_shift", [0.0, 1.0], ids=["shift0.0", "shift1.0"]) + @pytest.mark.parametrize("with_scale", [True, False], ids=["scale", "no-scale"]) + def test_gemma4_rmsnorm( + self, batch_size, hidden_size, dtype, scale_shift, with_scale + ): + x_list = [ + torch.randn([batch_size, hidden_size], dtype=dtype), + torch.randn([batch_size, 4, hidden_size], dtype=dtype), + ] - for x in x_list: - x = make_non_contiguous(x) - hidden_size = x.size(-1) - weight = torch.randn(hidden_size, dtype=dtype) - variance_epsilon = 1e-6 + for x in x_list: + x = make_non_contiguous(x) + weight = torch.randn(hidden_size, dtype=dtype) - out = torch.ops.sgl_kernel.gemma4_rmsnorm_cpu( - x, weight, variance_epsilon, scale_shift, with_scale - ) - ref_out = self._gemma4_rmsnorm_native( - x, weight, variance_epsilon, scale_shift, with_scale - ) + out = torch.ops.sgl_kernel.gemma4_rmsnorm_cpu( + x, weight, eps, scale_shift, with_scale + ) + ref_out = self._gemma4_rmsnorm_native( + x, weight, eps, scale_shift, with_scale + ) - atol = rtol = precision[ref_out.dtype] - torch.testing.assert_close(ref_out, out, atol=atol, rtol=rtol) + atol = rtol = precision[ref_out.dtype] + torch.testing.assert_close(ref_out, out, atol=atol, rtol=rtol) -class TestFusedRMSNormGated(CustomTestCase): - M = [4096, 1024] - N = [4096, 4096 + 13] - dtype = [torch.float16, torch.bfloat16] +class TestFusedRMSNormGated: def _forward_native( self, hidden_states: torch.Tensor, weight: torch.Tensor, - variance_epsilon: float = 1e-6, + variance_epsilon: float = eps, gate: Optional[torch.Tensor] = None, ) -> torch.Tensor: input_dtype = hidden_states.dtype @@ -280,37 +218,29 @@ class TestFusedRMSNormGated(CustomTestCase): return hidden_states.to(input_dtype) - def _norm_test(self, m, n, dtype): - - x = torch.randn([m, n], dtype=dtype) + @pytest.mark.parametrize("dtype", DTYPES, ids=DTYPE_IDS) + @pytest.mark.parametrize("hidden_size", [64, 1024 + 13]) + @pytest.mark.parametrize("batch_size", [32, 121]) + def test_fused_rmsnorm_gated(self, batch_size, hidden_size, dtype): + x = torch.randn([batch_size, hidden_size], dtype=dtype) x = make_non_contiguous(x) - batch_size = x.size(0) - hidden_size = x.size(-1) weight = torch.randn(hidden_size, dtype=dtype) - variance_epsilon = 1e-6 gate = torch.randn([batch_size, hidden_size], dtype=dtype) - out = torch.ops.sgl_kernel.fused_rmsnorm_gated_cpu( - x, weight, gate, variance_epsilon - ) - ref_out = self._forward_native(x, weight, variance_epsilon, gate) + out = torch.ops.sgl_kernel.fused_rmsnorm_gated_cpu(x, weight, gate, eps) + ref_out = self._forward_native(x, weight, eps, gate) - atol = rtol = precision[ref_out.dtype] * 2 + atol = rtol = precision[ref_out.dtype] torch.testing.assert_close(ref_out, out, atol=atol, rtol=rtol) - def test_norm(self): - for params in itertools.product(self.M, self.N, self.dtype): - with self.subTest(m=params[0], n=params[1], dtype=params[2]): - self._norm_test(*params) - -class TestLayerNorm(CustomTestCase): +class TestLayerNorm: def _forward_native( self, x: torch.Tensor, weight: torch.Tensor, - variance_epsilon: float, + variance_epsilon: float = eps, residual: Optional[torch.Tensor] = None, bias: Optional[torch.Tensor] = None, ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: @@ -328,109 +258,141 @@ class TestLayerNorm(CustomTestCase): x = x.to(orig_dtype) return x if residual is None else (x, residual) - @parametrize( - m=[4096, 1024], - n=[4096, 4109], - dtype=[torch.float16, torch.bfloat16], + @pytest.mark.parametrize("dtype", [torch.bfloat16], ids=["bfloat16"]) + @pytest.mark.parametrize("batch_size", [32, 121]) + @pytest.mark.parametrize("hidden_size", [128, 4096, 533]) + @pytest.mark.parametrize("has_bias", [False, True], ids=["no-bias", "bias"]) + def test_layernorm( + self, + batch_size: int, + hidden_size: int, + has_bias: bool, + dtype: torch.dtype, + ) -> None: + x_list = [ + torch.randn([batch_size, hidden_size], dtype=dtype), + torch.randn([batch_size, 3, hidden_size], dtype=dtype), + ] + + for x in x_list: + x = make_non_contiguous(x) + weight = torch.randn(hidden_size, dtype=dtype) + bias = torch.randn(hidden_size, dtype=dtype) if has_bias else None + + ln_out = torch.ops.sgl_kernel.layernorm_cpu(x, weight, bias, eps) + ref_ln_out = self._forward_native(x, weight, eps, residual=None, bias=bias) + + atol = rtol = precision[ref_ln_out.dtype] + torch.testing.assert_close(ln_out, ref_ln_out, atol=atol, rtol=rtol) + + residual = torch.randn(x.shape, dtype=dtype) + ref_residual = residual.clone() + + add_ln_out = torch.ops.sgl_kernel.fused_add_layernorm_cpu( + x, residual, weight, bias, eps + ) + ref_add_ln_out, ref_residual = self._forward_native( + x, weight, eps, residual=ref_residual, bias=bias + ) + + torch.testing.assert_close(add_ln_out, ref_add_ln_out, atol=atol, rtol=rtol) + torch.testing.assert_close(residual, ref_residual, atol=atol, rtol=rtol) + + +class TestFusedQKGemmaRMSNorm: + + def _gemma_rmsnorm_per_head_native( + self, + x: torch.Tensor, + weight: torch.Tensor, + head_dim: int, + variance_epsilon: float = eps, + ) -> torch.Tensor: + orig_dtype = x.dtype + x_f = x.to(torch.float32).reshape(-1, head_dim) + variance = x_f.pow(2).mean(dim=-1, keepdim=True) + x_f = x_f * torch.rsqrt(variance + variance_epsilon) + x_f = x_f * (1.0 + weight.to(torch.float32)) + return x_f.to(orig_dtype).reshape_as(x) + + @pytest.mark.parametrize("dtype", [torch.bfloat16], ids=["bfloat16"]) + @pytest.mark.parametrize( + "batch_size,num_head,num_head_kv,head_dim", + [ + (8, 4, 2, 128), + (17, 8, 2, 64), + (5, 3, 1, 96), + ], ) - def test_norm_input_2d(self, m: int, n: int, dtype: torch.dtype) -> None: - x = torch.randn([m, n], dtype=dtype) - x = make_non_contiguous(x) - hidden_size = x.size(-1) - weight = torch.randn(hidden_size, dtype=dtype) - bias = torch.randn(hidden_size, dtype=dtype) - variance_epsilon = 1e-6 + def test_fused_qk_gemma_rmsnorm( + self, batch_size: int, num_head: int, num_head_kv: int, head_dim: int, dtype + ): + q = torch.randn([batch_size, num_head * head_dim], dtype=dtype) + k = torch.randn([batch_size, num_head_kv * head_dim], dtype=dtype) - ln_out = torch.ops.sgl_kernel.layernorm_cpu(x, weight, None, variance_epsilon) - ref_ln_out = self._forward_native(x, weight, variance_epsilon) + # Keep last dim contiguous but make base storage non-contiguous to stress stride handling. + q = make_non_contiguous(q) + k = make_non_contiguous(k) - atol = rtol = precision[ref_ln_out.dtype] - torch.testing.assert_close(ln_out, ref_ln_out, atol=atol, rtol=rtol) + q_weight = torch.randn(head_dim, dtype=dtype) + k_weight = torch.randn(head_dim, dtype=dtype) - ln_out = torch.ops.sgl_kernel.layernorm_cpu(x, weight, bias, variance_epsilon) - ref_ln_out = self._forward_native( - x, weight, variance_epsilon, residual=None, bias=bias - ) - torch.testing.assert_close(ln_out, ref_ln_out, atol=atol, rtol=rtol) - - residual = torch.randn([m, hidden_size], dtype=dtype) - ref_residual = residual.clone() - - add_ln_out = torch.ops.sgl_kernel.fused_add_layernorm_cpu( - x, residual, weight, None, variance_epsilon - ) - ref_add_ln_out, ref_residual = self._forward_native( - x, weight, variance_epsilon, residual=ref_residual + q_out, k_out = torch.ops.sgl_kernel.fused_qk_gemma_rmsnorm_cpu( + q, k, q_weight, k_weight, eps, head_dim ) - torch.testing.assert_close(add_ln_out, ref_add_ln_out, atol=atol, rtol=rtol) - torch.testing.assert_close(residual, ref_residual, atol=atol, rtol=rtol) + ref_q_out = self._gemma_rmsnorm_per_head_native(q, q_weight, head_dim, eps) + ref_k_out = self._gemma_rmsnorm_per_head_native(k, k_weight, head_dim, eps) - residual = torch.randn([m, hidden_size], dtype=dtype) - ref_residual = residual.clone() + atol = rtol = precision[ref_q_out.dtype] + torch.testing.assert_close(q_out, ref_q_out, atol=atol, rtol=rtol) + torch.testing.assert_close(k_out, ref_k_out, atol=atol, rtol=rtol) - add_ln_out = torch.ops.sgl_kernel.fused_add_layernorm_cpu( - x, residual, weight, bias, variance_epsilon - ) - ref_add_ln_out, ref_residual = self._forward_native( - x, weight, variance_epsilon, residual=ref_residual, bias=bias - ) - - torch.testing.assert_close(add_ln_out, ref_add_ln_out, atol=atol, rtol=rtol) - torch.testing.assert_close(residual, ref_residual, atol=atol, rtol=rtol) - - @parametrize( - l=[4096, 1024], - m=[1, 4], - n=[4096, 4109, 2304], - dtype=[torch.float16, torch.bfloat16], + @pytest.mark.parametrize("dtype", [torch.bfloat16], ids=["bfloat16"]) + @pytest.mark.parametrize( + "batch_size,num_head,num_head_kv,head_dim", + [ + (8, 4, 2, 128), + (17, 8, 2, 64), + (5, 3, 1, 96), + ], ) - def test_norm_input_3d(self, l: int, m: int, n: int, dtype: torch.dtype) -> None: - x = torch.randn([l, m, n], dtype=dtype) - x = make_non_contiguous(x) - hidden_size = x.size(-1) - weight = torch.randn(hidden_size, dtype=dtype) - bias = torch.randn(hidden_size, dtype=dtype) - variance_epsilon = 1e-6 - - ln_out = torch.ops.sgl_kernel.layernorm_cpu(x, weight, None, variance_epsilon) - ref_ln_out = self._forward_native(x, weight, variance_epsilon) - - atol = rtol = precision[ref_ln_out.dtype] - torch.testing.assert_close(ln_out, ref_ln_out, atol=atol, rtol=rtol) - - ln_out = torch.ops.sgl_kernel.layernorm_cpu(x, weight, bias, variance_epsilon) - ref_ln_out = self._forward_native( - x, weight, variance_epsilon, residual=None, bias=bias + def test_fused_qk_gemma_rmsnorm_with_gate( + self, batch_size: int, num_head: int, num_head_kv: int, head_dim: int, dtype + ): + q = torch.randn([batch_size, num_head, head_dim], dtype=dtype) + gate = torch.randn([batch_size, num_head, head_dim], dtype=dtype) + q_gate = torch.cat((q, gate), dim=-1).reshape( + batch_size, num_head * head_dim * 2 ) - torch.testing.assert_close(ln_out, ref_ln_out, atol=atol, rtol=rtol) + k = torch.randn([batch_size, num_head_kv * head_dim], dtype=dtype) - residual = torch.randn([l, m, hidden_size], dtype=dtype) - ref_residual = residual.clone() + q_gate = make_non_contiguous(q_gate) + k = make_non_contiguous(k) - add_ln_out = torch.ops.sgl_kernel.fused_add_layernorm_cpu( - x, residual, weight, None, variance_epsilon - ) - ref_add_ln_out, ref_residual = self._forward_native( - x, weight, variance_epsilon, ref_residual + q_weight = torch.randn(head_dim, dtype=dtype) + k_weight = torch.randn(head_dim, dtype=dtype) + + q_out, k_out, gate_out = ( + torch.ops.sgl_kernel.fused_qk_gemma_rmsnorm_with_gate_cpu( + q_gate, k, q_weight, k_weight, eps, head_dim, num_head + ) ) - torch.testing.assert_close(add_ln_out, ref_add_ln_out, atol=atol, rtol=rtol) - torch.testing.assert_close(residual, ref_residual, atol=atol, rtol=rtol) + ref_q_out = self._gemma_rmsnorm_per_head_native(q, q_weight, head_dim, eps) + ref_k_out = self._gemma_rmsnorm_per_head_native(k, k_weight, head_dim, eps) - residual = torch.randn([l, m, hidden_size], dtype=dtype) - ref_residual = residual.clone() - - add_ln_out = torch.ops.sgl_kernel.fused_add_layernorm_cpu( - x, residual, weight, bias, variance_epsilon + atol = rtol = precision[ref_q_out.dtype] + torch.testing.assert_close( + q_out, ref_q_out.reshape(-1, head_dim), atol=atol, rtol=rtol ) - ref_add_ln_out, ref_residual = self._forward_native( - x, weight, variance_epsilon, residual=ref_residual, bias=bias + torch.testing.assert_close( + k_out, ref_k_out.reshape(-1, head_dim), atol=atol, rtol=rtol + ) + torch.testing.assert_close( + gate_out, gate.reshape(-1, head_dim), atol=atol, rtol=rtol ) - - torch.testing.assert_close(add_ln_out, ref_add_ln_out, atol=atol, rtol=rtol) - torch.testing.assert_close(residual, ref_residual, atol=atol, rtol=rtol) if __name__ == "__main__": - unittest.main() + sys.exit(pytest.main([__file__])) diff --git a/test/registered/cpu/utils.py b/test/registered/cpu/utils.py index 307cf8fb0..83dba491d 100644 --- a/test/registered/cpu/utils.py +++ b/test/registered/cpu/utils.py @@ -511,11 +511,14 @@ class MXFP4QuantizeUtil: def make_non_contiguous(x: torch.Tensor) -> torch.Tensor: - """ - Make a tensor non-contiguous by slicing it via last dimension. - """ + # Make a tensor non-contiguous without changing shape. + if not x.is_contiguous(): + return x + last_dim = x.shape[-1] - return x[..., : last_dim // 2] if x.is_contiguous() else x + expanded = torch.empty(*x.shape[:-1], last_dim + 32, dtype=x.dtype, device=x.device) + expanded[..., :last_dim].copy_(x) + return expanded.narrow(-1, 0, last_dim) def awq_reverse_reorder_int_tensor(int_tensor, bits: int):