diff --git a/python/sglang/srt/speculative/dflash_info_v2.py b/python/sglang/srt/speculative/dflash_info_v2.py index 0a22093d65e9..484d6e96f3d7 100644 --- a/python/sglang/srt/speculative/dflash_info_v2.py +++ b/python/sglang/srt/speculative/dflash_info_v2.py @@ -153,27 +153,16 @@ def prepare_for_decode(self, batch: ScheduleBatch): page_size=page_size, ) - max_top_k = 1 - uniform_top_k_value = None - uniform_top_k = True - nxt_kv_lens_sum = 0 - committed_seq_lens_sum = 0 - for i, (req, cur, nxt) in enumerate(zip(batch.reqs, cur_kv_lens, nxt_kv_lens)): - committed_len = int(req.kv.kv_committed_len) - committed_seq_lens_sum += committed_len - top_k = int(req.sampling_params.top_k) - - batch_seq_lens_cpu_t[i] = committed_len - cur_kv_lens_cpu_t[i] = cur - nxt_kv_lens_cpu_t[i] = nxt - - nxt_kv_lens_sum += nxt - if top_k > max_top_k: - max_top_k = top_k - if i == 0: - uniform_top_k_value = top_k - elif uniform_top_k and top_k != uniform_top_k_value: - uniform_top_k = False + committed_lens = [int(req.kv.kv_committed_len) for req in batch.reqs] + top_ks = [int(req.sampling_params.top_k) for req in batch.reqs] + batch_seq_lens_cpu_t.numpy()[:] = committed_lens + cur_kv_lens_cpu_t.numpy()[:] = cur_kv_lens + nxt_kv_lens_cpu_t.numpy()[:] = nxt_kv_lens + committed_seq_lens_sum = sum(committed_lens) + nxt_kv_lens_sum = sum(nxt_kv_lens) + max_top_k = max(top_ks) + uniform_top_k = top_ks.count(top_ks[0]) == bs + uniform_top_k_value = top_ks[0] self.max_top_k = max(max_top_k, 1) self.uniform_top_k_value = uniform_top_k_value if uniform_top_k else None