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

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

[V1][TPU] Integrate the new ragged paged attention kernel with vLLM v1 on TPU (#13379)

Signed-off-by: Xiongfei Wei <isaacwxf23@gmail.com> Signed-off-by: mgoin <mgoin64@gmail.com> Co-authored-by: mgoin <mgoin64@gmail.com>

iefgnoix committed Feb 28, 2025 at 10:01 UTC c3b6559a1019c73f5b2436d2d50cbaa3663e5382

6 files changed +353 -905

requirements-tpu.txt

+5 -6

@@ -17,9 +17,8 @@ ray[default]

17 --find-links https://storage.googleapis.com/libtpu-releases/index.html

18 --find-links https://storage.googleapis.com/jax-releases/jax_nightly_releases.html

19 --find-links https://storage.googleapis.com/jax-releases/jaxlib_nightly_releases.html

20 -torch @ https://download.pytorch.org/whl/nightly/cpu/torch-2.6.0.dev20241216%2Bcpu-cp39-cp39-linux_x86_64.whl ; python_version == "3.9"

21 -torch @ https://download.pytorch.org/whl/nightly/cpu/torch-2.6.0.dev20241216%2Bcpu-cp310-cp310-linux_x86_64.whl ; python_version == "3.10"

22 -torch @ https://download.pytorch.org/whl/nightly/cpu/torch-2.6.0.dev20241216%2Bcpu-cp311-cp311-linux_x86_64.whl ; python_version == "3.11"

23 -torch_xla[tpu, pallas] @ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.7.0.dev20250124-cp39-cp39-linux_x86_64.whl ; python_version == "3.9"

24 -torch_xla[tpu, pallas] @ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.7.0.dev20250124-cp310-cp310-linux_x86_64.whl ; python_version == "3.10"

25 -torch_xla[tpu, pallas] @ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.7.0.dev20250124-cp311-cp311-linux_x86_64.whl ; python_version == "3.11"

20 +

21 +torch==2.7.0.dev20250226+cpu

22 +torch_xla[tpu, pallas] @ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.7.0.dev20250226+cxx11-cp39-cp39-linux_x86_64.whl ; python_version == "3.9"

23 +torch_xla[tpu, pallas] @ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.7.0.dev20250226+cxx11-cp310-cp310-linux_x86_64.whl ; python_version == "3.10"

24 +torch_xla[tpu, pallas] @ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.7.0.dev20250226+cxx11-cp311-cp311-linux_x86_64.whl ; python_version == "3.11"

vllm/v1/attention/backends/pallas.py

+62 -218

@@ -4,13 +4,16 @@ from dataclasses import dataclass

4 from typing import Any, Dict, List, Optional, Tuple, Type

6 import torch

7 -import torch_xla.experimental.custom_kernel # Required to register custom ops.

7 +# Required to register custom ops.

8 +import torch_xla.experimental.custom_kernel # noqa: F401

10 from vllm.attention.backends.abstract import (AttentionBackend, AttentionImpl,

10 - AttentionLayer,

11 - AttentionMetadata, AttentionType)

11 + AttentionLayer, AttentionType)

12 from vllm.attention.backends.utils import CommonAttentionState

14 +NUM_QUERIES_PER_BLOCK = 16

15 +NUM_KV_PAGES_PER_BLOCK = 128

16 +

18 class PallasAttentionBackend(AttentionBackend):

@@ -47,47 +50,23 @@ class PallasAttentionBackend(AttentionBackend):

50 ) -> None:

51 raise RuntimeError("swap_blocks is not used for the TPU backend.")

50 - @torch.compile(backend="openxla")

51 - @staticmethod

52 - def copy_blocks(

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

54 - src_to_dists: Tuple[torch.Tensor, torch.Tensor],

55 - ) -> None:

56 - src_indices, dst_indices = src_to_dists

57 - for k_cache, v_cache in kv_caches:

58 - torch.ops.xla.dynamo_set_buffer_donor_(k_cache, True)

59 - k_cache[:, dst_indices] = k_cache[:, src_indices]

60 - torch.ops.xla.dynamo_set_buffer_donor_(v_cache, True)

61 - v_cache[:, dst_indices] = v_cache[:, src_indices]

62 -

54 @dataclass

65 -class PallasMetadata(AttentionMetadata):

66 -

67 - # Currently, input sequences can only contain all prefills

68 - # or all decoding.

69 - block_tables: Optional[torch.Tensor] = None

70 - context_lens: Optional[torch.Tensor] = None

71 - effective_query_lens: Optional[torch.Tensor] = None

72 -

73 - @property

74 - def prefill_metadata(self) -> Optional["PallasMetadata"]:

75 - if self.num_prefills == 0:

76 - return None

77 -

78 - assert self.num_decode_tokens == 0

79 - return self

80 -

81 - @property

82 - def decode_metadata(self) -> Optional["PallasMetadata"]:

83 - if self.num_decode_tokens == 0:

84 - return None

85 -

86 - assert self.num_prefills == 0

87 - assert self.num_prefill_tokens == 0

88 - assert self.block_tables is not None

89 - assert self.context_lens is not None

90 - return self

55 +class PallasMetadata:

56 + # NOTE(sang): Definition of context_len, query_len, and seq_len.

57 + # |---------- N-1 iteration --------|

58 + # |---------------- N iteration ---------------------|

59 + # |- tokenA -|......................|-- newTokens ---|

60 + # |---------- context_len ----------|

61 + # |-------------------- seq_len ---------------------|

62 + # |-- query_len ---|

63 +

64 + # Used in the PallasAttentionBackendImpl

65 + slot_mapping: torch.Tensor

66 + block_tables: torch.Tensor

67 + context_lens: torch.Tensor

68 + query_start_loc: torch.Tensor

69 + num_seqs: int

72 class PallasAttentionBackendImpl(AttentionImpl):

@@ -105,10 +84,13 @@ class PallasAttentionBackendImpl(AttentionImpl):

84 logits_soft_cap: Optional[float] = None,

85 attn_type: str = AttentionType.DECODER,

86 ) -> None:

87 + if blocksparse_params is not None:

88 + raise ValueError("Paged attention Pallas kernel does "

89 + "not support block-sparse attention.")

90 self.num_heads = num_heads

91 self.head_size = head_size

92 self.scale = float(scale)

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

93 + self.num_kv_heads = num_kv_heads

95 assert self.num_heads % self.num_kv_heads == 0

96 self.num_queries_per_kv = self.num_heads // self.num_kv_heads

@@ -126,25 +108,6 @@ class PallasAttentionBackendImpl(AttentionImpl):

108 raise NotImplementedError(

109 "Attention logits soft-capping is not supported.")

110

129 - if torch_xla.tpu.version() < 4:

130 - raise NotImplementedError("TPU version must be 4 or higher.")

131 -

132 - self.megacore_mode = None

133 - tpu_env = torch_xla.tpu.get_tpu_env()

134 - tpu_type = (tpu_env.get("ACCELERATOR_TYPE", None)

135 - or tpu_env.get("TYPE", None)

136 - or tpu_env.get("TPU_ACCELERATOR_TYPE", None))

137 - assert tpu_type is not None

138 - tpu_type = tpu_type.lower()

139 -

140 - if (("lite" not in tpu_type) and ("v6" not in tpu_type)):

141 - if self.num_kv_heads % 2 == 0:

142 - self.megacore_mode = "kv_head"

143 - else:

144 - # NOTE(woosuk): If the batch size is not a multiple of 2, the

145 - # megacore mode will be None.

146 - self.megacore_mode = "batch"

147 -

111 if attn_type != AttentionType.DECODER:

