ae88aada38 · tduncan/Nvidia-Optimization-Stack-VLLM

@@ -67,12 +67,19 @@ from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead

67 from vllm.model_executor.model_loader.weight_utils import default_weight_loader

68 from vllm.model_executor.models.module_mapping import MultiModelKeys

69 from vllm.multimodal import MULTIMODAL_REGISTRY

70 +from vllm.multimodal.evs import (

71 + compute_mrope_for_media,

72 + compute_retained_tokens_count,

73 + compute_retention_mask,

74 + recompute_mrope_positions,

75 +)

76 from vllm.multimodal.inputs import (

77 MultiModalDataDict,

78 MultiModalFeatureSpec,

79 MultiModalFieldConfig,

80 MultiModalKwargsItem,

81 MultiModalKwargsItems,

82 + PlaceholderRange,

83 VideoItem,

84 )

85 from vllm.multimodal.parse import ImageSize, MultiModalDataItems, MultiModalDataParser

@@ -92,6 +99,7 @@ from .interfaces import (

99 SupportsLoRA,

100 SupportsMRoPE,

101 SupportsMultiModal,

102 + SupportsMultiModalPruning,

103 SupportsPP,

104 _require_is_multimodal,

105 )

@@ -1043,13 +1051,39 @@ class Qwen3VLMultiModalProcessor(BaseMultiModalProcessor[Qwen3VLProcessingInfo])

1051 tokenizer.encode(f"<{curr_time:.1f} seconds>", add_special_tokens=False)

1052 for curr_time in timestamps

1053 ]

1046 - num_tokens_per_frame = int(grid_thw[1:].prod()) // merge_length

1054 + tokens_per_frame = int(grid_thw[1:].prod()) // merge_length

1055 + per_frame_token_counts = [tokens_per_frame for _ in frames_idx_token]

1056 +

1057 + video_pruning_rate = self.info.ctx.get_mm_config().video_pruning_rate

1058 + if video_pruning_rate is not None and video_pruning_rate > 0.0:

1059 + total_retained = compute_retained_tokens_count(

1060 + tokens_per_frame,

1061 + len(frames_idx_token),

1062 + video_pruning_rate,

1063 + )

1064 + if len(frames_idx_token) == 0:

1065 + per_frame_token_counts = []

1066 + elif len(frames_idx_token) == 1:

1067 + per_frame_token_counts = [tokens_per_frame]

1068 + else:

1069 + first_frame_tokens = tokens_per_frame

1070 + remaining_tokens = max(total_retained - first_frame_tokens, 0)

1071 + base = remaining_tokens // (len(frames_idx_token) - 1)

1072 + remainder = remaining_tokens % (len(frames_idx_token) - 1)

1073 + per_frame_token_counts = [first_frame_tokens]

1074 + for frame_idx in range(1, len(frames_idx_token)):

1075 + extra = base + (1 if (frame_idx - 1) < remainder else 0)

1076 + per_frame_token_counts.append(extra)

1077 +

1078 placeholder = []

1048 - for frame_idx in frames_idx_token:

1049 - placeholder.extend(frame_idx)

1079 + for frame_idx, timestamp_tokens in enumerate(frames_idx_token):

1080 + placeholder.extend(timestamp_tokens)

1081 + tokens_this_frame = per_frame_token_counts[

1082 + frame_idx if frame_idx < len(per_frame_token_counts) else -1

1083 + ]

1084 placeholder.extend(

1085 [vision_start_token_id]

1052 - + [video_token_id] * num_tokens_per_frame

1086 + + [video_token_id] * tokens_this_frame

1087 + [vision_end_token_id]

1088 )

1089 return PromptUpdateDetails.select_token_id(placeholder, video_token_id)

@@ -1190,6 +1224,7 @@ class Qwen3VLForConditionalGeneration(

1224 SupportsPP,

1225 SupportsMRoPE,

1226 SupportsEagle3,

1227 + SupportsMultiModalPruning,

1228 ):

