636efd10a5 · tduncan/Nvidia-Optimization-Stack-VLLM

@@ -1054,7 +1054,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1054

1055 def _get_encoder_seq_lens(

1056 self,

1057 - scheduler_output: "SchedulerOutput",

1057 + scheduled_encoder_inputs: dict[str, list[int]],

1058 kv_cache_spec: KVCacheSpec,

1059 num_reqs: int,

1060 ) -> np.ndarray | None:

@@ -1064,31 +1064,27 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1064 # Build encoder_seq_lens array mapping request indices to

1065 # encoder lengths for inputs scheduled in this batch

1066 encoder_seq_lens = np.zeros(num_reqs, dtype=np.int32)

1067 - for req_id in scheduler_output.scheduled_encoder_inputs:

1067 + for req_id in scheduled_encoder_inputs:

1068 req_index = self.input_batch.req_id_to_index[req_id]

1069 encoder_seq_lens[req_index] = self.max_encoder_len

1070

1071 return encoder_seq_lens

1072

1073 def _prepare_inputs(

1074 - self, scheduler_output: "SchedulerOutput"

1074 + self,

1075 + scheduler_output: "SchedulerOutput",

1076 + num_scheduled_tokens: np.ndarray,

1077 + max_num_scheduled_tokens: int,

1078 ) -> tuple[

1076 - PerLayerAttnMetadata,

1079 torch.Tensor,

1080 SpecDecodeMetadata | None,

1079 - np.ndarray,

1080 - CommonAttentionMetadata | None,

1081 - int,

1081 UBatchSlices | None,

1082 torch.Tensor | None,

1084 - bool,

1083 ]:

1084 """

1085 :return: tuple[

1088 - attn_metadata: layer-to-attention_metadata mapping,

1086 logits_indices, spec_decode_metadata,

1090 - num_scheduled_tokens, spec_decode_common_attn_metadata,

1091 - max_num_scheduled_tokens, use_cascade_attn

1087 + ubatch_slices, num_tokens_across_dp,

1088 ]

1089 """

1090 total_num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens

@@ -1100,12 +1096,6 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1096 # This way, we can overlap the copy with the following CPU operations.

1097 self.input_batch.block_table.commit_block_table(num_reqs)

1098

1103 - # Get the number of scheduled tokens for each request.

1104 - req_ids = self.input_batch.req_ids

1105 - tokens = [scheduler_output.num_scheduled_tokens[i] for i in req_ids]

1106 - num_scheduled_tokens = np.array(tokens, dtype=np.int32)

1107 - max_num_scheduled_tokens = max(tokens)

1108 -

1099 # Get request indices.

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

1101 req_indices = np.repeat(self.arange_np[:num_reqs], num_scheduled_tokens)

@@ -1232,8 +1222,6 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1222 # Fill unused with 0 for full cuda graph mode.

1223 self.seq_lens.np[num_reqs:].fill(0)

1224 self.seq_lens.copy_to_gpu()

1235 - seq_lens = self.seq_lens.gpu[:num_reqs]

1236 - max_seq_len = self.seq_lens.np[:num_reqs].max().item()

1225

1226 num_tokens = [self.requests[r].num_tokens for r in self.input_batch.req_ids]

1227 num_tokens_np = np.array(num_tokens, dtype=np.int32)

@@ -1305,11 +1293,46 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1293 self.num_decode_draft_tokens.np[num_reqs:].fill(-1)

1294 self.num_decode_draft_tokens.copy_to_gpu()

1295

1308 - logits_indices_padded = None

1309 - if self.cache_config.kv_sharing_fast_prefill:

