@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 82642d7d6c

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

Signed-off-by: yewentao256 <zhyanwentao@126.com> Signed-off-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com>

Wentao Ye committed Jul 29, 2026 at 16:45 UTC 82642d7d6c332fefdec06b260a246bd1a4ebf5ff

9 files changed +31 -29

csrc/libtorch_stable/layernorm_kernels.cu

+15 -11

@@ -110,7 +110,8 @@ fused_add_rms_norm_kernel(

110 const int64_t input_stride,

111 scalar_t* __restrict__ residual, // [..., hidden_size]

112 const scalar_t* __restrict__ weight, // [hidden_size], null if !HasWeight

113 - const float epsilon, const int num_tokens, const int hidden_size) {

113 + const float epsilon, const int num_tokens, const int hidden_size,

114 + const int64_t residual_stride) {

115 // Sanity checks on our vector struct and type-punned pointer arithmetic

116 static_assert(std::is_pod_v<_f16Vec<scalar_t, width>>);

117 static_assert(sizeof(_f16Vec<scalar_t, width>) == sizeof(scalar_t) * width);

@@ -130,7 +131,7 @@ fused_add_rms_norm_kernel(

131 reinterpret_cast<const _f16Vec<scalar_t, width>*>(weight);

132

133 for (int idx = threadIdx.x; idx < vec_hidden_size; idx += blockDim.x) {

133 - int id = blockIdx.x * vec_hidden_size + idx;

134 + int64_t id = blockIdx.x * residual_stride / width + idx;

135 int64_t strided_id = blockIdx.x * vec_input_stride + idx;

136 _f16Vec<scalar_t, width> temp = input_v[strided_id];

137 temp += residual_v[id];

@@ -148,7 +149,7 @@ fused_add_rms_norm_kernel(

149 __syncthreads();

150

151 for (int idx = threadIdx.x; idx < vec_hidden_size; idx += blockDim.x) {

151 - int id = blockIdx.x * vec_hidden_size + idx;

152 + int64_t id = blockIdx.x * residual_stride / width + idx;

153 int64_t strided_id = blockIdx.x * vec_input_stride + idx;

154 _f16Vec<scalar_t, width> res = residual_v[id];

155 _f16Vec<scalar_t, width> out;

@@ -182,16 +183,17 @@ fused_add_rms_norm_kernel(

183 const int64_t input_stride,

184 scalar_t* __restrict__ residual, // [..., hidden_size]

185 const scalar_t* __restrict__ weight, // [hidden_size], null if !HasWeight

185 - const float epsilon, const int num_tokens, const int hidden_size) {

186 + const float epsilon, const int num_tokens, const int hidden_size,

187 + const int64_t residual_stride) {

188 __shared__ float s_variance;

189 float variance = 0.0f;

190

191 for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) {

192 scalar_t z = input[blockIdx.x * input_stride + idx];

191 - z += residual[blockIdx.x * hidden_size + idx];

193 + z += residual[blockIdx.x * residual_stride + idx];

194 float x = (float)z;

195 variance += x * x;

194 - residual[blockIdx.x * hidden_size + idx] = z;

196 + residual[blockIdx.x * residual_stride + idx] = z;

197 }

198

199 using BlockReduce = cub::BlockReduce<float, 1024>;

@@ -204,7 +206,7 @@ fused_add_rms_norm_kernel(

206 __syncthreads();

207

208 for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) {

207 - float x = (float)residual[blockIdx.x * hidden_size + idx];

209 + float x = (float)residual[blockIdx.x * residual_stride + idx];

210 if constexpr (HasWeight) {

211 float w = (float)weight[idx];

212 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]

301 input.mutable_data_ptr<scalar_t>(), input_stride, \

302 residual.mutable_data_ptr<scalar_t>(), \

303 weight->const_data_ptr<scalar_t>(), epsilon, num_tokens, \

302 - hidden_size); \

304 + hidden_size, residual_stride); \

305 } else { \

306 vllm::fused_add_rms_norm_kernel<scalar_t, width, false> \

307 <<<grid, block, 0, stream>>>( \

308 input.mutable_data_ptr<scalar_t>(), input_stride, \

309 residual.mutable_data_ptr<scalar_t>(), nullptr, epsilon, \

308 - num_tokens, hidden_size); \

310 + num_tokens, hidden_size, residual_stride); \

311 } \

312 });