1229 packed_modules_mapping = {

1230 "qkv_proj": [

@@ -1232,6 +1267,11 @@ class Qwen3VLForConditionalGeneration(

1267 self.config = config

1268 self.multimodal_config = multimodal_config

1269 self.use_data_parallel = multimodal_config.mm_encoder_tp_mode == "data"

1270 + self.video_pruning_rate = multimodal_config.video_pruning_rate

1271 + self.is_multimodal_pruning_enabled = (

1272 + multimodal_config.is_multimodal_pruning_enabled()

1273 + )

1274 +

1275 if not multimodal_config.get_limit_per_prompt(

1276 "image"

1277 ) and not multimodal_config.get_limit_per_prompt("video"):

@@ -1418,6 +1458,109 @@ class Qwen3VLForConditionalGeneration(

1458 sizes = (grid_thw.prod(-1) // merge_size // merge_size).tolist()

1459 return video_embeds.split(sizes)

1460

1461 + def _postprocess_image_embeds_evs(

1462 + self,

1463 + image_embeds_split: tuple[torch.Tensor, ...],

1464 + image_input: Qwen2_5_VLImageInputs,

1465 + ) -> tuple[torch.Tensor, ...]:

1466 + """

1467 + Append mrope positions for each for images.

1468 + This is necessary to recover correct mrope

1469 + positions after video pruning

1470 +

1471 + Args:

1472 + image_embeds_split: Tuple of image embeddings for

1473 + each image item.

1474 + image_input: Image input data.

1475 +

1476 + Returns:

1477 + Tuple of image embeddings for each image item.

1478 + Resulting embeddings will have extra 4 channels for

1479 + computed mrope positions.

1480 + """

1481 + merge_size = self.visual.spatial_merge_size

1482 + grid_thw = image_input["image_grid_thw"]

1483 + grid_thw_list = grid_thw.tolist()

1484 + image_embeds_out = []

1485 + for emb, size in zip(image_embeds_split, grid_thw_list):

1486 + positions = compute_mrope_for_media(size, merge_size).to(emb.device)

1487 + emb = torch.cat([emb, positions], dim=1)

1488 + image_embeds_out.append(emb)

1489 + image_embeds_split = image_embeds_out

1490 + return tuple(image_embeds_split)

1491 +

1492 + def _postprocess_video_embeds_evs(

1493 + self,

1494 + video_embeds_split: tuple[torch.Tensor, ...],

1495 + video_input: Qwen2_5_VLVideoInputs,

1496 + ) -> tuple[torch.Tensor, ...]:

1497 + """

1498 + Prunes video embeddings via Efficient Video Sampling (EVS)

1499 + and then appends mrope positions for each retained embeddings

1500 +

1501 + Args:

1502 + video_embeds_split: Tuple of video embeddings for each video item.

1503 + video_input: Video input data.

1504 +

1505 + Returns:

1506 + Tuple of video embeddings for each video item.

1507 + Resulting embeddings will have extra 4 channels for

1508 + computed mrope positions.

1509 + """

1510 + grid_thw = video_input["video_grid_thw"]

1511 + assert grid_thw.ndim == 2

1512 + grid_thw_list = grid_thw.tolist()

1513 + merge_size = self.visual.spatial_merge_size

1514 +

1515 + # Cast to long to match the original code

1516 + # https://github.com/huggingface/transformers/blob/41980ce93e775f6c88500c51c8db7946fc6a2add/src/transformers/models/qwen2_5_vl/modular_qwen2_5_vl.py#L491 # noqa

1517 + second_per_grid_ts = video_input.get("second_per_grid_ts")

1518 + if second_per_grid_ts is None:

1519 + # For Qwen3-VL, second_per_grid_ts might not be available

1520 + # Use default value of 1.0 for each video

1521 + second_per_grid_ts = torch.ones(len(grid_thw_list), dtype=torch.long)

1522 + else:

1523 + second_per_grid_ts = second_per_grid_ts.long()

1524 + tokens_per_second = getattr(self.config.vision_config, "tokens_per_second", 1.0)

1525 +

1526 + video_embeds_out = []

1527 + for emb, size, video_second_per_grid_t in zip(

1528 + video_embeds_split, grid_thw_list, second_per_grid_ts

1529 + ):

1530 + # For each video, we compute retention mask using EVS

1531 + retention_mask = compute_retention_mask(

1532 + emb,

1533 + size,

1534 + spatial_merge_size=self.visual.spatial_merge_size,

1535 + q=self.video_pruning_rate,

1536 + )

1537 +

1538 + # Debug logging for EVS pruning

1539 + logger.debug(

1540 + "EVS: Video tokens pruned from %d to %d (T=%d,H=%d,W=%d, "

1541 + "pruning_rate=%.2f, reduction=%.1f%%)",

1542 + emb.shape[0],

1543 + retention_mask.sum().item(),

1544 + size[0],

1545 + size[1],

1546 + size[2],

1547 + self.video_pruning_rate,

1548 + (1 - retention_mask.float().mean().item()) * 100,

1549 + )

1550 +

1551 + positions = compute_mrope_for_media(

1552 + size,

1553 + merge_size,

1554 + tokens_per_second=tokens_per_second,

1555 + video_second_per_grid=video_second_per_grid_t.item(),

1556 + ).to(emb.device)

1557 +

1558 + emb = emb[retention_mask]

1559 + positions = positions[retention_mask]

1560 + emb = torch.cat([emb, positions], dim=1)

1561 + video_embeds_out.append(emb)

1562 + return tuple(video_embeds_out)

1563 +

1564 def _parse_and_validate_multimodal_inputs(self, **kwargs: object) -> dict:

1565 mm_input_by_modality = {}

1566 for input_key in kwargs:

@@ -1440,6 +1583,20 @@ class Qwen3VLForConditionalGeneration(

1583 def iter_mm_grid_hw(

1584 self, input_tokens: list[int], mm_features: list[MultiModalFeatureSpec]

1585 ) -> Iterator[tuple[int, int, int]]:

1586 + """

1587 + Iterate over multimodal features and yield grid information.

1588 +

1589 + For videos with EVS (Efficient Video Sampling) enabled, this function

1590 + computes the offset based on the pruned token count rather than relying

1591 + on input_tokens.index(), which would fail when tokens are pruned.

1592 +

1593 + Args:

1594 + input_tokens: List of token IDs in the prompt

1595 + mm_features: List of multimodal feature specifications

1596 +

1597 + Yields:

1598 + Tuple of (offset, grid_h, grid_w) for each frame/image

1599 + """

1600 video_token_id = self.config.video_token_id

1601 spatial_merge_size = self.config.vision_config.spatial_merge_size

1602 for mm_feature in sorted(mm_features, key=lambda f: f.mm_position.offset):

@@ -1452,42 +1609,289 @@ class Qwen3VLForConditionalGeneration(

1609 t, h, w = mm_feature.data["video_grid_thw"].data.tolist()

1610 llm_grid_h = h // spatial_merge_size

1611 llm_grid_w = w // spatial_merge_size

1455 - for _ in range(t):

1456 - offset = input_tokens.index(video_token_id, offset)

1457 - yield offset, llm_grid_h, llm_grid_w

1458 - offset += llm_grid_h * llm_grid_w

1612 +

1613 + # Check if EVS (Efficient Video Sampling) is enabled

1614 + is_evs_enabled = (

1615 + hasattr(self, "video_pruning_rate")

1616 + and self.video_pruning_rate is not None

1617 + and self.video_pruning_rate > 0.0

1618 + )

1619 +

1620 + if is_evs_enabled:

1621 + frame_offsets = self._extract_frame_offsets_from_mask(

1622 + mm_feature.mm_position, t

1623 + )

1624 + if frame_offsets is not None:

1625 + for rel_offset in frame_offsets:

1626 + yield offset + rel_offset, llm_grid_h, llm_grid_w

1627 + continue

1628 +

1629 + # If EVS is enabled but mask is missing, this indicates a bug

1630 + # in the prompt processing pipeline. The is_embed mask should

1631 + # always be present when video_pruning_rate > 0.

1632 + raise RuntimeError(

1633 + f"EVS is enabled (pruning_rate={self.video_pruning_rate}) "

1634 + "but is_embed mask is missing from mm_position. "

1635 + "This indicates a bug in prompt processing."

1636 + )

1637 + else:

1638 + # Non-EVS mode: Use original logic with input_tokens.index()

1639 + for _ in range(t):

1640 + offset = input_tokens.index(video_token_id, offset)

1641 + yield offset, llm_grid_h, llm_grid_w

1642 + offset += llm_grid_h * llm_grid_w

1643 else:

1644 raise ValueError(f"Unsupported modality: {mm_feature.modality}")

1645

1646 + def _get_evs_mask_segments(

1647 + self, mm_position: PlaceholderRange, expected_frames: int

1648 + ) -> list[torch.Tensor] | None:

1649 + """Extract contiguous segments from EVS is_embed mask.

1650 +

1651 + The EVS (Efficient Video Sampling) mask marks which placeholder

1652 + positions should be filled with video embeddings. This method splits

1653 + the mask into contiguous segments, where each segment represents one

1654 + retained frame.

1655 +

1656 + This is a pure function - it does not modify any state and always

1657 + returns the same output for the same input (idempotent).

1658 +

1659 + Args:

1660 + mm_position: MultiModal position containing the is_embed mask

1661 + expected_frames: Expected number of frame segments

1662 +

1663 + Returns:

1664 + List of tensors, each containing indices for one frame segment,

1665 + or None if EVS is not enabled or validation fails.

1666 + """

1667 + is_embed_mask = getattr(mm_position, "is_embed", None)

1668 + if is_embed_mask is None:

1669 + return None

1670 +

1671 + # Find all True positions in the mask

1672 + mask_tensor = torch.as_tensor(is_embed_mask, dtype=torch.bool).view(-1)

1673 + true_indices = torch.nonzero(mask_tensor, as_tuple=False).flatten()

1674 + if true_indices.numel() == 0:

1675 + return None

1676 +

1677 + # Split into contiguous segments (where diff > 1 indicates a gap)

1678 + if true_indices.numel() == 1:

1679 + segments = [true_indices]

1680 + else:

1681 + diffs = torch.diff(true_indices)

1682 + split_points = torch.nonzero(diffs != 1, as_tuple=False).flatten()

1683 + if split_points.numel() == 0:

1684 + segments = [true_indices]

1685 + else:

1686 + segments = torch.tensor_split(

1687 + true_indices, split_points.add(1).tolist()

1688 + )

1689 +

1690 + # Validate segment count matches expected frames

1691 + if len(segments) < expected_frames:

1692 + logger.debug(

1693 + "EVS mask segments (%d) do not match expected frames (%d)",

1694 + len(segments),

1695 + expected_frames,

1696 + )

1697 + return None

1698 +

1699 + return segments[:expected_frames]

1700 +

1701 + def _extract_frame_offsets_from_mask(

1702 + self, mm_position: PlaceholderRange, expected_frames: int

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

1704 + """Return relative offsets for each EVS-retained frame.

1705 +

1706 + The prompt processor stores a boolean mask inside ``mm_position`` that

1707 + marks which placeholder locations should be populated with video

1708 + embeddings. By splitting that mask into contiguous runs we can recover

1709 + the start of every retained frame without probing ``input_tokens``.

1710 +

1711 + Args:

1712 + mm_position: MultiModal position containing the is_embed mask

1713 + expected_frames: Expected number of frames

1714 +

1715 + Returns:

1716 + List of starting offsets (relative to mm_position) for each frame,

1717 + or None if EVS is not enabled.

1718 + """

1719 + segments = self._get_evs_mask_segments(mm_position, expected_frames)

1720 + if segments is None:

1721 + return None

1722 +

1723 + return [int(segment[0].item()) for segment in segments]

1724 +

1725 + def _get_actual_frame_token_counts(

1726 + self, mm_position: PlaceholderRange, expected_frames: int

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

1728 + """Return actual token count for each EVS-retained frame.

1729 +

1730 + This function calculates the actual number of tokens per frame by

1731 + analyzing the is_embed mask, accounting for EVS pruning. Each frame

1732 + may have a different token count due to content-aware pruning.

1733 +

1734 + Args:

1735 + mm_position: MultiModal position containing the is_embed mask

1736 + expected_frames: Expected number of frames

1737 +

1738 + Returns:

1739 + List of token counts for each frame, or None if EVS is not enabled.

1740 + """

1741 + segments = self._get_evs_mask_segments(mm_position, expected_frames)

1742 + if segments is None:

1743 + return None

1744 +

1745 + return [len(seg) for seg in segments]

1746 +

1747 + def recompute_mrope_positions(

1748 + self,

1749 + input_ids: list[int],

1750 + multimodal_embeddings: tuple[torch.Tensor, ...],

1751 + mrope_positions: torch.LongTensor,

1752 + num_computed_tokens: int,

1753 + ) -> tuple[tuple[torch.Tensor, ...], torch.Tensor, int]:

1754 + """

1755 + Update part of input mrope positions (starting with

1756 + num_computed_tokens index). Original mrope_positions are computed

1757 + for unpruned sequence and becomes incorrect once pruning occurs,

1758 + so once we prune media tokens we should reflect this in the

1759 + mrope_positions before we feed it to LLM.

1760 +

1761 + Args:

1762 + input_ids: (N,) All input tokens of the prompt (Containing

1763 + entire sequence).

1764 + multimodal_embeddings: Tuple of multimodal embeddings.

1765 + mrope_positions: Existing mrope positions (3, N) for entire

1766 + sequence

1767 + num_computed_tokens: A number of computed tokens so far.

1768 +

1769 + Returns:

1770 + Tuple of (multimodal_embeddings, mrope_positions,

1771 + mrope_position_delta).

1772 + """

1773 + image_token_id = self.config.image_token_id

1774 + video_token_id = self.config.video_token_id

1775 + vision_start_token_id = self.config.vision_start_token_id

1776 +

1777 + # Device

1778 + device = (

1779 + multimodal_embeddings[0].device

1780 + if len(multimodal_embeddings)

1781 + else mrope_positions.device

1782 + )

1783 +

1784 + # Tensors

1785 + input_ids_t = torch.as_tensor(input_ids, device=device, dtype=torch.long)

1786 +

1787 + mm_embeddings_out = [mm[:, :-4] for mm in multimodal_embeddings]

1788 + mm_embeddings_pos = [

1789 + mm[:, -4:].permute(1, 0).long() for mm in multimodal_embeddings

1790 + ]

1791 +

1792 + positions, mrope_positions_delta = recompute_mrope_positions(

1793 + input_ids_t,

1794 + mm_embeddings_pos,

1795 + mrope_positions,

1796 + num_computed_tokens,

1797 + vision_start_token_id,

1798 + image_token_id,

1799 + video_token_id,

1800 + )

1801 +

1802 + return tuple(mm_embeddings_out), positions, mrope_positions_delta

1803 +

1804 def get_mrope_input_positions(

1805 self,

1806 input_tokens: list[int],

1807 mm_features: list[MultiModalFeatureSpec],

1808 ) -> tuple[torch.Tensor, int]:

1809 + # Pre-collect actual frame token counts for EVS mode

1810 + frame_token_counts_map = {}

1811 + for mm_feature in mm_features:

1812 + if mm_feature.modality == "video":

1813 + is_evs_enabled = (

1814 + hasattr(self, "video_pruning_rate")

1815 + and self.video_pruning_rate is not None

1816 + and self.video_pruning_rate > 0.0

1817 + )

1818 + if is_evs_enabled:

1819 + t = mm_feature.data["video_grid_thw"].data.tolist()[0]

1820 + token_counts = self._get_actual_frame_token_counts(

1821 + mm_feature.mm_position, t

1822 + )

1823 + assert token_counts is not None, (

1824 + "EVS enabled but failed to extract frame token counts "

1825 + "from is_embed mask"

1826 + )

1827 + frame_token_counts_map[mm_feature.mm_position.offset] = token_counts

1828 +

1829 llm_pos_ids_list = []

1830 st = 0

1831 + frame_counts_idx = {}

1832 +

1833 for offset, llm_grid_h, llm_grid_w in self.iter_mm_grid_hw(

1834 input_tokens, mm_features

1835 ):

1836 text_len = offset - st

1837 st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0

1474 - llm_pos_ids_list.append(

1838 +

1839 + # Determine actual token count for this frame

1840 + base_offset = None

1841 + for feat_offset in frame_token_counts_map:

1842 + if offset >= feat_offset:

1843 + base_offset = feat_offset

1844 +

1845 + if base_offset is not None:

1846 + # EVS mode: use actual token count from is_embed mask

1847 + assert base_offset in frame_token_counts_map, (

1848 + f"Found base_offset {base_offset} but not in frame_token_counts_map"

1849 + )

1850 +

1851 + if base_offset not in frame_counts_idx:

1852 + frame_counts_idx[base_offset] = 0

1853 +

1854 + counts = frame_token_counts_map[base_offset]

1855 + idx = frame_counts_idx[base_offset]

1856 +

1857 + assert idx < len(counts), (

1858 + f"EVS frame index {idx} out of range (total frames: {len(counts)})"

1859 + )

1860 +

1861 + actual_frame_tokens = counts[idx]

1862 + frame_counts_idx[base_offset] += 1

1863 + else:

1864 + # Non-EVS mode (or image): use theoretical grid size

1865 + actual_frame_tokens = llm_grid_h * llm_grid_w

1866 +

1867 + # Add text segment

1868 + text_positions = (

1869 np.broadcast_to(np.arange(text_len), (3, text_len)) + st_idx

1870 )

1871 + llm_pos_ids_list.append(text_positions)

1872 + st_idx += text_len

1873

1874 + # Add frame segment with actual token count (not theoretical)

1875 grid_indices = np.indices((1, llm_grid_h, llm_grid_w)).reshape(3, -1)

1479 - llm_pos_ids_list.append(grid_indices + text_len + st_idx)

1480 - st = offset + llm_grid_h * llm_grid_w

1876 + # Only take the first actual_frame_tokens positions

1877 + frame_positions = grid_indices[:, :actual_frame_tokens] + st_idx

1878 + llm_pos_ids_list.append(frame_positions)

1879 +

1880 + # Update st using actual token count

1881 + st = offset + actual_frame_tokens

1882

1883 + # Handle final text segment

1884 if st < len(input_tokens):

1885 st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0

1886 text_len = len(input_tokens) - st

1485 - llm_pos_ids_list.append(

1887 + final_text_positions = (

1888 np.broadcast_to(np.arange(text_len), (3, text_len)) + st_idx

1889 )

1890 + llm_pos_ids_list.append(final_text_positions)

1891

1892 llm_positions = np.concatenate(llm_pos_ids_list, axis=1).reshape(3, -1)

1893 mrope_position_delta = (llm_positions.max() + 1 - len(input_tokens)).item()

1894 +

1895 return torch.from_numpy(llm_positions), mrope_position_delta

1896

1897 def get_language_model(self) -> torch.nn.Module:

@@ -1508,9 +1912,17 @@ class Qwen3VLForConditionalGeneration(

1912 multimodal_input = mm_input_by_modality[modality]

1913 if modality == "image":

1914 image_embeddings = self._process_image_input(multimodal_input)

1915 + if self.is_multimodal_pruning_enabled:

1916 + image_embeddings = self._postprocess_image_embeds_evs(

1917 + image_embeddings, multimodal_input

1918 + )

1919 multimodal_embeddings += tuple(image_embeddings)

1920 if modality == "video":

1921 video_embeddings = self._process_video_input(multimodal_input)

1922 + if self.is_multimodal_pruning_enabled:

1923 + video_embeddings = self._postprocess_video_embeds_evs(

1924 + video_embeddings, multimodal_input

1925 + )

1926 multimodal_embeddings += tuple(video_embeddings)

1927 return multimodal_embeddings

1928

[Feature]Add EVS (Efficient Video Sampling) Support for Qwen3-VL (#29752)