Skip to content

Commit 5d185ef

Browse files
authored
Fix prefill stats for dllm (sgl-project#18632)
1 parent dcc63dc commit 5d185ef

File tree

1 file changed

+12
-0
lines changed

1 file changed

+12
-0
lines changed

python/sglang/srt/dllm/mixin/scheduler.py

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -192,6 +192,18 @@ def _create_dllm_batch(
192192
new_batch.prepare_for_extend()
193193
new_batch.forward_mode = forward_mode
194194
new_batch.decoding_reqs = None
195+
196+
# Record prefill stats for logging after forward
197+
from sglang.srt.managers.scheduler_metrics_mixin import PrefillStats
198+
199+
new_batch.prefill_stats = PrefillStats(
200+
log_input_tokens=self.adder.log_input_tokens,
201+
log_hit_tokens=self.adder.log_hit_tokens,
202+
new_token_ratio=self.adder.new_token_ratio,
203+
running_bs=len(self.running_batch.reqs),
204+
num_new_seqs=len(can_run_list),
205+
)
206+
195207
return new_batch
196208

197209
def process_dllm_incoming_reqs(

0 commit comments

Comments
 (0)