313

@@ -314,13 +316,14 @@ void fused_add_rms_norm(torch::stable::Tensor& input, // [..., hidden_size]

316 std::optional<torch::stable::Tensor> weight,

317 double epsilon) {

318 STD_TORCH_CHECK(input.scalar_type() == residual.scalar_type());

317 - STD_TORCH_CHECK(residual.is_contiguous());

319 + STD_TORCH_CHECK(residual.stride(-1) == 1);

320 if (weight.has_value()) {

321 STD_TORCH_CHECK(weight->scalar_type() == input.scalar_type());

322 STD_TORCH_CHECK(weight->is_contiguous());

323 }

324 int hidden_size = input.size(-1);

325 int64_t input_stride = input.stride(-2);

326 + int64_t residual_stride = residual.stride(-2);

327 int num_tokens = input.numel() / hidden_size;

328

329 dim3 grid(num_tokens);

@@ -343,7 +346,8 @@ void fused_add_rms_norm(torch::stable::Tensor& input, // [..., hidden_size]

346 auto inp_ptr = reinterpret_cast<std::uintptr_t>(input.data_ptr());

347 auto res_ptr = reinterpret_cast<std::uintptr_t>(residual.data_ptr());

348 bool offsets_are_multiple_of_vector_width =

346 - hidden_size % vector_width == 0 && input_stride % vector_width == 0;

349 + hidden_size % vector_width == 0 && input_stride % vector_width == 0 &&

350 + residual_stride % vector_width == 0;

351 const bool has_weight = weight.has_value();

352 if (has_weight) {

353 auto wt_ptr = reinterpret_cast<std::uintptr_t>(weight->data_ptr());

tests/kernels/core/test_layernorm.py

+3 -1

@@ -60,7 +60,9 @@ def test_rms_norm(

60 x = x[..., :hidden_size]

61 assert x.is_contiguous() != strided_input

62 x *= scale

63 - residual = torch.randn_like(x) * scale if add_residual else None

63 + residual = x.new_empty_strided(x.size(), x.stride()) if add_residual else None

64 + if residual is not None:

65 + residual.normal_(std=scale)

67 # NOTE(woosuk): The reference implementation should be executed first

68 # because the custom kernel is in-place.

vllm/model_executor/models/apertus.py

+2 -2

@@ -211,8 +211,8 @@ class ApertusAttention(nn.Module):

211 ) -> torch.Tensor:

212 qkv, _ = self.qkv_proj(hidden_states)

213 q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)

214 - q = self.q_norm(q.contiguous().view(-1, self.head_dim)).view_as(q)

215 - k = self.k_norm(k.contiguous().view(-1, self.head_dim)).view_as(k)

214 + q = self.q_norm(q.view(-1, self.num_heads, self.head_dim)).view_as(q)

215 + k = self.k_norm(k.view(-1, self.num_kv_heads, self.head_dim)).view_as(k)

216 q, k = self.rotary_emb(positions, q, k)

217 attn_output = self.attn(q, k, v)

218 output, _ = self.o_proj(attn_output)

vllm/model_executor/models/deepseek_v2.py

-4

@@ -1481,8 +1481,6 @@ class DeepseekV2Model(nn.Module):

1481 hidden_states, residual = combined_states.split(

1482 [self.hidden_size, self.hidden_size], dim=-1

1483 )

1484 - # fused_add_rms_norm requires a contiguous residual

1485 - residual = residual.contiguous()

1484 if idx in self.aux_hidden_state_layers:

1485 aux_hidden_state = hidden_states + residual

1486 if aux_hidden_state.shape[0] != positions.shape[0]:

@@ -1507,8 +1505,6 @@ class DeepseekV2Model(nn.Module):

1505 hidden_states, residual = combined_states.split(

1506 [self.hidden_size, self.hidden_size], dim=-1

1507 )

1510 - # fused_add_rms_norm requires a contiguous residual

1511 - residual = residual.contiguous()

1508

1509 if self.end_layer in self.aux_hidden_state_layers:

1510 aux_hidden_states.append(hidden_states + residual)

vllm/model_executor/models/hunyuan_v1.py

+4 -4

@@ -238,10 +238,10 @@ class HunYuanAttention(nn.Module):

238 ori_k = k

239 if self.use_qk_norm:

240 q = self.query_layernorm(

241 - q.view(-1, self.num_heads, self.head_dim).contiguous()

241 + q.view(-1, self.num_heads, self.head_dim),

242 )

243 k = self.key_layernorm(

244 - k.view(-1, self.num_kv_heads, self.head_dim).contiguous()

244 + k.view(-1, self.num_kv_heads, self.head_dim),

245 )

246

247 attn_output = self.attn(q, k, v)

@@ -346,10 +346,10 @@ class HunYuanCrossAttention(nn.Module):

346 q, _ = self.rotary_emb(positions, q, k_tmp)

347 if self.use_qk_norm:

348 q = self.query_layernorm(

349 - q.view(-1, self.num_heads, self.head_dim).contiguous()

349 + q.view(-1, self.num_heads, self.head_dim),

350 )

351 k = self.key_layernorm(

352 - k.view(-1, self.num_kv_heads, self.head_dim).contiguous()

352 + k.view(-1, self.num_kv_heads, self.head_dim),

353 )

354

355 attn_output = self.attn(q, k, v)

vllm/model_executor/models/lfm2.py

+2 -2

@@ -169,8 +169,8 @@ class Lfm2Attention(nn.Module):

169 n_tokens, _ = hidden_states.shape

170 qkv, _ = self.qkv_proj(hidden_states)

171 q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)

172 - q = q.view(n_tokens, self.num_heads, self.head_dim).contiguous()

173 - k = k.view(n_tokens, self.num_kv_heads, self.head_dim).contiguous()

172 + q = q.view(n_tokens, self.num_heads, self.head_dim)

173 + k = k.view(n_tokens, self.num_kv_heads, self.head_dim)

174 q = self.q_layernorm(q)

175 k = self.k_layernorm(k)

176 q, k = self.rotary_emb(positions, q, k)

vllm/model_executor/models/lfm2_moe.py

+2 -2

@@ -255,8 +255,8 @@ class Lfm2MoeAttention(nn.Module):

255 n_tokens, _ = hidden_states.shape

256 qkv, _ = self.qkv_proj(hidden_states)

257 q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)

