@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / d8c6d7d6b5
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, 18:30 UTC

[V1][TPU] Support V1 Sampler for ragged attention (#14227)

Signed-off-by: NickLucche <nlucches@redhat.com>

Nicolò Lucchesi committed Mar 20, 2025 at 05:00 UTC d8c6d7d6b503000e46186d362626135887322006
6 files changed +535 -55
tests/v1/tpu/test_sampler.py new
+94
@@ -0,0 +1,94 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +import tempfile
3 +from time import time
4 +
5 +import pytest
6 +
7 +from vllm import LLM, envs
8 +from vllm.platforms import current_platform
9 +from vllm.sampling_params import SamplingParams
10 +
11 +if not envs.VLLM_USE_V1:
12 + pytest.skip(
13 + "Skipping V1 tests. Rerun with `VLLM_USE_V1=1` to test.",
14 + allow_module_level=True,
15 + )
16 +
17 +
18 +@pytest.mark.parametrize("model_name", ["D4nt3/Qwen2.5-two-layers"])
19 +@pytest.mark.skipif(not current_platform.is_tpu(),
20 + reason="This test needs a TPU")
21 +def test_sampler_compilation(model_name: str, monkeypatch):
22 + """
23 + Check that no recompilation happens despite changing sampling parameters.
24 + We can't read XLA metrics from the engine process, hence we measure time.
25 + """
26 + with tempfile.TemporaryDirectory() as temp_dir:
27 + monkeypatch.setenv("VLLM_XLA_CACHE_PATH", temp_dir)
28 + # Compiling model init may still take some time, enforce_eager to skip.
29 + llm = LLM(model_name,
30 + enforce_eager=True,
31 + max_num_seqs=16,
32 + max_model_len=1024,
33 + gpu_memory_utilization=0.5)
34 + prompts = [
35 + "A robot may not injure a human being",
36 + "It is only with the heart that one can see rightly;",
37 + ]
38 + # First inference should be slow
39 + sampling_params = SamplingParams(
40 + temperature=0.7,
41 + # top_p=0.6, # TODO too slow!
42 + # top_k=10,
43 + min_p=0.2,
44 + max_tokens=16)
45 + s = time()
46 + _ = llm.generate(prompts, sampling_params)
47 + run1 = time() - s
48 +
49 + # Second request with different params, but for which we
50 + # compiled for in previous eager iteration.
51 + sampling_params = SamplingParams(temperature=0.1,
52 + min_p=0.8,
53 + max_tokens=24)
54 + s = time()
55 + _ = llm.generate(prompts, sampling_params)
56 + run2 = time() - s
57 + # Much faster after compiling
58 + assert run1 * 0.1 > run2
59 + print("TIMES", run1, run2)
60 +
61 + # Third request with min_p set to "None". It will not trigger
62 + # recompilation as a default 0 value will be used.
63 + sampling_params = SamplingParams(max_tokens=24, temperature=0.0)
64 + s = time()
65 + _ = llm.generate(prompts, sampling_params)
66 + run3 = time() - s
67 + assert run1 * 0.1 > run3
68 + print("TIMES", run1, run3)
69 +
70 +
71 +@pytest.mark.parametrize("model_name", ["Qwen/Qwen2.5-1.5B-Instruct"])
72 +@pytest.mark.skipif(not current_platform.is_tpu(),
73 + reason="This test needs a TPU")
74 +def test_sampler_different(model_name: str):
75 + """
76 + Test significantly different sampling params to assert the model produces
77 + different results.
78 + """
79 + llm = LLM(
80 + model_name,
81 + enforce_eager=True,
82 + max_num_seqs=1,
83 + max_model_len=64,
84 + # TODO: setting to 0.5 or it will go OOM
85 + gpu_memory_utilization=0.5)
86 + prompts = [
87 + "Write a short story about a robot that dreams for the first time."
88 + ]
89 + sampling_params = SamplingParams(temperature=0.9, min_p=0.2, max_tokens=64)
90 + output = llm.generate(prompts, sampling_params)
91 +
92 + sampling_params = SamplingParams(temperature=0.1, min_p=0.8, max_tokens=64)
93 + output2 = llm.generate(prompts, sampling_params)
94 + assert output[0].outputs[0].text != output2[0].outputs[0].text
vllm/v1/sample/ops/topk_topp_sampler.py
+15 -1
@@ -65,6 +65,8 @@ class TopKTopPSampler(nn.Module):
65 "native implementation of top-p & top-k sampling. For the "
66 "best performance, please install FlashInfer.")
67 self.forward = self.forward_native
68 + elif current_platform.is_tpu():
69 + self.forward = self.forward_tpu
70 else:
71 self.forward = self.forward_native
72
@@ -96,6 +98,18 @@ class TopKTopPSampler(nn.Module):
98 return random_sample(probs, generators)
99 return flashinfer_sample(probs, k, p, generators)
100
101 + def forward_tpu(
102 + self,
103 + logits: torch.Tensor,
104 + generators: dict[int, torch.Generator],
105 + k: Optional[torch.Tensor],
106 + p: Optional[torch.Tensor],
107 + ) -> torch.Tensor:
108 + # TODO Placeholder for TPU optimized topk/p kernel
109 + # logits = apply_top_k_top_p(logits, k, p)
110 + probs = logits.softmax(dim=-1, dtype=torch.float32)
111 + return random_sample(probs, generators)
112 +
113
114 def apply_top_k_top_p(
115 logits: torch.Tensor,
@@ -112,7 +126,7 @@ def apply_top_k_top_p(
126
127 if k is not None:
128 # Apply top-k.
115 - top_k_mask = logits_sort.size(1) - k.to(torch.long)
129 + top_k_mask = logits_sort.size(1) - k.to(torch.long) # shape: B
130 # Get all the top_k values.
131 top_k_mask = logits_sort.gather(1, top_k_mask.unsqueeze(dim=1))
132 top_k_mask = logits_sort < top_k_mask
vllm/v1/sample/tpu/__init__.py
vllm/v1/sample/tpu/metadata.py new
+159
@@ -0,0 +1,159 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +from dataclasses import dataclass, field
3 +from typing import Optional
4 +
5 +import torch
6 +import torch_xla.core.xla_model as xm
7 +
8 +from vllm.v1.sample.metadata import SamplingMetadata
9 +
10 +
11 +@dataclass
12 +class TPUSupportedSamplingMetadata:
13 + # This class exposes a more xla-friendly interface than SamplingMetadata
14 + # on TPU, in particular all arguments should be traceable and no optionals
15 + # are allowed, to avoid graph recompilation on Nones.
16 + temperature: torch.Tensor
17 +
18 + min_p: torch.Tensor
19 + # Still too slow on forward_native!
20 + top_k: torch.Tensor = None
21 + top_p: torch.Tensor = None
22 +
23 + # XLA-unfriendly control flow in Sampler
24 + all_greedy: bool = False
25 + all_random: bool = False
26 + # Greedy sampling flag for compiling single xla graph.
27 + do_argmax: torch.Tensor = None
28 +
29 + # speculation not supported
30 + spec_token_ids = None
31 +
32 + # Generator not supported by xla
33 + generators: dict[int,
34 + torch.Generator] = field(default_factory=lambda: dict())
35 +
36 + # unsupported, you need to return an extra tensor of static size BxV
37 + max_num_logprobs = None
38 +
39 + # TODO No penalties for now
40 + no_penalties: bool = True
41 + prompt_token_ids = None
42 + frequency_penalties = None
43 + presence_penalties = None
44 + repetition_penalties = None
45 + # should use tensor
46 + output_token_ids: list[list[int]] = field(default_factory=lambda: list())
47 +
48 + min_tokens = None # impl is not vectorized
49 +
50 + logit_bias: list[Optional[dict[int, float]]] = field(
51 + default_factory=lambda: list())
52 +
53 + allowed_token_ids_mask = None
54 + bad_words_token_ids = None
55 + indices_do_sample: torch.Tensor = None
56 +
57 + def __post_init__(self):
58 + temp = self.temperature
59 + if self.indices_do_sample is None:
60 + self.indices_do_sample = torch.zeros(temp.shape[0],
61 + device=temp.device,
62 + dtype=torch.int32)
63 + if self.do_argmax is None:
64 + self.do_argmax = torch.tensor(0,
65 + dtype=torch.bool,
66 + device=temp.device)
67 +
68 + @classmethod
69 + def from_sampling_metadata(
70 + cls, metadata: SamplingMetadata,
71 + padded_do_sample_indices: torch.Tensor, num_do_sample: int,
72 + device: torch.device) -> "TPUSupportedSamplingMetadata":
73 + """
74 + Create an XLA-frienly SamplingMetadata structure. Do so by first
75 + instantiating an object with fixed-sized tensors and then writing the
76 + values in input `metadata`. Do that only for non-None values so that
77 + recompilation is not triggered for optional values (None/torch.Tensor).
78 +
79 + In order to handle different sizes for the params that range from 1 up
80 + to `max_num_seqs`, pad tensors to the closest pre-compiled shape.
81 + Same thing for `padded_do_sample_indices`, which contains the indices
82 + to be fed to the Sampler, padded to the closest pre-compiled shape.
83 +
84 + Eg. pad to 4 temperature: [0.7, 0.2]=>[0.7, 0.2, 0.0, 0.0]
85 + do_sample_indices: [4, 10]=>padded_do_sample_indices: [4, 10, 0, 0]
86 + """
87 + metadata = cls._validate_sampling_metadata(metadata)
88 + # NOTE we have to initialize default tensor-based params first and
89 + # skip None values altogether to produce the same xla graph.
90 + num_samples = len(padded_do_sample_indices)
91 + do_argmax = torch.tensor(metadata.all_greedy,
92 + dtype=torch.bool,
93 + device=device)
94 + new_metadata = cls.get_default_sampling_params(num_samples, device,
95 + indices_do_sample=\
96 + padded_do_sample_indices,
97 + do_argmax=do_argmax
98 + )
99 + supported_params = \
100 + TPUSupportedSamplingMetadata._get_default_params_values()
101 + # Copy input non-None values into `new_metadata` fixed-sized tensors.
102 + for p_name in supported_params:
103 + old_val = getattr(metadata, p_name)
104 + new_val = getattr(new_metadata, p_name)
105 + if isinstance(old_val, torch.Tensor):
106 + new_val[:num_do_sample] = old_val
107 + setattr(new_metadata, p_name, new_val)
108 +
109 + xm.mark_step()
110 + xm.wait_device_ops()
111 + return new_metadata
112 +
113 + @classmethod
114 + def get_default_sampling_params(
115 + cls,
116 + num_samples: int,
117 + device: torch.device,
118 + indices_do_sample=None,
119 + do_argmax=None) -> "TPUSupportedSamplingMetadata":
120 + # As sampling happens on a single traced graph, options
121 + # are "disabled" by having them evaluate to an Identity op.
122 + # Note that initialization is dependent on num_samples.
123 + sampling_metadata_disable_value = \
124 + TPUSupportedSamplingMetadata._get_default_params_values()
125 + init_kwargs = dict()
126 + for p_name, (default_val,
127 + dtype) in sampling_metadata_disable_value.items():
128 + default_tensor = torch.full((num_samples, ),
129 + default_val,
130 + dtype=dtype,
131 + device=device)
132 + init_kwargs[p_name] = default_tensor
133 +
134 + return cls(**init_kwargs,
135 + indices_do_sample=indices_do_sample,
136 + do_argmax=do_argmax)
137 +
138 + @staticmethod
139 + def _validate_sampling_metadata(
140 + sampling_metadata: SamplingMetadata) -> SamplingMetadata:
141 + if sampling_metadata.all_greedy:
142 + # Set to None since #13587. Make sure default isn't overruled.
143 + assert sampling_metadata.temperature is None
144 + return sampling_metadata
145 +
146 + @staticmethod
147 + def _get_default_params_values():
148 + return dict(
149 + # Since #13587 greedy sampling requires branching off which leads
150 + # to separate graphs. We set temp to noop and handle argmax here.
151 + temperature=(1.0, torch.float32),
152 + min_p=(0.0, torch.float32),
153 + # strictly disabled for now
154 + # top_k=(-1, torch.int32),
155 + # top_p=(0.0, torch.float32),
156 + # frequency_penalties=(0.0, torch.float32),
157 + # presence_penalties=(0.0, torch.float32),
158 + # repetition_penalties=(0.0, torch.float32),
159 + )
\ No newline at end of file
vllm/v1/sample/tpu/sampler.py new
+154
@@ -0,0 +1,154 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +"""Sampler layer implementing TPU supported operations."""
3 +
4 +import torch
5 +import torch.nn as nn
6 +
7 +from vllm.v1.outputs import LogprobsTensors, SamplerOutput
8 +from vllm.v1.sample.ops.topk_topp_sampler import TopKTopPSampler
9 +from vllm.v1.sample.tpu.metadata import TPUSupportedSamplingMetadata
10 +
11 +_SAMPLING_EPS = 1e-5
12 +
13 +
14 +class Sampler(nn.Module):
15 +
16 + def __init__(self):
17 + super().__init__()
18 + self.topk_topp_sampler = TopKTopPSampler()
19 +
20 + def forward(
21 + self,
22 + logits: torch.Tensor,
23 + sampling_metadata: TPUSupportedSamplingMetadata,
24 + ) -> SamplerOutput:
25 + # NOTE(woosuk): Use the original logits (before any penalties or
26 + # temperature scaling) for the top-k logprobs.
27 + # This is different from the V0 sampler, which uses the logits that
28 + # is used for sampling (after penalties and temperature scaling).
29 +
30 + # Use float32 for the logits.
31 + logits = logits.to(torch.float32)
32 + # Sample the next token.
33 + sampled = self.sample(logits, sampling_metadata)
34 +
35 + # Use int32 to reduce the tensor size.
36 + sampled = sampled.to(torch.int32)
37 +
38 + # These are GPU tensors.
39 + sampler_output = SamplerOutput(
40 + # The sampled tokens are expanded to 2D tensor with shape
41 + # [num_requests, 1], where each row represents one generated
42 + # token per request.
43 + sampled_token_ids=sampled.unsqueeze(-1),
44 + logprobs_tensors=None,
45 + )
46 + return sampler_output
47 +
48 + def apply_temperature(
49 + self,
50 + logits: torch.Tensor,
51 + temp: torch.Tensor,
52 + ) -> torch.Tensor:
53 + # Use in-place division to avoid creating a new tensor.
54 + return logits.div_(temp.unsqueeze(dim=1))
55 +
56 + def greedy_sample(self, logits: torch.Tensor) -> torch.Tensor:
57 + return logits.argmax(dim=-1).view(-1)
58 +
59 + def sample(
60 + self,
61 + logits: torch.Tensor,
62 + sampling_metadata: TPUSupportedSamplingMetadata,
63 + ) -> torch.Tensor:
64 + greedy_sampled = self.greedy_sample(logits)
65 +
66 + assert sampling_metadata.temperature is not None
67 +
68 + # Apply temperature.
69 + logits = self.apply_temperature(logits, sampling_metadata.temperature)
70 +
71 + # Apply min_p.
72 + if sampling_metadata.min_p is not None:
73 + logits = self.apply_min_p(logits, sampling_metadata.min_p)
74 +
75 + # Apply top_k and/or top_p.
76 + random_sampled = self.topk_topp_sampler(
77 + logits,
78 + sampling_metadata.generators,
79 + sampling_metadata.top_k,
80 + sampling_metadata.top_p,
81 + )
82 +
83 + sampled = torch.where(sampling_metadata.temperature < _SAMPLING_EPS,
84 + greedy_sampled, random_sampled)
85 + return sampled
86 +
87 + def compute_logprobs(self, logits: torch.Tensor) -> torch.Tensor:
88 + return logits.log_softmax(dim=-1, dtype=torch.float32)
89 +
90 + def gather_logprobs(
91 + self,
92 + logprobs: torch.Tensor,
93 + num_logprobs: int,
94 + token_ids: torch.Tensor,
95 + ) -> LogprobsTensors:
96 + """
97 + Gather logprobs for topk and sampled/prompt token.
98 +
99 + Args:
100 + logits: (num tokens) x (vocab) tensor
101 + num_logprobs: minimum number of logprobs to
102 + retain per token
103 + token_ids: prompt tokens (if prompt logprobs)
104 + or sampled tokens (if sampled
105 + logprobs); 1D token ID tensor
106 + with (num tokens) elements
107 +
108 + Returns:
109 + Top-k int indices tensor, (num tokens) x (num_logprobs + 1)
110 + Top-k float logprobs tensor, (num tokens) x (num_logprobs + 1)
111 + Sampled token rank tensor, (num tokens)
112 + """
113 + # Find the topK values.
114 + topk_logprobs, topk_indices = torch.topk(logprobs,
115 + num_logprobs,
116 + dim=-1)
117 +
118 + # Get with the logprob of the prompt or sampled token.
119 + token_ids = token_ids.unsqueeze(-1)
120 + token_logprobs = logprobs.gather(-1, token_ids)
121 +
122 + # Compute the ranks of the actual token.
123 + token_ranks = (logprobs >= token_logprobs).sum(-1)
124 +
125 + # Concatenate together with the topk.
126 + indices = torch.cat((token_ids, topk_indices), dim=1)
127 + logprobs = torch.cat((token_logprobs, topk_logprobs), dim=1)
128 +
129 + # Use int32 to reduce the tensor size.
130 + indices = indices.to(torch.int32)
131 +
132 + return LogprobsTensors(indices, logprobs, token_ranks)
133 +
134 + def apply_min_p(
135 + self,
136 + logits: torch.Tensor,
137 + min_p: torch.Tensor,
138 + ) -> torch.Tensor:
139 + """
140 + Filters logits using adaptive probability thresholding.
141 + """
142 + # Convert logits to probability distribution
143 + probability_values = torch.nn.functional.softmax(logits, dim=-1)
144 + # Calculate maximum probabilities per sequence
145 + max_probabilities = torch.amax(probability_values,
146 + dim=-1,
147 + keepdim=True)
148 + # Reshape min_p for broadcasting
149 + adjusted_min_p = min_p.unsqueeze(1) * max_probabilities
150 + # Identify valid tokens using threshold comparison
151 + valid_token_mask = probability_values >= adjusted_min_p
152 + # Apply mask using boolean indexing (xla friendly)
153 + logits.masked_fill_(~valid_token_mask, -float("inf"))
154 + return logits
vllm/v1/worker/tpu_model_runner.py
+113 -54
@@ -23,13 +23,16 @@ from vllm.multimodal.utils import group_mm_inputs_by_modality
23 from vllm.sampling_params import SamplingType
24 from vllm.sequence import IntermediateTensors
25 from vllm.utils import LayerBlockType, cdiv, is_pin_memory_available
26 -from vllm.v1.attention.backends.pallas import (PallasAttentionBackend,
26 +from vllm.v1.attention.backends.pallas import (NUM_KV_PAGES_PER_BLOCK,
27 + PallasAttentionBackend,
28 PallasMetadata)
29 from vllm.v1.core.encoder_cache_manager import compute_encoder_budget
30 from vllm.v1.kv_cache_interface import (FullAttentionSpec, KVCacheConfig,
31 KVCacheSpec)
32 from vllm.v1.outputs import (EMPTY_MODEL_RUNNER_OUTPUT, LogprobsTensors,
32 - ModelRunnerOutput)
33 + ModelRunnerOutput, SamplerOutput)
34 +from vllm.v1.sample.tpu.metadata import TPUSupportedSamplingMetadata
35 +from vllm.v1.sample.tpu.sampler import Sampler as TPUSampler
36 from vllm.v1.utils import bind_kv_cache
37 from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch
38
@@ -42,6 +45,8 @@ logger = init_logger(__name__)
45 # FIXME(woosuk): Find a more reliable way to prevent possible bugs.
46 _PAD_SLOT_ID = 1_000_000_000
47 INVALID_TOKEN_ID = -1
48 +# Smallest output size
49 +MIN_NUM_SEQS = 8
50
51
52 class TPUModelRunner:
@@ -138,8 +143,10 @@ class TPUModelRunner:
143 device="cpu")
144 self.slot_mapping_np = self.slot_mapping_cpu.numpy()
145
146 + padded_max_num_blocks_per_req = _get_padded_number(
147 + self.max_num_blocks_per_req, NUM_KV_PAGES_PER_BLOCK)
148 self.block_table_cpu = torch.zeros(
142 - (self.max_num_tokens, self.max_num_blocks_per_req),
149 + (self.max_num_tokens, padded_max_num_blocks_per_req),
150 dtype=self.input_batch.block_table.get_cpu_tensor().dtype,
151 device="cpu")
152
@@ -267,6 +274,9 @@ class TPUModelRunner:
274 req_data.num_computed_tokens)
275 self.input_batch.block_table.append_row(req_data.new_block_ids,
276 req_index)
277 + # Check if the batch has changed. If not, we can skip copying the
278 + # sampling metadata from CPU to GPU.
279 + batch_changed = len(removed_req_indices) > 0 or len(req_ids_to_add) > 0
280
281 # Add the new or resumed requests to the persistent batch.
282 # The smaller empty indices are filled first.
@@ -284,6 +294,10 @@ class TPUModelRunner:
294 # Condense the batched states if there are empty indices.
295 if removed_req_indices:
296 self.input_batch.condense(removed_req_indices)
297 +
298 + # TODO This slices tensors to copy to device, triggering recompilation.
299 + if batch_changed:
300 + self.input_batch.refresh_sampling_metadata()
301 return len(unscheduled_req_ids) > 0 or len(req_ids_to_add) > 0
302
303 def get_model(self) -> nn.Module:
@@ -447,6 +461,8 @@ class TPUModelRunner:
461 # TODO: Support prompt logprobs.
462 padded_num_reqs = _get_padded_num_reqs_with_upper_limit(
463 num_reqs, self.max_num_reqs)
464 + # Indices at which we sample (positions of last token in the sequence).
465 + # Padded to avoid recompiling when `num_reqs` varies.
466 logits_indices = self.query_start_loc_cpu[1:padded_num_reqs + 1] - 1
467 logits_indices = logits_indices.to(self.device)
468 return attn_metadata, logits_indices
@@ -576,7 +592,14 @@ class TPUModelRunner:
592 # then the embedding layer is not included in the CUDA graph.
593 input_ids = self.input_ids
594 inputs_embeds = None
579 -
595 + sampling_metadata = self.input_batch.sampling_metadata
596 + num_reqs = self.input_batch.num_reqs
597 + # NOTE (NickLucche) here we sync with TPU: if there's any shape
598 + # mismatch in pre-processing, it will trigger a small recompilation
599 + # of the code thus far. Forward graph remains untouched.
600 + tpu_sampling_metadata = TPUSupportedSamplingMetadata.\
601 + from_sampling_metadata(sampling_metadata, logits_indices,
602 + num_reqs, self.device)
603 # Run the decoder
604 with set_forward_context(attn_metadata, self.vllm_config):
605 hidden_states = self.model(
@@ -585,12 +608,13 @@ class TPUModelRunner:
608 kv_caches=self.kv_caches,
609 inputs_embeds=inputs_embeds,
610 )
588 - num_reqs = self.input_batch.num_reqs
589 - selected_token_ids = self.model.compute_logits(hidden_states,
590 - logits_indices, None)
611 + selected_token_ids = self.model.sample_from_hidden(
612 + hidden_states, tpu_sampling_metadata)
613 + # Remove padding on cpu and keep dynamic op outside of xla graph.
614 selected_token_ids = selected_token_ids.cpu()[:num_reqs]
615
593 - # Then, let's update the cache state.
616 + # Update the cache state concurrently. Code above will not block until
617 + # we use `selected_token_ids`. Add mark_step if post-processing changes
618 request_seq_lens: list[tuple[int, CachedRequestState, int]] = []
619 for i, req_id in zip(range(num_reqs), self.input_batch.req_ids):
620 assert req_id is not None
@@ -607,7 +631,6 @@ class TPUModelRunner:
631 # This relies on cuda-specific torch-internal impl details
632 generator.set_offset(generator.get_offset() - 4)
633
610 - # num_reqs entries should be non-None
634 assert all(
635 req_id is not None for req_id in
636 self.input_batch.req_ids[:num_reqs]), "req_ids contains None"
@@ -620,6 +643,7 @@ class TPUModelRunner:
643 max_gen_len = selected_token_ids.shape[-1]
644 if max_gen_len == 1:
645 valid_sampled_token_ids = selected_token_ids.tolist()
646 +
647 for i, req_state, seq_len in request_seq_lens:
648 token_id = valid_sampled_token_ids[i][0]
649 self.input_batch.token_ids_cpu[i, seq_len] = token_id
@@ -676,11 +700,8 @@ class TPUModelRunner:
700 fullgraph=True,
701 dynamic=False)
702
679 - def _dummy_run(
680 - self,
681 - kv_caches,
682 - num_tokens: int,
683 - ) -> None:
703 + @torch.no_grad()
704 + def _dummy_run(self, kv_caches, num_tokens: int) -> None:
705 if self.is_multimodal_model:
706 input_ids = None
707 inputs_embeds = torch.zeros((num_tokens, self.hidden_size),
@@ -729,32 +750,10 @@ class TPUModelRunner:
750 torch._dynamo.mark_dynamic(attn_metadata.slot_mapping, 0)
751
752 with set_forward_context(attn_metadata, self.vllm_config, 0):
732 - assert self.model is not None
733 - hidden_states = self.model(
734 - input_ids=input_ids,
735 - positions=position_ids,
736 - kv_caches=kv_caches,
737 - inputs_embeds=inputs_embeds,
738 - )
739 - num_reqs = _get_padded_num_reqs_with_upper_limit(
740 - 64, self.max_num_reqs)
741 - # NOTE(chengjiyao): In total, the compute_logits function utilizes a
742 - # compilation cache size of token_bucket_num multiplied by
743 - # req_bucket_num. This is acceptable, given the graph's relatively
744 - # small size.
745 - while True:
746 - logits_indices = torch.zeros(
747 - num_reqs,
748 - dtype=torch.int32,
749 - device=self.device,
750 - )
751 - torch._dynamo.mark_dynamic(hidden_states, 0)
752 - torch._dynamo.mark_dynamic(logits_indices, 0)
753 - self.model.compute_logits(hidden_states, logits_indices, None)
754 - if num_reqs >= self.max_num_reqs:
755 - break
756 - num_reqs = _get_padded_num_reqs_with_upper_limit(
757 - num_reqs + 1, self.max_num_reqs)
753 + self.model(input_ids=input_ids,
754 + positions=position_ids,
755 + kv_caches=kv_caches,
756 + inputs_embeds=inputs_embeds)
757
758 def capture_model(self) -> None:
759 """Compile the model."""
@@ -764,13 +763,51 @@ class TPUModelRunner:
763 start = time.perf_counter()
764 num_tokens = 16
765 while True:
767 - self._dummy_run(self.kv_caches, num_tokens)
766 logger.info(" -- num_tokens: %d", num_tokens)
767 + self._dummy_run(self.kv_caches, num_tokens)
768 xm.mark_step()
770 - xm.wait_device_ops()
769 if num_tokens >= self.max_num_tokens:
770 break
771 num_tokens *= 2
772 + xm.wait_device_ops()
773 + end = time.perf_counter()
774 + logger.info("Compilation finished in in %.2f [secs].", end - start)
775 +
776 + logger.info("Compiling sampling with different input shapes.")
777 + start = time.perf_counter()
778 + num_tokens = 16
779 + hsize = self.model_config.get_hidden_size()
780 + device = self.device
781 + # Compile sampling step for different model+sampler outputs in bucketed
782 + # n_tokens x max_num_reqs. Graph is really small so this is fine.
783 + while True:
784 + num_reqs_to_sample = MIN_NUM_SEQS
785 + dummy_hidden = torch.randn((num_tokens, hsize),
786 + device=device,
787 + dtype=torch.bfloat16)
788 + while True:
789 + # Default metadata is an all_greedy setup. But since the
790 + # `do_argmax` flag is a tensor, we still compile the full graph
791 + meta = self.input_batch.sampling_metadata
792 + indices = torch.zeros(
793 + num_reqs_to_sample,
794 + dtype=torch.int32,
795 + device=device,
796 + )
797 + sampling_meta = TPUSupportedSamplingMetadata.\
798 + from_sampling_metadata(meta, indices,
799 + num_reqs_to_sample, device)
800 + logger.info(" -- num_tokens: %d, num_seqs: %d", num_tokens,
801 + num_reqs_to_sample)
802 + self.model.sample_from_hidden(dummy_hidden, sampling_meta)
803 + xm.mark_step()
804 + if num_reqs_to_sample >= self.max_num_reqs:
805 + break
806 + num_reqs_to_sample *= 2
807 + if num_tokens >= self.max_num_tokens:
808 + break
809 + num_tokens *= 2
810 + xm.wait_device_ops()
811 end = time.perf_counter()
812 logger.info("Compilation finished in in %.2f [secs].", end - start)
813
@@ -818,6 +855,13 @@ class ModelWrapperV1(nn.Module):
855 def __init__(self, model: nn.Module):
856 super().__init__()
857 self.model = model
858 + self.sampler = TPUSampler()
859 +
860 + def sample(
861 + self, logits: torch.Tensor,
862 + sampling_metadata: TPUSupportedSamplingMetadata) -> SamplerOutput:
863 + sampler_out = self.sampler(logits, sampling_metadata)
864 + return sampler_out
865
866 def forward(
867 self,
@@ -826,7 +870,7 @@ class ModelWrapperV1(nn.Module):
870 kv_caches: list[tuple[torch.Tensor, torch.Tensor]],
871 inputs_embeds: Optional[torch.Tensor] = None,
872 ) -> torch.Tensor:
829 - """Executes the forward pass of the model and samples the next token.
873 + """Executes the forward pass of the model.
874
875 Args:
876 input_ids: The input token IDs of shape [num_tokens].
@@ -837,7 +881,6 @@ class ModelWrapperV1(nn.Module):
881 hidden_size]. It is used for multimodal models.
882 """
883
840 - assert self.model is not None
884 hidden_states = self.model(
885 input_ids=input_ids,
886 positions=positions,
@@ -846,17 +889,33 @@ class ModelWrapperV1(nn.Module):
889
890 return hidden_states
891
849 - @torch.compile(backend="openxla", fullgraph=True, dynamic=False)
850 - def compute_logits(
892 + def sample_from_hidden(
893 self,
894 hidden_states: torch.Tensor,
853 - logits_indices: torch.Tensor,
854 - sampling_metadata,
855 - ) -> Optional[torch.Tensor]:
856 - hidden_states = hidden_states[logits_indices]
857 - logits = self.model.compute_logits(hidden_states, sampling_metadata)
858 - selected_token_ids = torch.argmax(logits, dim=-1, keepdim=True)
859 - return selected_token_ids
895 + sampling_metadata: TPUSupportedSamplingMetadata,
896 + ) -> torch.Tensor:
897 + """
898 + Sample with xla-friendly function. This function is to be traced
899 + separately from `forward` for lighter compilation overhead.
900 + """
901 + # Tensor `sample_hidden_states` is of fixed pre-compiled size.
902 + sample_hidden_states = \
903 + hidden_states[sampling_metadata.indices_do_sample]
904 + logits = self.compute_logits(sample_hidden_states)
905 + # Greedy sampling can't be run without branching the graph on Sampler.
906 + # Therefore do_argmax/all_greedy is checked here in a xla-friendly way.
907 + # NOTE do_argmax is a scalar, this is just an optimized if/else.
908 + out_tokens = torch.where(sampling_metadata.do_argmax,
909 + torch.argmax(logits, dim=-1, keepdim=True),
910 + self.sample(logits, sampling_metadata)\
911 + .sampled_token_ids)
912 + return out_tokens
913 +
914 + def compute_logits(self,
915 + hidden_states: torch.Tensor) -> Optional[torch.Tensor]:
916 + # SamplingMetadata here for pruning output in LogitsProcessor, disabled
917 + logits = self.model.compute_logits(hidden_states, None)
918 + return logits
919
920 def get_multimodal_embeddings(self, *args, **kwargs):
921 return self.model.get_multimodal_embeddings(*args, **kwargs)
@@ -876,5 +935,5 @@ def _get_padded_token_len(x: int) -> int:
935
936
937 def _get_padded_num_reqs_with_upper_limit(x, upper_limit) -> int:
879 - res = 64 if x <= 64 else 1 << (x - 1).bit_length()
938 + res = MIN_NUM_SEQS if x <= MIN_NUM_SEQS else 1 << (x - 1).bit_length()
939 return min(res, upper_limit)