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%)
This commit is contained in:
2026-08-31 14:24:51 +08:00
parent 36e39496d4
commit 1cf7d6c76b
3 changed files with 96 additions and 17 deletions
+46 -1
View File
@@ -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