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, 14:45 UTC
[TPU][V1] Fix padding recompilation when `max-num-batched-tokens` is not even (#16726)
Signed-off-by: NickLucche <nlucches@redhat.com>
Nicolò Lucchesi committed
Apr 17, 2025 at 20:09 UTC
5989f4684d62d5cb1852624ce0fd04fc08dd239b
3 files changed
+17
-7
tests/v1/tpu/worker/test_tpu_model_runner.py
+8
@@ -294,11 +294,19 @@ def test_update_states_request_unscheduled(model_runner):
294
295
296
def test_get_paddings():
297
+ # Bucketed padding
298
min_token_size, max_token_size, padding_gap = 16, 512, 64
299
expected_paddings = [16, 32, 64, 128, 192, 256, 320, 384, 448, 512]
300
+ actual_paddings = _get_token_paddings(min_token_size, max_token_size,
301
+ padding_gap)
302
+
303
+ # Bucketed padding with max_token_size not a power of two.
304
+ max_token_size = 317
305
+ expected_paddings = [16, 32, 64, 128, 192, 256, 320]
306
actual_paddings = _get_token_paddings(min_token_size, max_token_size,
307
padding_gap)
308
assert actual_paddings == expected_paddings
309
+
310
# Exponential padding.
311
max_token_size, padding_gap = 1024, 0
312
expected_paddings = [16, 32, 64, 128, 256, 512, 1024]
vllm/v1/worker/tpu_model_runner.py
+7
-5
@@ -128,10 +128,16 @@ class TPUModelRunner:
128
self.block_size = cache_config.block_size
129
self.max_model_len = model_config.max_model_len
130
self.max_num_blocks_per_req = cdiv(self.max_model_len, self.block_size)
131
- self.max_num_tokens = scheduler_config.max_num_batched_tokens
131
# InputBatch needs to work with sampling tensors greater than padding
132
# to avoid dynamic shapes. Also, avoid suboptimal alignment.
133
self.max_num_reqs = max(scheduler_config.max_num_seqs, MIN_NUM_SEQS)
134
+ self.num_tokens_paddings = _get_token_paddings(
135
+ min_token_size=16,
136
+ max_token_size=scheduler_config.max_num_batched_tokens,
137
+ padding_gap=envs.VLLM_TPU_BUCKET_PADDING_GAP)
138
+ # In case `max_num_tokens < max(num_tokens_paddings)` use the actual
139
+ # padded max value to pre-allocate data structures and pre-compile.
140
+ self.max_num_tokens = self.num_tokens_paddings[-1]
141
142
# Model-related.
143
self.num_attn_layers = model_config.get_num_layers_by_block_type(
@@ -211,10 +217,6 @@ class TPUModelRunner:
217
# Range tensor with values [0 .. self.max_num_tokens - 1].
218
# Used to initialize positions / context_lens / seq_lens
219
self.arange_np = np.arange(self.max_num_tokens, dtype=np.int32)
214
- self.num_tokens_paddings = _get_token_paddings(
215
- min_token_size=16,
216
- max_token_size=self.max_num_tokens,
217
- padding_gap=envs.VLLM_TPU_BUCKET_PADDING_GAP)
220
self.num_reqs_paddings = _get_req_paddings(
221
min_req_size=MIN_NUM_SEQS, max_req_size=self.max_num_reqs)
222
vllm/v1/worker/tpu_worker.py
+2
-2
@@ -156,8 +156,8 @@ class TPUWorker:
156
self.vllm_config.compilation_config.static_forward_context,
157
runner_kv_caches)
158
159
- self.model_runner._dummy_run(
160
- self.scheduler_config.max_num_batched_tokens)
159
+ # `max_num_tokens >= max_num_batched_tokens` due to padding.
160
+ self.model_runner._dummy_run(self.model_runner.max_num_tokens)
161
162
# Synchronize before measuring the memory usage.
163
xm.wait_device_ops()