[perf] Add two stream norm for Olmo3 speedup 5% (#13681)

Signed-off-by: Vincent Zhong <207368749+vincentzed@users.noreply.github.com>
Co-authored-by: b8zhong <b8zhong@uwaterloo.ca>

python/sglang/srt/models/olmo2.py

@@ -43,9 +43,12 @@ from sglang.srt.layers.vocab_parallel_embedding import (
     ParallelLMHead,
     VocabParallelEmbedding,
 )
+from sglang.srt.model_executor.cuda_graph_runner import get_is_capture_mode
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch
 from sglang.srt.model_loader.weight_utils import default_weight_loader
-from sglang.srt.utils import add_prefix, make_layers
+from sglang.srt.utils import add_prefix, is_cuda, make_layers
+
+_is_cuda = is_cuda()
 
 
 # Aligned with HF's implementation, using sliding window inclusive with the last token
@@ -67,6 +70,7 @@ class Olmo2Attention(nn.Module):
         layer_id: int = 0,
         quant_config: Optional[QuantizationConfig] = None,
         prefix: str = "",
+        alt_stream: Optional[torch.cuda.Stream] = None,
     ):
         super().__init__()
         self.config = config
@@ -107,6 +111,7 @@ class Olmo2Attention(nn.Module):
             prefix=add_prefix("qkv_proj", prefix),
         )
         self.tp_rank = get_tensor_model_parallel_rank()
+        self.alt_stream = alt_stream
 
         self.k_norm = RMSNorm(
             self.total_num_kv_heads * self.head_dim,
@@ -161,8 +166,29 @@ class Olmo2Attention(nn.Module):
         if self.tp_size > 1:
             q = tensor_model_parallel_all_gather(q.contiguous())
             k = tensor_model_parallel_all_gather(k.contiguous())
-        q = self.q_norm.forward_native(q)
-        k = self.k_norm.forward_native(k)
+
+        if self.alt_stream is not None and get_is_capture_mode():
+            current_stream = torch.cuda.current_stream()
+            self.alt_stream.wait_stream(current_stream)
+
+            q_shape = q.shape
+            k_shape = k.shape
+
+            q_by_last = q.reshape(-1, q_shape[-1])
+            q_by_last = self.q_norm(q_by_last)
+
+            with torch.cuda.stream(self.alt_stream):
+                k_by_last = k.reshape(-1, k_shape[-1])
+                k_by_last = self.k_norm(k_by_last)
+
+            current_stream.wait_stream(self.alt_stream)
+
+            q = q_by_last.view(q_shape)
+            k = k_by_last.view(k_shape)
+        else:
+            q = self.q_norm.forward_native(q)
+            k = self.k_norm.forward_native(k)
+
         if self.tp_size > 1:
             splitter = partial(split_tensor_along_last_dim, num_partitions=self.tp_size)
             q = splitter(q)[self.tp_rank]
@@ -246,12 +272,18 @@ class Olmo2DecoderLayer(nn.Module):
         layer_id: int = 0,
         quant_config: Optional[QuantizationConfig] = None,
         prefix: str = "",
+        alt_stream: Optional[torch.cuda.Stream] = None,
     ):
         super().__init__()
         self.layer_id = layer_id
+        self.alt_stream = alt_stream
         # Attention block.
         self.self_attn = Olmo2Attention(
-            config, layer_id, quant_config, prefix=add_prefix("self_attn", prefix)
+            config,
+            layer_id,
+            quant_config,
+            prefix=add_prefix("self_attn", prefix),
+            alt_stream=alt_stream,
         )
 
         # MLP block.
@@ -293,9 +325,13 @@ class Olmo2Model(nn.Module):
         config: PretrainedConfig,
         quant_config: Optional[QuantizationConfig] = None,
         prefix: str = "",
+        alt_stream: Optional[torch.cuda.Stream] = None,
     ):
         super().__init__()
         self.config = config
+        if alt_stream is None and _is_cuda:
+            alt_stream = torch.cuda.Stream()
+        self.alt_stream = alt_stream
 
         self.embed_tokens = VocabParallelEmbedding(
             config.vocab_size,
@@ -309,6 +345,7 @@ class Olmo2Model(nn.Module):
                 layer_id=idx,
                 quant_config=quant_config,
                 prefix=prefix,
+                alt_stream=self.alt_stream,
             ),
             prefix=add_prefix("layers", prefix),
         )
@@ -357,11 +394,15 @@ class Olmo2ForCausalLM(nn.Module):
         config: PretrainedConfig,
         quant_config: Optional[QuantizationConfig] = None,
         prefix: str = "",
+        alt_stream: Optional[torch.cuda.Stream] = None,
     ):
         super().__init__()
         self.config = config
         self.model = Olmo2Model(
-            config, quant_config, prefix=add_prefix("model", prefix)
+            config,
+            quant_config,
+            prefix=add_prefix("model", prefix),
+            alt_stream=alt_stream,
         )
         if config.tie_word_embeddings:
             self.lm_head = self.model.embed_tokens