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 9, 20:45 UTC
[Model Runner V2] Minor refactor for `compute_slot_mappings` (#32794)
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
Woosuk Kwon committed
Jan 21, 2026 at 10:24 UTC
e1da249c931d65b0e9434028a136b52860384735
3 files changed
+33
-22
vllm/v1/worker/gpu/block_table.py
+22
-18
@@ -116,24 +116,26 @@ class BlockTables:
116
117
def compute_slot_mappings(
118
self,
119
+ idx_mapping: torch.Tensor,
120
query_start_loc: torch.Tensor,
121
positions: torch.Tensor,
122
) -> torch.Tensor:
122
- num_reqs = query_start_loc.shape[0] - 1
123
+ num_reqs = idx_mapping.shape[0]
124
num_tokens = positions.shape[0]
125
num_groups = self.num_kv_cache_groups
126
_compute_slot_mappings_kernel[(num_groups, num_reqs + 1)](
127
num_tokens,
128
self.max_num_batched_tokens,
129
+ idx_mapping,
130
query_start_loc,
131
positions,
130
- self.input_block_table_ptrs,
132
+ self.block_table_ptrs,
133
self.block_table_strides,
134
self.block_sizes_tensor,
135
self.slot_mappings,
136
self.slot_mappings.stride(0),
137
PAD_ID=PAD_SLOT_ID,
136
- BLOCK_SIZE=1024, # type: ignore
138
+ TRITON_BLOCK_SIZE=1024, # type: ignore
139
)
140
return self.slot_mappings[:, :num_tokens]
141
@@ -176,42 +178,44 @@ def _gather_block_tables_kernel(
178
def _compute_slot_mappings_kernel(
179
num_tokens,
180
max_num_tokens,
179
- cu_num_tokens, # [num_reqs + 1]
181
+ idx_mapping, # [num_reqs]
182
+ query_start_loc, # [num_reqs + 1]
183
pos, # [num_tokens]
184
block_table_ptrs, # [num_kv_cache_groups]
185
block_table_strides, # [num_kv_cache_groups]
183
- page_sizes, # [num_kv_cache_groups]
186
+ block_sizes, # [num_kv_cache_groups]
187
slot_mappings_ptr, # [num_kv_cache_groups, max_num_tokens]
188
slot_mappings_stride,
189
PAD_ID: tl.constexpr,
187
- BLOCK_SIZE: tl.constexpr,
190
+ TRITON_BLOCK_SIZE: tl.constexpr,
191
):
192
# kv cache group id
193
group_id = tl.program_id(0)
191
- req_idx = tl.program_id(1)
194
+ batch_idx = tl.program_id(1)
195
slot_mapping_ptr = slot_mappings_ptr + group_id * slot_mappings_stride
196
194
- if req_idx == tl.num_programs(1) - 1:
197
+ if batch_idx == tl.num_programs(1) - 1:
198
# Pad remaining slots to -1. This is needed for CUDA graphs.
196
- for i in range(num_tokens, max_num_tokens, BLOCK_SIZE):
197
- offset = i + tl.arange(0, BLOCK_SIZE)
199
+ for i in range(num_tokens, max_num_tokens, TRITON_BLOCK_SIZE):
200
+ offset = i + tl.arange(0, TRITON_BLOCK_SIZE)
201
tl.store(slot_mapping_ptr + offset, PAD_ID, mask=offset < max_num_tokens)
202
return
203
204
block_table_ptr = _load_ptr(block_table_ptrs + group_id, tl.int32)
205
block_table_stride = tl.load(block_table_strides + group_id)
203
- page_size = tl.load(page_sizes + group_id)
206
+ block_size = tl.load(block_sizes + group_id)
207
205
- start_idx = tl.load(cu_num_tokens + req_idx)
206
- end_idx = tl.load(cu_num_tokens + req_idx + 1)
207
- for i in range(start_idx, end_idx, BLOCK_SIZE):
208
- offset = i + tl.arange(0, BLOCK_SIZE)
208
+ req_state_idx = tl.load(idx_mapping + batch_idx)
209
+ start_idx = tl.load(query_start_loc + batch_idx)
210
+ end_idx = tl.load(query_start_loc + batch_idx + 1)
211
+ for i in range(start_idx, end_idx, TRITON_BLOCK_SIZE):
212
+ offset = i + tl.arange(0, TRITON_BLOCK_SIZE)
213
positions = tl.load(pos + offset, mask=offset < end_idx, other=0)
210
- block_indices = positions // page_size
214
+ block_indices = positions // block_size
215
block_numbers = tl.load(
212
- block_table_ptr + req_idx * block_table_stride + block_indices
216
+ block_table_ptr + req_state_idx * block_table_stride + block_indices
217
)
214
- slot_ids = block_numbers * page_size + positions % page_size
218
+ slot_ids = block_numbers * block_size + positions % block_size
219
tl.store(slot_mapping_ptr + offset, slot_ids, mask=offset < end_idx)
220
221
vllm/v1/worker/gpu/model_runner.py
+3
-1
@@ -607,7 +607,9 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
607
608
# Compute slot mappings: [num_kv_cache_groups, num_tokens]
609
slot_mappings = self.block_tables.compute_slot_mappings(
610
- query_start_loc, self.input_buffers.positions[:num_tokens]
610
+ idx_mapping,
611
+ query_start_loc,
612
+ self.input_buffers.positions[:num_tokens],
613
)
614
615
# Layer name -> attention metadata.
vllm/v1/worker/gpu/spec_decode/eagle.py
+8
-3
@@ -138,6 +138,7 @@ class EagleSpeculator:
138
) -> None:
139
pos = self.input_buffers.positions[:num_reqs]
140
query_start_loc = self.input_buffers.query_start_loc[: num_reqs + 1]
141
+ idx_mapping = self.idx_mapping[:num_reqs]
142
for step in range(1, self.num_speculative_steps):
143
# Run the eagle model.
144
last_hidden_states, hidden_states = self.run_model(
@@ -149,7 +150,7 @@ class EagleSpeculator:
150
# used for draft and target sampling.
151
draft_tokens = gumbel_sample(
152
logits,
152
- self.idx_mapping[:num_reqs],
153
+ idx_mapping,
154
self.temperature,
155
self.seeds,
156
pos + 1,
@@ -166,7 +167,9 @@ class EagleSpeculator:
167
self.hidden_states,
168
self.max_model_len,
169
)
169
- self.block_tables.compute_slot_mappings(query_start_loc, pos)
170
+ self.block_tables.compute_slot_mappings(
171
+ idx_mapping, query_start_loc, pos
172
+ )
173
174
def capture_model(self) -> None:
175
if self.num_speculative_steps == 1:
@@ -279,7 +282,9 @@ class EagleSpeculator:
282
self.max_num_reqs,
283
)
284
query_start_loc = self.input_buffers.query_start_loc[: num_reqs + 1]
282
- slot_mappings = self.block_tables.compute_slot_mappings(query_start_loc, pos)
285
+ slot_mappings = self.block_tables.compute_slot_mappings(
286
+ idx_mapping, query_start_loc, pos
287
+ )
288
289
cudagraph_size = self.cudagraph_manager.get_cudagraph_size(num_reqs)
290
if cudagraph_size is not None: