From c8492978a17c3a426b257d22660284e848d53dcd Mon Sep 17 00:00:00 2001 From: YAMY <74099316+YAMY1234@users.noreply.github.com> Date: Sat, 25 Oct 2025 12:28:26 -0700 Subject: [PATCH] Fix Illegal Instruction/IMA errors when using DP attention -- num_tokens_for_logprob calculation (#12115) --- python/sglang/srt/managers/scheduler.py | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 293851111..c1ae06612 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -2073,15 +2073,18 @@ class Scheduler( num_tokens_for_logprob = num_tokens else: num_tokens = local_batch.extend_num_tokens - num_tokens_for_logprob = sum( - [ + if local_batch.return_logprob: + num_tokens_for_logprob = sum( # We should have at least 1 token for sample in every case. max(extend_len - logprob_start_len, 1) for logprob_start_len, extend_len in zip( - local_batch.extend_logprob_start_lens, local_batch.extend_lens + local_batch.extend_logprob_start_lens, + local_batch.extend_lens, ) - ] - ) + ) + else: + # When return_logprob = False, only need last token per request + num_tokens_for_logprob = local_batch.batch_size() if local_batch is None or local_batch.forward_mode.is_decode_or_idle(): can_cuda_graph = 1