1310 - logits_indices_padded = self._prepare_kv_sharing_fast_prefill(

1311 - logits_indices

1296 + # Hot-Swap lora model

1297 + if self.lora_config:

1298 + assert (

1299 + np.sum(num_sampled_tokens)

1300 + <= self.vllm_config.scheduler_config.max_num_batched_tokens

1301 )

1302 + self.set_active_loras(

1303 + self.input_batch, num_scheduled_tokens, num_sampled_tokens

1304 + )

1305 +

1306 + return (

1307 + logits_indices,

1308 + spec_decode_metadata,

1309 + ubatch_slices,

1310 + num_tokens_across_dp,

1311 + )

1312 +

1313 + def _build_attention_metadata(

1314 + self,

1315 + total_num_scheduled_tokens: int,

1316 + max_num_scheduled_tokens: int,

1317 + num_reqs: int,

1318 + ubatch_slices: UBatchSlices | None = None,

1319 + logits_indices: torch.Tensor | None = None,

1320 + use_spec_decode: bool = False,

1321 + for_cudagraph_capture: bool = False,

1322 + scheduled_encoder_inputs: dict[str, list[int]] | None = None,

1323 + cascade_attn_prefix_lens: list[list[int]] | None = None,

1324 + ) -> tuple[PerLayerAttnMetadata, CommonAttentionMetadata | None]:

1325 + """

1326 + :return: tuple[attn_metadata, spec_decode_common_attn_metadata]

1327 + """

1328 + logits_indices_padded = None

1329 + num_logits_indices = 0

1330 + if logits_indices is not None:

1331 + num_logits_indices = logits_indices.size(0)

1332 + if self.cache_config.kv_sharing_fast_prefill:

1333 + logits_indices_padded = self._prepare_kv_sharing_fast_prefill(

1334 + logits_indices

1335 + )

1336

1337 # update seq_lens of decode reqs under DCP.

1338 if self.dcp_world_size > 1:

@@ -1324,15 +1347,28 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1347 attn_metadata: PerLayerAttnMetadata = {}

1348 if ubatch_slices is not None:

1349 attn_metadata = [dict() for _ in range(len(ubatch_slices))]

1327 - use_cascade_attn = False

1350

1329 - # Used in the below loop.

1351 + # Used in the below loop

1352 + query_start_loc = self.query_start_loc.gpu[: num_reqs + 1]

1353 query_start_loc_cpu = self.query_start_loc.cpu[: num_reqs + 1]

1354 + seq_lens = self.seq_lens.gpu[:num_reqs]

1355 seq_lens_cpu = self.seq_lens.cpu[:num_reqs]

1356 num_computed_tokens_cpu = self.input_batch.num_computed_tokens_cpu_tensor[

1357 :num_reqs

1358 ]

1359 + dcp_local_seq_lens = (

1360 + self.dcp_local_seq_lens.gpu[:num_reqs] if self.dcp_world_size > 1 else None

1361 + )

1362 spec_decode_common_attn_metadata = None

1363 +

1364 + if for_cudagraph_capture:

1365 + # For some attention backends (e.g. FA) with sliding window models we need

1366 + # to make sure the backend see a max_seq_len that is larger to the sliding

1367 + # window size when capturing to make sure the correct kernel is selected.

1368 + max_seq_len = self.max_model_len

1369 + else:

1370 + max_seq_len = self.seq_lens.np[:num_reqs].max().item()

1371 +

1372 if use_spec_decode:

1373 self.num_accepted_tokens.np[:num_reqs] = (

1374 self.input_batch.num_accepted_tokens_cpu[:num_reqs]

@@ -1342,14 +1378,16 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1378

1379 # Prepare the attention metadata for each KV cache group and make layers

1380 # in the same group share the same metadata.

1345 - for kv_cache_group_id, kv_cache_group_spec in enumerate(

1381 + for kv_cache_gid, kv_cache_group in enumerate(

1382 self.kv_cache_config.kv_cache_groups

1383 ):

1384 encoder_seq_lens = self._get_encoder_seq_lens(

1349 - scheduler_output, kv_cache_group_spec.kv_cache_spec, num_reqs

1385 + scheduled_encoder_inputs or {},

1386 + kv_cache_group.kv_cache_spec,

1387 + num_reqs,

1388 )

1389

1352 - if isinstance(kv_cache_group_spec.kv_cache_spec, EncoderOnlyAttentionSpec):

1390 + if isinstance(kv_cache_group.kv_cache_spec, EncoderOnlyAttentionSpec):

1391 # Encoder-only layers do not have KV cache, so we need to

1392 # create a dummy block table and slot mapping for them.

1393 blk_table_tensor = torch.zeros(

@@ -1362,18 +1400,14 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1400 dtype=torch.int64,

1401 device=self.device,

1402 )

1365 - num_common_prefix_blocks = 0

1403 else:

1367 - blk_table = self.input_batch.block_table[kv_cache_group_id]

1404 + blk_table = self.input_batch.block_table[kv_cache_gid]

1405 blk_table_tensor = blk_table.get_device_tensor(num_reqs)

1406 slot_mapping = blk_table.slot_mapping.gpu[:total_num_scheduled_tokens]

1407

1408 # Fill unused with -1. Needed for reshape_and_cache in full cuda

1409 # graph mode.

1410 blk_table.slot_mapping.gpu[total_num_scheduled_tokens:].fill_(-1)

1374 - num_common_prefix_blocks = scheduler_output.num_common_prefix_blocks[

1375 - kv_cache_group_id

1376 - ]

1411

1412 common_attn_metadata = CommonAttentionMetadata(

1413 query_start_loc=query_start_loc,

@@ -1388,35 +1422,26 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1422 block_table_tensor=blk_table_tensor,

1423 slot_mapping=slot_mapping,

1424 logits_indices_padded=logits_indices_padded,

1391 - num_logits_indices=logits_indices.size(0),

1425 + num_logits_indices=num_logits_indices,

1426 causal=True,

1427 encoder_seq_lens=encoder_seq_lens,

1394 - dcp_local_seq_lens=self.dcp_local_seq_lens.gpu[:num_reqs]

1395 - if self.dcp_world_size > 1

1396 - else None,

1428 + dcp_local_seq_lens=dcp_local_seq_lens,

1429 )

1430

1431 if self.speculative_config and spec_decode_common_attn_metadata is None:

1432 if isinstance(self.drafter, EagleProposer):

1401 - if (

1402 - self.drafter.attn_layer_names[0]

1403 - in kv_cache_group_spec.layer_names

1404 - ):

1433 + if self.drafter.attn_layer_names[0] in kv_cache_group.layer_names:

1434 spec_decode_common_attn_metadata = common_attn_metadata

1435 else:

1436 spec_decode_common_attn_metadata = common_attn_metadata

1437

1409 - for attn_group in self.attn_groups[kv_cache_group_id]:

1410 - # Prepare for cascade attention if enabled & beneficial.

1411 - common_prefix_len = 0

1438 + for attn_gid, attn_group in enumerate(self.attn_groups[kv_cache_gid]):

1439 + cascade_attn_prefix_len = (

1440 + cascade_attn_prefix_lens[kv_cache_gid][attn_gid]

1441 + if cascade_attn_prefix_lens

1442 + else 0

1443 + )

1444 builder = attn_group.get_metadata_builder()

1413 - if self.cascade_attn_enabled:

1414 - common_prefix_len = self._compute_cascade_attn_prefix_len(

1415 - num_scheduled_tokens,

1416 - num_common_prefix_blocks,

1417 - attn_group.kv_cache_spec,

1418 - builder,

1419 - )

1445

1446 extra_attn_metadata_args = {}

1447 if use_spec_decode and isinstance(builder, GDNAttentionMetadataBuilder):

@@ -1434,51 +1459,69 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1459 for ubid, common_attn_metadata in enumerate(

1460 common_attn_metadata_list

1461 ):

1437 - attn_metadata_i = attn_group.get_metadata_builder(

1438 - ubatch_id=ubid

1439 - ).build(

1440 - common_prefix_len=common_prefix_len,

1441 - common_attn_metadata=common_attn_metadata,

1442 - )

1443 - for layer_name in kv_cache_group_spec.layer_names:

1462 + builder = attn_group.get_metadata_builder(ubatch_id=ubid)

1463 + if for_cudagraph_capture:

1464 + attn_metadata_i = builder.build_for_cudagraph_capture(

1465 + common_attn_metadata

1466 + )

1467 + else:

1468 + attn_metadata_i = builder.build(

1469 + common_prefix_len=cascade_attn_prefix_len,

1470 + common_attn_metadata=common_attn_metadata,

1471 + )

1472 + for layer_name in kv_cache_group.layer_names:

1473 assert type(attn_metadata) is list

1474 attn_metadata[ubid][layer_name] = attn_metadata_i

1475 else:

1476 assert isinstance(attn_metadata, dict)

1448 - attn_metadata_i = builder.build(

1449 - common_prefix_len=common_prefix_len,

1450 - common_attn_metadata=common_attn_metadata,

1451 - **extra_attn_metadata_args,

1452 - )

1453 - use_cascade_attn |= getattr(attn_metadata_i, "use_cascade", False)

1477 + if for_cudagraph_capture:

1478 + attn_metadata_i = builder.build_for_cudagraph_capture(

1479 + common_attn_metadata

1480 + )

1481 + else:

1482 + attn_metadata_i = builder.build(

1483 + common_prefix_len=cascade_attn_prefix_len,

1484 + common_attn_metadata=common_attn_metadata,

1485 + **extra_attn_metadata_args,

1486 + )

1487 for layer_name in attn_group.layer_names:

1488 attn_metadata[layer_name] = attn_metadata_i

1489

1457 - # disable cascade attention when DBO

1458 - if ubatch_slices is not None:

1459 - use_cascade_attn = False

1490 + return attn_metadata, spec_decode_common_attn_metadata

1491

1461 - # Hot-Swap lora model

1462 - if self.lora_config:

1463 - assert (

1464 - np.sum(num_sampled_tokens)

1465 - <= self.vllm_config.scheduler_config.max_num_batched_tokens

1466 - )

1467 - self.set_active_loras(

1468 - self.input_batch, num_scheduled_tokens, num_sampled_tokens

1469 - )

1492 + def _compute_cascade_attn_prefix_lens(

1493 + self,

1494 + num_scheduled_tokens: np.ndarray,

1495 + num_common_prefix_blocks: list[int],

1496 + ) -> list[list[int]] | None:

1497 + """

1498 + :return: Optional[cascade_attn_prefix_lens]

1499 + cascade_attn_prefix_lens is 2D: ``[kv_cache_group_id][attn_group_idx]``,

1500 + None if we should not use cascade attention

1501 + """

1502

1471 - return (

1472 - attn_metadata,

1473 - logits_indices,

1474 - spec_decode_metadata,

1475 - num_scheduled_tokens,

1476 - spec_decode_common_attn_metadata,

1477 - max_num_scheduled_tokens,

1478 - ubatch_slices,

1479 - num_tokens_across_dp,

1480 - use_cascade_attn,

1481 - )

1503 + use_cascade_attn = False

1504 + num_kv_cache_groups = len(self.kv_cache_config.kv_cache_groups)

1505 + cascade_attn_prefix_lens: list[list[int]] = [

1506 + [] for _ in range(num_kv_cache_groups)

1507 + ]

1508 +

1509 + for kv_cache_gid in range(num_kv_cache_groups):

1510 + for attn_group in self.attn_groups[kv_cache_gid]:

1511 + if isinstance(attn_group.kv_cache_spec, EncoderOnlyAttentionSpec):

1512 + cascade_attn_prefix_len = 0

1513 + else:

1514 + # 0 if cascade attention should not be used

1515 + cascade_attn_prefix_len = self._compute_cascade_attn_prefix_len(

1516 + num_scheduled_tokens,

1517 + num_common_prefix_blocks[kv_cache_gid],

1518 + attn_group.kv_cache_spec,

1519 + attn_group.get_metadata_builder(),

1520 + )

1521 + cascade_attn_prefix_lens[kv_cache_gid].append(cascade_attn_prefix_len)

1522 + use_cascade_attn |= cascade_attn_prefix_len > 0

1523 +

1524 + return cascade_attn_prefix_lens if use_cascade_attn else None

1525

1526 def _compute_cascade_attn_prefix_len(

1527 self,

@@ -1504,6 +1547,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1547 Returns:

1548 int: Length of common prefix in tokens.

1549 """

1550 +

1551 common_prefix_len = num_common_prefix_blocks * kv_cache_spec.block_size

1552 if common_prefix_len == 0:

1553 # Common case.

@@ -2497,18 +2541,48 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

2541 "it when the requests need prompt logprobs"

2542 )

2543

2500 - # Prepare the decoder inputs.

2544 + num_reqs = self.input_batch.num_reqs

2545 + req_ids = self.input_batch.req_ids

2546 + tokens = [scheduler_output.num_scheduled_tokens[i] for i in req_ids]

2547 + num_scheduled_tokens_np = np.array(tokens, dtype=np.int32)

2548 + max_num_scheduled_tokens = int(num_scheduled_tokens_np.max())

2549 +

2550 (

2502 - attn_metadata,

2551 logits_indices,

2552 spec_decode_metadata,

2505 - num_scheduled_tokens_np,

2506 - spec_decode_common_attn_metadata,

2507 - max_query_len,

2553 ubatch_slices,

2554 num_tokens_across_dp,

2510 - use_cascade_attn,

2511 - ) = self._prepare_inputs(scheduler_output)

2555 + ) = self._prepare_inputs(

2556 + scheduler_output, num_scheduled_tokens_np, max_num_scheduled_tokens

2557 + )

2558 +

2559 + cascade_attn_prefix_lens = None

2560 + # Disable cascade attention when using microbatching (DBO)

2561 + if self.cascade_attn_enabled and ubatch_slices is None:

2562 + # Pre-compute cascade attention prefix lengths

2563 + # NOTE: Must be AFTER _prepare_inputs uses self.input_batch state

2564 + cascade_attn_prefix_lens = self._compute_cascade_attn_prefix_lens(

2565 + num_scheduled_tokens_np,

2566 + scheduler_output.num_common_prefix_blocks,

2567 + )

2568 +

2569 + # TODO(lucas): move cudagraph dispatching here:

2570 + # https://github.com/vllm-project/vllm/issues/23789

2571 +

2572 + total_num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens

2573 + use_spec_decode = len(scheduler_output.scheduled_spec_decode_tokens) > 0

2574 + attn_metadata, spec_decode_common_attn_metadata = (

2575 + self._build_attention_metadata(

2576 + total_num_scheduled_tokens=total_num_scheduled_tokens,

2577 + max_num_scheduled_tokens=max_num_scheduled_tokens,

2578 + num_reqs=num_reqs,

2579 + ubatch_slices=ubatch_slices,

2580 + logits_indices=logits_indices,

2581 + use_spec_decode=use_spec_decode,

2582 + scheduled_encoder_inputs=scheduler_output.scheduled_encoder_inputs,

2583 + cascade_attn_prefix_lens=cascade_attn_prefix_lens,

2584 + )

2585 + )

2586

2587 dp_rank = self.parallel_config.data_parallel_rank

2588 if ubatch_slices:

@@ -2532,16 +2606,19 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

2606 scheduler_output, num_input_tokens, intermediate_tensors

2607 )

2608

2535 - uniform_decode = (max_query_len == self.uniform_decode_query_len) and (

2536 - num_scheduled_tokens == self.input_batch.num_reqs * max_query_len

2537 - )

2609 + uniform_decode = (

2610 + max_num_scheduled_tokens == self.uniform_decode_query_len

2611 + ) and (num_scheduled_tokens == num_reqs * max_num_scheduled_tokens)

2612 batch_descriptor = BatchDescriptor(

2613 num_tokens=num_input_tokens,

2614 uniform_decode=uniform_decode,

2615 has_lora=len(self.input_batch.lora_id_to_lora_request) > 0,

2616 )

2617 cudagraph_runtime_mode, batch_descriptor = (

2544 - self.cudagraph_dispatcher.dispatch(batch_descriptor, use_cascade_attn)

2618 + self.cudagraph_dispatcher.dispatch(

2619 + batch_descriptor,

2620 + use_cascade_attn=cascade_attn_prefix_lens is not None,

2621 + )

2622 )

2623

2624 # Set cudagraph mode to none if calc_kv_scales is true.

@@ -3437,10 +3514,6 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

3514 # If force_attention is True, we always capture attention. Otherwise,

3515 # it only happens for cudagraph_runtime_mode=FULL.

3516 if force_attention or cudagraph_runtime_mode == CUDAGraphMode.FULL:

3440 - attn_metadata = {}

3441 - if ubatch_slices is not None:

3442 - attn_metadata = [dict() for _ in range(len(ubatch_slices))]

3443 -

3517 if create_mixed_batch:

3518 # In the mixed batch mode (used for FI warmup), we use

3519 # shorter sequence lengths to run faster.

@@ -3456,55 +3529,13 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

3529 self.query_start_loc.np[1 : num_reqs + 1] = cum_num_tokens

3530 self.query_start_loc.copy_to_gpu()

3531

3459 - for kv_cache_group_id, kv_cache_group_spec in enumerate(

3460 - self.kv_cache_config.kv_cache_groups

3461 - ):

3462 - common_attn_metadata = CommonAttentionMetadata(

3463 - query_start_loc=self.query_start_loc.gpu[: num_reqs + 1],

3464 - query_start_loc_cpu=self.query_start_loc.cpu[: num_reqs + 1],

3465 - seq_lens=self.seq_lens.gpu[:num_reqs],

3466 - seq_lens_cpu=self.seq_lens.cpu[:num_reqs],

3467 - num_computed_tokens_cpu=self.input_batch.num_computed_tokens_cpu_tensor[

3468 - :num_reqs

3469 - ],

3470 - num_reqs=num_reqs,

3471 - num_actual_tokens=num_tokens,

3472 - max_query_len=max_query_len,

3473 - max_seq_len=self.max_model_len,

3474 - block_table_tensor=self.input_batch.block_table[

3475 - kv_cache_group_id

3476 - ].get_device_tensor(num_reqs),

3477 - slot_mapping=self.input_batch.block_table[

3478 - kv_cache_group_id

3479 - ].slot_mapping.gpu[:num_tokens],

3480 - causal=True,

3481 - dcp_local_seq_lens=self.dcp_local_seq_lens.gpu[:num_reqs]

3482 - if self.dcp_world_size > 1

3483 - else None,

3484 - )

3485 - for attn_group in self.attn_groups[kv_cache_group_id]:

3486 - if ubatch_slices is not None:

3487 - common_attn_metadata_list = split_attn_metadata(

3488 - ubatch_slices, common_attn_metadata

3489 - )

3490 - for ubid, common_attn_metadata in enumerate(

3491 - common_attn_metadata_list

3492 - ):

3493 - assert common_attn_metadata.max_query_len == 1

3494 - attn_metadata_i = attn_group.get_metadata_builder(

3495 - ubatch_id=ubid

3496 - ).build_for_cudagraph_capture(common_attn_metadata)

3497 - for layer_name in attn_group.layer_names:

3498 - assert type(attn_metadata) is list

3499 - attn_metadata[ubid][layer_name] = attn_metadata_i

3500 - else:

3501 - assert type(attn_metadata) is dict

3502 - metadata_builder = attn_group.get_metadata_builder()

3503 - attn_metadata_i = metadata_builder.build_for_cudagraph_capture(

3504 - common_attn_metadata

3505 - )

3506 - for layer_name in attn_group.layer_names:

3507 - attn_metadata[layer_name] = attn_metadata_i

3532 + attn_metadata, _ = self._build_attention_metadata(

3533 + total_num_scheduled_tokens=num_tokens,

3534 + max_num_scheduled_tokens=max_query_len,

3535 + num_reqs=num_reqs,

3536 + ubatch_slices=ubatch_slices,

3537 + for_cudagraph_capture=True,

3538 + )

3539

3540 with self.maybe_dummy_run_with_lora(

3541 self.lora_config,

@@ -4478,9 +4509,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

4509 list[int]: List of kernel block sizes for each cache group.

4510 """

4511 kernel_block_sizes = []

4481 - for kv_cache_group_id, kv_cache_group in enumerate(

4482 - kv_cache_config.kv_cache_groups

4483 - ):

4512 + for kv_cache_gid, kv_cache_group in enumerate(kv_cache_config.kv_cache_groups):

4513 kv_cache_spec = kv_cache_group.kv_cache_spec

4514 if isinstance(kv_cache_spec, UniformTypeKVCacheSpecs):

4515 # All layers in the UniformTypeKVCacheSpecs have the same type,

@@ -4492,7 +4521,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

4521 # This is an attention backend that supports virtual

4522 # block splitting. Get the supported block sizes from

4523 # all backends in the group.

4495 - attn_groups = self.attn_groups[kv_cache_group_id]

4524 + attn_groups = self.attn_groups[kv_cache_gid]

4525 kv_manager_block_size = kv_cache_group.kv_cache_spec.block_size

4526 selected_kernel_size = self.select_common_block_size(

4527 kv_manager_block_size, attn_groups

[Core] Separate out attention metadata building logic from prepare inputs (#26764)