diff --git a/astrai/extension/attention_backend.py b/astrai/extension/attention_backend.py index d9df74a..d4221c2 100644 --- a/astrai/extension/attention_backend.py +++ b/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)) diff --git a/astrai/inference/core/cache.py b/astrai/inference/core/cache.py index 8f52e49..356bffb 100644 --- a/astrai/inference/core/cache.py +++ b/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 ---- diff --git a/astrai/inference/core/executor.py b/astrai/inference/core/executor.py index a7e5e40..c5b7fbe 100644 --- a/astrai/inference/core/executor.py +++ b/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( - [t.frequency_penalty for t in tasks], dtype=torch.float32, pin_memory=pin - ).to(device, non_blocking=True), + freq_penalties=freq_penalties, + has_freq=bool((freq_penalties != 0).any()), ) @@ -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 = [] diff --git a/astrai/inference/core/workspace.py b/astrai/inference/core/workspace.py index fe6e0eb..edd3959 100644 --- a/astrai/inference/core/workspace.py +++ b/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 - # 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). + # step, so the device buffer is pre-allocated (stable address for + # CUDA-graph capture) and filled via a host staging buffer. A + # double buffer keeps a copy in flight 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), + torch.empty((max_batch_size,), dtype=torch.long), + 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 - guarantees the copy never races the next call's host writes. + Host values are staged through the double buffer and copied into the + 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: