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 7, 06:30 UTC
[Model Runner V2] Misc code simplification (#35941)
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Nick Hill committed
Mar 4, 2026 at 15:26 UTC
a3299c3d1d6c260c35a866599bdf4d3e7b7d84dd
3 files changed
+7
-22
vllm/v1/worker/gpu/buffer_utils.py
-1
@@ -85,7 +85,6 @@ class UvaBackedTensor:
85
self, size: int | Sequence[int], dtype: torch.dtype, max_concurrency: int = 2
86
):
87
self.dtype = dtype
88
- self.max_concurrency = max_concurrency
88
89
# Source of truth
90
self.cpu = torch.zeros(size, dtype=dtype, device="cpu", pin_memory=False)
vllm/v1/worker/gpu/model_runner.py
+5
-10
@@ -96,11 +96,7 @@ logger = init_logger(__name__)
96
97
98
class GPUModelRunner(LoRAModelRunnerMixin):
99
- def __init__(
100
- self,
101
- vllm_config: VllmConfig,
102
- device: torch.device,
103
- ):
99
+ def __init__(self, vllm_config: VllmConfig, device: torch.device):
100
self.vllm_config = vllm_config
101
self.model_config = vllm_config.model_config
102
self.cache_config = vllm_config.cache_config
@@ -627,9 +623,10 @@ class GPUModelRunner(LoRAModelRunnerMixin):
623
num_reqs, dtype=torch.int32, device=self.device
624
)
625
else:
630
- num_draft_tokens = np.array(
631
- [len(draft_tokens.get(req_id, ())) for req_id in req_ids],
626
+ num_draft_tokens = np.fromiter(
627
+ (len(draft_tokens.get(req_id, ())) for req_id in req_ids),
628
dtype=np.int32,
629
+ count=num_reqs,
630
)
631
total_num_draft_tokens = int(num_draft_tokens.sum())
632
total_num_logits = num_reqs + total_num_draft_tokens
@@ -782,9 +779,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
779
780
if input_batch.num_draft_tokens == 0:
781
# No draft tokens (common case).
785
- num_sampled = torch.ones(
786
- input_batch.num_reqs, dtype=torch.int32, device=self.device
787
- )
782
+ num_sampled = input_batch.seq_lens.new_ones(input_batch.num_reqs)
783
else:
784
# Rejection sampling for spec decoding.
785
sampled_tokens, num_sampled = rejection_sample(
vllm/v1/worker/gpu/spec_decode/rejection_sample.py
+2
-11
@@ -48,17 +48,8 @@ def rejection_sample(
48
num_speculative_steps: int,
49
) -> tuple[torch.Tensor, torch.Tensor]:
50
num_reqs = cu_num_logits.shape[0] - 1
51
- sampled = torch.empty(
52
- num_reqs,
53
- num_speculative_steps + 1,
54
- dtype=target_sampled.dtype,
55
- device=target_sampled.device,
56
- )
57
- num_sampled = torch.empty(
58
- num_reqs,
59
- dtype=torch.int32,
60
- device=target_sampled.device,
61
- )
51
+ sampled = target_sampled.new_empty(num_reqs, num_speculative_steps + 1)
52
+ num_sampled = cu_num_logits.new_empty(num_reqs)
53
_rejection_sample_kernel[(num_reqs,)](
54
sampled,
55
sampled.stride(0),