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

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, 04:45 UTC

[TPU] Add TPU specific var VLLM_TPU_MOST_MODEL_LEN (#19919)

Signed-off-by: Chenyaaang <chenyangli@google.com>

Chenyaaang committed Jun 25, 2025 at 15:51 UTC 2d7620c3ebb3a3e0e600dd2781d7e5dfbd1c2382

5 files changed +185 -77

tests/v1/tpu/worker/test_tpu_model_runner.py

+14

@@ -587,3 +587,17 @@ def test_init_kv_cache_with_kv_sharing_valid():

587 assert len(kv_cache_config.kv_cache_groups[0].layer_names) == 2

588 assert kv_cache_config.kv_cache_groups[0].layer_names[0] == layer_0

589 assert kv_cache_config.kv_cache_groups[0].layer_names[1] == layer_1

590 +

591 +

592 +def test_most_model_len(monkeypatch: pytest.MonkeyPatch):

593 + monkeypatch.setenv("VLLM_TPU_MOST_MODEL_LEN", "2048")

594 + vllm_config = get_vllm_config()

595 + vllm_config.model_config.max_model_len = 32000

596 + vllm_config.scheduler_config.max_num_seqs = 1200

597 + model_runner = get_model_runner(vllm_config)

598 +

599 + # verify model runner will adjust num_reqs to avoid SMEM OOM.

600 + assert model_runner.num_reqs_most_model_len == 1200

601 + # num_page_per_req = 32k // 128

602 + # num_reqs = 1024 ** 2 // 2 // num_page_per_req // 4 = 524

603 + assert model_runner.num_reqs_max_model_len == 524

vllm/envs.py

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

119 VLLM_MARLIN_USE_ATOMIC_ADD: bool = False

120 VLLM_V0_USE_OUTLINES_CACHE: bool = False

121 VLLM_TPU_BUCKET_PADDING_GAP: int = 0

122 + VLLM_TPU_MOST_MODEL_LEN: Optional[int] = None

123 VLLM_USE_DEEP_GEMM: bool = False

124 VLLM_XGRAMMAR_CACHE_MB: int = 0

125 VLLM_MSGPACK_ZERO_COPY_THRESHOLD: int = 256

@@ -833,6 +834,8 @@ environment_variables: dict[str, Callable[[], Any]] = {

834 "VLLM_TPU_BUCKET_PADDING_GAP":

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

836 if "VLLM_TPU_BUCKET_PADDING_GAP" in os.environ else 0,

837 + "VLLM_TPU_MOST_MODEL_LEN":

838 + lambda: maybe_convert_int(os.environ.get("VLLM_TPU_MOST_MODEL_LEN", None)),

839

840 # Allow use of DeepGemm kernels for fused moe ops.

841 "VLLM_USE_DEEP_GEMM":

vllm/platforms/tpu.py

-10

@@ -122,16 +122,6 @@ class TpuPlatform(Platform):

122 PallasAttentionBackend)

123 cache_config.block_size = PallasAttentionBackend.get_page_size(

124 vllm_config) # type: ignore[assignment]

125 - min_page_size = PallasAttentionBackend.get_min_page_size(

126 - vllm_config)

127 - if min_page_size > cache_config.block_size:

128 - logger.warning(

129 - "Increase the page size from %s to %s to make sure there's"

130 - "no SMEM OOM",

131 - cache_config.block_size,

132 - min_page_size,

133 - )

134 - cache_config.block_size = min_page_size # type: ignore[assignment]

125

126 parallel_config = vllm_config.parallel_config

127 scheduler_config = vllm_config.scheduler_config

vllm/v1/attention/backends/pallas.py

@@ -71,6 +71,11 @@ class PallasAttentionBackend(AttentionBackend):

71 min_page_size = 1 << (min_page_size - 1).bit_length()

72 return min_page_size

74 + @staticmethod

75 + def get_max_num_seqs(model_len: int, page_size: int) -> int:

76 + num_page_per_req = cdiv(model_len, page_size)

77 + return 1024 * 1024 // 2 // num_page_per_req // 4

78 +

79 # TPU has limited SREGs (scalar registers), if page_size is too small, we

80 # can spill SREGs easily which leads to bad performance. The strategy we

81 # apply here is trying to split max-model-len to 16 pages which make the

vllm/v1/worker/tpu_model_runner.py

+163 -67

@@ -37,8 +37,8 @@ from vllm.v1.core.encoder_cache_manager import compute_encoder_budget

37 from vllm.v1.kv_cache_interface import (AttentionSpec, FullAttentionSpec,

38 KVCacheConfig, KVCacheSpec,

39 SlidingWindowSpec)

40 -from vllm.v1.outputs import (EMPTY_MODEL_RUNNER_OUTPUT, LogprobsTensors,

41 - ModelRunnerOutput)

40 +from vllm.v1.outputs import (EMPTY_MODEL_RUNNER_OUTPUT, LogprobsLists,

41 + LogprobsTensors, ModelRunnerOutput)

42 from vllm.v1.sample.tpu.metadata import TPUSupportedSamplingMetadata

43 from vllm.v1.sample.tpu.sampler import Sampler as TPUSampler

44 from vllm.v1.utils import bind_kv_cache

@@ -150,7 +150,11 @@ class TPUModelRunner(LoRAModelRunnerMixin):

150 self.sliding_window = model_config.get_sliding_window()

151 self.block_size = cache_config.block_size

152 self.max_model_len = model_config.max_model_len

153 + self.most_model_len = envs.VLLM_TPU_MOST_MODEL_LEN

154 self.max_num_blocks_per_req = cdiv(self.max_model_len, self.block_size)

155 + self.num_blocks_per_most_len_req = cdiv(

156 + self.most_model_len,

157 + self.block_size) if self.most_model_len is not None else None

158 # InputBatch needs to work with sampling tensors greater than padding

159 # to avoid dynamic shapes. Also, avoid suboptimal alignment.

160 self.max_num_reqs = max(scheduler_config.max_num_seqs, MIN_NUM_SEQS)

@@ -220,12 +224,19 @@ class TPUModelRunner(LoRAModelRunnerMixin):

224 dtype=torch.int32,

225 device="cpu")

