@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / f80ae5bdcf

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 15, 11:00 UTC

[Kernel] Use fused rmsnorm for some models like qwen3 series (#17735)

Signed-off-by: evian <eviantai@u.nus.edu> Co-authored-by: evian <eviantai@u.nus.edu>

Wanrui Dai committed May 7, 2025 at 14:10 UTC f80ae5bdcfa71ac56bae78ffb0d3ef85d49a67fb

7 files changed +19 -15

csrc/layernorm_kernels.cu

@@ -140,6 +140,10 @@ void rms_norm(torch::Tensor& out, // [..., hidden_size]

140 torch::Tensor& input, // [..., hidden_size]

141 torch::Tensor& weight, // [hidden_size]

142 double epsilon) {

143 + TORCH_CHECK(out.is_contiguous());

144 + TORCH_CHECK(input.is_contiguous());

145 + TORCH_CHECK(weight.is_contiguous());

146 +

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

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

149

vllm/_custom_ops.py

+3 -1

@@ -186,7 +186,9 @@ def batched_rotary_embedding(positions: torch.Tensor, query: torch.Tensor,

186 # layer norm ops

187 def rms_norm(out: torch.Tensor, input: torch.Tensor, weight: torch.Tensor,

188 epsilon: float) -> None:

189 - torch.ops._C.rms_norm(out, input, weight, epsilon)

189 + # TODO: Remove this contiguous call when the kernel is updated to support non-contiguous input

190 + input_contiguous = input.contiguous()

191 + torch.ops._C.rms_norm(out, input_contiguous, weight, epsilon)

192

193

194 def fused_add_rms_norm(input: torch.Tensor, residual: torch.Tensor,

vllm/model_executor/models/intern_vit.py

+4 -6

@@ -190,8 +190,8 @@ class InternParallelAttention(nn.Module):

190 if self.tp_size > 1:

191 q = tensor_model_parallel_all_gather(q.contiguous())

192 k = tensor_model_parallel_all_gather(k.contiguous())

193 - q = self.q_norm.forward_native(q)

194 - k = self.k_norm.forward_native(k)

193 + q = self.q_norm(q)

194 + k = self.k_norm(k)

195 if self.tp_size > 1:

196 splitter = partial(split_tensor_along_last_dim,

197 num_partitions=self.tp_size)

@@ -264,10 +264,8 @@ class InternSdpaAttention(nn.Module):

264

265 if self.qk_normalization:

266 B_, N_, H_, D_ = q.shape

267 - q = self.q_norm.forward_native(q.flatten(-2,

268 - -1)).view(B_, N_, H_, D_)

269 - k = self.k_norm.forward_native(k.flatten(-2,

270 - -1)).view(B_, N_, H_, D_)

267 + q = self.q_norm(q.flatten(-2, -1)).view(B_, N_, H_, D_)

268 + k = self.k_norm(k.flatten(-2, -1)).view(B_, N_, H_, D_)

269 q = q.transpose(1, 2)

270 k = k.transpose(1, 2)

271 v = v.transpose(1, 2)

vllm/model_executor/models/molmo.py

+2 -2

@@ -438,8 +438,8 @@ class MolmoAttention(nn.Module):

438 if self.tp_size > 1:

439 q = tensor_model_parallel_all_gather(q.contiguous())

440 k = tensor_model_parallel_all_gather(k.contiguous())

441 - q = self.q_norm.forward_native(q)

442 - k = self.k_norm.forward_native(k)

441 + q = self.q_norm(q)

442 + k = self.k_norm(k)

443 if self.tp_size > 1:

444 splitter = partial(split_tensor_along_last_dim,

445 num_partitions=self.tp_size)

vllm/model_executor/models/olmo2.py

+2 -2

@@ -139,8 +139,8 @@ class Olmo2Attention(nn.Module):

139 if self.tp_size > 1:

140 q = tensor_model_parallel_all_gather(q.contiguous())

141 k = tensor_model_parallel_all_gather(k.contiguous())

142 - q = self.q_norm.forward_native(q)

143 - k = self.k_norm.forward_native(k)

142 + q = self.q_norm(q)

143 + k = self.k_norm(k)

144 if self.tp_size > 1:

145 splitter = partial(split_tensor_along_last_dim,

146 num_partitions=self.tp_size)

vllm/model_executor/models/qwen3.py

+2 -2

@@ -133,11 +133,11 @@ class Qwen3Attention(nn.Module):

133 # Add qk-norm

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

135 self.head_dim)

136 - q_by_head = self.q_norm.forward_native(q_by_head)

136 + q_by_head = self.q_norm(q_by_head)

137 q = q_by_head.view(q.shape)

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

139 self.head_dim)

140 - k_by_head = self.k_norm.forward_native(k_by_head)

140 + k_by_head = self.k_norm(k_by_head)

141 k = k_by_head.view(k.shape)

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

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

vllm/model_executor/models/qwen3_moe.py

+2 -2

@@ -225,12 +225,12 @@ class Qwen3MoeAttention(nn.Module):

225 # Add qk-norm

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

227 self.head_dim)

228 - q_by_head = self.q_norm.forward_native(q_by_head)

228 + q_by_head = self.q_norm(q_by_head)

229 q = q_by_head.view(q.shape)

230

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

232 self.head_dim)

233 - k_by_head = self.k_norm.forward_native(k_by_head)

233 + k_by_head = self.k_norm(k_by_head)

234 k = k_by_head.view(k.shape)

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

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