From 282c46133f66d1ae9c2021cba04dc7541526b978 Mon Sep 17 00:00:00 2001 From: Liangsheng Yin Date: Sat, 30 May 2026 02:58:13 -0700 Subject: [PATCH] [Scheduler] Defer prefill input_ids H2D to forward stream, unify resolve via future_map (#25945) --- .../decode_schedule_batch_mixin.py | 8 +- python/sglang/srt/managers/overlap_utils.py | 85 ++++++++---------- python/sglang/srt/managers/schedule_batch.py | 71 ++++++++------- python/sglang/srt/managers/scheduler.py | 87 +++++++++++++------ .../sglang/srt/managers/scheduler_pp_mixin.py | 4 + .../sglang/srt/speculative/eagle_info_v2.py | 16 ++-- python/sglang/srt/speculative/eagle_utils.py | 2 +- python/sglang/srt/utils/common.py | 19 ++-- 8 files changed, 166 insertions(+), 126 deletions(-) diff --git a/python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py b/python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py index 5ab9edcb6..95e663128 100644 --- a/python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py +++ b/python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py @@ -151,7 +151,7 @@ class ScheduleBatchDisaggregationDecodeMixin: if spec_info is not None: self.spec_info = spec_info else: - # Non-spec: positive last token feeds decode directly. No FutureMap - # bootstrap needed (SB self-maintains seq_lens; resolve_future is - # a no-op on positive input_ids). - self.input_ids = last_tokens_tensor + # Non-spec: stash last token into the relay so the first DECODE's + # resolve_forward_inputs gathers it like any other decode iter. + future_map.stash(self.req_pool_indices, last_tokens_tensor) + self.input_ids = None diff --git a/python/sglang/srt/managers/overlap_utils.py b/python/sglang/srt/managers/overlap_utils.py index c4fff2fbd..23b920a77 100644 --- a/python/sglang/srt/managers/overlap_utils.py +++ b/python/sglang/srt/managers/overlap_utils.py @@ -1,10 +1,10 @@ from __future__ import annotations -import os from typing import TYPE_CHECKING, Optional, Sequence, Union import torch +from sglang.srt.environ import envs from sglang.srt.speculative.spec_utils import spec_need_hidden_states from sglang.srt.utils import is_cuda, is_hip, is_npu @@ -46,7 +46,7 @@ _is_npu = is_npu() # Token-buf consume tracking: init to -1, assert non-negative on gather, # write -1 back. Catches "gather without intermediate stash" bugs. CI enables # via the existing SGLANG_IS_IN_CI; off in production. -_DEBUG_ASSERT = os.getenv("SGLANG_IS_IN_CI", "").lower() == "true" +_DEBUG_ASSERT = envs.SGLANG_IS_IN_CI.get() @torch.compile(dynamic=True, disable=_is_npu) @@ -78,34 +78,42 @@ def _gather_spec_extras( return topk_p, topk_index, bonus_tokens, hidden_states -def _resolve_future_token_ids_native(input_ids, future_token_ids_map): - input_ids[:] = torch.where( - input_ids < 0, - future_token_ids_map[torch.clamp(-input_ids, min=0)], - input_ids, - ) +def resolve_forward_inputs(batch: ScheduleBatch, future_map: FutureMap) -> None: + """Materialize input_ids at forward entry. Two sources: + - Prefill: H2D copy from pinned CPU staging (prefill_input_ids_cpu). + - Decode/spec_v2: gather from FutureMap (last iter's sampled token). + """ + if batch.prefill_input_ids_cpu is not None: + prefill_gpu = batch.prefill_input_ids_cpu.to(batch.device, non_blocking=True) + if batch.mix_running_indices is not None: + decode_gpu = future_map.output_tokens_buf[batch.mix_running_indices] + if _DEBUG_ASSERT: + _assert_nonneg_and_invalidate( + decode_gpu, + future_map.output_tokens_buf, + batch.mix_running_indices, + ) + batch.input_ids = torch.cat([prefill_gpu, decode_gpu]) + else: + batch.input_ids = prefill_gpu + batch.prefill_input_ids_cpu = None + batch.mix_running_indices = None + elif batch.input_ids is None and future_map.spec_algo.is_none(): + batch.input_ids = future_map.output_tokens_buf[batch.req_pool_indices] + if _DEBUG_ASSERT: + _assert_nonneg_and_invalidate( + batch.input_ids, future_map.output_tokens_buf, batch.req_pool_indices + ) -if _is_cuda or _is_hip: - from sglang.jit_kernel.resolve_future_token_ids import ( - resolve_future_token_ids_cuda, - ) - - _resolve_future_token_ids = resolve_future_token_ids_cuda -else: - _resolve_future_token_ids = _resolve_future_token_ids_native + # spec_v1 (non-overlap spec) doesn't relay extras; only spec_v2 does. + if batch.is_spec_v2: + future_map._resolve_spec_extras(batch) class FutureMap: - """Cross-iter relay buffer for values the next iter's schedule cannot - compute locally (e.g. spec_v2 seq_lens after accept_lens, sampled tokens). - - Forward stream publishes into a buf; next iter's schedule pulls lazily. - Schedule-deterministic values (e.g. non-spec seq_lens via +1) stay - maintained by SB directly and do not need the relay. - - SB.seq_lens GPU is always a faithful seq_lens_cpu mirror; forward path - treats it as read-only, spec mutations land on forward_batch.seq_lens. + """Always-on pool-indexed relay for cross-iter values. Forward writes via + publish/stash; next iter reads via resolve_forward_inputs / resolve_seq_lens_cpu. """ def __init__( @@ -174,19 +182,6 @@ class FutureMap: device=self.device, ) - def resolve_future(self, batch: ScheduleBatch): - # seq_lens is already real on entry (SB +1 for non-spec; - # resolve_seq_lens_cpu pulled from buf for spec_v2). Only resolve - # input_ids tokens / spec extras here. - if self.spec_algo.is_none(): - _resolve_future_token_ids(batch.input_ids, self.output_tokens_buf) - if _DEBUG_ASSERT: - _assert_nonneg_and_invalidate( - batch.input_ids, self.output_tokens_buf, batch.req_pool_indices - ) - else: - self._resolve_spec_extras(batch) - def _resolve_spec_extras(self, batch: ScheduleBatch) -> None: draft_input: EagleDraftInput = batch.spec_info if draft_input is None: @@ -218,14 +213,6 @@ class FutureMap: draft_input.bonus_tokens, self.output_tokens_buf, indices ) - def set_input_ids_sentinel( - self, batch: ScheduleBatch, future_indices: torch.Tensor - ) -> None: - # Sentinel for the decode portion so mixed batches can cat extend - # (positive real tokens) + decode (negative sentinels) into one - # input_ids; resolve_future translates negatives via output_tokens_buf. - batch.input_ids = -future_indices - def resolve_seq_lens_cpu(self, batch: ScheduleBatch) -> None: # seq_lens_cpu may be needed on the host for kernel-launch prep (some backends). # Run this D2H on a standalone stream to avoid chain-blocking forward_n -> @@ -285,7 +272,11 @@ class FutureMap: if indices.shape[0] == 0: # DP idle: payload is empty stub; lazy-init shape peek would IndexError. return - if self.spec_algo.is_none(): + # Dispatch by payload type, not spec_algo: spec_v1 (non-overlap spec) + # also passes a token Tensor here. + # FIXME(lsyin): unify this relay path with a dataclass instead of the + # Tensor / EagleDraftInput type switch. + if isinstance(payload, torch.Tensor): self.output_tokens_buf[indices] = payload.to(torch.int64) return diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 53f593988..754793141 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -4,7 +4,7 @@ from sglang.srt.dllm.config import DllmConfig from sglang.srt.model_executor.forward_batch_info import ForwardBatch from sglang.srt.utils.common import ( ceil_align, - flatten_arrays_to_int64_tensor, + flatten_arrays_to_pinned_cpu, is_pin_memory_available, ) @@ -1525,6 +1525,9 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): # === GPU tensors crossing to ForwardBatch (clone targets for stream isolation) === # Batched arguments to model runner input_ids: torch.Tensor = None # shape: [b], int64 + # Staging consumed by resolve_forward_inputs (prefill H2D / mixed gather). + prefill_input_ids_cpu: Optional[torch.Tensor] = None + mix_running_indices: Optional[torch.Tensor] = None input_embeds: torch.Tensor = None # shape: [b, hidden_size], float32 # Token replacement embeddings and absolute positions (optional). @@ -1729,8 +1732,10 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): pt += req.extend_input_len - # Reassign - self.input_ids = flatten_arrays_to_int64_tensor(input_ids, self.device, _pin) + # Reassign: ED stripping rebuilds prefill_input_ids_cpu (CPU pinned); + # resolve_forward_inputs will H2D this on forward stream. self.input_ids + # stays None. + self.prefill_input_ids_cpu = flatten_arrays_to_pinned_cpu(input_ids, _pin) self.seq_lens = torch.tensor(seq_lens, dtype=torch.int64, pin_memory=_pin).to( self.device, non_blocking=True ) @@ -1833,7 +1838,9 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): ] _pin = is_pin_memory_available(self.device) - input_ids_tensor = flatten_arrays_to_int64_tensor(input_ids, self.device, _pin) + # Stay on pinned CPU; H2D is deferred to forward stream via + # resolve_forward_inputs. + pinned_input_ids = flatten_arrays_to_pinned_cpu(input_ids, _pin) seq_lens_tensor = torch.tensor(seq_lens, dtype=torch.int64, pin_memory=_pin).to( self.device, non_blocking=True ) @@ -2015,7 +2022,8 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): replace_embeds_tensor = None replace_positions_tensor = None - self.input_ids = input_ids_tensor + self.input_ids = None + self.prefill_input_ids_cpu = pinned_input_ids self.req_pool_indices = req_pool_indices_tensor self.req_pool_indices_cpu = req_pool_indices_cpu self.orig_seq_lens = orig_seq_lens_tensor @@ -2214,11 +2222,12 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): req.fill_ids = req.origin_input_ids + req.output_ids req.set_extend_input_len(1) - input_ids = torch.cat([self.input_ids, running_batch.input_ids]) + # Decode tokens of the running portion live in future_map.output_tokens_buf. + self.input_ids = None + self.mix_running_indices = running_batch.req_pool_indices out_cache_loc = torch.cat([self.out_cache_loc, running_batch.out_cache_loc]) self.merge_batch(running_batch) - self.input_ids = input_ids self.out_cache_loc = out_cache_loc # For overlap scheduler, the output_ids has one step delay @@ -2428,27 +2437,25 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): return if self.sampling_info.penalizer_orchestrator.is_required: - if self.enable_overlap: - # TODO: this can be slow, optimize this. - delayed_output_ids = torch.tensor( - [ - ( - req.output_ids[-1] - if len(req.output_ids) - else req.origin_input_ids[-1] - ) - for req in self.reqs - ], - dtype=torch.int64, - device=self.device, - ) - self.sampling_info.penalizer_orchestrator.cumulate_output_tokens( - delayed_output_ids - ) - else: - self.sampling_info.penalizer_orchestrator.cumulate_output_tokens( - self.input_ids - ) + # Under overlap batch.input_ids is just a placeholder here -- the + # real token is relayed via future_map and resolved at forward + # entry. So take the last output token from Req directly + # (origin_input_ids[-1] on the first decode, before any output). + latest_output_ids = torch.tensor( + [ + ( + req.output_ids[-1] + if len(req.output_ids) + else req.origin_input_ids[-1] + ) + for req in self.reqs + ], + dtype=torch.int64, + device=self.device, + ) + self.sampling_info.penalizer_orchestrator.cumulate_output_tokens( + latest_output_ids + ) # input_ids is set at end of previous run_batch (placeholder for # overlap; next_token_ids cast for non-overlap). @@ -2606,8 +2613,14 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): self.out_cache_loc = None # Sum is recomputed lazily by ForwardBatch.init_new. self.seq_lens_sum = None - if self.input_ids is not None: + # Cat only when both sides hold a real token tensor; otherwise drop to + # None and let resolve_forward_inputs rebuild from the merged + # req_pool_indices. Mismatch arises e.g. with spec_v1, which keeps its + # tensor while a relay-staged side is None -- there the worker rebuilds. + if self.input_ids is not None and other.input_ids is not None: self.input_ids = torch.cat([self.input_ids, other.input_ids]) + else: + self.input_ids = None # Optional under no-verify-sync; drop the mirror if either side absent. if self.seq_lens_cpu is None or other.seq_lens_cpu is None: self.seq_lens_cpu = None diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 08386e02f..14c5dac49 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -147,7 +147,10 @@ from sglang.srt.managers.io_struct import ( ) from sglang.srt.managers.load_snapshot import LoadSnapshot, create_load_snapshot_writer from sglang.srt.managers.multimodal_processor import get_mm_processor, import_processors -from sglang.srt.managers.overlap_utils import decide_needs_cpu_seq_lens +from sglang.srt.managers.overlap_utils import ( + decide_needs_cpu_seq_lens, + resolve_forward_inputs, +) from sglang.srt.managers.prefill_delayer import ( PrefillDelayer, PrefillDelayerSinglePassExecutor, @@ -1155,14 +1158,13 @@ class Scheduler( self.device_module = torch.get_device_module(self.device) if use_mlx(): - # MLX overlap scheduling uses mx.async_eval / mx.eval for - # synchronisation so no CUDA/MPS streams or FutureMap needed. + # MLX: no CUDA streams / FutureMap. self.future_map = None - # Empty result_queue is needed because idle-check references it - # when enable_overlap is True. self.result_queue: Deque = deque() return + # forward_stream_ctx / copy_stream are also used by PP (non-overlap) + # via scheduler_pp_mixin; init unconditionally to match main. self.forward_stream_ctx: CudaStreamContext = self.device_module.stream( self.forward_stream ) @@ -1171,10 +1173,7 @@ class Scheduler( self.copy_stream ) - if not self.enable_overlap: - self.future_map = None - return - + # FutureMap is always-on: input_ids relay used in both modes. # Workers not on BaseSpecWorker (e.g. FrozenKVMTPWorker) lack the # override; fall back to target-only so the helper still produces a # safe decision (no accidental opt-out for unaudited shapes). @@ -1192,6 +1191,10 @@ class Scheduler( self.req_to_token_pool, needs_cpu_seq_lens=needs_cpu_seq_lens, ) + + if not self.enable_overlap: + return + self.batch_record_buf = [None] * 2 self.batch_record_ct = 0 @@ -2371,9 +2374,12 @@ class Scheduler( batch.seq_lens_cpu = torch.tensor(seq_lens, dtype=torch.int64) batch.orig_seq_lens = torch.tensor(seq_lens, dtype=torch.int32, device=device) batch.seq_lens_sum = sum(seq_lens) - batch.input_ids = torch.tensor( + # Stash last token into relay; resolve_forward_inputs will gather. + last_tokens = torch.tensor( [r.output_ids[-1] for r in reqs], dtype=torch.int64, device=device ) + self.future_map.stash(batch.req_pool_indices, last_tokens) + batch.input_ids = None if batch.return_logprob: batch.top_logprobs_nums = [r.logprob.top_logprobs_num for r in reqs] @@ -2970,21 +2976,30 @@ class Scheduler( # no-ops (ForwardBatch.init_new lazily computes the sum). self.future_map.resolve_seq_lens_cpu(batch) - with self._overlap_forward_isolation(batch): - future_indices = batch.req_pool_indices + with self.forward_stream_ctx: + self.forward_stream.wait_stream(self.schedule_stream) + # resolve consumes SB staging (prefill_input_ids_cpu / + # mix_running_indices). Run OUTSIDE isolation so the + # snapshot captures the post-consume state — restoring + # post-forward must not un-consume staging. + resolve_forward_inputs(batch, self.future_map) - # Spec_v2 fires on_publish mid-worker (between verify and - # draft_extend) so schedule prep can overlap with draft_extend. - # Non-spec has no later work — scheduler publishes after return. - fwd_kwargs = ( - {"on_publish": partial(self.future_map.publish, future_indices)} - if batch.is_spec_v2 - else {} - ) + with self._overlap_forward_isolation(batch): + future_indices = batch.req_pool_indices + + # Spec_v2 fires on_publish mid-worker (between verify and + # draft_extend) so schedule prep can overlap with draft_extend. + # Non-spec has no later work — scheduler publishes after return. + fwd_kwargs = ( + { + "on_publish": partial( + self.future_map.publish, future_indices + ) + } + if batch.is_spec_v2 + else {} + ) - with self.forward_stream_ctx: - self.forward_stream.wait_stream(self.schedule_stream) - self.future_map.resolve_future(batch) # FIXME: pp is not compatible with overlap batch_result = self.model_worker.forward_batch_generation( batch, **fwd_kwargs @@ -3014,28 +3029,42 @@ class Scheduler( else: batch_result.future_indices = future_indices - self.future_map.set_input_ids_sentinel(batch, future_indices) + # Next-iter input_ids relayed via future_map. + batch.input_ids = None if batch.is_spec_v2: batch.spec_info = batch_result.next_draft_input batch.spec_info.future_indices = future_indices elif self.enable_pdmux and batch.forward_mode.is_split_prefill(): + resolve_forward_inputs(batch, self.future_map) batch_result = self.tp_worker.forward_batch_split_prefill(batch) if isinstance(batch_result.next_token_ids, torch.Tensor): - batch.input_ids = batch_result.next_token_ids.to(torch.int64) + self.future_map.stash( + batch.req_pool_indices, batch_result.next_token_ids + ) + batch.input_ids = None else: kwargs = ( {"pp_proxy_tensors": pp_proxy_tensors} if self.spec_algorithm.is_none() else {} ) + resolve_forward_inputs(batch, self.future_map) batch_result = self.model_worker.forward_batch_generation( batch, **kwargs ) - # PP intermediate ranks return None; DLLM returns a per-req list. - # Only the tensor case maps onto batch.input_ids as next-iter input. if isinstance(batch_result.next_token_ids, torch.Tensor): - batch.input_ids = batch_result.next_token_ids.to(torch.int64) + if self.spec_algorithm.is_none(): + # Non-spec: relay via future_map, gathered next iter. + self.future_map.stash( + batch.req_pool_indices, batch_result.next_token_ids + ) + batch.input_ids = None + else: + # Spec_v1 (non-overlap spec): worker shape doesn't match + # req_pool_indices; relay is unused (worker rebuilds input_ids + # inside verify). Keep pre-PR behavior. + batch.input_ids = batch_result.next_token_ids.to(torch.int64) self.update_cache_from_scheduler(batch, batch_result) # These 2 values are needed for processing the output, but the values can be @@ -3058,6 +3087,7 @@ class Scheduler( self.record_batch_in_overlap(batch) with self.forward_stream_ctx: self.forward_stream.wait_stream(self.schedule_stream) + resolve_forward_inputs(batch, self.future_map) pooler_output = self.tp_worker.forward_batch_embedding(batch) ret = EmbeddingBatchResult( embeddings=pooler_output.embeddings, @@ -3065,6 +3095,7 @@ class Scheduler( ) ret.copy_to_cpu() else: + resolve_forward_inputs(batch, self.future_map) pooler_output = self.tp_worker.forward_batch_embedding(batch) ret = EmbeddingBatchResult( embeddings=pooler_output.embeddings, diff --git a/python/sglang/srt/managers/scheduler_pp_mixin.py b/python/sglang/srt/managers/scheduler_pp_mixin.py index b802dda43..350790413 100644 --- a/python/sglang/srt/managers/scheduler_pp_mixin.py +++ b/python/sglang/srt/managers/scheduler_pp_mixin.py @@ -1091,6 +1091,10 @@ class SchedulerPPMixin: extend_logprob_start_len_per_req, ) = get_logprob_from_pp_outputs(pp_outputs) batch.input_ids = pp_outputs["next_token_ids"].to(torch.int64) + # PP rank 0 also relays into output_tokens_buf so the next iter's + # resolve_forward_inputs finds these tokens for the decode portion + # of mixed-chunk batches (which gather via mix_running_indices). + self.future_map.stash(batch.req_pool_indices, batch.input_ids) output_result = GenerationBatchResult( logits_output=logits_output, pp_hidden_states_proxy_tensors=None, diff --git a/python/sglang/srt/speculative/eagle_info_v2.py b/python/sglang/srt/speculative/eagle_info_v2.py index fb0c91b68..88e0bd03b 100644 --- a/python/sglang/srt/speculative/eagle_info_v2.py +++ b/python/sglang/srt/speculative/eagle_info_v2.py @@ -189,7 +189,7 @@ class EagleDraftInputV2Mixin: batch.out_cache_loc = torch.empty( (bs * topk * num_steps,), dtype=torch.int64, - device=batch.input_ids.device, + device=batch.device, ) # FIXME(lsyin): align with the default code path assign_draft_cache_locs_page_size_1[(bs,)]( @@ -294,7 +294,7 @@ class EagleVerifyInputV2Mixin: batch.model_config.vocab_size, "v2 prepare_for_verify input_ids", ) - device = batch.input_ids.device + device = batch.device batch.out_cache_loc = assign_extend_cache_locs_func( req_pool_indices=batch.req_pool_indices, req_to_token=req_to_token_pool.req_to_token, @@ -355,20 +355,16 @@ class EagleVerifyInputV2Mixin: Verify and find accepted tokens based on logits output and batch (which contains spec decoding information). """ + device = batch.device if batch.forward_mode.is_idle(): - predict = torch.empty(0, dtype=torch.int32, device=batch.input_ids.device) - num_correct_drafts = torch.empty( - 0, dtype=torch.int32, device=batch.input_ids.device - ) - accept_index = torch.empty( - 0, dtype=torch.int32, device=batch.input_ids.device - ) + predict = torch.empty(0, dtype=torch.int32, device=device) + num_correct_drafts = torch.empty(0, dtype=torch.int32, device=device) + accept_index = torch.empty(0, dtype=torch.int32, device=device) return predict, num_correct_drafts, accept_index bs = len(batch.seq_lens) sampling_info = batch.sampling_info next_token_logits = logits_output.next_token_logits - device = batch.input_ids.device # Apply penalty # This is a relaxed version of penalties for speculative decoding. diff --git a/python/sglang/srt/speculative/eagle_utils.py b/python/sglang/srt/speculative/eagle_utils.py index a350bb2c4..97b2f0fc3 100644 --- a/python/sglang/srt/speculative/eagle_utils.py +++ b/python/sglang/srt/speculative/eagle_utils.py @@ -61,7 +61,7 @@ def apply_eagle_prefill_input_rotation( return assert len(next_token_ids) == len(batch.seq_lens) extend_lens = torch.tensor( - batch.extend_lens, dtype=torch.int64, device=batch.input_ids.device + batch.extend_lens, dtype=torch.int64, device=batch.device ) seg_ends = extend_lens.cumsum(0) - 1 rotated = torch.empty_like(batch.input_ids) diff --git a/python/sglang/srt/utils/common.py b/python/sglang/srt/utils/common.py index f677b3ea5..0eff1209e 100644 --- a/python/sglang/srt/utils/common.py +++ b/python/sglang/srt/utils/common.py @@ -103,19 +103,24 @@ logger = logging.getLogger(__name__) torch_release = pkg_version.parse(torch.__version__).release -def flatten_arrays_to_int64_tensor( - parts: List[array[int]], device, pin: bool -) -> torch.Tensor: - """Flatten a list of array.array('q') buffers into one int64 tensor. +def flatten_arrays_to_pinned_cpu(parts: List[array[int]], pin: bool) -> torch.Tensor: + """Flatten array.array('q') buffers into one int64 CPU tensor. - Uses NumPy here to speed up the conversion by using memcpy - instead of a per-element PyLong-to-int64 walk. + NumPy memcpy instead of a per-element PyLong-to-int64 walk. Stays on + (optionally pinned) CPU; H2D is the caller's job. """ combined = np.concatenate([np.frombuffer(p, dtype=np.int64) for p in parts]) cpu_t = torch.from_numpy(combined) if pin: cpu_t = cpu_t.pin_memory() - return cpu_t.to(device, non_blocking=True) + return cpu_t + + +def flatten_arrays_to_int64_tensor( + parts: List[array[int]], device, pin: bool +) -> torch.Tensor: + """Flatten a list of array.array('q') buffers into one int64 tensor on `device`.""" + return flatten_arrays_to_pinned_cpu(parts, pin).to(device, non_blocking=True) # https://pytorch.org/docs/stable/notes/hip.html#checking-for-hip