diff --git a/astrai/inference/runtime/executor.py b/astrai/inference/runtime/executor.py index 7d4b611..612a26a 100644 --- a/astrai/inference/runtime/executor.py +++ b/astrai/inference/runtime/executor.py @@ -67,11 +67,15 @@ class DecodeSteadyState: When the same ordered task set decodes one token per step, sampling params and task signature are reused; only positions advance by 1. + ``last_tokens`` keeps that step's sampled ids on-device so the next + step with an unchanged signature can fill ``input_ids`` via a + device-to-device copy. """ task_sig: tuple positions: list[int] sampling_info: SamplingBatchInfo + last_tokens: Optional[Tensor] = None def _build_sampling_batch_info(tasks: List[Task], device) -> SamplingBatchInfo: @@ -250,6 +254,13 @@ class Executor: return_logprobs: bool = False, info: Optional[SamplingBatchInfo] = None, ): + """Sample from ``logits`` and return ``(host_payload, tokens)``. + + ``host_payload`` is the scheduler-facing list (token ids, or + ``(token_id, logprob)`` tuples with ``return_logprobs``); + ``tokens`` is the ``[B]`` device tensor that produced it, kept + for the steady-state decode fast path. + """ info = info or _build_sampling_batch_info(tasks, self.device) if info.has_freq: history_lists = [ @@ -284,14 +295,14 @@ class Executor: return_logprobs=return_logprobs, ) if not return_logprobs: - return result.tolist() + return result.tolist(), result tokens, logprobs = result tokens_list = tokens.tolist() logprobs_list = logprobs.tolist() for task, logprob in zip(tasks, logprobs_list): task.output_logprobs.append(float(logprob)) - return list(zip(tokens_list, logprobs_list)) + return list(zip(tokens_list, logprobs_list)), tokens def execute_prefill( self, @@ -336,7 +347,8 @@ class Executor: torch.arange(1, batch_sz + 1, device=self.device) * q_len - 1 ] - return tasks, self._sample_logits(logits, tasks, return_logprobs) + step_out, _ = self._sample_logits(logits, tasks, return_logprobs) + return tasks, step_out def execute_decode( self, tasks: List[Task], return_logprobs: bool = False @@ -360,24 +372,30 @@ class Executor: b = len(tasks) ws = self._workspace + task_ids = [t.task_id for t in tasks] + cur_positions = [t.next_pos for t in tasks] + task_sig = tuple(task_ids) # ---- pre-replay: update input buffers in-place ---- - input_ids = ws.fill_input_ids( - [t.output_ids[-1] if t.output_ids else t.prompt_ids[-1] for t in tasks] - ) - - task_ids = [t.task_id for t in tasks] - cur_positions = [t.next_pos for t in tasks] + # When the previous decode step ran this same ordered task set, its + # sampled tokens are still on-device and map 1:1 onto the current + # slots — fill input ids device-to-device. inference_mode guards + # the read because the source was produced under sampling's + # inference-mode context. + cached = self._decode_cache + sig_match = cached is not None and cached.task_sig == task_sig + if sig_match and cached.last_tokens is not None: + with torch.inference_mode(): + input_ids = ws.fill_input_ids_from_device(cached.last_tokens) + else: + input_ids = ws.fill_input_ids( + [t.output_ids[-1] if t.output_ids else t.prompt_ids[-1] for t in tasks] + ) kv_cache = self.task_cache.bind(task_ids, ws) - task_sig = tuple(task_ids) - reuse_decode_state = ( - self.task_cache.bind_was_steady - and self._decode_cache is not None - and self._decode_cache.task_sig == task_sig - ) + reuse_decode_state = self.task_cache.bind_was_steady and sig_match if reuse_decode_state: info = self._decode_cache.sampling_info ws.position_ids[:b] += 1 @@ -418,4 +436,8 @@ class Executor: ) logits = outputs["logits"] - return self._sample_logits(logits, tasks, return_logprobs, info=info) + step_out, tokens_dev = self._sample_logits( + logits, tasks, return_logprobs, info=info + ) + self._decode_cache.last_tokens = tokens_dev + return step_out diff --git a/astrai/inference/workspace.py b/astrai/inference/workspace.py index 6e4cc76..0c7a110 100644 --- a/astrai/inference/workspace.py +++ b/astrai/inference/workspace.py @@ -139,6 +139,18 @@ class InferenceWorkspace: self.input_ids[:b].copy_(pin[:b]) return self.input_ids[:b] + def fill_input_ids_from_device(self, tokens: Tensor) -> Tensor: + """Copy device-resident ``[B]`` token ids into the device buffer. + + Steady-state decode fast path: when the executor's cached task + signature still matches, the previous step's sampled tokens map + 1:1 onto the current slots, so the ids transfer device-to-device + instead of round-tripping through the host staging buffers. + """ + b = tokens.size(0) + self.input_ids[:b].copy_(tokens) + 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. diff --git a/tests/inference/test_scheduler.py b/tests/inference/test_scheduler.py index 81c3626..e25220e 100644 --- a/tests/inference/test_scheduler.py +++ b/tests/inference/test_scheduler.py @@ -393,7 +393,9 @@ def test_decode_does_not_reuse_previous_batch_state(): old_info = object() new_info = object() executor._decode_cache = DecodeSteadyState(("old",), [2], old_info) - executor._sample_logits = MagicMock(return_value=[3]) + executor._sample_logits = MagicMock( + return_value=([3], torch.tensor([3], dtype=torch.long)) + ) task = Task("new", list(range(8)), temperature=0) task.input_tokens = 8 @@ -412,3 +414,46 @@ def test_decode_does_not_reuse_previous_batch_state(): args, kwargs = executor._sample_logits.call_args assert args[1:] == ([task], False) assert kwargs["info"] is new_info + + +def test_decode_fills_input_ids_from_device_on_matching_signature(): + """Steady-state decode copies cached device tokens, skipping the host.""" + executor = object.__new__(Executor) + executor.device = torch.device("cpu") + executor.task_cache = MagicMock() + executor.task_cache.bind_was_steady = True + executor.task_cache.bind.return_value = MagicMock() + executor._graph_supported = False + executor._graph_ctx = SimpleNamespace(enabled=False) + + workspace = MagicMock() + workspace.position_ids = torch.tensor([2], dtype=torch.long) + workspace.fill_input_ids_from_device.return_value = torch.tensor( + [9], dtype=torch.long + ) + executor._workspace = workspace + executor.model = MagicMock( + return_value={"logits": torch.zeros(1, 1, 10, dtype=torch.float32)} + ) + + info = object() + tokens = torch.tensor([3], dtype=torch.long) + executor._decode_cache = DecodeSteadyState(("t1",), [2], info, last_tokens=tokens) + executor._sample_logits = MagicMock(return_value=([3], tokens)) + + task = Task("t1", list(range(8)), temperature=0) + task.input_tokens = 8 + task.output_ids = [7] + task.mark_prefill_done() + + with patch( + "astrai.inference.runtime.executor._build_sampling_batch_info", + return_value=info, + ): + assert executor.execute_decode([task]) == [3] + + workspace.fill_input_ids.assert_not_called() + workspace.fill_input_ids_from_device.assert_called_once_with(tokens) + assert workspace.position_ids.tolist() == [3] + assert executor._decode_cache.task_sig == ("t1",) + assert executor._decode_cache.last_tokens is tokens