chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

[CPU] support LayerNorm with 3D shape (#15075)

Browse Source 
Co-authored-by: Ma Mingfei <mingfei.ma@intel.com>
This commit is contained in:
Zaili Wang
2026-03-19 13:15:24 +08:00
committed by GitHub
parent dc6aa26ce9
commit 2f4babe32b
4 changed files with 196 additions and 54 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
python/sglang/srt/layers/layernorm.py
View File

@@ -421,8 +421,9 @@ class LayerNorm(MultiPlatformOp):
x: torch.Tensor,
) -> torch.Tensor:
if _is_cpu_amx_available:
bias_data = self.bias.data if self.use_bias else None
return torch.ops.sgl_kernel.layernorm_cpu(
x, self.weight.data, self.variance_epsilon
x, self.weight.data, bias_data, self.variance_epsilon
)
else:
return self.forward_native(x)

122
sgl-kernel/csrc/cpu/norm.cpp
View File

@@ -388,26 +388,32 @@ void fused_rmsnorm_gated_kernel_impl(
template <typename scalar_t>
void fused_add_layernorm_kernel_impl(
scalar_t* __restrict__ input,
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<scalar_t>;
using fVec = at::vec::Vectorized<float>;
constexpr int kVecSize = bVec::size();
at::parallel_for(0, batch_size, 0, [&](int64_t begin, int64_t end) {
int tid = at::get_thread_num();
float* __restrict__ buffer_ptr = buffer + tid * hidden_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__ input_ptr = input + i * input_strideN;
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 (residual != nullptr) {
if (has_residual) {
residual_ptr = residual + i * hidden_size;
}
@@ -422,7 +428,7 @@ void fused_add_layernorm_kernel_impl(
fVec x_fvec0, x_fvec1;
std::tie(x_fvec0, x_fvec1) = at::vec::convert_to_float(x_bvec);
if (residual_ptr != nullptr) {
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);
@@ -445,7 +451,7 @@ void fused_add_layernorm_kernel_impl(
#pragma GCC unroll 4
for (; d < hidden_size; ++d) {
float x_val = static_cast<float>(input_ptr[d]);
if (residual_ptr != nullptr) {
if (has_residual) {
float r_val = static_cast<float>(residual_ptr[d]);
x_val += r_val;
residual_ptr[d] = static_cast<scalar_t>(x_val);
@@ -475,7 +481,6 @@ void fused_add_layernorm_kernel_impl(
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);
@@ -483,14 +488,25 @@ void fused_add_layernorm_kernel_impl(
x_fvec0 = (x_fvec0 - mean_fvec) * scale_fvec * w_fvec0;
x_fvec1 = (x_fvec1 - mean_fvec) * scale_fvec * w_fvec1;
bVec x_bvec = convert_from_float_ext<scalar_t>(x_fvec0, x_fvec1);
x_bvec.store(input_ptr + d);
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<scalar_t>(x_fvec0, x_fvec1);
o_bvec.store(out_ptr + d);
}
#pragma GCC unroll 4
for (; d < hidden_size; ++d) {
float normalized = (buffer_ptr[d] - mean) * rsqrt_var;
float x_val = normalized * static_cast<float>(weight[d]);
input_ptr[d] = static_cast<scalar_t>(x_val);
if (has_bias) {
x_val += static_cast<float>(bias[d]);
}
out_ptr[d] = static_cast<scalar_t>(x_val);
}
}
});
@@ -543,33 +559,52 @@ at::Tensor rmsnorm_cpu(at::Tensor& input, at::Tensor& weight, double eps) {
return output;
}
// input : {batch_size, hidden_size}
// input : {batch_size, hidden_size} or {batch_size, seq_len, hidden_size}
// weight: {hidden_size}
void layernorm_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<at::Tensor>& bias, double eps) {
RECORD_FUNCTION("sgl-kernel::layernorm_cpu", std::vector<c10::IValue>({input, weight}));
CHECK_LAST_DIM_CONTIGUOUS_INPUT(input);
CHECK_INPUT(weight);
CHECK_DIM(2, input);
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));
int64_t batch_size = input.size(0);
int64_t hidden_size = input.size(1);
int64_t input_strideN = input.stride(0);
if (bias.has_value()) {
CHECK_DIM(1, bias.value());
CHECK_EQ(bias.value().size(0), weight.size(0));
}
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));
}
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<scalar_t>(
output.data_ptr<scalar_t>(),
input.data_ptr<scalar_t>(),
nullptr,
weight.data_ptr<scalar_t>(),
conditional_data_ptr<scalar_t>(bias),
buffer.data_ptr<float>(),
batch_size,
seq_len,
hidden_size,
input_strideN,
eps);
});
return output;
}
at::Tensor gemma_rmsnorm_cpu(at::Tensor& input, at::Tensor& weight, double eps) {
@@ -776,36 +811,65 @@ void gemma_fused_add_rmsnorm_cpu(at::Tensor& input, at::Tensor& residual, at::Te
});
}
// input : {batch_size, hidden_size}
// residual: {batch_size, hidden_size}
// input : {batch_size, hidden_size} or {batch_size, seq_len, hidden_size}
// residual: {batch_size, hidden_size} or {batch_size, seq_len, hidden_size}
// weight : {hidden_size}
void fused_add_layernorm_cpu(at::Tensor& input, at::Tensor& residual, at::Tensor& weight, double eps) {
// bias : {hidden_size}
at::Tensor fused_add_layernorm_cpu(
const at::Tensor& input,
at::Tensor& residual,
const at::Tensor& weight,
const std::optional<at::Tensor>& bias,
double eps) {
RECORD_FUNCTION("sgl-kernel::fused_add_layernorm_cpu", std::vector<c10::IValue>({input, residual, weight}));
CHECK_LAST_DIM_CONTIGUOUS_INPUT(input);
CHECK_INPUT(residual);
CHECK_INPUT(weight);
CHECK_DIM(2, input);
CHECK_DIM(2, residual);
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);
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));
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);
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));
}
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);
}
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<scalar_t>(
output.data_ptr<scalar_t>(),
input.data_ptr<scalar_t>(),
residual.data_ptr<scalar_t>(),
weight.data_ptr<scalar_t>(),
conditional_data_ptr<scalar_t>(bias),
buffer.data_ptr<float>(),
batch_size,
seq_len,
hidden_size,
input_strideN,
eps);
});
return output;
}

