perf: reduce remaining per-step allocations

- hoist prefill qo_indptr into the workspace so CudaBackend.fwd_prefill does not rebuild it per layer
- cache has_freq in SamplingBatchInfo to drop the per-step GPU any() sync
- drop pin_memory host staging for input_ids; sync copy suffices for a small batch
This commit is contained in:
2026-08-03 01:10:06 +08:00
parent a03504a280
commit d0e5d910de
4 changed files with 31 additions and 14 deletions
+7 -4
View File
@@ -28,10 +28,14 @@ class SamplingBatchInfo:
top_ks: Tensor # int32 [B]
top_ps: Tensor # float32 [B]
freq_penalties: Tensor # float32 [B]
has_freq: bool # any frequency_penalty != 0 (avoids per-step GPU .any())
def _build_sampling_batch_info(tasks: List[Task], device) -> SamplingBatchInfo:
pin = str(device).startswith("cuda")
freq_penalties = torch.tensor(
[t.frequency_penalty for t in tasks], dtype=torch.float32, pin_memory=pin
).to(device, non_blocking=True)
return SamplingBatchInfo(
temperatures=torch.tensor(
[t.temperature for t in tasks], dtype=torch.float32, pin_memory=pin
@@ -42,9 +46,8 @@ def _build_sampling_batch_info(tasks: List[Task], device) -> SamplingBatchInfo:
top_ps=torch.tensor(
[t.top_p for t in tasks], dtype=torch.float32, pin_memory=pin
).to(device, non_blocking=True),
freq_penalties=torch.tensor(
[t.frequency_penalty for t in tasks], dtype=torch.float32, pin_memory=pin
).to(device, non_blocking=True),
freq_penalties=freq_penalties,
has_freq=bool((freq_penalties != 0).any()),
)
@@ -164,7 +167,7 @@ class Executor:
total_len = max(t.next_pos for t in tasks) + 1
input_mask = self._workspace.decode_mask(position_ids, total_len)
has_freq = bool((info.freq_penalties != 0).any())
has_freq = info.has_freq
if has_freq:
history_lists = []
history_lens = []