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

test/registered/unit/model_executor/model_runner_components/test_cuda_graph_setup.py

@@ -48,6 +48,25 @@ def test_pipeline_attention_metadata_is_indexed_by_global_layer_id():
     assert companions[24] is companion24
 
 
+def test_reuse_tables_pass_through_but_distinct_duplicates_raise():
+    looped = SimpleNamespace(layer_id=1)
+    companion = object()
+    attention_in = [SimpleNamespace(layer_id=0), looped, looped]
+    companions_in = [None, companion, companion]
+
+    attention, companions = index_attention_layers_by_global_id(
+        attention_in, companions_in
+    )
+
+    assert attention is attention_in
+    assert companions is companions_in
+
+    with pytest.raises(ValueError, match="duplicate attention layer_id: 2"):
+        index_attention_layers_by_global_id(
+            [SimpleNamespace(layer_id=2), SimpleNamespace(layer_id=2)], [None, None]
+        )
+
+
 def test_model_runner_can_override_decode_graph_runner(monkeypatch):
     from sglang.srt.runtime_context import get_context
 

test/registered/unit/observability/test_req_time_stats.py

@@ -46,6 +46,26 @@ class TestSetstatePreservesUnsetTimeSentinels(CustomTestCase):
         self.assertAlmostEqual(hop2.wait_queue_entry_time, 123.456 - 9.0)
 
 
+class TestOutputMetaInfo(CustomTestCase):
+    def test_first_token_latency_requires_valid_timestamps(self):
+        for created, first, expected in (
+            (1.0, 1.5, 0.5),
+            (0.0, 1.5, None),
+            (1.0, 0.0, None),
+            (1.0, 1.0, None),
+            (1.0, 0.5, None),
+        ):
+            with self.subTest(created=created, first=first):
+                stats = rts.APIServerReqTimeStats()
+                stats.created_time = created
+                stats.first_token_time = first
+                meta_info = stats.convert_to_output_meta_info()
+                if expected is None:
+                    self.assertNotIn("first_token_latency", meta_info)
+                else:
+                    self.assertAlmostEqual(meta_info["first_token_latency"], expected)
+
+
 class TestConvertToGenAiSpanAttrs(CustomTestCase):
     def _stats_after_first_token(self) -> rts.APIServerReqTimeStats:
         stats = rts.APIServerReqTimeStats()