Simplify stream_output (#2398)

2024-12-08 12:27:13 -08:00
parent f62055b528
commit a6ca736c8e
9 changed files with 426 additions and 290 deletions
--- a/python/sglang/srt/model_executor/cuda_graph_runner.py
+++ b/python/sglang/srt/model_executor/cuda_graph_runner.py
@@ -400,9 +400,14 @@ class CudaGraphRunner:
                    forward_mode=ForwardMode.DECODE,
                    top_logprobs_nums=forward_batch.top_logprobs_nums,
                )
-                logits_output.output_top_logprobs = LogitsProcessor.get_top_logprobs(
+                (
+                    logits_output.output_top_logprobs_val,
+                    logits_output.output_top_logprobs_idx,
+                ) = LogitsProcessor.get_top_logprobs(
                    next_token_logprobs, logits_metadata
-                )[1]
+                )[
+                    2:4
+                ]
        else:
            logits_output = LogitsProcessorOutput(
                next_token_logits=next_token_logits,