From c2ab3713e95a8d2afc2222f8684f722648a185b9 Mon Sep 17 00:00:00 2001 From: Yi Zhong <207368749+vincentzed@users.noreply.github.com> Date: Sat, 31 Jan 2026 19:02:57 -0500 Subject: [PATCH] [Performance] Optimize Mllama LayerNorm -> Upd (#9725) --- python/sglang/srt/models/mllama.py | 31 ++++++++---------------------- 1 file changed, 8 insertions(+), 23 deletions(-) diff --git a/python/sglang/srt/models/mllama.py b/python/sglang/srt/models/mllama.py index 5be7cda58..e0d5713d7 100644 --- a/python/sglang/srt/models/mllama.py +++ b/python/sglang/srt/models/mllama.py @@ -475,23 +475,6 @@ class MllamaVisionModel(nn.Module): return hidden_state -class MllamaTextRMSNorm(nn.Module): - def __init__(self, hidden_size, eps=1e-6): - super().__init__() - self.weight = nn.Parameter(torch.ones(hidden_size)) - self.variance_epsilon = eps - - def forward(self, hidden_states): - input_dtype = hidden_states.dtype - hidden_states = hidden_states.to(torch.float32) - variance = hidden_states.pow(2).mean(-1, keepdim=True) - hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) - return self.weight * hidden_states.to(input_dtype) - - def extra_repr(self): - return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}" - - class MllamaTextCrossAttention(nn.Module): def __init__( self, @@ -534,10 +517,8 @@ class MllamaTextCrossAttention(nn.Module): quant_config=quant_config, prefix=add_prefix("o_proj", prefix), ) - # vllm.model_executor.layers.layernorm.RMSNorm has precision issue, - # use huggingface's instead - self.q_norm = MllamaTextRMSNorm(self.head_dim, eps=config.rms_norm_eps) - self.k_norm = MllamaTextRMSNorm(self.head_dim, eps=config.rms_norm_eps) + self.q_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) + self.k_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) self.scaling = self.head_dim**-0.5 self.attn = RadixAttention( @@ -572,9 +553,13 @@ class MllamaTextCrossAttention(nn.Module): ) k = k.view(-1, self.num_local_key_value_heads, self.head_dim) v = v.view(-1, self.num_local_key_value_heads, self.head_dim) - k = self.k_norm(k) + k = self.k_norm(k.reshape(-1, self.head_dim)).reshape( + -1, self.num_local_key_value_heads, self.head_dim + ) q = q.view(-1, self.num_local_heads, self.head_dim) - q = self.q_norm(q) + q = self.q_norm(q.reshape(-1, self.head_dim)).reshape( + -1, self.num_local_heads, self.head_dim + ) output = self.attn(q, k, v, forward_batch) out, _ = self.o_proj(output)