112 raise NotImplementedError("Encoder self-attention and "

113 "encoder/decoder cross-attention "

@@ -164,135 +127,47 @@ class PallasAttentionBackendImpl(AttentionImpl):

127 """Forward pass with Pallas attention.

128

129 Args:

167 - query: shape = [batch_size, seq_len, num_heads * head_size]

168 - key: shape = [batch_size, seq_len, num_kv_heads * head_size]

169 - value: shape = [batch_size, seq_len, num_kv_heads * head_size]

170 - kv_cache[0] = [num_kv_heads, num_blocks, block_size, head_size]

171 - kv_cache[1] = [num_kv_heads, num_blocks, block_size, head_size]

172 - NOTE: kv_cache[0] and kv_cache[1] will be an empty tensor

173 - with shape [0] for profiling run.

130 + query: shape = [num_tokens, num_heads * head_size]

131 + key: shape = [num_tokens, num_kv_heads * head_size]

132 + value: shape = [num_tokens, num_kv_heads * head_size]

133 + kv_cache = ([num_kv_heads, num_blocks, block_size, head_size],

134 + [num_kv_heads, num_blocks, block_size, head_size])

135 attn_metadata: Metadata for attention.

136 Returns:

176 - shape = [batch_size, seq_len, num_heads * head_size]

137 + shape = [num_tokens, num_heads * head_size]

138 """

178 -

179 - if attn_metadata is None:

139 + # For determine_available_memory case.

140 + if kv_cache[0].numel() == 0:

141 if output is None:

142 output = torch.ones_like(query)

143 return output

144

145 assert layer._k_scale_float == 1.0 and layer._v_scale_float == 1.0

185 - batch_size, seq_len, hidden_size = query.shape

186 - query = query.view(batch_size, seq_len, self.num_heads, self.head_size)

187 - key = key.view(batch_size, seq_len, self.num_kv_heads, self.head_size)

188 - value = value.view(batch_size, seq_len, self.num_kv_heads,

189 - self.head_size)

146 + num_tokens, hidden_size = query.shape

147 + query = query.view(num_tokens, self.num_heads, self.head_size)

148 + key = key.view(num_tokens, self.num_kv_heads, self.head_size)

149 + value = value.view(num_tokens, self.num_kv_heads, self.head_size)

150

151 + key_cache, value_cache = kv_cache

152 if kv_cache[0].numel() > 0:

153 slot_mapping = attn_metadata.slot_mapping

193 - key_cache, value_cache = kv_cache

154 write_to_kv_cache(key, value, key_cache, value_cache, slot_mapping)

155

156 query = query * self.scale

197 - if attn_metadata.num_prefills > 0:

198 - if attn_metadata.block_tables is None:

199 - # Prefill without paged KV cache.

200 - assert seq_len % 16 == 0, (

201 - "Pallas FlashAttention kernel requires seq_len to be a "

202 - f"multiple of 16 but got {seq_len}")

203 -

204 - # Handle GQA/MQA.

205 - if self.num_kv_heads != self.num_heads:

206 - key = key.repeat_interleave(self.num_queries_per_kv,

207 - dim=-2)

208 - key = key.view(batch_size, seq_len, self.num_heads,

209 - self.head_size)

210 - value = value.repeat_interleave(self.num_queries_per_kv,

211 - dim=-2)

212 - value = value.view(batch_size, seq_len, self.num_heads,

213 - self.head_size)

214 - # FlashAttention kernel requires the input shape to be

215 - # [batch_size, num_heads, seq_len, d_model]

216 - # while the input is [batch_size, seq_len, num_heads, d_model].

217 - # Permute the input to match the required format.

218 - output = torch.ops.xla.flash_attention(

219 - query.permute(0, 2, 1, 3),

220 - key.permute(0, 2, 1, 3),

221 - value.permute(0, 2, 1, 3),

222 - True,

223 - )

224 - output = output.permute(0, 2, 1, 3)

225 - else:

226 - # Prefill with paged KV cache.

227 - # TODO(woosuk): Tune the below knobs.

228 - num_kv_pages_per_compute_block = 16

229 - num_queries_per_compute_block = 16

230 - assert seq_len % num_queries_per_compute_block == 0

231 - output = torch.ops.xla.multi_queries_paged_attention(

232 - query,

233 - key_cache,

234 - value_cache,

235 - attn_metadata.context_lens,

236 - attn_metadata.block_tables,

237 - attn_metadata.effective_query_lens,

238 - num_kv_pages_per_compute_block,

239 - num_queries_per_compute_block,

240 - use_kernel=True,

241 - )

242 - else:

243 - # Decoding run.

244 - assert kv_cache[0].numel() > 0

245 - query = query.squeeze(dim=1)

246 - pages_per_compute_block = 16 # TODO(woosuk): Tune this value.

247 -

248 - assert attn_metadata.block_tables is not None

249 - assert attn_metadata.context_lens is not None

250 - # NOTE(woosuk): The PagedAttention Pallas kernel stores the entire

251 - # block table in SMEM. Therefore, if the block table is too large,

252 - # the kernel compilation will fail. To avoid this, we split the

253 - # batch dimension into smaller chunks and run the kernel multiple

254 - # times.

255 - MAX_SMEM_USAGE = 512 * 1024

256 - size_per_seq = 4 * attn_metadata.block_tables.shape[1]

257 - max_num_seq = MAX_SMEM_USAGE // size_per_seq

258 -

259 - if batch_size <= max_num_seq:

260 - output = paged_attention(

261 - query,

262 - key_cache,

263 - value_cache,

264 - attn_metadata.context_lens,

265 - attn_metadata.block_tables,

266 - pages_per_compute_block,

267 - self.megacore_mode,

268 - )

269 - else:

270 - chunk_size = max_num_seq

271 - # Make sure the chunk size is a multiple of 2.

272 - chunk_size = chunk_size // 2 * 2

273 - num_chunks = (batch_size + chunk_size - 1) // chunk_size

274 -

275 - output = torch.empty_like(query)

276 - for chunk_idx in range(num_chunks):

277 - chunk_start = chunk_idx * chunk_size

278 - chunk_end = chunk_start + chunk_size

279 - # NOTE(woosuk): We skip this line because it causes Dynamo

280 - # compilation error. Instead, we rely on the slice operation

281 - # to handle the out-of-bound case.

282 - # chunk_end = min(chunk_end, batch_size)

283 - chunk_output = paged_attention(

284 - query[chunk_start:chunk_end],

285 - key_cache,

286 - value_cache,

287 - attn_metadata.context_lens[chunk_start:chunk_end],

288 - attn_metadata.block_tables[chunk_start:chunk_end],

289 - pages_per_compute_block,

290 - self.megacore_mode,

291 - )

292 - output[chunk_start:chunk_end] = chunk_output

157 + output = torch.ops.xla.ragged_paged_attention(

158 + query,

159 + key_cache,

160 + value_cache,

161 + attn_metadata.context_lens,

162 + attn_metadata.block_tables,

163 + attn_metadata.query_start_loc,

164 + attn_metadata.num_seqs,

165 + num_kv_pages_per_block=NUM_KV_PAGES_PER_BLOCK,

166 + num_queries_per_block=NUM_QUERIES_PER_BLOCK,

167 + use_kernel=False,

168 + )

169

294 - # Reshape the output tensor.

295 - return output.reshape(batch_size, seq_len, hidden_size)

170 + return output.reshape(num_tokens, hidden_size)

171

172

173 def write_to_kv_cache(

@@ -302,52 +177,21 @@ def write_to_kv_cache(

177 value_cache: torch.Tensor,

178 slot_mapping: torch.Tensor,

179 ) -> None:

180 + """ Write the key and values to the KV cache.

181 +

182 + Args:

183 + key: shape = [num_tokens, num_kv_heads, head_size]

184 + value: shape = [num_tokens, num_kv_heads, head_size]

185 + k_cache = [num_kv_heads, num_blocks, block_size, head_size]

186 + v_cache = [num_kv_heads, num_blocks, block_size, head_size]

187 +

188 + """

189 torch.ops.xla.dynamo_set_buffer_donor_(key_cache, True)

190 torch.ops.xla.dynamo_set_buffer_donor_(value_cache, True)

191

308 - key = key.flatten(0, 2)

309 - value = value.flatten(0, 2)

192 + key = key.flatten(0, 1)

193 + value = value.flatten(0, 1)

194 key_cache = key_cache.flatten(0, 2)

195 value_cache = value_cache.flatten(0, 2)

196 key_cache.index_copy_(0, slot_mapping, key)

197 value_cache.index_copy_(0, slot_mapping, value)

314 -

315 -

316 -def paged_attention(

317 - query: torch.Tensor,

318 - key_cache: torch.Tensor,

319 - value_cache: torch.Tensor,

320 - context_lens: torch.Tensor,

321 - block_tables: torch.Tensor,

322 - pages_per_compute_block: int,

323 - megacore_mode: Optional[str],

324 -) -> torch.Tensor:

325 - batch_size = query.shape[0]

326 - if megacore_mode == "batch" and batch_size % 2 != 0:

327 - megacore_mode = None

328 - else:

329 - megacore_mode = megacore_mode

330 -

331 - # NOTE(woosuk): A temporary workaround to avoid the error:

332 - # "xla::paged_attention() Expected a value of type 'str' for

333 - # argument 'megacore_mode' but instead found type 'NoneType'."

334 - if megacore_mode is not None:

335 - output = torch.ops.xla.paged_attention(

336 - query,

337 - key_cache,

338 - value_cache,

339 - context_lens,

340 - block_tables,

341 - pages_per_compute_block,

342 - megacore_mode=megacore_mode,

343 - )

344 - else:

345 - output = torch.ops.xla.paged_attention(

346 - query,

347 - key_cache,

348 - value_cache,

349 - context_lens,

350 - block_tables,

351 - pages_per_compute_block,

352 - )

353 - return output

vllm/v1/outputs.py

+1 -1

@@ -79,4 +79,4 @@ class ModelRunnerOutput:

79 # [prompt_len, num_prompt_logprobs]

80 # [prompt_len, num_prompt_logprobs]

81 # [prompt_len]

82 - prompt_logprobs_dict: Dict[str, LogprobsTensors]

82 + prompt_logprobs_dict: Dict[str, Optional[LogprobsTensors]]

vllm/v1/worker/gpu_model_runner.py

+2 -2

@@ -1071,12 +1071,12 @@ class GPUModelRunner(LoRAModelRunnerMixin):

1071 self,

1072 hidden_states: torch.Tensor,

1073 scheduler_output: "SchedulerOutput",

1074 - ) -> Dict[str, LogprobsTensors]:

1074 + ) -> Dict[str, Optional[LogprobsTensors]]:

