perf: fill steady-state decode input ids via d2d copy

- add InferenceWorkspace.fill_input_ids_from_device copying device tokens straight into the fixed-address input_ids buffer
- cache each decode step's sampled tokens on-device in DecodeSteadyState.last_tokens; when the task signature is unchanged the next step reuses them, replacing the tolist -> python list -> elementwise host fill -> pageable h2d round-trip
- _sample_logits returns (host payload, device tokens); prefill discards the device tensor
- signature change (task join/leave/first decode) still takes the host path; both dispatch paths covered by tests

Benchmark: NVIDIA L20, BF16, 1B model + 0.11B test model (4 layers, hidden 512), contiguous KV cache, CUDA Graph, greedy, prompt 512, generation 256, engine decode via scripts/tools/benchmark.py (alternating A/B, 2-4 paired runs)
- 0.11B batch 32: 21429 -> 24415 tok/s mean (1.14x, +13.9%), 4/4 paired runs faster
- 1B batch 32: 4242 -> 4388 tok/s (1.034x, +3.4%), 7.54 -> 7.29 ms/step
- batch 1: no measurable change (<0.5%)

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(
+        # When the previous decode step ran this same ordered task set, its
-            [t.output_ids[-1] if t.output_ids else t.prompt_ids[-1] for t in tasks]
+        # 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
-        task_ids = [t.task_id for t in tasks]
+        # inference-mode context.
-        cur_positions = [t.next_pos for t in tasks]
+        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 sig_match
-        reuse_decode_state = (
-            self.task_cache.bind_was_steady
-            and self._decode_cache is not None
-            and self._decode_cache.task_sig == task_sig
-        )
         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

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.
 

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