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

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, 22:15 UTC

[Hardware][Gaudi][Feature] Support Contiguous Cache Fetch (#12139)

Signed-off-by: yuzhou <yuzhou@habana.ai> Signed-off-by: zhouyu5 <yu.zhou@intel.com> Co-authored-by: Cody Yu <hao.yu.cody@gmail.com>

Yu-Zhou committed Feb 19, 2025 at 11:40 UTC d0a7a2769d92619afdcdc3b91c78098eaa9e38c0

4 files changed +81 -48

vllm/attention/backends/hpu_attn.py

+1 -5

@@ -118,12 +118,8 @@ class HPUAttentionImpl(AttentionImpl, torch.nn.Module):

118 self.matmul_av = Matmul()

119 self.batch2block_matmul = Matmul()

120 self.block2batch_matmul = Matmul()

121 - # NOTE(kzawora): Contiguous PA is off until model runner supports it

121 self.k_cache = VLLMKVCache()

123 - self.k_cache.use_contiguous_pa = False

122 self.v_cache = VLLMKVCache()

125 - self.v_cache.use_contiguous_pa = False

126 - # NOTE(kzawora): Pipelined PA is off until model runner supports it

123 ops.pa_impl = ops.pa

124

125 self.num_kv_heads = num_heads if num_kv_heads is None else num_kv_heads

@@ -249,7 +245,7 @@ class HPUAttentionImpl(AttentionImpl, torch.nn.Module):

245 block_mapping=attn_metadata.block_mapping,

246 block_bias=attn_metadata.attn_bias,

247 block_scales=attn_metadata.block_scales,

252 - block_groups=None,

248 + block_groups=attn_metadata.block_groups,

249 scale=self.scale,

250 matmul_qk_op=self.matmul_qk,

251 matmul_av_op=self.matmul_av,

vllm/attention/ops/hpu_paged_attn.py

@@ -23,6 +23,7 @@ class HPUPagedAttentionMetadata:

23 block_indices: Optional[torch.Tensor]

24 block_offsets: Optional[torch.Tensor]

25 block_scales: Optional[torch.Tensor]

26 + block_groups: Optional[torch.Tensor]

29 class HPUPagedAttention:

vllm/envs.py

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

89 VLLM_RAY_PER_WORKER_GPUS: float = 1.0

90 VLLM_RAY_BUNDLE_INDICES: str = ""

91 VLLM_CUDART_SO_PATH: Optional[str] = None

92 + VLLM_USE_HPU_CONTIGUOUS_CACHE_FETCH: bool = True

95 def get_default_cache_root():

@@ -585,6 +586,13 @@ environment_variables: Dict[str, Callable[[], Any]] = {

586 # specify the path through environment variable VLLM_CUDART_SO_PATH.

587 "VLLM_CUDART_SO_PATH":

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

589 +

590 + # Contiguous cache fetching to avoid using costly gather operation on

591 + # Gaudi3. This is only applicable to HPU contiguous cache. If set to true,

592 + # contiguous cache fetch will be used.

593 + "VLLM_USE_HPU_CONTIGUOUS_CACHE_FETCH":

594 + lambda: os.environ.get("VLLM_CONTIGUOUS_PA", "true").lower() in

595 + ("1", "true"),

596 }

597

598 # end-env-vars-definition

vllm/worker/hpu_model_runner.py

+71 -43

@@ -25,9 +25,11 @@ import habana_frameworks.torch.internal.bridge_config as bc

25 import torch

26 import torch.nn as nn

27 from vllm_hpu_extension.ops import LoraMask as LoraMask

28 +from vllm_hpu_extension.ops import batch2block, block2batch

29 from vllm_hpu_extension.profiler import (HabanaHighLevelProfiler,

30 HabanaMemoryProfiler, format_bytes)

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.parallel_state import get_world_group

@@ -260,10 +262,19 @@ def setup_profiler():

262 return profiler

263

264

263 -def pad_list(list, k, v):

264 - target_len = round_up(len(list), k)

265 - padding = target_len - len(list)

266 - return list + [v] * padding

265 +def pad_list(input, k, v):

266 + input_len = len(input)

267 + target_len = round_up(input_len, k)

268 + padding = target_len - input_len

269 + return input + [v] * padding

270 +

271 +

272 +def gather_list(input, indices, v):

273 + return [input[i] if i is not None else v for i in indices]

274 +

275 +

276 +def flatten(in_list):

277 + return list(itertools.chain(*in_list))

278

279

280 def precompute_indices_and_offsets(block_size, slot_mapping, is_prompt):

@@ -334,13 +345,23 @@ class HpuModelAdapter:

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

346 attn_bias = (torch.zeros_like(mask, dtype=dtype).masked_fill_(

347 mask, -math.inf))

337 - block_mapping = torch.nn.functional.one_hot(metadata.block_mapping,

348 + block_mapping = torch.nn.functional.one_hot(metadata.block_groups,

349 num_classes=batch_size)

350 block_mapping = block_mapping.to(dtype)

351 metadata = metadata._replace(block_mapping=block_mapping,

352 attn_bias=attn_bias)

353 return metadata

354

355 + def _set_block_scales(self, metadata, device):

356 + block_mapping = metadata.block_mapping

357 + ones = torch.ones((block_mapping.size(0), ),

358 + device=device,

359 + dtype=block_mapping.dtype)

360 + sums = batch2block(block2batch(ones, block_mapping), block_mapping)

361 + block_scales = torch.reciprocal(torch.maximum(ones, sums))

362 + metadata = metadata._replace(block_scales=block_scales)

363 + return metadata

364 +

365 def _update_metadata(self, attn_metadata, batch_size, seq_len, device,

366 dtype):

367 if attn_metadata.is_prompt:

@@ -351,6 +372,7 @@ class HpuModelAdapter:

372 meta = attn_metadata

373 attn_metadata = self._set_block_mapping(meta, batch_size, device,

374 dtype)

375 + attn_metadata = self._set_block_scales(attn_metadata, device)

376 return attn_metadata

377

378 def forward(self, *args, **kwargs):

@@ -586,6 +608,7 @@ class HPUModelRunnerBase(ModelRunnerBase[TModelInputForHPU]):

608 self.bucketing_global_state = HPUBucketingGlobalState()

609 self._setup_buckets()

610 self._set_gc_threshold()

611 + self.use_contiguous_pa = envs.VLLM_USE_HPU_CONTIGUOUS_CACHE_FETCH

612

613 def _set_gc_threshold(self) -> None:

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

@@ -911,6 +934,7 @@ class HPUModelRunnerBase(ModelRunnerBase[TModelInputForHPU]):

934 block_indices=block_indices,

935 block_offsets=block_offsets,

936 block_scales=None,

937 + block_groups=None,

938 attn_bias=None,

939 seq_lens_tensor=seq_lens_tensor,

940 num_prefills=real_num_seqs,

@@ -1008,65 +1032,69 @@ class HPUModelRunnerBase(ModelRunnerBase[TModelInputForHPU]):

1032

1033 num_decode_tokens = sum(seq_lens)

1034

1011 - blocks_used = [len(bt) for bt in block_tables if bt]

1012 - block_list = []

1013 - block_scales = []

1014 - for i, bt in enumerate(block_tables):

1015 - block_list.extend(bt)

1016 - blocks_in_group = len(bt)

1017 - if blocks_in_group > 0:

1018 - scale = 1.0 / blocks_in_group

1019 - block_scales.extend([scale] * blocks_in_group)

1020 -

1021 - block_mapping_nested: List[List[int]] = [

1022 - [i] * b_u for i, b_u in enumerate(blocks_used)

1035 + last_block_usage = [

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

1037 ]

1024 - block_mapping: List[int] = list(

1025 - itertools.chain.from_iterable(block_mapping_nested))

1038 + block_groups = [[i] * len(bt) for i, bt in enumerate(block_tables)]

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

1040 + for bt, lbu in zip(block_tables, last_block_usage)

1041 + if bt]

1042 +

1043 + block_list = flatten(block_tables)

1044 + block_groups = flatten(block_groups)

1045 + block_usage = flatten(block_usage)

1046 +

1047 + assert len(block_list) == len(block_groups)

1048 + assert len(block_list) == len(block_usage)

1049 +

1050 + padding_fn = None

1051 + if self.use_contiguous_pa:

1052 + block_bucket_size = max(max(block_list) + 1, len(block_list))

1053 + block_bucket_size = find_bucket(

1054 + block_bucket_size,

1055 + self.bucketing_global_state.decode_block_bucket_cfg)

1056 + indices: List[Any]

1057 + indices = [None] * block_bucket_size

1058 + for i, bid in enumerate(block_list):

1059 + indices[bid] = i

1060 + padding_fn = lambda tensor, pad_value: gather_list(

1061 + tensor, indices, pad_value)

1062 + else:

1063 + block_bucket_size = find_bucket(

1064 + len(block_list),

1065 + self.bucketing_global_state.decode_block_bucket_cfg)

1066 + padding_fn = lambda tensor, pad_value: pad_list(

1067 + tensor, block_bucket_size, pad_value)

1068

1027 - last_block = [

1028 - sl % self.block_size + 1 for sl in itertools.chain(*slot_mapping)

1029 - ]

1030 - block_usage = [[self.block_size] * (b_u - 1) + [lb]

1031 - for b_u, lb in zip(blocks_used, last_block)]

1032 - block_usage = list(itertools.chain(*block_usage))

1033 -

1034 - block_bucket_size = find_bucket(

1035 - len(block_list),

1036 - self.bucketing_global_state.decode_block_bucket_cfg)

1037 - block_list = pad_list(block_list, block_bucket_size, _PAD_BLOCK_ID)

1038 - block_mapping = pad_list(block_mapping, block_bucket_size, -1)

1039 - block_usage = pad_list(block_usage, block_bucket_size, 1)

1040 - block_scales = pad_list(block_scales, block_bucket_size, 0.0)

1069 + block_list = padding_fn(block_list, _PAD_BLOCK_ID)

1070 + block_groups = padding_fn(block_groups, -1)

1071 + block_usage = padding_fn(block_usage, 1)

1072

1073 block_list = torch.tensor(block_list,

1074 dtype=torch.int,

1075 device=self.device)

1045 - block_mapping = torch.tensor(block_mapping,

1046 - dtype=torch.long,

1047 - device=self.device)

1076 + block_groups = torch.tensor(block_groups,

1077 + dtype=torch.int,

1078 + device=self.device)

1079 block_usage = torch.tensor(block_usage,

1080 dtype=self.model_config.dtype,

1081 device=self.device)

1051 -

1082 slot_mapping = torch.tensor(slot_mapping,

1083 dtype=torch.long,

1084 device=self.device)

1085

1086 block_indices, block_offsets = precompute_indices_and_offsets(

1087 self.block_size, slot_mapping, False)

1058 - block_scales = torch.tensor(block_scales,

1059 - dtype=self.model_config.dtype,

1060 - device=self.device)

1088

1089 attn_metadata = self.attn_backend.make_metadata(

1090 is_prompt=False,

1091 block_list=block_list,

1065 - block_mapping=block_mapping,

1092 + block_mapping=None,

1093 block_usage=block_usage,

1094 block_indices=block_indices,

1095 block_offsets=block_offsets,

1069 - block_scales=block_scales,

1096 + block_scales=None,

1097 + block_groups=block_groups,

1098 attn_bias=None,

1099 seq_lens_tensor=None,

1100 num_prefills=0,

@@ -1280,7 +1308,7 @@ class HPUModelRunnerBase(ModelRunnerBase[TModelInputForHPU]):

1308 attention_metadata = subtuple(metadata, 'TrimmedAttentionMetadata', [

1309 'attn_bias', 'seq_lens_tensor', 'block_list', 'block_mapping',

1310 'block_usage', 'slot_mapping', 'is_prompt', 'block_indices',

1283 - 'block_offsets', 'block_scales'

1311 + 'block_offsets', 'block_scales', 'block_groups'

1312 ])

1313 return attention_metadata

1314