@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / ac3cd6e83c

Fork 0

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, 17:15 UTC

[core] add bucket padding to tpu_model_runner (#14995)

Signed-off-by: Chenyaaang <llccyy1212@gmail.com> Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com> Co-authored-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>

Chenyaaang committed Mar 25, 2025 at 14:27 UTC ac3cd6e83c264255703651c085e567a684112938

3 files changed +63 -19

tests/v1/tpu/worker/test_tpu_model_runner.py

+21 -1

@@ -8,7 +8,9 @@ from vllm.sampling_params import SamplingParams

8 from vllm.v1.core.sched.output import (CachedRequestData, NewRequestData,

9 SchedulerOutput)

10 from vllm.v1.sample.metadata import SamplingMetadata

11 -from vllm.v1.worker.tpu_model_runner import TPUModelRunner

11 +from vllm.v1.worker.tpu_model_runner import (TPUModelRunner,

12 + _get_padded_token_len,

13 + _get_paddings)

15 # Mock torch_xla module since it may not be available in the test environments

16 torch_xla_patcher = mock.patch.dict(

@@ -305,3 +307,21 @@ def test_update_states_request_unscheduled(model_runner):

307

308 assert _is_req_added(model_runner, req_ids[1])

309 assert not _is_req_scheduled(model_runner, req_ids[1])

310 +

311 +

312 +def test_get_paddings():

313 + min_token_size, max_token_size, padding_gap = 16, 512, 64

314 + expected_paddings = [16, 32, 64, 128, 192, 256, 320, 384, 448, 512]

315 + actual_paddings = _get_paddings(min_token_size, max_token_size,

316 + padding_gap)

317 + assert actual_paddings == expected_paddings

318 +

319 +

320 +def test_get_padded_token_len():

321 + min_token_size, max_token_size, padding_gap = 16, 512, 64

322 + paddings = _get_paddings(min_token_size, max_token_size, padding_gap)

323 + assert _get_padded_token_len(paddings, 1) == 16

324 + assert _get_padded_token_len(paddings, 16) == 16

325 + assert _get_padded_token_len(paddings, 20) == 32

326 + assert _get_padded_token_len(paddings, 300) == 320

327 + assert _get_padded_token_len(paddings, 512) == 512

vllm/envs.py

@@ -97,6 +97,7 @@ if TYPE_CHECKING:

97 VLLM_MARLIN_USE_ATOMIC_ADD: bool = False

98 VLLM_V0_USE_OUTLINES_CACHE: bool = False

99 VLLM_TPU_DISABLE_TOPK_TOPP_OPTIMIZATION: bool = False

100 + VLLM_TPU_BUCKET_PADDING_GAP: int = 64

101

102

103 def get_default_cache_root():

@@ -627,6 +628,12 @@ environment_variables: dict[str, Callable[[], Any]] = {

628 "VLLM_TPU_DISABLE_TOPK_TOPP_OPTIMIZATION":

629 lambda: bool(int(os.environ["VLLM_TPU_DISABLE_TOPK_TOPP_OPTIMIZATION"]))

630 if "VLLM_TPU_DISABLE_TOPK_TOPP_OPTIMIZATION" in os.environ else None,

631 +

632 + # Gap between padding buckets for the forward pass. So we have

633 + # 8, we will run forward pass with [16, 24, 32, ...].

634 + "VLLM_TPU_BUCKET_PADDING_GAP":

635 + lambda: int(os.environ["VLLM_TPU_BUCKET_PADDING_GAP"])

636 + if "VLLM_TPU_BUCKET_PADDING_GAP" in os.environ else 64,

637 }

638

639 # end-env-vars-definition

vllm/v1/worker/tpu_model_runner.py

+35 -18

@@ -1,4 +1,5 @@

1 # SPDX-License-Identifier: Apache-2.0

2 +import bisect

3 import time

4 from typing import TYPE_CHECKING, Optional, cast

5 from unittest.mock import patch

@@ -170,6 +171,10 @@ class TPUModelRunner:

171 # Range tensor with values [0 .. self.max_num_tokens - 1].

172 # Used to initialize positions / context_lens / seq_lens

173 self.arange_np = np.arange(self.max_num_tokens, dtype=np.int32)

174 + self.num_tokens_paddings = _get_paddings(

175 + min_token_size=16,

176 + max_token_size=self.max_num_tokens,

177 + padding_gap=envs.VLLM_TPU_BUCKET_PADDING_GAP)

178

179 def _update_states(self, scheduler_output: "SchedulerOutput") -> bool:

180 """Update the cached states and the persistent batch with the scheduler

@@ -422,7 +427,7 @@ class TPUModelRunner:

427

428 # Do the padding and copy the tensors to the TPU.

429 padded_total_num_scheduled_tokens = _get_padded_token_len(

425 - total_num_scheduled_tokens)

430 + self.num_tokens_paddings, total_num_scheduled_tokens)

431 # Zero out to avoid spurious values from prev iteration (last cp chunk)

432 self.input_ids_cpu[

433 total_num_scheduled_tokens:padded_total_num_scheduled_tokens] = 0

@@ -573,7 +578,6 @@ class TPUModelRunner:

578

579 # Prepare inputs

580 attn_metadata, logits_indices = self._prepare_inputs(scheduler_output)

576 -

581 if self.is_multimodal_model:

582 # NOTE(woosuk): To unify token ids and soft tokens (vision

583 # embeddings), we always use embeddings (rather than token ids)

@@ -764,26 +768,21 @@ class TPUModelRunner:

768 logger.info("Compiling the model with different input shapes.")

769

770 start = time.perf_counter()

767 - num_tokens = 16

768 - while True:

771 + for num_tokens in self.num_tokens_paddings:

772 logger.info(" -- num_tokens: %d", num_tokens)

773 self._dummy_run(self.kv_caches, num_tokens)

774 xm.mark_step()

772 - if num_tokens >= self.max_num_tokens:

773 - break

774 - num_tokens *= 2

775 xm.wait_device_ops()

776 end = time.perf_counter()

777 logger.info("Compilation finished in in %.2f [secs].", end - start)

778

779 logger.info("Compiling sampling with different input shapes.")

780 start = time.perf_counter()

781 - num_tokens = 16

781 hsize = self.model_config.get_hidden_size()

782 device = self.device

783 # Compile sampling step for different model+sampler outputs in bucketed

784 # n_tokens x max_num_reqs. Graph is really small so this is fine.

786 - while True:

785 + for num_tokens in self.num_tokens_paddings:

786 num_reqs_to_sample = MIN_NUM_SEQS

787 dummy_hidden = torch.randn((num_tokens, hsize),

788 device=device,

@@ -805,9 +804,6 @@ class TPUModelRunner:

804 if num_reqs_to_sample >= self.max_num_reqs:

805 break

806 num_reqs_to_sample *= 2

808 - if num_tokens >= self.max_num_tokens:

809 - break

810 - num_tokens *= 2

807 xm.wait_device_ops()

808 end = time.perf_counter()

809 logger.info("Compilation finished in in %.2f [secs].", end - start)

@@ -939,12 +935,33 @@ def _get_padded_number(n: int, multiple: int) -> int:

935 return ((n + multiple - 1) // multiple) * multiple

936

937

942 -def _get_padded_token_len(x: int) -> int:

943 - if x <= 16:

944 - return 16

945 - return 1 << (x - 1).bit_length()

946 -

947 -

938 def _get_padded_num_reqs_with_upper_limit(x, upper_limit) -> int:

939 res = MIN_NUM_SEQS if x <= MIN_NUM_SEQS else 1 << (x - 1).bit_length()

940 return min(res, upper_limit)

941 +

942 +

943 +def _get_paddings(min_token_size: int, max_token_size: int,

944 + padding_gap: int) -> list[int]:

945 + """Generate a list of padding size, starting from min_token_size,

946 + ending with a number that can cover max_token_size

947 + first increase the size to twice,

948 + then increase the padding size by padding_gap.

949 + """

950 + paddings = []

951 + num = min_token_size

952 + while num <= padding_gap:

953 + paddings.append(num)

954 + num *= 2

955 + num //= 2

956 + while num < max_token_size:

957 + num += padding_gap

958 + paddings.append(num)

959 + return paddings

960 +

961 +

962 +def _get_padded_token_len(paddings: list[int], x: int) -> int:

963 + """Return the first element in paddings list greater or equal to x.

964 + """

965 + index = bisect.bisect_left(paddings, x)

966 + assert index < len(paddings)

967 + return paddings[index]