A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 07:15 UTC
create util function for batched arange (#18937)
Yu Guo committed
May 30, 2025 at 22:50 UTC
7782464a1714f6081ca06f47b75e824b14316c72
1 file changed
+35
-29
vllm/v1/worker/gpu_model_runner.py
+35
-29
@@ -500,6 +500,26 @@ class GPUModelRunner(LoRAModelRunnerMixin):
500
if batch_changed or batch_reordered:
501
self.input_batch.refresh_sampling_metadata()
502
503
+ def _get_cumsum_and_arange(
504
+ self,
505
+ num_tokens: np.ndarray,
506
+ cumsum_dtype: Optional[np.dtype] = None,
507
+ ) -> tuple[np.ndarray, np.ndarray]:
508
+ """Get the cumulative sum and batched arange of the given array.
509
+ # E.g., [2, 5, 3] -> ([2, 7, 10], [0, 1, 0, 1, 2, 3, 4, 0, 1, 2])
510
+ # Equivalent to but faster than:
511
+ # np.concatenate([np.arange(n) for n in num_tokens])
512
+ """
513
+ # Step 1. [2, 5, 3] -> [2, 7, 10]
514
+ cu_num_tokens = np.cumsum(num_tokens, dtype=cumsum_dtype)
515
+ total_num_tokens = cu_num_tokens[-1]
516
+ # Step 2. [2, 7, 10] -> [0, 0, 2, 2, 2, 2, 2, 7, 7, 7]
517
+ cumsums_offsets = np.repeat(cu_num_tokens - num_tokens, num_tokens)
518
+ # Step 3. [0, 1, 0, 1, 2, 3, 4, 0, 1, 2]
519
+ arange = self.arange_np[:total_num_tokens] - cumsums_offsets
520
+
521
+ return cu_num_tokens, arange
522
+
523
def _prepare_inputs(
524
self,
525
scheduler_output: "SchedulerOutput",
@@ -525,17 +545,10 @@ class GPUModelRunner(LoRAModelRunnerMixin):
545
req_indices = np.repeat(self.arange_np[:num_reqs],
546
num_scheduled_tokens)
547
528
- # Get batched arange.
529
- # E.g., [2, 5, 3] -> [0, 1, 0, 1, 2, 3, 4, 0, 1, 2]
530
- # Equivalent to but faster than:
531
- # np.concatenate([np.arange(n) for n in num_scheduled_tokens])
532
- # Step 1. [2, 5, 3] -> [2, 7, 10]
533
- cu_num_tokens = np.cumsum(num_scheduled_tokens)
534
- # Step 2. [2, 7, 10] -> [0, 0, 2, 2, 2, 2, 2, 7, 7, 7]
535
- cumsums_offsets = np.repeat(cu_num_tokens - num_scheduled_tokens,
536
- num_scheduled_tokens)
537
- # Step 3. [0, 1, 0, 1, 2, 3, 4, 0, 1, 2]
538
- arange = self.arange_np[:total_num_scheduled_tokens] - cumsums_offsets
548
+ # cu_num_tokens: [2, 5, 3] -> [2, 7, 10]
549
+ # arange: [0, 1, 0, 1, 2, 3, 4, 0, 1, 2]
550
+ cu_num_tokens, arange = self._get_cumsum_and_arange(
551
+ num_scheduled_tokens)
552
553
# Get positions.
554
positions_np = self.positions_np[:total_num_scheduled_tokens]
@@ -841,32 +854,25 @@ class GPUModelRunner(LoRAModelRunnerMixin):
854
# Compute the logits indices.
855
# [4, 1, 3, 1, 2]
856
num_sampled_tokens = num_draft_tokens + 1
844
- # Step 1. [4, 5, 8, 9, 11]
845
- cu_num_sampled_tokens = np.cumsum(num_sampled_tokens, dtype=np.int32)
846
- total_num_sampled_tokens = cu_num_sampled_tokens[-1]
847
- # Step 2. [0, 0, 0, 0, 4, 5, 5, 5, 8, 9, 9]
848
- cumsums_offsets = np.repeat(cu_num_sampled_tokens - num_sampled_tokens,
849
- num_sampled_tokens)
850
- # Step 3. [0, 1, 2, 3, 0, 0, 1, 2, 0, 0, 1]
851
- arange = self.arange_np[:total_num_sampled_tokens] - cumsums_offsets
852
- # Step 4. [0, 0, 0, 0, 103, 104, 104, 104, 206, 207, 207]
857
+
858
+ # Step 1. cu_num_sampled_tokens: [4, 5, 8, 9, 11]
859
+ # arange: [0, 1, 2, 3, 0, 0, 1, 2, 0, 0, 1]
860
+ cu_num_sampled_tokens, arange = self._get_cumsum_and_arange(
861
+ num_sampled_tokens, cumsum_dtype=np.int32)
862
+ # Step 2. [0, 0, 0, 0, 103, 104, 104, 104, 206, 207, 207]
863
logits_indices = np.repeat(
864
cu_num_scheduled_tokens - num_sampled_tokens, num_sampled_tokens)
855
- # Step 5. [0, 1, 2, 3, 103, 104, 105, 106, 206, 207, 208]
865
+ # Step 3. [0, 1, 2, 3, 103, 104, 105, 106, 206, 207, 208]
866
logits_indices += arange
867
868
# Compute the bonus logits indices.
869
bonus_logits_indices = cu_num_sampled_tokens - 1
870
871
# Compute the draft logits indices.
862
- # [3, 3, 5, 5, 6]
863
- cu_num_draft_tokens = np.cumsum(num_draft_tokens, dtype=np.int32)
864
- total_num_draft_tokens = cu_num_draft_tokens[-1]
865
- # [0, 0, 0, 3, 3, 5]
866
- cumsums_offsets = np.repeat(cu_num_draft_tokens - num_draft_tokens,
867
- num_draft_tokens)
868
- # [0, 1, 2, 0, 1, 0]
869
- arange = self.arange_np[:total_num_draft_tokens] - cumsums_offsets
872
+ # cu_num_draft_tokens: [3, 3, 5, 5, 6]
873
+ # arange: [0, 1, 2, 0, 1, 0]
874
+ cu_num_draft_tokens, arange = self._get_cumsum_and_arange(
875
+ num_draft_tokens, cumsum_dtype=np.int32)
876
# [0, 0, 0, 5, 5, 9]
877
target_logits_indices = np.repeat(
878
cu_num_sampled_tokens - num_sampled_tokens, num_draft_tokens)