[Scheduler] Defer prefill input_ids H2D to forward stream, unify resolve via future_map (#25945)

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

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
 

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

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,

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,

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.

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)

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