042da73244 · tduncan/Nvidia-Optimization-Stack-VLLM

@@ -1534,28 +1534,13 @@ class GPUModelRunner(

1534 """

1535 :return: tuple[attn_metadata, spec_decode_common_attn_metadata]

1536 """

1537 + # Attention metadata is not needed for attention free models

1538 + if len(self.kv_cache_config.kv_cache_groups) == 0:

1539 + return {}, None

1540 +

1541 num_tokens_padded = num_tokens_padded or num_tokens

1542 num_reqs_padded = num_reqs_padded or num_reqs

1539 -

1540 - logits_indices_padded = None

1541 - num_logits_indices = None

1542 - if logits_indices is not None:

1543 - num_logits_indices = logits_indices.size(0)

1544 - if self.cache_config.kv_sharing_fast_prefill:

1545 - logits_indices_padded = self._prepare_kv_sharing_fast_prefill(

1546 - logits_indices

1547 - )

1548 -

1549 - # update seq_lens of decode reqs under DCP.

1550 - if self.dcp_world_size > 1:

1551 - self.dcp_local_seq_lens.cpu[:num_reqs] = get_dcp_local_seq_lens(

1552 - self.seq_lens.cpu[:num_reqs],

1553 - self.dcp_world_size,

1554 - self.dcp_rank,

1555 - self.parallel_config.cp_kv_cache_interleave_size,

1556 - )

1557 - self.dcp_local_seq_lens.cpu[num_reqs:].fill_(0)

1558 - self.dcp_local_seq_lens.copy_to_gpu(num_reqs_padded)

1543 + assert num_reqs_padded is not None and num_tokens_padded is not None

1544

1545 attn_metadata: PerLayerAttnMetadata = {}

1546 if ubatch_slices is not None:

@@ -1576,36 +1561,12 @@ class GPUModelRunner(

1561 self.num_accepted_tokens.np[num_reqs:].fill(1)

1562 self.num_accepted_tokens.copy_to_gpu()

1563

1579 - # Used in the below loop, uses padded shapes

1580 - query_start_loc = self.query_start_loc.gpu[: num_reqs_padded + 1]

1581 - query_start_loc_cpu = self.query_start_loc.cpu[: num_reqs_padded + 1]

1582 - seq_lens = self.seq_lens.gpu[:num_reqs_padded]

1583 - seq_lens_cpu = self.seq_lens.cpu[:num_reqs_padded]

1584 - num_computed_tokens_cpu = self.input_batch.num_computed_tokens_cpu_tensor[

1585 - :num_reqs_padded

1586 - ]

1587 -

1588 - dcp_local_seq_lens, dcp_local_seq_lens_cpu = None, None

1589 - if self.dcp_world_size > 1:

1590 - dcp_local_seq_lens = self.dcp_local_seq_lens.gpu[:num_reqs_padded]

1591 - dcp_local_seq_lens_cpu = self.dcp_local_seq_lens.cpu[:num_reqs_padded]

1564 + kv_cache_groups = self.kv_cache_config.kv_cache_groups

1565

1593 - spec_decode_common_attn_metadata = None

1594 -

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

1596 - # in the same group share the same metadata.

1597 - for kv_cache_gid, kv_cache_group in enumerate(

1598 - self.kv_cache_config.kv_cache_groups

1599 - ):

1600 - encoder_seq_lens, encoder_seq_lens_cpu = self._get_encoder_seq_lens(

1601 - num_scheduled_tokens or {},

1602 - kv_cache_group.kv_cache_spec,

1603 - num_reqs_padded,

1604 - )

1605 -

1606 - if isinstance(kv_cache_group.kv_cache_spec, EncoderOnlyAttentionSpec):

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

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

1566 + def _get_block_table_and_slot_mapping(kv_cache_gid: int):

1567 + assert num_reqs_padded is not None and num_tokens_padded is not None

1568 + kv_cache_spec = kv_cache_groups[kv_cache_gid].kv_cache_spec

1569 + if isinstance(kv_cache_spec, EncoderOnlyAttentionSpec):

1570 blk_table_tensor = torch.zeros(

1571 (num_reqs_padded, 1),

1572 dtype=torch.int32,

@@ -1621,92 +1582,129 @@ class GPUModelRunner(

1582 blk_table_tensor = blk_table.get_device_tensor(num_reqs_padded)

1583 slot_mapping = blk_table.slot_mapping.gpu[:num_tokens_padded]

1584

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

1625 - # graph mode. `blk_table_tensor` -1 to match mamba PAD_SLOT_ID

1626 - slot_mapping[num_tokens:num_tokens_padded].fill_(-1)

1627 - blk_table_tensor[num_reqs:num_reqs_padded].fill_(-1)

1628 -

1629 - common_attn_metadata = CommonAttentionMetadata(

1630 - query_start_loc=query_start_loc,

1631 - query_start_loc_cpu=query_start_loc_cpu,

1632 - seq_lens=seq_lens,

1633 - _seq_lens_cpu=seq_lens_cpu,

1634 - _num_computed_tokens_cpu=num_computed_tokens_cpu,

1635 - num_actual_tokens=num_tokens_padded,

1636 - num_reqs=num_reqs_padded,

1637 - max_query_len=max_query_len,

1638 - max_seq_len=max_seq_len,

1639 - block_table_tensor=blk_table_tensor,

1640 - slot_mapping=slot_mapping,

1641 - logits_indices_padded=logits_indices_padded,

1642 - num_logits_indices=num_logits_indices,

1643 - causal=True,

1644 - encoder_seq_lens=encoder_seq_lens,

1645 - encoder_seq_lens_cpu=encoder_seq_lens_cpu,

1646 - dcp_local_seq_lens=dcp_local_seq_lens,

1647 - dcp_local_seq_lens_cpu=dcp_local_seq_lens_cpu,

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

1586 + # graph mode. `blk_table_tensor` -1 to match mamba PAD_SLOT_ID

1587 + slot_mapping[num_tokens:num_tokens_padded].fill_(-1)

1588 + blk_table_tensor[num_reqs:num_reqs_padded].fill_(-1)

1589 +

1590 + return blk_table_tensor, slot_mapping

1591 +

1592 + block_table_gid_0, slot_mapping_gid_0 = _get_block_table_and_slot_mapping(0)

1593 + cm_base = CommonAttentionMetadata(

1594 + query_start_loc=self.query_start_loc.gpu[: num_reqs_padded + 1],

1595 + query_start_loc_cpu=self.query_start_loc.cpu[: num_reqs_padded + 1],

1596 + seq_lens=self.seq_lens.gpu[:num_reqs_padded],

1597 + _seq_lens_cpu=self.seq_lens.cpu[:num_reqs_padded],

1598 + _num_computed_tokens_cpu=self.input_batch.num_computed_tokens_cpu_tensor[

1599 + :num_reqs_padded

1600 + ],

1601 + num_reqs=num_reqs_padded,

1602 + num_actual_tokens=num_tokens_padded,

1603 + max_query_len=max_query_len,

1604 + max_seq_len=max_seq_len,

1605 + block_table_tensor=block_table_gid_0,

1606 + slot_mapping=slot_mapping_gid_0,

1607 + causal=True,

1608 + )

1609 +

1610 + if self.dcp_world_size > 1:

1611 + self.dcp_local_seq_lens.cpu[:num_reqs] = get_dcp_local_seq_lens(

1612 + self.seq_lens.cpu[:num_reqs],

1613 + self.dcp_world_size,

1614 + self.dcp_rank,

1615 + self.parallel_config.cp_kv_cache_interleave_size,

1616 )

1617 + self.dcp_local_seq_lens.cpu[num_reqs:].fill_(0)

1618 + self.dcp_local_seq_lens.copy_to_gpu(num_reqs_padded)

1619 +

1620 + cm_base.dcp_local_seq_lens = self.dcp_local_seq_lens.gpu[:num_reqs_padded]

1621 + cm_base.dcp_local_seq_lens_cpu = self.dcp_local_seq_lens.cpu[

1622 + :num_reqs_padded

1623 + ]

1624 +

1625 + if logits_indices is not None and self.cache_config.kv_sharing_fast_prefill:

1626 + cm_base.num_logits_indices = logits_indices.size(0)

1627 + cm_base.logits_indices_padded = self._prepare_kv_sharing_fast_prefill(

1628 + logits_indices

1629 + )

1630 +

1631 + def _build_attn_group_metadata(

1632 + kv_cache_gid: int,

1633 + attn_gid: int,

1634 + common_attn_metadata: CommonAttentionMetadata,

1635 + ubid: int | None = None,

1636 + ) -> None:

1637 + attn_group = self.attn_groups[kv_cache_gid][attn_gid]

1638 + cascade_attn_prefix_len = (

1639 + cascade_attn_prefix_lens[kv_cache_gid][attn_gid]

1640 + if cascade_attn_prefix_lens

1641 + else 0

1642 + )

1643 +

1644 + builder = attn_group.get_metadata_builder(ubid or 0)

1645 + extra_attn_metadata_args = {}

1646 + if use_spec_decode and isinstance(builder, GDNAttentionMetadataBuilder):

1647 + assert ubid is None, "UBatching not supported with GDN yet"

1648 + extra_attn_metadata_args = dict(

1649 + num_accepted_tokens=self.num_accepted_tokens.gpu[:num_reqs_padded],

1650 + num_decode_draft_tokens_cpu=self.num_decode_draft_tokens.cpu[

1651 + :num_reqs_padded

1652 + ],

1653 + )

1654 +

1655 + if for_cudagraph_capture:

1656 + attn_metadata_i = builder.build_for_cudagraph_capture(

1657 + common_attn_metadata

1658 + )

1659 + else:

1660 + attn_metadata_i = builder.build(

1661 + common_prefix_len=cascade_attn_prefix_len,

1662 + common_attn_metadata=common_attn_metadata,

1663 + **extra_attn_metadata_args,

1664 + )

1665 +

1666 + if ubid is None:

1667 + assert isinstance(attn_metadata, dict)

1668 + attn_metadata_dict = attn_metadata

1669 + else:

1670 + assert isinstance(attn_metadata, list)

1671 + attn_metadata_dict = attn_metadata[ubid]

1672 +

1673 + for layer_name in attn_group.layer_names:

1674 + attn_metadata_dict[layer_name] = attn_metadata_i

1675 +

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

1677 + # in the same group share the same metadata.

1678 + spec_decode_common_attn_metadata = None

1679 + for kv_cache_gid, kv_cache_group in enumerate(kv_cache_groups):

1680 + cm = copy(cm_base) # shallow copy

1681 +

1682 + # Basically only the encoder seq_lens, block_table and slot_mapping change

1683 + # for each kv_cache_group.

1684 + cm.encoder_seq_lens, cm.encoder_seq_lens_cpu = self._get_encoder_seq_lens(

1685 + num_scheduled_tokens or {},

1686 + kv_cache_group.kv_cache_spec,

1687 + num_reqs_padded,

1688 + )

1689 + if kv_cache_gid > 0:

1690 + cm.block_table_tensor, cm.slot_mapping = (

1691 + _get_block_table_and_slot_mapping(kv_cache_gid)

1692 + )

1693

1694 if self.speculative_config and spec_decode_common_attn_metadata is None:

1695 if isinstance(self.drafter, EagleProposer):

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

1653 - spec_decode_common_attn_metadata = common_attn_metadata

1697 + spec_decode_common_attn_metadata = cm

1698 else:

1655 - spec_decode_common_attn_metadata = common_attn_metadata

1656 -

1657 - for attn_gid, attn_group in enumerate(self.attn_groups[kv_cache_gid]):

1658 - cascade_attn_prefix_len = (

1659 - cascade_attn_prefix_lens[kv_cache_gid][attn_gid]

1660 - if cascade_attn_prefix_lens

1661 - else 0

1662 - )

1663 - builder = attn_group.get_metadata_builder()

1664 -

1665 - extra_attn_metadata_args = {}

1666 - if use_spec_decode and isinstance(builder, GDNAttentionMetadataBuilder):

1667 - extra_attn_metadata_args = dict(

1668 - num_accepted_tokens=self.num_accepted_tokens.gpu[

1669 - :num_reqs_padded

1670 - ],

1671 - num_decode_draft_tokens_cpu=self.num_decode_draft_tokens.cpu[

1672 - :num_reqs_padded

1673 - ],

1674 - )

1699 + spec_decode_common_attn_metadata = cm

1700

1701 + for attn_gid in range(len(self.attn_groups[kv_cache_gid])):

1702 if ubatch_slices is not None:

1677 - common_attn_metadata_list = split_attn_metadata(

1678 - ubatch_slices, common_attn_metadata

1679 - )

1680 - for ubid, common_attn_metadata in enumerate(

1681 - common_attn_metadata_list

1682 - ):

1683 - builder = attn_group.get_metadata_builder(ubatch_id=ubid)

1684 - if for_cudagraph_capture:

1685 - attn_metadata_i = builder.build_for_cudagraph_capture(

1686 - common_attn_metadata

1687 - )

1688 - else:

1689 - attn_metadata_i = builder.build(

1690 - common_prefix_len=cascade_attn_prefix_len,

1691 - common_attn_metadata=common_attn_metadata,

1692 - )

1693 - for layer_name in kv_cache_group.layer_names:

1694 - assert type(attn_metadata) is list

1695 - attn_metadata[ubid][layer_name] = attn_metadata_i

1703 + for ubid, _cm in enumerate(split_attn_metadata(ubatch_slices, cm)):

1704 + _build_attn_group_metadata(kv_cache_gid, attn_gid, _cm, ubid)

1705 +

1706 else:

1697 - assert isinstance(attn_metadata, dict)

1698 - if for_cudagraph_capture:

1699 - attn_metadata_i = builder.build_for_cudagraph_capture(

1700 - common_attn_metadata

1701 - )

1702 - else:

1703 - attn_metadata_i = builder.build(

1704 - common_prefix_len=cascade_attn_prefix_len,

1705 - common_attn_metadata=common_attn_metadata,

1706 - **extra_attn_metadata_args,

1707 - )

1708 - for layer_name in attn_group.layer_names:

1709 - attn_metadata[layer_name] = attn_metadata_i

1707 + _build_attn_group_metadata(kv_cache_gid, attn_gid, cm)

1708

1709 if self.is_mm_prefix_lm:

1710 req_doc_ranges = {}

[Core] Refactor `_build_attention_metadata` (#29628)