From 06694071c6552383af9dce3f55e805d394ec4438 Mon Sep 17 00:00:00 2001 From: paulzhang-tm Date: Wed, 26 Aug 2026 17:51:31 -0400 Subject: [PATCH] [Spec] Avoid tensor scalar reads in spec decode allocation (#35377) --- python/sglang/srt/mem_cache/allocation.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/python/sglang/srt/mem_cache/allocation.py b/python/sglang/srt/mem_cache/allocation.py index 259be8837..af5038ec1 100644 --- a/python/sglang/srt/mem_cache/allocation.py +++ b/python/sglang/srt/mem_cache/allocation.py @@ -706,5 +706,6 @@ def alloc_for_spec_decode( len(reqs), ) - for i, req in enumerate(reqs): - req.kv.kv_allocated_len = max(req.kv.kv_allocated_len, int(nxt_kv_lens_cpu[i])) + nxt_kv_lens_list = nxt_kv_lens_cpu.tolist() + for req, nxt_kv_len in zip(reqs, nxt_kv_lens_list, strict=True): + req.kv.kv_allocated_len = max(req.kv.kv_allocated_len, nxt_kv_len)