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 12, 08:30 UTC
[Model Runner V2] Add sample/ directory and reorganize files (#29719)
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
Woosuk Kwon committed
Nov 29, 2025 at 00:41 UTC
6afc0ffaf61b49599dec25c30b28596a390ad2d9
10 files changed
+587
-570
vllm/v1/worker/gpu/model_runner.py
+11
-4
@@ -47,13 +47,18 @@ from vllm.v1.worker.gpu.input_batch import (
47
prepare_pos_seq_lens,
48
prepare_prefill_inputs,
49
)
50
-from vllm.v1.worker.gpu.sampler import Sampler, compute_prompt_logprobs
50
+from vllm.v1.worker.gpu.sample.logprob import compute_prompt_logprobs
51
+from vllm.v1.worker.gpu.sample.metadata import (
52
+ SamplingMetadata,
53
+ expand_sampling_metadata,
54
+)
55
+from vllm.v1.worker.gpu.sample.sampler import Sampler
56
from vllm.v1.worker.gpu.spec_decode import init_speculator
57
from vllm.v1.worker.gpu.spec_decode.rejection_sample import (
58
get_num_rejected,
59
rejection_sample,
60
)
56
-from vllm.v1.worker.gpu.states import RequestState, SamplingMetadata
61
+from vllm.v1.worker.gpu.states import RequestState
62
from vllm.v1.worker.gpu.structured_outputs import apply_grammar_bitmask
63
from vllm.v1.worker.kv_connector_model_runner_mixin import KVConnectorModelRunnerMixin
64
from vllm.v1.worker.lora_model_runner_mixin import LoRAModelRunnerMixin
@@ -890,8 +895,10 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
895
input_batch.idx_mapping, input_batch.idx_mapping_np, pos
896
)
897
if input_batch.num_draft_tokens > 0:
893
- sampling_metadata = self.req_states.expand_sampling_metadata(
894
- sampling_metadata, input_batch.cu_num_logits
898
+ sampling_metadata = expand_sampling_metadata(
899
+ sampling_metadata,
900
+ input_batch.cu_num_logits,
901
+ max_expand_len=self.num_speculative_steps + 1,
902
)
903
904
if self.lora_config:
vllm/v1/worker/gpu/sample/__init__.py
vllm/v1/worker/gpu/sample/gumbel.py
new
+100
@@ -0,0 +1,100 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import torch
4
+
5
+from vllm.triton_utils import tl, triton
6
+
7
+
8
+@triton.jit
9
+def _gumbel_sample_kernel(
10
+ local_argmax_ptr,
11
+ local_argmax_stride,
12
+ local_max_ptr,
13
+ local_max_stride,
14
+ logits_ptr,
15
+ logits_stride,
16
+ seeds_ptr,
17
+ pos_ptr,
18
+ temp_ptr,
19
+ vocab_size,
20
+ BLOCK_SIZE: tl.constexpr,
21
+ APPLY_TEMPERATURE: tl.constexpr,
22
+):
23
+ req_idx = tl.program_id(0)
24
+ block_idx = tl.program_id(1)
25
+ block = block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
26
+ mask = block < vocab_size
27
+ logits = tl.load(
28
+ logits_ptr + req_idx * logits_stride + block,
29
+ mask=mask,
30
+ other=float("-inf"),
31
+ )
32
+ logits = logits.to(tl.float32)
33
+
34
+ temp = tl.load(temp_ptr + req_idx).to(tl.float32)
35
+ if temp != 0.0:
36
+ # Calculate the seed for gumbel noise.
37
+ seed = tl.load(seeds_ptr + req_idx)
38
+ pos = tl.load(pos_ptr + req_idx)
39
+ gumbel_seed = tl.randint(seed, pos)
40
+
41
+ # Generate gumbel noise.
42
+ r = tl.rand(gumbel_seed, block).to(tl.float64)
43
+ gumbel_noise = -tl.log(-tl.log(r + 1e-20) + 1e-20)
44
+ gumbel_noise = gumbel_noise.to(tl.float32)
45
+
46
+ # Apply temperature.
47
+ if APPLY_TEMPERATURE:
48
+ # NOTE(woosuk): Use div_rn to match the behavior of torch.
49
+ logits = tl.div_rn(logits, temp)
50
+
51
+ # Apply gumbel noise.
52
+ logits = tl.where(mask, logits + gumbel_noise, float("-inf"))
53
+
54
+ idx = tl.argmax(logits, axis=0)
55
+ token_id = block_idx * BLOCK_SIZE + idx
56
+ value = tl.max(logits, axis=0)
57
+ tl.store(local_argmax_ptr + req_idx * local_argmax_stride + block_idx, token_id)
58
+ tl.store(local_max_ptr + req_idx * local_max_stride + block_idx, value)
59
+
60
+
61
+def gumbel_sample(
62
+ logits: torch.Tensor, # [num_reqs, vocab_size]
63
+ temperature: torch.Tensor, # [num_reqs]
64
+ seed: torch.Tensor, # [num_reqs]
65
+ pos: torch.Tensor, # [num_reqs]
66
+ apply_temperature: bool,
67
+) -> torch.Tensor:
68
+ num_reqs, vocab_size = logits.shape
69
+ BLOCK_SIZE = 1024
70
+ num_blocks = triton.cdiv(vocab_size, BLOCK_SIZE)
71
+ local_argmax = torch.empty(
72
+ num_reqs,
73
+ num_blocks,
74
+ dtype=torch.int64,
75
+ device=logits.device,
76
+ )
77
+ local_max = torch.empty(
78
+ num_reqs,
79
+ num_blocks,
80
+ dtype=torch.float32,
81
+ device=logits.device,
82
+ )
83
+ _gumbel_sample_kernel[(num_reqs, num_blocks)](
84
+ local_argmax,
85
+ local_argmax.stride(0),
86
+ local_max,
87
+ local_max.stride(0),
88
+ logits,
89
+ logits.stride(0),
90
+ seed,
91
+ pos,
92
+ temperature,
93
+ vocab_size,
94
+ BLOCK_SIZE=BLOCK_SIZE,
95
+ APPLY_TEMPERATURE=apply_temperature,
96
+ )
97
+ # NOTE(woosuk): Use int64 for later indexing.
98
+ max_block_idx = local_max.argmax(dim=-1, keepdim=True)
99
+ sampled = local_argmax.gather(dim=-1, index=max_block_idx).view(-1)
100
+ return sampled
vllm/v1/worker/gpu/sample/logprob.py
new
+167
@@ -0,0 +1,167 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from collections.abc import Callable
4
+
5
+import torch
6
+
7
+from vllm.triton_utils import tl, triton
8
+from vllm.v1.outputs import LogprobsTensors
9
+
10
+
11
+@triton.jit
12
+def _topk_log_softmax_kernel(
13
+ output_ptr,
14
+ logits_ptr,
15
+ logits_stride,
16
+ topk_ids_ptr,
17
+ topk,
18
+ vocab_size,
19
+ BLOCK_SIZE: tl.constexpr,
20
+ PADDED_TOPK: tl.constexpr,
21
+):
22
+ req_idx = tl.program_id(0)
23
+ row_ptr = logits_ptr + req_idx * logits_stride
24
+
25
+ max_val = float("-inf")
26
+ for i in range(0, vocab_size, BLOCK_SIZE):
27
+ block = i + tl.arange(0, BLOCK_SIZE)
28
+ logits = tl.load(row_ptr + block, mask=block < vocab_size, other=float("-inf"))
29
+ max_val = tl.max(tl.maximum(logits, max_val))
30
+ max_val = max_val.to(tl.float32) # type: ignore
31
+
32
+ se = 0.0
33
+ for i in range(0, vocab_size, BLOCK_SIZE):
34
+ block = i + tl.arange(0, BLOCK_SIZE)
35
+ logits = tl.load(row_ptr + block, mask=block < vocab_size, other=0.0)
36
+ # NOTE(woosuk): Make sure that logits and all following operations use FP32.
37
+ logits = logits.to(tl.float32)
38
+ e = tl.exp(logits - max_val)
39
+ e = tl.where(block < vocab_size, e, 0.0)
40
+ se += tl.sum(e)
41
+ lse = tl.log(se)
42
+
43
+ k_offset = tl.arange(0, PADDED_TOPK)
44
+ k_mask = k_offset < topk
45
+ topk_ids = tl.load(topk_ids_ptr + req_idx * topk + k_offset, mask=k_mask, other=0)
46
+
47
+ logits = tl.load(row_ptr + topk_ids, mask=k_mask)
48
+ logits = logits.to(tl.float32)
49
+ o = logits - max_val - lse
50
+ tl.store(output_ptr + req_idx * topk + k_offset, o, mask=k_mask)
51
+
52
+
53
+@triton.jit
54
+def _ranks_kernel(
55
+ output_ptr,
56
+ logits_ptr,
57
+ logits_stride,
58
+ token_ids_ptr,
59
+ vocab_size,
60
+ BLOCK_SIZE: tl.constexpr,
61
+):
62
+ req_idx = tl.program_id(0)
63
+ row_ptr = logits_ptr + req_idx * logits_stride
64
+
65
+ token_id = tl.load(token_ids_ptr + req_idx)
66
+ x = tl.load(row_ptr + token_id)
67
+
68
+ n = 0
69
+ for i in range(0, vocab_size, BLOCK_SIZE):
70
+ block = i + tl.arange(0, BLOCK_SIZE)
71
+ logits = tl.load(row_ptr + block, mask=block < vocab_size, other=float("-inf"))
72
+ n += tl.sum((logits > x).to(tl.int32))
73
+ tl.store(output_ptr + req_idx, n)
74
+
75
+
76
+def compute_token_logprobs(
77
+ logits: torch.Tensor,
78
+ token_ids: torch.Tensor,
79
+) -> torch.Tensor:
80
+ batch_size = logits.shape[0]
81
+ vocab_size = logits.shape[1]
82
+ token_ids = token_ids.to(torch.int64)
83
+ num_logprobs = token_ids.shape[1]
84
+ logprobs = torch.empty(
85
+ batch_size,
86
+ num_logprobs,
87
+ dtype=torch.float32,
88
+ device=logits.device,
89
+ )
90
+ _topk_log_softmax_kernel[(batch_size,)](
91
+ logprobs,
92
+ logits,
93
+ logits.stride(0),
94
+ token_ids,
95
+ num_logprobs,
96
+ vocab_size,
97
+ BLOCK_SIZE=1024, # type: ignore
98
+ PADDED_TOPK=triton.next_power_of_2(num_logprobs),
99
+ )
100
+ return logprobs
101
+
102
+
103
+def compute_topk_logprobs(
104
+ logits: torch.Tensor,
105
+ num_logprobs: int,
106
+ sampled_token_ids: torch.Tensor,
107
+) -> LogprobsTensors:
108
+ assert num_logprobs >= 0
109
+ batch_size, vocab_size = logits.shape
110
+ if num_logprobs == 0:
111
+ logprob_token_ids = sampled_token_ids.unsqueeze(-1)
112
+ else:
113
+ topk_indices = torch.topk(logits, num_logprobs, dim=-1).indices
114
+ logprob_token_ids = torch.cat(
115
+ (sampled_token_ids.unsqueeze(-1), topk_indices), dim=1
116
+ )
117
+
118
+ # NOTE(woosuk): Here, to save GPU memory, we do not materialize the full
119
+ # logprobs tensor. Instead, we only compute and return the logprobs of
120
+ # the topk + 1 tokens.
121
+ logprobs = compute_token_logprobs(logits, logprob_token_ids)
122
+ token_ranks = torch.empty(
123
+ batch_size,
124
+ dtype=torch.int64,
125
+ device=logits.device,
126
+ )
127
+ _ranks_kernel[(batch_size,)](
128
+ token_ranks,
129
+ logits,
130
+ logits.stride(0),
131
+ sampled_token_ids,
132
+ vocab_size,
133
+ BLOCK_SIZE=8192, # type: ignore
134
+ )
135
+ return LogprobsTensors(
136
+ logprob_token_ids=logprob_token_ids,
137
+ logprobs=logprobs,
138
+ selected_token_ranks=token_ranks,
139
+ )
140
+
141
+
142
+def compute_prompt_logprobs(
143
+ prompt_token_ids: torch.Tensor,
144
+ prompt_hidden_states: torch.Tensor,
145
+ logits_fn: Callable[[torch.Tensor], torch.Tensor],
146
+) -> tuple[torch.Tensor, torch.Tensor]:
147
+ # Since materializing the full prompt logits can take too much memory,
148
+ # we compute it in chunks.
149
+ CHUNK_SIZE = 1024
150
+ logprobs = []
151
+ ranks = []
152
+ prompt_token_ids = prompt_token_ids.to(torch.int64)
153
+ for start_idx in range(0, prompt_token_ids.shape[0], CHUNK_SIZE):
154
+ end_idx = start_idx + CHUNK_SIZE
155
+ # NOTE(woosuk): logits_fn can be slow because it involves all-gather.
156
+ prompt_logits = logits_fn(prompt_hidden_states[start_idx:end_idx])
157
+ prompt_logprobs = compute_topk_logprobs(
158
+ prompt_logits,
159
+ 0, # num_logprobs
160
+ prompt_token_ids[start_idx:end_idx],
161
+ )
162
+ logprobs.append(prompt_logprobs.logprobs)
163
+ ranks.append(prompt_logprobs.selected_token_ranks)
164
+
165
+ logprobs = torch.cat(logprobs, dim=0) if len(logprobs) > 1 else logprobs[0]
166
+ ranks = torch.cat(ranks, dim=0) if len(ranks) > 1 else ranks[0]
167
+ return logprobs, ranks
vllm/v1/worker/gpu/sample/metadata.py
new
+179
@@ -0,0 +1,179 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from dataclasses import dataclass
4
+
5
+import torch
6
+
7
+from vllm.triton_utils import tl, triton
8
+
9
+
10
+@dataclass
11
+class SamplingMetadata:
12
+ temperature: torch.Tensor
13
+
14
+ top_p: torch.Tensor | None
15
+ top_k: torch.Tensor | None
16
+
17
+ repetition_penalty: torch.Tensor
18
+ frequency_penalty: torch.Tensor
19
+ presence_penalty: torch.Tensor
20
+
21
+ seeds: torch.Tensor
22
+ pos: torch.Tensor
23
+
24
+ # None means no logprobs, 0 means sampled token logprobs only
25
+ max_num_logprobs: int | None
26
+
27
+ # For penalties
28
+ idx_mapping: torch.Tensor
29
+ prompt_bin_counts: torch.Tensor
30
+ output_bin_counts: torch.Tensor
31
+
32
+ @classmethod
33
+ def make_dummy(
34
+ cls,
35
+ num_reqs: int,
36
+ device: torch.device,
37
+ ) -> "SamplingMetadata":
38
+ assert num_reqs > 0
39
+ temperature = torch.zeros(num_reqs, dtype=torch.float32, device=device)
40
+ temperature[0] = 0.5
41
+ # TODO(woosuk): Use top-p and top-k for dummy sampler.
42
+ # Currently, they are disabled because of memory usage.
43
+ # top_p = torch.full((num_reqs,), 0.95, dtype=torch.float32, device=device)
44
+ # top_k = torch.full((num_reqs,), 20, dtype=torch.int32, device=device)
45
+ top_p = None
46
+ top_k = None
47
+ # NOTE(woosuk): We must set penalties to their default values to make sure
48
+ # the penalties kernel does not touch the placeholder bin_counts tensors.
49
+ repetition_penalty = torch.ones(num_reqs, dtype=torch.float32, device=device)
50
+ frequency_penalty = torch.zeros(num_reqs, dtype=torch.float32, device=device)
51
+ presence_penalty = torch.zeros(num_reqs, dtype=torch.float32, device=device)
52
+ seeds = torch.zeros(num_reqs, dtype=torch.int64, device=device)
53
+ pos = torch.zeros(num_reqs, dtype=torch.int64, device=device)
54
+ max_num_logprobs = 20
55
+
56
+ idx_mapping = torch.arange(num_reqs, dtype=torch.int32, device=device)
57
+ # NOTE(woosuk): These are placeholder tensors to avoid None checks in the
58
+ # penalties kernel. We use 2 instead of 1 as vocab_size to avoid Triton
59
+ # specialization and re-compilation at runtime.
60
+ prompt_bin_counts = torch.zeros(num_reqs, 2, dtype=torch.int32, device=device)
61
+ output_bin_counts = torch.zeros(num_reqs, 2, dtype=torch.int32, device=device)
62
+
63
+ return cls(
64
+ temperature=temperature,
65
+ top_p=top_p,
66
+ top_k=top_k,
67
+ repetition_penalty=repetition_penalty,
68
+ frequency_penalty=frequency_penalty,
69
+ presence_penalty=presence_penalty,
70
+ seeds=seeds,
71
+ pos=pos,
72
+ max_num_logprobs=max_num_logprobs,
73
+ idx_mapping=idx_mapping,
74
+ prompt_bin_counts=prompt_bin_counts,
75
+ output_bin_counts=output_bin_counts,
76
+ )
77
+
78
+
79
+# NOTE(woosuk): Re-compilation can happen at runtime since top_p and top_k can be None.
80
+@triton.jit
81
+def _expand_sampling_metadata_kernel(
82
+ temp_ptr,
83
+ expanded_temp_ptr,
84
+ top_p_ptr,
85
+ expanded_top_p_ptr,
86
+ top_k_ptr,
87
+ expanded_top_k_ptr,
88
+ rep_penalty_ptr,
89
+ expanded_rep_penalty_ptr,
90
+ freq_penalty_ptr,
91
+ expanded_freq_penalty_ptr,
92
+ pres_penalty_ptr,
93
+ expanded_pres_penalty_ptr,
94
+ seeds_ptr,
95
+ expanded_seeds_ptr,
96
+ cu_num_logits_ptr,
97
+ BLOCK_SIZE: tl.constexpr,
98
+):
99
+ req_idx = tl.program_id(0)
100
+ start_idx = tl.load(cu_num_logits_ptr + req_idx)
101
+ end_idx = tl.load(cu_num_logits_ptr + req_idx + 1)
102
+ num_tokens = end_idx - start_idx
103
+
104
+ block = tl.arange(0, BLOCK_SIZE)
105
+ mask = block < num_tokens
106
+
107
+ temp = tl.load(temp_ptr + req_idx)
108
+ tl.store(expanded_temp_ptr + start_idx + block, temp, mask=mask)
109
+
110
+ if top_p_ptr is not None:
111
+ top_p = tl.load(top_p_ptr + req_idx)
112
+ tl.store(expanded_top_p_ptr + start_idx + block, top_p, mask=mask)
113
+
114
+ if top_k_ptr is not None:
115
+ top_k = tl.load(top_k_ptr + req_idx)
116
+ tl.store(expanded_top_k_ptr + start_idx + block, top_k, mask=mask)
117
+
118
+ rep_penalty = tl.load(rep_penalty_ptr + req_idx)
119
+ tl.store(expanded_rep_penalty_ptr + start_idx + block, rep_penalty, mask=mask)
120
+
121
+ freq_penalty = tl.load(freq_penalty_ptr + req_idx)
122
+ tl.store(expanded_freq_penalty_ptr + start_idx + block, freq_penalty, mask=mask)
123
+
124
+ pres_penalty = tl.load(pres_penalty_ptr + req_idx)
125
+ tl.store(expanded_pres_penalty_ptr + start_idx + block, pres_penalty, mask=mask)
126
+
127
+ seed = tl.load(seeds_ptr + req_idx)
128
+ tl.store(expanded_seeds_ptr + start_idx + block, seed, mask=mask)
129
+
130
+
131
+def expand_sampling_metadata(
132
+ sampling_metadata: SamplingMetadata,
133
+ cu_num_logits: torch.Tensor,
134
+ max_expand_len: int,
135
+) -> SamplingMetadata:
136
+ total_num_logits = sampling_metadata.pos.shape[0]
137
+ create_empty = lambda x: x.new_empty(total_num_logits) if x is not None else None
138
+ expanded_temp = create_empty(sampling_metadata.temperature)
139
+ expanded_top_p = create_empty(sampling_metadata.top_p)
140
+ expanded_top_k = create_empty(sampling_metadata.top_k)
141
+ expanded_repetition_penalty = create_empty(sampling_metadata.repetition_penalty)
142
+ expanded_frequency_penalty = create_empty(sampling_metadata.frequency_penalty)
143
+ expanded_presence_penalty = create_empty(sampling_metadata.presence_penalty)
144
+ expanded_seeds = create_empty(sampling_metadata.seeds)
145
+
146
+ num_reqs = cu_num_logits.shape[0] - 1
147
+ _expand_sampling_metadata_kernel[(num_reqs,)](
148
+ sampling_metadata.temperature,
149
+ expanded_temp,
150
+ sampling_metadata.top_p,
151
+ expanded_top_p,
152
+ sampling_metadata.top_k,
153
+ expanded_top_k,
154
+ sampling_metadata.repetition_penalty,
155
+ expanded_repetition_penalty,
156
+ sampling_metadata.frequency_penalty,
157
+ expanded_frequency_penalty,
158
+ sampling_metadata.presence_penalty,
159
+ expanded_presence_penalty,
160
+ sampling_metadata.seeds,
161
+ expanded_seeds,
162
+ cu_num_logits,
163
+ BLOCK_SIZE=triton.next_power_of_2(max_expand_len),
164
+ )
165
+ return SamplingMetadata(
166
+ temperature=expanded_temp,
167
+ top_p=expanded_top_p,
168
+ top_k=expanded_top_k,
169
+ seeds=expanded_seeds,
170
+ repetition_penalty=expanded_repetition_penalty,
171
+ frequency_penalty=expanded_frequency_penalty,
172
+ presence_penalty=expanded_presence_penalty,
173
+ pos=sampling_metadata.pos,
174
+ max_num_logprobs=sampling_metadata.max_num_logprobs,
175
+ # TODO(woosuk): Support penalties with spec decoding.
176
+ idx_mapping=sampling_metadata.idx_mapping,
177
+ prompt_bin_counts=sampling_metadata.prompt_bin_counts,
178
+ output_bin_counts=sampling_metadata.output_bin_counts,
179
+ )
vllm/v1/worker/gpu/sample/penalties.py
renamed
+47
-1
@@ -3,7 +3,7 @@
3
import torch
4
5
from vllm.triton_utils import tl, triton
6
-from vllm.v1.worker.gpu.states import SamplingMetadata
6
+from vllm.v1.worker.gpu.sample.metadata import SamplingMetadata
7
8
9
@triton.jit
@@ -83,3 +83,49 @@ def apply_penalties(logits: torch.Tensor, sampling_metadata: SamplingMetadata) -
83
vocab_size,
84
BLOCK_SIZE=BLOCK_SIZE,
85
)
86
+
87
+
88
+@triton.jit(do_not_specialize=["prefill_len", "prompt_len"])
89
+def _bincount_kernel(
90
+ prefill_token_ids_ptr,
91
+ prefill_len,
92
+ prompt_len,
93
+ prompt_bin_counts_ptr,
94
+ output_bin_counts_ptr,
95
+ BLOCK_SIZE: tl.constexpr,
96
+):
97
+ block_idx = tl.program_id(0)
98
+ if block_idx * BLOCK_SIZE >= prefill_len:
99
+ return
100
+
101
+ block = block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
102
+ if block_idx * BLOCK_SIZE < prompt_len:
103
+ mask = block < prompt_len
104
+ prefill_tokens = tl.load(prefill_token_ids_ptr + block, mask=mask)
105
+ tl.atomic_add(prompt_bin_counts_ptr + prefill_tokens, 1, mask=mask)
106
+ if (block_idx + 1) * BLOCK_SIZE >= prompt_len:
107
+ mask = block < prefill_len
108
+ mask &= block >= prompt_len
109
+ prefill_tokens = tl.load(prefill_token_ids_ptr + block, mask=mask)
110
+ tl.atomic_add(output_bin_counts_ptr + prefill_tokens, 1, mask=mask)
111
+
112
+
113
+def bincount(
114
+ prefill_token_ids: torch.Tensor,
115
+ prefill_len: int,
116
+ prompt_len: int,
117
+ prompt_bin_counts: torch.Tensor,
118
+ output_bin_counts: torch.Tensor,
119
+) -> None:
120
+ prompt_bin_counts.zero_()
121
+ output_bin_counts.zero_()
122
+ BLOCK_SIZE = 1024
123
+ num_blocks = triton.cdiv(prefill_len, BLOCK_SIZE)
124
+ _bincount_kernel[(num_blocks,)](
125
+ prefill_token_ids,
126
+ prefill_len,
127
+ prompt_len,
128
+ prompt_bin_counts,
129
+ output_bin_counts,
130
+ BLOCK_SIZE=BLOCK_SIZE,
131
+ )
vllm/v1/worker/gpu/sample/sampler.py
new
+79
@@ -0,0 +1,79 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import torch
5
+
6
+from vllm.config.model import LogprobsMode
7
+from vllm.v1.outputs import SamplerOutput
8
+from vllm.v1.sample.ops.topk_topp_sampler import apply_top_k_top_p
9
+from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample
10
+from vllm.v1.worker.gpu.sample.logprob import compute_topk_logprobs
11
+from vllm.v1.worker.gpu.sample.metadata import SamplingMetadata
12
+from vllm.v1.worker.gpu.sample.penalties import apply_penalties
13
+
14
+
15
+class Sampler:
16
+ def __init__(
17
+ self,
18
+ logprobs_mode: LogprobsMode = "raw_logprobs",
19
+ ):
20
+ if logprobs_mode not in ["processed_logprobs", "raw_logprobs"]:
21
+ raise NotImplementedError(f"Unsupported logprobs_mode: {logprobs_mode}")
22
+ self.logprobs_mode = logprobs_mode
23
+
24
+ def __call__(
25
+ self,
26
+ logits: torch.Tensor,
27
+ sampling_metadata: SamplingMetadata,
28
+ ) -> SamplerOutput:
29
+ if sampling_metadata.max_num_logprobs is not None:
30
+ if self.logprobs_mode == "processed_logprobs":
31
+ sampled, logits = self.sample(
32
+ logits, sampling_metadata, return_logits=True
33
+ )
34
+ else:
35
+ assert self.logprobs_mode == "raw_logprobs"
36
+ sampled, _ = self.sample(logits, sampling_metadata, return_logits=False)
37
+
38
+ logprobs_tensors = compute_topk_logprobs(
39
+ logits,
40
+ sampling_metadata.max_num_logprobs,
41
+ sampled,
42
+ )
43
+ else:
44
+ sampled, _ = self.sample(logits, sampling_metadata, return_logits=False)
45
+ logprobs_tensors = None
46
+
47
+ # These are GPU tensors.
48
+ sampler_output = SamplerOutput(
49
+ # The sampled tokens are expanded to 2D tensor with shape
50
+ # [num_requests, 1], where each row represents one generated
51
+ # token per request.
52
+ sampled_token_ids=sampled.view(-1, 1),
53
+ logprobs_tensors=logprobs_tensors,
54
+ )
55
+ return sampler_output
56
+
57
+ def sample(
58
+ self,
59
+ logits: torch.Tensor,
60
+ sampling_metadata: SamplingMetadata,
61
+ return_logits: bool = False,
62
+ ) -> tuple[torch.Tensor, torch.Tensor | None]:
63
+ is_greedy = sampling_metadata.temperature == 0
64
+ temp = torch.where(is_greedy, 1.0, sampling_metadata.temperature)
65
+ logits = logits / temp.view(-1, 1)
66
+ logits = apply_top_k_top_p(
67
+ logits, sampling_metadata.top_k, sampling_metadata.top_p
68
+ )
69
+ # Apply penalties in place.
70
+ apply_penalties(logits, sampling_metadata)
71
+
72
+ sampled = gumbel_sample(
73
+ logits,
74
+ sampling_metadata.temperature,
75
+ sampling_metadata.seeds,
76
+ sampling_metadata.pos,
77
+ apply_temperature=False,
78
+ )
79
+ return sampled, logits if return_logits else None
vllm/v1/worker/gpu/sampler.py
deleted
-333
@@ -1,333 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-from collections.abc import Callable
4
-
5
-import torch
6
-
7
-from vllm.config.model import LogprobsMode
8
-from vllm.triton_utils import tl, triton
9
-from vllm.v1.outputs import LogprobsTensors, SamplerOutput
10
-from vllm.v1.sample.ops.topk_topp_sampler import apply_top_k_top_p
11
-from vllm.v1.worker.gpu.penalties import apply_penalties
12
-from vllm.v1.worker.gpu.states import SamplingMetadata
13
-
14
-
15
-class Sampler:
16
- def __init__(
17
- self,
18
- logprobs_mode: LogprobsMode = "raw_logprobs",
19
- ):
20
- if logprobs_mode not in ["processed_logprobs", "raw_logprobs"]:
21
- raise NotImplementedError(f"Unsupported logprobs_mode: {logprobs_mode}")
22
- self.logprobs_mode = logprobs_mode
23
-
24
- def __call__(
25
- self,
26
- logits: torch.Tensor,
27
- sampling_metadata: SamplingMetadata,
28
- ) -> SamplerOutput:
29
- if sampling_metadata.max_num_logprobs is not None:
30
- if self.logprobs_mode == "processed_logprobs":
31
- sampled, logits = self.sample(
32
- logits, sampling_metadata, return_logits=True
33
- )
34
- else:
35
- assert self.logprobs_mode == "raw_logprobs"
36
- sampled, _ = self.sample(logits, sampling_metadata, return_logits=False)
37
-
38
- logprobs_tensors = compute_topk_logprobs(
39
- logits,
40
- sampling_metadata.max_num_logprobs,
41
- sampled,
42
- )
43
- else:
44
- sampled, _ = self.sample(logits, sampling_metadata, return_logits=False)
45
- logprobs_tensors = None
46
-
47
- # These are GPU tensors.
48
- sampler_output = SamplerOutput(
49
- # The sampled tokens are expanded to 2D tensor with shape
50
- # [num_requests, 1], where each row represents one generated
51
- # token per request.
52
- sampled_token_ids=sampled.view(-1, 1),
53
- logprobs_tensors=logprobs_tensors,
54
- )
55
- return sampler_output
56
-
57
- def sample(
58
- self,
59
- logits: torch.Tensor,
60
- sampling_metadata: SamplingMetadata,
61
- return_logits: bool = False,
62
- ) -> tuple[torch.Tensor, torch.Tensor | None]:
63
- is_greedy = sampling_metadata.temperature == 0
64
- temp = torch.where(is_greedy, 1.0, sampling_metadata.temperature)
65
- logits = logits / temp.view(-1, 1)
66
- logits = apply_top_k_top_p(
67
- logits, sampling_metadata.top_k, sampling_metadata.top_p
68
- )
69
- # Apply penalties in place.
70
- apply_penalties(logits, sampling_metadata)
71
-
72
- sampled = gumbel_sample(
73
- logits,
74
- sampling_metadata.temperature,
75
- sampling_metadata.seeds,
76
- sampling_metadata.pos,
77
- apply_temperature=False,
78
- )
79
- return sampled, logits if return_logits else None
80
-
81
-
82
-@triton.jit
83
-def _gumbel_sample_kernel(
84
- local_argmax_ptr,
85
- local_argmax_stride,
86
- local_max_ptr,
87
- local_max_stride,
88
- logits_ptr,
89
- logits_stride,
90
- seeds_ptr,
91
- pos_ptr,
92
- temp_ptr,
93
- vocab_size,
94
- BLOCK_SIZE: tl.constexpr,
95
- APPLY_TEMPERATURE: tl.constexpr,
96
-):
97
- req_idx = tl.program_id(0)
98
- block_idx = tl.program_id(1)
99
- block = block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
100
- mask = block < vocab_size
101
- logits = tl.load(
102
- logits_ptr + req_idx * logits_stride + block,
103
- mask=mask,
104
- other=float("-inf"),
105
- )
106
- logits = logits.to(tl.float32)
107
-
108
- temp = tl.load(temp_ptr + req_idx).to(tl.float32)
109
- if temp != 0.0:
110
- # Calculate the seed for gumbel noise.
111
- seed = tl.load(seeds_ptr + req_idx)
112
- pos = tl.load(pos_ptr + req_idx)
113
- gumbel_seed = tl.randint(seed, pos)
114
-
115
- # Generate gumbel noise.
116
- r = tl.rand(gumbel_seed, block).to(tl.float64)
117
- gumbel_noise = -tl.log(-tl.log(r + 1e-20) + 1e-20)
118
- gumbel_noise = gumbel_noise.to(tl.float32)
119
-
120
- # Apply temperature.
121
- if APPLY_TEMPERATURE:
122
- # NOTE(woosuk): Use div_rn to match the behavior of torch.
123
- logits = tl.div_rn(logits, temp)
124
-
125
- # Apply gumbel noise.
126
- logits = tl.where(mask, logits + gumbel_noise, float("-inf"))
127
-
128
- idx = tl.argmax(logits, axis=0)
129
- token_id = block_idx * BLOCK_SIZE + idx
130
- value = tl.max(logits, axis=0)
131
- tl.store(local_argmax_ptr + req_idx * local_argmax_stride + block_idx, token_id)
132
- tl.store(local_max_ptr + req_idx * local_max_stride + block_idx, value)
133
-
134
-
135
-def gumbel_sample(
136
- logits: torch.Tensor, # [num_reqs, vocab_size]
137
- temperature: torch.Tensor, # [num_reqs]
138
- seed: torch.Tensor, # [num_reqs]
139
- pos: torch.Tensor, # [num_reqs]
140
- apply_temperature: bool,
141
-) -> torch.Tensor:
142
- num_reqs, vocab_size = logits.shape
143
- BLOCK_SIZE = 1024
144
- num_blocks = triton.cdiv(vocab_size, BLOCK_SIZE)
145
- local_argmax = torch.empty(
146
- num_reqs,
147
- num_blocks,
148
- dtype=torch.int64,
149
- device=logits.device,
150
- )
151
- local_max = torch.empty(
152
- num_reqs,
153
- num_blocks,
154
- dtype=torch.float32,
155
- device=logits.device,
156
- )
157
- _gumbel_sample_kernel[(num_reqs, num_blocks)](
158
- local_argmax,
159
- local_argmax.stride(0),
160
- local_max,
161
- local_max.stride(0),
162
- logits,
163
- logits.stride(0),
164
- seed,
165
- pos,
166
- temperature,
167
- vocab_size,
168
- BLOCK_SIZE=BLOCK_SIZE,
169
- APPLY_TEMPERATURE=apply_temperature,
170
- )
171
- # NOTE(woosuk): Use int64 for later indexing.
172
- max_block_idx = local_max.argmax(dim=-1, keepdim=True)
173
- sampled = local_argmax.gather(dim=-1, index=max_block_idx).view(-1)
174
- return sampled
175
-
176
-
177
-@triton.jit
178
-def _topk_log_softmax_kernel(
179
- output_ptr,
180
- logits_ptr,
181
- logits_stride,
182
- topk_ids_ptr,
183
- topk,
184
- vocab_size,
185
- BLOCK_SIZE: tl.constexpr,
186
- PADDED_TOPK: tl.constexpr,
187
-):
188
- req_idx = tl.program_id(0)
189
- row_ptr = logits_ptr + req_idx * logits_stride
190
-
191
- max_val = float("-inf")
192
- for i in range(0, vocab_size, BLOCK_SIZE):
193
- block = i + tl.arange(0, BLOCK_SIZE)
194
- logits = tl.load(row_ptr + block, mask=block < vocab_size, other=float("-inf"))
195
- max_val = tl.max(tl.maximum(logits, max_val))
196
- max_val = max_val.to(tl.float32) # type: ignore
197
-
198
- se = 0.0
199
- for i in range(0, vocab_size, BLOCK_SIZE):
200
- block = i + tl.arange(0, BLOCK_SIZE)
201
- logits = tl.load(row_ptr + block, mask=block < vocab_size, other=0.0)
202
- # NOTE(woosuk): Make sure that logits and all following operations use FP32.
203
- logits = logits.to(tl.float32)
204
- e = tl.exp(logits - max_val)
205
- e = tl.where(block < vocab_size, e, 0.0)
206
- se += tl.sum(e)
207
- lse = tl.log(se)
208
-
209
- k_offset = tl.arange(0, PADDED_TOPK)
210
- k_mask = k_offset < topk
211
- topk_ids = tl.load(topk_ids_ptr + req_idx * topk + k_offset, mask=k_mask, other=0)
212
-
213
- logits = tl.load(row_ptr + topk_ids, mask=k_mask)
214
- logits = logits.to(tl.float32)
215
- o = logits - max_val - lse
216
- tl.store(output_ptr + req_idx * topk + k_offset, o, mask=k_mask)
217
-
218
-
219
-@triton.jit
220
-def _ranks_kernel(
221
- output_ptr,
222
- logits_ptr,
223
- logits_stride,
224
- token_ids_ptr,
225
- vocab_size,
226
- BLOCK_SIZE: tl.constexpr,
227
-):
228
- req_idx = tl.program_id(0)
229
- row_ptr = logits_ptr + req_idx * logits_stride
230
-
231
- token_id = tl.load(token_ids_ptr + req_idx)
232
- x = tl.load(row_ptr + token_id)
233
-
234
- n = 0
235
- for i in range(0, vocab_size, BLOCK_SIZE):
236
- block = i + tl.arange(0, BLOCK_SIZE)
237
- logits = tl.load(row_ptr + block, mask=block < vocab_size, other=float("-inf"))
238
- n += tl.sum((logits > x).to(tl.int32))
239
- tl.store(output_ptr + req_idx, n)
240
-
241
-
242
-def compute_token_logprobs(
243
- logits: torch.Tensor,
244
- token_ids: torch.Tensor,
245
-) -> torch.Tensor:
246
- batch_size = logits.shape[0]
247
- vocab_size = logits.shape[1]
248
- token_ids = token_ids.to(torch.int64)
249
- num_logprobs = token_ids.shape[1]
250
- logprobs = torch.empty(
251
- batch_size,
252
- num_logprobs,
253
- dtype=torch.float32,
254
- device=logits.device,
255
- )
256
- _topk_log_softmax_kernel[(batch_size,)](
257
- logprobs,
258
- logits,
259
- logits.stride(0),
260
- token_ids,
261
- num_logprobs,
262
- vocab_size,
263
- BLOCK_SIZE=1024, # type: ignore
264
- PADDED_TOPK=triton.next_power_of_2(num_logprobs),
265
- )
266
- return logprobs
267
-
268
-
269
-def compute_topk_logprobs(
270
- logits: torch.Tensor,
271
- num_logprobs: int,
272
- sampled_token_ids: torch.Tensor,
273
-) -> LogprobsTensors:
274
- assert num_logprobs >= 0
275
- batch_size, vocab_size = logits.shape
276
- if num_logprobs == 0:
277
- logprob_token_ids = sampled_token_ids.unsqueeze(-1)
278
- else:
279
- topk_indices = torch.topk(logits, num_logprobs, dim=-1).indices
280
- logprob_token_ids = torch.cat(
281
- (sampled_token_ids.unsqueeze(-1), topk_indices), dim=1
282
- )
283
-
284
- # NOTE(woosuk): Here, to save GPU memory, we do not materialize the full
285
- # logprobs tensor. Instead, we only compute and return the logprobs of
286
- # the topk + 1 tokens.
287
- logprobs = compute_token_logprobs(logits, logprob_token_ids)
288
- token_ranks = torch.empty(
289
- batch_size,
290
- dtype=torch.int64,
291
- device=logits.device,
292
- )
293
- _ranks_kernel[(batch_size,)](
294
- token_ranks,
295
- logits,
296
- logits.stride(0),
297
- sampled_token_ids,
298
- vocab_size,
299
- BLOCK_SIZE=8192, # type: ignore
300
- )
301
- return LogprobsTensors(
302
- logprob_token_ids=logprob_token_ids,
303
- logprobs=logprobs,
304
- selected_token_ranks=token_ranks,
305
- )
306
-
307
-
308
-def compute_prompt_logprobs(
309
- prompt_token_ids: torch.Tensor,
310
- prompt_hidden_states: torch.Tensor,
311
- logits_fn: Callable[[torch.Tensor], torch.Tensor],
312
-) -> tuple[torch.Tensor, torch.Tensor]:
313
- # Since materializing the full prompt logits can take too much memory,
314
- # we compute it in chunks.
315
- CHUNK_SIZE = 1024
316
- logprobs = []
317
- ranks = []
318
- prompt_token_ids = prompt_token_ids.to(torch.int64)
319
- for start_idx in range(0, prompt_token_ids.shape[0], CHUNK_SIZE):
320
- end_idx = start_idx + CHUNK_SIZE
321
- # NOTE(woosuk): logits_fn can be slow because it involves all-gather.
322
- prompt_logits = logits_fn(prompt_hidden_states[start_idx:end_idx])
323
- prompt_logprobs = compute_topk_logprobs(
324
- prompt_logits,
325
- 0, # num_logprobs
326
- prompt_token_ids[start_idx:end_idx],
327
- )
328
- logprobs.append(prompt_logprobs.logprobs)
329
- ranks.append(prompt_logprobs.selected_token_ranks)
330
-
331
- logprobs = torch.cat(logprobs, dim=0) if len(logprobs) > 1 else logprobs[0]
332
- ranks = torch.cat(ranks, dim=0) if len(ranks) > 1 else ranks[0]
333
- return logprobs, ranks
vllm/v1/worker/gpu/spec_decode/eagle.py
+2
-2
@@ -18,9 +18,9 @@ from vllm.v1.kv_cache_interface import KVCacheConfig
18
from vllm.v1.worker.gpu.attn_utils import build_attn_metadata
19
from vllm.v1.worker.gpu.block_table import BlockTables
20
from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers
21
-from vllm.v1.worker.gpu.sampler import gumbel_sample
21
+from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample
22
+from vllm.v1.worker.gpu.sample.metadata import SamplingMetadata
23
from vllm.v1.worker.gpu.spec_decode.eagle_cudagraph import EagleCudaGraphManager
23
-from vllm.v1.worker.gpu.states import SamplingMetadata
24
25
logger = init_logger(__name__)
26
vllm/v1/worker/gpu/states.py
+2
-230
@@ -7,86 +7,18 @@ import torch
7
8
from vllm.lora.request import LoRARequest
9
from vllm.sampling_params import SamplingParams
10
-from vllm.triton_utils import tl, triton
10
from vllm.utils.platform_utils import is_uva_available
11
from vllm.utils.torch_utils import get_cuda_view_from_cpu_tensor
12
from vllm.v1.outputs import LogprobsTensors
13
from vllm.v1.utils import CpuGpuBuffer
14
+from vllm.v1.worker.gpu.sample.metadata import SamplingMetadata
15
+from vllm.v1.worker.gpu.sample.penalties import bincount
16
17
_NP_INT64_MIN = np.iinfo(np.int64).min
18
_NP_INT64_MAX = np.iinfo(np.int64).max
19
NO_LORA_ID = 0
20
21
21
-@dataclass
22
-class SamplingMetadata:
23
- temperature: torch.Tensor
24
-
25
- top_p: torch.Tensor | None
26
- top_k: torch.Tensor | None
27
-
28
- repetition_penalty: torch.Tensor
29
- frequency_penalty: torch.Tensor
30
- presence_penalty: torch.Tensor
31
-
32
- seeds: torch.Tensor
33
- pos: torch.Tensor
34
-
35
- # None means no logprobs, 0 means sampled token logprobs only
36
- max_num_logprobs: int | None
37
-
38
- # For penalties
39
- idx_mapping: torch.Tensor
40
- prompt_bin_counts: torch.Tensor
41
- output_bin_counts: torch.Tensor
42
-
43
- @classmethod
44
- def make_dummy(
45
- cls,
46
- num_reqs: int,
47
- device: torch.device,
48
- ) -> "SamplingMetadata":
49
- assert num_reqs > 0
50
- temperature = torch.zeros(num_reqs, dtype=torch.float32, device=device)
51
- temperature[0] = 0.5
52
- # TODO(woosuk): Use top-p and top-k for dummy sampler.
53
- # Currently, they are disabled because of memory usage.
54
- # top_p = torch.full((num_reqs,), 0.95, dtype=torch.float32, device=device)
55
- # top_k = torch.full((num_reqs,), 20, dtype=torch.int32, device=device)
56
- top_p = None
57
- top_k = None
58
- # NOTE(woosuk): We must set penalties to their default values to make sure
59
- # the penalties kernel does not touch the placeholder bin_counts tensors.
60
- repetition_penalty = torch.ones(num_reqs, dtype=torch.float32, device=device)
61
- frequency_penalty = torch.zeros(num_reqs, dtype=torch.float32, device=device)
62
- presence_penalty = torch.zeros(num_reqs, dtype=torch.float32, device=device)
63
- seeds = torch.zeros(num_reqs, dtype=torch.int64, device=device)
64
- pos = torch.zeros(num_reqs, dtype=torch.int64, device=device)
65
- max_num_logprobs = 20
66
-
67
- idx_mapping = torch.arange(num_reqs, dtype=torch.int32, device=device)
68
- # NOTE(woosuk): These are placeholder tensors to avoid None checks in the
69
- # penalties kernel. We use 2 instead of 1 as vocab_size to avoid Triton
70
- # specialization and re-compilation at runtime.
71
- prompt_bin_counts = torch.zeros(num_reqs, 2, dtype=torch.int32, device=device)
72
- output_bin_counts = torch.zeros(num_reqs, 2, dtype=torch.int32, device=device)
73
-
74
- return cls(
75
- temperature=temperature,
76
- top_p=top_p,
77
- top_k=top_k,
78
- repetition_penalty=repetition_penalty,
79
- frequency_penalty=frequency_penalty,
80
- presence_penalty=presence_penalty,
81
- seeds=seeds,
82
- pos=pos,
83
- max_num_logprobs=max_num_logprobs,
84
- idx_mapping=idx_mapping,
85
- prompt_bin_counts=prompt_bin_counts,
86
- output_bin_counts=output_bin_counts,
87
- )
88
-
89
-
22
class RequestState:
23
def __init__(
24
self,
@@ -311,17 +243,6 @@ class RequestState:
243
output_bin_counts=self.output_bin_counts,
244
)
245
314
- def expand_sampling_metadata(
315
- self,
316
- sampling_metadata: SamplingMetadata,
317
- cu_num_logits: torch.Tensor,
318
- ) -> SamplingMetadata:
319
- # For draft tokens, we need to expand the sampling param tensors as
320
- # each request samples multiple tokens in each step.
321
- return expand_sampling_metadata(
322
- sampling_metadata, cu_num_logits, self.num_speculative_steps
323
- )
324
-
246
def make_lora_inputs(
247
self,
248
req_ids: list[str],
@@ -376,158 +297,9 @@ class UvaBuffer:
297
self.gpu = get_cuda_view_from_cpu_tensor(self.cpu)
298
299
379
-# NOTE(woosuk): Re-compilation can happen at runtime since top_p and top_k can be None.
380
-@triton.jit
381
-def _expand_sampling_metadata_kernel(
382
- temp_ptr,
383
- expanded_temp_ptr,
384
- top_p_ptr,
385
- expanded_top_p_ptr,
386
- top_k_ptr,
387
- expanded_top_k_ptr,
388
- rep_penalty_ptr,
389
- expanded_rep_penalty_ptr,
390
- freq_penalty_ptr,
391
- expanded_freq_penalty_ptr,
392
- pres_penalty_ptr,
393
- expanded_pres_penalty_ptr,
394
- seeds_ptr,
395
- expanded_seeds_ptr,
396
- cu_num_logits_ptr,
397
- BLOCK_SIZE: tl.constexpr,
398
-):
399
- req_idx = tl.program_id(0)
400
- start_idx = tl.load(cu_num_logits_ptr + req_idx)
401
- end_idx = tl.load(cu_num_logits_ptr + req_idx + 1)
402
- num_tokens = end_idx - start_idx
403
-
404
- block = tl.arange(0, BLOCK_SIZE)
405
- mask = block < num_tokens
406
-
407
- temp = tl.load(temp_ptr + req_idx)
408
- tl.store(expanded_temp_ptr + start_idx + block, temp, mask=mask)
409
-
410
- if top_p_ptr is not None:
411
- top_p = tl.load(top_p_ptr + req_idx)
412
- tl.store(expanded_top_p_ptr + start_idx + block, top_p, mask=mask)
413
-
414
- if top_k_ptr is not None:
415
- top_k = tl.load(top_k_ptr + req_idx)
416
- tl.store(expanded_top_k_ptr + start_idx + block, top_k, mask=mask)
417
-
418
- rep_penalty = tl.load(rep_penalty_ptr + req_idx)
419
- tl.store(expanded_rep_penalty_ptr + start_idx + block, rep_penalty, mask=mask)
420
-
421
- freq_penalty = tl.load(freq_penalty_ptr + req_idx)
422
- tl.store(expanded_freq_penalty_ptr + start_idx + block, freq_penalty, mask=mask)
423
-
424
- pres_penalty = tl.load(pres_penalty_ptr + req_idx)
425
- tl.store(expanded_pres_penalty_ptr + start_idx + block, pres_penalty, mask=mask)
426
-
427
- seed = tl.load(seeds_ptr + req_idx)
428
- tl.store(expanded_seeds_ptr + start_idx + block, seed, mask=mask)
429
-
430
-
431
-def expand_sampling_metadata(
432
- sampling_metadata: SamplingMetadata,
433
- cu_num_logits: torch.Tensor,
434
- num_speculative_steps: int,
435
-) -> SamplingMetadata:
436
- total_num_logits = sampling_metadata.pos.shape[0]
437
- create_empty = lambda x: x.new_empty(total_num_logits) if x is not None else None
438
- expanded_temp = create_empty(sampling_metadata.temperature)
439
- expanded_top_p = create_empty(sampling_metadata.top_p)
440
- expanded_top_k = create_empty(sampling_metadata.top_k)
441
- expanded_repetition_penalty = create_empty(sampling_metadata.repetition_penalty)
442
- expanded_frequency_penalty = create_empty(sampling_metadata.frequency_penalty)
443
- expanded_presence_penalty = create_empty(sampling_metadata.presence_penalty)
444
- expanded_seeds = create_empty(sampling_metadata.seeds)
445
-
446
- num_reqs = cu_num_logits.shape[0] - 1
447
- _expand_sampling_metadata_kernel[(num_reqs,)](
448
- sampling_metadata.temperature,
449
- expanded_temp,
450
- sampling_metadata.top_p,
451
- expanded_top_p,
452
- sampling_metadata.top_k,
453
- expanded_top_k,
454
- sampling_metadata.repetition_penalty,
455
- expanded_repetition_penalty,
456
- sampling_metadata.frequency_penalty,
457
- expanded_frequency_penalty,
458
- sampling_metadata.presence_penalty,
459
- expanded_presence_penalty,
460
- sampling_metadata.seeds,
461
- expanded_seeds,
462
- cu_num_logits,
463
- BLOCK_SIZE=triton.next_power_of_2(num_speculative_steps + 1),
464
- )
465
- return SamplingMetadata(
466
- temperature=expanded_temp,
467
- top_p=expanded_top_p,
468
- top_k=expanded_top_k,
469
- seeds=expanded_seeds,
470
- repetition_penalty=expanded_repetition_penalty,
471
- frequency_penalty=expanded_frequency_penalty,
472
- presence_penalty=expanded_presence_penalty,
473
- pos=sampling_metadata.pos,
474
- max_num_logprobs=sampling_metadata.max_num_logprobs,
475
- # TODO(woosuk): Support penalties with spec decoding.
476
- idx_mapping=sampling_metadata.idx_mapping,
477
- prompt_bin_counts=sampling_metadata.prompt_bin_counts,
478
- output_bin_counts=sampling_metadata.output_bin_counts,
479
- )
480
-
481
-
300
def use_penalty(sampling_params: SamplingParams) -> bool:
301
return (
302
sampling_params.repetition_penalty != 1.0
303
or sampling_params.frequency_penalty != 0.0
304
or sampling_params.presence_penalty != 0.0
305
)
488
-
489
-
490
-@triton.jit(do_not_specialize=["prefill_len", "prompt_len"])
491
-def _bincount_kernel(
492
- prefill_token_ids_ptr,
493
- prefill_len,
494
- prompt_len,
495
- prompt_bin_counts_ptr,
496
- output_bin_counts_ptr,
497
- BLOCK_SIZE: tl.constexpr,
498
-):
499
- block_idx = tl.program_id(0)
500
- if block_idx * BLOCK_SIZE >= prefill_len:
501
- return
502
-
503
- block = block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
504
- if block_idx * BLOCK_SIZE < prompt_len:
505
- mask = block < prompt_len
506
- prefill_tokens = tl.load(prefill_token_ids_ptr + block, mask=mask)
507
- tl.atomic_add(prompt_bin_counts_ptr + prefill_tokens, 1, mask=mask)
508
- if (block_idx + 1) * BLOCK_SIZE >= prompt_len:
509
- mask = block < prefill_len
510
- mask &= block >= prompt_len
511
- prefill_tokens = tl.load(prefill_token_ids_ptr + block, mask=mask)
512
- tl.atomic_add(output_bin_counts_ptr + prefill_tokens, 1, mask=mask)
513
-
514
-
515
-def bincount(
516
- prefill_token_ids: torch.Tensor,
517
- prefill_len: int,
518
- prompt_len: int,
519
- prompt_bin_counts: torch.Tensor,
520
- output_bin_counts: torch.Tensor,
521
-) -> None:
522
- prompt_bin_counts.zero_()
523
- output_bin_counts.zero_()
524
- BLOCK_SIZE = 1024
525
- num_blocks = triton.cdiv(prefill_len, BLOCK_SIZE)
526
- _bincount_kernel[(num_blocks,)](
527
- prefill_token_ids,
528
- prefill_len,
529
- prompt_len,
530
- prompt_bin_counts,
531
- output_bin_counts,
532
- BLOCK_SIZE=BLOCK_SIZE,
533
- )