226 self.positions_np = self.positions_cpu.numpy()

223 -

227 self.block_table_cpu = torch.zeros(

228 (self.max_num_reqs, self.max_num_blocks_per_req),

229 dtype=torch.int32,

230 device="cpu")

228 -

231 + # adjust num_reqs to avoid SMEM OOM.

232 + self.num_reqs_most_model_len = min(

233 + PallasAttentionBackend.get_max_num_seqs(self.most_model_len,

234 + self.block_size),

235 + self.max_num_reqs) if self.most_model_len is not None else None

236 + self.num_reqs_max_model_len = min(

237 + PallasAttentionBackend.get_max_num_seqs(self.max_model_len,

238 + self.block_size),

239 + self.max_num_reqs)

240 self.query_start_loc_cpu = torch.zeros(self.max_num_tokens + 1,

241 dtype=torch.int32,

242 device="cpu",

@@ -515,25 +526,50 @@ class TPUModelRunner(LoRAModelRunnerMixin):

526

527 return kv_cache_spec

528

518 - def _prepare_inputs(self, scheduler_output: "SchedulerOutput"):

519 - total_num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens

520 - assert total_num_scheduled_tokens > 0

529 + def _prepare_inputs(self, scheduler_output: "SchedulerOutput",

530 + start_index: int):

531 + assert scheduler_output.total_num_scheduled_tokens > 0

532 num_reqs = self.input_batch.num_reqs

533 assert num_reqs > 0

534 + assert start_index < num_reqs

535

536 # Get the number of scheduled tokens for each request.

537 + use_max_model_len = self.most_model_len is None

538 num_scheduled_tokens_per_req = []

539 max_num_scheduled_tokens_all_reqs = 0

527 - for req_id in self.input_batch.req_ids[:num_reqs]:

540 + end_index = start_index

541 +

542 + # Use either most_model_len or max_model_len depending on request size.

543 + for i in range(start_index, num_reqs):

544 + req_id = self.input_batch.req_ids[i]

545 assert req_id is not None

546 num_tokens = scheduler_output.num_scheduled_tokens[req_id]

547 + if not use_max_model_len and num_tokens > self.most_model_len:

548 + use_max_model_len = True

549 num_scheduled_tokens_per_req.append(num_tokens)

531 - max_num_scheduled_tokens_all_reqs = max(

532 - max_num_scheduled_tokens_all_reqs, num_tokens)

550 + if use_max_model_len:

551 + if len(num_scheduled_tokens_per_req) > self.num_reqs_max_model_len:

552 + num_scheduled_tokens_per_req = \

553 + num_scheduled_tokens_per_req[:self.num_reqs_max_model_len]

554 + end_index = start_index + self.num_reqs_max_model_len

555 + else:

556 + end_index = num_reqs

557 + else:

558 + if len(num_scheduled_tokens_per_req

559 + ) > self.num_reqs_most_model_len:

560 + num_scheduled_tokens_per_req = \

561 + num_scheduled_tokens_per_req[:self.num_reqs_most_model_len]

562 + end_index = start_index + self.num_reqs_most_model_len

563 + else:

564 + end_index = num_reqs

565 + max_num_scheduled_tokens_all_reqs = max(num_scheduled_tokens_per_req)

566 num_scheduled_tokens_per_req = np.array(num_scheduled_tokens_per_req,

567 dtype=np.int32)

568 + total_num_scheduled_tokens = sum(num_scheduled_tokens_per_req)

569 assert max_num_scheduled_tokens_all_reqs > 0

570

571 + num_reqs = len(num_scheduled_tokens_per_req)

572 +

573 # Get request indices.

574 # E.g., [2, 5, 3] -> [0, 0, 1, 1, 1, 1, 1, 2, 2, 2]

575 # For each scheduled token, what are the corresponding req index.

@@ -615,13 +651,29 @@ class TPUModelRunner(LoRAModelRunnerMixin):

651 self.input_batch.block_table[0].

652 slot_mapping_cpu[:padded_total_num_scheduled_tokens].to(

653 self.device))

