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
+12
View File
@@ -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.