From 49424a2118087b04fc85d89eaa6825ffae924a7f Mon Sep 17 00:00:00 2001 From: Kirill Shkolenko Date: Tue, 30 Jun 2026 00:42:41 +0300 Subject: [PATCH 1/3] fix(sglang): treat streamed usage metadata as optional Signed-off-by: Kirill Shkolenko --- .../request_handlers/llm/decode_handler.py | 19 +++++++------- .../tests/test_sglang_decode_handler.py | 25 +++++++++++++++++++ 2 files changed, 35 insertions(+), 9 deletions(-) diff --git a/components/src/dynamo/sglang/request_handlers/llm/decode_handler.py b/components/src/dynamo/sglang/request_handlers/llm/decode_handler.py index 97c22b575f2e..d2ad6abee0a9 100644 --- a/components/src/dynamo/sglang/request_handlers/llm/decode_handler.py +++ b/components/src/dynamo/sglang/request_handlers/llm/decode_handler.py @@ -580,18 +580,19 @@ async def _process_token_stream( # as nvext.routed_experts); disaggregated_params stays KV-transfer only. out["engine_data"] = {"routed_experts": routed_experts} if finish_reason: - input_tokens = meta_info["prompt_tokens"] - completion_tokens = meta_info["completion_tokens"] - cached_tokens = meta_info["cached_tokens"] + input_tokens = meta_info.get("prompt_tokens") + completion_tokens = meta_info.get("completion_tokens") + cached_tokens = meta_info.get("cached_tokens") prefill_prompt_tokens_details = None if cached_tokens is not None and cached_tokens > 0: prefill_prompt_tokens_details = {"cached_tokens": cached_tokens} - out["completion_usage"] = { - "prompt_tokens": input_tokens, - "completion_tokens": completion_tokens, - "total_tokens": input_tokens + completion_tokens, - "prompt_tokens_details": prefill_prompt_tokens_details, - } + if input_tokens is not None and completion_tokens is not None: + out["completion_usage"] = { + "prompt_tokens": input_tokens, + "completion_tokens": completion_tokens, + "total_tokens": input_tokens + completion_tokens, + "prompt_tokens_details": prefill_prompt_tokens_details, + } if metadata_uploader is not None: try: await metadata_uploader.upload_choice(output_idx, meta_info) diff --git a/components/src/dynamo/sglang/tests/test_sglang_decode_handler.py b/components/src/dynamo/sglang/tests/test_sglang_decode_handler.py index 58a7aad72a8d..1aaa5ec260b1 100644 --- a/components/src/dynamo/sglang/tests/test_sglang_decode_handler.py +++ b/components/src/dynamo/sglang/tests/test_sglang_decode_handler.py @@ -456,6 +456,31 @@ async def test_metadata_upload_normalizes_numpy_values(tmp_path): assert uploaded_array["data"] == expected.tobytes() +@pytest.mark.asyncio +async def test_process_token_stream_treats_completion_usage_as_optional(): + handler = _new_decode_handler() + + chunks = await _collect( + handler._process_token_stream( + _stream( + [ + { + "index": 0, + "output_ids": [], + "meta_info": { + "id": "request-1", + "finish_reason": {"type": "stop"}, + }, + } + ] + ), + _Context(), + ) + ) + + assert chunks == [{"index": 0, "finish_reason": "stop", "token_ids": []}] + + @pytest.mark.asyncio async def test_process_token_stream_tracks_logprobs_per_choice_index(): handler = _new_decode_handler() From 5bda06d16037689509a60ade2d9d84259bbcfc44 Mon Sep 17 00:00:00 2001 From: Kirill Shkolenko Date: Tue, 30 Jun 2026 14:33:46 +0300 Subject: [PATCH 2/3] fix(sglang): omit empty prompt token details Signed-off-by: Kirill Shkolenko --- .../request_handlers/llm/decode_handler.py | 8 +++++-- .../tests/test_sglang_decode_handler.py | 24 ++++++++++++++++++- 2 files changed, 29 insertions(+), 3 deletions(-) diff --git a/components/src/dynamo/sglang/request_handlers/llm/decode_handler.py b/components/src/dynamo/sglang/request_handlers/llm/decode_handler.py index d2ad6abee0a9..ffb3dbc80dfc 100644 --- a/components/src/dynamo/sglang/request_handlers/llm/decode_handler.py +++ b/components/src/dynamo/sglang/request_handlers/llm/decode_handler.py @@ -587,12 +587,16 @@ async def _process_token_stream( if cached_tokens is not None and cached_tokens > 0: prefill_prompt_tokens_details = {"cached_tokens": cached_tokens} if input_tokens is not None and completion_tokens is not None: - out["completion_usage"] = { + completion_usage = { "prompt_tokens": input_tokens, "completion_tokens": completion_tokens, "total_tokens": input_tokens + completion_tokens, - "prompt_tokens_details": prefill_prompt_tokens_details, } + if prefill_prompt_tokens_details is not None: + completion_usage["prompt_tokens_details"] = ( + prefill_prompt_tokens_details + ) + out["completion_usage"] = completion_usage if metadata_uploader is not None: try: await metadata_uploader.upload_choice(output_idx, meta_info) diff --git a/components/src/dynamo/sglang/tests/test_sglang_decode_handler.py b/components/src/dynamo/sglang/tests/test_sglang_decode_handler.py index 1aaa5ec260b1..932f88e1eaf6 100644 --- a/components/src/dynamo/sglang/tests/test_sglang_decode_handler.py +++ b/components/src/dynamo/sglang/tests/test_sglang_decode_handler.py @@ -471,6 +471,16 @@ async def test_process_token_stream_treats_completion_usage_as_optional(): "id": "request-1", "finish_reason": {"type": "stop"}, }, + }, + { + "index": 1, + "output_ids": [], + "meta_info": { + "id": "request-1", + "finish_reason": {"type": "stop"}, + "prompt_tokens": 2, + "completion_tokens": 3, + }, } ] ), @@ -478,7 +488,19 @@ async def test_process_token_stream_treats_completion_usage_as_optional(): ) ) - assert chunks == [{"index": 0, "finish_reason": "stop", "token_ids": []}] + assert chunks == [ + {"index": 0, "finish_reason": "stop", "token_ids": []}, + { + "index": 1, + "finish_reason": "stop", + "token_ids": [], + "completion_usage": { + "prompt_tokens": 2, + "completion_tokens": 3, + "total_tokens": 5, + }, + }, + ] @pytest.mark.asyncio From 100e5b1ae6038ee58c980f417e35334690285697 Mon Sep 17 00:00:00 2001 From: Kirill Shkolenko Date: Thu, 2 Jul 2026 09:29:26 +0300 Subject: [PATCH 3/3] style(sglang): apply black formatting Signed-off-by: Kirill Shkolenko --- .../dynamo/sglang/request_handlers/llm/decode_handler.py | 6 +++--- .../src/dynamo/sglang/tests/test_sglang_decode_handler.py | 2 +- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/components/src/dynamo/sglang/request_handlers/llm/decode_handler.py b/components/src/dynamo/sglang/request_handlers/llm/decode_handler.py index ffb3dbc80dfc..23c04210acc0 100644 --- a/components/src/dynamo/sglang/request_handlers/llm/decode_handler.py +++ b/components/src/dynamo/sglang/request_handlers/llm/decode_handler.py @@ -593,9 +593,9 @@ async def _process_token_stream( "total_tokens": input_tokens + completion_tokens, } if prefill_prompt_tokens_details is not None: - completion_usage["prompt_tokens_details"] = ( - prefill_prompt_tokens_details - ) + completion_usage[ + "prompt_tokens_details" + ] = prefill_prompt_tokens_details out["completion_usage"] = completion_usage if metadata_uploader is not None: try: diff --git a/components/src/dynamo/sglang/tests/test_sglang_decode_handler.py b/components/src/dynamo/sglang/tests/test_sglang_decode_handler.py index 932f88e1eaf6..dcfa80a90766 100644 --- a/components/src/dynamo/sglang/tests/test_sglang_decode_handler.py +++ b/components/src/dynamo/sglang/tests/test_sglang_decode_handler.py @@ -481,7 +481,7 @@ async def test_process_token_stream_treats_completion_usage_as_optional(): "prompt_tokens": 2, "completion_tokens": 3, }, - } + }, ] ), _Context(),