From 4da5599e93b2d4b351a8934aa1a72ac5bc76d69e Mon Sep 17 00:00:00 2001 From: Lianmin Zheng Date: Tue, 15 Sep 2026 15:22:26 -0700 Subject: [PATCH] Fix first-token metadata and reused attention-layer indexing (#39328) Co-authored-by: Jinghui Zhang Co-authored-by: Lucia Fang <116399278+luccafong@users.noreply.github.com> --- .../cuda_graph_setup.py | 20 +++++++++++++++++-- .../srt/observability/req_time_stats.py | 2 ++ .../test_cuda_graph_setup.py | 19 ++++++++++++++++++ .../unit/observability/test_req_time_stats.py | 20 +++++++++++++++++++ 4 files changed, 59 insertions(+), 2 deletions(-) diff --git a/python/sglang/srt/model_executor/model_runner_components/cuda_graph_setup.py b/python/sglang/srt/model_executor/model_runner_components/cuda_graph_setup.py index d59182854..eda5b1099 100644 --- a/python/sglang/srt/model_executor/model_runner_components/cuda_graph_setup.py +++ b/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 diff --git a/python/sglang/srt/observability/req_time_stats.py b/python/sglang/srt/observability/req_time_stats.py index 4c174950a..4afd9d73b 100644 --- a/python/sglang/srt/observability/req_time_stats.py +++ b/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 diff --git a/test/registered/unit/model_executor/model_runner_components/test_cuda_graph_setup.py b/test/registered/unit/model_executor/model_runner_components/test_cuda_graph_setup.py index b886ef3fb..ac9aad7da 100644 --- a/test/registered/unit/model_executor/model_runner_components/test_cuda_graph_setup.py +++ b/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 diff --git a/test/registered/unit/observability/test_req_time_stats.py b/test/registered/unit/observability/test_req_time_stats.py index 40237b4b7..bfb5ddf4f 100644 --- a/test/registered/unit/observability/test_req_time_stats.py +++ b/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()