From a03504a280ae27280d74c632ea8fa52be8a0a22b Mon Sep 17 00:00:00 2001
From: ViperEkura <3081035982@qq.com>
Date: Mon, 3 Aug 2026 00:55:26 +0800
Subject: [PATCH] perf: preallocate inference decode buffers

- add InferenceWorkspace with fixed-shape per-step buffers (input_ids, decode mask, KV bind metadata) for CUDA-graph capture
- bind_tasks derives seq_lens from the pool's own _task_len tracking, dropping the seq_lens parameter
- update decode metadata in-place (position_ids, seq_lens, kv_indptr) instead of re-allocating per step
- task_extend advances _task_len in contiguous mode so the pool tracks current length
- skip log_softmax when logprobs are not requested
---
 astrai/inference/core/cache.py              | 146 +++++++++++---------
 astrai/inference/core/executor.py           |  35 +++--
 astrai/inference/core/workspace.py          | 107 ++++++++++++++
 astrai/inference/sample.py                  |   2 +-
 tests/extension/test_backend_equivalence.py |  27 ++--
 tests/inference/test_cache.py               |  22 ++-
 6 files changed, 240 insertions(+), 99 deletions(-)
 create mode 100644 astrai/inference/core/workspace.py

diff --git a/astrai/inference/core/cache.py b/astrai/inference/core/cache.py
index 9624216..8f52e49 100644
--- a/astrai/inference/core/cache.py
+++ b/astrai/inference/core/cache.py
@@ -20,6 +20,8 @@ from typing import Callable, Dict, List, Optional
 import torch
 from torch import Tensor
 
+from astrai.inference.core.workspace import InferenceWorkspace
+
 
 def page_hash(token_ids: List[int], page_idx: int, page_size: int) -> int:
     start = page_idx * page_size
@@ -217,24 +219,6 @@ class KVCache:
     kv_indptr: Optional[Tensor] = None
 
 