618 - block_tables = self.block_table_cpu[:self.max_num_reqs]

619 - block_tables[:num_reqs, :self.max_num_blocks_per_req] = (

620 - self.input_batch.block_table[0].get_cpu_tensor()[:num_reqs])

654 + if use_max_model_len:

655 + block_tables = self.block_table_cpu[:self.num_reqs_max_model_len, :

656 + self.max_num_blocks_per_req]

657 + block_tables[:num_reqs, :self.max_num_blocks_per_req] = (

658 + self.input_batch.block_table[0].get_cpu_tensor()[:num_reqs])

659 + query_start_loc = self.query_start_loc_cpu[:self.

660 + num_reqs_max_model_len +

661 + 1].to(self.device)

662 + seq_lens = self.seq_lens_cpu[:self.num_reqs_max_model_len].to(

663 + self.device)

664 + else:

665 + block_tables = self.block_table_cpu[:self.

666 + num_reqs_most_model_len, :self.

667 + num_blocks_per_most_len_req]

668 + block_tables[:num_reqs, :self.num_blocks_per_most_len_req] = (

669 + self.input_batch.block_table[0].get_cpu_tensor()

670 + [:num_reqs, :self.num_blocks_per_most_len_req])

671 + query_start_loc = self.query_start_loc_cpu[:self.

672 + num_reqs_most_model_len +

673 + 1].to(self.device)

674 + seq_lens = self.seq_lens_cpu[:self.num_reqs_most_model_len].to(

675 + self.device)

676 block_tables = block_tables.to(self.device)

622 - query_start_loc = self.query_start_loc_cpu[:self.max_num_reqs + 1].to(

623 - self.device)

624 - seq_lens = self.seq_lens_cpu[:self.max_num_reqs].to(self.device)

677

678 if self.lora_config is not None:

679 # We need to respect padding when activating LoRA adapters

@@ -672,7 +724,8 @@ class TPUModelRunner(LoRAModelRunnerMixin):

724 layer_name: attn_metadata

725 for layer_name in layer_names

726 }

675 - return per_layer_attn_metadata, logits_indices, padded_num_reqs

727 + return per_layer_attn_metadata, logits_indices, padded_num_reqs,\

728 + num_reqs, end_index

729

