[Perf] RMSNorm uncontiguous support, 1.2~3.1x kernel performance improvement (#49750)

Signed-off-by: yewentao256 <[email protected]>
Signed-off-by: Wentao Ye <[email protected]>

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<scalar_t, width>>);
   static_assert(sizeof(_f16Vec<scalar_t, width>) == sizeof(scalar_t) * width);
@@ -130,7 +131,7 @@ fused_add_rms_norm_kernel(
       reinterpret_cast<const _f16Vec<scalar_t, width>*>(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<scalar_t, width> 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<scalar_t, width> res = residual_v[id];
     _f16Vec<scalar_t, width> 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<float, 1024>;
@@ -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<scalar_t>(), input_stride,        \
                   residual.mutable_data_ptr<scalar_t>(),                   \
                   weight->const_data_ptr<scalar_t>(), epsilon, num_tokens, \
-                  hidden_size);                                            \
+                  hidden_size, residual_stride);                           \
         } else {                                                           \
           vllm::fused_add_rms_norm_kernel<scalar_t, width, false>          \
               <<<grid, block, 0, stream>>>(                                \
                   input.mutable_data_ptr<scalar_t>(), input_stride,        \
                   residual.mutable_data_ptr<scalar_t>(), 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<torch::stable::Tensor> 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<std::uintptr_t>(input.data_ptr());
   auto res_ptr = reinterpret_cast<std::uintptr_t>(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<std::uintptr_t>(weight->data_ptr());

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.

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)

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)

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)

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)

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)

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)

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)