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:
@@ -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.
|
||||
|
||||
|
||||
Reference in New Issue
Block a user