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

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, 01:00 UTC

[V0 deprecation] Remove V0 HPU backend (#21131)

Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>

Woosuk Kwon committed Jul 17, 2025 at 16:37 UTC 4de7146351d67be0010b7007ba4da48462962153

27 files changed +10 -3926

docker/Dockerfile.hpu deleted

-21

@@ -1,21 +0,0 @@

1 -FROM vault.habana.ai/gaudi-docker/1.20.1/ubuntu22.04/habanalabs/pytorch-installer-2.6.0:latest

2 -

3 -COPY ./ /workspace/vllm

4 -

5 -WORKDIR /workspace/vllm

6 -

7 -RUN pip install -v -r requirements/hpu.txt

8 -

9 -ENV no_proxy=localhost,127.0.0.1

10 -ENV PT_HPU_ENABLE_LAZY_COLLECTIVES=true

11 -

12 -RUN VLLM_TARGET_DEVICE=hpu python3 setup.py install

13 -

14 -# install development dependencies (for testing)

15 -RUN python3 -m pip install -e tests/vllm_test_utils

16 -

17 -WORKDIR /workspace/

18 -

19 -RUN ln -s /workspace/vllm/tests && ln -s /workspace/vllm/examples && ln -s /workspace/vllm/benchmarks

20 -

21 -ENTRYPOINT ["python3", "-m", "vllm.entrypoints.openai.api_server"]

requirements/hpu.txt deleted

-12

@@ -1,12 +0,0 @@

1 -# Common dependencies

2 --r common.txt

3 -

4 -# Dependencies for HPU code

5 -ray

6 -triton==3.1.0

7 -pandas

8 -numpy==1.26.4

9 -tabulate

10 -setuptools>=77.0.3,<80.0.0

11 -setuptools-scm>=8

12 -vllm-hpu-extension @ git+https://github.com/HabanaAI/vllm-hpu-extension.git@f1f6624

setup.py

+2 -34

@@ -410,29 +410,6 @@ class repackage_wheel(build_ext):

410 package_data[package_name].append(file_name)

411

412

413 -def _is_hpu() -> bool:

414 - # if VLLM_TARGET_DEVICE env var was set explicitly, skip HPU autodetection

415 - if os.getenv("VLLM_TARGET_DEVICE", None) == VLLM_TARGET_DEVICE:

416 - return VLLM_TARGET_DEVICE == "hpu"

417 -

418 - # if VLLM_TARGET_DEVICE was not set explicitly, check if hl-smi succeeds,

419 - # and if it doesn't, check if habanalabs driver is loaded

420 - is_hpu_available = False

421 - try:

422 - out = subprocess.run(["hl-smi"], capture_output=True, check=True)

423 - is_hpu_available = out.returncode == 0

424 - except (FileNotFoundError, PermissionError, subprocess.CalledProcessError):

425 - if sys.platform.startswith("linux"):

426 - try:

427 - output = subprocess.check_output(

428 - 'lsmod | grep habanalabs | wc -l', shell=True)

429 - is_hpu_available = int(output) > 0

430 - except (ValueError, FileNotFoundError, PermissionError,

431 - subprocess.CalledProcessError):

432 - pass

433 - return is_hpu_available

434 -

435 -

413 def _no_device() -> bool:

414 return VLLM_TARGET_DEVICE == "empty"

415

@@ -440,7 +417,7 @@ def _no_device() -> bool:

417 def _is_cuda() -> bool:

418 has_cuda = torch.version.cuda is not None

419 return (VLLM_TARGET_DEVICE == "cuda" and has_cuda

443 - and not (_is_neuron() or _is_tpu() or _is_hpu()))

420 + and not (_is_neuron() or _is_tpu()))

421

422

423 def _is_hip() -> bool:

@@ -573,12 +550,6 @@ def get_vllm_version() -> str:

550 if neuron_version != MAIN_CUDA_VERSION:

551 neuron_version_str = neuron_version.replace(".", "")[:3]

552 version += f"{sep}neuron{neuron_version_str}"

576 - elif _is_hpu():

577 - # Get the Intel Gaudi Software Suite version

578 - gaudi_sw_version = str(get_gaudi_sw_version())

579 - if gaudi_sw_version != MAIN_CUDA_VERSION:

580 - gaudi_sw_version = gaudi_sw_version.replace(".", "")[:3]

581 - version += f"{sep}gaudi{gaudi_sw_version}"

553 elif _is_tpu():

554 version += f"{sep}tpu"

555 elif _is_cpu():

@@ -625,8 +596,6 @@ def get_requirements() -> list[str]:

596 requirements = _read_requirements("rocm.txt")

597 elif _is_neuron():

598 requirements = _read_requirements("neuron.txt")

628 - elif _is_hpu():

629 - requirements = _read_requirements("hpu.txt")

599 elif _is_tpu():

600 requirements = _read_requirements("tpu.txt")

601 elif _is_cpu():

@@ -635,8 +604,7 @@ def get_requirements() -> list[str]:

604 requirements = _read_requirements("xpu.txt")

605 else:

606 raise ValueError(

638 - "Unsupported platform, please use CUDA, ROCm, Neuron, HPU, "

639 - "or CPU.")

607 + "Unsupported platform, please use CUDA, ROCm, Neuron, or CPU.")

608 return requirements

609

610

vllm/_custom_ops.py

+1 -2

@@ -13,8 +13,7 @@ from vllm.scalar_type import ScalarType

14 logger = init_logger(__name__)

16 -if not current_platform.is_tpu() and not current_platform.is_hpu()\

17 - and not current_platform.is_xpu():

16 +if not current_platform.is_tpu() and not current_platform.is_xpu():

17 try:

18 import vllm._C

19 except ImportError as e:

vllm/attention/backends/hpu_attn.py deleted

-319

@@ -1,319 +0,0 @@

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

2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 -

4 -###############################################################################

6 -###############################################################################

7 -

8 -from dataclasses import dataclass

9 -from typing import Any, Dict, List, Optional, Tuple, Type

10 -

11 -import torch

12 -import vllm_hpu_extension.kernels as kernels

13 -import vllm_hpu_extension.ops as ops

14 -from vllm_hpu_extension.flags import enabled_flags

15 -from vllm_hpu_extension.utils import Matmul, Softmax, VLLMKVCache

16 -

17 -from vllm.attention.backends.abstract import (AttentionBackend, AttentionImpl,

18 - AttentionLayer,

19 - AttentionMetadata, AttentionType,

20 - is_quantized_kv_cache)

21 -from vllm.attention.backends.utils import CommonAttentionState

22 -from vllm.attention.ops.hpu_paged_attn import (HPUPagedAttention,

23 - HPUPagedAttentionMetadata)

24 -from vllm.logger import init_logger

25 -

26 -logger = init_logger(__name__)

27 -

28 -

29 -class HPUAttentionBackend(AttentionBackend):

30 -

31 - @staticmethod

32 - def get_name() -> str:

33 - return "HPU_ATTN"

34 -

35 - @staticmethod

36 - def get_impl_cls() -> Type["HPUAttentionImpl"]:

37 - return HPUAttentionImpl

38 -

39 - @staticmethod

40 - def get_metadata_cls() -> Type["AttentionMetadata"]:

41 - return HPUAttentionMetadata

42 -

43 - @staticmethod

44 - def get_state_cls() -> Type["CommonAttentionState"]:

45 - return CommonAttentionState

46 -

47 - @staticmethod

48 - def get_kv_cache_shape(

49 - num_blocks: int,

50 - block_size: int,

51 - num_kv_heads: int,

52 - head_size: int,

53 - ) -> Tuple[int, ...]:

54 - return HPUPagedAttention.get_kv_cache_shape(num_blocks, block_size,

55 - num_kv_heads, head_size)

56 -

57 - @staticmethod

58 - def swap_blocks(

59 - src_kv_cache: torch.Tensor,

60 - dst_kv_cache: torch.Tensor,

61 - src_to_dsts: torch.Tensor,

62 - ) -> None:

63 - HPUPagedAttention.swap_blocks(src_kv_cache, dst_kv_cache, src_to_dsts)

64 -

65 - @staticmethod

66 - def copy_blocks(

67 - kv_caches: List[torch.Tensor],

68 - src_to_dsts: torch.Tensor,

69 - ) -> None:

70 - HPUPagedAttention.copy_blocks(kv_caches, src_to_dsts)

71 -

72 -

73 -@dataclass

74 -class HPUAttentionMetadata(HPUPagedAttentionMetadata, AttentionMetadata):

75 - """Metadata for HPUAttentionbackend."""

76 - # Currently, input sequences can only contain all prompts

77 - # or all decoding. True if all sequences are prompts.

78 - is_prompt: bool

79 - attn_bias: Optional[torch.Tensor]

80 - seq_lens_tensor: Optional[torch.Tensor]

81 - context_lens_tensor: Optional[torch.Tensor]

82 -

83 -

84 -class HPUAttentionImpl(AttentionImpl, torch.nn.Module):

85 - """

86 - If the input tensors contain prompt tokens, the layout is as follows:

87 - |<--------------- num_prefill_tokens ----------------->|

89 -

90 - Otherwise, the layout is as follows:

91 - |<----------------- num_decode_tokens ------------------>|

93 -

94 - Generation tokens can contain padding when cuda-graph is used.

95 - Currently, prompt tokens don't contain any padding.

96 -

97 - The prompts might have different lengths, while the generation tokens

98 - always have length 1.

99 - """

100 -

101 - def __init__(

102 - self,

103 - num_heads: int,

104 - head_size: int,

105 - scale: float,

106 - num_kv_heads: int,

107 - alibi_slopes: Optional[List[float]],

108 - sliding_window: Optional[int],

109 - kv_cache_dtype: str,

110 - blocksparse_params: Optional[Dict[str, Any]] = None,

111 - max_seq_len: int = 4096,

112 - attn_type: str = AttentionType.DECODER,

113 - kv_sharing_target_layer_name: Optional[str] = None,

114 - use_irope: bool = False,

115 - ) -> None:

116 - super(AttentionImpl, self).__init__()

117 - if kv_sharing_target_layer_name is not None:

118 - raise NotImplementedError("KV sharing is not supported in V0 "

119 - "HPU_ATTN backend.")

120 - if use_irope:

121 - logger.warning_once(

122 - "Using irope in HPU is not supported yet, it will fall back "

123 - "to global attention for long context.")

124 - self.kv_cache_dtype = kv_cache_dtype

125 - self.num_heads = num_heads

126 - self.head_size = head_size

127 - self.scale = float(scale)

128 - self.matmul_qk = Matmul()

129 - self.softmax = Softmax()

130 - self.matmul_av = Matmul()

131 - self.batch2block_matmul = Matmul()

132 - self.block2batch_matmul = Matmul()

133 - self.k_cache = VLLMKVCache()

134 - self.v_cache = VLLMKVCache()

135 - self.fused_scaled_dot_product_attention = kernels.fsdpa()

136 -

137 - self.prefill_impl = 'naive'

138 - if "flex_attention" in enabled_flags():

139 - self.prefill_impl = 'flex'

140 - if "fsdpa" in enabled_flags():

141 - assert alibi_slopes is None, \

142 - 'Prefill with FusedSDPA not supported with alibi slopes!'

143 - self.prefill_impl = 'fsdpa'

144 -

145 - self.num_kv_heads = num_heads if num_kv_heads is None else num_kv_heads

146 - self.sliding_window = sliding_window

147 - self.alibi_slopes = alibi_slopes

148 - if alibi_slopes is not None:

149 - alibi_slopes_tensor = torch.tensor(alibi_slopes,

150 - dtype=torch.bfloat16)

151 - self.alibi_slopes = alibi_slopes_tensor

152 - self.num_queries_per_kv = self.num_heads // self.num_kv_heads

153 -

154 - if self.prefill_impl == 'fsdpa':

155 - assert alibi_slopes is None, \

156 - 'Prefill with FusedSDPA not supported with alibi slopes!'

157 -

158 - supported_head_sizes = HPUPagedAttention.get_supported_head_sizes()

159 - if head_size not in supported_head_sizes:

160 - raise ValueError(

161 - f"Head size {head_size} is not supported by PagedAttention. "

162 - f"Supported head sizes are: {supported_head_sizes}.")

163 -

164 - self.attn_type = attn_type

165 - if self.attn_type != AttentionType.DECODER:

166 - raise NotImplementedError("Encoder self-attention and "

167 - "encoder/decoder cross-attention "

168 - "are not implemented for "

169 - "HPUAttentionImpl")

170 -

171 - if is_quantized_kv_cache(self.kv_cache_dtype):

172 - raise NotImplementedError(

173 - "HPUAttention with FP8 KV cache not yet supported")

174 -

175 - def forward(

176 - self,

177 - layer: AttentionLayer,

178 - query: torch.Tensor,

179 - key: torch.Tensor,

180 - value: torch.Tensor,

181 - kv_cache: torch.Tensor,

182 - attn_metadata: HPUAttentionMetadata,

183 - output: Optional[torch.Tensor] = None,

184 - output_scale: Optional[torch.Tensor] = None,

185 - ) -> torch.Tensor:

186 - """Forward pass with xFormers and PagedAttention.

187 -

188 - Args:

189 - query: shape = [num_tokens, num_heads * head_size]

190 - key: shape = [num_tokens, num_kv_heads * head_size]

191 - value: shape = [num_tokens, num_kv_heads * head_size]

192 - kv_cache = [2, num_blocks, block_size * num_kv_heads * head_size]

193 - attn_metadata: Metadata for attention.

194 - Returns:

195 - shape = [num_tokens, num_heads * head_size]

196 - """

197 - if output_scale is not None:

198 - raise NotImplementedError(

199 - "fused output quantization is not yet supported"

200 - " for HPUAttentionImpl")

201 -

202 - batch_size, seq_len, hidden_size = query.shape

203 - _, seq_len_kv, _ = key.shape

204 -

205 - key = key.view(-1, self.num_kv_heads, self.head_size)

206 - value = value.view(-1, self.num_kv_heads, self.head_size)

207 - block_indices = attn_metadata.block_indices

208 - block_offsets = attn_metadata.block_offsets

209 - key_cache = None

210 - value_cache = None

211 - if attn_metadata.is_prompt and self.attn_type \

212 - is not AttentionType.ENCODER_ONLY:

213 - key = key.unflatten(0, (block_indices.size(0), -1))

214 - value = value.unflatten(0, (block_indices.size(0), -1))

215 - if kv_cache is not None and isinstance(kv_cache, tuple):

216 - key_cache, value_cache = HPUPagedAttention.split_kv_cache(

217 - kv_cache, self.num_kv_heads, self.head_size)

218 -

219 - # Reshape the input keys and values and store them in the cache.

220 - # If kv_cache is not provided, the new key and value tensors are

221 - # not cached. This happens during the initial memory profiling run.

222 - key_cache = self.k_cache(key, key_cache, block_indices,

223 - block_offsets)

224 - value_cache = self.v_cache(value, value_cache, block_indices,

225 - block_offsets)

226 -

227 - if attn_metadata.is_prompt:

228 - # Prompt run.

229 - query_shape = (batch_size, seq_len, self.num_heads, self.head_size)

230 - kv_shape = (batch_size, seq_len_kv, self.num_kv_heads,

231 - self.head_size)

232 -

233 - attn_bias = attn_metadata.attn_bias

234 - if attn_bias is not None and self.alibi_slopes is not None:

235 - position_bias = _make_alibi_bias(self.alibi_slopes,

236 - self.num_kv_heads,

237 - attn_bias.dtype,

238 - attn_bias.shape[-1])

239 - attn_bias = attn_bias.tile((1, self.num_kv_heads, 1, 1))

240 - attn_bias.add_(position_bias)

241 -

242 - block_list = attn_metadata.block_list if attn_metadata \

243 - and attn_metadata.block_list is not None else None

244 -

245 - out = ops.prompt_attention(

246 - impl=self.prefill_impl,

247 - query=query.view(query_shape),

248 - key=key.view(kv_shape),

249 - value=value.view(kv_shape),

250 - is_causal=True,

251 - attn_bias=attn_bias,

252 - valid_seq_lengths=attn_metadata.seq_lens_tensor,

253 - **self.common_attention_args(block_list, key_cache,

254 - value_cache))

255 - output = out.reshape(batch_size, seq_len, hidden_size)

256 - else:

257 - # Decoding run.

258 - output = HPUPagedAttention.forward_decode(

259 - query=query,

260 - block_mapping=attn_metadata.block_mapping,

261 - block_bias=attn_metadata.attn_bias,

262 - block_groups=attn_metadata.block_groups,

263 - **self.common_attention_args(attn_metadata.block_list,

264 - key_cache, value_cache))

265 - # Reshape the output tensor.

266 - return output.view(batch_size, seq_len, hidden_size)

267 -

268 - def common_attention_args(self,

269 - block_list=None,

270 - key_cache=None,

271 - value_cache=None):

272 - fsdpa_op = self.fused_scaled_dot_product_attention.apply \

273 - if self.fused_scaled_dot_product_attention is not None else None

274 - return {

275 - 'scale': self.scale,

276 - 'matmul_qk_op': self.matmul_qk,

277 - 'matmul_av_op': self.matmul_av,

278 - 'batch2block_matmul_op': self.batch2block_matmul,

279 - 'block2batch_matmul_op': self.block2batch_matmul,

280 - 'fsdpa_op': fsdpa_op,

281 - 'keys_fetch_func': self.k_cache.fetch_from_cache,

282 - 'values_fetch_func': self.v_cache.fetch_from_cache,

283 - 'softmax_op': self.softmax,

284 - 'block_list': block_list,

285 - 'key_cache': key_cache,

286 - 'value_cache': value_cache,

287 - }

288 -

289 -

290 -def _make_alibi_bias(

291 - alibi_slopes: torch.Tensor,

292 - num_kv_heads: int,

293 - dtype: torch.dtype,

294 - seq_len: int,

295 -) -> torch.Tensor:

296 - bias = torch.arange(seq_len, dtype=dtype)

297 - # NOTE(zhuohan): HF uses

298 - # `bias = bias[None, :].repeat(seq_len, 1)`

299 - # here. We find that both biases give the same results, but

300 - # the bias below more accurately follows the original ALiBi

301 - # paper.

302 - # Calculate a matrix where each element represents ith element- jth

303 - # element.

304 - bias = bias[None, :] - bias[:, None]

305 -

306 - padded_len = (seq_len + 7) // 8 * 8

307 - num_heads = alibi_slopes.shape[0]

308 - bias = torch.empty(

309 - 1, # batch size

310 - num_heads,

311 - seq_len,

312 - padded_len,

313 - device=alibi_slopes.device,

314 - dtype=dtype,

315 - )[:, :, :, :seq_len].copy_(bias)

316 - bias.mul_(alibi_slopes[:, None, None])

317 - if num_heads != num_kv_heads:

