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

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

[Hardware][Intel-Gaudi] Support Automatic Prefix Caching on HPU (#17648)

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

Agata Dobrzyniewicz committed May 8, 2025 at 07:37 UTC 843b222723b659e4b80d71d3ffb4944266af1d74

3 files changed +146 -56

vllm/attention/backends/hpu_attn.py

+20 -12

@@ -57,16 +57,16 @@ class HPUAttentionBackend(AttentionBackend):

57 def swap_blocks(

58 src_kv_cache: torch.Tensor,

59 dst_kv_cache: torch.Tensor,

60 - src_to_dst: Dict[int, int],

60 + src_to_dsts: torch.Tensor,

61 ) -> None:

62 - HPUPagedAttention.swap_blocks(src_kv_cache, dst_kv_cache, src_to_dst)

62 + HPUPagedAttention.swap_blocks(src_kv_cache, dst_kv_cache, src_to_dsts)

64 @staticmethod

65 def copy_blocks(

66 kv_caches: List[torch.Tensor],

67 - src_to_dists: Dict[int, List[int]],

67 + src_to_dsts: torch.Tensor,

68 ) -> None:

69 - HPUPagedAttention.copy_blocks(kv_caches, src_to_dists)

69 + HPUPagedAttention.copy_blocks(kv_caches, src_to_dsts)

72 @dataclass

@@ -77,6 +77,7 @@ class HPUAttentionMetadata(HPUPagedAttentionMetadata, AttentionMetadata):

77 is_prompt: bool

78 attn_bias: Optional[torch.Tensor]

79 seq_lens_tensor: Optional[torch.Tensor]

80 + context_lens_tensor: Optional[torch.Tensor]

83 class HPUAttentionImpl(AttentionImpl, torch.nn.Module):

@@ -198,8 +199,7 @@ class HPUAttentionImpl(AttentionImpl, torch.nn.Module):

199 key_cache = None

200 value_cache = None

201 if attn_metadata.is_prompt and self.attn_type \

201 - is not AttentionType.ENCODER_ONLY \

202 - and attn_metadata.block_list is None:

202 + is not AttentionType.ENCODER_ONLY:

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

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

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

@@ -229,6 +229,9 @@ class HPUAttentionImpl(AttentionImpl, torch.nn.Module):

229 attn_bias = attn_bias.tile((1, self.num_kv_heads, 1, 1))

230 attn_bias.add_(position_bias)

231

232 + block_list = attn_metadata.block_list if attn_metadata \

233 + and attn_metadata.block_list is not None else None

234 +

235 out = ops.prompt_attention(

236 impl=self.prefill_impl,

237 query=query.view(query_shape),

@@ -237,23 +240,25 @@ class HPUAttentionImpl(AttentionImpl, torch.nn.Module):

240 is_causal=True,

241 attn_bias=attn_bias,

242 valid_seq_lengths=attn_metadata.seq_lens_tensor,

240 - **self.common_attention_args())

243 + **self.common_attention_args(block_list, key_cache,

244 + value_cache))

245 output = out.reshape(batch_size, seq_len, hidden_size)

246 else:

247 # Decoding run.

248 output = HPUPagedAttention.forward_decode(

249 query=query,

246 - key_cache=key_cache,

247 - value_cache=value_cache,

248 - block_list=attn_metadata.block_list,

250 block_mapping=attn_metadata.block_mapping,

251 block_bias=attn_metadata.attn_bias,

252 block_groups=attn_metadata.block_groups,

252 - **self.common_attention_args())

253 + **self.common_attention_args(attn_metadata.block_list,

254 + key_cache, value_cache))

255 # Reshape the output tensor.

256 return output.view(batch_size, seq_len, hidden_size)

257

256 - def common_attention_args(self):

258 + def common_attention_args(self,

259 + block_list=None,

260 + key_cache=None,

261 + value_cache=None):

262 fsdpa_op = self.fused_scaled_dot_product_attention.apply \

263 if self.fused_scaled_dot_product_attention is not None else None

264 return {

@@ -266,6 +271,9 @@ class HPUAttentionImpl(AttentionImpl, torch.nn.Module):

271 'keys_fetch_func': self.k_cache.fetch_from_cache,

272 'values_fetch_func': self.v_cache.fetch_from_cache,

273 'softmax_op': self.softmax,

274 + 'block_list': block_list,

275 + 'key_cache': key_cache,

276 + 'value_cache': value_cache,

277 }

278

279

vllm/attention/ops/hpu_paged_attn.py

+9 -27

@@ -5,7 +5,7 @@

5 ###############################################################################

7 from dataclasses import dataclass

8 -from typing import Dict, List, Optional, Tuple