258 - q = q.view(n_tokens, self.num_heads, self.head_dim).contiguous()

259 - k = k.view(n_tokens, self.num_kv_heads, self.head_dim).contiguous()

258 + q = q.view(n_tokens, self.num_heads, self.head_dim)

259 + k = k.view(n_tokens, self.num_kv_heads, self.head_dim)

260 q = self.q_layernorm(q)

261 k = self.k_layernorm(k)

262 q, k = self.rotary_emb(positions, q, k)

vllm/model_executor/models/minicpm3.py

+1 -1

@@ -146,7 +146,7 @@ class MiniCPM3Attention(nn.Module):

146 latent_cache, _ = self.kv_a_proj_with_mqa(hidden_states)

147 kv_a, _ = latent_cache.split([self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)

148 latent_cache = latent_cache.unsqueeze(1)

149 - kv_a = self.kv_a_layernorm(kv_a.contiguous())

149 + kv_a = self.kv_a_layernorm(kv_a)

150 kv, _ = self.kv_b_proj(kv_a)

151 kv = kv.view(-1, self.num_local_heads, self.qk_nope_head_dim + self.v_head_dim)

152 k_nope, v = kv.split([self.qk_nope_head_dim, self.v_head_dim], dim=-1)

vllm/model_executor/models/step3p5.py

+2 -2

@@ -271,11 +271,11 @@ class Step3p5Attention(nn.Module):

271 q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)

272 # Add qk-norm inline similar to Qwen3 MOE attention

273 q_by_head = q.view(*q.shape[:-1], q.shape[-1] // self.head_dim, self.head_dim)

274 - q_by_head = self.q_norm(q_by_head.contiguous())

274 + q_by_head = self.q_norm(q_by_head)

275 q = q_by_head.view(q.shape)

276

277 k_by_head = k.view(*k.shape[:-1], k.shape[-1] // self.head_dim, self.head_dim)

278 - k_by_head = self.k_norm(k_by_head.contiguous())

278 + k_by_head = self.k_norm(k_by_head)

279 k = k_by_head.view(k.shape)

280 if self.use_rope:

281 q, k = self.rotary_emb(positions, q, k)