16
sgl-kernel/csrc/cpu/torch_extension_cpu.cpp
View File

@@ -36,7 +36,8 @@ at::Tensor gemma_rmsnorm_cpu(at::Tensor& input, at::Tensor& weight, double eps);
at::Tensor gemma3_rmsnorm_cpu(at::Tensor& input, at::Tensor& weight, double eps);
// layernorm
void layernorm_cpu(at::Tensor& input, at::Tensor& weight, double eps);
at::Tensor
layernorm_cpu(const at::Tensor& input, const at::Tensor& weight, const std::optional<at::Tensor>& bias, double eps);
// qwen3_next_rmsnorm_gated
at::Tensor fused_rmsnorm_gated_cpu(at::Tensor& input, at::Tensor& weight, at::Tensor& gate, double eps);
@@ -46,7 +47,12 @@ void fused_add_rmsnorm_cpu(at::Tensor& input, at::Tensor& residual, at::Tensor&
void gemma_fused_add_rmsnorm_cpu(at::Tensor& input, at::Tensor& residual, at::Tensor& weight, double eps);
// fused_add_layernorm
void fused_add_layernorm_cpu(at::Tensor& input, at::Tensor& residual, at::Tensor& weight, double eps);
at::Tensor fused_add_layernorm_cpu(
const at::Tensor& input,
at::Tensor& residual,
const at::Tensor& weight,
const std::optional<at::Tensor>& bias,
double eps);
// topk
std::tuple<at::Tensor, at::Tensor>
@@ -364,7 +370,7 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) {
m.impl("gemma_rmsnorm_cpu", torch::kCPU, &gemma_rmsnorm_cpu);
m.def("gemma3_rmsnorm_cpu(Tensor input, Tensor weight, float eps) -> Tensor");
m.impl("gemma3_rmsnorm_cpu", torch::kCPU, &gemma3_rmsnorm_cpu);
m.def("layernorm_cpu(Tensor(a!) input, Tensor weight, float eps) -> ()");
m.def("layernorm_cpu(Tensor input, Tensor weight, Tensor? bias, float eps) -> Tensor");
m.impl("layernorm_cpu", torch::kCPU, &layernorm_cpu);
m.def("l2norm_cpu(Tensor input, float eps) -> Tensor");
m.impl("l2norm_cpu", torch::kCPU, &l2norm_cpu);
@@ -374,7 +380,9 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) {
m.impl("fused_add_rmsnorm_cpu", torch::kCPU, &fused_add_rmsnorm_cpu);
m.def("gemma_fused_add_rmsnorm_cpu(Tensor(a!) input, Tensor(a!) residual, Tensor weight, float eps) -> ()");
m.impl("gemma_fused_add_rmsnorm_cpu", torch::kCPU, &gemma_fused_add_rmsnorm_cpu);
m.def("fused_add_layernorm_cpu(Tensor(a!) input, Tensor(a!) residual, Tensor weight, float eps) -> ()");
m.def(
"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);
// topk

109
test/srt/cpu/test_norm.py
View File

@@ -215,6 +215,7 @@ class TestLayerNorm(CustomTestCase):
weight: torch.Tensor,
variance_epsilon: float,
residual: Optional[torch.Tensor] = None,
bias: Optional[torch.Tensor] = None,
) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
orig_dtype = x.dtype
x = x.to(torch.float32)
@@ -224,45 +225,113 @@ class TestLayerNorm(CustomTestCase):
variance, mean = torch.var_mean(x, dim=-1, keepdim=True, correction=0)
x = (x - mean) * torch.rsqrt(variance + variance_epsilon)
x = x.to(orig_dtype) * weight
if residual is None:
return x
else:
return x, residual
x = x * weight.to(torch.float32)
if bias is not None:
x = x + bias.to(torch.float32)
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: int, n: int, dtype: torch.dtype) -> None:
x_ln = torch.randn([m, n], dtype=dtype)
x_ln = make_non_contiguous(x_ln)
ref_x_ln = x_ln.clone()
hidden_size = x_ln.size(-1)
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
torch.ops.sgl_kernel.layernorm_cpu(x_ln, weight, variance_epsilon)
ref_ln_out = self._forward_native(ref_x_ln, weight, variance_epsilon)
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(x_ln, ref_ln_out, atol=atol, rtol=rtol)
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
)
torch.testing.assert_close(ln_out, ref_ln_out, atol=atol, rtol=rtol)
x_add_ln = torch.randn([m, n], dtype=dtype)
x_add_ln = make_non_contiguous(x_add_ln)
ref_x_add_ln = x_add_ln.clone()
residual = torch.randn([m, hidden_size], dtype=dtype)
ref_residual = residual.clone()
torch.ops.sgl_kernel.fused_add_layernorm_cpu(
x_add_ln, residual, weight, variance_epsilon
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(
ref_x_add_ln, weight, variance_epsilon, ref_residual
x, weight, variance_epsilon, residual=ref_residual
)
torch.testing.assert_close(x_add_ln, ref_add_ln_out, 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)
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, 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],
)
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
)
torch.testing.assert_close(ln_out, ref_ln_out, atol=atol, rtol=rtol)
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, None, variance_epsilon
)
ref_add_ln_out, ref_residual = self._forward_native(
x, weight, variance_epsilon, ref_residual
)
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)
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
)
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)