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

[Hardware][Intel-Gaudi] Update hpu-extension and update bucketing system for HPU device (#17186)

Signed-off-by: Agata Dobrzyniewicz <adobrzyniewicz@habana.ai>

Agata Dobrzyniewicz committed Apr 26, 2025 at 14:55 UTC c48334d405d6e0e106b4b6d922df8ff8ce58b64d
6 files changed +128 -335
requirements/hpu.txt
+1 -1
@@ -9,4 +9,4 @@ numpy==1.26.4
9 tabulate
10 setuptools>=61
11 setuptools-scm>=8
12 -vllm-hpu-extension @ git+https://github.com/HabanaAI/vllm-hpu-extension.git@4312768
12 +vllm-hpu-extension @ git+https://github.com/HabanaAI/vllm-hpu-extension.git@f1f6624
vllm/attention/backends/hpu_attn.py
+54 -52
@@ -4,14 +4,14 @@
4 # Copyright (C) 2024 Habana Labs, Ltd. an Intel Company
5 ###############################################################################
6
7 -import os
7 from dataclasses import dataclass
8 from typing import Any, Dict, List, Optional, Tuple, Type
9
10 import torch
11 +import vllm_hpu_extension.kernels as kernels
12 import vllm_hpu_extension.ops as ops
13 -from vllm_hpu_extension.utils import (Matmul, ModuleFusedSDPA, Softmax,
14 - VLLMKVCache)
13 +from vllm_hpu_extension.flags import enabled_flags
14 +from vllm_hpu_extension.utils import Matmul, Softmax, VLLMKVCache
15
16 from vllm.attention.backends.abstract import (AttentionBackend, AttentionImpl,
17 AttentionLayer,
@@ -126,7 +126,15 @@ class HPUAttentionImpl(AttentionImpl, torch.nn.Module):
126 self.block2batch_matmul = Matmul()
127 self.k_cache = VLLMKVCache()
128 self.v_cache = VLLMKVCache()
129 - ops.pa_impl = ops.pa
129 + self.fused_scaled_dot_product_attention = kernels.fsdpa()
130 +
131 + self.prefill_impl = 'naive'
132 + if "flex_attention" in enabled_flags():
133 + self.prefill_impl = 'flex'
134 + if "fsdpa" in enabled_flags():
135 + assert alibi_slopes is None, \
136 + 'Prefill with FusedSDPA not supported with alibi slopes!'
137 + self.prefill_impl = 'fsdpa'
138
139 self.num_kv_heads = num_heads if num_kv_heads is None else num_kv_heads
140 self.sliding_window = sliding_window
@@ -138,19 +146,9 @@ class HPUAttentionImpl(AttentionImpl, torch.nn.Module):
146 assert self.num_heads % self.num_kv_heads == 0
147 self.num_queries_per_kv = self.num_heads // self.num_kv_heads
148
141 - self.prefill_usefusedsdpa = os.getenv('VLLM_PROMPT_USE_FUSEDSDPA',
142 - '0').lower() in ['1', 'true']
143 - self.fused_scaled_dot_product_attention = None
144 - if self.prefill_usefusedsdpa:
149 + if self.prefill_impl == 'fsdpa':
150 assert alibi_slopes is None, \
151 'Prefill with FusedSDPA not supported with alibi slopes!'
147 - try:
148 - from habana_frameworks.torch.hpex.kernels import FusedSDPA
149 - self.fused_scaled_dot_product_attention = ModuleFusedSDPA(
150 - FusedSDPA)
151 - except ImportError:
152 - logger.warning("Could not import HPU FusedSDPA kernel. "
153 - "vLLM will use native implementation.")
152
153 supported_head_sizes = HPUPagedAttention.get_supported_head_sizes()
154 if head_size not in supported_head_sizes:
@@ -158,7 +156,8 @@ class HPUAttentionImpl(AttentionImpl, torch.nn.Module):
156 f"Head size {head_size} is not supported by PagedAttention. "
157 f"Supported head sizes are: {supported_head_sizes}.")
158
161 - if attn_type != AttentionType.DECODER:
159 + self.attn_type = attn_type
160 + if self.attn_type != AttentionType.DECODER:
161 raise NotImplementedError("Encoder self-attention and "
162 "encoder/decoder cross-attention "
163 "are not implemented for "
@@ -192,15 +191,18 @@ class HPUAttentionImpl(AttentionImpl, torch.nn.Module):
191 batch_size, seq_len, hidden_size = query.shape
192 _, seq_len_kv, _ = key.shape
193
195 - query = query.view(-1, self.num_heads, self.head_size)
194 key = key.view(-1, self.num_kv_heads, self.head_size)
195 value = value.view(-1, self.num_kv_heads, self.head_size)
196 block_indices = attn_metadata.block_indices
197 block_offsets = attn_metadata.block_offsets
200 - if attn_metadata.is_prompt:
198 + key_cache = None
199 + value_cache = None
200 + if attn_metadata.is_prompt and self.attn_type \
201 + is not AttentionType.ENCODER_ONLY \
202 + and attn_metadata.block_list is None:
203 key = key.unflatten(0, (block_indices.size(0), -1))
204 value = value.unflatten(0, (block_indices.size(0), -1))
203 - if kv_cache is not None:
205 + if kv_cache is not None and isinstance(kv_cache, tuple):
206 key_cache, value_cache = HPUPagedAttention.split_kv_cache(
207 kv_cache, self.num_kv_heads, self.head_size)
208
@@ -214,36 +216,28 @@ class HPUAttentionImpl(AttentionImpl, torch.nn.Module):
216
217 if attn_metadata.is_prompt:
218 # Prompt run.
217 - if not self.prefill_usefusedsdpa:
218 - # TODO: move this outside of model
219 - assert attn_metadata.attn_bias is not None, \
220 - 'attn_bias must be set before calling model.forward!'
221 - attn_bias = attn_metadata.attn_bias
222 - if self.alibi_slopes is not None:
223 - position_bias = _make_alibi_bias(self.alibi_slopes,
224 - self.num_kv_heads,
225 - attn_bias.dtype,
226 - attn_bias.shape[-1])
227 - attn_bias = attn_bias.tile((1, self.num_kv_heads, 1, 1))
228 - attn_bias.add_(position_bias)
229 - else:
230 - attn_bias = None
231 -
219 query_shape = (batch_size, seq_len, self.num_heads, self.head_size)
220 kv_shape = (batch_size, seq_len_kv, self.num_kv_heads,
221 self.head_size)
222 +
223 + attn_bias = attn_metadata.attn_bias
224 + if attn_bias is not None and self.alibi_slopes is not None:
225 + position_bias = _make_alibi_bias(self.alibi_slopes,
226 + self.num_kv_heads,
227 + attn_bias.dtype,
228 + attn_bias.shape[-1])
229 + attn_bias = attn_bias.tile((1, self.num_kv_heads, 1, 1))
230 + attn_bias.add_(position_bias)
231 +
232 out = ops.prompt_attention(
236 - query.view(query_shape),
237 - key.view(kv_shape),
238 - value.view(kv_shape),
233 + impl=self.prefill_impl,
234 + query=query.view(query_shape),
235 + key=key.view(kv_shape),
236 + value=value.view(kv_shape),
237 + is_causal=True,
238 attn_bias=attn_bias,
240 - p=0.0,
241 - scale=self.scale,
242 - matmul_qk_op=self.matmul_qk,
243 - softmax_op=self.softmax,
244 - matmul_av_op=self.matmul_av,
245 - fsdpa_op=self.fused_scaled_dot_product_attention,
246 - )
239 + valid_seq_lengths=attn_metadata.seq_lens_tensor,
240 + **self.common_attention_args())
241 output = out.reshape(batch_size, seq_len, hidden_size)
242 else:
243 # Decoding run.
@@ -254,18 +248,26 @@ class HPUAttentionImpl(AttentionImpl, torch.nn.Module):
248 block_list=attn_metadata.block_list,
249 block_mapping=attn_metadata.block_mapping,
250 block_bias=attn_metadata.attn_bias,
257 - block_scales=attn_metadata.block_scales,
251 block_groups=attn_metadata.block_groups,
259 - scale=self.scale,
260 - matmul_qk_op=self.matmul_qk,
261 - matmul_av_op=self.matmul_av,
262 - batch2block_matmul_op=self.batch2block_matmul,
263 - block2batch_matmul_op=self.block2batch_matmul,
264 - keys_fetch_func=self.k_cache.fetch_from_cache,
265 - values_fetch_func=self.v_cache.fetch_from_cache)
252 + **self.common_attention_args())
253 # Reshape the output tensor.
254 return output.view(batch_size, seq_len, hidden_size)
255
256 + def common_attention_args(self):
257 + fsdpa_op = self.fused_scaled_dot_product_attention.apply \
258 + if self.fused_scaled_dot_product_attention is not None else None
259 + return {
260 + 'scale': self.scale,
261 + 'matmul_qk_op': self.matmul_qk,
262 + 'matmul_av_op': self.matmul_av,
263 + 'batch2block_matmul_op': self.batch2block_matmul,
264 + 'block2batch_matmul_op': self.block2batch_matmul,
265 + 'fsdpa_op': fsdpa_op,
266 + 'keys_fetch_func': self.k_cache.fetch_from_cache,
267 + 'values_fetch_func': self.v_cache.fetch_from_cache,
268 + 'softmax_op': self.softmax,
269 + }
270 +
271
272 def _make_alibi_bias(
273 alibi_slopes: torch.Tensor,
vllm/attention/ops/hpu_paged_attn.py
-1
@@ -22,7 +22,6 @@ class HPUPagedAttentionMetadata:
22 block_usage: Optional[torch.Tensor]
23 block_indices: Optional[torch.Tensor]
24 block_offsets: Optional[torch.Tensor]
25 - block_scales: Optional[torch.Tensor]
25 block_groups: Optional[torch.Tensor]
26
27
vllm/model_executor/layers/layernorm.py
+2 -1
@@ -168,7 +168,8 @@ class RMSNorm(CustomOp):
168 x: torch.Tensor,
169 residual: Optional[torch.Tensor] = None,
170 ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
171 - from vllm_hpu_extension.ops import HPUFusedRMSNorm
171 + from vllm_hpu_extension.kernels import rms_norm
172 + HPUFusedRMSNorm = rms_norm()
173 if HPUFusedRMSNorm is None:
174 return self.forward_native(x, residual)
175 if residual is not None:
vllm/worker/hpu_model_runner.py
+70 -280
@@ -11,11 +11,9 @@ import functools
11 import gc
12 import itertools
13 import math
14 -import operator
14 import os
15 import time
16 from array import array
18 -from dataclasses import dataclass, field
17 from enum import IntEnum
18 from typing import (TYPE_CHECKING, Any, Callable, Dict, List, NamedTuple,
19 Optional, Set, Tuple, Type, TypeVar, Union)
@@ -24,8 +22,9 @@ import habana_frameworks.torch as htorch
22 import habana_frameworks.torch.internal.bridge_config as bc
23 import torch
24 import torch.nn as nn
25 +import vllm_hpu_extension.environment as environment
26 +from vllm_hpu_extension.bucketing.common import get_bucketing_context
27 from vllm_hpu_extension.ops import LoraMask as LoraMask
28 -from vllm_hpu_extension.ops import batch2block, block2batch
28 from vllm_hpu_extension.profiler import (HabanaHighLevelProfiler,
29 HabanaMemoryProfiler, format_bytes)
30
@@ -77,25 +76,6 @@ LORA_WARMUP_RANK = 8
76 DUMMY_TOKEN_ID = -1
77
78
80 -class Singleton(type):
81 - _instances: Dict[type, object] = {}
82 -
83 - def __call__(cls, *args, **kwargs):
84 - if cls not in cls._instances:
85 - cls._instances[cls] = super().__call__(*args, **kwargs)
86 - return cls._instances[cls]
87 -
88 -
89 -@dataclass
90 -class HPUBucketingGlobalState(metaclass=Singleton):
91 - prompt_bs_bucket_cfg: Tuple[int, int, int] = field(init=False)
92 - decode_bs_bucket_cfg: Tuple[int, int, int] = field(init=False)
93 - prompt_seq_bucket_cfg: Tuple[int, int, int] = field(init=False)
94 - decode_block_bucket_cfg: Tuple[int, int, int] = field(init=False)
95 - prompt_buckets: List[Tuple[int, int]] = field(init=False)
96 - decode_buckets: List[Tuple[int, int]] = field(init=False)
97 -
98 -
79 def subtuple(obj: object,
80 typename: str,
81 to_copy: List[str],
@@ -115,134 +95,10 @@ def subtuple(obj: object,
95 return _TYPE_CACHE[typename](**values)
96
97
118 -def read_bucket_settings(phase: str, dim: str, **defaults):
119 - """Read bucketing configuration from env variables.
120 -
121 - phase is either 'prompt' or 'decode'
122 - dim is either 'bs', 'seq' or 'block'
123 - param is either 'min', 'step' or 'max'
124 - example env variable: VLLM_DECODE_BS_BUCKET_STEP=128
125 - """
126 - params = ['min', 'step', 'max']
127 - env_vars = [f'VLLM_{phase}_{dim}_BUCKET_{p}'.upper() for p in params]
128 - default_values = [defaults[p] for p in params]
129 - values = [
130 - int(os.environ.get(e, d)) for e, d in zip(env_vars, default_values)
131 - ]
132 - for e, v, d in zip(env_vars, values, default_values):
133 - logger.info('%s=%s (default:%s)', e, v, d)
134 - return values
135 -
136 -
137 -def warmup_range(config: Tuple[int, int, int]):
138 - """Generate a warmup range.
139 -
140 - Start from bmin and multiply by 2 until you reach bstep.
141 - Then, increase the values in the range by the value of bstep until you
142 - reach bmax.
143 -
144 - Example:
145 - bmin = 2, bstep = 32, bmax = 64
146 - => ramp_up = (2, 4, 8, 16)
147 - => stable = (32, 64)
148 - => return ramp_up + stable => (2, 4, 8, 16, 32, 64)
149 - """
150 - bmin, bstep, bmax = config
151 - assert bmin <= bmax, ("Min. batch size cannot be greater than max. "
152 - "batch size. If you want to skip warmup, "
153 - "set VLLM_SKIP_WARMUP=true")
154 - base = itertools.repeat(2)
155 - ramp_up_acc = itertools.accumulate(base, func=operator.mul, initial=bmin)
156 - ramp_up_tw = itertools.takewhile(lambda x: x < bstep and x <= bmax, \
157 - ramp_up_acc)
158 - stable = range(bstep, bmax + 1, bstep)
159 - buckets = list(ramp_up_tw) + list(stable)
160 - return list(filter(lambda bucket: bucket >= bmin, buckets))
161 -
162 -
163 -def generate_prompt_buckets(bs_bucket_config,
164 - seq_bucket_config,
165 - max_num_batched_tokens=None):
166 - buckets = list(
167 - itertools.product(warmup_range(bs_bucket_config),
168 - warmup_range(seq_bucket_config)))
169 - if len(buckets) == 0:
170 - msg = ("No buckets could be captured with following config "
171 - f"(min, step, max_warmup): "
172 - f"bs:{bs_bucket_config}, "
173 - f"seq:{seq_bucket_config}")
174 - raise ValueError(msg)
175 -
176 - filtered_buckets = buckets
177 - if max_num_batched_tokens is not None:
178 - # Remove buckets exceeding batch token budget
179 - filtered_buckets = list(
180 - filter(
181 - lambda bucket: bucket[0] * bucket[1] <= max_num_batched_tokens,
182 - buckets))
183 -
184 - if len(filtered_buckets) == 0:
185 - # we can handle this if we ignore max_num_batched_tokens
186 - min_bucket_bs, min_bucket_seq = min(buckets,
187 - key=lambda b: (b[0] * b[1]))
188 - min_reqd_budget = min_bucket_bs * min_bucket_seq
189 - msg = (
190 - "The current bucketing configuration "
191 - f"(min, step, max_warmup): "
192 - f"bs:{bs_bucket_config}, "
193 - f"seq:{seq_bucket_config} cannot be used with specified "
194 - f"max_num_batched_tokens ({max_num_batched_tokens}), as the "
195 - f"smallest bucket ({min_reqd_budget}) would exceed token "
196 - "budget. Please increase max_num_batched_tokens or decrease "
197 - "bucket minimum Ignoring max_num_batched_tokens at risk of "
198 - "out-of-memory errors.")
199 - logger.error(msg)
200 - return list(
201 - sorted(buckets, key=lambda b: (b[0] * b[1], b[1], b[0]))), []
202 -
203 - captured_buckets = list(
204 - sorted(filtered_buckets, key=lambda b: (b[0] * b[1], b[1], b[0])))
205 - omitted_buckets = list(
206 - sorted([x for x in buckets if x not in filtered_buckets]))
207 - return captured_buckets, omitted_buckets
208 -
209 -
210 -def generate_decode_buckets(bs_bucket_config, blocks_bucket_config,
211 - max_blocks):
212 - buckets = []
213 - bs_buckets = warmup_range(bs_bucket_config)
214 - block_buckets = warmup_range(blocks_bucket_config)
215 - bmin, bstep, bmax = blocks_bucket_config
216 - last_bucket = round_up(max_blocks, bstep)
217 - for bs in bs_buckets:
218 - for blocks in block_buckets:
219 - if blocks < bs:
220 - continue
221 - if blocks > last_bucket:
222 - break
223 - buckets.append((bs, blocks))
224 - return list(sorted(buckets, key=lambda b: (b[0] * b[1], b[1], b[0])))
225 -
226 -
227 -def next_pow2(value: int, base: int):
228 - res = base
229 - while value > 1:
230 - value = (value + 1) // 2
231 - res *= 2
232 - return res
233 -
234 -
98 def round_up(value: int, k: int):
99 return (value + k - 1) // k * k
100
101
239 -def find_bucket(value: int, config: Tuple[int, int, int]):
240 - bmin, bstep, _ = config
241 - next_step = round_up(value, bstep)
242 - next_pow = next_pow2(value, bmin)
243 - return max(bmin, min(next_step, next_pow))
244 -
245 -
102 def align_workers(value, op):
103 group = get_world_group().cpu_group
104 world_size = torch.distributed.get_world_size()
@@ -406,16 +262,6 @@ class HpuModelAdapter:
262 attn_bias=attn_bias)
263 return metadata
264
409 - def _set_block_scales(self, metadata, device):
410 - block_mapping = metadata.block_mapping
411 - ones = torch.ones((block_mapping.size(0), ),
412 - device=device,
413 - dtype=block_mapping.dtype)
414 - sums = batch2block(block2batch(ones, block_mapping), block_mapping)
415 - block_scales = torch.reciprocal(torch.maximum(ones, sums))
416 - metadata = metadata._replace(block_scales=block_scales)
417 - return metadata
418 -
265 def _update_metadata(self, attn_metadata, batch_size, seq_len, device,
266 dtype):
267 if attn_metadata.is_prompt:
@@ -426,7 +272,6 @@ class HpuModelAdapter:
272 meta = attn_metadata
273 attn_metadata = self._set_block_mapping(meta, batch_size, device,
274 dtype)
429 - attn_metadata = self._set_block_scales(attn_metadata, device)
275 return attn_metadata
276
277 def forward(self, *args, **kwargs):
@@ -625,6 +470,7 @@ class HPUModelRunnerBase(ModelRunnerBase[TModelInputForHPU]):
470 return_hidden_states: bool = False,
471 ):
472 ModelRunnerBase.__init__(self, vllm_config=vllm_config)
473 + environment.set_model_config(self.model_config)
474 self.is_driver_worker = is_driver_worker
475 self.return_hidden_states = return_hidden_states
476
@@ -664,8 +510,13 @@ class HPUModelRunnerBase(ModelRunnerBase[TModelInputForHPU]):
510 self.profiler_counter_helper = HabanaProfilerCounterHelper()
511 self.seen_configs: set = set()
512 self._mem_margin: Optional[int] = None
667 - self.bucketing_global_state = HPUBucketingGlobalState()
668 - self._setup_buckets()
513 + HPUBucketingContext = get_bucketing_context()
514 + self.bucketing_ctx = HPUBucketingContext(self.max_num_seqs,
515 + self.max_num_prefill_seqs,
516 + self.block_size,
517 + self.max_num_batched_tokens,
518 + False, self.max_model_len)
519 + self.graphed_buckets: Set[Any] = set()
520 self._set_gc_threshold()
521 self.use_contiguous_pa = envs.VLLM_USE_HPU_CONTIGUOUS_CACHE_FETCH
522
@@ -773,6 +624,21 @@ class HPUModelRunnerBase(ModelRunnerBase[TModelInputForHPU]):
624 msg = f"Loading model weights took in total {m.get_summary_string()}"
625 logger.info(msg)
626
627 + def _add_dummy_seq(self, seq_group_metadata_list, is_prompt):
628 + real_batch_size = len(seq_group_metadata_list)
629 + batch_size_padded = self.bucketing_ctx.get_padded_batch_size(
630 + real_batch_size, is_prompt)
631 + batch_size_padding = batch_size_padded - real_batch_size
632 +
633 + seq_group_metadata_list = seq_group_metadata_list.copy()
634 +
635 + if batch_size_padding > 0:
636 + dummy_seq_group_metadata = self.create_dummy_seq_group_metadata(
637 + 0, 0, is_prompt)
638 + seq_group_metadata_list.extend(dummy_seq_group_metadata
639 + for _ in range(batch_size_padding))
640 + return seq_group_metadata_list, real_batch_size, batch_size_padded
641 +
642 def _maybe_wrap_in_hpu_graph(self, *args, **kwargs):
643 return htorch.hpu.wrap_in_hpu_graph(
644 HpuModelAdapter(*args, **kwargs), disable_tensor_cache=True
@@ -792,46 +658,6 @@ class HPUModelRunnerBase(ModelRunnerBase[TModelInputForHPU]):
658 def _is_valid_bucket(self, bucket):
659 return bucket[0] * bucket[1] <= self.max_num_batched_tokens
660
795 - def _setup_buckets(self) -> None:
796 - align_bs = lambda x: min(self.max_num_seqs, x)
797 - #FIXME: The default values should be max_model_len
798 - max_prompt_seq = 1024
799 - max_decode_seq = 2048
800 - self.bucketing_global_state.prompt_bs_bucket_cfg = read_bucket_settings(
801 - 'prompt',
802 - 'bs',
803 - min=1,
804 - step=align_bs(32),
805 - max=self.max_num_prefill_seqs)
806 - self.bucketing_global_state.decode_bs_bucket_cfg = read_bucket_settings(
807 - 'decode', 'bs', min=1, step=align_bs(32), max=self.max_num_seqs)
808 - self.bucketing_global_state.prompt_seq_bucket_cfg = \
809 - read_bucket_settings(
810 - 'prompt',
811 - 'seq',
812 - min=self.block_size,
813 - step=self.block_size,
814 - max=max_prompt_seq)
815 - self.bucketing_global_state.decode_block_bucket_cfg = \
816 - read_bucket_settings(
817 - 'decode',
818 - 'block',
819 - min=self.block_size,
820 - step=self.block_size,
821 - max=max(self.block_size,
822 - self.max_num_seqs * max_decode_seq // self.block_size))
823 - self.graphed_buckets: Set[Any] = set()
824 -
825 - msg = ("Prompt bucket config (min, step, max_warmup) "
826 - f"bs:{self.bucketing_global_state.prompt_bs_bucket_cfg}, "
827 - f"seq:{self.bucketing_global_state.prompt_seq_bucket_cfg}")
828 - logger.info(msg)
829 -
830 - msg = ("Decode bucket config (min, step, max_warmup) "
831 - f"bs:{self.bucketing_global_state.decode_bs_bucket_cfg}, "
832 - f"block:{self.bucketing_global_state.decode_block_bucket_cfg}")
833 - logger.info(msg)
834 -
661 def _prepare_prompt(
662 self,
663 seq_group_metadata_list: List[SequenceGroupMetadata],
@@ -947,8 +773,7 @@ class HPUModelRunnerBase(ModelRunnerBase[TModelInputForHPU]):
773 assert max_query_len > 0
774
775 max_prompt_len = max(
950 - find_bucket(max(seq_lens),
951 - self.bucketing_global_state.prompt_seq_bucket_cfg),
776 + self.bucketing_ctx.get_padded_prompt_seq_len(max_query_len),
777 self.block_size)
778
779 lora_ids: List[int] = []
@@ -997,7 +822,6 @@ class HPUModelRunnerBase(ModelRunnerBase[TModelInputForHPU]):
822 block_usage=None,
823 block_indices=block_indices,
824 block_offsets=block_offsets,
1000 - block_scales=None,
825 block_groups=None,
826 attn_bias=None,
827 seq_lens_tensor=seq_lens_tensor,
@@ -1124,9 +948,8 @@ class HPUModelRunnerBase(ModelRunnerBase[TModelInputForHPU]):
948 padding_fn = None
949 if self.use_contiguous_pa:
950 block_bucket_size = max(max(block_list) + 1, len(block_list))
1127 - block_bucket_size = find_bucket(
1128 - block_bucket_size,
1129 - self.bucketing_global_state.decode_block_bucket_cfg)
951 + block_bucket_size = self.bucketing_ctx.get_padded_decode_num_blocks(
952 + block_bucket_size)
953 indices: List[Any]
954 indices = [None] * block_bucket_size
955 for i, bid in enumerate(block_list):
@@ -1134,9 +957,9 @@ class HPUModelRunnerBase(ModelRunnerBase[TModelInputForHPU]):
957 padding_fn = lambda tensor, pad_value: gather_list(
958 tensor, indices, pad_value)
959 else:
1137 - block_bucket_size = find_bucket(
1138 - len(block_list),
1139 - self.bucketing_global_state.decode_block_bucket_cfg)
960 + block_bucket_size = \
961 + self.bucketing_ctx.get_padded_decode_num_blocks(
962 + len(block_list))
963 padding_fn = lambda tensor, pad_value: pad_list(
964 tensor, block_bucket_size, pad_value)
965
@@ -1167,7 +990,6 @@ class HPUModelRunnerBase(ModelRunnerBase[TModelInputForHPU]):
990 block_usage=block_usage,
991 block_indices=block_indices,
992 block_offsets=block_offsets,
1170 - block_scales=None,
993 block_groups=block_groups,
994 attn_bias=None,
995 seq_lens_tensor=None,
@@ -1210,17 +1032,8 @@ class HPUModelRunnerBase(ModelRunnerBase[TModelInputForHPU]):
1032 base_event_name = 'prompt' if is_prompt else 'decode'
1033 self.profiler.start('internal', base_event_name)
1034
1213 - real_batch_size = len(seq_group_metadata_list)
1214 - bucket_cfg = self.bucketing_global_state.prompt_bs_bucket_cfg \
1215 - if is_prompt else self.bucketing_global_state.decode_bs_bucket_cfg
1216 - batch_size_padded = find_bucket(real_batch_size, bucket_cfg)
1217 - batch_size_padding = batch_size_padded - real_batch_size
1218 - seq_group_metadata_list = seq_group_metadata_list.copy()
1219 - if batch_size_padding > 0:
1220 - dummy_seq_group_metadata = self.create_dummy_seq_group_metadata(
1221 - 0, 0, is_prompt)
1222 - seq_group_metadata_list.extend(dummy_seq_group_metadata
1223 - for _ in range(batch_size_padding))
1035 + seq_group_metadata_list, real_batch_size, batch_size_padded = (
1036 + self._add_dummy_seq(seq_group_metadata_list, is_prompt))
1037
1038 prefill_reqs = []
1039 decode_reqs = []
@@ -1382,7 +1195,7 @@ class HPUModelRunnerBase(ModelRunnerBase[TModelInputForHPU]):
1195 attention_metadata = subtuple(metadata, 'TrimmedAttentionMetadata', [
1196 'attn_bias', 'seq_lens_tensor', 'block_list', 'block_mapping',
1197 'block_usage', 'slot_mapping', 'is_prompt', 'block_indices',
1385 - 'block_offsets', 'block_scales', 'block_groups'
1198 + 'block_offsets', 'block_groups'
1199 ])
1200 return attention_metadata
1201
@@ -1420,16 +1233,18 @@ class HPUModelRunnerBase(ModelRunnerBase[TModelInputForHPU]):
1233 bind_kv_cache(
1234 self.vllm_config.compilation_config.static_forward_context,
1235 [kv_caches])
1423 - max_seq_len = self.bucketing_global_state.prompt_seq_bucket_cfg[-1]
1424 - max_batch_size = min(self.max_num_batched_tokens // max_seq_len,
1425 - self.scheduler_config.max_num_seqs)
1426 - self.warmup_scenario(max_batch_size, max_seq_len, True, False, True)
1236 + _, max_seq_len = self.bucketing_ctx.get_max_prompt_shape()
1237 + max_batch_size = min(self.max_num_seqs,
1238 + self.max_num_batched_tokens // max_seq_len)
1239 + self.warmup_scenario(max_batch_size, max_seq_len, True, kv_caches,
1240 + False, True)
1241 return
1242
1243 def warmup_scenario(self,
1244 batch_size,
1245 seq_len,
1246 is_prompt,
1247 + kv_caches,
1248 is_pt_profiler_run=False,
1249 is_lora_profile_run=False) -> None:
1250 use_graphs = self._use_graphs(batch_size, seq_len, is_prompt)
@@ -1565,16 +1380,17 @@ class HPUModelRunnerBase(ModelRunnerBase[TModelInputForHPU]):
1380 f"free_mem:{free_mem}")
1381 logger.info(msg)
1382
1568 - def warmup_all_buckets(self, buckets, is_prompt):
1383 + def warmup_all_buckets(self, buckets, is_prompt, kv_caches):
1384 for i, (batch_size, seq_len) in enumerate(reversed(buckets)):
1385 self.log_warmup('Prompt' if is_prompt else 'Decode', i,
1386 len(buckets), batch_size, seq_len)
1572 - self.warmup_scenario(batch_size, seq_len, is_prompt)
1387 + self.warmup_scenario(batch_size, seq_len, is_prompt, kv_caches)
1388
1389 def warmup_graphs(self,
1390 strategy,
1391 buckets,
1392 is_prompt,
1393 + kv_caches,
1394 available_mem,
1395 starting_mem=0,
1396 total_batch_seq=0.001):
@@ -1606,7 +1422,7 @@ class HPUModelRunnerBase(ModelRunnerBase[TModelInputForHPU]):
1422 self.graphed_buckets.add(graphed_bucket)
1423 self.log_warmup(phase, idx, num_candidates, batch_size, seq_len)
1424 with HabanaMemoryProfiler() as mem_prof:
1609 - self.warmup_scenario(batch_size, seq_len, is_prompt)
1425 + self.warmup_scenario(batch_size, seq_len, is_prompt, kv_caches)
1426 used_mem = align_workers(mem_prof.consumed_device_memory,
1427 torch.distributed.ReduceOp.MAX)
1428 available_mem -= used_mem
@@ -1630,50 +1446,21 @@ class HPUModelRunnerBase(ModelRunnerBase[TModelInputForHPU]):
1446
1447 @torch.inference_mode()
1448 def warmup_model(self, kv_caches: List[torch.Tensor]) -> None:
1449 + max_blocks = kv_caches[0][0].size(0)
1450 + self.bucketing_ctx.generate_decode_buckets(max_blocks)
1451 if profile := os.environ.get('VLLM_PT_PROFILE', None):
1452 phase, bs, seq_len, graph = profile.split('_')
1453 is_prompt = phase == 'prompt'
1454 graphs = graph == 't'
1455 if graphs:
1456 self.graphed_buckets.add((int(bs), int(seq_len), is_prompt))
1639 - self.warmup_scenario(int(bs), int(seq_len), is_prompt, True)
1457 + self.warmup_scenario(int(bs), int(seq_len), is_prompt, kv_caches,
1458 + True)
1459 raise AssertionError("Finished profiling")
1641 - if self.skip_warmup:
1642 - logger.info("Skipping warmup...")
1643 - return
1644 - self.profiler.start('internal', 'warmup')
1645 - max_blocks = kv_caches[0][0].size(0)
1646 -
1647 - self.bucketing_global_state.prompt_buckets, prompt_omitted_buckets = \
1648 - generate_prompt_buckets(
1649 - self.bucketing_global_state.prompt_bs_bucket_cfg,
1650 - self.bucketing_global_state.prompt_seq_bucket_cfg,
1651 - self.max_num_batched_tokens)
1652 -
1653 - msg = (f"Generated {len(self.bucketing_global_state.prompt_buckets)} "
1654 - f"prompt buckets [bs, seq]: \
1655 - {list(sorted(self.bucketing_global_state.prompt_buckets))}")
1656 - logger.info(msg)
1657 -
1658 - msg = (f"Omitted {len(prompt_omitted_buckets)} "
1659 - "prompt buckets due to exceeded token budget "
1660 - f"(max_num_batched_tokens={self.max_num_batched_tokens})")
1661 - logger.info(msg)
1662 -
1663 - msg = f"Omitted prompt buckets: {list(sorted(prompt_omitted_buckets))}"
1664 - logger.debug(msg)
1665 -
1666 - self.bucketing_global_state.decode_buckets = generate_decode_buckets(
1667 - self.bucketing_global_state.decode_bs_bucket_cfg,
1668 - self.bucketing_global_state.decode_block_bucket_cfg, max_blocks)
1669 - logger.info("Generated %d decode buckets [bs, total_blocks]: %s",
1670 - len(self.bucketing_global_state.decode_buckets),
1671 - list(sorted(self.bucketing_global_state.decode_buckets)))
1672 -
1460 if not htorch.utils.internal.is_lazy() and not self.enforce_eager:
1461 cache_size_limit = 1 + 3 * (
1675 - len(self.bucketing_global_state.prompt_buckets) +
1676 - len(self.bucketing_global_state.decode_buckets))
1462 + len(self.bucketing_ctx.prompt_buckets) +
1463 + len(self.bucketing_ctx.decode_buckets))
1464 torch._dynamo.config.cache_size_limit = max(
1465 cache_size_limit, torch._dynamo.config.cache_size_limit)
1466 # Multiply by 8 to follow the original default ratio between
@@ -1681,7 +1468,10 @@ class HPUModelRunnerBase(ModelRunnerBase[TModelInputForHPU]):
1468 torch._dynamo.config.accumulated_cache_size_limit = max(
1469 cache_size_limit * 8,
1470 torch._dynamo.config.accumulated_cache_size_limit)
1684 -
1471 + if self.skip_warmup:
1472 + logger.info("Skipping warmup...")
1473 + return
1474 + self.profiler.start('internal', 'warmup')
1475 start_mem = HabanaMemoryProfiler.current_device_memory_usage()
1476 start_time = time.perf_counter()
1477
@@ -1700,10 +1490,12 @@ class HPUModelRunnerBase(ModelRunnerBase[TModelInputForHPU]):
1490 'Please update Gaudi Software Suite.')
1491 with compile_only_mode_context(
1492 ) if can_use_compile_only_mode else contextlib.nullcontext():
1703 - self.warmup_all_buckets(self.bucketing_global_state.prompt_buckets,
1704 - True)
1705 - self.warmup_all_buckets(self.bucketing_global_state.decode_buckets,
1706 - False)
1493 + print("aa")
1494 + self.warmup_all_buckets(self.bucketing_ctx.prompt_buckets, True,
1495 + kv_caches)
1496 + print("bb")
1497 + self.warmup_all_buckets(self.bucketing_ctx.decode_buckets, False,
1498 + kv_caches)
1499
1500 if not self.enforce_eager and htorch.utils.internal.is_lazy():
1501 assert self.mem_margin is not None, \
@@ -1733,12 +1525,12 @@ class HPUModelRunnerBase(ModelRunnerBase[TModelInputForHPU]):
1525 'max_bs')
1526 mem_post_prompt, prompt_batch_seq, prompt_captured_all = \
1527 self.warmup_graphs(
1736 - prompt_strategy, self.bucketing_global_state.prompt_buckets,
1737 - True, prompt_available_memory)
1528 + prompt_strategy, self.bucketing_ctx.prompt_buckets,
1529 + True, kv_caches, prompt_available_memory)
1530 mem_post_decode, decode_batch_seq, decode_captured_all = \
1531 self.warmup_graphs(
1740 - decode_strategy, self.bucketing_global_state.decode_buckets,
1741 - False, decode_available_memory)
1532 + decode_strategy, self.bucketing_ctx.decode_buckets,
1533 + False, kv_caches, decode_available_memory)
1534
1535 # Not all prompt buckets were captured, but all decode buckets
1536 # were captured and we have some free graph-allocated space
@@ -1747,8 +1539,8 @@ class HPUModelRunnerBase(ModelRunnerBase[TModelInputForHPU]):
1539 and not prompt_captured_all and decode_captured_all):
1540 mem_post_prompt, _, prompt_captured_all = (
1541 self.warmup_graphs(
1750 - prompt_strategy,
1751 - self.bucketing_global_state.prompt_buckets, True,
1542 + prompt_strategy, self.bucketing_ctx.prompt_buckets,
1543 + True, kv_caches,
1544 graph_free_mem - mem_post_prompt - mem_post_decode,
1545 mem_post_prompt, prompt_batch_seq))
1546
@@ -1759,17 +1551,15 @@ class HPUModelRunnerBase(ModelRunnerBase[TModelInputForHPU]):
1551 and not decode_captured_all \
1552 and prompt_captured_all:
1553 mem_post_decode, _, _ = self.warmup_graphs(
1762 - decode_strategy,
1763 - self.bucketing_global_state.decode_buckets, False,
1554 + decode_strategy, self.bucketing_ctx.decode_buckets,
1555 + False, kv_caches,
1556 graph_free_mem - mem_post_prompt - mem_post_decode,
1557 mem_post_decode, decode_batch_seq)
1558
1559 self.log_graph_warmup_summary(
1768 - self.bucketing_global_state.prompt_buckets, True,
1769 - mem_post_prompt)
1560 + self.bucketing_ctx.prompt_buckets, True, mem_post_prompt)
1561 self.log_graph_warmup_summary(
1771 - self.bucketing_global_state.decode_buckets, False,
1772 - mem_post_decode)
1562 + self.bucketing_ctx.decode_buckets, False, mem_post_decode)
1563
1564 end_time = time.perf_counter()
1565 end_mem = HabanaMemoryProfiler.current_device_memory_usage()
vllm/worker/hpu_worker.py
+1
@@ -245,6 +245,7 @@ class HPUWorker(LocalOrDistributedWorkerBase):
245 cache_block_size)
246 num_hpu_blocks = max(num_hpu_blocks, 0)
247 num_cpu_blocks = max(num_cpu_blocks, 0)
248 + self.model_runner.bucketing_ctx.num_hpu_blocks = num_hpu_blocks
249
250 if self.model_runner.lora_manager:
251 self.model_runner.remove_all_loras()