-@dataclass
-class DecodeBindCache:
-    """Cached KV-addressing state for steady-state decode.
-
-    Valid for one ordered task set advancing every sequence by exactly one
-    token per step.  ``seq_lens`` is the Python mirror used to validate the
-    +1 progression without a GPU round-trip; on any task-set change or
-    non-monotonic seq_lens the whole entry is rebuilt.
-    """
-
-    sig: tuple
-    seq_lens: List[int]
-    req_pool_indices: Tensor
-    seq_lens_t: Tensor
-    kv_indptr: Tensor
-    inc: Tensor
-
-
 class PagePool:
     """Top-level KV cache manager.
 
@@ -305,11 +289,13 @@ class PagePool:
         self._task_pages: Dict[str, List[int]] = {}
         self._lock = threading.Lock()
 
-        # Single-slot incremental cache for steady-state decode: the same
-        # ordered task set advances every sequence by exactly one token per
-        # step, so seq_lens_t and kv_indptr can be updated in-place instead
-        # of re-allocating + re-cumsumming.  Any task-set change is a miss.
-        self._bind_cache: Optional[DecodeBindCache] = None
+        # Steady-state decode validation state: the ordered task set and its
+        # Python seq_lens mirror.  When the same set advances every sequence
+        # by exactly one token per step, bind_tasks updates the stable
+        # buffers in-place (+=1 / +=inc) instead of re-cumsumming.  Any
+        # task-set change is a miss and rebuilds.
+        self._bind_sig: Optional[tuple] = None
+        self._bind_seq_lens: Optional[List[int]] = None
 
     # ---- task lifecycle ----
 
@@ -395,33 +381,39 @@ class PagePool:
 
     def task_extend(self, task_id: str, pos: int) -> bool:
         req_idx = self._task_req.get(task_id)
-        if req_idx is None:
+        if req_idx is None or pos >= self.max_seq_len:
             return False
 
-        if self.contiguous:
-            return pos < self.max_seq_len
+        # Paged mode must also claim a physical slot for the new token;
+        # contiguous mode's block is pre-allocated so this is a no-op.
+        if not self.contiguous and not self._extend_slot(task_id, req_idx, pos):
+            return False
 
+        self._task_len[req_idx] = pos + 1
+        return True
+
+    def _extend_slot(self, task_id: str, req_idx: int, pos: int) -> bool:
+        """Allocate the physical slot for one extended token (paged mode)."""
         if self.page_size == 1:
             slots = self._alloc_tokens(1)
             if slots is None:
                 return False
             self._task_slots.setdefault(task_id, []).extend(slots)
             self._req_pool.req_to_token[req_idx, pos] = slots[0]
-        else:
-            page_idx = pos // self.page_size
-            existing = self._task_pages.get(task_id, [])
-            if page_idx >= len(existing):
-                p = self._alloc.alloc()
-                if p < 0:
-                    return False
-                existing.append(p)
-                self._task_pages[task_id] = existing
-            page_offset = pos % self.page_size
-            page = existing[page_idx]
-            token_slot = page * self.page_size + page_offset
-            self._req_pool.req_to_token[req_idx, pos] = token_slot
+            return True
 
-        self._task_len[req_idx] = pos + 1
+        page_idx = pos // self.page_size
+        existing = self._task_pages.get(task_id, [])
+        if page_idx >= len(existing):
+            p = self._alloc.alloc()
+            if p < 0:
+                return False
+            existing.append(p)
+            self._task_pages[task_id] = existing
+        page_offset = pos % self.page_size
+        page = existing[page_idx]
+        token_slot = page * self.page_size + page_offset
+        self._req_pool.req_to_token[req_idx, pos] = token_slot
         return True
 
     def task_cached(self, task_id: str) -> int:
@@ -442,43 +434,59 @@ class PagePool:
     def bind_tasks(
         self,
         task_ids: List[str],
-        seq_lens: List[int],
-        device: torch.device,
+        workspace: InferenceWorkspace,
+        device: Optional[torch.device] = None,
         start_pos: Optional[int] = None,
     ) -> KVCache:
+        if device is None:
+            device = workspace.device
         req_indices = [self._task_req[tid] for tid in task_ids]
+        # Per-request lengths come from the pool's own tracking (task_alloc
+        # sets len(prompt_ids); task_extend sets pos+1), so callers need not
+        # pass them.
+        seq_lens = [self._task_len[req_idx] for req_idx in req_indices]
+        b = len(task_ids)
         sig = tuple(task_ids)
 
-        cache = self._bind_cache
+        # Write into the caller's workspace buffers (fixed addresses, sized
+        # to max_batch/max_seq at init) — the sole owner of the per-step
+        # KV bind tensors.
+        rpi_buf = workspace.req_pool_indices
+        sl_buf = workspace.seq_lens
+        kvp_buf = workspace.kv_indptr
+        inc_buf = workspace.inc
+        ocl_buf = workspace.out_cache_loc
+
         incremental = (
             start_pos is None
-            and cache is not None
-            and cache.sig == sig
-            and len(cache.seq_lens) == len(seq_lens)
-            and all(s == p + 1 for s, p in zip(seq_lens, cache.seq_lens))
+            and self._bind_sig is not None
+            and self._bind_sig == sig
+            and self._bind_seq_lens is not None
+            and len(self._bind_seq_lens) == b
+            and all(s == p + 1 for s, p in zip(seq_lens, self._bind_seq_lens))
         )
         if incremental:
-            req_pool_indices = cache.req_pool_indices
-            seq_lens_t = cache.seq_lens_t + 1
-            kv_indptr = cache.kv_indptr + cache.inc
-            inc = cache.inc
+            # Steady-state decode: advance the stable buffers in-place.
+            # Normal-mode buffers keep ``+=`` legal regardless of whether
+            # this runs inside ``torch.inference_mode()``.
+            sl_buf[:b] += 1
+            kvp_buf[: b + 1] += inc_buf[: b + 1]
+            req_pool_indices = rpi_buf[:b]
+            seq_lens_t = sl_buf[:b]
+            kv_indptr = kvp_buf[: b + 1]
         else:
-            req_pool_indices = torch.tensor(
-                req_indices, dtype=torch.long, device=device
+            # Cold path: fill the stable buffers from fresh host tensors.
+            rpi_buf[:b].copy_(
+                torch.tensor(req_indices, dtype=torch.long, device=device)
             )
-            seq_lens_t = torch.tensor(seq_lens, dtype=torch.long, device=device)
-            kv_indptr = torch.zeros(len(seq_lens) + 1, dtype=torch.int32, device=device)
-            kv_indptr[1:] = seq_lens_t.cumsum(0).to(torch.int32)
-            inc = torch.arange(len(seq_lens) + 1, dtype=torch.int32, device=device)
-
-        self._bind_cache = DecodeBindCache(
-            sig=sig,
-            seq_lens=list(seq_lens),
-            req_pool_indices=req_pool_indices,
-            seq_lens_t=seq_lens_t,
-            kv_indptr=kv_indptr,
-            inc=inc,
-        )
+            sl_buf[:b].copy_(torch.tensor(seq_lens, dtype=torch.long, device=device))
+            kvp_buf[: b + 1].zero_()
+            kvp_buf[1 : b + 1] = sl_buf[:b].cumsum(0).to(torch.int32)
+            req_pool_indices = rpi_buf[:b]
+            seq_lens_t = sl_buf[:b]
+            kv_indptr = kvp_buf[: b + 1]
+        self._bind_sig = sig
+        self._bind_seq_lens = list(seq_lens)
 
         if start_pos is not None:
             seq_len = seq_lens[0]
@@ -487,9 +495,9 @@ class PagePool:
             ]
         else:
             write_pos = seq_lens_t - 1
-            out_cache_loc = self._req_pool.req_to_token[
-                req_pool_indices, write_pos
-            ].unsqueeze(-1)
+            loc = self._req_pool.req_to_token[req_pool_indices, write_pos].unsqueeze(-1)
+            ocl_buf[:b].copy_(loc)
+            out_cache_loc = ocl_buf[:b]
 
         return KVCache(
             k_buffer=self._storage.k_buffer,
diff --git a/astrai/inference/core/executor.py b/astrai/inference/core/executor.py
index d6ecf3d..a7e5e40 100644
--- a/astrai/inference/core/executor.py
+++ b/astrai/inference/core/executor.py
@@ -7,6 +7,7 @@ from torch import Tensor
 
 from astrai.inference.core.cache import PagePool
 from astrai.inference.core.task import Task
+from astrai.inference.core.workspace import InferenceWorkspace
 from astrai.inference.sample import sample
 from astrai.model.automodel import AutoModel
 from astrai.tokenize.tokenizer import AutoTokenizer
@@ -70,6 +71,17 @@ class Executor:
         # any task-set change is a cache miss.
         self._decode_cache: Optional[tuple] = None
 
+        # Pre-allocated fixed-shape buffers for the decode hot path
+        # (input_ids, decode mask, KV bind metadata).  Eagerly sized at init
+        # so the workspace is CUDA-graph-capture friendly — no allocation
+        # during capture.
+        self._workspace = InferenceWorkspace(
+            max_batch_size=kv_cache.max_batch_size,
+            max_seq_len=kv_cache.max_seq_len,
+            device=self.device,
+            dtype=self.dtype,
+        )
+
     def execute_prefill(self, tasks: List[Task], prompt_len: int, start_pos: int = 0):
         if start_pos >= prompt_len:
             return
@@ -99,7 +111,9 @@ class Executor:
                 input_mask=input_mask,
                 position_ids=position_ids,
                 kv_cache=self.kv_cache.bind_tasks(
-                    task_ids, [prompt_len] * batch_sz, self.device, start_pos=start_pos
+                    task_ids,
+                    self._workspace,
+                    start_pos=start_pos,
                 ),
             )
 
@@ -123,11 +137,9 @@ class Executor:
         if not tasks:
             return []
 
-        input_ids = torch.tensor(
-            [t.output_ids[-1] if t.output_ids else t.prompt_ids[-1] for t in tasks],
-            dtype=torch.long,
-            device=self.device,
-        )
+        input_ids = self._workspace.fill_input_ids(
+            [t.output_ids[-1] if t.output_ids else t.prompt_ids[-1] for t in tasks]
+        ).unsqueeze(1)
 
         task_ids = [t.task_id for t in tasks]
 
@@ -140,7 +152,7 @@ class Executor:
             and cur_positions == [p + 1 for p in cached[1]]
         ):
             _, _, info, position_ids = cached
-            position_ids = position_ids + 1
+            position_ids += 1
             self._decode_cache = (sig, cur_positions, info, position_ids)
         else:
             info = _build_sampling_batch_info(tasks, self.device)
@@ -150,9 +162,7 @@ class Executor:
             self._decode_cache = (sig, cur_positions, info, position_ids)
 
         total_len = max(t.next_pos for t in tasks) + 1
-        input_mask = position_ids[:, None, None] >= torch.arange(
-            total_len, device=self.device
-        )
+        input_mask = self._workspace.decode_mask(position_ids, total_len)
 
         has_freq = bool((info.freq_penalties != 0).any())
         if has_freq:
@@ -184,12 +194,11 @@ class Executor:
 
         with torch.inference_mode():
             outputs = self.model(
-                input_ids.unsqueeze(1),
+                input_ids,
                 input_mask=input_mask,
                 kv_cache=self.kv_cache.bind_tasks(
                     task_ids,
-                    [t.next_pos + 1 for t in tasks],
-                    self.device,
+                    self._workspace,
                 ),
                 position_ids=position_ids.unsqueeze(1),
             )
diff --git a/astrai/inference/core/workspace.py b/astrai/inference/core/workspace.py
new file mode 100644
index 0000000..fe6e0eb
--- /dev/null
+++ b/astrai/inference/core/workspace.py
@@ -0,0 +1,107 @@
+"""Pre-allocated buffers for the inference decode hot path.
+
+Mirrors SGLang's pre-allocated input buffers (``input_buffers.py``): tensors
+are sized once to the server's maximum dimensions and sliced to the live
+batch each step, so the per-token decode loop never calls
+``torch.empty``/``torch.zeros``/``torch.arange`` for the hot shapes.  Fills
+go through ``out=`` variants (``torch.ge``) which write into the stable
+buffers instead of allocating fresh results.
+
+All buffers are allocated eagerly at init (nothing is lazy), so the
+workspace is CUDA-graph-capture friendly: the decode step reads/writes
+fixed-address tensors with no allocation during capture.
+"""
+
+import torch
+from torch import Tensor
+
+
+class InferenceWorkspace:
+    """Reusable fixed-shape per-step buffers for decode.
+
+    Families of buffers, all sized to ``max_batch_size`` / ``max_seq_len``
+    and sliced via views each step:
+
+    - ``decode_mask``: a ``[B, 1, total_len]`` validity mask, the RHS
+      ``arange`` pre-computed so only a single ``torch.ge(out=)`` runs per
+      step.
+    - ``input_ids``: per-step token IDs filled from host (pinned, double-
+      buffered so an in-flight async H2D copy never races the next fill).
+    - KV-cache bind metadata (``req_pool_indices``, ``seq_lens``,
+      ``kv_indptr``, ``inc``, ``out_cache_loc``), written by
+      ``PagePool.bind_tasks`` when the Executor passes this workspace.
+
+    No re-allocation while the server's bounds are respected.
+    """
+
+    def __init__(
+        self,
+        max_batch_size: int,
+        max_seq_len: int,
+        device: torch.device,
+        dtype: torch.dtype,
+    ):
+        self.max_batch_size = max_batch_size
+        self.max_seq_len = max_seq_len
+        self.device = device
+        self.dtype = dtype
+
+        # ``position_ids[:, None, None] >= arange`` RHS, reused every step.
+        self.arange = torch.arange(max_seq_len, device=device)
+        # Decode validity mask: [max_batch, 1, max_seq_len] bool.
+        self.input_mask = torch.empty(
+            (max_batch_size, 1, max_seq_len), dtype=torch.bool, device=device
+        )
+
+        # Per-step token IDs.  Values come from host Python lists every
+        # step, so the device buffer is pre-allocated and filled via an
+        # async copy from a double-buffered pinned host buffer (stable
+        # address for CUDA-graph capture; alternating buffers keep an
+        # in-flight copy from being overwritten by the next fill).
+        self.input_ids = torch.empty((max_batch_size,), dtype=torch.long, device=device)
+        self._pin = [
+            torch.empty((max_batch_size,), dtype=torch.long, pin_memory=True),
+            torch.empty((max_batch_size,), dtype=torch.long, pin_memory=True),
+        ]
+        self._pin_idx = 0
+
+        # KV-cache bind metadata (fixed shape, written by ``PagePool.bind_tasks``
+        # when the Executor passes this workspace).  Stable addresses make the
+        # decode forward CUDA-graph capturable.
+        self.req_pool_indices = torch.empty(
+            (max_batch_size,), dtype=torch.long, device=device
+        )
+        self.seq_lens = torch.empty((max_batch_size,), dtype=torch.long, device=device)
+        self.kv_indptr = torch.empty(
+            (max_batch_size + 1,), dtype=torch.int32, device=device
+        )
+        self.inc = torch.arange(max_batch_size + 1, dtype=torch.int32, device=device)
+        self.out_cache_loc = torch.empty(
+            (max_batch_size, 1), dtype=torch.long, device=device
+        )
+
+    def fill_input_ids(self, ids: "list[int]") -> Tensor:
+        """Write ``ids`` into the device buffer and return ``[B]``.
+
+        Pinned host values are copied asynchronously; the double buffer
+        guarantees the copy never races the next call's host writes.
+        """
+        b = len(ids)
+        pin = self._pin[self._pin_idx]
+        self._pin_idx ^= 1
+        for i, v in enumerate(ids):
+            pin[i] = v
+        self.input_ids[:b].copy_(pin[:b], non_blocking=True)
+        return self.input_ids[:b]
+
+    def decode_mask(self, position_ids: Tensor, total_len: int) -> Tensor:
+        """Return the ``[B, 1, total_len]`` validity mask for this step.
+
+        Written into the pre-allocated buffer via ``torch.ge(out=)`` — no
+        new tensor is allocated.  ``position_ids`` is the current step's
+        ``[B]`` positions; ``total_len`` must not exceed ``max_seq_len``.
+        """
+        b = position_ids.size(0)
+        out = self.input_mask[:b, :, :total_len]
+        torch.ge(position_ids[:, None, None], self.arange[:total_len], out=out)
+        return out
diff --git a/astrai/inference/sample.py b/astrai/inference/sample.py
index 7e72657..36186ce 100644
--- a/astrai/inference/sample.py
+++ b/astrai/inference/sample.py
@@ -305,12 +305,12 @@ class SamplingPipeline(BaseSamplingStrategy):
             return tokens, chosen
 
         transformed = self.apply(logits, filter_value, input_ids, input_mask)
-        log_probs = torch.log_softmax(transformed.float(), dim=-1)
         tokens = torch.multinomial(
             torch.softmax(transformed, dim=-1), num_samples=1
         ).squeeze(-1)
         if not return_logprobs:
             return tokens
+        log_probs = torch.log_softmax(transformed.float(), dim=-1)
         chosen = torch.gather(log_probs, -1, tokens.unsqueeze(-1)).squeeze(-1)
         return tokens, chosen
 
diff --git a/tests/extension/test_backend_equivalence.py b/tests/extension/test_backend_equivalence.py
index 197f2be..acac283 100644
--- a/tests/extension/test_backend_equivalence.py
+++ b/tests/extension/test_backend_equivalence.py
@@ -8,9 +8,16 @@ import torch
 
 from astrai.extension import ATTN_BACKEND, attn_backend
 from astrai.inference.core.cache import PagePool
+from astrai.inference.core.workspace import InferenceWorkspace
 from tests.extension.conftest import D, skip_no_kernel
 
 
+def _ws(pool: PagePool) -> InferenceWorkspace:
+    return InferenceWorkspace(
+        pool.max_batch_size, pool.max_seq_len, pool.device, pool.dtype
+    )
+
+
 @skip_no_kernel
 def test_training_forward_matches_torch(cuda_model):
     """Training forward (kv_cache=None) should produce identical logits.
