Skip to content

[BugFix] Support rms_norm for stacked 2D weights on XPU #579

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
ccrhx4 wants to merge 2 commits into
base: main
Choose a base branch
from
+218 −15
Open

[BugFix] Support rms_norm for stacked 2D weights on XPU #579

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
292fb95
Fix rms_norm silently using layer-0 weight for stacked 2D weights on XPU
ccrhx4 Sep 7, 2026
8e2c0b0
Address PR review comments for rms_norm batched-weight fix
ccrhx4 Sep 8, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
102 changes: 87 additions & 15 deletions csrc/layernorm.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@ class rms_norm_kernel {
const int64_t input_shape_d2_, // input.size(-2)
const int64_t input_shape_d3_, // input.size(-3)
const scalar_t* weight_,
const int64_t weight_stride_, // 0 for 1D weight, else weight.stride(0)
const float epsilon_,
const int num_tokens_,
const int hidden_size_,
Expand All @@ -34,6 +35,7 @@ class rms_norm_kernel {
input_shape_d2(input_shape_d2_),
input_shape_d3(input_shape_d3_),
weight(weight_),
weight_stride(weight_stride_),
epsilon(epsilon_),
num_tokens(num_tokens_),
hidden_size(hidden_size_),
Expand All @@ -47,24 +49,32 @@ class rms_norm_kernel {
float variance = 0.0f;

const scalar_t* input_row;
// batch_idx selects the weight row for a 2D (stacked) weight; unused
// (stays 0) when weight is 1D since weight_stride is 0 in that case.
int batch_idx = 0;
if constexpr (NUM_DIMS == 2) {
// 2D for layernorm normal case [batch_size, hidden]
batch_idx = item_ct1.get_group(2);
input_row = input + item_ct1.get_group(2) * input_stride_d2;
} else if constexpr (NUM_DIMS == 3) {
// 3D for q/k norm [batch_size, num_heads, head_size]
int batch_idx = item_ct1.get_group(2) / input_shape_d2;
batch_idx = item_ct1.get_group(2) / input_shape_d2;
int head_idx = item_ct1.get_group(2) % input_shape_d2;
input_row =
input + batch_idx * input_stride_d3 + head_idx * input_stride_d2;
} else if constexpr (NUM_DIMS == 4) {
// 4D for transformers model_impl qk norm [batch, seq, head, head_dim]
int batch_idx = item_ct1.get_group(2) / (input_shape_d3 * input_shape_d2);
batch_idx = item_ct1.get_group(2) / (input_shape_d3 * input_shape_d2);
int remaining = item_ct1.get_group(2) % (input_shape_d3 * input_shape_d2);
int seq_idx = remaining / input_shape_d2;
int head_idx = remaining % input_shape_d2;
input_row = input + batch_idx * input_stride_d4 +
seq_idx * input_stride_d3 + head_idx * input_stride_d2;
}
const scalar_t* weight_row = nullptr;
if constexpr (HasWeight) {
weight_row = weight + batch_idx * weight_stride;
}

auto vec_op = [&variance](const vec_n_t<scalar_t, VEC_SIZE>& vec) {
#pragma unroll
Expand Down Expand Up @@ -96,7 +106,8 @@ class rms_norm_kernel {
scalar_t* out_row = out + item_ct1.get_group(2) * hidden_size;
auto* v_in =
reinterpret_cast<const vec_n_t<scalar_t, VEC_SIZE>*>(input_row);
auto* v_w = reinterpret_cast<const vec_n_t<scalar_t, VEC_SIZE>*>(weight);
auto* v_w =
reinterpret_cast<const vec_n_t<scalar_t, VEC_SIZE>*>(weight_row);
auto* v_out = reinterpret_cast<vec_n_t<scalar_t, VEC_SIZE>*>(out_row);
int64_t const out_num_vec_elems = hidden_size / VEC_SIZE;
float s_variance_val = *s_variance_ptr;
Expand Down Expand Up @@ -131,7 +142,9 @@ class rms_norm_kernel {
const int64_t input_stride_d4;
const int64_t input_shape_d2;
const int64_t input_shape_d3;
const scalar_t* __restrict__ weight; // [hidden_size]
const scalar_t* __restrict__ weight; // [hidden_size] or [num_rows,
// hidden_size]
const int64_t weight_stride; // 0 for 1D weight, else weight.stride(0)
const float epsilon;
const int num_tokens;
const int hidden_size;
Expand All @@ -151,6 +164,7 @@ class rms_norm_kernel<scalar_t, NUM_DIMS, 0, HasWeight> {
const int64_t input_shape_d2_, // input.size(-2)
const int64_t input_shape_d3_, // input.size(-3)
const scalar_t* weight_,
const int64_t weight_stride_, // 0 for 1D weight, else weight.stride(0)
const float epsilon_,
const int num_tokens_,
const int hidden_size_,
Expand All @@ -164,6 +178,7 @@ class rms_norm_kernel<scalar_t, NUM_DIMS, 0, HasWeight> {
input_shape_d2(input_shape_d2_),
input_shape_d3(input_shape_d3_),
weight(weight_),
weight_stride(weight_stride_),
epsilon(epsilon_),
num_tokens(num_tokens_),
hidden_size(hidden_size_),
Expand All @@ -177,24 +192,32 @@ class rms_norm_kernel<scalar_t, NUM_DIMS, 0, HasWeight> {
float variance = 0.0f;

const scalar_t* input_row;
// batch_idx selects the weight row for a 2D (stacked) weight; unused
// (stays 0) when weight is 1D since weight_stride is 0 in that case.
int batch_idx = 0;
if constexpr (NUM_DIMS == 2) {
// 2D for layernorm normal case [batch_size, hidden]
batch_idx = item_ct1.get_group(2);
input_row = input + item_ct1.get_group(2) * input_stride_d2;
} else if constexpr (NUM_DIMS == 3) {
// 3D for q/k norm [batch_size, num_heads, head_size]
int batch_idx = item_ct1.get_group(2) / input_shape_d2;
batch_idx = item_ct1.get_group(2) / input_shape_d2;
int head_idx = item_ct1.get_group(2) % input_shape_d2;
input_row =
input + batch_idx * input_stride_d3 + head_idx * input_stride_d2;
} else if constexpr (NUM_DIMS == 4) {
// 4D for transformers model_impl qk norm [batch, seq, head, head_dim]
int batch_idx = item_ct1.get_group(2) / (input_shape_d3 * input_shape_d2);
batch_idx = item_ct1.get_group(2) / (input_shape_d3 * input_shape_d2);
int remaining = item_ct1.get_group(2) % (input_shape_d3 * input_shape_d2);
int seq_idx = remaining / input_shape_d2;
int head_idx = remaining % input_shape_d2;
input_row = input + batch_idx * input_stride_d4 +
seq_idx * input_stride_d3 + head_idx * input_stride_d2;
}
const scalar_t* weight_row = nullptr;
if constexpr (HasWeight) {
weight_row = weight + batch_idx * weight_stride;
}

auto scalar_op = [&variance](const scalar_t& val) {
float x = static_cast<float>(val);
Expand Down Expand Up @@ -225,7 +248,7 @@ class rms_norm_kernel<scalar_t, NUM_DIMS, 0, HasWeight> {
float inv_rms = *s_variance_ptr;
if constexpr (HasWeight) {
out_row[idx] = static_cast<scalar_t>(
x * inv_rms * (static_cast<float>(weight[idx]) + weight_bias));
x * inv_rms * (static_cast<float>(weight_row[idx]) + weight_bias));
} else {
out_row[idx] = static_cast<scalar_t>(x * inv_rms);
}
Expand All @@ -240,7 +263,9 @@ class rms_norm_kernel<scalar_t, NUM_DIMS, 0, HasWeight> {
const int64_t input_stride_d4;
const int64_t input_shape_d2;
const int64_t input_shape_d3;
const scalar_t* __restrict__ weight; // [hidden_size]
const scalar_t* __restrict__ weight; // [hidden_size] or [num_rows,
// hidden_size]
const int64_t weight_stride; // 0 for 1D weight, else weight.stride(0)
const float epsilon;
const int num_tokens;
const int hidden_size;
Expand Down Expand Up @@ -269,6 +294,7 @@ class rms_norm_multi_row_kernel {
const int64_t input_shape_d2_,
const int64_t input_shape_d3_,
const scalar_t* weight_,
const int64_t weight_stride_, // 0 for 1D weight, else weight.stride(0)
const float epsilon_,
const int num_tokens_,
const int hidden_size_,
Expand All @@ -282,6 +308,7 @@ class rms_norm_multi_row_kernel {
input_shape_d2(input_shape_d2_),
input_shape_d3(input_shape_d3_),
weight(weight_),
weight_stride(weight_stride_),
epsilon(epsilon_),
num_tokens(num_tokens_),
hidden_size(hidden_size_),
Expand All @@ -304,21 +331,29 @@ class rms_norm_multi_row_kernel {
s_variance.template get_multi_ptr<sycl::access::decorated::no>().get();

const scalar_t* input_row;
// batch_idx selects the weight row for a 2D (stacked) weight; unused
// (stays 0) when weight is 1D since weight_stride is 0 in that case.
int batch_idx = 0;
if constexpr (NUM_DIMS == 2) {
batch_idx = global_row;
input_row = input + global_row * input_stride_d2;
} else if constexpr (NUM_DIMS == 3) {
int batch_idx = global_row / input_shape_d2;
batch_idx = global_row / input_shape_d2;
int head_idx = global_row % input_shape_d2;
input_row =
input + batch_idx * input_stride_d3 + head_idx * input_stride_d2;
} else if constexpr (NUM_DIMS == 4) {
int batch_idx = global_row / (input_shape_d3 * input_shape_d2);
batch_idx = global_row / (input_shape_d3 * input_shape_d2);
int remaining = global_row % (input_shape_d3 * input_shape_d2);
int seq_idx = remaining / input_shape_d2;
int head_idx = remaining % input_shape_d2;
input_row = input + batch_idx * input_stride_d4 +
seq_idx * input_stride_d3 + head_idx * input_stride_d2;
}
const scalar_t* weight_row = nullptr;
if constexpr (HasWeight) {
weight_row = weight + batch_idx * weight_stride;
}

float variance = 0.0f;
const int64_t num_vec_elems = hidden_size / VEC_SIZE;
Expand Down Expand Up @@ -359,7 +394,8 @@ class rms_norm_multi_row_kernel {
scalar_t* out_row = out + global_row * hidden_size;
auto* v_in =
reinterpret_cast<const vec_n_t<scalar_t, VEC_SIZE>*>(input_row);
auto* v_w = reinterpret_cast<const vec_n_t<scalar_t, VEC_SIZE>*>(weight);
auto* v_w =
reinterpret_cast<const vec_n_t<scalar_t, VEC_SIZE>*>(weight_row);
auto* v_out = reinterpret_cast<vec_n_t<scalar_t, VEC_SIZE>*>(out_row);
float s_var = s_variance_ptr[row_in_wg];

Expand Down Expand Up @@ -392,7 +428,9 @@ class rms_norm_multi_row_kernel {
const int64_t input_stride_d4;
const int64_t input_shape_d2;
const int64_t input_shape_d3;
const scalar_t* __restrict__ weight;
const scalar_t* __restrict__ weight; // [hidden_size] or [num_rows,
// hidden_size]
const int64_t weight_stride; // 0 for 1D weight, else weight.stride(0)
const float epsilon;
const int num_tokens;
const int hidden_size;
Expand All @@ -405,6 +443,7 @@ void call_rms_norm_kernel(
torch::Tensor& out,
torch::Tensor& input,
const scalar_t* weight_ptr,
int64_t weight_stride, // 0 for 1D weight, else weight.stride(0)
float epsilon,
float weight_bias = 0.0f) {
using sycl_t = typename vllm::xpu::SyclTypeTrait<scalar_t>::Type;
Expand Down Expand Up @@ -481,6 +520,7 @@ void call_rms_norm_kernel(
input_shape_d2,
input_shape_d3,
(const sycl_t*)weight_ptr,
weight_stride,
epsilon,
num_tokens,
hidden_size,
Expand All @@ -507,6 +547,7 @@ void call_rms_norm_kernel(
input_shape_d2,
input_shape_d3,
(const sycl_t*)weight_ptr,
weight_stride,
epsilon,
num_tokens,
hidden_size,
Expand All @@ -531,6 +572,7 @@ void call_rms_norm_kernel(
input_shape_d2,
input_shape_d3,
(const sycl_t*)weight_ptr,
weight_stride,
epsilon,
num_tokens,
hidden_size,
Expand Down Expand Up @@ -800,19 +842,44 @@ void rms_norm(
}
TORCH_CHECK(input.stride(-1) == 1);
const bool has_weight = weight.has_value();
// weight may be 1D `[hidden_size]` (the common case) or 2D
// `[input.size(0), hidden_size]` (a stacked per-outer-row weight, e.g. one
// row per speculative-decoding draft layer). weight_stride selects which
// row a given input row uses: 0 disables per-row selection (every row
// reads the same 1D weight), matching CUDA's rms_norm_kernel contract.
int64_t weight_stride = 0;
if (has_weight) {
TORCH_CHECK(weight->is_contiguous());
TORCH_CHECK(
weight->scalar_type() == input.scalar_type(),
"rms_norm: weight dtype must match input dtype");
if (weight->dim() == 1) {
TORCH_CHECK(
weight->size(0) == input.size(-1),
"rms_norm: 1D weight size must match hidden_size");
} else if (weight->dim() == 2) {
TORCH_CHECK(
weight->size(0) == input.size(0),
"rms_norm: 2D weight's outer dim must match input's outer "
"dim");
TORCH_CHECK(
weight->size(-1) == input.size(-1),
"rms_norm: 2D weight's hidden dim must match hidden_size");
weight_stride = weight->stride(0);
} else {
TORCH_CHECK(false, "rms_norm: weight must be 1D or 2D");
}
}
VLLM_DISPATCH_FLOATING_TYPES(
input.scalar_type(), "call_rms_norm_kernel", [&] {
const scalar_t* weight_ptr =
has_weight ? weight->data_ptr<scalar_t>() : nullptr;
if (has_weight) {
vllm::call_rms_norm_kernel<scalar_t, true>(
out, input, weight_ptr, epsilon);
out, input, weight_ptr, weight_stride, epsilon);
} else {
vllm::call_rms_norm_kernel<scalar_t, false>(
out, input, weight_ptr, epsilon);
out, input, weight_ptr, weight_stride, epsilon);
}
});
}
Expand Down Expand Up @@ -861,7 +928,12 @@ void gemma_rms_norm(
input.scalar_type(), "call_gemma_rms_norm_kernel", [&] {
const scalar_t* weight_ptr = weight.data_ptr<scalar_t>();
vllm::call_rms_norm_kernel<scalar_t, /*HasWeight=*/true>(
out, input, weight_ptr, epsilon, /*weight_bias=*/1.0f);
out,
input,
weight_ptr,
/*weight_stride=*/0,
epsilon,
/*weight_bias=*/1.0f);
});
}

Expand Down
Loading
Loading