8 +from typing import List, Optional, Tuple

10 import torch

11 from vllm_hpu_extension import cache_ops, ops

@@ -63,43 +63,25 @@ class HPUPagedAttention:

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

64 return ops.flat_pa(**kwargs)

66 - @staticmethod

67 - def forward_prefix(

68 - query: torch.Tensor,

69 - key: torch.Tensor,

70 - value: torch.Tensor,

71 - key_cache: torch.Tensor,

72 - value_cache: torch.Tensor,

73 - block_tables: torch.Tensor,

74 - subquery_start_loc: torch.Tensor,

75 - seq_lens_tensor: torch.Tensor,

76 - context_lens: torch.Tensor,

77 - max_query_len: int,

78 - alibi_slopes: Optional[torch.Tensor],

79 - sliding_window: Optional[int],

80 - ) -> torch.Tensor:

81 - raise NotImplementedError(

82 - "forward_prefix is not implemented for HPUPagedAttention")

83 -

66 @staticmethod

67 def swap_blocks(

86 - src_kv_cache: torch.Tensor,

87 - dst_kv_cache: torch.Tensor,

88 - src_to_dst: Dict[int, int],

68 + src_kv_cache: Tuple[torch.Tensor, torch.Tensor],

69 + dst_kv_cache: Tuple[torch.Tensor, torch.Tensor],

70 + src_to_dsts: torch.Tensor,

71 ) -> None:

72 src_key_cache = src_kv_cache[0]

73 dst_key_cache = dst_kv_cache[0]

92 - cache_ops.swap_blocks(src_key_cache, dst_key_cache, src_to_dst)

74 + cache_ops.swap_blocks(src_key_cache, dst_key_cache, src_to_dsts)

76 src_value_cache = src_kv_cache[1]

77 dst_value_cache = dst_kv_cache[1]

96 - cache_ops.swap_blocks(src_value_cache, dst_value_cache, src_to_dst)

78 + cache_ops.swap_blocks(src_value_cache, dst_value_cache, src_to_dsts)

80 @staticmethod

81 def copy_blocks(

100 - kv_caches: List[torch.Tensor],

101 - src_to_dists: Dict[int, List[int]],

82 + kv_caches: List[Tuple[torch.Tensor, torch.Tensor]],

83 + src_to_dsts: torch.Tensor,

84 ) -> None:

85 key_caches = [kv_cache[0] for kv_cache in kv_caches]

86 value_caches = [kv_cache[1] for kv_cache in kv_caches]

105 - cache_ops.copy_blocks(key_caches, value_caches, src_to_dists)

87 + cache_ops.copy_blocks(key_caches, value_caches, src_to_dsts)

vllm/worker/hpu_model_runner.py

+117 -17

@@ -14,7 +14,7 @@ import math

14 import os

15 import time

16 from array import array

17 -from enum import IntEnum

17 +from enum import Enum, IntEnum

18 from typing import (TYPE_CHECKING, Any, Callable, Dict, List, NamedTuple,

19 Optional, Set, Tuple, Type, TypeVar, Union)

@@ -75,6 +75,12 @@ LORA_WARMUP_RANK = 8

75 DUMMY_TOKEN_ID = -1

78 +class PhaseType(Enum):

79 + PREFILL = 'prefill'

80 + PREFIX_PREFILL = 'prefix_prefill'

81 + DECODE = 'decode'

82 +

83 +

84 def subtuple(obj: object,

85 typename: str,

86 to_copy: List[str],

@@ -213,20 +219,40 @@ class HpuModelAdapter:

219

220 def _set_attn_bias(self, attn_metadata, batch_size, seq_len, device,

221 dtype):

216 - prefill_metadata = attn_metadata

217 - if prefill_metadata is None or self.prefill_use_fusedsdpa:

222 + if (attn_metadata is None

223 + or (self.prefill_use_fusedsdpa \

224 + and attn_metadata.block_list is None)

225 + or not attn_metadata.is_prompt):

226 return attn_metadata

227

228 + prefill_metadata = attn_metadata

229 +

230 seq_lens_t = prefill_metadata.seq_lens_tensor

231 + context_lens_t = prefill_metadata.context_lens_tensor

232 + query_lens_t = seq_lens_t - context_lens_t

233 +

234 + block_list = attn_metadata.block_list

235 + max_context_len = (block_list.size(-1) //

236 + batch_size if block_list is not None else 0)

237 + max_context_len = max_context_len * self.block_size

238 + past_mask = torch.arange(0,

239 + max_context_len,

240 + dtype=torch.int32,

241 + device=device)

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

243 + context_lens_t.view(-1, 1)).view(batch_size, 1, -1).expand(

244 + batch_size, seq_len, -1).view(batch_size, 1, seq_len, -1))

