perf: reduce remaining per-step allocations

- hoist prefill qo_indptr into the workspace so CudaBackend.fwd_prefill does not rebuild it per layer
- cache has_freq in SamplingBatchInfo to drop the per-step GPU any() sync
- drop pin_memory host staging for input_ids; sync copy suffices for a small batch

astrai/extension/attention_backend.py

@@ -376,7 +376,7 @@ class CudaBackend(AttentionBackend):
         q_len = q.size(1)
 
         kv_indptr = kv_cache.kv_indptr
-        qo_indptr = torch.arange(b + 1, dtype=torch.int32, device=q.device) * q_len
+        qo_indptr = kv_cache.qo_indptr
 
         q_flat = q.reshape(b * q_len, q.size(2), q.size(3))
 

astrai/inference/core/cache.py

@@ -217,6 +217,7 @@ class KVCache:
     out_cache_loc: Tensor
     max_len: int = 0
     kv_indptr: Optional[Tensor] = None
+    qo_indptr: Optional[Tensor] = None
 
 
 class PagePool:
@@ -493,11 +494,19 @@ class PagePool:
             out_cache_loc = self._req_pool.req_to_token[
                 req_pool_indices, start_pos:seq_len
             ]
+            # Ragged query segmentation for the prefill kernel, computed once
+            # (was rebuilt per layer in CudaBackend.fwd_prefill).
+            q_len = seq_len - start_pos
+            workspace.qo_indptr[: b + 1].copy_(
+                torch.arange(b + 1, dtype=torch.int32, device=device) * q_len
+            )
+            qo_indptr = workspace.qo_indptr[: b + 1]
         else:
             write_pos = seq_lens_t - 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]
+            qo_indptr = None
 
         return KVCache(
             k_buffer=self._storage.k_buffer,
@@ -508,6 +517,7 @@ class PagePool:
             out_cache_loc=out_cache_loc,
             max_len=max(seq_lens),
             kv_indptr=kv_indptr,
+            qo_indptr=qo_indptr,
         )
 
     # ---- internals ----

astrai/inference/core/executor.py

@@ -28,10 +28,14 @@ class SamplingBatchInfo:
     top_ks: Tensor  # int32  [B]
     top_ps: Tensor  # float32 [B]
     freq_penalties: Tensor  # float32 [B]
+    has_freq: bool  # any frequency_penalty != 0 (avoids per-step GPU .any())
 
 
 def _build_sampling_batch_info(tasks: List[Task], device) -> SamplingBatchInfo:
     pin = str(device).startswith("cuda")
+    freq_penalties = torch.tensor(
+        [t.frequency_penalty for t in tasks], dtype=torch.float32, pin_memory=pin
+    ).to(device, non_blocking=True)
     return SamplingBatchInfo(
         temperatures=torch.tensor(
             [t.temperature for t in tasks], dtype=torch.float32, pin_memory=pin
@@ -42,9 +46,8 @@ def _build_sampling_batch_info(tasks: List[Task], device) -> SamplingBatchInfo:
         top_ps=torch.tensor(
             [t.top_p for t in tasks], dtype=torch.float32, pin_memory=pin
         ).to(device, non_blocking=True),
-        freq_penalties=torch.tensor(
+        freq_penalties=freq_penalties,
-            [t.frequency_penalty for t in tasks], dtype=torch.float32, pin_memory=pin
+        has_freq=bool((freq_penalties != 0).any()),
-        ).to(device, non_blocking=True),
     )
 
 
@@ -164,7 +167,7 @@ class Executor:
         total_len = max(t.next_pos for t in tasks) + 1
         input_mask = self._workspace.decode_mask(position_ids, total_len)
 
-        has_freq = bool((info.freq_penalties != 0).any())
+        has_freq = info.has_freq
         if has_freq:
             history_lists = []
             history_lens = []

astrai/inference/core/workspace.py

@@ -54,14 +54,14 @@ class InferenceWorkspace:
         )
 
         # Per-step token IDs.  Values come from host Python lists every
-        # step, so the device buffer is pre-allocated and filled via an
+        # step, so the device buffer is pre-allocated (stable address for
-        # async copy from a double-buffered pinned host buffer (stable
+        # CUDA-graph capture) and filled via a host staging buffer.  A
-        # address for CUDA-graph capture; alternating buffers keep an
+        # double buffer keeps a copy in flight from being overwritten by
-        # in-flight copy from being overwritten by the next fill).
+        # 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),
-            torch.empty((max_batch_size,), dtype=torch.long, pin_memory=True),
+            torch.empty((max_batch_size,), dtype=torch.long),
         ]
         self._pin_idx = 0
 
@@ -75,6 +75,9 @@ class InferenceWorkspace:
         self.kv_indptr = torch.empty(
             (max_batch_size + 1,), dtype=torch.int32, device=device
         )
+        self.qo_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
@@ -83,15 +86,16 @@ class InferenceWorkspace:
     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
+        Host values are staged through the double buffer and copied into the
-        guarantees the copy never races the next call's host writes.
+        stable device buffer (``copy_`` without pinning is synchronous, so
+        the alternating buffers guard against an in-flight transfer).
         """
         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)
+        self.input_ids[:b].copy_(pin[:b])
         return self.input_ids[:b]
 
     def decode_mask(self, position_ids: Tensor, total_len: int) -> Tensor: