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

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

[Hardware][Intel-Gaudi] Multi-step scheduling implementation for HPU (#12779)

Signed-off-by: Tomasz Zielinski <tomasz.zielinski@intel.com>

Tomasz Zielinski committed Apr 11, 2025 at 16:38 UTC 34b2cf3b3366443e6b68a43572aee1800f6f2b90

3 files changed +439 -116

vllm/platforms/hpu.py

+3 -3

@@ -46,15 +46,15 @@ class HpuPlatform(Platform):

46 def check_and_update_config(cls, vllm_config: VllmConfig) -> None:

48 scheduler_config = vllm_config.scheduler_config

49 + parallel_config = vllm_config.parallel_config

50 if scheduler_config.is_multi_step:

50 - raise NotImplementedError(

51 - "Multi-step execution is not implemented for HPU")

51 + parallel_config.worker_cls = \

52 + "vllm.worker.multi_step_hpu_worker.MultiStepHPUWorker"

54 if vllm_config.speculative_config is not None:

55 raise NotImplementedError(

56 "Speculative decoding is not implemented for HPU")

57 - parallel_config = vllm_config.parallel_config

58 if parallel_config.worker_cls == "auto":

59 parallel_config.worker_cls = "vllm.worker.hpu_worker.HPUWorker"

vllm/worker/hpu_model_runner.py

+314 -113

@@ -32,6 +32,7 @@ from vllm_hpu_extension.profiler import (HabanaHighLevelProfiler,

32 import vllm.envs as envs

33 from vllm.attention import AttentionMetadata, get_attn_backend

34 from vllm.config import DeviceConfig, VllmConfig

35 +from vllm.distributed import broadcast_tensor_dict

36 from vllm.distributed.parallel_state import get_world_group

37 from vllm.forward_context import set_forward_context

38 from vllm.logger import init_logger

@@ -44,11 +45,13 @@ from vllm.model_executor.layers.sampler import SamplerOutput

45 from vllm.model_executor.layers.vocab_parallel_embedding import (

46 VocabParallelEmbedding)

47 from vllm.model_executor.model_loader import get_model

48 +from vllm.model_executor.sampling_metadata import SequenceGroupToSample

49 from vllm.multimodal import (MULTIMODAL_REGISTRY, BatchedTensorInputs,

50 MultiModalKwargs)

51 from vllm.sampling_params import SamplingParams

50 -from vllm.sequence import (IntermediateTensors, SequenceData,

51 - SequenceGroupMetadata)

52 +from vllm.sequence import (CompletionSequenceGroupOutput, IntermediateTensors,

53 + Logprob, SequenceData, SequenceGroupMetadata,

54 + SequenceOutput)

55 from vllm.utils import (bind_kv_cache, is_pin_memory_available,

56 make_tensor_with_pad)

57 from vllm.worker.model_runner_base import (

@@ -100,7 +103,10 @@ def subtuple(obj: object,

103 if to_override is None:

104 to_override = {}

105 fields = set(to_copy) | set(to_override.keys())

103 - values = {f: to_override.get(f, getattr(obj, f)) for f in fields}

106 + if type(obj) is dict:

107 + values = {key: obj[key] for key in fields if key in obj}

108 + else:

109 + values = {f: to_override.get(f, getattr(obj, f)) for f in fields}

110 if typename not in _TYPE_CACHE:

111 _TYPE_CACHE[typename] = collections.namedtuple(typename,

112 ' '.join(fields))

@@ -533,6 +539,8 @@ class ModelInputForHPU(ModelRunnerInputBase):

539 virtual_engine: int = 0

540 lora_ids: Optional[List[int]] = None

541 async_callback: Optional[Callable] = None

542 + is_first_multi_step: bool = True

543 + is_last_step: bool = True

544

545 def as_broadcastable_tensor_dict(self) -> Dict[str, Any]:

546 tensor_dict = {

@@ -545,6 +553,8 @@ class ModelInputForHPU(ModelRunnerInputBase):

553 "batch_size_padded": self.batch_size_padded,

554 "virtual_engine": self.virtual_engine,

555 "lora_ids": self.lora_ids,

556 + "is_first_multi_step": self.is_first_multi_step,

557 + "is_last_step": self.is_last_step,

558 }

559 _add_attn_metadata_broadcastable_dict(tensor_dict, self.attn_metadata)

560 return tensor_dict

@@ -656,6 +666,9 @@ class HPUModelRunnerBase(ModelRunnerBase[TModelInputForHPU]):

666 self._set_gc_threshold()

667 self.use_contiguous_pa = envs.VLLM_USE_HPU_CONTIGUOUS_CACHE_FETCH

668

669 + # For multi-step scheduling

670 + self.cached_step_outputs: List[torch.Tensor] = []

671 +

672 def _set_gc_threshold(self) -> None:

673 # Read https://docs.python.org/3/library/gc.html#gc.set_threshold

674 # for comprehensive description of gc generations.

@@ -1005,6 +1018,7 @@ class HPUModelRunnerBase(ModelRunnerBase[TModelInputForHPU]):

1018 def _prepare_decode(

1019 self,

1020 seq_group_metadata_list: List[SequenceGroupMetadata],

1021 + output=None,

1022 ) -> PrepareDecodeMetadata:

1023 input_tokens: List[List[int]] = []

1024 input_positions: List[List[int]] = []

@@ -1035,8 +1049,9 @@ class HPUModelRunnerBase(ModelRunnerBase[TModelInputForHPU]):

1049

1050 for seq_id in seq_ids:

1051 seq_data = seq_group_metadata.seq_data[seq_id]

1038 - generation_token = seq_data.get_last_token_id()

1039 - input_tokens.append([generation_token])

1052 + if output is None:

1053 + generation_token = seq_data.get_last_token_id()

1054 + input_tokens.append([generation_token])

1055

1056 seq_len = seq_data.get_len()

1057 position = seq_len - 1

@@ -1047,6 +1062,9 @@ class HPUModelRunnerBase(ModelRunnerBase[TModelInputForHPU]):

1062 seq_lens.append(seq_len)

1063

1064 block_table = seq_group_metadata.block_tables[seq_id]

1065 + num_fully_occupied_blocks = position // self.block_size

1066 + block_table = block_table[:num_fully_occupied_blocks + 1]

1067 +

1068 if len(block_table) == 0:

1069 block_number = _PAD_BLOCK_ID

1070 else:

@@ -1066,9 +1084,14 @@ class HPUModelRunnerBase(ModelRunnerBase[TModelInputForHPU]):

1084 block_table = block_table[-sliding_window_blocks:]

1085 block_tables.append(block_table)

1086

1069 - input_tokens = torch.tensor(input_tokens,

1070 - dtype=torch.long,

1071 - device=self.device)

1087 + if output is None:

1088 + input_tokens = torch.tensor(input_tokens,

1089 + dtype=torch.long,

1090 + device=self.device)

1091 + else:

1092 + real_batch_size = len(seq_group_metadata_list)

1093 + input_tokens = output[:real_batch_size]

1094 +

1095 input_positions = torch.tensor(input_positions,

1096 dtype=torch.long,

1097 device=self.device)

@@ -1462,7 +1485,27 @@ class HPUModelRunnerBase(ModelRunnerBase[TModelInputForHPU]):

1485 profiler.start()

1486 for _ in range(times):

1487 inputs = self.prepare_model_input(seqs)

1465 - self.execute_model(inputs, None, warmup_mode=True)

1488 + is_single_step = \

1489 + self.vllm_config.scheduler_config.num_scheduler_steps == 1

1490 + if is_prompt or is_single_step:

1491 + self.execute_model(inputs, None, warmup_mode=True)

1492 + else: # decode with multi-step

1493 + inputs = dataclasses.replace(inputs,

1494 + is_first_multi_step=True,

1495 + is_last_step=False)

1496 + self.execute_model(inputs,

1497 + None,

1498 + warmup_mode=True,

1499 + num_steps=2,

1500 + seqs=seqs)

1501 + inputs = dataclasses.replace(inputs,

1502 + is_first_multi_step=False,

1503 + is_last_step=True)

1504 + self.execute_model(inputs,

1505 + None,

1506 + warmup_mode=True,

1507 + num_steps=2,

1508 + seqs=seqs)

1509 torch.hpu.synchronize()

1510 if profiler:

1511 profiler.step()

@@ -1985,115 +2028,273 @@ class HPUModelRunner(HPUModelRunnerBase[ModelInputForHPUWithSamplingMetadata]):

2028 intermediate_tensors: Optional[IntermediateTensors] = None,

2029 num_steps: int = 1,

2030 warmup_mode=False,

2031 + seqs=None,

2032 ) -> Optional[Union[List[SamplerOutput], IntermediateTensors]]:

1989 - if num_steps > 1:

1990 - raise ValueError(

1991 - "num_steps > 1 is not supported in HPUModelRunner")

2033 + if not model_input.is_first_multi_step:

2034 + if not model_input.is_last_step:

2035 + # not first or last multi-step

2036 + return []

2037 + # last multi-step

2038 + output = self._decode_sampler_outputs(

2039 + model_input) if self.is_driver_worker else []

2040 + torch.hpu.synchronize()

2041 + if model_input.is_first_multi_step:

2042 + # first multi-step

2043 + if self.lora_config:

2044 + assert model_input.lora_requests is not None

2045 + assert model_input.lora_mapping is not None

2046 + self.set_active_loras(model_input.lora_requests,

2047 + model_input.lora_mapping)

2048 + input_tokens = model_input.input_tokens

2049 + input_positions = model_input.input_positions

2050 + attn_metadata = model_input.attn_metadata

2051 + sampling_metadata = model_input.sampling_metadata

2052 + real_batch_size = model_input.real_batch_size

2053 + batch_size_padded = model_input.batch_size_padded

2054 + assert input_tokens is not None

2055 + assert input_positions is not None

2056 + assert sampling_metadata is not None

2057 + assert attn_metadata is not None

2058 + is_prompt = attn_metadata.is_prompt

2059 + assert is_prompt is not None

2060 + batch_size = input_tokens.size(0)

2061 + seq_len = self._seq_len(attn_metadata)

2062 + use_graphs = self._use_graphs(batch_size, seq_len, is_prompt)

2063 + self._check_config(batch_size, seq_len, is_prompt, warmup_mode)

2064 +

2065 + lora_mask: torch.Tensor = None

2066 + lora_logits_mask: torch.Tensor = None

2067 + if self.lora_config:

2068 + assert model_input.lora_ids is not None

2069 + lora_mask, lora_logits_mask = self.create_lora_mask(

2070 + input_tokens, model_input.lora_ids,

2071 + attn_metadata.is_prompt)

2072 +

2073 + execute_model_kwargs = {

2074 + "input_ids": input_tokens,

2075 + "positions": input_positions,

2076 + "attn_metadata": self.trim_attn_metadata(attn_metadata),

2077 + "intermediate_tensors": intermediate_tensors,

2078 + "lora_mask": lora_mask,

2079 + "virtual_engine": model_input.virtual_engine,

2080 + **(model_input.multi_modal_kwargs or {}),

2081 + }

2082 + if htorch.utils.internal.is_lazy():

2083 + execute_model_kwargs.update(

2084 + {"bypass_hpu_graphs": not use_graphs})

2085

1993 - if self.lora_config:

1994 - assert model_input.lora_requests is not None

1995 - assert model_input.lora_mapping is not None

1996 - self.set_active_loras(model_input.lora_requests,

1997 - model_input.lora_mapping)

1998 - input_tokens = model_input.input_tokens

1999 - input_positions = model_input.input_positions

2000 - attn_metadata = model_input.attn_metadata

2001 - sampling_metadata = model_input.sampling_metadata

2002 - real_batch_size = model_input.real_batch_size

2003 - batch_size_padded = model_input.batch_size_padded

2004 - assert input_tokens is not None

2005 - assert input_positions is not None

2006 - assert sampling_metadata is not None

2007 - assert attn_metadata is not None

2008 - is_prompt = attn_metadata.is_prompt

2009 - assert is_prompt is not None

2010 - batch_size = input_tokens.size(0)

2011 - seq_len = self._seq_len(attn_metadata)

2012 - use_graphs = self._use_graphs(batch_size, seq_len, is_prompt)

2013 - self._check_config(batch_size, seq_len, is_prompt, warmup_mode)

2086 + htorch.core.mark_step()

2087 + if self.is_driver_worker:

2088 + model_event_name = ("model_"

2089 + f"{'prompt' if is_prompt else 'decode'}_"

2090 + f"bs{batch_size}_"

2091 + f"seq{seq_len}_"

2092 + f"graphs{'T' if use_graphs else 'F'}")

2093 + else:

2094 + model_event_name = 'model_executable'

2095 + if num_steps > 1:

2096 + # in case of multi-step scheduling

2097 + # we only want to pythonize in the last step

2098 + sampling_metadata.skip_sampler_cpu_output = True

2099 + self.model.model.sampler.include_gpu_probs_tensor = True

2100 + cache_orig_output_tokens_len: List[Dict] = []

2101 +

2102 + def try_revert_dummy_output_tokens():

2103 + if len(cache_orig_output_tokens_len) > 0:

2104 + # Reuse the original output token ids length

2105 + for i, seq_group_metadata in enumerate(

2106 + seq_group_metadata_list):

2107 + for j, data in seq_group_metadata.seq_data.items():

2108 + orig_output_tokens_len = \

2109 + cache_orig_output_tokens_len[i][j]

2110 + data.output_token_ids = \

2111 + data.output_token_ids[:orig_output_tokens_len]

2112 +

2113 + for i in range(num_steps):

2114 + if i != 0 and not self.is_driver_worker:

2115 + broadcast_data = broadcast_tensor_dict(src=0)

2116 + if 'early_exit' in broadcast_data and broadcast_data[

2117 + 'early_exit']:

2118 + return [output] if num_steps == 1 else []

2119 + execute_model_kwargs.update({

2120 + "input_ids":

2121 + broadcast_data["input_ids"],

2122 + "positions":

2123 + broadcast_data["positions"],

2124 + "attn_metadata":

2125 + self.trim_attn_metadata(

2126 + broadcast_data["attn_metadata"])

2127 + })

2128 + with self.profiler.record_event('internal', model_event_name):

2129 + hidden_states = self.model.forward(

2130 + **execute_model_kwargs,

2131 + selected_token_indices=sampling_metadata.

2132 + selected_token_indices)

2133 +

2134 + if self.lora_config:

2135 + LoraMask.setLoraMask(

2136 + lora_logits_mask.index_select(

2137 + 0, sampling_metadata.selected_token_indices))

2138 +

2139 + # Compute the logits.

2140 + with self.profiler.record_event(

2141 + 'internal',

2142 + ('compute_logits_'

2143 + f'{"prompt" if is_prompt else "decode"}_bs'

2144 + f'{batch_size}_'

2145 + f'seq{seq_len}')):

2146 + if num_steps == 1:

2147 + sampling_metadata.selected_token_indices = None

2148 + logits = self.model.compute_logits(hidden_states,

2149 + sampling_metadata)

2150 + htorch.core.mark_step()

2151 + # Only perform sampling in the driver worker.

2152 + if not self.is_driver_worker:

2153 + continue

2154

2015 - lora_mask: torch.Tensor = None

2016 - lora_logits_mask: torch.Tensor = None

2017 - if self.lora_config:

2018 - assert model_input.lora_ids is not None

2019 - lora_mask, lora_logits_mask = self.create_lora_mask(

2020 - input_tokens, model_input.lora_ids, attn_metadata.is_prompt)

2021 -

2022 - execute_model_kwargs = {

2023 - "input_ids": input_tokens,

2024 - "positions": input_positions,

2025 - "attn_metadata": self.trim_attn_metadata(attn_metadata),

2026 - "intermediate_tensors": intermediate_tensors,

2027 - "lora_mask": lora_mask,

2028 - "virtual_engine": model_input.virtual_engine,

2029 - **(model_input.multi_modal_kwargs or {}),

2030 - }

2031 - if htorch.utils.internal.is_lazy():

2032 - execute_model_kwargs.update({"bypass_hpu_graphs": not use_graphs})

2033 -

2034 - htorch.core.mark_step()

2035 - if self.is_driver_worker:

2036 - model_event_name = ("model_"

2037 - f"{'prompt' if is_prompt else 'decode'}_"

2038 - f"bs{batch_size}_"

2039 - f"seq{seq_len}_"

2040 - f"graphs{'T' if use_graphs else 'F'}")

2155 + if model_input.async_callback is not None:

2156 + model_input.async_callback()

2157 + # Sample the next token.

2158 + with self.profiler.record_event(

2159 + 'internal', ('sample_'

2160 + f'{"prompt" if is_prompt else "decode"}_'

2161 + f'bs{batch_size}_'

2162 + f'seq{seq_len}')):

2163 + output = self.model.sample(

2164 + logits=logits,

2165 + sampling_metadata=sampling_metadata,

2166 + )

2167 + if num_steps > 1:

2168 + output = output.sampled_token_ids

2169 + self.cached_step_outputs.append(

2170 + output.detach().clone())

2171 + htorch.core.mark_step()

2172 + if i < num_steps - 1:

2173 + if i == 0:

2174 + if model_input.async_callback is not None:

2175 + ctx = model_input.async_callback.keywords[ # type: ignore

2176 + "ctx"]

2177 + seq_group_metadata_list = \

2178 + ctx.seq_group_metadata_list

2179 + elif seqs is not None:

2180 + seq_group_metadata_list = seqs

2181 + else:

2182 + raise RuntimeError(

2183 + "seq_group_metadata_list is uninitialized")

2184 + for i, seq_group_metadata in enumerate(

2185 + seq_group_metadata_list):

2186 + # Skip empty steps

2187 + seq_group_metadata.state.current_step += (

2188 + num_steps - 2)

2189 + # Cache the original output token ids

2190 + cache_orig_output_tokens_len.append({})

2191 + for j, data in seq_group_metadata.seq_data.items():

2192 + cache_orig_output_tokens_len[i][j] = \

2193 + len(data.output_token_ids)

2194 + for seq_group_metadata in seq_group_metadata_list:

2195 + for data in seq_group_metadata.seq_data.values():

2196 + max_output_len = sampling_metadata.seq_groups[

2197 + 0].sampling_params.max_tokens

2198 + if len(data.output_token_ids) < max_output_len - 1:

2199 + # add a place holder for prepare_decode

2200 + # arbitrary value, this could be any token

2201 + dummy_token = (540, )

2202 + data.output_token_ids += (dummy_token)

2203 + else:

2204 + broadcast_tensor_dict({'early_exit': True},

2205 + src=0)

2206 + if num_steps == 1:

2207 + return [output]

2208 + else:

2209 + try_revert_dummy_output_tokens()

2210 + return []

2211 +

2212 + result = self._prepare_decode(seq_group_metadata_list,

2213 + output=output)

2214 + execute_model_kwargs.update({

2215 + "input_ids":

2216 + result.input_tokens,

2217 + "positions":

2218 + result.input_positions,

2219 + "attn_metadata":

2220 + self.trim_attn_metadata(result.attn_metadata)

2221 + })

2222 + model_kwargs_broadcast_data = {

2223 + "input_ids": result.input_tokens,

2224 + "positions": result.input_positions,

2225 + "attn_metadata": vars(result.attn_metadata)

2226 + }

2227 + broadcast_tensor_dict(model_kwargs_broadcast_data, src=0)

2228 + else:

2229 + try_revert_dummy_output_tokens()

2230 +

2231 + if self.is_driver_worker and self.profiler.enabled:

2232 + # Stop recording 'execute_model' event

2233 + self.profiler.end()

2234 + event_end = self.profiler.get_timestamp_us()

2235 + counters = self.profiler_counter_helper.get_counter_dict(

2236 + cache_config=self.cache_config,

2237 + duration=event_end - self.event_start,

2238 + seq_len=seq_len,

2239 + batch_size_padded=batch_size_padded,

2240 + real_batch_size=real_batch_size,

2241 + is_prompt=is_prompt)

2242 + self.profiler.record_counter(self.event_start, counters)

2243 + if num_steps == 1:

2244 + return [output] if self.is_driver_worker else []

2245 + else:

2246 + return []

2247 + return output if type(output) is list else [output]

2248 +

2249 + def _decode_sampler_outputs(self, model_input):

2250 + use_async_out_proc = model_input.async_callback is not None

2251 + sampler_outputs = []

2252 + num_outputs = len(self.cached_step_outputs)

2253 + for i in range(num_outputs):

2254 + next_token_ids = self.cached_step_outputs.pop(0)

2255 + next_token_ids = next_token_ids.cpu().tolist()

2256 + sampler_output = self._make_decode_output(

2257 + next_token_ids, model_input.sampling_metadata.seq_groups)

2258 + sampler_outputs.append(sampler_output)

2259 +

2260 + if i < num_outputs - 1 and use_async_out_proc:

2261 + assert model_input.async_callback is not None

2262 + ctx = model_input.async_callback.keywords[ # type: ignore

2263 + "ctx"]

2264 + ctx.append_output(

2265 + outputs=[sampler_output],

2266 + seq_group_metadata_list=ctx.seq_group_metadata_list,

2267 + scheduler_outputs=ctx.scheduler_outputs,

2268 + is_async=False,

2269 + is_last_step=False,

2270 + is_first_step_output=False)

2271 + model_input.async_callback()

2272 +

2273 + if use_async_out_proc:

2274 + return [sampler_outputs[-1]]

2275 else:

2042 - model_event_name = 'model_executable'

2043 - with self.profiler.record_event('internal', model_event_name):

2044 - hidden_states = self.model.forward(

2045 - **execute_model_kwargs,

2046 - selected_token_indices=sampling_metadata.selected_token_indices

2047 - )

2276 + return sampler_outputs

2277

2049 - if self.lora_config:

2050 - LoraMask.setLoraMask(

2051 - lora_logits_mask.index_select(

2052 - 0, sampling_metadata.selected_token_indices))

2053 -

2054 - # Compute the logits.

2055 - with self.profiler.record_event(

2056 - 'internal', ('compute_logits_'

2057 - f'{"prompt" if is_prompt else "decode"}_bs'

2058 - f'{batch_size}_'

2059 - f'seq{seq_len}')):

2060 - sampling_metadata.selected_token_indices = None

2061 - logits = self.model.compute_logits(hidden_states,

2062 - sampling_metadata)

2063 - htorch.core.mark_step()

2064 - # Only perform sampling in the driver worker.

2065 - if not self.is_driver_worker:

2066 - return []

2067 -

2068 - if model_input.async_callback is not None:

2069 - model_input.async_callback()

2070 -

2071 - # Sample the next token.

2072 - with self.profiler.record_event(

2073 - 'internal', ('sample_'

2074 - f'{"prompt" if is_prompt else "decode"}_'

2075 - f'bs{batch_size}_'

2076 - f'seq{seq_len}')):

2077 - output = self.model.sample(

2078 - logits=logits,

2079 - sampling_metadata=sampling_metadata,

2080 - )

2081 - output.outputs = output.outputs[:real_batch_size]

2082 - htorch.core.mark_step()

2083 -

2084 - if self.is_driver_worker and self.profiler.enabled:

2085 - # Stop recording 'execute_model' event

2086 - self.profiler.end()

2087 - event_end = self.profiler.get_timestamp_us()

2088 - counters = self.profiler_counter_helper.get_counter_dict(

2089 - cache_config=self.cache_config,

2090 - duration=event_end - self.event_start,

2091 - seq_len=seq_len,

2092 - batch_size_padded=batch_size_padded,

2093 - real_batch_size=real_batch_size,

2094 - is_prompt=is_prompt)

2095 - self.profiler.record_counter(self.event_start, counters)

2096 - return [output]

2278 + def _make_decode_output(

2279 + self,

2280 + next_token_ids: List[List[int]],

2281 + seq_groups: List[SequenceGroupToSample],

2282 + ) -> SamplerOutput:

2283 + zero_logprob = Logprob(0.0)

2284 + sampler_outputs = []

2285 + batch_idx = 0

2286 + for seq_group in seq_groups:

2287 + seq_ids = seq_group.seq_ids

2288 + seq_outputs = []

2289 + for seq_id in seq_ids:

2290 + next_token_id = next_token_ids[batch_idx][0]

2291 + seq_outputs.append(

2292 + SequenceOutput(seq_id, next_token_id,

2293 + {next_token_id: zero_logprob}))

2294 + batch_idx += 1

2295 + sampler_outputs.append(

2296 + CompletionSequenceGroupOutput(seq_outputs, None))

2297 + return SamplerOutput(sampler_outputs)

2298

2299 def shutdown_inc(self):

2300 can_finalize_inc = False

vllm/worker/multi_step_hpu_worker.py new

+122

@@ -0,0 +1,122 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +

3 +###############################################################################

5 +###############################################################################

6 +

7 +import dataclasses

8 +from typing import Dict, Optional, Tuple

9 +

10 +import torch

11 +

12 +from vllm.distributed import broadcast_tensor_dict

13 +from vllm.sequence import ExecuteModelRequest

14 +from vllm.worker.hpu_model_runner import ModelInputForHPU

15 +from vllm.worker.hpu_worker import HPUWorker

16 +from vllm.worker.worker_base import WorkerInput

17 +

18 +

19 +class MultiStepHPUWorker(HPUWorker):

20 +

21 + def __init__(self, *args, **kwargs):

22 + super().__init__(*args, **kwargs)

23 + self.cached_model_input: Optional[ModelInputForHPU] = None

24 +

25 + def _get_driver_input_and_broadcast(

26 + self, execute_model_req: ExecuteModelRequest

27 + ) -> Tuple[ModelInputForHPU, WorkerInput, Dict[str, torch.Tensor]]:

28 + """

29 + Get the driver input and broadcast it to other workers.

30 + """

31 + assert self.is_driver_worker

32 + assert execute_model_req.virtual_engine == 0

33 +

34 + is_first_multi_step = execute_model_req.is_first_multi_step

35 + is_last_step = execute_model_req.is_last_step

36 +

37 + if is_first_multi_step:

38 + # on first step we prepare the worker input and model input normally

39 + worker_input: WorkerInput = self.prepare_worker_input(

40 + execute_model_req=execute_model_req)

41 + worker_input = dataclasses.replace(

42 + worker_input,

43 + num_steps=execute_model_req.num_lookahead_slots + 1)

44 + model_input: ModelInputForHPU = (

45 + self.model_runner.prepare_model_input(

46 + execute_model_req.seq_group_metadata_list,

47 + execute_model_req.virtual_engine,

48 + execute_model_req.finished_requests_ids))

49 +

50 + if execute_model_req.async_callback:

51 + model_input = dataclasses.replace(

52 + model_input,

53 + async_callback=execute_model_req.async_callback)

54 + else:

55 + # on subsequent steps we reuse the worker input and model input

56 + assert self.cached_model_input is not None

57 + model_input = self.cached_model_input

58 + worker_input = WorkerInput()

59 +

60 + model_input = dataclasses.replace(

61 + model_input,

62 + is_first_multi_step=is_first_multi_step,

63 + is_last_step=is_last_step)

64 +

65 + if self.do_metadata_broadcast:

66 + if is_first_multi_step:

67 + broadcast_data = worker_input.as_broadcastable_tensor_dict()

68 + broadcast_data.update(

69 + model_input.as_broadcastable_tensor_dict())

70 + broadcast_tensor_dict(broadcast_data, src=0)

71 + else:

72 + broadcast_data = {

73 + "is_first_multi_step": is_first_multi_step,

74 + "is_last_step": is_last_step,

75 + }

76 + broadcast_tensor_dict(broadcast_data, src=0)

77 +

78 + # Returning empty dict here to keep this compatible with

79 + # `LocalOrDistributedWorkerBase._get_driver_input_and_broadcast`

80 + return model_input, worker_input, {}

81 +

82 + def prepare_input(

83 + self,

84 + execute_model_req: Optional[ExecuteModelRequest] = None,

85 + ) -> Optional[Tuple[ModelInputForHPU, WorkerInput, Dict[str,

86 + torch.Tensor]]]:

87 + if self.is_driver_worker:

88 + if execute_model_req is None:

89 + if self.do_metadata_broadcast:

90 + # This signals that there's no more requests to process for

91 + # now. All workers are running infinite loop with

92 + # broadcast_tensor_dict, and it stops the loop when the

93 + # driver broadcasts an empty input. Send an empty input to

94 + # notify all other workers to stop their execution loop.

95 + broadcast_tensor_dict({}, src=0)

96 + return None

97 + model_input, worker_input, _ = self._get_driver_input_and_broadcast(

98 + execute_model_req)

99 + if model_input.is_first_multi_step:

100 + self.cached_model_input = model_input

101 + return model_input, worker_input, {}

102 + else:

103 + broadcast_data = broadcast_tensor_dict(src=0)

104 + if not broadcast_data:

105 + return None

106 +

107 + if len(broadcast_data) == 2:

108 + assert self.cached_model_input is not None

109 + self.cached_model_input = dataclasses.replace(

110 + self.cached_model_input,

111 + is_first_multi_step=broadcast_data["is_first_multi_step"],

112 + is_last_step=broadcast_data["is_last_step"])

113 + empty_worker_input = WorkerInput()

114 + return self.cached_model_input, empty_worker_input, {}

115 +

116 + worker_input = WorkerInput.from_broadcasted_tensor_dict(

117 + broadcast_data)

118 + model_input = (

119 + self.model_runner.

120 + make_model_input_from_broadcasted_tensor_dict(broadcast_data))

121 + self.cached_model_input = model_input

122 + return model_input, worker_input, {}