245 +

246 len_mask = (torch.arange(0, seq_len, device=device,

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

223 - seq_lens_t.unsqueeze(-1)).view(

248 + query_lens_t.unsqueeze(-1)).view(

249 batch_size, 1, 1, seq_len))

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

251 device=device,

252 dtype=torch.bool),

253 diagonal=1)

254 mask = causal_mask.logical_or(len_mask)

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

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

257 mask, -math.inf))

258 attn_metadata = prefill_metadata._replace(attn_bias=attn_bias)

@@ -517,6 +543,11 @@ class HPUModelRunnerBase(ModelRunnerBase[TModelInputForHPU]):

543 False, self.max_model_len)

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

545 self._set_gc_threshold()

546 + if self.vllm_config.cache_config.enable_prefix_caching:

547 + os.environ.setdefault("VLLM_CONTIGUOUS_PA", "False")

548 + assert os.environ.get(

549 + "VLLM_CONTIGUOUS_PA",

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

551 self.use_contiguous_pa = envs.VLLM_USE_HPU_CONTIGUOUS_CACHE_FETCH

552

553 # For multi-step scheduling

@@ -702,6 +733,10 @@ class HPUModelRunnerBase(ModelRunnerBase[TModelInputForHPU]):

733 computed_block_nums) > 0 and self.sliding_window is None:

734 # Prefix is not supported with sliding_window

735 context_len = len(computed_block_nums) * self.block_size

736 + if context_len == seq_len \

737 + and self.vllm_config.cache_config.enable_prefix_caching:

738 + # Fully cached prompt - compute only last token

739 + context_len = context_len - 1

740 prompt_tokens = prompt_tokens[context_len:]

741 prefix_block_tables.append(computed_block_nums)

742 elif self.scheduler_config.chunked_prefill_enabled:

@@ -779,12 +814,33 @@ class HPUModelRunnerBase(ModelRunnerBase[TModelInputForHPU]):

814 if lora_id > 0:

815 lora_requests.add(seq_group_metadata.lora_request)

816

782 - lora_index_mapping += [lora_id] * (max_prompt_len - context_len)

817 + lora_index_mapping += [lora_id] * max_prompt_len

818 lora_prompt_mapping.extend(

819 [lora_id] *

785 - (max_prompt_len - context_len

820 + (max_prompt_len

821 if seq_group_metadata.sampling_params.prompt_logprobs else 1))

822

823 + if any(context_lens):

824 + assert not self.scheduler_config.chunked_prefill_enabled

825 + # prefix caching

826 +

827 + max_num_block = max(len(bt) for bt in prefix_block_tables)

828 + prefix_block_list = list(

829 + itertools.chain.from_iterable(

830 + bt if len(bt) == max_num_block else bt +

831 + ([_PAD_BLOCK_ID] * (max_num_block - len(bt)))

832 + for bt in prefix_block_tables))

833 +

834 + pad_len = len(prefix_block_list)

835 + prefix_block_list = pad_list(prefix_block_list, pad_len,

836 + _PAD_BLOCK_ID)

837 +

838 + prefix_block_list_tensor = torch.tensor(prefix_block_list,

839 + dtype=torch.long,

840 + device=self.device)

841 + else:

842 + prefix_block_list_tensor = None

843 +

844 input_tokens = make_tensor_with_pad(input_tokens,

845 max_len=max_prompt_len,

846 pad=0,

@@ -807,11 +863,15 @@ class HPUModelRunnerBase(ModelRunnerBase[TModelInputForHPU]):

863 dtype=torch.long,

864 device=self.device)

865

866 + context_lens_tensor = torch.tensor(context_lens,

867 + dtype=torch.long,

868 + device=self.device)

869 +

870 block_indices, block_offsets = precompute_indices_and_offsets(

871 self.block_size, slot_mapping, True)

872 attn_metadata = self.attn_backend.make_metadata(

873 is_prompt=True,

814 - block_list=None,

874 + block_list=prefix_block_list_tensor,

875 block_mapping=None,

876 block_usage=None,

877 block_indices=block_indices,

@@ -819,6 +879,7 @@ class HPUModelRunnerBase(ModelRunnerBase[TModelInputForHPU]):

879 block_groups=None,

880 attn_bias=None,

881 seq_lens_tensor=seq_lens_tensor,

882 + context_lens_tensor=context_lens_tensor,

883 num_prefills=real_num_seqs,

884 num_prefill_tokens=sum_query_len,

885 num_decode_tokens=0,

@@ -987,6 +1048,7 @@ class HPUModelRunnerBase(ModelRunnerBase[TModelInputForHPU]):

1048 block_groups=block_groups,

1049 attn_bias=None,

1050 seq_lens_tensor=None,

1051 + context_lens_tensor=None,

1052 num_prefills=0,

1053 num_prefill_tokens=0,

1054 num_decode_tokens=num_decode_tokens,

@@ -1091,7 +1153,7 @@ class HPUModelRunnerBase(ModelRunnerBase[TModelInputForHPU]):

1153 # FIXME: We need to adjust selected_token_indices to accommodate

1154 # for padding

1155 max_len = input_tokens.size(1)

1094 - paddings = [max_len - s for s in seq_lens]

1156 + paddings = [max_len - q for q in query_lens]

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

1158 paddings = list(itertools.accumulate(paddings))

1159 paddings_prompt_logprobs = []

@@ -1187,9 +1249,17 @@ class HPUModelRunnerBase(ModelRunnerBase[TModelInputForHPU]):

1249 # input_hash(123) != input_hash(321)

1250 # input_hash("abc") != input_hash("cba")

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

1190 - 'attn_bias', 'seq_lens_tensor', 'block_list', 'block_mapping',

1191 - 'block_usage', 'slot_mapping', 'is_prompt', 'block_indices',

1192 - 'block_offsets', 'block_groups'

1252 + 'attn_bias',

1253 + 'seq_lens_tensor',

1254 + 'context_lens_tensor',

1255 + 'block_list',

1256 + 'block_mapping',

1257 + 'block_usage',

1258 + 'slot_mapping',

1259 + 'is_prompt',

1260 + 'block_indices',

1261 + 'block_offsets',

1262 + 'block_groups',

1263 ])

