Fix first-token metadata and reused attention-layer indexing (#39328)

Co-authored-by: Jinghui Zhang <jinghui@meta.com>
Co-authored-by: Lucia Fang <116399278+luccafong@users.noreply.github.com>

python/sglang/srt/model_executor/model_runner_components/cuda_graph_setup.py

@@ -95,7 +95,12 @@ def index_attention_layers_by_global_id(
     mha_companion_layers: list[Any],
     layer_model=None,
 ) -> tuple[list[Any], list[Any]]:
-    """Pad PP-local attention metadata so global layer_id remains a valid index."""
+    """Pad PP-local attention metadata so global layer_id remains a valid index.
+
+    Models that re-execute layers pre-expand these lists into position-indexed
+    lookup tables (the same layer at several positions); such tables are
+    returned unchanged.
+    """
     if len(attention_layers) != len(mha_companion_layers):
         raise ValueError("attention and MHA companion metadata must be parallel")
     populated = [layer for layer in attention_layers if layer is not None]
@@ -109,16 +114,27 @@ def index_attention_layers_by_global_id(
     max_layer_id = max(int(layer.layer_id) for layer in populated)
     indexed_attention = [None] * (max_layer_id + 1)
     indexed_companions = [None] * (max_layer_id + 1)
+    has_reused_layers = False
     for attention, companion in zip(attention_layers, mha_companion_layers):
         if attention is None:
             if companion is not None:
                 raise ValueError("MHA companion has no primary attention layer")
             continue
         layer_id = int(attention.layer_id)
-        if layer_id < 0 or indexed_attention[layer_id] is not None:
+        if layer_id < 0:
             raise ValueError(f"invalid or duplicate attention layer_id: {layer_id}")
+        if indexed_attention[layer_id] is not None:
+            if (
+                indexed_attention[layer_id] is not attention
+                or indexed_companions[layer_id] is not companion
+            ):
+                raise ValueError(f"invalid or duplicate attention layer_id: {layer_id}")
+            has_reused_layers = True
+            continue
         indexed_attention[layer_id] = attention
         indexed_companions[layer_id] = companion
+    if has_reused_layers:
+        return attention_layers, mha_companion_layers
     return indexed_attention, indexed_companions
 
 

python/sglang/srt/observability/req_time_stats.py

@@ -499,6 +499,8 @@ class APIServerReqTimeStats(ReqTimeStatsBase):
             meta_info["request_received_ts"] = convert_time_to_realtime(
                 self.created_time
             )
+        if self.created_time > 0.0 and self.first_token_time > self.created_time:
+            meta_info["first_token_latency"] = self.get_first_token_latency()
         if self.api_server_dispatch_finish_time > 0.0:
             meta_info["api_server_dispatch_finish_ts"] = convert_time_to_realtime(
                 self.api_server_dispatch_finish_time