318 - bias = bias.unflatten(1, (num_kv_heads, num_heads // num_kv_heads))

319 - return bias

vllm/attention/ops/hpu_paged_attn.py deleted

-88

@@ -1,88 +0,0 @@

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

2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 -

4 -###############################################################################

6 -###############################################################################

7 -

8 -from dataclasses import dataclass

9 -from typing import List, Optional, Tuple

10 -

11 -import torch

12 -from vllm_hpu_extension import cache_ops, ops

13 -

14 -# Should be the same as PARTITION_SIZE in `paged_attention_v2_launcher`.

15 -_PARTITION_SIZE = 512

16 -

17 -

18 -@dataclass

19 -class HPUPagedAttentionMetadata:

20 - """Metadata for PagedAttention."""

21 - block_list: Optional[torch.Tensor]

22 - block_mapping: Optional[torch.Tensor]

23 - block_usage: Optional[torch.Tensor]

24 - block_indices: Optional[torch.Tensor]

25 - block_offsets: Optional[torch.Tensor]

26 - block_groups: Optional[torch.Tensor]

27 -

28 -

29 -class HPUPagedAttention:

30 -

31 - @staticmethod

32 - def get_supported_head_sizes() -> List[int]:

33 - return [64, 80, 96, 112, 128, 256]

34 -

35 - @staticmethod

36 - def get_kv_cache_shape(

37 - num_blocks: int,

38 - block_size: int,

39 - num_kv_heads: int,

40 - head_size: int,

41 - ) -> Tuple[int, ...]:

42 - return (num_blocks, block_size, num_kv_heads, head_size)

43 -

44 - @staticmethod

45 - def split_kv_cache(

46 - kv_cache: torch.Tensor,

47 - num_kv_heads: int,

48 - head_size: int,

49 - ) -> Tuple[torch.Tensor, torch.Tensor]:

50 - key_cache = kv_cache[0]

51 - value_cache = kv_cache[1]

52 - return key_cache, value_cache

53 -

54 - @staticmethod

55 - def write_to_paged_cache(key: torch.Tensor, value: torch.Tensor,

56 - key_cache: torch.Tensor,

57 - value_cache: torch.Tensor,

58 - slot_mapping: torch.Tensor, kv_cache_dtype: str,

59 - is_prompt: bool) -> None:

60 - cache_ops.reshape_and_cache(key, value, key_cache, value_cache,

61 - slot_mapping, kv_cache_dtype, is_prompt)

62 -

63 - @staticmethod

64 - def forward_decode(**kwargs) -> torch.Tensor:

65 - return ops.flat_pa(**kwargs)

66 -

67 - @staticmethod

68 - def swap_blocks(

69 - src_kv_cache: Tuple[torch.Tensor, torch.Tensor],

70 - dst_kv_cache: Tuple[torch.Tensor, torch.Tensor],

71 - src_to_dsts: torch.Tensor,

72 - ) -> None:

73 - src_key_cache = src_kv_cache[0]

74 - dst_key_cache = dst_kv_cache[0]

75 - cache_ops.swap_blocks(src_key_cache, dst_key_cache, src_to_dsts)

76 -

77 - src_value_cache = src_kv_cache[1]

78 - dst_value_cache = dst_kv_cache[1]

79 - cache_ops.swap_blocks(src_value_cache, dst_value_cache, src_to_dsts)

80 -

81 - @staticmethod

82 - def copy_blocks(

83 - kv_caches: List[Tuple[torch.Tensor, torch.Tensor]],

84 - src_to_dsts: torch.Tensor,

85 - ) -> None:

86 - key_caches = [kv_cache[0] for kv_cache in kv_caches]

87 - value_caches = [kv_cache[1] for kv_cache in kv_caches]

88 - cache_ops.copy_blocks(key_caches, value_caches, src_to_dsts)

vllm/config.py

+1 -1

@@ -2452,7 +2452,7 @@ class SchedulerConfig:

2452 return self.num_scheduler_steps > 1

2453

2454

2455 -Device = Literal["auto", "cuda", "neuron", "cpu", "tpu", "xpu", "hpu"]

2455 +Device = Literal["auto", "cuda", "neuron", "cpu", "tpu", "xpu"]

2456

2457

2458 @config

vllm/core/block/cpu_gpu_block_allocator.py

+1 -3

@@ -7,7 +7,6 @@ from vllm.core.block.interfaces import (Block, BlockAllocator, BlockId,

7 DeviceAwareBlockAllocator)

8 from vllm.core.block.naive_block import NaiveBlock, NaiveBlockAllocator

9 from vllm.core.block.prefix_caching_block import PrefixCachingBlockAllocator

10 -from vllm.platforms import current_platform

10 from vllm.utils import Device

@@ -56,8 +55,7 @@ class CpuGpuBlockAllocator(DeviceAwareBlockAllocator):

55 - The block IDs are assigned contiguously, with GPU block IDs coming

56 before CPU block IDs.

57 """

59 - # For HPU, block id 0 is used only for padding

60 - reserved_blocks = 1 if current_platform.is_hpu() else 0

58 + reserved_blocks = 0

59 block_ids = list(

60 range(reserved_blocks, num_gpu_blocks + num_cpu_blocks))

61 num_gpu_blocks -= reserved_blocks

vllm/distributed/device_communicators/hpu_communicator.py deleted

-46

@@ -1,46 +0,0 @@

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

2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 -

4 -import torch

5 -import torch.distributed as dist

6 -

7 -from vllm.platforms import current_platform

8 -

9 -from .base_device_communicator import DeviceCommunicatorBase

10 -

11 -if current_platform.is_hpu():

12 - import habana_frameworks.torch as htorch # noqa: F401

13 -

14 -

15 -class HpuCommunicator(DeviceCommunicatorBase):

16 -

17 - def all_reduce(self, input_: torch.Tensor) -> torch.Tensor:

18 - # FIXME(kzawora): this is a workaround for a bug in Habana PT bridge

19 - # occurring when PT_HPU_ENABLE_LAZY_COLLECTIVES=true env var is used

20 - # (which is required for tensor parallel HPUGraph inference)

21 - htorch.core.mark_step()

22 - dist.all_reduce(input_, group=self.device_group)

23 - return input_

24 -

25 - def all_gather(self, input_: torch.Tensor, dim: int = -1) -> torch.Tensor:

26 - world_size = self.world_size

27 - if dim < 0:

28 - # Convert negative dim to positive.

29 - dim += input_.dim()

30 - input_size = input_.size()

31 - # Allocate output tensor.

32 - output_tensor = torch.empty((world_size, ) + input_size,

33 - dtype=input_.dtype,

34 - device=input_.device)

35 - # All-gather.

36 - htorch.core.mark_step()

37 - dist.all_gather_into_tensor(output_tensor,

38 - input_,

39 - group=self.device_group)

40 - # Reshape

41 - output_tensor = output_tensor.movedim(0, dim)

42 - output_tensor = output_tensor.reshape(input_size[:dim] +

43 - (world_size *

44 - input_size[dim], ) +

45 - input_size[dim + 1:])

46 - return output_tensor

vllm/engine/arg_utils.py

+2 -3

@@ -1365,9 +1365,8 @@ class EngineArgs:

1365 supported = False

1366 if current_platform.is_rocm() or (

1367 current_platform.is_cuda()

1368 - and current_platform.is_device_capability(100)) or (

1369 - current_platform.device_name

1370 - == "hpu"): # handle hpu also for OOT platform

1368 + and current_platform.is_device_capability(100)

1369 + ): # handle hpu also for OOT platform

1370 supported = True

1371 elif fp8_attention and will_use_fa:

1372 from vllm.attention.utils.fa_utils import (

vllm/envs.py

-15

@@ -106,8 +106,6 @@ if TYPE_CHECKING:

106 VLLM_RAY_PER_WORKER_GPUS: float = 1.0

107 VLLM_RAY_BUNDLE_INDICES: str = ""

108 VLLM_CUDART_SO_PATH: Optional[str] = None

109 - VLLM_USE_HPU_CONTIGUOUS_CACHE_FETCH: bool = True

110 - VLLM_HPU_USE_DELAYED_SAMPLING: bool = False

109 VLLM_DP_RANK: int = 0

110 VLLM_DP_RANK_LOCAL: int = -1

111 VLLM_DP_SIZE: int = 1

@@ -780,19 +778,6 @@ environment_variables: dict[str, Callable[[], Any]] = {

778 "VLLM_CUDART_SO_PATH":

779 lambda: os.getenv("VLLM_CUDART_SO_PATH", None),

780

783 - # Contiguous cache fetching to avoid using costly gather operation on

784 - # Gaudi3. This is only applicable to HPU contiguous cache. If set to true,

785 - # contiguous cache fetch will be used.

786 - "VLLM_USE_HPU_CONTIGUOUS_CACHE_FETCH":

787 - lambda: os.environ.get("VLLM_CONTIGUOUS_PA", "true").lower() in

788 - ("1", "true"),

789 -

790 - # Use delayed sampling for HPU to reduce host cpu overhead

791 - # between each step.

792 - "VLLM_HPU_USE_DELAYED_SAMPLING":

793 - lambda: os.environ.get("VLLM_DELAYED_SAMPLING", "false").lower() in

794 - ("1", "true"),

795 -

781 # Rank of the process in the data parallel setting

782 "VLLM_DP_RANK":

783 lambda: int(os.getenv("VLLM_DP_RANK", "0")),

vllm/lora/layers.py

-4

@@ -1164,10 +1164,6 @@ class LogitsProcessorWithLoRA(BaseLayerWithLoRA):

1164 posinf=pos_inf,

1165 neginf=neg_inf))

1166

1167 - # HPU needs special handling to prune out dummy samples.

1168 - if current_platform.is_hpu():

1169 - lora_logits = lora_logits[:logits.shape[0], :]

1170 -

1167 logits[:,

1168 self.base_layer.org_vocab_size:self.base_layer.org_vocab_size +

1169 lora_logits.shape[1]] = lora_logits

vllm/lora/punica_wrapper/punica_hpu.py deleted

-145

@@ -1,145 +0,0 @@

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

2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 -

4 -from typing import TYPE_CHECKING, Optional, Union, final

5 -

6 -import torch

7 -from vllm_hpu_extension.ops import (dispatch_bgmv_embedding,

8 - dispatch_bgmv_linear)

9 -

10 -from .punica_base import PunicaWrapperBase

11 -from .utils import convert_mapping

12 -

13 -if TYPE_CHECKING:

14 - # avoid circuit import

15 - from vllm.lora.layers import LoRAMapping

16 - from vllm.lora.models import LongContextLoRAContext

17 -

18 -

19 -@final

20 -class PunicaWrapperHPU(PunicaWrapperBase):

21 -

22 - def __init__(self, max_num_batched_tokens: int, max_batches: int,

23 - device: Union[torch.device, str], **kwargs):

24 - # Increasing max_num_batched_tokens by 3x to handle increase in

25 - # tensor size due to padding.

26 - PunicaWrapperBase.__init__(self, 3 * max_num_batched_tokens,

27 - max_batches, device)

28 -

29 - def _update_base_metadata(

30 - self,

31 - mapping: "LoRAMapping",

32 - lora_index_to_id: list[Optional[int]],

33 - max_loras: int,

34 - vocab_size: int,

35 - extra_vocab_size: int,

36 - long_lora_context: Optional["LongContextLoRAContext"] = None,

37 - ):

38 - (

39 - base_indices,

40 - sampler_indices,

41 - sampler_indices_padded,

42 - embeddings_indices,

43 - long_lora_offsets_tensor,

44 - indices_len,

45 - ) = convert_mapping(mapping, lora_index_to_id, max_loras, vocab_size,

46 - extra_vocab_size, self.device, None)

47 - # Updating each element in `long_lora_offsets` with `lora_offset` slows

48 - # down perf in HPU due to a series of `strided_insert` ops during lazy

49 - # graph accumulation. Hence HPU appends `lora_offset` to a list and

50 - # converts it to a tensor only after it is ready.

51 - if long_lora_context:

52 - index_mapping_indices: list[int] = list(

53 - mapping.index_mapping).copy()

54 - long_lora_offsets: list[int] = []

55 - for i in range(len(index_mapping_indices)):

56 - lora_offset: int = long_lora_context.offsets_by_lora_id.get(

57 - index_mapping_indices[i], 0)

58 - long_lora_offsets.append(lora_offset)

59 - long_lora_offsets_tensor = torch.tensor(long_lora_offsets,

60 - device=self.device,

61 - dtype=torch.long)

62 - indices_len[-1] = long_lora_offsets_tensor.shape[-1]

63 -

64 - self._token_lora_indices[:base_indices.shape[0]].copy_(base_indices)

65 - self._sampler_indices[:sampler_indices.shape[0]].copy_(sampler_indices)

66 - self._sampler_indices_padded[:sampler_indices_padded.shape[0]].copy_(

67 - sampler_indices_padded)

68 - self._embeddings_indices[:embeddings_indices.

69 - shape[0], :embeddings_indices.shape[1]].copy_(

70 - embeddings_indices)

71 - if long_lora_offsets_tensor is not None:

72 - self._long_lora_indices[:long_lora_offsets_tensor.shape[0]].copy_(

73 - long_lora_offsets_tensor)

74 - else:

75 - self._long_lora_indices.zero_()

76 - self.indices_len[:] = indices_len

77 -

78 - def add_lora_embedding(self,

79 - y: torch.Tensor,

80 - x: torch.Tensor,

81 - lora_b_stacked: torch.Tensor,

82 - add_inputs: bool = True,

83 - **kwargs) -> None:

84 - dispatch_bgmv_embedding(y, x, lora_b_stacked, 0)

85 -

86 - def add_lora_linear(self,

87 - y: torch.Tensor,

88 - x: torch.Tensor,

89 - lora_a_stacked: tuple[torch.Tensor, ...],

90 - lora_b_stacked: tuple[torch.Tensor, ...],

91 - lora_bias_stacked: Optional[tuple[torch.Tensor, ...]],

92 - scale: float,

93 - output_slices: tuple[int, ...],

94 - *,

95 - buffer: Optional[tuple[torch.Tensor, ...]] = None,

96 - **kwargs) -> None:

97 - y_org = y

98 - x = x.view(-1, x.shape[-1])

99 - y = y.view(-1, y.shape[-1])

100 - offset_left = 0

101 -

102 - for slice_idx in range(len(output_slices)):

103 - dispatch_bgmv_linear(

104 - y[:, offset_left:offset_left + output_slices[slice_idx]], x,

105 - lora_a_stacked[slice_idx], lora_b_stacked[slice_idx], 0, scale)

106 - offset_left += output_slices[slice_idx]

107 - y = y.view_as(y_org)

108 -

109 - def add_lora_logits(self,

110 - y: torch.Tensor,

111 - x: torch.Tensor,

112 - lora_a_stacked: torch.Tensor,

113 - lora_b_stacked: torch.Tensor,

114 - scale,

115 - *,

116 - buffer: Optional[torch.Tensor] = None,

117 - **kwargs) -> None:

118 - y_org = y

119 - y = y.view(-1, y.shape[-1])

120 - x = x.view(-1, x.shape[-1])

121 - dispatch_bgmv_linear(y, x, lora_a_stacked, lora_b_stacked, 0, scale)

122 - y = y.view_as(y_org)

123 -

124 - def add_shrink(

125 - self,

126 - y: Union[tuple[torch.Tensor, ...], torch.Tensor],

127 - x: torch.Tensor,

128 - lora_a_stacked: tuple[torch.Tensor, ...],

129 - scale: float,

130 - **kwargs,

131 - ) -> None:

132 - raise NotImplementedError

133 -

134 - def add_expand(

135 - self,

136 - y: torch.Tensor,

137 - x: Union[tuple[torch.Tensor, ...], torch.Tensor],

138 - lora_b_stacked: tuple[torch.Tensor, ...],

139 - lora_bias_stacked: Optional[tuple[torch.Tensor, ...]],

140 - output_slices: tuple[int, ...],

141 - offset_start: int = 0,

142 - add_inputs=True,

143 - **kwargs,

144 - ) -> None:

145 - raise NotImplementedError

vllm/model_executor/custom_op.py

-7

@@ -73,11 +73,6 @@ class CustomOp(nn.Module):

73 # NOTE(woosuk): This is a placeholder for future extensions.

74 return self.forward_native(*args, **kwargs)

76 - def forward_hpu(self, *args, **kwargs):

77 - # By default, we assume that Gaudi ops are compatible with the

78 - # PyTorch-native implementation.

79 - return self.forward_native(*args, **kwargs)

80 -

76 def forward_neuron(self, *args, **kwargs):

77 # By default, we assume that Neuron ops are compatible with the

78 # PyTorch-native implementation.

@@ -106,8 +101,6 @@ class CustomOp(nn.Module):

101 return self.forward_hip

102 elif current_platform.is_cpu():

103 return self.forward_cpu

109 - elif current_platform.is_hpu():

110 - return self.forward_hpu

104 elif current_platform.is_tpu():

105 return self.forward_tpu

106 elif current_platform.is_xpu():

vllm/model_executor/layers/fused_moe/layer.py

-36

@@ -475,39 +475,6 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):

475 activation,

476 )

477

478 - def forward_hpu(

479 - self,

480 - layer: torch.nn.Module,

481 - x: torch.Tensor,

482 - use_grouped_topk: bool,

483 - top_k: int,

484 - router_logits: torch.Tensor,

485 - renormalize: bool,

486 - topk_group: Optional[int] = None,

487 - num_expert_group: Optional[int] = None,

488 - global_num_experts: int = -1,

489 - expert_map: Optional[torch.Tensor] = None,

490 - custom_routing_function: Optional[Callable] = None,

491 - scoring_func: str = "softmax",

492 - e_score_correction_bias: Optional[torch.Tensor] = None,

493 - apply_router_weight_on_input: bool = False,

494 - activation: str = "silu",

495 - ) -> torch.Tensor:

496 - assert not use_grouped_topk

497 - assert num_expert_group is None

498 - assert topk_group is None

499 - assert custom_routing_function is None

500 - assert layer is not None

501 - assert apply_router_weight_on_input is False

502 - if scoring_func != "softmax":

503 - raise NotImplementedError(

504 - "Only softmax scoring function is supported for HPU.")

505 - if e_score_correction_bias is not None:

506 - raise NotImplementedError(

507 - "Expert score correction bias is not supported for HPU.")

508 - return layer.hpu_fused_moe(x, layer.w13_weight, layer.w2_weight,

509 - router_logits, top_k)

510 -

478 def forward_tpu(

479 self,

480 layer: torch.nn.Module,

@@ -716,9 +683,6 @@ class FusedMoE(torch.nn.Module):

683 if self.scoring_func != "softmax" and not self.use_grouped_topk:

684 raise ValueError("Only softmax scoring function is supported for "

685 "non-grouped topk.")

719 - if current_platform.is_hpu():

720 - from vllm_hpu_extension.ops import DynamicFusedMOE

721 - self.hpu_fused_moe = DynamicFusedMOE(self.global_num_experts)

686

687 if vllm_config.model_config is not None:

688 model_dtype = vllm_config.model_config.dtype

vllm/model_executor/layers/layernorm.py

-20

@@ -170,26 +170,6 @@ class RMSNorm(CustomOp):

170 else:

171 return norm_func(x, self.weight.data, self.variance_epsilon)

172

173 - def forward_hpu(

174 - self,

175 - x: torch.Tensor,

176 - residual: Optional[torch.Tensor] = None,

177 - ) -> Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:

178 - from vllm_hpu_extension.kernels import rms_norm

179 - HPUFusedRMSNorm = rms_norm()

180 - if HPUFusedRMSNorm is None:

181 - return self.forward_native(x, residual)

182 - if residual is not None:

183 - orig_shape = x.shape

184 - residual += x.view(residual.shape)

185 - # Note: HPUFusedRMSNorm requires 3D tensors as inputs

186 - x = HPUFusedRMSNorm.apply(residual, self.weight,

187 - self.variance_epsilon)

188 - return x.view(orig_shape), residual

189 -

190 - x = HPUFusedRMSNorm.apply(x, self.weight, self.variance_epsilon)

191 - return x

192 -

173 def forward_xpu(

174 self,

175 x: torch.Tensor,

vllm/model_executor/layers/rotary_embedding.py

-58

@@ -229,64 +229,6 @@ class RotaryEmbedding(CustomOp):

229 self.cos_sin_cache, self.is_neox_style)

230 return query, key

231

232 - def forward_hpu(

233 - self,

234 - positions: torch.Tensor,

235 - query: torch.Tensor,

236 - key: Optional[torch.Tensor] = None,

237 - offsets: Optional[torch.Tensor] = None,

238 - ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:

239 - from habana_frameworks.torch.hpex.kernels import (

240 - RotaryPosEmbeddingMode, apply_rotary_pos_emb)

241 - if offsets is not None:

242 - offsets = offsets.view(positions.shape[0], -1)

243 - positions = positions + offsets

244 - positions = positions.flatten()

245 - num_tokens = positions.shape[0]

246 - cos_sin = self.cos_sin_cache.index_select(0, positions).view(

247 - num_tokens, 1, -1)

248 - cos, sin = cos_sin.chunk(2, dim=-1)

249 - # HPU RoPE kernel requires hidden dimension for cos and sin to be equal

250 - # to query hidden dimension, so the original tensors need to be

251 - # expanded

252 - # GPT-NeoX kernel requires position_ids = None, offset, mode = BLOCKWISE

253 - # and expansion of cos/sin tensors via concatenation

254 - # GPT-J kernel requires position_ids = None, offset = 0, mode = PAIRWISE

255 - # and expansion of cos/sin tensors via repeat_interleave

256 - rope_mode: RotaryPosEmbeddingMode

257 - if self.is_neox_style:

258 - rope_mode = RotaryPosEmbeddingMode.BLOCKWISE

259 - cos = torch.cat((cos, cos), dim=-1)

260 - sin = torch.cat((sin, sin), dim=-1)

261 - else:

262 - rope_mode = RotaryPosEmbeddingMode.PAIRWISE

263 - sin = torch.repeat_interleave(sin,

264 - 2,

265 - dim=-1,

266 - output_size=cos_sin.shape[-1])

267 - cos = torch.repeat_interleave(cos,

268 - 2,

269 - dim=-1,

270 - output_size=cos_sin.shape[-1])

271 -

272 - query_shape = query.shape

273 - query = query.view(num_tokens, -1, self.head_size)

274 - query_rot = query[..., :self.rotary_dim]

275 - query_pass = query[..., self.rotary_dim:]

276 - query_rot = apply_rotary_pos_emb(query_rot, cos, sin, None, 0,

277 - rope_mode)

278 - query = torch.cat((query_rot, query_pass), dim=-1).reshape(query_shape)

279 -

280 - if key is not None:

281 - key_shape = key.shape

282 - key = key.view(num_tokens, -1, self.head_size)

283 - key_rot = key[..., :self.rotary_dim]

284 - key_pass = key[..., self.rotary_dim:]

285 - key_rot = apply_rotary_pos_emb(key_rot, cos, sin, None, 0,

286 - rope_mode)

287 - key = torch.cat((key_rot, key_pass), dim=-1).reshape(key_shape)

288 - return query, key

289 -

232 def forward_neuron(

233 self,

234 positions: torch.Tensor,

vllm/model_executor/layers/vocab_parallel_embedding.py

+2 -14

@@ -388,20 +388,8 @@ class VocabParallelEmbedding(torch.nn.Module):

388

389 # Copy the data. Select chunk corresponding to current shard.

390 loaded_weight = loaded_weight.narrow(output_dim, start_idx, shard_size)

391 -

392 - if current_platform.is_hpu():

393 - # FIXME(kzawora): Weight copy with slicing bugs out on Gaudi here,

394 - # so we're using a workaround. Remove this when fixed in

395 - # HPU PT bridge.

396 - padded_weight = torch.cat([

397 - loaded_weight,

398 - torch.zeros(param.shape[0] - loaded_weight.shape[0],

399 - *loaded_weight.shape[1:])

400 - ])

401 - param.data.copy_(padded_weight)

402 - else:

403 - param[:loaded_weight.shape[0]].data.copy_(loaded_weight)

404 - param[loaded_weight.shape[0]:].data.fill_(0)

391 + param[:loaded_weight.shape[0]].data.copy_(loaded_weight)

392 + param[loaded_weight.shape[0]:].data.fill_(0)

393

394 def forward(self, input_):

395 if self.tp_size > 1:

vllm/model_executor/model_loader/bitsandbytes_loader.py

+1 -10

@@ -199,10 +199,6 @@ class BitsAndBytesModelLoader(BaseModelLoader):

199

200 if self.pre_quant:

201 if self.load_8bit:

202 - if current_platform.is_hpu():

203 - raise ValueError(

204 - "currently hpu supports 4bit quantization only")

205 -

202 return self._quantized_8bit_generator(

203 hf_weights_files, use_safetensors,

204 quant_state_dict), quant_state_dict

@@ -306,10 +302,6 @@ class BitsAndBytesModelLoader(BaseModelLoader):

302 in temp_state_dict):

303 quant_state = _parse_quant_state(mapped_weight_name,

304 temp_state_dict)

309 - if current_platform.is_hpu():

310 - assert quant_state.quant_type == "nf4", (

311 - "currently hpu supports nf4 quant_type only")

312 -

305 quant_state_dict[mapped_weight_name] = quant_state

306 yield org_weight_name, weight_tensor

307 else:

@@ -380,8 +372,7 @@ class BitsAndBytesModelLoader(BaseModelLoader):

372 ...]

373

374 # bitsandbytes requires data in GPU

383 - if (weight_sub_tensor.is_cuda

384 - or weight_sub_tensor.device.type == "hpu"):

375 + if weight_sub_tensor.is_cuda:

376 loaded_weight = weight_sub_tensor

377 else:

378 loaded_weight = weight_sub_tensor.to(

vllm/model_executor/model_loader/default_loader.py

-10

@@ -218,16 +218,6 @@ class DefaultModelLoader(BaseModelLoader):

218

219 weights_iterator = _xla_weights_iterator(weights_iterator)

220

221 - elif current_platform.is_hpu():

222 - import habana_frameworks.torch.core as htcore

223 -

224 - def _hpu_weights_iterator(iterator: Generator):

225 - for weights in iterator:

226 - yield weights

227 - htcore.mark_step()

228 -

229 - weights_iterator = _hpu_weights_iterator(weights_iterator)

230 -

221 if self.counter_before_loading_weights == 0.0:

222 self.counter_before_loading_weights = time.perf_counter()

223 # Apply the prefix.

vllm/platforms/__init__.py

-18

@@ -116,23 +116,6 @@ def rocm_platform_plugin() -> Optional[str]:

116 return "vllm.platforms.rocm.RocmPlatform" if is_rocm else None

117

118

119 -def hpu_platform_plugin() -> Optional[str]:

120 - is_hpu = False

121 - logger.debug("Checking if HPU platform is available.")

122 - try:

123 - from importlib import util

124 - is_hpu = util.find_spec('habana_frameworks') is not None

125 - if is_hpu:

126 - logger.debug("Confirmed HPU platform is available.")

127 - else:

128 - logger.debug("HPU platform is not available because "

129 - "habana_frameworks is not found.")

130 - except Exception as e:

131 - logger.debug("HPU platform is not available because: %s", str(e))

132 -

133 - return "vllm.platforms.hpu.HpuPlatform" if is_hpu else None

134 -

135 -

119 def xpu_platform_plugin() -> Optional[str]:

120 is_xpu = False

121 logger.debug("Checking if XPU platform is available.")

@@ -208,7 +191,6 @@ builtin_platform_plugins = {

191 'tpu': tpu_platform_plugin,

192 'cuda': cuda_platform_plugin,

193 'rocm': rocm_platform_plugin,

211 - 'hpu': hpu_platform_plugin,

194 'xpu': xpu_platform_plugin,

195 'cpu': cpu_platform_plugin,

196 'neuron': neuron_platform_plugin,

vllm/platforms/hpu.py deleted

-114

@@ -1,114 +0,0 @@

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

2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 -

4 -import os

5 -from typing import TYPE_CHECKING, Optional

6 -

7 -import torch

8 -

9 -from vllm import envs

10 -from vllm.logger import init_logger

11 -from vllm.utils import DEFAULT_MAX_NUM_BATCHED_TOKENS

12 -

13 -from .interface import Platform, PlatformEnum, _Backend

14 -

15 -if TYPE_CHECKING:

16 - from vllm.config import VllmConfig

17 -else:

18 - VllmConfig = None

19 -

20 -logger = init_logger(__name__)

21 -

22 -

23 -class HpuPlatform(Platform):

24 - _enum = PlatformEnum.HPU

25 - device_name: str = "hpu"

26 - device_type: str = "hpu"

27 - dispatch_key: str = "HPU"

28 - ray_device_key: str = "HPU"

29 - dist_backend: str = "hccl"

30 - device_control_env_var: str = "HABANA_VISIBLE_MODULES"

31 -

32 - @classmethod

33 - def get_attn_backend_cls(cls, selected_backend: _Backend, head_size: int,

34 - dtype: torch.dtype, kv_cache_dtype: Optional[str],

35 - block_size: int, use_v1: bool,

36 - use_mla: bool) -> str:

37 - logger.info("Using HPUAttention backend.")

38 - return "vllm.attention.backends.hpu_attn.HPUAttentionBackend"

39 -

40 - @classmethod

41 - def is_async_output_supported(cls, enforce_eager: Optional[bool]) -> bool:

42 - return True

43 -

44 - @classmethod

45 - def inference_mode(cls):

46 - return torch.no_grad()

47 -

48 - @classmethod

49 - def set_device(cls, device: torch.device) -> None:

50 - """

51 - Set the device for the current platform.

52 - """

53 - torch.hpu.set_device(device)

54 -

55 - @classmethod

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

57 -

58 - scheduler_config = vllm_config.scheduler_config

59 - parallel_config = vllm_config.parallel_config

60 - if scheduler_config.is_multi_step:

61 - parallel_config.worker_cls = \

62 - "vllm.worker.multi_step_hpu_worker.MultiStepHPUWorker"

63 -

64 - if vllm_config.speculative_config is not None:

65 - raise NotImplementedError(

66 - "Speculative decoding is not implemented for HPU")

67 -

68 - if parallel_config.worker_cls == "auto":

69 - parallel_config.worker_cls = "vllm.worker.hpu_worker.HPUWorker"

70 -

71 - # NOTE(kzawora): default block size for Gaudi should be 128

72 - # smaller sizes still work, but very inefficiently

73 - cache_config = vllm_config.cache_config

74 - if cache_config and cache_config.block_size is None:

75 - cache_config.block_size = 128

76 - if (parallel_config.distributed_executor_backend == 'mp'

77 - and envs.VLLM_WORKER_MULTIPROC_METHOD == 'fork'):

78 - if os.environ.get("VLLM_WORKER_MULTIPROC_METHOD",

79 - None) is not None:

80 - logger.warning("On HPU, VLLM_WORKER_MULTIPROC_METHOD=fork "

81 - "might cause application hangs on exit. Using "

82 - "VLLM_WORKER_MULTIPROC_METHOD=fork anyway, "

83 - "as it was explicitly requested.")

84 - else:

85 - logger.warning(

86 - "On HPU, VLLM_WORKER_MULTIPROC_METHOD=fork "

87 - "might cause application hangs on exit. Setting "

88 - "VLLM_WORKER_MULTIPROC_METHOD to 'spawn'. "

89 - "To override that behavior, please set "

90 - "VLLM_WORKER_MULTIPROC_METHOD=fork explicitly.")

91 - os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"

92 -

93 - if vllm_config.model_config and vllm_config.model_config.use_mla:

94 - logger.info(

95 - "MLA is enabled on a non-GPU platform; forcing chunked "

96 - "prefill and prefix caching to be disabled.")

97 - vllm_config.scheduler_config.enable_chunked_prefill = False

98 - vllm_config.scheduler_config.chunked_prefill_enabled = False

99 - vllm_config.scheduler_config.max_num_batched_tokens = max(

100 - vllm_config.scheduler_config.max_model_len,

101 - DEFAULT_MAX_NUM_BATCHED_TOKENS)

102 -

103 - @classmethod

104 - def is_pin_memory_available(cls):

105 - logger.warning("Pin memory is not supported on HPU.")

106 - return False

107 -

108 - @classmethod

109 - def get_punica_wrapper(cls) -> str:

110 - return "vllm.lora.punica_wrapper.punica_hpu.PunicaWrapperHPU"

111 -

112 - @classmethod

113 - def get_device_communicator_cls(cls) -> str:

114 - return "vllm.distributed.device_communicators.hpu_communicator.HpuCommunicator" # noqa

vllm/platforms/interface.py

-5

@@ -54,7 +54,6 @@ class _Backend(enum.Enum):

54 FLASHMLA_VLLM_V1 = enum.auto()

55 FLASHMLA = enum.auto() # Supported by V1

56 CUTLASS_MLA_VLLM_V1 = enum.auto()

57 - HPU_ATTN = enum.auto()

57 PALLAS = enum.auto()

58 PALLAS_VLLM_V1 = enum.auto()

59 IPEX = enum.auto()

@@ -69,7 +68,6 @@ class PlatformEnum(enum.Enum):

68 CUDA = enum.auto()

69 ROCM = enum.auto()

70 TPU = enum.auto()

72 - HPU = enum.auto()

71 XPU = enum.auto()

72 CPU = enum.auto()

73 NEURON = enum.auto()

@@ -154,9 +152,6 @@ class Platform:

152 def is_tpu(self) -> bool:

153 return self._enum == PlatformEnum.TPU

154

157 - def is_hpu(self) -> bool:

158 - return self._enum == PlatformEnum.HPU

159 -

155 def is_xpu(self) -> bool:

156 return self._enum == PlatformEnum.XPU

157

vllm/plugins/__init__.py

-13

@@ -2,7 +2,6 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

4 import logging

5 -import os

5 from typing import Any, Callable

7 import torch

@@ -75,18 +74,6 @@ def load_general_plugins():

74 if current_platform.is_xpu():

75 # see https://github.com/pytorch/pytorch/blob/43c5f59/torch/_dynamo/config.py#L158

76 torch._dynamo.config.disable = True

78 - elif current_platform.is_hpu():

79 - # NOTE(kzawora): PT HPU lazy backend (PT_HPU_LAZY_MODE = 1)

80 - # does not support torch.compile

81 - # Eager backend (PT_HPU_LAZY_MODE = 0) must be selected for

82 - # torch.compile support

83 - is_lazy = os.environ.get('PT_HPU_LAZY_MODE', '1') == '1'

84 - if is_lazy:

85 - torch._dynamo.config.disable = True

86 - # NOTE(kzawora) multi-HPU inference with HPUGraphs (lazy-only)

87 - # requires enabling lazy collectives

88 - # see https://docs.habana.ai/en/latest/PyTorch/Inference_on_PyTorch/Inference_Using_HPU_Graphs.html # noqa: E501

89 - os.environ['PT_HPU_ENABLE_LAZY_COLLECTIVES'] = 'true'

78 plugins = load_plugins_by_group(group=DEFAULT_PLUGINS_GROUP)

79 # general plugins, we only need to execute the loaded functions

vllm/worker/hpu_model_runner.py deleted

-2320

@@ -1,2320 +0,0 @@

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

2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 -

4 -###############################################################################

6 -###############################################################################

7 -

8 -import collections

9 -import contextlib

10 -import dataclasses

11 -import functools

12 -import gc

13 -import itertools

14 -import math

15 -import os

16 -import time

17 -from array import array

18 -from enum import Enum, IntEnum

19 -from typing import (TYPE_CHECKING, Any, Callable, Dict, List, NamedTuple,

20 - Optional, Set, Tuple, Type, TypeVar, Union)

21 -

22 -import habana_frameworks.torch as htorch

23 -import habana_frameworks.torch.internal.bridge_config as bc

24 -import torch

25 -import torch.nn as nn

26 -import vllm_hpu_extension.environment as environment

27 -from vllm_hpu_extension.bucketing.common import get_bucketing_context

28 -from vllm_hpu_extension.ops import LoraMask as LoraMask

29 -from vllm_hpu_extension.profiler import (HabanaHighLevelProfiler,

30 - HabanaMemoryProfiler, format_bytes)

31 -

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

39 -from vllm.lora.layers import LoRAMapping

40 -from vllm.lora.request import LoRARequest

41 -from vllm.lora.worker_manager import LRUCacheWorkerLoRAManager

42 -from vllm.model_executor import SamplingMetadata

43 -from vllm.model_executor.layers.layernorm import RMSNorm

44 -from vllm.model_executor.layers.sampler import SamplerOutput, get_sampler

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 BatchedTensorInputs, MultiModalKwargs

50 -from vllm.sampling_params import SamplingParams

51 -from vllm.sequence import (CompletionSequenceGroupOutput, IntermediateTensors,

52 - Logprob, SequenceData, SequenceGroupMetadata,

53 - SequenceOutput)

54 -from vllm.utils import (bind_kv_cache, is_pin_memory_available,

55 - make_tensor_with_pad)

56 -from vllm.worker.model_runner_base import (

57 - ModelRunnerBase, ModelRunnerInputBase,

58 - _add_attn_metadata_broadcastable_dict,

59 - _add_sampling_metadata_broadcastable_dict,

60 - _init_attn_metadata_from_tensor_dict,

61 - _init_sampling_metadata_from_tensor_dict)

62 -

63 -if TYPE_CHECKING:

64 - from vllm.attention.backends.abstract import AttentionBackend

65 -

66 -logger = init_logger(__name__)

67 -

68 -_TYPE_CACHE = {}

69 -# These values are assumed to be zero in several places.

70 -# Use caution when updating them!

71 -_PAD_SLOT_ID = 0

72 -_PAD_BLOCK_ID = 0

73 -

74 -LORA_WARMUP_RANK = 8

75 -

76 -DUMMY_TOKEN_ID = -1

77 -

78 -

79 -class PhaseType(Enum):

80 - PREFILL = 'prefill'

81 - PREFIX_PREFILL = 'prefix_prefill'

82 - DECODE = 'decode'

83 -

84 -

85 -def subtuple(obj: object,

86 - typename: str,

87 - to_copy: List[str],

88 - to_override: Optional[Dict[str, object]] = None):

89 - if obj is None:

90 - return None

91 - if to_override is None:

92 - to_override = {}

93 - fields = set(to_copy) | set(to_override.keys())

94 - if type(obj) is dict:

95 - values = {key: obj[key] for key in fields if key in obj}

96 - else:

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

98 - if typename not in _TYPE_CACHE:

99 - _TYPE_CACHE[typename] = collections.namedtuple(typename,

100 - ' '.join(fields))

101 - return _TYPE_CACHE[typename](**values)

102 -

103 -

104 -def round_up(value: int, k: int):

105 - return (value + k - 1) // k * k

106 -

107 -

108 -def align_workers(value, op):

109 - group = get_world_group().cpu_group

110 - world_size = torch.distributed.get_world_size()

111 - if world_size <= 1:

112 - return value

113 - value_t = torch.tensor(value, device='cpu')

114 - torch.distributed.all_reduce(value_t, op=op, group=group)

115 - return value_t.item()

116 -

117 -

118 -def setup_profiler():

119 - schedule = torch.profiler.schedule(wait=0, warmup=2, active=1, repeat=1)

120 - DEVICE = 'hpu'

121 - activities = [torch.profiler.ProfilerActivity.CPU]

122 - activities.extend([torch.profiler.ProfilerActivity.HPU] if DEVICE ==

123 - 'hpu' else [])

124 - #from habana_frameworks.torch.activity_profiler import DebugActivity

125 - #debug_activities=[DebugActivity.BRIDGE_FUNCTION_CALLS]

126 -

127 - profiler = torch.profiler.profile(

128 - schedule=schedule,

129 - activities=activities,

130 - #debug_activities=debug_activities,

131 - on_trace_ready=torch.profiler.tensorboard_trace_handler('.',

132 - use_gzip=True),

133 - record_shapes=False,

134 - with_stack=True)

135 - return profiler

136 -

137 -

138 -def pad_list(input, k, v):

139 - input_len = len(input)

140 - target_len = round_up(input_len, k)

141 - padding = target_len - input_len

142 - return input + [v] * padding

143 -

144 -

145 -def gather_list(input, indices, v):

146 - return [input[i] if i is not None else v for i in indices]

147 -

148 -

149 -def flatten(in_list):

150 - return list(itertools.chain(*in_list))

151 -

152 -

153 -def precompute_indices_and_offsets(block_size, slot_mapping, is_prompt):

154 - slot_mapping = slot_mapping.flatten()

155 - indices = torch.div(slot_mapping, block_size, rounding_mode="floor")

156 - if is_prompt:

157 - indices = indices.unflatten(0, (-1, block_size))[:, 0]

158 - offsets = None

159 - else:

160 - offsets = torch.fmod(slot_mapping, block_size)

161 - return indices, offsets

162 -

163 -

164 -def modify_decoder_layer(module: torch.nn.Module, suffix="DecoderLayer"):

165 - if module.__class__.__name__.endswith(suffix):

166 -

167 - def forward_hook(module, args, output):

168 - htorch.core.mark_step()

169 - return output

170 -

171 - module.register_forward_hook(forward_hook)

172 -

173 - for child_name, child_module in module.named_children():

174 - modify_decoder_layer(child_module)

175 -

176 -

177 -class HpuModelAdapter:

178 -

179 - def __init__(self, model, vllm_config):

180 - self.model = model

181 - self.sampler = get_sampler()

182 - self.prefill_use_fusedsdpa = os.getenv('VLLM_PROMPT_USE_FUSEDSDPA',

183 - '0').lower() in ['1', 'true']

184 - self.vllm_config = vllm_config

185 - self.block_size = vllm_config.cache_config.block_size

186 - self.dtype = vllm_config.model_config.dtype

187 - enforce_eager = vllm_config.model_config.enforce_eager

188 -

189 - if not htorch.utils.internal.is_lazy() and not enforce_eager:

190 - if os.getenv('VLLM_REGIONAL_COMPILATION',

191 - 'true').lower() == 'true':

192 - self.regional_compilation_layers_list = [

193 - RMSNorm, VocabParallelEmbedding

194 - ]

195 - self._regional_compilation(self.model)

196 - else:

197 - self.model = torch.compile(self.model,

198 - backend='hpu_backend',

199 - dynamic=False)

200 -

201 - def _regional_compilation(self,

202 - module,

203 - parent_module=None,

204 - module_name=None):

205 - if isinstance(module, torch.nn.ModuleList):

206 - for children_name, children_module in module.named_children():

207 - self._compile_region(module, children_name, children_module)

208 - elif any(

209 - isinstance(module, layer)

210 - for layer in self.regional_compilation_layers_list):

211 - self._compile_region(parent_module, module_name, module)

212 - else:

213 - for children_name, children_module in module.named_children():

214 - self._regional_compilation(children_module, module,

215 - children_name)

216 -

217 - def _compile_region(self, model, name, module):

218 - module = torch.compile(module, backend='hpu_backend', dynamic=False)

219 - setattr(model, name, module)

220 -

221 - def _set_attn_bias(self, attn_metadata, batch_size, seq_len, device,

222 - dtype):

223 - if (attn_metadata is None

224 - or (self.prefill_use_fusedsdpa \

225 - and attn_metadata.block_list is None)

226 - or not attn_metadata.is_prompt):

227 - return attn_metadata

228 -

229 - prefill_metadata = attn_metadata

230 -

231 - seq_lens_t = prefill_metadata.seq_lens_tensor

232 - context_lens_t = prefill_metadata.context_lens_tensor

233 - query_lens_t = seq_lens_t - context_lens_t

234 -

235 - block_list = attn_metadata.block_list

236 - max_context_len = (block_list.size(-1) //

237 - batch_size if block_list is not None else 0)

238 - max_context_len = max_context_len * self.block_size

239 - past_mask = torch.arange(0,

240 - max_context_len,

241 - dtype=torch.int32,

242 - device=device)

243 - past_mask = (past_mask.view(1, -1).expand(batch_size, -1).ge(

244 - context_lens_t.view(-1, 1)).view(batch_size, 1, -1).expand(

245 - batch_size, seq_len, -1).view(batch_size, 1, seq_len, -1))

246 -

247 - len_mask = (torch.arange(0, seq_len, device=device,

248 - dtype=torch.int32).view(1, seq_len).ge(

249 - query_lens_t.unsqueeze(-1)).view(

250 - batch_size, 1, 1, seq_len))

251 - causal_mask = torch.triu(torch.ones((batch_size, 1, seq_len, seq_len),

252 - device=device,

253 - dtype=torch.bool),

254 - diagonal=1)

255 - mask = causal_mask.logical_or(len_mask)

256 - mask = torch.concat((past_mask, mask), dim=-1)

257 - attn_bias = (torch.zeros_like(mask, dtype=dtype).masked_fill_(

258 - mask, -math.inf))

259 - attn_metadata = prefill_metadata._replace(attn_bias=attn_bias)

260 - return attn_metadata

261 -

262 - def _set_block_mapping(self, metadata, batch_size, device, dtype):

263 - mask = torch.arange(0,

264 - self.block_size,

265 - device=device,

266 - dtype=torch.int32).unsqueeze(0)

267 - mask = mask >= metadata.block_usage.unsqueeze(-1)

268 - attn_bias = (torch.zeros_like(mask, dtype=dtype).masked_fill_(

269 - mask, -math.inf))

270 - if os.environ.get('VLLM_USE_FAKE_HPU',

271 - '0') == '0' and htorch.utils.internal.is_lazy():

272 - block_mapping = torch.nn.functional.one_hot(metadata.block_groups,

273 - num_classes=batch_size)

274 - else:

275 - # Unfortunately one_hot on CPU/torch.compile mode/eager mode

276 - # doesn't handle out of bounds classes so we need to convert

277 - # all negative values to 0 (block_mapping) or bs (block_groups)

278 - block_groups = metadata.block_groups.to(torch.long)

279 - block_mapping = torch.nn.functional.relu(block_groups)

280 - block_mapping = torch.nn.functional.one_hot(block_mapping,

281 - num_classes=batch_size)

282 - oob_values = block_groups.lt(0)

283 - block_mapping.masked_fill_(oob_values.unsqueeze(-1), 0)

284 - block_groups.masked_fill_(oob_values, batch_size)

285 - metadata = metadata._replace(block_groups=block_groups)

286 - block_mapping = block_mapping.to(dtype)

287 - metadata = metadata._replace(block_mapping=block_mapping,

288 - attn_bias=attn_bias)

289 - return metadata

290 -

291 - def _update_metadata(self, attn_metadata, batch_size, seq_len, device,

292 - dtype):

293 - if attn_metadata.is_prompt:

294 - meta = attn_metadata

295 - attn_metadata = self._set_attn_bias(meta, batch_size, seq_len,

296 - device, dtype)

297 - else:

298 - meta = attn_metadata

299 - attn_metadata = self._set_block_mapping(meta, batch_size, device,

300 - dtype)

301 - return attn_metadata

302 -

303 - def forward(self, *args, **kwargs):

304 - kwargs = kwargs.copy()

305 - selected_token_indices = kwargs.pop('selected_token_indices')

306 - if 'warmup_mode' in kwargs:

307 - kwargs.pop('warmup_mode')

308 - virtual_engine = 0

309 - if 'virtual_engine' in kwargs:

310 - virtual_engine = kwargs.pop('virtual_engine')

311 - input_ids = kwargs['input_ids']

312 - attn_metadata = self._update_metadata(kwargs.pop('attn_metadata'),

313 - input_ids.size(0),

314 - input_ids.size(1),

315 - input_ids.device, self.dtype)

316 - LoraMask.setLoraMask(kwargs.pop('lora_mask'))

317 - with set_forward_context(attn_metadata, self.vllm_config,

318 - virtual_engine):

319 - hidden_states = self.model(*args, **kwargs)

320 - hidden_states = hidden_states.view(-1, hidden_states.shape[-1])

321 - hidden_states = hidden_states.index_select(0,

322 - selected_token_indices)

323 - return hidden_states

324 -

325 - def compute_logits(self, *args, **kwargs):

326 - return self.model.compute_logits(*args, **kwargs)

327 -

328 - def sample(self, *args, **kwargs):

329 - return self.sampler(*args, **kwargs)

330 -

331 -

332 -class PreparePromptMetadata(NamedTuple):

333 - input_tokens: torch.Tensor

334 - input_positions: List[List[int]]

335 - attn_metadata: Optional[AttentionMetadata]

336 - seq_lens: List[int]

337 - query_lens: List[int]

338 - lora_index_mapping: List[List[int]]

339 - lora_prompt_mapping: List[List[int]]

340 - lora_requests: Set[LoRARequest]

341 - multi_modal_kwargs: Optional[Dict[str, BatchedTensorInputs]]

342 - slot_mapping: List[List[int]]

343 - lora_ids: List[int]

344 -

345 - @classmethod

346 - def empty(cls):

347 - return PreparePromptMetadata(input_tokens=[],

348 - input_positions=[],

349 - attn_metadata=None,

350 - seq_lens=[],

351 - query_lens=[],

352 - lora_index_mapping=[],

353 - lora_prompt_mapping=[],

354 - lora_requests=set(),

355 - multi_modal_kwargs=None,

356 - slot_mapping=[],

357 - lora_ids=[])

358 -

359 -

360 -class PrepareDecodeMetadata(NamedTuple):

361 - input_tokens: torch.Tensor

362 - input_positions: List[List[int]]

363 - attn_metadata: Optional[AttentionMetadata]

364 - lora_index_mapping: List[List[int]]

365 - lora_prompt_mapping: List[List[int]]

366 - lora_requests: Set[LoRARequest]

367 - slot_mapping: List[List[int]]

368 - lora_ids: List[int]

369 -

370 - @classmethod

371 - def empty(cls):

372 - return PrepareDecodeMetadata(input_tokens=[],

373 - input_positions=[],

374 - attn_metadata=None,

375 - lora_index_mapping=[],

376 - lora_prompt_mapping=[],

377 - lora_requests=set(),

378 - slot_mapping=[],

379 - lora_ids=[])

380 -

381 -

382 -# How batches are constructed.

383 -class BatchType(IntEnum):

384 - # Every batch is prefill.

385 - PREFILL = 0

386 - # Every batch is decode.

387 - DECODE = 1

388 - # Batch is a mixture of prefill and decode.

389 - MIXED = 2

390 -

391 -

392 -TModelInputForHPU = TypeVar('TModelInputForHPU', bound="ModelInputForHPU")

393 -

394 -

395 -@dataclasses.dataclass(frozen=True)

396 -class ModelInputForHPU(ModelRunnerInputBase):

397 - """

398 - This base class contains metadata needed for the base model forward pass

399 - but not metadata for possible additional steps, e.g., sampling. Model

400 - runners that run additional steps should subclass this method to add

401 - additional fields.

402 - """

403 - input_tokens: Optional[torch.Tensor] = None

404 - input_positions: Optional[torch.Tensor] = None

405 - seq_lens: Optional[List[int]] = None

406 - query_lens: Optional[List[int]] = None

407 - lora_mapping: Optional["LoRAMapping"] = None

408 - lora_requests: Optional[Set[LoRARequest]] = None

409 - attn_metadata: Optional["AttentionMetadata"] = None

410 - multi_modal_kwargs: Optional[Dict[str, torch.Tensor]] = None

411 - real_batch_size: Optional[int] = None

412 - batch_size_padded: Optional[int] = None

413 - virtual_engine: int = 0

414 - lora_ids: Optional[List[int]] = None

415 - async_callback: Optional[Callable] = None

416 - is_first_multi_step: bool = True

417 - is_last_step: bool = True

418 -

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

420 - tensor_dict = {

421 - "input_tokens": self.input_tokens,

422 - "input_positions": self.input_positions,

423 - "lora_requests": self.lora_requests,

424 - "lora_mapping": self.lora_mapping,

425 - "multi_modal_kwargs": self.multi_modal_kwargs,

426 - "real_batch_size": self.real_batch_size,

427 - "batch_size_padded": self.batch_size_padded,

428 - "virtual_engine": self.virtual_engine,

429 - "lora_ids": self.lora_ids,

430 - "is_first_multi_step": self.is_first_multi_step,

431 - "is_last_step": self.is_last_step,

432 - }

433 - _add_attn_metadata_broadcastable_dict(tensor_dict, self.attn_metadata)

434 - return tensor_dict

435 -

436 - @classmethod

437 - def from_broadcasted_tensor_dict(

438 - cls: Type[TModelInputForHPU],

439 - tensor_dict: Dict[str, Any],

440 - attn_backend: Optional["AttentionBackend"] = None,

441 - ) -> TModelInputForHPU:

442 - if attn_backend is not None:

443 - tensor_dict = _init_attn_metadata_from_tensor_dict(

444 - attn_backend, tensor_dict)

445 - return cls(**tensor_dict)

446 -

447 -

448 -@dataclasses.dataclass(frozen=True)

449 -class ModelInputForHPUWithSamplingMetadata(ModelInputForHPU):

450 - """

451 - Used by the ModelRunner.

452 - """

453 - sampling_metadata: Optional["SamplingMetadata"] = None

454 - # Used for speculative decoding. We do not broadcast it because it is only

455 - # used by the driver worker.

456 - is_prompt: Optional[bool] = None

457 -

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

459 - tensor_dict = {

460 - "input_tokens": self.input_tokens,

461 - "input_positions": self.input_positions,

462 - "lora_requests": self.lora_requests,

463 - "lora_mapping": self.lora_mapping,

464 - "multi_modal_kwargs": self.multi_modal_kwargs,

465 - "lora_ids": self.lora_ids,

466 - }

467 - _add_attn_metadata_broadcastable_dict(tensor_dict, self.attn_metadata)

468 - _add_sampling_metadata_broadcastable_dict(tensor_dict,

469 - self.sampling_metadata)

470 - return tensor_dict

471 -

472 - @classmethod

473 - def from_broadcasted_tensor_dict(

474 - cls,

475 - tensor_dict: Dict[str, Any],

476 - attn_backend: Optional["AttentionBackend"] = None,

477 - ) -> "ModelInputForHPUWithSamplingMetadata":

478 - tensor_dict = _init_sampling_metadata_from_tensor_dict(tensor_dict)

479 - # FIXME(kzawora): this fails for whatever reason - why?

480 - if attn_backend is not None:

481 - tensor_dict = _init_attn_metadata_from_tensor_dict(

482 - attn_backend, tensor_dict)

483 - return cls(**tensor_dict)

484 -

485 -

486 -class HPUModelRunnerBase(ModelRunnerBase[TModelInputForHPU]):

487 - """

488 - Helper class for shared methods between GPU model runners.

489 - """

490 - _model_input_cls: Type[TModelInputForHPU]

491 -

492 - def __init__(

493 - self,

494 - vllm_config: VllmConfig,

495 - is_driver_worker: bool = False,

496 - return_hidden_states: bool = False,

497 - ):

498 - ModelRunnerBase.__init__(self, vllm_config=vllm_config)

499 - environment.set_model_config(self.model_config)

500 - self.is_driver_worker = is_driver_worker

501 - self.return_hidden_states = return_hidden_states

502 -

503 - self.sliding_window = (self.model_config.get_sliding_window()

504 - if self.model_config is not None else None)

505 - self.device_config = (self.device_config if self.device_config

506 - is not None else DeviceConfig())

507 - self.device = self.device_config.device

508 - self.enforce_eager = self.model_config.enforce_eager

509 - self.max_num_seqs = self.scheduler_config.max_num_seqs

510 - # NOTE(kzawora): Change that to scheduler_config.max_num_prefill_seqs

511 - # once padding-aware scheduling gets merged

512 - self.max_num_prefill_seqs = 64

513 - self.max_model_len = self.scheduler_config.max_model_len

514 - self.max_num_batched_tokens = \

515 - self.scheduler_config.max_num_batched_tokens

516 - self.block_size = self.cache_config.block_size

517 -

518 - self.pin_memory = is_pin_memory_available()

519 - self.kv_cache_dtype = self.cache_config.cache_dtype

520 -

521 - self.attn_backend = get_attn_backend(

522 - self.model_config.get_head_size(),

523 - self.model_config.dtype,

524 - self.kv_cache_dtype,

525 - self.block_size,

526 - self.model_config.is_attention_free,

527 - )

528 -

529 - # Lazy initialization

530 - self.lora_manager: LRUCacheWorkerLoRAManager = None

531 - self.model: torch.nn.Module = None

532 - self.inc_initialized_successfully = False

533 -

534 - # Profiler stats

535 - self.profiler = HabanaHighLevelProfiler()

536 - self.profiler_counter_helper = HabanaProfilerCounterHelper()

537 - self.seen_configs: set = set()

538 - self._mem_margin: Optional[int] = None

539 - HPUBucketingContext = get_bucketing_context()

540 - self.bucketing_ctx = HPUBucketingContext(self.max_num_seqs,

541 - self.max_num_prefill_seqs,

542 - self.block_size,

543 - self.max_num_batched_tokens,

544 - False, self.max_model_len)

545 - self.graphed_buckets: Set[Any] = set()

546 - self._set_gc_threshold()

547 - if self.vllm_config.cache_config.enable_prefix_caching:

548 - os.environ.setdefault("VLLM_CONTIGUOUS_PA", "False")

549 - assert os.environ.get(

550 - "VLLM_CONTIGUOUS_PA",

551 - "").lower() != "true", "Contiguous PA doesn't support APC"

552 - self.use_contiguous_pa = envs.VLLM_USE_HPU_CONTIGUOUS_CACHE_FETCH

553 -

554 - # For multi-step scheduling

555 - self.cached_step_outputs: List[torch.Tensor] = []

556 - # For delayed sampling

557 - self.cached_step_inputs: List[

558 - ModelInputForHPUWithSamplingMetadata] = []

559 -

560 - def _set_gc_threshold(self) -> None:

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

562 - # for comprehensive description of gc generations.

563 - # We can either use VLLM_GC_THR_GEN[0-2] (this has higher priority)

564 - # to set particular generation threshold or use simpler

565 - # VLLM_GC_THR_MULTIPLIER to multiply default values.

566 - default_gc_thrs = list(gc.get_threshold())

567 - requested_gc_thrs = [0] * len(default_gc_thrs)

568 - for i in range(len(default_gc_thrs)):

569 - requested_gc_thrs[i] = int(

570 - os.environ.get(f'VLLM_GC_THR_GEN{i}', default_gc_thrs[i]))

571 - if requested_gc_thrs == default_gc_thrs:

572 - gc_thr_multiplier = int(os.environ.get('VLLM_GC_THR_MULTIPLIER',

573 - 2))

574 - requested_gc_thrs = [

575 - t * gc_thr_multiplier for t in default_gc_thrs

576 - ]

577 - gc.set_threshold(*requested_gc_thrs)

578 -

579 - self.skip_warmup = os.environ.get('VLLM_SKIP_WARMUP',

580 - 'false').lower() == 'true'

581 -

582 - def load_model(self) -> None:

583 - import habana_frameworks.torch.core as htcore

584 - if self.model_config.quantization == 'inc' or \

585 - self.model_config.quantization == 'fp8':

586 - htcore.hpu_set_env()

587 - with HabanaMemoryProfiler() as m:

588 - with HabanaMemoryProfiler() as m_getmodel:

589 - self.model = get_model(vllm_config=self.vllm_config)

590 - msg = ("Pre-loading model weights on "

591 - f"{next(self.model.parameters()).device} "

592 - f"took {m_getmodel.get_summary_string()}")

593 - logger.info(msg)

594 -

595 - if self.lora_config:

596 - assert hasattr(self.model, "embedding_modules"

597 - ), "Model does not have embedding_modules"

598 - assert hasattr(

599 - self.model, "embedding_padding_modules"

600 - ), "Model does not have embedding_padding_modules"

601 - assert not self.lora_config.bias_enabled, \

602 - "Bias support in LoRA is not enabled in HPU yet."

603 - assert not self.lora_config.fully_sharded_loras, \

604 - "Fully sharded LoRAs is not enabled in HPU yet."

605 -

606 - # Use get_text_config() in case of multimodal models

607 - text_config = self.model_config.hf_config.get_text_config()

608 -

609 - self.lora_manager = LRUCacheWorkerLoRAManager(

610 - self.scheduler_config.max_num_seqs,

611 - self.scheduler_config.max_num_batched_tokens,

612 - self.vocab_size,

613 - self.lora_config,

614 - self.device,

615 - self.model.embedding_modules,

616 - self.model.embedding_padding_modules,

617 - max_position_embeddings=text_config.

618 - max_position_embeddings,

619 - )

620 - self.model = self.lora_manager.create_lora_manager(self.model)

621 -

622 - if self.model_config.quantization == 'inc':

623 - logger.info("Preparing model with INC..")

624 - with HabanaMemoryProfiler() as m_inc:

625 - from neural_compressor.torch.quantization import (

626 - FP8Config, convert, prepare)

627 - config = FP8Config.from_json_file(

628 - os.getenv("QUANT_CONFIG", ""))

629 - if config.measure:

630 - self.model = prepare(self.model, config)

631 - elif config.quantize:

632 - self.model = convert(self.model, config)

633 - htcore.hpu_initialize(self.model,

634 - mark_only_scales_as_const=True)

635 - self.inc_initialized_successfully = True

636 - logger.info("Preparing model with INC took %s",

637 - m_inc.get_summary_string())

638 - else:

639 - self.model = self.model.to("hpu")

640 - htcore.mark_step()

641 - modify_decoder_layer(self.model)

642 - torch.hpu.synchronize()

643 -

644 - with HabanaMemoryProfiler() as m_wrap:

645 - self.model = _maybe_wrap_in_hpu_graph(

646 - self.model, vllm_config=self.vllm_config)

647 - msg = f"Wrapping in HPU Graph took {m_wrap.get_summary_string()}"

648 - logger.info(msg)

649 -

650 - self.model_memory_usage = m.consumed_device_memory

651 - msg = f"Loading model weights took in total {m.get_summary_string()}"

652 - logger.info(msg)

653 -

654 - def _add_dummy_seq(self, seq_group_metadata_list, is_prompt):

655 - real_batch_size = len(seq_group_metadata_list)

656 - batch_size_padded = self.bucketing_ctx.get_padded_batch_size(

657 - real_batch_size, is_prompt)

658 - batch_size_padding = batch_size_padded - real_batch_size

659 -

660 - seq_group_metadata_list = seq_group_metadata_list.copy()

661 -

662 - if batch_size_padding > 0:

663 - dummy_seq_group_metadata = self.create_dummy_seq_group_metadata(

664 - 0, 0, is_prompt)

665 - seq_group_metadata_list.extend(dummy_seq_group_metadata

666 - for _ in range(batch_size_padding))

667 - return seq_group_metadata_list, real_batch_size, batch_size_padded

668 -

669 - def _maybe_wrap_in_hpu_graph(self, *args, **kwargs):

670 - return htorch.hpu.wrap_in_hpu_graph(

671 - HpuModelAdapter(*args, **kwargs), disable_tensor_cache=True

672 - ) if htorch.utils.internal.is_lazy() else HpuModelAdapter(

673 - *args, **kwargs)

674 -

675 - def get_model(self) -> nn.Module:

676 - return self.model

677 -

678 - def _use_graphs(self, batch_size, seq_len, is_prompt):

679 - if self.enforce_eager:

680 - return False

681 - if self.skip_warmup:

682 - return True

683 - return (batch_size, seq_len, is_prompt) in self.graphed_buckets

684 -

685 - def _is_valid_bucket(self, bucket):

686 - return bucket[0] * bucket[1] <= self.max_num_batched_tokens

687 -

688 - def _prepare_prompt(

689 - self,

690 - seq_group_metadata_list: List[SequenceGroupMetadata],

691 - ) -> PreparePromptMetadata:

692 - input_tokens: List[List[int]] = []

693 - input_positions: List[List[int]] = []

694 - slot_mapping: List[List[int]] = []

695 - lora_index_mapping: List[List[int]] = []

696 - lora_prompt_mapping: List[List[int]] = []

697 - lora_requests: Set[LoRARequest] = set()

698 -

699 - seq_lens: List[int] = []

700 - context_lens: List[int] = []

701 - query_lens: List[int] = []

702 - prefix_block_tables: List[List[int]] = []

703 - multi_modal_kwargs_list: List[MultiModalKwargs] = []

704 -

705 - if len(seq_group_metadata_list) == 0:

706 - return PreparePromptMetadata.empty()

707 -

708 - for seq_group_metadata in seq_group_metadata_list:

709 - assert seq_group_metadata.is_prompt

710 - seq_ids = list(seq_group_metadata.seq_data.keys())

711 - assert len(seq_ids) == 1

712 - seq_id = seq_ids[0]

713 -

714 - computed_block_nums = seq_group_metadata.computed_block_nums

715 - if (self.scheduler_config is not None

716 - and self.scheduler_config.chunked_prefill_enabled

717 - and not (computed_block_nums is None

718 - or computed_block_nums == [])):

719 - raise RuntimeError(

720 - "chunked prefill cannot be used with prefix caching "

721 - "now.")

722 -

723 - token_chunk_size = seq_group_metadata.token_chunk_size

724 - seq_data = seq_group_metadata.seq_data[seq_id]

725 - context_len = seq_data.get_num_computed_tokens()

726 - # We should use get_len here because in case of preemption

727 - # it contains output tokens.

728 - seq_len = min(seq_data.get_len(), context_len + token_chunk_size)

729 - prompt_tokens = seq_data.get_token_ids()[context_len:seq_len]

730 - seq_lens.append(seq_len)

731 -

732 - # NOTE: This only works for oooooooxxx style attention.

733 - if computed_block_nums is not None and len(

734 - computed_block_nums) > 0 and self.sliding_window is None:

735 - # Prefix is not supported with sliding_window

736 - context_len = len(computed_block_nums) * self.block_size

737 - if context_len == seq_len \

738 - and self.vllm_config.cache_config.enable_prefix_caching:

739 - # Fully cached prompt - compute only last token

740 - context_len = context_len - 1

741 - prompt_tokens = prompt_tokens[context_len:]

742 - prefix_block_tables.append(computed_block_nums)

743 - elif self.scheduler_config.chunked_prefill_enabled:

744 - if seq_group_metadata.block_tables is not None:

745 - # Prefill has chunked before.

746 - block_table = seq_group_metadata.block_tables[seq_id]

747 - prefix_block_tables.append(block_table)

748 - else:

749 - # The first prefill.

750 - prefix_block_tables.append([])

751 - else:

752 - prefix_block_tables.append([])

753 - # Right now, prefill start is always 0. However, this

754 - # assumption can be changed once chunked prefill is introduced.

755 - assert context_len == 0

756 -

757 - # actual prompt lens

758 - context_lens.append(context_len)

759 - query_lens.append(seq_len - context_len)

760 - input_tokens.append(prompt_tokens)

761 - # NOTE(woosuk): Here we assume that the first token in the prompt

762 - # is always the first token in the sequence.

763 - input_positions.append(list(range(context_len, seq_len)))

764 -

765 - mm_kwargs = seq_group_metadata.multi_modal_data

766 - if mm_kwargs:

767 - multi_modal_kwargs_list.append(mm_kwargs)

768 -

769 - if seq_group_metadata.block_tables is None:

770 - # During memory profiling, the block tables are not initialized

771 - # yet. In this case, we just use a dummy slot mapping.

772 - slot_mapping.append([_PAD_SLOT_ID] * seq_len)

773 - continue

774 -

775 - # Compute the slot mapping.

776 - slot_mapping.append([])

777 - block_table = seq_group_metadata.block_tables[seq_id]

778 -

779 - # Mask the [0, start_idx) tokens of the prompt with _PAD_SLOT_ID,

780 - # where start_idx is max(0, seq_len - sliding_window).

781 - # For example, if the prompt len is 10, sliding window is 8, and

782 - # block size is 4, the first two tokens are masked and the slot

783 - # mapping will be [-1, -1, 2, 3, 4, 5, 6, 7, 0, 1].

784 - start_idx = 0

785 - if self.sliding_window is not None:

786 - assert context_len == 0, (

787 - "Prefix caching is currently not supported with "

788 - "sliding window attention")

789 - start_idx = max(0, seq_len - self.sliding_window)

790 - for i in range(context_len, seq_len):

791 - if i < start_idx:

792 - slot_mapping[-1].append(_PAD_SLOT_ID)

793 - continue

794 -

795 - block_number = block_table[i // self.block_size]

796 - block_offset = i % self.block_size

797 - slot = block_number * self.block_size + block_offset

798 - slot_mapping[-1].append(slot)

799 -

800 - max_query_len = max(query_lens)

801 - sum_query_len = sum(query_lens)

802 - real_num_seqs = len(query_lens)

803 - assert max_query_len > 0

804 -

805 - max_prompt_len = max(

806 - self.bucketing_ctx.get_padded_prompt_seq_len(max_query_len),

807 - self.block_size)

808 -

809 - lora_ids: List[int] = []

810 - for seq_group_metadata, context_len in zip(seq_group_metadata_list,

811 - context_lens):

812 - lora_id = seq_group_metadata.lora_int_id

813 - lora_ids.append(lora_id)

814 -

815 - if lora_id > 0:

816 - lora_requests.add(seq_group_metadata.lora_request)

817 -

818 - lora_index_mapping += [lora_id] * max_prompt_len

819 - lora_prompt_mapping.extend(

820 - [lora_id] *

821 - (max_prompt_len

822 - if seq_group_metadata.sampling_params.prompt_logprobs else 1))

823 -

824 - if any(context_lens):

825 - assert not self.scheduler_config.chunked_prefill_enabled

826 - # prefix caching

827 -

828 - max_num_block = max(len(bt) for bt in prefix_block_tables)

829 - prefix_block_list = list(

830 - itertools.chain.from_iterable(

831 - bt if len(bt) == max_num_block else bt +

832 - ([_PAD_BLOCK_ID] * (max_num_block - len(bt)))

833 - for bt in prefix_block_tables))

834 -

835 - pad_len = len(prefix_block_list)

836 - prefix_block_list = pad_list(prefix_block_list, pad_len,

837 - _PAD_BLOCK_ID)

838 -

839 - prefix_block_list_tensor = torch.tensor(prefix_block_list,

840 - dtype=torch.long,

841 - device=self.device)

842 - else:

843 - prefix_block_list_tensor = None

844 -

845 - input_tokens = make_tensor_with_pad(input_tokens,

846 - max_len=max_prompt_len,

847 - pad=0,

848 - dtype=torch.long,

849 - device=self.device)

850 -

851 - input_positions = make_tensor_with_pad(input_positions,

852 - max_len=max_prompt_len,

853 - pad=0,

854 - dtype=torch.long,

855 - device=self.device)

856 -

857 - slot_mapping = make_tensor_with_pad(slot_mapping,

858 - max_len=max_prompt_len,

859 - pad=_PAD_SLOT_ID,

860 - dtype=torch.long,

861 - device=self.device)

862 -

863 - seq_lens_tensor = torch.tensor(seq_lens,

864 - dtype=torch.long,

865 - device=self.device)

866 -

867 - context_lens_tensor = torch.tensor(context_lens,

868 - dtype=torch.long,

869 - device=self.device)

870 -

871 - block_indices, block_offsets = precompute_indices_and_offsets(

872 - self.block_size, slot_mapping, True)

873 - attn_metadata = self.attn_backend.make_metadata(

874 - is_prompt=True,

875 - block_list=prefix_block_list_tensor,

876 - block_mapping=None,

877 - block_usage=None,

878 - block_indices=block_indices,

879 - block_offsets=block_offsets,

880 - block_groups=None,

881 - attn_bias=None,

882 - seq_lens_tensor=seq_lens_tensor,

883 - context_lens_tensor=context_lens_tensor,

884 - num_prefills=real_num_seqs,

885 - num_prefill_tokens=sum_query_len,

886 - num_decode_tokens=0,

887 - slot_mapping=slot_mapping,

888 - multi_modal_placeholder_index_maps=

889 - None, # FIXME(kzawora): multi-modality will not work here

890 - enable_kv_scales_calculation=False,

891 - )

892 - multi_modal_kwargs = MultiModalKwargs.batch(multi_modal_kwargs_list)

893 -

894 - return PreparePromptMetadata(input_tokens=input_tokens,

895 - input_positions=input_positions,

896 - attn_metadata=attn_metadata,

897 - seq_lens=seq_lens,

898 - query_lens=query_lens,

899 - lora_index_mapping=lora_index_mapping,

900 - lora_prompt_mapping=lora_prompt_mapping,

901 - lora_requests=lora_requests,

902 - multi_modal_kwargs=multi_modal_kwargs,

903 - slot_mapping=slot_mapping,

904 - lora_ids=lora_ids)

905 -

906 - def _prepare_decode(

907 - self,

908 - seq_group_metadata_list: List[SequenceGroupMetadata],

909 - output=None,

910 - ) -> PrepareDecodeMetadata:

911 - input_tokens: List[List[int]] = []

912 - input_positions: List[List[int]] = []

913 - slot_mapping: List[List[int]] = []

914 - seq_lens: List[int] = []

915 - block_tables: List[List[int]] = []

916 - lora_index_mapping: List[List[int]] = []

917 - lora_prompt_mapping: List[List[int]] = []

918 - lora_requests: Set[LoRARequest] = set()

919 -

920 - if len(seq_group_metadata_list) == 0:

921 - return PrepareDecodeMetadata.empty()

922 - lora_ids: List[int] = []

923 -

924 - dummy_slots = itertools.cycle(

925 - range(_PAD_SLOT_ID, _PAD_SLOT_ID + self.block_size))

926 -

927 - for seq_group_metadata in seq_group_metadata_list:

928 - assert not seq_group_metadata.is_prompt

929 - assert seq_group_metadata.token_chunk_size == 1

930 -

931 - seq_ids = list(seq_group_metadata.seq_data.keys())

932 - lora_id = seq_group_metadata.lora_int_id

933 - lora_ids.append(lora_id)

934 -

935 - if lora_id > 0:

936 - lora_requests.add(seq_group_metadata.lora_request)

937 -

938 - for seq_id in seq_ids:

939 - seq_data = seq_group_metadata.seq_data[seq_id]

940 - if output is None:

941 - generation_token = seq_data.get_last_token_id()

942 - input_tokens.append([generation_token])

943 -

944 - seq_len = seq_data.get_len()

945 - position = seq_len - 1

946 - input_positions.append([position])

947 -

948 - seq_len = seq_len if self.sliding_window is None else min(

949 - seq_len, self.sliding_window)

950 - seq_lens.append(seq_len)

951 -

952 - block_table = seq_group_metadata.block_tables[seq_id]

953 - num_fully_occupied_blocks = position // self.block_size

954 - block_table = block_table[:num_fully_occupied_blocks + 1]

955 -

956 - if len(block_table) == 0:

957 - block_number = _PAD_BLOCK_ID

958 - else:

959 - block_number = block_table[position // self.block_size]

960 - if block_number == _PAD_BLOCK_ID:

961 - slot = next(dummy_slots)

962 - else:

963 - block_offset = position % self.block_size

964 - slot = block_number * self.block_size + block_offset

965 - slot_mapping.append([slot])

966 - lora_index_mapping.append(lora_id)

967 - lora_prompt_mapping.append(lora_id)

968 -

969 - if self.sliding_window is not None:

970 - sliding_window_blocks = (self.sliding_window //

971 - self.block_size)

972 - block_table = block_table[-sliding_window_blocks:]

973 - block_tables.append(block_table)

974 -

975 - if output is None:

976 - input_tokens = torch.tensor(input_tokens,

977 - dtype=torch.long,

978 - device=self.device)

979 - else:

980 - real_batch_size = len(seq_group_metadata_list)

981 - input_tokens = output[:real_batch_size]

982 -

983 - input_positions = torch.tensor(input_positions,

984 - dtype=torch.long,

985 - device=self.device)

986 -

987 - num_decode_tokens = sum(seq_lens)

988 -

989 - last_block_usage = [

990 - slot[0] % self.block_size + 1 for slot in slot_mapping

991 - ]

992 - block_groups = [[i] * len(bt) for i, bt in enumerate(block_tables)]

993 - block_usage = [[self.block_size] * (len(bt) - 1) + [lbu]

994 - for bt, lbu in zip(block_tables, last_block_usage)

995 - if bt]

996 -

997 - block_list = flatten(block_tables)

998 - block_groups = flatten(block_groups)

999 - block_usage = flatten(block_usage)

1000 -

1001 - assert len(block_list) == len(block_groups)

1002 - assert len(block_list) == len(block_usage)

1003 -

1004 - padding_fn = None

1005 - if self.use_contiguous_pa:

1006 - block_bucket_size = max(max(block_list) + 1, len(block_list))

1007 - block_bucket_size = self.bucketing_ctx.get_padded_decode_num_blocks(

1008 - block_bucket_size)

1009 - indices: List[Any]

1010 - indices = [None] * block_bucket_size

1011 - for i, bid in enumerate(block_list):

1012 - indices[bid] = i

1013 - padding_fn = lambda tensor, pad_value: gather_list(

1014 - tensor, indices, pad_value)

1015 - else:

1016 - block_bucket_size = \

1017 - self.bucketing_ctx.get_padded_decode_num_blocks(

1018 - len(block_list))

1019 - padding_fn = lambda tensor, pad_value: pad_list(

1020 - tensor, block_bucket_size, pad_value)

1021 -

1022 - block_list = padding_fn(block_list, _PAD_BLOCK_ID)

1023 - block_groups = padding_fn(block_groups, -1)

1024 - block_usage = padding_fn(block_usage, 1)

1025 -

1026 - block_list = torch.tensor(block_list,

1027 - dtype=torch.int,

1028 - device=self.device)

1029 - block_groups = torch.tensor(block_groups,

1030 - dtype=torch.int,

1031 - device=self.device)

1032 - block_usage = torch.tensor(block_usage,

1033 - dtype=self.model_config.dtype,

1034 - device=self.device)

1035 - slot_mapping = torch.tensor(slot_mapping,

1036 - dtype=torch.long,

1037 - device=self.device)

1038 -

1039 - block_indices, block_offsets = precompute_indices_and_offsets(

1040 - self.block_size, slot_mapping, False)

1041 -

1042 - attn_metadata = self.attn_backend.make_metadata(

1043 - is_prompt=False,

1044 - block_list=block_list,

1045 - block_mapping=None,

1046 - block_usage=block_usage,

1047 - block_indices=block_indices,

1048 - block_offsets=block_offsets,

1049 - block_groups=block_groups,

1050 - attn_bias=None,

1051 - seq_lens_tensor=None,

1052 - context_lens_tensor=None,

1053 - num_prefills=0,

1054 - num_prefill_tokens=0,

1055 - num_decode_tokens=num_decode_tokens,

1056 - slot_mapping=slot_mapping,

1057 - multi_modal_placeholder_index_maps=None,

1058 - enable_kv_scales_calculation=False,

1059 - )

1060 - return PrepareDecodeMetadata(input_tokens=input_tokens,

1061 - input_positions=input_positions,

1062 - attn_metadata=attn_metadata,

1063 - lora_index_mapping=lora_index_mapping,

1064 - lora_prompt_mapping=lora_prompt_mapping,

1065 - lora_requests=lora_requests,

1066 - slot_mapping=slot_mapping,

1067 - lora_ids=lora_ids)

1068 -

1069 - def prepare_input_tensors(

1070 - self,

1071 - seq_group_metadata_list: List[SequenceGroupMetadata],

1072 - ) -> Tuple[TModelInputForHPU, SamplingMetadata]:

1073 - if len(seq_group_metadata_list) == 0:

1074 - return self._model_input_cls(), None

1075 -

1076 - input_tokens = None

1077 - input_positions = None

1078 - lora_mapping = None

1079 - lora_requests = None

1080 - multi_modal_kwargs = None

1081 - batch_type = None

1082 - seq_lens = None

1083 - query_lens = None

1084 - real_batch_size = None

1085 - batch_size_padded = None

1086 -

1087 - self.event_start = self.profiler.get_timestamp_us()

1088 - is_prompt = seq_group_metadata_list[0].is_prompt

1089 - base_event_name = 'prompt' if is_prompt else 'decode'

1090 - self.profiler.start('internal', base_event_name)

1091 -

1092 - seq_group_metadata_list, real_batch_size, batch_size_padded = (

1093 - self._add_dummy_seq(seq_group_metadata_list, is_prompt))

1094 -

1095 - prefill_reqs = []

1096 - decode_reqs = []

1097 - for seq_group_meta in seq_group_metadata_list:

1098 - if seq_group_meta.is_prompt:

1099 - prefill_reqs.append(seq_group_meta)

1100 - else:

1101 - decode_reqs.append(seq_group_meta)

1102 -

1103 - # Prepare input tensors.

1104 - (

1105 - input_tokens,

1106 - input_positions,

1107 - prefill_attn_metadata,

1108 - seq_lens,

1109 - query_lens,

1110 - lora_index_mapping,

1111 - lora_prompt_mapping,

1112 - lora_requests,

1113 - multi_modal_kwargs,

1114 - slot_mapping,

1115 - lora_ids,

1116 - ) = self._prepare_prompt(prefill_reqs)

1117 - (

1118 - decode_input_tokens,

1119 - decode_input_positions,

1120 - decode_attn_metadata,

1121 - decode_lora_index_mapping,

1122 - decode_lora_prompt_mapping,

1123 - decode_lora_requests,

1124 - decode_slot_mapping,

1125 - decode_lora_ids,

1126 - ) = self._prepare_decode(decode_reqs)

1127 - sampling_metadata = SamplingMetadata.prepare(seq_group_metadata_list,

1128 - seq_lens, query_lens,

1129 - self.device,

1130 - self.pin_memory)

1131 -

1132 - if not self.scheduler_config.chunked_prefill_enabled:

1133 - assert (len(prefill_reqs) and len(decode_reqs)) == 0

1134 -

1135 - num_prefills = len(seq_lens)

1136 - num_prefill_tokens = len(input_tokens)

1137 - num_decode_tokens = len(decode_input_tokens)

1138 -

1139 - # NOTE(kzawora): Here we diverge from GPU code - we don't

1140 - # support mixed batches, so we either use decode or prefill

1141 - # inputs, without coalescing.

1142 - assert (num_prefills == 0 and num_decode_tokens > 0) or (

1143 - num_prefills > 0

1144 - and num_decode_tokens == 0), "HPU does not support mixed batches!"

1145 - if num_decode_tokens > 0:

1146 - input_tokens = decode_input_tokens

1147 - input_positions = decode_input_positions

1148 - slot_mapping = decode_slot_mapping

1149 - lora_index_mapping = decode_lora_index_mapping

1150 - lora_prompt_mapping = decode_lora_prompt_mapping

1151 - lora_requests = decode_lora_requests

1152 - lora_ids = decode_lora_ids

1153 -

1154 - # FIXME: We need to adjust selected_token_indices to accommodate

1155 - # for padding

1156 - max_len = input_tokens.size(1)

1157 - paddings = [max_len - q for q in query_lens]

1158 - paddings = [0] + paddings[:-1]

1159 - paddings = list(itertools.accumulate(paddings))

1160 - paddings_prompt_logprobs = []

1161 - for i, seq_group_metadata in enumerate(seq_group_metadata_list):

1162 - if seq_group_metadata.sampling_params.prompt_logprobs is not None \

1163 - and seq_group_metadata.is_prompt:

1164 - paddings_prompt_logprobs += ([paddings[i]] * seq_lens[i])

1165 - paddings = torch.tensor(

1166 - paddings_prompt_logprobs if paddings_prompt_logprobs else paddings,

1167 - dtype=sampling_metadata.selected_token_indices.dtype,

1168 - device=sampling_metadata.selected_token_indices.device)

1169 - sampling_metadata.selected_token_indices.add_(paddings)

1170 -

1171 - if self.lora_config:

1172 - lora_mapping = LoRAMapping(

1173 - **dict(index_mapping=lora_index_mapping,

1174 - prompt_mapping=lora_prompt_mapping,

1175 - is_prefill=(num_prefills > 0)))

1176 - else:

1177 - lora_mapping = None

1178 -

1179 - if (prefill_attn_metadata is not None

1180 - and decode_attn_metadata is not None):

1181 - batch_type = BatchType.MIXED

1182 - raise NotImplementedError("Mixed batch is not supported on HPU")

1183 - elif prefill_attn_metadata is not None:

1184 - batch_type = BatchType.PREFILL

1185 - else:

1186 - batch_type = BatchType.DECODE

1187 -

1188 - metadata_dict = {

1189 - "input_tokens": input_tokens,

1190 - "input_positions": input_positions,

1191 - "selected_token_indices": sampling_metadata.selected_token_indices,

1192 - "lora_requests": lora_requests,

1193 - "lora_mapping": lora_mapping,

1194 - "multi_modal_kwargs": multi_modal_kwargs,

1195 - "num_prefill_tokens": num_prefill_tokens,

1196 - "num_decode_tokens": num_decode_tokens,

1197 - "slot_mapping": slot_mapping,

1198 - "num_prefills": num_prefills,

1199 - "batch_type": batch_type,

1200 - "seq_lens": seq_lens,

1201 - "query_lens": query_lens

1202 - }

1203 - if prefill_attn_metadata is not None:

1204 - metadata_dict.update(prefill_attn_metadata.asdict_zerocopy())

1205 - else:

1206 - assert decode_attn_metadata is not None

1207 - metadata_dict.update(decode_attn_metadata.asdict_zerocopy())

1208 -

1209 - attn_metadata = prefill_attn_metadata if \

1210 - prefill_attn_metadata is not None else decode_attn_metadata

1211 -

1212 - return self._model_input_cls(input_tokens=input_tokens,

1213 - seq_lens=seq_lens,

1214 - query_lens=query_lens,

1215 - input_positions=input_positions,

1216 - attn_metadata=attn_metadata,

1217 - lora_requests=lora_requests,

1218 - lora_mapping=lora_mapping,

1219 - multi_modal_kwargs=multi_modal_kwargs,

1220 - real_batch_size=real_batch_size,

1221 - batch_size_padded=batch_size_padded,

1222 - lora_ids=lora_ids), \

1223 - sampling_metadata

1224 -

1225 - def _seq_len(self, attn_metadata):

1226 - if attn_metadata.num_prefills != 0:

1227 - return attn_metadata.slot_mapping.size(1)

1228 - else:

1229 - return attn_metadata.block_list.numel()

1230 -

1231 - def trim_attn_metadata(self, metadata: AttentionMetadata) -> object:

1232 - # NOTE(kzawora): To anyone working on this in the future:

1233 - # Trimming metadata is required when using HPUGraphs.

1234 - # Attention metadata is going to be hashed by PT bridge, and

1235 - # appropriate HPUGraphs will be matched based on all inputs' hash.

1236 -

1237 - # Before you put more keys in here, make sure you know their

1238 - # value type and make sure you know how it's going to be hashed.

1239 - # You can find that information in input_hash function

1240 - # in habana_frameworks/torch/hpu/graphs.py. You can also hash

1241 - # it manually with torch.hpu.graphs.input_hash(attention_metadata)

1242 -

1243 - # If you use primitive types here - they will get hashed based

1244 - # on their value. You *will* get lots of excessive graph captures

1245 - # (and an OOM eventually) if you decide to put something like

1246 - # seq_len int here.

1247 - # If you absolutely need a scalar, put it in a tensor. Tensors

1248 - # get hashed using their metadata, not their values:

1249 - # input_hash(torch.tensor(123)) == input_hash(torch.tensor(321))

1250 - # input_hash(123) != input_hash(321)

1251 - # input_hash("abc") != input_hash("cba")

1252 - attention_metadata = subtuple(metadata, 'TrimmedAttentionMetadata', [

1253 - 'attn_bias',

1254 - 'seq_lens_tensor',

1255 - 'context_lens_tensor',

1256 - 'block_list',

1257 - 'block_mapping',

1258 - 'block_usage',

1259 - 'slot_mapping',

1260 - 'is_prompt',

1261 - 'block_indices',

1262 - 'block_offsets',

1263 - 'block_groups',

1264 - ])

1265 - return attention_metadata

1266 -

1267 - def create_dummy_seq_group_metadata(self,

1268 - group_id,

1269 - seq_len,

1270 - is_prompt,

1271 - lora_request=None):

1272 - sampling_params = SamplingParams(temperature=0)

1273 - num_blocks = math.ceil(seq_len / self.block_size)

1274 - seq_len = max(seq_len, 1)

1275 - if is_prompt:

1276 - input_len = seq_len

1277 - output_len = 0

1278 - block_tables = None

1279 - else:

1280 - input_len = seq_len - 1

1281 - output_len = 1

1282 - block_tables = {group_id: [_PAD_BLOCK_ID] * num_blocks}

1283 - prompt_token_ids = [0] * input_len

1284 - output_token_ids = [1] * output_len

1285 - prompt_token_ids_array = array('l', prompt_token_ids) # noqa: F821

1286 - seq_data = SequenceData(prompt_token_ids_array)

1287 - seq_data.output_token_ids = output_token_ids

1288 - return SequenceGroupMetadata(request_id=str(group_id),

1289 - is_prompt=(output_len == 0),

1290 - seq_data={group_id: seq_data},

1291 - sampling_params=sampling_params,

1292 - block_tables=block_tables,

1293 - lora_request=lora_request)

1294 -

1295 - def profile_run(self) -> None:

1296 - num_layers = self.model_config.get_num_layers(self.parallel_config)

1297 - kv_caches = [None] * num_layers

1298 - bind_kv_cache(

1299 - self.vllm_config.compilation_config.static_forward_context,

1300 - [kv_caches])

1301 - _, max_seq_len = self.bucketing_ctx.get_max_prompt_shape()

1302 - max_batch_size = min(self.max_num_seqs,

1303 - self.max_num_batched_tokens // max_seq_len)

1304 - self.warmup_scenario(max_batch_size, max_seq_len, True, kv_caches,

1305 - False, True)

1306 - return

1307 -

1308 - def warmup_scenario(self,

1309 - batch_size,

1310 - seq_len,

1311 - is_prompt,

1312 - kv_caches,

1313 - is_pt_profiler_run=False,

1314 - is_lora_profile_run=False) -> None:

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

1316 - scenario_name = ("warmup_"

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

1318 - f"bs{batch_size}_"

1319 - f"seq{seq_len}_"

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

1321 - # This represents the maximum number of different requests

1322 - # that will have unique loras, an therefore the max amount of memory

1323 - # consumption create dummy lora request copies from the lora request

1324 - # passed in, which contains a lora from the lora warmup path.

1325 - dummy_lora_requests: List[LoRARequest] = []

1326 - dummy_lora_requests_per_seq: List[LoRARequest] = []

1327 - if self.lora_config and is_lora_profile_run:

1328 - assert self.lora_manager is not None

1329 - with self.lora_manager.dummy_lora_cache():

1330 - for idx in range(self.lora_config.max_loras):

1331 - lora_id = idx + 1

1332 - dummy_lora_request = LoRARequest(

1333 - lora_name=f"warmup_{lora_id}",

1334 - lora_int_id=lora_id,

1335 - lora_local_path="/not/a/real/path",

1336 - )

1337 - self.lora_manager.add_dummy_lora(dummy_lora_request,

1338 - rank=LORA_WARMUP_RANK)

1339 - dummy_lora_requests.append(dummy_lora_request)

1340 - dummy_lora_requests_per_seq = [

1341 - dummy_lora_requests[idx % len(dummy_lora_requests)]

1342 - for idx in range(batch_size)

1343 - ]

1344 - self.profiler.start('internal', scenario_name)

1345 - times = 3 if use_graphs or is_pt_profiler_run else 1

1346 - if is_prompt:

1347 - seqs = [

1348 - self.create_dummy_seq_group_metadata(

1349 - i,

1350 - seq_len,

1351 - is_prompt,

1352 - lora_request=dummy_lora_requests_per_seq[i]

1353 - if dummy_lora_requests_per_seq else None)

1354 - for i in range(batch_size)

1355 - ]

1356 - else:

1357 - # FIXME: seq_len is actually number of blocks

1358 - blocks = [seq_len // batch_size for _ in range(batch_size)]

1359 - blocks[0] += seq_len % batch_size

1360 - seqs = [

1361 - self.create_dummy_seq_group_metadata(

1362 - i,

1363 - b * self.block_size - 1,

1364 - is_prompt,

1365 - lora_request=dummy_lora_requests_per_seq[i]

1366 - if dummy_lora_requests_per_seq else None)

1367 - for i, b in enumerate(blocks)

1368 - ]

1369 - torch.hpu.synchronize()

1370 - profiler = None

1371 - if is_pt_profiler_run and self.is_driver_worker:

1372 - profiler = setup_profiler()

1373 - profiler.start()

1374 - for _ in range(times):

1375 - inputs = self.prepare_model_input(seqs)

1376 - is_single_step = \

1377 - self.vllm_config.scheduler_config.num_scheduler_steps == 1

1378 - if is_prompt or is_single_step:

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

1380 - else: # decode with multi-step

1381 - inputs = dataclasses.replace(inputs,

1382 - is_first_multi_step=True,

1383 - is_last_step=False)

1384 - self.execute_model(inputs,

1385 - None,

1386 - warmup_mode=True,

1387 - num_steps=2,

1388 - seqs=seqs)

1389 - inputs = dataclasses.replace(inputs,

1390 - is_first_multi_step=False,

1391 - is_last_step=True)

1392 - self.execute_model(inputs,

1393 - None,

1394 - warmup_mode=True,

1395 - num_steps=2,

1396 - seqs=seqs)

1397 - torch.hpu.synchronize()

1398 - if profiler:

1399 - profiler.step()

1400 - if profiler:

1401 - profiler.stop()

1402 - self.profiler.end()

1403 - gc.collect()

1404 -

1405 - def remove_all_loras(self):

1406 - if not self.lora_manager:

1407 - raise RuntimeError("LoRA is not enabled.")

1408 - self.lora_manager.remove_all_adapters()

1409 -

1410 - def set_active_loras(self, lora_requests: Set[LoRARequest],

1411 - lora_mapping: LoRAMapping) -> None:

1412 - if not self.lora_manager:

1413 - raise RuntimeError("LoRA is not enabled.")

1414 - self.lora_manager.set_active_adapters(lora_requests, lora_mapping)

1415 -

1416 - def add_lora(self, lora_request: LoRARequest) -> bool:

1417 - if not self.lora_manager:

1418 - raise RuntimeError("LoRA is not enabled.")

1419 - return self.lora_manager.add_adapter(lora_request)

1420 -

1421 - def remove_lora(self, lora_id: int) -> bool:

1422 - if not self.lora_manager:

1423 - raise RuntimeError("LoRA is not enabled.")

1424 - return self.lora_manager.remove_adapter(lora_id)

1425 -

1426 - def pin_lora(self, lora_id: int) -> bool:

1427 - if not self.lora_manager:

1428 - raise RuntimeError("LoRA is not enabled.")

1429 - return self.lora_manager.pin_adapter(lora_id)

1430 -

1431 - def list_loras(self) -> Set[int]:

1432 - if not self.lora_manager:

1433 - raise RuntimeError("LoRA is not enabled.")

1434 - return self.lora_manager.list_adapters()

1435 -

1436 - def log_warmup(self, phase, i, max_i, batch_size, seq_len):

1437 - free_mem = format_bytes(

1438 - HabanaMemoryProfiler.current_free_device_memory())

1439 - dim = "num_blocks"

1440 - if phase == "Prompt":

1441 - dim = "seq_len"

1442 - msg = (f"[Warmup][{phase}][{i+1}/{max_i}] "

1443 - f"batch_size:{batch_size} "

1444 - f"{dim}:{seq_len} "

1445 - f"free_mem:{free_mem}")

1446 - logger.info(msg)

1447 -

1448 - def warmup_all_buckets(self, buckets, is_prompt, kv_caches):

1449 - for i, (batch_size, seq_len) in enumerate(reversed(buckets)):

1450 - self.log_warmup('Prompt' if is_prompt else 'Decode', i,

1451 - len(buckets), batch_size, seq_len)

1452 - self.warmup_scenario(batch_size, seq_len, is_prompt, kv_caches)

1453 -

1454 - def warmup_graphs(self,

1455 - strategy,

1456 - buckets,

1457 - is_prompt,

1458 - kv_caches,

1459 - available_mem,

1460 - starting_mem=0,

1461 - total_batch_seq=0.001):

1462 - total_mem = starting_mem

1463 - idx = 0

1464 - phase = f'Graph/{"Prompt" if is_prompt else "Decode"}'

1465 - num_candidates = len(buckets)

1466 - ordering : Union[Callable[[Any], Tuple[Any, Any]], \

1467 - Callable[[Any], Tuple[Any, Any, Any]]]

1468 - if strategy == 'min_tokens':

1469 - ordering = lambda b: (b[0] * b[1], b[1], b[0])

1470 - elif strategy == 'max_bs':

1471 - ordering = lambda b: (-b[0], b[1])

1472 - else:

1473 - raise NotImplementedError(

1474 - f'Unsupported graph allocation strategy: {strategy}')

1475 - buckets = list(sorted(buckets, key=ordering))

1476 - captured_all = True

1477 - for idx, (batch_size, seq_len) in enumerate(buckets):

1478 - # Graph memory usage is proportional to seq dimension in a batch

1479 - batch_seq = batch_size * seq_len if is_prompt else batch_size

1480 - mem_estimate = batch_seq / total_batch_seq * total_mem

1481 - if mem_estimate >= available_mem:

1482 - captured_all = False

1483 - continue

1484 - graphed_bucket = (batch_size, seq_len, is_prompt)

1485 - if graphed_bucket in self.graphed_buckets:

1486 - continue

1487 - self.graphed_buckets.add(graphed_bucket)

1488 - self.log_warmup(phase, idx, num_candidates, batch_size, seq_len)

1489 - with HabanaMemoryProfiler() as mem_prof:

1490 - self.warmup_scenario(batch_size, seq_len, is_prompt, kv_caches)

1491 - used_mem = align_workers(mem_prof.consumed_device_memory,

1492 - torch.distributed.ReduceOp.MAX)

1493 - available_mem -= used_mem

1494 - total_mem += used_mem

1495 - total_batch_seq += batch_seq

1496 -

1497 - return total_mem, total_batch_seq, captured_all

1498 -

1499 - def log_graph_warmup_summary(self, buckets, is_prompt, total_mem):

1500 - num_candidates = len(buckets)

1501 - phase = f'Graph/{"Prompt" if is_prompt else "Decode"}'

1502 - graphed = list(c[:2] for c in self.graphed_buckets

1503 - if c[2] == is_prompt)

1504 - if num_candidates == 0:

1505 - num_candidates = 1

1506 - msg = (f'{phase} captured:{len(graphed)} '

1507 - f'({100 * len(graphed) / num_candidates:.1f}%) '

1508 - f'used_mem:{format_bytes(total_mem)} '

1509 - f'buckets:{sorted(list(graphed))}')

1510 - logger.info(msg)

1511 -

1512 - @torch.inference_mode()

1513 - def warmup_model(self, kv_caches: List[torch.Tensor]) -> None:

1514 - max_blocks = kv_caches[0][0].size(0)

1515 - self.bucketing_ctx.generate_decode_buckets(max_blocks)

1516 - if profile := os.environ.get('VLLM_PT_PROFILE', None):

1517 - phase, bs, seq_len, graph = profile.split('_')

1518 - is_prompt = phase == 'prompt'

1519 - graphs = graph == 't'

1520 - if graphs:

1521 - self.graphed_buckets.add((int(bs), int(seq_len), is_prompt))

1522 - self.warmup_scenario(int(bs), int(seq_len), is_prompt, kv_caches,

1523 - True)

1524 - raise AssertionError("Finished profiling")

1525 - if not htorch.utils.internal.is_lazy() and not self.enforce_eager:

1526 - cache_size_limit = 1 + 3 * (

1527 - len(self.bucketing_ctx.prompt_buckets) +

1528 - len(self.bucketing_ctx.decode_buckets))

1529 - torch._dynamo.config.cache_size_limit = max(

1530 - cache_size_limit, torch._dynamo.config.cache_size_limit)

1531 - # Multiply by 8 to follow the original default ratio between

1532 - # the cache_size_limit and accumulated_cache_size_limit

1533 - torch._dynamo.config.accumulated_cache_size_limit = max(

1534 - cache_size_limit * 8,

1535 - torch._dynamo.config.accumulated_cache_size_limit)

1536 - if self.skip_warmup:

1537 - logger.info("Skipping warmup...")

1538 - return

1539 - self.profiler.start('internal', 'warmup')

1540 - start_mem = HabanaMemoryProfiler.current_device_memory_usage()

1541 - start_time = time.perf_counter()

1542 -

1543 - compile_only_mode_context = functools.partial(bc.env_setting,

1544 - "PT_COMPILE_ONLY_MODE",

1545 - True)

1546 - can_use_compile_only_mode = True

1547 - try:

1548 - with compile_only_mode_context():

1549 - pass

1550 - logger.debug("Using PT_COMPILE_ONLY_MODE.")

1551 - except KeyError:

1552 - can_use_compile_only_mode = False

1553 - logger.warning('Cannot use PT_COMPILE_ONLY_MODE. '

1554 - 'Warmup time will be negatively impacted. '

1555 - 'Please update Gaudi Software Suite.')

1556 - with compile_only_mode_context(

1557 - ) if can_use_compile_only_mode else contextlib.nullcontext():

1558 - self.warmup_all_buckets(self.bucketing_ctx.prompt_buckets, True,

1559 - kv_caches)

1560 - self.warmup_all_buckets(self.bucketing_ctx.decode_buckets, False,

1561 - kv_caches)

1562 -

1563 - if not self.enforce_eager and htorch.utils.internal.is_lazy():

1564 - assert self.mem_margin is not None, \

1565 - ("HabanaWorker.determine_num_available_blocks needs "

1566 - "to be called before warming up the model.")

1567 - free_mem = HabanaMemoryProfiler.current_free_device_memory()

1568 - graph_free_mem = free_mem - self.mem_margin

1569 - graph_free_mem = align_workers(graph_free_mem,

1570 - torch.distributed.ReduceOp.MIN)

1571 - prompt_graph_mem_ratio = float(

1572 - os.environ.get('VLLM_GRAPH_PROMPT_RATIO', '0.3'))

1573 - prompt_available_memory = (prompt_graph_mem_ratio *

1574 - graph_free_mem)

1575 - decode_available_memory = (graph_free_mem -

1576 - prompt_available_memory)

1577 - msg = (

1578 - f"Using {format_bytes(graph_free_mem)}"

1579 - f"/{format_bytes(free_mem)} "

1580 - "of free device memory for HPUGraphs, "

1581 - f"{format_bytes(prompt_available_memory)} for prompt and "

1582 - f"{format_bytes(decode_available_memory)} for decode "

1583 - f"(VLLM_GRAPH_PROMPT_RATIO={prompt_graph_mem_ratio})")

1584 - logger.info(msg)

1585 - prompt_strategy = os.environ.get('VLLM_GRAPH_PROMPT_STRATEGY',

1586 - 'min_tokens')

1587 - decode_strategy = os.environ.get('VLLM_GRAPH_DECODE_STRATEGY',

1588 - 'max_bs')

1589 - mem_post_prompt, prompt_batch_seq, prompt_captured_all = \

1590 - self.warmup_graphs(

1591 - prompt_strategy, self.bucketing_ctx.prompt_buckets,

1592 - True, kv_caches, prompt_available_memory)

1593 - mem_post_decode, decode_batch_seq, decode_captured_all = \

1594 - self.warmup_graphs(

1595 - decode_strategy, self.bucketing_ctx.decode_buckets,

1596 - False, kv_caches, decode_available_memory)

1597 -

1598 - # Not all prompt buckets were captured, but all decode buckets

1599 - # were captured and we have some free graph-allocated space

1600 - # left. Let's try to use it for capturing more prompt buckets.

1601 - if (mem_post_decode + mem_post_prompt < graph_free_mem

1602 - and not prompt_captured_all and decode_captured_all):

1603 - mem_post_prompt, _, prompt_captured_all = (

1604 - self.warmup_graphs(

1605 - prompt_strategy, self.bucketing_ctx.prompt_buckets,

1606 - True, kv_caches,

1607 - graph_free_mem - mem_post_prompt - mem_post_decode,

1608 - mem_post_prompt, prompt_batch_seq))

1609 -

1610 - # Not all decode buckets were captured, but all prompt buckets

1611 - # were captured and we have some free graph-allocated space

1612 - # left. Let's try to use it for capturing more decode buckets.

1613 - if mem_post_decode + mem_post_prompt < graph_free_mem \

1614 - and not decode_captured_all \

1615 - and prompt_captured_all:

1616 - mem_post_decode, _, _ = self.warmup_graphs(

1617 - decode_strategy, self.bucketing_ctx.decode_buckets,

1618 - False, kv_caches,

1619 - graph_free_mem - mem_post_prompt - mem_post_decode,

1620 - mem_post_decode, decode_batch_seq)

1621 -

1622 - self.log_graph_warmup_summary(

1623 - self.bucketing_ctx.prompt_buckets, True, mem_post_prompt)

1624 - self.log_graph_warmup_summary(

1625 - self.bucketing_ctx.decode_buckets, False, mem_post_decode)

1626 -

1627 - end_time = time.perf_counter()

1628 - end_mem = HabanaMemoryProfiler.current_device_memory_usage()

1629 - elapsed_time = end_time - start_time

1630 - msg = (

1631 - f"Warmup finished in {elapsed_time:.0f} secs, "

1632 - f"allocated {format_bytes(end_mem - start_mem)} of device memory")

1633 - logger.info(msg)

1634 - self.profiler.end()

1635 -

1636 - @property

1637 - def vocab_size(self) -> int:

1638 - return self.model_config.get_vocab_size()

1639 -

1640 - @property

1641 - def mem_margin(self) -> Optional[int]:

1642 - return self._mem_margin

1643 -

1644 - @mem_margin.setter

1645 - def mem_margin(self, value):

1646 - self._mem_margin = value

1647 -

1648 -

1649 -def _maybe_wrap_in_hpu_graph(*args, **kwargs):

1650 - return htorch.hpu.wrap_in_hpu_graph(

1651 - HpuModelAdapter(*args, **kwargs), disable_tensor_cache=True

1652 - ) if htorch.utils.internal.is_lazy() else HpuModelAdapter(*args, **kwargs)

1653 -

1654 -

1655 -class HabanaProfilerCounterHelper:

1656 -

1657 - def __init__(self):

1658 - self.niter = 0

1659 - self.average_real_throughput = None

1660 - self.logged_once = False

1661 - self.real_seq_lens = []

1662 - self.prompt_seq_lens = []

1663 -

1664 - def capture_seq_group_metadata_stats(self, seq_group_metadata_list):

1665 - self.real_seq_lens = [

1666 - len(seq_data.prompt_token_ids) + len(seq_data.output_token_ids)

1667 - for seq_group_metadata in seq_group_metadata_list

1668 - for seq_data in seq_group_metadata.seq_data.values()

1669 - ]

1670 - self.prompt_seq_lens = [

1671 - len(seq_data.prompt_token_ids)

1672 - for seq_group_metadata in seq_group_metadata_list

1673 - for seq_data in seq_group_metadata.seq_data.values()

1674 - ]

1675 -

1676 - def get_counter_dict(self, cache_config, duration, seq_len,

1677 - batch_size_padded, real_batch_size, is_prompt):

1678 - throughput = batch_size_padded / (duration / 1e6)

1679 - throughput_effective = real_batch_size / (duration / 1e6)

1680 -

1681 - real_max_seq_len = max(self.real_seq_lens)

1682 - real_num_tokens = sum(self.real_seq_lens)

1683 - padded_num_tokens = batch_size_padded * seq_len

1684 - batch_token_utilization = real_num_tokens / padded_num_tokens

1685 - if self.average_real_throughput is None:

1686 - self.average_real_throughput = throughput_effective

1687 - else: # https://www.heikohoffmann.de/htmlthesis/node134.html

1688 - self.average_real_throughput = self.average_real_throughput + 1 / (

1689 - self.niter + 1) * (throughput_effective -

1690 - self.average_real_throughput)

1691 - phase = "prompt" if is_prompt else "decode"

1692 - counters = {

1693 - f'{phase}_bucket_batch_size': batch_size_padded,

1694 - f'{phase}_batch_size': real_batch_size,

1695 - f'{phase}_bucket_seq_len': seq_len,

1696 - f'{phase}_seq_len': real_max_seq_len,

1697 - f'{phase}_bucket_gen_throughput': throughput,

1698 - f'{phase}_real_gen_throughput': throughput_effective,

1699 - f'{phase}_batch_token_utilization': batch_token_utilization,

1700 - 'average_real_throughput': self.average_real_throughput,

1701 - 'engine_iteration': self.niter,

1702 - }

1703 - self.niter += 1

1704 - if is_prompt:

1705 - prompt_bucket_in_throughput = (seq_len * batch_size_padded) / (

1706 - duration / 1e6)

1707 - prompt_real_in_throughput = sum(

1708 - self.prompt_seq_lens) / (duration / 1e6)

1709 - counters[

1710 - f'{phase}_bucket_in_throughput'] = prompt_bucket_in_throughput

1711 - counters[f'{phase}_real_in_throughput'] = prompt_real_in_throughput

1712 -

1713 - # KV cache might not be created yet (e.g. for profiling run)

1714 - if cache_config.num_gpu_blocks is not None and \

1715 - cache_config.num_gpu_blocks != 0:

1716 - cache_num_blocks_used = [

1717 - math.ceil(sl / cache_config.block_size)

1718 - for sl in self.real_seq_lens

1719 - ]

1720 - cache_total_num_blocks_used = sum(cache_num_blocks_used)

1721 - num_cache_blocks = cache_config.num_gpu_blocks

1722 - cache_total_num_free_blocks = \

1723 - num_cache_blocks - cache_total_num_blocks_used

1724 - cache_computed_utilization = \

1725 - cache_total_num_blocks_used / num_cache_blocks

1726 - max_blocks_per_seq = math.ceil(seq_len / cache_config.block_size)

1727 - batch_block_utilization = cache_total_num_blocks_used / (

1728 - batch_size_padded * max_blocks_per_seq)

1729 - counters['cache_num_blocks_used'] = cache_total_num_blocks_used

1730 - counters['cache_num_free_blocks'] = cache_total_num_free_blocks

1731 - counters['cache_computed_utilization'] = cache_computed_utilization

1732 - counters[

1733 - f'{phase}_batch_block_utilization'] = batch_block_utilization

1734 - if not self.logged_once:

1735 - counters['const_cache_num_blocks'] = cache_config.num_gpu_blocks

1736 - counters[

1737 - 'const_gpu_memory_utilization'] = \

1738 - cache_config.gpu_memory_utilization

1739 - counters['const_block_size'] = cache_config.block_size

1740 - self.logged_once = True

1741 - return counters

1742 -

1743 -

1744 -def unwrap_model(model):

1745 - if isinstance(model, torch._dynamo.eval_frame.OptimizedModule):

1746 - return unwrap_model(model._orig_mod)

1747 - else:

1748 - model = list(vars(model)['_modules'].values())[0]

1749 - modules = list(vars(model)['_modules'].values())

1750 - return modules

1751 -

1752 -

1753 -class HPUModelRunner(HPUModelRunnerBase[ModelInputForHPUWithSamplingMetadata]):

1754 - """

1755 - GPU model runner with sampling step.

1756 - """

1757 - _model_input_cls: Type[ModelInputForHPUWithSamplingMetadata] = (

1758 - ModelInputForHPUWithSamplingMetadata)

1759 -

1760 - def make_model_input_from_broadcasted_tensor_dict(

1761 - self,

1762 - tensor_dict: Dict[str, Any],

1763 - ) -> ModelInputForHPUWithSamplingMetadata:

1764 - return (

1765 - ModelInputForHPUWithSamplingMetadata.from_broadcasted_tensor_dict(

1766 - tensor_dict,

1767 - attn_backend=self.attn_backend,

1768 - ))

1769 -

1770 - @torch.inference_mode()

1771 - def prepare_model_input(

1772 - self,

1773 - seq_group_metadata_list: List[SequenceGroupMetadata],

1774 - virtual_engine: int = 0,

1775 - finished_requests_ids: Optional[List[str]] = None

1776 - ) -> ModelInputForHPUWithSamplingMetadata:

1777 - """Prepare the model input based on a given sequence group, including

1778 - metadata for the sampling step.

1779 - The API assumes seq_group_metadata_list is sorted by prefill -> decode.

1780 - The result tensors and data structure also batches input in prefill

1781 - -> decode order. For example,

1782 - - input_tokens[:num_prefill_tokens] contains prefill tokens.

1783 - - input_tokens[num_prefill_tokens:] contains decode tokens.

1784 - If cuda graph is required, this API automatically pads inputs.

1785 - """

1786 - with self.profiler.record_event('internal', 'prepare_input_tensors'):

1787 - assert seq_group_metadata_list is not None

1788 - if self.profiler.enabled:

1789 - self.profiler_counter_helper.capture_seq_group_metadata_stats(

1790 - seq_group_metadata_list=seq_group_metadata_list)

1791 - model_input, sampling_metadata = self.prepare_input_tensors(

1792 - seq_group_metadata_list)

1793 - assert model_input.attn_metadata is not None

1794 - is_prompt = model_input.attn_metadata.is_prompt

1795 -

1796 - return dataclasses.replace(model_input,

1797 - sampling_metadata=sampling_metadata,

1798 - is_prompt=is_prompt,

1799 - virtual_engine=virtual_engine)

1800 -

1801 - def finish_measurements(self):

1802 - from neural_compressor.torch.quantization import finalize_calibration

1803 - finalize_calibration(self.model.model)

1804 -

1805 - def _num_blocks(self, attn_metadata):

1806 - if attn_metadata.block_list is None:

1807 - return 0

1808 - return attn_metadata.block_list.numel()

1809 -

1810 - def _phase(self, attn_metadata):

1811 - phase_type: PhaseType

1812 - is_prompt = attn_metadata.is_prompt

1813 - is_prefix_prefill = is_prompt and attn_metadata.block_list is not None

1814 - if is_prompt and is_prefix_prefill:

1815 - phase_type = PhaseType.PREFIX_PREFILL

1816 - elif is_prompt and not is_prefix_prefill:

1817 - phase_type = PhaseType.PREFILL

1818 - elif not is_prompt:

1819 - phase_type = PhaseType.DECODE

1820 - else:

1821 - raise ValueError("Unrecognized pass type, likely due to malformed "

1822 - "attention metadata")

1823 - return phase_type

1824 -

1825 - def _check_config(self, batch_size, seq_len, attn_metadata, warmup_mode):

1826 - is_prefix_caching = self.vllm_config.cache_config.enable_prefix_caching

1827 - cfg: Optional[tuple] = None

1828 - assert cfg is None, "Configs changed between 2D and 3D"

1829 - if is_prefix_caching:

1830 - phase = self._phase(attn_metadata)

1831 - num_blocks = self._num_blocks(attn_metadata)

1832 - cfg = (batch_size, seq_len, num_blocks, phase)

1833 - else:

1834 - phase = 'prompt' if attn_metadata.is_prompt else 'decode'

1835 - cfg = (batch_size, seq_len, phase)

1836 - seen = cfg in self.seen_configs

1837 - self.seen_configs.add(cfg)

1838 - if not seen and not warmup_mode:

1839 - logger.warning("Configuration: %s was not warmed-up!",

1840 - (phase.value, batch_size, seq_len,

1841 - num_blocks) if is_prefix_caching else

1842 - (phase, batch_size, seq_len))

1843 -

1844 - def create_lora_mask(self, input_tokens: torch.Tensor, lora_ids: List[int],

1845 - is_prompt: bool):

1846 - '''

1847 - This is a helper function to create the mask for lora computations.

1848 - Lora Mask is needed to ensure we match the correct lora weights for the

1849 - for the request.

1850 - For Prompt phase we have

1851 - lora_mask with shape (batch_size * seq_len, max_loras * max_rank)

1852 - lora_logits_mask with shape (batch_size, max_loras * max_rank)

1853 - For Decode phase we have both

1854 - lora_mask and lora_logits_mask with shape

1855 - (batch_size, max_loras * max_rank)

1856 - '''

1857 - lora_mask: torch.Tensor = None

1858 - lora_logits_mask: torch.Tensor = None

1859 - lora_index = 0

1860 -

1861 - if self.lora_config:

1862 - if is_prompt:

1863 - lora_mask = torch.zeros(

1864 - input_tokens.shape[0] * input_tokens.shape[1],

1865 - (self.lora_config.max_loras) *\

1866 - self.lora_config.max_lora_rank,

1867 - dtype=self.lora_config.lora_dtype)

1868 - lora_logits_mask = torch.zeros(

1869 - input_tokens.shape[0], (self.lora_config.max_loras) *

1870 - self.lora_config.max_lora_rank,

1871 - dtype=self.lora_config.lora_dtype)

1872 -

1873 - ones = torch.ones(input_tokens.shape[1],

1874 - self.lora_config.max_lora_rank,

1875 - dtype=self.lora_config.lora_dtype)

1876 - logit_ones = torch.ones(1,

1877 - self.lora_config.max_lora_rank,

1878 - dtype=self.lora_config.lora_dtype)

1879 -

1880 - for i in range(len(lora_ids)):

1881 - if lora_ids[i] == 0:

1882 - continue

1883 - lora_index = self.lora_manager._adapter_manager.\

1884 - lora_index_to_id.index(lora_ids[i])

1885 - start_row = i * input_tokens.shape[1]

1886 - end_row = start_row + input_tokens.shape[1]

1887 - start_col = lora_index * self.lora_config.max_lora_rank

1888 - end_col = start_col + self.lora_config.max_lora_rank

1889 - lora_mask[start_row:end_row, start_col:end_col] = ones

1890 - lora_logits_mask[i, start_col:end_col] = logit_ones

1891 - lora_mask = lora_mask.to('hpu')

1892 - lora_logits_mask = lora_logits_mask.to('hpu')

1893 - else:

1894 - lora_mask = torch.zeros(input_tokens.shape[0],

1895 - (self.lora_config.max_loras) *

1896 - self.lora_config.max_lora_rank,

1897 - dtype=self.lora_config.lora_dtype)

1898 - ones = torch.ones(1,

1899 - self.lora_config.max_lora_rank,

1900 - dtype=self.lora_config.lora_dtype)

1901 - for i in range(len(lora_ids)):

1902 - if lora_ids[i] == 0:

1903 - continue

1904 - lora_index = self.lora_manager._adapter_manager.\

1905 - lora_index_to_id.index(lora_ids[i])

1906 - start_pos = lora_index * self.lora_config.max_lora_rank

1907 - end_pos = start_pos + self.lora_config.max_lora_rank

1908 - lora_mask[i, start_pos:end_pos] = ones

1909 - lora_mask = lora_mask.to('hpu')

1910 - lora_logits_mask = lora_mask

1911 -

1912 - return lora_mask, lora_logits_mask

1913 -

1914 - def _get_seq_ids(self, model_input):

1915 - return ([

1916 - sg.seq_ids[0] for sg in model_input.sampling_metadata.seq_groups

1917 - ])

1918 -

1919 - def _pad_to_max_num_seqs(self, tensor, value):

1920 - padding_needed = self.max_num_seqs - tensor.size(0)

1921 - if padding_needed:

1922 - padding = torch.full((padding_needed, *tensor.shape[1:]),

1923 - value,

1924 - device=tensor.device,

1925 - dtype=tensor.dtype)

1926 - tensor = torch.cat([tensor, padding])

1927 - return tensor

1928 -

1929 - @torch.inference_mode()

1930 - def execute_model(

1931 - self,

1932 - model_input: ModelInputForHPUWithSamplingMetadata,

1933 - kv_caches: List[torch.Tensor],

1934 - intermediate_tensors: Optional[IntermediateTensors] = None,

1935 - num_steps: int = 1,

1936 - warmup_mode=False,

1937 - seqs=None,

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

1939 - VLLM_DELAYED_SAMPLING = envs.VLLM_HPU_USE_DELAYED_SAMPLING

1940 - use_delayed_sampling = VLLM_DELAYED_SAMPLING and not warmup_mode

1941 - assert not (use_delayed_sampling and num_steps != 1), \

1942 - 'Delayed sampling is not compatible with MSS!'

1943 - assert model_input.input_tokens is not None

1944 - if use_delayed_sampling and not model_input.is_prompt and \

1945 - self.is_driver_worker:

1946 - num_cached = len(self.cached_step_outputs)

1947 - assert num_cached > 0

1948 - cur_seq_ids = self._get_seq_ids(model_input)

1949 - cur_seq_id_pos = {

1950 - sid: idx

1951 - for idx, sid in enumerate(cur_seq_ids) if sid >= 0

1952 - }

1953 - htorch.core.mark_step()

1954 - for i in range(num_cached):

1955 - prev_seq_ids = self._get_seq_ids(self.cached_step_inputs[i])

1956 - target_indices = [

1957 - cur_seq_id_pos.get(psi, -1) for psi in prev_seq_ids

1958 - ]

1959 - padding = self.cached_step_outputs[i].size(0) - len(

1960 - target_indices)

1961 - target_indices.extend([-1] * padding)

1962 - target_indices = torch.tensor(

1963 - target_indices,

1964 - device=model_input.input_tokens.device,

1965 - dtype=model_input.input_tokens.dtype)

1966 - model_input.input_tokens.index_copy_(

1967 - 0, target_indices, self.cached_step_outputs[i])

1968 - htorch.core.mark_step()

1969 -

1970 - if not model_input.is_first_multi_step:

1971 - if not model_input.is_last_step:

1972 - # not first or last multi-step

1973 - return []

1974 - # last multi-step

1975 - output = self._decode_sampler_outputs(

1976 - model_input) if self.is_driver_worker else []

1977 - torch.hpu.synchronize()

1978 - if model_input.is_first_multi_step:

1979 - # first multi-step

1980 - if self.lora_config:

1981 - assert model_input.lora_requests is not None

1982 - assert model_input.lora_mapping is not None

1983 - self.set_active_loras(model_input.lora_requests,

1984 - model_input.lora_mapping)

1985 - # Rank!=0 workers has is_prompt==None

1986 - if use_delayed_sampling and not model_input.is_prompt and \

1987 - model_input.input_tokens.size(1) == 1:

1988 - if self.is_driver_worker:

1989 - model_kwargs_broadcast_data = {

1990 - "input_tokens": model_input.input_tokens

1991 - }

1992 - broadcast_tensor_dict(model_kwargs_broadcast_data, src=0)

1993 - input_tokens = model_input.input_tokens

1994 -

1995 - else:

1996 - model_kwargs_broadcast_data = broadcast_tensor_dict(src=0)

1997 - input_tokens = model_kwargs_broadcast_data["input_tokens"]

1998 - else:

1999 - input_tokens = model_input.input_tokens

2000 - input_positions = model_input.input_positions

2001 - attn_metadata = model_input.attn_metadata

2002 - sampling_metadata = model_input.sampling_metadata

2003 - real_batch_size = model_input.real_batch_size

2004 - batch_size_padded = model_input.batch_size_padded

2005 - assert input_tokens is not None

2006 - assert input_positions is not None

2007 - assert sampling_metadata is not None

2008 - assert attn_metadata is not None

2009 - is_prompt = attn_metadata.is_prompt

2010 - assert is_prompt is not None

2011 - batch_size = input_tokens.size(0)

2012 - seq_len = self._seq_len(attn_metadata)

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

2014 - self._check_config(batch_size, seq_len, attn_metadata, warmup_mode)

2015 -

2016 - lora_mask: torch.Tensor = None

2017 - lora_logits_mask: torch.Tensor = None

2018 - if self.lora_config:

2019 - assert model_input.lora_ids is not None

2020 - lora_mask, lora_logits_mask = self.create_lora_mask(

2021 - input_tokens, model_input.lora_ids,

2022 - attn_metadata.is_prompt)

2023 -

2024 - execute_model_kwargs = {

2025 - "input_ids": input_tokens,

2026 - "positions": input_positions,

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

2028 - "intermediate_tensors": intermediate_tensors,

2029 - "lora_mask": lora_mask,

2030 - "virtual_engine": model_input.virtual_engine,

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

2032 - }

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

2034 - execute_model_kwargs.update(

2035 - {"bypass_hpu_graphs": not use_graphs})

2036 -

2037 - htorch.core.mark_step()

2038 - if self.is_driver_worker:

2039 - model_event_name = ("model_"

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

2041 - f"bs{batch_size}_"

2042 - f"seq{seq_len}_"

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

2044 - else:

2045 - model_event_name = 'model_executable'

2046 - if num_steps > 1 or use_delayed_sampling:

2047 - # in case of multi-step scheduling

2048 - # we only want to pythonize in the last step

2049 - sampling_metadata.skip_sampler_cpu_output = True

2050 - self.model.sampler.include_gpu_probs_tensor = True

2051 - cache_orig_output_tokens_len: List[Dict] = []

2052 -

2053 - def try_revert_dummy_output_tokens():

2054 - if len(cache_orig_output_tokens_len) > 0:

2055 - # Reuse the original output token ids length

2056 - for i, seq_group_metadata in enumerate(

2057 - seq_group_metadata_list):

2058 - for j, data in seq_group_metadata.seq_data.items():

2059 - orig_output_tokens_len = \

2060 - cache_orig_output_tokens_len[i][j]

2061 - data.output_token_ids = \

2062 - data.output_token_ids[:orig_output_tokens_len]

2063 -

2064 - for i in range(num_steps):

2065 - if i != 0 and not self.is_driver_worker:

2066 - broadcast_data = broadcast_tensor_dict(src=0)

2067 - if 'early_exit' in broadcast_data and broadcast_data[

2068 - 'early_exit']:

2069 - return [output] if num_steps == 1 else []

2070 - execute_model_kwargs.update({

2071 - "input_ids":

2072 - broadcast_data["input_ids"],

2073 - "positions":

2074 - broadcast_data["positions"],

2075 - "attn_metadata":

2076 - self.trim_attn_metadata(

2077 - broadcast_data["attn_metadata"])

2078 - })

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

2080 - hidden_states = self.model.forward(

2081 - **execute_model_kwargs,

2082 - selected_token_indices=sampling_metadata.

2083 - selected_token_indices)

2084 -

2085 - if self.lora_config:

2086 - LoraMask.setLoraMask(

2087 - lora_logits_mask.index_select(

2088 - 0, sampling_metadata.selected_token_indices))

2089 -

2090 - # Compute the logits.

2091 - with self.profiler.record_event(

2092 - 'internal',

2093 - ('compute_logits_'

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

2095 - f'{batch_size}_'

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

2097 - if num_steps == 1:

2098 - sampling_metadata.selected_token_indices = None

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

2100 - sampling_metadata)

2101 - htorch.core.mark_step()

2102 - # Only perform sampling in the driver worker.

2103 - if not self.is_driver_worker:

2104 - continue

2105 -

2106 - if use_delayed_sampling:

2107 - fake_output = self._delayed_sampler_outputs(model_input)

2108 -

2109 - with self.profiler.record_event(

2110 - 'internal', ('sample_'

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

2112 - f'bs{batch_size}_'

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

2114 - output = self.model.sample(

2115 - logits=logits,

2116 - sampling_metadata=sampling_metadata,

2117 - )

2118 - if num_steps > 1:

2119 - output = output.sampled_token_ids

2120 - self.cached_step_outputs.append(output)

2121 - if use_delayed_sampling and self.is_driver_worker:

2122 - self._patch_prev_output()

2123 - output = self._pad_to_max_num_seqs(

2124 - output.sampled_token_ids, DUMMY_TOKEN_ID)

2125 - self.cached_step_outputs.append(output)

2126 - self.cached_step_inputs.append(model_input)

2127 - htorch.core.mark_step()

2128 - if model_input.async_callback is not None:

2129 - model_input.async_callback()

2130 - if i < num_steps - 1:

2131 - if i == 0:

2132 - if model_input.async_callback is not None:

2133 - ctx = model_input.async_callback.keywords[ # type: ignore

2134 - "ctx"]

2135 - seq_group_metadata_list = \

2136 - ctx.seq_group_metadata_list

2137 - elif seqs is not None:

2138 - seq_group_metadata_list = seqs

2139 - else:

2140 - raise RuntimeError(

2141 - "seq_group_metadata_list is uninitialized")

2142 - for i, seq_group_metadata in enumerate(

2143 - seq_group_metadata_list):

2144 - # Skip empty steps

2145 - seq_group_metadata.state.current_step += (

2146 - num_steps - 2)

2147 - # Cache the original output token ids

2148 - cache_orig_output_tokens_len.append({})

2149 - for j, data in seq_group_metadata.seq_data.items():

2150 - cache_orig_output_tokens_len[i][j] = \

2151 - len(data.output_token_ids)

2152 - for seq_group_metadata in seq_group_metadata_list:

2153 - for data in seq_group_metadata.seq_data.values():

2154 - max_output_len = sampling_metadata.seq_groups[

2155 - 0].sampling_params.max_tokens

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

2157 - # add a place holder for prepare_decode

2158 - # arbitrary value, this could be any token

2159 - dummy_token = (540, )

2160 - data.output_token_ids += (dummy_token)

2161 - else:

2162 - broadcast_tensor_dict({'early_exit': True},

2163 - src=0)

2164 - if num_steps == 1:

2165 - return [output]

2166 - else:

2167 - try_revert_dummy_output_tokens()

2168 - return []

2169 -

2170 - result = self._prepare_decode(seq_group_metadata_list,

2171 - output=output)

2172 - execute_model_kwargs.update({

2173 - "input_ids":

2174 - result.input_tokens,

2175 - "positions":

2176 - result.input_positions,

2177 - "attn_metadata":

2178 - self.trim_attn_metadata(result.attn_metadata)

2179 - })

2180 - model_kwargs_broadcast_data = {

2181 - "input_ids": result.input_tokens,

2182 - "positions": result.input_positions,

2183 - "attn_metadata": vars(result.attn_metadata)

2184 - }

2185 - broadcast_tensor_dict(model_kwargs_broadcast_data, src=0)

2186 - else:

2187 - try_revert_dummy_output_tokens()

2188 -

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

2190 - # Stop recording 'execute_model' event

2191 - self.profiler.end()

2192 - event_end = self.profiler.get_timestamp_us()

2193 - counters = self.profiler_counter_helper.get_counter_dict(

2194 - cache_config=self.cache_config,

2195 - duration=event_end - self.event_start,

2196 - seq_len=seq_len,

2197 - batch_size_padded=batch_size_padded,

2198 - real_batch_size=real_batch_size,

2199 - is_prompt=is_prompt)

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

2201 - if num_steps == 1:

2202 - if self.return_hidden_states:

2203 - # we only need to pass hidden states of most recent token

2204 - assert model_input.sampling_metadata is not None

2205 - if model_input.is_prompt:

2206 - output.prefill_hidden_states = hidden_states

2207 - output.hidden_states = hidden_states

2208 - if use_delayed_sampling:

2209 - if self.is_driver_worker:

2210 - return [fake_output]

2211 - else:

2212 - return []

2213 -

2214 - return [output] if self.is_driver_worker else []

2215 - else:

2216 - return []

2217 - return output if type(output) is list else [output]

2218 -

2219 - def _delayed_sampler_outputs(self, model_input):

2220 - next_token_ids = [[DUMMY_TOKEN_ID]] * len(

2221 - model_input.sampling_metadata.seq_groups)

2222 - sampler_output = self._make_decode_output(

2223 - next_token_ids, model_input.sampling_metadata.seq_groups)

2224 - return sampler_output

2225 -

2226 - def _decode_sampler_outputs(self, model_input):

2227 - use_async_out_proc = model_input.async_callback is not None

2228 - sampler_outputs = []

2229 - num_outputs = len(self.cached_step_outputs)

2230 - for i in range(num_outputs):

2231 - next_token_ids = self.cached_step_outputs.pop(0)

2232 - next_token_ids = next_token_ids.cpu().tolist()

2233 - sampler_output = self._make_decode_output(

2234 - next_token_ids, model_input.sampling_metadata.seq_groups)

2235 - sampler_outputs.append(sampler_output)

2236 -

2237 - if i < num_outputs - 1 and use_async_out_proc:

2238 - assert model_input.async_callback is not None

2239 - ctx = model_input.async_callback.keywords[ # type: ignore

2240 - "ctx"]

2241 - ctx.append_output(

2242 - outputs=[sampler_output],

2243 - seq_group_metadata_list=ctx.seq_group_metadata_list,

2244 - scheduler_outputs=ctx.scheduler_outputs,

2245 - is_async=False,

2246 - is_last_step=False,

2247 - is_first_step_output=False)

2248 - model_input.async_callback()

2249 -

2250 - if use_async_out_proc:

2251 - return [sampler_outputs[-1]]

2252 - else:

2253 - return sampler_outputs

2254 -

2255 - def _make_decode_output(

2256 - self,

2257 - next_token_ids: List[List[int]],

2258 - seq_groups: List[SequenceGroupToSample],

2259 - ) -> SamplerOutput:

2260 - zero_logprob = Logprob(0.0)

2261 - sampler_outputs = []

2262 - batch_idx = 0

2263 - for seq_group in seq_groups:

2264 - seq_ids = seq_group.seq_ids

2265 - seq_outputs = []

2266 - for seq_id in seq_ids:

2267 - next_token_id = next_token_ids[batch_idx][0]

2268 - seq_outputs.append(

2269 - SequenceOutput(seq_id, next_token_id,

2270 - {next_token_id: zero_logprob}))

2271 - batch_idx += 1

2272 - sampler_outputs.append(

2273 - CompletionSequenceGroupOutput(seq_outputs, None))

2274 - return SamplerOutput(sampler_outputs)

2275 -

2276 - def shutdown_inc(self):

2277 - can_finalize_inc = False

2278 - from contextlib import suppress

2279 - with suppress(AttributeError):

2280 - can_finalize_inc = (self.model_config.quantization == 'inc') and \

2281 - (self.model.model is not None) and \

2282 - self.inc_initialized_successfully and \

2283 - not getattr(self, "_is_inc_finalized", False)

2284 - if can_finalize_inc:

2285 - from neural_compressor.torch.quantization import (

2286 - finalize_calibration)

2287 - finalize_calibration(self.model.model)

2288 - self._is_inc_finalized = True

2289 -

2290 - def __del__(self):

2291 - self.shutdown_inc()

2292 -

2293 - def _patch_prev_output(self):

2294 - assert len(self.cached_step_inputs) == len(self.cached_step_outputs), \

2295 - f'''Inputs and outputs are out of sync!

2296 - {len(self.cached_step_inputs)} vs {len(self.cached_step_outputs)}'''

2297 - if len(self.cached_step_inputs) == 0:

2298 - return

2299 - model_input = self.cached_step_inputs.pop(0)

2300 - delayed_output = self.cached_step_outputs.pop(0).cpu().squeeze(

2301 - -1).tolist()

2302 - ctx = model_input.async_callback.keywords["ctx"] # type: ignore

2303 - # If there's no output to patch with, which is usually the case when

2304 - # we're starting a new request after all requests are completed.

2305 - if len(ctx.output_queue) == 0:

2306 - return

2307 - assert len(

2308 - ctx.output_queue) == 1, 'There should be exactly 1 output waiting!'

2309 - output_data = ctx.output_queue[0]

2310 - assert len(output_data.outputs) == 1

2311 - for fake_out, real_out in zip(output_data.outputs[0], delayed_output):

2312 - fake_out.samples[0].output_token = real_out

2313 - for sg, real_out in zip(output_data.seq_group_metadata_list,

2314 - delayed_output):

2315 - assert len(sg.seq_data) == 1

2316 - seq_data = list(sg.seq_data.values())[0]

2317 - # This is a hack. Assigning output_token_ids triggers

2318 - # a cache recomputation and we only need to update the last token

2319 - seq_data.output_token_ids_array[-1] = real_out

2320 - seq_data._cached_all_token_ids[-1] = real_out

vllm/worker/hpu_worker.py deleted

-485

@@ -1,485 +0,0 @@

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

2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 -

4 -###############################################################################

6 -###############################################################################

7 -

8 -import contextlib

9 -import gc

10 -import os

11 -from typing import List, Optional, Set, Tuple, Type

12 -

13 -import habana_frameworks.torch as htorch # noqa:F401

14 -import torch

15 -import torch.distributed

16 -from vllm_hpu_extension.profiler import HabanaMemoryProfiler, format_bytes

17 -

18 -import vllm.envs as envs

19 -from vllm.config import ParallelConfig, VllmConfig

20 -from vllm.distributed import (ensure_model_parallel_initialized,

21 - init_distributed_environment)

22 -from vllm.logger import init_logger

23 -from vllm.lora.request import LoRARequest

24 -from vllm.model_executor import set_random_seed

25 -from vllm.model_executor.layers.sampler import SamplerOutput

26 -from vllm.platforms import current_platform

27 -from vllm.prompt_adapter.request import PromptAdapterRequest

28 -from vllm.sequence import ExecuteModelRequest

29 -from vllm.utils import bind_kv_cache

30 -from vllm.worker.cache_engine import CacheEngine

31 -from vllm.worker.hpu_model_runner import HPUModelRunner

32 -from vllm.worker.model_runner_base import ModelRunnerBase

33 -from vllm.worker.worker_base import (LocalOrDistributedWorkerBase, WorkerBase,

34 - WorkerInput)

35 -

36 -logger = init_logger(__name__)

37 -

38 -

39 -class HPUWorker(LocalOrDistributedWorkerBase):

40 - """A worker class that executes (a partition of) the model on a HPU.

41 -

42 - Each worker is associated with a single HPU. The worker is responsible for

43 - maintaining the KV cache and executing the model on the HPU. In case of

44 - distributed inference, each worker is assigned a partition of the model.

45 - """

46 -

47 - def __init__(

48 - self,

49 - vllm_config: VllmConfig,

50 - local_rank: int,

51 - rank: int,

52 - distributed_init_method: str,

53 - is_driver_worker: bool = False,

54 - model_runner_cls: Optional[Type[ModelRunnerBase]] = None,

55 - ) -> None:

56 - WorkerBase.__init__(self, vllm_config=vllm_config)

57 - self.parallel_config.rank = rank

58 - self.local_rank = local_rank

59 - self.rank = rank

60 - self.distributed_init_method = distributed_init_method

61 - self.is_driver_worker = is_driver_worker

62 - if self.is_driver_worker:

63 - assert self.rank == 0, "The driver worker must have rank 0."

64 -

65 - if self.model_config.trust_remote_code:

66 - # note: lazy import to avoid importing torch before initializing

67 - from vllm.utils import init_cached_hf_modules

68 - init_cached_hf_modules()

69 -

70 - self.model_runner: HPUModelRunner = HPUModelRunner(

71 - vllm_config=vllm_config, is_driver_worker=is_driver_worker)

72 - # Uninitialized cache engine. Will be initialized by

73 - # initialize_cache.

74 - self.cache_engine: List[HPUCacheEngine]

75 - # Initialize gpu_cache as pooling models don't initialize kv_caches

76 - self.hpu_cache: Optional[List[List[torch.Tensor]]] = None

77 - # Torch profiler. Enabled and configured through env vars:

78 - # VLLM_TORCH_PROFILER_DIR=/path/to/save/trace

79 - if envs.VLLM_TORCH_PROFILER_DIR:

80 - torch_profiler_trace_dir = envs.VLLM_TORCH_PROFILER_DIR

81 - logger.info("Profiling enabled. Traces will be saved to: %s",

82 - torch_profiler_trace_dir)

83 - self.profiler = torch.profiler.profile(

84 - activities=[

85 - torch.profiler.ProfilerActivity.CPU,

86 - torch.profiler.ProfilerActivity.HPU,

87 - ],

88 - with_stack=True,

89 - on_trace_ready=torch.profiler.tensorboard_trace_handler(

90 - torch_profiler_trace_dir, use_gzip=True))

91 - else:

92 - self.profiler = None

93 -

94 - def start_profile(self):

95 - if self.profiler is None:

96 - raise RuntimeError("Profiler is not enabled.")

97 - self.profiler.start()

98 -

99 - def stop_profile(self):

100 - if self.profiler is None:

101 - raise RuntimeError("Profiler is not enabled.")

102 - self.profiler.stop()

103 -

104 - def _set_env_vars(self):

105 - local_rank = self.local_rank

106 - if self.parallel_config.world_size == 1:

107 - local_rank = -1

108 - import os

109 - os.environ["LOCAL_RANK"] = str(local_rank)

110 - os.environ["ID"] = str(local_rank)

111 - os.environ["WORLD_SIZE"] = str(self.parallel_config.world_size)

112 - os.environ["RANK"] = str(self.rank)

113 -

114 - def init_device(self) -> None:

115 - if self.device_config.device.type == "hpu":

116 - self.device = torch.device("hpu")

117 - torch.hpu.set_device(self.device)

118 - else:

119 - raise RuntimeError(

120 - f"Not support device type: {self.device_config.device}")

121 - # Initialize the distributed environment.

122 - if self.model_config.quantization == 'inc':

123 - self._set_env_vars()

124 - init_worker_distributed_environment(self.parallel_config, self.rank,

125 - self.distributed_init_method,

126 - self.local_rank)

127 - # Set random seed.

128 - set_random_seed(self.model_config.seed)

129 -

130 - def load_model(self):

131 - self.model_runner.load_model()

132 -

133 - def execute_model(

134 - self,

135 - execute_model_req: Optional[ExecuteModelRequest] = None,

136 - ) -> Optional[List[SamplerOutput]]:

137 - # VLLM_HPU_LOG_STEP_GRAPH_COMPILATION - will log graph compilations per engine step, only when there was any - highly recommended to use alongside PT_HPU_METRICS_GC_DETAILS! # noqa:E501

138 - # VLLM_HPU_LOG_STEP_GRAPH_COMPILATION_ALL - will log graph compilations per engine step, always, even if there were none # noqa:E501

139 - # VLLM_HPU_LOG_STEP_CPU_FALLBACKS - will log cpu fallbacks per engine step, only when there was any # noqa:E501

140 - # VLLM_HPU_LOG_STEP_CPU_FALLBACKS_ALL - will log cpu fallbacks per engine step, always, even if there were none # noqa:E501

141 - log_graph_compilation_all = os.environ.get(

142 - 'VLLM_HPU_LOG_STEP_GRAPH_COMPILATION_ALL', '0') != '0'

143 - log_graph_compilation = os.environ.get(

144 - 'VLLM_HPU_LOG_STEP_GRAPH_COMPILATION',

145 - '0') != '0' or log_graph_compilation_all

146 - log_cpu_fallbacks_all = os.environ.get(

147 - 'VLLM_HPU_LOG_STEP_CPU_FALLBACKS_ALL', '0') != '0'

148 - log_cpu_fallbacks = os.environ.get('VLLM_HPU_LOG_STEP_CPU_FALLBACKS',

149 - '0') != '0' or log_cpu_fallbacks_all

150 - if (log_graph_compilation or log_cpu_fallbacks) and \

151 - execute_model_req is not None:

152 - from habana_frameworks.torch.hpu.metrics import metric_localcontext

153 - seq_group_metadata_list = execute_model_req.seq_group_metadata_list

154 - is_prompt = any([

155 - seq_group_metadata.is_prompt

156 - for seq_group_metadata in seq_group_metadata_list

157 - ])

158 - max_context_len = max([

159 - max([

160 - len(v.prompt_token_ids) + len(v.output_token_ids)

161 - for v in seq_group_metadata.seq_data.values()

162 - ]) for seq_group_metadata in seq_group_metadata_list

163 - ]) # whoa, that's some spicy stuff right here

164 - max_num_blocks = (

165 - (max_context_len - 1) // self.cache_config.block_size) + 1

166 - input_stats = (f'is_prompt: {is_prompt}, '

167 - f'num_seqs: {len(seq_group_metadata_list)}, '

168 - f'max_context_len: {max_context_len}, '

169 - f'max_num_blocks {max_num_blocks}')

170 - gc_ctx = metric_localcontext(

171 - "graph_compilation"

172 - ) if log_graph_compilation else contextlib.nullcontext()

173 - cpu_fallback_ctx = metric_localcontext(

174 - "cpu_fallback"

175 - ) if log_cpu_fallbacks else contextlib.nullcontext()

176 - with gc_ctx as gc_local_metric, \

177 - cpu_fallback_ctx as cpu_fallback_local_metric:

178 - output = LocalOrDistributedWorkerBase.execute_model(

179 - self, execute_model_req)

180 - if (log_graph_compilation and gc_local_metric.stats()[0][1]

181 - > 0) or log_graph_compilation_all:

182 - msg = ("VLLM_HPU_STEP_GRAPH_COMPILATION: "

183 - f"{gc_local_metric.stats()}, {input_stats}")

184 - logger.warning(msg)

185 - if (log_cpu_fallbacks and cpu_fallback_local_metric.stats()[0][1]

186 - > 0) or log_cpu_fallbacks_all:

187 - msg = ("VLLM_HPU_STEP_CPU_FALLBACK: "

188 - f"{cpu_fallback_local_metric.stats()}, {input_stats}")

189 - logger.warning(msg)

190 -

191 - return output

192 -

193 - output = LocalOrDistributedWorkerBase.execute_model(

194 - self, execute_model_req)

195 - return output

196 -

197 - @torch.inference_mode()

198 - def determine_num_available_blocks(self) -> Tuple[int, int]:

199 - """Profiles the peak memory usage of the model to determine how many

200 - KV blocks may be allocated without OOMs.

201 -

202 - The engine will first conduct a profiling of the existing memory usage.

203 - Then, it calculate the maximum possible number of GPU and CPU blocks

204 - that can be allocated with the remaining free memory.

205 -

206 - Tip:

207 - You may limit the usage of GPU memory

208 - by adjusting the `gpu_memory_utilization` parameter.

209 - """

210 - # Profile the memory usage of the model and get the maximum number of

211 - # cache blocks that can be allocated with the remaining free memory.

212 -

213 - # Execute a forward pass with dummy inputs to profile the memory usage

214 - # of the model.

215 - with HabanaMemoryProfiler() as m:

216 - self.model_runner.profile_run()

217 - torch.hpu.synchronize()

218 - msg = ("Model profiling run "

219 - f"took {m.get_summary_string()}")

220 - logger.info(msg)

221 - # At this point we should've allocated the maximum workspace for all

222 - # recipes we will use the extra memory for graphs/blocks

223 - free_hpu_memory = torch.hpu.mem_get_info()[0]

224 -

225 - cache_block_size = self.get_cache_block_size_bytes()

226 - graph_reserved_mem = (float(

227 - os.environ.get('VLLM_GRAPH_RESERVED_MEM', '0.1'))

228 - if not self.model_config.enforce_eager else 0)

229 - graph_headroom = 1 - graph_reserved_mem

230 - available_hpu_memory = free_hpu_memory * \

231 - self.cache_config.gpu_memory_utilization

232 - hpu_memory_margin = free_hpu_memory * (

233 - 1 - self.cache_config.gpu_memory_utilization)

234 - self.model_runner.mem_margin = hpu_memory_margin

235 - cache_size_bytes = available_hpu_memory * graph_headroom

236 - graph_headroom_bytes = available_hpu_memory * (1 - graph_headroom)

237 - msg = (

238 - f"Free device memory: {format_bytes(free_hpu_memory)}, "

239 - f"{format_bytes(available_hpu_memory)} usable "

240 - f"(gpu_memory_utilization={self.cache_config.gpu_memory_utilization}),"

241 - f" {format_bytes(graph_headroom_bytes)} reserved for HPUGraphs "

242 - f"(VLLM_GRAPH_RESERVED_MEM={graph_reserved_mem}), "

243 - f"{format_bytes(cache_size_bytes)} reserved for KV cache")

244 - logger.info(msg)

245 - num_hpu_blocks = int(cache_size_bytes // cache_block_size)

246 - num_cpu_blocks = int(self.cache_config.swap_space_bytes //

247 - cache_block_size)

248 - num_hpu_blocks = max(num_hpu_blocks, 0)

249 - num_cpu_blocks = max(num_cpu_blocks, 0)

250 - self.model_runner.bucketing_ctx.num_hpu_blocks = num_hpu_blocks

251 -

252 - if self.model_runner.lora_manager:

253 - self.model_runner.remove_all_loras()

254 -

255 - gc.collect()

256 - return num_hpu_blocks, num_cpu_blocks

257 -

258 - def initialize_cache(self, num_gpu_blocks: int,

259 - num_cpu_blocks: int) -> None:

260 - """Allocate GPU and CPU KV cache with the specified number of blocks.

261 -

262 - This also warms up the model, which may record CUDA graphs.

263 - """

264 - raise_if_cache_size_invalid(

265 - num_gpu_blocks, self.cache_config.block_size,

266 - self.model_config.max_model_len,

267 - self.parallel_config.pipeline_parallel_size)

268 -

269 - self.cache_config.num_gpu_blocks = num_gpu_blocks

270 - self.cache_config.num_cpu_blocks = num_cpu_blocks

271 -

272 - with HabanaMemoryProfiler() as m:

273 - self._init_cache_engine()

274 - torch.hpu.synchronize()

275 - msg = ("Initializing cache engine "

276 - f"took {m.get_summary_string()}")

277 - logger.info(msg)

278 - self._warm_up_model()

279 -

280 - def _init_cache_engine(self):

281 - assert self.cache_config.num_gpu_blocks is not None

282 - self.cache_engine = [

283 - HPUCacheEngine(self.cache_config, self.model_config,

284 - self.parallel_config, self.device_config)

285 - for _ in range(self.parallel_config.pipeline_parallel_size)

286 - ]

287 - self.hpu_cache = [

288 - self.cache_engine[ve].gpu_cache

289 - for ve in range(self.parallel_config.pipeline_parallel_size)

290 - ]

291 - bind_kv_cache(self.compilation_config.static_forward_context,

292 - self.hpu_cache)

293 -

294 - def _warm_up_model(self) -> None:

295 - # NOTE(kzawora): We should use virtual engine index here

296 - # for pipeline parallelism. Using 0 for now.

297 - assert self.hpu_cache is not None

298 - self.model_runner.warmup_model(self.hpu_cache[0])

299 - # Reset the seed to ensure that the random state is not affected by

300 - # the model initialization and profiling.

301 - set_random_seed(self.model_config.seed)

302 -

303 - def finish_measurements(self):

304 - self.model_runner.finish_measurements()

305 -

306 - @property

307 - def do_metadata_broadcast(self) -> bool:

308 - return self.parallel_config.tensor_parallel_size > 1

309 -

310 - @property

311 - def kv_cache(self) -> Optional[List[List[torch.Tensor]]]:

312 - return self.hpu_cache

313 -

314 - @torch.inference_mode()

315 - def prepare_worker_input(

316 - self, execute_model_req: ExecuteModelRequest) -> WorkerInput:

317 - virtual_engine = execute_model_req.virtual_engine

318 - num_seq_groups = len(execute_model_req.seq_group_metadata_list)

319 - # `blocks_to_swap_in` and `blocks_to_swap_out` are cpu tensors.

320 - # they contain parameters to launch cudamemcpyasync.

321 - blocks_to_swap_in = torch.tensor(execute_model_req.blocks_to_swap_in,

322 - device="cpu",

323 - dtype=torch.int64).view(-1, 2)

324 - blocks_to_swap_out = torch.tensor(execute_model_req.blocks_to_swap_out,

325 - device="cpu",

326 - dtype=torch.int64).view(-1, 2)

327 - # `blocks_to_copy` is a gpu tensor. The src and tgt of

328 - # blocks to copy are in the same device, and `blocks_to_copy`

329 - # can be used directly within cuda kernels.

330 - blocks_to_copy = torch.tensor(execute_model_req.blocks_to_copy,

331 - device=self.device,

332 - dtype=torch.int64).view(-1, 2)

333 -

334 - return WorkerInput(

335 - num_seq_groups=num_seq_groups,

336 - blocks_to_swap_in=blocks_to_swap_in,

337 - blocks_to_swap_out=blocks_to_swap_out,

338 - blocks_to_copy=blocks_to_copy,

339 - virtual_engine=virtual_engine,

340 - )

341 -

342 - @torch.inference_mode()

343 - def execute_worker(self, worker_input: WorkerInput) -> None:

344 - virtual_engine = worker_input.virtual_engine

345 - # Issue cache operations.

346 - if (worker_input.blocks_to_swap_in is not None

347 - and worker_input.blocks_to_swap_in.numel() > 0):

348 - self.cache_engine[virtual_engine].swap_in(

349 - worker_input.blocks_to_swap_in)

350 - if (worker_input.blocks_to_swap_out is not None

351 - and worker_input.blocks_to_swap_out.numel() > 0):

352 - self.cache_engine[virtual_engine].swap_out(

353 - worker_input.blocks_to_swap_out)

354 - if (worker_input.blocks_to_copy is not None

355 - and worker_input.blocks_to_copy.numel() > 0):

356 - self.cache_engine[virtual_engine].copy(worker_input.blocks_to_copy)

357 -

358 - def add_lora(self, lora_request: LoRARequest) -> bool:

359 - return self.model_runner.add_lora(lora_request)

360 -

361 - def remove_lora(self, lora_id: int) -> bool:

362 - return self.model_runner.remove_lora(lora_id)

363 -

364 - def pin_lora(self, lora_id: int) -> bool:

365 - return self.model_runner.pin_lora(lora_id)

366 -

367 - def list_loras(self) -> Set[int]:

368 - return self.model_runner.list_loras()

369 -

370 - def add_prompt_adapter(

371 - self, prompt_adapter_request: PromptAdapterRequest) -> bool:

372 - raise NotImplementedError(

373 - "Prompt Adapter is not implemented for HPU backend.")

374 -

375 - def remove_prompt_adapter(self, prompt_adapter_id: int) -> bool:

376 - raise NotImplementedError(

377 - "Prompt Adapter is not implemented for HPU backend.")

378 -

379 - def pin_prompt_adapter(self, prompt_adapter_id: int) -> bool:

380 - raise NotImplementedError(

381 - "Prompt Adapter is not implemented for HPU backend.")

382 -

383 - def list_prompt_adapters(self) -> Set[int]:

384 - raise NotImplementedError(

385 - "Prompt Adapter is not implemented for HPU backend.")

386 -

387 - def shutdown_inc(self):

388 - self.model_runner.shutdown_inc()

389 -

390 - @property

391 - def max_model_len(self) -> int:

392 - return self.model_config.max_model_len

393 -

394 - @property

395 - def vocab_size(self) -> int:

396 - return self.model_runner.vocab_size

397 -

398 - def get_cache_block_size_bytes(self) -> int:

399 - """Get the size of the KV cache block size in bytes.

400 - """

401 - return HPUCacheEngine.get_cache_block_size(self.cache_config,

402 - self.model_config,

403 - self.parallel_config)

404 -

405 -

406 -def init_worker_distributed_environment(

407 - parallel_config: ParallelConfig,

408 - rank: int,

409 - distributed_init_method: Optional[str] = None,

410 - local_rank: int = -1,

411 -) -> None:

412 - """Initialize the distributed environment."""

413 - init_distributed_environment(parallel_config.world_size,

414 - rank,

415 - distributed_init_method,

416 - local_rank,

417 - backend=current_platform.dist_backend)

418 -

419 - ensure_model_parallel_initialized(parallel_config.tensor_parallel_size,

420 - parallel_config.pipeline_parallel_size)

421 -

422 - if torch.distributed.is_initialized():

423 - torch_world_size = torch.distributed.get_world_size()

424 - if torch_world_size != parallel_config.world_size:

425 - raise RuntimeError(

426 - "torch.distributed is already initialized but the torch world "

427 - "size does not match parallel_config.world_size "

428 - f"({torch_world_size} vs. {parallel_config.world_size}).")

429 - elif not distributed_init_method:

430 - raise ValueError(

431 - "distributed_init_method must be set if torch.distributed "

432 - "is not already initialized")

433 - else:

434 - torch.distributed.init_process_group(

435 - backend="hccl",

436 - world_size=parallel_config.world_size,

437 - rank=rank,

438 - init_method=distributed_init_method,

439 - )

440 -

441 - # A small all_reduce for warmup & checking conformance.

442 - dummy_tensor_hpu = torch.ones(1).to('hpu')

443 - torch.distributed.all_reduce(dummy_tensor_hpu)

444 - assert dummy_tensor_hpu.item() == parallel_config.world_size

445 - ensure_model_parallel_initialized(parallel_config.tensor_parallel_size,

446 - parallel_config.pipeline_parallel_size)

447 -

448 -

449 -def raise_if_cache_size_invalid(num_gpu_blocks, block_size, max_model_len,

450 - pipeline_parallel_size) -> None:

451 - if num_gpu_blocks <= 0:

452 - raise ValueError("No available memory for the cache blocks. "

453 - "Try increasing `gpu_memory_utilization` when "

454 - "initializing the engine.")

455 - max_seq_len = block_size * (num_gpu_blocks // pipeline_parallel_size)

456 - if max_model_len > max_seq_len:

457 - raise ValueError(

458 - f"The model's max seq len ({max_model_len}) "

459 - "is larger than the maximum number of tokens that can be "

460 - f"stored in KV cache ({max_seq_len}). Try increasing "

461 - "`gpu_memory_utilization` or decreasing `max_model_len` when "

462 - "initializing the engine.")

463 -

464 -

465 -class HPUCacheEngine(CacheEngine):

466 -

467 - def _allocate_kv_cache(

468 - self,

469 - num_blocks: int,

470 - device: str,

471 - ) -> List[Tuple[torch.Tensor, torch.Tensor]]:

472 - """Allocates KV cache on the specified device."""

473 - kv_cache_shape = self.attn_backend.get_kv_cache_shape(

474 - num_blocks, self.block_size, self.num_kv_heads, self.head_size)

475 - kv_cache: List[Tuple[torch.Tensor, torch.Tensor]] = []

476 - for _ in range(self.num_attention_layers):

477 - key_cache = torch.zeros(kv_cache_shape,

478 - dtype=self.dtype,

479 - device=device)

480 - value_cache = torch.zeros(kv_cache_shape,

481 - dtype=self.dtype,

482 - device=device)

483 - kv_layer = (key_cache, value_cache)

484 - kv_cache.append(kv_layer)

485 - return kv_cache

vllm/worker/multi_step_hpu_worker.py deleted

-123

@@ -1,123 +0,0 @@

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

2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 -

4 -###############################################################################

6 -###############################################################################

7 -

8 -import dataclasses

9 -from typing import Dict, Optional, Tuple

10 -

11 -import torch

12 -

13 -from vllm.distributed import broadcast_tensor_dict

14 -from vllm.sequence import ExecuteModelRequest

15 -from vllm.worker.hpu_model_runner import ModelInputForHPU

16 -from vllm.worker.hpu_worker import HPUWorker

17 -from vllm.worker.worker_base import WorkerInput

18 -

19 -

20 -class MultiStepHPUWorker(HPUWorker):

21 -

22 - def __init__(self, *args, **kwargs):

23 - super().__init__(*args, **kwargs)

24 - self.cached_model_input: Optional[ModelInputForHPU] = None

25 -

26 - def _get_driver_input_and_broadcast(

27 - self, execute_model_req: ExecuteModelRequest

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

29 - """

30 - Get the driver input and broadcast it to other workers.

31 - """

32 - assert self.is_driver_worker

33 - assert execute_model_req.virtual_engine == 0

34 -

35 - is_first_multi_step = execute_model_req.is_first_multi_step

36 - is_last_step = execute_model_req.is_last_step

37 -

38 - if is_first_multi_step:

39 - # on first step we prepare the worker input and model input normally

40 - worker_input: WorkerInput = self.prepare_worker_input(

41 - execute_model_req=execute_model_req)

42 - worker_input = dataclasses.replace(

43 - worker_input,

44 - num_steps=execute_model_req.num_lookahead_slots + 1)

45 - model_input: ModelInputForHPU = (

46 - self.model_runner.prepare_model_input(

47 - execute_model_req.seq_group_metadata_list,

48 - execute_model_req.virtual_engine,

49 - execute_model_req.finished_requests_ids))

50 -

51 - if execute_model_req.async_callback:

52 - model_input = dataclasses.replace(

53 - model_input,

54 - async_callback=execute_model_req.async_callback)

55 - else:

56 - # on subsequent steps we reuse the worker input and model input

57 - assert self.cached_model_input is not None

58 - model_input = self.cached_model_input

59 - worker_input = WorkerInput()

60 -

61 - model_input = dataclasses.replace(

62 - model_input,

63 - is_first_multi_step=is_first_multi_step,

64 - is_last_step=is_last_step)

65 -

66 - if self.do_metadata_broadcast:

67 - if is_first_multi_step:

68 - broadcast_data = worker_input.as_broadcastable_tensor_dict()

69 - broadcast_data.update(

70 - model_input.as_broadcastable_tensor_dict())

71 - broadcast_tensor_dict(broadcast_data, src=0)

72 - else:

73 - broadcast_data = {

74 - "is_first_multi_step": is_first_multi_step,

75 - "is_last_step": is_last_step,

76 - }

77 - broadcast_tensor_dict(broadcast_data, src=0)

78 -

79 - # Returning empty dict here to keep this compatible with

80 - # `LocalOrDistributedWorkerBase._get_driver_input_and_broadcast`

81 - return model_input, worker_input, {}

82 -

83 - def prepare_input(

84 - self,

85 - execute_model_req: Optional[ExecuteModelRequest] = None,

86 - ) -> Optional[Tuple[ModelInputForHPU, WorkerInput, Dict[str,

87 - torch.Tensor]]]:

88 - if self.is_driver_worker:

89 - if execute_model_req is None:

90 - if self.do_metadata_broadcast:

91 - # This signals that there's no more requests to process for

92 - # now. All workers are running infinite loop with

93 - # broadcast_tensor_dict, and it stops the loop when the

94 - # driver broadcasts an empty input. Send an empty input to

95 - # notify all other workers to stop their execution loop.

96 - broadcast_tensor_dict({}, src=0)

97 - return None

98 - model_input, worker_input, _ = self._get_driver_input_and_broadcast(

99 - execute_model_req)

100 - if model_input.is_first_multi_step:

101 - self.cached_model_input = model_input

102 - return model_input, worker_input, {}

103 - else:

104 - broadcast_data = broadcast_tensor_dict(src=0)

105 - if not broadcast_data:

106 - return None

107 -

108 - if len(broadcast_data) == 2:

109 - assert self.cached_model_input is not None

110 - self.cached_model_input = dataclasses.replace(

111 - self.cached_model_input,

112 - is_first_multi_step=broadcast_data["is_first_multi_step"],

113 - is_last_step=broadcast_data["is_last_step"])

114 - empty_worker_input = WorkerInput()

115 - return self.cached_model_input, empty_worker_input, {}

116 -

117 - worker_input = WorkerInput.from_broadcasted_tensor_dict(

118 - broadcast_data)

119 - model_input = (

120 - self.model_runner.

121 - make_model_input_from_broadcasted_tensor_dict(broadcast_data))

122 - self.cached_model_input = model_input

123 - return model_input, worker_input, {}