1264 return attention_metadata

1265

@@ -1733,14 +1803,44 @@ class HPUModelRunner(HPUModelRunnerBase[ModelInputForHPUWithSamplingMetadata]):

1803 from neural_compressor.torch.quantization import finalize_calibration

1804 finalize_calibration(self.model.model)

1805

1736 - def _check_config(self, batch_size, seq_len, is_prompt, warmup_mode):

1737 - cfg = (batch_size, seq_len, is_prompt)

1806 + def _num_blocks(self, attn_metadata):

1807 + if attn_metadata.block_list is None:

1808 + return 0

1809 + return attn_metadata.block_list.numel()

1810 +

1811 + def _phase(self, attn_metadata):

1812 + phase_type: PhaseType

1813 + is_prompt = attn_metadata.is_prompt

1814 + is_prefix_prefill = is_prompt and attn_metadata.block_list is not None

1815 + if is_prompt and is_prefix_prefill:

1816 + phase_type = PhaseType.PREFIX_PREFILL

1817 + elif is_prompt and not is_prefix_prefill:

1818 + phase_type = PhaseType.PREFILL

1819 + elif not is_prompt:

1820 + phase_type = PhaseType.DECODE

1821 + else:

1822 + raise ValueError("Unrecognized pass type, likely due to malformed "

1823 + "attention metadata")

1824 + return phase_type

1825 +

1826 + def _check_config(self, batch_size, seq_len, attn_metadata, warmup_mode):

1827 + is_prefix_caching = self.vllm_config.cache_config.enable_prefix_caching

1828 + cfg: Optional[tuple] = None

1829 + assert cfg is None, "Configs changed between 2D and 3D"

1830 + if is_prefix_caching:

1831 + phase = self._phase(attn_metadata)

1832 + num_blocks = self._num_blocks(attn_metadata)

1833 + cfg = (batch_size, seq_len, num_blocks, phase)

1834 + else:

1835 + phase = 'prompt' if attn_metadata.is_prompt else 'decode'

1836 + cfg = (batch_size, seq_len, phase)

1837 seen = cfg in self.seen_configs

1838 self.seen_configs.add(cfg)

1839 if not seen and not warmup_mode:

1741 - phase = 'prompt' if is_prompt else 'decode'

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

1743 - phase, batch_size, seq_len)

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

1841 + (phase.value, batch_size, seq_len,

1842 + num_blocks) if is_prefix_caching else

1843 + (phase, batch_size, seq_len))

1844

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

1846 is_prompt: bool):

@@ -1912,7 +2012,7 @@ class HPUModelRunner(HPUModelRunnerBase[ModelInputForHPUWithSamplingMetadata]):

2012 batch_size = input_tokens.size(0)

2013 seq_len = self._seq_len(attn_metadata)

2014 use_graphs = self._use_graphs(batch_size, seq_len, is_prompt)

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

2015 + self._check_config(batch_size, seq_len, attn_metadata, warmup_mode)

2016

2017 lora_mask: torch.Tensor = None

2018 lora_logits_mask: torch.Tensor = None