[mem_cache] Build empty-prefix last_loc sentinel on-device to avoid per-call H2D sync (#32575)

Signed-off-by: xingyuliu <charlotteliu12x@gmail.com>
2026-08-03 18:55:03 -07:00
parent a84e70eb1e
commit 572924634b
1 changed files with 1 additions and 1 deletions
@@ -352,7 +352,7 @@ def alloc_for_extend(
    else:
        # Paged allocation - build last_loc
        last_loc = [
-            (t[-1:] if len(t) > 0 else torch.tensor([-1], device=batch.device))
+            (t[-1:] if len(t) > 0 else torch.full((1,), -1, device=batch.device))
            for t in prefix_tensors
        ]
        out_cache_loc = alloc_paged_token_slots_extend(