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

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

[v1] Pass BlockTable and KVCacheSpec to AttentionMetadataBuilders (#17483)

Signed-off-by: Chen Zhang <zhangch99@outlook.com>

Chen Zhang committed May 11, 2025 at 07:12 UTC 950751a9870f040ad0f1a876a8c8a03ecb97000b

11 files changed +132 -68

tests/v1/worker/test_gpu_input_batch.py

@@ -221,6 +221,7 @@ def test_sampling_metadata_in_input_batch(device: str, batch_size: int):

221 max_num_reqs=batch_size,

222 max_model_len=1024,

223 max_num_blocks_per_req=10,

224 + max_num_batched_tokens=1024,

225 device=torch.device(device),

226 pin_memory=is_pin_memory_available(),

227 vocab_size=1024,

@@ -310,6 +311,7 @@ def test_swap_states_in_input_batch(device: str, batch_size: int,

311 max_num_reqs=batch_size,

312 max_model_len=1024,

313 max_num_blocks_per_req=10,

314 + max_num_batched_tokens=1024,

315 device=torch.device(device),

316 pin_memory=is_pin_memory_available(),

317 vocab_size=1024,

@@ -318,6 +320,7 @@ def test_swap_states_in_input_batch(device: str, batch_size: int,

320 max_num_reqs=batch_size,

321 max_model_len=1024,

322 max_num_blocks_per_req=10,

323 + max_num_batched_tokens=1024,

324 device=torch.device(device),

325 pin_memory=is_pin_memory_available(),

326 vocab_size=1024,

tests/v1/worker/test_gpu_model_runner.py

+20 -1

@@ -1,14 +1,31 @@

1 # SPDX-License-Identifier: Apache-2.0

2 +import weakref

3 +

4 import pytest

5 +import torch

7 from vllm.config import CacheConfig, ModelConfig, SchedulerConfig, VllmConfig

8 from vllm.sampling_params import SamplingParams

9 from vllm.v1.core.sched.output import (CachedRequestData, NewRequestData,

10 SchedulerOutput)

11 +from vllm.v1.kv_cache_interface import FullAttentionSpec

12 from vllm.v1.sample.metadata import SamplingMetadata

13 from vllm.v1.worker.gpu_model_runner import GPUModelRunner

16 +def initialize_kv_cache(runner: GPUModelRunner):

17 + """

18 + Only perform necessary steps in GPUModelRunner.initialize_kv_cache()

19 + """

20 + kv_cache_spec = FullAttentionSpec(block_size=16,

21 + num_kv_heads=1,

22 + head_size=64,

23 + dtype=torch.float16,

24 + use_mla=False)

25 + runner.attn_metadata_builder = runner.attn_backend.get_builder_cls()(

26 + weakref.proxy(runner), kv_cache_spec, runner.input_batch.block_table)

27 +

28 +

29 @pytest.fixture

30 def model_runner():

31 scheduler_config = SchedulerConfig(

@@ -38,7 +55,9 @@ def model_runner():

55 )

57 device = "cuda"

41 - return GPUModelRunner(vllm_config, device)

58 + runner = GPUModelRunner(vllm_config, device)

59 + initialize_kv_cache(runner)

60 + return runner

63 def _schedule_new_request(*req_ids: str) -> SchedulerOutput:

vllm/v1/attention/backends/flash_attn.py

+30 -17

@@ -19,6 +19,8 @@ from vllm.logger import init_logger

19 from vllm.platforms import current_platform

20 from vllm.utils import cdiv

21 from vllm.v1.attention.backends.utils import CommonAttentionMetadata

22 +from vllm.v1.kv_cache_interface import AttentionSpec

23 +from vllm.v1.worker.block_table import BlockTable

25 if TYPE_CHECKING:

26 from vllm.v1.core.sched.output import SchedulerOutput

@@ -167,7 +169,7 @@ def make_local_attention_virtual_batches(

169 query_start_loc_np: np.ndarray,

170 seq_lens_np: np.ndarray,

171 block_table: torch.Tensor,

170 - page_size: int = 0,

172 + block_size: int = 0,

173 ) -> tuple[np.ndarray, np.ndarray, np.ndarray, torch.Tensor]:

174 q_seqlens = query_start_loc_np[1:] - query_start_loc_np[:-1]

175 actual_batch_size = seq_lens_np.shape[0]

@@ -238,14 +240,14 @@ def make_local_attention_virtual_batches(

240 # For the example the local attention blocks start at:

241 # _b0_ _____b1_____ _b2_

242 # k_seqstarts_absolute = [0, 4, 4, 8, 12, 16, 4, 8]

241 - block_starts = k_seqstarts_absolute // page_size

242 - assert attn_chunk_size % page_size == 0, \

243 + block_starts = k_seqstarts_absolute // block_size

244 + assert attn_chunk_size % block_size == 0, \

245 f"attn_chunk_size {attn_chunk_size} is not " \

244 - f"divisible by page_size {page_size}"

245 - pages_per_local_batch = attn_chunk_size // page_size

246 + f"divisible by block_size {block_size}"

247 + pages_per_local_batch = attn_chunk_size // block_size

248

249 # Create a block_table for the local attention blocks

248 - # For out example if we have a block-table like (assuming page_size=2):

250 + # For out example if we have a block-table like (assuming block_size=2):

251 # block_table = [

252 # [ 0, 1, 2, 3, 4, 5, 6, 7, 8, 9], < batch 0

253 # [10, 11, 12, 13, 14, 15, 16, 17, 18, 19], < batch 1

@@ -289,7 +291,8 @@ def _get_sliding_window_configs(

291

292 class FlashAttentionMetadataBuilder:

293

292 - def __init__(self, runner: "GPUModelRunner"):

294 + def __init__(self, runner: "GPUModelRunner", kv_cache_spec: AttentionSpec,

295 + block_table: BlockTable):

296 model_config = runner.model_config

297 compilation_config = runner.vllm_config.compilation_config

298

@@ -299,7 +302,9 @@ class FlashAttentionMetadataBuilder:

302 self.num_heads_kv = model_config.get_num_kv_heads(

303 runner.parallel_config)

304 self.headdim = model_config.get_head_size()

302 - self.page_size = self.runner.block_size

305 + self.block_size = kv_cache_spec.block_size

306 + self.kv_cache_spec = kv_cache_spec

307 + self.block_table = block_table

308

309 if get_flash_attn_version() == 3:

310 self.aot_schedule = not compilation_config.full_cuda_graph

@@ -323,9 +328,17 @@ class FlashAttentionMetadataBuilder:

328 max_seq_len = self.runner.seq_lens_np[:num_reqs].max()

329 query_start_loc = common_attn_metadata.query_start_loc

330 seq_lens = common_attn_metadata.seq_lens

326 - block_table = (

327 - self.runner.input_batch.block_table.get_device_tensor()[:num_reqs])

328 - slot_mapping = self.runner.slot_mapping[:num_actual_tokens]

331 + block_table = self.block_table

332 + block_table_tensor = block_table.get_device_tensor()[:num_reqs]

333 +

334 + block_table.slot_mapping[:num_actual_tokens].copy_(

335 + block_table.slot_mapping_cpu[:num_actual_tokens],

336 + non_blocking=True)

337 + # Fill unused with -1. Needed for reshape_and_cache in full cuda graph

338 + # mode.

339 + block_table.slot_mapping[num_actual_tokens:].fill_(-1)

340 +

341 + slot_mapping = block_table.slot_mapping[:num_actual_tokens]

342

343 if self.aot_sliding_window is None:

344 self.aot_sliding_window = (-1, -1)

@@ -354,7 +367,7 @@ class FlashAttentionMetadataBuilder:

367 num_heads_q=self.num_heads_q,

368 num_heads_kv=self.num_heads_kv,

369 headdim=self.headdim,

357 - page_size=self.page_size,

370 + page_size=self.block_size,

371 cu_seqlens_q=cu_query_lens,

372 causal=causal,

373 window_size=self.aot_sliding_window,

@@ -365,12 +378,12 @@ class FlashAttentionMetadataBuilder:

378 local_attn_metadata = None

379 if self.runner.attention_chunk_size is not None:

380 seqlens_q_local_np, virt_q_cu_seqlens_np, virt_k_seqlens_np, \

368 - virt_block_table = make_local_attention_virtual_batches(

381 + virt_block_table_tensor = make_local_attention_virtual_batches(

382 self.runner.attention_chunk_size,

383 self.runner.query_start_loc_np[:num_reqs + 1],

384 self.runner.seq_lens_np[:num_reqs],

372 - block_table,

373 - self.runner.block_size,

385 + block_table_tensor,

386 + self.block_size,

387 )

388 local_query_start_loc = torch.from_numpy(virt_q_cu_seqlens_np).to(

389 self.runner.device, non_blocking=True)

@@ -389,7 +402,7 @@ class FlashAttentionMetadataBuilder:

402 local_attn_metadata = FlashAttentionMetadata.LocalAttentionMetadata(

403 local_query_start_loc=local_query_start_loc,

404 local_seqused_k=local_seqused_k,

392 - local_block_table=virt_block_table,

405 + local_block_table=virt_block_table_tensor,

406 local_max_query_len=local_max_query_len,

407 local_max_seq_len=local_max_seq_len,

408 local_scheduler_metadata=local_scheduler_metadata,

@@ -440,7 +453,7 @@ class FlashAttentionMetadataBuilder:

453 query_start_loc=query_start_loc,

454 max_seq_len=max_seq_len,

455 seq_lens=seq_lens,

443 - block_table=block_table,

456 + block_table=block_table_tensor,

457 slot_mapping=slot_mapping,

458 use_cascade=use_cascade,

459 common_prefix_len=common_prefix_len,

vllm/v1/attention/backends/flashinfer.py

+20 -15

@@ -19,6 +19,8 @@ from vllm.config import (VllmConfig, get_current_vllm_config,

19 from vllm.logger import init_logger

20 from vllm.v1.attention.backends.flash_attn import use_cascade_attention

21 from vllm.v1.attention.backends.utils import CommonAttentionMetadata

22 +from vllm.v1.kv_cache_interface import AttentionSpec

23 +from vllm.v1.worker.block_table import BlockTable

25 if TYPE_CHECKING:

26 from vllm.v1.core.sched.output import SchedulerOutput

@@ -202,7 +204,8 @@ class FlashInferMetadata:

204

205 class FlashInferMetadataBuilder:

206

205 - def __init__(self, runner: GPUModelRunner):

207 + def __init__(self, runner: GPUModelRunner, kv_cache_spec: AttentionSpec,

208 + block_table: BlockTable):

209 self.runner = runner

210 self._workspace_buffer = None

211 self._prefill_wrapper = None # Wrapper for prefill/append

@@ -213,6 +216,8 @@ class FlashInferMetadataBuilder:

216 self.global_hyperparameters: Optional[PerLayerParameters] = None

217

218 self.vllm_config = get_current_vllm_config()

219 + self.kv_cache_spec = kv_cache_spec

220 + self.block_table = block_table

221

222 def reorder_batch(self, input_batch: InputBatch,

223 scheduler_output: SchedulerOutput) -> bool:

@@ -400,13 +405,12 @@ class FlashInferMetadataBuilder:

405 assert self._num_decodes + self._num_prefills == num_reqs

406 assert (self._num_decode_tokens +

407 self._num_prefill_tokens == num_actual_tokens)

403 - page_size = self.runner.block_size

408 + page_size = self.kv_cache_spec.block_size

409 device = self.runner.device

410 qo_indptr = common_attn_metadata.query_start_loc

411 seq_lens = common_attn_metadata.seq_lens

407 - block_table = (

408 - self.runner.input_batch.block_table.get_device_tensor()[:num_reqs])

409 - slot_mapping = self.runner.slot_mapping_cpu[:num_actual_tokens].to(

412 + block_table_tensor = self.block_table.get_device_tensor()[:num_reqs]

413 + slot_mapping = self.block_table.slot_mapping_cpu[:num_actual_tokens].to(

414 self.runner.device, non_blocking=True).long()

415

416 block_table_bounds = (seq_lens + page_size - 1) // page_size

@@ -422,12 +426,13 @@ class FlashInferMetadataBuilder:

426 shared_kv_page_indptr = torch.tensor([0, num_common_kv_blocks],

427 dtype=torch.int32,

428 device=device)

425 - shared_kv_page_indices = block_table[0, :num_common_kv_blocks]

429 + shared_kv_page_indices = block_table_tensor[

430 + 0, :num_common_kv_blocks]

431 shared_kv_last_page_len = torch.tensor([page_size],

432 dtype=torch.int32,

433 device=device)

434 # Remove the blocks of the shared prefix from all requests.

430 - block_table = block_table[:, num_common_kv_blocks:]

435 + block_table_tensor = block_table_tensor[:, num_common_kv_blocks:]

436 block_table_bounds -= num_common_kv_blocks

437 else:

438 shared_qo_indptr = None

@@ -435,11 +440,11 @@ class FlashInferMetadataBuilder:

440 shared_kv_page_indices = None

441 shared_kv_last_page_len = None

442

438 - mask = (torch.arange(block_table.size(1),

439 - dtype=block_table.dtype,

440 - device=block_table.device).unsqueeze(0)

443 + mask = (torch.arange(block_table_tensor.size(1),

444 + dtype=block_table_tensor.dtype,

445 + device=block_table_tensor.device).unsqueeze(0)

446 < block_table_bounds.unsqueeze(1))

442 - paged_kv_indices = block_table[mask]

447 + paged_kv_indices = block_table_tensor[mask]

448

449 paged_kv_indptr = torch.cat([

450 torch.zeros(1,

@@ -459,10 +464,10 @@ class FlashInferMetadataBuilder:

464 paged_kv_indices=paged_kv_indices,

465 paged_kv_last_page_len=paged_kv_last_page_len,

466 num_qo_heads=self.runner.num_query_heads,

462 - num_kv_heads=self.runner.num_kv_heads,

463 - head_dim=self.runner.head_size,

467 + num_kv_heads=self.kv_cache_spec.num_kv_heads,

468 + head_dim=self.kv_cache_spec.head_size,

469 page_size=page_size,

465 - data_type=self.runner.kv_cache_dtype,

470 + data_type=self.kv_cache_spec.dtype,

471 q_data_type=self.runner.dtype,

472 slot_mapping=slot_mapping,

473 num_decodes=self._num_decodes,

@@ -481,7 +486,7 @@ class FlashInferMetadataBuilder:

486 return attn_metadata

487

488 def use_cascade_attention(self, *args, **kwargs) -> bool:

484 - if self.runner.kv_cache_dtype != self.runner.model_config.dtype:

489 + if self.kv_cache_spec.dtype != self.runner.model_config.dtype:

490 # TODO: The cascade wrapper currently does not support setting

491 # kv cache dtype to something different from query dtype.

492 return False

vllm/v1/attention/backends/mla/common.py

+15 -8

@@ -207,6 +207,8 @@ from vllm.model_executor.layers.linear import (ColumnParallelLinear,

207 from vllm.platforms import current_platform

208 from vllm.utils import cdiv, round_down

209 from vllm.v1.attention.backends.utils import CommonAttentionMetadata

210 +from vllm.v1.kv_cache_interface import AttentionSpec

211 +from vllm.v1.worker.block_table import BlockTable

212

213 try:

214 from vllm.vllm_flash_attn import flash_attn_varlen_func

@@ -334,6 +336,8 @@ class MLACommonMetadataBuilder(Generic[M]):

336

337 def __init__(self,

338 runner: "GPUModelRunner",

339 + kv_cache_spec: AttentionSpec,

340 + block_table: BlockTable,

341 metadata_cls: Optional[type[M]] = None):

342 self.metadata_cls = metadata_cls \

343 if metadata_cls is not None else MLACommonMetadata

@@ -346,10 +350,11 @@ class MLACommonMetadataBuilder(Generic[M]):

350 runner.parallel_config)

351 self.mla_dims = get_mla_dims(model_config)

352 self.aot_schedule = is_vllm_fa and (get_flash_attn_version() == 3)

353 + self.kv_cache_spec = kv_cache_spec

354

355 # Dont try to access the runner on AMD

356 if self.aot_schedule:

352 - self.page_size = self.runner.block_size

357 + self.page_size = self.kv_cache_spec.block_size

358

359 if self.chunked_prefill_enabled:

360 self.chunked_prefill_workspace_size = min(

@@ -375,6 +380,7 @@ class MLACommonMetadataBuilder(Generic[M]):

380 dtype=model_config.dtype,

381 device=runner.device,

382 )

383 + self.block_table = block_table

384

385 def reorder_batch(self, input_batch: "InputBatch",

386 scheduler_output: "SchedulerOutput") -> bool:

@@ -436,9 +442,10 @@ class MLACommonMetadataBuilder(Generic[M]):

442

443 return modified_batch

444

439 - def _build_decode(self, block_table: torch.Tensor, seq_lens: torch.Tensor):

445 + def _build_decode(self, block_table_tensor: torch.Tensor,

446 + seq_lens: torch.Tensor):

447 return MLACommonDecodeMetadata(

441 - block_table=block_table,

448 + block_table=block_table_tensor,

449 seq_lens=seq_lens,

450 )

451

@@ -451,9 +458,9 @@ class MLACommonMetadataBuilder(Generic[M]):

458 # function. We should avoid GPU -> CPU sync as much as possible because

459 # it blocks on all previous kernels.

460 device = self.runner.device

454 - block_table = (

455 - self.runner.input_batch.block_table.get_device_tensor()[:num_reqs])

456 - slot_mapping = self.runner.slot_mapping_cpu[:num_actual_tokens].to(

461 + block_table = self.block_table

462 + block_table_tensor = block_table.get_device_tensor()[:num_reqs]

463 + slot_mapping = block_table.slot_mapping_cpu[:num_actual_tokens].to(

464 device, non_blocking=True).long()

465

466 query_start_loc = common_attn_metadata.query_start_loc

@@ -530,7 +537,7 @@ class MLACommonMetadataBuilder(Generic[M]):

537 self.chunked_prefill_workspace_size

538

539 prefill_metadata = MLACommonPrefillMetadata(

533 - block_table=block_table[reqs_start:, ...],

540 + block_table=block_table_tensor[reqs_start:, ...],

541 query_start_loc=prefill_query_start_loc,

542 max_query_len=max_query_len,

543 chunked_context=chunked_context_metadata,

@@ -539,7 +546,7 @@ class MLACommonMetadataBuilder(Generic[M]):

546 decode_metadata = None

547 if self._num_decodes > 0:

548 decode_metadata = self._build_decode(

542 - block_table=block_table[:self._num_decodes, ...],

549 + block_table_tensor=block_table_tensor[:self._num_decodes, ...],

550 seq_lens=seq_lens[:self._num_decodes],

551 )

552

vllm/v1/attention/backends/mla/flashmla.py

+7 -4

@@ -16,6 +16,8 @@ from vllm.v1.attention.backends.mla.common import (MLACommonBackend,

16 MLACommonImpl,

17 MLACommonMetadata,

18 MLACommonMetadataBuilder)

19 +from vllm.v1.kv_cache_interface import AttentionSpec

20 +from vllm.v1.worker.block_table import BlockTable

22 logger = init_logger(__name__)

@@ -52,13 +54,14 @@ class FlashMLAMetadata(MLACommonMetadata[FlashMLADecodeMetadata]):

55 class FlashMLAMetadataBuilder(MLACommonMetadataBuilder[FlashMLAMetadata]):

55 - def __init__(self, runner):

56 - super().__init__(runner)

57 + def __init__(self, runner, kv_cache_spec: AttentionSpec,

58 + block_table: BlockTable):

59 + super().__init__(runner, kv_cache_spec, block_table)

61 self.num_q_heads = self.runner.model_config.get_num_attention_heads(

62 self.runner.parallel_config)

61 - def _build_decode(self, block_table: torch.Tensor,

64 + def _build_decode(self, block_table_tensor: torch.Tensor,

65 seq_lens: torch.Tensor) -> FlashMLADecodeMetadata:

66 tile_scheduler_metadata, num_splits = \

67 get_mla_metadata(

@@ -68,7 +71,7 @@ class FlashMLAMetadataBuilder(MLACommonMetadataBuilder[FlashMLAMetadata]):

71 )

73 return FlashMLADecodeMetadata(

71 - block_table=block_table,

74 + block_table=block_table_tensor,

75 seq_lens=seq_lens,

76 tile_scheduler_metadata=tile_scheduler_metadata,

77 num_splits=num_splits,

vllm/v1/attention/backends/mla/rocm_aiter_mla.py

+5 -2

@@ -14,6 +14,8 @@ from vllm.v1.attention.backends.mla.common import (MLACommonBackend,

14 MLACommonImpl,

15 MLACommonMetadata,

16 MLACommonMetadataBuilder)

17 +from vllm.v1.kv_cache_interface import AttentionSpec

18 +from vllm.v1.worker.block_table import BlockTable

20 # yapf: enable

@@ -59,8 +61,9 @@ class AiterMLAMetadata(MLACommonMetadata[AiterMLADecodeMetadata]):

62 class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):

62 - def __init__(self, runner):

63 - super().__init__(runner)

64 + def __init__(self, runner, kv_cache_spec: AttentionSpec,

65 + block_table: BlockTable):

66 + super().__init__(runner, kv_cache_spec, block_table)

67 max_model_len = self.runner.model_config.max_model_len

68 assert max_model_len == 32768,\

69 "AITER MLA requires max_model_len=32768"

vllm/v1/worker/block_table.py

+11

@@ -14,11 +14,13 @@ class BlockTable:

14 self,

15 max_num_reqs: int,

16 max_num_blocks_per_req: int,

17 + max_num_batched_tokens: int,

18 pin_memory: bool,

19 device: torch.device,

20 ):

21 self.max_num_reqs = max_num_reqs

22 self.max_num_blocks_per_req = max_num_blocks_per_req

23 + self.max_num_batched_tokens = max_num_batched_tokens

24 self.pin_memory = pin_memory

25 self.device = device

@@ -36,6 +38,15 @@ class BlockTable:

38 self.block_table_np = self.block_table_cpu.numpy()

39 self.num_blocks_per_row = np.zeros(max_num_reqs, dtype=np.int32)

41 + self.slot_mapping_cpu = torch.zeros(self.max_num_batched_tokens,

42 + dtype=torch.int64,

43 + device="cpu",

44 + pin_memory=self.pin_memory)

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

46 + self.slot_mapping = torch.zeros(self.max_num_batched_tokens,

47 + dtype=torch.int64,

48 + device=self.device)

49 +

50 def append_row(

51 self,

52 block_ids: list[int],

vllm/v1/worker/gpu_input_batch.py

@@ -59,6 +59,7 @@ class InputBatch:

59 max_num_reqs: int,

60 max_model_len: int,

61 max_num_blocks_per_req: int,

62 + max_num_batched_tokens: int,

63 device: torch.device,

64 pin_memory: bool,

65 vocab_size: int,

@@ -66,6 +67,7 @@ class InputBatch:

67 self.max_num_reqs = max_num_reqs

68 self.max_model_len = max_model_len

69 self.max_num_blocks_per_req = max_num_blocks_per_req

70 + self.max_num_batched_tokens = max_num_batched_tokens

71 self.device = device

72 self.pin_memory = pin_memory

73 self.vocab_size = vocab_size

@@ -100,6 +102,7 @@ class InputBatch:

102 self.block_table = BlockTable(

103 max_num_reqs=max_num_reqs,

104 max_num_blocks_per_req=max_num_blocks_per_req,

105 + max_num_batched_tokens=max_num_batched_tokens,

106 pin_memory=pin_memory,

107 device=device,

108 )

vllm/v1/worker/gpu_model_runner.py

+9 -12

@@ -150,8 +150,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):

150 f"FA3. Current attention backend is {attn_backend_name}, "

151 f"FlashAttention version is {flash_attn_version}.")

152

153 - self.attn_metadata_builder = self.attn_backend.get_builder_cls()(

154 - weakref.proxy(self))

153 self.cascade_attn_enabled = not self.model_config.disable_cascade_attn

154

155 # Multi-modal data support

@@ -174,6 +172,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):

172 # Initialize in initialize_kv_cache

173 self.kv_caches: list[torch.Tensor] = []

174 # self.kv_cache_config: KVCacheConfig

175 + # self.attn_metadata_builder: type[AttentionMetadataBuilder]

176

177 # req_id -> (input_id -> encoder_output)

178 self.encoder_cache: dict[str, dict[int, torch.Tensor]] = {}

@@ -203,6 +202,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):

202 max_num_reqs=self.max_num_reqs,

203 max_model_len=self.max_model_len,

204 max_num_blocks_per_req=self.max_num_blocks_per_req,

205 + max_num_batched_tokens=self.max_num_tokens,

206 device=self.device,

207 pin_memory=self.pin_memory,

208 vocab_size=model_config.get_vocab_size(),

@@ -291,11 +291,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):

291 device="cpu",

292 pin_memory=self.pin_memory)

293 self.positions_np = self.positions_cpu.numpy()

294 - self.slot_mapping_cpu = torch.zeros(self.max_num_tokens,

295 - dtype=torch.int64,

296 - device="cpu",

297 - pin_memory=self.pin_memory)

298 - self.slot_mapping_np = self.slot_mapping_cpu.numpy()

294 self.query_start_loc_cpu = torch.zeros(self.max_num_reqs + 1,

295 dtype=torch.int32,

296 device="cpu",

@@ -586,7 +581,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):

581 block_offsets = positions_np % self.block_size

582 np.add(block_numbers * self.block_size,

583 block_offsets,

589 - out=self.slot_mapping_np[:total_num_scheduled_tokens])

584 + out=self.input_batch.block_table.

585 + slot_mapping_np[:total_num_scheduled_tokens])

586

587 # Prepare the attention metadata.

588 self.query_start_loc_np[0] = 0

@@ -614,12 +610,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):

610 self.query_start_loc_cpu[:num_reqs + 1], non_blocking=True)

611 self.seq_lens[:num_reqs].copy_(self.seq_lens_cpu[:num_reqs],

612 non_blocking=True)

617 - self.slot_mapping[:total_num_scheduled_tokens].copy_(

618 - self.slot_mapping_cpu[:total_num_scheduled_tokens],

619 - non_blocking=True)

613

614 # Fill unused with -1. Needed for reshape_and_cache

622 - self.slot_mapping[total_num_scheduled_tokens:].fill_(-1)

615 self.seq_lens[num_reqs:].fill_(0)

616 self.query_start_loc[num_reqs + 1:].fill_(-1)

617

@@ -1821,6 +1813,11 @@ class GPUModelRunner(LoRAModelRunnerMixin):

1813 self.vllm_config.compilation_config.static_forward_context,

1814 self.kv_caches)

1815

1816 + self.attn_metadata_builder = self.attn_backend.get_builder_cls()(

1817 + weakref.proxy(self),

1818 + kv_cache_config.kv_cache_groups[0].kv_cache_spec,

1819 + self.input_batch.block_table)

1820 +

1821 def get_kv_cache_spec(self) -> dict[str, KVCacheSpec]:

1822 """

1823 Generates the KVCacheSpec by parsing the kv cache format from each

vllm/v1/worker/tpu_model_runner.py

+9 -9

@@ -179,6 +179,7 @@ class TPUModelRunner(LoRAModelRunnerMixin):

179 max_num_reqs=self.max_num_reqs,

180 max_model_len=self.max_model_len,

181 max_num_blocks_per_req=self.max_num_blocks_per_req,

182 + max_num_batched_tokens=self.max_num_tokens,

183 device=self.device,

184 pin_memory=self.pin_memory,

185 vocab_size=self.vocab_size,

@@ -197,10 +198,6 @@ class TPUModelRunner(LoRAModelRunnerMixin):

198 device="cpu")

199 self.positions_np = self.positions_cpu.numpy()

200

200 - self.slot_mapping_cpu = torch.zeros(self.max_num_tokens,

201 - dtype=torch.int64,

202 - device="cpu")

203 - self.slot_mapping_np = self.slot_mapping_cpu.numpy()

201 self.block_table_cpu = torch.zeros(

202 (self.max_num_reqs, self.max_num_blocks_per_req),

203 dtype=self.input_batch.block_table.get_cpu_tensor().dtype,

@@ -533,7 +530,8 @@ class TPUModelRunner(LoRAModelRunnerMixin):

530 block_offsets = positions_np % self.block_size

531 np.add(block_numbers * self.block_size,

532 block_offsets,

536 - out=self.slot_mapping_np[:total_num_scheduled_tokens])

533 + out=self.input_batch.block_table.

534 + slot_mapping_cpu[:total_num_scheduled_tokens])

535

536 # Prepare the attention metadata.

537 self.query_start_loc_np[0] = 0

@@ -557,10 +555,12 @@ class TPUModelRunner(LoRAModelRunnerMixin):

555 self.position_ids = self.positions_cpu[:

556 padded_total_num_scheduled_tokens].to(

557 self.device)

560 - self.slot_mapping_cpu[total_num_scheduled_tokens:] = _PAD_SLOT_ID

561 - slot_mapping = self.slot_mapping_cpu[:

562 - padded_total_num_scheduled_tokens].to(

563 - self.device)

558 + self.input_batch.block_table.slot_mapping_cpu[

559 + total_num_scheduled_tokens:] = _PAD_SLOT_ID

560 + slot_mapping = (

561 + self.input_batch.block_table.

562 + slot_mapping_cpu[:padded_total_num_scheduled_tokens].to(

563 + self.device))

564 block_tables = self.block_table_cpu[:self.max_num_reqs]

565 block_tables[:num_reqs, :self.max_num_blocks_per_req] = (

566 self.input_batch.block_table.get_cpu_tensor()[:num_reqs])