[Spec] Avoid tensor scalar reads in spec decode allocation (#35377)
This commit is contained in:
@@ -706,5 +706,6 @@ def alloc_for_spec_decode(
|
|||||||
len(reqs),
|
len(reqs),
|
||||||
)
|
)
|
||||||
|
|
||||||
for i, req in enumerate(reqs):
|
nxt_kv_lens_list = nxt_kv_lens_cpu.tolist()
|
||||||
req.kv.kv_allocated_len = max(req.kv.kv_allocated_len, int(nxt_kv_lens_cpu[i]))
|
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)
|
||||||
|
|||||||
Reference in New Issue
Block a user