From 82642d7d6c332fefdec06b260a246bd1a4ebf5ff Mon Sep 17 00:00:00 2001 From: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Date: Wed, 29 Jul 2026 16:45:38 -0400 Subject: [PATCH] [Perf] RMSNorm uncontiguous support, 1.2~3.1x kernel performance improvement (#49750) Signed-off-by: yewentao256 Signed-off-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com> --- csrc/libtorch_stable/layernorm_kernels.cu | 26 +++++++++++++---------- tests/kernels/core/test_layernorm.py | 4 +++- vllm/model_executor/models/apertus.py | 4 ++-- vllm/model_executor/models/deepseek_v2.py | 4 ---- vllm/model_executor/models/hunyuan_v1.py | 8 +++---- vllm/model_executor/models/lfm2.py | 4 ++-- vllm/model_executor/models/lfm2_moe.py | 4 ++-- vllm/model_executor/models/minicpm3.py | 2 +- vllm/model_executor/models/step3p5.py | 4 ++-- 9 files changed, 31 insertions(+), 29 deletions(-) diff --git a/csrc/libtorch_stable/layernorm_kernels.cu b/csrc/libtorch_stable/layernorm_kernels.cu index 7a1051d2c00..b342c59f180 100644 --- a/csrc/libtorch_stable/layernorm_kernels.cu +++ b/csrc/libtorch_stable/layernorm_kernels.cu @@ -110,7 +110,8 @@ fused_add_rms_norm_kernel( const int64_t input_stride, scalar_t* __restrict__ residual, // [..., hidden_size] const scalar_t* __restrict__ weight, // [hidden_size], null if !HasWeight - const float epsilon, const int num_tokens, const int hidden_size) { + const float epsilon, const int num_tokens, const int hidden_size, + const int64_t residual_stride) { // Sanity checks on our vector struct and type-punned pointer arithmetic static_assert(std::is_pod_v<_f16Vec>); static_assert(sizeof(_f16Vec) == sizeof(scalar_t) * width); @@ -130,7 +131,7 @@ fused_add_rms_norm_kernel( reinterpret_cast*>(weight); for (int idx = threadIdx.x; idx < vec_hidden_size; idx += blockDim.x) { - int id = blockIdx.x * vec_hidden_size + idx; + int64_t id = blockIdx.x * residual_stride / width + idx; int64_t strided_id = blockIdx.x * vec_input_stride + idx; _f16Vec temp = input_v[strided_id]; temp += residual_v[id]; @@ -148,7 +149,7 @@ fused_add_rms_norm_kernel( __syncthreads(); for (int idx = threadIdx.x; idx < vec_hidden_size; idx += blockDim.x) { - int id = blockIdx.x * vec_hidden_size + idx; + int64_t id = blockIdx.x * residual_stride / width + idx; int64_t strided_id = blockIdx.x * vec_input_stride + idx; _f16Vec res = residual_v[id]; _f16Vec out; @@ -182,16 +183,17 @@ fused_add_rms_norm_kernel( const int64_t input_stride, scalar_t* __restrict__ residual, // [..., hidden_size] const scalar_t* __restrict__ weight, // [hidden_size], null if !HasWeight - const float epsilon, const int num_tokens, const int hidden_size) { + const float epsilon, const int num_tokens, const int hidden_size, + const int64_t residual_stride) { __shared__ float s_variance; float variance = 0.0f; for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) { scalar_t z = input[blockIdx.x * input_stride + idx]; - z += residual[blockIdx.x * hidden_size + idx]; + z += residual[blockIdx.x * residual_stride + idx]; float x = (float)z; variance += x * x; - residual[blockIdx.x * hidden_size + idx] = z; + residual[blockIdx.x * residual_stride + idx] = z; } using BlockReduce = cub::BlockReduce; @@ -204,7 +206,7 @@ fused_add_rms_norm_kernel( __syncthreads(); for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) { - float x = (float)residual[blockIdx.x * hidden_size + idx]; + float x = (float)residual[blockIdx.x * residual_stride + idx]; if constexpr (HasWeight) { float w = (float)weight[idx]; input[blockIdx.x * input_stride + idx] = (scalar_t)(x * s_variance * w); @@ -299,13 +301,13 @@ void rms_norm(torch::stable::Tensor& out, // [..., hidden_size] input.mutable_data_ptr(), input_stride, \ residual.mutable_data_ptr(), \ weight->const_data_ptr(), epsilon, num_tokens, \ - hidden_size); \ + hidden_size, residual_stride); \ } else { \ vllm::fused_add_rms_norm_kernel \ <<>>( \ input.mutable_data_ptr(), input_stride, \ residual.mutable_data_ptr(), nullptr, epsilon, \ - num_tokens, hidden_size); \ + num_tokens, hidden_size, residual_stride); \ } \ }); @@ -314,13 +316,14 @@ void fused_add_rms_norm(torch::stable::Tensor& input, // [..., hidden_size] std::optional weight, double epsilon) { STD_TORCH_CHECK(input.scalar_type() == residual.scalar_type()); - STD_TORCH_CHECK(residual.is_contiguous()); + STD_TORCH_CHECK(residual.stride(-1) == 1); if (weight.has_value()) { STD_TORCH_CHECK(weight->scalar_type() == input.scalar_type()); STD_TORCH_CHECK(weight->is_contiguous()); } int hidden_size = input.size(-1); int64_t input_stride = input.stride(-2); + int64_t residual_stride = residual.stride(-2); int num_tokens = input.numel() / hidden_size; dim3 grid(num_tokens); @@ -343,7 +346,8 @@ void fused_add_rms_norm(torch::stable::Tensor& input, // [..., hidden_size] auto inp_ptr = reinterpret_cast(input.data_ptr()); auto res_ptr = reinterpret_cast(residual.data_ptr()); bool offsets_are_multiple_of_vector_width = - hidden_size % vector_width == 0 && input_stride % vector_width == 0; + hidden_size % vector_width == 0 && input_stride % vector_width == 0 && + residual_stride % vector_width == 0; const bool has_weight = weight.has_value(); if (has_weight) { auto wt_ptr = reinterpret_cast(weight->data_ptr()); diff --git a/tests/kernels/core/test_layernorm.py b/tests/kernels/core/test_layernorm.py index 6e546f154c2..1a6c893ad6b 100644 --- a/tests/kernels/core/test_layernorm.py +++ b/tests/kernels/core/test_layernorm.py @@ -60,7 +60,9 @@ def test_rms_norm( x = x[..., :hidden_size] assert x.is_contiguous() != strided_input x *= scale - residual = torch.randn_like(x) * scale if add_residual else None + residual = x.new_empty_strided(x.size(), x.stride()) if add_residual else None + if residual is not None: + residual.normal_(std=scale) # NOTE(woosuk): The reference implementation should be executed first # because the custom kernel is in-place. diff --git a/vllm/model_executor/models/apertus.py b/vllm/model_executor/models/apertus.py index b997e153a99..74b17129c98 100644 --- a/vllm/model_executor/models/apertus.py +++ b/vllm/model_executor/models/apertus.py @@ -211,8 +211,8 @@ class ApertusAttention(nn.Module): ) -> torch.Tensor: qkv, _ = self.qkv_proj(hidden_states) q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) - q = self.q_norm(q.contiguous().view(-1, self.head_dim)).view_as(q) - k = self.k_norm(k.contiguous().view(-1, self.head_dim)).view_as(k) + q = self.q_norm(q.view(-1, self.num_heads, self.head_dim)).view_as(q) + k = self.k_norm(k.view(-1, self.num_kv_heads, self.head_dim)).view_as(k) q, k = self.rotary_emb(positions, q, k) attn_output = self.attn(q, k, v) output, _ = self.o_proj(attn_output) diff --git a/vllm/model_executor/models/deepseek_v2.py b/vllm/model_executor/models/deepseek_v2.py index bf67e040a15..66068dcca68 100644 --- a/vllm/model_executor/models/deepseek_v2.py +++ b/vllm/model_executor/models/deepseek_v2.py @@ -1481,8 +1481,6 @@ class DeepseekV2Model(nn.Module): hidden_states, residual = combined_states.split( [self.hidden_size, self.hidden_size], dim=-1 ) - # fused_add_rms_norm requires a contiguous residual - residual = residual.contiguous() if idx in self.aux_hidden_state_layers: aux_hidden_state = hidden_states + residual if aux_hidden_state.shape[0] != positions.shape[0]: @@ -1507,8 +1505,6 @@ class DeepseekV2Model(nn.Module): hidden_states, residual = combined_states.split( [self.hidden_size, self.hidden_size], dim=-1 ) - # fused_add_rms_norm requires a contiguous residual - residual = residual.contiguous() if self.end_layer in self.aux_hidden_state_layers: aux_hidden_states.append(hidden_states + residual) diff --git a/vllm/model_executor/models/hunyuan_v1.py b/vllm/model_executor/models/hunyuan_v1.py index 4f70a966289..d646b132589 100644 --- a/vllm/model_executor/models/hunyuan_v1.py +++ b/vllm/model_executor/models/hunyuan_v1.py @@ -238,10 +238,10 @@ class HunYuanAttention(nn.Module): ori_k = k if self.use_qk_norm: q = self.query_layernorm( - q.view(-1, self.num_heads, self.head_dim).contiguous() + q.view(-1, self.num_heads, self.head_dim), ) k = self.key_layernorm( - k.view(-1, self.num_kv_heads, self.head_dim).contiguous() + k.view(-1, self.num_kv_heads, self.head_dim), ) attn_output = self.attn(q, k, v) @@ -346,10 +346,10 @@ class HunYuanCrossAttention(nn.Module): q, _ = self.rotary_emb(positions, q, k_tmp) if self.use_qk_norm: q = self.query_layernorm( - q.view(-1, self.num_heads, self.head_dim).contiguous() + q.view(-1, self.num_heads, self.head_dim), ) k = self.key_layernorm( - k.view(-1, self.num_kv_heads, self.head_dim).contiguous() + k.view(-1, self.num_kv_heads, self.head_dim), ) attn_output = self.attn(q, k, v) diff --git a/vllm/model_executor/models/lfm2.py b/vllm/model_executor/models/lfm2.py index a8387a33641..e18edb0b36e 100644 --- a/vllm/model_executor/models/lfm2.py +++ b/vllm/model_executor/models/lfm2.py @@ -169,8 +169,8 @@ class Lfm2Attention(nn.Module): n_tokens, _ = hidden_states.shape qkv, _ = self.qkv_proj(hidden_states) q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) - q = q.view(n_tokens, self.num_heads, self.head_dim).contiguous() - k = k.view(n_tokens, self.num_kv_heads, self.head_dim).contiguous() + q = q.view(n_tokens, self.num_heads, self.head_dim) + k = k.view(n_tokens, self.num_kv_heads, self.head_dim) q = self.q_layernorm(q) k = self.k_layernorm(k) q, k = self.rotary_emb(positions, q, k) diff --git a/vllm/model_executor/models/lfm2_moe.py b/vllm/model_executor/models/lfm2_moe.py index abdf12c7651..9d1f7bd50c7 100644 --- a/vllm/model_executor/models/lfm2_moe.py +++ b/vllm/model_executor/models/lfm2_moe.py @@ -255,8 +255,8 @@ class Lfm2MoeAttention(nn.Module): n_tokens, _ = hidden_states.shape qkv, _ = self.qkv_proj(hidden_states) q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) - q = q.view(n_tokens, self.num_heads, self.head_dim).contiguous() - k = k.view(n_tokens, self.num_kv_heads, self.head_dim).contiguous() + q = q.view(n_tokens, self.num_heads, self.head_dim) + k = k.view(n_tokens, self.num_kv_heads, self.head_dim) q = self.q_layernorm(q) k = self.k_layernorm(k) q, k = self.rotary_emb(positions, q, k) diff --git a/vllm/model_executor/models/minicpm3.py b/vllm/model_executor/models/minicpm3.py index e61e9d06103..30a9cd34ed0 100644 --- a/vllm/model_executor/models/minicpm3.py +++ b/vllm/model_executor/models/minicpm3.py @@ -146,7 +146,7 @@ class MiniCPM3Attention(nn.Module): latent_cache, _ = self.kv_a_proj_with_mqa(hidden_states) kv_a, _ = latent_cache.split([self.kv_lora_rank, self.qk_rope_head_dim], dim=-1) latent_cache = latent_cache.unsqueeze(1) - kv_a = self.kv_a_layernorm(kv_a.contiguous()) + kv_a = self.kv_a_layernorm(kv_a) kv, _ = self.kv_b_proj(kv_a) kv = kv.view(-1, self.num_local_heads, self.qk_nope_head_dim + self.v_head_dim) k_nope, v = kv.split([self.qk_nope_head_dim, self.v_head_dim], dim=-1) diff --git a/vllm/model_executor/models/step3p5.py b/vllm/model_executor/models/step3p5.py index 07a25d23c8c..1f4fbecce56 100644 --- a/vllm/model_executor/models/step3p5.py +++ b/vllm/model_executor/models/step3p5.py @@ -271,11 +271,11 @@ class Step3p5Attention(nn.Module): q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) # Add qk-norm inline similar to Qwen3 MOE attention q_by_head = q.view(*q.shape[:-1], q.shape[-1] // self.head_dim, self.head_dim) - q_by_head = self.q_norm(q_by_head.contiguous()) + q_by_head = self.q_norm(q_by_head) q = q_by_head.view(q.shape) k_by_head = k.view(*k.shape[:-1], k.shape[-1] // self.head_dim, self.head_dim) - k_by_head = self.k_norm(k_by_head.contiguous()) + k_by_head = self.k_norm(k_by_head) k = k_by_head.view(k.shape) if self.use_rope: q, k = self.rotary_emb(positions, q, k)