@@ -62,11 +69,10 @@ def test_prefill_with_kv_cache_matches_torch(cuda_model):
         dtype=torch.bfloat16,
     )
 
+    ws = _ws(cache)
     cache.task_alloc("t1", prompt_ids[0])
     cache.task_alloc("t2", prompt_ids[1])
-    kv1 = cache.bind_tasks(
-        ["t1", "t2"], [len(prompt_ids[0]), len(prompt_ids[1])], device, start_pos=0
-    )
+    kv1 = cache.bind_tasks(["t1", "t2"], ws, start_pos=0)
     with torch.inference_mode():
         out_torch = model(
             input_ids, input_mask=input_mask, kv_cache=kv1, position_ids=position_ids
@@ -76,9 +82,7 @@ def test_prefill_with_kv_cache_matches_torch(cuda_model):
     cache.task_free("t2")
     cache.task_alloc("t1", prompt_ids[0])
     cache.task_alloc("t2", prompt_ids[1])
-    kv2 = cache.bind_tasks(
-        ["t1", "t2"], [len(prompt_ids[0]), len(prompt_ids[1])], device, start_pos=0
-    )
+    kv2 = cache.bind_tasks(["t1", "t2"], ws, start_pos=0)
     with attn_backend(ATTN_BACKEND.CUDA):
         with torch.inference_mode():
             out_cuda = model(
@@ -129,11 +133,10 @@ def test_decode_mixed_seq_lens_matches_torch(cuda_model):
         input_mask[i, : len(p)] = True
         position_ids[i, : len(p)] = torch.arange(len(p), device=device)
 
+    ws = _ws(cache)
     cache.task_alloc("t1", prompt_ids[0])
     cache.task_alloc("t2", prompt_ids[1])
-    kv = cache.bind_tasks(
-        ["t1", "t2"], [len(prompt_ids[0]), len(prompt_ids[1])], device, start_pos=0
-    )
+    kv = cache.bind_tasks(["t1", "t2"], ws, start_pos=0)
     with torch.inference_mode():
         model(input_ids, input_mask=input_mask, kv_cache=kv, position_ids=position_ids)
 
@@ -143,13 +146,15 @@ def test_decode_mixed_seq_lens_matches_torch(cuda_model):
     total_len = 9
     dec_mask = dec_pos[:, None, None] >= torch.arange(total_len, device=device)
 
-    kv_t = cache.bind_tasks(["t1", "t2"], [9, 7], device)
+    cache.task_extend("t1", 8)
+    cache.task_extend("t2", 6)
+    kv_t = cache.bind_tasks(["t1", "t2"], ws)
     with torch.inference_mode():
         out_torch = model(
             dec_ids, input_mask=dec_mask, kv_cache=kv_t, position_ids=dec_pos
         )
 
-    kv_c = cache.bind_tasks(["t1", "t2"], [9, 7], device)
+    kv_c = cache.bind_tasks(["t1", "t2"], ws)
     with attn_backend(ATTN_BACKEND.CUDA):
         with torch.inference_mode():
             out_cuda = model(
diff --git a/tests/inference/test_cache.py b/tests/inference/test_cache.py
index 6cd08db..125cf7d 100644
--- a/tests/inference/test_cache.py
+++ b/tests/inference/test_cache.py
@@ -10,6 +10,15 @@ from astrai.inference import (
     ReqToTokenPool,
     page_hash,
 )
+from astrai.inference.core.workspace import InferenceWorkspace
+
+
+def _ws(pool: PagePool) -> InferenceWorkspace:
+    """Workspace sized to the pool (bind_tasks requires it)."""
+    return InferenceWorkspace(
+        pool.max_batch_size, pool.max_seq_len, pool.device, pool.dtype
+    )
+
 
 # ---- page_hash ----
 
@@ -216,7 +225,7 @@ def test_page_pool_contiguous_bind_tasks_prefill():
     pool = _make_contiguous_pool()
     pool.task_alloc("t1", list(range(10)))
     pool.task_alloc("t2", list(range(10)))
-    kv = pool.bind_tasks(["t1", "t2"], [10, 10], torch.device("cpu"), start_pos=0)
+    kv = pool.bind_tasks(["t1", "t2"], _ws(pool), start_pos=0)
     assert kv.out_cache_loc.shape == (2, 10)
     assert kv.seq_lens.tolist() == [10, 10]
     assert kv.req_pool_indices.shape == (2,)
@@ -226,7 +235,10 @@ def test_page_pool_contiguous_bind_tasks_decode():
     pool = _make_contiguous_pool()
     pool.task_alloc("t1", list(range(10)))
     pool.task_alloc("t2", list(range(8)))
-    kv = pool.bind_tasks(["t1", "t2"], [11, 9], torch.device("cpu"))
+    # Simulate one decode extension so seq_lens advance to 11 and 9.
+    assert pool.task_extend("t1", 10)
+    assert pool.task_extend("t2", 8)
+    kv = pool.bind_tasks(["t1", "t2"], _ws(pool))
     assert kv.out_cache_loc.shape == (2, 1)
     assert kv.seq_lens.tolist() == [11, 9]
 
@@ -236,7 +248,7 @@ def test_page_pool_contiguous_bind_roundtrip():
     pool = _make_contiguous_pool(n_layers=1, n_kv_heads=2, head_dim=4)
     pool.task_alloc("t1", list(range(4)))
 
-    kv = pool.bind_tasks(["t1"], [4], torch.device("cpu"), start_pos=0)
+    kv = pool.bind_tasks(["t1"], _ws(pool), start_pos=0)
     k = torch.randn(1, 4, 2, 4)
     v = torch.randn(1, 4, 2, 4)
     kv.k_buffer[0, kv.out_cache_loc] = k
@@ -298,7 +310,7 @@ def test_page_pool_paged_bind_roundtrip():
     pool = _make_paged_pool(n_layers=1, n_kv_heads=2, head_dim=4)
     pool.task_alloc("t1", list(range(4)))
 
-    kv = pool.bind_tasks(["t1"], [4], torch.device("cpu"), start_pos=0)
+    kv = pool.bind_tasks(["t1"], _ws(pool), start_pos=0)
     k = torch.randn(1, 4, 2, 4)
     v = torch.randn(1, 4, 2, 4)
     kv.k_buffer[0, kv.out_cache_loc] = k
@@ -349,7 +361,7 @@ def test_page_pool_paged_ps64_bind_roundtrip():
     prompt = list(range(128))
     pool.task_alloc("t1", prompt)
 
-    kv = pool.bind_tasks(["t1"], [128], torch.device("cpu"), start_pos=0)
+    kv = pool.bind_tasks(["t1"], _ws(pool), start_pos=0)
     k = torch.randn(1, 128, 2, 4)
     v = torch.randn(1, 128, 2, 4)
     kv.k_buffer[0, kv.out_cache_loc] = k