730 def _scatter_placeholders(

731 self,

@@ -847,52 +900,84 @@ class TPUModelRunner(LoRAModelRunnerMixin):

900 else:

901 mm_embeds = []

902 xm.mark_step()

850 - # Prepare inputs

851 - attn_metadata, logits_indices, padded_num_reqs = self._prepare_inputs(

852 - scheduler_output)

853 - input_ids, inputs_embeds = self._get_model_inputs(

854 - self.input_ids, mm_embeds)

855 - xm.mark_step()

856 - num_reqs = self.input_batch.num_reqs

857 - # Run the decoder

858 - with set_forward_context(

859 - attn_metadata,

860 - self.vllm_config,

861 - num_tokens=scheduler_output.total_num_scheduled_tokens):

862 - hidden_states = self.model(

863 - input_ids=input_ids,

864 - positions=self.position_ids,

865 - inputs_embeds=inputs_embeds,

866 - )

867 - hidden_states = self.select_hidden_states(hidden_states,

868 - logits_indices)

869 - logits = self.compute_logits(hidden_states)

870 - tpu_sampling_metadata = TPUSupportedSamplingMetadata.\

871 - from_input_batch(self.input_batch, padded_num_reqs, self.device)

872 - if scheduler_output.grammar_bitmask is not None:

873 - require_struct_decoding, grammar_bitmask_padded, arange = \

874 - self.prepare_structured_decoding_input(logits, scheduler_output)

875 - logits = self.structured_decode(require_struct_decoding,

876 - grammar_bitmask_padded, logits,

877 - arange)

878 - selected_token_ids = self.sample_from_logits_func(

879 - logits, tpu_sampling_metadata)

880 - # NOTE (NickLucche) Use the original logits (before any penalties or

881 - # temperature scaling) for the top-k logprobs. We can't enforce it due

882 - # to recompilations outside torch.compiled code, so just make sure

883 - # `sample_from_logits` does not modify the logits in-place.

884 - logprobs = self.gather_logprobs(logits, selected_token_ids) \

885 - if tpu_sampling_metadata.logprobs else None

886 -

887 - # Remove padding on cpu and keep dynamic op outside of xla graph.

888 - selected_token_ids = selected_token_ids.cpu()[:num_reqs]

889 - logprobs_lists = logprobs.tolists() \

890 - if tpu_sampling_metadata.logprobs else None

903 + # Prepare inputs, the requests might be splitted into multiple

904 + # executions, combine the result of each execution.

905 + start_index = 0

906 + combined_selected_tokens: list[torch.Tensor] = []

907 + combined_logprobs: list[LogprobsLists] = []

908 + while start_index < self.input_batch.num_reqs:

909 + attn_metadata, logits_indices, padded_num_reqs, num_reqs,\

910 + end_index = self._prepare_inputs(scheduler_output, start_index)

911 + input_ids, inputs_embeds = self._get_model_inputs(

912 + self.input_ids, mm_embeds)

913 + xm.mark_step()

914 + # Run the decoder

915 + with set_forward_context(

916 + attn_metadata,

917 + self.vllm_config,

918 + num_tokens=scheduler_output.total_num_scheduled_tokens):

919 + hidden_states = self.model(

920 + input_ids=input_ids,

921 + positions=self.position_ids,

922 + inputs_embeds=inputs_embeds,

923 + )

924 + hidden_states = self.select_hidden_states(hidden_states,

925 + logits_indices)

926 + logits = self.compute_logits(hidden_states)

927 + tpu_sampling_metadata = TPUSupportedSamplingMetadata.\

928 + from_input_batch(self.input_batch, padded_num_reqs, self.device)

929 + if scheduler_output.grammar_bitmask is not None:

930 + require_struct_decoding, grammar_bitmask_padded, arange = \

931 + self.prepare_structured_decoding_input(logits,

932 + scheduler_output)

933 + logits = self.structured_decode(require_struct_decoding,

934 + grammar_bitmask_padded, logits,

935 + arange)

936 + selected_token_ids = self.sample_from_logits_func(

937 + logits, tpu_sampling_metadata)

938 + # NOTE (NickLucche) Use the original logits (before any penalties or

939 + # temperature scaling) for the top-k logprobs. We can't enforce it

940 + # due to recompilations outside torch.compiled code, so just make

941 + # sure `sample_from_logits` does not modify the logits in-place.

942 + logprobs = self.gather_logprobs(logits, selected_token_ids) \

943 + if tpu_sampling_metadata.logprobs else None

944 +

945 + # Remove padding on cpu and keep dynamic op outside of xla graph.

946 + selected_token_ids = selected_token_ids.cpu()[:num_reqs]

947 +

948 + combined_selected_tokens.append(selected_token_ids)

949 + if tpu_sampling_metadata.logprobs:

950 + combined_logprobs.append(logprobs.tolists())

951 +

952 + start_index = end_index

953 +

954 + selected_token_ids = torch.cat(combined_selected_tokens, dim=0)

955 + if tpu_sampling_metadata.logprobs:

956 +

957 + def concat_lists(input_lists):

958 + result = []

959 + for input_list in input_lists:

960 + result.extend(input_list)

961 + return result

962 +

963 + logprobs_lists = LogprobsLists(logprob_token_ids=concat_lists(

964 + [lp.logprob_token_ids for lp in combined_logprobs]),

965 + logprobs=concat_lists([

966 + lp.logprobs

967 + for lp in combined_logprobs

968 + ]),

969 + sampled_token_ranks=concat_lists([

970 + lp.sampled_token_ranks

971 + for lp in combined_logprobs

972 + ]))

973 + else:

974 + logprobs_lists = None

975

976 # Update the cache state concurrently. Code above will not block until

977 # we use `selected_token_ids`. Add mark_step if post-processing changes

978 request_seq_lens: list[tuple[int, CachedRequestState, int]] = []

979 discard_sampled_tokens_req_indices = []

980 + num_reqs = self.input_batch.num_reqs

981 for i, req_id in zip(range(num_reqs), self.input_batch.req_ids):

982 assert req_id is not None

983 req_state = self.requests[req_id]

@@ -1020,7 +1105,8 @@ class TPUModelRunner(LoRAModelRunnerMixin):

1105 self.sampler = TPUSampler()

1106

1107 @torch.no_grad()

1023 - def _dummy_run(self, num_tokens: int) -> None:

1108 + def _dummy_run(self, num_tokens: int, num_reqs: int,

1109 + num_blocks: int) -> None:

1110 if self.is_multimodal_model:

1111 input_ids = None

1112 inputs_embeds = torch.zeros((num_tokens, self.hidden_size),

@@ -1030,20 +1116,19 @@ class TPUModelRunner(LoRAModelRunnerMixin):

1116 input_ids = torch.zeros((num_tokens),

1117 dtype=torch.int32).to(self.device)

1118 inputs_embeds = None

1033 - actual_num_reqs = min(num_tokens, self.max_num_reqs)

1119 + actual_num_reqs = min(num_tokens, num_reqs)

1120 position_ids = torch.zeros(num_tokens,

1121 dtype=torch.int32).to(self.device)

1122 slot_mapping = torch.zeros(num_tokens,

1123 dtype=torch.int64).to(self.device)

1038 - block_tables = torch.zeros(

1039 - (self.max_num_reqs, self.block_table_cpu.shape[1]),

1040 - dtype=torch.int32).to(self.device)

1041 - query_lens = [1] * self.max_num_reqs

1124 + block_tables = torch.zeros((num_reqs, num_blocks),

1125 + dtype=torch.int32).to(self.device)

1126 + query_lens = [1] * num_reqs

1127 query_start_loc = torch.cumsum(torch.tensor([0] + query_lens,

1128 dtype=torch.int32),

1129 dim=0,

1130 dtype=torch.int32).to(self.device)

1046 - context_lens = torch.ones((self.max_num_reqs, ),

1131 + context_lens = torch.ones((num_reqs, ),

1132 dtype=torch.int32).to(self.device)

1133 num_seqs = torch.tensor([actual_num_reqs],

1134 dtype=torch.int32).to(self.device)

@@ -1061,6 +1146,9 @@ class TPUModelRunner(LoRAModelRunnerMixin):

1146 torch._dynamo.mark_dynamic(input_ids, 0)

1147 torch._dynamo.mark_dynamic(position_ids, 0)

1148 torch._dynamo.mark_dynamic(attn_metadata.slot_mapping, 0)

1149 + torch._dynamo.mark_dynamic(attn_metadata.block_tables, (0, 1))

1150 + torch._dynamo.mark_dynamic(attn_metadata.context_lens, 0)

1151 + torch._dynamo.mark_dynamic(attn_metadata.query_start_loc, 0)

1152

1153 layer_names = get_layers_from_vllm_config(self.vllm_config,

1154 Attention).keys()

@@ -1152,7 +1240,11 @@ class TPUModelRunner(LoRAModelRunnerMixin):

1240 start = time.perf_counter()

1241 for num_tokens in self.num_tokens_paddings:

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

1155 - self._dummy_run(num_tokens)

1243 + self._dummy_run(num_tokens, self.num_reqs_max_model_len,

1244 + self.max_num_blocks_per_req)

1245 + if self.most_model_len is not None:

1246 + self._dummy_run(num_tokens, self.num_reqs_most_model_len,

1247 + self.num_blocks_per_most_len_req)

1248 xm.wait_device_ops()

1249 end = time.perf_counter()

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

@@ -1341,7 +1433,11 @@ class TPUModelRunner(LoRAModelRunnerMixin):

1433 self.encoder_cache["tmp"] = dict(enumerate(dummy_encoder_outputs))

1434

1435 # Trigger compilation for general shape.

1344 - self._dummy_run(num_tokens)

1436 + self._dummy_run(num_tokens, self.num_reqs_max_model_len,

1437 + self.max_num_blocks_per_req)

1438 + if self.most_model_len is not None:

1439 + self._dummy_run(num_tokens, self.num_reqs_most_model_len,

1440 + self.num_blocks_per_most_len_req)

1441

1442 xm.mark_step()

1443 xm.wait_device_ops()