1075 num_prompt_logprobs_dict = self.input_batch.num_prompt_logprobs

1076 if not num_prompt_logprobs_dict:

1077 return {}

1078

1079 - prompt_logprobs_dict: Dict[str, LogprobsTensors] = {}

1079 + prompt_logprobs_dict: Dict[str, Optional[LogprobsTensors]] = {}

1080

1081 # Since prompt logprobs are a rare feature, prioritize simple,

1082 # maintainable loop over optimal performance.

vllm/v1/worker/tpu_model_runner.py

+281 -674

@@ -1,8 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

2 -import enum

2 import time

4 -from dataclasses import dataclass

5 -from typing import TYPE_CHECKING, Dict, List, Optional, Tuple

3 +from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, cast

4 from unittest.mock import patch

6 import numpy as np

@@ -21,7 +19,9 @@ from vllm.logger import init_logger

19 from vllm.model_executor.model_loader import get_model

20 from vllm.sampling_params import SamplingType

21 from vllm.utils import LayerBlockType, cdiv, is_pin_memory_available

24 -from vllm.v1.attention.backends.pallas import (PallasAttentionBackend,

22 +from vllm.v1.attention.backends.pallas import (NUM_KV_PAGES_PER_BLOCK,

23 + NUM_QUERIES_PER_BLOCK,

24 + PallasAttentionBackend,

25 PallasMetadata)

26 from vllm.v1.kv_cache_interface import (FullAttentionSpec, KVCacheConfig,

27 KVCacheSpec)

@@ -37,36 +37,7 @@ logger = init_logger(__name__)

37 # Here we utilize the behavior that out-of-bound index is ignored.

38 # FIXME(woosuk): Find a more reliable way to prevent possible bugs.

39 _PAD_SLOT_ID = 1_000_000_000

40 -

41 -

42 -class ExecutionMode(enum.Enum):

43 - PREFILL = enum.auto()

44 - DECODE = enum.auto()

45 - PREFIX_PREFILL = enum.auto()

46 -

47 - def is_prefill(self) -> bool:

48 - return self in (ExecutionMode.PREFILL, ExecutionMode.PREFIX_PREFILL)

49 -

50 -

51 -@dataclass

52 -class PromptDecodeInfo:

53 - prompt_req_ids: List[str]

54 - decode_req_ids: List[str]

55 - prompt_scheduled_tokens: List[int]

56 -

57 -

58 -@dataclass

59 -class PromptData:

60 - input_tokens: torch.Tensor

61 - input_positions: torch.Tensor

62 - attn_metadata: PallasMetadata

63 -

64 -

65 -@dataclass

66 -class DecodeData:

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

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

69 - attn_metadata: Optional[PallasMetadata] = None

40 +INVALID_TOKEN_ID = -1

43 class TPUModelRunner:

@@ -113,8 +84,6 @@ class TPUModelRunner:

84 self.head_size = model_config.get_head_size()

85 self.hidden_size = model_config.get_hidden_size()

116 - self.model: Optional[nn.Module] = None

117 -

87 # Persistent batch.

88 self.input_batch = InputBatch(

89 max_num_reqs=self.max_num_reqs,

@@ -134,50 +103,48 @@ class TPUModelRunner:

103 # KV caches for forward pass

104 self.kv_caches: List[Tuple[torch.Tensor, torch.Tensor]] = []

105

137 - # Cached torch/numpy tensors

138 - self.num_swaps = 2

139 - self.cur_swap_id = 0

140 - self.input_ids_cpu = []

141 - self.input_ids_np = []

142 - self.input_positions_cpu = []

143 - self.input_positions_np = []

144 - self.slot_mapping_cpu = []

145 - self.slot_mapping_np = []

146 - self.prompt_context_lens_cpu = []

147 - self.prompt_effective_query_lens_cpu = []

148 - self.decode_context_lens_cpu = []

149 - self.decode_context_lens_np = []

150 - for _ in range(self.num_swaps):

151 - self.input_ids_cpu.append(

152 - torch.empty(self.max_num_tokens,

153 - dtype=torch.int32,

154 - device="cpu"))

155 - self.input_ids_np.append(self.input_ids_cpu[-1].numpy())

156 -

157 - self.input_positions_cpu.append(

158 - torch.empty(self.max_num_tokens,

159 - dtype=torch.int32,

160 - device="cpu"))

161 - self.input_positions_np.append(

162 - self.input_positions_cpu[-1].numpy())

163 -

164 - self.slot_mapping_cpu.append(

165 - torch.empty(self.max_num_tokens,

166 - dtype=torch.int64,

167 - device="cpu"))

168 - self.slot_mapping_np.append(self.slot_mapping_cpu[-1].numpy())

169 -

170 - self.prompt_context_lens_cpu.append(

171 - torch.empty((1), dtype=torch.int32, device="cpu"))

172 - self.prompt_effective_query_lens_cpu.append(

173 - torch.empty((1), dtype=torch.int32, device="cpu"))

174 -

175 - self.decode_context_lens_cpu.append(

176 - torch.empty(self.max_num_tokens,

177 - dtype=torch.int32,

178 - device="cpu"))

179 - self.decode_context_lens_np.append(

180 - self.decode_context_lens_cpu[-1].numpy())

106 + # Cached torch/numpy tensor

107 + # The pytorch tensor and numpy array share the same buffer.

108 + # Sometimes the numpy op is faster so we create both.

109 + self.input_ids_cpu = torch.zeros(self.max_num_tokens,

110 + dtype=torch.int32,

111 + device="cpu")

112 + self.input_ids_np = self.input_ids_cpu.numpy()

113 +

114 + self.positions_cpu = torch.zeros(self.max_num_tokens,

115 + dtype=torch.int32,

116 + device="cpu")

117 + self.positions_np = self.positions_cpu.numpy()

118 +

119 + self.slot_mapping_cpu = torch.zeros(self.max_num_tokens,

120 + dtype=torch.int64,

121 + device="cpu")

122 + self.slot_mapping_np = self.slot_mapping_cpu.numpy()

123 +

124 + # self.input_batch.block_table has a shape of [max_num_reqs,

125 + # max_num_blocks_per_req]. To reduce the number of recompilation,

126 + # we want the block_table.shape[0] to be num_tokens.

127 + # To make the block_table to be compatible with the paged attention

128 + # kernel, we want the block_table[1] to be multiple of

129 + # NUM_KV_PAGES_PER_BLOCK.

130 + padded_max_num_blocks_per_req = _get_padded_number(

131 + self.max_num_blocks_per_req, NUM_KV_PAGES_PER_BLOCK)

132 + self.block_table_cpu = torch.zeros(

133 + (self.max_num_tokens, padded_max_num_blocks_per_req),

134 + dtype=self.input_batch.block_table.get_cpu_tensor().dtype,

135 + device="cpu")

136 +

137 + self.query_start_loc_cpu = torch.zeros(self.max_num_tokens + 1,

138 + dtype=torch.int32,

139 + device="cpu",

140 + pin_memory=self.pin_memory)

141 + self.query_start_loc_np = self.query_start_loc_cpu.numpy()

142 +

143 + self.seq_lens_cpu = torch.zeros(self.max_num_tokens,

144 + dtype=torch.int32,

145 + device="cpu",

146 + pin_memory=self.pin_memory)

147 + self.seq_lens_np = self.seq_lens_cpu.numpy()

148

149 # Range tensor with values [0 .. self.max_num_tokens - 1].

150 # Used to initialize positions / context_lens / seq_lens

@@ -191,7 +158,7 @@ class TPUModelRunner:

158 the input GPU tensors for the model.

159

160 Returns:

194 - True if there is a new/resumed/paused/finished request in the batch.

161 + True if there is a new/resumed/paused/finished request.

162 If False, we can skip copying SamplingMetadata to the GPU.

163 """

164 # Remove finished requests from the cached states.

@@ -303,9 +270,6 @@ class TPUModelRunner:

270 self.input_batch.condense(removed_req_indices)

271 return len(unscheduled_req_ids) > 0 or len(req_ids_to_add) > 0

272

306 - def swap_step(self):

307 - self.cur_swap_id = (self.cur_swap_id + 1) % self.num_swaps

308 -

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

274 assert self.model is not None

275 return self.model

@@ -345,238 +309,124 @@ class TPUModelRunner:

309

310 return kv_cache_spec

311

348 - def _get_prompts_and_decodes(

349 - self,

350 - scheduler_output: "SchedulerOutput",

351 - ) -> PromptDecodeInfo:

312 + def _prepare_inputs(self, scheduler_output: "SchedulerOutput"):

313 total_num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens

314 assert total_num_scheduled_tokens > 0

315 num_reqs = self.input_batch.num_reqs

316 assert num_reqs > 0

317

357 - # Traverse decodes first

358 - decode_req_ids = []

359 - for i in range(num_reqs):

360 - req_id = self.input_batch.req_ids[i]

361 - assert req_id is not None

362 -

363 - num_computed_tokens = self.input_batch.num_computed_tokens_cpu[i]

364 - num_prompt_tokens = self.input_batch.num_prompt_tokens[i]

365 - num_scheduled_tokens = scheduler_output.num_scheduled_tokens[

366 - req_id]

367 -

368 - if num_computed_tokens < num_prompt_tokens:

369 - # This is prompt

370 - break

371 -

372 - # This is decode

373 - assert num_scheduled_tokens == 1

374 - decode_req_ids.append(req_id)

375 -

376 - # Traverse prompts

377 - prompt_req_ids = []

378 - prompt_scheduled_tokens = []

379 - for i in range(len(decode_req_ids), num_reqs):

380 - req_id = self.input_batch.req_ids[i]

318 + # Get the number of scheduled tokens for each request.

319 + num_scheduled_tokens_per_req = []

320 + max_num_scheduled_tokens_all_reqs = 0

321 + for req_id in self.input_batch.req_ids[:num_reqs]:

322 assert req_id is not None

382 -

383 - num_computed_tokens = self.input_batch.num_computed_tokens_cpu[i]

384 - num_prompt_tokens = self.input_batch.num_prompt_tokens[i]

385 - num_scheduled_tokens = scheduler_output.num_scheduled_tokens[

386 - req_id]

387 -

388 - # Must be prompt

389 - assert num_computed_tokens < num_prompt_tokens

390 -

391 - prompt_req_ids.append(req_id)

392 - prompt_scheduled_tokens.append(num_scheduled_tokens)

393 -

394 - return PromptDecodeInfo(prompt_req_ids, decode_req_ids,

395 - prompt_scheduled_tokens)

396 -

397 - def _prepare_prompt(self, req_index: int,

398 - num_scheduled_tokens: int) -> PromptData:

399 - num_computed_tokens = self.input_batch.num_computed_tokens_cpu[

400 - req_index]

401 - num_prompt_tokens = self.input_batch.num_prompt_tokens[req_index]

402 -

403 - # Must be prompt

404 - assert num_computed_tokens < num_prompt_tokens

405 -

406 - # Prompt len

407 - prompt_len = num_scheduled_tokens

408 - padded_prompt_len = _get_padded_prompt_len(prompt_len)

409 - assert padded_prompt_len <= self.max_model_len

410 -

411 - # Seq len

412 - seq_len = num_computed_tokens + prompt_len

413 - padded_seq_len = num_computed_tokens + padded_prompt_len

414 -

415 - # Input tokens

416 - input_tokens_cpu = self.input_batch.token_ids_cpu_tensor[

417 - req_index, num_computed_tokens:padded_seq_len]

418 - input_tokens_cpu[prompt_len:] = 0

419 -

420 - # Input positions

421 - input_positions_np = self.input_positions_np[

422 - self.cur_swap_id][:padded_prompt_len]

423 - np.add(num_computed_tokens,

424 - self.arange_np[:padded_prompt_len],

425 - out=input_positions_np)

426 - input_positions_np[prompt_len:] = 0

427 -

428 - # Slot mapping

429 - block_table_np = \

430 - self.input_batch.block_table.get_numpy_array()

431 - block_numbers_np = block_table_np[req_index, input_positions_np //

432 - self.block_size]

433 - block_offsets_np = input_positions_np % self.block_size

434 -

435 - slot_mapping_np = self.slot_mapping_np[

436 - self.cur_swap_id][:padded_prompt_len]

437 - np.add(block_numbers_np * self.block_size,

438 - block_offsets_np,

439 - out=slot_mapping_np)

440 - slot_mapping_np[prompt_len:] = _PAD_SLOT_ID

441 -

442 - # Block table

443 - block_table_cpu = None

444 - if num_computed_tokens > 0:

445 - block_table_cpu = self.input_batch.block_table.get_cpu_tensor()

446 - block_table_cpu = block_table_cpu[req_index]

447 -

448 - # Context len

449 - self.prompt_context_lens_cpu[self.cur_swap_id][0] = 0

450 - if num_computed_tokens > 0:

451 - self.prompt_context_lens_cpu[self.cur_swap_id][0] = seq_len

452 -

453 - # Effective query len

454 - self.prompt_effective_query_lens_cpu[self.cur_swap_id][0] = prompt_len

455 -

456 - # Get final tensors

457 - input_tokens = input_tokens_cpu.reshape(1, -1).to(self.device)

458 - input_positions = self.input_positions_cpu[

459 - self.cur_swap_id][:padded_prompt_len].reshape(1,

460 - -1).to(self.device)

461 - slot_mapping = self.slot_mapping_cpu[

462 - self.cur_swap_id][:padded_prompt_len].reshape(1,

463 - -1).to(self.device)

464 - block_table = block_table_cpu.reshape(1, -1).to(

465 - self.device) if block_table_cpu is not None else None

466 -

467 - context_lens = self.prompt_context_lens_cpu[self.cur_swap_id].to(

468 - self.device)

469 - effective_query_lens = self.prompt_effective_query_lens_cpu[

470 - self.cur_swap_id].to(self.device)

471 -

472 - self.swap_step()

473 -

474 - # Attn metadata

475 - attn_metadata = PallasMetadata(

476 - num_prefills=1,

477 - num_prefill_tokens=0, # NOTE: This is not used.

478 - num_decode_tokens=0,

479 - slot_mapping=slot_mapping,

480 - multi_modal_placeholder_index_maps=None,

481 - enable_kv_scales_calculation=True,

482 - block_tables=block_table,

483 - context_lens=context_lens,

484 - effective_query_lens=effective_query_lens,

485 - )

486 -

487 - return PromptData(input_tokens, input_positions, attn_metadata)

488 -

489 - def _prepare_decode(

490 - self,

491 - decode_req_ids: List[str],

492 - ) -> DecodeData:

493 - # Batch size

494 - batch_size = len(decode_req_ids)

495 - padded_batch_size = _get_padded_batch_size(batch_size)

496 - assert padded_batch_size <= self.max_model_len

497 -

498 - # Init [0 .. batch_size - 1]

499 - req_indices_np = self.arange_np[:padded_batch_size]

500 -

501 - # Input positions

502 - input_positions_np = self.input_positions_np[

503 - self.cur_swap_id][:padded_batch_size]

504 - np.add(self.input_batch.num_computed_tokens_cpu[:padded_batch_size],

505 - 0,

506 - out=input_positions_np)

507 - input_positions_np[batch_size:] = 0

508 - input_positions_cpu = self.input_positions_cpu[

509 - self.cur_swap_id][:padded_batch_size]

510 -

511 - # Input tokens

512 - token_indices_np = (

513 - input_positions_np +

514 - req_indices_np * self.input_batch.token_ids_cpu.shape[1])

515 - input_tokens_cpu = self.input_ids_cpu[

516 - self.cur_swap_id][:padded_batch_size]

323 + num_tokens = scheduler_output.num_scheduled_tokens[req_id]

324 + num_scheduled_tokens_per_req.append(num_tokens)

325 + max_num_scheduled_tokens_all_reqs = max(

326 + max_num_scheduled_tokens_all_reqs, num_tokens)

327 + num_scheduled_tokens_per_req = np.array(num_scheduled_tokens_per_req,

328 + dtype=np.int32)

329 + assert max_num_scheduled_tokens_all_reqs > 0

330 +

331 + # Get request indices.

332 + # E.g., [2, 5, 3] -> [0, 0, 1, 1, 1, 1, 1, 2, 2, 2]

333 + # For each scheduled token, what are the corresponding req index.

334 + req_indices = np.repeat(self.arange_np[:num_reqs],

335 + num_scheduled_tokens_per_req)

336 +

337 + # Get batched arange.

338 + # E.g., [2, 5, 3] -> [0, 1, 0, 1, 2, 3, 4, 0, 1, 2]

339 + # For each scheduled token, what is its position in corresponding req.

340 + arange = np.concatenate(

341 + [self.arange_np[:n] for n in num_scheduled_tokens_per_req])

342 +

343 + # Get positions.

344 + positions_np = self.positions_np[:total_num_scheduled_tokens]

345 + np.add(self.input_batch.num_computed_tokens_cpu[req_indices],

346 + arange,

347 + out=positions_np)

348 +

349 + # Get token indices.

350 + # E.g., [0, 1, 0, 1, 2, 3, 4, 0, 1, 2]

351 + # -> [0, 1, M, M + 1, M + 2, M + 3, M + 4, 2 * M, 2 * M + 1, 2 * M + 2]

352 + # where M is the max_model_len.

353 + token_indices = (positions_np +

354 + req_indices * self.input_batch.token_ids_cpu.shape[1])

355 +

356 + # NOTE(woosuk): We use torch.index_select instead of np.take here

357 + # because torch.index_select is much faster than np.take for large

358 + # tensors.

359 torch.index_select(self.input_batch.token_ids_cpu_tensor.flatten(),

360 0,

519 - torch.from_numpy(token_indices_np),

520 - out=input_tokens_cpu)

521 - input_tokens_cpu[batch_size:] = 0

522 -

523 - # Slot mapping

524 - block_table_indices_np = (

525 - req_indices_np * self.max_num_blocks_per_req +

526 - input_positions_np // self.block_size)

527 -

361 + torch.from_numpy(token_indices),

362 + out=self.input_ids_cpu[:total_num_scheduled_tokens])

363 +

364 + # Calculate the slot mapping.

365 + # E.g., [0, 1, 0, 1, 2, 3, 4, 0, 1, 2]

366 + # -> [0, 0, K, K, K + 1, K + 1, K + 2, 2 * K, 2 * K, 2 * K + 1]

367 + # where K is the max_num_blocks_per_req and the block size is 2.

368 + # NOTE(woosuk): We can't simply use `token_indices // block_size` here

369 + # because M (max_model_len) is not necessarily divisible by block_size.

370 + # req_indices: # E.g., [2, 5, 3] -> [0, 0, 1, 1, 1, 1, 1, 2, 2, 2]

371 + block_table_indices = (req_indices * self.max_num_blocks_per_req +

372 + positions_np // self.block_size)

373 + # NOTE(woosuk): We use torch.index_select instead of np.take here

374 + # because torch.index_select is much faster than np.take for large

375 + # tensors.

376 block_table_cpu = self.input_batch.block_table.get_cpu_tensor()

377 + block_numbers = block_table_cpu.flatten()[block_table_indices].numpy()

378 + block_offsets = positions_np % self.block_size

379 + np.add(block_numbers * self.block_size,

380 + block_offsets,

381 + out=self.slot_mapping_np[:total_num_scheduled_tokens])

382 +

383 + # Prepare the attention metadata.

384 + self.query_start_loc_np[0] = 0

385 + np.cumsum(num_scheduled_tokens_per_req,

386 + out=self.query_start_loc_np[1:num_reqs + 1])

387 +

388 + self.seq_lens_np[:num_reqs] = (

389 + self.input_batch.num_computed_tokens_cpu[:num_reqs] +

390 + num_scheduled_tokens_per_req)

391 +

392 + # Do the padding and copy the tensors to the TPU.

393 + padded_total_num_scheduled_tokens = _get_padded_number(

394 + total_num_scheduled_tokens, NUM_QUERIES_PER_BLOCK)

395 + self.input_ids = self.input_ids_cpu[:

396 + padded_total_num_scheduled_tokens].to(

397 + self.device)

398 + self.position_ids = self.positions_cpu[:

399 + padded_total_num_scheduled_tokens].to(

400 + self.device)

401 + self.slot_mapping_cpu[total_num_scheduled_tokens:] = _PAD_SLOT_ID

402 + slot_mapping = self.slot_mapping_cpu[:

403 + padded_total_num_scheduled_tokens].to(

404 + self.device)

405 + padded_block_table = self.block_table_cpu[:

406 + padded_total_num_scheduled_tokens]

407 + padded_block_table[:num_reqs, :self.max_num_blocks_per_req] = (

408 + self.input_batch.block_table.get_cpu_tensor()[:num_reqs])

409 + padded_block_table = padded_block_table.to(self.device)

410 + query_start_loc = self.query_start_loc_cpu[:

411 + padded_total_num_scheduled_tokens

412 + + 1].to(self.device)

413 + seq_lens = self.seq_lens_cpu[:padded_total_num_scheduled_tokens].to(

414 + self.device)

415

530 - block_numbers_np = block_table_cpu.flatten(

531 - )[block_table_indices_np].numpy()

532 -

533 - block_offsets_np = input_positions_np % self.block_size

534 -

535 - slot_mapping_np = self.slot_mapping_np[

536 - self.cur_swap_id][:padded_batch_size]

537 - np.add(block_numbers_np * self.block_size,

538 - block_offsets_np,

539 - out=slot_mapping_np)

540 - slot_mapping_np[batch_size:] = _PAD_SLOT_ID

541 -

542 - block_table_cpu = block_table_cpu[:padded_batch_size]

543 -

544 - # Context lens

545 - context_lens_np = self.decode_context_lens_np[

546 - self.cur_swap_id][:padded_batch_size]

547 - np.add(self.input_batch.num_computed_tokens_cpu[:padded_batch_size],

548 - 1,

549 - out=context_lens_np)

550 - context_lens_np[batch_size:] = 0

551 -

552 - # Get final tensors

553 - input_tokens = input_tokens_cpu.reshape(-1, 1).to(self.device)

554 - input_positions = input_positions_cpu.reshape(-1, 1).to(self.device)

555 - slot_mapping = self.slot_mapping_cpu[

556 - self.cur_swap_id][:padded_batch_size].reshape(-1,

557 - 1).to(self.device)

558 - block_table = block_table_cpu.to(self.device)

559 - context_lens = self.decode_context_lens_cpu[

560 - self.cur_swap_id][:padded_batch_size].to(self.device)

561 -

562 - self.swap_step()

563 -

564 - # Attn metadata

416 attn_metadata = PallasMetadata(

566 - num_prefills=0,

567 - num_prefill_tokens=0,

568 - num_decode_tokens=padded_batch_size,

417 slot_mapping=slot_mapping,

570 - multi_modal_placeholder_index_maps=None,

571 - enable_kv_scales_calculation=True,

572 - block_tables=block_table,

573 - context_lens=context_lens,

574 - effective_query_lens=None,

418 + block_tables=padded_block_table,

419 + context_lens=seq_lens,

420 + query_start_loc=query_start_loc,

421 + num_seqs=num_reqs,

422 )

576 -

577 - return DecodeData(input_tokens=input_tokens,

578 - input_positions=input_positions,

579 - attn_metadata=attn_metadata)

423 + # NOTE(woosuk): Due to chunked prefills, there can be at most 1 partial

424 + # request in the batch. While we should not sample any token from this

425 + # partial request, we do so for simplicity. We will ignore the sampled

426 + # token from the partial request.

427 + # TODO: Support prompt logprobs.

428 + logits_indices = query_start_loc[1:] - 1

429 + return attn_metadata, logits_indices

430

431 @torch.no_grad()

432 def execute_model(

@@ -586,118 +436,81 @@ class TPUModelRunner:

436 # Update cached state

437 self._update_states(scheduler_output)

438

589 - # If necessary, swap decodes/prompts to have all decodes on the start

590 - ensure_decodes_first(self.input_batch)

591 -

592 - # Prepare prompts/decodes info

593 - pd_info = self._get_prompts_and_decodes(scheduler_output)

594 -

595 - # Init

596 - num_prompts = len(pd_info.prompt_req_ids)

597 - num_decodes = len(pd_info.decode_req_ids)

598 - decode_data = None

599 - sampled_token_ids = [0] * self.input_batch.num_reqs

600 -

601 - # Run each prompt individually

602 - is_first = True

603 - for i in range(num_prompts):

604 - req_id = pd_info.prompt_req_ids[i]

605 - req_index = num_decodes + i

606 - assert req_index == self.input_batch.req_id_to_index[

607 - req_id] # TODO: Remove

608 - req_state = self.requests[req_id]

609 - num_scheduled_tokens = pd_info.prompt_scheduled_tokens[i]

610 - prompt_len = num_scheduled_tokens

611 - seq_len = req_state.num_computed_tokens + num_scheduled_tokens

612 -

613 - # Prepare first prompt

614 - if is_first:

615 - prompt_data = self._prepare_prompt(req_index,

616 - num_scheduled_tokens)

617 - is_first = False

618 -

619 - # Run forward pass

620 - with set_forward_context(prompt_data.attn_metadata,

621 - self.vllm_config):

622 - assert self.model is not None

623 - selected_token_ids = self.model(prompt_data.input_tokens,

624 - prompt_data.input_positions,

625 - self.kv_caches)

626 -

627 - # In parallel to TPU execution, prepare the next iteration

628 - if i < num_prompts - 1:

629 - # There is next prompt => prepare it

630 - prompt_data = self._prepare_prompt(

631 - req_index + 1, pd_info.prompt_scheduled_tokens[i + 1])

632 - elif i == num_prompts - 1 and num_decodes > 0:

633 - # There is next decode => prepare it

634 - decode_data = self._prepare_decode(pd_info.decode_req_ids)

635 -

636 - # Update cached state (if prompt is fully done)

637 - if seq_len >= len(req_state.prompt_token_ids):

638 - # Transfer sampled tokens from TPU to CPU

639 - selected_token_ids_cpu = selected_token_ids.cpu()

640 -

641 - # Get output token

642 - token_id = selected_token_ids_cpu[prompt_len - 1].item()

643 - sampled_token_ids[req_index] = token_id

644 -

645 - # Add output token to the request

646 - self.input_batch.token_ids_cpu[req_index, seq_len] = token_id

647 - self.input_batch.num_tokens[req_index] += 1

648 - req_state.output_token_ids.append(token_id)

439 + # Prepare inputs

440 + attn_metadata, logits_indices = self._prepare_inputs(scheduler_output)

441 + total_num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens

442

650 - # Run decodes (a single batch)

651 - if num_decodes > 0:

652 -

653 - # Prepare decode (if was not yet prepared)

654 - if decode_data is None:

655 - decode_data = self._prepare_decode(pd_info.decode_req_ids)

656 -

657 - # Run forward pass

658 - with set_forward_context(decode_data.attn_metadata,

659 - self.vllm_config):

660 - assert self.model is not None

661 - selected_token_ids = self.model(decode_data.input_tokens,

662 - decode_data.input_positions,

663 - self.kv_caches)

664 -

665 - # Transfer sampled tokens from TPU to CPU

666 - decode_token_ids_cpu = selected_token_ids.cpu()

667 - # Convert to list

668 - decode_token_ids_list = decode_token_ids_cpu.tolist()

669 -

670 - # Update cached state for each decode request

671 - for i in range(num_decodes):

672 - req_id = pd_info.decode_req_ids[i]

673 - req_index = i

674 - assert req_index == self.input_batch.req_id_to_index[

675 - req_id] # TODO: Remove

676 - req_state = self.requests[req_id]

677 - seq_len = req_state.num_computed_tokens + 1

678 -

679 - token_id = decode_token_ids_list[i]

680 - sampled_token_ids[req_index] = token_id

681 -

682 - self.input_batch.token_ids_cpu[req_index, seq_len] = token_id

683 - self.input_batch.num_tokens[req_index] += 1

684 - req_state.output_token_ids.append(token_id)

443 + # Run the decoder

444 + with set_forward_context(attn_metadata, self.vllm_config):

445 + hidden_states = self.model(

446 + token_ids=self.input_ids,

447 + position_ids=self.position_ids,

448 + kv_caches=self.kv_caches,

449 + )

450 + hidden_states = hidden_states[:total_num_scheduled_tokens]

451 + num_reqs = self.input_batch.num_reqs

452 + logits_indices = logits_indices[:num_reqs]

453 + hidden_states = hidden_states[logits_indices]

454 + logits = self.model.compute_logits(hidden_states, None)

455 + selected_token_ids = torch.argmax(logits, dim=-1, keepdim=True)

456 +

457 + # Then, let's update the cache state.

458 + request_seq_lens: List[Tuple[int, CachedRequestState, int]] = []

459 + for i, req_id in zip(range(num_reqs), self.input_batch.req_ids):

460 + assert req_id is not None

461 + req_state = self.requests[req_id]

462 + seq_len = (req_state.num_computed_tokens +

463 + scheduler_output.num_scheduled_tokens[req_id])

464 + if seq_len >= req_state.num_tokens:

465 + request_seq_lens.append((i, req_state, seq_len))

466 + else:

467 + # Ignore the sampled token from the partial request.

468 + # Rewind the generator state as if the token was not sampled.

469 + generator = self.input_batch.generators.get(i)

470 + if generator is not None:

471 + # This relies on cuda-specific torch-internal impl details

472 + generator.set_offset(generator.get_offset() - 4)

473 +

474 + # num_reqs entries should be non-None

475 + assert all(

476 + req_id is not None for req_id in

477 + self.input_batch.req_ids[:num_reqs]), "req_ids contains None"

478 + req_ids = cast(List[str], self.input_batch.req_ids[:num_reqs])

479

686 - # Create output.

687 - all_req_ids = pd_info.decode_req_ids + pd_info.prompt_req_ids

480 prompt_logprobs_dict: Dict[str, Optional[LogprobsTensors]] = {}

689 - for req_id in all_req_ids:

481 + for req_id in self.input_batch.req_ids[:num_reqs]:

482 prompt_logprobs_dict[req_id] = None

483

484 + max_gen_len = selected_token_ids.shape[-1]

485 + if max_gen_len == 1:

486 + valid_sampled_token_ids = selected_token_ids.tolist()

487 + for i, req_state, seq_len in request_seq_lens:

488 + token_id = valid_sampled_token_ids[i][0]

489 + self.input_batch.token_ids_cpu[i, seq_len] = token_id

490 + req_state.output_token_ids.append(token_id)

491 + self.input_batch.num_tokens[i] += 1

492 + else:

493 + valid_mask = selected_token_ids != INVALID_TOKEN_ID

494 + gen_lens = valid_mask.sum(dim=1).tolist()

495 + valid_sampled_token_ids = [

496 + seq.tolist()

497 + for seq in selected_token_ids[valid_mask].split(gen_lens)

498 + ]

499 + self.input_batch.num_tokens[:num_reqs] += gen_lens

500 + for i, req_state, seq_len in request_seq_lens:

501 + target_slice = slice(seq_len - gen_lens[i] + 1, seq_len + 1)

502 + self.input_batch.token_ids_cpu[

503 + i, target_slice] = valid_sampled_token_ids[i]

504 + req_state.output_token_ids.extend(valid_sampled_token_ids[i])

505 +

506 model_runner_output = ModelRunnerOutput(

693 - req_ids=all_req_ids,

507 + req_ids=req_ids,

508 req_id_to_index=self.input_batch.req_id_to_index,

695 - sampled_token_ids=[[token_id] for token_id in sampled_token_ids],

509 + sampled_token_ids=valid_sampled_token_ids,

510 spec_token_ids=None,

511 logprobs=None,

698 - prompt_logprobs_dict=prompt_logprobs_dict, # type: ignore[arg-type]

512 + prompt_logprobs_dict=prompt_logprobs_dict,

513 )

700 -

514 return model_runner_output

515

516 def load_model(self) -> None:

@@ -731,185 +544,63 @@ class TPUModelRunner:

544 self,

545 kv_caches,

546 num_tokens: int,

734 - seq_len: Optional[int] = None,

735 - exec_mode: Optional[ExecutionMode] = None,

547 ) -> None:

737 - assert seq_len is not None

738 - assert exec_mode is not None

739 -

740 - exec_mode = ExecutionMode(exec_mode)

741 - if exec_mode.is_prefill():

742 - seq_len = (seq_len + 15) // 16 * 16

743 - token_ids = torch.zeros((num_tokens, seq_len),

744 - dtype=torch.int32,

745 - device=self.device)

746 - position_ids = torch.zeros((num_tokens, seq_len),

747 - dtype=torch.int32,

748 - device=self.device)

749 - slot_mapping = torch.zeros((num_tokens, seq_len),

750 - dtype=torch.int64,

751 - device=self.device)

752 - if exec_mode == ExecutionMode.PREFILL:

753 - attn_metadata = PallasMetadata(

754 - num_prefills=num_tokens,

755 - num_prefill_tokens=num_tokens * seq_len,

756 - num_decode_tokens=0,

757 - slot_mapping=slot_mapping,

758 - multi_modal_placeholder_index_maps=None,

759 - enable_kv_scales_calculation=True,

760 - block_tables=None,

761 - context_lens=None,

762 - effective_query_lens=None,

763 - )

764 -

765 - else:

766 - context_lens = torch.ones((num_tokens, ),

767 - dtype=torch.int32,

768 - device=self.device)

769 -

770 - block_tables = torch.zeros(

771 - (num_tokens, self.max_num_blocks_per_req),

772 - dtype=torch.int32,

773 - device=self.device)

774 -

775 - effective_query_lens = torch.ones_like(context_lens)

776 -

777 - attn_metadata = PallasMetadata(

778 - num_prefills=num_tokens,

779 - num_prefill_tokens=num_tokens * seq_len,

780 - num_decode_tokens=0,

781 - slot_mapping=slot_mapping,

782 - multi_modal_placeholder_index_maps=None,

783 - enable_kv_scales_calculation=True,

784 - block_tables=block_tables,

785 - context_lens=context_lens,

786 - effective_query_lens=effective_query_lens,

787 - )

788 - else:

789 - assert seq_len == 1

790 - token_ids = torch.zeros((num_tokens, seq_len),

791 - dtype=torch.int32,

792 - device=self.device)

793 - position_ids = torch.zeros((num_tokens, seq_len),

794 - dtype=torch.int32,

795 - device=self.device)

796 - slot_mapping = torch.zeros((num_tokens, seq_len),

797 - dtype=torch.int64,

798 - device=self.device)

799 - block_tables = torch.zeros(

800 - (num_tokens, self.max_num_blocks_per_req),

801 - dtype=torch.int32,

802 - device=self.device)

803 - context_lens = torch.ones((num_tokens, ),

804 - dtype=torch.int32,

805 - device=self.device)

806 - attn_metadata = PallasMetadata(

807 - num_prefills=0,

808 - num_prefill_tokens=0,

809 - num_decode_tokens=num_tokens * seq_len,

810 - slot_mapping=slot_mapping,

811 - multi_modal_placeholder_index_maps=None,

812 - enable_kv_scales_calculation=True,

813 - block_tables=block_tables,

814 - context_lens=context_lens,

815 - )

548 + input_ids = torch.zeros(num_tokens,

549 + dtype=torch.int32,

550 + device=self.device)

551 + position_ids = torch.zeros(num_tokens,

552 + dtype=torch.int32,

553 + device=self.device)

554 + slot_mapping = torch.zeros(num_tokens,

555 + dtype=torch.int64,

556 + device=self.device)

557 + block_tables = torch.zeros((num_tokens, self.block_table_cpu.shape[1]),

558 + dtype=torch.int32,

559 + device=self.device)

560 + query_lens = [1] * num_tokens

561 + query_start_loc = torch.cumsum(torch.tensor([0] + query_lens,

562 + dtype=torch.int32),

563 + dim=0,

564 + dtype=torch.int32).to(self.device)

565 + context_lens = torch.ones((num_tokens, ),

566 + dtype=torch.int32,

567 + device=self.device)

568 + attn_metadata = PallasMetadata(

569 + slot_mapping=slot_mapping,

570 + block_tables=block_tables,

571 + context_lens=context_lens,

572 + query_start_loc=query_start_loc,

573 + num_seqs=num_tokens,

574 + )

575

817 - # NOTE(woosuk): There are two stages of compilation: torch.compile and

818 - # XLA compilation. Using `mark_dynamic` can reduce the torch.compile

819 - # overhead by reusing the FX graph for different shapes.

820 - # However, the XLA graph will still require static shapes and needs to

821 - # be re-compiled for every different shapes. This overhead is inevitable

822 - # in the first run, but can be skipped afterwards as we cache the XLA

823 - # graphs in the disk (VLLM_XLA_CACHE_PATH).

824 - if exec_mode.is_prefill():

825 - # Prefll

826 - torch._dynamo.mark_dynamic(token_ids, 1)

827 - torch._dynamo.mark_dynamic(position_ids, 1)

828 - torch._dynamo.mark_dynamic(attn_metadata.slot_mapping, 1)

829 - else:

830 - # Decode

831 - torch._dynamo.mark_dynamic(token_ids, 0)

832 - torch._dynamo.mark_dynamic(position_ids, 0)

833 - torch._dynamo.mark_dynamic(attn_metadata.slot_mapping, 0)

834 - torch._dynamo.mark_dynamic(attn_metadata.context_lens, 0)

835 - torch._dynamo.mark_dynamic(attn_metadata.block_tables, 0)

576 + torch._dynamo.mark_dynamic(input_ids, 0)

577 + torch._dynamo.mark_dynamic(position_ids, 0)

578 + torch._dynamo.mark_dynamic(attn_metadata.slot_mapping, 0)

579 + torch._dynamo.mark_dynamic(attn_metadata.block_tables, 0)

580 + torch._dynamo.mark_dynamic(attn_metadata.query_start_loc, 0)

581 + torch._dynamo.mark_dynamic(attn_metadata.context_lens, 0)

582

583 with set_forward_context(attn_metadata, self.vllm_config, 0):

584 assert self.model is not None

839 - self.model(token_ids, position_ids, kv_caches)

585 + self.model(input_ids, position_ids, kv_caches)

586

587 def capture_model(self) -> None:

588 """Compile the model."""

589

844 - # Prefill

845 - logger.info(

846 - "Compiling the model with different input shapes for prefill:")

847 - start = time.time()

848 - for batch_size in [1]:

849 - seq_len = 16

850 - while seq_len <= self.model_config.max_model_len:

851 - self.dummy_run(self.kv_caches,

852 - batch_size,

853 - seq_len,

854 - exec_mode=ExecutionMode.PREFILL)

855 - xm.wait_device_ops()

856 - logger.info(" batch_size: %d, seq_len: %d", batch_size,

857 - seq_len)

858 - num_tokens = batch_size * seq_len

859 - if num_tokens >= self.scheduler_config.max_num_batched_tokens:

860 - break

861 - seq_len = seq_len * 2

862 -

863 - end = time.time()

864 - logger.info(" -- Compilation for prefill done in %.2f [secs].",

865 - end - start)

866 -

867 - # Prefix prefill

868 - if self.scheduler_config.enable_chunked_prefill:

869 - logger.info("Compiling the model with different input shapes for "

870 - "prefix prefill:")

871 - start = time.time()

872 - for batch_size in [1]:

873 - seq_len = 16

874 - while seq_len <= self.model_config.max_model_len:

875 - self.dummy_run(self.kv_caches,

876 - batch_size,

877 - seq_len,

878 - exec_mode=ExecutionMode.PREFIX_PREFILL)

879 - xm.wait_device_ops()

880 - logger.info(" batch_size: %d, seq_len: %d", batch_size,

881 - seq_len)

882 - num_tokens = batch_size * seq_len

883 - if (num_tokens

884 - >= self.scheduler_config.max_num_batched_tokens):

885 - break

886 - seq_len = seq_len * 2

887 - end = time.time()

888 - logger.info(

889 - " -- Compilation for prefix prefill done in %.2f [secs].",

890 - end - start)

891 -

892 - # Decode

893 - logger.info(

894 - "Compiling the model with different input shapes for decode:")

895 - start = time.time()

896 - seq_len = 1

897 - batch_size = 8 # Must be in sync with _get_padded_batch_size()

590 + logger.info("Compiling the model with different input shapes.")

591 +

592 + start = time.perf_counter()

593 + num_tokens = 16

594 while True:

899 - self.dummy_run(self.kv_caches,

900 - batch_size,

901 - seq_len,

902 - exec_mode=ExecutionMode.DECODE)

595 + self.dummy_run(self.kv_caches, num_tokens)

596 + logger.info(" -- num_tokens: %d", num_tokens)

597 + xm.mark_step()

598 xm.wait_device_ops()

904 - logger.info(" batch_size: %d, seq_len: %d", batch_size, seq_len)

905 -

906 - if batch_size >= self.scheduler_config.max_num_seqs:

599 + if num_tokens >= self.scheduler_config.max_num_batched_tokens:

600 break

908 - batch_size = batch_size + 16 if batch_size >= 16 else batch_size * 2

909 -

910 - end = time.time()

911 - logger.info(" -- Compilation for decode done in %.2f [secs].",

912 - end - start)

601 + num_tokens *= 2

602 + end = time.perf_counter()

603 + logger.info("Compilation finished in in %.2f [secs].", end - start)

604

605 def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None:

606 """

@@ -965,12 +656,8 @@ class ModelWrapperV1(nn.Module):

656 """Executes the forward pass of the model and samples the next token.

657

658 Args:

968 - token_ids: The input token IDs of shape [batch_size, seq_len].

969 - position_ids: The input position IDs of shape [batch_size, seq_len].

970 - input_lens: The actual input lengths of shape [batch_size].

971 - t: The sampling temperature of shape [batch_size].

972 - p: The top-p probability of shape [batch_size].

973 - num_samples: Number of samples to draw from each logits vector.

659 + token_ids: The input token IDs of shape [num_tokens].

660 + position_ids: The input position IDs of shape [num_tokens].

661 kv_caches: The key and value caches. They can be None during the

662 memory profiling at initialization.

663 """

@@ -982,6 +669,7 @@ class ModelWrapperV1(nn.Module):

669 # [num_kv_heads, num_blocks, block_size, head_size]. To make it

670 # work, we need to flatten the first three dimensions and modify

671 # the slot_mapping accordingly.

672 + # kv_caches: List[Tuple[torch.Tensor, torch.Tensor]]

673 num_kv_heads, num_blocks, block_size, _ = kv_caches[0][0].shape

674 slot_mapping = attn_metadata.slot_mapping

675 slot_mapping = slot_mapping.flatten()

@@ -997,103 +685,22 @@ class ModelWrapperV1(nn.Module):

685 attn_metadata.slot_mapping = slot_mapping

686

687 assert self.model is not None

1000 - hidden_states = self.model(token_ids, position_ids)

1001 -

1002 - hidden_states = hidden_states.flatten(0, 1)

1003 - logits = self.model.compute_logits(hidden_states, None)

1004 -

1005 - # Greedy sampling.

1006 - argmax_token_ids = torch.argmax(logits, dim=-1, keepdim=True)

1007 - argmax_token_ids = argmax_token_ids.squeeze(dim=-1)

1008 - return argmax_token_ids

1009 -

1010 -

1011 -def swap_positions(b: InputBatch, id_1, id_2):

1012 - assert id_1 != id_2

1013 - req_id_1 = b.req_ids[id_1]

1014 - req_id_2 = b.req_ids[id_2]

1015 - assert req_id_1 is not None

1016 - assert req_id_2 is not None

1017 - assert id_1 == b.req_id_to_index[req_id_1]

1018 - assert id_2 == b.req_id_to_index[req_id_2]

1019 -

1020 - b.req_ids[id_1], b.req_ids[id_2] = b.req_ids[id_2], b.req_ids[id_1]

1021 - b.req_id_to_index[req_id_1], b.req_id_to_index[

1022 - req_id_2] = b.req_id_to_index[req_id_2], b.req_id_to_index[req_id_1]

1023 -

1024 - ids = [id_1, id_2]

1025 - rev_ids = [id_2, id_1]

1026 - b.num_tokens[ids] = b.num_tokens[rev_ids]

1027 - b.token_ids_cpu[ids] = b.token_ids_cpu[rev_ids]

1028 - b.num_prompt_tokens[ids] = b.num_prompt_tokens[rev_ids]

1029 - b.num_computed_tokens_cpu[ids] = b.num_computed_tokens_cpu[rev_ids]

1030 -

1031 - b.block_table.swap_row(id_1, id_2)

1032 -

1033 - b.temperature_cpu[ids] = b.temperature_cpu[rev_ids]

1034 - b.top_p_cpu[ids] = b.top_p_cpu[rev_ids]

1035 - b.top_k_cpu[ids] = b.top_k_cpu[rev_ids]

1036 - b.frequency_penalties_cpu[ids] = b.frequency_penalties_cpu[rev_ids]

1037 - b.presence_penalties_cpu[ids] = b.presence_penalties_cpu[rev_ids]

1038 - b.repetition_penalties_cpu[ids] = b.repetition_penalties_cpu[rev_ids]

1039 -

1040 - b.min_tokens[id_1], b.min_tokens[id_2] = b.min_tokens[id_2], b.min_tokens[

1041 - id_1]

1042 -

1043 - gen_1 = b.generators.pop(id_1, None)

1044 - gen_2 = b.generators.pop(id_2, None)

1045 - if gen_1 is not None:

1046 - b.generators[id_2] = gen_1

1047 - if gen_2 is not None:

1048 - b.generators[id_1] = gen_2

1049 -

1050 -

1051 -def ensure_decodes_first(b: InputBatch):

1052 - num_reqs = b.num_reqs

1053 - while True:

1054 - # Find the first prompt index

1055 - first_prompt_index = None

1056 - for i in range(num_reqs):

1057 - if b.num_computed_tokens_cpu[i] < b.num_prompt_tokens[i]:

1058 - first_prompt_index = i

1059 - break

1060 - if first_prompt_index is None:

1061 - break

1062 -

1063 - # Find the last decode index

1064 - last_decode_index = None

1065 - for i in reversed(range(num_reqs)):

1066 - if b.num_computed_tokens_cpu[i] >= b.num_prompt_tokens[i]:

1067 - last_decode_index = i

1068 - break

1069 - if last_decode_index is None:

1070 - break

1071 -

1072 - # Sanity

1073 - assert first_prompt_index != last_decode_index

1074 -

1075 - # Check if done

1076 - if first_prompt_index > last_decode_index:

1077 - break

1078 -

1079 - # Swap

1080 - swap_positions(b, first_prompt_index, last_decode_index)

688 + hidden_states = self.model(

689 + token_ids,

690 + position_ids,

691 + kv_caches,

692 + )

693

694 + return hidden_states

695

1083 -def _get_padded_prompt_len(x: int) -> int:

1084 - # NOTE(woosuk): The pallas FlashAttention kernel requires the sequence

1085 - # length to be a multiple of 16. We pad the prompt length to the nearest

1086 - # multiple of 16. This is also good for performance.

1087 - if x <= 16:

1088 - return 16

1089 - return 1 << (x - 1).bit_length()

696 + def compute_logits(

697 + self,

698 + hidden_states: torch.Tensor,

699 + sampling_metadata,

700 + ) -> Optional[torch.Tensor]:

701 + logits = self.model.compute_logits(hidden_states, sampling_metadata)

702 + return logits

703

704

1092 -def _get_padded_batch_size(batch_size: int) -> int:

1093 - # The GMM Pallas kernel requires num_tokens * topk to be a multiple of 16.

1094 - # To meet this requirement in the simplest way, we set the minimal batch

1095 - # size to 8.

1096 - if batch_size <= 8:

1097 - return 8

1098 - else:

1099 - return ((batch_size + 15) // 16) * 16

705 +def _get_padded_number(n: int, multiple: int) -> int:

706 + return ((n + multiple - 1) // multiple) * multiple

vllm/v1/worker/tpu_worker.py

+2 -4

@@ -21,7 +21,7 @@ from vllm.v1.kv_cache_interface import (FullAttentionSpec, KVCacheConfig,

21 KVCacheSpec)

22 from vllm.v1.outputs import ModelRunnerOutput

23 from vllm.v1.utils import bind_kv_cache

24 -from vllm.v1.worker.tpu_model_runner import ExecutionMode, TPUModelRunner

24 +from vllm.v1.worker.tpu_model_runner import TPUModelRunner

26 logger = init_logger(__name__)

@@ -126,9 +126,7 @@ class TPUWorker:

126

127 self.model_runner.dummy_run(

128 runner_kv_caches,

129 - num_tokens=1,

130 - seq_len=self.scheduler_config.max_num_batched_tokens,

131 - exec_mode=ExecutionMode.PREFILL,

129 + num_tokens=self.scheduler_config.max_num_batched_tokens,

130 )

131

132 # Synchronize before measuring the memory usage.